Merge branch 'master' into tgi133

amzn-choeric · web-flow · commit 3f3f947ab5f3 · 2023-12-21T15:06:54.000-05:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,16 @@
 # Changelog
 
+## v2.202.0 (2023-12-21)
+
+### Features
+
+ * support remote debug for sagemaker training job
+
+### Bug Fixes and Other Changes
+
+ * update image_uri_configs  12-21-2023 08:32:41 PST
+ * Update tblib constraint
+
 ## v2.201.0 (2023-12-20)
 
 ### Features
diff --git a/VERSION b/VERSION
@@ -1 +1 @@
-2.201.1.dev0
+2.202.1.dev0
diff --git a/setup.py b/setup.py
@@ -63,7 +63,7 @@ def read_requirements(filename):
     "PyYAML~=6.0",
     "jsonschema",
     "platformdirs",
-    "tblib==1.7.0",
+    "tblib>=1.7.0,<3",
     "urllib3<1.27",
     "uvicorn==0.22.0",
     "fastapi==0.95.2",
diff --git a/src/sagemaker/automl/automl.py b/src/sagemaker/automl/automl.py
@@ -332,7 +332,7 @@ def attach(cls, auto_ml_job_name, sagemaker_session=None):
             total_job_runtime_in_seconds=auto_ml_job_desc.get("AutoMLJobConfig", {})
             .get("CompletionCriteria", {})
             .get("MaxAutoMLJobRuntimeInSeconds"),
-            job_objective=auto_ml_job_desc.get("AutoMLJobObjective", {}).get("MetricName"),
+            job_objective=auto_ml_job_desc.get("AutoMLJobObjective", {}),
             generate_candidate_definitions_only=auto_ml_job_desc.get(
                 "GenerateCandidateDefinitionsOnly", False
             ),
diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py
@@ -178,6 +178,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
+        enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -540,6 +541,8 @@ def __init__(
                 to Amazon S3 without compression after training finishes.
             enable_infra_check (bool or PipelineVariable): Optional.
                 Specifies whether it is running Sagemaker built-in infra check jobs.
+            enable_remote_debug (bool or PipelineVariable): Optional.
+                Specifies whether RemoteDebug is enabled for the training job
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -777,6 +780,8 @@ def __init__(
 
         self.tensorboard_app = TensorBoardApp(region=self.sagemaker_session.boto_region_name)
 
+        self._enable_remote_debug = enable_remote_debug
+
     @abstractmethod
     def training_image_uri(self):
         """Return the Docker image to use for training.
@@ -1958,6 +1963,11 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
             max_wait = job_details.get("StoppingCondition", {}).get("MaxWaitTimeInSeconds")
             if max_wait:
                 init_params["max_wait"] = max_wait
+
+        if "RemoteDebugConfig" in job_details:
+            init_params["enable_remote_debug"] = job_details["RemoteDebugConfig"].get(
+                "EnableRemoteDebug"
+            )
         return init_params
 
     def _get_instance_type(self):
@@ -2292,6 +2302,32 @@ def update_profiler(
 
         _TrainingJob.update(self, profiler_rule_configs, profiler_config_request_dict)
 
+    def get_remote_debug_config(self):
+        """dict: Return the configuration of RemoteDebug"""
+        return (
+            None
+            if self._enable_remote_debug is None
+            else {"EnableRemoteDebug": self._enable_remote_debug}
+        )
+
+    def enable_remote_debug(self):
+        """Enable remote debug for a training job."""
+        self._update_remote_debug(True)
+
+    def disable_remote_debug(self):
+        """Disable remote debug for a training job."""
+        self._update_remote_debug(False)
+
+    def _update_remote_debug(self, enable_remote_debug: bool):
+        """Update to enable or disable remote debug for a training job.
+
+        This method updates the ``_enable_remote_debug`` parameter
+        and enables or disables remote debug for a training job
+        """
+        self._ensure_latest_training_job()
+        _TrainingJob.update(self, remote_debug_config={"EnableRemoteDebug": enable_remote_debug})
+        self._enable_remote_debug = enable_remote_debug
+
     def get_app_url(
         self,
         app_type,
@@ -2520,6 +2556,9 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
         if estimator.profiler_config:
             train_args["profiler_config"] = estimator.profiler_config._to_request_dict()
 
+        if estimator.get_remote_debug_config() is not None:
+            train_args["remote_debug_config"] = estimator.get_remote_debug_config()
+
         return train_args
 
     @classmethod
@@ -2549,7 +2588,12 @@ def _is_local_channel(cls, input_uri):
 
     @classmethod
     def update(
-        cls, estimator, profiler_rule_configs=None, profiler_config=None, resource_config=None
+        cls,
+        estimator,
+        profiler_rule_configs=None,
+        profiler_config=None,
+        resource_config=None,
+        remote_debug_config=None,
     ):
         """Update a running Amazon SageMaker training job.
 
@@ -2562,20 +2606,31 @@ def update(
             resource_config (dict): Configuration of the resources for the training job. You can
                 update the keep-alive period if the warm pool status is `Available`. No other fields
                 can be updated. (default: None).
+            remote_debug_config (dict): Configuration for RemoteDebug. (default: ``None``)
+                The dict can contain 'EnableRemoteDebug'(bool).
+                For example,
+
+                .. code:: python
+
+                    remote_debug_config = {
+                        "EnableRemoteDebug": True,
+                    } (default: None).
 
         Returns:
             sagemaker.estimator._TrainingJob: Constructed object that captures
             all information about the updated training job.
         """
         update_args = cls._get_update_args(
-            estimator, profiler_rule_configs, profiler_config, resource_config
+            estimator, profiler_rule_configs, profiler_config, resource_config, remote_debug_config
         )
         estimator.sagemaker_session.update_training_job(**update_args)
 
         return estimator.latest_training_job
 
     @classmethod
-    def _get_update_args(cls, estimator, profiler_rule_configs, profiler_config, resource_config):
+    def _get_update_args(
+        cls, estimator, profiler_rule_configs, profiler_config, resource_config, remote_debug_config
+    ):
         """Constructs a dict of arguments for updating an Amazon SageMaker training job.
 
         Args:
@@ -2596,6 +2651,7 @@ def _get_update_args(cls, estimator, profiler_rule_configs, profiler_config, res
         update_args.update(build_dict("profiler_rule_configs", profiler_rule_configs))
         update_args.update(build_dict("profiler_config", profiler_config))
         update_args.update(build_dict("resource_config", resource_config))
+        update_args.update(build_dict("remote_debug_config", remote_debug_config))
 
         return update_args
 
@@ -2694,6 +2750,7 @@ def __init__(
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: bool = False,
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
+        enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -3055,6 +3112,8 @@ def __init__(
                 to Amazon S3 without compression after training finishes.
             enable_infra_check (bool or PipelineVariable): Optional.
                 Specifies whether it is running Sagemaker built-in infra check jobs.
+            enable_remote_debug (bool or PipelineVariable): Optional.
+                Specifies whether RemoteDebug is enabled for the training job
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -3106,6 +3165,7 @@ def __init__(
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
+            enable_remote_debug=enable_remote_debug,
             **kwargs,
         )
 
diff --git a/src/sagemaker/image_uri_config/spark.json b/src/sagemaker/image_uri_config/spark.json
@@ -21,6 +21,7 @@
                     "ap-southeast-3": "800295151634",
                     "ap-southeast-4": "819679513684",
                     "ca-central-1": "446299261295",
+                    "ca-west-1": "000907499111",
                     "cn-north-1": "671472414489",
                     "cn-northwest-1": "844356804704",
                     "eu-central-1": "906073651304",
@@ -61,6 +62,7 @@
                     "ap-southeast-3": "800295151634",
                     "ap-southeast-4": "819679513684",
                     "ca-central-1": "446299261295",
+                    "ca-west-1": "000907499111",
                     "cn-north-1": "671472414489",
                     "cn-northwest-1": "844356804704",
                     "eu-central-1": "906073651304",
@@ -101,6 +103,7 @@
                     "ap-southeast-3": "800295151634",
                     "ap-southeast-4": "819679513684",
                     "ca-central-1": "446299261295",
+                    "ca-west-1": "000907499111",
                     "cn-north-1": "671472414489",
                     "cn-northwest-1": "844356804704",
                     "eu-central-1": "906073651304",
@@ -141,6 +144,7 @@
                     "ap-southeast-3": "800295151634",
                     "ap-southeast-4": "819679513684",
                     "ca-central-1": "446299261295",
+                    "ca-west-1": "000907499111",
                     "cn-north-1": "671472414489",
                     "cn-northwest-1": "844356804704",
                     "eu-central-1": "906073651304",
@@ -181,6 +185,7 @@
                     "ap-southeast-3": "800295151634",
                     "ap-southeast-4": "819679513684",
                     "ca-central-1": "446299261295",
+                    "ca-west-1": "000907499111",
                     "cn-north-1": "671472414489",
                     "cn-northwest-1": "844356804704",
                     "eu-central-1": "906073651304",
diff --git a/src/sagemaker/jumpstart/estimator.py b/src/sagemaker/jumpstart/estimator.py
@@ -106,6 +106,7 @@ def __init__(
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
+        enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
     ):
         """Initializes a ``JumpStartEstimator``.
 
@@ -495,6 +496,8 @@ def __init__(
                 a training job.
             disable_output_compression (Optional[bool]): When set to true, Model is uploaded
                 to Amazon S3 without compression after training finishes.
+            enable_remote_debug (bool or PipelineVariable): Optional.
+                Specifies whether RemoteDebug is enabled for the training job
 
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
@@ -569,6 +572,7 @@ def _is_valid_model_id_hook():
             container_arguments=container_arguments,
             disable_output_compression=disable_output_compression,
             enable_infra_check=enable_infra_check,
+            enable_remote_debug=enable_remote_debug,
         )
 
         self.model_id = estimator_init_kwargs.model_id
diff --git a/src/sagemaker/jumpstart/factory/estimator.py b/src/sagemaker/jumpstart/factory/estimator.py
@@ -127,6 +127,7 @@ def get_init_kwargs(
     container_arguments: Optional[List[str]] = None,
     disable_output_compression: Optional[bool] = None,
     enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
+    enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
 ) -> JumpStartEstimatorInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Estimator` object."""
 
@@ -183,6 +184,7 @@ def get_init_kwargs(
         container_arguments=container_arguments,
         disable_output_compression=disable_output_compression,
         enable_infra_check=enable_infra_check,
+        enable_remote_debug=enable_remote_debug,
     )
 
     estimator_init_kwargs = _add_model_version_to_kwargs(estimator_init_kwargs)
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -1280,6 +1280,7 @@ class JumpStartEstimatorInitKwargs(JumpStartKwargs):
         "container_arguments",
         "disable_output_compression",
         "enable_infra_check",
+        "enable_remote_debug",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1344,6 +1345,7 @@ def __init__(
         container_arguments: Optional[List[str]] = None,
         disable_output_compression: Optional[bool] = None,
         enable_infra_check: Optional[Union[bool, PipelineVariable]] = None,
+        enable_remote_debug: Optional[Union[bool, PipelineVariable]] = None,
     ) -> None:
         """Instantiates JumpStartEstimatorInitKwargs object."""
 
@@ -1401,6 +1403,7 @@ def __init__(
         self.container_arguments = container_arguments
         self.disable_output_compression = disable_output_compression
         self.enable_infra_check = enable_infra_check
+        self.enable_remote_debug = enable_remote_debug
 
 
 class JumpStartEstimatorFitKwargs(JumpStartKwargs):
diff --git a/src/sagemaker/serve/utils/telemetry_logger.py b/src/sagemaker/serve/utils/telemetry_logger.py
@@ -59,21 +59,35 @@ def wrapper(self, *args, **kwargs):
             caught_ex = None
 
             image_uri_tail = self.image_uri.split("/")[1]
-            extra = f"{func_name}&{MODEL_SERVER_TO_CODE[str(self.model_server)]}&{image_uri_tail}"
+            extra = (
+                f"{func_name}"
+                f"&x-modelServer={MODEL_SERVER_TO_CODE[str(self.model_server)]}"
+                f"&x-imageTag={image_uri_tail}"
+            )
 
             if self.model_server == ModelServer.DJL_SERVING or self.model_server == ModelServer.TGI:
-                extra += f"&{self.model}"
+                extra += f"&x-modelName={self.model}"
 
             try:
                 response = func(self, *args, **kwargs)
                 if not self.serve_settings.telemetry_opt_out:
                     _send_telemetry(
-                        "1", MODE_TO_CODE[str(self.mode)], self.sagemaker_session, None, extra
+                        "1",
+                        MODE_TO_CODE[str(self.mode)],
+                        self.sagemaker_session,
+                        None,
+                        None,
+                        extra,
                     )
             except ModelBuilderException as e:
                 if not self.serve_settings.telemetry_opt_out:
                     _send_telemetry(
-                        "0", MODE_TO_CODE[str(self.mode)], self.sagemaker_session, str(e), extra
+                        "0",
+                        MODE_TO_CODE[str(self.mode)],
+                        self.sagemaker_session,
+                        str(e),
+                        e.__class__.__name__,
+                        extra,
                     )
                 caught_ex = e
             except Exception as e:  # pylint: disable=W0703
@@ -93,13 +107,22 @@ def _send_telemetry(
     mode: int,
     session: Session,
     failure_reason: str = None,
+    failure_type: str = None,
     extra_info: str = None,
 ) -> None:
     """Make GET request to an empty object in S3 bucket"""
     try:
         accountId = _get_accountId(session)
         region = _get_region_or_default(session)
-        url = _construct_url(accountId, str(mode), status, failure_reason, extra_info, region)
+        url = _construct_url(
+            accountId,
+            str(mode),
+            status,
+            failure_reason,
+            failure_type,
+            extra_info,
+            region,
+        )
         _requests_helper(url, 2)
         logger.debug("ModelBuilder metrics emitted.")
     except Exception:  # pylint: disable=W0703
@@ -111,6 +134,7 @@ def _construct_url(
     mode: str,
     status: str,
     failure_reason: str,
+    failure_type: str,
     extra_info: str,
     region: str,
 ) -> str:
@@ -124,6 +148,7 @@ def _construct_url(
     )
     if failure_reason:
         base_url += f"&x-failureReason={failure_reason}"
+        base_url += f"&x-failureType={failure_type}"
     if extra_info:
         base_url += f"&x-extra={extra_info}"
     return base_url
diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py
diff --git a/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py b/tests/unit/sagemaker/serve/utils/test_telemetry_logger.py
diff --git a/tests/unit/sagemaker/workflow/test_automl_step.py b/tests/unit/sagemaker/workflow/test_automl_step.py
diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
