Merge master

Author: Captainia

src/sagemaker/jumpstart/estimator.py

@@ -112,6 +112,7 @@ def __init__(
         disable_output_compression: Optional[bool] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         config_name: Optional[str] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
     ):
         """Initializes a ``JumpStartEstimator``.
 
@@ -505,6 +506,8 @@ def __init__(
                 Specifies whether RemoteDebug is enabled for the training job
             config_name (Optional[str]):
                 Name of the training configuration to apply to the Estimator. (Default: None).
+            enable_session_tag_chaining (bool or PipelineVariable): Optional.
+                Specifies whether SessionTagChaining is enabled for the training job
 
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
@@ -584,6 +587,7 @@ def _validate_model_id_and_get_type_hook():
             enable_infra_check=enable_infra_check,
             enable_remote_debug=enable_remote_debug,
             config_name=config_name,
+            enable_session_tag_chaining=enable_session_tag_chaining,
         )
 
         self.model_id = estimator_init_kwargs.model_id

src/sagemaker/jumpstart/factory/estimator.py

@@ -131,6 +131,7 @@ def get_init_kwargs(
     enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
     enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
     config_name: Optional[str] = None,
+    enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -190,6 +191,7 @@ def get_init_kwargs(
         enable_infra_check=enable_infra_check,
         enable_remote_debug=enable_remote_debug,
         config_name=config_name,
+        enable_session_tag_chaining=enable_session_tag_chaining,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)

src/sagemaker/jumpstart/session_utils.py

@@ -219,7 +219,7 @@ def get_model_info_from_training_job(
         model_id,
         inferred_model_version,
         inference_config_name,
-        trainig_config_name,
+        training_config_name,
     ) = get_jumpstart_model_info_from_resource_arn(training_job_arn, sagemaker_session)
 
     model_version = inferred_model_version or None
@@ -231,4 +231,4 @@ def get_model_info_from_training_job(
             "for this training job."
         )
 
-    return model_id, model_version, inference_config_name, trainig_config_name
+    return model_id, model_version, inference_config_name, training_config_name

src/sagemaker/jumpstart/types.py

@@ -1078,7 +1078,7 @@ class JumpStartMetadataConfig(JumpStartDataHolderType):
         "resolved_metadata_config",
         "config_name",
         "default_inference_config",
-        "default_incremental_trainig_config",
+        "default_incremental_training_config",
         "supported_inference_configs",
         "supported_incremental_training_configs",
     ]
@@ -1114,7 +1114,7 @@ def __init__(
         self.resolved_metadata_config: Optional[Dict[str, Any]] = None
         self.config_name: Optional[str] = config_name
         self.default_inference_config: Optional[str] = config.get("default_inference_config")
-        self.default_incremental_trainig_config: Optional[str] = config.get(
+        self.default_incremental_training_config: Optional[str] = config.get(
             "default_incremental_training_config"
         )
         self.supported_inference_configs: Optional[List[str]] = config.get(
@@ -1775,6 +1775,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "enable_infra_check",
         "enable_remote_debug",
         "config_name",
+        "enable_session_tag_chaining",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1844,6 +1845,7 @@ def __init__(
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         config_name: Optional[str] = None,
+        enable_session_tag_chaining: Optional[Union[bool, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1904,6 +1906,7 @@ def __init__(
         self.enable_infra_check = enable_infra_check
         self.enable_remote_debug = enable_remote_debug
         self.config_name = config_name
+        self.enable_session_tag_chaining = enable_session_tag_chaining
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

tests/unit/sagemaker/jumpstart/test_types.py

@@ -1055,6 +1055,14 @@ def test_inference_configs_parsing():
     )
     assert list(config.config_components.keys()) == ["neuron-inference"]
 
+    spec = {
+        **BASE_SPEC,
+        **INFERENCE_CONFIGS,
+        **INFERENCE_CONFIG_RANKINGS,
+        "unrecognized-field": "blah",  # New fields in base metadata fields should be ignored
+    }
+    specs1 = JumpStartModelSpecs(spec)
+
 
 def test_set_inference_configs():
     spec = {**BASE_SPEC, **INFERENCE_CONFIGS, **INFERENCE_CONFIG_RANKINGS}