feature: support intelligent defaults config for pathways

Author: Namrata Madan

.flake8

@@ -3,3 +3,4 @@ application_import_names = sagemaker, tests
 import-order-style = google
 per-file-ignores =
     tests/unit/test_tuner.py: F405
+    src/sagemaker/config/config_schema.py: E501

src/sagemaker/config/config_schema.py

@@ -44,6 +44,15 @@
 SAGEMAKER = "SageMaker"
 PYTHON_SDK = "PythonSDK"
 MODULES = "Modules"
+REMOTE_FUNCTION = "RemoteFunction"
+DEPENDENCIES = "Dependencies"
+ENVIRONMENT_VARIABLES = "EnvironmentVariables"
+IMAGE_URI = "ImageUri"
+INSTANCE_TYPE = "InstanceType"
+S3_KMS_KEY_ID = "S3KmsKeyId"
+S3_ROOT_URI = "S3RootUri"
+SOURCE_DIR = "SourceDir"
+JOB_CONDA_ENV = "JobCondaEnvironment"
 OFFLINE_STORE_CONFIG = "OfflineStoreConfig"
 ONLINE_STORE_CONFIG = "OnlineStoreConfig"
 S3_STORAGE_CONFIG = "S3StorageConfig"
@@ -246,6 +255,7 @@ def _simple_path(*args: str):
 )
 
 
