swap back to nvenc/HEVC and redefine hosts accordingly

encode_VOD_pyro.py

@@ -11,10 +11,34 @@ import shutil
 import time
 
 # Distributed mode requires tqdm and ffmpeg_distributed.py (SSH, Unix select.poll); on Windows use WSL or Linux.
-DISTRIBUTED_HOSTS_DEFAULT = ["Pyro", "RenderScrap", "root@GuiltsCurse", "PostIrony", "root@Godzilla"]
-DISTRIBUTED_REMOTE_ARGS_DEFAULT = "-c:v libaom-av1 -crf 0 -b:v 9000k -maxrate 9000k -bufsize 18000k -cpu-used 8 -row-mt 1 -an"
+# Workers = (ssh_host, gpu_index). Unraid (GuiltsCurse, Godzilla) excluded; RenderScrap has 2 GPUs.
+DISTRIBUTED_WORKERS_DEFAULT = [
+    ("Pyro", 0),
+    ("RenderScrap", 0),
+    ("RenderScrap", 1),
+    ("PostIrony", 0),
+]
+DISTRIBUTED_REMOTE_ARGS_DEFAULT = "-c:v hevc_nvenc -preset p7 -tune hq -rc vbr -rc-lookahead 32 -spatial-aq 1 -aq-strength 15 -cq 0 -b:v 9000k -maxrate 9000k -bufsize 18000k -an"
 DISTRIBUTED_SEGMENT_SECONDS = 60
 
+
+def _parse_workers_env(s):
+    """Parse DISTRIBUTED_WORKERS e.g. 'Pyro:0,RenderScrap:0,RenderScrap:1,PostIrony:0' -> [(host, gpu_id), ...]."""
+    out = []
+    for part in (s or "").strip().split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if ":" in part:
+            host, gpu = part.rsplit(":", 1)
+            try:
+                out.append((host.strip(), int(gpu.strip())))
+            except ValueError:
+                pass
+        else:
+            out.append((part, 0))
+    return out
+
 # ANSI color codes
 class Colors:
     PURPLE = '\033[95m'
@@ -342,9 +366,9 @@ def encode_dvr(input_file, output_dir, gpu):
                       f"{Colors.RED}Unexpected error encoding {input_path}: {e}{Colors.ENDC}")
 
 
-def encode_dvr_distributed(input_file, output_dir, hosts, segment_seconds=60, remote_args=None, concat_args="-c:a copy", probe_host=None, probe_path=None, remote_ffmpeg_path=None):
-    """Encode one file using ffmpeg_distributed (split -> farm -> concat). Segment temp dirs go under script dir/tmp/.
-    If probe_host and probe_path are set, ffprobe runs there (faster when input is on NAS)."""
+def encode_dvr_distributed(input_file, output_dir, workers, segment_seconds=60, remote_args=None, concat_args="-c:a copy", probe_host=None, probe_path=None, remote_ffmpeg_path=None):
+    """Encode one file using ffmpeg_distributed (split -> farm -> concat). workers = [(host, gpu_id), ...].
+    Segment temp dirs go under script dir/tmp/. If probe_host and probe_path are set, ffprobe runs there (faster when input is on NAS)."""
     input_path = Path(input_file).resolve()
     output_path = (Path(output_dir) / f"{input_path.stem}{input_path.suffix}").resolve()
     if output_path.exists():
