Leonid Pershin
241ada62cb
- Added a `--update` option to the `up` command in the CLI for forced updates during warm ACTIVE states. - Improved the idle-killer logic with new functions to check for reusable credentials and refresh settings on existing remote credentials. - Enhanced the `seed_civitai` function to implement caching for model jobs, optimizing the provisioning process. - Updated the installation script for Ollama to prevent unnecessary restarts when the service configuration has not changed. - Added diagnostics for model job caching and improved error handling in various functions.
312 lines
11 KiB
Python
312 lines
11 KiB
Python
"""Create OpenStack application credential and arm idle-killer on the VM."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import secrets
|
||
import time
|
||
from collections.abc import Callable
|
||
from importlib.resources import files
|
||
|
||
from gpu_rent.config import Config
|
||
from gpu_rent.errors import CloudError
|
||
from gpu_rent.ssh_ops import put_text, run_ssh
|
||
|
||
Log = Callable[[str], None]
|
||
DATA = "/mnt/swarm_data"
|
||
CREDS_REMOTE = "/root/.gpu-rent/idle-killer.json"
|
||
SCRIPT_REMOTE = "/usr/local/lib/gpu-rent/idle_killer.py"
|
||
UNIT = "gpu-rent-idle-killer"
|
||
CRED_NAME_PREFIX = "gpu-rent-idle-killer"
|
||
|
||
|
||
def _pkg_text(name: str) -> str:
|
||
return files("gpu_rent.remote").joinpath(name).read_text(encoding="utf-8")
|
||
|
||
|
||
def revoke_old_credentials(conn, log: Log) -> None:
|
||
user_id = getattr(conn, "current_user_id", None)
|
||
if not user_id:
|
||
return
|
||
try:
|
||
for ac in conn.identity.application_credentials(user=user_id):
|
||
name = getattr(ac, "name", "") or ""
|
||
if name.startswith(CRED_NAME_PREFIX):
|
||
try:
|
||
conn.identity.delete_application_credential(user_id, ac.id)
|
||
log(f"отозван старый app cred {name}")
|
||
except Exception as exc:
|
||
log(f"не отозвать {name}: {exc}")
|
||
except Exception as exc:
|
||
log(f"list application_credentials: {exc}")
|
||
|
||
|
||
def create_application_credential(conn, cfg: Config, server_id: str, log: Log) -> dict:
|
||
user_id = getattr(conn, "current_user_id", None)
|
||
if not user_id:
|
||
raise CloudError("нет current_user_id — idle-killer без app cred")
|
||
revoke_old_credentials(conn, log)
|
||
secret = secrets.token_urlsafe(32)
|
||
name = f"{CRED_NAME_PREFIX}-{server_id[:8]}"
|
||
# Only this compute — never /servers/* (would allow deleting any VM in the project).
|
||
access_rules = [
|
||
{
|
||
"service": "compute",
|
||
"method": "DELETE",
|
||
"path": f"/v2.1/servers/{server_id}",
|
||
},
|
||
{
|
||
"service": "compute",
|
||
"method": "GET",
|
||
"path": f"/v2.1/servers/{server_id}",
|
||
},
|
||
]
|
||
try:
|
||
ac = conn.identity.create_application_credential(
|
||
user=user_id,
|
||
name=name,
|
||
secret=secret,
|
||
description="gpu-rent idle-killer: delete this compute only",
|
||
access_rules=access_rules,
|
||
)
|
||
except Exception as exc:
|
||
raise CloudError(
|
||
f"не создать application credential с access_rules (только этот server): {exc}. "
|
||
"Без узких правил idle-killer не вооружаем (fail closed). "
|
||
"Нужны права identity:application_credential_create на сервисного пользователя."
|
||
) from exc
|
||
ac_id = getattr(ac, "id", None) or (ac.get("id") if isinstance(ac, dict) else None)
|
||
ac_secret = getattr(ac, "secret", None) or secret
|
||
if not ac_id:
|
||
raise CloudError("application credential создан без id")
|
||
log(f"application credential {name} (DELETE/GET только {server_id[:12]}…)")
|
||
return {
|
||
"auth_url": cfg.os_auth_url,
|
||
"project_id": cfg.os_project_id,
|
||
"region_name": cfg.os_region_name,
|
||
"application_credential_id": ac_id,
|
||
"application_credential_secret": ac_secret,
|
||
"server_id": server_id,
|
||
"idle_minutes": cfg.idle_minutes,
|
||
"grace_minutes": cfg.idle_grace_minutes,
|
||
"grace_from": int(time.time()),
|
||
"swarm_url": "http://127.0.0.1:7801",
|
||
}
|
||
|
||
|
||
def install_units(cfg: Config, host: str, log: Log) -> None:
|
||
script = _pkg_text("idle_killer.py")
|
||
run_ssh(cfg, host, "sudo -n mkdir -p /usr/local/lib/gpu-rent /root/.gpu-rent", check=False)
|
||
put_text(cfg, host, "/tmp/gpu-rent-idle_killer.py", script)
|
||
run_ssh(
|
||
cfg,
|
||
host,
|
||
f"sudo -n mv /tmp/gpu-rent-idle_killer.py {SCRIPT_REMOTE} && "
|
||
f"sudo -n chmod 755 {SCRIPT_REMOTE}",
|
||
)
|
||
service = f"""[Unit]
|
||
Description=gpu-rent idle-killer (one shot)
|
||
After=network-online.target
|
||
|
||
[Service]
|
||
Type=oneshot
|
||
ExecStart=/usr/bin/python3 {SCRIPT_REMOTE}
|
||
Nice=10
|
||
"""
|
||
timer = f"""[Unit]
|
||
Description=gpu-rent idle-killer every minute
|
||
|
||
[Timer]
|
||
OnBootSec=2min
|
||
OnUnitActiveSec=1min
|
||
AccuracySec=15s
|
||
Unit={UNIT}.service
|
||
|
||
[Install]
|
||
WantedBy=timers.target
|
||
"""
|
||
put_text(cfg, host, f"/tmp/{UNIT}.service", service)
|
||
put_text(cfg, host, f"/tmp/{UNIT}.timer", timer)
|
||
run_ssh(
|
||
cfg,
|
||
host,
|
||
f"sudo -n mv /tmp/{UNIT}.service /etc/systemd/system/{UNIT}.service && "
|
||
f"sudo -n mv /tmp/{UNIT}.timer /etc/systemd/system/{UNIT}.timer && "
|
||
"sudo -n systemctl daemon-reload && "
|
||
f"sudo -n systemctl enable --now {UNIT}.timer",
|
||
)
|
||
log(f"systemd timer {UNIT}.timer")
|
||
|
||
|
||
def _remote_killer_reusable(cfg: Config, host: str, server_id: str) -> bool:
|
||
"""True if VM already has matching idle-killer creds + active timer."""
|
||
cmd = (
|
||
f"sid=$(cat {DATA}/.gpu-rent-server-id 2>/dev/null | tr -d '[:space:]'); "
|
||
f"echo sid=$sid; echo want={server_id}; "
|
||
f"test -f {DATA}/.gpu-rent-killer-creds-ok && echo creds=yes || echo creds=no; "
|
||
f"test -f {DATA}/.gpu-rent-killer-armed && echo armed=yes || echo armed=no; "
|
||
f"sudo -n test -f {CREDS_REMOTE} && echo rootcreds=yes || echo rootcreds=no; "
|
||
f"echo timer=$(systemctl is-active {UNIT}.timer 2>/dev/null || echo inactive)"
|
||
)
|
||
try:
|
||
out = run_ssh(cfg, host, cmd, check=False, timeout=25)
|
||
except Exception:
|
||
return False
|
||
parsed: dict[str, str] = {}
|
||
for line in out.splitlines():
|
||
if "=" in line:
|
||
k, v = line.strip().split("=", 1)
|
||
parsed[k] = v
|
||
return (
|
||
parsed.get("sid") == server_id
|
||
and parsed.get("want") == server_id
|
||
and parsed.get("creds") == "yes"
|
||
and parsed.get("armed") == "yes"
|
||
and parsed.get("rootcreds") == "yes"
|
||
and parsed.get("timer") == "active"
|
||
)
|
||
|
||
|
||
def _refresh_killer_grace(cfg: Config, host: str) -> None:
|
||
"""Bump grace_from / idle settings on existing remote creds; clear idle-since."""
|
||
run_ssh(
|
||
cfg,
|
||
host,
|
||
"sudo -n python3 -c \""
|
||
"import json,time; from pathlib import Path; "
|
||
f"p=Path('{CREDS_REMOTE}'); d=json.loads(p.read_text()); "
|
||
f"d['idle_minutes']={int(cfg.idle_minutes)}; "
|
||
f"d['grace_minutes']={int(cfg.idle_grace_minutes)}; "
|
||
"d['grace_from']=int(time.time()); "
|
||
"p.write_text(json.dumps(d,indent=2)+chr(10)); p.chmod(0o600)"
|
||
"\" && "
|
||
f"sudo -n rm -f {DATA}/.gpu-rent-idle-since && "
|
||
f"date -u +%Y-%m-%dT%H:%M:%SZ | sudo -n tee {DATA}/.gpu-rent-killer-armed >/dev/null && "
|
||
f"sudo -n chmod 644 {DATA}/.gpu-rent-killer-armed "
|
||
f"{DATA}/.gpu-rent-killer-creds-ok {DATA}/.gpu-rent-server-id 2>/dev/null || true",
|
||
check=False,
|
||
timeout=30,
|
||
)
|
||
|
||
|
||
def arm_idle_killer(
|
||
cfg: Config,
|
||
host: str,
|
||
conn,
|
||
server_id: str,
|
||
log: Log,
|
||
) -> None:
|
||
if not server_id:
|
||
raise CloudError("idle-killer: нет server_id")
|
||
if _remote_killer_reusable(cfg, host, server_id):
|
||
_refresh_killer_grace(cfg, host)
|
||
log(
|
||
f"idle-killer: reuse app cred (server {server_id[:12]}…) — "
|
||
f"льгота {cfg.idle_grace_minutes} мин, потом {cfg.idle_minutes} мин → delete"
|
||
)
|
||
return
|
||
creds = create_application_credential(conn, cfg, server_id, log)
|
||
put_text(
|
||
cfg,
|
||
host,
|
||
"/tmp/gpu-rent-idle-killer.json",
|
||
json.dumps(creds, indent=2) + "\n",
|
||
mode=0o600,
|
||
)
|
||
run_ssh(
|
||
cfg,
|
||
host,
|
||
"sudo -n mkdir -p /root/.gpu-rent && "
|
||
f"sudo -n mv /tmp/gpu-rent-idle-killer.json {CREDS_REMOTE} && "
|
||
f"sudo -n chmod 600 {CREDS_REMOTE}",
|
||
)
|
||
put_text(cfg, host, f"{DATA}/.gpu-rent-server-id", server_id + "\n")
|
||
put_text(
|
||
cfg,
|
||
host,
|
||
f"{DATA}/.gpu-rent-killer-armed",
|
||
time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + "\n",
|
||
)
|
||
put_text(cfg, host, f"{DATA}/.gpu-rent-killer-creds-ok", "1\n")
|
||
run_ssh(
|
||
cfg,
|
||
host,
|
||
f"sudo -n rm -f {DATA}/.gpu-rent-idle-since && "
|
||
f"sudo -n chmod 644 {DATA}/.gpu-rent-server-id "
|
||
f"{DATA}/.gpu-rent-killer-armed {DATA}/.gpu-rent-killer-creds-ok",
|
||
check=False,
|
||
)
|
||
install_units(cfg, host, log)
|
||
log(
|
||
f"idle-killer вооружён: льгота {cfg.idle_grace_minutes} мин, "
|
||
f"потом {cfg.idle_minutes} мин пустой очереди → delete compute"
|
||
)
|
||
|
||
|
||
def killer_status_lines(cfg: Config, host: str) -> list[str]:
|
||
"""Best-effort SSH snapshot for `gpu-rent status`."""
|
||
lines: list[str] = []
|
||
script = f"""
|
||
python3 - <<'PY'
|
||
from pathlib import Path
|
||
import time
|
||
data = Path("{DATA}")
|
||
now = time.time()
|
||
armed = (data / ".gpu-rent-killer-armed").is_file()
|
||
creds = (data / ".gpu-rent-killer-creds-ok").is_file()
|
||
print("armed", "yes" if armed else "no")
|
||
print("creds", "yes" if creds else "no")
|
||
hold = data / ".gpu-rent-hold-until"
|
||
if hold.is_file():
|
||
try:
|
||
ts = float(hold.read_text().strip().split()[0])
|
||
left = int(ts - now)
|
||
print("hold", left if left > 0 else 0)
|
||
except Exception:
|
||
print("hold", "bad")
|
||
else:
|
||
print("hold", "none")
|
||
idle = data / ".gpu-rent-idle-since"
|
||
if idle.is_file():
|
||
try:
|
||
ts = float(idle.read_text().strip().split()[0])
|
||
print("idle_for", int(now - ts))
|
||
except Exception:
|
||
print("idle_for", "bad")
|
||
else:
|
||
print("idle_for", "none")
|
||
PY
|
||
"""
|
||
try:
|
||
from gpu_rent.ssh_ops import run_ssh as _ssh
|
||
|
||
out = _ssh(cfg, host, script.strip(), check=False, timeout=20)
|
||
except Exception as exc:
|
||
return [f"ssh fail: {exc}"]
|
||
parsed = {}
|
||
for line in out.splitlines():
|
||
parts = line.strip().split(None, 1)
|
||
if len(parts) == 2:
|
||
parsed[parts[0]] = parts[1]
|
||
if parsed.get("armed") != "yes":
|
||
lines.append("не вооружён")
|
||
elif parsed.get("creds") != "yes":
|
||
lines.append("слеп (нет кредов)")
|
||
else:
|
||
lines.append("armed")
|
||
hold = parsed.get("hold")
|
||
if hold and hold not in {"none", "bad", "0"}:
|
||
try:
|
||
sec = int(hold)
|
||
lines.append(f"hold ещё {sec // 60}m")
|
||
except ValueError:
|
||
lines.append(f"hold={hold}")
|
||
idle_for = parsed.get("idle_for")
|
||
if idle_for and idle_for not in {"none", "bad"}:
|
||
try:
|
||
sec = int(idle_for)
|
||
lines.append(f"пустая очередь {sec // 60}m {sec % 60}s / {cfg.idle_minutes}m")
|
||
except ValueError:
|
||
pass
|
||
return lines or ["неизвестно"]
|