Nova training support

* feature: Added Amazon Nova training support for ModelTrainer and Estimator

Author: benieric

.gitignore

@@ -37,4 +37,5 @@ src/sagemaker/modules/train/container_drivers/sourcecode.json
 src/sagemaker/modules/train/container_drivers/distributed.json
 tests/data/**/_repack_model.py
 tests/data/experiment/sagemaker-dev-1.0.tar.gz
-src/sagemaker/serve/tmp_workspace
+src/sagemaker/serve/tmp_workspace
+test-examples

src/sagemaker/estimator.py

@@ -905,6 +905,30 @@ def _json_encode_hyperparameters(hyperparameters: Dict[str, Any]) -> Dict[str, A
             }
         return hyperparameters
 
+    @staticmethod
+    def _nova_encode_hyperparameters(hyperparameters: Dict[str, Any]) -> Dict[str, Any]:
+        """Applies JSON encoding for Nova job hyperparameters, preserving string values.
+
+        For Nova jobs, string values should not be JSON-encoded.
+
+        Args:
+            hyperparameters (dict): Dictionary of hyperparameters.
+
+        Returns:
+            dict: Dictionary with encoded hyperparameters.
+        """
+        current_hyperparameters = hyperparameters
+        if current_hyperparameters is not None:
+            hyperparameters = {}
+            for k, v in current_hyperparameters.items():
+                if is_pipeline_variable(v):
+                    hyperparameters[str(k)] = v.to_string()
+                elif isinstance(v, str):
+                    hyperparameters[str(k)] = v
+                else:
+                    hyperparameters[str(k)] = json.dumps(v)
+        return hyperparameters
+
     def _prepare_for_training(self, job_name=None):
         """Set any values in the estimator that need to be set before training.
 
@@ -938,7 +962,11 @@ def _prepare_for_training(self, job_name=None):
             self.source_dir = updated_paths["source_dir"]
             self.dependencies = updated_paths["dependencies"]
 
-        if self.source_dir or self.entry_point or self.dependencies:
+        if (
+            self.source_dir
+            or self.entry_point
+            or (self.dependencies and len(self.dependencies) > 0)
+        ):
             # validate source dir will raise a ValueError if there is something wrong with
             # the source directory. We are intentionally not handling it because this is a
             # critical error.
@@ -3579,7 +3607,11 @@ def __init__(
             git_config=git_config,
             enable_network_isolation=enable_network_isolation,
         )
-        if not is_pipeline_variable(entry_point) and entry_point.startswith("s3://"):
+        if (
+            not is_pipeline_variable(entry_point)
+            and entry_point is not None
+            and entry_point.startswith("s3://")
+        ):
             raise ValueError(
                 "Invalid entry point script: {}. Must be a path to a local file.".format(
                     entry_point
@@ -3599,6 +3631,7 @@ def __init__(
         self.checkpoint_s3_uri = checkpoint_s3_uri
         self.checkpoint_local_path = checkpoint_local_path
         self.enable_sagemaker_metrics = enable_sagemaker_metrics
+        self.is_nova_job = kwargs.get("is_nova_job", False)
 
     def _prepare_for_training(self, job_name=None):
         """Set hyperparameters needed for training. This method will also validate ``source_dir``.
@@ -3713,7 +3746,10 @@ def _model_entry_point(self):
 
     def set_hyperparameters(self, **kwargs):
         """Escapes the dict argument as JSON, updates the private hyperparameter attribute."""
-        self._hyperparameters.update(EstimatorBase._json_encode_hyperparameters(kwargs))
+        if self.is_nova_job:
+            self._hyperparameters.update(EstimatorBase._nova_encode_hyperparameters(kwargs))
+        else:
+            self._hyperparameters.update(EstimatorBase._json_encode_hyperparameters(kwargs))
 
     def hyperparameters(self):
         """Returns the hyperparameters as a dictionary to use for training.
@@ -3724,7 +3760,10 @@ def hyperparameters(self):
         Returns:
             dict[str, str]: The hyperparameters.
         """
-        return EstimatorBase._json_encode_hyperparameters(self._hyperparameters)
+        if self.is_nova_job:
+            return EstimatorBase._nova_encode_hyperparameters(self._hyperparameters)
+        else:
+            return EstimatorBase._json_encode_hyperparameters(self._hyperparameters)
 
     @classmethod
     def _prepare_init_params_from_job_description(cls, job_details, model_channel_name=None):

src/sagemaker/fw_utils.py

@@ -1063,7 +1063,7 @@ def validate_torch_distributed_distribution(
             )
 
     # Check entry point type
