feat: jumpstart instance specific hyperparameters (#4180)

Author: evakravi

src/sagemaker/hyperparameters.py

@@ -31,6 +31,7 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
+    instance_type: Optional[str] = None,
     include_container_hyperparameters: bool = False,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -45,6 +46,8 @@ def retrieve_default(
             retrieve the default hyperparameters. (Default: None).
         model_version (str): The version of the model for which to retrieve the
             default hyperparameters. (Default: None).
+        instance_type (str): An instance type to optionally supply in order to get hyperparameters
+            specific for the instance type.
         include_container_hyperparameters (bool): ``True`` if the container hyperparameters
             should be returned. Container hyperparameters are not used to tune
             the specific algorithm. They are used by SageMaker Training jobs to set up
@@ -75,12 +78,13 @@ def retrieve_default(
         )
 
     return artifacts._retrieve_default_hyperparameters(
-        model_id,
-        model_version,
-        region,
-        include_container_hyperparameters,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,
+        instance_type=instance_type,
+        region=region,
+        include_container_hyperparameters=include_container_hyperparameters,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
     )
 

src/sagemaker/jumpstart/artifacts/hyperparameters.py

@@ -35,6 +35,7 @@ def _retrieve_default_hyperparameters(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    instance_type: Optional[str] = None,
 ):
     """Retrieves the training hyperparameters for the model matching the given arguments.
 
@@ -63,6 +64,8 @@ def _retrieve_default_hyperparameters(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        instance_type (str): An instance type to optionally supply in order to get hyperparameters
+            specific for the instance type.
     Returns:
         dict: the hyperparameters to use for the model.
     """
@@ -86,4 +89,19 @@ def _retrieve_default_hyperparameters(
             include_container_hyperparameters and hyperparameter.scope == VariableScope.CONTAINER
         ) or hyperparameter.scope == VariableScope.ALGORITHM:
             default_hyperparameters[hyperparameter.name] = str(hyperparameter.default)
+
+    instance_specific_hyperparameters = (
+        model_specs.training_instance_type_variants.get_instance_specific_hyperparameters(
+            instance_type
+        )
+        if instance_type
+        and getattr(model_specs, "training_instance_type_variants", None) is not None
+        else []
+    )
+
+    for instance_specific_hyperparameter in instance_specific_hyperparameters:
+        default_hyperparameters[instance_specific_hyperparameter.name] = str(
+            instance_specific_hyperparameter.default
+        )
+
     return default_hyperparameters

src/sagemaker/jumpstart/factory/estimator.py

@@ -599,6 +599,7 @@ def _add_hyperparameters_to_kwargs(
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
         sagemaker_session=kwargs.sagemaker_session,
+        instance_type=kwargs.instance_type,
     )
 
     for key, value in default_hyperparameters.items():

src/sagemaker/jumpstart/types.py

@@ -403,6 +403,50 @@ def to_json(self) -> Dict[str, Any]:
         json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
         return json_obj
 
+    def get_instance_specific_hyperparameters(
+        self, instance_type: str
+    ) -> List[JumpStartHyperparameter]:
+        """Returns instance specific hyperparameters.
+
+        Returns empty list if a model, instance type tuple does not have specific
+        hyperparameters.
+        """
+
+        if self.variants is None:
+            return []
+
+        instance_specific_hyperparameters: List[JumpStartHyperparameter] = [
+            JumpStartHyperparameter(json)
+            for json in self.variants.get(instance_type, {})
+            .get("properties", {})
+            .get("hyperparameters", [])
+        ]
+
+        instance_type_family = get_instance_type_family(instance_type)
+
+        instance_family_hyperparameters: List[JumpStartHyperparameter] = [
+            JumpStartHyperparameter(json)
+            for json in (
+                self.variants.get(instance_type_family, {})
+                .get("properties", {})
+                .get("hyperparameters", [])
+                if instance_type_family not in {"", None}
+                else []
+            )
+        ]
+
+        instance_specific_hyperparameter_names = {
+            hyperparameter.name for hyperparameter in instance_specific_hyperparameters
+        }
+
+        hyperparams_to_return = deepcopy(instance_specific_hyperparameters)
+
+        for hyperparameter in instance_family_hyperparameters:
+            if hyperparameter.name not in instance_specific_hyperparameter_names:
+                hyperparams_to_return.append(hyperparameter)
+
+        return hyperparams_to_return
+
     def get_instance_specific_environment_variables(self, instance_type: str) -> Dict[str, str]:
         """Returns instance specific environment variables.
 
@@ -867,10 +911,10 @@ def __init__(
         self.instance_type = instance_type
         self.region = region
         self.image_uri = image_uri
-        self.model_data = model_data
+        self.model_data = deepcopy(model_data)
         self.source_dir = source_dir
         self.entry_point = entry_point
-        self.env = env
+        self.env = deepcopy(env)
         self.predictor_cls = predictor_cls
         self.role = role
         self.name = name
@@ -963,7 +1007,7 @@ def __init__(
         self.deserializer = deserializer
         self.accelerator_type = accelerator_type
         self.endpoint_name = endpoint_name
-        self.tags = tags
+        self.tags = deepcopy(tags)
         self.kms_key = kms_key
         self.wait = wait
         self.data_capture_config = data_capture_config
@@ -1111,8 +1155,8 @@ def __init__(
         self.model_uri = model_uri
         self.source_dir = source_dir
         self.entry_point = entry_point
-        self.hyperparameters = hyperparameters
-        self.metric_definitions = metric_definitions
+        self.hyperparameters = deepcopy(hyperparameters)
+        self.metric_definitions = deepcopy(metric_definitions)
         self.role = role
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
         self.volume_size = volume_size
@@ -1123,7 +1167,7 @@ def __init__(
         self.output_kms_key = output_kms_key
         self.base_job_name = base_job_name
         self.sagemaker_session = sagemaker_session
-        self.tags = tags
+        self.tags = deepcopy(tags)
         self.subnets = subnets
         self.security_group_ids = security_group_ids
         self.model_channel_name = model_channel_name
@@ -1139,7 +1183,7 @@ def __init__(
         self.enable_sagemaker_metrics = enable_sagemaker_metrics
         self.profiler_config = profiler_config
         self.disable_profiler = disable_profiler
-        self.environment = environment
+        self.environment = deepcopy(environment)
         self.max_retry_attempts = max_retry_attempts
         self.git_config = git_config
         self.container_log_level = container_log_level
@@ -1319,13 +1363,13 @@ def __init__(
         self.image_uri = image_uri
         self.source_dir = source_dir
         self.entry_point = entry_point
-        self.env = env
+        self.env = deepcopy(env)
         self.predictor_cls = predictor_cls
         self.serializer = serializer
         self.deserializer = deserializer
         self.accelerator_type = accelerator_type
         self.endpoint_name = endpoint_name
-        self.tags = tags
+        self.tags = deepcopy(tags)
         self.kms_key = kms_key
         self.wait = wait
         self.data_capture_config = data_capture_config

tests/unit/sagemaker/hyperparameters/jumpstart/test_default.py

@@ -19,7 +19,7 @@
 
 from sagemaker import hyperparameters
 
-from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
+from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
 
 
 mock_client = boto3.client("s3")
@@ -116,3 +116,74 @@ def test_jumpstart_default_hyperparameters(patched_get_model_specs):
         hyperparameters.retrieve_default(
             model_id=model_id,
         )
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_sdk_hyperparameters_instance_type_overrides(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_special_model_spec
+
+    model_id = "variant-model"
+    region = "us-west-2"
+
+    # assert that we can add hyperparameters to default
+    vars = hyperparameters.retrieve_default(
+        region=region,
+        model_id=model_id,
+        model_version="*",
+        sagemaker_session=mock_session,
+        instance_type="ml.p2.48xlarge",
+    )
+    assert vars == {
+        "adam-learning-rate": "0.05",
+        "batch-size": "4",
+        "epochs": "3",
+        "num_bag_sets": "5",
+        "num_stack_levels": "6",
+        "refit_full": "False",
+        "sagemaker_container_log_level": "20",
+        "sagemaker_program": "transfer_learning.py",
+        "sagemaker_submit_directory": "/opt/ml/input/data/code/sourcedir.tar.gz",
+        "save_space": "False",
+        "set_best_to_refit_full": "False",
+        "verbosity": "2",
+    }
+
+    # assert that we can override default environment variables (instance family + instance type
+    # specific)
+    vars = hyperparameters.retrieve_default(
+        region=region,
+        model_id=model_id,
+        model_version="*",
+        sagemaker_session=mock_session,
+        instance_type="ml.p2.12xlarge",
+    )
+    assert vars == {
+        "adam-learning-rate": "0.05",
+        "batch-size": "1",
+        "epochs": "3",
+        "num_bag_sets": "1",
+        "num_stack_levels": "0",
+        "refit_full": "False",
+        "eval_metric": "auto",
+        "num_bag_folds": "0",
+        "presets": "medium_quality",
+        "auto_stack": "False",
+        "sagemaker_container_log_level": "20",
+        "sagemaker_program": "transfer_learning.py",
+        "sagemaker_submit_directory": "/opt/ml/input/data/code/sourcedir.tar.gz",
+        "save_space": "False",
+        "set_best_to_refit_full": "False",
+        "verbosity": "2",
+    }
+
+    # assert that we can return default hyperparameters for unrecognized instance
+    vars = hyperparameters.retrieve_default(
+        region=region,
+        model_id=model_id,
+        model_version="*",
+        sagemaker_session=mock_session,
+        instance_type="ml.p9999.48xlarge",
+    )
+
+    assert vars == {"epochs": "3", "adam-learning-rate": "0.05", "batch-size": "4"}