Skip to content

fix(generate): don't kill the model worker when cancelling a job that already ended - #341

Open
kevin9327 wants to merge 1 commit into
lightningpixel:devfrom
kevin9327:fix/cancel-finished-job-keeps-worker
Open

kevin9327 wants to merge 1 commit into
lightningpixel:devfrom
kevin9327:fix/cancel-finished-job-keeps-worker

Conversation

@kevin9327

Copy link
Copy Markdown
Contributor

What

Cancelling a generation job that has already ended still kills the model worker:

  • If another generation is running, it fails with Subprocess died during generation.
  • If nothing is running, the warm model is thrown away and the next generation has to load it from scratch.

This is easy to hit from the headless surface: python tools/modly-cli/agent.py legacy cancel <job_id> or workflow-run cancel <run_id> on a job that finished, failed, or was already cancelled. It also happens in the app when Cancel is pressed just as a generation completes, because the poll loop posts the cancel on its next tick, after the backend has marked the job done.

Why it triggers

cancel_job (/generate/cancel/{id}) and its sibling cancel_run (/workflow-runs/{id}/cancel) only update the status while the job is pending or running. The subprocess kill after that runs unconditionally:

if job.status in ("pending", "running"):
    job.status = "cancelled"
    _completed_at[job_id] = time.monotonic()
# Kill the active generator subprocess immediately so inference stops now.
try:
    gen = generator_registry._generators.get(generator_registry._active_id)
    if gen is not None and hasattr(gen, "_proc") and gen._proc and gen._proc.poll() is None:
        gen._proc.kill()

ExtensionProcess keeps its subprocess alive between generations to hold the loaded model, so gen._proc.poll() is None is true for an idle worker too. For an ended job, that subprocess isn't running the job being cancelled.

Fix

In both endpoints, return before the kill when the job has already ended. Cancelling a pending or running job is unchanged. The response stays {"cancelled": true} in every case, so clients see no difference.

Verification

New CancelEndedJobTests in api/tests/test_workflow_runs_lifecycle.py, next to the existing cancel_run test. The registry's active worker is a fake with a live _proc.

  • test_cancelling_a_finished_job_leaves_the_active_worker_alone (cancel_job, status done): fails before, passes after.
  • test_cancelling_a_failed_run_leaves_the_active_worker_alone (cancel_run, status error): fails before, passes after.
  • test_cancelling_a_running_job_still_stops_the_worker and test_cancelling_a_running_run_still_stops_the_worker: pass before and after. A running job is still marked cancelled, and its worker is still killed, detached and marked unloaded. That proves the change doesn't widen into leaving a real cancellation running.

The ended-job tests also assert the job's status is unchanged (done / error), which holds before and after.

Fail-before output, with the two routers reverted and the test kept. It is verbatim except that the local checkout path prefix is replaced with <repo>, and the registry's import-time path banner that follows it is omitted:

test_cancelling_a_failed_run_leaves_the_active_worker_alone (tests.test_workflow_runs_lifecycle.CancelEndedJobTests.test_cancelling_a_failed_run_leaves_the_active_worker_alone) ... FAIL
test_cancelling_a_finished_job_leaves_the_active_worker_alone (tests.test_workflow_runs_lifecycle.CancelEndedJobTests.test_cancelling_a_finished_job_leaves_the_active_worker_alone) ... FAIL
test_cancelling_a_running_job_still_stops_the_worker (tests.test_workflow_runs_lifecycle.CancelEndedJobTests.test_cancelling_a_running_job_still_stops_the_worker) ... ok
test_cancelling_a_running_run_still_stops_the_worker (tests.test_workflow_runs_lifecycle.CancelEndedJobTests.test_cancelling_a_running_run_still_stops_the_worker) ... ok
test_cancel_run_records_completion_so_it_can_be_purged (tests.test_workflow_runs_lifecycle.WorkflowRunJobLifecycleTests.test_cancel_run_records_completion_so_it_can_be_purged) ... ok
test_create_run_purges_terminal_jobs_past_ttl (tests.test_workflow_runs_lifecycle.WorkflowRunJobLifecycleTests.test_create_run_purges_terminal_jobs_past_ttl) ... ok

======================================================================
FAIL: test_cancelling_a_failed_run_leaves_the_active_worker_alone (tests.test_workflow_runs_lifecycle.CancelEndedJobTests.test_cancelling_a_failed_run_leaves_the_active_worker_alone)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "<repo>\api\tests\test_workflow_runs_lifecycle.py", line 163, in test_cancelling_a_failed_run_leaves_the_active_worker_alone
    self.assertFalse(proc.killed)
AssertionError: True is not false

======================================================================
FAIL: test_cancelling_a_finished_job_leaves_the_active_worker_alone (tests.test_workflow_runs_lifecycle.CancelEndedJobTests.test_cancelling_a_finished_job_leaves_the_active_worker_alone)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "<repo>\api\tests\test_workflow_runs_lifecycle.py", line 152, in test_cancelling_a_finished_job_leaves_the_active_worker_alone
    self.assertFalse(proc.killed)
AssertionError: True is not false

----------------------------------------------------------------------
Ran 6 tests in 0.004s

FAILED (failures=2)

After the fix, python -m unittest tests.test_workflow_runs_lifecycle passes all 6 tests (Ran 6 tests in 0.021s / OK).

Whole suite, python -m unittest discover -s tests in api/ (venv with fastapi + python-multipart + httpx):

dev this branch
Whole suite Ran 93 tests / OK (skipped=3) Ran 97 tests / OK (skipped=3)
test_workflow_runs_lifecycle.py 2 pass 6 pass

The 3 skips were already there on dev.

Lint: the repo configures no Python linter, and ESLint ignores api/**. As a spot check, ruff check --isolated reports the same counts on dev and this branch: generation.py 10 → 10, workflow_runs.py 8 → 8, test_workflow_runs_lifecycle.py 4 → 4, so nothing new. No TypeScript touched.

🤖 Generated with Claude Code

… already ended

cancel_job and cancel_run kill the active generator's subprocess so
inference stops at once, but they did so whatever the job's status. For
a job that had already finished, failed or been cancelled, that
subprocess belongs to whatever is generating now, or holds the warm
model: cancelling a finished job (the CLI's `legacy cancel` /
`workflow-run cancel`, or pressing Cancel as a generation completes)
failed the other generation with "Subprocess died during generation" or
forced the next one to reload the model from scratch.

Only kill the subprocess when the job being cancelled was still pending
or running; cancelling an ended job keeps returning {"cancelled": true}.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@iammojogo-sudo

Copy link
Copy Markdown
Contributor

Question: Im curious as if this also has an effect of the ability to free memory after a cancel/kill. Has anyone experience python separation from Modly processes when doing this? If so, that could be tied into this kind of fix.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants