sagemaker 2.232.1__tar.gz → 2.232.2__tar.gz

{sagemaker-2.232.1 → sagemaker-2.232.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: sagemaker
-Version: 2.232.1
+Version: 2.232.2
 Summary: Open source library for training and deploying models on Amazon SageMaker.
 Project-URL: Homepage, https://github.com/aws/sagemaker-python-sdk
 Author: Amazon Web Services
@@ -34,6 +34,7 @@ Requires-Dist: psutil
 Requires-Dist: pyyaml~=6.0
 Requires-Dist: requests
 Requires-Dist: sagemaker-core<2.0.0,>=1.0.0
+Requires-Dist: sagemaker-mlflow
 Requires-Dist: schema
 Requires-Dist: smdebug-rulesconfig==1.0.1
 Requires-Dist: tblib<4,>=1.7.0

sagemaker-2.232.2/VERSION

@@ -0,0 +1 @@
+2.232.2

{sagemaker-2.232.1 → sagemaker-2.232.2}/pyproject.toml

@@ -48,6 +48,7 @@ dependencies = [
   "PyYAML~=6.0",
   "requests",
   "sagemaker-core>=1.0.0,<2.0.0",
+  "sagemaker-mlflow",
   "schema",
   "smdebug_rulesconfig==1.0.1",
   "tblib>=1.7.0,<4",

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/estimator.py

@@ -107,6 +107,8 @@ from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
 
+from sagemaker.mlflow.forward_sagemaker_metrics import log_sagemaker_job_to_mlflow
+
 logger = logging.getLogger(__name__)
 
 
@@ -1366,8 +1368,14 @@ class EstimatorBase(with_metaclass(ABCMeta, object)):  # pylint: disable=too-man
         experiment_config = check_and_get_run_experiment_config(experiment_config)
         self.latest_training_job = _TrainingJob.start_new(self, inputs, experiment_config)
         self.jobs.append(self.latest_training_job)
+        forward_to_mlflow_tracking_server = False
+        if os.environ.get("MLFLOW_TRACKING_URI") and self.enable_network_isolation():
+            wait = True
+            forward_to_mlflow_tracking_server = True
         if wait:
             self.latest_training_job.wait(logs=logs)
+        if forward_to_mlflow_tracking_server:
+            log_sagemaker_job_to_mlflow(self.latest_training_job.name)
 
     def _compilation_job_name(self):
         """Placeholder docstring"""

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/huggingface/model.py

@@ -334,11 +334,7 @@ class HuggingFaceModel(FrameworkModel):
             container_startup_health_check_timeout=container_startup_health_check_timeout,
             inference_recommendation_id=inference_recommendation_id,
             explainer_config=explainer_config,
-            endpoint_logging=kwargs.get("endpoint_logging", False),
-            endpoint_type=kwargs.get("endpoint_type", None),
-            resources=kwargs.get("resources", None),
-            managed_instance_scaling=kwargs.get("managed_instance_scaling", None),
-            routing_config=kwargs.get("routing_config", None),
+            **kwargs,
         )
 
     def register(

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/image_uri_config/pytorch.json

@@ -1407,7 +1407,8 @@
             "2.0": "2.0.1",
             "2.1": "2.1.0",
             "2.2": "2.2.0",
-            "2.3": "2.3.0"
+            "2.3": "2.3.0",
+            "2.4": "2.4.0"
         },
         "versions": {
             "0.4.0": {
@@ -2429,6 +2430,47 @@
                     "us-west-2": "763104351884"
                 },
                 "repository": "pytorch-training"
+            },
+            "2.4.0": {
+                "py_versions": [
+                    "py311"
+                ],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ap-southeast-4": "457447274322",
+                    "ca-central-1": "763104351884",
+                    "ca-west-1": "204538143572",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "me-central-1": "914824155844",
+                    "me-south-1": "217643126080",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884"
+                },
+                "repository": "pytorch-training"
             }
         }
     }

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/jumpstart/cache.py

@@ -293,7 +293,8 @@ class JumpStartModelsCache:
             raise KeyError(error_msg)
 
         error_msg = f"Unable to find model manifest for '{model_id}' with version '{version}'. "
-        error_msg += f"Visit {MODEL_ID_LIST_WEB_URL} for updated list of models. "
+        error_msg += "Specify a different model ID or try a different AWS Region. "
+        error_msg += f"For a list of available models, see {MODEL_ID_LIST_WEB_URL}. "
 
         other_model_id_version = None
         if model_type == JumpStartModelType.OPEN_WEIGHTS:

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/jumpstart/exceptions.py

