Skip to content

Commit 7f1b17b

Browse files
committed
fix: pipeline variable kms key
1 parent e0102e7 commit 7f1b17b

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

src/sagemaker/estimator.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -614,16 +614,24 @@ def __init__(
614614
self.output_kms_key = resolve_value_from_config(
615615
output_kms_key, TRAINING_JOB_KMS_KEY_ID_PATH, sagemaker_session=self.sagemaker_session
616616
)
617+
use_volume_kms_config: bool = False
617618
if instance_type is None or isinstance(instance_type, str):
618619
instance_type_for_volume_kms = instance_type
619620
elif isinstance(instance_type, ParameterString):
620621
instance_type_for_volume_kms = instance_type.default_value
621622
else:
622-
raise ValueError(f"Bad value for instance type: '{instance_type}'")
623+
# Always attach a KMS key if an instance type that isn't a string or parameter
624+
# string is used.
625+
use_volume_kms_config = True
626+
if not isinstance(instance_type, PipelineVariable):
627+
logger.warning("Unrecognized value for instance type: '%s'", instance_type)
623628

624629
# KMS can only be attached to supported instances
625630
use_volume_kms_config = (
626-
(instance_type_for_volume_kms and instance_supports_kms(instance_type_for_volume_kms))
631+
use_volume_kms_config
632+
or (
633+
instance_type_for_volume_kms and instance_supports_kms(instance_type_for_volume_kms)
634+
)
627635
or instance_groups is not None
628636
and any(
629637
[

0 commit comments

Comments
 (0)