Enhance Ollama model management and performance tuning
- Updated the `provision_llm` function to utilize the `/api/tags` endpoint for verifying available models, improving accuracy in model management. - Introduced a new `already_have_ollama_tag` function to ensure exact tag matching, preventing mismatches during model checks. - Enhanced the `pull_stream` function to require a successful status from the API before proceeding, ensuring reliable model downloads. - Added logic to handle unwritten blob files, improving the robustness of the model pulling process. - Updated documentation and tests to reflect these changes, ensuring clarity and reliability in Ollama model operations.
This commit is contained in:
@@ -127,6 +127,17 @@ def parse_ollama_models(path: Path) -> list[OllamaModelEntry]:
|
||||
return out
|
||||
|
||||
|
||||
def already_have_ollama_tag(have: set[str], wanted: str) -> bool:
|
||||
"""Exact tag match only — qwen2.5:3b must not satisfy qwen2.5:7b."""
|
||||
if wanted in have:
|
||||
return True
|
||||
if ":" not in wanted and f"{wanted}:latest" in have:
|
||||
return True
|
||||
if wanted.endswith(":latest") and wanted.rsplit(":", 1)[0] in have:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def write_ollama_models_preset(path: Path, preset: str) -> None:
|
||||
key = (preset or "recommended").strip().lower()
|
||||
if key not in OLLAMA_PRESETS:
|
||||
|
||||
+76
-29
@@ -923,8 +923,49 @@ def ensure_swarmui_running(cfg: Config, host: str, log: Log, restart: bool) -> N
|
||||
run_ssh(cfg, host, "sudo -n systemctl enable swarmui", check=False)
|
||||
|
||||
|
||||
_OLLAMA_TAGS_PY = r"""
|
||||
import json, urllib.request
|
||||
try:
|
||||
with urllib.request.urlopen("http://127.0.0.1:11434/api/tags", timeout=8) as r:
|
||||
data = json.loads(r.read().decode())
|
||||
except Exception as exc:
|
||||
print("ERR " + str(exc)[:200])
|
||||
raise SystemExit(0)
|
||||
for m in data.get("models") or []:
|
||||
if isinstance(m, dict):
|
||||
for key in ("name", "model"):
|
||||
n = m.get(key)
|
||||
if n:
|
||||
print(n)
|
||||
elif isinstance(m, str) and m.strip():
|
||||
print(m.strip())
|
||||
"""
|
||||
|
||||
|
||||
def _ollama_api_tags(cfg: Config, host: str) -> set[str]:
|
||||
"""Names from Ollama /api/tags (same source as Assistent / verify)."""
|
||||
out = run_ssh(
|
||||
cfg,
|
||||
host,
|
||||
"python3 - <<'PY'\n" + _OLLAMA_TAGS_PY + "\nPY",
|
||||
check=False,
|
||||
timeout=20,
|
||||
)
|
||||
names: set[str] = set()
|
||||
for ln in out.splitlines():
|
||||
s = ln.strip()
|
||||
if not s or s.startswith("ERR "):
|
||||
continue
|
||||
names.add(s)
|
||||
return names
|
||||
|
||||
|
||||
def provision_llm(cfg: Config, host: str, log: Log) -> None:
|
||||
from gpu_rent.llm_runtime import normalize_runtime, parse_ollama_models
|
||||
from gpu_rent.llm_runtime import (
|
||||
already_have_ollama_tag,
|
||||
normalize_runtime,
|
||||
parse_ollama_models,
|
||||
)
|
||||
from gpu_rent.ssh_ops import run_script_sudo
|
||||
from gpu_rent.state import load_state, save_state
|
||||
|
||||
@@ -950,6 +991,7 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None:
|
||||
st.notes.pop("llm_error", None)
|
||||
save_state(st)
|
||||
return
|
||||
still: list[str] = []
|
||||
if runtime == "ollama":
|
||||
_stop_units("gpu-rent-llamacpp")
|
||||
log("LLM: ставим/запускаем Ollama")
|
||||
@@ -967,48 +1009,53 @@ def provision_llm(cfg: Config, host: str, log: Log) -> None:
|
||||
if defaults:
|
||||
log(f"Ollama preferred: {defaults[0]}")
|
||||
names = [e.name for e in entries]
|
||||
still: list[str] = []
|
||||
if not names:
|
||||
log("ollama-models.yaml пуст — pull skip")
|
||||
else:
|
||||
# Fast path: all tags already present — skip upload/pull script.
|
||||
listed = run_ssh(
|
||||
cfg,
|
||||
host,
|
||||
"ollama list 2>/dev/null | awk 'NR>1 {print $1}' || true",
|
||||
check=False,
|
||||
timeout=30,
|
||||
)
|
||||
have = {ln.strip() for ln in listed.splitlines() if ln.strip()}
|
||||
missing = []
|
||||
for name in names:
|
||||
if name in have or (
|
||||
":" not in name and f"{name}:latest" in have
|
||||
) or (
|
||||
name.endswith(":latest") and name.rsplit(":", 1)[0] in have
|
||||
):
|
||||
continue
|
||||
missing.append(name)
|
||||
have = _ollama_api_tags(cfg, host)
|
||||
missing = [n for n in names if not already_have_ollama_tag(have, n)]
|
||||
if not missing:
|
||||
log(f"ollama pull: skip — все {len(names)} уже есть")
|
||||
log(f"ollama pull: skip — /api/tags уже {sorted(have)}")
|
||||
else:
|
||||
put_text(
|
||||
cfg, host, "/tmp/gpu-rent-ollama-models.json", json.dumps(missing, indent=2)
|
||||
)
|
||||
log(f"Ollama: pull {len(missing)} из манифеста (нет: {len(missing)})")
|
||||
run_python(
|
||||
cfg,
|
||||
host,
|
||||
_pkg_text("ollama_pull.py"),
|
||||
remote_path="/tmp/gpu-rent-ollama_pull.py",
|
||||
timeout=7200,
|
||||
log=log,
|
||||
log(
|
||||
f"Ollama: pull {len(missing)} из манифеста "
|
||||
f"(/api/tags={len(have)})"
|
||||
)
|
||||
try:
|
||||
run_python(
|
||||
cfg,
|
||||
host,
|
||||
_pkg_text("ollama_pull.py"),
|
||||
remote_path="/tmp/gpu-rent-ollama_pull.py",
|
||||
timeout=7200,
|
||||
log=log,
|
||||
)
|
||||
except Exception as exc:
|
||||
log(f"⚠ Ollama pull: {exc}")
|
||||
have = _ollama_api_tags(cfg, host)
|
||||
still = [n for n in names if not already_have_ollama_tag(have, n)]
|
||||
if still:
|
||||
log(
|
||||
"⚠ Ollama /api/tags без "
|
||||
+ ", ".join(still[:5])
|
||||
+ f" (есть: {sorted(have) or 'пусто'}). "
|
||||
"SwarmUI ок — GPU не гасим; Assistent будет пустой."
|
||||
)
|
||||
else:
|
||||
raise CloudError(f"неизвестный LLM_RUNTIME={runtime!r}")
|
||||
st = load_state()
|
||||
st.notes = dict(st.notes or {})
|
||||
st.notes["llm_runtime"] = runtime
|
||||
st.notes.pop("llm_error", None)
|
||||
if still:
|
||||
st.notes["llm_error"] = (
|
||||
"нет в /api/tags: " + ", ".join(still[:5])
|
||||
)[:500]
|
||||
else:
|
||||
st.notes.pop("llm_error", None)
|
||||
save_state(st)
|
||||
|
||||
|
||||
|
||||
+68
-4
@@ -84,7 +84,8 @@ print("WAIT timeout-slice")
|
||||
|
||||
# One-shot probe of configured stack endpoints on the VM (JSON line).
|
||||
_REMOTE_STACK_PROBE = r'''
|
||||
import json, urllib.error, urllib.request, subprocess
|
||||
import json, time, urllib.error, urllib.request, subprocess
|
||||
from pathlib import Path
|
||||
|
||||
def http_ok(url, timeout=4.0):
|
||||
try:
|
||||
@@ -107,9 +108,19 @@ def unit_active(name):
|
||||
except Exception:
|
||||
return "unknown"
|
||||
|
||||
def pulling_age():
|
||||
p = Path("/mnt/swarm_data/.gpu-rent-ollama-pulling")
|
||||
try:
|
||||
if p.is_file():
|
||||
return max(0.0, time.time() - p.stat().st_mtime)
|
||||
except OSError:
|
||||
return None
|
||||
return None
|
||||
|
||||
checks = []
|
||||
want_swarm = WANT_SWARM
|
||||
want_ollama = WANT_OLLAMA
|
||||
want_ollama_models = WANT_OLLAMA_MODELS
|
||||
|
||||
if want_swarm:
|
||||
ok, detail = http_ok("http://127.0.0.1:7801/")
|
||||
@@ -134,9 +145,11 @@ if want_swarm:
|
||||
"ok": ok,
|
||||
"detail": detail,
|
||||
"unit": unit_active("swarmui"),
|
||||
"retry": not ok,
|
||||
})
|
||||
|
||||
if want_ollama:
|
||||
retry = True
|
||||
try:
|
||||
req = urllib.request.Request("http://127.0.0.1:11434/api/tags", method="GET")
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
@@ -154,15 +167,27 @@ if want_ollama:
|
||||
preview = ", ".join(names[:3])
|
||||
extra = "" if len(names) <= 3 else f" +{len(names) - 3}"
|
||||
ok, detail = True, f"{len(names)} models ({preview}{extra})"
|
||||
retry = False
|
||||
elif want_ollama_models:
|
||||
age = pulling_age()
|
||||
if age is not None and age < 2700:
|
||||
ok, detail = False, f"Ollama up, 0 models — pull идёт ({int(age)}s)"
|
||||
retry = True
|
||||
else:
|
||||
ok, detail = True, "WARN 0 models — Assistent empty (GPU не гасим)"
|
||||
retry = False
|
||||
else:
|
||||
ok, detail = False, "Ollama up, 0 models — Assistent dropdown empty; ollama pull"
|
||||
ok, detail = True, "0 models (манифест пуст)"
|
||||
retry = False
|
||||
except Exception as exc:
|
||||
ok, detail = False, str(exc)[:160]
|
||||
retry = True
|
||||
checks.append({
|
||||
"name": "ollama",
|
||||
"ok": ok,
|
||||
"detail": detail,
|
||||
"unit": unit_active("gpu-rent-ollama"),
|
||||
"retry": retry,
|
||||
})
|
||||
|
||||
print(json.dumps({"checks": checks}, ensure_ascii=False))
|
||||
@@ -175,6 +200,7 @@ class ServiceCheck:
|
||||
ok: bool
|
||||
detail: str
|
||||
where: str = "vm" # vm | local
|
||||
retry: bool = True
|
||||
|
||||
|
||||
Log = Callable[[str], None]
|
||||
@@ -343,10 +369,29 @@ def _expected_services(cfg: Config) -> tuple[bool, bool]:
|
||||
return swarm, rt == "ollama"
|
||||
|
||||
|
||||
def _want_ollama_models(cfg: Config) -> bool:
|
||||
"""True when ollama-models.yaml lists tags that must appear in /api/tags."""
|
||||
if normalize_runtime(getattr(cfg, "llm_runtime", "none")) != "ollama":
|
||||
return False
|
||||
from gpu_rent.llm_runtime import parse_ollama_models
|
||||
|
||||
path = getattr(cfg, "ollama_models_manifest", None)
|
||||
if path is None:
|
||||
return False
|
||||
try:
|
||||
return bool(parse_ollama_models(path))
|
||||
except (OSError, ValueError):
|
||||
return False
|
||||
|
||||
|
||||
def _probe_vm_once(cfg: Config, host: str) -> list[ServiceCheck]:
|
||||
want_swarm, want_ollama = _expected_services(cfg)
|
||||
script = (
|
||||
_REMOTE_STACK_PROBE.replace("WANT_SWARM", "True" if want_swarm else "False")
|
||||
.replace(
|
||||
"WANT_OLLAMA_MODELS",
|
||||
"True" if (want_ollama and _want_ollama_models(cfg)) else "False",
|
||||
)
|
||||
.replace("WANT_OLLAMA", "True" if want_ollama else "False")
|
||||
)
|
||||
out = run_ssh(
|
||||
@@ -380,7 +425,13 @@ def _probe_vm_once(cfg: Config, host: str) -> list[ServiceCheck]:
|
||||
if unit and unit != "unknown":
|
||||
detail = f"{detail}; unit={unit}"
|
||||
checks.append(
|
||||
ServiceCheck(name=name, ok=bool(item.get("ok")), detail=detail, where="vm")
|
||||
ServiceCheck(
|
||||
name=name,
|
||||
ok=bool(item.get("ok")),
|
||||
detail=detail,
|
||||
where="vm",
|
||||
retry=bool(item.get("retry", True)),
|
||||
)
|
||||
)
|
||||
return checks
|
||||
|
||||
@@ -417,9 +468,22 @@ def verify_stack_on_vm(
|
||||
last = [ServiceCheck("ssh", False, str(exc)[:200], "vm")]
|
||||
if last and all(c.ok for c in last):
|
||||
for c in last:
|
||||
log(f" [ok] {c.name}: {c.detail}")
|
||||
mark = "warn" if c.detail.startswith("WARN") else "ok"
|
||||
log(f" [{mark}] {c.name}: {c.detail}")
|
||||
log("проверка VM: всё отвечает")
|
||||
return last
|
||||
stuck = [c for c in last if not c.ok and not c.retry]
|
||||
if stuck:
|
||||
for c in last:
|
||||
mark = "ok" if c.ok else "FAIL"
|
||||
log(f" [{mark}] {c.name}: {c.detail}")
|
||||
if raise_on_fail:
|
||||
failed = [c.name for c in stuck]
|
||||
raise CloudError(
|
||||
f"{', '.join(failed)} не готов и ждать бесполезно: "
|
||||
f"{stuck[0].detail}. GPU жив — gpu-rent logs / повторный up (pull)"
|
||||
)
|
||||
return last
|
||||
bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
|
||||
wait.tick(f" … ещё нет: {bad}")
|
||||
time.sleep(poll_every)
|
||||
|
||||
@@ -109,6 +109,38 @@ else
|
||||
apt-get install -y -qq git python3 python3-venv python3-pip python3-full curl wget ca-certificates
|
||||
fi
|
||||
|
||||
# Triton/sage JIT: python3.12-dev even on light (warm VMs skipped full apt).
|
||||
ensure_triton_build_deps() {
|
||||
if [[ -f /usr/include/python3.12/Python.h ]] && command -v gcc >/dev/null 2>&1; then
|
||||
return 0
|
||||
fi
|
||||
log "ставим gcc + python3.12-dev (Triton cuda_utils JIT)"
|
||||
apt-get update -qq
|
||||
apt-get install -y -qq build-essential python3-dev python3.12-dev
|
||||
}
|
||||
|
||||
ensure_libcuda_for_triton() {
|
||||
# Triton links -l:libcuda.so.1 -L/lib/x86_64-linux-gnu; Ubuntu driver is under /usr/lib.
|
||||
local dst=/lib/x86_64-linux-gnu/libcuda.so.1
|
||||
if [[ -e "$dst" ]]; then
|
||||
return 0
|
||||
fi
|
||||
local src
|
||||
for src in /usr/lib/x86_64-linux-gnu/libcuda.so.1 /usr/lib/x86_64-linux-gnu/nvidia/current/libcuda.so.1; do
|
||||
if [[ -e "$src" ]]; then
|
||||
mkdir -p /lib/x86_64-linux-gnu
|
||||
ln -sf "$src" "$dst"
|
||||
log "symlink $src → $dst (Triton gcc)"
|
||||
return 0
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
ensure_triton_build_deps
|
||||
ensure_libcuda_for_triton
|
||||
mkdir -p /var/tmp/gpu-rent-triton
|
||||
chown "${SWARM_USER}:${SWARM_USER}" /var/tmp/gpu-rent-triton 2>/dev/null || true
|
||||
|
||||
ensure_data_mount
|
||||
|
||||
mkdir -p \
|
||||
@@ -257,6 +289,13 @@ Environment=HOME=/home/${SWARM_USER}
|
||||
Environment=DOTNET_ROOT=/home/${SWARM_USER}/.dotnet
|
||||
Environment=DOTNET_CLI_HOME=/home/${SWARM_USER}
|
||||
Environment=PATH=/home/${SWARM_USER}/.dotnet:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
|
||||
Environment=TMPDIR=/tmp
|
||||
Environment=TEMP=/tmp
|
||||
Environment=TMP=/tmp
|
||||
Environment=TRITON_CACHE_DIR=/var/tmp/gpu-rent-triton
|
||||
Environment=TRITON_HOME=/var/tmp/gpu-rent-triton
|
||||
Environment=LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/lib/x86_64-linux-gnu
|
||||
Environment=LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu
|
||||
ExecStart=${SWARM_ROOT}/launch-linux.sh --launch_mode none --host 127.0.0.1 --port 7801
|
||||
Restart=on-failure
|
||||
RestartSec=8
|
||||
|
||||
@@ -2,10 +2,12 @@
|
||||
"""Pull Ollama models from a JSON list. Stdlib only. Runs on the VM.
|
||||
|
||||
Uses POST /api/pull with stream JSON for completed/total + speed lines.
|
||||
Stream end without status=success is NOT ok (cached layers ≠ registered model).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
@@ -16,6 +18,7 @@ from pathlib import Path
|
||||
JOBS = Path("/tmp/gpu-rent-ollama-models.json")
|
||||
MARKER = Path("/mnt/swarm_data/.gpu-rent-ollama-pulling")
|
||||
OLLAMA = "http://127.0.0.1:11434"
|
||||
MODELS_DIR = Path(os.environ.get("OLLAMA_MODELS") or "/mnt/swarm_data/ollama")
|
||||
|
||||
|
||||
def fmt_bytes(n: float) -> str:
|
||||
@@ -42,8 +45,12 @@ def progress_line(label: str, done: int, total: int | None, speed: float, *, wid
|
||||
|
||||
def listed() -> set[str]:
|
||||
"""Exact tags from `ollama list` (NAME column), e.g. qwen2.5:7b."""
|
||||
env = os.environ.copy()
|
||||
env.setdefault("OLLAMA_HOST", "127.0.0.1:11434")
|
||||
try:
|
||||
out = subprocess.check_output(["ollama", "list"], text=True, stderr=subprocess.DEVNULL)
|
||||
out = subprocess.check_output(
|
||||
["ollama", "list"], text=True, stderr=subprocess.DEVNULL, env=env
|
||||
)
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
return set()
|
||||
names: set[str] = set()
|
||||
@@ -56,6 +63,61 @@ def listed() -> set[str]:
|
||||
return names
|
||||
|
||||
|
||||
def api_tags() -> set[str]:
|
||||
"""Names from GET /api/tags (`name` and `model`)."""
|
||||
try:
|
||||
with urllib.request.urlopen(f"{OLLAMA}/api/tags", timeout=8) as resp:
|
||||
data = json.loads(resp.read().decode())
|
||||
except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError, json.JSONDecodeError):
|
||||
return set()
|
||||
names: set[str] = set()
|
||||
for m in data.get("models") or []:
|
||||
if isinstance(m, dict):
|
||||
for key in ("name", "model"):
|
||||
val = m.get(key)
|
||||
if val:
|
||||
names.add(str(val))
|
||||
elif isinstance(m, str) and m.strip():
|
||||
names.add(m.strip())
|
||||
return names
|
||||
|
||||
|
||||
def is_unwritten_blob(path: Path, *, sample: int = 64) -> bool:
|
||||
"""True if the file is sparse/truncated zeros (Ollama skip-download bug)."""
|
||||
try:
|
||||
size = path.stat().st_size
|
||||
except OSError:
|
||||
return False
|
||||
if size <= 0:
|
||||
return False
|
||||
n = min(sample, size)
|
||||
try:
|
||||
with path.open("rb") as fh:
|
||||
head = fh.read(n)
|
||||
except OSError:
|
||||
return False
|
||||
return bool(head) and head == b"\x00" * len(head)
|
||||
|
||||
|
||||
def purge_nul_blobs(models_dir: Path) -> int:
|
||||
"""Delete unwritten blob files so the next pull re-downloads layers."""
|
||||
blobs = models_dir / "blobs"
|
||||
if not blobs.is_dir():
|
||||
return 0
|
||||
n = 0
|
||||
for path in blobs.iterdir():
|
||||
if not path.is_file() or not is_unwritten_blob(path):
|
||||
continue
|
||||
size = path.stat().st_size
|
||||
try:
|
||||
path.unlink()
|
||||
except OSError:
|
||||
continue
|
||||
n += 1
|
||||
print(f"удалён пустой blob {path.name} ({size}B)", flush=True)
|
||||
return n
|
||||
|
||||
|
||||
def already_have(have: set[str], wanted: str) -> bool:
|
||||
"""Exact tag match only — qwen2.5:3b must not satisfy qwen2.5:7b."""
|
||||
if wanted in have:
|
||||
@@ -67,8 +129,38 @@ def already_have(have: set[str], wanted: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def pull_stream(name: str, label: str) -> None:
|
||||
body = json.dumps({"name": name, "stream": True}).encode("utf-8")
|
||||
def wait_in_tags(name: str, *, timeout: float) -> set[str]:
|
||||
deadline = time.monotonic() + max(timeout, 0.0)
|
||||
have: set[str] = set()
|
||||
while True:
|
||||
have = api_tags() | listed()
|
||||
if already_have(have, name):
|
||||
return have
|
||||
if time.monotonic() >= deadline:
|
||||
return have
|
||||
time.sleep(1.0)
|
||||
|
||||
|
||||
def pull_cli(name: str, label: str) -> None:
|
||||
env = os.environ.copy()
|
||||
env["OLLAMA_HOST"] = "127.0.0.1:11434"
|
||||
proc = subprocess.Popen(
|
||||
["ollama", "pull", name],
|
||||
env=env,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
)
|
||||
assert proc.stdout is not None
|
||||
for line in proc.stdout:
|
||||
print(f"{label} {line.rstrip()}", flush=True)
|
||||
code = proc.wait()
|
||||
if code != 0:
|
||||
raise RuntimeError(f"ollama pull CLI exit {code}")
|
||||
|
||||
|
||||
def pull_stream(name: str, label: str, *, tags_wait: float = 45.0) -> None:
|
||||
body = json.dumps({"model": name, "name": name, "stream": True}).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
f"{OLLAMA}/api/pull",
|
||||
data=body,
|
||||
@@ -78,6 +170,8 @@ def pull_stream(name: str, label: str) -> None:
|
||||
t0 = time.monotonic()
|
||||
last_print = 0.0
|
||||
last_touch = 0.0
|
||||
last_status = ""
|
||||
got_success = False
|
||||
with urllib.request.urlopen(req, timeout=7200) as resp:
|
||||
while True:
|
||||
raw = resp.readline()
|
||||
@@ -90,10 +184,10 @@ def pull_stream(name: str, label: str) -> None:
|
||||
if ev.get("error"):
|
||||
raise RuntimeError(str(ev["error"]))
|
||||
status = str(ev.get("status") or "")
|
||||
last_status = status or last_status
|
||||
total = int(ev.get("total") or 0) or None
|
||||
done = int(ev.get("completed") or 0)
|
||||
now = time.monotonic()
|
||||
# Keep idle-killer marker fresh for long pulls (max-age 3h).
|
||||
if now - last_touch >= 60.0:
|
||||
try:
|
||||
MARKER.write_text(f"{int(time.time())}\n", encoding="utf-8")
|
||||
@@ -112,8 +206,34 @@ def pull_stream(name: str, label: str) -> None:
|
||||
print(f"{label} {status}", flush=True)
|
||||
last_print = now
|
||||
if status == "success":
|
||||
got_success = True
|
||||
break
|
||||
print(f"{label} ok", flush=True)
|
||||
if not got_success:
|
||||
print(
|
||||
f"{label} HTTP stream без success (last={last_status!r}) — ollama pull CLI",
|
||||
flush=True,
|
||||
)
|
||||
pull_cli(name, label)
|
||||
have = wait_in_tags(name, timeout=tags_wait)
|
||||
if already_have(have, name):
|
||||
print(f"{label} ok (в /api/tags)", flush=True)
|
||||
return
|
||||
purged = purge_nul_blobs(MODELS_DIR)
|
||||
if purged:
|
||||
print(
|
||||
f"{label} {purged} blob(ов) из нулей (pull skip по размеру) — ollama pull CLI",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
print(f"{label} нет в /api/tags — ollama pull CLI", flush=True)
|
||||
pull_cli(name, label)
|
||||
have = wait_in_tags(name, timeout=tags_wait)
|
||||
if already_have(have, name):
|
||||
print(f"{label} ok (в /api/tags)", flush=True)
|
||||
return
|
||||
raise RuntimeError(
|
||||
f"после pull нет в /api/tags (есть: {sorted(have) or 'пусто'})"
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
@@ -124,7 +244,7 @@ def main() -> int:
|
||||
if not isinstance(models, list) or not models:
|
||||
print("ollama pull: пустой список — skip")
|
||||
return 0
|
||||
have = listed()
|
||||
have = api_tags() | listed()
|
||||
MARKER.parent.mkdir(parents=True, exist_ok=True)
|
||||
MARKER.write_text(f"{int(time.time())}\n", encoding="utf-8")
|
||||
failed = 0
|
||||
@@ -141,7 +261,13 @@ def main() -> int:
|
||||
try:
|
||||
pull_stream(name, label=f"{prefix} {name}")
|
||||
have.add(name)
|
||||
except (urllib.error.URLError, urllib.error.HTTPError, OSError, TimeoutError, RuntimeError) as exc:
|
||||
except (
|
||||
urllib.error.URLError,
|
||||
urllib.error.HTTPError,
|
||||
OSError,
|
||||
TimeoutError,
|
||||
RuntimeError,
|
||||
) as exc:
|
||||
failed += 1
|
||||
print(f"FAIL pull {name}: {exc}")
|
||||
finally:
|
||||
|
||||
@@ -243,7 +243,70 @@ def patch_backends_extra_args(extra: str) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def sage_already_importable(py: Path) -> bool:
|
||||
def strip_sage_extra_args() -> bool:
|
||||
"""Remove --use-sage-attention so Comfy does not JIT-crash without python3-dev."""
|
||||
if not BACKENDS.is_file():
|
||||
return False
|
||||
text = BACKENDS.read_text(encoding="utf-8")
|
||||
out: list[str] = []
|
||||
n = 0
|
||||
for line in text.splitlines():
|
||||
m = re.match(r"^(\s*ExtraArgs:\s*)(.*)$", line)
|
||||
if not m or "--use-sage-attention" not in m.group(2):
|
||||
out.append(line)
|
||||
continue
|
||||
rest = re.sub(r"(^|\s)--use-sage-attention\b", "", m.group(2)).strip()
|
||||
out.append(f"{m.group(1)}{rest}")
|
||||
n += 1
|
||||
if not n:
|
||||
return False
|
||||
BACKENDS.write_text("\n".join(out) + "\n", encoding="utf-8")
|
||||
print("stripped --use-sage-attention (triton JIT failed)")
|
||||
return True
|
||||
|
||||
|
||||
_TRITON_JIT_PROBE = r"""
|
||||
import os
|
||||
os.environ.setdefault("TRITON_CACHE_DIR", "/var/tmp/gpu-rent-triton")
|
||||
os.environ.setdefault("TMPDIR", "/tmp")
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
@triton.jit
|
||||
def _add(x_ptr, y_ptr, n):
|
||||
i = tl.program_id(0)
|
||||
if i < n:
|
||||
tl.store(y_ptr + i, tl.load(x_ptr + i) + 1)
|
||||
|
||||
x = torch.ones(32, device="cuda", dtype=torch.float16)
|
||||
y = torch.empty_like(x)
|
||||
_add[(32,)](x, y, 32)
|
||||
torch.cuda.synchronize()
|
||||
print("triton JIT ok")
|
||||
"""
|
||||
|
||||
|
||||
def triton_jit_ok(py: Path) -> bool:
|
||||
"""import sageattention is not enough — first gen compiles cuda_utils.c."""
|
||||
env = _pip_env()
|
||||
env["TRITON_CACHE_DIR"] = "/var/tmp/gpu-rent-triton"
|
||||
env["TMPDIR"] = "/tmp"
|
||||
try:
|
||||
out = subprocess.check_output(
|
||||
[str(py), "-c", _TRITON_JIT_PROBE],
|
||||
text=True,
|
||||
stderr=subprocess.STDOUT,
|
||||
timeout=180,
|
||||
env=env,
|
||||
)
|
||||
print((out or "").strip()[-400:] or "triton JIT ok")
|
||||
return True
|
||||
except (subprocess.CalledProcessError, OSError, subprocess.TimeoutExpired) as exc:
|
||||
body = getattr(exc, "output", None) or str(exc)
|
||||
print(f"WARN triton JIT failed: {body[-1500:]}")
|
||||
print("нужны python3.12-dev + gcc; ExtraArgs без --use-sage-attention")
|
||||
return False
|
||||
code, _ = _run(
|
||||
[str(py), "-c", "import triton, sageattention"],
|
||||
timeout=60,
|
||||
@@ -269,6 +332,14 @@ def _pip_env() -> dict[str, str]:
|
||||
return env
|
||||
|
||||
|
||||
def sage_already_importable(py: Path) -> bool:
|
||||
code, _ = _run(
|
||||
[str(py), "-c", "import triton, sageattention"],
|
||||
timeout=60,
|
||||
)
|
||||
return code == 0
|
||||
|
||||
|
||||
def pip_install_sage(py: Path) -> bool:
|
||||
"""Install via ``python -m pip`` (never exec venv/bin/pip directly)."""
|
||||
if sage_already_importable(py):
|
||||
@@ -299,16 +370,20 @@ def main() -> int:
|
||||
except json.JSONDecodeError:
|
||||
prev = {}
|
||||
same_gpu = bool(prev.get("uuid") and prev.get("uuid") == plan["uuid"])
|
||||
# Always repair FDS corruption / relative StartScript even when tune is a no-op.
|
||||
fds_fixed = sanitize_backends_fds()
|
||||
if ensure_absolute_start_script():
|
||||
fds_fixed = True
|
||||
if same_gpu and prev.get("extra_args") == plan["extra_args"]:
|
||||
if prev.get("pip_ok") or not plan["use_sage"]:
|
||||
jit_ok_prev = bool(prev.get("jit_ok"))
|
||||
if (prev.get("pip_ok") or not plan["use_sage"]) and (
|
||||
not plan["use_sage"] or jit_ok_prev
|
||||
):
|
||||
print(f"perf tune already applied for {plan['name']} ({plan['tier']})")
|
||||
if fds_fixed:
|
||||
print("RESTART_SWARMUI=1")
|
||||
return 0
|
||||
if plan["use_sage"] and prev.get("pip_ok") and not jit_ok_prev:
|
||||
print("perf tune: retry triton JIT (cuda_utils)")
|
||||
|
||||
if same_gpu and plan["use_sage"] and prev.get("pip_ok") is False:
|
||||
print("perf tune: retry (previous pip_ok=false — sage/triton ещё не встали)")
|
||||
@@ -316,16 +391,22 @@ def main() -> int:
|
||||
print(f"perf tune: {plan['name']} tier={plan['tier']} sage={plan['use_sage']}")
|
||||
restarted_needed = bool(fds_fixed)
|
||||
pip_ok = not plan["use_sage"]
|
||||
jit_ok = not plan["use_sage"]
|
||||
if plan["use_sage"]:
|
||||
py = find_comfy_python()
|
||||
if py:
|
||||
pip_ok = pip_install_sage(py)
|
||||
jit_ok = bool(pip_ok and triton_jit_ok(py))
|
||||
else:
|
||||
print("Comfy venv python not found yet — will retry next up")
|
||||
pip_ok = False
|
||||
# Only patch ExtraArgs when wheels installed — otherwise Comfy may break.
|
||||
if pip_ok and patch_backends_extra_args(plan["extra_args"]):
|
||||
restarted_needed = True
|
||||
jit_ok = False
|
||||
if pip_ok and jit_ok:
|
||||
if patch_backends_extra_args(plan["extra_args"]):
|
||||
restarted_needed = True
|
||||
else:
|
||||
if strip_sage_extra_args():
|
||||
restarted_needed = True
|
||||
if ensure_absolute_start_script():
|
||||
restarted_needed = True
|
||||
|
||||
@@ -333,8 +414,9 @@ def main() -> int:
|
||||
"uuid": plan["uuid"],
|
||||
"name": plan["name"],
|
||||
"tier": plan["tier"],
|
||||
"extra_args": plan["extra_args"] if pip_ok else "",
|
||||
"extra_args": plan["extra_args"] if (pip_ok and jit_ok) else "",
|
||||
"pip_ok": pip_ok,
|
||||
"jit_ok": jit_ok,
|
||||
"restart_needed": restarted_needed,
|
||||
}
|
||||
MARKER.write_text(json.dumps(marker, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
Reference in New Issue
Block a user