feat: meta llama fine tuning (#4021)

* feat: meta llama fine tuning

* chore: simplify code

* chore: rename function to _retrieve_model_package_model_artifact_s3_uri

* chore: add unit tests, fixes

* fix: env var collisions, unit test

* chore: de morgans law simplification

* chore: address PR comments

* chore: improve docstring

* fix: better metadata for unit tests

* fix: flake8 issues

Author: evakravi

src/sagemaker/jumpstart/artifacts/__init__.py

@@ -11,16 +11,21 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """This module imports all JumpStart artifact functions from the respective sub-module."""
-from sagemaker.jumpstart.artifacts.prepack import _model_supports_prepacked_inference  # noqa: F401
 from sagemaker.jumpstart.artifacts.resource_names import (  # noqa: F401
     _retrieve_resource_name_base,
 )
 from sagemaker.jumpstart.artifacts.incremental_training import (  # noqa: F401
     _model_supports_incremental_training,
 )
 from sagemaker.jumpstart.artifacts.image_uris import _retrieve_image_uri  # noqa: F401
-from sagemaker.jumpstart.artifacts.script_uris import _retrieve_script_uri  # noqa: F401
-from sagemaker.jumpstart.artifacts.model_uris import _retrieve_model_uri  # noqa: F401
+from sagemaker.jumpstart.artifacts.script_uris import (  # noqa: F401
+    _retrieve_script_uri,
+    _model_supports_inference_script_uri,
+)
+from sagemaker.jumpstart.artifacts.model_uris import (  # noqa: F401
+    _retrieve_model_uri,
+    _model_supports_training_model_uri,
+)
 from sagemaker.jumpstart.artifacts.hyperparameters import (  # noqa: F401
     _retrieve_default_hyperparameters,
 )
@@ -52,4 +57,7 @@
     _retrieve_supported_accept_types,
     _retrieve_supported_content_types,
 )
-from sagemaker.jumpstart.artifacts.model_packages import _retrieve_model_package_arn  # noqa: F401
+from sagemaker.jumpstart.artifacts.model_packages import (  # noqa: F401
+    _retrieve_model_package_arn,
+    _retrieve_model_package_model_artifact_s3_uri,
+)

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -75,3 +75,61 @@ def _retrieve_model_package_arn(
         return regional_arn
 
     raise NotImplementedError(f"Model Package ARN not supported for scope: '{scope}'")
+
+
+def _retrieve_model_package_model_artifact_s3_uri(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    scope: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+) -> Optional[str]:
+    """Retrieves s3 artifact uri associated with model package.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the model package artifact.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            model package artifact.
+        region (Optional[str]): Region for which to retrieve the model package artifact.
+            (Default: None).
+        scope (Optional[str]): Scope for which to retrieve the model package artifact.
+            (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+
+    Returns:
+        str: the model package artifact uri to use for the model or None.
+
+    Raises:
+        NotImplementedError: If an unsupported script is used.
+    """
+
+    if scope == JumpStartScriptScope.TRAINING:
+
+        if region is None:
+            region = JUMPSTART_DEFAULT_REGION_NAME
+
+        model_specs = verify_model_region_and_return_specs(
+            model_id=model_id,
+            version=model_version,
+            scope=scope,
+            region=region,
+            tolerate_vulnerable_model=tolerate_vulnerable_model,
+            tolerate_deprecated_model=tolerate_deprecated_model,
+        )
+
+        if model_specs.training_model_package_artifact_uris is None:
+            return None
+
+        model_s3_uri = model_specs.training_model_package_artifact_uris.get(region)
+
+        return model_s3_uri
+
+    raise NotImplementedError(f"Model Package Artifact URI not supported for scope: '{scope}'")

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -92,3 +92,45 @@ def _retrieve_model_uri(
     model_s3_uri = f"s3://{bucket}/{model_artifact_key}"
 
     return model_s3_uri
+
+
+def _model_supports_training_model_uri(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+) -> bool:
+    """Returns True if the model supports training with model uri field.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the support status for model uri with training.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            support status for model uri with training.
+        region (Optional[str]): Region for which to retrieve the
+            support status for model uri with training.
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+    Returns:
+        bool: the support status for model uri with training.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        scope=JumpStartScriptScope.TRAINING,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+    )
+
+    return model_specs.use_training_model_artifact()

src/sagemaker/jumpstart/artifacts/prepack.py

@@ -90,3 +90,45 @@ def _retrieve_script_uri(
     script_s3_uri = f"s3://{bucket}/{model_script_key}"
 
     return script_s3_uri
+
+
+def _model_supports_inference_script_uri(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+) -> bool:
+    """Returns True if the model supports inference with script uri field.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the support status for script uri with inference.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            support status for script uri with inference.
+        region (Optional[str]): Region for which to retrieve the
+            support status for script uri with inference.
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+    Returns:
+        bool: the support status for script uri with inference.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        scope=JumpStartScriptScope.INFERENCE,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+    )
+
+    return model_specs.use_inference_script_uri()

src/sagemaker/jumpstart/artifacts/script_uris.py

@@ -147,6 +147,7 @@
 
 JUMPSTART_RESOURCE_BASE_NAME = "sagemaker-jumpstart"
 
+SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY = "SageMakerGatedModelS3Uri"
 
 CONTENT_TYPE_TO_SERIALIZER_TYPE_MAP: Dict[MIMEType, SerializerType] = {
     MIMEType.X_IMAGE: SerializerType.RAW_BYTES,

src/sagemaker/jumpstart/constants.py

@@ -23,7 +23,10 @@
     model_uris,
     script_uris,
 )
-from sagemaker.jumpstart.artifacts import _model_supports_incremental_training
+from sagemaker.jumpstart.artifacts import (
+    _model_supports_incremental_training,
+    _retrieve_model_package_model_artifact_s3_uri,
+)
 from sagemaker.jumpstart.artifacts.resource_names import _retrieve_resource_name_base
 from sagemaker.session import Session
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
@@ -37,11 +40,13 @@
 from sagemaker.jumpstart.artifacts import (
     _retrieve_estimator_init_kwargs,
     _retrieve_estimator_fit_kwargs,
+    _model_supports_training_model_uri,
 )
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
     JUMPSTART_LOGGER,
     TRAINING_ENTRY_POINT_SCRIPT_NAME,
+    SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY,
 )
 from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.factory import model
