Merge branch 'master' into feat/py-typed

Author: sage-maker

doc/overview.rst

@@ -1958,15 +1958,15 @@ Make sure to have a Compose Version compatible with your Docker Engine installat
 Local mode configuration
 ========================
 
-The local mode uses a YAML configuration file located at ``~/.sagemaker/config.yaml`` to define the default values that are automatically passed to the ``config`` attribute of ``LocalSession``. This is an example of the configuration, for the full schema, see `sagemaker.config.config_schema.SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA <https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/config/config_schema.py>`_.
+The local mode uses a YAML configuration file located at ``${user_config_directory}/sagemaker/config.yaml`` to define the default values that are automatically passed to the ``config`` attribute of ``LocalSession``. This is an example of the configuration, for the full schema, see `sagemaker.config.config_schema.SAGEMAKER_PYTHON_SDK_LOCAL_MODE_CONFIG_SCHEMA <https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/config/config_schema.py>`_.
 
 .. code:: yaml
 
     local:
         local_code: true # Using everything locally
         region_name: "us-west-2" # Name of the region
         container_config: # Additional docker container config
-            shm_size: "128M
+            shm_size: "128M"
 
 If you want to keep everything local, and not use Amazon S3 either, you can enable "local code" in one of two ways:
 

src/sagemaker/estimator.py