@@ -366,11 +390,11 @@ def encode_dvr_distributed(input_file, output_dir, hosts, segment_seconds=60, re
     try:
         os.chdir(output_dir)
         from ffmpeg_distributed import encode as distributed_encode
-        safe_log_info(f"Distributed encode: {input_path} -> {output_path} (hosts: {hosts})")
-        print(f"{Colors.BLUE}Distributed encode (AV1): {input_path.name}{Colors.ENDC}")
+        safe_log_info(f"Distributed encode: {input_path} -> {output_path} (workers: {workers})")
+        print(f"{Colors.BLUE}Distributed encode (HEVC): {input_path.name}{Colors.ENDC}")
         remote_ffmpeg = remote_ffmpeg_path or os.environ.get("DISTRIBUTED_REMOTE_FFMPEG_PATH")
         distributed_encode(
-            hosts,
+            workers,
             str(input_path),
             str(output_path),
             segment_seconds=segment_seconds,
@@ -399,13 +423,14 @@ if __name__ == "__main__":
     output_dir = "output"
     os.makedirs(output_dir, exist_ok=True)
 
-    hosts_str = os.environ.get("DISTRIBUTED_HOSTS")
-    if hosts_str:
-        hosts = [h.strip() for h in hosts_str.split(",") if h.strip()]
+    workers_str = os.environ.get("DISTRIBUTED_WORKERS")
+    if workers_str:
+        workers = _parse_workers_env(workers_str)
     else:
-        hosts = DISTRIBUTED_HOSTS_DEFAULT
-        print(f"{Colors.BLUE}Using hosts: {', '.join(hosts)}{Colors.ENDC}")
-    safe_log_info(f"Distributed mode; hosts: {hosts}")
+        workers = DISTRIBUTED_WORKERS_DEFAULT
+    workers_desc = ", ".join(f"{h}:gpu{g}" for h, g in workers)
+    print(f"{Colors.BLUE}Using workers: {workers_desc}{Colors.ENDC}")
+    safe_log_info(f"Distributed mode; workers: {workers}")
 
     files = [f for f in os.listdir(input_dir) if f.endswith(('.mp4', '.DVR.mp4'))]
     total_files = len(files)
@@ -419,4 +444,4 @@ if __name__ == "__main__":
             input_file = os.path.join(input_dir, file)
             safe_log_info(f"Processing file {i}/{total_files}: {file}")
             print(f"\n{Colors.BLUE}Processing file {i}/{total_files}: {file}{Colors.ENDC}")
-            encode_dvr_distributed(input_file, output_dir, hosts, segment_seconds=DISTRIBUTED_SEGMENT_SECONDS) 
+            encode_dvr_distributed(input_file, output_dir, workers, segment_seconds=DISTRIBUTED_SEGMENT_SECONDS) 

ffmpeg_distributed.py

@@ -257,21 +257,23 @@ HOST_COLORS = ['\033[94m', '\033[92m', '\033[93m', '\033[95m', '\033[96m', '\033
 RESET = '\033[0m'
 
 class TaskThread(Thread):
-    def __init__(self, host: str, source_file: str, task_queue: SimpleQueue, bar_pos: int, remote_ffmpeg_path: str = None):
+    def __init__(self, host: str, gpu_id: int, source_file: str, task_queue: SimpleQueue, bar_pos: int, remote_ffmpeg_path: str = None):
         super().__init__()
         self._should_stop = False
         self._host = host
+        self._gpu_id = gpu_id
+        self._host_desc = f"{host}:gpu{gpu_id}"
         self._bar_pos = bar_pos
         self._remote_ffmpeg_path = remote_ffmpeg_path
         self._source_file = source_file
         self._task_queue = task_queue
         self._ffmpeg = None
-        self._bar = TqdmAbsolute(desc=host, position=bar_pos)
+        self._bar = TqdmAbsolute(desc=self._host_desc, position=bar_pos)
         self._current_file = None
 
     def _host_tag(self):
         c = HOST_COLORS[self._bar_pos % len(HOST_COLORS)]
-        return f'{c}{self._host}{RESET}'
+        return f'{c}{self._host_desc}{RESET}'
 
     def stop(self):
         self._should_stop = True
@@ -287,7 +289,7 @@ class TaskThread(Thread):
         last_log = [0.0]  # mutable for progress heartbeat
         def upd(frames, fps, t, duration, speed):
             self._bar.total = duration or 999
-            self._bar.desc = self._host + ': ' + (self._current_file or '')
+            self._bar.desc = self._host_desc + ': ' + (self._current_file or '')
             self._bar.update(t)
             if duration and duration > 0 and (time() - last_log[0]) >= 30:
                 tqdm.write(f'  {self._host_tag()}: {self._current_file} {t:.0f}s / {duration:.0f}s ({speed:.1f}x)', file=stderr)
@@ -307,8 +309,9 @@ class TaskThread(Thread):
                 ffmpeg_bin = (self._remote_ffmpeg_path or 'ffmpeg') if self._host != 'localhost' else 'ffmpeg'
                 encoder_cmd = [
                     ffmpeg_bin, '-f', 'mpegts', '-i', 'pipe:',
+                    '-gpu', str(self._gpu_id),
                     *task.ffmpeg_args,
-                    '-f', 'mpegts', 'pipe:1'
+                    '-f', 'mp4', '-movflags', 'frag_keyframe+empty_moov', 'pipe:1'
                 ]
                 if self._host == 'localhost' and sys_platform != 'win32':
                     encoder_cmd = ['nice', '-n10', 'ionice', '-c3'] + encoder_cmd
@@ -351,7 +354,7 @@ class TaskThread(Thread):
             pass
         self._bar.close()
 
-def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds: float = 60, remote_args: str = '', concat_args: str = '', tmp_dir: str = None, keep_tmp=False, resume=False, copy_input=False, probe_host: str = None, probe_path: str = None, remote_ffmpeg_path: str = None):
+def encode(workers: List[Tuple[str, int]], input_file: str, output_file: str, segment_seconds: float = 60, remote_args: str = '', concat_args: str = '', tmp_dir: str = None, keep_tmp=False, resume=False, copy_input=False, probe_host: str = None, probe_path: str = None, remote_ffmpeg_path: str = None):
     input_file = abspath(expanduser(input_file))
     output_file = abspath(expanduser(output_file))
     tmp_dir = tmp_dir or 'ffmpeg_segments_'+md5(input_file.encode()).hexdigest()
@@ -377,7 +380,7 @@ def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds:
     removed = 0
     for i, (start_sec, end_sec) in enumerate(segments):
         duration_sec = end_sec - start_sec
-        output_path = f'{tmp_dir}/{i:08d}.ts'
+        output_path = f'{tmp_dir}/{i:08d}.mp4'
         if isfile(output_path):
             try:
                 if getsize(output_path) < MIN_SEGMENT_BYTES:
@@ -402,9 +405,9 @@ def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds:
     stderr.flush()
     dprint(f'Segments: {len(segments)} total, {n_tasks} tasks queued')
 
-    tqdm.write(f'[3/4] Encoding segments on {len(hosts)} host(s)...', file=stderr)
+    tqdm.write(f'[3/4] Encoding segments on {len(workers)} worker(s)...', file=stderr)
     stderr.flush()
-    threads = [TaskThread(host, input_file, task_queue, pos, remote_ffmpeg_path) for pos, host in enumerate(hosts, 0)]
+    threads = [TaskThread(host, gpu_id, input_file, task_queue, pos, remote_ffmpeg_path) for pos, (host, gpu_id) in enumerate(workers, 0)]
 
     def stop_all():
         tqdm.write('Stopping all workers (killing ffmpeg/SSH on each host)...', file=stderr)
@@ -470,7 +473,7 @@ def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds:
 
     list_path = f'{tmp_dir}/output_segments.txt'
     with open(list_path, 'w') as f:
-        f.write('\n'.join([f"file '{fpath}'" for fpath in sorted(glob(f'{tmp_dir}/*.ts'))]))
+        f.write('\n'.join([f"file '{fpath}'" for fpath in sorted(glob(f'{tmp_dir}/*.mp4'))]))
 
     tqdm.write('[4/4] Concatenating segments and muxing with audio...', file=stderr)
     concat_extra = ['-stats_period', '5'] if verbose else []
@@ -497,15 +500,33 @@ def encode(hosts: List[str], input_file: str, output_file: str, segment_seconds:
     if not keep_tmp:
         rmtree(tmp_dir)
 
+def _parse_workers(host_specs):
+    """Parse list of 'host' or 'host:gpu' into [(host, gpu_id), ...]."""
+    workers = []
+    for spec in host_specs or []:
+        spec = (spec or "").strip()
+        if not spec:
+            continue
+        if ":" in spec:
+            host, gpu = spec.rsplit(":", 1)
+            try:
+                workers.append((host.strip(), int(gpu.strip())))
+            except ValueError:
+                workers.append((spec, 0))
+        else:
+            workers.append((spec, 0))
+    return workers
+
+
 if __name__ == '__main__':
     import argparse
-    parser = argparse.ArgumentParser(description='Splits a file into segments and processes them on multiple hosts in parallel using ffmpeg over SSH.')
+    parser = argparse.ArgumentParser(description='Splits a file into segments and processes them on multiple workers (host:gpu) in parallel using ffmpeg over SSH.')
     parser.add_argument('input_file', help='File to encode.')
     parser.add_argument('output_file', help='Path to encoded output file.')
-    parser.add_argument('remote_args', help='Arguments to pass to the remote ffmpeg instances. For example: "-c:v libx264 -crf 23 -preset fast"')
-    parser.add_argument('concat_args', default='', help='Arguments to pass to the local ffmpeg concatenating the processed video segments and muxing it with the original audio/subs/metadata. Mainly useful for audio encoding options, or "-an" to get rid of it.')
+    parser.add_argument('remote_args', help='Arguments to pass to the remote ffmpeg instances (e.g. NVENC: -c:v hevc_nvenc -preset p7 ...). -gpu is added per worker.')
+    parser.add_argument('concat_args', default='', help='Arguments to pass to the local ffmpeg concatenating the processed video segments and muxing it with the original audio/subs/metadata.')
     parser.add_argument('-s', '--segment-length', type=float, default=10, help='Segment length in seconds.')
-    parser.add_argument('-H', '--host', action='append', help='SSH hostname(s) to encode on. Use "localhost" to include the machine you\'re running this from. Can include username.', required=True)
+    parser.add_argument('-H', '--host', action='append', help='Worker as host or host:gpu (e.g. -H Pyro:0 -H RenderScrap:0 -H RenderScrap:1).', required=True)
     parser.add_argument('-k', '--keep-tmp', action='store_true', help='Keep temporary segment files instead of deleting them on successful exit.')
     parser.add_argument('-r', '--resume', action='store_true', help='Don\'t split the input file again, keep existing segments and only process the missing ones.')
     parser.add_argument('-t', '--tmp-dir', default=None, help='Directory to use for temporary files. Should not already exist and will be deleted afterwards.')
@@ -513,8 +534,11 @@ if __name__ == '__main__':
     parser.add_argument('-P', '--probe-host', default=None, help='SSH host to run ffprobe on (file must be at --probe-path there). Speeds up [1/4] when input is on slow/NAS path.')
     parser.add_argument('--probe-path', default=None, help='Path to input file as seen on --probe-host (required if -P set).')
     args = parser.parse_args()
+    workers = _parse_workers(args.host)
+    if not workers:
+        parser.error('At least one worker required (e.g. -H Pyro:0 -H RenderScrap:0)')
     encode(
-        args.host,
+        workers,
         args.input_file,
         args.output_file,
         segment_seconds=args.segment_length,