add metrics client to session object (aws#745)

Co-authored-by: Dewen Qi <[email protected]>
Co-authored-by: Dana Benson <[email protected]>
Co-authored-by: Dana Benson <[email protected]>
Co-authored-by: qidewenwhen <[email protected]>

Author: 3 people

src/sagemaker/session.py

@@ -89,6 +89,7 @@ def __init__(
         sagemaker_featurestore_runtime_client=None,
         default_bucket=None,
         settings=SessionSettings(),
+        sagemaker_metrics_client=None,
     ):
         """Initialize a SageMaker ``Session``.
 
@@ -116,6 +117,10 @@ def __init__(
                 Example: "sagemaker-my-custom-bucket".
             settings (sagemaker.session_settings.SessionSettings): Optional. Set of optional
                 parameters to apply to the session.
+            sagemaker_metrics_client (boto3.SageMakerMetrics.Client):
+                Client which makes SageMaker Metrics related calls to Amazon SageMaker
+                (default: None). If not provided, one will be created using
+                this instance's ``boto_session``.
         """
         self._default_bucket = None
         self._default_bucket_name_override = default_bucket
@@ -130,6 +135,7 @@ def __init__(
             sagemaker_client=sagemaker_client,
             sagemaker_runtime_client=sagemaker_runtime_client,
             sagemaker_featurestore_runtime_client=sagemaker_featurestore_runtime_client,
+            sagemaker_metrics_client=sagemaker_metrics_client,
         )
 
     def _initialize(
@@ -138,6 +144,7 @@ def _initialize(
         sagemaker_client,
         sagemaker_runtime_client,
         sagemaker_featurestore_runtime_client,
+        sagemaker_metrics_client,
     ):
         """Initialize this SageMaker Session.
 
@@ -172,6 +179,12 @@ def _initialize(
                 "sagemaker-featurestore-runtime"
             )
 
+        if sagemaker_metrics_client:
+            self.sagemaker_metrics_client = sagemaker_metrics_client
+        else:
+            self.sagemaker_metrics_client = self.boto_session.client("sagemaker-metrics")
+        prepend_user_agent(self.sagemaker_metrics_client)
+
         self.local_mode = False
 
     @property

tests/data/experiment/docker/Dockerfile

@@ -0,0 +1,50 @@
+# 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.
+
+FROM ubuntu:22.04
+
+ARG script
+ARG library
+ARG botomodel
+ARG metricsmodel
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        python3 \
+        curl \
+        ca-certificates \
+        awscli \
+        && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \
+    python3 get-pip.py && \
+    rm get-pip.py
+
+
+WORKDIR /root
+
+COPY $library .
+
+# use a custom model
+# TODO: coumment out these four lines once the new API changes for Run release
+COPY $botomodel sagemaker-2017-07-24.normal.json
+RUN aws configure add-model --service-name sagemaker --service-model file://sagemaker-2017-07-24.normal.json
+COPY $metricsmodel sagemaker-metrics-2022-09-30.normal.json
+RUN aws configure add-model --service-name sagemaker-metrics --service-model file://sagemaker-metrics-2022-09-30.normal.json
+
+RUN python3 -m pip install $(basename $library)
+
+COPY $script script.py
+
+ENTRYPOINT ["python3", "./script.py"]

tests/data/experiment/docker/scripts/train_job_script_for_run_clz.py

@@ -0,0 +1,55 @@
+# 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 script file runs on SageMaker training job"""
+from __future__ import absolute_import
+
+import logging
+import time
+
+import os
+
+import boto3
+
+from sagemaker import Session
+from sagemaker.experiments.run import Run
+
+for key, value in os.environ.items():
+    logging.info("OS env var - {}: {}".format(key, value))
+
+boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
+sagemaker_session = Session(boto_session=boto_session)
+
+with Run.init(
+    experiment_name="my-train-job-exp-in-script",
+    run_name="my-train-job-run-in-script",
+    sagemaker_session=sagemaker_session,
+) as run:
+    logging.info(f"Run name: {run.run_name}")
+    logging.info(f"Experiment name: {run.experiment_name}")
+    logging.info(f"Trial component name: {run._trial_component.trial_component_name}")
+    run.log_parameter("p1", 1.0)
+    run.log_parameter("p2", 2.0)
+    if "TRAINING_JOB_ARN" in os.environ:
+        for i in range(2):
+            run.log_metric("A", i)
+        for i in range(2):
+            run.log_metric("B", i)
+        for i in range(2):
+            run.log_metric("C", i)
+        for i in range(2):
+            time.sleep(0.003)
+            run.log_metric("D", i)
+        for i in range(2):
+            time.sleep(0.003)
+            run.log_metric("E", i)
+        time.sleep(15)