Test for SIGTERM

feature: Test for SIGTERM

Author: bstriner

CONTRIBUTING.md

@@ -37,7 +37,7 @@ GitHub provides additional document on [forking a repository](https://help.githu
 ### Running the unit tests
 
 1. Install tox using `pip install tox`
-1. Install coverage using `pip install .[test]`
+1. Install coverage using `pip install ".[test]"`
 1. cd into the sagemaker-training-toolkit folder: `cd sagemaker-training-toolkit`
 1. Run the following tox command and verify that all code checks and unit tests pass: `tox test/unit`
 

src/sagemaker_training/process.py

@@ -137,7 +137,7 @@ async def run_async(cmd, processes_per_host, env, cwd, stderr, **kwargs):
         cmd, env=env, cwd=cwd, stdout=PIPE, stderr=stderr, **kwargs
     )
 
-    with capture_signal(signal.SIGTERM, lambda signalnum, frame: proc.send_signal(signalnum)):
+    with capture_signal(signal.SIGTERM, lambda signalnum,  *_: proc.send_signal(signalnum)):
         output = await asyncio.gather(
             watch(proc.stdout, processes_per_host), watch(proc.stderr, processes_per_host)
         )
@@ -218,7 +218,7 @@ def check_error(cmd, error_class, processes_per_host, cwd=None, capture_error=Tr
         process = subprocess.Popen(
             cmd, env=os.environ, cwd=cwd or environment.code_dir, stderr=stderr, **kwargs
         )
-        with capture_signal(signal.SIGTERM, lambda signalnum, frame: process.send_signal(signalnum)):
+        with capture_signal(signal.SIGTERM, lambda signalnum, *_: process.send_signal(signalnum)):
             return_code = process.wait()
     if return_code:
         extra_info = None

test/unit/_test_process_helper.py

@@ -0,0 +1,24 @@
+"""
+Helper script for testing signal handling
+
+- If it receives SIGTERM, immediately exit "21"
+- If it doesn't receive a signal, sleep for 3 seconds then exit "-1"
+"""
+
+import signal
+import time
+
+
+def signal_handler(signalnum, *_):
+    assert signalnum == signal.SIGTERM
+    exit(21)
+
+
+def main():
+    signal.signal(signal.SIGTERM, signal_handler)
+    time.sleep(3)
+    exit(-1)
+
+
+if __name__ == "__main__":
+    main()

test/unit/test_process.py

@@ -15,11 +15,14 @@
 import asyncio
 import os
 import sys
-
+import time
+import multiprocessing
 from mock import ANY, MagicMock, patch
 import pytest
 
 from sagemaker_training import environment, errors, process
+from sagemaker_training.process import capture_signal
+import signal
 
 
 class AsyncMock(MagicMock):
@@ -175,3 +178,32 @@ def test_run_python(log, async_shell, async_gather, entry_point_type_script, eve
         stdout=asyncio.subprocess.PIPE,
     )
     log.assert_called_with(cmd, {})
+
+
+def _sleep_subprocess(capture_error):
+    with pytest.raises(errors.ExecuteUserScriptError) as error:
+        process.check_error(
+            [
+                sys.executable,
+                os.path.abspath(os.path.join(__file__, "../_test_process_helper.py"))
+            ],
+            errors.ExecuteUserScriptError,
+            1,
+            capture_error=capture_error
+        )
+    assert int(error.value.return_code) == 21
+    exit(42)
+
+
+@pytest.mark.skipif(sys.version_info != (3, 7), reason="requires python3.7")
+@pytest.mark.parametrize("capture_error", [True, False])
+def test_check_error_signal(capture_error):
+    proc = multiprocessing.Process(
+        target=_sleep_subprocess,
+        args=(capture_error,)
+    )
+    proc.start()
+    time.sleep(1)
+    os.kill(proc.pid, signal.SIGTERM)
+    proc.join(1)
+    assert int(proc.exitcode) == 42