+<<<<<<< HEAD
 SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {
     "$schema": "https://json-schema.org/draft/2020-12/schema",
     TYPE: OBJECT,

src/sagemaker/remote_function/job.py

@@ -17,8 +17,10 @@
 import re
 from typing import Dict, List, Tuple
 
+from sagemaker.config import config_schema
+from sagemaker.config.config_factory import SageMakerConfigFactory
 from sagemaker.session import get_execution_role, _logs_for_job, Session
-from sagemaker.utils import name_from_base, base_name_from_image
+from sagemaker.utils import name_from_base
 from sagemaker.s3 import s3_path_join
 from sagemaker import vpc_utils
 from sagemaker.remote_function.core.stored_function import StoredFunction
@@ -32,7 +34,6 @@
 logger = logging_config.get_logger()
 
 
-# TODO: extend this class to load job settings from the configuration files.
 class _JobSettings:
     """Helper class that processes the job settings.
 
@@ -62,46 +63,84 @@ def __init__(
         volume_kms_key: str = None,
         volume_size: int = 30,
     ):
+
+        self.sagemaker_config = SageMakerConfigFactory.build_sagemaker_config(
+            additional_override_config_location=os.getcwd()
+        )
         self.sagemaker_session = sagemaker_session or Session()
 
-        self.environment_variables = environment_variables
+        self.environment_variables = self._get_from_config(
+            environment_variables, config_schema.ENVIRONMENT_VARIABLES
+        )
 
-        self.image_uri = image_uri or _JobSettings._get_default_image_uri()
-        self.dependencies = dependencies
+        # TODO: provide default image uri if not set
+        self.image_uri = self._get_from_config(image_uri, config_schema.IMAGE_URI, required=True)
+        self.dependencies = self._get_from_config(dependencies, config_schema.DEPENDENCIES)
 
-        self.instance_type = instance_type
+        self.instance_type = self._get_from_config(
+            instance_type, config_schema.INSTANCE_TYPE, required=True
+        )
         self.instance_count = instance_count
         self.volume_size = volume_size
         self.max_runtime_in_seconds = max_runtime_in_seconds
         self.max_retry_attempts = max_retry_attempts
         self.keep_alive_period_in_seconds = keep_alive_period_in_seconds
-        self.source_dir = source_dir
-        self.job_conda_env = job_conda_env
+        self.source_dir = self._get_from_config(source_dir, config_schema.SOURCE_DIR)
+        self.job_conda_env = self._get_from_config(job_conda_env, config_schema.JOB_CONDA_ENV)
 
-        if role is not None:
-            self.role = self.sagemaker_session.expand_role(role)
+        _role = self._get_from_config(role, config_schema.ROLE_ARN)
+        if _role:
+            self.role = self.sagemaker_session.expand_role(_role)
         else:
             self.role = get_execution_role(self.sagemaker_session)
 
-        self.s3_root_uri = s3_root_uri or os.path.join(
-            "s3://",
-            self.sagemaker_session.default_bucket(),
-            base_name_from_image(self.image_uri),
+        self.s3_root_uri = self._get_from_config(
+            s3_root_uri,
+            config_schema.S3_ROOT_URI,
+            default=os.path.join("s3://", self.sagemaker_session.default_bucket()),
         )
-        self.s3_kms_key = s3_kms_key
-        self.volume_kms_key = volume_kms_key
+        self.s3_kms_key = self._get_from_config(s3_kms_key, config_schema.S3_KMS_KEY_ID)
+        self.volume_kms_key = self._get_from_config(volume_kms_key, config_schema.VOLUME_KMS_KEY_ID)
 
-        vpc_config = vpc_utils.to_dict(subnets=subnets, security_group_ids=security_group_ids)
+        _subnets = self._get_from_config(subnets, config_schema.SUBNETS)
+        _security_group_ids = self._get_from_config(
+            security_group_ids, config_schema.SECURITY_GROUP_IDS
+        )
+        vpc_config = vpc_utils.to_dict(subnets=_subnets, security_group_ids=_security_group_ids)
         self.vpc_config = vpc_utils.sanitize(vpc_config)
 
-        self.tags = [] if tags is None else [{"Key": k, "Value": v} for k, v in tags]
-
-    @staticmethod
-    def _get_default_image_uri():
-        """Get the default image uri"""
+        self.tags = self._get_from_config(
+            tags,
+            config_schema.TAGS,
+            transform=lambda x: [next(iter(tuple(e.items()))) for e in x],
+            default=[],
+        )
 
-        # TODO: provide default image uri if not set
-        raise ValueError("image_uri must be set")
+    def _get_from_config(
+        self,
+        override_value,
+        sagemaker_config_key,
+        transform=lambda x: x,
+        default=None,
+        required=False,
+    ):
+        """Get default value from sagemaker config."""
+        if override_value:
+            return override_value
+        config_value = self.sagemaker_config.get_config_value(
+            "{}.{}.{}.{}.{}".format(
+                config_schema.SAGEMAKER,
+                config_schema.PYTHON_SDK,
+                config_schema.MODULES,
+                config_schema.REMOTE_FUNCTION,
+                sagemaker_config_key,
+            )
+        )
+        if config_value:
+            return transform(config_value)
+        if not default and required:
+            raise ValueError(f"{sagemaker_config_key} is a required parameter!")
+        return default
 
 
 class _Job:

tests/data/config/config.yaml

@@ -1,5 +1,21 @@
 SchemaVersion: '1.0'
 SageMaker:
+  PythonSDK:
+    Modules:
+      RemoteFunction:
+        Dependencies: "path/to/requirements.txt"
+        EnvironmentVariables: {"EnvVarKey": "EnvVarValue"}
+        ImageUri: "366666666666.dkr.ecr.us-west-2.amazonaws.com/my-image:latest"
+        InstanceType: "ml.m5.large"
+        JobCondaEnvironment: "my_conda_env"
+        RoleArn: "arn:aws:iam::366666666666:role/IMRole"
+        S3KmsKeyId: "somekmskeyid"
+        S3RootUri: "s3://bucket/key"
+        SecurityGroupIds: ["sg123"]
+        SourceDir: "../mymodule"
+        Subnets: ["subnet-1234"]
+        Tags: [{"someTagKey": "someTagValue"}]
+        VolumeKmsKeyId: "somekmskeyid"
   FeatureGroup:
     OnlineStoreConfig:
       SecurityConfig:

tests/data/remote_function/config.yaml

@@ -0,0 +1,15 @@
+SchemaVersion: '1.0'
+SageMaker:
+  PythonSDK:
+    Modules:
+      RemoteFunction:
+        Dependencies: "path/to/requirements.txt"
+        EnvironmentVariables: {"EnvVarKey": "EnvVarValue"}
+        InstanceType: "ml.m5.large"
+        JobCondaEnvironment: "my_conda_env"
+        S3KmsKeyId: "someS3KmsKey"
+        SecurityGroupIds: ["sg123"]
+        SourceDir: "../mymodule"
+        Subnets: ["subnet-1234"]
+        Tags: [{"someTagKey": "someTagValue"}, {"someTagKey2": "someTagValue2"}]
+        VolumeKmsKeyId: "someVolumeKmsKey"

tests/unit/sagemaker/config/conftest.py

@@ -22,17 +22,32 @@ def base_config_with_schema():
     return {"SchemaVersion": "1.0"}
 
 
+<<<<<<< HEAD
 @pytest.fixture()
 def valid_vpc_config():
     return {"SecurityGroupIds": ["sg123"], "Subnets": ["subnet-1234"]}
+=======
+@pytest.fixture(scope="module")
+def valid_vpc_subnet():
+    return "subnet-1234"
+
+
+@pytest.fixture(scope="module")
+def valid_vpc_security_group():
+    return "sg123"
+
+
+@pytest.fixture(scope="module")
+def valid_vpc_config(valid_vpc_security_group, valid_vpc_subnet):
+    return {"SecurityGroupIds": [valid_vpc_security_group], "Subnets": [valid_vpc_subnet]}
+>>>>>>> 740cd77d (feature: support intelligent defaults config for pathways)
 
 
 @pytest.fixture()
 def valid_iam_role_arn():
     return "arn:aws:iam::555555555555:role/IMRole"
 
 
-@pytest.fixture()
 def valid_feature_group_config(valid_iam_role_arn):
     security_storage_config = {"KmsKeyId": "kmskeyid1"}
     s3_storage_config = {"KmsKeyId": "kmskeyid2"}

tests/unit/sagemaker/config/test_config_factory.py

@@ -70,13 +70,14 @@ def test_default_config_file_when_directory_is_provided_as_the_path(
 
 
 def test_default_config_file_when_path_is_provided_as_environment_variable(
-    get_data_dir, valid_config_with_all_the_scopes, base_config_with_schema
+    get_data_dir, valid_config_with_all_the_scopes, base_config_with_schema, monkeypatch
 ):
-    os.environ["SAGEMAKER_DEFAULT_CONFIG_OVERRIDE"] = get_data_dir
+    monkeypatch.setenv("SAGEMAKER_DEFAULT_CONFIG_OVERRIDE", get_data_dir)
     # This will try to load config.yaml file from that directory if present.
     expected_config = base_config_with_schema
     expected_config["SageMaker"] = valid_config_with_all_the_scopes
     assert expected_config == SageMakerConfigFactory.build_sagemaker_config().get_config()
+    monkeypatch.delenv("SAGEMAKER_DEFAULT_CONFIG_OVERRIDE")
 
 
 def test_merge_behavior_when_additional_config_file_is_not_found(

tests/unit/sagemaker/remote_function/test_job.py

@@ -12,8 +12,11 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import os
+
 import pytest
 from mock import patch, Mock, ANY
+from tests.unit import DATA_DIR
 from sagemaker.remote_function.job import _JobSettings, _Job
 
 
@@ -80,19 +83,38 @@ def job_function(a, b=1, *, c, d=3):
 
 @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_default_settings(*args):
+def test_sagemaker_config_job_settings(get_execution_role, session, monkeypatch):
+    monkeypatch.setenv(
+        "SAGEMAKER_DEFAULT_CONFIG_OVERRIDE", os.path.join(DATA_DIR, "remote_function")
+    )
+
     job_settings = _JobSettings(image_uri="image_uri")
     assert job_settings.image_uri == "image_uri"
-    assert job_settings.s3_root_uri == f"s3://{BUCKET}/image_uri"
+    assert job_settings.s3_root_uri == f"s3://{BUCKET}"
     assert job_settings.role == DEFAULT_ROLE_ARN
-    assert job_settings.tags == []
+    assert job_settings.tags == [("someTagKey", "someTagValue"), ("someTagKey2", "someTagValue2")]
+    assert job_settings.vpc_config == {"Subnets": ["subnet-1234"], "SecurityGroupIds": ["sg123"]}
+    assert job_settings.dependencies == "path/to/requirements.txt"
+    assert job_settings.environment_variables == {"EnvVarKey": "EnvVarValue"}
+    assert job_settings.job_conda_env == "my_conda_env"
+    assert job_settings.source_dir == "../mymodule"
+    assert job_settings.volume_kms_key == "someVolumeKmsKey"
+    assert job_settings.s3_kms_key == "someS3KmsKey"
+    assert job_settings.instance_type == "ml.m5.large"
+
+    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_fails_on_missing_image_uri(*args):
-    with pytest.raises(ValueError):
-        _JobSettings(image_uri=None)
+def test_sagemaker_config_job_settings_missing_image_uri(get_execution_role, session, monkeypatch):
+    monkeypatch.setenv(
+        "SAGEMAKER_DEFAULT_CONFIG_OVERRIDE", os.path.join(DATA_DIR, "remote_function")
+    )
+
+    with pytest.raises(ValueError, match="ImageUri is a required parameter!"):
+        _JobSettings()
+    monkeypatch.delenv("SAGEMAKER_DEFAULT_CONFIG_OVERRIDE")
 
 
 @patch("sagemaker.remote_function.job.RuntimeEnvironmentManager")
@@ -101,7 +123,11 @@ def test_fails_on_missing_image_uri(*args):
 def test_start(session, mock_stored_function, mock_runtime_manager):
 
     job_settings = _JobSettings(
-        image_uri=IMAGE, s3_root_uri=S3_URI, role=ROLE_ARN, source_dir=PATH_TO_SRC_DIR
+        image_uri=IMAGE,
+        s3_root_uri=S3_URI,
+        role=ROLE_ARN,
+        source_dir=PATH_TO_SRC_DIR,
+        instance_type="ml.m5.large",
     )
 
     job = _Job.start(job_settings, job_function, func_args=(1, 2), func_kwargs={"c": 3, "d": 4})
@@ -139,7 +165,10 @@ def test_start(session, mock_stored_function, mock_runtime_manager):
             ],
         ),
         ResourceConfig=dict(
-            VolumeSizeInGB=30, InstanceCount=1, InstanceType=None, KeepAlivePeriodInSeconds=0
+            VolumeSizeInGB=30,
+            InstanceCount=1,
+            InstanceType="ml.m5.large",
+            KeepAlivePeriodInSeconds=0,
         ),
     )
 
@@ -221,7 +250,11 @@ def test_start_with_complete_job_settings(session, mock_stored_function, mock_ru
 def test_describe(session, *args):
 
     job_settings = _JobSettings(
-        image_uri=IMAGE, s3_root_uri=S3_URI, role=ROLE_ARN, source_dir=PATH_TO_SRC_DIR
+        image_uri=IMAGE,
+        s3_root_uri=S3_URI,
+        role=ROLE_ARN,
+        source_dir=PATH_TO_SRC_DIR,
+        instance_type="ml.m5.large",
     )
     job = _Job.start(job_settings, job_function, func_args=(1, 2), func_kwargs={"c": 3, "d": 4})
 
@@ -236,7 +269,11 @@ def test_describe(session, *args):
 def test_stop(session, *args):
 
     job_settings = _JobSettings(
-        image_uri=IMAGE, s3_root_uri=S3_URI, role=ROLE_ARN, source_dir=PATH_TO_SRC_DIR
+        image_uri=IMAGE,
+        s3_root_uri=S3_URI,
+        role=ROLE_ARN,
+        source_dir=PATH_TO_SRC_DIR,
+        instance_type="ml.m5.large",
     )
     job = _Job.start(job_settings, job_function, func_args=(1, 2), func_kwargs={"c": 3, "d": 4})
 
@@ -253,7 +290,11 @@ def test_stop(session, *args):
 def test_wait(session, mock_stored_function, mock_logs_for_job):
 
     job_settings = _JobSettings(
-        image_uri=IMAGE, s3_root_uri=S3_URI, role=ROLE_ARN, source_dir=PATH_TO_SRC_DIR
+        image_uri=IMAGE,
+        s3_root_uri=S3_URI,
+        role=ROLE_ARN,
+        source_dir=PATH_TO_SRC_DIR,
+        instance_type="ml.m5.large",
     )
     job = _Job.start(job_settings, job_function, func_args=(1, 2), func_kwargs={"c": 3, "d": 4})