feature: add integ tests for remote_function auto_capture functionality

Author: Rohan Gujarathi

tests/integ/sagemaker/remote_function/conftest.py

@@ -18,6 +18,7 @@
 import shutil
 import pytest
 import docker
+import re
 
 from sagemaker.utils import sagemaker_timestamp, _tmpdir, sts_regional_endpoint
 
@@ -57,25 +58,30 @@
     "ENV SAGEMAKER_JOB_CONDA_ENV=default_env\n"
 )
 
-CONDA_YML_FILE_TEMPLATE = (
-    "name: integ_test_env\n"
-    "channels:\n"
-    "  - defaults\n"
-    "dependencies:\n"
-    "  - scipy=1.7.3\n"
-    "  - pip:\n"
-    "    - /sagemaker-{sagemaker_version}.tar.gz\n"
-    "prefix: /opt/conda/bin/conda\n"
+AUTO_CAPTURE_CLIENT_DOCKER_TEMPLATE = (
+    "FROM public.ecr.aws/docker/library/python:{py_version}-slim\n\n"
+    'SHELL ["/bin/bash", "-c"]\n'
+    "RUN apt-get update -y \
+        && apt-get install -y unzip curl\n\n"
+    "RUN curl -L -O 'https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh' \
+        && bash Mambaforge-Linux-x86_64.sh -b -p '/opt/conda' \
+        && /opt/conda/bin/conda init bash\n\n"
+    "ENV PATH $PATH:/opt/conda/bin\n"
+    "COPY {source_archive} ./\n"
+    "RUN mamba create -n auto_capture_client python={py_version} -y \
+        && mamba run -n auto_capture_client pip install '{source_archive}' awscli boto3\n"
+    "COPY test_auto_capture.py .\n"
+    "CMD [\"mamba\", \"run\", \"-n\", \"auto_capture_client\", \"python\", \"test_auto_capture.py\"]\n"
 )
 
-CONDA_YML_FILE_WITH_SM_FROM_INPUT_CHANNEL = (
+CONDA_YML_FILE_TEMPLATE = (
     "name: integ_test_env\n"
     "channels:\n"
     "  - defaults\n"
     "dependencies:\n"
     "  - scipy=1.7.3\n"
     "  - pip:\n"
-    "    - sagemaker-2.132.1.dev0-py2.py3-none-any.whl\n"
+    "    - /sagemaker-{sagemaker_version}.tar.gz\n"
     "prefix: /opt/conda/bin/conda\n"
 )
 
@@ -99,6 +105,12 @@ def dummy_container_with_conda(sagemaker_session):
     return ecr_uri
 
 
+@pytest.fixture(scope="package")
+def auto_capture_test_container(sagemaker_session):
+    ecr_uri = _build_container_locally("3.10", AUTO_CAPTURE_CLIENT_DOCKER_TEMPLATE)
+    return ecr_uri
+
+
 @pytest.fixture(scope="package")
 def conda_env_yml():
     """Write conda yml file needed for tests"""
@@ -116,22 +128,6 @@ def conda_env_yml():
         os.remove(conda_yml_file_name)
 
 
-@pytest.fixture(scope="package")
-def conda_yml_file_sm_from_input_channel():
-    """Write conda yml file needed for tests"""
-
-    conda_yml_file_name = "conda_env_sm_from_input_channel.yml"
-    conda_file_path = os.path.join(os.getcwd(), conda_yml_file_name)
-
-    with open(conda_file_path, "w") as yml_file:
-        yml_file.writelines(CONDA_YML_FILE_WITH_SM_FROM_INPUT_CHANNEL)
-    yield conda_file_path
-
-    # cleanup
-    if os.path.isfile(conda_yml_file_name):
-        os.remove(conda_yml_file_name)
-
-
 def _build_container(sagemaker_session, py_version, docker_templete):
     """Build a dummy test container locally and push a container to an ecr repo"""
 
@@ -178,6 +174,23 @@ def _build_container(sagemaker_session, py_version, docker_templete):
     return ecr_image
 
 
+def _build_container_locally(py_version, docker_templete):
+    with _tmpdir() as tmpdir:
+        print("building docker image locally in ", tmpdir)
+        print("building source archive...")
+        source_archive = _generate_sdk_tar_with_public_version(tmpdir)
+        _move_auto_capture_test_file(tmpdir)
+        with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
+            file.writelines(
+                docker_templete.format(py_version=py_version, source_archive=source_archive)
+            )
+
+        docker_client = docker.from_env()
+
+        print("building docker image...")
+        image, build_logs = docker_client.images.build(path=tmpdir, tag=REPO_NAME, rm=True)
+        return image.id
+
 def _is_repository_exists(ecr_client, repo_name):
     try:
         ecr_client.describe_repositories(repositoryNames=[repo_name])
@@ -212,3 +225,30 @@ def _generate_and_move_sagemaker_sdk_tar(destination_folder):
     shutil.copy2(source_path, destination_path)
 
     return source_archive
+
+
+def _generate_sdk_tar_with_public_version(destination_folder):
+    with open(os.path.join(os.getcwd(), "VERSION"), "r+") as version_file:
+        dev_sagemaker_version = version_file.readline().strip()
+        public_sagemaker_version = re.sub("1.dev0", "0", dev_sagemaker_version)
+        version_file.seek(0)
+        version_file.write(public_sagemaker_version)
+        version_file.truncate()
+    shutil.rmtree("dist")
+    
+    source_archive = _generate_and_move_sagemaker_sdk_tar(destination_folder)
+
+    with open(os.path.join(os.getcwd(), "VERSION"), "r+") as version_file:
+        version_file.seek(0)
+        version_file.write(dev_sagemaker_version)
+        version_file.truncate()
+    shutil.rmtree("dist")
+
+    return source_archive
+
+
+def _move_auto_capture_test_file(destination_folder):
+    test_file_name = "test_auto_capture.py"
+    source_path = os.path.join(os.getcwd(), "tests", "integ", "sagemaker", "remote_function", test_file_name)
+    destination_path = os.path.join(destination_folder, test_file_name)
+    shutil.copy2(source_path, destination_path)

tests/integ/sagemaker/remote_function/test_auto_capture.py

@@ -0,0 +1,40 @@
+# 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 __future__ import absolute_import
+
+import os
+import pandas as pd
+
+from sagemaker.remote_function import remote
+
+
+@remote(
+    role="SageMakerRole",
+    instance_type="ml.m5.xlarge",
+    dependencies="auto_capture",
+)
+def multiply(dataframe: pd.DataFrame, factor: float):
+    return dataframe * factor
+
+
+df = pd.DataFrame(
+    {
+        "A": [14, 4, 5, 4, 1],
+        "B": [5, 2, 54, 3, 2],
+        "C": [20, 20, 7, 3, 8],
+        "D": [14, 3, 6, 2, 6],
+    }
+)
+
+if __name__ == "__main__":
+    multiply(df, 10.0)

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -20,6 +20,8 @@
 import random
 import string
 import pandas as pd
+import subprocess
+import shlex
 from sagemaker.experiments.run import Run, load_run
 from tests.integ.sagemaker.experiments.helpers import cleanup_exp_resources
 from sagemaker.experiments.trial_component import _TrialComponent
@@ -596,3 +598,20 @@ def get_file_content(file_names):
     with pytest.raises(RuntimeEnvironmentError) as e:
         get_file_content(["test_file_1", "test_file_2", "test_file_3"])
         assert "line 2: bws: command not found" in str(e)
+
+
+def test_decorator_auto_capture(sagemaker_session, auto_capture_test_container):
+    creds = sagemaker_session.boto_session.get_credentials()
+    region = sagemaker_session.boto_session.region_name
+    env = {
+            "AWS_ACCESS_KEY_ID": str(creds.access_key),
+            "AWS_SECRET_ACCESS_KEY": str(creds.secret_key),
+            "AWS_SESSION_TOKEN": str(creds.token),
+        }
+    cmd = (f"docker run -e AWS_ACCESS_KEY_ID={env['AWS_ACCESS_KEY_ID']} "
+    f"-e AWS_SECRET_ACCESS_KEY={env['AWS_SECRET_ACCESS_KEY']} "
+    f"-e AWS_SESSION_TOKEN={env['AWS_SESSION_TOKEN']} "
+    f"-e AWS_DEFAULT_REGION={region} "
+    f"-it {auto_capture_test_container}")
+    result = subprocess.check_output(shlex.split(cmd), stderr=subprocess.STDOUT).decode("utf-8")
+    print(result)