pathway: Update remote function test setup

* Update to remote function test image setup.
* Add basic test scenarios. Happy cases are enabled and can run successfully

Author: Zhankuil

tests/data/remote_function/containers/without_error/Dockerfile

@@ -19,55 +19,78 @@
 import pytest
 import docker
 
-from tests.integ import DATA_DIR
 from sagemaker import utils
-from sagemaker.utils import sagemaker_timestamp
+from sagemaker.utils import sagemaker_timestamp, _tmpdir
 
 REPO_ACCOUNT_ID = "033110030271"
 
+REPO_NAME = "remote-function-dummy-container"
 
-@pytest.fixture(scope="module")
+DOCKERFILE_TEMPLATE = (
+    "FROM public.ecr.aws/docker/library/python:{py_version}-slim\n\n"
+    "WORKDIR /opt/ml/remote_function/\n"
+    "COPY {source_archive} ./\n"
+    "RUN pip3 install '{source_archive}[remote_function]'\n"
+    "RUN rm {source_archive}\n"
+)
+
+
+@pytest.fixture(scope="package")
 def dummy_container_without_error(sagemaker_session):
-    repository_name = "remote-function-dummy-container"
-    docker_file_path = os.path.join(DATA_DIR, "remote_function", "containers", "without_error")
-    ecr_uri = _build_container(sagemaker_session, repository_name, docker_file_path)
+    # TODO: the python version should be dynamically specified instead of hardcoding
+    ecr_uri = _build_container(sagemaker_session, "3.10")
     return ecr_uri
 
 
-def _build_container(sagemaker_session, repository_name, docker_file_path):
-    """Build a dummy test container locally and push a container to an ecr repo"""
+@pytest.fixture(scope="package")
+def dummy_container_incompatible_python_runtime(sagemaker_session):
+    ecr_uri = _build_container(sagemaker_session, "3.7")
+    return ecr_uri
+
 
-    _generate_and_move_sagemaker_sdk_tar(docker_file_path)
+def _build_container(sagemaker_session, py_version):
+    """Build a dummy test container locally and push a container to an ecr repo"""
 
-    image_tag = sagemaker_timestamp()
+    region = sagemaker_session.boto_region_name
+    image_tag = f"{py_version.replace('.', '-')}-{sagemaker_timestamp()}"
     ecr_client = sagemaker_session.boto_session.client("ecr")
     username, password = _ecr_login(ecr_client)
 
-    docker_client = docker.from_env()
-    # build docker locally
-    image, _ = docker_client.images.build(path=docker_file_path, tag=repository_name, rm=True)
+    with _tmpdir() as tmpdir:
+        print("building docker image locally in ", tmpdir)
+        print("building source archive...")
+        source_archive = _generate_and_move_sagemaker_sdk_tar(tmpdir)
+        with open(os.path.join(tmpdir, "Dockerfile"), "w") as file:
+            file.writelines(
+                DOCKERFILE_TEMPLATE.format(py_version=py_version, source_archive=source_archive)
+            )
 
-    region = sagemaker_session.boto_region_name
-    ecr_image = ""
-    if _is_repository_exists(ecr_client, repository_name):
+        docker_client = docker.from_env()
+
+        print("building docker image...")
+        image, build_logs = docker_client.images.build(path=tmpdir, tag=REPO_NAME, rm=True)
+
+    if _is_repository_exists(ecr_client, REPO_NAME):
         sts_client = sagemaker_session.boto_session.client(
             "sts", region_name=region, endpoint_url=utils.sts_regional_endpoint(region)
         )
         account_id = sts_client.get_caller_identity()["Account"]
-        # When the test is run locally, repo will exists in same account whose credentials are used to run the test
+        # When the test is run locally, repo will exist in same account whose credentials are used to run the test
         ecr_image = _ecr_image_uri(
-            account_id, sagemaker_session.boto_region_name, repository_name, image_tag
+            account_id, sagemaker_session.boto_region_name, REPO_NAME, image_tag
         )
     else:
         ecr_image = _ecr_image_uri(
             REPO_ACCOUNT_ID,
             sagemaker_session.boto_region_name,
-            repository_name,
+            REPO_NAME,
             image_tag,
         )
 
+    print("pushing image...")
     image.tag(ecr_image, tag=image_tag)
     docker_client.images.push(ecr_image, auth_config={"username": username, "password": password})
+
     return ecr_image
 
 
