Pathways cleanup before GA (aws#908)

Co-authored-by: Rohan Gujarathi <[email protected]>

Author: 2 people

requirements/extras/remote_function_requirements.txt

@@ -49,6 +49,7 @@ def read_requirements(filename):
 required_packages = [
     "attrs>=20.3.0,<23",
     "boto3>=1.26.28,<2.0",
+    "cloudpickle==2.2.1",
     "google-pasta",
     "numpy>=1.9.0,<2.0",
     "protobuf>=3.1,<4.0",
@@ -62,14 +63,14 @@ def read_requirements(filename):
     "PyYAML==5.4.1",
     "jsonschema",
     "platformdirs",
+    "tblib==1.7.0",
 ]
 
 # Specific use case dependencies
 # Keep format of *_requirements.txt to be tracked by dependabot
 extras = {
     "local": read_requirements("requirements/extras/local_requirements.txt"),
     "scipy": read_requirements("requirements/extras/scipy_requirements.txt"),
-    "remote_function": read_requirements("requirements/extras/remote_function_requirements.txt"),
 }
 # Meta dependency groups
 extras["all"] = [item for group in extras.values() for item in group]

setup.py

@@ -0,0 +1,36 @@
+{
+    "versions": {
+        "1.0": {
+            "registries": {
+                "us-east-2": "429704687514",
+                "me-south-1": "117516905037",
+                "us-west-2": "236514542706",
+                "ca-central-1": "310906938811",
+                "ap-east-1": "493642496378",
+                "us-east-1": "081325390199",
+                "ap-northeast-2": "806072073708",
+                "eu-west-2": "712779665605",
+                "ap-southeast-2": "52832661640",
+                "cn-northwest-1": "390780980154",
+                "eu-north-1": "243637512696",
+                "cn-north-1": "390048526115",
+                "ap-south-1": "394103062818",
+                "eu-west-3": "615547856133",
+                "ap-southeast-3": "276181064229",
+                "af-south-1": "559312083959",
+                "eu-west-1": "470317259841",
+                "eu-central-1": "936697816551",
+                "sa-east-1": "782484402741",
+                "ap-northeast-3": "792733760839",
+                "eu-south-1": "592751261982",
+                "ap-northeast-1": "102112518831",
+                "us-west-1": "742091327244",
+                "ap-southeast-1": "492261229750",
+                "me-central-1": "103105715889",
+                "us-gov-east-1": "107072934176",
+                "us-gov-west-1": "107173498710"
+            },
+            "repository": "sagemaker-base-python"
+        }
+    }
+}

src/sagemaker/image_uri_config/sagemaker-base-python.json

@@ -663,3 +663,29 @@ def get_training_image_uri(
         container_version=container_version,
         training_compiler_config=compiler_config,
     )
+
+
+def get_base_python_image_uri(region, py_version="310") -> str:
+    """Retrieves the image URI for base python image.
+
+    Args:
+        region (str): The AWS region to use for image URI.
+        py_version (str): The python version to use for the image. Can be 310 or 38
+        Default to 310
+
+    Returns:
+        str: The image URI string.
+    """
+
+    framework = "sagemaker-base-python"
+    version = "1.0"
+    hostname = utils._botocore_resolver().construct_endpoint("ecr", region)["hostname"]
+    config = config_for_framework(framework)
+    version_config = config["versions"][_version_for_config(version, config)]
+
+    registry = _registry_from_region(region, version_config["registries"])
+
+    repo = version_config["repository"] + "-" + py_version
+    repo_and_tag = repo + ":" + version
+
+    return ECR_URI_TEMPLATE.format(registry=registry, hostname=hostname, repository=repo_and_tag)

src/sagemaker/image_uris.py

@@ -41,6 +41,7 @@
 )
 from sagemaker.experiments._run_context import _RunContext
 from sagemaker.experiments.run import Run
+from sagemaker.image_uris import get_base_python_image_uri
 from sagemaker.session import get_execution_role, _logs_for_job, Session
 from sagemaker.utils import name_from_base, _tmpdir, resolve_value_from_config
 from sagemaker.s3 import s3_path_join, S3Uploader
@@ -60,16 +61,11 @@
 # training channel names
 RUNTIME_SCRIPTS_CHANNEL_NAME = "sagemaker_remote_function_bootstrap"
 REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
-SAGEMAKER_WHL_CHANNEL_NAME = "sagemaker_whl_file"
 
 # run context dictionary keys
 KEY_EXPERIMENT_NAME = "experiment_name"
 KEY_RUN_NAME = "run_name"
 
-SAGEMAKER_SDK_WHL_FILE = (
-    "s3://sagemaker-pathways/beta/pysdk/sagemaker-2.132.1.dev0-py2.py3-none-any.whl"
-)
-
 JOBS_CONTAINER_ENTRYPOINT = [
     "/bin/bash",
     f"/opt/ml/input/data/{RUNTIME_SCRIPTS_CHANNEL_NAME}/{ENTRYPOINT_SCRIPT_NAME}",
@@ -280,22 +276,18 @@ def _get_default_image(session):
         ):
             return os.environ["SAGEMAKER_INTERNAL_IMAGE_URI"]
 
-        py_major_version = sys.version_info[0]
-        py_minor_version = sys.version_info[1]
+        py_version = str(sys.version_info[0]) + str(sys.version_info[1])
 
-        # TODO:Add Support for 3.8
-        if py_major_version != 3 or py_minor_version != 10:
-            raise ValueError("Use supported Python version or provide compatible ImageUri.")
+        if py_version not in ["310", "38"]:
+            raise ValueError(
+                "Default image is supported only for Python versions 3.8 and 3.10. If you "
+                "are using any other python version, you must provide a compatible image_uri."
+            )
 
-        # TODO: Support only supported by Studio
         region = session.boto_region_name
+        image_uri = get_base_python_image_uri(region=region, py_version=py_version)
 
-        # TODO: Remove beta image and use public base python
-        beta_image = (
-            f"581474259216.dkr.ecr.{region}.amazonaws.com/"
-            f"sagemaker-pathways-beta:basepy_3_10_latest"
-        )
-        return beta_image
+        return image_uri
 
 
 class _Job:
@@ -394,19 +386,6 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs, run_info=Non
                 )
             )
 
-        # temporary solution for public beta to make sagemaker installer available
-        # in the images, this should be removed before pathways GA.
-        input_data_config.append(
-            dict(
-                ChannelName=SAGEMAKER_WHL_CHANNEL_NAME,
-                DataSource={
-                    "S3DataSource": {
-                        "S3Uri": SAGEMAKER_SDK_WHL_FILE,
-                        "S3DataType": "S3Prefix",
-                    }
-                },
-            )
-        )
         request_dict["InputDataConfig"] = input_data_config
 
         output_config = {"S3OutputPath": s3_base_uri}

src/sagemaker/remote_function/job.py

@@ -33,8 +33,6 @@
 REMOTE_FUNCTION_WORKSPACE = "sm_rf_user_ws"
 BASE_CHANNEL_PATH = "/opt/ml/input/data"
 FAILURE_REASON_PATH = "/opt/ml/output/failure"
-SAGEMAKER_WHL_FILE_NAME = "sagemaker-2.132.1.dev0-py2.py3-none-any.whl"
-SAGEMAKER_WHL_CHANNEL = "sagemaker_whl_file"
 PRE_EXECUTION_SCRIPT_NAME = "pre_exec.sh"
 
 
