feature: Auto-capture image URI (aws#861)

* pathways:add default image uri config

* pathways: Add unit test for studio image uri

* pathways: Add unit test for studio image uri

* change: Move pathways image config to job

---------

Co-authored-by: Mourya Baddam <[email protected]>

Author: 2 people

src/sagemaker/remote_function/job.py

@@ -16,6 +16,7 @@
 import os
 import re
 import shutil
+import sys
 from typing import Dict, List, Tuple
 
 from sagemaker.config import config_schema
@@ -120,8 +121,12 @@ def __init__(
             environment_variables, config_schema.ENVIRONMENT_VARIABLES
         )
 
-        # TODO: provide default image uri if not set
-        self.image_uri = self._get_from_config(image_uri, config_schema.IMAGE_URI, required=True)
+        _image_uri = self._get_from_config(image_uri, config_schema.IMAGE_URI)
+        if _image_uri:
+            self.image_uri = _image_uri
+        else:
+            self.image_uri = self._get_default_image(self.sagemaker_session)
+
         self.dependencies = self._get_from_config(dependencies, config_schema.DEPENDENCIES)
 
         self.instance_type = self._get_from_config(
@@ -171,7 +176,7 @@ def _get_from_config(
         default=None,
         required=False,
     ):
-        """Get default value from sagemaker config."""
+        """Get value from sagemaker config."""
         if override_value:
             return override_value
         config_value = self.sagemaker_config.get_config_value(
@@ -189,6 +194,33 @@ def _get_from_config(
             raise ValueError(f"{sagemaker_config_key} is a required parameter!")
         return default
 
+    @staticmethod
+    def _get_default_image(session):
+        """Return Studio notebook image, if in Studio env. Else, base python"""
+
+        if (
+            "SAGEMAKER_INTERNAL_IMAGE_URI" in os.environ
+            and os.environ["SAGEMAKER_INTERNAL_IMAGE_URI"]
+        ):
+            return os.environ["SAGEMAKER_INTERNAL_IMAGE_URI"]
+
+        py_major_version = sys.version_info[0]
+        py_minor_version = 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.")
+
+        # TODO: Support only supported by Studio
+        region = session.boto_region_name
+
+        # 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
+
 
 class _Job:
     """Helper class that interacts with the SageMaker training service."""

tests/unit/sagemaker/remote_function/test_job.py

@@ -13,6 +13,7 @@
 from __future__ import absolute_import
 
 import os
+import sys
 
 import pytest
 from mock import patch, Mock, ANY
@@ -123,8 +124,35 @@ def test_sagemaker_config_job_settings_missing_image_uri(get_execution_role, ses
         "SAGEMAKER_DEFAULT_CONFIG_OVERRIDE", os.path.join(DATA_DIR, "remote_function")
     )
 
-    with pytest.raises(ValueError, match="ImageUri is a required parameter!"):
-        _JobSettings()
+    py_major_version = sys.version_info[0]
+    py_minor_version = sys.version_info[1]
+    if py_major_version != 3 or py_minor_version != 10:
+        with pytest.raises(
+            ValueError, match="Use supported Python version or provide compatible ImageUri."
+        ):
+            _JobSettings()
+    else:
+        job_settings = _JobSettings()
+        assert (
+            job_settings.image_uri
+            == f"581474259216.dkr.ecr.{TEST_REGION}.amazonaws.com/sagemaker-pathways-beta:basepy_3_10_latest"
+        )
+
+    monkeypatch.delenv("SAGEMAKER_DEFAULT_CONFIG_OVERRIDE")
+
+
+@patch("sagemaker.remote_function.job.Session", return_value=mock_session())
+@patch("sagemaker.remote_function.job.get_execution_role", return_value=DEFAULT_ROLE_ARN)
+def test_sagemaker_config_job_settings_studio_image_uri(get_execution_role, session, monkeypatch):
+    monkeypatch.setenv(
+        "SAGEMAKER_DEFAULT_CONFIG_OVERRIDE", os.path.join(DATA_DIR, "remote_function")
+    )
+    monkeypatch.setenv("SAGEMAKER_INTERNAL_IMAGE_URI", "studio_image_uri")
+
+    job_settings = _JobSettings()
+    assert job_settings.image_uri == "studio_image_uri"
+
+    monkeypatch.delenv("SAGEMAKER_INTERNAL_IMAGE_URI")
     monkeypatch.delenv("SAGEMAKER_DEFAULT_CONFIG_OVERRIDE")