diff --git a/app/ui/features/generate/index.js b/app/ui/features/generate/index.js index c08d3908..7644f074 100644 --- a/app/ui/features/generate/index.js +++ b/app/ui/features/generate/index.js @@ -2,6 +2,7 @@ import { composer as c } from "../../composer/composer.js"; import { isGenerative, isVideoClip, genUnitId, unitRoot, effFrames, effFps } from "../../shell/scenes.js"; import { QUEUED, RUNNING } from "../../shell/run.js"; +import { jobsRunning, onJobs, stopJobs } from "../../shell/renders.js"; import { inPlace } from "../../shell/place.js"; import { buildInputs, rolesAt } from "./inputs.js"; import { offer } from "../../shell/actions.js"; @@ -201,15 +202,18 @@ export default { let drawn = ""; function draw() { const open = p.project; - const picked = open ? pickedScenes().length : 0, key = `${busy}|${held}|${g.state().phase}|${Boolean(open)}|${picked}`; + const picked = open ? pickedScenes().length : 0, key = `${busy}|${held}|${g.state().phase}|${Boolean(open)}|${picked}|${jobsRunning()}`; if (key === drawn) return; // progress ticks must not rebuild the buttons drawn = key; const working = busy || held || g.state().phase === QUEUED || g.state().phase === RUNNING; // also a run this page did not start const go = c.button.sm({ label: working ? "Generating…" : "▶ Generate", tone: "primary", disabled: working || !open, title: "Generate every scene on the timeline, except left-out scenes, video clips and Auto Montage pieces. An earlier render of a scene stays as a take.", onClick: () => runUnits(unitsOf(p.scenes)) }); const one = c.button.sm({ label: picked > 1 ? `Selected (${picked})` : "Selected", title: "Generate only the picked scenes", disabled: working || !picked, onClick: () => runUnits(unitsOf(pickedScenes())) }); - const stop = c.button.sm({ label: "■ Stop", tone: "danger", title: "Stop the run that is going", onClick: () => { stopped = true; g.cancel(); } }); - stop.node.hidden = !working; + const stop = c.button.sm({ label: "■ Stop", tone: "danger", title: "Stop the run that is going, and any Upscale or Interpolate in progress", onClick: () => { + stopJobs(app.api.stopJob); + if (working) { stopped = true; g.cancel(); } + } }); + stop.node.hidden = !working && !jobsRunning(); put(go.node, one.node, stop.node); } // A run found after a reload belongs to this page's last session: its result still goes on its clips. @@ -244,11 +248,11 @@ export default { }); }); draw(); - const offRun = g.subscribe(draw); + const offRun = g.subscribe(draw), offJobs = onJobs(draw); // the Stop button also shows while an Upscale or Interpolate runs const offApp = app.on((what) => { if (what === "generate.scene" && p.selected) runUnits(unitsOf([p.selected])); else if (what === "generate.selected" && pickedScenes().length) runUnits(unitsOf(pickedScenes())); else draw(); }); const offers = [offer(app, { id: "generate-all", label: "Generate", icon: "▶", run: () => { if (p.project) runUnits(unitsOf(p.scenes)); } }), offer(app, { id: "generate-selected", label: "Generate selected", icon: "▶", run: () => { if (pickedScenes().length) runUnits(unitsOf(pickedScenes())); } })]; - return () => { offRun(); offApp(); offers.forEach((f) => f()); }; + return () => { offRun(); offJobs(); offApp(); offers.forEach((f) => f()); }; }, }; diff --git a/app/ui/features/interpolate/index.js b/app/ui/features/interpolate/index.js index 3111c1d3..c5605f60 100644 --- a/app/ui/features/interpolate/index.js +++ b/app/ui/features/interpolate/index.js @@ -48,7 +48,7 @@ export default { const swapped = await p.editFor(pid, (pr) => swap(pr, media, out) > 0); if (!swapped) c.toast.warn({ text: "Interpolated, but that render was replaced meanwhile, so nothing was swapped." }); else c.toast.good({ text: `Interpolated: the render now plays at ${opts().multiplier} times the frames.` }); - } catch (err) { c.toast.warn({ text: `Interpolate failed: ${err.message}` }); } + } catch (err) { c.toast[err.stopped ? "info" : "warn"]({ text: err.stopped ? "Interpolate stopped. The render stays as it was." : `Interpolate failed: ${err.message}` }); } finally { busy.delete(key); draw(); } } diff --git a/app/ui/features/interpolate/job.test.js b/app/ui/features/interpolate/job.test.js index 3c0a9ab1..09d1ca5e 100644 --- a/app/ui/features/interpolate/job.test.js +++ b/app/ui/features/interpolate/job.test.js @@ -1,7 +1,7 @@ import test from "node:test"; import assert from "node:assert/strict"; import { interpolate, modeLabel } from "./job.js"; -import { derived, gaveUp, onGaveUp, swap, wentThrough } from "../../shell/renders.js"; +import { derived, gaveUp, onGaveUp, swap, wentThrough, waitJob, jobsRunning, onJobs, stopJobs, JobStopped } from "../../shell/renders.js"; const m = (f, sub = "") => ({ filename: f, subfolder: sub, type: "output" }); @@ -46,3 +46,20 @@ test("a step that gave up is announced, so a step waiting behind it can go on wi gaveUp(m("r.mp4"), "upscale"); assert.deepEqual(heard, [["r.mp4", "upscale"]]); }); + +test("Stop: a job being waited for can be stopped, ends as stopped (not failed, not restarted), and the Stop button hears about it", async () => { + const stopped = [], heard = []; + const off = onJobs(() => heard.push(jobsRunning())); + let polls = 0; + const read = async () => (++polls < 2 ? { state: "running" } : { state: "stopped" }); + const waiting = waitJob(read, "j1", async () => { if (polls === 0) await stopJobs((id) => stopped.push(id)); }); + await assert.rejects(waiting, (err) => err instanceof JobStopped && err.stopped === true && err.message === "Stopped."); + assert.deepEqual(stopped, ["j1"]); + assert.deepEqual(heard, [1, 0]); // shown while it waits, hidden when it ends + assert.equal(jobsRunning(), 0); + off(); +}); + +test("Stop: a stopped job that is read once more is still there (no 404 turned into 'ComfyUI restarted')", async () => { + await assert.rejects(waitJob(async () => ({ state: "stopped" }), "j2", async () => {}), (err) => err.stopped === true); +}); diff --git a/app/ui/features/upscale/index.js b/app/ui/features/upscale/index.js index 54ba3899..79536688 100644 --- a/app/ui/features/upscale/index.js +++ b/app/ui/features/upscale/index.js @@ -53,7 +53,10 @@ export default { const swapped = await p.editFor(pid, (pr) => swap(pr, media, out) > 0); if (!swapped) c.toast.warn({ text: "Upscaled, but that render was replaced meanwhile, so nothing was swapped." }); else c.toast.good({ text: "Upscaled: the render now plays the larger video." }); - } catch (err) { gaveUp(media, "upscale"); c.toast.warn({ text: `Upscale failed: ${err.message}` }); } + } catch (err) { + if (err.stopped) c.toast.info({ text: "Upscale stopped. The render stays as it was." }); + else { gaveUp(media, "upscale"); c.toast.warn({ text: `Upscale failed: ${err.message}` }); } + } finally { busy.delete(key); draw(); } } diff --git a/app/ui/shell/api.js b/app/ui/shell/api.js index be8a0185..b7fb2373 100644 --- a/app/ui/shell/api.js +++ b/app/ui/shell/api.js @@ -138,6 +138,8 @@ export const api = { /** The upscale is a background job of FunUI, not a ComfyUI prompt (that would make ComfyUI drop the workflow's loaded models). -> {job_id}. */ upscale: (media, model, rtx) => call("POST", "/api/m/render/upscale", { media: { filename: media.filename, subfolder: media.subfolder || "", type: media.type === "temp" ? "temp" : "output" }, model, ...(rtx ? { rtx_options: rtx } : {}) }), upscaleJob: (id) => call("GET", `/api/m/render/upscale/${encodeURIComponent(id)}`), + /** The Stop button for Upscale and Interpolate jobs (one job list serves both). -> {stopped} */ + stopJob: (id) => call("POST", `/api/m/render/jobs/${encodeURIComponent(id)}/stop`), frameGeneration: () => call("GET", "/api/m/interpolate/frame_generation"), /** Interpolation is a background job of FunUI, not a ComfyUI prompt (that would make ComfyUI drop the workflow's loaded models). -> {job_id}. */ interpolate: (media, options) => call("POST", "/api/m/interpolate/run", { media: { filename: media.filename, subfolder: media.subfolder || "", type: media.type === "temp" ? "temp" : "output" }, options }), diff --git a/app/ui/shell/renders.js b/app/ui/shell/renders.js index c73f37ce..3fdfd73b 100644 --- a/app/ui/shell/renders.js +++ b/app/ui/shell/renders.js @@ -2,21 +2,37 @@ // scene that plays it, and remembering which steps a file has been through so two "always" modes can never feed each other forever. Pure of the page. export const keyOf = (m) => `${m.type || "output"}/${m.subfolder || ""}/${m.filename}`; -/** Wait for a FunUI background job. `read(id)` -> {state: running | done | error, media, detail}. -> the media, or throws a sentence. `sleep` is injectable so a test need not wait. */ +const active = new Set(); // ids of the Upscale / Interpolate jobs this page is waiting for +const listeners = []; +const changed = () => { for (const fn of listeners) { try { fn(); } catch { /* one listener must not stop the others */ } } }; +/** How many jobs the Stop button can stop right now. */ +export const jobsRunning = () => active.size; +/** -> a function that stops listening. */ +export const onJobs = (fn) => { listeners.push(fn); return () => { const i = listeners.indexOf(fn); if (i >= 0) listeners.splice(i, 1); }; }; +/** The Stop button: ask every job this page is waiting for to stop. `stop(id)` is the API call. */ +export const stopJobs = (stop) => Promise.all([...active].map((id) => Promise.resolve(stop(id)).catch(() => {}))); +/** What waitJob throws for a job that was stopped: not a failure, so callers say "stopped" and not "failed". */ +export class JobStopped extends Error { constructor() { super("Stopped."); this.stopped = true; } } + +/** Wait for a FunUI background job. `read(id)` -> {state: running | done | error | stopped, media, detail}. -> the media, or throws a sentence (a JobStopped when it was stopped). `sleep` is injectable so a test need not wait. */ export async function waitJob(read, id, sleep = (ms) => new Promise((r) => setTimeout(r, ms))) { let lost = 0; - for (;;) { - await sleep(1000); - let r; - try { r = await read(id); } catch (err) { - if (err.status === 404) throw new Error("ComfyUI restarted while it was working"); - if (++lost < 30) continue; // a dropped reply is not a failed job - throw new Error("lost contact with ComfyUI"); + active.add(id); changed(); + try { + for (;;) { + await sleep(1000); + let r; + try { r = await read(id); } catch (err) { + if (err.status === 404) throw new Error("ComfyUI restarted while it was working"); + if (++lost < 30) continue; // a dropped reply is not a failed job + throw new Error("lost contact with ComfyUI"); + } + lost = 0; + if (r.state === "stopped") throw new JobStopped(); + if (r.state === "error") throw new Error(r.detail || "failed inside ComfyUI"); + if (r.state === "done") return r.media; } - lost = 0; - if (r.state === "error") throw new Error(r.detail || "failed inside ComfyUI"); - if (r.state === "done") return r.media; - } + } finally { active.delete(id); changed(); } } /** One file backs a whole chain run, so every scene render and ghost that plays `from` gets `to`. -> how many were swapped. */ diff --git a/core/jobs.py b/core/jobs.py index ff73b181..9a243195 100644 --- a/core/jobs.py +++ b/core/jobs.py @@ -8,12 +8,32 @@ from collections import OrderedDict MAX_JOBS = 50 # finished jobs are kept this long for the poll that reads them _jobs: "OrderedDict[str, dict]" = OrderedDict() +_stop: set = set() # ids asked to stop; the job notices between chunks + + +class Stopped(Exception): + """Raised inside a job that was asked to stop; its partial output is removed by whoever was writing it.""" + + +def stop(job: str) -> bool: + """Ask a queued or running job to stop. False when there is no such job or it already ended.""" + entry = _jobs.get(job) + if not entry or entry.get("state") not in ("queued", "running"): + return False + _stop.add(job) + return True + + +def check(job: str) -> None: + """Raises Stopped once `stop(job)` was called: the hook a long job calls between its chunks.""" + if job in _stop: + raise Stopped() def new() -> str: job = uuid.uuid4().hex _jobs[job] = {"state": "queued"} - for old in [k for k, v in _jobs.items() if v.get("state") in ("done", "error")]: + for old in [k for k, v in _jobs.items() if v.get("state") in ("done", "error", "stopped")]: if len(_jobs) <= MAX_JOBS: break del _jobs[old] # only finished ones: a running job's poll must keep finding it @@ -24,14 +44,22 @@ def get(job: str) -> dict | None: return _jobs.get(job) -async def run(job: str, fn, *args, error) -> None: +async def run(job: str, fn, *args, error, stoppable: bool = False) -> None: """Run the blocking `fn(*args)` off the event loop and record how it ended. `error` is the exception class whose message is meant for the person; anything else is reported plainly.""" + if job in _stop: # stopped while it was still queued + _stop.discard(job) + _jobs[job] = {"state": "stopped"} + return _jobs[job] = {"state": "running"} try: - result = await asyncio.to_thread(fn, *args) + result = await asyncio.to_thread(fn, *args, **({"interrupted": lambda: check(job)} if stoppable else {})) _jobs[job] = {"state": "done", **result} + except Stopped: + _jobs[job] = {"state": "stopped"} except error as exc: _jobs[job] = {"state": "error", "detail": str(exc)} except Exception as exc: # noqa: BLE001 _jobs[job] = {"state": "error", "detail": f"{type(exc).__name__}: {exc}"} + finally: + _stop.discard(job) diff --git a/modules/output/interpolate/__init__.py b/modules/output/interpolate/__init__.py index 7749bc6a..ecfd18b0 100644 --- a/modules/output/interpolate/__init__.py +++ b/modules/output/interpolate/__init__.py @@ -43,7 +43,7 @@ def routes(table, base, web): except ValueError as exc: return bad(str(exc)) job = jobs.new() - asyncio.create_task(jobs.run(job, nodes.interpolate, media, body.get("options") or {}, error=nodes.InterpolateError)) + asyncio.create_task(jobs.run(job, nodes.interpolate, media, body.get("options") or {}, error=nodes.InterpolateError, stoppable=True)) return web.json_response({"job_id": job}) @table.get(base + "/jobs/{job}") diff --git a/modules/output/render/__init__.py b/modules/output/render/__init__.py index ca28da41..fa77d72f 100644 --- a/modules/output/render/__init__.py +++ b/modules/output/render/__init__.py @@ -102,9 +102,14 @@ def routes(table, base, web): if not isinstance(media, dict) or not isinstance(media.get("filename"), str) or not media["filename"] or not isinstance(model, str) or not model: return bad("Send {media: {filename, subfolder, type}, model, rtx_options}.") job = jobs.new() - asyncio.create_task(jobs.run(job, upscale.run, media, model, body.get("rtx_options") or "{}", error=(RuntimeError, ValueError, files.ClipError))) + asyncio.create_task(jobs.run(job, upscale.run, media, model, body.get("rtx_options") or "{}", error=(RuntimeError, ValueError, files.ClipError), stoppable=True)) return web.json_response({"job_id": job}) + @table.post(base + "/jobs/{job}/stop") + async def _stop_job(req): + """Stop an Upscale or Interpolate job (the page's Stop button). Both ride this module's job list.""" + return web.json_response({"stopped": jobs.stop(req.match_info["job"])}) + @table.get(base + "/upscale/{job}") async def _upscale_job(req): job = jobs.get(req.match_info["job"]) diff --git a/modules/output/render/tests/test_jobs_routes.py b/modules/output/render/tests/test_jobs_routes.py index e2202232..f2a313a5 100644 --- a/modules/output/render/tests/test_jobs_routes.py +++ b/modules/output/render/tests/test_jobs_routes.py @@ -149,3 +149,63 @@ def test_render_keeps_the_real_frame_rate_of_interpolated_clips_only_when_asked( job = finish(server, f"{BASE}/projects/{proj.id}/render", queued["job_id"]) assert job["state"] == "done" assert files.frame_rate(str(comfy.temp / job["media"]["filename"])) == expected + + +def test_a_running_job_asked_to_stop_ends_stopped_and_stays_pollable(monkeypatch): + import asyncio + import threading + monkeypatch.setattr(jobs, "MAX_JOBS", 3) + started, seen = threading.Event(), [] + + def long_job(interrupted): + started.set() + while True: + interrupted() # between chunks, as stream_file does + seen.append(1) + threading.Event().wait(0.01) + + async def go(): + job = jobs.new() + task = asyncio.create_task(jobs.run(job, long_job, error=RuntimeError, stoppable=True)) + while not started.is_set(): + await asyncio.sleep(0.01) + assert jobs.stop(job) is True + await task + return job + + job = asyncio.run(go()) + assert jobs.get(job) == {"state": "stopped"} # a poll finds it: no 404, so no "ComfyUI restarted" + assert jobs.stop(job) is False and jobs.stop("nope") is False + + +def test_a_job_stopped_before_it_started_never_runs(): + import asyncio + ran = [] + job = jobs.new() + assert jobs.stop(job) is True + asyncio.run(jobs.run(job, lambda interrupted: ran.append(1), error=RuntimeError, stoppable=True)) + assert not ran and jobs.get(job) == {"state": "stopped"} + + +@needs_ffmpeg +def test_a_stopped_stream_leaves_no_half_written_file(comfy): + from core import video_pipe + src = comfy.output / "a.mp4" + make_clip(src, seconds=1) + dst = comfy.output / "out.mp4" + calls = [] + + def stop_after_first_chunk(): + calls.append(1) + if len(calls) > 1: + raise jobs.Stopped() + with pytest.raises(jobs.Stopped): + video_pipe.stream_file(str(src), str(dst), lambda x: x, stop_after_first_chunk) + assert not dst.exists() + + +def test_the_stop_route_answers_whether_it_stopped_anything(comfy, server): + job = jobs.new() + status, body = call(server, "POST", f"{BASE}/jobs/{job}/stop") + assert status == 200 and body == {"stopped": True} + assert call(server, "POST", f"{BASE}/jobs/nope/stop")[1] == {"stopped": False}