@@ -187,6 +192,7 @@ def get_init_kwargs(
     estimator_init_kwargs = _add_metric_definitions_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_estimator_extra_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_role_to_kwargs(estimator_init_kwargs)
+    estimator_init_kwargs = _add_env_to_kwargs(estimator_init_kwargs)
 
     return estimator_init_kwargs
 
@@ -446,32 +452,39 @@ def _add_image_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartE
 def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets model uri in kwargs based on default or override, returns full kwargs."""
 
-    default_model_uri = model_uris.retrieve(
-        model_scope=JumpStartScriptScope.TRAINING,
+    if _model_supports_training_model_uri(
         model_id=kwargs.model_id,
         model_version=kwargs.model_version,
+        region=kwargs.region,
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
-    )
-
-    if (
-        kwargs.model_uri is not None
-        and kwargs.model_uri != default_model_uri
-        and not _model_supports_incremental_training(
+    ):
+        default_model_uri = model_uris.retrieve(
+            model_scope=JumpStartScriptScope.TRAINING,
             model_id=kwargs.model_id,
             model_version=kwargs.model_version,
-            region=kwargs.region,
             tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
             tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
         )
-    ):
-        JUMPSTART_LOGGER.warning(
-            "'%s' does not support incremental training but is being trained with"
-            " non-default model artifact.",
-            kwargs.model_id,
-        )
 
-    kwargs.model_uri = kwargs.model_uri or default_model_uri
+        if (
+            kwargs.model_uri is not None
+            and kwargs.model_uri != default_model_uri
+            and not _model_supports_incremental_training(
+                model_id=kwargs.model_id,
+                model_version=kwargs.model_version,
+                region=kwargs.region,
+                tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+                tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+            )
+        ):
+            JUMPSTART_LOGGER.warning(
+                "'%s' does not support incremental training but is being trained with"
+                " non-default model artifact.",
+                kwargs.model_id,
+            )
+
+        kwargs.model_uri = kwargs.model_uri or default_model_uri
 
     return kwargs
 
@@ -501,6 +514,31 @@ def _add_source_dir_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStart
     return kwargs
 
 
+def _add_env_to_kwargs(
+    kwargs: JumpStartEstimatorInitKwargs,
+) -> JumpStartEstimatorInitKwargs:
+    """Sets environment in kwargs based on default or override, returns full kwargs."""
+
+    model_package_artifact_uri = _retrieve_model_package_model_artifact_s3_uri(
+        model_id=kwargs.model_id,
+        model_version=kwargs.model_version,
+        region=kwargs.region,
+        scope=JumpStartScriptScope.TRAINING,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+    )
+
+    if model_package_artifact_uri:
+        if kwargs.environment is None:
+            kwargs.environment = {}
+        kwargs.environment = {
+            **{SAGEMAKER_GATED_MODEL_S3_URI_TRAINING_ENV_VAR_KEY: model_package_artifact_uri},
+            **kwargs.environment,
+        }
+
+    return kwargs
+
+
 def _add_entry_point_to_kwargs(
     kwargs: JumpStartEstimatorInitKwargs,
 ) -> JumpStartEstimatorInitKwargs: