Actively running jobs respond to the kill switch by checking the switch's status in the monitor function. The monitor function, up to this point, has been for gathering memory usage statistics via a timer thread. A second timer thread now monitors the kill switch.

Author: AlexTate

cwltool/job.py

@@ -507,6 +507,7 @@ def process_monitor(self, sproc: "subprocess.Popen[str]", kill_switch: threading
         memory_usage: MutableSequence[Optional[int]] = [None]
 
         mem_tm: "Optional[Timer]" = None
+        ks_tm: "Optional[Timer]" = None
 
         def get_tree_mem_usage(memory_usage: MutableSequence[Optional[int]]) -> None:
             nonlocal mem_tm
@@ -528,10 +529,27 @@ def get_tree_mem_usage(memory_usage: MutableSequence[Optional[int]]) -> None:
                 if mem_tm is not None:
                     mem_tm.cancel()
 
+        def monitor_kill_switch() -> None:
+            nonlocal ks_tm
+            if kill_switch.is_set():
+                _logger.error("[job %s] terminating by kill switch", self.name)
+                if sproc.stdin: sproc.stdin.close()
+                sproc.terminate()
+            else:
+                ks_tm = Timer(interval=1, function=monitor_kill_switch)
+                ks_tm.daemon = True
+                ks_tm.start()
+
+        ks_tm = Timer(interval=1, function=monitor_kill_switch)
+        ks_tm.daemon = True
+        ks_tm.start()
+
         mem_tm = Timer(interval=1, function=get_tree_mem_usage, args=(memory_usage,))
         mem_tm.daemon = True
         mem_tm.start()
+
         sproc.wait()
+        ks_tm.cancel()
         mem_tm.cancel()
         if memory_usage[0] is not None:
             _logger.info(
@@ -845,20 +863,48 @@ def docker_monitor(
         process: "subprocess.Popen[str]",
         kill_switch: threading.Event,
     ) -> None:
-        """Record memory usage of the running Docker container."""
+        """Record memory usage of the running Docker container. Terminate if kill_switch is activated."""
+
+        ks_tm: "Optional[Timer]" = None
+        cid: Optional[str] = None
+
+        def monitor_kill_switch() -> None:
+            nonlocal ks_tm
+            if kill_switch.is_set():
+                _logger.error("[job %s] terminating by kill switch", self.name)
+                if process.stdin:
+                    process.stdin.close()
+                if cid is not None:
+                    kill_proc = subprocess.Popen(  # nosec
+                        [docker_exe, "kill", cid], shell=False  # nosec
+                    )
+                    try:
+                        kill_proc.wait(timeout=10)
+                    except subprocess.TimeoutExpired:
+                        kill_proc.kill()
+                process.terminate()  # Always terminate, even if we tried with the cidfile
+            else:
+                ks_tm = Timer(interval=1, function=monitor_kill_switch)
+                ks_tm.daemon = True
+                ks_tm.start()
+
+        ks_tm = Timer(interval=1, function=monitor_kill_switch)
+        ks_tm.daemon = True
+        ks_tm.start()
+
         # Todo: consider switching to `docker create` / `docker start`
         # instead of `docker run` as `docker create` outputs the container ID
         # to stdout, but the container is frozen, thus allowing us to start the
         # monitoring process without dealing with the cidfile or too-fast
         # container execution
-        cid: Optional[str] = None
         while cid is None:
             time.sleep(1)
             # This is needed to avoid a race condition where the job
             # was so fast that it already finished when it arrives here
             if process.returncode is None:
                 process.poll()
             if process.returncode is not None:
+                ks_tm.cancel()
                 if cleanup_cidfile:
                     try:
                         os.remove(cidfile)
@@ -890,6 +936,9 @@ def docker_monitor(
         except OSError as exc:
             _logger.warning("Ignored error with %s stats: %s", docker_exe, exc)
             return
+        finally:
+            ks_tm.cancel()
+
         max_mem_percent: float = 0.0
         mem_percent: float = 0.0
         with open(stats_file_name) as stats:
@@ -924,7 +973,7 @@ def _job_popen(
     job_script_contents: Optional[str] = None,
     timelimit: Optional[int] = None,
     name: Optional[str] = None,
-    monitor_function: Optional[Callable[["subprocess.Popen[str]"], None]] = None,
+    monitor_function: Optional[Callable[["subprocess.Popen[str]", "threading.Event"], None]] = None,
     default_stdout: Optional[Union[IO[bytes], TextIO]] = None,
     default_stderr: Optional[Union[IO[bytes], TextIO]] = None,
 ) -> int:
@@ -979,7 +1028,7 @@ def terminate():  # type: () -> None
             tm.daemon = True
             tm.start()
         if monitor_function:
-            monitor_function(sproc)
+            monitor_function(sproc, kill_switch)
         rcode = sproc.wait()
 
         if tm is not None:
@@ -1055,7 +1104,7 @@ def terminate():  # type: () -> None
                 tm.daemon = True
                 tm.start()
             if monitor_function:
-                monitor_function(sproc)
+                monitor_function(sproc, kill_switch)
 
             rcode = sproc.wait()