studio / renderer /ffmpeg /runner.py
Ava2lon's picture
Upload 170 files
345855e verified
Raw
History Blame Contribute Delete
3.08 kB
from __future__ import annotations
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable
try:
import psutil
except Exception: # pragma: no cover - optional runtime dependency fallback
psutil = None
class FFmpegError(RuntimeError):
def __init__(self, message: str, command: list[str], stderr: str = "") -> None:
super().__init__(message)
self.command = command
self.stderr = stderr
@dataclass
class CommandResult:
command: list[str]
returncode: int
duration_seconds: float
stdout: str = ""
stderr: str = ""
metrics: dict[str, float | int] = field(default_factory=dict)
class FFmpegRunner:
def __init__(
self,
timeout_seconds: int = 900,
log: Callable[[str], None] | None = None,
on_command: Callable[[list[str]], None] | None = None,
) -> None:
self.timeout_seconds = timeout_seconds
self.log = log or (lambda _: None)
self.on_command = on_command or (lambda _: None)
def run(self, command: list[str], cwd: Path | None = None) -> CommandResult:
started = time.time()
self.on_command(command)
self.log("$ " + " ".join(command))
process = subprocess.Popen(
command,
cwd=str(cwd) if cwd else None,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
encoding="utf-8",
errors="replace",
)
peak_rss = 0
cpu_percent = 0.0
proc = psutil.Process(process.pid) if psutil else None
try:
stdout, stderr = process.communicate(timeout=self.timeout_seconds)
if proc:
try:
peak_rss = max(peak_rss, proc.memory_info().rss)
cpu_percent = proc.cpu_percent(interval=None)
except Exception:
pass
except subprocess.TimeoutExpired as exc:
self._kill_process(process)
stdout, stderr = process.communicate()
raise FFmpegError(f"FFmpeg timed out after {self.timeout_seconds}s", command, stderr) from exc
duration = time.time() - started
result = CommandResult(
command=command,
returncode=process.returncode,
duration_seconds=duration,
stdout=stdout,
stderr=stderr,
metrics={"duration_seconds": duration, "peak_rss_bytes": peak_rss, "cpu_percent": cpu_percent},
)
if process.returncode != 0:
raise FFmpegError("FFmpeg failed", command, stderr)
return result
@staticmethod
def _kill_process(process: subprocess.Popen[str]) -> None:
if psutil:
try:
parent = psutil.Process(process.pid)
for child in parent.children(recursive=True):
child.kill()
parent.kill()
return
except Exception:
pass
process.kill()