change: Enable Experiment integ test on beta clients (aws#3590)

Author: qidewenwhen

tests/conftest.py

@@ -22,7 +22,7 @@
 from botocore.config import Config
 from packaging.version import Version
 
-from sagemaker import Session, image_uris, utils
+from sagemaker import Session, image_uris, utils, get_execution_role
 from sagemaker.local import LocalSession
 from sagemaker.workflow.pipeline_context import PipelineSession, LocalPipelineSession
 
@@ -93,6 +93,7 @@ def pytest_addoption(parser):
     parser.addoption("--sagemaker-client-config", action="store", default=None)
     parser.addoption("--sagemaker-runtime-config", action="store", default=None)
     parser.addoption("--boto-config", action="store", default=None)
+    parser.addoption("--sagemaker-metrics-config", action="store", default=None)
 
 
 def pytest_configure(config):
@@ -115,6 +116,12 @@ def sagemaker_runtime_config(request):
     return json.loads(config) if config else None
 
 
+@pytest.fixture(scope="session")
+def sagemaker_metrics_config(request):
+    config = request.config.getoption("--sagemaker-metrics-config")
+    return json.loads(config) if config else None
+
+
 @pytest.fixture(scope="session")
 def boto_session(request):
     config = request.config.getoption("--boto-config")
@@ -135,7 +142,9 @@ def region(boto_session):
 
 
 @pytest.fixture(scope="session")
-def sagemaker_session(sagemaker_client_config, sagemaker_runtime_config, boto_session):
+def sagemaker_session(
+    sagemaker_client_config, sagemaker_runtime_config, boto_session, sagemaker_metrics_config
+):
     sagemaker_client_config.setdefault("config", Config(retries=dict(max_attempts=10)))
     sagemaker_client = (
         boto_session.client("sagemaker", **sagemaker_client_config)
@@ -147,11 +156,17 @@ def sagemaker_session(sagemaker_client_config, sagemaker_runtime_config, boto_se
         if sagemaker_runtime_config
         else None
     )
+    metrics_client = (
+        boto_session.client("sagemaker-metrics", **sagemaker_metrics_config)
+        if sagemaker_metrics_config
+        else None
+    )
 
     return Session(
         boto_session=boto_session,
         sagemaker_client=sagemaker_client,
         sagemaker_runtime_client=runtime_client,
+        sagemaker_metrics_client=metrics_client,
     )
 
 
@@ -170,6 +185,11 @@ def local_pipeline_session(boto_session):
     return LocalPipelineSession(boto_session=boto_session)
 
 
+@pytest.fixture(scope="session")
+def execution_role(sagemaker_session):
+    return get_execution_role(sagemaker_session)
+
+
 @pytest.fixture(scope="module")
 def custom_bucket_name(boto_session):
     region = boto_session.region_name

tests/data/experiment/inference.py

@@ -10,6 +10,7 @@
 #  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.
+import json
 import logging
 import os
 import pickle as pkl
@@ -24,11 +25,32 @@
 sdk_file = f"{code_dir}/{sdk_name}"
 os.system(f"pip install {sdk_file}")
 
+
+def _get_client_config_in_dict(cfg_in_str) -> dict:
+    return json.loads(cfg_in_str) if cfg_in_str else None
+
+
 from sagemaker.session import Session
 from sagemaker.experiments import load_run
 
 boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-sagemaker_session = Session(boto_session=boto_session)
+
+sagemaker_client_config = _get_client_config_in_dict(os.environ.get("SM_CLIENT_CONFIG", None))
+sagemaker_metrics_config = _get_client_config_in_dict(os.environ.get("SM_METRICS_CONFIG", None))
+sagemaker_client = (
+    boto_session.client("sagemaker", **sagemaker_client_config) if sagemaker_client_config else None
+)
+metrics_client = (
+    boto_session.client("sagemaker-metrics", **sagemaker_metrics_config)
+    if sagemaker_metrics_config
+    else None
+)
+
+sagemaker_session = Session(
+    boto_session=boto_session,
+    sagemaker_client=sagemaker_client,
+    sagemaker_metrics_client=metrics_client,
+)
 
 
 def model_fn(model_dir):

tests/data/experiment/process_job_script_for_run_clz.py

@@ -13,6 +13,7 @@
 """This script file runs on SageMaker processing job"""
 from __future__ import absolute_import
 
+import json
 import logging
 import os
 import boto3
@@ -25,8 +26,28 @@
 from sagemaker.experiments import load_run
 
 
+def _get_client_config_in_dict(cfg_in_str) -> dict:
+    return json.loads(cfg_in_str) if cfg_in_str else None
+
+
 boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-sagemaker_session = Session(boto_session=boto_session)
+
+sagemaker_client_config = _get_client_config_in_dict(os.environ.get("SM_CLIENT_CONFIG", None))
+sagemaker_metrics_config = _get_client_config_in_dict(os.environ.get("SM_METRICS_CONFIG", None))
+sagemaker_client = (
+    boto_session.client("sagemaker", **sagemaker_client_config) if sagemaker_client_config else None
+)
+metrics_client = (
+    boto_session.client("sagemaker-metrics", **sagemaker_metrics_config)
+    if sagemaker_metrics_config
+    else None
+)
+
+sagemaker_session = Session(
+    boto_session=boto_session,
+    sagemaker_client=sagemaker_client,
+    sagemaker_metrics_client=metrics_client,
+)
 
 
 with load_run(sagemaker_session=sagemaker_session) as run:

tests/data/experiment/train_job_script_for_run_clz.py

@@ -13,6 +13,7 @@
 """This script file runs on SageMaker training job"""
 from __future__ import absolute_import
 
+import json
 import logging
 import time
 import os
@@ -24,8 +25,29 @@
 from sagemaker import Session
 from sagemaker.experiments import load_run, Run
 
+
+def _get_client_config_in_dict(cfg_in_str) -> dict:
+    return json.loads(cfg_in_str) if cfg_in_str else None
+
+
 boto_session = boto3.Session(region_name=os.environ["AWS_REGION"])
-sagemaker_session = Session(boto_session=boto_session)
+
+sagemaker_client_config = _get_client_config_in_dict(os.environ.get("SM_CLIENT_CONFIG", None))
+sagemaker_metrics_config = _get_client_config_in_dict(os.environ.get("SM_METRICS_CONFIG", None))
+sagemaker_client = (
+    boto_session.client("sagemaker", **sagemaker_client_config) if sagemaker_client_config else None
+)
+metrics_client = (
+    boto_session.client("sagemaker-metrics", **sagemaker_metrics_config)
+    if sagemaker_metrics_config
+    else None
+)
+
+sagemaker_session = Session(
+    boto_session=boto_session,
+    sagemaker_client=sagemaker_client,
+    sagemaker_metrics_client=metrics_client,
+)
 
 if os.environ["RUN_OPERATION"] == "init":
     logging.info("Initializing a Run")