Merge branch 'master' into chore/move-jumpstart-region-definitions-to-json-file

Author: benieric

doc/requirements.txt

@@ -2,7 +2,7 @@ sphinx==5.1.1
 sphinx-rtd-theme==0.5.0
 docutils==0.15.2
 packaging==20.9
-jinja2==3.1.4
+jinja2==3.1.6
 schema==0.7.5
 accelerate>=0.24.1,<=0.27.0
 graphene<4.0

requirements/extras/test_requirements.txt

@@ -20,7 +20,7 @@ attrs>=23.1.0,<24
 fabric==2.6.0
 requests==2.32.2
 sagemaker-experiments==0.1.35
-Jinja2==3.1.4
+Jinja2==3.1.6
 pyvis==0.2.1
 pandas==1.4.4
 scikit-learn==1.3.0
@@ -33,7 +33,7 @@ pillow>=10.0.1,<=11
 opentelemetry-proto==1.27.0
 protobuf==4.25.5
 tensorboard>=2.9.0,<=2.15.2
-transformers==4.46.1
+transformers==4.48.0
 sentencepiece==0.1.99
 # https://github.com/triton-inference-server/server/issues/6246
 tritonclient[http]<2.37.0

src/sagemaker/image_uri_config/huggingface-llm-neuronx.json

