Source code for remote_compression.host

"""Execution hosts: this machine, or a remote server behind an SSH session."""

import logging
import re
import shutil
import stat as stat_module
import subprocess
import threading
import time
from collections import deque
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
from uuid import uuid4

import paramiko

from remote_compression.command import FORBIDDEN_CHARS
from remote_compression.probe import probe
from remote_compression.ssh import (
    RemoteExecError,
    SSHConnectionError,
    SSHSession,
    is_localhost,
)

logger = logging.getLogger("rcomp")

#: Remote working directory, relative to the SFTP home (never absolutized:
#: OpenSSH on Windows reports homes as ``/C:/...``, invalid in a cmd argv).
WORKDIR = ".rcomp"
#: Seconds granted to get the connection back mid-batch.
RECONNECT_WINDOW = 300.0
#: A remote output whose size stalls this long is finished (or dead).
SETTLE_TIME = 60.0
#: Orphan polling interval.
POLL = 10.0
#: Seconds granted to a remote output to appear after a reconnection.
GRACE = 30.0

#: Encoder names safe to embed in a command line; anything else skips the
#: bootstrap encoder check (and will fail loudly at encode time instead).
_ENCODER_NAME = re.compile(r"[A-Za-z0-9_.-]+")

#: ffmpeg's message for an unknown encoder (exit codes vary across versions).
_UNKNOWN_ENCODER = "is not recognized"


