pathway: Fix race conditions in the RemoteExecutor

There is a race condition which may cause the polling worker to abort
early: it checks the number of running jobs as the exist criteria, however
it is possible that the running job list is empty because the submit worker
is slow to put the pending job to the running job list.

Author: Zhankuil

src/sagemaker/remote_function/client.py

@@ -14,7 +14,7 @@
 from __future__ import absolute_import
 
 from concurrent.futures import ThreadPoolExecutor
-import queue
+from collections import deque
 import time
 import threading
 from typing import Dict, List, Tuple, Any
@@ -171,26 +171,36 @@ def __init__(self, future, job_settings: _JobSettings, func, func_args, func_kwa
 
 def _submit_worker(executor):
     """Background worker that submits job requests."""
+
+    def has_work_to_do():
+        return (
+            len(executor._pending_request_queue) > 0
+            and len(executor._running_jobs) < executor.max_parallel_job
+        )
+
     try:
         while True:
-            request = executor._pending_request_queue.get(block=True)
+            with executor._state_condition:
+                executor._state_condition.wait_for(has_work_to_do)
+                request = executor._pending_request_queue[0]
 
             if request is None:
+                with executor._state_condition:
+                    # remove the anchor from the pending queue
+                    executor._pending_request_queue.popleft()
                 return
 
-            executor._semaphore.acquire(blocking=True)
-
             time.sleep(_API_CALL_LIMIT["SubmittingIntervalInSecs"])
             # submit a new job
             job = request.future._start_and_notify(
                 request.job_settings, request.func, request.args, request.kwargs
             )
 
-            if job is None:
-                # job fails to submit
-                executor._semaphore.release()
-            else:
-                executor._running_jobs[job.job_name] = job
+            with executor._state_condition:
+                if job:
+                    executor._running_jobs[job.job_name] = job
+                # remove the request from the pending queue
+                executor._pending_request_queue.popleft()
     except Exception:  # pylint: disable=broad-except
         logger.exception("Error occurred while submitting CreateTrainingJob requests.")
 
@@ -199,8 +209,12 @@ def _polling_worker(executor):
     """Background worker that polls the status of the running jobs."""
     try:
         while True:
-            if executor._shutdown and len(executor._running_jobs) == 0:
-                return
+            with executor._state_condition:
+                if (
+                    executor._shutdown
+                    and len(executor._running_jobs) + len(executor._pending_request_queue) == 0
+                ):
+                    return
 
             time.sleep(
                 max(
@@ -211,16 +225,17 @@ def _polling_worker(executor):
             )
 
             # check if running jobs are terminated
-            for job_name in executor._running_jobs.keys():
+            for job_name in list(executor._running_jobs.keys()):
                 try:
                     time.sleep(_API_CALL_LIMIT["PollingIntervalInSecs"])
                     if executor._running_jobs[job_name].describe()["TrainingJobStatus"] in [
                         "Completed",
                         "Failed",
                         "Stopped",
                     ]:
-                        del executor._running_jobs[job_name]
-                        executor._semaphore.release()
+                        with executor._state_condition:
+                            del executor._running_jobs[job_name]
+                            executor._state_condition.notify_all()
                 except Exception as e:  # pylint: disable=broad-except
                     if (
                         not isinstance(e, ClientError)
@@ -231,8 +246,9 @@ def _polling_worker(executor):
                         logger.exception(
                             "Error occurred while checking the status of job %s", job_name
                         )
-                        del executor._running_jobs[job_name]
-                        executor._semaphore.release()
+                        with executor._state_condition:
+                            del executor._running_jobs[job_name]
+                            executor._state_condition.notify_all()
     except Exception:  # pylint: disable=broad-except
         logger.exception("Error occurred while monitoring the job statuses.")
 
@@ -331,14 +347,14 @@ def __init__(
             volume_size=volume_size,
         )
 
-        self._pending_request_queue = queue.SimpleQueue()
-        self._semaphore = threading.BoundedSemaphore(self.max_parallel_job)
+        self._state_condition = threading.Condition()
+        self._pending_request_queue = deque()
         # For thread safety, see
         # https://web.archive.org/web/20201108091210/http://effbot.org/pyfaq/what-kinds-of-global-value-mutation-are-thread-safe.htm
         self._running_jobs = dict()
+        self._shutdown = False
 
         self._workers: ThreadPoolExecutor = None
-        self._shutdown = False
 
     def submit(self, func, *args, **kwargs):
         """Execute the input function as a SageMaker job asynchronously.
@@ -353,24 +369,30 @@ def submit(self, func, *args, **kwargs):
 
         self._validate_submit_args(func, *args, **kwargs)
 
-        future = Future()
-        self._pending_request_queue.put(
-            _SubmitRequest(future, self.job_settings, func, args, kwargs)
-        )
+        with self._state_condition:
+            future = Future()
+            self._pending_request_queue.append(
+                _SubmitRequest(future, self.job_settings, func, args, kwargs)
+            )
+
+            if self._workers is None:
+                self._workers = ThreadPoolExecutor(2)
+                self._workers.submit(_submit_worker, self)
+                self._workers.submit(_polling_worker, self)
 
-        if self._workers is None:
-            self._workers = ThreadPoolExecutor(2)
-            self._workers.submit(_submit_worker, self)
-            self._workers.submit(_polling_worker, self)
+            self._state_condition.notify_all()
 
         return future
 
     def shutdown(self):
         """Prevent more function executions to be submitted to this executor."""
-        self._shutdown = True
+        with self._state_condition:
+            self._shutdown = True
+
+            # give a signal to the submitting worker so that it doesn't block on empty queue forever
+            self._pending_request_queue.append(None)
 
-        # give a signal to the submitting worker so that it doesn't block on empty queue forever
-        self._pending_request_queue.put(None)
+            self._state_condition.notify_all()
 
         if self._workers is not None:
             self._workers.shutdown(wait=True)

src/sagemaker/remote_function/core/stored_function.py

@@ -56,12 +56,20 @@ def save(self, func, source_dir=None, *args, **kwargs):
         """
         if source_dir:
             self._zip_and_upload_source_dir(source_dir)
+
+        logger.info(
+            f"Serializing function code to {s3_path_join(self.s3_base_uri, 'function.pkl')}"
+        )
         serialization.serialize_func_to_s3(
             func,
             self.sagemaker_session,
             s3_path_join(self.s3_base_uri, "function.pkl"),
             self.s3_kms_key,
         )
+
+        logger.info(
+            f"Serializing function arguments to {s3_path_join(self.s3_base_uri, 'arguments.pkl')}"
+        )
         serialization.serialize_obj_to_s3(
             (args, kwargs),
             self.sagemaker_session,

src/sagemaker/remote_function/job.py

@@ -22,11 +22,15 @@
 from sagemaker import vpc_utils
 from sagemaker.remote_function.core.stored_function import StoredFunction
 from sagemaker.remote_function.core.runtime_environment import RuntimeEnvironmentManager
+from sagemaker.remote_function import logging_config
 
 
 JOBS_CONTAINER_ENTRYPOINT = ["invoke-remote-function"]
 
 
+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.
@@ -191,6 +195,7 @@ def start(job_settings: _JobSettings, func, func_args, func_kwargs):
         if job_settings.environment_variables:
             request_dict["Environment"] = job_settings.environment_variables
 
+        logger.info("Creating job: %s", job_name)
         job_settings.sagemaker_session.sagemaker_client.create_training_job(**request_dict)
 
         return _Job(job_name, job_settings)

tests/unit/sagemaker/remote_function/test_client.py

@@ -55,6 +55,13 @@ def describe_training_job_response(job_status):
 }
 
 
+def create_mock_job(job_name, describe_return):
+    mock_job = Mock()
+    mock_job.describe.return_value = describe_return
+    mock_job.job_name = job_name
+    return mock_job
+
+
 def job_function(a, b=1, *, c, d=3):
     return a * b * c * d
 
@@ -179,37 +186,49 @@ def test_executor_submit_after_shutdown(*args):
         e.submit(job_function, 1, 2, c=3, d=4)
 
 
+@pytest.mark.parametrize("parallelism", [1, 2, 3, 4])
 @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")
-def test_executor_submit_happy_case(mock_start, *args):
-    mock_job = Mock()
-    mock_job.describe.return_value = COMPLETED_TRAINING_JOB
-    mock_job.job_name = TRAINING_JOB_NAME
-    mock_start.return_value = mock_job
-
-    with RemoteExecutor(max_parallel_job=2, s3_root_uri="s3://bucket/") as e:
+def test_executor_submit_happy_case(mock_start, MockJobSetting, parallelism):
+    mock_job_1 = create_mock_job("job_1", COMPLETED_TRAINING_JOB)
+    mock_job_2 = create_mock_job("job_2", COMPLETED_TRAINING_JOB)
+    mock_job_3 = create_mock_job("job_3", COMPLETED_TRAINING_JOB)
+    mock_job_4 = create_mock_job("job_4", COMPLETED_TRAINING_JOB)
+    mock_start.side_effect = [mock_job_1, mock_job_2, mock_job_3, mock_job_4]
+
+    with RemoteExecutor(max_parallel_job=parallelism, s3_root_uri="s3://bucket/") as e:
         future_1 = e.submit(job_function, 1, 2, c=3, d=4)
         future_2 = e.submit(job_function, 5, 6, c=7, d=8)
+        future_3 = e.submit(job_function, 9, 10, c=11, d=12)
+        future_4 = e.submit(job_function, 13, 14, c=15, d=16)
 
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (1, 2), {"c": 3, "d": 4}),
             call(ANY, job_function, (5, 6), {"c": 7, "d": 8}),
+            call(ANY, job_function, (9, 10), {"c": 11, "d": 12}),
+            call(ANY, job_function, (13, 14), {"c": 15, "d": 16}),
         ]
     )
