fix: remote function include_local_workdir default value (#1342)

Co-authored-by: Namrata Madan <[email protected]>

Author: 2 people

src/sagemaker/remote_function/client.py

@@ -65,7 +65,7 @@ def remote(
     pre_execution_script: str = None,
     environment_variables: Dict[str, str] = None,
     image_uri: str = None,
-    include_local_workdir: bool = False,
+    include_local_workdir: bool = None,
     custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
     instance_count: int = 1,
     instance_type: str = None,
@@ -495,7 +495,7 @@ def __init__(
         pre_execution_script: str = None,
         environment_variables: Dict[str, str] = None,
         image_uri: str = None,
-        include_local_workdir: bool = False,
+        include_local_workdir: bool = None,
         custom_file_filter: Optional[Union[Callable[[str, List], List], CustomFileFilter]] = None,
         instance_count: int = 1,
         instance_type: str = None,

src/sagemaker/workflow/pipeline.py

@@ -971,8 +971,8 @@ def result(self, step_name: str):
         try:
             self.wait()
         except WaiterError as e:
-            if "Waiter encountered a terminal failure state" in str(e):
-                pass
+            if "Waiter encountered a terminal failure state" not in str(e):
+                raise
         step = next(filter(lambda x: x["StepName"] == step_name, self.list_steps()), None)
         if not step:
             raise ValueError(f"Invalid step name {step_name}")

tests/unit/sagemaker/remote_function/test_client.py

@@ -12,11 +12,14 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import os
 import threading
 import time
 
 import pytest
 from mock import MagicMock, patch, Mock, ANY, call
+
+from sagemaker.config import load_sagemaker_config
 from sagemaker.exceptions import UnexpectedStatusException
 
 from botocore.exceptions import ClientError
@@ -36,6 +39,7 @@
     RuntimeEnvironmentError,
 )
 from sagemaker.remote_function.job import _RunInfo
+from tests.unit import DATA_DIR
 
 from tests.unit.sagemaker.experiments.helpers import (
     mock_tc_load_or_create_func,
@@ -50,6 +54,7 @@
 EXPECTED_JOB_RESULT = [1, 2, 3]
 PATH_TO_SRC_DIR = "path/to/src/dir"
 HMAC_KEY = "some-hmac-key"
+ROLE_ARN = "arn:aws:iam::555555555555:role/my_execution_role_arn"
 
 
 def describe_training_job_response(job_status):
@@ -175,6 +180,37 @@ def square(x):
     assert mock_job_settings.call_args.kwargs["image_uri"] == IMAGE
 
 
+@patch(
+    "sagemaker.remote_function.client.serialization.deserialize_obj_from_s3",
+    return_value=EXPECTED_JOB_RESULT,
+)
+@patch("sagemaker.remote_function.client._Job.start")
+@patch("sagemaker.remote_function.job.Session")
+def test_decorator_with_config_file(session, mock_start, mock_deserialize_obj_from_s3):
+    session().get_caller_identity_arn = lambda: ROLE_ARN
+    session().sagemaker_config = load_sagemaker_config(
+        additional_config_paths=[os.path.join(DATA_DIR, "remote_function")]
+    )
+
+    mock_job = Mock(job_name=TRAINING_JOB_NAME)
+    mock_job.describe.return_value = COMPLETED_TRAINING_JOB
+
+    mock_start.return_value = mock_job
+
+    @remote(image_uri=IMAGE, s3_root_uri=S3_URI)
+    def square(x):
+        return x * x
+
+    result = square(5)
+    assert result == EXPECTED_JOB_RESULT
+    assert square.job_settings.image_uri == IMAGE
+    assert square.job_settings.s3_root_uri == S3_URI
+    # assert values are read from sagemaker defaults config file
+    assert square.job_settings.include_local_workdir is True
+    assert square.job_settings.custom_file_filter.ignore_name_patterns == ["data", "test"]
+    assert square.job_settings.s3_kms_key == "someS3KmsKey"
+
+
 @patch(
     "sagemaker.remote_function.client.serialization.deserialize_obj_from_s3",
     return_value=EXPECTED_JOB_RESULT,
@@ -609,6 +645,37 @@ def test_executor_submit_happy_case(mock_start, mock_job_settings, parallelism):
     assert future_4.done()
 
 
+@patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
+@patch("sagemaker.remote_function.client._Job.start")
+@patch("sagemaker.session.Session")
+def test_executor_submit_with_config_file(session, mock_start):
+    session().get_caller_identity_arn = lambda: ROLE_ARN
+    session().sagemaker_config = load_sagemaker_config(
+        additional_config_paths=[os.path.join(DATA_DIR, "remote_function")]
+    )
+
+    mock_job = create_mock_job("job_1", COMPLETED_TRAINING_JOB)
+    mock_start.side_effect = [mock_job]
+
+    with RemoteExecutor(
+        max_parallel_jobs=1,
+        s3_root_uri="s3://bucket/",
+        image_uri=IMAGE,
+        sagemaker_session=session(),
+    ) as e:
+        future = e.submit(job_function, 1, 2, c=3, d=4)
+
+        # assert values are read from sagemaker defaults config file
+        assert e.job_settings.include_local_workdir is True
+        assert e.job_settings.custom_file_filter.ignore_name_patterns == ["data", "test"]
+        assert e.job_settings.s3_kms_key == "someS3KmsKey"
+
+    mock_start.assert_called_with(ANY, job_function, (1, 2), {"c": 3, "d": 4}, None)
+    mock_job.describe.assert_called()
+
+    assert future.done()
+
+
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)
 @patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")