Merge branch 'master' into master

zhaoqizqwang · web-flow · commit 8bf0026738fd · 2025-02-11T13:37:58.000-05:00
diff --git a/doc/frameworks/tensorflow/using_tf.rst b/doc/frameworks/tensorflow/using_tf.rst
@@ -246,7 +246,7 @@ Training with parameter servers
 
 If you specify parameter_server as the value of the distribution parameter, the container launches a parameter server
 thread on each instance in the training cluster, and then executes your training code. You can find more information on
-TensorFlow distributed training at `TensorFlow docs <https://www.tensorflow.org/deploy/distributed>`__.
+TensorFlow distributed training at `TensorFlow docs <https://www.tensorflow.org/guide/distributed_training>`__.
 To enable parameter server training:
 
 .. code:: python
diff --git a/doc/overview.rst b/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:
 
diff --git a/src/sagemaker/estimator.py b/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.
 
diff --git a/src/sagemaker/feature_store/feature_group.py b/src/sagemaker/feature_store/feature_group.py
@@ -631,7 +631,7 @@ def __str__(self) -> str:
 class FeatureGroup:
     """FeatureGroup definition.
 
-    This class instantiates a FeatureGroup object that comprises of a name for the FeatureGroup,
+    This class instantiates a FeatureGroup object that comprises a name for the FeatureGroup,
     session instance, and a list of feature definition objects i.e., FeatureDefinition.
 
     Attributes:
diff --git a/src/sagemaker/fw_utils.py b/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),
         )
 
diff --git a/src/sagemaker/huggingface/estimator.py b/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
diff --git a/src/sagemaker/jumpstart/cache.py b/src/sagemaker/jumpstart/cache.py
@@ -263,7 +263,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
         ]
@@ -541,9 +541,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:
diff --git a/src/sagemaker/jumpstart/estimator.py b/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).
diff --git a/src/sagemaker/jumpstart/model.py b/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).
diff --git a/src/sagemaker/jumpstart/utils.py b/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)
diff --git a/src/sagemaker/model.py b/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
 
diff --git a/src/sagemaker/mxnet/estimator.py b/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
diff --git a/src/sagemaker/processing.py b/src/sagemaker/processing.py
@@ -1416,7 +1416,7 @@ class RunArgs(object):
 class FeatureStoreOutput(ApiObject):
     """Configuration for processing job outputs in Amazon SageMaker Feature Store."""
 
-    feature_group_name = None
+    feature_group_name: Optional[str] = None
 
 
 class FrameworkProcessor(ScriptProcessor):
diff --git a/src/sagemaker/pytorch/estimator.py b/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
diff --git a/src/sagemaker/rl/estimator.py b/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
diff --git a/src/sagemaker/sklearn/estimator.py b/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
diff --git a/src/sagemaker/xgboost/estimator.py b/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
diff --git a/tests/integ/sagemaker/serve/test_base_model_builder_deploy.py b/tests/integ/sagemaker/serve/test_base_model_builder_deploy.py
@@ -12,38 +12,72 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import pytest
-
-from sagemaker import get_execution_role
-from sklearn.datasets import load_iris
-from sklearn.model_selection import train_test_split
-
 import os
+import uuid
+from typing import Generator
 
+import numpy as np
+import pandas as pd
+import pytest
+from sagemaker_core.main.resources import TrainingJob
 from sagemaker_core.main.shapes import (
     AlgorithmSpecification,
     Channel,
     DataSource,
-    S3DataSource,
     OutputDataConfig,
     ResourceConfig,
+    S3DataSource,
     StoppingCondition,
 )
-import uuid
-from sagemaker.serve.builder.model_builder import ModelBuilder
-import pandas as pd
-import numpy as np
-from sagemaker.serve import InferenceSpec, SchemaBuilder
-from sagemaker_core.main.resources import TrainingJob
+from sklearn.datasets import load_iris
+from sklearn.model_selection import train_test_split
 from xgboost import XGBClassifier
 
-from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
-
-from sagemaker.s3_utils import s3_path_join
+from sagemaker import get_execution_role
 from sagemaker.async_inference import AsyncInferenceConfig
+from sagemaker.s3_utils import s3_path_join
+from sagemaker.serve import InferenceSpec, SchemaBuilder
+from sagemaker.serve.builder.model_builder import ModelBuilder
+from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
 from tests.integ.utils import cleanup_model_resources
 
 
+@pytest.fixture(autouse=True)
+def cleanup_endpoints(mb_sagemaker_session) -> Generator[None, None, None]:
+    """Clean up any existing endpoints before and after tests."""
+    sagemaker_client = mb_sagemaker_session.sagemaker_client
+
+    # Pre-test cleanup
+    try:
+        endpoints = sagemaker_client.list_endpoints()
+        for endpoint in endpoints["Endpoints"]:
+            try:
+                sagemaker_client.delete_endpoint(EndpointName=endpoint["EndpointName"])
+                sagemaker_client.delete_endpoint_config(
+                    EndpointConfigName=endpoint["EndpointConfigName"]
+                )
+            except Exception as e:
+                print(f"Error cleaning up endpoint {endpoint['EndpointName']}: {e}")
+    except Exception as e:
+        print(f"Error listing endpoints: {e}")
+
+    yield
+
+    # Post-test cleanup
+    try:
+        endpoints = sagemaker_client.list_endpoints()
+        for endpoint in endpoints["Endpoints"]:
+            try:
+                sagemaker_client.delete_endpoint(EndpointName=endpoint["EndpointName"])
+                sagemaker_client.delete_endpoint_config(
+                    EndpointConfigName=endpoint["EndpointConfigName"]
+                )
+            except Exception as e:
+                print(f"Error cleaning up endpoint {endpoint['EndpointName']}: {e}")
+    except Exception as e:
+        print(f"Error listing endpoints: {e}")
+
+
 @pytest.fixture(scope="module")
 def xgboost_model_builder(mb_sagemaker_session):
     sagemaker_session = mb_sagemaker_session
diff --git a/tests/unit/sagemaker/jumpstart/test_cache.py b/tests/unit/sagemaker/jumpstart/test_cache.py
diff --git a/tests/unit/sagemaker/jumpstart/test_utils.py b/tests/unit/sagemaker/jumpstart/test_utils.py

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