@@ -90,6 +90,7 @@ def remote(
     spark_config: SparkConfig = None,
     use_spot_instances=False,
     max_wait_time_in_seconds=None,
+    disable_output_compression: bool = False,
     use_torchrun: bool = False,
     use_mpirun: bool = False,
     nproc_per_node: Optional[int] = None,
@@ -283,13 +284,16 @@ def remote(
           After this amount of time Amazon SageMaker will stop waiting for managed spot training
           job to complete. Defaults to ``None``.
 
+        disable_output_compression (bool): Optional. When set to true, Model is uploaded to
+          Amazon S3 without compression after training finishes.
+
         use_torchrun (bool): Specifies whether to use torchrun for distributed training.
           Defaults to ``False``.
 
         use_mpirun (bool): Specifies whether to use mpirun for distributed training.
           Defaults to ``False``.
 
-        nproc_per_node (Optional int): Specifies the number of processes per node for
+        nproc_per_node (int): Optional. Specifies the number of processes per node for
           distributed training. Defaults to ``None``.
           This is defined automatically configured on the instance type.
     """
@@ -324,6 +328,7 @@ def _remote(func):
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            disable_output_compression=disable_output_compression,
             use_torchrun=use_torchrun,
             use_mpirun=use_mpirun,
             nproc_per_node=nproc_per_node,
@@ -543,6 +548,7 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        disable_output_compression: bool = False,
         use_torchrun: bool = False,
         use_mpirun: bool = False,
         nproc_per_node: Optional[int] = None,
@@ -736,13 +742,16 @@ def __init__(
               After this amount of time Amazon SageMaker will stop waiting for managed spot training
               job to complete. Defaults to ``None``.
 
+            disable_output_compression (bool): Optional. When set to true, Model is uploaded to
+              Amazon S3 without compression after training finishes.
+
             use_torchrun (bool): Specifies whether to use torchrun for distributed training.
               Defaults to ``False``.
 
             use_mpirun (bool): Specifies whether to use mpirun for distributed training.
               Defaults to ``False``.
 
-            nproc_per_node (Optional int): Specifies the number of processes per node for
+            nproc_per_node (int): Optional. Specifies the number of processes per node for
               distributed training. Defaults to ``None``.
               This is defined automatically configured on the instance type.
         """
@@ -790,6 +799,7 @@ def __init__(
             spark_config=spark_config,
             use_spot_instances=use_spot_instances,
             max_wait_time_in_seconds=max_wait_time_in_seconds,
+            disable_output_compression=disable_output_compression,
             use_torchrun=use_torchrun,
             use_mpirun=use_mpirun,
             nproc_per_node=nproc_per_node,

src/sagemaker/image_uri_config/huggingface-llm.json

@@ -373,6 +373,7 @@ def __init__(
         spark_config: SparkConfig = None,
         use_spot_instances=False,
         max_wait_time_in_seconds=None,
+        disable_output_compression: bool = False,
         use_torchrun: bool = False,
         use_mpirun: bool = False,
         nproc_per_node: Optional[int] = None,
@@ -558,13 +559,16 @@ def __init__(
               After this amount of time Amazon SageMaker will stop waiting for managed spot
               training job to complete. Defaults to ``None``.
 
+            disable_output_compression (bool): Optional. When set to true, Model is uploaded to
+              Amazon S3 without compression after training finishes.
+
             use_torchrun (bool): Specifies whether to use torchrun for distributed training.
               Defaults to ``False``.
 
             use_mpirun (bool): Specifies whether to use mpirun for distributed training.
               Defaults to ``False``.
 
-            nproc_per_node (Optional int): Specifies the number of processes per node for
+            nproc_per_node (int): Optional. Specifies the number of processes per node for
               distributed training. Defaults to ``None``.
               This is defined automatically configured on the instance type.
         """
@@ -725,6 +729,7 @@ def __init__(
         tags = format_tags(tags)
         self.tags = self.sagemaker_session._append_sagemaker_config_tags(tags, REMOTE_FUNCTION_TAGS)
 
+        self.disable_output_compression = disable_output_compression
         self.use_torchrun = use_torchrun
         self.use_mpirun = use_mpirun
         self.nproc_per_node = nproc_per_node
@@ -954,6 +959,8 @@ def compile(
             output_config = {"S3OutputPath": s3_base_uri}
         if job_settings.s3_kms_key is not None:
             output_config["KmsKeyId"] = job_settings.s3_kms_key
+        if job_settings.disable_output_compression:
+            output_config["CompressionType"] = "NONE"
         request_dict["OutputDataConfig"] = output_config
 
         container_args = ["--s3_base_uri", s3_base_uri]

src/sagemaker/remote_function/client.py

@@ -35,7 +35,7 @@
 SPARK_PATH = os.path.join(DATA_DIR, "spark")
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="module", autouse=True)
 def build_jar():
     jar_file_path = os.path.join(SPARK_PATH, "code", "java", "hello-java-spark")
     # compile java file
@@ -207,12 +207,10 @@ def configuration() -> list:
 
 
 def test_sagemaker_pyspark_v3(
-    spark_v3_py_processor, spark_v3_jar_processor, sagemaker_session, configuration, build_jar
+    spark_v3_py_processor, spark_v3_jar_processor, sagemaker_session, configuration
 ):
     test_sagemaker_pyspark_multinode(spark_v3_py_processor, sagemaker_session, configuration)
-    test_sagemaker_java_jar_multinode(
-        spark_v3_jar_processor, sagemaker_session, configuration, build_jar
-    )
+    test_sagemaker_java_jar_multinode(spark_v3_jar_processor, sagemaker_session, configuration)
 
 
 def test_sagemaker_pyspark_multinode(spark_py_processor, sagemaker_session, configuration):
@@ -280,9 +278,7 @@ def test_sagemaker_pyspark_multinode(spark_py_processor, sagemaker_session, conf
     assert len(output_contents) != 0
 
 
-def test_sagemaker_java_jar_multinode(
-    spark_jar_processor, sagemaker_session, configuration, build_jar
-):
+def test_sagemaker_java_jar_multinode(spark_jar_processor, sagemaker_session, configuration):
     """Test SparkJarProcessor using Java application jar"""
     bucket = spark_jar_processor.sagemaker_session.default_bucket()
     with open(os.path.join(SPARK_PATH, "files", "data.jsonl")) as data:

src/sagemaker/remote_function/job.py

@@ -907,6 +907,7 @@ def test_remote_decorator_fields_consistency(get_execution_role, session):
         "use_spot_instances",
         "max_wait_time_in_seconds",
         "custom_file_filter",
+        "disable_output_compression",
         "use_torchrun",
         "use_mpirun",
         "nproc_per_node",

tests/integ/test_spark_processing.py

@@ -568,11 +568,12 @@ def test_clarify_model_monitor():
 
     # The subclass should has monitoring_type() defined
     # noinspection PyAbstractClass
-    class DummyClarifyModelMonitoir(ClarifyModelMonitor):
+    class DummyClarifyModelMonitor(ClarifyModelMonitor):
+        _TEST_CLASS = True
         pass
 
     with pytest.raises(TypeError):
-        DummyClarifyModelMonitoir.monitoring_type()
+        DummyClarifyModelMonitor.monitoring_type()
 
 
 def test_clarify_model_monitor_invalid_update(clarify_model_monitors):
@@ -593,6 +594,8 @@ def test_clarify_model_monitor_invalid_attach(sagemaker_session):
     )
     # attach, invalid monitoring type
     for clarify_model_monitor_cls in ClarifyModelMonitor.__subclasses__():
+        if hasattr(clarify_model_monitor_cls, "_TEST_CLASS"):
+            continue
         with pytest.raises(TypeError):
             clarify_model_monitor_cls.attach(SCHEDULE_NAME, sagemaker_session)
 

tests/unit/sagemaker/feature_store/feature_processor/test_feature_scheduler.py

@@ -1504,6 +1504,7 @@ def test_consistency_between_remote_and_step_decorator():
         "s3_kms_key",
         "s3_root_uri",
         "sagemaker_session",
+        "disable_output_compression",
         "use_torchrun",
         "use_mpirun",
         "nproc_per_node",