-    if not entry_point.endswith(".py"):
+    if entry_point is not None and not entry_point.endswith(".py"):
         err_msg += (
             "Unsupported entry point type for the distribution torch_distributed.\n"
             "Only python programs (*.py) are supported."

src/sagemaker/modules/constants.py

@@ -25,6 +25,10 @@
     os.path.dirname(os.path.abspath(__file__)), "train/container_drivers"
 )
 
+SM_RECIPE = "recipe"
+SM_RECIPE_YAML = "recipe.yaml"
+SM_RECIPE_CONTAINER_PATH = f"/opt/ml/input/data/recipe/{SM_RECIPE_YAML}"
+
 SOURCE_CODE_JSON = "sourcecode.json"
 DISTRIBUTED_JSON = "distributed.json"
 TRAIN_SCRIPT = "sm_train.sh"

src/sagemaker/modules/train/model_trainer.py

@@ -85,6 +85,9 @@
     SM_CODE_CONTAINER_PATH,
     SM_DRIVERS,
     SM_DRIVERS_LOCAL_PATH,
+    SM_RECIPE,
+    SM_RECIPE_YAML,
+    SM_RECIPE_CONTAINER_PATH,
     TRAIN_SCRIPT,
     DEFAULT_CONTAINER_ENTRYPOINT,
     DEFAULT_CONTAINER_ARGUMENTS,
@@ -100,7 +103,12 @@
 from sagemaker.telemetry.telemetry_logging import _telemetry_emitter
 from sagemaker.telemetry.constants import Feature
 from sagemaker.modules import logger
-from sagemaker.modules.train.sm_recipes.utils import _get_args_from_recipe, _determine_device_type
+from sagemaker.modules.train.sm_recipes.utils import (
+    _get_args_from_recipe,
+    _determine_device_type,
+    _is_nova_recipe,
+    _load_base_recipe,
+)
 
 
 class Mode(Enum):
@@ -242,6 +250,7 @@ class ModelTrainer(BaseModel):
     _remote_debug_config: Optional[RemoteDebugConfig] = PrivateAttr(default=None)
     _metric_definitions: Optional[List[MetricDefinition]] = PrivateAttr(default=None)
 
+    _is_nova_recipe: Optional[bool] = PrivateAttr(default=None)
     _temp_recipe_train_dir: Optional[TemporaryDirectory] = PrivateAttr(default=None)
 
     CONFIGURABLE_ATTRIBUTES: ClassVar[List[str]] = [
@@ -449,6 +458,33 @@ def _validate_source_code(self, source_code: Optional[SourceCode]):
                                 + "Must be a valid file within the 'source_dir'.",
                             )
 
+    @staticmethod
+    def _validate_and_load_hyperparameters_file(hyperparameters_file: str) -> Dict[str, Any]:
+        """Validate the hyperparameters file."""
+        if not os.path.exists(hyperparameters_file):
+            raise ValueError(f"Hyperparameters file not found: {hyperparameters_file}")
+        logger.info(f"Loading hyperparameters from file: {hyperparameters_file}")
+        with open(hyperparameters_file, "r") as f:
+            contents = f.read()
+            try:
+                hyperparameters = json.loads(contents)
+                logger.debug("Hyperparameters loaded as JSON")
+                return hyperparameters
+            except json.JSONDecodeError:
+                try:
+                    logger.info(f"contents: {contents}")
+                    hyperparameters = yaml.safe_load(contents)
+                    if not isinstance(hyperparameters, dict):
+                        raise ValueError("YAML contents must be a valid mapping")
+                    logger.info(f"hyperparameters: {hyperparameters}")
+                    logger.debug("Hyperparameters loaded as YAML")
+                    return hyperparameters
+                except (yaml.YAMLError, ValueError):
+                    raise ValueError(
+                        f"Invalid hyperparameters file: {hyperparameters_file}. "
+                        "Must be a valid JSON or YAML file."
+                    )
+
     def model_post_init(self, __context: Any):
         """Post init method to perform custom validation and set default values."""
         self._validate_training_image_and_algorithm_name(self.training_image, self.algorithm_name)
@@ -510,27 +546,9 @@ def model_post_init(self, __context: Any):
             )
 
         if self.hyperparameters and isinstance(self.hyperparameters, str):
