livekit-cameras/run.py

"""Orchestrator: discover cameras and run one publisher worker each.

Usage:
    python run.py                 # use .env (or environment variables)
    python run.py --cameras 1,5   # publish only cameras 1 and 5
    python run.py --once          # run, print a status report, exit (for CI)
"""

from __future__ import annotations

import argparse
import asyncio
import contextlib
import json
import logging
import os
import signal
import subprocess
import sys
import time
from pathlib import Path

from config import load_config, _load_dotenv
from discovery import discover_cameras, select_cameras
from tokens import make_token

log = logging.getLogger("cameras.main")


def worker_args_for(cfg, cam) -> list[str]:
    return [
        "--camera-json", json.dumps(cam.__dict__),
        "--room", cfg.room,
        "--url", cfg.url,
        "--token", make_token(cfg.url, cfg.api_key, cfg.api_secret,
                              cfg.room, cam.identity),
        "--width", str(cfg.width),
        "--height", str(cfg.height),
        "--fps", str(cfg.fps),
        "--bitrate", str(cfg.video_bitrate),
        "--video-codec", cfg.video_codec,
        "--video-encoder", cfg.video_encoder,
        "--vaapi-device", cfg.vaapi_device,
        "--audio-rate", str(cfg.audio_rate),
        "--enhance-mode", cfg.enhance_mode,
        "--enhance-model", cfg.enhance_model,
        "--enhance-chunk-s", str(cfg.enhance_chunk_s),
        "--enhance-hop-s", str(cfg.enhance_hop_s),
        "--beamform-mode", cfg.beamform_mode,
        "--torch-num-threads", str(cfg.torch_num_threads),
        "--tdoa-every", str(cfg.beamform_tdoa_every),
    ]


class Supervisor:
    def __init__(self, cfg, cameras):
        self.cfg = cfg
        self.cameras = cameras
        self.procs: dict[str, subprocess.Popen] = {}
        self.last_exit: dict[str, int | None] = {}
        self.restarts: dict[str, int] = {}

    def spawn(self, cam) -> None:
        args = worker_args_for(self.cfg, cam)
        env = dict(os.environ)
        env.setdefault("LIBVA_DRIVER_NAME", "iHD")
        env.setdefault("LIBVA_DRM_DEVICE", self.cfg.vaapi_device)
        p = subprocess.Popen(
            [sys.executable, str(Path(__file__).parent / "run_worker.py"),
             *args],
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            text=True, bufsize=1, env=env,
        )
        self.procs[cam.identity] = p
        self.restarts[cam.identity] = self.restarts.get(cam.identity, 0) + 1

    def pump(self) -> None:
        """Read one line from each running worker; print them."""
        for ident, p in list(self.procs.items()):
            line = p.stdout.readline() if p.stdout else ""
            if line:
                sys.stdout.write(f"[{ident}] {line}")
                sys.stdout.flush()
            if p.poll() is not None and not line:
                pass  # EOF handled by poll below

    def reap(self) -> None:
        for ident, p in list(self.procs.items()):
            rc = p.poll()
            if rc is None:
                continue
            self.last_exit[ident] = rc
            if rc != 0:
                log.warning("%s exited rc=%s (restart %d/%d)", ident, rc,
                            self.restarts.get(ident, 0), 10)
                if self.restarts.get(ident, 0) < 10:
                    time.sleep(1.0)
                    # find the camera object again
                    cam = next(c for c in self.cameras
                               if c.identity == ident)
                    self.spawn(cam)
                else:
                    log.error("%s gave up after 10 restarts", ident)
            else:
                log.info("%s exited cleanly", ident)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--cameras", default=None,
                    help="subset spec, e.g. 'all', '1,5', '0-19'")
    ap.add_argument("--once", action="store_true",
                    help="spawn workers, report status after "
                         "PUBLISH_TIMEOUT_S, then exit")
    args = ap.parse_args()

    cfg = load_config()
    logging.basicConfig(
        level=getattr(logging, cfg.log_level, logging.INFO),
        format="%(asctime)s [main] %(levelname)s %(message)s")

    if args.cameras:
        from config import _parse_cameras
        cfg.cameras = _parse_cameras(args.cameras)

    if not cfg.url or not cfg.api_key or not cfg.api_secret:
        log.error("LIVEKIT_URL / LIVEKIT_API_KEY / LIVEKIT_API_SECRET "
                  "not set (create .env from .env.example)")
        return 2

    cameras = discover_cameras()
    log.info("discovered %d cameras:", len(cameras))
    for c in cameras:
        log.info("  %s  video=%s  audio=hw:%s  usb=%s",
                 c.identity, c.video_device,
                 c.audio_card if c.audio_card is not None else "-",
                 c.usb_port or "-")

    cameras = select_cameras(cameras, cfg.cameras)
    if not cameras:
        log.error("no cameras match --cameras %s", args.cameras)
        return 2

    sup = Supervisor(cfg, cameras)
    for cam in cameras:
        sup.spawn(cam)
    log.info("spawned %d workers", len(cameras))

    if args.once:
        deadline = time.monotonic() + cfg.publish_timeout_s
        while time.monotonic() < deadline:
            sup.pump()
            sup.reap()
            time.sleep(0.5)
        # final report
        up = [i for i, p in sup.procs.items() if p.poll() is None]
        log.info("STATUS: %d/%d workers alive", len(up), len(sup.procs))
        for ident in up:
            pass
        # stop everything
        for p in sup.procs.values():
            p.send_signal(signal.SIGTERM)
        for p in sup.procs.values():
            with contextlib.suppress(Exception):
                p.wait(timeout=10)
        return 0 if len(up) == len(sup.procs) else 1

    try:
        while True:
            sup.pump()
            sup.reap()
            time.sleep(0.2)
    except KeyboardInterrupt:
        log.info("shutting down ...")
        for p in sup.procs.values():
            if p.poll() is None:
                p.send_signal(signal.SIGTERM)
        for p in sup.procs.values():
            with contextlib.suppress(Exception):
                p.wait(timeout=10)
        for p in sup.procs.values():
            if p.poll() is None:
                p.kill()
    return 0


if __name__ == "__main__":
    sys.exit(main())