+    mock_job_1.describe.assert_called()
+    mock_job_2.describe.assert_called()
+    mock_job_3.describe.assert_called()
+    mock_job_4.describe.assert_called()
+
     assert future_1.done()
     assert future_2.done()
-    mock_job.describe.assert_called()
+    assert future_3.done()
+    assert future_4.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")
 def test_executor_submit_enforcing_max_parallel_jobs(mock_start, *args):
-    mock_job = Mock()
-    mock_job.describe.return_value = INPROGRESS_TRAINING_JOB
-    mock_start.return_value = mock_job
+    mock_job_1 = create_mock_job("job_1", INPROGRESS_TRAINING_JOB)
+    mock_job_2 = create_mock_job("job_2", INPROGRESS_TRAINING_JOB)
+    mock_start.side_effect = [mock_job_1, mock_job_2]
 
     e = RemoteExecutor(max_parallel_job=1, s3_root_uri="s3://bucket/")
     future_1 = e.submit(job_function, 1, 2, c=3, d=4)
@@ -221,10 +240,15 @@ def test_executor_submit_enforcing_max_parallel_jobs(mock_start, *args):
     assert not future_2.running()
     mock_start.assert_called_with(ANY, job_function, (1, 2), {"c": 3, "d": 4})
 
-    mock_job.describe.return_value = COMPLETED_TRAINING_JOB
+    mock_job_1.describe.return_value = COMPLETED_TRAINING_JOB
+    mock_job_2.describe.return_value = COMPLETED_TRAINING_JOB
+
     e.shutdown()
 
     mock_start.assert_called_with(ANY, job_function, (5, 6), {"c": 7, "d": 8})