-            if not os.path.exists(self.hyperparameters):
-                raise ValueError(f"Hyperparameters file not found: {self.hyperparameters}")
-            logger.info(f"Loading hyperparameters from file: {self.hyperparameters}")
-            with open(self.hyperparameters, "r") as f:
-                contents = f.read()
-                try:
-                    self.hyperparameters = json.loads(contents)
-                    logger.debug("Hyperparameters loaded as JSON")
-                except json.JSONDecodeError:
-                    try:
-                        logger.info(f"contents: {contents}")
-                        self.hyperparameters = yaml.safe_load(contents)
-                        if not isinstance(self.hyperparameters, dict):
-                            raise ValueError("YAML contents must be a valid mapping")
-                        logger.info(f"hyperparameters: {self.hyperparameters}")
-                        logger.debug("Hyperparameters loaded as YAML")
-                    except (yaml.YAMLError, ValueError):
-                        raise ValueError(
-                            f"Invalid hyperparameters file: {self.hyperparameters}. "
-                            "Must be a valid JSON or YAML file."
-                        )
+            self.hyperparameters = self._validate_and_load_hyperparameters_file(
+                self.hyperparameters
+            )
 
         if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB:
             if self.output_data_config is None:
@@ -613,6 +631,22 @@ def train(
 
             final_input_data_config = list(existing_channels.values()) + new_channels
 
+        if self._is_nova_recipe:
+            for input_data in final_input_data_config:
+                if input_data.channel_name == SM_RECIPE:
+                    raise ValueError(
+                        "Cannot use reserved channel name 'recipe' as an input channel name "
+                        " for Nova Recipe"
+                    )
+            recipe_file_path = os.path.join(self._temp_recipe_train_dir.name, SM_RECIPE_YAML)
+            recipe_channel = self.create_input_data_channel(
+                channel_name=SM_RECIPE,
+                data_source=recipe_file_path,
+                key_prefix=input_data_key_prefix,
+            )
+            final_input_data_config.append(recipe_channel)
+            self.hyperparameters.update({"sagemaker_recipe_local_path": SM_RECIPE_CONTAINER_PATH})
+
         if final_input_data_config:
             final_input_data_config = self._get_input_data_config(
                 final_input_data_config, input_data_key_prefix
@@ -1005,6 +1039,7 @@ def from_recipe(
         checkpoint_config: Optional[shapes.CheckpointConfig] = None,
         training_input_mode: Optional[str] = "File",
         environment: Optional[Dict[str, str]] = None,
+        hyperparameters: Optional[Union[Dict[str, Any], str]] = {},
         tags: Optional[List[Tag]] = None,
         sagemaker_session: Optional[Session] = None,
         role: Optional[str] = None,
@@ -1101,14 +1136,21 @@ def from_recipe(
         """
         if compute.instance_type is None:
             raise ValueError(
-                "Must set ``instance_type`` in compute_config when using training recipes."
+                "Must set ``instance_type`` in ``compute`` input when using training recipes."
             )
         device_type = _determine_device_type(compute.instance_type)
-        if device_type == "cpu":
+        recipe = _load_base_recipe(
+            training_recipe=training_recipe, recipe_overrides=recipe_overrides
+        )
+        is_nova = _is_nova_recipe(recipe=recipe)
+
+        if device_type == "cpu" and not is_nova:
             raise ValueError(
-                "Training recipes are not supported for CPU instances. "
+                "Training recipe is not supported for CPU instances. "
                 + "Please provide a GPU or Tranium instance type."
             )
+        if training_image is None and is_nova:
+            raise ValueError("training_image must be provided when using recipe for Nova.")
 
         if training_image_config and training_image is None:
             raise ValueError("training_image must be provided when using training_image_config.")
@@ -1126,15 +1168,27 @@ def from_recipe(
         # - distributed
         # - compute
         # - hyperparameters
-        model_trainer_args, recipe_train_dir = _get_args_from_recipe(
-            training_recipe=training_recipe,
+        model_trainer_args, tmp_dir = _get_args_from_recipe(
+            training_recipe=recipe,
             recipe_overrides=recipe_overrides,
             requirements=requirements,
             compute=compute,
             region_name=sagemaker_session.boto_region_name,
+            role=role,
         )
         if training_image is not None:
             model_trainer_args["training_image"] = training_image
+        if hyperparameters and not is_nova:
+            logger.warning(
+                "Hyperparameters are not supported for general training recipes. "
+                + "Ignoring hyperparameters input."
+            )
+        if is_nova:
+            if hyperparameters and isinstance(hyperparameters, str):
+                hyperparameters = cls._validate_and_load_hyperparameters_file(hyperparameters)
+                model_trainer_args["hyperparameters"].update(hyperparameters)
+            elif hyperparameters and isinstance(hyperparameters, dict):
+                model_trainer_args["hyperparameters"].update(hyperparameters)
 
         model_trainer = cls(
             sagemaker_session=sagemaker_session,
@@ -1151,8 +1205,8 @@ def from_recipe(
             tags=tags,
             **model_trainer_args,
         )
-
-        model_trainer._temp_recipe_train_dir = recipe_train_dir
+        model_trainer._is_nova_recipe = is_nova
+        model_trainer._temp_recipe_train_dir = tmp_dir
         return model_trainer
 
     def with_tensorboard_output_config(