[docs] class BootstrapError(RuntimeError): """A prerequisite is missing (ffmpeg, workspace...); the batch cannot start."""
class _ConnectionLost(Exception): """Internal: the connection died during a phase of a remote compression.""" def __init__(self, phase, cause): super().__init__(f"connection lost during {phase}: {cause}") self.phase = phase self.cause = cause
[docs] @dataclass(frozen=True, slots=True) class HostResult: """ Outcome of one compression attempt on a host. Attributes ---------- ok: :class:`bool` True iff the target file was produced and retrieved. exit_status: :class:`int`, optional ffmpeg exit status when known. stderr_tail: :class:`str` Bounded tail of ffmpeg stderr (diagnosis). reason: :class:`str` Failure category: ``ffmpeg-failed``, ``upload-failed``, ``download-failed`` or ``exec-timeout``. Empty on success. """ ok: bool exit_status: int | None = None stderr_tail: str = "" reason: str = ""
[docs] @dataclass(frozen=True, slots=True) class WorkspaceEntry: """A regular file in the remote workspace.""" name: str size: int mtime: float
[docs] def age(self, now=None): """:class:`datetime.timedelta`: age of the file (relative to local `now`).""" now = time.time() if now is None else now return timedelta(seconds=max(0.0, now - self.mtime))
[docs] @dataclass(slots=True) class WorkspaceStatus: """Content of the remote workspace.""" entries: list @property def count(self): """:class:`int`: number of files.""" return len(self.entries) @property def total_bytes(self): """:class:`int`: cumulated size.""" return sum(e.size for e in self.entries)
[docs] @dataclass(slots=True) class PurgeReport: """Outcome of a workspace purge.""" removed: list failed: list dry_run: bool @property def freed_bytes(self): """:class:`int`: cumulated size of the removed files.""" return sum(e.size for e in self.removed)
def _safe_suffix(path, fallback): """Lower-cased suffix of `path` if safe for a remote command line, else `fallback`.""" suffix = Path(path).suffix.lower() if suffix.isascii() and not (FORBIDDEN_CHARS & set(suffix)) and " " not in suffix: return suffix return fallback
[docs] def open_host(settings, config_path=None): """ Build the right host for the given settings. Parameters ---------- settings: :class:`~remote_compression.settings.Settings` Resolved settings (`hostname`, `auto_local`, `auto_purge_days`). config_path: :class:`str`, optional Alternative openSSH config file (mostly for tests). Returns ------- :class:`LocalHost` or :class:`RemoteHost` Not yet connected: use it as a context manager, whose ``__enter__`` performs connection and fail-fast bootstrap checks. """ hostname = settings.hostname if hostname.lower() == "local": return LocalHost(encoder=settings.codec) if settings.auto_local and is_localhost(hostname, config_path): logger.info("Host '%s' resolves to this machine: switching to local mode.", hostname) return LocalHost(encoder=settings.codec) return RemoteHost( hostname, config_path=config_path, auto_purge_days=settings.auto_purge_days, encoder=settings.codec, remote_ffmpeg=settings.remote_ffmpeg, workdir=settings.remote_workdir, workdir_sftp=settings.remote_workdir_sftp, )
[docs] class LocalHost: """ Compression on this very machine. Parameters ---------- encoder: :class:`str`, optional Encoder whose availability is verified at bootstrap. """ hostname = "local" is_remote = False def __init__(self, encoder=None): self.encoder = encoder def __enter__(self): self.bootstrap() return self def __exit__(self, exc_type, exc_val, exc_tb): pass
[docs] def bootstrap(self): """ Raises ------ BootstrapError If ffmpeg is not in the local PATH, or lacks the requested encoder. """ if shutil.which("ffmpeg") is None: msg = "ffmpeg not found on this machine. Install it or add it to the PATH." raise BootstrapError(msg) if self.encoder and _ENCODER_NAME.fullmatch(self.encoder): result = subprocess.run( ["ffmpeg", "-hide_banner", "-h", f"encoder={self.encoder}"], capture_output=True, check=False, ) output = (result.stdout + result.stderr).decode("utf-8", errors="replace") if result.returncode != 0 or _UNKNOWN_ENCODER in output: msg = ( f"Encoder '{self.encoder}' is not available in the local ffmpeg build. " "Pick another codec (-C) or install a fuller ffmpeg." ) raise BootstrapError(msg)
[docs] def compress_file(self, source, target, plan, on_progress=None): """ Run ffmpeg locally. Parameters ---------- source: :class:`~pathlib.Path` Input video. target: :class:`~pathlib.Path` Output video (the temporary ``comp_*`` file). plan: :class:`~remote_compression.command.FfmpegPlan` Command to run. on_progress: callable, optional Receives decoded chunks of the ffmpeg ``-progress`` stream. Returns ------- :class:`HostResult` """ argv = plan.argv(source, target) logger.debug("Running locally: %s", subprocess.list2cmdline(argv)) proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE) stderr_chunks = deque(maxlen=64) def drain_stderr(): for chunk in iter(lambda: proc.stderr.read(32768), b""): stderr_chunks.append(chunk) thread = threading.Thread(target=drain_stderr, daemon=True) thread.start() try: for raw_line in proc.stdout: if on_progress is not None: on_progress(raw_line.decode("utf-8", errors="replace")) finally: exit_status = proc.wait() thread.join(timeout=5) stderr_tail = b"".join(stderr_chunks)[-4096:].decode("utf-8", errors="replace") if exit_status != 0: return HostResult(ok=False, exit_status=exit_status, stderr_tail=stderr_tail, reason="ffmpeg-failed") return HostResult(ok=True, exit_status=0, stderr_tail=stderr_tail)
[docs] class RemoteHost: """ Compression on a remote server, over a single SSH session. All file operations go through SFTP (platform-agnostic: Linux and Windows servers behave the same); only ``ffmpeg -version`` and the ffmpeg command itself are executed remotely, with a syntax common to sh and cmd.exe. Parameters ---------- hostname: :class:`str` SSH alias of the server. config_path: :class:`str`, optional Alternative openSSH config file. auto_purge_days: :class:`float`, optional Age threshold of the opportunistic purge run at connection time (None disables it). encoder: :class:`str`, optional Encoder whose availability is verified at bootstrap. remote_ffmpeg: :class:`str`, optional Explicit path of the ffmpeg binary on the server (space-free ASCII; for hosts whose non-interactive PATH misses ffmpeg). workdir: :class:`str`, optional Workspace directory as seen by the shell running ffmpeg (default: ``.rcomp``, relative to the login home). workdir_sftp: :class:`str`, optional The same directory as seen by the SFTP channel, when the two views differ (Synology chroots SFTP: its ``/home`` is the login home, so ``/home/.rcomp`` pairs with the default `workdir`). Defaults to `workdir`. session: :class:`~remote_compression.ssh.SSHSession`, optional Injected session (tests). """ is_remote = True def __init__(self, hostname, config_path=None, auto_purge_days=7.0, encoder=None, remote_ffmpeg=None, workdir=None, workdir_sftp=None, session=None): self.hostname = hostname self.session = session if session is not None else SSHSession(hostname, config_path=config_path) self.auto_purge_days = auto_purge_days self.encoder = encoder self.remote_ffmpeg = remote_ffmpeg self.workdir = workdir if workdir is not None else WORKDIR self.workdir_sftp = workdir_sftp if workdir_sftp is not None else self.workdir # -- lifecycle ----------------------------------------------------------- def __enter__(self): self.session.open() self.bootstrap() return self def __exit__(self, exc_type, exc_val, exc_tb): self.session.close()
[docs] def bootstrap(self): """Fail-fast checks: workspace, ffmpeg, then best-effort auto-purge.""" self.ensure_workspace() self.check_ffmpeg() if self.auto_purge_days is not None: try: report = self.purge_workspace(older_than=timedelta(days=self.auto_purge_days)) if report.removed: logger.info( "Auto-purged %d stale file(s), %d bytes, from %s:%s.", len(report.removed), report.freed_bytes, self.hostname, self.workdir_sftp, ) except Exception as e: # noqa: BLE001 - opportunistic purge is best-effort by design logger.warning("Opportunistic purge of %s:%s failed: %s", self.hostname, self.workdir_sftp, e)
[docs] def ensure_workspace(self): """ Create the remote workspace when missing (pure SFTP). Raises ------ BootstrapError If it can neither be accessed nor created. """ sftp = self.session.sftp try: sftp.stat(self.workdir_sftp) except FileNotFoundError: try: sftp.mkdir(self.workdir_sftp) logger.info("Created workspace %s on %s.", self.workdir_sftp, self.hostname) except OSError as e: msg = ( f"Cannot create the workspace {self.workdir_sftp} on {self.hostname}: {e}. " "The parent folder must exist and be writable over SFTP. On Synology, " "SFTP is chrooted to the shared folders and its /home is the login home: " 'set remote_workdir_sftp = "/home/.rcomp" in the configuration.' ) raise BootstrapError(msg) from e except OSError as e: msg = f"Cannot access the workspace {self.workdir_sftp} on {self.hostname}: {e}" raise BootstrapError(msg) from e
[docs] def check_ffmpeg(self): """ Raises ------ BootstrapError If ffmpeg fails on the server, or lacks the requested encoder. """ binary = self.remote_ffmpeg or "ffmpeg" result = self.session.run(f"{binary} -version", inactivity_timeout=30.0) if not result.ok: detail = result.stderr_tail.strip() or result.stdout_tail.strip() or "no output" msg = ( f"ffmpeg not found on {self.hostname} (tried '{binary}', exit {result.exit_status}): " f"{detail}. Install ffmpeg on the server, add it to its PATH, or set " "remote_ffmpeg in the configuration." ) raise BootstrapError(msg) if self.encoder and _ENCODER_NAME.fullmatch(self.encoder): result = self.session.run( f"{binary} -hide_banner -h encoder={self.encoder}", inactivity_timeout=30.0 ) if not result.ok or _UNKNOWN_ENCODER in result.stdout_tail + result.stderr_tail: msg = ( f"Encoder '{self.encoder}' is not available in the ffmpeg build on " f"{self.hostname}. Pick another codec (-C) or install a fuller ffmpeg there." ) raise BootstrapError(msg)
# -- workspace inspection --------------------------------------------------
[docs] def workspace_status(self): """ Returns ------- :class:`WorkspaceStatus` or None Regular files of the workspace (sub-directories and symlinks are ignored), sorted by mtime; None when the workspace does not exist. """ try: attrs = self.session.sftp.listdir_attr(self.workdir_sftp) except FileNotFoundError: return None entries = [ WorkspaceEntry(name=a.filename, size=a.st_size or 0, mtime=a.st_mtime or 0.0) for a in attrs if stat_module.S_ISREG(a.st_mode or 0) ] entries.sort(key=lambda e: e.mtime) return WorkspaceStatus(entries=entries)
[docs] def purge_workspace(self, older_than=timedelta(0), dry_run=False): """ Remove workspace files older than a threshold. Parameters ---------- older_than: :class:`datetime.timedelta` Minimal age (local clock vs remote mtimes; skew is negligible at the thresholds in use). Zero purges everything. dry_run: :class:`bool` List without removing. Returns ------- :class:`PurgeReport` Individual removal failures are collected, never raised. """ status = self.workspace_status() if status is None: return PurgeReport(removed=[], failed=[], dry_run=dry_run) now = time.time() threshold = older_than.total_seconds() removed, failed = [], [] for entry in status.entries: if now - entry.mtime <= threshold: continue if dry_run: removed.append(entry) continue try: self.session.sftp.remove(f"{self.workdir_sftp}/{entry.name}") removed.append(entry) except OSError as e: failed.append((entry.name, str(e))) return PurgeReport(removed=removed, failed=failed, dry_run=dry_run)
# -- compression -----------------------------------------------------------
[docs] def compress_file(self, source, target, plan, on_progress=None): """ Upload, encode remotely, download. On a connection loss, reconnects patiently then tries to *re-attach* to the orphan encode (ffmpeg may well survive the disconnection on the server) instead of blindly relaunching. At most one relaunch, always with fresh temporary names — a zombie may still be writing to the old target. Parameters ---------- source: :class:`~pathlib.Path` Local input video. target: :class:`~pathlib.Path` Local output (the temporary ``comp_*`` file). plan: :class:`~remote_compression.command.FfmpegPlan` Command to run. on_progress: callable, optional Receives decoded chunks of the ffmpeg ``-progress`` stream. Returns ------- :class:`HostResult` Raises ------ SSHConnectionError When the connection cannot be recovered, or is lost twice on the same file: the batch cannot reasonably continue. """ for attempt in range(2): # same physical files, two path views: SFTP operations vs the shell # command (they differ on chrooted-SFTP servers like Synology) src_name = f"{uuid4().hex}{_safe_suffix(source, '.src')}" tgt_name = f"{uuid4().hex}{_safe_suffix(target, '.mkv')}" r_source = f"{self.workdir_sftp}/{src_name}" r_target = f"{self.workdir_sftp}/{tgt_name}" exec_paths = (f"{self.workdir}/{src_name}", f"{self.workdir}/{tgt_name}") try: return self._attempt(source, target, plan, r_source, r_target, exec_paths, on_progress) except _ConnectionLost as lost: logger.warning( "Connection to %s lost while processing %s (%s phase).", self.hostname, Path(source).name, lost.phase, ) self._patient_reconnect() if lost.phase == "run" and self._reattach(r_target, target): return HostResult(ok=True, exit_status=0) self._cleanup(r_source, r_target) if attempt == 1: msg = ( f"Lost the connection to {self.hostname} twice on " f"{Path(source).name}; giving up on the batch." ) raise SSHConnectionError(msg) from lost.cause return None # pragma: no cover - unreachable
def _attempt(self, source, target, plan, r_source, r_target, exec_paths, on_progress): """One full upload / encode / download cycle.""" session = self.session try: session.ensure_connected() except SSHConnectionError as e: raise _ConnectionLost("connect", e) from e try: session.sftp.put(str(source), r_source) except Exception as e: # noqa: BLE001 - sorted into connection-loss vs file failure self._raise_if_lost(e, "put") return HostResult(ok=False, stderr_tail=str(e), reason="upload-failed") try: result = session.run( plan.command(*exec_paths, binary=self.remote_ffmpeg), on_stdout=on_progress ) except RemoteExecError as e: self._cleanup(r_source, r_target) return HostResult(ok=False, stderr_tail=str(e), reason="exec-timeout") except SSHConnectionError as e: raise _ConnectionLost("run", e) from e try: if result.ok and self._remote_size(r_target) > 0: session.sftp.get(r_target, str(target)) return HostResult(ok=True, exit_status=result.exit_status, stderr_tail=result.stderr_tail) detail = result.stderr_tail or result.stdout_tail return HostResult(ok=False, exit_status=result.exit_status, stderr_tail=detail, reason="ffmpeg-failed") except Exception as e: # noqa: BLE001 - sorted into connection-loss vs file failure self._raise_if_lost(e, "get") return HostResult(ok=False, exit_status=result.exit_status, stderr_tail=str(e), reason="download-failed") finally: self._cleanup(r_source, r_target) def _raise_if_lost(self, exc, phase): """Translate an exception into :class:`_ConnectionLost` when the transport died.""" if isinstance(exc, _ConnectionLost): raise exc lost_types = (paramiko.ssh_exception.SSHException, SSHConnectionError, EOFError, ConnectionError) if isinstance(exc, lost_types) or not self.session.is_alive(): raise _ConnectionLost(phase, exc) from exc def _remote_size(self, r_path): try: return self.session.sftp.stat(r_path).st_size or 0 except FileNotFoundError: return 0 def _cleanup(self, *r_paths): """Best-effort removal of remote temporaries (auto-purge catches leftovers).""" for r_path in r_paths: try: self.session.sftp.remove(r_path) except Exception: # noqa: BLE001, S110 - cleanup is best-effort by design pass def _patient_reconnect(self): """ Reconnect with a long patience window (a wifi/VPN blip mid-batch should not kill hours of work). Raises ------ SSHConnectionError When :data:`RECONNECT_WINDOW` is exhausted. """ deadline = time.monotonic() + RECONNECT_WINDOW delay = 2.0 while True: try: self.session.reconnect() return except SSHConnectionError as e: if time.monotonic() + delay > deadline: msg = f"Could not reconnect to {self.hostname} within {RECONNECT_WINDOW:.0f} s." raise SSHConnectionError(msg) from e logger.warning("Reconnection to %s failed, retrying in %.0f s.", self.hostname, delay) time.sleep(delay) delay = min(delay * 2, 30.0) def _reattach(self, r_target, target): """ Try to recover the output of an orphan encode after a reconnection. Polls the remote target: a growing size means the orphan is still encoding (wait); a size stable for :data:`SETTLE_TIME` means it finished or died — since its exit status is unknowable, the file is downloaded and validated with a *local* ffprobe. Returns ------- :class:`bool` True when a valid output was recovered into `target`. Raises ------ SSHConnectionError On a second connection loss (no third chance). """ logger.info("Checking whether the remote encode of %s survived...", Path(target).name) last_size = -1 last_change = time.monotonic() absent_deadline = time.monotonic() + GRACE while True: try: size = self._remote_size(r_target) except Exception as e: # noqa: BLE001 - sorted into connection-loss vs give-up self._raise_if_lost_for_good(e) return False now = time.monotonic() if size <= 0: if now > absent_deadline: logger.info("No remote output: the orphan encode died early. Relaunching.") return False elif size != last_size: last_size = size last_change = now logger.debug("Orphan encode still writing (%d bytes).", size) elif now - last_change > SETTLE_TIME: break time.sleep(POLL) try: self.session.sftp.get(r_target, str(target)) except Exception as e: # noqa: BLE001 - sorted into connection-loss vs give-up self._raise_if_lost_for_good(e) return False finally: self._cleanup(r_target) if probe(target) is None: Path(target).unlink(missing_ok=True) logger.info("Recovered remote output is invalid; relaunching.") return False logger.info("Recovered the orphan encode output (%d bytes).", last_size) return True def _raise_if_lost_for_good(self, exc): """During re-attachment, a second connection loss aborts the batch.""" lost_types = (paramiko.ssh_exception.SSHException, SSHConnectionError, EOFError, ConnectionError) if isinstance(exc, lost_types) or not self.session.is_alive(): msg = f"Lost the connection to {self.hostname} again while re-attaching; giving up on the batch." raise SSHConnectionError(msg) from exc