LLMVideoPipeline/pipeline/remote.py

"""
Remote PC manager with SSH, SCP, and HTTP API helpers.
"""

import subprocess
import requests
import logging
from typing import Dict, Any, Optional, Tuple

logger = logging.getLogger("YTFactory.Remote")

class RemotePC:
    def __init__(self, name: str, host: str, ssh_user: str, services: Optional[Dict[str, Any]] = None):
        self.name = name
        self.host = host
        self.ssh_user = ssh_user
        self.services = services or {}
        self.ssh_base = [
            "ssh",
            "-o", "StrictHostKeyChecking=no",
            "-o", "ConnectTimeout=5",
            "-o", "BatchMode=yes",
            f"{ssh_user}@{host}"
        ]

    def check(self, port: int) -> bool:
        """Check if service is responding on port."""
        for path in ("/health", "/system_stats", "/v1/models", "/"):
            try:
                r = requests.get(f"http://{self.host}:{port}{path}", timeout=3)
                if r.status_code in (200, 204, 404):  # 404 still means server is alive
                    return True
            except Exception:
                pass
        return False

    def post_json(self, port: int, endpoint: str, payload: Dict[str, Any], timeout: int = 180) -> Optional[Dict[str, Any]]:
        """Send JSON POST request to service."""
        url = f"http://{self.host}:{port}{endpoint}"
        try:
            r = requests.post(url, json=payload, timeout=timeout)
            r.raise_for_status()
            return r.json()
        except Exception as e:
            logger.debug(f"POST {self.name}:{port}{endpoint} failed: {e}")
            return None

    def post_binary(self, port: int, endpoint: str, payload: Dict[str, Any], timeout: int = 180) -> Optional[bytes]:
        """Send JSON POST request expecting binary audio/image output."""
        url = f"http://{self.host}:{port}{endpoint}"
        try:
            r = requests.post(url, json=payload, timeout=timeout)
            r.raise_for_status()
            return r.content
        except Exception as e:
            logger.debug(f"POST binary {self.name}:{port}{endpoint} failed: {e}")
            return None

    def get_json(self, port: int, endpoint: str, timeout: int = 30) -> Optional[Dict[str, Any]]:
        """Send GET request expecting JSON."""
        url = f"http://{self.host}:{port}{endpoint}"
        try:
            r = requests.get(url, timeout=timeout)
            r.raise_for_status()
            return r.json()
        except Exception as e:
            logger.debug(f"GET {self.name}:{port}{endpoint} failed: {e}")
            return None

    def run(self, cmd: str, timeout: int = 600) -> Tuple[int, str, str]:
        """Execute command over SSH."""
        try:
            r = subprocess.run(
                self.ssh_base + [cmd],
                capture_output=True,
                text=True,
                timeout=timeout
            )
            return r.returncode, r.stdout, r.stderr
        except subprocess.TimeoutExpired:
            logger.error(f"SSH command timed out on {self.name}: {cmd[:80]}")
            return -1, "", "TimeoutExpired"
        except Exception as e:
            logger.error(f"SSH command failed on {self.name}: {e}")
            return -1, "", str(e)

    def scp_to(self, local_path: str, remote_path: str) -> bool:
        """Upload file via SCP."""
        cmd = [
            "scp",
            "-o", "StrictHostKeyChecking=no",
            "-o", "ConnectTimeout=10",
            local_path,
            f"{self.ssh_user}@{self.host}:{remote_path}"
        ]
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            logger.error(f"SCP to {self.name}:{remote_path} failed: {r.stderr}")
            return False
        return True

    def scp_from(self, remote_path: str, local_path: str) -> bool:
        """Download file via SCP."""
        cmd = [
            "scp",
            "-o", "StrictHostKeyChecking=no",
            "-o", "ConnectTimeout=10",
            f"{self.ssh_user}@{self.host}:{remote_path}",
            local_path
        ]
        r = subprocess.run(cmd, capture_output=True, text=True)
        if r.returncode != 0:
            logger.error(f"SCP from {self.name}:{remote_path} failed: {r.stderr}")
            return False
        return True