fix: support legacy training models, fix cache override for unsupported files

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/types.py

@@ -1940,9 +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."""
+        # 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."""