Merge branch 'master' into feat/jumpstart-instance-specific-hyperparameters

Author: evakravi

CHANGELOG.md

@@ -1,5 +1,26 @@
 # Changelog
 
+## v2.192.0 (2023-10-11)
+
+### Features
+
+ * jumpstart estimator enable infra check flag
+ * jumpstart default payloads
+ * allow non-python files in job dependencies
+ * allow configuring docker container in local mode
+
+### Bug Fixes and Other Changes
+
+ * js tagging s3 prefix
+ * Batch transform: Add support for split_type == "None" in local mode
+ * use correct line endings and s3 uris on windows
+ * Fixed bug in _create_training_details
+ * DJL Neuronx 0.24.0
+
+### Documentation Changes
+
+ * Include FeatureGroup's load_feature_definitions API documentation
+
 ## v2.191.0 (2023-10-05)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.191.1.dev0
+2.192.1.dev0

src/sagemaker/config/config_schema.py

@@ -106,6 +106,7 @@
 LOCAL_CODE = "local_code"
 SERVING_PORT = "serving_port"
 CONTAINER_CONFIG = "container_config"
+CONTAINER_ROOT = "container_root"
 REGION_NAME = "region_name"
 
 
@@ -1090,6 +1091,9 @@ def _simple_path(*args: str):
                 SERVING_PORT: {
                     TYPE: "integer",
                 },
+                CONTAINER_ROOT: {
+                    TYPE: "string",
+                },
                 CONTAINER_CONFIG: {
                     TYPE: OBJECT,
                 },

src/sagemaker/estimator.py