@@ -32,8 +32,8 @@ NO_AVAILABLE_RESOURCE_REQUIREMENT_RECOMMENDATION_ERROR_MSG = (
 )
 
 INVALID_MODEL_ID_ERROR_MSG = (
-    "Invalid model ID: '{model_id}'. Please visit "
-    f"{MODEL_ID_LIST_WEB_URL} for a list of valid model IDs. "
+    "Invalid model ID: '{model_id}'. Specify a different model ID or try a different AWS Region. "
+    f"For a list of available models, see {MODEL_ID_LIST_WEB_URL}. "
     "The module `sagemaker.jumpstart.notebook_utils` contains utilities for "
     "fetching model IDs. We recommend upgrading to the latest version of sagemaker "
     "to get access to the most models."

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/jumpstart/utils.py

@@ -50,7 +50,12 @@ from sagemaker.jumpstart.types import (
 )
 from sagemaker.session import Session
 from sagemaker.config import load_sagemaker_config
-from sagemaker.utils import resolve_value_from_config, TagsDict, get_instance_rate_per_hour
+from sagemaker.utils import (
+    resolve_value_from_config,
+    TagsDict,
+    get_instance_rate_per_hour,
+    get_domain_for_region,
+)
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.user_agent import get_user_agent_extra_suffix
 
@@ -553,7 +558,7 @@ def get_eula_message(model_specs: JumpStartModelSpecs, region: str) -> str:
     return (
         f"Model '{model_specs.model_id}' requires accepting end-user license agreement (EULA). "
         f"See https://{get_jumpstart_content_bucket(region=region)}.s3.{region}."
-        f"amazonaws.com{'.cn' if region.startswith('cn-') else ''}"
+        f"{get_domain_for_region(region)}"
         f"/{model_specs.hosting_eula_key} for terms of use."
     )
 

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/local/image.py

@@ -30,7 +30,6 @@ import sys
 import tarfile
 import tempfile
 
-from distutils.spawn import find_executable
 from threading import Thread
 from typing import Dict, List
 from six.moves.urllib.parse import urlparse
@@ -170,7 +169,7 @@ class _SageMakerContainer(object):
             compose_cmd_prefix.extend(["docker", "compose"])
             return compose_cmd_prefix
 
-        if find_executable("docker-compose") is not None:
+        if shutil.which("docker-compose") is not None:
             logger.info("'Docker Compose' found using Docker Compose CLI.")
             compose_cmd_prefix.extend(["docker-compose"])
             return compose_cmd_prefix

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/local/utils.py

@@ -21,7 +21,6 @@ import json
 import re
 import errno
 
-from distutils.dir_util import copy_tree
 from six.moves.urllib.parse import urlparse
 
 from sagemaker import s3
@@ -102,7 +101,7 @@ def move_to_destination(source, destination, job_name, sagemaker_session, prefix
 
 
 def recursive_copy(source, destination):
-    """A wrapper around distutils.dir_util.copy_tree.
+    """A wrapper around shutil.copy_tree.
 
     This won't throw any exception when the source directory does not exist.
 
@@ -111,7 +110,7 @@ def recursive_copy(source, destination):
         destination (str): destination path
     """
     if os.path.isdir(source):
-        copy_tree(source, destination)
+        shutil.copytree(source, destination, dirs_exist_ok=True)
 
 
 def kill_child_processes(pid):

sagemaker-2.232.2/src/sagemaker/mlflow/forward_sagemaker_metrics.py

@@ -0,0 +1,311 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+"""This module contains code related to forwarding SageMaker TrainingJob Metrics to MLflow."""
+
+from __future__ import absolute_import
+
+import os
+import platform
+import re
+from typing import Set, Tuple, List, Dict, Generator
+import boto3
+import mlflow
+from mlflow import MlflowClient
+from mlflow.entities import Metric, Param, RunTag
+
+from packaging import version
+
+
+def encode(name: str, existing_names: Set[str]) -> str:
+    """Encode a string to comply with MLflow naming restrictions and ensure uniqueness.
+
+    Args:
+        name (str): The original string to be encoded.
+        existing_names (Set[str]): Set of existing encoded names to avoid collisions.
+
+    Returns:
+        str: The encoded string if changes were necessary, otherwise the original string.
+    """
+
+    def encode_char(match):
+        return f"_{ord(match.group(0)):02x}_"
+
+    # Check if we're on Mac/Unix and using MLflow 2.16.0 or greater
+    is_unix = platform.system() != "Windows"
+    mlflow_version = version.parse(mlflow.__version__)
+    allow_colon = is_unix and mlflow_version >= version.parse("2.16.0")
+
+    if allow_colon:
+        pattern = r"[^\w\-./:\s]"
+    else:
+        pattern = r"[^\w\-./\s]"
+
+    encoded = re.sub(pattern, encode_char, name)
+    base_name = encoded[:240]  # Leave room for potential suffix to accommodate duplicates
+
+    if base_name in existing_names:
+        suffix = 1
+        # Edge case where even with suffix space there is a collision
+        # we will override one of the keys.
+        while f"{base_name}_{suffix}" in existing_names:
+            suffix += 1
+        encoded = f"{base_name}_{suffix}"
+
+    # Max length is 250 for mlflow metric/params
+    encoded = encoded[:250]
+
+    existing_names.add(encoded)
+    return encoded
+
+
+def decode(encoded_metric_name: str) -> str:
+    """Decodes an encoded metric name by replacing hexadecimal representations with ASCII
+
+    This function reverses the encoding process by converting hexadecimal codes
+    back to their original characters. It looks for patterns of the form "_XX_"
+    where XX is a two-digit hexadecimal code, and replaces them with the
+    corresponding ASCII character.
+
+    Args:
+        encoded_metric_name (str): The encoded metric name to be decoded.
+
+    Returns:
+        str: The decoded metric name with hexadecimal codes replaced by their
+             corresponding characters.
+
+    Example:
+        >>> decode("loss_3a_val")
+        "loss:val"
+    """
+
+    def replace_code(match):
+        code = match.group(1)
+        return chr(int(code, 16))
+
+    # Replace encoded characters
+    decoded = re.sub(r"_([0-9a-f]{2})_", replace_code, encoded_metric_name)
+
+    return decoded
+
+
+def get_training_job_details(job_arn: str) -> dict:
+    """Retrieve details of a SageMaker training job.
+
+    Args:
+        job_arn (str): The ARN of the SageMaker training job.
+
+    Returns:
+        dict: A dictionary containing the details of the training job.
+
+    Raises:
+        boto3.exceptions.Boto3Error: If there's an issue with the AWS API call.
+    """
+    sagemaker_client = boto3.client("sagemaker")
+    job_name = job_arn.split("/")[-1]
+    return sagemaker_client.describe_training_job(TrainingJobName=job_name)
+
+
+def create_metric_queries(job_arn: str, metric_definitions: list) -> list:
+    """Create metric queries for SageMaker metrics.
+
+    Args:
+        job_arn (str): The ARN of the SageMaker training job.
+        metric_definitions (list): List of metric definitions from the training job.
+
+    Returns:
+        list: A list of dictionaries, each representing a metric query.
+    """
+    metric_queries = []
+    for metric in metric_definitions:
+        query = {
+            "MetricName": metric["Name"],
+            "XAxisType": "Timestamp",
+            "MetricStat": "Avg",
+            "Period": "OneMinute",
+            "ResourceArn": job_arn,
+        }
+        metric_queries.append(query)
+    return metric_queries
+
+
+def get_metric_data(metric_queries: list) -> dict:
+    """Retrieve metric data from SageMaker.
+
+    Args:
+        metric_queries (list): A list of metric queries.
+
+    Returns:
+        dict: A dictionary containing the metric data results.
+
+    Raises:
+        boto3.exceptions.Boto3Error: If there's an issue with the AWS API call.
+    """
+    sagemaker_metrics_client = boto3.client("sagemaker-metrics")
+    metric_data = sagemaker_metrics_client.batch_get_metrics(MetricQueries=metric_queries)
+    return metric_data
+
+
+def prepare_mlflow_metrics(
+    metric_queries: list, metric_results: list
+) -> Tuple[List[Metric], Dict[str, str]]:
+    """Prepare metrics for MLflow logging, encoding metric names if necessary.
+
+    Args:
+        metric_queries (list): The original metric queries sent to SageMaker.
+        metric_results (list): The metric results from SageMaker batch_get_metrics.
+
+    Returns:
+        Tuple[List[Metric], Dict[str, str]]:
+            - A list of Metric objects with encoded names (if necessary)
+            - A mapping of encoded to original names for metrics (only for encoded metrics)
+    """
+    mlflow_metrics = []
+    metric_name_mapping = {}
+    existing_names = set()
+
+    for query, result in zip(metric_queries, metric_results):
+        if result["Status"] == "Complete":
+            metric_name = query["MetricName"]
+            encoded_name = encode(metric_name, existing_names)
+            metric_name_mapping[encoded_name] = metric_name
+
+            for step, (timestamp, value) in enumerate(
+                zip(result["XAxisValues"], result["MetricValues"])
+            ):
+                metric = Metric(key=encoded_name, value=value, timestamp=timestamp, step=step)
+                mlflow_metrics.append(metric)
+
+    return mlflow_metrics, metric_name_mapping
+
+
+def prepare_mlflow_params(hyperparameters: Dict[str, str]) -> Tuple[List[Param], Dict[str, str]]:
+    """Prepare hyperparameters for MLflow logging, encoding parameter names if necessary.
+
+    Args:
+        hyperparameters (Dict[str, str]): The hyperparameters from the SageMaker job.
+
+    Returns:
+        Tuple[List[Param], Dict[str, str]]:
+            - A list of Param objects with encoded names (if necessary)
+            - A mapping of encoded to original names for
+                hyperparameters (only for encoded parameters)
+    """
+    mlflow_params = []
+    param_name_mapping = {}
+    existing_names = set()
+
+    for key, value in hyperparameters.items():
+        encoded_key = encode(key, existing_names)
+        param_name_mapping[encoded_key] = key
+        mlflow_params.append(Param(encoded_key, str(value)))
+
+    return mlflow_params, param_name_mapping
+
+
+def batch_items(items: list, batch_size: int) -> Generator:
+    """Yield successive batch_size chunks from items.
+
+    Args:
+        items (list): The list of items to be batched.
+        batch_size (int): The size of each batch.
+
+    Yields:
+        list: A batch of items.
+    """
+    for i in range(0, len(items), batch_size):
+        yield items[i : i + batch_size]
+
+
+def log_to_mlflow(metrics: list, params: list, tags: dict) -> None:
+    """Log metrics, parameters, and tags to MLflow.
+
+    Args:
+        metrics (list): List of metrics to log.
+        params (list): List of parameters to log.
+        tags (dict): Dictionary of tags to set.
+
+    Raises:
+        mlflow.exceptions.MlflowException: If there's an issue with MLflow logging.
+    """
+    client = MlflowClient()
+
+    experiment_name = os.getenv("MLFLOW_EXPERIMENT_NAME")
+    if experiment_name is None or experiment_name.strip() == "":
+        experiment_name = "Default"
+        print("MLFLOW_EXPERIMENT_NAME not set. Using Default")
+
+    experiment = client.get_experiment_by_name(experiment_name)
+    if experiment is None:
+        experiment_id = client.create_experiment(experiment_name)
+    else:
+        experiment_id = experiment.experiment_id
+
+    run = client.create_run(experiment_id)
+
+    for metric_batch in batch_items(metrics, 1000):
+        client.log_batch(
+            run.info.run_id,
+            metrics=metric_batch,
+        )
+    for param_batch in batch_items(params, 1000):
+        client.log_batch(run.info.run_id, params=param_batch)
+
+    tag_items = list(tags.items())
+    for tag_batch in batch_items(tag_items, 1000):
+        tag_objects = [RunTag(key, str(value)) for key, value in tag_batch]
+        client.log_batch(run.info.run_id, tags=tag_objects)
+    client.set_terminated(run.info.run_id)
+
+
+def log_sagemaker_job_to_mlflow(training_job_arn: str) -> None:
+    """Retrieve SageMaker metrics and hyperparameters and log them to MLflow.
+
+    Args:
+        training_job_arn (str): The ARN of the SageMaker training job.
+
+    Raises:
+        Exception: If there's any error during the process.
+    """
+    # Get training job details
+    mlflow.set_tracking_uri(os.getenv("MLFLOW_TRACKING_URI"))
+    job_details = get_training_job_details(training_job_arn)
+
+    # Extract hyperparameters and metric definitions
+    hyperparameters = job_details["HyperParameters"]
+    metric_definitions = job_details["AlgorithmSpecification"]["MetricDefinitions"]
+
+    # Create and get metric queries
+    metric_queries = create_metric_queries(job_details["TrainingJobArn"], metric_definitions)
+    metric_data = get_metric_data(metric_queries)
+
+    # Create a mapping of encoded to original metric names
+    # Prepare data for MLflow
+    mlflow_metrics, metric_name_mapping = prepare_mlflow_metrics(
+        metric_queries, metric_data["MetricQueryResults"]
+    )
+
+    # Create a mapping of encoded to original hyperparameter names
+    # Prepare data for MLflow
+    mlflow_params, param_name_mapping = prepare_mlflow_params(hyperparameters)
+
+    mlflow_tags = {
+        "training_job_arn": training_job_arn,
+        "metric_name_mapping": str(metric_name_mapping),
+        "param_name_mapping": str(param_name_mapping),
+    }
+
+    # Log to MLflow
+    log_to_mlflow(mlflow_metrics, mlflow_params, mlflow_tags)
+    print(f"Logged {len(mlflow_metrics)} metric datapoints to MLflow")
+    print(f"Logged {len(mlflow_params)} hyperparameters to MLflow")

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/utils.py

@@ -52,6 +52,15 @@ from sagemaker.session_settings import SessionSettings
 from sagemaker.workflow import is_pipeline_variable, is_pipeline_parameter_string
 from sagemaker.workflow.entities import PipelineVariable
 
+ALTERNATE_DOMAINS = {
+    "cn-north-1": "amazonaws.com.cn",
+    "cn-northwest-1": "amazonaws.com.cn",
+    "us-iso-east-1": "c2s.ic.gov",
+    "us-isob-east-1": "sc2s.sgov.gov",
+    "us-isof-south-1": "csp.hci.ic.gov",
+    "us-isof-east-1": "csp.hci.ic.gov",
+}
+
 ECR_URI_PATTERN = r"^(\d+)(\.)dkr(\.)ecr(\.)(.+)(\.)(.*)(/)(.*:.*)$"
 MODEL_PACKAGE_ARN_PATTERN = (
     r"arn:aws([a-z\-]*)?:sagemaker:([a-z0-9\-]*):([0-9]{12}):model-package/(.*)"
@@ -1905,3 +1914,12 @@ def remove_tag_with_key(key: str, tags: Optional[Tags]) -> Optional[Tags]:
     if len(updated_tags) == 1:
         return updated_tags[0]
     return updated_tags
+
+
+def get_domain_for_region(region: str) -> str:
+    """Returns the domain for the given region.
+
+    Args:
+        region (str): AWS region name.
+    """
+    return ALTERNATE_DOMAINS.get(region, "amazonaws.com")

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/workflow/_repack_model.py

@@ -27,13 +27,6 @@ import tempfile
 # is unpacked for inference, the custom entry point will be used.
 # Reference: https://docs.aws.amazon.com/sagemaker/latest/dg/amazon-sagemaker-toolkits.html
 
-# distutils.dir_util.copy_tree works way better than the half-baked
-# shutil.copytree which bombs on previously existing target dirs...
-# alas ... https://bugs.python.org/issue10948
-# we'll go ahead and use the copy_tree function anyways because this
-# repacking is some short-lived hackery, right??
-from distutils.dir_util import copy_tree
-
 from os.path import abspath, realpath, dirname, normpath, join as joinpath
 
 logger = logging.getLogger(__name__)
@@ -188,7 +181,7 @@ def repack(inference_script, model_archive, dependencies=None, source_dir=None):
 
         # copy the "src" dir, which includes the previous training job's model and the
         # custom inference script, to the output of this training job
-        copy_tree(src_dir, "/opt/ml/model")
+        shutil.copytree(src_dir, "/opt/ml/model", dirs_exist_ok=True)
 
 
 if __name__ == "__main__":  # pragma: no cover

{sagemaker-2.232.1 → sagemaker-2.232.2}/src/sagemaker/workflow/_utils.py

@@ -46,7 +46,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-FRAMEWORK_VERSION = "0.23-1"
+FRAMEWORK_VERSION = "1.2-1"
 INSTANCE_TYPE = "ml.m5.large"
 REPACK_SCRIPT = "_repack_model.py"
 REPACK_SCRIPT_LAUNCHER = "_repack_script_launcher.sh"

sagemaker-2.232.1/VERSION

@@ -1 +0,0 @@
-2.232.1