52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
import asyncio
|
|
import logging
|
|
import shlex
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
|
|
log = logging.getLogger("llamacpp_restart")
|
|
|
|
|
|
@dataclass
|
|
class RestartPlan:
|
|
method: str
|
|
command: Optional[str]
|
|
url: Optional[str]
|
|
allowed_container: Optional[str] = None
|
|
|
|
|
|
async def trigger_restart(plan: RestartPlan, payload: Optional[dict] = None) -> None:
|
|
if plan.method == "none":
|
|
log.warning("Restart requested but restart method is none")
|
|
return
|
|
if plan.method == "http":
|
|
if not plan.url:
|
|
raise RuntimeError("restart url is required for http method")
|
|
async with httpx.AsyncClient(timeout=60) as client:
|
|
resp = await client.post(plan.url, json=payload or {})
|
|
resp.raise_for_status()
|
|
return
|
|
if plan.method == "docker":
|
|
if not plan.command:
|
|
raise RuntimeError("restart command must include container id or name for docker method")
|
|
if plan.allowed_container and plan.command != plan.allowed_container:
|
|
raise RuntimeError("docker restart command not allowed for non-target container")
|
|
async with httpx.AsyncClient(transport=httpx.AsyncHTTPTransport(uds="/var/run/docker.sock"), timeout=30) as client:
|
|
resp = await client.post(f"http://docker/containers/{plan.command}/restart")
|
|
resp.raise_for_status()
|
|
return
|
|
if plan.method == "shell":
|
|
if not plan.command:
|
|
raise RuntimeError("restart command is required for shell method")
|
|
cmd = plan.command
|
|
args = shlex.split(cmd)
|
|
proc = await asyncio.create_subprocess_exec(*args)
|
|
code = await proc.wait()
|
|
if code != 0:
|
|
raise RuntimeError(f"restart command failed with exit code {code}")
|
|
return
|
|
raise RuntimeError(f"unknown restart method {plan.method}")
|