feat: add build/deploy support for fine-tuned JS models (#1473)

* feat: add support for fine-tuned JS models

* Refactor

* Refactor

* Refactor

* Refactor

* pylint

* pylint

---------

Co-authored-by: Jacky Lee <[email protected]>

Author: grenmester

src/sagemaker/enums.py

@@ -46,3 +46,5 @@ class Tag(str, Enum):
     """Enum class for tag keys to apply to models."""
 
     OPTIMIZATION_JOB_NAME = "sagemaker-sdk:optimization-job-name"
+    FINE_TUNING_MODEL_PATH = "sagemaker-sdk:fine-tuning-model-path"
+    FINE_TUNING_JOB_NAME = "sagemaker-sdk:fine-tuning-job-name"

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -14,11 +14,15 @@
 from __future__ import absolute_import
 
 import copy
+import re
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
 from typing import Type, Any, List, Dict, Optional
 import logging
 
+from botocore.exceptions import ClientError
+
+from sagemaker.enums import Tag
 from sagemaker.jumpstart import enums
 from sagemaker.jumpstart.utils import verify_model_region_and_return_specs, get_eula_message
 from sagemaker.model import Model
@@ -105,6 +109,7 @@ def __init__(self):
         self.nb_instance_type = None
         self.ram_usage_model_load = None
         self.jumpstart = None
+        self.model_metadata = None
 
     @abstractmethod
     def _prepare_for_mode(self):
@@ -520,6 +525,54 @@ def list_deployment_configs(self) -> List[Dict[str, Any]]:
 
         return self.pysdk_model.list_deployment_configs()
 
+    def _is_fine_tuned_model(self) -> bool:
+        """Checks whether a fine-tuned model exists."""
+        return self.model_metadata and (
+            self.model_metadata.get("FINE_TUNING_MODEL_PATH")
+            or self.model_metadata.get("FINE_TUNING_JOB_NAME")
+        )
+
+    def _update_model_data_for_fine_tuned_model(self, pysdk_model: Type[Model]) -> Type[Model]:
+        """Set the model path and data and add fine-tuning tags for the model."""
+        # TODO: determine precedence of FINE_TUNING_MODEL_PATH and FINE_TUNING_JOB_NAME
+        if fine_tuning_model_path := self.model_metadata.get("FINE_TUNING_MODEL_PATH"):
+            if not re.match("^(https|s3)://([^/]+)/?(.*)$", fine_tuning_model_path):
+                raise ValueError(
+                    f"Invalid path for FINE_TUNING_MODEL_PATH: {fine_tuning_model_path}."
+                )
+            pysdk_model.model_data["S3DataSource"]["S3Uri"] = fine_tuning_model_path
+            pysdk_model.add_tags(
+                {"key": Tag.FINE_TUNING_MODEL_PATH, "value": fine_tuning_model_path}
+            )
+            return pysdk_model
+
+        if fine_tuning_job_name := self.model_metadata.get("FINE_TUNING_JOB_NAME"):
+            try:
+                response = self.sagemaker_session.sagemaker_client.describe_training_job(
+                    TrainingJobName=fine_tuning_job_name
+                )
+                fine_tuning_model_path = response["OutputDataConfig"]["S3OutputPath"]
+                pysdk_model.model_data["S3DataSource"]["S3Uri"] = fine_tuning_model_path
+                pysdk_model.model_data["S3DataSource"]["CompressionType"] = response[
+                    "OutputDataConfig"
+                ]["CompressionType"]
+                pysdk_model.add_tags(
+                    [
+                        {"key": Tag.FINE_TUNING_JOB_NAME, "value": fine_tuning_job_name},
+                        {"key": Tag.FINE_TUNING_MODEL_PATH, "value": fine_tuning_model_path},
+                    ]
+                )
+                return pysdk_model
+            except ClientError:
+                raise ValueError(
+                    f"Invalid job name for FINE_TUNING_JOB_NAME: {fine_tuning_job_name}."
+                )
+
+        raise ValueError(
+            "Input model not found. Please provide either `model_path`, or "
+            "`FINE_TUNING_MODEL_PATH` or `FINE_TUNING_JOB_NAME` under `model_metadata`."
+        )
+
     def _build_for_jumpstart(self):
         """Placeholder docstring"""
         if hasattr(self, "pysdk_model") and self.pysdk_model is not None:
@@ -534,6 +587,9 @@ def _build_for_jumpstart(self):
 
         logger.info("JumpStart ID %s is packaged with Image URI: %s", self.model, image_uri)
 
+        if self._is_fine_tuned_model():
+            pysdk_model = self._update_model_data_for_fine_tuned_model(pysdk_model)
+
         if self._is_gated_model(pysdk_model) and self.mode != Mode.SAGEMAKER_ENDPOINT:
             raise ValueError(
                 "JumpStart Gated Models are only supported in SAGEMAKER_ENDPOINT mode."
@@ -714,7 +770,7 @@ def _optimize_for_jumpstart(
             **create_optimization_job_args
         )
 
-    def _is_gated_model(self, model) -> bool:
+    def _is_gated_model(self, model: Model) -> bool:
         """Determine if ``this`` Model is Gated
 
         Args:

src/sagemaker/serve/builder/model_builder.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """Holds the ModelBuilder class and the ModelServer enum."""
 from __future__ import absolute_import
+
 import uuid
 from typing import Any, Type, List, Dict, Optional, Union
 from dataclasses import dataclass, field
@@ -278,8 +279,9 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
         default=None,
         metadata={
             "help": "Define the model metadata to override, currently supports `HF_TASK`, "
-            "`MLFLOW_MODEL_PATH`. HF_TASK should be set for new models without task metadata in "
-            "the Hub, Adding unsupported task types will throw an exception"
+            "`MLFLOW_MODEL_PATH`, `FINE_TUNING_MODEL_PATH`, and `FINE_TUNING_JOB_NAME`. HF_TASK "
+            "should be set for new models without task metadata in the Hub, Adding unsupported "
+            "task types will throw an exception."
         },
     )
 
@@ -739,8 +741,8 @@ def build(  # pylint: disable=R0911
         )
 
         self.serve_settings = self._get_serve_setting()
-
         self._is_custom_image_uri = self.image_uri is not None
+
         self._is_mlflow_model = self._check_if_input_is_mlflow_model()
         if self._is_mlflow_model:
             logger.warning(
@@ -925,7 +927,7 @@ def _try_fetch_gpu_info(self):
                 f"Unable to determine single GPU size for instance: [{self.instance_type}]"
             )
 
-    def optimize(self, *args, **kwargs) -> Type[Model]:
+    def optimize(self, *args, **kwargs) -> Model:
         """Runs a model optimization job.
 
         Args:
@@ -948,7 +950,7 @@ def optimize(self, *args, **kwargs) -> Type[Model]:
                 function creates one using the default AWS configuration chain.
 
         Returns:
-            Type[Model]: A deployable ``Model`` object.
+            Model: A deployable ``Model`` object.
         """
         # need to get telemetry_opt_out info before telemetry decorator is called
         self.serve_settings = self._get_serve_setting()
@@ -972,7 +974,7 @@ def _model_builder_optimize_wrapper(
         kms_key: Optional[str] = None,
         max_runtime_in_sec: Optional[int] = None,
         sagemaker_session: Optional[Session] = None,
-    ) -> Type[Model]:
+    ) -> Model:
         """Runs a model optimization job.
 
         Args:
@@ -1002,7 +1004,7 @@ def _model_builder_optimize_wrapper(
                 function creates one using the default AWS configuration chain.
 
         Returns:
-            Type[Model]: A deployable ``Model`` object.
+            Model: A deployable ``Model`` object.
         """
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
         self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)