fix: include model channel for gated uncompressed models

Author: evakravi

src/sagemaker/jumpstart/cache.py

@@ -372,10 +372,18 @@ def _get_json_file(
         object and None when reading from the local file system.
         """
         if self._is_local_metadata_mode():
-            file_content, etag = self._get_json_file_from_local_override(key, filetype), None
-        else:
-            file_content, etag = self._get_json_file_and_etag_from_s3(key)
-        return file_content, etag
+            if filetype in {
+                JumpStartS3FileType.OPEN_WEIGHT_MANIFEST,
+                JumpStartS3FileType.OPEN_WEIGHT_SPECS,
+            }:
+                return self._get_json_file_from_local_override(key, filetype), None
+            else:
+                JUMPSTART_LOGGER.warning(
+                    "Local metadata mode is enabled, but the file type %s is not supported "
+                    "for local override. Falling back to s3.",
+                    filetype,
+                )
+        return self._get_json_file_and_etag_from_s3(key)
 
     def _get_json_md5_hash(self, key: str):
         """Retrieves md5 object hash for s3 objects, using `s3.head_object`.

src/sagemaker/jumpstart/factory/estimator.py

@@ -632,13 +632,7 @@ def _add_model_reference_arn_to_kwargs(
 
 def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets model uri in kwargs based on default or override, returns full kwargs."""
-    # hub_arn is by default None unless the user specifies the hub_name
-    # If no hub_name is specified, it is assumed the public hub
-    is_private_hub = JUMPSTART_MODEL_HUB_NAME not in kwargs.hub_arn if kwargs.hub_arn else False
-    if (
-        _model_supports_training_model_uri(**get_model_info_default_kwargs(kwargs))
-        or is_private_hub
-    ):
+    if _model_supports_training_model_uri(**get_model_info_default_kwargs(kwargs)):
         default_model_uri = model_uris.retrieve(
             model_scope=JumpStartScriptScope.TRAINING,
             instance_type=kwargs.instance_type,

src/sagemaker/jumpstart/types.py

@@ -1940,12 +1940,20 @@ def use_inference_script_uri(self) -> bool:
 
     def use_training_model_artifact(self) -> bool:
         """Returns True if the model should use a model uri when kicking off training job."""
-        # gated model never use training model artifact
-        if self.gated_bucket:
+        # old models with this environment variable present don't use model channel
+        if any(
+            self.training_instance_type_variants.get_instance_specific_gated_model_key_env_var_value(
+                instance_type
+            )
+            for instance_type in self.supported_training_instance_types
+        ):
+            return False
+
+        # even older models with training model package artifact uris present also don't use model channel
+        if len(self.training_model_package_artifact_uris or {}) > 0:
             return False
 
-        # otherwise, return true is a training model package is not set
-        return len(self.training_model_package_artifact_uris or {}) == 0
+        return getattr(self, "training_artifact_key", None) is not None
 
     def is_gated_model(self) -> bool:
         """Returns True if the model has a EULA key or the model bucket is gated."""