diff --git a/src/campaign/README.md b/src/campaign/README.md index 7e129c9..eea9d34 100644 --- a/src/campaign/README.md +++ b/src/campaign/README.md @@ -3,7 +3,7 @@ RADICAL asyncflow-native orchestrator for multi-workflow HPC campaigns. Runs concurrent replicas of heterogeneous workflows inside a single `asyncio` event loop backed by `radical.asyncflow`, with priority-based scheduling, -sliding-window concurrency caps, resource-pool gating, and workflow-driven +sliding-window concurrency caps, resource-pool gating, and adaptive cascading dependency signalling. --- @@ -29,19 +29,21 @@ All user workflows subclass `BaseWorkflow`. class BaseWorkflow: workflow_id: str = "base" # unique prefix for replica IDs - def __init__(self, config, on_ready, asyncflow, policies, engine_dragon): ... + def __init__(self, config, _cm, _group_name, asyncflow, policies, engine_dragon): ... - async def run(self, replica_id: str): ... # entry point (override run OR start) - async def on_replica_done(self, replica_id, cm, final_state): ... # optional hook - async def _signal_ready(self): ... # call from run() to unblock dependents + async def run(self, replica_id: str): ... # entry point (override run OR start) + async def on_replica_done(self, replica_id, cm, final_state): ... # optional hook + async def _signal_done(self): ... # broadcast signal to all dependent groups + async def _trigger_dependent(self, name, replicas=1): ... # explicit activation of a named group ``` -The CM injects five objects at construction time: +The CM injects six objects at construction time: | Injected attribute | Type | Purpose | |--------------------|------|---------| | `self.config` | `dict` | per-group config section (CM scheduling keys stripped) | -| `self._on_ready` | async callable | calls `cm.signal_ready(group_name)` | +| `self._cm` | `AsyncCampaignManager` | reference to the running CM (`None` in unit tests without a CM) | +| `self._group_name` | `str` | name of this replica's group (used by `_signal_done`) | | `self.asyncflow` | `WorkflowEngine` | shared `radical.asyncflow` engine | | `self.policies` | `list[Policy]` | one Dragon `Policy` per assigned GPU (empty on concurrent backend) | | `self.engine_dragon` | backend handle | Dragon backend; `None` on concurrent | @@ -66,13 +68,13 @@ has: | Field | Meaning | |-------|---------| -| `replicas` | total replicas to complete | +| `replicas` | total replicas to complete (omit / set to 0 for dependent groups) | | `max_replicas` | sliding-window concurrency cap (defaults to `replicas` if 0) | | `min_replicas` | minimum guaranteed concurrent slots (Pass 1 of scheduler) | | `priority` | higher → scheduled first | | `required_cpus` | CPU cores reserved from the pool while a replica runs | | `required_gpus` | GPU slots reserved from the pool while a replica runs | -| `dependencies` | list of group names that must signal ready first | +| `dependencies` | upstream groups; used to route `_signal_done()` and gate scheduling | | `dependency_threshold` | count-based fallback: N finished replicas in a dep group counts as "ready" (default 1) | ### ResourcePool @@ -104,10 +106,80 @@ replica finishes. --- +## Adaptive cascading dependency model + +### Two group modes + +A workflow group is **independent** or **dependent**, controlled entirely by +the config — no workflow code changes required to switch between modes. + +**Independent** — `replicas: N` present; group starts immediately on `cm.start()`. + +**Dependent** — `replicas` omitted (defaults to 0); group stays inactive until an +upstream replica signals the CM. Each signal adds more replicas to the queue; +signals can repeat throughout the lifetime of the upstream run. + +### Two signalling methods + +#### `_signal_done()` — broadcast, topology-driven + +```python +await self._signal_done() +``` + +Called from `run()` to indicate that this iteration has produced output. +The CM auto-routes the signal to **every group** that lists the caller's group +in its `dependencies` config field, adding +1 replica to each. The caller +does not need to know downstream group names — the pipeline topology lives +entirely in the config. + +Use this for **data-driven fan-out**: one upstream replica fires once per +result, and the CM decides which downstream groups get a new replica based on +the config graph. + +``` +md ──signal_done()──► CM routes ──► miniapps (+1 replica per signal) +``` + +#### `_trigger_dependent(name, replicas=N)` — explicit, named + +```python +await self._trigger_dependent("downstream_group", replicas=1) +``` + +Called from `run()` or `on_replica_done()` when the upstream workflow +decides—based on its own logic—to start a specific number of downstream +replicas. Each call is additive: calling it again queues more replicas. +If the group was already marked done, it is re-opened for scheduling. + +Use this when the **calling workflow knows** the target name and controls +exactly how many replicas to spawn per event (e.g. one inference result +triggers exactly one downstream job). + +``` +inference ──_trigger_dependent("dummy", replicas=1)──► dummy (+1 per result) +``` + +### Scheduler re-runs on every signal + +Every call to `_signal_done()` or `_trigger_dependent()` increments the +target group's `replicas` counter and immediately re-runs the two-pass +scheduler. If resources are available the new replica starts at once; +otherwise it queues until resources free up. + +### Campaign completion + +Groups registered with `replicas=0` (dependent groups that were never +triggered) are **excluded** from the all-done check. The campaign completes +when all groups that were actually triggered have finished, plus all +independent groups are done. + +--- + ## Scheduling model The CM runs a **two-pass greedy scheduler** on every state change (replica -start, replica finish, `signal_ready`, `add_replicas`): +start, replica finish, `signal_done`, `trigger_dependent`): 1. **Pass 1** — guarantee `min_replicas` concurrent slots for all eligible groups, highest priority first. @@ -120,43 +192,76 @@ and a WARNING is emitted. A group is **eligible** when every dependency group is **ready**: -- **Workflow-driven** (preferred): the dependency called - `await self._signal_ready()` at any point during execution. This fires - immediately, regardless of how many replicas have finished. +- **Workflow-driven** (preferred): a dependency group called `_signal_done()` + at any point during execution (`group.ready = True`). - **Count-based fallback**: `dep.finished_replicas >= dep_threshold` (default 1). --- ## Authoring a workflow +### Independent workflow + ```python from src.campaign import BaseWorkflow -class MyWorkflow(BaseWorkflow): - workflow_id = "my_wf" +class SimWorkflow(BaseWorkflow): + workflow_id = "sim" async def run(self, replica_id: str) -> None: # self.config — dict forwarded from YAML workflow section # self.asyncflow — shared WorkflowEngine # self.policies — Dragon Policy list (empty on concurrent backend) - await do_simulation(self.asyncflow, self.config) + result = await do_simulation(self.asyncflow, self.config) + + # Signal the CM every time a result is ready. + # CM auto-routes +1 replica to every group in config's dependencies. + await self._signal_done() +``` + +### Dependent workflow (topology-driven via _signal_done) - # Unblock dependent groups immediately (does not wait for run() to return). - await self._signal_ready() +No changes needed in the dependent workflow itself — it just runs normally. +The CM starts it when an upstream `_signal_done()` fires. - await do_training(self.asyncflow, self.config) +```yaml +# config.yaml +workflows: + sim: + replicas: 4 # independent: starts immediately + ... + + analysis: + dependencies: [sim] # dependent: starts at replicas=0; sim's _signal_done() adds replicas + ... # no "replicas:" key — the count comes from signals at runtime +``` + +### Dependent workflow (explicit via _trigger_dependent) + +Use when this workflow decides the count and the target name based on its +execution logic (e.g. a quality filter on results). + +```python +class InferenceWorkflow(BaseWorkflow): + workflow_id = "inference" + + async def run(self, replica_id: str) -> None: + results = await run_inference(self.asyncflow, self.config) + for r in results: + if r.quality > THRESHOLD: + # Explicitly queue 1 more replica of the downstream group. + await self._trigger_dependent("downstream", replicas=1) async def on_replica_done(self, replica_id, cm, final_state): - # Optional: called after run() returns or raises. - # final_state is "done" or "failed". - if final_state == "done": - await cm.add_replicas("downstream_group", n=1) + # on_replica_done fires after run() returns; useful for teardown + # that should happen once per replica (e.g. releasing shared services). + ... ``` Rules: - Define **either** `run()` or `start()` — not both. -- `_signal_ready()` is idempotent at the CM level; call it as many times as - needed, only the first call has effect. +- Both `_signal_done()` and `_trigger_dependent()` are no-ops when no CM was + injected (safe to call in unit tests). - `on_replica_done` may be `async def` or `def`; the CM handles both. - Do **not** call `asyncflow.shutdown()` from within a replica — the engine is owned by the caller and shut down after `cm.close()`. @@ -168,11 +273,11 @@ Rules: ```python from src.campaign import AsyncCampaignManager -WORKFLOW_REGISTRY = {"my_wf": MyWorkflow, "downstream": DownstreamWorkflow} +WORKFLOW_REGISTRY = {"sim": SimWorkflow, "analysis": AnalysisWorkflow} cm = AsyncCampaignManager.from_config(config, WORKFLOW_REGISTRY) -await cm.start() # schedules all groups with satisfied dependencies -await cm.wait() # blocks until every group is done +await cm.start() # schedules all groups with replicas > 0 +await cm.wait() # blocks until every triggered group is done await cm.close() # releases CM resources (does NOT shut down asyncflow) # caller shuts down asyncflow separately, after telemetry is stopped ``` @@ -201,7 +306,6 @@ await asyncflow.shutdown() ```yaml # ── Cluster resource budget ────────────────────────────────────────────────── -# Set either to 0 to disable tracking (unlimited). resources: total_cpus: 128 total_gpus: 4 @@ -210,26 +314,57 @@ resources: engine: dragon # "dragon" or "concurrent" (falls back to concurrent if Dragon unavailable) # ── Workflow groups ────────────────────────────────────────────────────────── +# +# Two modes — controlled by whether 'replicas' is present: +# +# Independent (replicas: N): +# Group starts immediately on cm.start(). +# +# Dependent (no replicas / replicas: 0): +# Group starts at 0 replicas; stays inactive until an upstream replica +# calls _signal_done() or _trigger_dependent(). Each call is additive — +# the upstream workflow decides when and how many replicas to add based on +# its own execution logic. Calls can repeat across the lifetime of one +# upstream replica (e.g. once per iteration, once per result). +# +# To switch a dependent group to independent: add 'replicas: N' and remove +# 'dependencies'. No workflow code needs to change. + workflows: - ddsim: - replicas: 8 - min_replicas: 2 # guaranteed concurrent minimum - max_replicas: 4 # sliding-window cap - priority: 5 - required_cpus: 20 # cores held while one replica runs - required_gpus: 0 # ddsim is CPU-only - dependencies: [] - dependency_threshold: 1 # unused when no deps + md: + replicas: 2 # independent: starts immediately + min_replicas: 1 + max_replicas: 2 + priority: 10 + required_cpus: 4 + required_gpus: 1 + # Each iteration calls _signal_done() → CM routes +1 replica to miniapps. + + miniapps: + priority: 8 + min_replicas: 1 + max_replicas: 2 + required_cpus: 4 + required_gpus: 1 + dependencies: [md] # dependent: no replicas key → starts at 0 + # md's _signal_done() adds replicas at runtime inference: - replicas: 16 - min_replicas: 1 - max_replicas: 4 - priority: 10 # higher → scheduled before ddsim overflow slots - required_cpus: 32 - required_gpus: 1 # one GPU per inference replica - dependencies: [ddsim] - dependency_threshold: 1 + replicas: 8 # independent + min_replicas: 1 + max_replicas: 4 + priority: 6 + required_cpus: 4 + required_gpus: 1 + # on_replica_done calls _trigger_dependent("dummy", replicas=1) per result. + + dummy: + priority: 5 + min_replicas: 2 + max_replicas: 4 + required_cpus: 4 + required_gpus: 0 + dependencies: [inference] # dependent: inference triggers via _trigger_dependent ``` Config keys consumed by the CM and stripped before forwarding to `workflow.config`: @@ -249,11 +384,12 @@ min_replicas max_replicas required_cpus required_gpus |--------|-------------| | `from_config(config, registry, asyncflow=None, engine_dragon=None)` | Build from YAML config dict + `{name: cls}` registry | | `register_group(name, cls, ...)` | Register a workflow group | -| `start()` | Schedule all eligible groups; creates the shared asyncflow engine if not pre-built | -| `wait(timeout=None)` | Async-block until all groups complete; returns `True` on success | +| `start()` | Schedule all groups with `replicas > 0`; creates the shared asyncflow engine if not pre-built | +| `wait(timeout=None)` | Async-block until all triggered groups complete; returns `True` on success | | `close()` | Release CM resources (does NOT shut down asyncflow) | -| `signal_ready(group_name)` | Mark a group ready; unblock its dependents | -| `add_replicas(group_name, n)` | Dynamically extend a group (capped at `configured_replicas`) | +| `signal_done(group_name)` | Called by `_signal_done()`; adds +1 replica to every group that lists `group_name` in `dependencies` | +| `trigger_dependent(name, replicas, config=None)` | Called by `_trigger_dependent()`; adds `replicas` to the named group and re-opens it if done | +| `add_replicas(group_name, n)` | Dynamically extend a group up to its `configured_replicas` cap | | `status()` | Snapshot dict of all group states + `"resources"` key | | `stats()` | Per-group `WorkflowStats(replicas_started, replicas_finished)` | @@ -261,7 +397,7 @@ min_replicas max_replicas required_cpus required_gpus | Parameter | Default | Meaning | |-----------|---------|---------| -| `replicas` | `1` | Total replicas | +| `replicas` | `1` | Total replicas (0 for dependent groups) | | `min_replicas` | `0` | Guaranteed concurrent minimum | | `max_replicas` | `0` | Sliding-window cap (0 → equals `replicas`) | | `priority` | `0` | Scheduling priority (higher = first) | @@ -282,9 +418,11 @@ plain blocking calls. Same `from_config` / `register_group` / `start` / |--------------------|-------------| | `workflow_id` | class-level string; used as replica ID prefix | | `config` | dict forwarded from the group's config section (CM keys stripped) | +| `_cm` | reference to the running `AsyncCampaignManager` (`None` if no CM) | +| `_group_name` | name of this group in the CM (used by `_signal_done`) | | `asyncflow` | shared `WorkflowEngine` | | `policies` | list of Dragon `Policy` objects for assigned GPUs (empty on concurrent) | | `engine_dragon` | Dragon backend handle (`None` on concurrent) | -| `_on_ready` | injected async callable; invoke via `_signal_ready()` | -| `_signal_ready()` | fires the on-ready callback; no-op if not injected | +| `_signal_done()` | broadcast signal to CM; adds +1 replica to all downstream groups; no-op without a CM | +| `_trigger_dependent(name, replicas)` | explicitly queue N replicas of a named group; no-op without a CM | | `on_replica_done(replica_id, cm, state)` | post-replica hook; override as needed | diff --git a/src/campaign/campaign_manager.py b/src/campaign/campaign_manager.py index bb21943..ad2e4cb 100644 --- a/src/campaign/campaign_manager.py +++ b/src/campaign/campaign_manager.py @@ -14,9 +14,10 @@ Pass 2 — fill remaining capacity up to ``max_replicas`` (highest priority first). -A group becomes *eligible* when each dependency group has either explicitly -called ``signal_ready()`` (workflow-driven) or has ``dep_threshold`` or more -finished replicas (count-based fallback, default 1). +A group becomes *eligible* either when a parent workflow calls +``_trigger_dependent()`` (explicit activation) or when each dependency group +has ``dep_threshold`` or more finished replicas / has called ``_signal_done()`` +(count-based fallback, default 1). Usage ----- @@ -122,14 +123,16 @@ class BaseWorkflow: def __init__( self, config: Optional[dict] = None, - on_ready: Optional[object] = None, + _cm: Optional[object] = None, + _group_name: Optional[str] = None, asyncflow: Optional[object] = None, policies: Optional[list] = None, engine_dragon: Optional[object] = None, ) -> None: self.config = config - # Async callable injected by AsyncCampaignManager; invoke via _signal_ready(). - self._on_ready = on_ready + # AsyncCampaignManager reference injected at construction. + self._cm = _cm + self._group_name = _group_name # Shared WorkflowEngine injected by AsyncCampaignManager (optional). self.asyncflow = asyncflow # One Dragon Policy per assigned GPU, injected by AsyncCampaignManager. @@ -138,20 +141,34 @@ def __init__( # Dragon backend handle for routing tasks to specific backends. self.engine_dragon: Optional[object] = engine_dragon - async def _signal_ready(self) -> None: + async def _trigger_dependent( + self, + name: str, + replicas: int = 1, + **kwargs, + ) -> None: + """ + Tell the CM to activate a dependent workflow group with *replicas* replicas. + + The group must already be registered (via config or register_group) with + ``replicas=0``. Calling this mid-run is the canonical way to start a + workflow that depends on data produced by the current workflow. + No-op when no CM was injected. """ - Signal the CM that this workflow has produced enough data for its - dependents to start. + if self._cm is not None: + await self._cm.trigger_dependent(name, replicas=replicas, **kwargs) - Calls the ``on_ready`` coroutine injected by the CM at construction. - No-op if no callback was provided. + async def _signal_done(self) -> None: """ - import asyncio + Signal the CM that this workflow has finished producing data for + its dependents (count-based dependency fallback). - if self._on_ready is not None: - result = self._on_ready() - if asyncio.iscoroutine(result): - await result + Marks this group ``ready=True`` in the CM, which unblocks any groups + that list this one as a dependency with ``dep_threshold > finished_replicas``. + No-op when no CM was injected. + """ + if self._cm is not None and self._group_name is not None: + await self._cm.signal_done(self._group_name) def run(self, replica_id: str) -> None: """ @@ -387,10 +404,14 @@ def from_config( cm._log.warning(f"from_config: no class registered for {name!r} — skipping") continue + # Groups with dependencies default to replicas=0 — they remain + # inactive until a parent workflow calls _trigger_dependent(). + has_deps = bool(wf_cfg.get("dependencies", [])) + default_replicas = 0 if has_deps else 1 cm.register_group( name=name, workflow_class=wf_class, - replicas=int(wf_cfg.get("replicas", 1)), + replicas=int(wf_cfg.get("replicas", default_replicas)), dependencies=list(wf_cfg.get("dependencies", [])), dep_threshold=int(wf_cfg.get("dependency_threshold", 1)), priority=int(wf_cfg.get("priority", 0)), @@ -637,22 +658,65 @@ async def close(self) -> None: # Monitoring # ------------------------------------------------------------------ - async def signal_ready(self, group_name: str) -> None: + async def signal_done(self, group_name: str) -> None: """ - Mark *group_name* as having produced enough data for its dependents. + Signal that *group_name* has produced output and downstream work should run. - Called by a workflow via ``self._signal_ready()`` (injected at construction). - Once set, the flag overrides the ``dep_threshold`` replica-count check so - dependent groups are unblocked immediately. + Called by a workflow via ``self._signal_done()`` — can be called multiple + times per replica (e.g. once per iteration). Each call queues **1 more + replica** in every group that lists *group_name* in its ``dependencies``. + The CM routes the signal to the right downstream groups automatically, so + the calling workflow does not need to know their names. - Idempotent — subsequent calls for the same group are no-ops. + Also sets ``group.ready = True`` to satisfy count-based dep checks. """ async with self._lock: group = self._groups.get(group_name) - if group is None or group.ready: + if group is None: return group.ready = True - self._log.info(f"Group {group_name!r} signaled ready — unblocking dependents") + dependents = [g for g in self._groups.values() if group_name in g.dependencies] + for dep in dependents: + dep.replicas += 1 + dep.configured_replicas += 1 + if dep.status == "done": + dep.status = "pending" + if dependents: + self._log.info( + f"{group_name!r} signaled done → +1 replica for {[d.name for d in dependents]}" + ) + await self._schedule() + + async def trigger_dependent( + self, + name: str, + replicas: int = 1, + config: Optional[dict] = None, + ) -> None: + """ + Queue *replicas* more runs of the dependent group *name*. + + Called by a parent workflow (via ``self._trigger_dependent()``) each + time its execution logic decides to launch downstream work. May be + called multiple times — each call adds *replicas* to the group's + total and re-opens the group for scheduling if it had already finished. + + The group must be pre-registered (via config or ``register_group``). + """ + async with self._lock: + group = self._groups.get(name) + if group is None: + self._log.warning(f"trigger_dependent: group {name!r} not registered — ignoring") + return + group.replicas += replicas + group.configured_replicas += replicas + if group.status == "done": + group.status = "pending" + if config: + group.group_config = {**(group.group_config or {}), **config} + self._log.info( + f"trigger_dependent: {name!r} +{replicas} replicas (total={group.replicas})" + ) await self._schedule() async def add_replicas(self, group_name: str, n: int = 1) -> None: @@ -739,7 +803,11 @@ def _can_start_locked(self, group: _GroupInfo) -> bool: return False if group.started_count >= group.replicas: return False - if group.running_count >= group.max_replicas: + # max_replicas == 0 means "no explicit cap — use replicas count". + # This handles dependent groups registered with replicas=0 that later + # receive replicas via trigger_dependent() or signal_done(). + effective_max = group.max_replicas if group.max_replicas > 0 else group.replicas + if group.running_count >= effective_max: return False if not self._deps_satisfied_locked(group): return False @@ -816,7 +884,7 @@ def _schedule_locked(self) -> list[tuple["_GroupInfo", int]]: for g in eligible: if ( g.started_count < g.replicas - and g.running_count < g.max_replicas + and g.running_count < (g.max_replicas if g.max_replicas > 0 else g.replicas) and self._deps_satisfied_locked(g) and not self._resources.can_fit(g.required_cpus, g.required_gpus) ): @@ -901,7 +969,8 @@ async def _run_replica(self, group: _GroupInfo, replica_idx: int) -> None: } wf = group.workflow_class( config=replica_config, - on_ready=lambda: self.signal_ready(group.name), + _cm=self, + _group_name=group.name, asyncflow=self._asyncflow, policies=policies, engine_dragon=self._engine_dragon, @@ -1005,7 +1074,13 @@ async def _on_replica_finished(self, group: _GroupInfo, replica_id: str) -> None await self._schedule() async with self._lock: - all_done = bool(self._groups) and all(g.status == "done" for g in self._groups.values()) + # Groups registered with replicas=0 (triggered dependents not yet activated) + # are excluded from the completion check — they only count once triggered. + all_done = ( + bool(self._groups) + and all(g.status == "done" or g.replicas == 0 for g in self._groups.values()) + and any(g.replicas > 0 for g in self._groups.values()) + ) if all_done: self._all_done.set() @@ -1131,10 +1206,12 @@ def from_config( wf_class = workflow_registry.get(name) if wf_class is None: continue + has_deps = bool(wf_cfg.get("dependencies", [])) + default_replicas = 0 if has_deps else 1 cm.register_group( name=name, workflow_class=wf_class, - replicas=int(wf_cfg.get("replicas", 1)), + replicas=int(wf_cfg.get("replicas", default_replicas)), dependencies=list(wf_cfg.get("dependencies", [])), dep_threshold=int(wf_cfg.get("dependency_threshold", 1)), priority=int(wf_cfg.get("priority", 0)), diff --git a/tests/test_campaign_manager.py b/tests/test_campaign_manager.py index 63403af..292e3cf 100644 --- a/tests/test_campaign_manager.py +++ b/tests/test_campaign_manager.py @@ -55,16 +55,27 @@ async def run(self, replica_id: str) -> None: RecordingWorkflow.ran.append(replica_id) -class SignalWorkflow(BaseWorkflow): - """Fires _signal_ready() immediately then finishes after a tiny sleep.""" +class SignalDoneWorkflow(BaseWorkflow): + """Fires _signal_done() immediately then finishes after a tiny sleep.""" - workflow_id = "signal" + workflow_id = "signal_done" async def run(self, replica_id: str) -> None: - await self._signal_ready() + await self._signal_done() await asyncio.sleep(0.01) +class TriggerWorkflow(BaseWorkflow): + """Triggers a dependent group named 'downstream' then finishes.""" + + workflow_id = "trigger" + dependent_name: str = "downstream" + dependent_replicas: int = 1 + + async def run(self, replica_id: str) -> None: + await self._trigger_dependent(self.dependent_name, replicas=self.dependent_replicas) + + class HookWorkflow(BaseWorkflow): """Records (replica_id, final_state) tuples in on_replica_done.""" @@ -149,25 +160,25 @@ async def _fake_init(): class TestBaseWorkflow: - async def test_signal_ready_no_callback_is_noop(self): + async def test_signal_done_no_cm_is_noop(self): wf = NullWorkflow() - await wf._signal_ready() # must not raise + await wf._signal_done() # must not raise - async def test_signal_ready_calls_sync_callback(self): - called = [] - wf = NullWorkflow(on_ready=lambda: called.append(1)) - await wf._signal_ready() - assert called == [1] - - async def test_signal_ready_awaits_async_callback(self): - called = [] + async def test_trigger_dependent_no_cm_is_noop(self): + wf = NullWorkflow() + await wf._trigger_dependent("some_group", replicas=2) # must not raise - async def cb(): - called.append(1) + async def test_signal_done_calls_cm(self): + cm_mock = AsyncMock() + wf = NullWorkflow(_cm=cm_mock, _group_name="mygroup") + await wf._signal_done() + cm_mock.signal_done.assert_awaited_once_with("mygroup") - wf = NullWorkflow(on_ready=cb) - await wf._signal_ready() - assert called == [1] + async def test_trigger_dependent_calls_cm(self): + cm_mock = AsyncMock() + wf = NullWorkflow(_cm=cm_mock) + await wf._trigger_dependent("dep", replicas=3) + cm_mock.trigger_dependent.assert_awaited_once_with("dep", replicas=3) def test_base_run_raises_not_implemented(self): wf = BaseWorkflow() @@ -258,9 +269,9 @@ async def run(self, replica_id: str) -> None: b_idx = next(i for i, (wf, _) in enumerate(order) if wf == "B") assert all(wf == "A" for wf, _ in order[:b_idx]) - async def test_dependency_via_signal_ready(self, acm): - """_signal_ready() unblocks B even before all of A's replicas finish.""" - acm.register_group("a", SignalWorkflow, replicas=1) + async def test_dependency_via_signal_done(self, acm): + """_signal_done() unblocks B even before all of A's replicas finish.""" + acm.register_group("a", SignalDoneWorkflow, replicas=1) acm.register_group( "b", NullWorkflow, @@ -275,6 +286,25 @@ async def test_dependency_via_signal_ready(self, acm): assert s["groups"]["a"]["ready"] is True assert s["groups"]["b"]["status"] == "done" + async def test_trigger_dependent_activates_group(self, acm): + """Parent workflow calls _trigger_dependent to start a replicas=0 group.""" + acm.register_group("upstream", TriggerWorkflow, replicas=1) + acm.register_group("downstream", RecordingWorkflow, replicas=0) + await acm.start() + assert await acm.wait(timeout=3.0) + + s = acm.status() + assert s["groups"]["downstream"]["status"] == "done" + assert "downstream_0" in RecordingWorkflow.ran + + async def test_untriggered_group_does_not_block_completion(self, acm): + """A replicas=0 group that is never triggered must not prevent _all_done.""" + acm.register_group("a", NullWorkflow, replicas=1) + acm.register_group("never_triggered", NullWorkflow, replicas=0) + await acm.start() + assert await acm.wait(timeout=3.0) + assert acm.status()["groups"]["a"]["status"] == "done" + async def test_on_replica_done_hook_called(self, acm): acm.register_group("a", HookWorkflow, replicas=2) await acm.start() @@ -328,7 +358,8 @@ async def test_from_config_registers_groups(self): config = { "workflows": { "x": {"replicas": 2, "max_replicas": 1, "priority": 3}, - "y": {"replicas": 1, "dependencies": ["x"], "dependency_threshold": 2}, + # y has dependencies → replicas defaults to 0 (triggered group) + "y": {"dependencies": ["x"], "dependency_threshold": 2}, } } cm = AsyncCampaignManager.from_config(config, {"x": NullWorkflow, "y": NullWorkflow}) @@ -336,6 +367,7 @@ async def test_from_config_registers_groups(self): assert s["x"]["replicas_total"] == 2 assert s["x"]["max_replicas"] == 1 assert s["x"]["priority"] == 3 + assert s["y"]["replicas_total"] == 0 # triggered group: not yet activated assert s["y"]["dependencies"] == ["x"] assert s["y"]["dep_threshold"] == 2 @@ -443,15 +475,20 @@ def test_from_config_registers_groups(self): config = { "workflows": { "alpha": {"replicas": 3, "max_replicas": 2, "priority": 7}, + # beta has dependencies → replicas defaults to 0 (triggered group) + "beta": {"dependencies": ["alpha"]}, } } - cm = CampaignManager.from_config(config, {"alpha": SyncRecordingWorkflow}) + cm = CampaignManager.from_config( + config, {"alpha": SyncRecordingWorkflow, "beta": SyncRecordingWorkflow} + ) s = cm.status()["groups"] cm.close() assert "alpha" in s assert s["alpha"]["replicas_total"] == 3 assert s["alpha"]["max_replicas"] == 2 assert s["alpha"]["priority"] == 7 + assert s["beta"]["replicas_total"] == 0 # triggered group: not yet activated def test_unknown_group_skipped_in_from_config(self): config = {"workflows": {"ghost": {"replicas": 1}}} diff --git a/workflows/run_campaign/config.yaml b/workflows/run_campaign/config.yaml index 8b4a5ad..57d3124 100644 --- a/workflows/run_campaign/config.yaml +++ b/workflows/run_campaign/config.yaml @@ -28,23 +28,58 @@ workflow_registry: # ── Workflow Groups ─────────────────────────────────────────────────────────── # Each key must match a name in the workflow_registry above. # -# replicas : total replicas to run -# min_replicas : guaranteed concurrent slots -# max_replicas : sliding-window cap on concurrent replicas -# priority : higher → scheduled before lower-priority groups -# dependencies : groups that must signal ready before this starts -# dependency_threshold : fallback: unblock after N finished replicas -# required_cpus/gpus: per-replica resource reservation -# config_file : path (relative to this file) to the workflow's own -# parameter file — merged into the config passed to the -# workflow constructor +# Two modes — controlled by whether 'dependencies' is set: +# +# Independent (no dependencies): +# replicas: N — group starts immediately on cm.start() +# +# Dependent (has dependencies): +# omit 'replicas' — group starts at 0 and stays inactive until an upstream +# replica calls await self._trigger_dependent("name", replicas=N) +# The upstream workflow decides *when* and *how many* based on its own +# execution logic (e.g. 1 downstream run per result produced, or N based +# on a quality threshold). Each call adds N more replicas to the queue; +# calls can repeat across the lifetime of the upstream run. +# +# To switch a dependent group to independent: add 'replicas: N' and remove +# 'dependencies'. No workflow code needs to change — the pipeline topology +# lives entirely in this config. +# +# Field reference +# --------------- +# replicas : total replicas for independent groups; omit for dependent +# min_replicas : guaranteed concurrent slots (scheduler pass 1) +# max_replicas : sliding-window concurrency cap (scheduler pass 2) +# priority : higher → scheduled first when resources are contested +# dependencies : upstream groups this workflow depends on; also used by +# the CM to route _signal_done() — when group X signals +# done, every group listing X here gets +1 replica queued +# dependency_threshold : N finished upstream replicas satisfies dep check when +# using _signal_done() fallback (default 1) +# required_cpus/gpus : per-replica resource reservation (ResourcePool) +# config_file : path to the workflow's own YAML merged into the +# constructor config (${VAR} expanded at load time) workflows: - # ── DDSim (simulation + training loop) ───────────────────────────────────── + # ── ESM2 Inference ───────────────────────────────────────────────────────── + # Independent: starts immediately on cm.start(). + # on_replica_done triggers 1 'dummy' replica per successful inference result. + inference: + priority: 6 + replicas: 8 + min_replicas: 1 + max_replicas: 4 + required_cpus: 4 + required_gpus: 1 + config_file: "${INF_DIR}/config.yaml" + + # ── DDSim ────────────────────────────────────────────────────────────────── + # Dependent on inference: stays at replicas=0 until InferenceWorkflow calls + # _trigger_dependent("dummy", replicas=1) for each successful result. + # To run independently: add 'replicas: N' and remove 'dependencies'. dummy: priority: 5 - replicas: 8 min_replicas: 2 max_replicas: 4 required_cpus: 4 @@ -52,36 +87,27 @@ workflows: dependencies: [inference] config_file: "${DUMMY_DIR}/config.yaml" - # ── ESM2 Inference ────────────────────────────────────────────────────────── - inference: - priority: 6 - replicas: 8 - min_replicas: 1 - max_replicas: 4 - required_cpus: 4 - required_gpus: 1 - dependencies: [] - dependency_threshold: 1 - config_file: "${INF_DIR}/config.yaml" + # ── DDMd (MD simulations) ────────────────────────────────────────────────── + # Independent: starts immediately, runs in parallel with inference. + # Each completed iteration calls _signal_done() → CM queues 1 miniapps replica. + md: + priority: 10 + replicas: 2 + min_replicas: 1 + max_replicas: 1 + required_cpus: 4 + required_gpus: 1 + config_file: "${MD_HOME}/config.yaml" - # ── MiniApps Workflow ──────────────────────────────────────────────────────── + # ── MiniApps ─────────────────────────────────────────────────────────────── + # Dependent on md: stays at replicas=0 until DDMdWrapperWorkflow calls + # _trigger_dependent("miniapps", replicas=N) based on MD results. + # To run independently: add 'replicas: N' and remove 'dependencies'. miniapps: priority: 8 - replicas: 4 min_replicas: 1 - max_replicas: 2 + max_replicas: 1 required_cpus: 4 required_gpus: 1 dependencies: [md] config_file: "${MINAPPS_DIR}/config.yaml" - - # ── DDMd Workflow (real MD simulations) ────────────────────────────────────── - md: - priority: 10 - replicas: 4 - min_replicas: 2 - max_replicas: 4 - required_cpus: 4 - required_gpus: 1 - dependencies: [] - config_file: "${MD_HOME}/config.yaml" diff --git a/workflows/run_campaign/ddmd_workflow.py b/workflows/run_campaign/ddmd_workflow.py index eb97ccc..235eb30 100644 --- a/workflows/run_campaign/ddmd_workflow.py +++ b/workflows/run_campaign/ddmd_workflow.py @@ -64,7 +64,7 @@ async def run(self, replica_id: str) -> None: asyncflow=asyncflow, config=replica_config_path, name=name, - on_ready=self._on_ready, + on_ready=lambda: self._signal_done(), policies=self.policies, engine_dragon=self.engine_dragon, ) diff --git a/workflows/run_campaign/dummy_workflow.py b/workflows/run_campaign/dummy_workflow.py index b82ae81..cc0c4c7 100644 --- a/workflows/run_campaign/dummy_workflow.py +++ b/workflows/run_campaign/dummy_workflow.py @@ -57,7 +57,8 @@ async def run(self, replica_id: str) -> None: name=name, asyncflow=asyncflow, home_dir=str(home_base), - on_ready=self._on_ready, + _cm=self._cm, + _group_name=self._group_name, policies=self.policies, ) except Exception: diff --git a/workflows/run_campaign/inference_workflow.py b/workflows/run_campaign/inference_workflow.py index 85abfc9..1aa8939 100644 --- a/workflows/run_campaign/inference_workflow.py +++ b/workflows/run_campaign/inference_workflow.py @@ -192,7 +192,7 @@ def _reset_service_queues(svc) -> None: # ------------------------------------------------------------------ # async def on_replica_done(self, replica_id: str, cm, final_state: str) -> None: - """Queue a DDSim replica per completed inference run; teardown on last.""" + """Queue 1 downstream replica per successful inference; teardown on last.""" status = cm.status() g = status["groups"].get("inference", {}) finished = g.get("replicas_finished", 0) + 1 @@ -205,18 +205,9 @@ async def on_replica_done(self, replica_id: str, cm, final_state: str) -> None: ) if final_state == "done": - InferenceWorkflow._log.info( - "queuing 1 ddsim replica", - component="workflow", - task_name=replica_id, - ) - await cm.add_replicas("dummy", n=1) - else: - InferenceWorkflow._log.warning( - f"skipping ddsim replica (state={final_state})", - component="workflow", - task_name=replica_id, - ) + # Each successful inference result triggers 1 downstream replica. + # The count is determined here by execution logic, not by config. + await self._trigger_dependent("dummy", replicas=1) if finished >= total: InferenceWorkflow._log.info( diff --git a/workflows/run_campaign/miniapps_workflow.py b/workflows/run_campaign/miniapps_workflow.py index 6ffd54f..849a6c6 100644 --- a/workflows/run_campaign/miniapps_workflow.py +++ b/workflows/run_campaign/miniapps_workflow.py @@ -67,7 +67,8 @@ async def run(self, replica_id: str) -> None: asyncflow=asyncflow, home_dir=str(home_base), name=name, - on_ready=self._on_ready, + _cm=self._cm, + _group_name=self._group_name, policies=self.policies, ) except Exception: diff --git a/workflows/run_campaign/plot_cm_timeline.py b/workflows/run_campaign/plot_cm_timeline.py index 67a5fe1..9d4023e 100644 --- a/workflows/run_campaign/plot_cm_timeline.py +++ b/workflows/run_campaign/plot_cm_timeline.py @@ -3,7 +3,7 @@ Plot replica execution timeline from a campaign SLURM log. Usage: - python plot_replicas.py slurm-17715157.out [--out timeline.png] + python plot_cm_timeline.py slurm-XXXXXX.out [--out timeline.png] """ import argparse @@ -23,6 +23,7 @@ matplotlib.use("Agg") import matplotlib.gridspec as gridspec +import matplotlib.lines as mlines import matplotlib.patches as mpatches import matplotlib.pyplot as plt @@ -35,13 +36,16 @@ r"min=(\d+) max=(\d+) deps=\[([^\]]*)\] dep_threshold=(\d+) " r"resources=\(cpus=(\d+), gpus=(\d+)\)" ) -_GPU_ASSIGN_RE = re.compile(r"GPU assign: '(\w+)' \u2192 GPU\(s\) \[([^\]]*)\]") -# "resources: cpus=USED/TOTAL gpus=USED/TOTAL" from scheduler log +_GPU_ASSIGN_RE = re.compile(r"GPU assign: '(\w+)' → GPU\(s\) \[([^\]]*)\]") _USAGE_RE = re.compile(r"resources: cpus=(\d+)/(\d+)\s+gpus=(\d+)/(\d+)") -# "available: cpus=AVAIL/TOTAL gpus=AVAIL/TOTAL" from replica-finished log _AVAIL_RE = re.compile(r"available: cpus=(\d+)/(\d+)\s+gpus=(\d+)/(\d+)") _TOTAL_RES_RE = re.compile(r"Resource pool: total_cpus=(\d+)\s+total_gpus=(\d+)") -_READY_RE = re.compile(r"Group '(\w+)' signaled ready") + +# Signal events — new dependency model +# "signal_done": 'md' signaled done → +1 replica for ['miniapps'] +# "trigger_dep": trigger_dependent: 'dummy' +1 replicas (total=3) +_SIGNAL_DONE_RE = re.compile(r"'(\w+)' signaled done.*?\+(\d+) replica.*?\['([\w,\s]*)'\]") +_TRIGGER_DEP_RE = re.compile(r"trigger_dependent: '(\w+)' \+(\d+) replicas \(total=(\d+)\)") GROUP_COLORS = { "inference": "#4C72B0", @@ -55,12 +59,16 @@ def parse_log(path: str): - """Parse SLURM log; return spans, group_meta, resource_timeline, gpu_assignments.""" + """Parse SLURM log; return spans, group_meta, resource_timeline, + gpu_assignments, signal_events.""" starts: dict[str, datetime] = {} spans = [] group_meta = {} gpu_assignments = {} - resource_timeline = [] # (elapsed_s, used_cpus, total_cpus, used_gpus, total_gpus) + resource_timeline = [] + # (elapsed_s, source_group, target_groups, n_replicas, kind) + # kind: "signal_done" | "trigger_dep" + signal_events = [] t0_dt = None total_cpus = total_gpus = 0 @@ -74,7 +82,6 @@ def parse_log(path: str): with open(path) as fh: for raw in fh: - # Group registration lines (no timestamp required) if m := _GROUP_RE.search(raw): name = m.group(1) deps_raw = m.group(6) @@ -102,6 +109,7 @@ def parse_log(path: str): dt = datetime.strptime(f"{date_ref} {ts_m.group(1)}", "%Y-%m-%d %H:%M:%S.%f") if t0_dt is None: t0_dt = dt + elapsed = (dt - t0_dt).total_seconds() if m := _GPU_ASSIGN_RE.search(raw): rid = m.group(1) @@ -109,18 +117,38 @@ def parse_log(path: str): gpu_ids = [int(x) for x in gpu_str.split(",") if x.strip()] if gpu_str else [] gpu_assignments[rid] = gpu_ids - # Resource usage from scheduler summary line if m := _USAGE_RE.search(raw): - uc, tc, ug, tg = int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4)) - elapsed = (dt - t0_dt).total_seconds() if t0_dt else 0 + uc, tc, ug, tg = ( + int(m.group(1)), + int(m.group(2)), + int(m.group(3)), + int(m.group(4)), + ) resource_timeline.append((elapsed, uc, tc, ug, tg)) - - # Resource available from replica-finished line (convert to used) elif m := _AVAIL_RE.search(raw): - ac, tc, ag, tg = int(m.group(1)), int(m.group(2)), int(m.group(3)), int(m.group(4)) - elapsed = (dt - t0_dt).total_seconds() if t0_dt else 0 + ac, tc, ag, tg = ( + int(m.group(1)), + int(m.group(2)), + int(m.group(3)), + int(m.group(4)), + ) resource_timeline.append((elapsed, tc - ac, tc, tg - ag, tg)) + # Signal events — new dependency model + if m := _SIGNAL_DONE_RE.search(raw): + src = m.group(1) + n = int(m.group(2)) + targets_raw = m.group(3) + targets = [t.strip().strip("'\"") for t in targets_raw.split(",") if t.strip()] + for tgt in targets: + signal_events.append((elapsed, src, tgt, n, "signal_done")) + + if m := _TRIGGER_DEP_RE.search(raw): + tgt = m.group(1) + n = int(m.group(2)) + # Source unknown from log line — mark as "trigger_dep" + signal_events.append((elapsed, None, tgt, n, "trigger_dep")) + if m := _START_RE.search(raw): starts[m.group(1)] = dt elif m := _FINISH_RE.search(raw): @@ -138,14 +166,31 @@ def parse_log(path: str): group = rid.rsplit("_", 1)[0] spans.append((rid, group, start, start, None)) - # Sort resource timeline and deduplicate by elapsed resource_timeline.sort(key=lambda x: x[0]) + signal_events.sort(key=lambda x: x[0]) t0 = min(s[2] for s in spans) if spans else t0_dt - return spans, group_meta, resource_timeline, gpu_assignments, t0, (total_cpus, total_gpus) + return ( + spans, + group_meta, + resource_timeline, + gpu_assignments, + signal_events, + t0, + (total_cpus, total_gpus), + ) -def plot(spans, group_meta, resource_timeline, gpu_assignments, t0, total_resources, out_path): +def plot( + spans, + group_meta, + resource_timeline, + gpu_assignments, + signal_events, + t0, + total_resources, + out_path, +): if not spans: print("No replica events found.", file=sys.stderr) return @@ -188,7 +233,7 @@ def sort_key(s): # ── Gantt chart ────────────────────────────────────────────────────────── yticks, ylabels = [], [] - group_row_ranges = {} # group -> (first_row, last_row) + group_row_ranges = {} prev_group = None for row, (rid, group, start, end, ok) in enumerate(spans): @@ -201,18 +246,15 @@ def sort_key(s): lw = 1.5 if ok is False else 0 alpha = 0.45 if ok is None else 0.88 - # Draw separator line between groups if group != prev_group and prev_group is not None: ax_gantt.axhline(row - 0.5, color="grey", lw=0.6, alpha=0.5, linestyle="--") prev_group = group - # Track row ranges per group if group not in group_row_ranges: group_row_ranges[group] = [row, row] else: group_row_ranges[group][1] = row - # Alternating background per group g_idx = GROUP_ORDER.index(group) if group in GROUP_ORDER else len(GROUP_ORDER) if g_idx % 2 == 0: ax_gantt.axhspan(row - 0.5, row + 0.5, color="grey", alpha=0.04, linewidth=0) @@ -228,7 +270,6 @@ def sort_key(s): alpha=alpha, ) - # Annotate bar with GPU assignment or CPU count gpu_ids = gpu_assignments.get(rid, []) meta = group_meta.get(group, {}) if gpu_ids: @@ -254,14 +295,15 @@ def sort_key(s): yticks.append(row) ylabels.append(rid) - # Group section labels on the right of the Gantt + # Group section labels on the right for group, (r0, r1) in group_row_ranges.items(): meta = group_meta.get(group, {}) mid = (r0 + r1) / 2 pri = meta.get("priority", "?") cpus = meta.get("cpus", 0) gpus = meta.get("gpus", 0) - info = f"priority={pri}\ncpu={cpus} gpu={gpus}" + mode = "dep." if meta.get("deps") else "indep." + info = f"priority={pri}\ncpu={cpus} gpu={gpus}\n{mode}" ax_gantt.text( 1.002, 1.0 - (mid + 0.5) / n_rows, @@ -274,48 +316,155 @@ def sort_key(s): fontweight="bold", ) - # Build a lookup: group -> sorted list of (row, start, end, ok) for its spans - group_spans = {} + # Build group_spans lookup: group -> sorted list of (row, start_dt, end_dt, ok) + group_spans: dict[str, list] = {} for row, (_, group, start, end, ok) in enumerate(spans): group_spans.setdefault(group, []).append((row, start, end, ok)) - # Dependency arrows: - # tail — midpoint of first dep-group replica bar (where signal_ready likely fires) - # head — start of first dependent-group replica bar - # vline — trigger moment (start of first dependent replica) + # ── Dependency signal arrows ───────────────────────────────────────────── + # Each signal event gets its own arrow: signal point → triggered replica start. + consumed_tgt_rows: set[int] = set() + consumed_src_rows: set[int] = set() + # Round-robin counters for signal_done — distributes signals evenly across + # parallel replicas when the log doesn't record which replica sent each signal. + signal_done_rr: dict[str, int] = {} + # available target replicas per group, sorted by start time + available: dict[str, list] = { + g: sorted(sl, key=lambda s: s[1]) for g, sl in group_spans.items() + } + + for sig_elapsed, src_group, tgt_group, _n, kind in signal_events: + if tgt_group not in available: + continue + + # Find the earliest unconsumed target replica that starts at or after signal + tgt_span = None + for s in available[tgt_group]: + if s[0] not in consumed_tgt_rows and (s[1] - t0).total_seconds() >= sig_elapsed - 0.5: + tgt_span = s + break + if tgt_span is None: + continue + consumed_tgt_rows.add(tgt_span[0]) + + tgt_row = tgt_span[0] + tgt_start_t = (tgt_span[1] - t0).total_seconds() + + # Determine the row to use for the signal source. + # For trigger_dep events src_group is None — infer from tgt_group's own + # dependency list (what tgt_group depends ON, not who depends on it). + resolved_src = src_group + if not resolved_src: + tgt_deps = group_meta.get(tgt_group, {}).get("deps", []) + resolved_src = tgt_deps[0] if tgt_deps else None + + if resolved_src and resolved_src in group_row_ranges: + r0, r1 = group_row_ranges[resolved_src] + src_spans = group_spans.get(resolved_src, []) + + if kind == "trigger_dep": + # Signal fires from on_replica_done — the source replica may not + # yet have its "finished" line in the log. Find the closest + # unconsumed source replica by finish time, without a direction + # constraint. + candidates = [s for s in src_spans if s[0] not in consumed_src_rows] + if candidates: + best = min( + candidates, + key=lambda s: abs((s[2] - t0).total_seconds() - sig_elapsed), + ) + src_row = best[0] + consumed_src_rows.add(best[0]) + else: + src_row = (r0 + r1) / 2 + else: + # signal_done fires from within run() — multiple replicas may be + # running in parallel and the log doesn't record which sent it. + # Distribute signals round-robin across the running replicas. + running = [ + s + for s in src_spans + if (s[1] - t0).total_seconds() + <= sig_elapsed + <= (s[2] - t0).total_seconds() + 0.5 + ] + if running: + rr_idx = signal_done_rr.get(resolved_src, 0) + src_row = running[rr_idx % len(running)][0] + signal_done_rr[resolved_src] = rr_idx + 1 + else: + src_row = (r0 + r1) / 2 + + src_color = GROUP_COLORS.get(resolved_src, DEFAULT_COLOR) + else: + src_row = tgt_row - 1.5 + src_color = GROUP_COLORS.get(tgt_group, DEFAULT_COLOR) + + # Draw diamond marker at signal point on source row + ax_gantt.plot( + sig_elapsed, + src_row, + "D", + markersize=5, + color=src_color, + zorder=5, + markeredgecolor="white", + markeredgewidth=0.5, + ) + + # Draw arrow from signal diamond to triggered replica start + ax_gantt.annotate( + "", + xy=(tgt_start_t, tgt_row), + xytext=(sig_elapsed, src_row), + arrowprops=dict( + arrowstyle="->", + color="#555555", + lw=1.1, + connectionstyle="arc3,rad=0.25", + ), + annotation_clip=False, + ) + + # Fall back to a single structural arrow for deps with no logged signals + # (e.g. log truncated or dep_threshold path) + drawn_dep_pairs: set[tuple[str, str]] = set() + for _, src, tgt, _, kind in signal_events: + if src: + drawn_dep_pairs.add((src, tgt)) + elif kind == "trigger_dep": + # src is None for trigger_dep log lines — resolve from tgt's dep list + tgt_deps = group_meta.get(tgt, {}).get("deps", []) + if tgt_deps: + drawn_dep_pairs.add((tgt_deps[0], tgt)) for group, _ in group_row_ranges.items(): meta = group_meta.get(group, {}) for dep_name in meta.get("deps", []): + if (dep_name, group) in drawn_dep_pairs: + continue if dep_name not in group_spans or group not in group_spans: continue - - dep_first = group_spans[dep_name][0] # (row, start, end, ok) of first dep replica - grp_first = group_spans[group][ - 0 - ] # (row, start, end, ok) of first replica of this group - + dep_first = group_spans[dep_name][0] + grp_first = group_spans[group][0] dep_row, dep_start, dep_end, _ = dep_first - grp_row, grp_start, _grp_end, _ = grp_first - - dep_bar_mid_t = (dep_start - t0).total_seconds() + grp_row, grp_start, _, _ = grp_first + dep_mid_t = (dep_start - t0).total_seconds() if dep_end != dep_start: - dep_bar_mid_t += ((dep_end - dep_start).total_seconds()) / 2 - - grp_bar_start_t = (grp_start - t0).total_seconds() - + dep_mid_t += (dep_end - dep_start).total_seconds() / 2 + grp_start_t = (grp_start - t0).total_seconds() ax_gantt.annotate( "", - xy=(grp_bar_start_t, grp_row), # head: start of first dependent replica - xytext=(dep_bar_mid_t, dep_row), # tail: mid of first dep replica bar + xy=(grp_start_t, grp_row), + xytext=(dep_mid_t, dep_row), arrowprops=dict( arrowstyle="->", - color="#555555", - lw=1.3, + color="#888888", + lw=1.0, connectionstyle="arc3,rad=0.35", + linestyle="dashed", ), annotation_clip=False, ) - ax_gantt.axvline(grp_bar_start_t, color="#555555", lw=0.6, linestyle=":", alpha=0.5) ax_gantt.set_yticks(yticks) ax_gantt.set_yticklabels(ylabels, fontsize=7) @@ -328,6 +477,15 @@ def sort_key(s): legend_patches += [ mpatches.Patch(facecolor="white", edgecolor="red", linewidth=1.2, label="error"), mpatches.Patch(color="grey", alpha=0.45, label="still running"), + mlines.Line2D( + [0], + [0], + marker="D", + color="w", + markerfacecolor="#666666", + markersize=6, + label="signal / trigger", + ), ] ax_gantt.legend(handles=legend_patches, loc="lower right", fontsize=7, framealpha=0.8) @@ -340,19 +498,12 @@ def sort_key(s): g for g in group_meta if g not in GROUP_ORDER ] - col_labels = [ - "Work\nflow", - "Priority", - "CPUs", - "GPUs", - "min/max\nreplicas", - "Total\nreplicas", - "Deps", - ] + col_labels = ["Workflow", "Priority", "CPUs", "GPUs", "min/max", "Mode", "Deps"] rows_data, row_colors = [], [] for gname in info_groups: m = group_meta[gname] deps = ", ".join(m.get("deps", [])) or "—" + mode = "dep." if m.get("deps") else "indep." rows_data.append( [ gname, @@ -360,7 +511,7 @@ def sort_key(s): str(m.get("cpus", 0)), str(m.get("gpus", 0)), f"{m.get('min', 0)}/{m.get('max', 0)}", - str(m.get("replicas", "?")), + mode, deps, ] ) @@ -378,25 +529,30 @@ def sort_key(s): tbl.set_fontsize(7.5) tbl.scale(1.0, 1.5) - # Style header row for j in range(len(col_labels)): tbl[0, j].set_facecolor("#333333") tbl[0, j].set_text_props(color="white", fontweight="bold") - # Dependency chain text - dep_chains = [] + # Signal-based dependency graph + dep_lines = [] + # Collect unique dep relationships with signal counts + sig_counts: dict[tuple[str, str], int] = {} + for _, src, tgt, n, _kind in signal_events: + if src: + sig_counts[(src, tgt)] = sig_counts.get((src, tgt), 0) + n + for gname in info_groups: m = group_meta[gname] - deps = m.get("deps", []) - if deps: - thr = m.get("dep_threshold", 1) - dep_chains.append(f" {', '.join(deps)} →(≥{thr}) {gname}") + for dep_name in m.get("deps", []): + count = sig_counts.get((dep_name, gname), 0) + count_str = f" ×{count}" if count else "" + dep_lines.append(f" {dep_name} —signals→ {gname}{count_str}") - if dep_chains: - dep_str = "Dependency graph:\n" + "\n".join(dep_chains) + if dep_lines: + dep_str = "Dependency graph (signals):\n" + "\n".join(dep_lines) ax_info.text( 0.5, - 0.2, + 0.22, dep_str, transform=ax_info.transAxes, va="bottom", @@ -406,21 +562,26 @@ def sort_key(s): bbox=dict(boxstyle="round,pad=0.5", facecolor="#f0f4ff", edgecolor="#aabbdd"), ) - # Scheduling logic note sched_note = ( + "Dependency model:\n" + " Independent: starts on cm.start()\n" + " Dependent: waits for upstream signal\n" + " ◆ _signal_done() → +1 replica per\n" + " downstream group in dependencies\n" + " ◆ _trigger_dependent() → explicit N\n" + "\n" "Scheduler:\n" - " Pass 1: guarantee min replicas (by priority)\n" - "Pass 2: fill up to max replicas (by priority)\n" - "Blocked by: resource limits & unmet deps" + " Pass 1: guarantee min replicas (priority)\n" + " Pass 2: fill up to max replicas (priority)" ) ax_info.text( 0.5, - 0.5, + 0.52, sched_note, transform=ax_info.transAxes, va="bottom", ha="center", - fontsize=9, + fontsize=8, bbox=dict(boxstyle="round,pad=0.5", facecolor="#fffbe6", edgecolor="#ccaa00"), ) @@ -450,6 +611,11 @@ def sort_key(s): ax_res.tick_params(axis="y", labelcolor="#4C72B0", labelsize=7) ax_res.set_ylim(bottom=0) + # Mark signal events on resource plot + for sig_elapsed, src, tgt, _n, _kind in signal_events: + color = GROUP_COLORS.get(src or tgt, "#888888") + ax_res.axvline(sig_elapsed, color=color, lw=0.7, alpha=0.5, linestyle=":") + ax_cpu = ax_res.twinx() ax_cpu.step(times, used_cpus, where="post", color="#DD8452", lw=1.8, label="CPUs used") ax_cpu.fill_between(times, used_cpus, step="post", color="#DD8452", alpha=0.12) @@ -482,7 +648,7 @@ def sort_key(s): def parse_config(path: str) -> dict: - """Load group_meta from a campaign config.yaml (authoritative source).""" + """Load group_meta from a campaign config.yaml.""" if not _HAVE_YAML: print("PyYAML not installed — falling back to log-parsed group metadata", file=sys.stderr) return {} @@ -490,8 +656,10 @@ def parse_config(path: str) -> dict: cfg = yaml.safe_load(fh) group_meta = {} for name, wf in cfg.get("workflows", {}).items(): + has_deps = bool(wf.get("dependencies", [])) + default_replicas = 0 if has_deps else 1 group_meta[name] = { - "replicas": int(wf.get("replicas", 1)), + "replicas": int(wf.get("replicas", default_replicas)), "priority": int(wf.get("priority", 0)), "min": int(wf.get("min_replicas", 0)), "max": int(wf.get("max_replicas", 0)), @@ -504,38 +672,40 @@ def parse_config(path: str) -> dict: def _default_out(log_path: str) -> str: - stem = Path(log_path).stem # e.g. "slurm-17716610" + stem = Path(log_path).stem m = re.search(r"(\d+)", stem) run_num = m.group(1) if m else stem - return f"replica_timeline_{run_num}.png" + return f"cm_timeline_{run_num}.png" def main(): - parser = argparse.ArgumentParser() + parser = argparse.ArgumentParser(description="Plot campaign manager replica timeline") parser.add_argument("log", help="SLURM output file") parser.add_argument( "--config", default=None, - help="Campaign config.yaml (authoritative source for group metadata). " - "Auto-detected as config.yaml next to the log file if not given.", - ) - parser.add_argument( - "--out", default=None, help="Output PNG (default: replica_timeline_.png)" + help="Campaign config.yaml (auto-detected as config.yaml next to log if not given)", ) + parser.add_argument("--out", default=None, help="Output PNG (default: cm_timeline_.png)") args = parser.parse_args() if args.out is None: args.out = _default_out(args.log) - # Auto-detect config.yaml next to the log file if args.config is None: candidate = Path(args.log).parent / "config.yaml" if candidate.exists(): args.config = str(candidate) - spans, group_meta_log, resource_timeline, gpu_assignments, t0, total_resources = parse_log( - args.log - ) + ( + spans, + group_meta_log, + resource_timeline, + gpu_assignments, + signal_events, + t0, + total_resources, + ) = parse_log(args.log) if args.config: group_meta = parse_config(args.config) @@ -548,12 +718,20 @@ def main(): f"Parsed {len(spans)} replica spans, " f"{len(group_meta)} groups, " f"{len(resource_timeline)} resource events, " - f"{len(gpu_assignments)} GPU assignments" + f"{len(gpu_assignments)} GPU assignments, " + f"{len(signal_events)} signal events" + ) + plot( + spans, + group_meta, + resource_timeline, + gpu_assignments, + signal_events, + t0, + total_resources, + args.out, ) - plot(spans, group_meta, resource_timeline, gpu_assignments, t0, total_resources, args.out) if __name__ == "__main__": main() - -# python plot_replicas.py slurm-17715157.out --out replica_timeline.png diff --git a/workflows/run_campaign/plots/adaptive_cm_timeline.png b/workflows/run_campaign/plots/adaptive_cm_timeline.png new file mode 100644 index 0000000..ace6c86 Binary files /dev/null and b/workflows/run_campaign/plots/adaptive_cm_timeline.png differ diff --git a/workflows/run_campaign/cm_timeline.png b/workflows/run_campaign/plots/cm_timeline.png similarity index 100% rename from workflows/run_campaign/cm_timeline.png rename to workflows/run_campaign/plots/cm_timeline.png