Combined review: WebSocket execution stack - #480
Closed
hermabr wants to merge 19 commits into
Closed
Conversation
- Replace the FastAPI/HTTP API with a versioned WebSocket worker protocol - Remove HTTP client dependencies and coordinate local worker pools directly
- Use ExecutionCoordinator directly for worker backend coordination - Update protocol discriminators and tests
- Simplify hello messages and server connection handling - Remove obsolete protocol version validation test
lease_job now blocks on the coordinator's condition until a job is available or the run is done, deleting the "wait" token and the server's 1s wake poll (and with it a missed-wakeup race). The lock and wake primitives merge into a single Condition notified by mutators. The strictly sequential worker channel drops its envelopes: workers send a bare JobResult instead of ResultMessage, hellos validate directly, and worker_message_adapter goes away. lease_job also stops releasing prior leases per request - connection teardown (worker_lost) is the single release path. client.py folds into its only caller, worker_loop. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Idle worker handlers block in lease_job without touching the socket, so closing their connections never woke them: an interrupted run waited out the 10s drain timeout and leaked non-daemon handler threads that hang interpreter exit. Server shutdown now marks the coordinator finished first, which wakes every waiting handler, and the run loop closes the server before joining worker pools so pool stops only reap already-exiting workers. With that, the explicit stop/assign envelopes are unnecessary: the server sends Job frames directly and a normal WebSocket close is the stop signal, removing AssignMessage, StopMessage, their kind discriminators, and the server message adapter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Only one job can be outstanding per connection and the server retains the Job it sent until the matching result arrives, so per-member UUID lease IDs carry no information. Key the coordinator's running table by object id (unique among running jobs because nodes are popped from ready) and send Job.artifacts as a plain list of ArtifactSpec. This deletes JobMember, RunningJob.lease_id, FailedJob.lease_id, UUID generation, and the lease-centric log details. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The dynamic scaler (scale thread, polling interval, failed-thread bookkeeping, restart budget, idle timeouts forcing workers to exit and be recreated) was built for the HTTP era when idle workers polled. An idle WebSocket worker is just a blocked thread, so start max_workers threads once, keep them connected with no idle timeout, and let the server closing the connection end them. A crashed worker thread now fails the run immediately: job failures are already caught inside the worker loop, so a thread-level crash means a furu bug, not a job problem. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Three failure budgets overlapped: coordinator per-object retries, worker consecutive-job failures, and pool failed-worker restarts. The worker threshold (5) rarely fired under the coordinator's retry default (3), and when it did the worker exited zero, Slurm recorded a normal completion, and the pool replaced it without consuming any restart budget. The Slurm restart budget itself could never govern anything because the scale loop fails the whole run on the first failed Slurm state. Keep two policies: the coordinator owns job failure and retry; the pool owns worker/allocation failure. Remove max_consecutive_failures (worker loop, CLI flag, sbatch script plumbing), the unused Slurm restart budget, the now-unused worker.max_failed_restarts config, and the never-read SlurmWorkerPool._server_url/_auth_token fields. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Aggregate the full WebSocket execution stack into one PR against
mainso the complete change can be reviewed together. This branch contains the exact stacked history through #479, without extra code changes or synthetic merge commits.Includes:
Testing
uv run pytest -n 0 tests/test_worker_execution_coordinator.py tests/test_slurm_backend.py tests/test_worker_subprocess.py tests/test_config.py tests/test_provenance_wiring.py— 145 passeduv run ruff check .— passeduv run ty check— passed