Fix/batch placement#54
Conversation
Dragon's dragon.workflows.batch.Manager builds a fixed worker Pool at startup with one Policy per worker (random GPU per worker). Per-task ProcessTemplate.policy is never consulted — placement is decided once at Manager init and tasks land on whatever pool worker happens to grab them off work_q, ignoring any HOST_NAME / HOST_ID requested by the caller. At Perlmutter scale (16 nodes x 4 GPU) this manifested as 1..7 processes per node instead of the requested 4. Detect host-pinned tasks in _build_task_sync (Priority 0) and bypass the batch pool: nproc=1 tasks (the common matey-style case) spawn a dragon.native.process.Process(..., policy=...) directly — one GS round-trip per task; nproc>1 / multi-template tasks fall back to a one-shot ProcessGroup with ignore_error_on_exit=True. Both kinds are tracked in _pinned_tasks (slot.kind discriminator), driven by a throttled (50 ms) _sweep_pinned step inside _monitor_loop, and feed the existing _deliver_batch path so downstream callbacks are unchanged. A background _pinned_reaper_loop calls pg.close() so the monitor never blocks on dragon's close patience. The unpinned batch.process/job/function path is byte-identical to before — no throughput regression for non-pinned tasks. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… contract The feature/edge cherry-pick reshaped _deliver_batch's success branch to surface the executable shell-redirect file paths, but switched from dev's guarded `if stdout: task_desc["stdout"] = stdout` to an unconditional assignment that wrote `None` whenever no path / no captured output / no preset value was available. This broke test_v3_deliver_batch_success_stores_value_and_fires_done, which asserts `"stdout" not in task_desc` for tasks that produced no output. Compute the same 3-way preferred value (redirect file path, then dragon-captured text, then preset), but guard the assignment with `is not None` so success with truly no output leaves the keys absent. Failure branch is unchanged — its diagnostic path always has a stderr worth surfacing. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for host-pinned tasks in the Dragon backend, enabling tasks with specific placement policies to bypass the batch worker pool and run as independent processes or process groups. The implementation includes a throttled monitoring sweep, a background reaper thread for resource cleanup, and updated submission and cancellation logic. Feedback from the review highlights the need for improved thread safety when initializing the reaper thread and incrementing task counters, as well as a recommendation to use public properties instead of protected attributes for state monitoring.
| # kind == "pg" | ||
| pg = slot["pg"] | ||
| try: | ||
| state = pg._state |
There was a problem hiding this comment.
Accessing the protected attribute _state is discouraged when a public property is available. In Dragon's ProcessGroup, the state property provides access to the internal state object.
| state = pg._state | |
| state = pg.state |
| if (self._pinned_reaper_thread is not None | ||
| and self._pinned_reaper_thread.is_alive()): | ||
| return | ||
| self._pinned_reaper_thread = threading.Thread( | ||
| target=self._pinned_reaper_loop, | ||
| name="dragon_pinned_reaper", | ||
| daemon=True, | ||
| ) | ||
| self._pinned_reaper_thread.start() |
There was a problem hiding this comment.
The _ensure_pinned_reaper method is not thread-safe. Since it can be called from multiple threads during task submission (via asyncio.to_thread), there is a risk of starting multiple reaper threads if they concurrently check _pinned_reaper_thread. This initialization should be protected by a lock.
| if (self._pinned_reaper_thread is not None | |
| and self._pinned_reaper_thread.is_alive()): | |
| return | |
| self._pinned_reaper_thread = threading.Thread( | |
| target=self._pinned_reaper_loop, | |
| name="dragon_pinned_reaper", | |
| daemon=True, | |
| ) | |
| self._pinned_reaper_thread.start() | |
| with self._pinned_lock: | |
| if (self._pinned_reaper_thread is not None | |
| and self._pinned_reaper_thread.is_alive()): | |
| return | |
| self._pinned_reaper_thread = threading.Thread( | |
| target=self._pinned_reaper_loop, | |
| name="dragon_pinned_reaper", | |
| daemon=True, | |
| ) | |
| self._pinned_reaper_thread.start() |
| n_total = self._pinned_submitted_total = ( | ||
| getattr(self, "_pinned_submitted_total", 0) + 1) |
There was a problem hiding this comment.
The increment of _pinned_submitted_total is not thread-safe. Because _submit_pinned is called within a thread pool during task submission, concurrent calls can lead to race conditions on this counter. It should be updated within the _pinned_lock block.
with self._pinned_lock:
self._pinned_submitted_total = (
getattr(self, "_pinned_submitted_total", 0) + 1)
n_total = self._pinned_submitted_totalThere was a problem hiding this comment.
Pull request overview
This PR adds a separate execution path in DragonExecutionBackendV3 for tasks whose ProcessTemplate.policy pins placement to a specific host, bypassing Dragon Batch’s pre-sized worker pool so that per-task placement policies are respected (and improving throughput in the author’s benchmark).
Changes:
- Detect host-pinned
Policyon user-suppliedprocess_template(s)and route those tasks through a pinned submission path instead ofdragon.workflows.batch.Batch. - Add pinned-task lifecycle tracking (
_pinned_tasks), polling (_sweep_pinned), and background ProcessGroup reaping to avoid blocking the monitor loop. - Update cancellation/shutdown handling for pinned tasks and adjust
_deliver_batchto avoid writingstdout=None/stderr=Noneon successful completion.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| task = self._task_registry[uid]["description"] | ||
|
|
||
| # Pinned tasks own their own dragon resource (Process or | ||
| # ProcessGroup) and can be cancelled at the dragon level; batch | ||
| # tasks have no per-task cancel hook in dragon.workflows.batch, | ||
| # so the asyncflow is just notified. | ||
| with self._pinned_lock: | ||
| slot = self._pinned_tasks.pop(uid, None) | ||
| if slot is not None: | ||
| kind = slot.get("kind", "pg") | ||
| if kind == "process": | ||
| try: | ||
| slot["process"].terminate() | ||
| except Exception as e: | ||
| self.logger.debug("pinned cancel terminate %s: %s", | ||
| uid, e) | ||
| else: | ||
| try: | ||
| slot["pg"].stop(patience=2.0) | ||
| except Exception as e: | ||
| self.logger.debug("pinned cancel stop %s: %s", uid, e) | ||
| try: | ||
| self._pinned_reap_q.put_nowait(slot["pg"]) | ||
| except Exception: | ||
| pass | ||
|
|
||
| self._callback_func(task, "CANCELED") | ||
| self._cancelled_tasks.add(uid) |
| def _submit_pinned(self, uid: str, pinned_templates: list, | ||
| name: str, timeout: float) -> None: | ||
| """Submit a pinned task. | ||
|
|
||
| Single-template, single-process tasks (the common case for | ||
| host+gpu pinning) are spawned via ``dragon.native.process.Process``: | ||
| one GS round-trip with the policy attached, no Manager process, | ||
| no PGSetProperties / PGAddProcessTemplates handshake — roughly | ||
| two orders of magnitude cheaper than ``ProcessGroup.init()``. | ||
|
|
||
| Multi-process or multi-template tasks fall back to a one-shot | ||
| ``ProcessGroup`` (slower but supports per-rank Policy). | ||
|
|
||
| Returns None either way — pinned tasks are tracked in | ||
| ``self._pinned_tasks`` rather than ``self._monitored_batches``. | ||
| """ | ||
| # Single template + nproc == 1 is the cheap path (Process API). | ||
| if len(pinned_templates) == 1 and pinned_templates[0][0] == 1: | ||
| self._submit_pinned_process(uid, pinned_templates[0][1]) | ||
| else: | ||
| self._submit_pinned_pg(uid, pinned_templates, name, timeout) | ||
| n_total = self._pinned_submitted_total = ( | ||
| getattr(self, "_pinned_submitted_total", 0) + 1) | ||
| if n_total == 1 or n_total % 256 == 0: | ||
| with self._pinned_lock: | ||
| in_flight = len(self._pinned_tasks) | ||
| self.logger.info( | ||
| "pinned submit %d in-flight=%d uid=%s", | ||
| n_total, in_flight, uid) | ||
| return None | ||
|
|
||
| def _submit_pinned_process(self, uid: str, template) -> None: | ||
| """Cheap path: single dragon Process with explicit Policy.""" | ||
| proc = Process( | ||
| target = template.target, | ||
| args = template.args, | ||
| kwargs = template.kwargs, | ||
| cwd = template.cwd, | ||
| env = template.env, | ||
| stdin = template.stdin, | ||
| stdout = template.stdout, | ||
| stderr = template.stderr, | ||
| policy = template.policy, | ||
| options= template.options, | ||
| ) | ||
| proc.start() |
| # kind == "pg" | ||
| pg = slot["pg"] | ||
| try: | ||
| state = pg._state | ||
| except Exception as e: | ||
| self.logger.debug("pinned %s state query failed: %s", | ||
| uid, e) | ||
| continue | ||
| if state.name not in ("IDLE", "ERROR"): | ||
| continue |
| # Detect host-pinned templates upfront. dragon.workflows.batch's | ||
| # worker pool ignores per-task ProcessTemplate.policy (see | ||
| # _is_pinned_policy), so any task that asks for a specific host | ||
| # bypasses Batch and runs as its own ProcessGroup whose member | ||
| # template carries the user's Policy. | ||
| pinned_templates = None # list[(nproc, ProcessTemplate)] or None | ||
| if process_templates_config is not None: | ||
| if any(_is_pinned_policy(tc.get("policy")) | ||
| for _, tc in process_templates_config): | ||
| pinned_templates = [ | ||
| (nranks, | ||
| ProcessTemplate( | ||
| target, **{**tc, "args": task_args, | ||
| "kwargs": task_kwargs})) | ||
| for nranks, tc in process_templates_config | ||
| ] | ||
| elif process_template_config is not None: | ||
| if _is_pinned_policy(process_template_config.get("policy")): | ||
| pinned_templates = [( | ||
| 1, | ||
| ProcessTemplate( | ||
| target, **{**process_template_config, | ||
| "args": task_args, | ||
| "kwargs": task_kwargs}), | ||
| )] | ||
| # Auto-build (mpi/function/executable) paths attach no policy and | ||
| # are therefore never pinned. | ||
|
|
||
| # Single decision tree - no redundant checks | ||
| if pinned_templates is not None: | ||
| # Priority 0: pinned task — bypass Batch's worker pool. | ||
| batch_task = self._submit_pinned(uid, pinned_templates, | ||
| name, timeout) | ||
| execution_mode = "pinned" | ||
|
|
…G race _submit_pinned_process referenced template.kwargs, but ProcessTemplate folds kwargs into argdata for python targets and exposes no .kwargs attribute (and kwargs is meaningless for executable targets). Branch on template.is_python and recover via get_original_python_parameters() for python tasks; use empty dict otherwise. For fast-completing pinned tasks (e.g. /bin/true), _sweep_pinned delivered DONE via call_soon_threadsafe while _build_chunk was still running in the worker thread; the post-build dispatch loop in submit_tasks then fired RUNNING for every task, clobbering the already- delivered terminal state. Tasks ended up stuck in RUNNING after their future had already resolved. Fix: fire RUNNING from inside _build_task_sync (pinned branch) before _submit_pinned starts the process, and skip the dispatch-loop RUNNING for pinned tasks. FIFO ordering of call_soon_threadsafe guarantees RUNNING precedes DONE on the loop. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Closing — superseded by Dragon 0.14.0. This PR added a host-pinned bypass to work around a bug in Dragon's batch worker pool: I checked the current Dragon release (0.14.0, which
So the bypass this branch implemented is no longer needed on 0.14.0 — per-task host pinning is honoured natively. (Verified by code inspection of the 0.14.0 wheel; worth a quick on-cluster confirmation, but the workaround should not be carried forward.) The per-rank function stdout/stderr capture and the submit chunking this branch also carried are tracked separately in #51 and #53. Closing. |
This PR creates a secondary code path for tasks which are pinned to a specific host and gpu. Tasks which are not pinned re handled by BATCH as before.
Surprisingly, this code path improves throughput (for my specific test case):
bypassis using this PR. But more importantly, the placement policies are now respected.