@@ -387,8 +387,8 @@ def __init__(
             source_dir (str or PipelineVariable): The absolute, relative, or S3 URI Path to
                 a directory with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. The structure within this directory is preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
+                point to a file with name ``sourcedir.tar.gz``. The structure within this directory
+                is preserved when training on Amazon SageMaker. If 'git_config' is provided,
                 'source_dir' should be a relative location to a directory in the Git
                 repo.
                 With the following GitHub repo directory structure:
@@ -3421,8 +3421,8 @@ def __init__(
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI)
                 to a directory with any other training source code dependencies aside from
                 the entry point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                are preserved when training on Amazon SageMaker. If 'git_config' is provided,
                 'source_dir' should be a relative location to a directory in the Git
                 repo.
 

src/sagemaker/fw_utils.py

@@ -252,7 +252,7 @@ def validate_source_code_input_against_pipeline_variables(
         logger.warning(
             "The source_dir is a pipeline variable: %s. During pipeline execution, "
             "the interpreted value of source_dir has to be an S3 URI and "
-            "must point to a tar.gz file",
+            "must point to a file with name ``sourcedir.tar.gz``",
             type(source_dir),
         )
 

src/sagemaker/huggingface/estimator.py

@@ -84,8 +84,8 @@ def __init__(
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to a
                 directory with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker.
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory are
+                preserved when training on Amazon SageMaker.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None). The hyperparameters are made
                 accessible as a dict[str, str] to the training code on

src/sagemaker/jumpstart/cache.py

@@ -262,7 +262,7 @@ def _model_id_retrieval_function(
             return JumpStartVersionedModelId(model_id, sm_compatible_model_version)
 
         versions_incompatible_with_sagemaker = [
-            Version(header.version)
+            header.version
             for header in manifest.values()  # type: ignore
             if header.model_id == model_id
         ]
@@ -540,9 +540,7 @@ def _select_version(
         """
 
         if version_str == "*":
-            if len(available_versions) == 0:
-                return None
-            return str(max(available_versions))
+            return utils.get_latest_version(available_versions)
 
         if model_type == JumpStartModelType.PROPRIETARY:
             if "*" in version_str:

src/sagemaker/jumpstart/estimator.py

@@ -350,8 +350,8 @@ def __init__(
             source_dir (Optional[Union[str, PipelineVariable]]): The absolute, relative, or
                 S3 URI Path to a directory with any other training source code dependencies
                 aside from the entry point file. If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory is preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                is preserved when training on Amazon SageMaker. If 'git_config' is provided,
                 'source_dir' should be a relative location to a directory in the Git
                 repo.
                 (Default: None).
@@ -947,8 +947,8 @@ def deploy(
             source_dir (Optional[str]): The absolute, relative, or S3 URI Path to a directory
                 with any other training source code dependencies aside from the entry
                 point file (Default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory is preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory is
+                preserved when training on Amazon SageMaker. If 'git_config' is provided,
                 'source_dir' should be a relative location to a directory in the Git repo.
                 If the directory points to S3, no code is uploaded and the S3 location
                 is used instead. (Default: None).

src/sagemaker/jumpstart/model.py

@@ -178,8 +178,8 @@ def __init__(
             source_dir (Optional[str]): The absolute, relative, or S3 URI Path to a directory
                 with any other training source code dependencies aside from the entry
                 point file (Default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory is preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory is
+                preserved when training on Amazon SageMaker. If 'git_config' is provided,
                 'source_dir' should be a relative location to a directory in the Git repo.
                 If the directory points to S3, no code is uploaded and the S3 location
                 is used instead. (Default: None).

src/sagemaker/jumpstart/utils.py

@@ -21,7 +21,7 @@
 from urllib.parse import urlparse
 import boto3
 from botocore.exceptions import ClientError
-from packaging.version import Version
+from packaging.version import Version, InvalidVersion
 import botocore
 from sagemaker_core.shapes import ModelAccessConfig
 import sagemaker
@@ -1630,3 +1630,11 @@ def get_draft_model_content_bucket(provider: Dict, region: str) -> str:
             return get_jumpstart_gated_content_bucket(region=region)
         return get_jumpstart_content_bucket(region=region)
     return neo_bucket
+
+
+def get_latest_version(versions: List[str]) -> Optional[str]:
+    """Returns the latest version using sem-ver when possible."""
+    try:
+        return None if not versions else max(versions, key=Version)
+    except InvalidVersion:
+        return max(versions)

src/sagemaker/model.py

@@ -215,8 +215,8 @@ def __init__(
             source_dir (str): The absolute, relative, or S3 URI Path to a directory
                 with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory is preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                is preserved when training on Amazon SageMaker. If 'git_config' is provided,
                 'source_dir' should be a relative location to a directory in the Git repo.
                 If the directory points to S3, no code is uploaded and the S3 location
                 is used instead.
@@ -1996,11 +1996,11 @@ def __init__(
             source_dir (str): Path (absolute, relative or an S3 URI) to a directory
                 with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker. If 'git_config' is provided,
-                'source_dir' should be a relative location to a directory in the Git repo.
-                If the directory points to S3, no code will be uploaded and the S3 location
-                will be used instead.
+                point to a file with name ``sourcedir.tar.gz``. Structure within this
+                directory are preserved when training on Amazon SageMaker. If 'git_config'
+                is provided, 'source_dir' should be a relative location to a directory in the
+                Git repo. If the directory points to S3, no code will be uploaded and the S3
+                location will be used instead.
 
                 .. admonition:: Example
 

src/sagemaker/mxnet/estimator.py

@@ -84,8 +84,8 @@ def __init__(
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to
                 a directory with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker.
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                are preserved when training on Amazon SageMaker.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None). The hyperparameters are made
                 accessible as a dict[str, str] to the training code on

src/sagemaker/pytorch/estimator.py

@@ -182,8 +182,8 @@ def __init__(
                 unless ``image_uri`` is provided.
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to
                 a directory with any other training source code dependencies aside from the entry
-                point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
+                point file (default: None). If ``source_dir`` is an S3 URI, it must point to a
+                file with name ``sourcedir.tar.gz``. Structure within this directory are preserved
                 when training on Amazon SageMaker. Must be a local path when using training_recipe.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None). The hyperparameters are made

src/sagemaker/rl/estimator.py

@@ -120,8 +120,8 @@ def __init__(
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI)
                 to a directory with any other training source code dependencies aside from
                 the entry point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker.
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                are preserved when training on Amazon SageMaker.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None). The hyperparameters are made
                 accessible as a dict[str, str] to the training code on

src/sagemaker/sklearn/estimator.py

@@ -83,8 +83,8 @@ def __init__(
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to
                 a directory with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker.
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                are preserved when training on Amazon SageMaker.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None). The hyperparameters are made
                 accessible as a dict[str, str] to the training code on

src/sagemaker/xgboost/estimator.py

@@ -78,8 +78,8 @@ def __init__(
             source_dir (str or PipelineVariable): Path (absolute, relative or an S3 URI) to
                 a directory with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
-                point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker.
+                point to a file with name ``sourcedir.tar.gz``. Structure within this directory
+                are preserved when training on Amazon SageMaker.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None).
                 The hyperparameters are made accessible as a dict[str, str] to the training code

tests/unit/sagemaker/jumpstart/test_cache.py

@@ -22,7 +22,10 @@
 from mock.mock import MagicMock
 import pytest
 from mock import patch
+from packaging.version import Version
 
+
+from sagemaker.jumpstart import utils
 from sagemaker.jumpstart.cache import (
     JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY,
     JUMPSTART_DEFAULT_PROPRIETARY_MANIFEST_KEY,
@@ -33,6 +36,7 @@
     ENV_VARIABLE_JUMPSTART_SPECS_LOCAL_ROOT_DIR_OVERRIDE,
 )
 from sagemaker.jumpstart.types import (
+    JumpStartCachedContentValue,
     JumpStartModelHeader,
     JumpStartModelSpecs,
     JumpStartVersionedModelId,
@@ -1119,3 +1123,124 @@ def test_jumpstart_local_metadata_override_specs_not_exist_both_directories(
             ),
         ]
     )
+
+
+@patch.object(JumpStartModelsCache, "_retrieval_function")
+def test_jumpstart_cache_handles_versioning_correctly_for_open_source_weights(
+    retrieval_function: Mock,
+):
+    sm_version = Version(utils.get_sagemaker_version())
+    new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
+    print(str(new_sm_version))
+    versions = ["1.0.0", "2.9.1", "2.16.0"]
+    manifest = [
+        {
+            "model_id": "test-model",
+            "version": version,
+            "min_version": "2.49.0",
+            "spec_key": "spec_key",
+        }
+        for version in versions
+    ]
+
+    manifest.append(
+        {
+            "model_id": "test-model",
+            "version": "3.0.0",
+            "min_version": str(new_sm_version),
+            "spec_key": "spec_key",
+        }
+    )
+
+    manifest_dict = {}
+    for header in manifest:
+        header_obj = JumpStartModelHeader(header)
+        manifest_dict[JumpStartVersionedModelId(header_obj.model_id, header_obj.version)] = (
+            header_obj
+        )
+    retrieval_function.return_value = JumpStartCachedContentValue(formatted_content=manifest_dict)
+    key = JumpStartVersionedModelId("test-model", "*")
+
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    result = cache._get_open_weight_manifest_key_from_model_id(key=key, value=None)
+
+    assert_key = JumpStartVersionedModelId("test-model", "2.16.0")
+
+    assert result == assert_key
+
+
+@patch.object(JumpStartModelsCache, "_retrieval_function")
+def test_jumpstart_cache_handles_versioning_correctly_for_proprietary_weights(
+    retrieval_function: Mock,
+):
+    sm_version = Version(utils.get_sagemaker_version())
+    new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
+    print(str(new_sm_version))
+    versions = ["1.0.0", "2.9.1", "2.16.0"]
+    manifest = [
+        {
+            "model_id": "test-model",
+            "version": version,
+            "min_version": "2.49.0",
+            "spec_key": "spec_key",
+        }
+        for version in versions
+    ]
+
+    manifest.append(
+        {
+            "model_id": "test-model",
+            "version": "3.0.0",
+            "min_version": str(new_sm_version),
+            "spec_key": "spec_key",
+        }
+    )
+
+    manifest_dict = {}
+    for header in manifest:
+        header_obj = JumpStartModelHeader(header)
+        manifest_dict[JumpStartVersionedModelId(header_obj.model_id, header_obj.version)] = (
+            header_obj
+        )
+    retrieval_function.return_value = JumpStartCachedContentValue(formatted_content=manifest_dict)
+    key = JumpStartVersionedModelId("test-model", "*")
+
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    result = cache._get_proprietary_manifest_key_from_model_id(key=key, value=None)
+
+    assert_key = JumpStartVersionedModelId("test-model", "2.16.0")
+
+    assert result == assert_key
+
+
+@patch.object(JumpStartModelsCache, "_retrieval_function")
+def test_jumpstart_cache_handles_versioning_correctly_non_sem_ver(retrieval_function: Mock):
+    sm_version = Version(utils.get_sagemaker_version())
+    new_sm_version = Version(str(sm_version.major + 1) + ".0.0")
+    print(str(new_sm_version))
+    versions = ["abc", "2.9.1", "2.16.0"]
+    manifest = [
+        {
+            "model_id": "test-model",
+            "version": version,
+            "min_version": "2.49.0",
+            "spec_key": "spec_key",
+        }
+        for version in versions
+    ]
+
+    manifest_dict = {}
+    for header in manifest:
+        header_obj = JumpStartModelHeader(header)
+        manifest_dict[JumpStartVersionedModelId(header_obj.model_id, header_obj.version)] = (
+            header_obj
+        )
+    retrieval_function.return_value = JumpStartCachedContentValue(formatted_content=manifest_dict)
+    key = JumpStartVersionedModelId("test-model", "*")
+
+    cache = JumpStartModelsCache(s3_bucket_name="some_bucket")
+    result = cache._get_open_weight_manifest_key_from_model_id(key=key, value=None)
+
+    assert_key = JumpStartVersionedModelId("test-model", "abc")
+
+    assert result == assert_key