@@ -3082,6 +3082,7 @@ def __init__(
             hyperparameters=hyperparameters,
             instance_groups=instance_groups,
             training_repository_access_mode=training_repository_access_mode,
+            enable_infra_check=enable_infra_check,
             training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,  # noqa: E501 # pylint: disable=line-too-long
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,

src/sagemaker/jumpstart/estimator.py

@@ -102,6 +102,7 @@ def __init__(
         instance_groups: Optional[List[InstanceGroup]] = None,
         training_repository_access_mode: Optional[Union[str, PipelineVariable]] = None,
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
@@ -485,6 +486,8 @@ def __init__(
                 private Docker registry where your training image is hosted (Default: None).
                 When it's set to None, SageMaker will not do authentication before pulling the image
                 in the private Docker registry. (Default: None).
+            enable_infra_check (Optional[Union[bool, PipelineVariable]]):
+                Specifies whether it is running Sagemaker built-in infra check jobs.
             container_entry_point (Optional[List[str]]): The entrypoint script for a Docker
                 container used to run a training job. This script takes precedence over
                 the default train processing instructions.
@@ -565,6 +568,7 @@ def _is_valid_model_id_hook():
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
+            enable_infra_check=enable_infra_check,
         )
 
         self.model_id = estimator_init_kwargs.model_id

src/sagemaker/jumpstart/factory/estimator.py

@@ -123,6 +123,7 @@ def get_init_kwargs(
     container_entry_point: Optional[List[str]] = None,
     container_arguments: Optional[List[str]] = None,
     disable_output_compression: Optional[bool] = None,
+    enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -178,6 +179,7 @@ def get_init_kwargs(
         container_entry_point=container_entry_point,
         container_arguments=container_arguments,
         disable_output_compression=disable_output_compression,
+        enable_infra_check=enable_infra_check,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)

src/sagemaker/jumpstart/payload_utils.py

@@ -37,12 +37,12 @@ class PayloadSerializer:
 
     def __init__(
         self,
-        bucket: str = get_jumpstart_content_bucket(),
+        bucket: Optional[str] = None,
         region: str = JUMPSTART_DEFAULT_REGION_NAME,
         s3_client: Optional[boto3.client] = None,
     ) -> None:
         """Initializes PayloadSerializer object."""
-        self.bucket = bucket
+        self.bucket = bucket or get_jumpstart_content_bucket()
         self.region = region
         self.s3_client = s3_client
 

src/sagemaker/jumpstart/types.py

@@ -18,6 +18,7 @@
 from sagemaker.utils import get_instance_type_family
 
 from sagemaker.session import Session
+from sagemaker.workflow.entities import PipelineVariable
 
 
 class JumpStartDataHolderType:
@@ -1055,6 +1056,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "container_entry_point",
         "container_arguments",
         "disable_output_compression",
+        "enable_infra_check",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1118,6 +1120,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
+        enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1174,6 +1177,7 @@ def __init__(
         self.container_entry_point = container_entry_point
         self.container_arguments = container_arguments
         self.disable_output_compression = disable_output_compression
+        self.enable_infra_check = enable_infra_check
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):

src/sagemaker/jumpstart/utils.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 import logging
 import os
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlparse
 from packaging.version import Version
 import sagemaker
@@ -277,7 +277,7 @@ def get_jumpstart_base_name_if_jumpstart_model(
 
 def add_jumpstart_tags(
     tags: Optional[List[Dict[str, str]]] = None,
-    inference_model_uri: Optional[str] = None,
+    inference_model_uri: Optional[Union[str, dict]] = None,
     inference_script_uri: Optional[str] = None,
     training_model_uri: Optional[str] = None,
     training_script_uri: Optional[str] = None,
@@ -289,7 +289,7 @@ def add_jumpstart_tags(
     Args:
         tags (Optional[List[Dict[str,str]]): Current tags for JumpStart inference
             or training job. (Default: None).
-        inference_model_uri (Optional[str]): S3 URI for inference model artifact.
+        inference_model_uri (Optional[Union[dict, str]]): S3 URI for inference model artifact.
             (Default: None).
         inference_script_uri (Optional[str]): S3 URI for inference script tarball.
             (Default: None).
@@ -302,6 +302,10 @@ def add_jumpstart_tags(
         "The URI (%s) is a pipeline variable which is only interpreted at execution time. "
         "As a result, the JumpStart resources will not be tagged."
     )
+
+    if isinstance(inference_model_uri, dict):
+        inference_model_uri = inference_model_uri.get("S3DataSource", {}).get("S3Uri", None)
+
     if inference_model_uri:
         if is_pipeline_variable(inference_model_uri):
             logging.warning(warn_msg, "inference_model_uri")

src/sagemaker/model.py

@@ -1345,7 +1345,9 @@ def deploy(
 
         tags = add_jumpstart_tags(
             tags=tags,
-            inference_model_uri=self.model_data if isinstance(self.model_data, str) else None,
+            inference_model_uri=self.model_data
+            if isinstance(self.model_data, (str, dict))
+            else None,
             inference_script_uri=self.source_dir,
         )
 

tests/integ/test_huggingface_torch_distributed.py

@@ -32,7 +32,7 @@ def test_huggingface_torch_distributed_g5_glue(
             transformers_version=huggingface_training_latest_version,
             pytorch_version=huggingface_training_pytorch_latest_version,
             instance_count=1,
-            instance_type="ml.g5.8xlarge",
+            instance_type="ml.g5.4xlarge",
             hyperparameters={
                 "model_name_or_path": "distilbert-base-cased",
                 "task_name": "wnli",

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -677,7 +677,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
         Please add the new argument to the skip set below,
         and reach out to JumpStart team."""
 
-        init_args_to_skip: Set[str] = set(["kwargs", "enable_infra_check"])
+        init_args_to_skip: Set[str] = set(["kwargs"])
         fit_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])
 

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -331,7 +331,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
 
         """If you add arguments to <Model constructor>, this test will fail.
         Please add the new argument to the skip set below,
-        and cut a ticket sev-3 to JumpStart team: AWS > SageMaker > JumpStart"""
+        and reach out to JumpStart team."""
 
         init_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -216,6 +216,34 @@ def test_add_jumpstart_tags_inference():
         inference_script_uri=inference_script_uri,
     ) == [{"Key": JumpStartTag.INFERENCE_MODEL_URI.value, "Value": inference_model_uri}]
 
+    tags = []
+    inference_model_uri = {"S3DataSource": {"S3Uri": random_jumpstart_s3_uri("random_key")}}
+    inference_script_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            "Key": JumpStartTag.INFERENCE_MODEL_URI.value,
+            "Value": inference_model_uri["S3DataSource"]["S3Uri"],
+        }
+    ]
+
+    tags = []
+    inference_model_uri = {"S3DataSource": {"S3Uri": random_jumpstart_s3_uri("random_key/prefix/")}}
+    inference_script_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            "Key": JumpStartTag.INFERENCE_MODEL_URI.value,
+            "Value": inference_model_uri["S3DataSource"]["S3Uri"],
+        }
+    ]
+
     tags = [{"Key": "some", "Value": "tag"}]
     inference_model_uri = random_jumpstart_s3_uri("random_key")
     inference_script_uri = "dfsdfs"