adjust stopping behavior

ffmpeg_distributed.py

@@ -10,7 +10,7 @@ from threading import Thread
 from time import sleep, time
 from typing import List, NamedTuple, Callable, Union, Tuple
 from signal import signal, SIGINT
-from sys import exit, stderr, platform as sys_platform
+from sys import exit, stderr, stdin, platform as sys_platform
 from shlex import split, join
 from hashlib import md5
 from time import strptime
@@ -161,6 +161,11 @@ class FFMPEGProc:
 
     def stop(self):
         self._should_stop = True
+        if getattr(self, '_proc', None) is not None and self._proc.poll() is None:
+            try:
+                self._proc.terminate()
+            except OSError:
+                pass
 
     def _read_stderr_loop(self, stderr_lines: list):
         """Read stderr in a loop (used on Windows where select.poll is unavailable)."""
@@ -270,6 +275,11 @@ class TaskThread(Thread):
 
     def stop(self):
         self._should_stop = True
+        if getattr(self, '_reader_proc', None) is not None and self._reader_proc.poll() is None:
+            try:
+                self._reader_proc.terminate()
+            except OSError:
+                pass
         if self._ffmpeg:
             self._ffmpeg.stop()
 
@@ -305,21 +315,24 @@ class TaskThread(Thread):
                 if self._host != 'localhost':
                     encoder_cmd = ['ssh', '-o', 'ConnectTimeout=15', self._host, join(encoder_cmd)]
 
-                reader_proc = Popen(reader_cmd, stdout=PIPE, stderr=DEVNULL)
+                self._reader_proc = Popen(reader_cmd, stdout=PIPE, stderr=DEVNULL)
                 ret = -1
                 try:
                     with open(task.output_file, 'wb') as outfile:
                         self._ffmpeg = FFMPEGProc(
-                            encoder_cmd, stdin=reader_proc.stdout, stdout=outfile,
+                            encoder_cmd, stdin=self._reader_proc.stdout, stdout=outfile,
                             update_callback=upd, binary_io=True
                         )
                         ret = self._ffmpeg.run()
                 finally:
-                    try:
-                        reader_proc.wait(timeout=2)
-                    except TimeoutExpired:
-                        reader_proc.terminate()
-                        reader_proc.wait(timeout=5)
+                    rp = getattr(self, '_reader_proc', None)
+                    if rp is not None:
+                        try:
+                            rp.wait(timeout=2)
+                        except TimeoutExpired:
+                            rp.terminate()
+                            rp.wait(timeout=5)
+                        self._reader_proc = None
 
                 if ret != 0:
                     try:
@@ -393,20 +406,67 @@ def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds:
     stderr.flush()
     threads = [TaskThread(host, input_file, task_queue, pos, remote_ffmpeg_path) for pos, host in enumerate(hosts, 0)]
 
-    def sigint(sig, stack):
-        print('Got SIGINT, stopping...')
+    def stop_all():
+        tqdm.write('Stopping all workers (killing ffmpeg/SSH on each host)...', file=stderr)
+        stderr.flush()
         for thread in threads:
             thread.stop()
         for thread in threads:
-            thread.join()
+            thread.join(timeout=5)
         exit(1)
 
-    signal(SIGINT, sigint)
+    stop_file = abspath('.encode_stop')
+
+    def stdin_stop_listener():
+        try:
+            while True:
+                line = stdin.readline()
+                if not line:
+                    break
+                if line.strip().lower() == 'stop':
+                    stop_all()
+                    break
+        except (OSError, EOFError):
+            pass
+
+    def stop_file_poller():
+        try:
+            while True:
+                sleep(1.5)
+                if isfile(stop_file):
+                    try:
+                        unlink(stop_file)
+                    except OSError:
+                        pass
+                    stop_all()
+                    break
+        except OSError:
+            pass
+
+    def sigint(sig, stack):
+        stop_all()
+
+    try:
+        signal(SIGINT, sigint)
+    except (ValueError, OSError):
+        pass  # SIGINT not available on this platform
+
+    if getattr(stdin, 'isatty', lambda: False)():
+        tqdm.write("To cancel: type 'stop' + Enter, or create .encode_stop (e.g. from another terminal: echo. > .encode_stop)", file=stderr)
+    else:
+        tqdm.write("To cancel: create .encode_stop in current dir (e.g. from another terminal: echo. > .encode_stop)", file=stderr)
+    stderr.flush()
+    if getattr(stdin, 'isatty', lambda: False)():
+        Thread(target=stdin_stop_listener, daemon=True).start()
+    Thread(target=stop_file_poller, daemon=True).start()
 
     for thread in threads:
         thread.start()
-    for thread in threads:
-        thread.join()
+    try:
+        for thread in threads:
+            thread.join()
+    except KeyboardInterrupt:
+        stop_all()
 
     list_path = f'{tmp_dir}/output_segments.txt'
     with open(list_path, 'w') as f: