Repository navigation
feat(aorlta): two runs on one node: engine ports chosen free, GPU and directory claims between jobs - #617
Draft
vivekkhandelwal1 wants to merge 6 commits into
Draft
feat(aorlta): two runs on one node: engine ports chosen free, GPU and directory claims between jobs#617vivekkhandelwal1 wants to merge 6 commits into
vivekkhandelwal1 wants to merge 6 commits into
Conversation
… and directory claims between jobs - rollout.port_base (null by default): each engine start picks a free port window outside the ephemeral range and away from this process's other engines, and retries on a taken port; an int keeps the fixed layout (port_base + 2000 i, 20000 before). - The run job claims its GPUs and its two /dev/shm directories under /dev/shm/aorlta-claims and stops (exit 89) when a job Slurm still runs holds one; left-over containers, claims and images are removed only when their job has ended; an image another job may use is kept. sbatch --exclusive=user -w <node> lets your own jobs share a node.
…d-node test_run_job.py: both sides added tests at its end; both kept. Also fixed in the shared-node tests: a slice that did not parse, the missing pytest import, bash run by absolute path when a test empties PATH, and the no-squeue case given a PATH without the host's squeue (the helper rewrites its fake squeue on every call).
…are a node A partition with OverSubscribe=EXCLUSIVE never puts two jobs on a node, whatever --exclusive=user asks. The comment no longer contains the directive's text, which the score test reads the header up to.
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Unresolved directory-claim and cleanup flaws allow concurrent jobs to delete other jobs’ live files through path aliases, overlapping directories, or repeated cleanup.
5 open issues
What changed in this PR
This PR prepares AORLTA runs to share a node by choosing engine ports automatically and coordinating GPU, directory, and container ownership between jobs. Simultaneous two-job execution remains unverified on GPU hardware.
Changes:
- Adds automatic port-window selection and startup retries, with an optional fixed port base.
- Adds node-local resource claims and Slurm-aware container and image cleanup.
- Records engine startup details and expands port, ownership, and cleanup tests.
| File | Description |
|---|---|
aorlta/tests/test_config.py |
Tests automatic and fixed port configuration. |
aorlta/tests/rollout/test_ports.py |
Tests port layouts, availability, and window allocation. |
aorlta/tests/rollout/test_engine_proc.py |
Covers port selection, retries, and engine restarts. |
aorlta/tests/rollout/test_engine_pool.py |
Tests startup and restart notifications. |
aorlta/tests/orchestrator/test_wiring.py |
Tests port configuration and startup records. |
aorlta/tests/orchestrator/test_score.py |
Isolates resource claims in the job test. |
aorlta/tests/orchestrator/test_run_job.py |
Tests claims, Slurm state handling, and cleanup. |
aorlta/src/aorlta/rollout/ports.py |
Implements port-window selection and availability checks. |
aorlta/src/aorlta/rollout/fake_engine.py |
Adds optional socket binding for port tests. |
aorlta/src/aorlta/rollout/engine_proc.py |
Selects ports and retries conflicting engine starts. |
aorlta/src/aorlta/rollout/engine_pool.py |
Reports successful starts through a callback. |
aorlta/src/aorlta/orchestrator/wiring.py |
Applies port settings and writes engine startup records. |
aorlta/src/aorlta/config.py |
Adds and validates rollout.port_base. |
aorlta/jobs/aorlta-run.sbatch |
Adds resource claims, sharing guidance, and job-aware cleanup. |
Review effort: Balanced
…ndows never fragment what is left A window at any base in a span could leave no room: with Random(105), seven windows left none for an eighth on an idle host (2 of 300 seeds). Bases are now AUTO_LOW + 2000 n inside the spans (12 with the default ephemeral range), tried in a random order.
… overlap; cleanup runs once; a kept image is removed by the last job - A directory claim's key is its realpath -m, so /a/run and /a/run/ (or a symlink alias) are one claim. - A live job's directory inside or around the one claimed stops the job (exit 89): either cleanup removes its directory recursively. - The TERM trap's exit ran cleanup a second time, after the claims were released, so it could remove directories another job had claimed in between; cleanup now runs once. - An image a job built or pulled and kept for another job is recorded under the claims directory and tried again by every later cleanup, so the last job to leave removes it; images no job recorded are never removed. - The score test reads the job's header up to its first #SBATCH directive, not the first mention of one.
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Critical path-handling flaws can delete live jobs’ sockets or the shared claim registry, enabling conflicting GPU claims, while claim permissions also block cross-user stale-claim recovery.
3 open issues
5 resolved since last review
Review effort: Balanced
| } | ||
| canon_dir() { realpath -m -- "$1"; } # one spelling per directory: symlinks resolved, no trailing slash, no . or .. | ||
| dir_key() { echo "dir-$(canon_dir "$1" | tr -d '\n' | sha256sum | cut -c1-16)"; } | ||
| claim_dir() { local d; d=$(canon_dir "$1"); claim "directory $d" "$(dir_key "$d")"; } |
| # the GPUs it uses on this node, which are the engines', the trainer ranks' own (with trainer.devices null they share | ||
| # the engines') and the checker servers'. The fakes use no GPU. | ||
| claim_dir "$SYNC_DIR/$RUN_ID" | ||
| [ -n "$SOCK_DIR" ] && claim_dir "$SOCK_DIR" |
| fi | ||
| stamp "taking over $key from job $owner, which has ended" | ||
| fi | ||
| echo "$SLURM_JOB_ID $what" > "$f" && CLAIMED+=("$f") || { echo "[guard] cannot write $f" >&2; exec {fd}>&-; exit 89; } |
This branch has not been deployed
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
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.
Lets two AORLTA runs share a node. Engine ports were fixed (20000 + 2000 i), so a second run's engines collided with the first's.
rollout.port_base(null by default): each engine start picks a free port window, outside the kernel's ephemeral range and away from this process's other engines, testing the same ports TokenSpeed derives and binds. A start that fails on a taken port tries another window (3 attempts). An int keeps the fixed layout./dev/shmdirectories under/dev/shm/aorlta-claims. It stops with exit 89 when a job Slurm still runs holds one, and takes over a claim whose job has ended.sbatch --exclusive=user -w <node>, disjoint GPUs, distinct run ids and socket directories), on a partition that allows node sharing.Tests: the claim, stale-container, image and cleanup logic driven as shell functions against fake
squeue/docker; port windows, retries and the fixed layout; the job's GPU list from real configs.On GPU: a single run of this branch on one MI350X node (Qwen3.5-4B, arithmetic score) chose its engine ports automatically (window at 11193), claimed GPU 0, scored 256 of 256 (243 correct), and left no containers. Not yet verified: two jobs on one node. The cluster partition I have to use for GPU checks never places two jobs on a node (OverSubscribe=EXCLUSIVE), so that check needs a partition that allows sharing.
Stacked on #613. Draft.