@@ -109,7 +107,6 @@ def _bootstrap_runtime_environment(
             break
 
     if dependencies_file:
-        _create_sm_installer_symlinks()
         RuntimeEnvironmentManager().bootstrap(
             local_dependencies_file=dependencies_file,
             conda_env=conda_env,
@@ -123,29 +120,6 @@ def _bootstrap_runtime_environment(
         )
 
 
-def _create_sm_installer_symlinks():
-    """Creates a symlink for sagemaker installer
-
-    This function create a symlink of sagemaker whl such that public beta users won't need
-    to provide whl file path in their dependencies file. This logic should be removed before GA
-    """
-
-    whl_file_path = os.path.join(BASE_CHANNEL_PATH, SAGEMAKER_WHL_CHANNEL, SAGEMAKER_WHL_FILE_NAME)
-
-    destination_path1 = os.path.join(
-        BASE_CHANNEL_PATH, REMOTE_FUNCTION_WORKSPACE, SAGEMAKER_WHL_FILE_NAME
-    )
-    os.symlink(whl_file_path, destination_path1)
-    logger.info("Successfully created symlink from '%s' to '%s'.", destination_path1, whl_file_path)
-
-    destination_path2 = os.path.join(os.getcwd(), SAGEMAKER_WHL_FILE_NAME)
-    if not os.path.exists(destination_path2):
-        logger.info(
-            "Successfully created symlink from '%s' to '%s'.", destination_path2, whl_file_path
-        )
-        os.symlink(whl_file_path, destination_path2)
-
-
 def _write_failure_reason_file(failure_msg):
     """Create a file 'failure' with failure reason written if bootstrap runtime env failed.
 

src/sagemaker/remote_function/runtime_environment/bootstrap_runtime_environment.py

@@ -20,7 +20,6 @@
 import shlex
 import os
 import subprocess
-import tempfile
 import time
 
 
@@ -127,14 +126,10 @@ def _capture_from_local_runtime(self) -> str:
             )
 
         local_dependencies_path = os.path.join(os.getcwd(), "env_snapshot.yml")
-        with tempfile.NamedTemporaryFile(suffix=".yml", prefix="tmp_export") as tmp_file:
-            if conda_env_name is not None:
-                self._export_conda_env_from_env_name(conda_env_name, tmp_file.name)
-            else:
-                self._export_conda_env_from_prefix(conda_env_prefix, tmp_file.name)
-            data = self._replace_sagemaker_in_conda_env_yml(tmp_file.read())
-            with open(local_dependencies_path, "wb") as file:
-                file.write(data)
+        if conda_env_name is not None:
+            self._export_conda_env_from_env_name(conda_env_name, local_dependencies_path)
+        else:
+            self._export_conda_env_from_prefix(conda_env_prefix, local_dependencies_path)
 
         return local_dependencies_path
 
@@ -263,14 +258,6 @@ def _export_conda_env_from_prefix(self, prefix, local_path):
         _run_shell_cmd(cmd)
         logger.info("Conda environment %s exported successfully", prefix)
 
-    def _replace_sagemaker_in_conda_env_yml(self, data):
-        """Replace the sagemaker in conda yml file with local sagemaker installer"""
-
-        sagemaker_with_version = b"sagemaker==2.132.1.dev0"
-        sagemaker_whl_file = b"sagemaker-2.132.1.dev0-py2.py3-none-any.whl[remote_function]"
-
-        return data.replace(sagemaker_with_version, sagemaker_whl_file)
-
     def _write_conda_env_to_file(self, env_name):
         """Writes conda env to the text file"""
 

src/sagemaker/remote_function/runtime_environment/runtime_environment_manager.py

@@ -34,7 +34,7 @@
         && unzip awscliv2.zip \
         && ./aws/install\n\n"
     "COPY {source_archive} ./\n"
-    "RUN pip3 install '{source_archive}[remote_function]'\n"
+    "RUN pip3 install '{source_archive}'\n"
     "RUN rm {source_archive}\n"
 )
 
@@ -51,10 +51,10 @@
     "RUN mamba create -n integ_test_env python={py_version} -y \
         && mamba create -n default_env python={py_version} -y\n"
     "COPY {source_archive} ./\n"
-    "RUN pip install '{source_archive}[remote_function]' \
-        && mamba run -n base pip install '{source_archive}[remote_function]' \
-        && mamba run -n default_env pip install '{source_archive}[remote_function]' \
-        && mamba run -n integ_test_env pip install '{source_archive}[remote_function]'\n"
+    "RUN pip install '{source_archive}' \
+        && mamba run -n base pip install '{source_archive}' \
+        && mamba run -n default_env pip install '{source_archive}' \
+        && mamba run -n integ_test_env pip install '{source_archive}'\n"
     "ENV SHELL=/bin/bash\n"
     "ENV SAGEMAKER_JOB_CONDA_ENV=default_env\n"
 )
@@ -66,7 +66,7 @@
     "dependencies:\n"
     "  - scipy=1.7.3\n"
     "  - pip:\n"
-    "    - /opt/ml/remote_function/sagemaker-{sagemaker_version}.tar.gz[remote_function]\n"
+    "    - /opt/ml/remote_function/sagemaker-{sagemaker_version}.tar.gz\n"
     "prefix: /opt/conda/bin/conda\n"
 )
 
@@ -77,7 +77,7 @@
     "dependencies:\n"
     "  - scipy=1.7.3\n"
     "  - pip:\n"
-    "    - sagemaker-2.132.1.dev0-py2.py3-none-any.whl[remote_function]\n"
+    "    - sagemaker-2.132.1.dev0-py2.py3-none-any.whl\n"
     "prefix: /opt/conda/bin/conda\n"
 )
 

tests/integ/sagemaker/remote_function/conftest.py

@@ -290,31 +290,6 @@ def cuberoot(x):
         assert cuberoot(27) == 3
 
 
-# TODO: it pulls the SDK from our beta release bucket, which is outdated already.
-# It will be cleaned up before GA release.
-@pytest.mark.skip
-def test_additional_dependencies_sm_from_input_channel(
-    sagemaker_session,
-    dummy_container_with_conda,
-    cpu_instance_type,
-    conda_yml_file_sm_from_input_channel,
-):
-    @remote(
-        role=ROLE,
-        image_uri=dummy_container_with_conda,
-        dependencies=conda_yml_file_sm_from_input_channel,
-        instance_type=cpu_instance_type,
-        sagemaker_session=sagemaker_session,
-        job_conda_env="integ_test_env",
-    )
-    def cuberoot(x):
-        from scipy.special import cbrt
-
-        return cbrt(x)
-
-    assert cuberoot(27) == 3
-
-
 def test_with_non_existent_dependencies(
     sagemaker_session, dummy_container_without_error, cpu_instance_type
 ):

tests/integ/sagemaker/remote_function/test_decorator.py

@@ -78,3 +78,9 @@ def djl_framework_uri(repo, account, djl_version, primary_framework, region=REGI
     domain = ALTERNATE_DOMAINS.get(region, DOMAIN)
     tag = f"{djl_version}-{primary_framework}"
     return IMAGE_URI_FORMAT.format(account, region, domain, repo, tag)
+
+
+def base_python_uri(repo, account, region=REGION):
+    domain = ALTERNATE_DOMAINS.get(region, DOMAIN)
+    tag = "1.0"
+    return IMAGE_URI_FORMAT.format(account, region, domain, repo, tag)