@@ -99,7 +122,9 @@ def _generate_and_move_sagemaker_sdk_tar(destination_folder):
     """
     subprocess.run("python3 setup.py sdist", shell=True)
     dist_dir = "dist"
-    for item in os.listdir(dist_dir):
-        source = os.path.join(dist_dir, item)
-        destination = os.path.join(destination_folder, item)
-        shutil.copy2(source, destination)
+    source_archive = os.listdir(dist_dir)[0]
+    source_path = os.path.join(dist_dir, source_archive)
+    destination_path = os.path.join(destination_folder, source_archive)
+    shutil.copy2(source_path, destination_path)
+
+    return source_archive

tests/data/remote_function/containers/without_error/serve

@@ -15,38 +15,136 @@
 
 import pytest
 
-import sagemaker.exceptions
-from sagemaker import image_uris
 from sagemaker.remote_function import remote
 
+from tests.integ.kms_utils import get_or_create_kms_key
+
 ROLE = "SageMakerRole"
 
 
 @pytest.fixture(scope="module")
-def image_uri(
-    sklearn_latest_version,
-    sklearn_latest_py_version,
-    cpu_instance_type,
-    sagemaker_session,
+def s3_kms_key(sagemaker_session):
+    return get_or_create_kms_key(sagemaker_session=sagemaker_session)
+
+
+def test_decorator(sagemaker_session, dummy_container_without_error, cpu_instance_type):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+        keep_alive_period_in_seconds=30,
+    )
+    def divide(x, y):
+        return x / y
+
+    assert divide(10, 2) == 5
+    assert divide(20, 2) == 10
+
+
+@pytest.mark.skip
+def test_decorated_function_raises_exception(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
+):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+    )
+    def divide(x, y):
+        return x / y
+
+    with pytest.raises(ZeroDivisionError):
+        divide(10, 0)
+
+
+@pytest.mark.skip
+def test_remote_python_runtime_is_incompatible(
+    sagemaker_session, dummy_container_incompatible_python_runtime, cpu_instance_type
+):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_incompatible_python_runtime,
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+    )
+    def divide(x, y):
+        return x / y
+
+    # TODO: should raise serialization error
+    with pytest.raises(RuntimeError):
+        divide(10, 2)
+
+
+@pytest.mark.skip
+def test_advanced_job_setting(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type, s3_kms_key
+):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        instance_type=cpu_instance_type,
+        # TODO: add VPC settings
+        s3_kms_key=s3_kms_key,
+        sagemaker_session=sagemaker_session,
+    )
+    def divide(x, y):
+        return x / y
+
+    assert divide(10, 2) == 5
+
+
+@pytest.mark.skip
+def test_with_additional_dependencies(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
+):
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        dependencies="./requirements.txt",
+        instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
+    )
+    def divide(x, y):
+        return x / y
+
+    assert divide(10, 2) == 5
+
+
+@pytest.mark.skip
+def test_with_non_existent_dependencies(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):
-    return image_uris.retrieve(
-        "sklearn",
-        sagemaker_session.boto_region_name,
-        version=sklearn_latest_version,
-        py_version=sklearn_latest_py_version,
+    @remote(
+        role=ROLE,
+        image_uri=dummy_container_without_error,
+        dependencies="./requirements.txt",
         instance_type=cpu_instance_type,
+        sagemaker_session=sagemaker_session,
     )
+    def divide(x, y):
+        return x / y
+
+    # TODO: this should raise RuntimeEnvironmentError
+    with pytest.raises(RuntimeError):
+        divide(10, 2)
 
 
-def test_decorator(sagemaker_session, image_uri, cpu_instance_type):
+@pytest.mark.skip
+def test_with_incompatible_dependencies(
+    sagemaker_session, dummy_container_without_error, cpu_instance_type
+):
     @remote(
         role=ROLE,
-        image_uri=image_uri,
+        image_uri=dummy_container_without_error,
+        dependencies="./requirements.txt",
         instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,
     )
-    def square(x):
-        return x * x
+    def divide(x, y):
+        return x / y
 
-    with pytest.raises(sagemaker.exceptions.UnexpectedStatusException):
-        square(10)
+    # TODO: this should raise DeserializationError
+    with pytest.raises(RuntimeError):
+        divide(10, 2)

tests/data/remote_function/containers/without_error/train

@@ -12,42 +12,29 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import pytest
 
-import sagemaker.exceptions
-from sagemaker import image_uris
 from sagemaker.remote_function import RemoteExecutor
 
 ROLE = "SageMakerRole"
 
 
-@pytest.fixture(scope="module")
-def image_uri(
-    sklearn_latest_version,
-    sklearn_latest_py_version,
-    cpu_instance_type,
-    sagemaker_session,
-):
-    return image_uris.retrieve(
-        "sklearn",
-        sagemaker_session.boto_region_name,
-        version=sklearn_latest_version,
-        py_version=sklearn_latest_py_version,
-        instance_type=cpu_instance_type,
-    )
-
-
-def test_executor(sagemaker_session, image_uri, cpu_instance_type):
+def test_executor(sagemaker_session, dummy_container_without_error, cpu_instance_type):
     def square(x):
         return x * x
 
+    def cube(x):
+        return x * x * x
+
     with RemoteExecutor(
+        max_parallel_job=1,
         role=ROLE,
-        image_uri=image_uri,
+        image_uri=dummy_container_without_error,
         instance_type=cpu_instance_type,
         sagemaker_session=sagemaker_session,
+        keep_alive_period_in_seconds=30,
     ) as e:
-        future = e.submit(square, 10)
+        future_1 = e.submit(square, 10)
+        future_2 = e.submit(cube, 10)
 
-    with pytest.raises(sagemaker.exceptions.UnexpectedStatusException):
-        future.result()
+    assert future_1.result() == 100
+    assert future_2.result() == 1000