+    mock_job_1.describe.assert_called()
+    mock_job_2.describe.assert_called()
+
     assert future_1.done()
     assert future_2.done()
 
@@ -252,9 +276,9 @@ def test_executor_fails_to_start_job(mock_start, *args):
 @patch("sagemaker.remote_function.client._JobSettings")
 @patch("sagemaker.remote_function.client._Job.start")
 def test_executor_submit_and_cancel(mock_start, *args):
-    mock_job = Mock()
-    mock_job.describe.return_value = INPROGRESS_TRAINING_JOB
-    mock_start.return_value = mock_job
+    mock_job_1 = create_mock_job("job_1", INPROGRESS_TRAINING_JOB)
+    mock_job_2 = create_mock_job("job_2", INPROGRESS_TRAINING_JOB)
+    mock_start.side_effect = [mock_job_1, mock_job_2]
 
     e = RemoteExecutor(max_parallel_job=1, s3_root_uri="s3://bucket/")
 
@@ -266,12 +290,14 @@ def test_executor_submit_and_cancel(mock_start, *args):
     future_2.cancel()
 
     # let the first job complete
-    mock_job.describe.return_value = COMPLETED_TRAINING_JOB
+    mock_job_1.describe.return_value = COMPLETED_TRAINING_JOB
     e.shutdown()
 
+    mock_start.assert_called_once_with(ANY, job_function, (1, 2), {"c": 3, "d": 4})
+    mock_job_1.describe.assert_called()
+
     assert future_1.done()
     assert future_2.cancelled()
-    mock_start.assert_called_once_with(ANY, job_function, (1, 2), {"c": 3, "d": 4})
 
 
 @patch("sagemaker.remote_function.client._API_CALL_LIMIT", new=API_CALL_LIMIT)