fix: make RemoteExecutor context manager non-blocking on pending futures (#3822)

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

Author: nmadan

src/sagemaker/remote_function/client.py

@@ -731,7 +731,7 @@ def map(self, func, *iterables):
         futures = map(self.submit, itertools.repeat(func), *iterables)
         return [future.result() for future in futures]
 
-    def shutdown(self):
+    def shutdown(self, wait=True):
         """Prevent more function executions to be submitted to this executor."""
         with self._state_condition:
             self._shutdown = True
@@ -742,15 +742,15 @@ def shutdown(self):
             self._state_condition.notify_all()
 
         if self._workers is not None:
-            self._workers.shutdown(wait=True)
+            self._workers.shutdown(wait)
 
     def __enter__(self):
         """Create an executor instance and return it"""
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Make sure the executor instance is shutdown."""
-        self.shutdown()
+        self.shutdown(wait=False)
         return False
 
     @staticmethod

tests/unit/sagemaker/remote_function/test_client.py

@@ -509,6 +509,11 @@ def test_executor_submit_happy_case(mock_start, mock_job_settings, parallelism):
         future_3 = e.submit(job_function, 9, 10, c=11, d=12)
         future_4 = e.submit(job_function, 13, 14, c=15, d=16)
 
+    future_1.wait()
+    future_2.wait()
+    future_3.wait()
+    future_4.wait()
+
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (1, 2), {"c": 3, "d": 4}, None),
@@ -517,10 +522,6 @@ def test_executor_submit_happy_case(mock_start, mock_job_settings, parallelism):
             call(ANY, job_function, (13, 14), {"c": 15, "d": 16}, None),
         ]
     )
-    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()
@@ -545,14 +546,15 @@ def test_executor_submit_with_run(mock_start, mock_job_settings, run_obj):
             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_1.wait()
+    future_2.wait()
+
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (1, 2), {"c": 3, "d": 4}, run_info),
             call(ANY, job_function, (5, 6), {"c": 7, "d": 8}, run_info),
         ]
     )
-    mock_job_1.describe.assert_called()
-    mock_job_2.describe.assert_called()
 
     assert future_1.done()
     assert future_2.done()
@@ -562,14 +564,15 @@ def test_executor_submit_with_run(mock_start, mock_job_settings, run_obj):
             future_3 = e.submit(job_function, 9, 10, c=11, d=12)
             future_4 = e.submit(job_function, 13, 14, c=15, d=16)
 
+    future_3.wait()
+    future_4.wait()
+
     mock_start.assert_has_calls(
         [
             call(ANY, job_function, (9, 10), {"c": 11, "d": 12}, run_info),
             call(ANY, job_function, (13, 14), {"c": 15, "d": 16}, run_info),
         ]
     )
-    mock_job_3.describe.assert_called()
-    mock_job_4.describe.assert_called()
 
     assert future_3.done()
     assert future_4.done()
@@ -621,7 +624,7 @@ def test_executor_fails_to_start_job(mock_start, *args):
 
     with pytest.raises(TypeError):
         future_1.result()
-    print(future_2._state)
+    future_2.wait()
     assert future_2.done()
 
 
@@ -678,6 +681,8 @@ def test_executor_describe_job_throttled_temporarily(mock_start, *args):
         # submit second job
         future_2 = e.submit(job_function, 5, 6, c=7, d=8)
 
+    future_1.wait()
+    future_2.wait()
     assert future_1.done()
     assert future_2.done()
 
@@ -697,9 +702,9 @@ def test_executor_describe_job_failed_permanently(mock_start, *args):
         future_2 = e.submit(job_function, 5, 6, c=7, d=8)
 
     with pytest.raises(RuntimeError):
-        future_1.done()
+        future_1.result()
     with pytest.raises(RuntimeError):
-        future_2.done()
+        future_2.result()
 
 
 @pytest.mark.parametrize(