115 lines
4.2 KiB
Python
115 lines
4.2 KiB
Python
"""
|
|
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
|