From e1dfdd80d99669ee92a356fdf85804ca20991175 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:41:33 +0300 Subject: [PATCH] feat: the slow-torch chip switches torch to its CUDA 13 build in place: downloads first, installs from those files, proves it imports, else puts the old build back; restarts Co-Authored-By: Claude Opus 5.5 --- app/ui/features/generate/warn.js | 3 +- app/ui/features/generate/warn.test.js | 6 + app/ui/shell/api.js | 2 + app/ui/shell/maintenance.js | 19 +++- core/routes.py | 42 +++++++ core/tests/test_torch_swap.py | 158 ++++++++++++++++++++++++++ core/torch_swap.py | 122 ++++++++++++++++++++ 7 files changed, 350 insertions(+), 2 deletions(-) create mode 100644 core/tests/test_torch_swap.py create mode 100644 core/torch_swap.py diff --git a/app/ui/features/generate/warn.js b/app/ui/features/generate/warn.js index 597c14a0..1cdbbe1b 100644 --- a/app/ui/features/generate/warn.js +++ b/app/ui/features/generate/warn.js @@ -22,7 +22,8 @@ export default { title: "SLA and ComfyUI's Model Sparse Attention node are both in this pipeline. Both sparsify attention; remove one." }).node; twice.hidden = true; // A torch build that halves int8 speed on this GPU (cu12x on Blackwell). Asked once: the build cannot change while ComfyUI runs. - const slow = c.button.sm({ label: "⚠ Slow torch for int8", tone: "ghost", disabled: true, title: "" }).node; + const slow = c.button.sm({ label: "⚠ Slow torch for int8", tone: "ghost", title: "", + onClick: () => app.maintenance?.cuda13(slow.title) }).node; slow.hidden = true; host.append(chip, plain, keyless, twice, slow); Promise.resolve().then(() => app.api?.torchBuild?.()).then((r) => { diff --git a/app/ui/features/generate/warn.test.js b/app/ui/features/generate/warn.test.js index 62fd4c14..8de5a2d4 100644 --- a/app/ui/features/generate/warn.test.js +++ b/app/ui/features/generate/warn.test.js @@ -46,6 +46,12 @@ test("the slow-torch chip shows the server's sentence when this build runs int8 await tick(); assert.equal(slowChip(host).hidden, false); assert.match(slowChip(host).title, /CUDA 12\.8/); + let asked = null; + host = document.createElement("div"); + warn.setup({ host, app: { pipeline, maintenance: { cuda13: (why) => { asked = why; } }, api: { torchBuild: async () => ({ slow_int8: "torch is built for CUDA 12.8: reinstall" }) } } }); + await tick(); + slowChip(host).click(); + assert.match(asked, /CUDA 12\.8/, "the chip offers the switch, with the reason"); host = document.createElement("div"); warn.setup({ host, app: { pipeline, api: { torchBuild: async () => ({ slow_int8: null }) } } }); await tick(); diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index 4160892f..51f764ab 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -118,6 +118,8 @@ export const api = { pack: (action, body) => call("POST", `/api/packs/${action}`, body || {}), git: (action, body) => call("POST", `/api/git/${action}`, body || {}), health: () => call("GET", "/api/health"), + torchSwap: () => call("POST", "/api/torch/cuda13"), + torchSwapState: () => call("GET", "/api/torch/cuda13"), gitFull: () => call("GET", "/api/git/status"), gitStatus: () => call("GET", "/api/git/status?remote=0"), media: () => call("GET", "/api/media"), diff --git a/app/ui/shell/maintenance.js b/app/ui/shell/maintenance.js index b21d752c..8a7d8d70 100644 --- a/app/ui/shell/maintenance.js +++ b/app/ui/shell/maintenance.js @@ -32,7 +32,7 @@ export function createMaintenance({ api, flush }) { const card = c.modal.generic({ title: "Working", size: "sm", closeOnOutside: false, closeOnEsc: false, body: c.region.stack({ gap: "md", children: [c.progress.indeterminate({ label: "Working" }), text] }) }); card.setText = (t) => text.setText ? text.setText(t) : (text.node.textContent = t); try { - const res = await action(); + const res = await action(card); if (res && res.restarting === false) { // nothing restarted: say what is true, and stay on the page card.close(); return c.modal.dialogue({ title: "Not restarted", message: `${res.blocked || (said ? said(res) : "Nothing changed.")} ${res.blocked ? "Restart ComfyUI yourself when the generation has finished." : ""}`.trim(), confirmLabel: "OK", cancelLabel: "Close" }).result; @@ -79,6 +79,23 @@ export function createMaintenance({ api, flush }) { if (!(await ask("Roll back", `Roll back to ${String(target.commit).slice(0, 8)}${target.subject ? ` ("${target.subject}")` : ""}? If that update changed requirements.txt, dependencies are not reinstalled: run pip by hand afterwards if things do not load.`))) return; return run(`Rolling back to ${String(target.commit).slice(0, 8)}…`, () => api.git("rollback"), (r) => `Rolled back ${r.before} → ${r.after}.`); }, + /** Switch torch to its CUDA 13 build (same versions), then restart. `why` is the warning that offered it. */ + async cuda13(why) { + if (!(await ask("Switch torch to CUDA 13", `${why || ""}\n\nDownload the CUDA 13 build of the same torch, torchvision and torchaudio versions (several GB), install it, check it imports, and restart ComfyUI. If it fails, the old build is put back. This can take 10+ minutes.`.trim()))) return; + // Started, then polled: the download outlives one request through a rental's proxy. A server that comes back + // with no job and no result has restarted on its own after a success -- run() then waits for it and reloads. + const action = async (card) => { + await api.torchSwap(); + for (;;) { + await wait(1000); + let s; try { s = await api.torchSwapState(); } catch { continue; } // down while restarting: keep asking + if (s.step) card.setText(s.step); + if (s.result) return s.result; + if (!s.running) return { restarting: true }; + } + }; + return run("Switching torch to CUDA 13…", action, (r) => r.message || ""); + }, async restart() { if (!(await ask("Restart ComfyUI", "Restart now? The server is down for 10–40 s and any running generation is lost. This page reloads when it is back."))) return; const action = () => api.git("restart"); action.dropsConnection = true; diff --git a/core/routes.py b/core/routes.py index 593b83cf..f897bd33 100644 --- a/core/routes.py +++ b/core/routes.py @@ -17,6 +17,7 @@ from . import (backend_log, config, control as control_mod, graph as graph_mod, story, sysinfo, temp_files, + torch_swap, readiness, update as update_mod, workflow_import, registry as registry_mod, serve as static, widgets) from .contract import CONTRACT_VERSION @@ -573,6 +574,47 @@ def register(routes, prefix=None): _schedule_restart() return web.json_response({"restarting": True}) + # The torch switch runs as a task the client polls: a multi-GB download outlives the ~100s a + # rental's proxy keeps one request open. Shares the git lock and its restart rules. + torch_job = {"running": False, "result": None, "task": None} + + async def _torch_run(): + global _pending_restart + async with _git_lock: + try: + result = await asyncio.to_thread(torch_swap.swap) + except Exception as exc: # noqa: BLE001 + log.broke("torch switch", exc, doing="swap") + result = {"installed": False, "message": f"{type(exc).__name__}: {exc}"} + if not result.get("installed"): + result = {"restarting": False, **result} + elif _generation_running(): + _pending_restart = True + result = {"restarting": False, "blocked": "The new torch is installed, but a generation started: restart once it finishes.", **result} + else: + _pending_restart = False + result = {"restarting": True, **result} + # Written BEFORE the restart is scheduled, so a poll can read it while the server is still up. + torch_job.update(running=False, result=result) + if result["restarting"]: + _schedule_restart() + + @routes.post(P + "/api/torch/cuda13") + async def _torch_cuda13(_req): + """Start switching torch to the CUDA 13 build of the same versions; GET says how it goes.""" + if _generation_running(): + return web.json_response({"detail": "A generation is running. Switch torch when it has finished."}, status=409) + if torch_job["running"] or _git_lock.locked(): + return web.json_response({"detail": "An update or a torch switch is already running."}, status=409) + if _pending_restart: + return web.json_response({"detail": "An earlier change is waiting on a restart. Restart first."}, status=409) + torch_job.update(running=True, result=None, task=asyncio.ensure_future(_torch_run())) + return web.json_response({"started": True}) + + @routes.get(P + "/api/torch/cuda13") + async def _torch_cuda13_state(_req): + return web.json_response({"running": torch_job["running"], "step": torch_swap.step, "result": torch_job["result"]}) + @routes.get(P + "/api/git/status") async def _git_status(req): # `?remote=0` answers from the checkout alone. The fetch is the only part diff --git a/core/tests/test_torch_swap.py b/core/tests/test_torch_swap.py new file mode 100644 index 00000000..8944e713 --- /dev/null +++ b/core/tests/test_torch_swap.py @@ -0,0 +1,158 @@ +"""Switching torch to its CUDA 13 build: download first, install from those files, prove it imports, +or put the old build back. pip is faked: a real run downloads several GB.""" + +import subprocess +import time +from collections import namedtuple + +import pytest + +from core import torch_swap + +Disk = namedtuple("Disk", "total used free") + + +@pytest.fixture +def fake(monkeypatch): + """pip calls recorded; each verb answers from `codes`; versions from `have`; imports from `cuda`.""" + state = {"calls": [], "codes": {}, "cuda": ["13.0"], + "have": {"torch": "2.10.0+cu128", "torchvision": "0.25.0+cu128", "torchaudio": "2.10.0+cu128"}} + + def pip(args, timeout): + state["calls"].append(args) + out = "torch==2.10.0+cu128\nnumpy==2.1.0\nnvidia-cublas-cu12==12.8\ntriton==3.6.0\nfoo @ file:///x\n" if args[0] == "freeze" else "" + return subprocess.CompletedProcess(args, state["codes"].get(args[0], 0), out, "pip said no") + + monkeypatch.setattr(torch_swap, "_pip", pip) + monkeypatch.setattr(torch_swap, "_version", lambda n: state["have"].get(n)) + def imports(names): + state["imported"] = names + return (state["cuda"].pop(0) if state["cuda"] else None), "ImportError: boom" + monkeypatch.setattr(torch_swap, "_imports", imports) + monkeypatch.setattr(torch_swap.shutil, "disk_usage", lambda p: Disk(0, 0, 100 * 1024 ** 3)) + import torch + monkeypatch.setattr(torch.version, "cuda", "12.8") + return state + + +def verbs(state): + return [c[0] for c in state["calls"]] + + +def test_success_downloads_then_installs_the_tagged_pins_from_those_files_only(fake): + r = torch_swap.swap() + assert r["installed"] is True and "13.0" in r["message"] + assert verbs(fake) == ["freeze", "download", "install"] + download, install = fake["calls"][1], fake["calls"][2] + assert "torch==2.10.0+cu130" in download, "a bare torch==2.10.0 is satisfied by +cu128 and pip does nothing" + assert "torchaudio==2.10.0+cu130" in install and "--no-index" in install, "install must not reach the network" + assert torch_swap.step is None + + +def test_constraints_hold_everything_but_torch_triton_and_the_cuda_runtime(fake, monkeypatch, tmp_path): + seen = {} + real = torch_swap._constraints + + def spy(path): + real(path) + seen["text"] = open(path).read() + monkeypatch.setattr(torch_swap, "_constraints", spy) + torch_swap.swap() + assert seen["text"].split() == ["numpy==2.1.0"] + + +def test_a_failed_download_changes_nothing(fake): + fake["codes"]["download"] = 1 + r = torch_swap.swap() + assert r["installed"] is False and "nothing was changed" in r["message"] + assert "install" not in verbs(fake) + + +def test_a_build_that_does_not_import_is_rolled_back_to_the_old_tagged_build(fake): + fake["cuda"] = [None, "12.8"] # new torch fails to import; the restored one works + r = torch_swap.swap() + assert r["installed"] is False and "old one" in r["message"] and "boom" in r["message"] + undo = fake["calls"][-1] + assert "torch==2.10.0+cu128" in undo and "https://download.pytorch.org/whl/cu128" in undo + + +def test_a_failed_rollback_says_it_plainly_with_the_command_to_fix_it(fake): + fake["cuda"] = [None, None] + r = torch_swap.swap() + assert r.get("broken") is True and "--force-reinstall" in r["message"] and "torch==2.10.0+cu128" in r["message"] + + +@pytest.mark.parametrize("torch_v,why", [("2.10.0+cu130", "already"), ("2.10.0+cpu", "not a CUDA build")]) +def test_refuses_what_it_should_not_touch(fake, torch_v, why): + fake["have"]["torch"] = torch_v + r = torch_swap.swap() + assert r["installed"] is False and why in r["message"] and fake["calls"] == [] + + +def test_refuses_without_disk_room(fake, monkeypatch): + monkeypatch.setattr(torch_swap.shutil, "disk_usage", lambda p: Disk(0, 0, 1024 ** 3)) + r = torch_swap.swap() + assert r["installed"] is False and "disk" in r["message"] and fake["calls"] == [] + + +def test_untagged_pypi_torch_rolls_back_to_pypi(fake): + fake["have"] = {"torch": "2.10.0", "torchvision": "0.25.0"} + fake["cuda"] = [None, "12.8"] + torch_swap.swap() + assert fake["imported"] == ["torch", "torchvision"], "a missing torchaudio is not a failed import" + undo = fake["calls"][-1] + assert "torch==2.10.0" in undo and not any("download.pytorch.org" in a for a in undo) + + +# ---- the route: the git lock's rules, and a restart only when the new build is on disk ---- + +from core.tests.test_routes_update import _request, server # noqa: E402,F401 +from core import routes # noqa: E402 + + +def _finish(server): + for _ in range(100): + _, state = _request(server["port"], "GET", "/funpack/api/torch/cuda13") + if not state["running"]: + return state + time.sleep(0.05) + raise AssertionError("the switch never finished") + + +def test_route_starts_a_job_and_restarts_only_after_an_install(server, monkeypatch): + monkeypatch.setattr(routes, "_generation_running", lambda: False) + monkeypatch.setattr(torch_swap, "swap", lambda: {"installed": False, "message": "download failed"}) + status, body = _request(server["port"], "POST", "/funpack/api/torch/cuda13") + assert status == 200 and body == {"started": True}, "answers at once: the download outlives one proxied request" + state = _finish(server) + assert state["result"]["restarting"] is False and state["result"]["message"] == "download failed" + assert server["restarts"] == [] + monkeypatch.setattr(torch_swap, "swap", lambda: {"installed": True, "message": "ok"}) + _request(server["port"], "POST", "/funpack/api/torch/cuda13") + assert _finish(server)["result"]["restarting"] is True + + +def test_a_second_start_while_one_runs_is_refused(server, monkeypatch): + import threading + gate = threading.Event() + monkeypatch.setattr(routes, "_generation_running", lambda: False) + monkeypatch.setattr(torch_swap, "swap", lambda: (gate.wait(5), {"installed": False, "message": "x"})[1]) + _request(server["port"], "POST", "/funpack/api/torch/cuda13") + status, body = _request(server["port"], "POST", "/funpack/api/torch/cuda13") + gate.set() + assert status == 409 and "already running" in body["detail"] + _finish(server) + + +def test_route_refuses_during_a_generation(server, monkeypatch): + monkeypatch.setattr(routes, "_generation_running", lambda: True) + monkeypatch.setattr(torch_swap, "swap", lambda: (_ for _ in ()).throw(AssertionError("must not run"))) + status, body = _request(server["port"], "POST", "/funpack/api/torch/cuda13") + assert status == 409 and "generation" in body["detail"] + + +def test_a_torch_without_cuda_is_refused_before_any_download(fake, monkeypatch): + import torch + monkeypatch.setattr(torch.version, "cuda", None) + r = torch_swap.swap() + assert r["installed"] is False and "no CUDA" in r["message"] and fake["calls"] == [] diff --git a/core/torch_swap.py b/core/torch_swap.py new file mode 100644 index 00000000..482f8345 --- /dev/null +++ b/core/torch_swap.py @@ -0,0 +1,122 @@ +"""Switch this ComfyUI's torch to the CUDA 13 build of the SAME versions -- or leave it as it was. + +Why: on a Blackwell rental, torch 2.10+cu128 sampled int8 models ~2x slower than bf16; the +same torch from the cu130 index made int8 faster than bf16 (2026-10-10). + +Order is the safety. Everything is DOWNLOADED first (a failed download changes nothing), then +installed from those files only, then a fresh interpreter must import all three and report +CUDA 13. If any of that fails the old build is put back. A torch that does not import means a +ComfyUI that does not start, and then there is no UI left to repair it from. + +Pins carry the local tag (`torch==2.10.0+cu130`): a bare `torch==2.10.0` is already satisfied +by `2.10.0+cu128`, so pip does nothing -- which is exactly how the hand-typed version failed. +""" + +import shutil +import subprocess +import sys +import tempfile +from importlib import metadata + +PACKAGES = ("torch", "torchvision", "torchaudio") +INDEX = "https://download.pytorch.org/whl/{tag}" +PYPI = "https://pypi.org/simple" +TARGET = "cu130" +NEED_FREE = 12 * 1024 ** 3 # the cu13 runtime wheels alone are several GB +# Compiled against one torch: worth checking after the swap, not touched by it. +COMPILED = ("xformers", "sageattn3", "flash_attn", "flash-attn") + +step = None # what is happening now, for the UI's progress card + + +def _version(name): + try: + return metadata.version(name) + except metadata.PackageNotFoundError: + return None + + +def _pip(args, timeout): + return subprocess.run([sys.executable, "-m", "pip", *args, "--disable-pip-version-check"], + capture_output=True, text=True, timeout=timeout) + + +def _tail(proc): + return "\n".join(((proc.stderr or "") + (proc.stdout or "")).strip().splitlines()[-6:]) + + +def _pins(have, tag): + """`torch==2.10.0+cu130` for each installed package; no tag means PyPI's own build.""" + return [f"{p}=={v.split('+')[0]}{'+' + tag if tag else ''}" for p, v in have.items()] + + +def _constraints(path): + """Everything else stays at the version it is now. torch's own family and the CUDA runtime + wheels are left free: they are what changes.""" + proc = _pip(["freeze", "--exclude-editable"], 120) + free = PACKAGES + ("triton",) + lines = [l for l in proc.stdout.splitlines() + if "==" in l and not l.split("==")[0].lower().startswith(free + ("nvidia-",))] + with open(path, "w") as f: + f.write("\n".join(lines) + "\n") + + +def _imports(names): + """CUDA version a FRESH interpreter's torch reports, or the error. This process keeps the old + torch in memory until the restart, so asking here would always say the old answer.""" + proc = subprocess.run([sys.executable, "-c", f"import {', '.join(names)}; print(torch.version.cuda)"], + capture_output=True, text=True, timeout=300) + return (proc.stdout.strip().splitlines() or [""])[-1] if proc.returncode == 0 else None, _tail(proc) + + +def swap(timeout=3600): + """-> {"installed": bool (the new build is on disk), "message": str, "rebuild": [names]}.""" + global step + have = {p: _version(p) for p in PACKAGES if _version(p)} + if "torch" not in have: + return {"installed": False, "message": "torch is not installed in this Python, so there is nothing to switch."} + import torch + if not getattr(torch.version, "cuda", None): + return {"installed": False, "message": "This torch has no CUDA (a Mac, or a CPU build): there is no CUDA 13 build to switch to."} + torch_v = have["torch"] + old_tag = torch_v.split("+")[1] if "+" in torch_v else None + if old_tag == TARGET or (old_tag or "").startswith("cu13"): + return {"installed": False, "message": f"torch {torch_v} is already a CUDA 13 build."} + if old_tag and not old_tag.startswith("cu"): + return {"installed": False, "message": f"torch {torch_v} is not a CUDA build ({old_tag}); switching it to CUDA 13 is not something this can judge."} + work = tempfile.mkdtemp(prefix="funpack_torch_") + try: + if shutil.disk_usage(work).free < NEED_FREE: + return {"installed": False, "message": f"Not enough free disk for the download (needs ~{NEED_FREE // 1024 ** 3} GB). Nothing was changed."} + constraints = f"{work}/constraints.txt" + _constraints(constraints) + new = _pins(have, TARGET) + step = f"Downloading {', '.join(new)} (several GB)…" + proc = _pip(["download", "-d", f"{work}/wheels", "--index-url", INDEX.format(tag=TARGET), + "--extra-index-url", PYPI, "-c", constraints, *new], timeout) + if proc.returncode != 0: + return {"installed": False, "message": f"The download failed, so nothing was changed.\n{_tail(proc)}"} + step = "Installing the downloaded build…" + proc = _pip(["install", "--no-index", "--find-links", f"{work}/wheels", "-c", constraints, *new], timeout) + cuda, err = _imports(list(have)) if proc.returncode == 0 else (None, _tail(proc)) + if cuda and cuda.startswith("13"): + rebuild = [n for n in COMPILED if _version(n)] + return {"installed": True, "rebuild": rebuild, + "message": f"torch {', '.join(new)} installed and imports with CUDA {cuda}." + + (f" Check these after the restart (built for the old torch): {', '.join(rebuild)}." if rebuild else "")} + step = "That did not work; putting the old build back…" + back = _pins(have, old_tag) if old_tag else [f"{p}=={v}" for p, v in have.items()] + index = ["--index-url", INDEX.format(tag=old_tag), "--extra-index-url", PYPI] if old_tag else [] + undo = _pip(["install", *index, "-c", constraints, *back], timeout) + restored, _ = _imports(list(have)) + if undo.returncode == 0 and restored: + return {"installed": False, "message": f"The CUDA 13 build failed to install or import, so the old one ({torch_v}) is back.\n{err}"} + return {"installed": False, "broken": True, + "message": (f"The CUDA 13 build failed AND putting {torch_v} back failed. ComfyUI will probably not start " + f"after a restart. Fix it in a terminal:\n {sys.executable} -m pip install --force-reinstall " + f"{' '.join(back)} {' '.join(index)}\n{err}\n{_tail(undo)}")} + except subprocess.TimeoutExpired: + return {"installed": False, "message": "pip took too long and was stopped. Check the terminal: torch may be half-installed."} + finally: + step = None + shutil.rmtree(work, ignore_errors=True)