diff --git a/.gitignore b/.gitignore index 2acfac3..f2b3470 100644 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,11 @@ outputs/ # Large generated files — not tracked in git *.bak *.pptx -workflows/run_campaign/dreamer_campaign/*.json -workflows/run_campaign/dreamer_campaign/log -workflows/run_campaign/dreamer_campaign/plots/ \ No newline at end of file + +# Runtime output directories (campaign artifacts) +workflows/esm2_inference/ESM2 +workflows/esm2_inference/cache/ +workflows/*/telemetry*/ +workflows/*/*.npy +workflows/*/ddict_* +workflows/sgdes/mayv_output/ \ No newline at end of file diff --git a/CLAUDE.md b/CLAUDE.md index b36347e..3cfa262 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -2,12 +2,16 @@ This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. +> **Note**: The `AsyncCampaignManager` and all campaign orchestration code has been moved to +> `/scratch/bblj/mgoliyad1/campaign_manager`. This repo now contains only the +> multi-GPU inference service framework and SGDES workflow. + ## Quick Start ### Installation ```bash -# Core installation +# Core installation (inference + utils) pip install -e . # With ESM2 model support (torch + transformers) @@ -16,11 +20,11 @@ pip install -e ".[esm2]" # With Dragon/RADICAL HPC support pip install -e ".[dragon]" -# Full dev setup (recommended) +# Full dev setup pip install -e ".[esm2,dragon,dev]" ``` -Requires **Python ≥ 3.10** (union type syntax, structural pattern matching). +Requires **Python ≥ 3.10**. ### Common Commands @@ -32,8 +36,8 @@ pytest pytest --cov=src --cov-report=html # Run a single test file or test -pytest tests/test_campaign_manager.py -pytest tests/test_campaign_manager.py::TestClass::test_method +pytest tests/test_inference_service.py +pytest tests/test_inference_service.py::TestClass::test_method # Lint and format check ruff check . @@ -41,282 +45,119 @@ ruff format --check . # Auto-format code ruff format . - -# Run linting and format with auto-fixes ruff check . --fix -ruff format . ``` ### Key Entry Points -- **Multi-workflow campaign**: `python workflows/run_campaign/esm2_ddsim_campaign/run_campaing.py --config workflows/run_campaign/esm2_ddsim_campaign/config.yaml` *(note: `run_campaing.py` — intentional typo in filename)* -- **Dreamer campaign (single run)**: `python workflows/run_campaign/dreamer_campaign/run_campaign.py --config workflows/run_campaign/dreamer_campaign/config.yaml` -- **Dreamer benchmark (multi-config, N runs)**: `python workflows/run_campaign/dreamer_campaign/benchmark.py --config workflows/run_campaign/dreamer_campaign/config.yaml --runs 5 --out benchmark_results.json` -- **ESM2 inference (standalone)**: `python workflows/esm2_inference/run_esm2_infern.py --config workflows/esm2_inference/config.yaml --mode local` +- **ESM2 inference (standalone)**: `python workflows/esm2_inference/run_esm2_infern.py --config workflows/esm2_inference/inference.yaml --mode local` +- **SGDES workflow**: `python workflows/sgdes/run_workflow.py` +- **SLURM submission (Delta GPU)**: `sbatch workflows/esm2_inference/delta_gpu_sbatch.sh` --- ## Architecture Overview -SPHERICAL is an **async-native HPC workflow orchestrator** built on `radical.asyncflow`. The core innovation is **AsyncCampaignManager**, which orchestrates multiple heterogeneous workflow groups (replicas) inside a single Python asyncio event loop with sophisticated multi-stage dependency signalling, adaptive resource scheduling, and optional adaptive batching. +SPHERICAL provides a **multi-GPU inference service framework** for ESM2 protein language model embeddings, +built on `radical.asyncflow`. Workflows submit sequences to inference services running on one or more GPUs; +results feed downstream molecular dynamics or ML pipelines. ### High-Level Design ``` -AsyncCampaignManager (campaign_manager.py) -├── SchedulerMixin (scheduler.py) -│ └── Two-pass greedy scheduler: guarantee min_replicas, fill to max_replicas -├── ExecutorMixin (executor.py) -│ └── Replica lifecycle: launch, monitor, completion, GPU assignment -├── MonitorMixin (monitor_mixin.py) -│ └── Periodic health checks + drift detection -│ -├── BaseWorkflow (base_workflow.py) -│ └── User-defined workflow classes subclass this; override run() or start() -│ -└── Optional Features (feature flags in config) - ├── BackpressureNegotiator (backpressure.py) — per-edge queue depth controller - ├── Sharder (sharder.py) — batches upstream triggers before downstream dispatch - ├── Monitor (monitor.py) — detects pass-through & budget burn drift - ├── Bandit (bandit.py) — Thompson-sampling arm selection for scheduling & sharding - └── CandidateLog (candidate_log.py) — tracks upstream results for sharder ranking +src/inference/ +├── esm2_service/ +│ ├── esm2_service.py ESM2InferenceService — loads model, runs batched inference +│ └── esm2_client.py ESM2Client — submits sequences, collects embeddings +├── inference_service.py BaseInferenceService — async queue-based service skeleton +├── inference_client.py BaseInferenceClient — client protocol +├── orchestrator.py start_services() / start_services_local() — spin up N service instances +├── server.py aiohttp HTTP server wrapping a service +├── dragon_launcher.py Dragon backend launcher for HPC nodes +├── utils.py shared helpers (queue drain, batch sizing, etc.) +└── kill_all.py emergency cleanup for zombie processes + +src/utils/ +├── logger.py colored structured logging with metrics recording +└── workflow.py _expand_env(), load_config(), find_gpus(), make_policies() ``` ### Core Concepts -#### AsyncCampaignManager - -Orchestrates workflow groups with dependencies and resource constraints: - -- **Groups**: Named pools of replicas of the same workflow class. Each group has: - - `replicas`: total count (0 = dependent, wait for trigger) - - `min_replicas` / `max_replicas`: concurrent caps - - `priority`: scheduling priority (higher = first) - - `required_cpus` / `required_gpus`: per-replica resource reservation - - `dependencies`: upstream groups that must signal before this group starts - -- **Two signalling modes**: - - `_signal_done()`: broadcast to all downstream groups (topology-driven) - - `_trigger_dependent(name, replicas=N)`: explicit queue N replicas to a named group - -- **Scheduler**: Runs on every state change (replica finish, signal received). Two-pass greedy: - 1. Pass 1: guarantee `min_replicas` for all eligible groups (highest priority first) - 2. Pass 2: fill remaining capacity up to `max_replicas` (highest priority first) - -A group is **eligible** when its dependencies are **ready**: - - Workflow-driven: dependency called `_signal_done()` (sets `group.ready = True`) - - Count-based fallback: `dep.finished_replicas >= dep.dependency_threshold` - -#### BaseWorkflow - -All user workflows subclass `BaseWorkflow`. The CM injects six objects at construction: - -| Attribute | Type | Purpose | -|-----------|------|---------| -| `config` | dict | per-group config (CM scheduling keys stripped) | -| `_cm` | AsyncCampaignManager | reference to running CM (`None` in unit tests) | -| `_group_name` | str | name of this group (used by `_signal_done()`) | -| `asyncflow` | WorkflowEngine | shared radical.asyncflow engine | -| `policies` | list[Policy] | Dragon `Policy` per assigned GPU (empty on concurrent) | -| `engine_dragon` | object | Dragon backend handle (`None` on concurrent) | - -When GPUs are assigned, two extra config keys are injected: -- `assigned_gpu_ids`: list of GPU IDs for this replica -- `group_gpu_ids`: all GPUs held by the group right now - -**Workflow entry points**: Define **either** `async def run()` or `def start()`, not both. The CM detects which is overridden and raises `ValueError` if both or neither are defined. Async coroutines are awaited directly; sync functions run via `asyncio.to_thread`. +#### ESM2 Inference Service -Optional hook: `on_replica_done(replica_id, cm, final_state)` — called after entry-point returns/raises; can be async or sync. - -#### ResourcePool - -Tracks available CPU cores and GPU slots: +Each `ESM2InferenceService` instance owns one GPU. The orchestrator starts N instances (one per GPU): ```python -pool = ResourcePool(total_cpus=128, total_gpus=4) -pool.can_fit(cpus=4, gpus=1) # → True/False -pool.allocate(cpus=4, gpus=1) -pool.release(cpus=4, gpus=1) -pool.usage_str() # → "cpus=20/128 gpus=1/4" +handles = await start_services_local(config, ESM2InferenceService) +# handles[i].service → ESM2InferenceService on GPU i +# handles[i].endpoint → HTTP URL (server mode) +# handles[i].close() → teardown ``` -Setting total to 0 disables tracking (unlimited). - ---- - -## Campaign Configuration - -All campaigns use YAML config files with two sections: +Services expose three asyncio queues: +- `input_queue`: caller puts `(batch_id, sequences)` +- `processed_queue`: service puts `(batch_id, embeddings)` after inference +- `work_queue`: internal task tracking (joined for backpressure) -### Resources & Engine +#### ESM2 Client -```yaml -engine: concurrent # or "dragon" for HPC -resources: - total_cpus: 128 - total_gpus: 4 -``` +`ESM2Client` serialises access to one service via an `asyncio.Lock`. It: +1. Drains any in-flight work (`work_queue.join()`, `processed_queue.join()`) +2. Resets queue state for a clean run +3. Submits sequences → collects embeddings +4. Writes outputs to `config["output_dir"]` -### Workflow Groups +#### Server Mode -```yaml -workflows: - sim: - replicas: 8 # independent: starts immediately - min_replicas: 2 - max_replicas: 4 - priority: 10 - required_cpus: 4 - required_gpus: 1 - # other keys forwarded to workflow.config - - analysis: - priority: 8 - min_replicas: 1 - max_replicas: 4 - required_cpus: 4 - required_gpus: 1 - dependencies: [sim] # dependent: starts at 0 replicas - # sim's _signal_done() adds replicas at runtime -``` +`server.py` wraps a service with an aiohttp HTTP interface. Clients POST sequences and GET results. +Used for cross-node communication in Dragon/HPC deployments. -**Key insight**: Switch a group between independent/dependent modes purely through config, no workflow code changes needed. +#### Resource Helpers (`src/utils/workflow.py`) -### Optional Feature Flags - -```yaml -cm: - features: - backpressure: true # hysteresis queue depth controller per edge - sharder: true # adaptive batch dispatch from trigger buffer - monitor: true # periodic health checks + drift alerts - bandit: true # Thompson-sampling cross-stage optimization - monitor_interval_s: 30 # tick interval for monitor - telemetry: - collect_telemetry: true - telemetry_dir: telemetry-results - workflow_registry: - sim: my_module.SimWorkflow - analysis: my_module.AnalysisWorkflow -``` +- `_expand_env(value)` — expands `${VAR}` in config strings +- `load_config(path)` — loads YAML with env expansion +- `find_gpus(node)` — enumerates GPUs on a Dragon node +- `make_policies(gpu_ids)` — builds Dragon `Policy` objects for GPU affinity --- -## Optional Features Deep Dive - -### Backpressure (backpressure.py) - -Per-edge hysteresis state machine that throttles downstream queue depth: - -```yaml -workflows: - downstream: - backpressure_high: 200 # queue ≥ 200 → THROTTLE (block new starts) - backpressure_low: 100 # queue ≤ 100 → WIDEN (dispatch more) -``` - -Three states (HOLD → THROTTLE → WIDEN → HOLD): -- **HOLD**: normal, neither throttling nor widening -- **THROTTLE**: queue too deep, sharder dispatch returns 0 -- **WIDEN**: queue drained, dispatch multiplier increases - -### Sharder (sharder.py) - -Buffers upstream trigger signals and batch-dispatches downstream, with optional priority ranking: - -```yaml -workflows: - downstream: - sharding: - target_size: 100 # nominal batch size - min_size: 10 - max_size: 200 - stratify: soft # soft | strict | off - use_bandit: true # Thompson-sampling BP multiplier selection - dispatch_cap: 500 # max replicas to dispatch (drop low-priority candidates) -``` - -**Stratify modes**: -- `off`: dispatch exactly 1 trigger per cycle -- `soft`: adaptive sizing with tail dispatch (partial batches acceptable) -- `strict`: hold buffer until target_size or upstream done (chemical diversity, etc.) - -**Candidate ranking** (via ProfileWeights): -- Score, surrogate prediction, uncertainty, age, diversity (scaffold novelty) - -### Monitor (monitor.py, monitor_mixin.py) - -Periodic health checks and drift detection: - -```yaml -cm: - features: - monitor: true - monitor_interval_s: 30 - replan: - budget_burn_deviation_pct: 20 # alert if spend > expected + 20% - pass_through_deviation_pct: 25 # alert if pass-through ratio deviates 25% - surrogate_recall_floor: 0.90 # alert if surrogate recall < 90% -``` +## Key File Organization -Two monitoring paths: -1. **Reactive** (per replica finish) — low-latency drift check -2. **Periodic** (background task) — full health table, stall detection +### Inference Framework (`src/inference/`) -### Bandit (bandit.py) +| File | Purpose | +|------|---------| +| `esm2_service/esm2_service.py` | ESM2 model loading + batched GPU inference | +| `esm2_service/esm2_client.py` | HTTP / in-process client for ESM2 service | +| `inference_service.py` | Abstract async queue-based service | +| `inference_client.py` | Abstract client protocol | +| `orchestrator.py` | `start_services()` / `start_services_local()` | +| `server.py` | aiohttp HTTP server | +| `dragon_launcher.py` | Dragon HPC backend launcher | +| `utils.py` | Queue helpers, batch sizing | +| `kill_all.py` | Process cleanup | -Thompson-sampling multi-armed bandit for optimization. Two use cases: +### Utilities (`src/utils/`) -**Shard optimizer**: arms = multiplier factors [0.5, 0.75, 1.0, 1.25, 1.5]; reward = throughput +| File | Purpose | +|------|---------| +| `logger.py` | Colored structured logging, metrics recording | +| `workflow.py` | Config loading, GPU enumeration, Dragon policies | -**Scheduling bandit**: arms = cross-stage priority; reward = downstream BP state quality +### Workflows ---- +- **`workflows/esm2_inference/`**: Standalone ESM2 inference workflow + - `run_esm2_infern.py`: entry point (local or server mode) + - `inference.yaml`: config (model path, GPU count, output dir) + - `delta_gpu_sbatch.sh`: SLURM script for Delta HPC -## Key File Organization +- **`workflows/sgdes/`**: SGDES protein engineering workflow + - `sgdes_workflow.py`: main workflow class + - `run_workflow.py`: entry point + - `config.yaml`: campaign configuration -### Campaign Manager Core - -- **campaign_manager.py**: Main class; constructor, config loading, group registration -- **base_workflow.py**: User-defined workflow base class -- **types.py**: `_GroupInfo`, `ResourcePool`, `WorkflowStats` data structures -- **scheduler.py**: SchedulerMixin — two-pass scheduling logic -- **executor.py**: ExecutorMixin — replica launch/completion/GPU assignment -- **monitor_mixin.py**: MonitorMixin — periodic health checks -- **gpu.py**: `detect_gpus()` (CUDA/nvidia-smi probe) and `find_gpus()` (Dragon node enumeration) — both degrade gracefully without Dragon/CUDA -- **sync_wrapper.py**: `CampaignManager` — synchronous wrapper around AsyncCampaignManager (thin thread-based bridge) - -### Optional Features - -- **backpressure.py**: `BackpressureNegotiator` — hysteresis state machine -- **sharder.py**: `Sharder` — buffering and batch dispatch with priority ranking -- **candidate_log.py**: `CandidateLog`, `CandidateHistory` — tracks upstream results -- **monitor.py**: `Monitor`, `DriftEvent` — drift detection logic -- **bandit.py**: `Bandit`, `SchedulingBandit` — Thompson-sampling optimization -- **profiles.py**: `ProfileWeights`, `PROFILES` — candidate ranking profiles -- **metrics.py**: `CampaignMetrics` — in-process event recording (timing, BP transitions, etc.) - -### Utilities - -- **src/utils/logger.py**: Colored structured logging with metrics recording -- **src/utils/workflow.py**: `_expand_env()`, `load_config()`, `find_gpus()`, `make_policies()` -- **src/inference/**: Multi-GPU inference service framework (ESM2 embeddings, HTTP server/client) - -### Examples & Workflows - -- **workflows/run_campaign/esm2_ddsim_campaign/**: Real multi-workflow campaign (DDMd sim + ESM2 inference) - - `run_campaing.py`: main entry point - - `ddmd_workflow.py`: wraps DeepDriveSim DDMd pipeline - - `inference_workflow.py`: ESM2 client workflow - - `config.yaml`: multi-stage campaign config - -- **workflows/run_campaign/dreamer_campaign/**: Emulation campaign using radical.dreamer - - `dreamer_workflow.py`: simulates task execution in-process - - `config.yaml`: plan-based campaign (cm-prototype schema) - - `benchmark.py`: runner with profiling - -- **workflows/esm2_inference/**: Standalone ESM2 service - - `run_esm2_infern.py`: launches inference server with worker pools per GPU +- **`workflows/plot_telemetry.sh`**: plots asyncflow JSONL telemetry to PNG --- @@ -324,24 +165,26 @@ Thompson-sampling multi-armed bandit for optimization. Two use cases: ### Environment Variable Expansion -All config files support `${VAR}` and `$VAR` shell-style references, expanded at load time by `_expand_env()` in `src/utils/workflow.py`: +All config files support `${VAR}` and `$VAR` references, expanded at load time by `_expand_env()`: ```yaml service_python: "${VE_HOME}/esm2/bin/python" outdir: "${SPHERICAL_DIR}/workflows/sgdes/output" +model_path: "${HF_HOME}/models/esm2_t33_650M_UR50D" ``` -Unset variables are preserved as literal strings (fail fast with clear `FileNotFoundError`). - -### Config File Merging +Unset variables raise `KeyError` at load time. -When a workflow entry has `config_file`, that YAML is loaded and merged (scheduling params in the main config take precedence): +### Inference Config Example ```yaml -workflows: - inference: - config_file: inference_specific.yaml # loaded and merged - replicas: 4 # overrides any value in the file +mode: local # "local" (in-process) or "server" (HTTP) +num_services: 2 # one per GPU +model_name: esm2_t33_650M_UR50D +batch_size: 32 +output_dir: outputs/embeddings +metrics_dir: outputs/metrics +stub_sleep_s: 0.1 # debug: bypass real inference, sleep instead ``` --- @@ -350,30 +193,26 @@ workflows: ### Test Organization -- **tests/test_campaign_manager.py**: Core CM logic (scheduling, execution, hooks) -- **tests/test_inference_service.py**: Multi-GPU inference orchestration -- **tests/test_server.py**: aiohttp server endpoints -- **tests/test_client.py**: HTTP client interface -- **tests/test_sgdes_workflow.py**: SGDES protein engineering workflow -- **tests/test_logger.py**: Structured logging utilities -- **tests/test_utils.py**: Config loading and GPU helpers +| File | What it tests | +|------|--------------| +| `tests/test_inference_service.py` | Multi-GPU service orchestration | +| `tests/test_server.py` | aiohttp server endpoints | +| `tests/test_client.py` | HTTP client interface | +| `tests/test_logger.py` | Structured logging utilities | +| `tests/test_utils.py` | Config loading and GPU helpers | +| `tests/test_sgdes_workflow.py` | SGDES protein engineering workflow | ### Test Markers ```bash -# Run only fast tests (skip slow) -pytest -m "not slow" - -# Run only integration tests -pytest -m integration - -# Run only GPU tests (if available) -pytest -m gpu +pytest -m "not slow" # skip slow integration tests +pytest -m integration # only integration tests +pytest -m gpu # only GPU tests (if CUDA available) ``` ### Async Tests -All async tests use `anyio` (not `pytest-asyncio` directly). The standard pattern used throughout the test suite: +All async tests use `anyio`: ```python import pytest @@ -387,161 +226,29 @@ async def test_something(): ... ``` -`asyncio_mode = "auto"` in `pyproject.toml` applies to `pytest-asyncio`; the tests themselves rely on `anyio` with the `anyio_backend` fixture pinning execution to asyncio. - ---- - -## Key Design Patterns - -### Workflow Authoring Pattern - -```python -from src.campaign import BaseWorkflow - -class MyWorkflow(BaseWorkflow): - workflow_id = "my_wf" - - async def run(self, replica_id: str) -> None: - # Do work - result = await compute(self.asyncflow, self.config) - - # Signal downstream groups - await self._signal_done() # broadcast to all dependents - # OR - await self._trigger_dependent("specific_group", replicas=1) # explicit - - async def on_replica_done(self, replica_id, cm, final_state): - if final_state == "done": - # cleanup - pass -``` - -### Campaign Runner Pattern - -```python -from src.campaign import AsyncCampaignManager - -WORKFLOW_REGISTRY = { - "workflow1": Workflow1, - "workflow2": Workflow2, -} - -# Option 1: from config -cm = AsyncCampaignManager.from_config(config, WORKFLOW_REGISTRY) - -# Option 2: manual registration -cm = AsyncCampaignManager(engine="concurrent", total_cpus=128, total_gpus=4) -cm.register_group("wf1", Workflow1, replicas=4, ...) -cm.register_group("wf2", Workflow2, dependencies=["wf1"], ...) - -# Run -await cm.start() -await cm.wait() -await cm.close() - -# Caller is responsible for asyncflow lifecycle -asyncflow = await WorkflowEngine.create(backend) -# ... start telemetry if needed -cm = AsyncCampaignManager.from_config(config, registry, asyncflow=asyncflow) -# ... run campaign -await telemetry.stop() -await asyncflow.shutdown() -``` - -### GPU Assignment - -When `required_gpus > 0`: -1. CM pops GPU IDs from a global free list (FIFO) -2. Injects `assigned_gpu_ids` and `group_gpu_ids` into replica config -3. Builds Dragon `Policy(HOST_NAME, gpu_affinity=[...])` and injects as `self.policies[0]` -4. Returns IDs to free list when replica finishes - ---- - -## Important Implementation Notes - -### No-op Signals in Unit Tests - -Both `_signal_done()` and `_trigger_dependent()` are no-ops when `_cm is None`, making workflows safe to unit test without a CM: - -```python -# Unit test — no CM injected -wf = MyWorkflow(config={...}) -await wf.run("test_0") # signals are no-ops -``` - -### Two Signalling Methods Are Mutually Exclusive per Workflow - -- If a workflow calls `_signal_done()`, downstream groups are determined entirely by config `dependencies` -- If a workflow calls `_trigger_dependent()`, it explicitly decides which group gets replicas -- Mixing both on the same workflow is allowed but unconventional — `_signal_done()` is simpler for data-driven fan-out - -### Scheduler Re-runs on Every State Change - -Every replica completion or signal call triggers `_schedule()`, which: -1. Acquires the lock -2. Runs `_schedule_locked()` (two-pass greedy with dependency + resource checks) -3. Fires resulting replica tasks outside the lock - -This keeps scheduling immediate and fair across groups. - -### Campaign Completion Logic - -The CM completes when: -- All groups with `replicas > 0` are finished -- All sharder buffers are empty -- Groups that were never triggered (dependent groups with 0 finished replicas) are excluded from the check - -This allows purely dependent groups to remain inactive without stalling the campaign. - ---- - -## Performance Tuning - -### Concurrency Caps - -- `max_replicas`: sliding-window concurrency cap per group -- `min_replicas`: guaranteed concurrent slots (priority-ordered across groups in Pass 1) -- If a group has slots but cannot be satisfied by resources, a WARNING is logged - -### Backpressure Tuning - -High `backpressure_high` + low `backpressure_low` gap = frequent oscillation. Recommend: -- `high_water ≈ 1.5 × downstream_total` -- `low_water ≈ 0.5 × downstream_total` - -### Sharder Target Size - -- Too small (e.g., 1): no batching benefit, frequent dispatch overhead -- Too large: causes queue buildup and backpressure throttling -- Recommend: 5–20% of downstream group's total replicas, tuned via A/B testing - -### Monitor Interval - -- Too small (< 10s): log spam, overhead -- Too large (> 120s): miss transient drifts -- Recommend: 20–60s for typical campaigns; shorter (5–15s) for debugging - --- ## Deployment -### Local Testing (Concurrent Backend) +### Local Mode (no GPU required) ```bash -python run_campaing.py --config config.yaml +python workflows/esm2_inference/run_esm2_infern.py \ + --config workflows/esm2_inference/inference.yaml \ + --mode local ``` -Uses `radical.asyncflow` ConcurrentExecutionBackend (pure asyncio, no MPI). +### Delta HPC (SLURM + Dragon) -### HPC Deployment (Dragon Backend) +Set `SBATCH_ACCOUNT` and `HF_TOKEN` before submitting: ```bash -dragon -m workflows/run_campaign/esm2_ddsim_campaign/run_campaing.py \ - --config workflows/run_campaign/esm2_ddsim_campaign/config.yaml +export SBATCH_ACCOUNT=your-project-id +export HF_TOKEN=hf_... +sbatch workflows/esm2_inference/delta_gpu_sbatch.sh ``` -Sets `engine: dragon` in config; CM detects and uses DragonExecutionBackendV3. +The batch script reads `$SBATCH_ACCOUNT` for the `#SBATCH -A` directive. ### Telemetry Visualization @@ -550,17 +257,3 @@ bash workflows/plot_telemetry.sh \ workflows/sgdes/telemetry_output/out.jsonl \ --out-dir plots/sgdes ``` - -Plots asyncflow native JSONL telemetry to workflow dashboard PNG. - -### Campaign Timeline Visualization - -```bash -python workflows/run_campaign/plot_cm_timeline.py \ - slurm-17715157.out \ - --config workflows/run_campaign/config.yaml \ - --out replica_timeline.png -``` - -Parses SLURM log; produces Gantt chart (replicas + resource utilization) + config summary table. - diff --git a/README.md b/README.md index eede016..9622ddb 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,7 @@ HPC workflow orchestration framework for multi-GPU protein inference and enginee ## Features - **AsyncCampaignManager** — async-native orchestrator for concurrent multi-workflow campaigns with priority scheduling, resource pools, and dependency signalling +- **Adaptive Optimization Layers** — opt-in, config-driven: quality routing (Sharder), flow control (Backpressure), surrogate-gated Triage (RUN/DISCARD/ADVANCE), and a BudgetController that keeps spend on plan; drift-driven Replanning. Cross-stage scheduling priority is driven by the **ADR agent layer** (rule / bandit / LLM policies), not an in-CM bandit - **Multi-GPU Inference** — worker pool per GPU with automatic load balancing; aiohttp HTTP server/client - **ESM2 Inference Workflow** — standalone or campaign-embedded ESM2-650M embedding service - **SGDES Workflow** — Structure-Guided Deep Evolution Solver for iterative protein sequence optimisation @@ -20,8 +21,6 @@ HPC workflow orchestration framework for multi-GPU protein inference and enginee ``` spherical/ ├── src/ -│ ├── campaign/ # AsyncCampaignManager + BaseWorkflow + ResourcePool -│ │ └── campaign_manager.py │ ├── inference/ # InferenceService base, orchestrator, server │ │ ├── esm2_service/ # ESM2InferenceService + ESM2Client │ │ ├── inference_client.py @@ -36,18 +35,20 @@ spherical/ │ ├── esm2_inference/ # Standalone ESM2 inference runner │ │ ├── run_esm2_infern.py │ │ └── config.yaml -│ ├── run_campaign/ # Multi-workflow campaign (DDSim + Inference) -│ │ ├── run_campaing.py -│ │ ├── inference_workflow.py -│ │ ├── ddmd_workflow.py -│ │ ├── plot_cm_timeline.py.py # Gantt timeline + resource chart from SLURM log -│ │ └── config.yaml +│ ├── run_campaign/ # Multi-workflow campaigns +│ │ ├── plot_cm_timeline.py # Gantt timeline + resource chart from SLURM log +│ │ ├── esm2_ddsim_campaign/ # real HPC campaign: ESM2 inference + DeepDriveSim (Dragon/GPU) +│ │ │ ├── run_campaing.py · config.yaml · gpu_sbatch.sh +│ │ │ └── inference_workflow.py · ddmd_workflow.py · miniapps_workflow.py · dummy_workflow.py +│ │ └── dreamer_campaign/ # in-process emulation (radical.dreamer) for benchmarking +│ │ ├── run_campaign.py · config*.yaml +│ │ ├── benchmark.py · benchmark_adr.py # feature-flag + ADR-policy benchmarks +│ │ └── plot_optimizations.py · plot_policy_comparison.py · plot_deadline_yield.py │ └── sgdes/ # SGDES protein engineering │ ├── run_workflow.py │ ├── sgdes_workflow.py │ └── config.yaml └── tests/ - ├── test_campaign_manager.py ├── test_inference_service.py ├── test_client.py ├── test_server.py @@ -107,8 +108,17 @@ service_python: "${VE_HOME}/esm2/bin/python" # resolved at load time ### Multi-workflow Campaign +Two campaigns ship under `workflows/run_campaign/`: + ```bash -python workflows/run_campaign/run_campaing.py --config workflows/run_campaign/config.yaml +# Real HPC campaign (ESM2 inference + DeepDriveSim) — Dragon backend, real GPUs: +cd workflows/run_campaign/esm2_ddsim_campaign +dragon run_campaing.py --config config.yaml +# local smoke test (no Dragon): python run_campaing.py --config config.yaml --engine concurrent + +# Emulated campaign (radical.dreamer, in-process) — for benchmarking scheduling policies: +cd workflows/run_campaign/dreamer_campaign +python run_campaign.py --config config.yaml --policy rule # none | rule | bandit | llm ``` Config structure: @@ -118,19 +128,25 @@ resources: total_cpus: 128 total_gpus: 4 +# Optional ADR agent layer — drives cross-stage scheduling priority each tick. +cm: + adr: + policy: rule # none | rule | bandit | llm (override with --policy) + tick_s: 2.0 + workflows: ddsim: replicas: 8 - min_replicas: 2 - max_replicas: 4 + concurrency_floor: 2 + concurrency_cap: 4 priority: 5 required_cpus: 20 dependencies: [] inference: replicas: 16 - min_replicas: 1 - max_replicas: 4 + concurrency_floor: 1 + concurrency_cap: 4 priority: 10 required_cpus: 32 required_gpus: 1 @@ -147,40 +163,6 @@ See [workflows/sgdes/README.md](workflows/sgdes/README.md) for full setup, confi --- -## Campaign Manager - -`AsyncCampaignManager` orchestrates heterogeneous workflow groups inside a single `asyncio` event loop. - -### Authoring a workflow - -```python -from src.campaign import BaseWorkflow - -class MyWorkflow(BaseWorkflow): - workflow_id = "my_wf" - - async def run(self, replica_id: str) -> None: - await do_work(self.asyncflow, self.config) - await self._signal_ready() # unblock dependent groups immediately - - async def on_replica_done(self, replica_id, cm, final_state): - if final_state == "done": - await cm.add_replicas("downstream", n=1) -``` - -### Runner pattern - -```python -cm = AsyncCampaignManager.from_config(config, WORKFLOW_REGISTRY) -await cm.start() -await cm.wait() -await cm.close() -``` - -See [src/campaign/README.md](src/campaign/README.md) for full API reference, scheduler details, and a live run trace. - ---- - ## Extending for New Model Types Subclass `InferenceService` from `src.inference.inference_service`: @@ -250,13 +232,13 @@ bash workflows/plot_telemetry.sh \ ### Campaign Manager replica timeline -`workflows/run_campaign/plot_cm_timeline.py.py` parses a SLURM output log and +`workflows/run_campaign/plot_cm_timeline.py` parses a SLURM output log and produces a Gantt chart of replica execution spans with a resource utilization panel (GPU/CPU in use over time) and a campaign config summary table. ```bash -python workflows/run_campaign/plot_cm_timeline.py.py slurm-.out \ - [--config workflows/run_campaign/config.yaml] \ +python workflows/run_campaign/plot_cm_timeline.py slurm-.out \ + [--config workflows/run_campaign/esm2_ddsim_campaign/config.yaml] \ [--out timeline.png] ``` @@ -272,11 +254,110 @@ from the log lines. **Example**: ```bash -python workflows/run_campaign/plot_cm_timeline.py.py \ +python workflows/run_campaign/plot_cm_timeline.py \ workflows/run_campaign/slurm-17715157.out \ --out replica_timeline.png ``` +### Dreamer campaign timeline (with simulation stats) + +`workflows/run_campaign/dreamer_campaign/plot_dreamer_timeline.py` is a +Dreamer-specific superset of the timeline above: it produces the same Gantt + +resource-utilization rows **plus** a third row of emulation metrics (simulated +makespan per replica, task-ops box plots from the `dreamer-profiles/*.json`, +and a per-workflow stats table). + +```bash +python workflows/run_campaign/dreamer_campaign/plot_dreamer_timeline.py \ + [--profiles-dir dreamer-profiles/] \ + [--config workflows/run_campaign/dreamer_campaign/config.yaml] \ + [--out dreamer_timeline.png] +``` + +The profiles directory is auto-detected next to the log when `--profiles-dir` +is omitted. Use `plot_cm_timeline.py` for non-Dreamer campaigns. + +### Benchmark optimization plots + +`workflows/run_campaign/dreamer_campaign/plot_optimizations.py` reads the +`benchmark_results.json` produced by `benchmark.py` and writes 7 comparison +plots (wall time, pipeline Gantt, cascade funnel, GPU utilization, shard +dispatch, bandit convergence, time-to-target) — one per optimization axis. + +```bash +# 1. produce the results (N runs per configuration) +python workflows/run_campaign/dreamer_campaign/benchmark.py \ + --config workflows/run_campaign/dreamer_campaign/config.yaml \ + --runs 5 --out benchmark_results.json + +# 2. render the plots +python workflows/run_campaign/dreamer_campaign/plot_optimizations.py \ + [--results benchmark_results.json] \ + [--out-dir plots/optimizations] +``` + +Config display names are mapped via `CFG_DISPLAY` and workflow stage labels via +`DISPLAY` at the top of the script; both default to the antigen-cascade names. + +### Budget-control illustration + +`workflows/run_campaign/dreamer_campaign/plot_budget_control.py` renders the +score-cutoff adaptation and burn-ratio convergence for the `budget_control` +benchmark case (a 2-panel figure) from the same `benchmark_results.json`. + +```bash +python workflows/run_campaign/dreamer_campaign/plot_budget_control.py \ + [--results benchmark_results.json] \ + [--out plots/diagrams/budget_control_illustration.png] +``` + +### ADR scheduling-policy comparison + +The dreamer runner can drive scheduling from a swappable `radical.adr` policy +(`--policy {none|rule|bandit|llm}`) and record each decision cycle to JSONL with +`--record`. `plot_policy_comparison.py` then plots the policies side by side — +assigned priority per workflow over cycles — so the rule/llm stable downstream-first +ladder contrasts visually with the bandit's still-exploring (reshuffling) priorities. + +```bash +cd workflows/run_campaign/dreamer_campaign + +# run the same campaign under each policy, recording decisions +python run_campaign.py --policy rule --record +python run_campaign.py --policy bandit --record +python run_campaign.py --policy llm --record # needs OPENROUTER_API_KEY + +# plot them together +python plot_policy_comparison.py \ + adr-decisions-rule.jsonl adr-decisions-bandit.jsonl adr-decisions-llm.jsonl \ + --out plots/policy_comparison.png +``` + +Requires `pip install -e ".[adr]"` (the LLM policy also needs `".[llm]"`). The +policy and recording can also be set in `config.yaml` under `cm.adr`. + +**Batch benchmark (all policies in one job).** `benchmark_adr.py` runs every +policy N times (same metrics shape as `benchmark.py`), writing one results JSON +plus per-cycle decision logs under `adr-logs/`: + +```bash +python workflows/run_campaign/dreamer_campaign/benchmark_adr.py \ + --runs 5 --out benchmark_adr_results.json + # or restrict: --policies none rule bandit +``` + +Cross-stage scheduling priority is owned entirely by the ADR policy (the CM has +no in-loop scheduling bandit); `--policy bandit` runs the same Thompson-sampling +bandit wrapped as an ADR agent. + +`benchmark_adr.py` also supports a **deadline-yield** objective (`--mode +deadline-yield --deadline 60`): instead of time-to-N-leads, it measures how many +terminal leads each policy produces within a fixed wall-clock window (higher is +better — the realistic HPC framing). `plot_deadline_yield.py` renders the +leads-per-policy figure with per-run spread. For the full analysis of when each +policy wins and why downstream-first is hard to beat, see +[docs/scheduling_policy_comparison.md](docs/scheduling_policy_comparison.md). + --- ## Development diff --git a/pyproject.toml b/pyproject.toml index 8d570a2..ef1d41c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,13 +5,13 @@ build-backend = "setuptools.build_meta" [project] name = "spherical" version = "0.1.0" -description = "Multi-GPU Inference Service Framework with Worker Pool Management" +description = "Multi-GPU Inference Service Framework — ESM2 inference service and SGDES workflow" authors = [ - {name = "Masha", email = "masha@example.com"}, + {name = "Masha", email = "mg2347@soe.rutgers.edu"}, ] maintainers = [ - {name = "Masha", email = "masha@example.com"}, + {name = "Masha", email = "mg2347@soe.rutgers.edu"}, ] readme = "README.md" requires-python = ">=3.10" @@ -29,23 +29,23 @@ Homepage = "https://github.com/masha/spherical" Issues = "https://github.com/masha/spherical/issues" [project.optional-dependencies] -# ESM2 model support +# ESM2 model support (torch + transformers) esm2 = [ "torch>=2.0.0,<2.4.0", "transformers>=4.30.0", "numpy>=1.24.0,<2.0.0", ] -# Dragon/RADICAL support +# Dragon/RADICAL HPC backend dragon = [ "dragonhpc>=0.13.2", "rhapsody-py>=0.2.0", "nvidia-ml-py" ] -# SGDES workflow (examples/sgdes) — PyPI-available deps only. +# SGDES workflow — PyPI-available deps only. # PyTorch+CUDA and foldseek/seqkit binaries must be installed separately; -# see examples/sgdes/requirements.txt and delta_env_setup.sh / bridges2_env_setup.sh. +# see workflows/sgdes/requirements.txt and delta_env_setup.sh / bridges2_env_setup.sh. sgdes = [ "biopython>=1.81", "scikit-learn>=1.4.0", @@ -89,15 +89,9 @@ doc = [ "mkdocstrings[python]>=0.24.0", ] -# Plotting/metrics visualization -plotting = [ - "matplotlib>=3.7.0", - "numpy>=1.24.0,<2.0.0", -] - [tool.setuptools.packages.find] where = ["."] -include = ["src*", "examples*"] +include = ["src*"] [tool.pyright] pythonVersion = "3.10" @@ -120,7 +114,7 @@ indent-style = "space" [tool.pytest.ini_options] minversion = "7.0" testpaths = ["tests"] -asyncio_mode = "auto" +asyncio_mode = "strict" markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", "integration: marks tests as integration tests", @@ -129,7 +123,7 @@ markers = [ [tool.coverage.run] source = ["src"] -omit = ["tests/*", "example/*"] +omit = ["tests/*"] [tool.coverage.report] exclude_lines = [ diff --git a/src/campaign/README.md b/src/campaign/README.md deleted file mode 100644 index eea9d34..0000000 --- a/src/campaign/README.md +++ /dev/null @@ -1,428 +0,0 @@ -# Campaign Manager - -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 adaptive cascading -dependency signalling. - ---- - -## Module layout - -``` -src/campaign/ -├── campaign_manager.py # AsyncCampaignManager, CampaignManager, BaseWorkflow, -│ # ResourcePool, WorkflowStats — all in one file -└── __init__.py # re-exports all five public names -``` - ---- - -## Core concepts - -### BaseWorkflow - -All user workflows subclass `BaseWorkflow`. - -```python -class BaseWorkflow: - workflow_id: str = "base" # unique prefix for replica IDs - - 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_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 six objects at construction time: - -| Injected attribute | Type | Purpose | -|--------------------|------|---------| -| `self.config` | `dict` | per-group config section (CM scheduling keys stripped) | -| `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 | - -When GPUs are assigned, the CM also injects two extra keys into `config`: - -| Config key | Value | -|------------|-------| -| `assigned_gpu_ids` | list of GPU IDs assigned to this replica | -| `group_gpu_ids` | all GPU IDs held by the group right now (useful for multi-GPU service init) | - -### Workflow entry point - -Define **either** `run()` or `start()` — not both. The CM detects which one -is overridden at `register_group` time and raises `ValueError` if both or -neither are defined. - -### Workflow groups - -A **group** is a named pool of replicas of the same workflow class. Each group -has: - -| Field | Meaning | -|-------|---------| -| `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` | 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 - -The CM maintains a single `ResourcePool` tracking CPU cores and GPU slots. -Setting a total to `0` disables tracking for that type (unlimited). - -``` -ResourcePool(total_cpus=128, total_gpus=4) - available_cpus=108 available_gpus=3 ← after some allocations -``` - -| Method | Description | -|--------|-------------| -| `can_fit(cpus, gpus)` | `True` when the requested amounts are currently available | -| `allocate(cpus, gpus)` | Decrement available counts (on replica start) | -| `release(cpus, gpus)` | Increment available counts (on replica finish) | -| `usage_str()` | `"cpus=20/128 gpus=1/4"` (used/total, tracked types only) | -| `available_str()` | `"cpus=108/128 gpus=3/4"` | -| `as_dict()` | Full snapshot included in `cm.status()["resources"]` | - -### GPU assignment - -When `required_gpus > 0`, the CM pops GPU IDs from a global free list (FIFO) -and injects them as `assigned_gpu_ids` and `group_gpu_ids` into the replica's -`config`. A Dragon `Policy(HOST_NAME, gpu_affinity=[...])` is also built and -injected as `self.policies[0]`. IDs are returned to the free list when the -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_done`, `trigger_dependent`): - -1. **Pass 1** — guarantee `min_replicas` concurrent slots for all eligible - groups, highest priority first. -2. **Pass 2** — fill remaining capacity up to `max_replicas`, highest priority - first. - -Each pass also gates on `ResourcePool.can_fit()`: a group that has slots under -`max_replicas` but cannot be satisfied by the current resource pool is skipped -and a WARNING is emitted. - -A group is **eligible** when every dependency group is **ready**: - -- **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 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) - 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) - -No changes needed in the dependent workflow itself — it just runs normally. -The CM starts it when an upstream `_signal_done()` fires. - -```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): - # 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. -- 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()`. - ---- - -## Runner pattern - -```python -from src.campaign import AsyncCampaignManager - -WORKFLOW_REGISTRY = {"sim": SimWorkflow, "analysis": AnalysisWorkflow} - -cm = AsyncCampaignManager.from_config(config, WORKFLOW_REGISTRY) -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 -``` - -### Pre-built engine (recommended for telemetry) - -When the caller builds the asyncflow engine itself (to start telemetry before -the CM runs), pass it at construction time: - -```python -asyncflow = await WorkflowEngine.create(backend) -telemetry = await asyncflow.start_telemetry(...) - -cm = AsyncCampaignManager.from_config(config, WORKFLOW_REGISTRY, asyncflow=asyncflow) -await cm.start() -await cm.wait() -await cm.close() - -await telemetry.stop() -await asyncflow.shutdown() -``` - ---- - -## Configuration - -```yaml -# ── Cluster resource budget ────────────────────────────────────────────────── -resources: - total_cpus: 128 - total_gpus: 4 - -# ── Execution backend ──────────────────────────────────────────────────────── -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: - 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: 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`: - -``` -replicas dependencies dependency_threshold priority -min_replicas max_replicas required_cpus required_gpus -``` - ---- - -## API reference - -### `AsyncCampaignManager` - -| Method | Description | -|--------|-------------| -| `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 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_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)` | - -`register_group` key parameters: - -| Parameter | Default | Meaning | -|-----------|---------|---------| -| `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) | -| `required_cpus` | `0` | CPU cores reserved per running replica | -| `required_gpus` | `0` | GPU slots reserved per running replica | -| `dep_threshold` | `1` | Finished-replica count fallback for dependency readiness | - -### `CampaignManager` (sync wrapper) - -Thin synchronous wrapper around `AsyncCampaignManager`. Runs a dedicated -event loop in a background thread so callers without an async context can use -plain blocking calls. Same `from_config` / `register_group` / `start` / -`wait` / `close` / `status` / `stats` API. - -### `BaseWorkflow` - -| Attribute / method | Description | -|--------------------|-------------| -| `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) | -| `_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/__init__.py b/src/campaign/__init__.py deleted file mode 100644 index fdc6623..0000000 --- a/src/campaign/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Campaign management for multi-workflow orchestration.""" - -from .campaign_manager import AsyncCampaignManager -from .base_workflow import BaseWorkflow -from .sync_wrapper import CampaignManager -from .types import ResourcePool, WorkflowStats -from .backpressure import BackpressureNegotiator, BPState -from .candidate_log import CandidateLog, CandidateHistory, StageResult -from .monitor import Monitor, DriftEvent, DriftKind -from .profiles import ProfileWeights, PROFILES, get_profile -from .sharder import Sharder, ShardingSpec -from .bandit import Bandit, BanditArm, shard_bandit, resource_bandit, SchedulingBandit, scheduling_bandit - -__all__ = [ - "AsyncCampaignManager", - "CampaignManager", - "BaseWorkflow", - "ResourcePool", - "WorkflowStats", - "BackpressureNegotiator", - "BPState", - "CandidateLog", - "CandidateHistory", - "StageResult", - "Monitor", - "DriftEvent", - "DriftKind", - "ProfileWeights", - "PROFILES", - "get_profile", - "Sharder", - "ShardingSpec", - "Bandit", - "BanditArm", - "shard_bandit", - "resource_bandit", - "SchedulingBandit", - "scheduling_bandit", -] diff --git a/src/campaign/backpressure.py b/src/campaign/backpressure.py deleted file mode 100644 index ff75a1b..0000000 --- a/src/campaign/backpressure.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Per-edge backpressure hysteresis controller. - -Ported from cm-prototype/src/cm/components/backpressure.py. - -Controls how many replicas are allowed to queue up for a downstream stage. -When the queue depth (triggered but not yet started) crosses high_water, -the controller enters THROTTLE state and the scheduler will not start new -downstream replicas. When depth drops back below low_water it enters WIDEN. - -The hysteresis gap between high_water and low_water prevents oscillation. -""" - -from dataclasses import dataclass -from enum import Enum - - -class BPState(Enum): - HOLD = "hold" # normal — neither throttling nor widening - THROTTLE = "throttle" # queue too deep — block new starts - WIDEN = "widen" # queue drained — allow new starts freely - - -@dataclass -class BackpressureNegotiator: - """Hysteresis state machine for one downstream stage queue.""" - - edge_name: str - high_water: int - low_water: int - state: BPState = BPState.HOLD - - def __post_init__(self) -> None: - if self.low_water >= self.high_water: - raise ValueError( - f"BackpressureNegotiator [{self.edge_name}]: " - f"low_water ({self.low_water}) must be < high_water ({self.high_water})" - ) - - def step(self, depth: int) -> BPState: - """Update state given current queue depth; return new state. - - Pure 3-state function — state is determined solely by depth: - depth ≥ high_water → THROTTLE (queue too deep, stop dispatching) - depth ≤ low_water → WIDEN (queue drained, dispatch more) - otherwise → HOLD (queue balanced, dispatch at normal rate) - - This makes HOLD reachable from both sides: queue rising through the - (low_water, high_water) band gives HOLD on the way to THROTTLE, and - queue draining gives HOLD on the way back to WIDEN. - """ - if depth >= self.high_water: - self.state = BPState.THROTTLE - elif depth <= self.low_water: - self.state = BPState.WIDEN - else: - self.state = BPState.HOLD - return self.state - - def __repr__(self) -> str: - return (f"BackpressureNegotiator({self.edge_name!r}, " - f"hi={self.high_water}, lo={self.low_water}, state={self.state.value})") diff --git a/src/campaign/bandit.py b/src/campaign/bandit.py deleted file mode 100644 index 18171f1..0000000 --- a/src/campaign/bandit.py +++ /dev/null @@ -1,258 +0,0 @@ -""" -Thompson-sampling multi-armed bandit for campaign optimization. - -Two planned use cases (integration wired later via feature flag): - - Shard optimizer - Arms: multiplier factors applied to ShardingSpec.target_size - e.g. [0.50, 0.75, 1.00, 1.25, 1.50] - Reward: throughput — replicas dispatched per second after the shard; - normalized to [0, 1] relative to a rolling max - - Resource optimizer - Arms: per-stage budget reallocation factors - e.g. {"s1": 0.9, "s2": 1.1, ...} encoded as discrete options - Reward: pass-through efficiency — observed / expected trigger fraction; - clamped to [0, 1] - -Algorithm ---------- -Beta-Bernoulli Thompson sampling with continuous reward: - - Prior: Beta(α=1, β=1) — uniform, no preference - Update: α += reward (reward ∈ [0, 1]) - β += 1 - reward - Select: sample each arm from Beta(α, β); - choose arm with the highest sample - -The continuous update degrades gracefully: reward=1.0 is a pure success, -reward=0.0 is a pure failure, values in between are fractional credit. - -Usage ------ - # Create a bandit with discrete multiplier arms - b = Bandit(arms=[0.5, 0.75, 1.0, 1.25, 1.5], seed=42) - - # At each decision point, select an arm - factor = b.select() - - # After observing the outcome, update with a normalized reward - b.update(factor, reward=0.8) - - # Inspect current estimates - print(b.summary()) # {0.5: 0.52, 0.75: 0.61, 1.0: 0.78, ...} - print(b.best()) # 1.0 (arm with highest mean so far) -""" - -import random -from dataclasses import dataclass, field -from typing import Any, Optional - - -@dataclass -class BanditArm: - """One arm of a Beta-Bernoulli bandit.""" - - label: Any - alpha: float = 1.0 # successes + prior - beta: float = 1.0 # failures + prior - - def sample(self, rng: random.Random) -> float: - """Draw a Thompson sample from Beta(alpha, beta).""" - return rng.betavariate(self.alpha, self.beta) - - def update(self, reward: float) -> None: - """Update with *reward* ∈ [0, 1]. Values outside are clamped.""" - reward = max(0.0, min(1.0, reward)) - self.alpha += reward - self.beta += 1.0 - reward - - def reset(self) -> None: - """Return to uninformative uniform prior.""" - self.alpha = 1.0 - self.beta = 1.0 - - @property - def mean(self) -> float: - """Current posterior mean estimate.""" - return self.alpha / (self.alpha + self.beta) - - @property - def pulls(self) -> int: - """Effective number of updates (alpha + beta - 2 initial prior units).""" - return max(0, round(self.alpha + self.beta - 2)) - - def __repr__(self) -> str: - return ( - f"BanditArm({self.label!r} " - f"mean={self.mean:.3f} pulls={self.pulls} " - f"α={self.alpha:.2f} β={self.beta:.2f})" - ) - - -class Bandit: - """Multi-armed bandit with Thompson sampling over a fixed discrete action set. - - Parameters - ---------- - arms: - Iterable of labels (any hashable value) representing the discrete actions. - seed: - Optional RNG seed for reproducibility. - """ - - def __init__(self, arms: list[Any], seed: Optional[int] = None) -> None: - if not arms: - raise ValueError("Bandit requires at least one arm") - self._rng = random.Random(seed) - self._arms: dict[Any, BanditArm] = { - label: BanditArm(label=label) for label in arms - } - - # ── Decision ────────────────────────────────────────────────────────────── - - def select(self) -> Any: - """Thompson sampling: return the label of the arm with the highest sample.""" - return max(self._arms.values(), key=lambda a: a.sample(self._rng)).label - - # ── Learning ────────────────────────────────────────────────────────────── - - def update(self, arm_label: Any, reward: float) -> None: - """Update *arm_label* with *reward* ∈ [0, 1]. - - Unknown labels are silently ignored so callers don't need to guard. - """ - arm = self._arms.get(arm_label) - if arm is not None: - arm.update(reward) - - def reset(self, arm_label: Optional[Any] = None) -> None: - """Reset one arm (or all arms if *arm_label* is None) to the uniform prior.""" - targets = [self._arms[arm_label]] if arm_label is not None else self._arms.values() - for arm in targets: - arm.reset() - - # ── Inspection ──────────────────────────────────────────────────────────── - - def best(self) -> Any: - """Return the label of the arm with the highest posterior mean.""" - return max(self._arms.values(), key=lambda a: a.mean).label - - def summary(self) -> dict[Any, float]: - """Posterior mean estimate for each arm — useful for logging.""" - return {label: arm.mean for label, arm in self._arms.items()} - - def arms(self) -> list[BanditArm]: - """All arms, sorted by label (for deterministic logging).""" - try: - return sorted(self._arms.values(), key=lambda a: a.label) - except TypeError: - return list(self._arms.values()) - - def __repr__(self) -> str: - arm_str = " ".join(repr(a) for a in self.arms()) - return f"Bandit(best={self.best()!r} [{arm_str}])" - - -# ── Preconfigured factory functions ─────────────────────────────────────────── - -def shard_bandit(seed: Optional[int] = None) -> Bandit: - """Bandit for shard-size multiplier selection. - - Arms represent scale factors applied to ShardingSpec.target_size. - Replaces the hardcoded 0.5× / 1.5× multipliers in Sharder._adaptive_size(). - """ - return Bandit(arms=[0.50, 0.75, 1.00, 1.25, 1.50], seed=seed) - - -def resource_bandit(stage_ids: list[str], seed: Optional[int] = None) -> Bandit: - """Bandit for per-stage budget reallocation. - - Each arm is a tuple of (stage_id, factor) pairs encoded as a frozenset, - representing a candidate budget allocation across stages. - In practice the caller constructs the arms based on the plan's - per_stage_band_pct constraints. - """ - return Bandit(arms=stage_ids, seed=seed) - - -class SchedulingBandit: - """Thompson-sampling bandit for cross-stage scheduling priority. - - One BanditArm per pipeline stage. When multiple stages are eligible - simultaneously, rank() returns them sorted by Thompson-sampled Beta value - so the scheduler tries the most-promising stage first. - - Reward signal (fed at replica completion via update()): - WIDEN → 0.8 downstream hungry — this stage's output is needed, keep going - HOLD → 0.7 balanced — good scheduling rate - THROTTLE → 0.2 downstream flooded — back off this stage - none → 0.5 terminal stage or no BP tracking — neutral - - stage_priors: optional per-stage (alpha, beta) warm-start values. Use to - give CPU-only source stages a head start so they are not starved during the - cold-start window before the bandit has accumulated enough observations. - Example: {"s1_ligand_filter": (2.0, 1.0)} → initial mean 0.67 vs 0.5 default. - """ - - def __init__( - self, - stage_names: list[str], - seed: Optional[int] = None, - stage_priors: Optional[dict[str, tuple[float, float]]] = None, - ) -> None: - self._arms: dict[str, BanditArm] = {} - for n in stage_names: - arm = BanditArm(label=n) - if stage_priors and n in stage_priors: - arm.alpha, arm.beta = stage_priors[n] - self._arms[n] = arm - self._rng = random.Random(seed) - - def rank(self, eligible: list) -> list: - """Return eligible groups sorted by Thompson-sampled priority (highest first). - - Groups not in the bandit's arm set (e.g. added dynamically) fall back - to a neutral 0.5 sample so they're still scheduled fairly. - """ - if len(eligible) <= 1: - return eligible - return sorted( - eligible, - key=lambda g: ( - self._arms[g.name].sample(self._rng) - if g.name in self._arms else 0.5 - ), - reverse=True, - ) - - def update(self, stage_name: str, reward: float) -> None: - """Update the arm for *stage_name* with *reward* ∈ [0, 1].""" - arm = self._arms.get(stage_name) - if arm is not None: - arm.update(reward) - - def summary(self) -> dict[str, float]: - """Posterior mean per stage — for logging.""" - return {name: arm.mean for name, arm in self._arms.items()} - - def best(self) -> Optional[str]: - """Stage name with highest posterior mean.""" - if not self._arms: - return None - return max(self._arms, key=lambda n: self._arms[n].mean) - - def __repr__(self) -> str: - arms_str = " ".join( - f"{n}:{arm.mean:.3f}" for n, arm in self._arms.items() - ) - return f"SchedulingBandit(best={self.best()!r} [{arms_str}])" - - -def scheduling_bandit( - stage_names: list[str], - seed: Optional[int] = None, - stage_priors: Optional[dict[str, tuple[float, float]]] = None, -) -> SchedulingBandit: - """Factory for the cross-stage scheduling bandit.""" - return SchedulingBandit(stage_names=stage_names, seed=seed, stage_priors=stage_priors) diff --git a/src/campaign/base_workflow.py b/src/campaign/base_workflow.py deleted file mode 100644 index b6ebf0f..0000000 --- a/src/campaign/base_workflow.py +++ /dev/null @@ -1,83 +0,0 @@ -""" -BaseWorkflow — base class for all campaign workflow implementations. - -Subclass contract ------------------ -- ``workflow_id`` (class attr, str): unique prefix for replica IDs. -- ``run(replica_id)`` **or** ``start(replica_id)``: execute the workflow. - Exactly one must be defined. Async coroutines are awaited directly; - sync functions are run via ``asyncio.to_thread``. -- ``on_replica_done(replica_id, cm, final_state)`` (optional): hook called - by the CM after the entry-point returns or raises. - ``final_state`` is ``"done"`` or ``"failed"``. -""" - -from typing import TYPE_CHECKING, Optional - -if TYPE_CHECKING: - from .campaign_manager import AsyncCampaignManager - - -class BaseWorkflow: - workflow_id: str = "base" - - def __init__( - self, - config: Optional[dict] = None, - _cm: Optional["AsyncCampaignManager"] = None, - _group_name: Optional[str] = None, - asyncflow: Optional[object] = None, - policies: Optional[list] = None, - engine_dragon: Optional[object] = None, - ) -> None: - self.config = config - self._cm = _cm - self._group_name = _group_name - self.asyncflow = asyncflow - # One Dragon Policy per assigned GPU, injected by AsyncCampaignManager. - self.policies: list = policies or [] - self.engine_dragon: Optional[object] = engine_dragon - - async def _trigger_dependent( - self, - name: str, - replicas: int = 1, - **kwargs, - ) -> None: - """Tell the CM to activate a dependent workflow group.""" - if self._cm is not None: - await self._cm.trigger_dependent(name, replicas=replicas, **kwargs) - - async def _trigger_batch( - self, - name: str, - candidates: list[dict], - ) -> None: - """Add multiple candidates to the sharder buffer in one scheduler cycle. - - Fills the buffer with N candidates before dispatch() runs, enabling - meaningful priority ranking. Each dict must contain ``candidate_id`` - and may include ``score``, ``scaffold_class``, ``surrogate_pred``, - ``surrogate_unc``. - """ - if self._cm is not None: - await self._cm.trigger_batch(name, candidates) - - async def _signal_done(self) -> None: - """Signal the CM that this workflow has finished producing data.""" - 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: - """Execute the workflow for one replica. Override in subclasses.""" - raise NotImplementedError( - f"{type(self).__name__}.run() not implemented (replica_id={replica_id!r})" - ) - - def on_replica_done( - self, - replica_id: str, - cm: "AsyncCampaignManager", - final_state: str, - ) -> None: - """Hook called after this replica's entry-point finishes. No-op by default.""" diff --git a/src/campaign/campaign_manager.py b/src/campaign/campaign_manager.py deleted file mode 100644 index 483e9ec..0000000 --- a/src/campaign/campaign_manager.py +++ /dev/null @@ -1,655 +0,0 @@ -""" -AsyncCampaignManager — async-native campaign orchestrator. - -Each workflow replica is an asyncio Task. Supports both -``async def run(replica_id)`` and sync ``def run(replica_id)`` entry points -(sync ones run via ``asyncio.to_thread``). - -Scheduling model ----------------- -Two-pass greedy scheduler on every state change (see scheduler.py): - Pass 1 — guarantee ``min_replicas`` for all eligible groups (highest priority). - Pass 2 — fill remaining capacity up to ``max_replicas`` (highest priority). - -A group becomes eligible either via ``trigger_dependent()`` (explicit) or when -each dependency has ``dep_threshold`` finished replicas (count-based fallback). - -Usage ------ - cm = AsyncCampaignManager.from_config(config, WORKFLOW_REGISTRY) - await cm.start() - await cm.wait() - await cm.close() -""" - -import asyncio -import itertools -from typing import Optional - -from ..utils.logger import Logger -from .backpressure import BackpressureNegotiator, BPState # noqa: F401 (re-exported) -from .metrics import CampaignMetrics -from .bandit import Bandit, BanditArm, shard_bandit, resource_bandit, SchedulingBandit, scheduling_bandit # noqa: F401 -from .base_workflow import BaseWorkflow -from .candidate_log import CandidateLog, CandidateHistory, StageResult # noqa: F401 -from .executor import ExecutorMixin -from .monitor import Monitor, DriftKind # noqa: F401 (re-exported) -from .monitor_mixin import MonitorMixin -from .scheduler import SchedulerMixin -from .sharder import Sharder, ShardingSpec -from .types import _GroupInfo, ResourcePool, WorkflowStats - - -class AsyncCampaignManager(SchedulerMixin, ExecutorMixin, MonitorMixin): - """ - Async campaign manager that orchestrates multiple replicas of one - or more :class:`BaseWorkflow` subclasses. - """ - - def __init__( - self, - max_workers: Optional[int] = None, - engine: str = "concurrent", - total_cpus: int = 0, - total_gpus: int = 0, - num_workers: Optional[int] = None, - debug: bool = False, - asyncflow=None, - engine_dragon=None, - features: Optional[dict] = None, - ) -> None: - self._log = Logger(name="AsyncCampaignManager", use_colors=True) - self._seq = itertools.count() - self._lock = asyncio.Lock() - self._engine_type = engine - self._num_workers = num_workers - self._debug = debug - self._asyncflow = asyncflow - self._engine_dragon = engine_dragon - self._gpu_pool: list[tuple[str, int]] = [] - self._free_gpu_ids: list[int] = [] - self._replica_gpu_assignments: dict[str, list[int]] = {} - self._resources = ResourcePool(total_cpus=total_cpus, total_gpus=total_gpus) - - self._groups: dict[str, _GroupInfo] = {} - self._stats: dict[str, WorkflowStats] = {} - self._all_done = asyncio.Event() - - self._features: dict[str, bool] = features or {} - self._bp: dict[str, BackpressureNegotiator] = {} - self._sharders: dict[str, Sharder] = {} - self._monitor: Optional[Monitor] = None - self._monitor_interval_s: float = 30.0 # overwritten by from_config - self._monitor_task: Optional[asyncio.Task] = None - self._scheduling_bandit: Optional[SchedulingBandit] = None - self._candidate_log: Optional[CandidateLog] = None - self._cand_seq: itertools.count = itertools.count() - self._replica_candidate_assignments: dict[str, str] = {} - self._metrics: CampaignMetrics = CampaignMetrics() - - feat_summary = ", ".join(f"{k}={'on' if v else 'off'}" for k, v in self._features.items()) - self._log.info( - f"AsyncCampaignManager initialised (engine={engine})" - + (f" features: [{feat_summary}]" if feat_summary else "") - ) - - # ------------------------------------------------------------------ - # Alternative constructor - # ------------------------------------------------------------------ - - @classmethod - def from_config( - cls, - config: dict, - workflow_registry: dict[str, type[BaseWorkflow]], - asyncflow=None, - engine_dragon=None, - ) -> "AsyncCampaignManager": - """Build an AsyncCampaignManager from a config dict + workflow registry.""" - res_cfg = config.get("resources", {}) - num_workers = config.get("num_workers") - features = config.get("features", {}) - - cm = cls( - max_workers=config.get("max_workers"), - engine=config.get("engine", "concurrent"), - total_cpus=int(res_cfg.get("total_cpus", 0)), - total_gpus=int(res_cfg.get("total_gpus", 0)), - num_workers=int(num_workers) if num_workers is not None else None, - debug=bool(config.get("debug", False)), - asyncflow=asyncflow, - engine_dragon=engine_dragon, - features=dict(features) if features else {}, - ) - - _cm_keys = { - "replicas", "dependencies", "dependency_threshold", - "min_replicas", "max_replicas", "priority", "required_cpus", "required_gpus", - "concurrency_cap", - "sharding", - } - - for name, wf_cfg in config.get("workflows", {}).items(): - wf_class = workflow_registry.get(name) - if wf_class is None: - cm._log.warning(f"from_config: no class registered for {name!r} — skipping") - continue - - has_deps = bool(wf_cfg.get("dependencies", [])) - default_replicas = 0 if has_deps else 1 - max_replicas = int(wf_cfg.get("max_replicas") or - wf_cfg.get("concurrency_cap") or 0) - cm.register_group( - name=name, - workflow_class=wf_class, - replicas=int(wf_cfg.get("replicas", default_replicas)), - dependencies=list(wf_cfg.get("dependencies", [])), - dep_threshold=int(wf_cfg.get("dependency_threshold", 1)), - min_replicas=int(wf_cfg.get("min_replicas", 0)), - max_replicas=max_replicas, - priority=int(wf_cfg.get("priority", 0)), - required_cpus=int(wf_cfg.get("required_cpus", 0)), - required_gpus=int(wf_cfg.get("required_gpus", 0)), - config={k: v for k, v in wf_cfg.items() if k not in _cm_keys} or None, - ) - - # ── Feature: Backpressure ───────────────────────────────────────────── - if features.get("backpressure"): - for name, wf_cfg in config.get("workflows", {}).items(): - hi = int(wf_cfg.get("backpressure_high") or 0) - lo = int(wf_cfg.get("backpressure_low") or 0) - if hi > 0 and lo > 0 and hi > lo: - cm._bp[name] = BackpressureNegotiator( - edge_name=f"*_to_{name}", high_water=hi, low_water=lo, - ) - cm._log.info(f"Backpressure [{name}]: high_water={hi} low_water={lo}") - - # ── Feature: Sharder ───────────────────────────────────────────────── - if features.get("sharder"): - for name, wf_cfg in config.get("workflows", {}).items(): - sh_raw = wf_cfg.get("sharding") - if sh_raw and isinstance(sh_raw, dict): - spec = ShardingSpec.from_dict(sh_raw) - sharder = Sharder(name=name, spec=spec) - sharder._log_fn = cm._log.info - sharder._metrics_fn = lambda sid, n, sc, pr, _name=name: \ - cm._metrics.record_shard(_name, sid, n, sc, pr) - cm._sharders[name] = sharder - cm._log.info( - f"Sharder [{name}]: target={spec.target_size} " - f"[{spec.min_size}, {spec.max_size}] stratify={spec.stratify}" - ) - - # ── Candidate log (always enabled when any sharder exists) ─────────── - if any(wf_cfg.get("sharding") for wf_cfg in config.get("workflows", {}).values()): - cm._candidate_log = CandidateLog() - cm._log.info("CandidateLog enabled") - - # ── Feature: Monitor ────────────────────────────────────────────────── - if features.get("monitor"): - replan = config.get("replan", {}) - cm._monitor = Monitor( - burn_dev_pct=float(replan.get("budget_burn_deviation_pct", 20.0)), - passthrough_dev_pct=float(replan.get("pass_through_deviation_pct", 25.0)), - recall_floor=float(replan.get("surrogate_recall_floor", 0.90)), - breaches_to_escalate=2, - ) - cm._monitor_interval_s = float( - config.get("cm", {}).get("monitor_interval_s", 30.0) - ) - cm._log.info( - f"Monitor enabled: pass_through_dev={cm._monitor.passthrough_dev_pct}% " - f"budget_dev={cm._monitor.burn_dev_pct}% " - f"escalate_after={cm._monitor.breaches_to_escalate} breaches " - f"interval={cm._monitor_interval_s}s" - ) - - # ── Feature: Scheduling bandit ──────────────────────────────────────── - stage_names = list(config.get("workflows", {}).keys()) - if features.get("bandit") and len(stage_names) > 1: - bandit_seed = config.get("bandit", {}).get("seeds", {}).get("bandit") - # Warm-start: downstream stages get higher initial priority so the bandit - # minimises time-to-target (first N terminal-stage completions). - # Without this, the default FIFO order (insertion = upstream-first) runs - # s1 at full capacity before feeding s4/s5, delaying the first leads. - # - # Priority by pipeline depth (deepest = highest priority): - # depth-0 (source) → Beta(1, 1) mean≈0.50 (lowest — runs last when competing) - # depth-1 → Beta(2, 1) mean≈0.67 - # depth-2 → Beta(3, 1) mean≈0.75 - # depth-3 → Beta(4, 1) mean≈0.80 - # depth-4+ (terminal)→ Beta(5, 1) mean≈0.83 (highest — reaches target fastest) - # Priors are weak (~1-5 effective observations) and quickly overridden by - # the utilisation-based reward signal (see executor.py). - def _dep_depth(name: str, visited: frozenset = frozenset()) -> int: - if name in visited: - return 0 - deps = [d for d in cm._groups[name].dependencies if d in cm._groups] - return 0 if not deps else 1 + max( - _dep_depth(d, visited | {name}) for d in deps - ) - - depths = {n: _dep_depth(n) for n in stage_names if n in cm._groups} - max_alpha = 5 # terminal stage gets Beta(5,1) mean=0.83 - - stage_priors: dict[str, tuple[float, float]] = {} - for name in stage_names: - if name not in cm._groups: - continue - d = depths.get(name, 0) - alpha = float(min(max_alpha, d + 1)) # deeper = higher priority - stage_priors[name] = (alpha, 1.0) - - cm._scheduling_bandit = scheduling_bandit( - stage_names, seed=bandit_seed, stage_priors=stage_priors or None - ) - warm_str = " ".join( - f"{n}=Beta({a:.0f},1)" for n, (a, _) in stage_priors.items() - ) - cm._log.info( - f"SchedulingBandit enabled: {len(stage_names)} stages " - f"[{' '.join(stage_names)}]" - + (f" warm-start: {warm_str}" if warm_str else "") - ) - - return cm - - # ------------------------------------------------------------------ - # Group registration - # ------------------------------------------------------------------ - - @staticmethod - def _resolve_entry_point(workflow_class: type[BaseWorkflow]) -> str: - has_run = workflow_class.run is not BaseWorkflow.run - has_start = "start" in workflow_class.__dict__ or ( - hasattr(workflow_class, "start") - and workflow_class.start is not getattr(BaseWorkflow, "start", None) - ) - if has_run and has_start: - raise ValueError( - f"{workflow_class.__name__} defines both 'run' and 'start' — " - "choose exactly one as the workflow entry point" - ) - if not has_run and not has_start: - raise ValueError(f"{workflow_class.__name__} must define either 'run' or 'start'") - return "run" if has_run else "start" - - def register_group( - self, - name: str, - workflow_class: type[BaseWorkflow], - replicas: int = 1, - dependencies: Optional[list[str]] = None, - dep_threshold: int = 1, - min_replicas: int = 0, - max_replicas: int = 0, - priority: int = 0, - required_cpus: int = 0, - required_gpus: int = 0, - config: Optional[dict] = None, - ) -> None: - entry_point = self._resolve_entry_point(workflow_class) - effective_max = max_replicas if max_replicas > 0 else replicas - - self._groups[name] = _GroupInfo( - name=name, - workflow_class=workflow_class, - replicas=replicas, - dependencies=list(dependencies or []), - group_config=config, - configured_replicas=replicas, - min_replicas=min_replicas, - max_replicas=effective_max, - priority=priority, - required_cpus=required_cpus, - required_gpus=required_gpus, - dep_threshold=dep_threshold, - entry_point=entry_point, - ) - self._stats[name] = WorkflowStats() - self._log.info( - f"Registered group {name!r}: replicas={replicas} " - f"min={min_replicas} max={effective_max} " - f"deps={dependencies or []} dep_threshold={dep_threshold} " - f"resources=(cpus={required_cpus}, gpus={required_gpus})" - ) - - # ------------------------------------------------------------------ - # Lifecycle - # ------------------------------------------------------------------ - - async def _setup_resources(self) -> None: - """Sync CM resource state against the pre-built asyncflow engine. - - Called once from start(). The caller (run_campaign.py) is responsible - for creating the backend and WorkflowEngine before passing asyncflow= - to from_config() / __init__. This method only does CM-side setup: - debug logging, GPU pool discovery, and ResourcePool cap correction. - - Raises RuntimeError if asyncflow was not provided. - """ - if self._asyncflow is None: - raise RuntimeError( - "asyncflow engine not provided — create the backend and " - "WorkflowEngine in your run script and pass asyncflow= to from_config()" - ) - - if self._debug: - try: - from rhapsody import enable_logging - enable_logging(level="DEBUG") - import logging as _logging - _logging.getLogger("radical.asyncflow").setLevel(_logging.WARNING) - _logging.getLogger("asyncio").setLevel(_logging.WARNING) - self._log.warning("rhapsody.enable_logging active") - except ImportError: - self._log.warning("rhapsody.enable_logging not available — skipping") - - from .gpu import find_gpus, detect_gpus - - if self._engine_type == "dragon": - self._gpu_pool = find_gpus() - self._free_gpu_ids = [gid for _, gid in self._gpu_pool] - actual_gpus = len(self._free_gpu_ids) - if actual_gpus != self._resources.total_gpus: - self._log.warning( - f"config total_gpus={self._resources.total_gpus} " - f"!= discovered GPUs={actual_gpus} — " - f"capping ResourcePool to {actual_gpus}" - ) - self._resources.total_gpus = actual_gpus - self._resources.available_gpus = actual_gpus - self._log.info( - f"GPU pool: {len(self._gpu_pool)} GPU(s) — " - + (", ".join(f"{h}:{g}" for h, g in self._gpu_pool) or "none found") - ) - else: - # Concurrent mode: auto-detect CUDA GPUs for assignment tracking. - if not self._free_gpu_ids: - n = detect_gpus() - self._free_gpu_ids = list(range(n)) - if n: - self._log.info(f"Concurrent mode: auto-detected {n} GPU(s) for assignment") - - self._log.info(f"CM ready (engine={self._engine_type})") - - async def start(self) -> None: - """Kick off the campaign — schedule all eligible groups.""" - if not self._groups: - self._all_done.set() - return - await self._setup_resources() - res = self._resources - if res.total_cpus > 0 or res.total_gpus > 0: - self._log.info( - f"Resource pool: total_cpus={res.total_cpus} total_gpus={res.total_gpus}" - ) - if self._monitor is not None: - self._monitor_task = self._start_monitor_loop(self._monitor_interval_s) - await self._schedule() - - async def wait(self, timeout: Optional[float] = None) -> bool: - """Block (async) until all workflow groups have finished.""" - if timeout is not None: - try: - await asyncio.wait_for(asyncio.shield(self._all_done.wait()), timeout=timeout) - return True - except asyncio.TimeoutError: - return False - await self._all_done.wait() - return True - - async def close(self) -> None: - """Release CM resources (asyncflow shutdown left to caller).""" - if self._monitor_task and not self._monitor_task.done(): - self._monitor_task.cancel() - try: - await self._monitor_task - except asyncio.CancelledError: - pass - self._asyncflow = None - self._metrics.finish() - self._log.info("AsyncCampaignManager closed") - - # ------------------------------------------------------------------ - # Public control API - # ------------------------------------------------------------------ - - async def signal_done(self, group_name: str) -> None: - """Signal that *group_name* has produced output; queue 1 replica in each dependent.""" - async with self._lock: - group = self._groups.get(group_name) - if group is None: - return - group.ready = True - 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, - candidate_id: Optional[str] = None, - score: float = 0.0, - surrogate_pred: float = 0.0, - surrogate_unc: float = 0.0, - scaffold_class: str = "", - source_stage: str = "", - ) -> None: - """Queue replicas of the dependent group *name*. - - When candidate signals are supplied (candidate_id, score, …) the call - is treated as a single candidate trigger: - 1. The result is recorded in CandidateLog for *source_stage*. - 2. threshold_top_fraction of *source_stage* gates whether the candidate - enters the sharder buffer. - 3. The sharder ranks the buffer by profile-weighted priority on dispatch. - - When candidate_id is None the call is a count-based trigger (legacy API): - *replicas* anonymous entries are added to the buffer with score=0. - Routes directly to group.replicas when no sharder is registered. - """ - 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 - if config: - group.group_config = {**(group.group_config or {}), **config} - - sharder = self._sharders.get(name) - if sharder is not None: - if candidate_id is not None: - # ── Candidate-aware single-trigger path ────────────────── - enqueue_time = None - if self._candidate_log and source_stage: - result = self._candidate_log.record( - candidate_id, source_stage, score, - surrogate_pred, surrogate_unc, scaffold_class, - ) - enqueue_time = self._candidate_log.get(candidate_id).enqueue_time - src_group = self._groups.get(source_stage) - top_frac = float( - (src_group.group_config or {}).get("threshold_top_fraction", 1.0) - ) if src_group else 1.0 - if not self._candidate_log.passes_threshold( - candidate_id, source_stage, top_frac - ): - cutoff = self._candidate_log.threshold_cutoff(source_stage, top_frac) - self._log.info( - f" Filtered {candidate_id!r} at {source_stage!r}: " - f"score={score:.4f} < cutoff={cutoff:.4f} " - f"(top-{top_frac:.0%})" - ) - result.decision = "filtered" - return - result.decision = "passed" - sharder.receive( - candidate_id=candidate_id, - score=score, - surrogate_pred=surrogate_pred, - surrogate_unc=surrogate_unc, - scaffold_class=scaffold_class, - enqueue_time=enqueue_time, - ) - self._log.info( - f"trigger_dependent: {name!r} candidate={candidate_id!r} " - f"score={score:.4f} → shard buffer (buffered={sharder.buffered})" - ) - else: - # ── Anonymous count-based path (legacy) ────────────────── - for _ in range(replicas): - anon_id = f"cand_{name}_{next(self._cand_seq):06d}" - sharder.receive( - candidate_id=anon_id, - score=score, - surrogate_pred=surrogate_pred, - surrogate_unc=surrogate_unc, - scaffold_class=scaffold_class, - ) - self._log.info( - f"trigger_dependent: {name!r} +{replicas} anonymous → shard buffer " - f"(buffered={sharder.buffered})" - ) - else: - group.replicas += replicas - group.configured_replicas += replicas - if group.status == "done": - group.status = "pending" - self._log.info( - f"trigger_dependent: {name!r} +{replicas} replicas (total={group.replicas})" - ) - await self._schedule() - - async def trigger_candidate( - self, - name: str, - candidate_id: str, - score: float = 0.0, - surrogate_pred: float = 0.0, - surrogate_unc: float = 0.0, - scaffold_class: str = "", - source_stage: str = "", - config: Optional[dict] = None, - ) -> None: - """Convenience wrapper for a single named-candidate trigger. - - Equivalent to trigger_dependent(name, replicas=1, candidate_id=...). - Workflows prefer this over trigger_dependent when they have scored results. - """ - await self.trigger_dependent( - name=name, - replicas=1, - config=config, - candidate_id=candidate_id, - score=score, - surrogate_pred=surrogate_pred, - surrogate_unc=surrogate_unc, - scaffold_class=scaffold_class, - source_stage=source_stage, - ) - - async def trigger_batch( - self, - name: str, - candidates: list[dict], - ) -> None: - """Add N candidates to the sharder buffer in one lock acquisition. - - Each dict in *candidates* must contain ``candidate_id`` and may include - ``score``, ``scaffold_class``, ``surrogate_pred``, ``surrogate_unc``. - - Unlike N sequential trigger_dependent() calls (each of which calls - _schedule() after releasing the lock), this method holds the lock through - all sharder.receive() calls and calls _schedule() exactly once. The - sharder buffer therefore accumulates N candidates before the first - dispatch() runs, enabling meaningful priority ranking across the batch. - - Routes directly to group.replicas when no sharder is registered. - """ - async with self._lock: - group = self._groups.get(name) - if group is None: - self._log.warning(f"trigger_batch: group {name!r} not registered — ignoring") - return - sharder = self._sharders.get(name) - n_added = 0 - for cand in candidates: - candidate_id = str(cand["candidate_id"]) - score = float(cand.get("score", 0.0)) - scaffold_class = str(cand.get("scaffold_class", "")) - surrogate_pred = float(cand.get("surrogate_pred", 0.0)) - surrogate_unc = float(cand.get("surrogate_unc", 0.0)) - if sharder is not None: - sharder.receive( - candidate_id=candidate_id, - score=score, - surrogate_pred=surrogate_pred, - surrogate_unc=surrogate_unc, - scaffold_class=scaffold_class, - ) - else: - group.replicas += 1 - group.configured_replicas += 1 - if group.status == "done": - group.status = "pending" - n_added += 1 - self._log.info( - f"trigger_batch: {name!r} +{n_added} candidates " - f"(buffered={sharder.buffered if sharder else '—'})" - ) - await self._schedule() - - # ------------------------------------------------------------------ - # Status / stats - # ------------------------------------------------------------------ - - def status(self) -> dict: - return { - "resources": self._resources.as_dict(), - "groups": { - name: { - "status": g.status, - "replicas_total": g.replicas, - "replicas_configured": g.configured_replicas, - "replicas_started": g.started_count, - "replicas_running": g.running_count, - "replicas_finished": g.finished_replicas, - "min_replicas": g.min_replicas, - "max_replicas": g.max_replicas, - "required_cpus": g.required_cpus, - "required_gpus": g.required_gpus, - "dep_threshold": g.dep_threshold, - "ready": g.ready, - "dependencies": g.dependencies, - } - for name, g in self._groups.items() - }, - } - - def stats(self) -> dict[str, WorkflowStats]: - return {name: WorkflowStats(**vars(s)) for name, s in self._stats.items()} - - def metrics(self) -> CampaignMetrics: - """Return the live metrics recorder for this campaign run.""" - return self._metrics - - # ------------------------------------------------------------------ - # Internal — schedule dispatch (called outside the lock) - # ------------------------------------------------------------------ - - async def _schedule(self) -> None: - async with self._lock: - to_start = self._schedule_locked() - for group, replica_idx in to_start: - asyncio.get_running_loop().create_task(self._run_replica(group, replica_idx)) diff --git a/src/campaign/candidate_log.py b/src/campaign/candidate_log.py deleted file mode 100644 index ce50f9c..0000000 --- a/src/campaign/candidate_log.py +++ /dev/null @@ -1,161 +0,0 @@ -""" -CandidateLog — per-candidate result history across pipeline stages. - -Workflows call trigger_dependent() with a score; the CM records it here. -Two purposes: - 1. threshold_top_fraction gating — streaming quantile cutoff decides whether - a candidate advances to the next stage. - 2. Sharder priority signals — score / surrogate / uncertainty / enqueue_time - stored here are read by the sharder's profile-based priority scorer. - -Not thread-safe by design: all access goes through the CM's asyncio lock. -""" -from __future__ import annotations - -import time -from dataclasses import dataclass, field -from typing import Optional - -import numpy as np - - -@dataclass -class StageResult: - """One per-stage record in a candidate's history.""" - stage_id: str - score: float - surrogate_pred: float = 0.0 - surrogate_unc: float = 0.0 - scaffold_class: str = "" - decision: str = "" # "passed" | "filtered" | "" - timestamp: float = field(default_factory=time.time) - - -@dataclass -class CandidateHistory: - """Accumulated per-stage results for one candidate.""" - candidate_id: str - scaffold_class: str = "" - enqueue_time: float = field(default_factory=time.time) - results: list[StageResult] = field(default_factory=list) - - @property - def latest_score(self) -> float: - return self.results[-1].score if self.results else 0.0 - - @property - def latest_surrogate_pred(self) -> float: - return self.results[-1].surrogate_pred if self.results else 0.0 - - @property - def latest_surrogate_unc(self) -> float: - return self.results[-1].surrogate_unc if self.results else 0.0 - - def score_at(self, stage_id: str) -> Optional[float]: - """Most recent score recorded for stage_id, or None.""" - for r in reversed(self.results): - if r.stage_id == stage_id: - return r.score - return None - - -class CandidateLog: - """In-memory registry of candidate histories.""" - - def __init__(self) -> None: - self._histories: dict[str, CandidateHistory] = {} - self._stage_scores: dict[str, list[float]] = {} # stage_id → all scores seen - - # ── Write ───────────────────────────────────────────────────────────────── - - def register( - self, - candidate_id: str, - scaffold_class: str = "", - enqueue_time: Optional[float] = None, - ) -> CandidateHistory: - """Create a history record for a new candidate (idempotent).""" - if candidate_id not in self._histories: - self._histories[candidate_id] = CandidateHistory( - candidate_id=candidate_id, - scaffold_class=scaffold_class, - enqueue_time=enqueue_time if enqueue_time is not None else time.time(), - ) - return self._histories[candidate_id] - - def record( - self, - candidate_id: str, - stage_id: str, - score: float, - surrogate_pred: float = 0.0, - surrogate_unc: float = 0.0, - scaffold_class: str = "", - decision: str = "", - ) -> StageResult: - """Append a stage result. Auto-registers the candidate if unknown.""" - if candidate_id not in self._histories: - self.register(candidate_id, scaffold_class=scaffold_class) - result = StageResult( - stage_id=stage_id, - score=score, - surrogate_pred=surrogate_pred, - surrogate_unc=surrogate_unc, - scaffold_class=scaffold_class, - decision=decision, - ) - self._histories[candidate_id].results.append(result) - self._stage_scores.setdefault(stage_id, []).append(score) - return result - - # ── Read ────────────────────────────────────────────────────────────────── - - def get(self, candidate_id: str) -> Optional[CandidateHistory]: - return self._histories.get(candidate_id) - - def threshold_cutoff(self, stage_id: str, top_fraction: float) -> float: - """Running quantile score cutoff for the top_fraction at stage_id. - - Returns -inf when fewer than 2 scores recorded so early candidates - always pass (the distribution isn't established yet). - """ - scores = self._stage_scores.get(stage_id, []) - if len(scores) < 2: - return float("-inf") - arr = np.asarray(scores, dtype=float) - quantile = max(0.0, min(1.0, 1.0 - top_fraction)) - return float(np.quantile(arr, quantile)) - - def passes_threshold( - self, - candidate_id: str, - stage_id: str, - top_fraction: float, - ) -> bool: - """True if the candidate's score at stage_id is in the top_fraction. - - Always returns True when top_fraction >= 1.0 or no score is recorded. - """ - if top_fraction >= 1.0: - return True - history = self._histories.get(candidate_id) - if history is None: - return True - score = history.score_at(stage_id) - if score is None: - return True - return score >= self.threshold_cutoff(stage_id, top_fraction) - - def stage_summary(self, stage_id: str) -> dict: - """Score distribution stats for one stage — for monitor logging.""" - scores = self._stage_scores.get(stage_id, []) - if not scores: - return {"n": 0} - arr = np.asarray(scores, dtype=float) - return { - "n": len(scores), - "mean": round(float(arr.mean()), 4), - "p50": round(float(np.median(arr)), 4), - "p90": round(float(np.quantile(arr, 0.90)), 4), - "max": round(float(arr.max()), 4), - } diff --git a/src/campaign/executor.py b/src/campaign/executor.py deleted file mode 100644 index 347ce17..0000000 --- a/src/campaign/executor.py +++ /dev/null @@ -1,291 +0,0 @@ -""" -ExecutorMixin — replica launch, completion, and monitor logic for AsyncCampaignManager. - -Mixed into AsyncCampaignManager; all methods use ``self`` to access shared -state (``_groups``, ``_resources``, ``_sharders``, ``_monitor``, ``_log``). -""" - -import asyncio -from typing import TYPE_CHECKING - -from .backpressure import BPState -from .gpu import make_policies -from .types import _GroupInfo - -if TYPE_CHECKING: - from .base_workflow import BaseWorkflow - - -def _campaign_complete(groups: dict, sharders: dict) -> bool: - """True when every group has finished all its replicas and all sharder buffers are empty. - - Deliberately does NOT use group.status so it works even when the - deps_done status-propagation chain stalls (e.g. a downstream stage - finishes all replicas before its upstream is marked 'done'). - """ - if not groups: - return False - if not any(g.replicas > 0 for g in groups.values()): - return False # nothing has started yet - for g in groups.values(): - if g.replicas == 0: - continue # not yet activated - if g.running_count > 0: - return False - if g.finished_replicas < g.replicas: - return False - if any(s.buffered > 0 for s in sharders.values()): - return False - return True - - -class ExecutorMixin: - - async def _run_replica(self, group: _GroupInfo, replica_idx: int) -> None: - """Execute one replica of a workflow group.""" - replica_id = f"{group.name}_{replica_idx}" - final_state = "done" - - gpu_ids = self._replica_gpu_assignments.get(replica_id, []) - policies = make_policies(self._gpu_pool, gpu_ids) - - res_tag = "" - if group.required_cpus > 0 or group.required_gpus > 0: - res_tag = f" [cpus={group.required_cpus} gpus={group.required_gpus}]" - if gpu_ids: - host = self._gpu_pool[0][0] if self._gpu_pool else "?" - res_tag += f" [gpu_affinity={gpu_ids} host={host}]" - self._log.info(f" starting replica {replica_id!r}{res_tag}") - - # Build per-replica config: start from group config, layer in GPU and candidate info. - replica_config = group.group_config - if gpu_ids: - replica_config = { - **(replica_config or {}), - "assigned_gpu_ids": gpu_ids, - "group_gpu_ids": list(group.running_gpu_ids), - } - candidate_id = self._replica_candidate_assignments.pop(replica_id, None) - score = None - if candidate_id and self._candidate_log: - h = self._candidate_log.get(candidate_id) - if h: - score = h.latest_score - replica_config = { - **(replica_config or {}), - "candidate_id": candidate_id, - "candidate_score": h.latest_score, # upstream quality score - "candidate_surr": h.latest_surrogate_pred, # surrogate model prediction - "candidate_surr_unc": h.latest_surrogate_unc, # surrogate uncertainty - "candidate_scaffold": h.scaffold_class, # chemical scaffold class - } - else: - replica_config = {**(replica_config or {}), "candidate_id": candidate_id} - self._metrics.record_replica_start(group.name, replica_id, candidate_id=candidate_id, score=score) - wf = group.workflow_class( - config=replica_config, - _cm=self, - _group_name=group.name, - asyncflow=self._asyncflow, - policies=policies, - engine_dragon=self._engine_dragon, - ) - - entry = getattr(wf, group.entry_point) - try: - if asyncio.iscoroutinefunction(entry): - await entry(replica_id) - else: - await asyncio.to_thread(entry, replica_id) - except asyncio.CancelledError: - raise - except BaseException as exc: - self._log.error(f"Replica {replica_id!r} raised: {type(exc).__name__}: {exc}") - final_state = "failed" - - await self._handle_replica_done(wf, group, replica_id, replica_idx, final_state) - - async def _handle_replica_done( - self, - wf: "BaseWorkflow", - group: _GroupInfo, - replica_id: str, - replica_idx: int, - final_state: str, - ) -> None: - """Call workflow hook, then update group state and re-schedule.""" - try: - hook = wf.on_replica_done - if asyncio.iscoroutinefunction(hook): - await hook(replica_id, self, final_state) - else: - hook(replica_id, self, final_state) - except Exception as exc: - self._log.error(f"Replica {replica_id!r} on_replica_done raised: {exc}") - - self._metrics.record_replica_finish(group.name, replica_id, final_state) - await self._on_replica_finished(group, replica_id) - - async def _on_replica_finished(self, group: _GroupInfo, replica_id: str) -> None: - """Update group counters, notify sharders, run monitor, then re-schedule.""" - group_done = False - async with self._lock: - group.finished_replicas += 1 - group.running_count -= 1 - self._resources.release(group.required_cpus, group.required_gpus) - self._stats[group.name].replicas_finished = group.finished_replicas - - if group.finished_replicas >= group.replicas: - # Dependent groups receive triggers incrementally while their - # upstream runs, so finished==replicas fires spuriously after - # every single completion (e.g. 1/1 when only 1 trigger has - # arrived and more are still coming). Only mark truly done - # when all upstream dependencies are also done — i.e., no - # more triggers can arrive from them. - deps_done = not group.dependencies or all( - self._groups.get(d) is not None - and self._groups[d].status == "done" - for d in group.dependencies - ) - if deps_done: - group.status = "done" - group_done = True - - # Update scheduling bandit: reward for this group based on downstream BP. - if self._scheduling_bandit is not None: - downstream_name = (group.group_config or {}).get("trigger_downstream") - bp_ctrl = self._bp.get(downstream_name) if downstream_name else None - if bp_ctrl is not None and bp_ctrl.state in (BPState.THROTTLE, BPState.WIDEN): - # Use BP state only for the extreme cases where it carries a clear - # directional signal: THROTTLE means this stage is flooding its - # downstream (back off), WIDEN means downstream is starved (run more). - sched_reward = 0.8 if bp_ctrl.state == BPState.WIDEN else 0.2 - else: - # BP HOLD (healthy) or no BP at all: use downstream utilisation as a - # fine-grained reward signal. This lets the bandit differentiate - # stages even when BP never fires (all high_waters are above peak queue). - # Terminal stage (no downstream) → max reward; every finish directly - # counts toward the campaign target. - if downstream_name is None: - sched_reward = 1.0 - else: - downstream_grp = self._groups.get(downstream_name) - if downstream_grp is not None and downstream_grp.max_replicas > 0: - util = downstream_grp.running_count / downstream_grp.max_replicas - sched_reward = max(0.2, 1.0 - 0.5 * util) - else: - sched_reward = 0.5 - self._scheduling_bandit.update(group.name, sched_reward) - - freed_gpu_ids = self._replica_gpu_assignments.pop(replica_id, []) - self._free_gpu_ids.extend(freed_gpu_ids) - for gid in freed_gpu_ids: - try: - group.running_gpu_ids.remove(gid) - except ValueError: - pass - - if freed_gpu_ids: - if self._replica_gpu_assignments: - asgn_str = ", ".join( - f"{rid}→{gids}" - for rid, gids in sorted(self._replica_gpu_assignments.items()) - ) - self._log.info( - f" GPU freed: {replica_id!r} released {freed_gpu_ids}" - f" | active: [{asgn_str}]" - f" | free: {sorted(self._free_gpu_ids)}" - ) - else: - self._log.info( - f" GPU freed: {replica_id!r} released {freed_gpu_ids}" - f" | active: (none)" - f" | free: {sorted(self._free_gpu_ids)}" - ) - - release_tag = "" - if group.required_cpus > 0 or group.required_gpus > 0: - release_tag = ( - f" | released cpus={group.required_cpus} gpus={group.required_gpus}" - f" | available: {self._resources.available_str()}" - ) - if freed_gpu_ids: - release_tag += ( - f" [freed gpu_affinity={freed_gpu_ids} | free_gpus={sorted(self._free_gpu_ids)}]" - ) - self._log.info(f"Replica {replica_id!r} finished{release_tag}") - - if group_done: - self._log.info(f"Workflow group {group.name!r} completed") - # Notify downstream sharders: no more triggers from this group, - # so strict-stratify partial tails are safe to flush. - for sh_name, sharder in self._sharders.items(): - sh_group = self._groups.get(sh_name) - if sh_group and group.name in sh_group.dependencies: - sharder.mark_upstream_done() - self._log.info( - f"Sharder [{sh_name}]: upstream {group.name!r} done " - f"— partial tail ({sharder.buffered}) will flush next cycle" - ) - - self._log.info( - f"_on_replica_finished: {group.name!r} - " - f"finished_replicas={group.finished_replicas}/{group.replicas}" - ) - - # ── Monitor: pass-through and budget drift checks ───────────────────── - if self._monitor and group.finished_replicas > 0: - grp_cfg = group.group_config or {} - trigger_name = grp_cfg.get("trigger_downstream") - expected_frac = float(grp_cfg.get("trigger_fraction", 1.0)) - _MIN_PASSTHROUGH_SAMPLE = 10 - if (trigger_name and trigger_name in self._groups and expected_frac < 1.0 - and group.finished_replicas >= _MIN_PASSTHROUGH_SAMPLE): - downstream_total = self._groups[trigger_name].replicas - observed_frac = downstream_total / group.finished_replicas - ev = self._monitor.check_passthrough(group.name, observed_frac, expected_frac) - if ev: - tag = " [ESCALATING]" if self._monitor.is_escalating(ev) else "" - self._log.warning( - f"Monitor [{group.name}] pass_through drift{tag}: " - f"observed={observed_frac:.3f} expected={expected_frac:.3f}" - f" dev={ev.deviation_pct:.1f}% breach={ev.breach_count}" - ) - - budget = float(grp_cfg.get("budget_node_hours") or 0) - if budget > 0 and group.replicas > 0: - pilot = grp_cfg.get("pilot", {}) - nodes = int(pilot.get("nodes", 1)) - walltime_h = float(pilot.get("walltime_h", 1)) - spent_actual = nodes * walltime_h * group.finished_replicas / group.replicas - expected_so_far = budget * group.finished_replicas / group.replicas - ev = self._monitor.check_budget(group.name, spent_actual, expected_so_far) - if ev: - self._log.warning( - f"Monitor [{group.name}] budget drift: " - f"spent={spent_actual:.1f} expected={expected_so_far:.1f} node-hours" - f" dev={ev.deviation_pct:.1f}%" - ) - - # ── Early termination: downstream_input_target ──────────────────────── - # Check BEFORE scheduling so that when the target is hit, _schedule_locked - # sees _all_done=True and returns [] immediately — no new replicas start. - if not self._all_done.is_set(): - for gname, g in self._groups.items(): - target = int((g.group_config or {}).get("downstream_input_target") or 0) - if target > 0 and g.finished_replicas >= target: - self._all_done.set() - self._log.info( - f"Campaign target reached: {gname!r} finished " - f"{g.finished_replicas}/{target} replicas — stopping early" - ) - return # _schedule_locked will be a no-op for all future calls - - await self._schedule() - - async with self._lock: - all_done = _campaign_complete(self._groups, self._sharders) - - if all_done: - self._all_done.set() - self._log.info("All campaign workflow groups finished") diff --git a/src/campaign/gpu.py b/src/campaign/gpu.py deleted file mode 100644 index 8dc3245..0000000 --- a/src/campaign/gpu.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Dragon GPU discovery and policy helpers. - -Both functions degrade gracefully when Dragon is not installed or when -running under the ConcurrentExecutionBackend (local testing). -""" - - -def detect_gpus() -> int: - """Count CUDA-visible GPUs for concurrent-mode assignment tracking. Returns 0 when none found.""" - try: - import torch - return torch.cuda.device_count() - except Exception: - pass - try: - import subprocess - out = subprocess.check_output( - ["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"], text=True - ) - return len(out.strip().splitlines()) - except Exception: - return 0 - - -def find_gpus() -> list[tuple[str, int]]: - """Return [(hostname, gpu_id), ...] for every GPU visible to Dragon.""" - try: - from dragon.native.machine import Node, System - - gpus = [] - for huid in System().nodes: - node = Node(huid) - for gpu_id in node.gpus: - gpus.append((node.hostname, gpu_id)) - return gpus - except Exception: - return [] - - -def make_policies(gpu_pool: list[tuple[str, int]], gpu_ids: list[int]) -> list: - """Build a single Dragon Policy covering all assigned GPU IDs. - - Returns a list with exactly one Policy whose gpu_affinity lists every - assigned GPU. Returns an empty list when gpu_ids is empty or Dragon is - not available. - """ - if not gpu_ids or not gpu_pool: - return [] - try: - from dragon.infrastructure.policy import Policy - - hostname = gpu_pool[0][0] - return [ - Policy( - placement=Policy.Placement.HOST_NAME, - host_name=hostname, - gpu_affinity=list(gpu_ids), - ) - ] - except Exception: - return [] diff --git a/src/campaign/metrics.py b/src/campaign/metrics.py deleted file mode 100644 index cc2df51..0000000 --- a/src/campaign/metrics.py +++ /dev/null @@ -1,230 +0,0 @@ -""" -CampaignMetrics — lightweight in-process event recorder. - -Records timestamped events during a campaign run: - - ReplicaEvent per-replica start/finish with timing - - BPEvent backpressure state transitions per edge - - ShardEvent sharder dispatch metadata (shard quality) - - SchedulingEvent scheduling decisions (which stage chosen, bandit scores) - -All timestamps are wall-clock seconds via time.time(). -to_dict() serialises to JSON-compatible dicts for benchmark aggregation. -""" -import time -from dataclasses import dataclass, field -from typing import Optional - - -@dataclass -class ReplicaEvent: - group: str - replica_id: str - event: str # "start" | "finish" | "failed" - timestamp: float - candidate_id: Optional[str] = None - score: Optional[float] = None - duration_s: Optional[float] = None # set on finish event - - -@dataclass -class BPEvent: - group: str - old_state: str - new_state: str - queue_depth: int - timestamp: float - - -@dataclass -class ShardEvent: - group: str - shard_id: int - n: int - scores: list[float] # raw candidate scores in dispatch order - priorities: list[float] # profile priority scores in dispatch order - timestamp: float - - @property - def mean_score(self) -> float: - return sum(self.scores) / len(self.scores) if self.scores else 0.0 - - @property - def mean_priority(self) -> float: - return sum(self.priorities) / len(self.priorities) if self.priorities else 0.0 - - -@dataclass -class SchedulingEvent: - chosen_groups: list[str] # groups started this cycle (in order) - eligible_groups: list[str] # all eligible groups before selection - bandit_scores: dict[str, float] # Thompson sample per group (empty if no bandit) - timestamp: float - - -class CampaignMetrics: - """Accumulates events during one campaign run.""" - - def __init__(self) -> None: - self.start_time: float = time.time() - self.end_time: Optional[float] = None - self.replica_events: list[ReplicaEvent] = [] - self.bp_events: list[BPEvent] = [] - self.shard_events: list[ShardEvent] = [] - self.scheduling_events: list[SchedulingEvent] = [] - self._replica_starts: dict[str, float] = {} # replica_id → start time - - # ── Writers ─────────────────────────────────────────────────────────────── - - def record_replica_start( - self, - group: str, - replica_id: str, - candidate_id: Optional[str] = None, - score: Optional[float] = None, - ) -> None: - t = time.time() - self._replica_starts[replica_id] = t - self.replica_events.append(ReplicaEvent( - group=group, replica_id=replica_id, event="start", - timestamp=t, candidate_id=candidate_id, score=score, - )) - - def record_replica_finish( - self, - group: str, - replica_id: str, - final_state: str, # "done" | "failed" - ) -> None: - t = time.time() - start = self._replica_starts.pop(replica_id, t) - self.replica_events.append(ReplicaEvent( - group=group, replica_id=replica_id, - event="finish" if final_state == "done" else "failed", - timestamp=t, duration_s=t - start, - )) - - def record_bp_transition( - self, - group: str, - old_state: str, - new_state: str, - queue_depth: int, - ) -> None: - self.bp_events.append(BPEvent( - group=group, old_state=old_state.upper(), new_state=new_state.upper(), - queue_depth=queue_depth, timestamp=time.time(), - )) - - def record_shard( - self, - group: str, - shard_id: int, - n: int, - scores: list[float], - priorities: list[float], - ) -> None: - self.shard_events.append(ShardEvent( - group=group, shard_id=shard_id, n=n, - scores=scores, priorities=priorities, timestamp=time.time(), - )) - - def record_scheduling( - self, - chosen_groups: list[str], - eligible_groups: list[str], - bandit_scores: dict[str, float], - ) -> None: - self.scheduling_events.append(SchedulingEvent( - chosen_groups=chosen_groups, - eligible_groups=eligible_groups, - bandit_scores=bandit_scores, - timestamp=time.time(), - )) - - def finish(self) -> None: - self.end_time = time.time() - - # ── Summary ─────────────────────────────────────────────────────────────── - - @property - def wall_time_s(self) -> float: - end = self.end_time or time.time() - return end - self.start_time - - def group_stats(self) -> dict: - """Per-group throughput and timing summary.""" - from collections import defaultdict - starts: dict[str, list[float]] = defaultdict(list) - finishes: dict[str, list[float]] = defaultdict(list) - durations: dict[str, list[float]] = defaultdict(list) - for ev in self.replica_events: - if ev.event == "start": - starts[ev.group].append(ev.timestamp) - elif ev.event in ("finish", "failed"): - finishes[ev.group].append(ev.timestamp) - if ev.duration_s is not None: - durations[ev.group].append(ev.duration_s) - groups = set(starts) | set(finishes) - out = {} - for g in groups: - s_times = sorted(starts.get(g, [])) - f_times = sorted(finishes.get(g, [])) - durs = durations.get(g, []) - span = (max(f_times) - min(s_times)) if s_times and f_times else 0.0 - out[g] = { - "n_started": len(s_times), - "n_finished": len(f_times), - "first_start": min(s_times) - self.start_time if s_times else None, - "last_finish": max(f_times) - self.start_time if f_times else None, - "span_s": span, - "mean_dur_s": sum(durs) / len(durs) if durs else None, - "throughput_rps": len(f_times) / span if span > 0 else None, - } - return out - - def bp_state_fractions(self) -> dict: - """Per-group fraction of inter-event time spent in each BP state.""" - from collections import defaultdict - durations: dict[str, dict[str, float]] = defaultdict(lambda: defaultdict(float)) - # track last transition time and state - last: dict[str, tuple[float, str]] = {} - for ev in sorted(self.bp_events, key=lambda e: e.timestamp): - if ev.group in last: - prev_t, prev_state = last[ev.group] - durations[ev.group][prev_state] += ev.timestamp - prev_t - last[ev.group] = (ev.timestamp, ev.new_state) - # close open intervals - end = self.end_time or time.time() - for g, (t, state) in last.items(): - durations[g][state] += end - t - result = {} - for g, d in durations.items(): - total = sum(d.values()) or 1.0 - result[g] = {k: v / total for k, v in d.items()} - return result - - def to_dict(self) -> dict: - """Full serialization for JSON storage.""" - return { - "wall_time_s": self.wall_time_s, - "group_stats": self.group_stats(), - "bp_fractions": self.bp_state_fractions(), - "shard_events": [ - {"group": e.group, "shard_id": e.shard_id, "n": e.n, - "mean_score": e.mean_score, "mean_priority": e.mean_priority, - "scores": e.scores, "timestamp": e.timestamp - self.start_time} - for e in self.shard_events - ], - "scheduling_events": [ - {"chosen": e.chosen_groups, "eligible": e.eligible_groups, - "bandit": e.bandit_scores, - "timestamp": e.timestamp - self.start_time} - for e in self.scheduling_events - ], - "replica_events": [ - {"group": e.group, "replica_id": e.replica_id, "event": e.event, - "t": e.timestamp - self.start_time, - "dur": e.duration_s, "score": e.score} - for e in self.replica_events - ], - } diff --git a/src/campaign/monitor.py b/src/campaign/monitor.py deleted file mode 100644 index 60ceee3..0000000 --- a/src/campaign/monitor.py +++ /dev/null @@ -1,109 +0,0 @@ -""" -Drift detection monitor for campaign stage deviations. - -Ported from cm-prototype/src/cm/components/monitor.py. - -Tracks two signals per stage: - pass_through — fraction of upstream completions that triggered downstream - vs. the plan's threshold_top_fraction / trigger_fraction - budget_burn — node-hours spent vs. proportional plan budget - -Each check returns a DriftEvent when the deviation exceeds the configured -percentage threshold. Consecutive breaches increment a counter; once it -reaches breaches_to_escalate the monitor flags escalation. -""" - -from dataclasses import dataclass, field -from enum import Enum -from typing import Optional - - -class DriftKind(Enum): - BUDGET_BURN = "budget_burn" - PASS_THROUGH = "pass_through" - SURROGATE_RECALL = "surrogate_recall" - - -@dataclass -class DriftEvent: - kind: DriftKind - stage_id: str - observed: float - expected: float - deviation_pct: float - breach_count: int = 1 - - -@dataclass -class Monitor: - """Detect plan vs. actual deviations and escalate after N consecutive breaches.""" - - burn_dev_pct: float = 20.0 # % deviation allowed for budget burn - passthrough_dev_pct: float = 25.0 # % deviation allowed for pass-through rate - recall_floor: float = 0.90 # minimum surrogate recall before alert - breaches_to_escalate: int = 2 # consecutive breaches before escalation - - # internal breach counters: (stage_id, DriftKind) → count - _counts: dict = field(default_factory=dict) - - def check_passthrough( - self, stage_id: str, observed: float, expected: float - ) -> Optional[DriftEvent]: - """Fire when |observed - expected| / expected > passthrough_dev_pct.""" - return self._check( - stage_id, DriftKind.PASS_THROUGH, - observed, expected, self.passthrough_dev_pct, - ) - - def check_budget( - self, stage_id: str, spent: float, expected: float - ) -> Optional[DriftEvent]: - """Fire when |spent - expected| / expected > burn_dev_pct.""" - return self._check( - stage_id, DriftKind.BUDGET_BURN, - spent, expected, self.burn_dev_pct, - ) - - def check_recall( - self, stage_id: str, recall: float - ) -> Optional[DriftEvent]: - """Fire when surrogate recall drops below recall_floor.""" - if recall >= self.recall_floor: - self._reset(stage_id, DriftKind.SURROGATE_RECALL) - return None - dev_pct = (self.recall_floor - recall) / self.recall_floor * 100 - key = (stage_id, DriftKind.SURROGATE_RECALL) - self._counts[key] = self._counts.get(key, 0) + 1 - return DriftEvent(DriftKind.SURROGATE_RECALL, stage_id, - recall, self.recall_floor, dev_pct, self._counts[key]) - - def is_escalating(self, ev: DriftEvent) -> bool: - """True when the breach count has reached the escalation threshold.""" - return ev.breach_count >= self.breaches_to_escalate - - def reset(self, stage_id: str, kind: DriftKind) -> None: - self._reset(stage_id, kind) - - # ── internal ───────────────────────────────────────────────────────────── - - def _check( - self, - stage_id: str, - kind: DriftKind, - observed: float, - expected: float, - threshold_pct: float, - ) -> Optional[DriftEvent]: - if expected <= 0: - return None - dev_pct = abs(observed - expected) / expected * 100 - key = (stage_id, kind) - if dev_pct > threshold_pct: - self._counts[key] = self._counts.get(key, 0) + 1 - return DriftEvent(kind, stage_id, observed, expected, - dev_pct, self._counts[key]) - self._reset(stage_id, kind) - return None - - def _reset(self, stage_id: str, kind: DriftKind) -> None: - self._counts.pop((stage_id, kind), None) diff --git a/src/campaign/monitor_mixin.py b/src/campaign/monitor_mixin.py deleted file mode 100644 index dbfb176..0000000 --- a/src/campaign/monitor_mixin.py +++ /dev/null @@ -1,136 +0,0 @@ -""" -MonitorMixin — active periodic monitoring loop for AsyncCampaignManager. - -Mixed into AsyncCampaignManager alongside SchedulerMixin and ExecutorMixin. - -Two monitoring paths --------------------- -Reactive (executor.py): fires after every replica completion — low-latency, - per-event drift check. -Periodic (this module): background task at monitor_interval_s cadence — - logs a full health table, catches stalls (stages not - finishing), and uses a richer pass-through calculation - that includes the shard buffer. - -Pass-through formula (periodic) --------------------------------- - observed = (downstream.replicas + sharder.buffered) / upstream.finished - -Including the buffer gives the "true" pipeline pass-through rate. Without it, -strict-stratify batching makes the ratio appear 0 until the first full batch -dispatches — a false positive in the reactive path. -""" - -import asyncio -from typing import Optional - -from .executor import _campaign_complete - - -class MonitorMixin: - - def _start_monitor_loop(self, interval_s: float) -> "asyncio.Task": - """Spawn the background monitor loop; return the Task.""" - self._monitor_interval_s: float = interval_s - task = asyncio.get_running_loop().create_task(self._run_monitor_loop()) - return task - - async def _run_monitor_loop(self) -> None: - """Periodic health check — runs until all campaign groups are done.""" - while not self._all_done.is_set(): - try: - await asyncio.wait_for( - asyncio.shield(self._all_done.wait()), - timeout=self._monitor_interval_s, - ) - break # campaign finished while we were waiting - except asyncio.TimeoutError: - pass # normal — interval elapsed, run a tick - await self._tick_monitor() - - async def _tick_monitor(self) -> None: - """Snapshot all groups and run health checks. Acquires the lock.""" - async with self._lock: - bandit_info = "" - if self._scheduling_bandit is not None: - bsum = self._scheduling_bandit.summary() - best = self._scheduling_bandit.best() - bandit_info = ( - f" sched_bandit best={best!r} " - + " ".join(f"{n}:{v:.2f}" for n, v in bsum.items()) - ) - self._log.info(f"── Monitor tick ───{bandit_info}") - for name, g in self._groups.items(): - if g.replicas == 0: - continue # not yet activated - if g.replicas > 0 and g.finished_replicas >= g.replicas and g.running_count == 0: - continue # group fully finished — skip to avoid log spam - - completion_pct = ( - g.finished_replicas / g.replicas * 100 if g.replicas else 0 - ) - sharder = self._sharders.get(name) - extra = "" - if sharder and sharder.buffered: - extra += f" buffered={sharder.buffered}" - if sharder: - bsum = sharder.bandit_summary() - if bsum is not None: - best = max(bsum, key=bsum.get) - extra += ( - f" shard_bandit_best={best:.2f}×" - f" [{' '.join(f'{k:.2f}:{v:.2f}' for k, v in bsum.items())}]" - ) - self._log.info( - f" {name}: {g.finished_replicas}/{g.replicas} done " - f"({completion_pct:.0f}%) running={g.running_count}{extra}" - ) - - if not self._monitor or g.finished_replicas == 0: - continue - - grp_cfg = g.group_config or {} - - # ── Pass-through: include shard buffer in downstream count ───── - # Skip until enough upstream completions for a stable ratio. - _MIN_PASSTHROUGH_SAMPLE = 10 - trigger_name = grp_cfg.get("trigger_downstream") - expected_frac = float(grp_cfg.get("trigger_fraction", 1.0)) - if (trigger_name and trigger_name in self._groups and expected_frac < 1.0 - and g.finished_replicas >= _MIN_PASSTHROUGH_SAMPLE): - ds = self._groups[trigger_name] - sharder = self._sharders.get(trigger_name) - buffered = sharder.buffered if sharder else 0 - observed_frac = (ds.replicas + buffered) / g.finished_replicas - ev = self._monitor.check_passthrough( - name, observed_frac, expected_frac - ) - if ev: - tag = " [ESCALATING]" if self._monitor.is_escalating(ev) else "" - self._log.warning( - f" Monitor [{name}] pass_through drift{tag}: " - f"observed={observed_frac:.3f} expected={expected_frac:.3f}" - f" dev={ev.deviation_pct:.1f}% breach={ev.breach_count}" - ) - - # ── Budget burn ─────────────────────────────────────────────── - budget = float(grp_cfg.get("budget_node_hours") or 0) - if budget > 0 and g.replicas > 0: - pilot = grp_cfg.get("pilot", {}) - nodes = int(pilot.get("nodes", 1)) - walltime_h = float(pilot.get("walltime_h", 1)) - spent_actual = nodes * walltime_h * g.finished_replicas / g.replicas - expected_so_far = budget * g.finished_replicas / g.replicas - ev = self._monitor.check_budget(name, spent_actual, expected_so_far) - if ev: - self._log.warning( - f" Monitor [{name}] budget drift: " - f"spent={spent_actual:.1f} expected={expected_so_far:.1f} node-hours" - f" dev={ev.deviation_pct:.1f}%" - ) - - # Fallback: if everything is done but _all_done was never set - # (status-propagation chain stalled), detect it here. - if _campaign_complete(self._groups, self._sharders): - self._all_done.set() - self._log.info("Monitor tick: all groups done — signalling campaign complete") diff --git a/src/campaign/profiles.py b/src/campaign/profiles.py deleted file mode 100644 index f5fc748..0000000 --- a/src/campaign/profiles.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -Named priority profiles for candidate dispatch ordering. - -Each profile is a weight vector over five signals: - score — most recent upstream stage score (primary quality signal) - surrogate — surrogate model prediction - uncertainty — surrogate uncertainty (drives active learning) - age — time since enqueue (anti-starvation bonus) - diversity — novelty bonus for under-represented scaffold classes (MMR-style) - -Higher weight → that signal contributes more to the candidate's priority score -when the sharder ranks its buffer for dispatch. - -Ported from cm-prototype/src/cm/plan/profiles.py; weights are unchanged. -""" -from __future__ import annotations - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class ProfileWeights: - score: float - surrogate: float - uncertainty: float - age: float - diversity: float - - def as_dict(self) -> dict[str, float]: - return { - "score": self.score, - "surrogate": self.surrogate, - "uncertainty": self.uncertainty, - "age": self.age, - "diversity": self.diversity, - } - - -PROFILES: dict[str, ProfileWeights] = { - # Greedy: rank purely by score; tiny age bonus prevents starvation. - "pure_promise": ProfileWeights(score=1.0, surrogate=0.0, uncertainty=0.0, age=0.05, diversity=0.0), - # Active learning: maximize uncertainty reduction; ignore score. - "active_learning": ProfileWeights(score=0.0, surrogate=0.0, uncertainty=1.0, age=0.05, diversity=0.0), - # Balanced: score + surrogate + uncertainty (good for exploration with a model). - "explore_exploit": ProfileWeights(score=0.5, surrogate=0.3, uncertainty=0.4, age=0.05, diversity=0.0), - # Score-weighted with scaffold diversity to avoid chemical echo chambers. - "diverse_top": ProfileWeights(score=0.6, surrogate=0.0, uncertainty=0.1, age=0.05, diversity=0.3), - # Pure diversity: round-robin across scaffold classes. - "round_robin": ProfileWeights(score=0.0, surrogate=0.0, uncertainty=0.0, age=0.05, diversity=1.0), -} - - -def get_profile(name: str) -> ProfileWeights: - if name not in PROFILES: - raise KeyError(f"unknown profile {name!r}; known: {sorted(PROFILES)}") - return PROFILES[name] diff --git a/src/campaign/scheduler.py b/src/campaign/scheduler.py deleted file mode 100644 index eeb900e..0000000 --- a/src/campaign/scheduler.py +++ /dev/null @@ -1,297 +0,0 @@ -""" -SchedulerMixin — two-pass greedy scheduling logic for AsyncCampaignManager. - -Mixed into AsyncCampaignManager; all methods use ``self`` to access shared -state (``_groups``, ``_resources``, ``_bp``, ``_sharders``, ``_log``). - -Scheduling model ----------------- -Every state change (replica finished, trigger received) calls ``_schedule``, -which holds the lock, calls ``_schedule_locked``, then fires the resulting -tasks outside the lock. - -Pass 1 — guarantee ``min_replicas`` for all eligible groups (highest priority). -Pass 2 — fill remaining capacity up to ``max_replicas`` (highest priority). - -A group is eligible when its dependencies are satisfied and it has replicas -waiting to be started. -""" - -from .backpressure import BPState -from .types import _GroupInfo - - -class SchedulerMixin: - - # ------------------------------------------------------------------ - # Dependency and resource checks (must be called under self._lock) - # ------------------------------------------------------------------ - - def _deps_satisfied_locked(self, group: _GroupInfo) -> bool: - """True when every dependency group is considered ready. - - Ready means any of: - - dep.status == "done" (all replicas finished — always satisfies, regardless of threshold) - - explicit ``signal_ready()`` call (workflow-driven), OR - - ``dep_threshold`` or more finished replicas (count-based fallback). - - The status=="done" path enables true sequential waterfall when dep_threshold - is set very high: the downstream stage is blocked until the upstream group - fully completes every replica, not just the first dep_threshold ones. - """ - for dep_name in group.dependencies: - dep = self._groups.get(dep_name) - if dep is None: - return False - if dep.status == "done": - continue # fully completed upstream always satisfies dependency - if not dep.ready and dep.finished_replicas < group.dep_threshold: - return False - return True - - def _can_start_locked(self, group: _GroupInfo) -> bool: - """True if one more replica of *group* can be started right now.""" - if group.status == "done": - return False - if group.started_count >= group.replicas: - return False - # max_replicas == 0 means "no explicit cap — use replicas count". - 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 - if not self._resources.can_fit(group.required_cpus, group.required_gpus): - return False - return True - - def _allocate_locked(self, group: _GroupInfo) -> int: - """Record one replica start for *group*; update counters; return replica idx.""" - idx = group.started_count - group.started_count += 1 - group.running_count += 1 - self._resources.allocate(group.required_cpus, group.required_gpus) - self._stats[group.name].replicas_started = group.started_count - replica_id = f"{group.name}_{idx}" - # Assign the next pending candidate ID to this replica (FIFO from shard dispatch). - if group._pending_candidates: - self._replica_candidate_assignments[replica_id] = group._pending_candidates.popleft() - gpu_ids = [ - self._free_gpu_ids.pop(0) - for _ in range(group.required_gpus) - if self._free_gpu_ids - ] - self._replica_gpu_assignments[replica_id] = gpu_ids - group.running_gpu_ids.extend(gpu_ids) - if gpu_ids: - self._log.info( - f" GPU assign: {replica_id!r} → GPU(s) {gpu_ids}" - f" | free: {sorted(self._free_gpu_ids)}" - ) - return idx - - # ------------------------------------------------------------------ - # Sharder flush (must be called under self._lock) - # ------------------------------------------------------------------ - - def _flush_sharders_locked(self) -> None: - """Drain shard buffers into their groups' runnable replica queues. - - dispatch() returns a priority-ordered list of candidate IDs. Each ID - is appended to group._pending_candidates; _allocate_locked pops them - FIFO so replica_idx → candidate_id mapping is stable. - """ - for name, sharder in self._sharders.items(): - if sharder.buffered <= 0: - continue - g = self._groups.get(name) - if g is None: - continue - bp = self._bp.get(name) - cap = g.max_replicas if g.max_replicas > 0 else max(g.replicas, 1) - occupancy = min(1.0, g.running_count / cap) - # Collect scaffold classes of currently-running replicas for diversity scoring. - running_scaffolds: set[str] = set() - if self._candidate_log: - for rid, cid in self._replica_candidate_assignments.items(): - if rid.startswith(f"{name}_"): - h = self._candidate_log.get(cid) - if h and h.scaffold_class: - running_scaffolds.add(h.scaffold_class) - dispatched = sharder.dispatch( - bp=bp, occupancy=occupancy, running_scaffolds=running_scaffolds or None - ) - n = len(dispatched) - if n > 0: - g.replicas += n - g.configured_replicas += n - g._pending_candidates.extend(dispatched) - if g.status == "done": - g.status = "pending" - self._log.info( - f"Sharder [{name}]: dispatched {n} → runnable " - f"(buffered={sharder.buffered} queue_depth={g.replicas - g.started_count})" - ) - elif sharder.buffered > 0: - self._log.info( - f"Sharder [{name}]: holding {sharder.buffered} " - f"(bp={bp.state.value if bp else 'none'} occupancy={occupancy:.2f})" - ) - - # ------------------------------------------------------------------ - # Main scheduler (must be called under self._lock) - # ------------------------------------------------------------------ - - def _schedule_locked(self) -> list[tuple[_GroupInfo, int]]: - """Two-pass greedy scheduler. Must be called under ``self._lock``. - - Returns a list of (group, replica_idx) pairs to start. - """ - to_start: list[tuple[_GroupInfo, int]] = [] - - # Stop scheduling immediately after early termination or natural completion. - if self._all_done.is_set(): - return to_start - - # ── Sharder: flush buffers into runnable queues ────────────────────── - if self._sharders: - self._flush_sharders_locked() - - # ── Backpressure: refresh state for all controlled groups ──────────── - if self._features.get("backpressure"): - for bp_name, bp_ctrl in self._bp.items(): - if bp_name in self._groups: - g = self._groups[bp_name] - queue_depth = max(0, g.replicas - g.started_count) - old_state = bp_ctrl.state - bp_ctrl.step(queue_depth) - if bp_ctrl.state != old_state: - # Suppress the trivial HOLD→WIDEN at startup (empty queue - # always triggers this; it carries no actionable information). - startup_widen = ( - old_state.value == "hold" - and bp_ctrl.state.value == "widen" - and queue_depth == 0 - ) - if not startup_widen: - level = ( - self._log.warning - if bp_ctrl.state == BPState.THROTTLE - else self._log.info - ) - level( - f"Backpressure [{bp_name}]: " - f"{old_state.value} → {bp_ctrl.state.value}" - f" (queue_depth={queue_depth})" - ) - self._metrics.record_bp_transition( - bp_name, old_state.value, bp_ctrl.state.value, queue_depth - ) - - eligible = [ - g - for g in self._groups.values() - if g.status != "done" - and g.started_count < g.replicas - and self._deps_satisfied_locked(g) - ] - - for g in eligible: - if g.status == "pending": - g.status = "running" - self._log.info(f"Group {g.name!r} is now eligible — status → running") - - if self._scheduling_bandit is not None: - eligible = self._scheduling_bandit.rank(eligible) - else: - # Sort by group priority (higher = scheduled first). - # stable sort: equal-priority groups keep registration order (FIFO). - eligible = sorted(eligible, key=lambda g: -g.priority) - - # Pass 1: guarantee min_replicas. - for g in eligible: - deficit = g.min_replicas - g.running_count - for _ in range(deficit): - if not self._can_start_locked(g): - break - idx = self._allocate_locked(g) - to_start.append((g, idx)) - - # Pass 2: fill remaining capacity up to max_replicas. - for g in eligible: - while self._can_start_locked(g): - idx = self._allocate_locked(g) - to_start.append((g, idx)) - - # Warn about groups stalled on resources. - for g in eligible: - if ( - g.started_count < g.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) - ): - self._log.warning( - f"Group {g.name!r} stalled — waiting for resources " - f"(needs cpus={g.required_cpus} gpus={g.required_gpus} " - f"available: {self._resources.available_str()})" - ) - - if to_start: - bandit_scores: dict = {} - if self._scheduling_bandit is not None: - bandit_scores = { - g.name: self._scheduling_bandit._arms[g.name].sample( - self._scheduling_bandit._rng - ) - for g in eligible if g.name in self._scheduling_bandit._arms - } - self._metrics.record_scheduling( - chosen_groups=[g.name for g, _ in to_start], - eligible_groups=[g.name for g in eligible], - bandit_scores=bandit_scores, - ) - - def _gpu_tag(g, idx): - ids = self._replica_gpu_assignments.get(f"{g.name}_{idx}", []) - return f"gpu={ids}" if ids else "" - - summary = ", ".join( - f"{g.name}_{idx}" + (f"[{_gpu_tag(g, idx)}]" if _gpu_tag(g, idx) else "") - for g, idx in to_start - ) - _used: set[str] = set() - _abbrevs: dict[str, str] = {} - for g in self._groups.values(): - ch = next( - (c.upper() for c in g.name if c.upper() not in _used), - chr(ord("A") + len(_abbrevs)), - ) - _abbrevs[g.name] = ch - _used.add(ch) - viz = "".join(_abbrevs[g.name] * g.running_count for g in self._groups.values()) - buf_str = {n: s.buffered for n, s in self._sharders.items() if s.buffered} - col = max(len(g.name) for g in self._groups.values()) + 2 - group_lines = "\n".join( - f" {g.name:<{col}} run={g.running_count:<3} " - f"done={g.finished_replicas}/{g.replicas}" - + (f" buf={buf_str[g.name]}" if g.name in buf_str else "") - for g in self._groups.values() - ) - res_line = f" {self._resources.usage_str()}" - bandit_line = "" - if self._scheduling_bandit is not None: - bsum = self._scheduling_bandit.summary() - best = self._scheduling_bandit.best() - bandit_line = ( - "\n sched_bandit_best=" + repr(best) + " " - + " ".join(f"{n}:{v:.2f}" for n, v in bsum.items()) - ) - self._log.info( - f"Scheduling: [{summary}] viz=[{viz}]\n" - + group_lines + "\n" - + res_line - + bandit_line - ) - - return to_start diff --git a/src/campaign/sharder.py b/src/campaign/sharder.py deleted file mode 100644 index 9cb742e..0000000 --- a/src/campaign/sharder.py +++ /dev/null @@ -1,479 +0,0 @@ -""" -Per-stage sharder: buffers upstream triggers and batch-dispatches downstream. - -The sharder sits between the upstream producer (trigger_dependent signals) and -the downstream execution queue (group.replicas counter). Incoming candidates -accumulate in a buffer; each scheduling cycle the CM calls dispatch() to move -a priority-ordered batch from the buffer into the runnable queue. - -Priority scoring ----------------- -Each candidate entry carries (score, surrogate_pred, surrogate_unc, -scaffold_class, enqueue_time). dispatch() ranks buffered candidates under the -active ProfileWeights before selecting the top-N to release. - -Five signals, all percentile-ranked within the buffer: - score — upstream stage score - surrogate — surrogate model prediction - uncertainty — surrogate uncertainty - age — time since enqueue (anti-starvation) - diversity — novelty: less-common scaffold class in buffer → higher score; - scaffolds already running in the downstream group score 0 - -threshold_top_fraction ----------------------- -Set top_fraction < 1.0 on a stage to gate candidates at trigger time. -Only candidates whose score is in the top fraction of all scores seen so far -at that stage are accepted into the buffer. Gating is handled by the CM -(trigger_dependent) before receive() is called; the sharder itself does not -re-check the threshold. - -Backpressure and stratify semantics are unchanged from the integer-buffer version. - -Bandit mode (use_bandit: true) -------------------------------- -Thompson-sampling bandit selects the BP multiplier from -[0.50, 0.75, 1.00, 1.25, 1.50]. Feedback is the BP state of the PREVIOUS -dispatch cycle. In strict mode the selected factor is clamped to ≥ 1.0 so -the bandit's reward is consistent with what the sharder actually dispatched -(strict mode never sends fewer than target_size). -""" - -import time -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Callable, Optional - -import numpy as np - -from .backpressure import BackpressureNegotiator, BPState - -if TYPE_CHECKING: # pragma: no cover - from .bandit import Bandit - from .profiles import ProfileWeights - - -# ── Candidate entry ──────────────────────────────────────────────────────────── - -@dataclass -class _CandidateEntry: - candidate_id: str - score: float = 0.0 - surrogate_pred: float = 0.0 - surrogate_unc: float = 0.0 - scaffold_class: str = "" - enqueue_time: float = 0.0 # set by Sharder.receive() via its now_fn - - -# ── Percentile-rank helper ───────────────────────────────────────────────────── - -def _percentile_rank(values: list[float]) -> list[float]: - """Map values to percentile ranks in [0, 1]. Stable for ties. - - Special case: a single value returns [1.0] so a one-element buffer - doesn't zero out the score/surrogate/uncertainty contributions to - priority. (np.linspace(0, 1, 1) is [0.0], which would otherwise - suppress real signal magnitude on tail-of-campaign single dispatches.) - """ - if not values: - return [] - if len(values) == 1: - return [1.0] - arr = np.asarray(values, dtype=float) - order = np.argsort(arr, kind="stable") - ranks = np.empty_like(order, dtype=float) - ranks[order] = np.linspace(0.0, 1.0, len(arr)) - return ranks.tolist() - - -# ── ShardingSpec ─────────────────────────────────────────────────────────────── - -@dataclass -class ShardingSpec: - """Plan-level sharding bounds and dispatch strategy for one stage.""" - target_size: int = 50 - min_size: int = 1 - max_size: int = 200 - stratify: str = "soft" # off | soft | strict - top_fraction: float = 1.0 # threshold_top_fraction gate (1.0 = no gate) - profile: str = "diverse_top" - use_bandit: bool = False - bandit_seed: Optional[int] = None - - @classmethod - def from_dict(cls, d: dict) -> "ShardingSpec": - return cls( - target_size = int(d.get("target_size", 50)), - min_size = int(d.get("min_size", 1)), - max_size = int(d.get("max_size", 200)), - stratify = str(d.get("stratify", "soft")), - top_fraction = float(d.get("top_fraction", 1.0)), - profile = str(d.get("profile", "diverse_top")), - use_bandit = bool(d.get("use_bandit", False)), - bandit_seed = d.get("bandit_seed"), - ) - - -# ── Sharder ──────────────────────────────────────────────────────────────────── - -@dataclass -class Sharder: - """Buffer + priority-dispatch controller for one downstream stage.""" - name: str - spec: ShardingSpec - - _buf: list = field(default_factory=list, init=False) - _upstream_done: bool = field(default=False, init=False) - _last_factor: Optional[float] = field(default=None, init=False) - _last_occupancy: float = field(default=0.0, init=False) - _shard_seq: int = field(default=0, init=False) - _bandit: Optional["Bandit"] = field(default=None, init=False) - _profile: Optional["ProfileWeights"] = field(default=None, init=False) - _log_fn: Optional[Callable[[str], None]] = field(default=None, init=False) - _metrics_fn: Optional[Callable[[int, int, list, list], None]] = field(default=None, init=False) - _now_fn: Callable[[], float] = field(default=time.time, init=False) - - def _log(self, msg: str) -> None: - if self._log_fn is not None: - self._log_fn(msg) - - def __post_init__(self) -> None: - if self.spec.use_bandit: - from .bandit import shard_bandit - self._bandit = shard_bandit(seed=self.spec.bandit_seed) - from .profiles import get_profile - self._profile = get_profile(self.spec.profile) - - # ── Configuration ───────────────────────────────────────────────────────── - - def set_profile(self, profile_name: str) -> None: - """Switch the active priority profile (e.g., on surrogate redeploy - or when the downstream input target nears completion). - """ - from .profiles import get_profile - self.spec.profile = profile_name - self._profile = get_profile(profile_name) - - def set_now_fn(self, fn: Callable[[], float]) -> None: - """Inject a clock (default time.time) so the age signal and enqueue - timestamps work under simulated time (e.g., SimPy env.now).""" - self._now_fn = fn - - # ── Producer side ───────────────────────────────────────────────────────── - - def receive( - self, - candidate_id: str, - score: float = 0.0, - surrogate_pred: float = 0.0, - surrogate_unc: float = 0.0, - scaffold_class: str = "", - enqueue_time: Optional[float] = None, - ) -> None: - """Accept one candidate into the buffer.""" - self._buf.append(_CandidateEntry( - candidate_id=candidate_id, - score=score, - surrogate_pred=surrogate_pred, - surrogate_unc=surrogate_unc, - scaffold_class=scaffold_class, - enqueue_time=enqueue_time if enqueue_time is not None else self._now_fn(), - )) - - def mark_upstream_done(self) -> None: - """Signal that no more triggers will arrive. - - Behavior on the next dispatch(): - - strict mode: flush ALL remaining candidates (priority-ordered) so - the stage doesn't deadlock waiting for a full target_size batch - that will never arrive. - - soft / off mode: no effect — the campaign terminates via natural - completion or _all_done elsewhere in the CM; flushing here would - release low-priority candidates that should never run. - """ - self._upstream_done = True - - def clear(self) -> None: - """Drop all buffered candidates. Used when a stage is cancelled or - when the CM needs to reset state (e.g., after a drain → resume).""" - n = len(self._buf) - self._buf.clear() - self._log(f" Sharder [{self.name}] cleared: dropped {n} buffered") - - @property - def buffered(self) -> int: - return len(self._buf) - - # ── Priority scoring ─────────────────────────────────────────────────────── - - def _score_entries( - self, - entries: list[_CandidateEntry], - now: float, - running_scaffolds: Optional[set] = None, - ) -> list[float]: - """Compute one priority score per entry under the active profile. - - Signals are percentile-ranked within the buffer so all profiles operate - on a common [0, 1] scale regardless of raw value magnitudes. - Diversity: scaffold classes already running downstream score 0; - within the buffer, less common scaffold class → higher diversity. - """ - if not entries: - return [] - w = self._profile - already = running_scaffolds or set() - - score_n = _percentile_rank([e.score for e in entries]) - surr_n = _percentile_rank([e.surrogate_pred for e in entries]) - unc_n = _percentile_rank([e.surrogate_unc for e in entries]) - - ages = [max(0.0, now - e.enqueue_time) for e in entries] - age_max = max(ages) if max(ages) > 0 else 1.0 - age_n = [a / age_max for a in ages] - - counts: dict[str, int] = {} - for e in entries: - counts[e.scaffold_class] = counts.get(e.scaffold_class, 0) + 1 - max_count = max(counts.values()) if counts else 1 - diversity = [ - 0.0 if e.scaffold_class in already - else 1.0 - (counts[e.scaffold_class] - 1) / max_count - for e in entries - ] - - return [ - w.score * score_n[i] - + w.surrogate * surr_n[i] - + w.uncertainty * unc_n[i] - + w.age * age_n[i] - + w.diversity * diversity[i] - for i in range(len(entries)) - ] - - # ── Bandit helpers ───────────────────────────────────────────────────────── - - def _update_bandit(self, bp_state: Optional[BPState]) -> None: - """Feed reward for the previous dispatch factor. - - Reward formula: - THROTTLE → 0.1 + 0.3 × occupancy - (previous factor flooded the queue; partial credit for high - occupancy so we don't converge to the smallest factor just - to avoid any THROTTLE.) - HOLD → 0.8 (queue depth healthy, flat — no occupancy signal) - WIDEN → 0.2 + 0.7 × occupancy - (high occupancy means the previous dispatch kept the stage - well-fed despite the queue draining → good factor choice; - low occupancy means the stage was starved.) - - Using occupancy rather than raw WIDEN/HOLD lets the bandit converge - even when BP stays WIDEN permanently (queue drains between dispatches), - which happens when makespan < dispatch interval. - """ - if self._bandit is None or self._last_factor is None: - return - if bp_state == BPState.THROTTLE: - reward = 0.1 + 0.3 * self._last_occupancy - elif bp_state == BPState.HOLD: - reward = 0.8 - else: # WIDEN or no BP - reward = 0.2 + 0.7 * self._last_occupancy - prev_factor = self._last_factor - self._bandit.update(prev_factor, reward) - self._last_factor = None - # Defer string formatting until we know logging is wired up — large - # bandit summaries are expensive to stringify in hot dispatch loops. - if self._log_fn is not None: - bsum = self._bandit.summary() - best = max(bsum, key=bsum.get) - reward_src = ( - f"throttle+occ={self._last_occupancy:.2f}→{reward:.2f}" if bp_state == BPState.THROTTLE - else f"hold→0.80" if bp_state == BPState.HOLD - else f"occ={self._last_occupancy:.2f}→{reward:.2f}" - ) - self._log( - f" Sharder [{self.name}] shard_bandit update: " - f"factor={prev_factor:.2f} reward={reward_src} " - f"best={best:.2f}× [{' '.join(f'{k:.2f}:{v:.2f}' for k, v in bsum.items())}]" - ) - - def _bp_factor(self, bp_state: Optional[BPState]) -> float: - """Select dispatch multiplier from bandit or fixed BP→factor mapping. - - In strict mode the factor is floored at 1.0 so the bandit's selected - arm matches what dispatch() actually emits (strict mode never sends - fewer than target_size). Without this clamp, bandit selections of - 0.5 or 0.75 would be silently promoted to 1.0 by the strict-mode - floor in dispatch(), and the bandit's posterior would inflate the - value of small arms that had no real effect. - """ - if self._bandit is not None: - factor = self._bandit.select() - if self.spec.stratify == "strict": - factor = max(factor, 1.0) - self._last_factor = factor - return factor - # Fixed mapping - if bp_state == BPState.THROTTLE: - return 1.0 if self.spec.stratify == "strict" else 0.5 - if bp_state == BPState.WIDEN: - return 1.5 - return 1.0 - - def bandit_summary(self) -> Optional[dict]: - if self._bandit is None: - return None - return self._bandit.summary() - - # ── Consumer side ────────────────────────────────────────────────────────── - - def adaptive_size( - self, - bp: "BackpressureNegotiator | None", - occupancy: float, - ) -> int: - """Compute dispatch batch size for this scheduling cycle. - - Note: this method has a side effect when use_bandit=True — it consumes - one bandit arm selection via _bp_factor. Call it only from dispatch() - or from tests that intend to advance the bandit. - """ - sh = self.spec - bp_state = bp.state if bp is not None else None - buf_len = len(self._buf) - - factor = self._bp_factor(bp_state) - size_bp = max(sh.min_size, min(sh.max_size, int(sh.target_size * factor))) - if self._bandit is not None: - bp_tag = f"×{factor:.2f}(bandit) target={sh.target_size}→{size_bp}" - else: - bp_tag = (f"bp={bp_state.value if bp_state else 'none'} " - f"×{factor:.2f}(fixed) target={sh.target_size}→{size_bp}") - - if occupancy > 0.85: - size_occ = max(sh.min_size, int(size_bp * 0.75)) - occ_tag = f"occ={occupancy:.2f}(high ×0.75) {size_bp}→{size_occ}" - elif occupancy < 0.40: - size_occ = min(sh.max_size, int(size_bp * 1.25)) - occ_tag = f"occ={occupancy:.2f}(low ×1.25) {size_bp}→{size_occ}" - else: - size_occ = size_bp - occ_tag = f"occ={occupancy:.2f}(ok)" - - if 0 < buf_len < size_occ and sh.stratify != "strict": - size_final = max(sh.min_size, buf_len) - tail_tag = f"tail(buf={buf_len}<{size_occ})→{size_final}" - else: - size_final = size_occ - tail_tag = "" - - size_final = max(sh.min_size, min(sh.max_size, size_final)) - - if self._log_fn is not None: - parts = [bp_tag, occ_tag] - if tail_tag: - parts.append(tail_tag) - self._log( - f" Sharder [{self.name}] adaptive: buf={buf_len} | " - + " | ".join(parts) - + f" | size={size_final}" - ) - return size_final - - def dispatch( - self, - bp: "BackpressureNegotiator | None", - occupancy: float, - running_scaffolds: Optional[set] = None, - ) -> list[str]: - """Dispatch one priority-ordered shard from the buffer. - - Returns a list of candidate_ids (highest priority first). - Returns [] when: - - buffer is empty - - BP state is THROTTLE - - stratify=strict and buffer < target_size (unless upstream is done) - - Candidates are ranked by the active profile's weight vector before - selection; the top-N by priority score are dispatched. - """ - if not self._buf: - return [] - - bp_state = bp.state if bp is not None else None - # Store occupancy so _update_bandit can use it as the reward signal - # for the *previous* factor (occupancy at this dispatch = outcome of prev dispatch). - self._last_occupancy = occupancy - self._update_bandit(bp_state) - - # ── Upstream-done fast-path (strict only) ────────────────────────────── - # STRICT mode: flush ALL remaining candidates in priority order so we - # don't deadlock on a partial tail smaller than target_size. - # SOFT/OFF: no effect — the campaign terminates via natural completion - # or _all_done elsewhere; flushing here would release low-priority - # candidates that should never run. - if self._upstream_done and self.spec.stratify == "strict": - now = self._now_fn() - scores = self._score_entries(self._buf, now, running_scaffolds) - order = sorted(range(len(self._buf)), key=lambda i: -scores[i]) - dispatched = [self._buf[i].candidate_id for i in order] - self._buf.clear() - self._shard_seq += 1 - self._log( - f" Sharder [{self.name}] [upstream-done strict-flush] " - f"shard={self._shard_seq}: dispatched {len(dispatched)} " - f"(priority-ordered) buffered=0 remaining" - ) - return dispatched - - if bp_state == BPState.THROTTLE: - return [] - if (self.spec.stratify == "strict" - and len(self._buf) < self.spec.target_size): - return [] - - if self.spec.stratify == "off": - n = 1 - else: - n = self.adaptive_size(bp, occupancy) - # In strict mode floor at target_size so dispatch never sends - # fewer than a full batch. _bp_factor already clamps the bandit's - # selection to ≥ 1.0 in strict mode; this is defense-in-depth for - # the fixed-mapping path and any future factor changes. - if self.spec.stratify == "strict": - n = max(n, self.spec.target_size) - n = min(n, len(self._buf)) - - # Rank buffer by priority, take top-n - now = self._now_fn() - scores = self._score_entries(self._buf, now, running_scaffolds) - order = sorted(range(len(self._buf)), key=lambda i: -scores[i]) - top_n = order[:n] - - top_entries = [self._buf[i] for i in top_n] - top_scores = [scores[i] for i in top_n] - dispatched = [e.candidate_id for e in top_entries] - - dispatched_set = set(top_n) - self._buf = [e for i, e in enumerate(self._buf) if i not in dispatched_set] - - self._shard_seq += 1 - if self._metrics_fn is not None: - self._metrics_fn( - self._shard_seq, - n, - [e.score for e in top_entries], - top_scores, - ) - if self._log_fn is not None: - cand_str = " ".join( - f"{e.candidate_id}(s={e.score:.3f} p={ps:.3f} sc={e.scaffold_class})" - for e, ps in zip(top_entries, top_scores) - ) - self._log( - f" Sharder [{self.name}] shard={self._shard_seq}: " - f"dispatched {n} profile={self.spec.profile} " - f"buffered={len(self._buf)} remaining\n {cand_str}" - ) - - return dispatched diff --git a/src/campaign/sync_wrapper.py b/src/campaign/sync_wrapper.py deleted file mode 100644 index aeefbf2..0000000 --- a/src/campaign/sync_wrapper.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -CampaignManager — synchronous shim around AsyncCampaignManager. - -Runs a dedicated event loop in a background thread so callers without an -async context can orchestrate workflows with plain blocking calls. -""" - -import asyncio -from typing import Optional - -from .base_workflow import BaseWorkflow -from .campaign_manager import AsyncCampaignManager -from .types import WorkflowStats - - -class CampaignManager: - """Synchronous campaign manager — thin wrapper around AsyncCampaignManager.""" - - def __init__( - self, - max_workers: Optional[int] = None, - engine: str = "concurrent", - total_cpus: int = 0, - total_gpus: int = 0, - num_workers: Optional[int] = None, - debug: bool = False, - asyncflow=None, - engine_dragon=None, - features: Optional[dict] = None, - ) -> None: - import threading - - self._acm = AsyncCampaignManager( - max_workers=max_workers, - engine=engine, - total_cpus=total_cpus, - total_gpus=total_gpus, - num_workers=num_workers, - debug=debug, - asyncflow=asyncflow, - engine_dragon=engine_dragon, - features=features, - ) - - async def _noop_init() -> None: - pass - - self._acm._setup_resources = _noop_init - - self._loop = asyncio.new_event_loop() - self._thread = threading.Thread( - target=self._loop.run_forever, daemon=True, name="CampaignManagerLoop" - ) - self._thread.start() - - def register_group(self, *args, **kwargs) -> None: - self._acm.register_group(*args, **kwargs) - - def start(self) -> None: - future = asyncio.run_coroutine_threadsafe(self._acm.start(), self._loop) - future.result() - - def wait(self, timeout: Optional[float] = None) -> bool: - future = asyncio.run_coroutine_threadsafe(self._acm.wait(timeout=timeout), self._loop) - outer_timeout = (timeout + 2.0) if timeout is not None else None - try: - return bool(future.result(timeout=outer_timeout)) - except Exception: - return False - - def close(self) -> None: - try: - future = asyncio.run_coroutine_threadsafe(self._acm.close(), self._loop) - future.result(timeout=5.0) - except Exception: - pass - self._loop.call_soon_threadsafe(self._loop.stop) - self._thread.join(timeout=5.0) - - def status(self) -> dict: - return self._acm.status() - - def stats(self) -> dict[str, WorkflowStats]: - return self._acm.stats() - - @classmethod - def from_config( - cls, - config: dict, - workflow_registry: dict[str, type[BaseWorkflow]], - **kwargs, - ) -> "CampaignManager": - res_cfg = config.get("resources", {}) - num_workers = config.get("num_workers") - cm = cls( - max_workers=config.get("max_workers"), - engine=config.get("engine", "concurrent"), - total_cpus=int(res_cfg.get("total_cpus", 0)), - total_gpus=int(res_cfg.get("total_gpus", 0)), - num_workers=int(num_workers) if num_workers is not None else None, - debug=bool(config.get("debug", False)), - **kwargs, - ) - - _cm_keys = { - "replicas", "dependencies", "dependency_threshold", "priority", - "min_replicas", "max_replicas", "required_cpus", "required_gpus", - "concurrency_cap", - } - - for name, wf_cfg in config.get("workflows", {}).items(): - 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 - max_replicas = int(wf_cfg.get("max_replicas") or - wf_cfg.get("concurrency_cap") or 0) - cm.register_group( - name=name, - workflow_class=wf_class, - 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)), - min_replicas=int(wf_cfg.get("min_replicas", 0)), - max_replicas=max_replicas, - required_cpus=int(wf_cfg.get("required_cpus", 0)), - required_gpus=int(wf_cfg.get("required_gpus", 0)), - config={k: v for k, v in wf_cfg.items() if k not in _cm_keys} or None, - ) - - return cm diff --git a/src/campaign/types.py b/src/campaign/types.py deleted file mode 100644 index f166ef6..0000000 --- a/src/campaign/types.py +++ /dev/null @@ -1,110 +0,0 @@ -""" -Shared data types for the campaign manager. - - _GroupInfo — internal per-group runtime state (not public API) - ResourcePool — CPU/GPU availability tracker - WorkflowStats — public per-group statistics snapshot -""" - -from collections import deque -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Optional - -if TYPE_CHECKING: - from .base_workflow import BaseWorkflow - - -@dataclass -class _GroupInfo: - name: str - workflow_class: "type[BaseWorkflow]" - replicas: int - dependencies: list[str] - group_config: Optional[dict] - configured_replicas: int = 0 - min_replicas: int = 0 - max_replicas: int = 0 - priority: int = 0 - required_cpus: int = 0 - required_gpus: int = 0 - dep_threshold: int = 1 - entry_point: str = "run" - status: str = "pending" - started_count: int = 0 - running_count: int = 0 - finished_replicas: int = 0 - # Set to True when the workflow explicitly signals it has produced enough - # data (via cm.signal_ready). Takes precedence over dep_threshold check. - ready: bool = False - # GPU IDs currently held by all running replicas of this group. - # Populated by _allocate_locked; cleared by _on_replica_finished. - running_gpu_ids: list[int] = field(default_factory=list) - # Candidate IDs waiting to be assigned to the next replica that starts. - # Populated by _flush_sharders_locked when dispatch returns candidate IDs; - # consumed FIFO by _allocate_locked so replica_idx → candidate_id is stable. - _pending_candidates: deque = field(default_factory=deque, repr=False) - - -@dataclass -class ResourcePool: - """ - Tracks available CPU cores and GPU slots for the campaign. - - Both counters are optional: a value of 0 disables tracking for that - resource type (unlimited). - """ - - total_cpus: int = 0 - total_gpus: int = 0 - available_cpus: int = field(default=0, init=False) - available_gpus: int = field(default=0, init=False) - - def __post_init__(self) -> None: - self.available_cpus = self.total_cpus - self.available_gpus = self.total_gpus - - def can_fit(self, cpus: int, gpus: int) -> bool: - if self.total_cpus > 0 and cpus > self.available_cpus: - return False - if self.total_gpus > 0 and gpus > self.available_gpus: - return False - return True - - def allocate(self, cpus: int, gpus: int) -> None: - self.available_cpus -= cpus - self.available_gpus -= gpus - - def release(self, cpus: int, gpus: int) -> None: - self.available_cpus += cpus - self.available_gpus += gpus - - def usage_str(self) -> str: - parts = [] - if self.total_cpus > 0: - parts.append(f"cpus={self.total_cpus - self.available_cpus}/{self.total_cpus}") - if self.total_gpus > 0: - parts.append(f"gpus={self.total_gpus - self.available_gpus}/{self.total_gpus}") - return " ".join(parts) if parts else "—" - - def available_str(self) -> str: - parts = [] - if self.total_cpus > 0: - parts.append(f"cpus={self.available_cpus}/{self.total_cpus}") - if self.total_gpus > 0: - parts.append(f"gpus={self.available_gpus}/{self.total_gpus}") - return " ".join(parts) if parts else "—" - - def as_dict(self) -> dict: - return { - "total_cpus": self.total_cpus, - "available_cpus": self.available_cpus, - "total_gpus": self.total_gpus, - "available_gpus": self.available_gpus, - } - - -@dataclass -class WorkflowStats: - """Cumulative statistics for one workflow group.""" - replicas_started: int = 0 - replicas_finished: int = 0 diff --git a/tests/test_campaign_manager.py b/tests/test_campaign_manager.py deleted file mode 100644 index 1634810..0000000 --- a/tests/test_campaign_manager.py +++ /dev/null @@ -1,699 +0,0 @@ -"""Tests for CampaignManager and AsyncCampaignManager.""" - -import asyncio -from unittest.mock import AsyncMock - -import pytest - -from src.campaign import ( - AsyncCampaignManager, - BaseWorkflow, - CampaignManager, - ResourcePool, - WorkflowStats, -) - -pytestmark = pytest.mark.anyio - - -@pytest.fixture -def anyio_backend(): - """Run all async tests with the asyncio backend only.""" - return "asyncio" - - -# --------------------------------------------------------------------------- -# Workflow stubs -# --------------------------------------------------------------------------- - - -class NullWorkflow(BaseWorkflow): - """No-op async workflow.""" - - workflow_id = "null" - - async def run(self, replica_id: str) -> None: - pass - - -class SleepWorkflow(BaseWorkflow): - """Sleeps briefly so concurrency effects are observable.""" - - workflow_id = "sleep" - - async def run(self, replica_id: str) -> None: - await asyncio.sleep(0.02) - - -class RecordingWorkflow(BaseWorkflow): - """Appends each replica_id it runs to a class-level list.""" - - workflow_id = "recording" - ran: list = [] - - async def run(self, replica_id: str) -> None: - RecordingWorkflow.ran.append(replica_id) - - -class SignalDoneWorkflow(BaseWorkflow): - """Fires _signal_done() immediately then finishes after a tiny sleep.""" - - workflow_id = "signal_done" - - async def run(self, replica_id: str) -> None: - 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.""" - - workflow_id = "hook" - calls: list = [] - - async def run(self, replica_id: str) -> None: - pass - - async def on_replica_done(self, replica_id, cm, final_state): - HookWorkflow.calls.append((replica_id, final_state)) - - -class FailingHookWorkflow(BaseWorkflow): - """Raises during run; records (replica_id, final_state) in on_replica_done.""" - - workflow_id = "failing_hook" - calls: list = [] - - async def run(self, replica_id: str) -> None: - raise RuntimeError("deliberate failure") - - async def on_replica_done(self, replica_id, cm, final_state): - FailingHookWorkflow.calls.append((replica_id, final_state)) - - -# Sync variants for CampaignManager (thread-pool) tests - - -class SyncRecordingWorkflow(BaseWorkflow): - workflow_id = "sync_rec" - ran: list = [] - - def run(self, replica_id: str) -> None: - SyncRecordingWorkflow.ran.append(replica_id) - - -class SyncHookWorkflow(BaseWorkflow): - workflow_id = "sync_hook" - calls: list = [] - - def run(self, replica_id: str) -> None: - pass - - def on_replica_done(self, replica_id, cm, final_state): - SyncHookWorkflow.calls.append((replica_id, final_state)) - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture(autouse=True) -def reset_class_state(): - """Clear class-level recording lists before every test.""" - RecordingWorkflow.ran = [] - HookWorkflow.calls = [] - FailingHookWorkflow.calls = [] - SyncRecordingWorkflow.ran = [] - SyncHookWorkflow.calls = [] - yield - - -@pytest.fixture -async def acm(): - """AsyncCampaignManager with asyncflow initialization mocked out.""" - cm = AsyncCampaignManager() - mock_af = AsyncMock() - - async def _fake_init(): - cm._asyncflow = mock_af - - cm._init_asyncflow = _fake_init - yield cm - await cm.close() - - -# --------------------------------------------------------------------------- -# BaseWorkflow -# --------------------------------------------------------------------------- - - -class TestBaseWorkflow: - async def test_signal_done_no_cm_is_noop(self): - wf = NullWorkflow() - await wf._signal_done() # must not raise - - 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 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") - - 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() - with pytest.raises(NotImplementedError): - wf.run("r0") - - def test_resolve_entry_point_both_raises(self): - class BothWorkflow(BaseWorkflow): - workflow_id = "both" - - def run(self, replica_id): - pass - - def start(self, replica_id): - pass - - with pytest.raises(ValueError, match="both 'run' and 'start'"): - AsyncCampaignManager._resolve_entry_point(BothWorkflow) - - def test_resolve_entry_point_neither_raises(self): - class NeitherWorkflow(BaseWorkflow): - workflow_id = "neither" - - with pytest.raises(ValueError, match="must define either"): - AsyncCampaignManager._resolve_entry_point(NeitherWorkflow) - - -# --------------------------------------------------------------------------- -# AsyncCampaignManager -# --------------------------------------------------------------------------- - - -class TestAsyncCampaignManager: - async def test_single_replica_completes(self, acm): - acm.register_group("a", NullWorkflow, replicas=1) - await acm.start() - assert await acm.wait(timeout=3.0) - s = acm.status() - assert s["groups"]["a"]["status"] == "done" - assert s["groups"]["a"]["replicas_finished"] == 1 - - async def test_all_replicas_run(self, acm): - acm.register_group("a", RecordingWorkflow, replicas=4, max_replicas=4) - await acm.start() - assert await acm.wait(timeout=3.0) - assert sorted(RecordingWorkflow.ran) == ["a_0", "a_1", "a_2", "a_3"] - - async def test_max_replicas_cap_respected(self, acm): - """Concurrent running count must never exceed max_replicas.""" - peak = [] - - class PeakObserver(BaseWorkflow): - workflow_id = "peak" - _active = 0 - - async def run(self, replica_id: str) -> None: - PeakObserver._active += 1 - peak.append(PeakObserver._active) - await asyncio.sleep(0.02) - PeakObserver._active -= 1 - - acm.register_group("a", PeakObserver, replicas=6, max_replicas=2) - await acm.start() - assert await acm.wait(timeout=5.0) - assert max(peak) <= 2 - - async def test_dependency_count_based(self, acm): - """Group B must not start until A has dep_threshold finished replicas.""" - order = [] - - class A(BaseWorkflow): - workflow_id = "A" - - async def run(self, replica_id: str) -> None: - order.append(("A", replica_id)) - - class B(BaseWorkflow): - workflow_id = "B" - - async def run(self, replica_id: str) -> None: - order.append(("B", replica_id)) - - acm.register_group("a", A, replicas=2) - acm.register_group("b", B, replicas=1, dependencies=["a"], dep_threshold=2) - await acm.start() - assert await acm.wait(timeout=3.0) - - 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_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, - replicas=1, - dependencies=["a"], - dep_threshold=999, # count-based fallback would never fire - ) - await acm.start() - assert await acm.wait(timeout=3.0) - - s = acm.status() - 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() - assert await acm.wait(timeout=3.0) - assert len(HookWorkflow.calls) == 2 - assert {rid for rid, _ in HookWorkflow.calls} == {"a_0", "a_1"} - assert all(st == "done" for _, st in HookWorkflow.calls) - - async def test_run_exception_marks_replica_failed(self, acm): - """An exception in run() sets final_state="failed"; campaign still completes.""" - acm.register_group("a", FailingHookWorkflow, replicas=2) - await acm.start() - assert await acm.wait(timeout=3.0) - assert len(FailingHookWorkflow.calls) == 2 - assert all(st == "failed" for _, st in FailingHookWorkflow.calls) - assert acm.status()["groups"]["a"]["status"] == "done" - - async def test_status_transitions_pending_running_done(self, acm): - """Status progresses: pending before start → running during → done after.""" - acm.register_group("a", SleepWorkflow, replicas=1) - assert acm.status()["groups"]["a"]["status"] == "pending" - await acm.start() - await asyncio.sleep(0.005) # yield to let the replica task begin - assert acm.status()["groups"]["a"]["status"] == "running" - assert await acm.wait(timeout=3.0) - assert acm.status()["groups"]["a"]["status"] == "done" - - async def test_empty_campaign_finishes_immediately(self, acm): - await acm.start() - assert await acm.wait(timeout=1.0) - - async def test_status_snapshot_fields(self, acm): - acm.register_group("a", NullWorkflow, replicas=2, max_replicas=1) - s = acm.status()["groups"]["a"] - assert s["status"] == "pending" - assert s["replicas_total"] == 2 - assert s["max_replicas"] == 1 - assert s["dependencies"] == [] - - async def test_stats_reflect_finished_count(self, acm): - acm.register_group("a", NullWorkflow, replicas=3) - await acm.start() - assert await acm.wait(timeout=3.0) - st = acm.stats() - assert st["a"].replicas_started == 3 - assert st["a"].replicas_finished == 3 - assert isinstance(st["a"], WorkflowStats) - - async def test_from_config_registers_groups(self): - config = { - "workflows": { - "x": {"replicas": 2, "max_replicas": 1}, - # y has dependencies → replicas defaults to 0 (triggered group) - "y": {"dependencies": ["x"], "dependency_threshold": 2}, - } - } - cm = AsyncCampaignManager.from_config(config, {"x": NullWorkflow, "y": NullWorkflow}) - s = cm.status()["groups"] - assert s["x"]["replicas_total"] == 2 - assert s["x"]["max_replicas"] == 1 - assert s["y"]["replicas_total"] == 0 # triggered group: not yet activated - assert s["y"]["dependencies"] == ["x"] - assert s["y"]["dep_threshold"] == 2 - - async def test_unknown_group_skipped_in_from_config(self): - config = {"workflows": {"unknown": {"replicas": 1}}} - cm = AsyncCampaignManager.from_config(config, {}) # empty registry - assert "unknown" not in cm.status()["groups"] - - -# --------------------------------------------------------------------------- -# CampaignManager (sync / thread-pool) -# --------------------------------------------------------------------------- - - -class TestCampaignManager: - @pytest.fixture - def cm(self): - manager = CampaignManager() - yield manager - manager.close() - - def test_single_replica_runs(self, cm): - cm.register_group("a", SyncRecordingWorkflow, replicas=1) - cm.start() - assert cm.wait(timeout=5.0) - assert SyncRecordingWorkflow.ran == ["a_0"] - - def test_multiple_replicas_all_run(self, cm): - cm.register_group("a", SyncRecordingWorkflow, replicas=3) - cm.start() - assert cm.wait(timeout=5.0) - assert sorted(SyncRecordingWorkflow.ran) == ["a_0", "a_1", "a_2"] - - def test_sliding_window_max_replicas(self, cm): - cm.register_group("a", SyncRecordingWorkflow, replicas=4, max_replicas=2) - cm.start() - assert cm.wait(timeout=5.0) - assert sorted(SyncRecordingWorkflow.ran) == ["a_0", "a_1", "a_2", "a_3"] - - def test_dependency_respected(self, cm): - """Group B must start only after group A completes.""" - order = [] - - class A(BaseWorkflow): - workflow_id = "A" - - def run(self, replica_id: str) -> None: - order.append(("A", replica_id)) - - class B(BaseWorkflow): - workflow_id = "B" - - def run(self, replica_id: str) -> None: - order.append(("B", replica_id)) - - cm.register_group("a", A, replicas=2) - cm.register_group("b", B, replicas=1, dependencies=["a"]) - cm.start() - assert cm.wait(timeout=5.0) - - b_idx = next(i for i, (wf, _) in enumerate(order) if wf == "B") - assert all(wf == "A" for wf, _ in order[:b_idx]) - - def test_on_replica_done_hook_called(self, cm): - cm.register_group("a", SyncHookWorkflow, replicas=2) - cm.start() - assert cm.wait(timeout=5.0) - assert len(SyncHookWorkflow.calls) == 2 - assert {rid for rid, _ in SyncHookWorkflow.calls} == {"a_0", "a_1"} - - def test_status_snapshot_fields(self, cm): - cm.register_group("a", SyncRecordingWorkflow, replicas=1, max_replicas=1) - s = cm.status()["groups"]["a"] - assert s["status"] == "pending" - assert s["replicas_total"] == 1 - assert s["max_replicas"] == 1 - - def test_stats_reflect_finished_count(self, cm): - cm.register_group("a", SyncRecordingWorkflow, replicas=3) - cm.start() - assert cm.wait(timeout=5.0) - st = cm.stats() - assert st["a"].replicas_finished == 3 - assert isinstance(st["a"], WorkflowStats) - - def test_from_config_registers_groups(self): - config = { - "workflows": { - "alpha": {"replicas": 3, "max_replicas": 2}, - # beta has dependencies → replicas defaults to 0 (triggered group) - "beta": {"dependencies": ["alpha"]}, - } - } - 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["beta"]["replicas_total"] == 0 # triggered group: not yet activated - - def test_unknown_group_skipped_in_from_config(self): - config = {"workflows": {"ghost": {"replicas": 1}}} - cm = CampaignManager.from_config(config, {}) - cm.close() - assert "ghost" not in cm.status()["groups"] - - -# --------------------------------------------------------------------------- -# ResourcePool unit tests -# --------------------------------------------------------------------------- - - -class TestResourcePool: - def test_initial_available_equals_total(self): - rp = ResourcePool(total_cpus=16, total_gpus=4) - assert rp.available_cpus == 16 - assert rp.available_gpus == 4 - - def test_can_fit_within_budget(self): - rp = ResourcePool(total_cpus=8, total_gpus=2) - assert rp.can_fit(8, 2) - assert rp.can_fit(1, 0) - assert rp.can_fit(0, 1) - - def test_cannot_fit_over_budget(self): - rp = ResourcePool(total_cpus=4, total_gpus=1) - assert not rp.can_fit(5, 0) - assert not rp.can_fit(0, 2) - - def test_zero_total_means_unlimited(self): - rp = ResourcePool(total_cpus=0, total_gpus=0) - assert rp.can_fit(9999, 9999) - - def test_allocate_decrements_available(self): - rp = ResourcePool(total_cpus=16, total_gpus=4) - rp.allocate(4, 1) - assert rp.available_cpus == 12 - assert rp.available_gpus == 3 - - def test_release_increments_available(self): - rp = ResourcePool(total_cpus=16, total_gpus=4) - rp.allocate(4, 1) - rp.release(4, 1) - assert rp.available_cpus == 16 - assert rp.available_gpus == 4 - - def test_as_dict_keys(self): - rp = ResourcePool(total_cpus=8, total_gpus=2) - d = rp.as_dict() - assert set(d) == {"total_cpus", "available_cpus", "total_gpus", "available_gpus"} - - def test_usage_str_tracks_used(self): - rp = ResourcePool(total_cpus=8, total_gpus=4) - rp.allocate(3, 2) - s = rp.usage_str() - assert "3/8" in s - assert "2/4" in s - - def test_available_str_tracks_free(self): - rp = ResourcePool(total_cpus=8, total_gpus=4) - rp.allocate(3, 2) - s = rp.available_str() - assert "5/8" in s - assert "2/4" in s - - def test_unlimited_usage_str_returns_dash(self): - rp = ResourcePool(total_cpus=0, total_gpus=0) - assert rp.usage_str() == "—" - - -# --------------------------------------------------------------------------- -# Resource-aware scheduling tests (AsyncCampaignManager) -# --------------------------------------------------------------------------- - - -class TestAsyncCampaignManagerResources: - @pytest.fixture - async def racm(self): - """AsyncCampaignManager with 4 CPUs and 2 GPUs, asyncflow mocked.""" - cm = AsyncCampaignManager(total_cpus=4, total_gpus=2) - mock_af = AsyncMock() - - async def _fake_init(): - cm._asyncflow = mock_af - - cm._init_asyncflow = _fake_init - yield cm - await cm.close() - - async def test_resource_limits_concurrency(self, racm): - """With 2 GPUs and 1 GPU/replica, at most 2 replicas run concurrently.""" - peak = [] - - class GpuWorkflow(BaseWorkflow): - workflow_id = "gpu" - _active = 0 - - async def run(self, replica_id: str) -> None: - GpuWorkflow._active += 1 - peak.append(GpuWorkflow._active) - await asyncio.sleep(0.02) - GpuWorkflow._active -= 1 - - racm.register_group("g", GpuWorkflow, replicas=6, max_replicas=6, required_gpus=1) - await racm.start() - assert await racm.wait(timeout=5.0) - assert max(peak) <= 2 # only 2 GPUs available - - async def test_resources_released_after_replica(self, racm): - """Available resources return to full after all replicas complete.""" - racm.register_group("g", NullWorkflow, replicas=2, required_cpus=2, required_gpus=1) - await racm.start() - assert await racm.wait(timeout=3.0) - s = racm.status()["resources"] - assert s["available_cpus"] == 4 # total_cpus restored - assert s["available_gpus"] == 2 # total_gpus restored - - async def test_status_includes_resource_snapshot(self, racm): - racm.register_group("g", NullWorkflow, replicas=1, required_cpus=2, required_gpus=1) - s = racm.status() - assert "resources" in s - assert s["resources"]["total_cpus"] == 4 - assert s["resources"]["total_gpus"] == 2 - assert s["resources"]["available_cpus"] == 4 - assert s["resources"]["available_gpus"] == 2 - - async def test_from_config_parses_resources(self): - config = { - "resources": {"total_cpus": 64, "total_gpus": 8}, - "workflows": { - "a": {"replicas": 1, "required_cpus": 4, "required_gpus": 2}, - }, - } - cm = AsyncCampaignManager.from_config(config, {"a": NullWorkflow}) - s = cm.status() - assert s["resources"]["total_cpus"] == 64 - assert s["resources"]["total_gpus"] == 8 - assert s["groups"]["a"]["required_cpus"] == 4 - assert s["groups"]["a"]["required_gpus"] == 2 - - async def test_resource_constrained_scheduling(self, racm): - """Both groups run to completion despite resource contention.""" - started_order = [] - - class TrackWorkflow(BaseWorkflow): - workflow_id = "track" - - async def run(self, replica_id: str) -> None: - started_order.append(replica_id) - await asyncio.sleep(0.01) - - racm.register_group("lo", TrackWorkflow, replicas=2, required_gpus=1) - racm.register_group("hi", TrackWorkflow, replicas=2, required_gpus=1) - await racm.start() - assert await racm.wait(timeout=3.0) - # All 4 replicas should complete - assert len(started_order) == 4 - assert sum(1 for r in started_order if r.startswith("lo")) == 2 - assert sum(1 for r in started_order if r.startswith("hi")) == 2 - - -# --------------------------------------------------------------------------- -# Resource-aware scheduling tests (CampaignManager sync) -# --------------------------------------------------------------------------- - - -class TestCampaignManagerResources: - @pytest.fixture - def rcm(self): - cm = CampaignManager(total_cpus=4, total_gpus=2) - yield cm - cm.close() - - def test_resource_limits_concurrency(self, rcm): - """With 2 GPUs and 1 GPU/replica, at most 2 run concurrently.""" - import threading - - peak = [] - lock = threading.Lock() - - class GpuWorkflow(BaseWorkflow): - workflow_id = "gpu" - _active = 0 - - def run(self, replica_id: str) -> None: - with lock: - GpuWorkflow._active += 1 - peak.append(GpuWorkflow._active) - import time - - time.sleep(0.02) - with lock: - GpuWorkflow._active -= 1 - - rcm.register_group("g", GpuWorkflow, replicas=6, max_replicas=6, required_gpus=1) - rcm.start() - assert rcm.wait(timeout=5.0) - assert max(peak) <= 2 - - def test_resources_released_after_replica(self, rcm): - rcm.register_group("g", SyncRecordingWorkflow, replicas=2, required_cpus=2, required_gpus=1) - rcm.start() - assert rcm.wait(timeout=3.0) - s = rcm.status()["resources"] - assert s["available_cpus"] == 4 - assert s["available_gpus"] == 2 - - def test_status_includes_resource_snapshot(self, rcm): - rcm.register_group("g", SyncRecordingWorkflow, replicas=1, required_cpus=1, required_gpus=0) - s = rcm.status() - assert "resources" in s - assert s["resources"]["total_cpus"] == 4 - assert s["resources"]["total_gpus"] == 2 - - def test_from_config_parses_resources(self): - config = { - "resources": {"total_cpus": 32, "total_gpus": 4}, - "workflows": { - "a": {"replicas": 1, "required_cpus": 8, "required_gpus": 1}, - }, - } - cm = CampaignManager.from_config(config, {"a": SyncRecordingWorkflow}) - s = cm.status() - cm.close() - assert s["resources"]["total_cpus"] == 32 - assert s["resources"]["total_gpus"] == 4 - assert s["groups"]["a"]["required_cpus"] == 8 - assert s["groups"]["a"]["required_gpus"] == 1 diff --git a/workflows/esm2_inference/config.yaml b/workflows/esm2_inference/config.yaml index 98b11a6..580236c 100644 --- a/workflows/esm2_inference/config.yaml +++ b/workflows/esm2_inference/config.yaml @@ -11,7 +11,7 @@ cache_dir: "cache" # ----------------------------------------------------------------------------- # GPU / Service Configuration # ----------------------------------------------------------------------------- -num_services: 1 # Match max_replicas so each service gets its own CM-assigned GPU +num_services: 1 # Match concurrency_cap so each service gets its own CM-assigned GPU num_gpus_per_service: 1 num_cpus_per_service: 64 num_workers_per_gpu: 4 diff --git a/workflows/esm2_inference/delta_cpu_batch.sh b/workflows/esm2_inference/delta_cpu_batch.sh new file mode 100755 index 0000000..568ab69 --- /dev/null +++ b/workflows/esm2_inference/delta_cpu_batch.sh @@ -0,0 +1,55 @@ +#!/bin/sh -l +# +# SPHERICAL ESM2/DDSim Campaign — SLURM CPU batch script (concurrent backend) +# +# Runs the campaign without Dragon/GPU — useful for functional testing and +# development. Workflows execute via the asyncflow ConcurrentExecutionBackend. +# +# Account: set SBATCH_ACCOUNT=-delta-cpu before calling sbatch +#SBATCH --partition=cpu +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=64 +#SBATCH --time=00:30:00 +#SBATCH --job-name=campaign-cpu +#xSBATCH --mail-user=${USER}@institution.edu +#SBATCH --mail-type=ALL + +# ── Environment ─────────────────────────────────────────────────────────────── +if [ -z "${SBATCH_ACCOUNT:-}${SLURM_JOB_ACCOUNT:-}" ]; then + echo "WARNING: SBATCH_ACCOUNT is not set — job may be charged to default account." + echo " Set it with: export SBATCH_ACCOUNT=-delta-cpu" +fi +echo "Account: ${SLURM_JOB_ACCOUNT:-unknown}" + +if [ -z "${SCRATCH:-}" ]; then + echo "ERROR: SCRATCH is not set." + echo " export SCRATCH=/scratch/ && sbatch delta_cpu_batch.sh" + exit 1 +fi + +# ── Project paths (adjust base dirs if layout differs) ─────────────────────── +export DDSIM_DIR="${DDSIM_DIR:-${SCRATCH}/${USER}/DeepDriveSim}" +export SPHERICAL_DIR="${SPHERICAL_DIR:-${SCRATCH}/${USER}/SPHERICAL}" +export VE_HOME=/u/${USER}/ve + +export MD_DIR=${DDSIM_DIR}/workflows/ddmd_workflow +export MINAPPS_DIR=${DDSIM_DIR}/workflows/miniapps_workflow +export DUMMY_DIR=${DDSIM_DIR}/workflows/dummy_workflow +export INF_DIR=${SPHERICAL_DIR}/workflows/esm2_inference + +export MD_HOME=${DDSIM_DIR}/workflows/ddmd_workflow +export MD_INPUT=${MD_HOME}/data + +export WORK_DIR=${SPHERICAL_DIR}/workflows/esm2_inference + +cd ${WORK_DIR} + +# ── Clean previous run artifacts ───────────────────────────────────────────── +rm -rf DDMD* telemetry-results asyncflow.session* + +# ── Activate campaign environment ───────────────────────────────────────────── +source ${VE_HOME}/esm2_ddsim_campaign/bin/activate + +# ── Launch (concurrent backend, no Dragon) ─────────────────────────────────── +python run_esm2_infern.py diff --git a/workflows/esm2_inference/delta_gpu_batch.sh b/workflows/esm2_inference/delta_gpu_batch.sh new file mode 100644 index 0000000..4767719 --- /dev/null +++ b/workflows/esm2_inference/delta_gpu_batch.sh @@ -0,0 +1,88 @@ +#!/bin/sh -l +# +# SPHERICAL ESM2/DDSim Campaign — SLURM GPU batch script (Dragon backend) +# +# GPU stress scenario: inference (priority 9) competes with md (priority 4) for +# 4 A40 GPUs. inference.cap=4 × required_gpus=1 fills the entire pool under the +# 'none' baseline, starving Pipeline B (md → miniapps) completely. +# +# --gpus=4 matches config.yaml resources.total_gpus=4 +# +# Compare policies by changing cm.adr.policy in config.yaml: +# none → miniapps=0 (inference monopolizes all 4 GPUs) +# rule → miniapps>0 (ADR boosts md into pass-2 GPU allocation) +# bandit → miniapps>0 (converges within 3–5 cycles) +# llm → miniapps>0 (if HF_TOKEN is valid) +# +# Account: set SBATCH_ACCOUNT=-delta-gpu before calling sbatch +#SBATCH --partition=gpuA40x4 +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=4 +#SBATCH --cpus-per-task=16 +#SBATCH --gpus=4 +#SBATCH --exclusive +#SBATCH --time=01:30:00 +#SBATCH --job-name=esm2 +#SBATCH --mail-user=${USER}@institution.edu +#SBATCH --mail-type=ALL + +# ── System library paths (Delta-specific) ──────────────────────────────────── +export CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_x86_64/25.3/cuda/12.8 +export MPI_LIB=/opt/cray/pe/mpich/8.1.32/ofi/gnu/11.2/lib-abi-mpich +export FAB_LIB=/opt/cray/libfabric/1.22.0/lib64 +export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${MPI_LIB}:${FAB_LIB}:${LD_LIBRARY_PATH} + +export TF_FORCE_GPU_ALLOW_GROWTH=true +export JAX_PLATFORMS=cpu +export TF_CPP_MIN_LOG_LEVEL=3 # suppress TF/XLA C++ log noise (cuInit probe at import time) + +# ── Environment ─────────────────────────────────────────────────────────────── +if [ -z "${SBATCH_ACCOUNT:-}${SLURM_JOB_ACCOUNT:-}" ]; then + echo "WARNING: SBATCH_ACCOUNT is not set — job may be charged to default account." + echo " Set it with: export SBATCH_ACCOUNT=-delta-gpu" +fi +echo "Account: ${SLURM_JOB_ACCOUNT:-unknown}" + +if [ -z "${SCRATCH:-}" ]; then + echo "ERROR: SCRATCH is not set." + echo " export SCRATCH=/scratch/ && sbatch delta_gpu_batch.sh" + exit 1 +fi + +# ── Project paths (adjust base dirs if layout differs) ─────────────────────── +export DDSIM_DIR="${DDSIM_DIR:-${SCRATCH}/${USER}/DeepDriveSim}" +export SPHERICAL_DIR="${SPHERICAL_DIR:-${SCRATCH}/${USER}/SPHERICAL}" +export VE_HOME=/u/${USER}/ve + +export MD_DIR=${DDSIM_DIR}/workflows/ddmd_workflow +export MINAPPS_DIR=${DDSIM_DIR}/workflows/miniapps_workflow +export DUMMY_DIR=${DDSIM_DIR}/workflows/dummy_workflow +export INF_DIR=${SPHERICAL_DIR}/workflows/esm2_inference + +export MD_HOME=${DDSIM_DIR}/workflows/ddmd_workflow +export MD_INPUT=${MD_HOME}/data +export SGDES_DIR="${SGDES_DIR:-${SCRATCH}/${USER}/SGDES}" + +export WORK_DIR=${SPHERICAL_DIR}/workflows/esm2_inference + +cd ${WORK_DIR} + +# ── Clean previous run artifacts ───────────────────────────────────────────── +rm -rf DDMD* telemetry-results nvml-telemetry asyncflow.session* + +# ── Activate campaign environment and configure Dragon ─────────────────────── +source ${VE_HOME}/esm2/bin/activate +dragon-config add --ofi-runtime-lib=${FAB_LIB} + +# ── Launch ─────────────────────────────────────────────────────────────────── +GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-4} +export TOTAL_GPUS=$(( SLURM_NNODES * GPUS_PER_NODE )) +echo "Nodes: ${SLURM_NNODES} GPUs/node: ${GPUS_PER_NODE} Total GPUs: ${TOTAL_GPUS}" + +if [ "${SLURM_NNODES}" -gt 1 ]; then + dragon -m run_esm2_infern.py +else + dragon -s run_esm2_infern.py +fi + +echo "=== Campaign done: $(date) ===" diff --git a/workflows/esm2_inference/delta_gpu_sbatch.sh b/workflows/esm2_inference/delta_gpu_sbatch.sh deleted file mode 100755 index 66f8dcb..0000000 --- a/workflows/esm2_inference/delta_gpu_sbatch.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/bin/sh -l - -#SBATCH -A ***-delta-gpu -#SBATCH --partition=gpuA40x4 -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=64 -#SBATCH --gpus-per-node=4 -#SBATCH --time=00:30:00 -#SBATCH --job-name=esm2_inf -#SBATCH --mail-user=mariya.goliyad@rutgers.edu -#SBATCH --mail-type=ALL - -export CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_x86_64/25.3/cuda/12.8 -export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH - -export TF_FORCE_GPU_ALLOW_GROWTH=true - -export SPHERICAL_DIR=/scratch/***/${USER}/SPHERICAL -export WORK_DIR=${SPHERICAL_DIR}/workflows/esm2_inference - -export VE_HOME=/u/${USER}/ve -cd ${WORK_DIR} - -# ── Clean previous run artifacts ────────────────────────────────────────────── -rm -rf data/outputs_test - -source ${VE_HOME}/esm2/bin/activate -dragon-config add --ofi-runtime-lib=/opt/cray/libfabric/1.22.0/lib64 - -# Compute total GPUs and choose single- vs multi-node Dragon launch. -GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-1} -export TOTAL_GPUS=$(( SLURM_NNODES * GPUS_PER_NODE )) -echo "Nodes: ${SLURM_NNODES} GPUs/node: ${GPUS_PER_NODE} Total GPUs: ${TOTAL_GPUS}" - -if [ "${SLURM_NNODES}" -gt 1 ]; then - dragon -m run_esm2_infern.py --config_file config.yaml -else - dragon -s run_esm2_infern.py --config_file config.yaml -fi diff --git a/workflows/run_campaign/dreamer_campaign/benchmark.py b/workflows/run_campaign/dreamer_campaign/benchmark.py deleted file mode 100644 index 90123de..0000000 --- a/workflows/run_campaign/dreamer_campaign/benchmark.py +++ /dev/null @@ -1,258 +0,0 @@ -#!/usr/bin/env python3 -""" -Benchmark runner — measures performance across feature-flag configurations. - -Each configuration is a dict of feature overrides applied on top of the -base config.yaml. For each configuration the campaign is run N_RUNS times -(different random seeds) and metrics are aggregated. - -Results are written to benchmark_results.json for consumption by -plot_optimizations.py. - -Usage: - python benchmark.py [--config config.yaml] [--runs 3] [--out benchmark_results.json] -""" - -import argparse -import asyncio -import copy -import json -import sys -import time -from pathlib import Path - -import yaml - -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -# ── Benchmark configurations ────────────────────────────────────────────────── - -CONFIGURATIONS: dict[str, dict] = { - # ─── Dumb waterfall baseline ────────────────────────────────────────────── - # True sequential pipeline: each stage starts ONLY after ALL replicas of the - # previous stage have finished (dep_threshold_override=9999999 forces the - # scheduler to wait for upstream.status=="done", not just the first replica). - # Sharder OFF → FIFO arrival order (random quality). No priorities → s1 - # monopolises all GPUs. Underutilises resources at every stage transition. - "baseline": { - "features": {"backpressure": False, "monitor": False, "sharder": False, "bandit": False}, - # "dep_threshold_override": 9999999, - "stage_replicas_overrides": { - "s2_ml_affinity": {"min_replicas": 2}, - "s3_docking": {"min_replicas": 1}, - "s4_md_refinement": {"min_replicas": 1}, - "s5_fep_ranking": {"min_replicas": 1}, - }, - }, - - # ─── Smart sharding axis ───────────────────────────────────────────────── - # Demonstrates the CANDIDATE-ROUTING + PARTIAL PIPELINE benefit. - # Sharder ON, stratify=soft: adaptive batch dispatch ranked by score. - # NO static stage priorities (no bandit) — but min_replicas guarantees a - # concurrency floor for every downstream stage via Pass 1 of the scheduler. - # This forces PARTIAL overlap between stages without requiring the bandit - # to learn it. Combines quality filtering with basic pipeline configuration. - "sharding+bp": { - "features": {"backpressure": True, "monitor": False, "sharder": True, "bandit": False}, - "sharding_overrides": {"stratify": "soft", "use_bandit": True, - "min_size": 1, "target_size": 8, "max_size": 32}, - # min_replicas guarantees Pass-1 concurrency floor for downstream stages: - # scheduler always reserves this many GPU slots even while s1 is running. - "stage_replicas_overrides": { - "s2_ml_affinity": {"min_replicas": 2}, - "s3_docking": {"min_replicas": 1}, - "s4_md_refinement": {"min_replicas": 1}, - "s5_fep_ranking": {"min_replicas": 1}, - }, - }, - - # ─── Scheduling bandit axis ─────────────────────────────────────────────── - # Demonstrates the GPU-ALLOCATION benefit in isolation. - # No sharding: FIFO dispatch (sharder=False), random quality order. - # Cross-stage Thompson-sampling bandit LEARNS downstream-first GPU allocation. - # No static priorities (bandit must learn them). No BP. - "scheduling_bandit": { - "features": {"backpressure": False, "monitor": False, "sharder": False, "bandit": True}, - }, - - # ─── Combined ───────────────────────────────────────────────────────────── - # Both axes together: adaptive soft sharding (stratify=soft + shard_bandit + - # BP) AND cross-stage scheduling bandit. - # NO static priorities — the cross-stage bandit learns optimal GPU allocation - # via Thompson sampling with depth-based warm-start priors (s5=Beta(5,1), - # s1=Beta(1,1)). Static priorities would pre-answer what the bandit is - # supposed to discover, hiding whether it adds value beyond the static ordering. - "all_optimizations": { - "features": {"backpressure": True, "monitor": False, "sharder": True, "bandit": True}, - "sharding_overrides": {"stratify": "soft", "use_bandit": True, - "min_size": 1, "target_size": 8, "max_size": 32}, - }, -} - - -def _apply_config_override(base: dict, override: dict) -> dict: - """Deep-merge override into a copy of base config.""" - cfg = copy.deepcopy(base) - # Feature flags - if "features" in override: - cfg.setdefault("cm", {}).setdefault("features", {}).update(override["features"]) - cfg.setdefault("features", {}).update(override["features"]) - # Sharding overrides: applied to all stages that have a sharding block - if "sharding_overrides" in override: - sh_ov = override["sharding_overrides"] - for stage in cfg.get("stages", []): - if "sharding" in stage: - stage["sharding"].update(sh_ov) - # stage_replicas_overrides: per-stage min_replicas / max_replicas overrides. - # Used to guarantee a minimum concurrency floor for downstream stages even without - # a scheduling bandit — Pass 1 of the scheduler ensures min_replicas is always - # satisfied first, forcing some GPU sharing across stages. - if "stage_replicas_overrides" in override: - for stage in cfg.get("stages", []): - sid = stage["id"] - if sid in override["stage_replicas_overrides"]: - stage.update(override["stage_replicas_overrides"][sid]) - - # dep_threshold_override: sets dependency_threshold for every DEPENDENT stage to a - # very large value so it only becomes eligible when upstream.status == "done" — - # not after the first upstream replica finishes. Creates a true sequential - # waterfall: stage N+1 waits for ALL of stage N to complete before starting. - if "dep_threshold_override" in override: - dt = int(override["dep_threshold_override"]) - stage_ids_set = {s["id"] for s in cfg.get("stages", [])} - for stage in cfg.get("stages", []): - if stage.get("upstream", "") in stage_ids_set: - stage["dependency_threshold"] = dt - - # Stage priority overrides: sets scheduler priority per stage (higher = scheduled first). - # Used to give downstream stages static priority without a scheduling bandit. - if "stage_priority_overrides" in override: - pri_ov = override["stage_priority_overrides"] - for stage in cfg.get("stages", []): - if stage["id"] in pri_ov: - stage["priority"] = pri_ov[stage["id"]] - # Dreamer overrides: applied to all stages' dreamer block (flat key update). - # Used to set trigger_mode and other dreamer simulation parameters. - if "dreamer_overrides" in override: - dr_ov = override["dreamer_overrides"] - for stage in cfg.get("stages", []): - stage.setdefault("dreamer", {}).update(dr_ov) - return cfg - - -async def _run_once(config: dict, seed_offset: int) -> dict: - """Run one campaign with the given config and return its metrics dict.""" - import random as _random - # Fix the global random state so score-cascade outcomes are identical across - # configs within the same run index. Without this, sequential config runs - # consume different random numbers from a shared state, making comparisons - # unfair (different random realizations of the score cascade). - _random.seed(seed_offset + 1337) - - from src.campaign import AsyncCampaignManager as CampaignManager - from src.inference.utils import load_config - import importlib - - # Reset DreamerWorkflow class-level state so trigger counts don't bleed - # across benchmark runs (class vars persist for the lifetime of the process). - sys.path.insert(0, str(Path(__file__).parent)) - from dreamer_workflow import DreamerWorkflow - DreamerWorkflow._group_state = {} - DreamerWorkflow._trigger_lock = None - - # Translate plan format - if "stages" in config: - from run_campaign import _build_from_plan, _build_registry - cm_cfg = config.get("cm", {}) - config["workflows"] = _build_from_plan(config) - for key in ("engine", "resources", "telemetry", "workflow_registry", "features"): - if key in cm_cfg and key not in config: - config[key] = cm_cfg[key] - config["debug"] = bool(cm_cfg.get("debug", False)) - - # Bump seeds for reproducible variance across runs - if "provenance" in config: - for k in config["provenance"].get("seeds", {}): - config["provenance"]["seeds"][k] += seed_offset - - from radical.asyncflow import WorkflowEngine - from rhapsody.backends import ConcurrentExecutionBackend - backend = await ConcurrentExecutionBackend() - asyncflow = await WorkflowEngine.create(backend) - - registry = _build_registry(config) - cm = CampaignManager.from_config(config, registry, asyncflow=asyncflow) - - # Hard timeout: guards against any stall in cm.wait(). - # All stages GPU-bound; baseline ~1200s/run. Optimised runs can take - # longer during bandit warm-up (before it learns to keep s2 running). - RUN_TIMEOUT_S = 5400 - - try: - await cm.start() - finished = await cm.wait(timeout=RUN_TIMEOUT_S) - if not finished: - raise TimeoutError( - f"Campaign did not finish within {RUN_TIMEOUT_S}s " - f"(likely a runaway dreamer replica)" - ) - finally: - await cm.close() - await asyncflow.shutdown() - - m = cm.metrics().to_dict() - - # ── Time-to-target: seconds until the Nth terminal-stage replica finishes ── - # Extracted from replica_events so plot_optimizations can draw the step curve. - _TARGET_STAGE = "s5_fep_ranking" - _TARGET_N = 5 - s5_finishes = sorted( - e["t"] for e in m.get("replica_events", []) - if e["group"] == _TARGET_STAGE and e["event"] == "finish" - ) - m["time_to_target_s"] = s5_finishes[_TARGET_N - 1] if len(s5_finishes) >= _TARGET_N else None - return m - - -async def run_benchmark( - config_path: str, - n_runs: int, - out_path: str, -) -> None: - with open(config_path) as f: - base_config = yaml.safe_load(f) - - results: dict = {} - for cfg_name, override in CONFIGURATIONS.items(): - print(f"\n{'='*60}") - print(f"Configuration: {cfg_name}") - print(f"{'='*60}") - cfg_results = [] - for run_idx in range(n_runs): - print(f" Run {run_idx + 1}/{n_runs}...", end=" ", flush=True) - cfg = _apply_config_override(base_config, override) - t0 = time.time() - try: - metrics = await _run_once(cfg, seed_offset=run_idx * 100) - elapsed = time.time() - t0 - print(f"done in {elapsed:.1f}s (campaign wall_time={metrics['wall_time_s']:.1f}s)") - cfg_results.append(metrics) - except Exception as exc: - print(f"FAILED: {exc}") - cfg_results.append({"error": str(exc), "wall_time_s": None}) - results[cfg_name] = cfg_results - - with open(out_path, "w") as f: - json.dump(results, f, indent=2) - print(f"\nResults written to {out_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--config", default="config.yaml") - parser.add_argument("--runs", type=int, default=3) - parser.add_argument("--out", default="benchmark_results.json") - args = parser.parse_args() - asyncio.run(run_benchmark(args.config, args.runs, args.out)) - -#python benchmark.py --runs 3 --out benchmark_results.json \ No newline at end of file diff --git a/workflows/run_campaign/dreamer_campaign/config.yaml b/workflows/run_campaign/dreamer_campaign/config.yaml deleted file mode 100644 index 60e02ce..0000000 --- a/workflows/run_campaign/dreamer_campaign/config.yaml +++ /dev/null @@ -1,365 +0,0 @@ -# ============================================================================= -# SPHERICAL Dreamer Campaign Plan -# -# Format: cm-prototype cm-plan/1.0 + SPHERICAL extensions -# Mirrors /u/mgoliyad1/cm-prototype/plans/example_campaign.yaml exactly. -# Two SPHERICAL-specific additions: -# cm: — Campaign Manager runtime settings (engine, resources, registry) -# dreamer: — radical.dreamer emulation parameters (one block per stage) -# -# Prototype fields per stage (unchanged): -# id, upstream, downstream, threshold_top_fraction, budget_node_hours, -# pilot, concurrency_cap, surrogate, downstream_input_target, variant -# -# SPHERICAL-only additions per stage: -# sharding — adaptive batch dispatch from the trigger buffer (see below) -# dreamer — radical.dreamer emulation block (num_cores, perf_dist, …) -# -# ── sharding block ──────────────────────────────────────────────────────────── -# Controls how trigger_dependent() signals are batched before entering the -# runnable queue. Only meaningful for dependent (non-root) stages. -# -# target_size: int Target batch size per dispatch cycle. -# Actual size is modulated by BP state and pilot occupancy. -# -# min_size: int Hard lower bound on dispatch size (never dispatch fewer). -# -# max_size: int Hard upper bound on dispatch size (never dispatch more). -# -# stratify: str Batching strategy — one of: -# off Bypass adaptive_size entirely; always dispatch exactly 1 -# trigger per scheduling cycle. Useful when no batching -# is desired but backpressure control is still wanted. -# soft Adaptive sizing (BP + occupancy modulation) with tail -# dispatch: if the buffer holds fewer than target_size -# triggers, dispatch whatever is buffered immediately. -# Best for stages where partial batches are acceptable. -# strict Adaptive sizing, but HOLD the buffer until it accumulates -# at least target_size triggers before dispatching. -# A partial tail is only flushed when the upstream group -# is truly done (no more triggers will arrive). -# Best for stratified sampling stages (docking, MD) where -# chemical diversity within a batch matters. -# -# use_bandit: bool Replace the fixed BP multipliers (0.5× THROTTLE / -# 1.0× HOLD / 1.5× WIDEN) with a Thompson-sampling bandit -# that learns which multiplier leads to healthy queue depth. -# Arms: [0.50, 0.75, 1.00, 1.25, 1.50]. -# Reward: HOLD→0.8, WIDEN→0.3, THROTTLE→0.1. -# Useful for A/B comparison: set true on one stage, false -# on a parallel stage with identical sharding parameters. -# -# bandit_seed: int Optional RNG seed for reproducible bandit exploration. -# -# ── edges / backpressure block ──────────────────────────────────────────────── -# Each edge can carry a backpressure stanza that controls the downstream stage's -# queue depth via a hysteresis state machine: -# -# high_water: int Queue depth (triggered but not yet started) at which the -# controller enters THROTTLE state. While THROTTLE, the -# sharder's dispatch() returns 0 — the producer side is -# gated, so the consumer (executor) is never blocked. -# -# low_water: int Queue depth at which the controller leaves THROTTLE and -# enters WIDEN state (dispatch multiplier increases). -# Must be strictly less than high_water. -# -# States: HOLD (normal) → THROTTLE (queue too deep) → WIDEN (queue drained) -# -# ── edges / profile ─────────────────────────────────────────────────────────── -# profile maps to dreamer schedule_strategy and early_binding in run_campaign.py: -# round_robin random strategy, early binding — diversity at intake -# diverse_top smallest_to_fastest, early bind — score + coverage -# explore_exploit largest_to_fastest, late bind — score + uncertainty -# pure_promise largest_to_fastest, late bind — greedy score-only -# -# ── cm / features block ─────────────────────────────────────────────────────── -# Feature flags for A/B comparison — run twice with a flag flipped to isolate -# the effect of each feature: -# backpressure Per-edge hysteresis queue depth controller. -# Disabling removes all THROTTLE/WIDEN state transitions. -# monitor Periodic health table + pass-through and budget drift alerts. -# Disabling suppresses all Monitor tick logs and drift warnings. -# sharder Adaptive batch dispatch from the trigger buffer. -# Disabling means every trigger_dependent() goes directly into -# the runnable queue (no buffering, no BP gate). -# bandit Thompson-sampling cross-stage scheduling bandit. -# Learns which stage to prioritize when multiple are eligible -# simultaneously. Reward: downstream BP state after each replica. -# -# ── dreamer / perf_dist and ops_dist ───────────────────────────────────────── -# Both distributions share the same schema: -# name: uniform | normal -# mean: float Centre of the distribution. -# var_spatial: float Variance across cores/tasks (fixed at run start). -# var_temporal: float Variance across time steps (optional; normal only). -# ============================================================================= - -plan_id: vax-funnel-dreamer/v1 -campaign_id: vax-funnel-dreamer -schema_version: cm-plan/1.0 -created_at: "2026-05-07T00:00:00Z" - -# ── Campaign contract ───────────────────────────────────────────────────────── -contract: - total_budget_node_hours: 100000 - deadline_hours: 504 - min_acceptable_yield: 50 - facility_caps: - frontier: 70000 - perlmutter: 30000 - -# ── Stages ──────────────────────────────────────────────────────────────────── -stages: - - # ── S1: Ligand property filter ───────────────────────────────────────────── - # Prototype: Perlmutter GPU, 300 nodes × 24 h = 7 200 node-hours - # GPU-bound: fast GPU-accelerated graph-NN filter (all stages GPU-bound). - # s1 alone demands 30 GPUs > 24 available → immediately GPU-limited at 24 concurrent. - # s1 runtime (debug): 500 × 2s ÷ 24 = 41.7s — long enough for s2→s5 pipeline to start - # and overlap with s1 for ~26s, giving the scheduling bandit rich multi-stage decisions. - # Pipeline latency s1→s5: 2+1+3+5+4=15s; baseline target=5 at ~57s; bandit target at ~15s. - # Peak GPU demand: s1(24)+s2(16)+s3(12)+s4(16)+s5(6) = 74 vs 24 → 3× oversubscription. - - id: s1_ligand_filter - upstream: library - downstream: s2_ml_affinity - threshold_top_fraction: 1.0 # score_threshold in dreamer block gates s1→s2; CM filter disabled - budget_node_hours: 7200 - pilot: { facility: perlmutter, partition: gpu, nodes: 300, walltime_h: 24 } - concurrency_cap: 30 - min_replicas: 4 # guaranteed floor so s1 isn't starved by downstream priorities - - dreamer: - use_stub: true - simulated_duration: 0.5 # DEBUG: 500×0.5s÷24GPU=10.4s s1 phase; pipeline latency=6s - simulated_jitter: 0.05 - # Score cascade: s1 generates initial random score in [0,1]. - # Only candidates with score ≥ 0.60 (top 40%) enter the s2 sharder. - # Sharding dispatches the highest-scored ones first; baseline (no sharder) dispatches FIFO. - score_noise: 0.0 # root stage — raw screen score, no refinement - score_threshold: 0.6 # top 40% pass → ~2000 enter s2 buffer - - # ── S2: ML affinity prediction ───────────────────────────────────────────── - # Prototype: Perlmutter GPU, 100 nodes × 48 h = 20 000 node-hours - # GPU-bound: ESM2/protein-ligand ML model on GPU. Drain rate = 16 / 10 s = 1.6/s. - # s2 runs concurrently with s3/s4/s5 → multi-stage GPU contention for scheduling bandit. - - id: s2_ml_affinity - upstream: s1_ligand_filter - downstream: s3_docking - threshold_top_fraction: 0.14 - budget_node_hours: 20000 - pilot: { facility: perlmutter, partition: gpu, nodes: 100, walltime_h: 48 } - surrogate: - model_ref: surrogate-s2-v1 - uncertainty_cutoff: 0.20 - cutoff_nudge_bounds: [0.10, 0.40] - threshold_top_fraction: 1.0 # CM filter disabled; dreamer score_threshold=0.35 gates s2→s3 - concurrency_cap: 16 - # All triggers dispatched — optimisations manage scheduling, not work reduction. - # BP throttles queue depth to prevent runaway cascade; bandit allocates GPUs. - sharding: { target_size: 80, min_size: 8, max_size: 160, stratify: soft, use_bandit: true, bandit_seed: 42 } - - dreamer: - use_stub: true - simulated_duration: 0.5 # DEBUG: was 5.0s - simulated_jitter: 0.05 - score_noise: 0.05 # small noise: high-scored candidates stay high across stages - score_threshold: 0.65 # ~75% of s2 inputs pass → more cascade material - - # ── S3: High-precision docking ───────────────────────────────────────────── - # Prototype: Frontier GPU, 200 nodes × 72 h = 30 000 node-hours - # GPU-bound: AutoDock-GPU / Glide GPU. s2+s3+s4+s5 compete for 24 GPU slots: - # s2=16 GPUs, s3=12 GPUs, s4=16 GPUs, s5=6 GPUs → peak demand 50 > 24 (2× oversubscription). - # Scheduling bandit arbitrates the fierce 4-way contention over ~1300s overlap window. - - id: s3_docking - upstream: s2_ml_affinity - downstream: s4_md_refinement - variant: full - threshold_top_fraction: 0.50 - budget_node_hours: 30000 - pilot: { facility: frontier, partition: gpu, nodes: 200, walltime_h: 72 } - surrogate: - model_ref: surrogate-s3-v1 - uncertainty_cutoff: 0.25 - cutoff_nudge_bounds: [0.15, 0.40] - threshold_top_fraction: 1.0 # CM filter disabled; dreamer score_threshold=0.40 gates s3→s4 - concurrency_cap: 12 - sharding: { target_size: 20, min_size: 4, max_size: 40, stratify: soft, use_bandit: true, profile: diverse_top } - - dreamer: - use_stub: true - simulated_duration: 1.0 # DEBUG: was 15.0s - simulated_jitter: 0.1 - score_noise: 0.05 - score_threshold: 0.70 # ~85% of s3 inputs pass - - # ── S4: MD refinement ────────────────────────────────────────────────────── - # Prototype: Frontier MPI+GPU, 400 nodes × 168 h = 30 000 node-hours, cap=200 - # Dependent on s3; diverse_top edge. - - id: s4_md_refinement - upstream: s3_docking - downstream: s5_fep_ranking - threshold_top_fraction: 0.60 - budget_node_hours: 30000 - pilot: { facility: frontier, partition: gpu, nodes: 400, walltime_h: 168 } - # DEBUG: changed mpi+gpu→gpu (1 GPU/replica instead of 2) so s4 can start when a - # single GPU is freed — required_gpus=2 deadlocked the pipeline while s1 ran. - threshold_top_fraction: 1.0 # CM filter disabled; dreamer score_threshold=0.45 gates s4→s5 - concurrency_cap: 8 - sharding: { target_size: 12, min_size: 4, max_size: 24, stratify: soft, use_bandit: true } - - dreamer: - use_stub: true - simulated_duration: 2.0 # DEBUG: was 30.0s - simulated_jitter: 0.2 - score_noise: 0.05 - score_threshold: 0.75 # ~80% of s4 inputs pass - - # ── S5: FEP ranking ──────────────────────────────────────────────────────── - # Prototype: Frontier LargeMem GPU, 200 nodes × 240 h = 25 000 node-hours, cap=50 - # Dependent on s4; pure_promise edge (greedy, score-only). Terminal stage. - - id: s5_fep_ranking - upstream: s4_md_refinement - downstream: final_lead_set - threshold_top_fraction: 0.60 - budget_node_hours: 25000 - pilot: { facility: frontier, partition: largemem, nodes: 200, walltime_h: 240 } - threshold_top_fraction: 1.0 # CM filter disabled; all s4 pass (score_threshold=0.0 in dreamer) - downstream_input_target: 5 # stop when 5 s5 leads produced - concurrency_cap: 6 - sharding: { target_size: 6, min_size: 2, max_size: 12, stratify: soft, use_bandit: true } - - dreamer: - use_stub: true - simulated_duration: 2.0 # DEBUG: was 22.0s - simulated_jitter: 0.4 - score_noise: 0.05 - score_threshold: 0.80 # terminal gate; ~3% of s1 candidates expected to reach here - -# ── Edges ───────────────────────────────────────────────────────────────────── -# Unchanged from prototype. -# profile → dreamer schedule_strategy via _PROFILE_STRATEGY in run_campaign.py -edges: - # Backpressure values are expressed in local emulation scale - # (proportional to each downstream stage's expected replica total). - # Queue depth = replicas triggered but not yet started. - # Throttle when depth ≥ high_water; widen when depth ≤ low_water. - - name: lib_to_s1 - upstream: library - downstream: s1_ligand_filter # total=2000 - profile: round_robin - backpressure: { high_water: 2250, low_water: 1000 } - - name: s1_to_s2 - upstream: s1_ligand_filter - downstream: s2_ml_affinity # total≈325 (500 × 0.65 — score_threshold=0.35 in dreamer) - profile: diverse_top - # Debug-scale BP: set high_water above total expected so throttle only fires on genuine - # overflow; sharding benefit comes from priority ranking, not queue control in debug runs. - backpressure: { high_water: 500, low_water: 200 } - - name: s2_to_s3 - upstream: s2_ml_affinity - downstream: s3_docking # total≈195 (325 × 0.60 — score_threshold=0.40) - profile: explore_exploit - backpressure: { high_water: 300, low_water: 120 } - - name: s3_to_s4 - upstream: s3_docking - downstream: s4_md_refinement # total≈107 (195 × 0.55 — score_threshold=0.45) - profile: diverse_top - backpressure: { high_water: 200, low_water: 80 } - - name: s4_to_s5 - upstream: s4_md_refinement - downstream: s5_fep_ranking # total≈107 (all pass, score_threshold=0.0) - profile: pure_promise - backpressure: { high_water: 200, low_water: 80 } - -# ── Bandit budget reallocation ──────────────────────────────────────────────── -# Unchanged from prototype. -bandit: - enabled: true - update_interval_hours: 2 - per_stage_band_pct: - s1_ligand_filter: 0.05 - s2_ml_affinity: 0.10 - s3_docking: 0.15 - s4_md_refinement: 0.10 - s5_fep_ranking: 0.05 - -# ── Replanning triggers ─────────────────────────────────────────────────────── -# Unchanged from prototype. -replan: - cadence_hours: 6 - budget_burn_deviation_pct: 80 # pilot costs < budget ceiling by design (safety margin); - # largest gap is s2 (50×48=2400 vs 10000 nh = 76%) — alert - # only on genuine overspend above this threshold - pass_through_deviation_pct: 10 # tight: fires when strict-sharder batching delays S3 - surrogate_recall_floor: 0.90 - deadline_slip_hours: 24 - facility_outage_hours: 4 - -# ── Provenance ──────────────────────────────────────────────────────────────── -provenance: - cm_code_commit: dreamer-emulation - inputs_hash: "" - outputs_hash: "" - seeds: { sharder: 19, bandit: 42 } - containers: - s1_ligand_filter: "ghcr.io/demo/ligand-filter@sha256:demo" - s2_ml_affinity: "ghcr.io/demo/ml-affinity@sha256:demo" - s3_docking: "ghcr.io/demo/docking@sha256:demo" - s4_md_refinement: "ghcr.io/demo/md@sha256:demo" - s5_fep_ranking: "ghcr.io/demo/fep@sha256:demo" - -# ── Debug / emulation overrides ────────────────────────────────────────────── -# SPHERICAL-specific parameters not in the cm-prototype schema. -# Controls the local radical.dreamer emulation; has no effect on real HPC runs. -debug: - - # Total replicas for independent (root) stages. - # Dependent stages grow via trigger_dependent() signals from upstream. - stage_replicas: - s1_ligand_filter: 10000 # 10000×0.5s÷24GPU=208s s1 phase; longer overlap for bandit convergence - - # Fraction of stage-N total that becomes stage-N+1 total. - # Values mirror threshold_top_fraction from the stage specs. - # Trigger fractions replaced by score-based gating in dreamer_workflow. - # Each stage triggers downstream only when output_score >= score_threshold. - # Expected funnel: 200 s1 → ~130 s2 → ~78 s3 → ~43 s4 → ~24 s5 (stops at target=5) - - -# ── SPHERICAL Campaign Manager runtime ─────────────────────────────────────── -# Not part of the cm-prototype schema. Hoisted to top-level by the plan -# translator in run_campaign.py when a "stages" key is detected. -cm: - # engine: concurrent | dragon - # concurrent — asyncio ConcurrentExecutionBackend (local, no MPI) - # dragon — radical.asyncflow DragonExecutionBackendV3 (HPC, requires Dragon) - engine: concurrent - debug: false - - # ── Optional feature flags ──────────────────────────────────────────────── - # Set to true to enable; false to run baseline without the feature. - # Enables A/B comparison: run twice, once with false and once with true. - features: - backpressure: true # per-edge hysteresis queue depth controller - monitor: true # pass-through + budget drift detection alerts - sharder: true # adaptive batch dispatch from trigger buffer - bandit: true # Thompson-sampling cross-stage scheduling bandit - - # Periodic monitor tick interval (seconds). - # Set short for local emulation so ticks appear during the ~90s campaign run. - monitor_interval_s: 8 - - resources: - total_cpus: 1024 # 30×16(s1)+16×4(s2)+12×4(s3)+8×16(s4)+6×8(s5)=768 → 1024 with headroom - total_gpus: 24 # s1(30×1)+s2(16×1)+s3(12×1)+s4(8×2)+s5(6×1)=74 > 24 → 3× oversubscription; all stages GPU-bound - - telemetry: - collect_telemetry: true - telemetry_dir: telemetry-results - - workflow_registry: - s1_ligand_filter: dreamer_workflow.DreamerWorkflow - s2_ml_affinity: dreamer_workflow.DreamerWorkflow - s3_docking: dreamer_workflow.DreamerWorkflow - s4_md_refinement: dreamer_workflow.DreamerWorkflow - s5_fep_ranking: dreamer_workflow.DreamerWorkflow diff --git a/workflows/run_campaign/dreamer_campaign/delta_cpu_sbatch.sh b/workflows/run_campaign/dreamer_campaign/delta_cpu_sbatch.sh deleted file mode 100644 index 766c799..0000000 --- a/workflows/run_campaign/dreamer_campaign/delta_cpu_sbatch.sh +++ /dev/null @@ -1,56 +0,0 @@ -#!/bin/sh -l -# -# SPHERICAL Dreamer Campaign — Delta CPU benchmark -# -# All stages GPU-bound: s1=2s(3000 reps), s2=5s, s3=15s, s4=30s, s5=22s. -# s1 GPU-limited at 24 concurrent → 250s (4.2min); all stages overlap (3× GPU oversubscription). -# Expected runtimes per run: baseline ~25min, optimised ~15min. -# Total benchmark (5 runs × 4 configs): ~5.4h → 7h walltime gives 1.6h safety margin. -# -# Submit: sbatch delta_cpu_sbatch.sh -# Logs: slurm-.out (stdout+stderr, streamed live) -# -#SBATCH -A bblj-delta-cpu -#SBATCH --partition=cpu -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=64 -#SBATCH --time=05:00:00 -#SBATCH --job-name=dreamer_bench -#SBATCH --mail-user=mariya.goliyad@rutgers.edu -#SBATCH --mail-type=ALL -#SBATCH --output=slurm-%j.out -#SBATCH --error=slurm-%j.out - -# ── Paths ───────────────────────────────────────────────────────────────────── -export SPHERICAL_DIR="/scratch/bblj/${USER}/SPHERICAL" -export DREAMER_DIR="/scratch/bblj/${USER}/radical.dreamer" -export ENV_DIR="/u/${USER}/ve/dreamer_campaign" - -# ── Activate venv ───────────────────────────────────────────────────────────── -source "${ENV_DIR}/bin/activate" - -# ── Run ─────────────────────────────────────────────────────────────────────── -CAMPAIGN_DIR="${SPHERICAL_DIR}/workflows/run_campaign/dreamer_campaign" -cd "${CAMPAIGN_DIR}" - -# Clean stale artifacts from previous runs -rm -rf dreamer-profiles telemetry-results - -echo "=== Dreamer benchmark: $(date) ===" -echo " Node: ${SLURMD_NODENAME} CPUs: ${SLURM_CPUS_PER_TASK}" -echo " Config: config.yaml Runs: 5" -echo " Stage durations: s1=2s(1500 reps) s2=5s s3=15s s4=30s s5=22s (all GPU-bound)" -echo " Workload: s2=787, s3=110, s4=55, s5=33 replicas (identical across all configs)" -echo " Expected: ~10min/run, total ~3.5h (5 runs x 4 configs)" - -python benchmark.py --config config.yaml --runs 5 --out benchmark_results.json - -echo "=== Benchmark done: $(date) ===" - -# Regenerate plots -python plot_optimizations.py --results benchmark_results.json --out-dir plots/optimizations - -echo "=== Plots written: $(date) ===" - -rm -rf asyncflow.session* \ No newline at end of file diff --git a/workflows/run_campaign/dreamer_campaign/delta_gpu_sbatch.sh b/workflows/run_campaign/dreamer_campaign/delta_gpu_sbatch.sh deleted file mode 100755 index 89d1906..0000000 --- a/workflows/run_campaign/dreamer_campaign/delta_gpu_sbatch.sh +++ /dev/null @@ -1,31 +0,0 @@ -#!/bin/sh -l - -#SBATCH -A bblj-delta-gpu -#SBATCH --partition=gpuA100x4 -#SBATCH --nodes=1 -#SBATCH --gpus-per-node=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=64 -#SBATCH --time=06:00:00 -#SBATCH --job-name=dreamer_bench -#SBATCH --mail-user=mariya.goliyad@rutgers.edu -#SBATCH --mail-type=ALL - -export SPHERICAL_DIR="/scratch/bblj/${USER}/SPHERICAL" -export DREAMER_DIR="/scratch/bblj/${USER}/radical.dreamer" -export ENV_DIR="/u/${USER}/ve/dreamer_campaign" - -source "${ENV_DIR}/bin/activate" - -CAMPAIGN_DIR="${SPHERICAL_DIR}/workflows/run_campaign/dreamer_campaign" -cd "${CAMPAIGN_DIR}" -rm -rf dreamer-profiles telemetry-results - -echo "=== Dreamer benchmark: $(date) === Node: ${SLURMD_NODENAME}" -echo "Config: 10000 s1, target=20 s5, 5 runs x 4 configs" - -python benchmark.py --config config.yaml --runs 5 --out benchmark_results.json - -python plot_optimizations.py --results benchmark_results.json --out-dir plots/optimizations - -echo "=== Done: $(date) ===" diff --git a/workflows/run_campaign/dreamer_campaign/dreamer_workflow.py b/workflows/run_campaign/dreamer_campaign/dreamer_workflow.py deleted file mode 100644 index 9d1a7bc..0000000 --- a/workflows/run_campaign/dreamer_campaign/dreamer_workflow.py +++ /dev/null @@ -1,179 +0,0 @@ -""" -DreamerWorkflow — emulates workflow tasks using radical.dreamer in peer mode. - -Score cascade model -------------------- -Each stage computes an output score from its upstream input score and fires a -downstream trigger only if the score clears a per-stage threshold. This models -a real drug-discovery funnel where each expensive stage refines the quality -estimate and only the most promising candidates proceed. - - s1 (root): output_score = uniform [0, 1] — initial ligand screen - s2–s5: output_score = input_score + N(0, score_noise) - only triggers downstream if output_score >= score_threshold - -Three outputs are forwarded to the downstream candidate metadata: - score — this stage's refined quality estimate - surrogate_pred — cheap prediction of what the NEXT stage will produce - surrogate_unc — model uncertainty (high = explore this candidate) - -The sharder uses all three (weighted by the active profile) to dispatch the -highest-value candidates first, so the pipeline converges on good leads faster -than FIFO (baseline) dispatch. - -Config keys (per stage in config.yaml, under dreamer:): - use_stub / simulated_duration / simulated_jitter — timing emulation - score_noise float Std of Gaussian noise added to input score (default 0.05) - score_threshold float Min output score to trigger downstream (default 0.0) - surr_decay float Correlation factor for surrogate prediction (default 0.90) - surr_noise float Noise on surrogate prediction (default 0.05) - -Config keys forwarded by run_campaign.py: - trigger_downstream — downstream group name - candidate_score — upstream score (None for root stage) - candidate_surr — upstream surrogate prediction - candidate_surr_unc — upstream surrogate uncertainty - candidate_scaffold — upstream scaffold class -""" - -import asyncio -import hashlib -import json -import os -import random -import sys -from pathlib import Path -from typing import ClassVar, Optional - -# Scaffold alphabet for diversity signal (8 classes, assigned by hash of candidate_id) -_SCAFFOLDS = ["scaf_A", "scaf_B", "scaf_C", "scaf_D", - "scaf_E", "scaf_F", "scaf_G", "scaf_H"] - -_dreamer_dir = os.environ.get("DREAMER_DIR") -if _dreamer_dir: - _src = str(Path(_dreamer_dir) / "src") - if _src not in sys.path: - sys.path.insert(0, _src) - -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from src.campaign import BaseWorkflow # noqa: E402 - -try: - from radical.dreamer import Resource, Workload - from radical.dreamer.configs import ScheduleConfig - from radical.dreamer.managers.ext.schedule import Schedule - from radical.dreamer.managers.resource import ResourceManager - _DREAMER_AVAILABLE = True -except ImportError: - _DREAMER_AVAILABLE = False - - -class DreamerWorkflow(BaseWorkflow): - """Async wrapper that uses radical.dreamer (peer mode) to emulate tasks.""" - - workflow_id = "dreamer" - - # ── Workflow entry point ────────────────────────────────────────────────── - - async def run(self, replica_id: str) -> None: - cfg = self.config or {} - await asyncio.to_thread(self._run_simulation, replica_id, cfg) - - # ── Simulation (runs in a thread pool worker) ───────────────────────────── - - @staticmethod - def _run_simulation(replica_id: str, cfg: dict) -> dict: - if not _DREAMER_AVAILABLE or cfg.get("use_stub"): - import time - dur = float(cfg.get("simulated_duration", 1.0)) - jitter = float(cfg.get("simulated_jitter", 0.1)) - time.sleep(dur + random.uniform(0.0, jitter)) - return {"stub": True} - - num_cores = int(cfg.get("num_cores", 32)) - perf_dist = dict(cfg.get("perf_dist", {"name": "uniform", "mean": 16.0, "var_spatial": 2.0})) - num_tasks = int(cfg.get("num_tasks", 64)) - ops_dist = dict(cfg.get("ops_dist", {"mean": 512.0})) - strategy = str(cfg.get("schedule_strategy", "smallest_to_fastest")) - early_binding = bool(cfg.get("early_binding", True)) - - resource = Resource(num_cores=num_cores, perf_dist=perf_dist) - workload = Workload(num_tasks=num_tasks, ops_dist=ops_dist) - schedule = Schedule(cfg=ScheduleConfig(from_dict={ - "strategy": strategy, "early_binding": early_binding, "is_adaptive": False, - })) - ResourceManager.processing(resource=resource, workload=workload, schedule=schedule) - return {"stub": False} - - # ── Scaffold helper ─────────────────────────────────────────────────────── - - @staticmethod - def _scaffold_for(candidate_id: str) -> str: - """Deterministic scaffold class from candidate id hash.""" - return _SCAFFOLDS[int(hashlib.md5(candidate_id.encode()).hexdigest()[:2], 16) % len(_SCAFFOLDS)] - - # ── Score cascade ───────────────────────────────────────────────────────── - - async def on_replica_done(self, replica_id: str, cm, final_state: str) -> None: - """Compute output score, apply threshold gate, trigger downstream. - - Root stage (s1): generates an initial random score in [0, 1]. - Downstream stages: refine the upstream score with Gaussian noise, - modelling each stage as a progressively more accurate quality estimate. - - Only candidates that clear score_threshold trigger the next stage. - The downstream candidate receives output_score, surrogate_pred, and - surrogate_unc so the sharder can dispatch highest-value candidates first. - """ - cfg = self.config or {} - trigger = cfg.get("trigger_downstream") - if not trigger or final_state != "done": - return - - # ── Read upstream context ───────────────────────────────────────────── - input_score = cfg.get("candidate_score") # None for root stage (s1) - scaffold = cfg.get("candidate_scaffold") or self._scaffold_for(replica_id) - - # ── Compute this stage's output score ───────────────────────────────── - noise_std = float(cfg.get("score_noise", 0.05)) - - if input_score is None: - # Root stage: initial screen produces a random quality score. - output_score = random.random() - else: - # Downstream: refine upstream estimate with stage-specific noise. - # Noise models imperfect correlation between successive assays. - output_score = max(0.0, min(1.0, - input_score + random.gauss(0.0, noise_std) - )) - - # ── Threshold gate ──────────────────────────────────────────────────── - threshold = float(cfg.get("score_threshold", 0.0)) - if output_score < threshold: - return # candidate does not proceed to next stage - - # ── Surrogate outputs for sharder priority ranking ──────────────────── - # surr_pred: predict what the NEXT stage will produce. - # Modelled as a noisy, slightly-decayed version of the current score. - surr_decay = float(cfg.get("surr_decay", 0.90)) - surr_noise = float(cfg.get("surr_noise", 0.05)) - surr_pred = max(0.0, min(1.0, - output_score * surr_decay + random.gauss(0.0, surr_noise) - )) - # surr_unc: uncertainty decreases for high-scoring candidates - # (the model is more confident about good leads). - surr_unc = 0.4 * (1.0 - output_score) - - # ── Trigger downstream with full candidate metadata ─────────────────── - cand_id = f"{replica_id}_d" - await self._trigger_dependent( - trigger, - replicas=1, - candidate_id=cand_id, - score=output_score, - surrogate_pred=surr_pred, - surrogate_unc=surr_unc, - scaffold_class=scaffold, - source_stage=self._group_name, - ) diff --git a/workflows/run_campaign/dreamer_campaign/env_setup.sh b/workflows/run_campaign/dreamer_campaign/env_setup.sh deleted file mode 100644 index 2d975ea..0000000 --- a/workflows/run_campaign/dreamer_campaign/env_setup.sh +++ /dev/null @@ -1,157 +0,0 @@ -#!/bin/bash -# ============================================================================= -# SPHERICAL Dreamer Campaign — environment setup -# -# Creates a Python venv with SPHERICAL, radical.dreamer, and the async backend. -# -# Usage: -# bash env_setup.sh [--env-dir DIR] [--spherical-dir DIR] [--dreamer-dir DIR] -# -# Defaults: -# ENV_DIR = /u/$USER/ve/dreamer_campaign -# SPHERICAL_DIR = /scratch/bblj/$USER/SPHERICAL -# DREAMER_DIR = /scratch/bblj/$USER/radical.dreamer -# ============================================================================= -if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then - set -euo pipefail -fi - -ENV_DIR="${ENV_DIR:-/u/${USER}/ve/dreamer_campaign}" -SPHERICAL_DIR="${SPHERICAL_DIR:-/scratch/bblj/${USER}/SPHERICAL}" -DREAMER_DIR="${DREAMER_DIR:-/scratch/bblj/${USER}/radical.dreamer}" - -while [[ $# -gt 0 ]]; do - case $1 in - --env-dir) ENV_DIR="$2"; shift 2 ;; - --spherical-dir) SPHERICAL_DIR="$2"; shift 2 ;; - --dreamer-dir) DREAMER_DIR="$2"; shift 2 ;; - *) echo "Unknown argument: $1"; exit 1 ;; - esac -done - -echo "=================================================================" -echo " ENV_DIR = ${ENV_DIR}" -echo " SPHERICAL_DIR = ${SPHERICAL_DIR}" -echo " DREAMER_DIR = ${DREAMER_DIR}" -echo "=================================================================" - -# ── 0. Clone repositories ───────────────────────────────────────────────────── -echo "" -echo "── Step 0: Checking repositories ──" - -if [ ! -d "${SPHERICAL_DIR}/.git" ]; then - echo "Cloning SPHERICAL → ${SPHERICAL_DIR}" - git clone git@github.com:radical-collaboration/SPHERICAL.git "${SPHERICAL_DIR}" -else - echo "SPHERICAL already cloned at ${SPHERICAL_DIR}" -fi - -if [ ! -d "${DREAMER_DIR}/.git" ]; then - echo "Cloning radical.dreamer → ${DREAMER_DIR}" - git clone https://github.com/radical-cybertools/radical.dreamer.git "${DREAMER_DIR}" -else - echo "radical.dreamer already cloned at ${DREAMER_DIR}" -fi - -# ── 1. Create venv ──────────────────────────────────────────────────────────── -echo "" -echo "── Step 1: Creating venv ──" - -BASE_PY=$(command -v python3.11 2>/dev/null || true) - -if [ -z "${BASE_PY}" ]; then - module load cray-python/3.11.7 2>/dev/null || true - BASE_PY=$(command -v python3.11 2>/dev/null || true) -fi - -if [ -z "${BASE_PY}" ]; then - module load anaconda3 2>/dev/null || true - BASE_PY=$(command -v python3.10 2>/dev/null || true) -fi - -if [ -z "${BASE_PY}" ]; then - BASE_PY=$(command -v python3 2>/dev/null || true) -fi - -if [ -z "${BASE_PY}" ]; then - echo "ERROR: no Python 3.10+ interpreter found." - exit 1 -fi - -PY_VERSION=$("${BASE_PY}" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") -PY="${ENV_DIR}/bin/python${PY_VERSION}" -PIP="${ENV_DIR}/bin/pip" -echo "Using Python: ${BASE_PY} ($(${BASE_PY} --version))" - -if [ ! -x "${PY}" ]; then - echo "Creating venv at ${ENV_DIR}..." - "${BASE_PY}" -m venv "${ENV_DIR}" -else - echo "venv already exists at ${ENV_DIR}" -fi - -ln -sf "${ENV_DIR}/bin/python${PY_VERSION}" "${ENV_DIR}/bin/python" 2>/dev/null || true -ln -sf "${ENV_DIR}/bin/python${PY_VERSION}" "${ENV_DIR}/bin/python3" 2>/dev/null || true - -# ── 2. Bootstrap pip ────────────────────────────────────────────────────────── -echo "" -echo "── Step 2: Bootstrapping pip ──" -"${PY}" -m pip install -q --upgrade pip wheel -"${PIP}" install -q --force-reinstall "setuptools<71" - -# ── 3. Async backend ───────────────────────────────────────────────────────── -echo "" -echo "── Step 3: Async backend (rhapsody + radical.asyncflow) ──" -"${PIP}" install -q \ - "rhapsody-py>=0.2.0" \ - "radical.asyncflow>=0.3.1" \ - "pyyaml" \ - "numpy>=1.26.3,<2.0.0" - -# ── 4. radical.dreamer ──────────────────────────────────────────────────────── -echo "" -echo "── Step 4: radical.dreamer ──" -"${PIP}" install -q -e "${DREAMER_DIR}" - -# ── 5. SPHERICAL ───────────────────────────────────────────────────────────── -echo "" -echo "── Step 5: SPHERICAL ──" -"${PIP}" install -q -e "${SPHERICAL_DIR}" - -# ── 6. Dreamer campaign requirements ───────────────────────────────────────── -echo "" -echo "── Step 6: dreamer_campaign requirements ──" -CAMP_DIR="${SPHERICAL_DIR}/workflows/run_campaign/dreamer_campaign" -if [ -f "${CAMP_DIR}/requirements.txt" ]; then - "${PIP}" install -q -r "${CAMP_DIR}/requirements.txt" -fi - -# ── 7. Verify ──────────────────────────────────────────────────────────────── -echo "" -echo "── Verifying installation ──" -_check() { - local label="$1"; shift - if out=$("$@" 2>&1); then - echo " ${label}: OK (${out})" - else - echo " WARNING: ${label} failed" - echo " ${out}" | head -3 - fi -} - -_check "radical.asyncflow" "${PY}" -c "import radical.asyncflow; print('ok')" -_check "rhapsody" "${PY}" -c "import rhapsody; print('ok')" -_check "radical.dreamer" "${PY}" -c "import radical.dreamer; print('ok')" -_check "spherical" "${PY}" -c "import src.campaign; print('ok')" - -echo "" -echo "=================================================================" -echo "Setup complete." -echo "" -echo "Activate with:" -echo " source ${ENV_DIR}/bin/activate" -echo "" -echo "Run the campaign:" -echo " cd ${SPHERICAL_DIR}/workflows/run_campaign/dreamer_campaign" -echo " python run_campaign.py --config config.yaml" -echo "=================================================================" diff --git a/workflows/run_campaign/dreamer_campaign/make_presentation.py b/workflows/run_campaign/dreamer_campaign/make_presentation.py deleted file mode 100644 index 3ef7b86..0000000 --- a/workflows/run_campaign/dreamer_campaign/make_presentation.py +++ /dev/null @@ -1,1431 +0,0 @@ -#!/usr/bin/env python3 -""" -Generate SPHERICAL benchmark PowerPoint presentation. - -Usage: - python make_presentation.py [--out spherical_benchmark.pptx] -""" - -import argparse -from pathlib import Path - -import matplotlib -matplotlib.use("Agg") -import matplotlib.patches as mpatches -import matplotlib.pyplot as plt -import numpy as np -from matplotlib.patches import FancyBboxPatch - -from pptx import Presentation -from pptx.util import Inches, Pt -from pptx.dml.color import RGBColor -from pptx.enum.text import PP_ALIGN - -# ── Colour constants ────────────────────────────────────────────────────────── - -WHITE = RGBColor(0xFF, 0xFF, 0xFF) -BLACK = RGBColor(0x00, 0x00, 0x00) -DARK_BG = RGBColor(0x1A, 0x1A, 0x2E) -ACCENT = RGBColor(0x0F, 0x3C, 0x78) -LIGHT_PANEL = RGBColor(0xF0, 0xF4, 0xFF) -GRAY_TEXT = RGBColor(0x55, 0x55, 0x55) -BASELINE_C = RGBColor(0x9E, 0x9E, 0x9E) -SHARD_C = RGBColor(0x4C, 0xAF, 0x50) -BANDIT_C = RGBColor(0x9C, 0x27, 0xB0) -ALLOPT_C = RGBColor(0xF4, 0x43, 0x36) -HIGHLIGHT = RGBColor(0xFF, 0xC1, 0x07) -TEAL_C = RGBColor(0x00, 0x83, 0x8F) -ORANGE_C = RGBColor(0xE6, 0x51, 0x00) - -SLIDE_W = Inches(13.33) -SLIDE_H = Inches(7.5) - - -# ── Low-level helpers ───────────────────────────────────────────────────────── - -def _bg(slide, color: RGBColor): - fill = slide.background.fill - fill.solid() - fill.fore_color.rgb = color - - -def _box(slide, l, t, w, h, text="", font_size=18, bold=False, - color=WHITE, bg=None, align=PP_ALIGN.LEFT, - font_name="Calibri", italic=False, wrap=True): - txBox = slide.shapes.add_textbox(l, t, w, h) - tf = txBox.text_frame - tf.word_wrap = wrap - p = tf.paragraphs[0] - p.alignment = align - run = p.add_run() - run.text = text - run.font.size = Pt(font_size) - run.font.bold = bold - run.font.italic = italic - run.font.color.rgb = color - run.font.name = font_name - if bg is not None: - txBox.fill.solid() - txBox.fill.fore_color.rgb = bg - return txBox - - -def _rect(slide, l, t, w, h, fill_color: RGBColor, line_color=None, line_width=0): - shape = slide.shapes.add_shape(1, l, t, w, h) - shape.fill.solid() - shape.fill.fore_color.rgb = fill_color - if line_color: - shape.line.color.rgb = line_color - shape.line.width = Pt(line_width) - else: - shape.line.fill.background() - return shape - - -def _img(slide, path, l, t, w, h=None): - if h is not None: - slide.shapes.add_picture(str(path), l, t, w, h) - else: - slide.shapes.add_picture(str(path), l, t, w) - - -def _title_bar(slide, title: str, subtitle: str = ""): - _rect(slide, Inches(0), Inches(0), SLIDE_W, Inches(1.1), ACCENT) - _box(slide, Inches(0.25), Inches(0.08), Inches(12.5), Inches(0.6), - title, font_size=28, bold=True, color=WHITE, align=PP_ALIGN.LEFT) - if subtitle: - _box(slide, Inches(0.25), Inches(0.65), Inches(12.5), Inches(0.35), - subtitle, font_size=13, color=RGBColor(0xBB, 0xCC, 0xFF), - align=PP_ALIGN.LEFT) - - -def _stat_box(slide, l, t, w, h, value, label, val_color=HIGHLIGHT, - lbl_color=WHITE, bg=ACCENT): - _rect(slide, l, t, w, h, bg) - _box(slide, l, t + Inches(0.05), w, Inches(0.55), - value, font_size=36, bold=True, color=val_color, align=PP_ALIGN.CENTER) - _box(slide, l, t + Inches(0.6), w, Inches(0.35), - label, font_size=11, color=lbl_color, align=PP_ALIGN.CENTER) - - -def _section_panel(slide, l, t, w, h, title, bullets, title_bg, title_color=WHITE, - body_bg=None, bullet_color=None, title_size=11, bullet_size=10): - """Titled panel with bullet items.""" - if body_bg is None: - body_bg = RGBColor(0xF8, 0xF8, 0xF8) - if bullet_color is None: - bullet_color = BLACK - _rect(slide, l, t, w, Inches(0.32), title_bg) - _box(slide, l + Inches(0.06), t + Inches(0.02), w - Inches(0.12), Inches(0.30), - title, font_size=title_size, bold=True, color=title_color) - _rect(slide, l, t + Inches(0.32), w, h - Inches(0.32), body_bg) - y = t + Inches(0.36) - per = (h - Inches(0.40)) / max(len(bullets), 1) - for b in bullets: - _box(slide, l + Inches(0.1), y, w - Inches(0.15), per, - f"• {b}", font_size=bullet_size, color=bullet_color) - y += per - - -# ── Slide builders ──────────────────────────────────────────────────────────── - -PLOT_DIR = Path(__file__).parent / "plots" / "optimizations" - - -def slide_title(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, DARK_BG) - _rect(slide, Inches(0), Inches(0), SLIDE_W, Inches(2.8), ACCENT) - _box(slide, Inches(0.5), Inches(0.3), Inches(12), Inches(0.7), - "SPHERICAL", font_size=18, bold=True, - color=RGBColor(0xBB, 0xCC, 0xFF), align=PP_ALIGN.LEFT) - _box(slide, Inches(0.5), Inches(0.85), Inches(12), Inches(1.1), - "Adaptive HPC Campaign Optimisation", font_size=40, bold=True, - color=WHITE, align=PP_ALIGN.LEFT) - _box(slide, Inches(0.5), Inches(1.9), Inches(12), Inches(0.6), - "Benchmark Results: 4-Configuration Drug-Discovery Pipeline Study", - font_size=18, color=RGBColor(0xBB, 0xCC, 0xFF), align=PP_ALIGN.LEFT) - - _rect(slide, Inches(0.5), Inches(3.1), Inches(5.5), Inches(2.0), - RGBColor(0x0A, 0x2A, 0x52)) - _box(slide, Inches(0.6), Inches(3.2), Inches(5.2), Inches(0.5), - "Key Result", font_size=14, bold=True, color=HIGHLIGHT, align=PP_ALIGN.LEFT) - _box(slide, Inches(0.6), Inches(3.6), Inches(5.2), Inches(0.8), - "10.9× faster time-to-target", font_size=28, bold=True, - color=WHITE, align=PP_ALIGN.LEFT) - _box(slide, Inches(0.6), Inches(4.25), Inches(5.2), Inches(0.7), - "82.6 s → 7.6 s to find 5 drug leads\nfrom 10,000 candidate ligands", - font_size=13, color=RGBColor(0xBB, 0xCC, 0xFF), align=PP_ALIGN.LEFT) - - for i, (label, col) in enumerate([ - ("Quality Routing", SHARD_C), - ("Adaptive Scheduling", BANDIT_C), - ("Combined", ALLOPT_C), - ]): - x = Inches(6.4 + i * 2.2) - _rect(slide, x, Inches(3.1), Inches(2.0), Inches(0.5), col) - _box(slide, x, Inches(3.15), Inches(2.0), Inches(0.45), - label, font_size=12, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - - _box(slide, Inches(0.5), Inches(6.9), Inches(12), Inches(0.4), - "5 independent runs per configuration · 10,000 s1 ligands · target = 5 s5 FEP completions", - font_size=10, color=RGBColor(0x77, 0x88, 0xAA), align=PP_ALIGN.CENTER) - - -def slide_pipeline_overview(prs): - """5-stage drug-discovery pipeline — no redundant optimization-axis preview.""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Drug-Discovery Pipeline", - "5-stage cascade — each stage refines quality and filters candidates") - - stages = [ - ("S1\nLigand Filter", "~3,200 start\n(10,000 queued)", "#42a5f5", "score > 0.60"), - ("S2\nML Affinity", "~313 enter", "#66bb6a", "score > 0.65"), - ("S3\nDocking", "~77 enter", "#ffa726", "score > 0.70"), - ("S4\nMD Refinement", "~21 enter", "#ef5350", "score > 0.75"), - ("S5\nFEP Ranking", "5 hit target", "#ab47bc", "score > 0.80"), - ] - bw = Inches(2.35) - gap = Inches(0.14) - for i, (name, count, col, filt) in enumerate(stages): - x = Inches(0.25) + i * (bw + gap) - c = RGBColor(int(col[1:3], 16), int(col[3:5], 16), int(col[5:7], 16)) - _rect(slide, x, Inches(1.5), bw, Inches(2.6), c) - _box(slide, x, Inches(1.58), bw, Inches(0.8), - name, font_size=15, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, x, Inches(2.4), bw, Inches(0.55), - count, font_size=11, color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, x, Inches(2.95), bw, Inches(0.85), - filt, font_size=10, italic=True, - color=RGBColor(0xEE, 0xEE, 0xEE), align=PP_ALIGN.CENTER) - if i < 4: - ax = x + bw - _box(slide, ax, Inches(2.35), gap + Inches(0.05), Inches(0.45), - "▶", font_size=22, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - - # Footnote: where counts come from - _box(slide, Inches(0.25), Inches(4.1), Inches(12.8), Inches(0.22), - "† Counts are averages from baseline benchmark runs (5 independent runs). " - "Campaign stops when 5th s5 hit is found — most of the 10,000 s1 candidates " - "never execute (early termination). Counts vary by configuration; see Cascade Funnel slide.", - font_size=8.5, italic=True, color=GRAY_TEXT) - - # Resource model description - _rect(slide, Inches(0.25), Inches(4.35), Inches(12.85), Inches(2.8), - RGBColor(0xF3, 0xF4, 0xFF)) - _box(slide, Inches(0.4), Inches(4.42), Inches(12.5), Inches(0.38), - "How SPHERICAL executes this pipeline", font_size=14, bold=True, - color=ACCENT) - - cols = [ - ("Each stage = a workflow group", - ["Replicas run in parallel within each group", - "GPU slots assigned per replica (required_gpus)", - "min_replicas ensures downstream stages always have slots", - "max_replicas caps concurrent usage per group"]), - ("Dependencies drive execution order", - ["Downstream groups start when upstream signals done", - "_signal_done() or _trigger_dependent() from workflow code", - "Sharder buffers upstream results before dispatching", - "Backpressure prevents queue flooding"]), - ("Adaptive resource allocation", - ["Thompson-sampling bandit allocates GPUs cross-stage", - "Warm-start priors favour terminal stages (s5 > s1)", - "Learns from backpressure feedback each scheduling cycle", - "Combined: routes best candidates to best-resourced stage"]), - ] - for ci, (title, bullets) in enumerate(cols): - x = Inches(0.4 + ci * 4.3) - _box(slide, x, Inches(4.85), Inches(4.1), Inches(0.35), - title, font_size=11, bold=True, color=ACCENT) - for bi, b in enumerate(bullets): - _box(slide, x + Inches(0.1), Inches(5.25 + bi * 0.4), Inches(4.0), Inches(0.38), - f"• {b}", font_size=9.5, color=GRAY_TEXT) - - -def slide_spherical_architecture(prs, diag_dir: Path): - """Merged architecture slide: CM class hierarchy + scheduler description.""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "SPHERICAL — System Architecture", - "AsyncCampaignManager: mixin-based design with optional feature flags") - - _img(slide, diag_dir / "cm_architecture.png", - Inches(0.15), Inches(1.1), Inches(8.7), Inches(5.9)) - - # Right: Scheduler + key design notes - _rect(slide, Inches(9.05), Inches(1.1), Inches(4.1), Inches(5.9), - RGBColor(0xF5, 0xF5, 0xF5)) - _box(slide, Inches(9.15), Inches(1.15), Inches(3.9), Inches(0.38), - "Scheduling Algorithm", font_size=13, bold=True, color=BLACK) - - sched_items = [ - (RGBColor(0x2E, 0x7D, 0x32), "Pass 1 — Fairness", - "For every eligible group: allocate until running == min_replicas. " - "Highest priority first. Prevents s1 from monopolising all GPUs."), - (ORANGE_C, "Pass 2 — Throughput", - "After min_replicas satisfied: fill remaining capacity up to max_replicas. " - "Highest priority (or bandit-ranked) stage gets extras first."), - (BANDIT_C, "Bandit override", - "When bandit=true: Thompson-sample Beta arm per stage to replace " - "static priority sort. Learns downstream-first allocation."), - (RGBColor(0x01, 0x57, 0x9B), "Dependency eligibility", - "Group eligible when: deps called _signal_done() OR " - "dep.finished_replicas ≥ dep_threshold."), - ] - for i, (col, title, body) in enumerate(sched_items): - y = Inches(1.6 + i * 1.35) - _rect(slide, Inches(9.05), y, Inches(4.1), Inches(0.3), col) - _box(slide, Inches(9.1), y + Inches(0.02), Inches(4.0), Inches(0.28), - title, font_size=10, bold=True, color=WHITE) - _box(slide, Inches(9.1), y + Inches(0.33), Inches(4.0), Inches(0.88), - body, font_size=9, color=GRAY_TEXT) - - _box(slide, Inches(0.15), Inches(7.1), Inches(13.1), Inches(0.3), - "Feature flags: cm.features.sharder / backpressure / bandit / monitor — all disabled by default", - font_size=9, italic=True, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - - -def slide_stage_profiles(prs, diag_dir: Path): - """NEW: Candidate ranking profiles used by the sharder.""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Candidate Ranking Profiles", - "The sharder scores candidates as: priority = Σ(weight × signal) — choose profile per campaign stage") - - _img(slide, diag_dir / "profiles.png", - Inches(0.15), Inches(1.1), Inches(8.0), Inches(5.9)) - - # Right panel: when to use each profile - _rect(slide, Inches(8.35), Inches(1.1), Inches(4.8), Inches(5.9), - RGBColor(0xF3, 0xF4, 0xFF)) - _box(slide, Inches(8.45), Inches(1.15), Inches(4.6), Inches(0.38), - "When to use each profile", font_size=13, bold=True, color=ACCENT) - - profiles_guide = [ - (SHARD_C, "pure_promise", - "Best when upstream scores are reliable.\nPure quality routing — top candidates only.\nUsed in this benchmark (sharding+bp config)."), - (BANDIT_C, "active_learning", - "Early campaign: model needs diverse data.\nMaximises uncertainty reduction.\nSacrifices short-term quality for model accuracy."), - (TEAL_C, "explore_exploit", - "Mid-campaign with calibrated surrogate.\nBalances score, model prediction, uncertainty.\nDefault for most drug-discovery campaigns."), - (ALLOPT_C, "diverse_top", - "Late campaign: avoid chemical echoes.\nQuality-weighted + scaffold novelty bonus.\nPrevents converging on one chemical series."), - (BASELINE_C, "round_robin", - "Initial screening / diversity mandate.\nRound-robin across scaffold classes.\nIgnores quality — maximises chemical diversity."), - ] - for i, (col, name, desc) in enumerate(profiles_guide): - y = Inches(1.6 + i * 1.07) - _rect(slide, Inches(8.35), y, Inches(4.8), Inches(0.28), col) - _box(slide, Inches(8.42), y + Inches(0.02), Inches(4.6), Inches(0.26), - name, font_size=10, bold=True, color=WHITE) - _box(slide, Inches(8.42), y + Inches(0.30), Inches(4.65), Inches(0.72), - desc, font_size=9, color=GRAY_TEXT) - - _box(slide, Inches(0.15), Inches(7.1), Inches(13.1), Inches(0.3), - "Profiles are configured per-group in the YAML. The sharder evaluates all buffered candidates each dispatch cycle.", - font_size=9, italic=True, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - - -def slide_benchmark_design(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Benchmark Design", - "Same workload, different optimisation features — 5 independent runs per config") - - configs = [ - ("baseline", BASELINE_C, "No features", - ["Pipelined FIFO — no quality routing", "No resource allocation learning", - "s1 monopolises GPUs; downstream starved", "Result: 82.6 s ± 19.3 s"]), - ("sharding+bp", SHARD_C, "Quality Routing", - ["Sharder: highest-score candidates dispatched first", - "Backpressure: THROTTLE/WIDEN queue control", - "min_replicas floor keeps s2–s5 slots open", - "Result: 17.0 s ± 2.5 s (4.9× faster)"]), - ("scheduling_bandit", BANDIT_C, "Adaptive Scheduling", - ["Thompson-sampling bandit per stage", - "Learns downstream-first GPU allocation", - "No quality routing — FIFO dispatch", - "Result: 26.7 s ± 5.8 s (3.1× faster)"]), - ("all_optimizations", ALLOPT_C, "Both Combined", - ["Quality routing + adaptive scheduling", - "Bandit + sharder + backpressure all active", - "Most consistent: σ=0.9 s on 7.6 s mean", - "Result: 7.6 s ± 0.9 s (10.9× faster)"]), - ] - - for i, (name, color, tag, bullets) in enumerate(configs): - x = Inches(0.25 + i * 3.27) - _rect(slide, x, Inches(1.2), Inches(3.1), Inches(0.45), color) - _box(slide, x, Inches(1.22), Inches(3.1), Inches(0.42), - f"{name}", font_size=13, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - _rect(slide, x, Inches(1.65), Inches(3.1), Inches(0.3), - RGBColor(0xEE, 0xEE, 0xEE)) - _box(slide, x, Inches(1.66), Inches(3.1), Inches(0.3), - tag, font_size=11, italic=True, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - _rect(slide, x, Inches(1.95), Inches(3.1), Inches(2.5), - RGBColor(0xF8, 0xF8, 0xF8)) - for j, b in enumerate(bullets): - bold_last = (j == len(bullets) - 1) - col = color if bold_last else BLACK - _box(slide, x + Inches(0.05), Inches(2.0 + j * 0.58), - Inches(3.0), Inches(0.55), - f"{'→' if bold_last else '•'} {b}", - font_size=11, bold=bold_last, color=col) - - _rect(slide, Inches(0.25), Inches(4.6), Inches(12.8), Inches(0.65), - RGBColor(0xE3, 0xF2, 0xFD)) - _box(slide, Inches(0.35), Inches(4.65), Inches(12.5), Inches(0.55), - "Setup: 10,000 s1 candidates · early termination when s5_fep_ranking hits 5 completions " - "· same random seed per run index across all configs · concurrent asyncio backend (no real GPU hardware)", - font_size=11, color=RGBColor(0x0D, 0x47, 0xA1)) - - _box(slide, Inches(0.25), Inches(5.35), Inches(12.8), Inches(1.95), - "Note: 'baseline' is pipelined FIFO (not a true sequential waterfall) — stages start as " - "soon as the first upstream replica completes. This makes baseline HARDER to beat than a " - "pure waterfall, so the measured speedups are conservative.", - font_size=11, italic=True, color=GRAY_TEXT) - - -def slide_cascade_funnel(prs): - """Standalone cascade funnel — shows pipeline compute cost per config.""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Cascade Funnel — Total Pipeline Work per Configuration", - "How many replicas were launched at each stage to find 5 s5 FEP hits") - - _img(slide, PLOT_DIR / "3_cascade_funnel.png", - Inches(0.15), Inches(1.15), Inches(8.7), Inches(5.0)) - - # Right: key numbers - _rect(slide, Inches(9.05), Inches(1.15), Inches(4.1), Inches(5.0), - RGBColor(0xF5, 0xF5, 0xF5)) - _box(slide, Inches(9.15), Inches(1.2), Inches(3.9), Inches(0.38), - "Total replicas launched", font_size=12, bold=True, color=BLACK) - - funnel_stats = [ - (BASELINE_C, "baseline", "3,622 replicas\ns1 monopolises pipeline"), - (SHARD_C, "sharding+bp", "726 replicas\n5× less compute"), - (BANDIT_C, "scheduling_bandit", "1,109 replicas\n3.3× less compute"), - (ALLOPT_C, "all_optimizations", "119 replicas\n30× less s1 work"), - ] - for i, (col, name, stat) in enumerate(funnel_stats): - y = Inches(1.65 + i * 1.1) - _rect(slide, Inches(9.05), y, Inches(4.1), Inches(1.0), col) - _box(slide, Inches(9.12), y + Inches(0.04), Inches(3.9), Inches(0.34), - name, font_size=11, bold=True, color=WHITE) - _box(slide, Inches(9.12), y + Inches(0.40), Inches(3.9), Inches(0.52), - stat, font_size=13, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - - _box(slide, Inches(0.15), Inches(6.2), Inches(13.0), Inches(1.0), - "Left: stacked bars show absolute replica counts per config — s1 dominates baseline. " - "Right: log scale reveals all 5 stages. Sharding dispatches only the top-scoring ~20% of s1 " - "results downstream — drastically shrinking every subsequent stage.", - font_size=9.5, italic=True, color=GRAY_TEXT) - - -def slide_main_result(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Main Result — Wall Time to Target", - "Time from campaign start until 5th s5_fep_ranking completion · 5 runs per config") - - _img(slide, PLOT_DIR / "1_wall_time.png", - Inches(0.2), Inches(1.15), Inches(7.8), Inches(5.5)) - - stats = [ - ("82.6 s", "baseline\n(±19.3 s)", BASELINE_C), - ("17.0 s", "sharding+bp\n4.9× faster", SHARD_C), - ("26.7 s", "scheduling_bandit\n3.1× faster", BANDIT_C), - ("7.6 s", "all_optimizations\n10.9× faster", ALLOPT_C), - ] - for i, (val, lbl, col) in enumerate(stats): - y = Inches(1.2 + i * 1.5) - _rect(slide, Inches(8.3), y, Inches(4.8), Inches(1.3), col) - _box(slide, Inches(8.3), y + Inches(0.08), Inches(4.8), Inches(0.65), - val, font_size=38, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, Inches(8.3), y + Inches(0.72), Inches(4.8), Inches(0.5), - lbl, font_size=12, color=WHITE, align=PP_ALIGN.CENTER) - - _box(slide, Inches(8.3), Inches(7.1), Inches(4.8), Inches(0.3), - "Lower is better · white dots = individual runs", - font_size=9, italic=True, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - - -def slide_sharding_bp(prs, diag_dir: Path = None): - """Optimisation 1 — fully pptx-native layout (no embedded image).""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Optimisation 1 — Quality Routing", - "Sharder + BackpressureNegotiator + Shard Bandit · 17.0 s ± 2.5 s (4.9× faster)") - - # ── Left column: visual diagrams ────────────────────────────────────── - - # --- Data flow section label --- - _rect(slide, Inches(0.15), Inches(1.18), Inches(8.55), Inches(0.28), - RGBColor(0xE8, 0xF5, 0xE9)) - _box(slide, Inches(0.22), Inches(1.20), Inches(8.4), Inches(0.25), - "SHARDER — upstream trigger → buffer → rank → dispatch", - font_size=9.5, bold=True, color=RGBColor(0x1B, 0x5E, 0x20)) - - # Data flow boxes (4 boxes + arrows) - flow_y = Inches(1.52) - flow_h = Inches(1.6) - box_w = Inches(1.88) - arr_w = Inches(0.36) - gap = arr_w - flow_items = [ - (Inches(0.15), "S1 Replica\nDone", SHARD_C, - "trigger(cid,\nscore, surr,\nuncertainty,\nscaffold)"), - (Inches(0.15) + box_w + gap, "BUFFER", RGBColor(0x1B, 0x5E, 0x20), - "all candidates\nenqueued\nsorted by\npriority score"), - (Inches(0.15) + 2*(box_w + gap), "Ranking\nEngine", BANDIT_C, - "priority =\nΣ weight\n× signal\n(profile)"), - (Inches(0.15) + 3*(box_w + gap), "S2 Queue", RGBColor(0x01, 0x57, 0x9B), - "highest-score\ncandidates\ndispatched\nfirst"), - ] - for i, (x, name, col, sub) in enumerate(flow_items): - _rect(slide, x, flow_y, box_w, flow_h, col) - _box(slide, x + Inches(0.05), flow_y + Inches(0.07), - box_w - Inches(0.1), Inches(0.44), - name, font_size=11, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, x + Inches(0.05), flow_y + Inches(0.52), - box_w - Inches(0.1), Inches(0.98), - sub, font_size=9, color=RGBColor(0xDD, 0xFF, 0xDD), align=PP_ALIGN.CENTER) - if i < 3: - arr_x = x + box_w - _box(slide, arr_x, flow_y + Inches(0.65), arr_w, Inches(0.35), - "▶", font_size=18, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - - # Adaptive batch sizing note (below flow) - _rect(slide, Inches(0.15), Inches(3.18), Inches(8.55), Inches(0.28), - RGBColor(0x00, 0x60, 0x64)) - _box(slide, Inches(0.22), Inches(3.20), Inches(8.4), Inches(0.25), - "Batch sizing: stratify=soft → adaptive; stratify=strict → hold until full batch; " - "shard bandit learns multiplier arms [0.5 0.75 1.0 1.25 1.5]", - font_size=9, color=WHITE) - - # --- Backpressure state machine --- - _rect(slide, Inches(0.15), Inches(3.55), Inches(8.55), Inches(0.28), - RGBColor(0xFF, 0xE0, 0xB2)) - _box(slide, Inches(0.22), Inches(3.57), Inches(8.4), Inches(0.25), - "BACKPRESSURE NEGOTIATOR — hysteresis state machine controlling dispatch rate", - font_size=9.5, bold=True, color=RGBColor(0xE6, 0x51, 0x00)) - - # Three state boxes - BP_Y = Inches(3.9) - BP_H = Inches(1.55) - BP_W = Inches(2.35) - states = [ - (Inches(0.15), "HOLD", "normal operation", - "dispatch proceeds\nat normal rate", RGBColor(0x43, 0xA0, 0x47)), - (Inches(2.97), "THROTTLE", "queue ≥ high_water", - "dispatch = 0\npipeline paused", RGBColor(0xE5, 0x39, 0x35)), - (Inches(5.8), "WIDEN", "queue ≤ low_water", - "dispatch × mult\nqueue drained", RGBColor(0x1E, 0x88, 0xE5)), - ] - for x, title, cond, action, col in states: - _rect(slide, x, BP_Y, BP_W, BP_H, col) - _box(slide, x + Inches(0.05), BP_Y + Inches(0.06), - BP_W - Inches(0.1), Inches(0.38), - title, font_size=14, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, x + Inches(0.05), BP_Y + Inches(0.44), - BP_W - Inches(0.1), Inches(0.3), - cond, font_size=8.5, italic=True, - color=RGBColor(0xFF, 0xFF, 0xCC), align=PP_ALIGN.CENTER) - _box(slide, x + Inches(0.05), BP_Y + Inches(0.78), - BP_W - Inches(0.1), Inches(0.6), - action, font_size=10, color=WHITE, align=PP_ALIGN.CENTER) - - # Transition arrows between states (text-based) - _box(slide, Inches(2.52), BP_Y + Inches(0.58), Inches(0.45), Inches(0.4), - "▶", font_size=18, color=RGBColor(0xE5, 0x39, 0x35), align=PP_ALIGN.CENTER) - _box(slide, Inches(5.35), BP_Y + Inches(0.58), Inches(0.45), Inches(0.4), - "▶", font_size=18, color=RGBColor(0x1E, 0x88, 0xE5), align=PP_ALIGN.CENTER) - - # Return path label - _rect(slide, Inches(0.15), BP_Y + BP_H + Inches(0.05), - Inches(8.55), Inches(0.3), RGBColor(0x43, 0xA0, 0x47)) - _box(slide, Inches(0.22), BP_Y + BP_H + Inches(0.07), - Inches(8.4), Inches(0.25), - "◀─── when queue returns to normal range, state resets to HOLD ───────────────────────", - font_size=9, color=WHITE) - - # Config params - _box(slide, Inches(0.15), BP_Y + BP_H + Inches(0.42), Inches(8.55), Inches(0.28), - "Config: backpressure_high (high_water mark) · backpressure_low (low_water mark) " - "· queue_depth = group.replicas − group.started_count", - font_size=8.5, italic=True, color=GRAY_TEXT) - - # ── Right column: bullet descriptions ───────────────────────────────── - RW = Inches(4.3) - RX = Inches(8.95) - - _section_panel( - slide, RX, Inches(1.18), RW, Inches(1.85), - "Sharder (sharder.py)", - ["Buffer upstream trigger results (score, surrogate, uncertainty, scaffold)", - "Rank candidates: priority = Σ(weight × signal) via ProfileWeights", - "Dispatch best candidates to s2 first (stratify=soft: adaptive batch size)", - "Flush buffer when upstream stage completes"], - title_bg=SHARD_C, - body_bg=RGBColor(0xE8, 0xF5, 0xE9), - bullet_color=RGBColor(0x1B, 0x5E, 0x20), - bullet_size=9.5, - ) - - _section_panel( - slide, RX, Inches(3.13), RW, Inches(1.95), - "BackpressureNegotiator (backpressure.py)", - ["HOLD → queue depth between thresholds → dispatch normal", - "THROTTLE → depth ≥ high_water → pause dispatch (return 0)", - "WIDEN → depth ≤ low_water → dispatch × multiplier (> 1)", - "Hysteresis prevents rapid oscillation between states", - "Config: backpressure_high / backpressure_low per group"], - title_bg=RGBColor(0xE6, 0x51, 0x00), - body_bg=RGBColor(0xFF, 0xF3, 0xE0), - bullet_color=RGBColor(0x7F, 0x3B, 0x00), - bullet_size=9.5, - ) - - _section_panel( - slide, RX, Inches(5.18), RW, Inches(1.42), - "Shard Bandit (optional, bandit.py)", - ["Arms = [0.5, 0.75, 1.0, 1.25, 1.5] (dispatch multipliers)", - "Thompson-samples arm to adjust batch size each dispatch cycle", - "Reward = throughput improvement over last window", - "Learns optimal batch size for current pipeline state"], - title_bg=BANDIT_C, - body_bg=RGBColor(0xF3, 0xE5, 0xF5), - bullet_color=RGBColor(0x4A, 0x14, 0x8C), - bullet_size=9.5, - ) - - # Result callout - _rect(slide, RX, Inches(6.7), RW, Inches(0.68), SHARD_C) - _box(slide, RX + Inches(0.08), Inches(6.73), RW - Inches(0.15), Inches(0.28), - "Result: 17.0 s ± 2.5 s · 4.9× faster wall time", - font_size=11, bold=True, color=WHITE) - _box(slide, RX + Inches(0.08), Inches(7.01), RW - Inches(0.15), Inches(0.28), - "5× fewer total replicas launched (3,622 → 726)", font_size=11, color=WHITE) - - -def slide_bandit(prs): - """Optimisation 2 — Scheduling Bandit: GPU utilization + algorithm description.""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Optimisation 2 — Adaptive Scheduling (Thompson-sampling Bandit)", - "Learns to allocate freed GPU slots to the highest-value stage · 26.7 s ± 5.8 s (3.1× faster)") - - # Left: GPU utilization plot — full height - _img(slide, PLOT_DIR / "4_gpu_utilization.png", - Inches(0.15), Inches(1.15), Inches(8.6), Inches(5.85)) - - # Right: Algorithm description - RW = Inches(4.3) - RX = Inches(8.95) - - _rect(slide, RX, Inches(1.15), RW, Inches(0.35), BANDIT_C) - _box(slide, RX + Inches(0.08), Inches(1.17), RW - Inches(0.15), Inches(0.32), - "Thompson Sampling Algorithm", font_size=11, bold=True, color=WHITE) - - algo_steps = [ - "1. GPU slot freed → collect eligible stages", - "2. Sample θᵢ ~ Beta(αᵢ, βᵢ) for each stage", - "3. Assign slot to stage with highest θ", - "4. Replica runs → measure downstream BP state", - "5. Compute reward r ∈ [0,1] (see below)", - "6. Update posterior: αᵢ += r, βᵢ += (1-r)", - ] - _rect(slide, RX, Inches(1.5), RW, Inches(2.2), RGBColor(0xF3, 0xE5, 0xF5)) - for j, step in enumerate(algo_steps): - _box(slide, RX + Inches(0.1), Inches(1.53 + j * 0.35), RW - Inches(0.15), Inches(0.34), - step, font_size=9.5, color=RGBColor(0x4A, 0x14, 0x8C)) - - # Warm-start priors - _rect(slide, RX, Inches(3.8), RW, Inches(0.3), RGBColor(0x4A, 0x14, 0x8C)) - _box(slide, RX + Inches(0.08), Inches(3.82), RW - Inches(0.15), Inches(0.28), - "Warm-start priors", font_size=10, bold=True, color=WHITE) - _rect(slide, RX, Inches(4.1), RW, Inches(0.95), RGBColor(0xED, 0xE7, 0xF6)) - priors = [ - ("s5 FEP Ranking", "Beta(5,1)", BANDIT_C), - ("s4 MD Refine", "Beta(4,1)", RGBColor(0xEF, 0x53, 0x50)), - ("s3 Docking", "Beta(3,1)", RGBColor(0xFF, 0xA7, 0x26)), - ("s2 ML Affinity", "Beta(2,1)", RGBColor(0x66, 0xBB, 0x6A)), - ("s1 Ligand Filter", "Beta(1,1) neutral", BASELINE_C), - ] - for j, (stage, prior, col) in enumerate(priors): - x_off = j * (RW / 5) - _box(slide, RX + x_off, Inches(4.13), RW / 5, Inches(0.42), - f"{prior}\n{stage.split()[0]}", font_size=8, bold=False, - color=col, align=PP_ALIGN.CENTER) - - # Reward signal - _rect(slide, RX, Inches(5.15), RW, Inches(0.3), RGBColor(0x4A, 0x14, 0x8C)) - _box(slide, RX + Inches(0.08), Inches(5.17), RW - Inches(0.15), Inches(0.28), - "Reward signal (from BP state)", font_size=10, bold=True, color=WHITE) - _rect(slide, RX, Inches(5.45), RW, Inches(0.85), RGBColor(0xED, 0xE7, 0xF6)) - for j, (label, r, col) in enumerate([ - ("THROTTLE (queue flooded)", "r = 0.2", RGBColor(0xC6, 0x28, 0x28)), - ("HOLD (queue healthy)", "r = 0.5-0.8", RGBColor(0x2E, 0x7D, 0x32)), - ("WIDEN (queue drained)", "r = 0.8", RGBColor(0x15, 0x65, 0xC0)), - ]): - _box(slide, RX + Inches(0.1), Inches(5.48 + j * 0.28), RW - Inches(0.2), Inches(0.27), - f"• {label} → {r}", font_size=9.5, color=col) - - # Result callout - _rect(slide, RX, Inches(6.4), RW, Inches(0.6), BANDIT_C) - _box(slide, RX + Inches(0.08), Inches(6.43), RW - Inches(0.15), Inches(0.27), - "s5 first start: baseline 14.5 s → bandit 6.9 s (2.1×)", font_size=10, - bold=True, color=WHITE) - _box(slide, RX + Inches(0.08), Inches(6.7), RW - Inches(0.15), Inches(0.27), - "Wall time 3.1× faster · 3.3× fewer total replicas", font_size=10, color=WHITE) - - -def slide_all_opt(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Optimisation 3 — Combined (all_optimizations)", - "Quality routing + adaptive scheduling — both axes active simultaneously") - - _img(slide, PLOT_DIR / "7_time_to_target.png", - Inches(0.15), Inches(1.15), Inches(8.0), Inches(5.1)) - - _rect(slide, Inches(8.3), Inches(1.15), Inches(4.8), Inches(5.1), - RGBColor(0xFF, 0xEB, 0xEE)) - _box(slide, Inches(8.4), Inches(1.2), Inches(4.6), Inches(0.4), - "Why combined > each alone", font_size=14, bold=True, - color=RGBColor(0xB7, 0x1C, 0x1C)) - - for j, txt in enumerate([ - "Sharding routes high-quality candidates\n to s2–s5 FIRST", - "Bandit allocates s5 GPUs from t~1 s\n (warm-start prior)", - "Together: s5 gets the BEST candidates\n with the MOST GPU resources", - "s5 hits target with only 10 s5 starts\n from just 119 s1 replicas", - "15× less total compute vs baseline", - ]): - _box(slide, Inches(8.4), Inches(1.65 + j * 0.72), Inches(4.6), Inches(0.7), - f"• {txt}", font_size=11, color=RGBColor(0xB7, 0x1C, 0x1C)) - - _rect(slide, Inches(8.3), Inches(5.4), Inches(4.8), Inches(0.85), ALLOPT_C) - _box(slide, Inches(8.35), Inches(5.42), Inches(4.7), Inches(0.42), - "10.9× faster · 15× less compute", font_size=22, bold=True, - color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, Inches(8.35), Inches(5.82), Inches(4.7), Inches(0.38), - "7.6 s ± 0.9 s (most consistent of all configs)", - font_size=11, color=WHITE, align=PP_ALIGN.CENTER) - - _box(slide, Inches(0.15), Inches(6.3), Inches(13.0), Inches(0.85), - "Time-to-target step curves: faint lines = all 5 runs per config; bold = median run. " - "▼ markers show when each config crosses N=5. " - "Expected independent speedup = 4.9×3.1=15.2×; actual 10.9× shows moderate overlap " - "(both optimisations reduce wasted compute, so savings partially overlap).", - font_size=9, italic=True, color=GRAY_TEXT) - - -def slide_gantt(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Pipeline Stage Overlap — Gantt View", - "Average first-start to last-finish per stage; more overlap = better pipeline utilisation") - - _img(slide, PLOT_DIR / "2_pipeline_gantt.png", - Inches(0.15), Inches(1.15), Inches(9.5), Inches(5.8)) - - _rect(slide, Inches(9.8), Inches(1.15), Inches(3.3), Inches(5.8), - RGBColor(0xF5, 0xF5, 0xF5)) - _box(slide, Inches(9.9), Inches(1.2), Inches(3.1), Inches(0.4), - "What to look for", font_size=13, bold=True, color=BLACK) - - insights = [ - (BASELINE_C, "baseline", - "All stages sequential — s1 finishes before s2 fills up."), - (SHARD_C, "sharding+bp", - "s2–s5 bars start within seconds of s1 due to min_replicas floor."), - (BANDIT_C, "scheduling_bandit", - "s3/s4/s5 overlap deeply with s1 — bandit feeds terminal stages while s1 still runs."), - (ALLOPT_C, "all_optimizations", - "All 5 bars nearly co-incident; campaign ends at t=7.6 s."), - ] - for j, (col, name, desc) in enumerate(insights): - y = Inches(1.7 + j * 1.3) - _rect(slide, Inches(9.85), y, Inches(0.18), Inches(0.28), col) - _box(slide, Inches(10.1), y - Inches(0.02), Inches(2.9), Inches(0.32), - name, font_size=11, bold=True, color=col) - _box(slide, Inches(10.0), y + Inches(0.3), Inches(3.0), Inches(0.75), - desc, font_size=10, color=GRAY_TEXT) - - -def slide_planner_execution(prs, diag_dir: Path): - """NEW: SPHERICAL execution model — from YAML config to running replicas.""" - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "SPHERICAL — Execution Model", - "From campaign YAML config to asyncflow engine: how the planner and executor relate") - - _img(slide, diag_dir / "planner.png", - Inches(0.15), Inches(1.1), Inches(8.7), Inches(5.9)) - - # Right: key concepts - _rect(slide, Inches(9.05), Inches(1.1), Inches(4.1), Inches(5.9), - RGBColor(0xF5, 0xF5, 0xF5)) - _box(slide, Inches(9.15), Inches(1.15), Inches(3.9), Inches(0.38), - "Execution Plan Concepts", font_size=13, bold=True, color=BLACK) - - plan_items = [ - (RGBColor(0x01, 0x57, 0x9B), "YAML Config = Plan", - "Each workflow group is a named pool of replicas. " - "Dependencies define the DAG. " - "Resources (CPUs/GPUs) cap concurrent execution."), - (SHARD_C, "Dynamic plan update", - "_trigger_dependent(name, N) lets an upstream workflow " - "add N replicas to a downstream group at runtime — " - "the plan adapts to intermediate results."), - (BANDIT_C, "Scheduler as planner", - "Every state change re-runs _schedule_locked(). " - "The two-pass algorithm is the 'planner' that decides " - "which groups get resources each cycle."), - (ALLOPT_C, "asyncflow = executor", - "CM creates asyncio tasks; asyncflow engine manages " - "the event loop, task lifecycle, and backend " - "(concurrent or Dragon HPC)."), - ] - for i, (col, title, body) in enumerate(plan_items): - y = Inches(1.6 + i * 1.35) - _rect(slide, Inches(9.05), y, Inches(4.1), Inches(0.3), col) - _box(slide, Inches(9.1), y + Inches(0.02), Inches(4.0), Inches(0.28), - title, font_size=10, bold=True, color=WHITE) - _box(slide, Inches(9.1), y + Inches(0.33), Inches(4.0), Inches(0.88), - body, font_size=9, color=GRAY_TEXT) - - _box(slide, Inches(0.15), Inches(7.1), Inches(13.1), Inches(0.3), - "Sync wrapper (CampaignManager) provides a blocking API for non-async callers; " - "AsyncCampaignManager is the native async class.", - font_size=9, italic=True, color=GRAY_TEXT, align=PP_ALIGN.CENTER) - - -def slide_methodology(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, WHITE) - _title_bar(slide, "Methodology — What's Real vs What to Watch", - "Verified findings and known caveats") - - confirmed = [ - "10.9× wall-time speedup is correctly measured (time from start to 5th s5 finish)", - "Cascade funnel reduction (15× less total work) uses n_started — accurate", - "Same random seed per run index ensures consistent score distributions across configs", - "All 5 runs per config completed successfully (no timeouts or failures)", - "all_optimizations has the lowest run-to-run variance (σ/mean = 12% vs 23% for baseline)", - ] - caveats = [ - "Baseline label says 'waterfall' but dep_threshold_override is commented out — " - "baseline is actually pipelined FIFO, making speedups MORE conservative", - "all_optimizations over-provisions s5: ~10 s5 replicas start to find 5 hits " - "(bandit warm-start Beta(5,1) is aggressive — 2× s5 waste)", - "BP fractions show 100% WIDEN for all runs — queue never hit high-water mark " - "(BP controlled dispatch rate but never fully throttled in these short runs)", - "Runs use asyncio concurrent backend (no real GPU hardware) — timing models " - "stage durations with simulated sleep + jitter, not actual compute", - ] - - _rect(slide, Inches(0.25), Inches(1.15), Inches(6.2), Inches(5.5), - RGBColor(0xE8, 0xF5, 0xE9)) - _box(slide, Inches(0.35), Inches(1.2), Inches(5.9), Inches(0.4), - "✓ Confirmed — results are real", font_size=13, bold=True, - color=RGBColor(0x1B, 0x5E, 0x20)) - for j, txt in enumerate(confirmed): - _box(slide, Inches(0.4), Inches(1.65 + j * 0.95), Inches(5.9), Inches(0.9), - f"✓ {txt}", font_size=10.5, color=RGBColor(0x1B, 0x5E, 0x20)) - - _rect(slide, Inches(6.7), Inches(1.15), Inches(6.4), Inches(5.5), - RGBColor(0xFF, 0xF9, 0xC4)) - _box(slide, Inches(6.8), Inches(1.2), Inches(6.1), Inches(0.4), - "⚠ Caveats — known limitations", font_size=13, bold=True, - color=RGBColor(0xE6, 0x5C, 0x00)) - for j, txt in enumerate(caveats): - _box(slide, Inches(6.85), Inches(1.65 + j * 1.22), Inches(6.1), Inches(1.1), - f"⚠ {txt}", font_size=10.5, color=RGBColor(0x7F, 0x3B, 0x00)) - - -def slide_summary(prs): - slide = prs.slides.add_slide(prs.slide_layouts[6]) - _bg(slide, DARK_BG) - _rect(slide, Inches(0), Inches(0), SLIDE_W, Inches(1.1), ACCENT) - _box(slide, Inches(0.3), Inches(0.12), Inches(12.5), Inches(0.85), - "Summary", font_size=32, bold=True, color=WHITE) - - stats = [ - ("10.9×", "wall-time speedup\nall_optimizations vs baseline", ALLOPT_C), - ("15×", "less total compute\n(replicas launched)", SHARD_C), - ("7.6 s", "median time-to-5-hits\nall_optimizations (±0.9 s)", ALLOPT_C), - ("4.9× / 3.1×", "sharding+bp / bandit\nindividual gains", BANDIT_C), - ] - for i, (val, lbl, col) in enumerate(stats): - x = Inches(0.25 + i * 3.27) - _rect(slide, x, Inches(1.25), Inches(3.1), Inches(1.7), col) - _box(slide, x, Inches(1.3), Inches(3.1), Inches(0.9), - val, font_size=34, bold=True, color=WHITE, align=PP_ALIGN.CENTER) - _box(slide, x, Inches(2.15), Inches(3.1), Inches(0.65), - lbl, font_size=11, color=WHITE, align=PP_ALIGN.CENTER) - - takeaways = [ - ("Quality routing wins on compute", SHARD_C, - "Routing high-score candidates first cuts total pipeline work by 5×. " - "The cascade stays narrow — only the best leads reach expensive downstream stages."), - ("Bandit wins on latency", BANDIT_C, - "Thompson-sampling allocation gets s5 slots filled 2× earlier than baseline. " - "Even without quality filtering, earlier resource allocation cuts wall time 3×."), - ("Combined is super-additive", ALLOPT_C, - "Best candidates reach a well-resourced s5 simultaneously. " - "Result: 10.9× speedup and 12% run-to-run variance — the most reliable configuration."), - ] - for i, (title, col, body) in enumerate(takeaways): - x = Inches(0.25 + i * 4.37) - _rect(slide, x, Inches(3.1), Inches(4.1), Inches(0.38), col) - _box(slide, x + Inches(0.05), Inches(3.12), Inches(4.0), Inches(0.35), - title, font_size=12, bold=True, color=WHITE) - _rect(slide, x, Inches(3.48), Inches(4.1), Inches(2.5), - RGBColor(0x0A, 0x2A, 0x52)) - _box(slide, x + Inches(0.08), Inches(3.52), Inches(3.95), Inches(2.42), - body, font_size=11, color=RGBColor(0xCC, 0xDD, 0xFF)) - - _box(slide, Inches(0.3), Inches(6.1), Inches(12.5), Inches(0.3), - "Experiment: 10,000 s1 ligands · target = 5 s5 FEP completions · 5 independent runs · asyncio concurrent backend", - font_size=9, color=RGBColor(0x77, 0x88, 0xAA), align=PP_ALIGN.CENTER) - - -# ── Architecture diagram generators ────────────────────────────────────────── - -DIAG_DIR = Path(__file__).parent / "plots" / "diagrams" - -_BG = "#F4F6FB" -_NAVY = "#1A237E" -_BLUE = "#1565C0" -_GRN = "#2E7D32" -_PRP = "#6A1B9A" -_ORG = "#E65100" -_RED = "#B71C1C" -_GRY = "#37474F" -_LBL = "#546E7A" -_TEAL = "#00838F" - - -def _fbox(ax, x, y, w, h, label, sublabel="", fc="#1565C0", tc="white", - fs=11, sfs=8.5, radius=0.05, lw=1.5): - ax.add_patch(FancyBboxPatch((x, y), w, h, - boxstyle=f"round,pad={radius}", - facecolor=fc, edgecolor="white", linewidth=lw, zorder=3)) - ly = y + h * (0.62 if sublabel else 0.5) - ax.text(x + w / 2, ly, label, ha="center", va="center", - fontsize=fs, fontweight="bold", color=tc, zorder=4) - if sublabel: - ax.text(x + w / 2, y + h * 0.28, sublabel, ha="center", va="center", - fontsize=sfs, color=tc, alpha=0.88, zorder=4, fontstyle="italic") - - -def _arrow(ax, x0, y0, x1, y1, color="#555", lw=1.5, style="->"): - ax.annotate("", xy=(x1, y1), xytext=(x0, y0), - arrowprops=dict(arrowstyle=style, color=color, - lw=lw, connectionstyle="arc3,rad=0")) - - -def _label(ax, x, y, text, fs=9, color="#333", ha="center", va="center", - bold=False, italic=False): - ax.text(x, y, text, ha=ha, va=va, fontsize=fs, color=color, - fontweight="bold" if bold else "normal", - fontstyle="italic" if italic else "normal", zorder=5) - - -def _diag_save(fig, path, facecolor=_BG): - fig.patch.set_facecolor(facecolor) - plt.savefig(path, dpi=150, bbox_inches="tight", facecolor=facecolor) - plt.close() - - -# ── Diagram 1: CM Architecture ──────────────────────────────────────────────── - -def make_cm_arch_diagram(path: Path) -> None: - fig, ax = plt.subplots(figsize=(14, 8.5)) - ax.set_xlim(0, 14); ax.set_ylim(0, 8.5); ax.axis("off") - - _fbox(ax, 0.3, 7.0, 13.4, 1.2, - "AsyncCampaignManager", - "from_config(yaml, registry, asyncflow) · start() · wait() · close() · metrics()", - fc=_NAVY, fs=18, sfs=10) - - mixin_specs = [ - ("SchedulerMixin", "Two-pass greedy scheduler\nPass 1: guarantee min_replicas\nPass 2: fill to max_replicas\nPriority / bandit ordering", _GRN, 0.3), - ("ExecutorMixin", "Replica lifecycle\nLaunch → monitor → complete\nGPU ID assignment\nEarly termination", _ORG, 4.85), - ("MonitorMixin", "Periodic health checks\nDrift detection\nStall alerting\nBP transitions", _PRP, 9.4), - ] - for name, desc, fc, x in mixin_specs: - _fbox(ax, x, 4.8, 4.2, 2.0, name, desc, fc=fc, fs=13, sfs=9.5) - _arrow(ax, x + 2.1, 7.0, x + 2.1, 6.8, color="white") - - struct_specs = [ - ("_GroupInfo", "replicas · min/max_replicas\npriority · dependencies\nstatus · running_count", "#01579B"), - ("ResourcePool", "total_cpus · total_gpus\ncan_fit() / allocate()\nrelease()", "#01579B"), - ("CampaignMetrics", "replica_events\nscheduling_events\nbp_fractions · shard_events", "#01579B"), - ("BaseWorkflow", "_signal_done()\n_trigger_dependent()\nrun() or start()", _GRY), - ] - for i, (name, desc, fc) in enumerate(struct_specs): - x = 0.3 + i * 3.42 - _fbox(ax, x, 2.7, 3.1, 1.85, name, desc, fc=fc, fs=10.5, sfs=8.5) - if i < 3: - _arrow(ax, x + 1.55, 4.8, x + 1.55, 4.55, color="#aaa") - - ax.text(7.0, 2.55, "Core data structures", ha="center", va="center", - fontsize=9, style="italic", color=_LBL) - - feat_specs = [ - ("Sharder", "Buffer → Rank → Dispatch\nAdaptive batch sizing\nShard bandit arm"), - ("BackpressureNegotiator", "HOLD / THROTTLE / WIDEN\nHysteresis queue control\nDispatch multiplier"), - ("SchedulingBandit", "Thompson sampling\nCross-stage GPU allocation\nBeta arm per stage"), - ("CandidateLog", "Score + surrogate history\nScaffold class diversity\nPriority ranking"), - ] - for i, (name, desc) in enumerate(feat_specs): - x = 0.3 + i * 3.42 - _fbox(ax, x, 0.4, 3.1, 1.9, name, desc, fc=_GRN, fs=10.5, sfs=8.5) - - ax.text(7.0, 0.22, "Optional features — enabled via feature flags in config YAML", - ha="center", va="center", fontsize=9, style="italic", color=_GRN) - - for i, flag in enumerate(["sharder=true", "backpressure=true", "bandit=true", ""]): - if flag: - ax.text(0.3 + i * 3.42 + 1.55, 2.45, flag, - ha="center", va="center", fontsize=7.5, color=_GRN, style="italic", - bbox=dict(boxstyle="round,pad=0.2", fc="#E8F5E9", ec=_GRN, lw=0.8)) - - ax.set_title("AsyncCampaignManager — Class Architecture", fontsize=14, - fontweight="bold", color=_NAVY, pad=6) - _diag_save(fig, path) - - -# ── Diagram 2: Sharder + Backpressure ───────────────────────────────────────── - -def make_sharder_diagram(path: Path) -> None: - fig, ax = plt.subplots(figsize=(14, 8)) - ax.set_xlim(0, 14); ax.set_ylim(0, 8); ax.axis("off") - - _fbox(ax, 0.2, 4.5, 2.5, 2.4, - "Upstream\nStage (s1)", - "on_replica_done()\n→ score, surr_pred,\n surr_unc computed\n→ _trigger_dependent()", fc=_BLUE, fs=12, sfs=9) - - _arrow(ax, 2.7, 5.7, 3.2, 5.7, color=_GRN, lw=2) - ax.text(2.95, 5.95, "trigger(candidate_id,\n score, surr_pred,\n surr_unc, scaffold)", - ha="center", va="bottom", fontsize=7.5, color=_GRN) - - buf_fc = "#E8F5E9" - ax.add_patch(FancyBboxPatch((3.2, 3.5), 2.8, 4.2, - boxstyle="round,pad=0.12", - facecolor=buf_fc, edgecolor=_GRN, linewidth=2, zorder=2)) - ax.text(4.6, 7.4, "BUFFER", ha="center", va="center", - fontsize=13, fontweight="bold", color=_GRN, zorder=5) - ax.text(4.6, 7.05, "candidate queue", ha="center", va="center", - fontsize=9, color=_GRN, style="italic", zorder=5) - - for yi, (score, lbl) in enumerate([ - (0.94, "0.94"), (0.88, "0.88"), (0.81, "0.81"), - (0.75, "0.75"), (0.68, "0.68"), (0.61, "0.61"), - ]): - y = 6.55 - yi * 0.48 - c = plt.cm.RdYlGn(score) - ax.add_patch(plt.Circle((3.95, y), 0.17, color=c, zorder=6)) - ax.text(4.25, y, f"score={lbl}", va="center", fontsize=7.5, color=_GRY, zorder=6) - - _fbox(ax, 6.3, 4.5, 3.5, 2.4, - "Ranking Engine", - "priority =\nw_score × score\n+ w_surr × surr_pred\n+ w_unc × surr_unc\n+ w_age × age", - fc=_PRP, fs=12, sfs=8.5) - _arrow(ax, 6.0, 5.7, 6.3, 5.7, color=_PRP, lw=2) - ax.text(6.15, 5.95, "dispatch()", ha="center", va="bottom", fontsize=8, color=_PRP) - - _fbox(ax, 10.1, 4.5, 3.0, 2.4, - "Downstream\nStage (s2)", - "receives candidates\nin priority order\n(highest score first)\nvia _pending_candidates", - fc=_BLUE, fs=12, sfs=9) - _arrow(ax, 9.8, 5.7, 10.1, 5.7, color=_NAVY, lw=2) - ax.text(9.95, 5.95, "priority-ranked\nreplicas", ha="center", va="bottom", - fontsize=7.5, color=_NAVY) - - _fbox(ax, 6.3, 2.3, 3.5, 1.9, - "Adaptive Batch Sizing", - "target_size = base × BP_multiplier\nstratify=soft: tail dispatch allowed\nstratify=strict: hold until full\nshard bandit learns multiplier", - fc="#006064", fs=11, sfs=8) - - ax.add_patch(FancyBboxPatch((0.2, 0.3), 5.6, 3.8, - boxstyle="round,pad=0.1", - facecolor="#FFF9C4", edgecolor="#F57F17", linewidth=2, zorder=2)) - ax.text(3.0, 3.8, "BackpressureNegotiator", ha="center", va="center", - fontsize=12, fontweight="bold", color="#E65100", zorder=5) - - states = [("HOLD\n(normal)", 1.0, 2.6, "#43A047"), - ("THROTTLE\n(queue too deep)", 3.0, 2.6, "#E53935"), - ("WIDEN\n(queue drained)", 5.0, 2.6, "#1E88E5")] - for lbl, x, y, c in states: - ax.add_patch(plt.Circle((x, y), 0.55, color=c, zorder=4)) - ax.text(x, y, lbl, ha="center", va="center", fontsize=7.5, - fontweight="bold", color="white", zorder=5) - - for (x0, y0), (x1, y1), lbl, c in [ - ((1.55, 2.85), (2.45, 2.85), "q ≥ high_water", "#E53935"), - ((3.55, 2.35), (4.45, 2.35), "q ≤ low_water", "#1E88E5"), - ((4.45, 2.85), (2.6, 2.85), "q in range → HOLD", "#43A047"), - ]: - _arrow(ax, x0, y0, x1, y1, color=c) - ax.text((x0 + x1) / 2, (y0 + y1) / 2 + 0.18, lbl, - ha="center", fontsize=7, color=c) - - for x, lbl in [(1.0, "dispatch\nnormal"), (3.0, "dispatch\n= 0"), (5.0, "dispatch\n× mult")]: - ax.text(x, 1.8, lbl, ha="center", va="center", fontsize=7.5, color=_GRY, - bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="#ccc", lw=0.8)) - - ax.text(3.0, 0.6, "queue_depth = group.replicas − group.started_count", - ha="center", fontsize=8, color=_GRY, style="italic") - - _arrow(ax, 5.8, 2.6, 6.3, 2.8, color="#F57F17", lw=1.5) - ax.text(6.1, 2.85, "BP state\n→ multiplier", ha="center", fontsize=7.5, color="#E65100") - - ax.set_title("Sharder Module — Buffering, Priority Ranking, and Dispatch Control", - fontsize=13, fontweight="bold", color=_NAVY, pad=6) - _diag_save(fig, path) - - -# ── Diagram 3: Thompson-Sampling Bandit ─────────────────────────────────────── - -def make_bandit_diagram(path: Path) -> None: - fig = plt.figure(figsize=(14, 9)) - fig.patch.set_facecolor(_BG) - - fig.text(0.5, 0.97, "Scheduling Bandit — Thompson Sampling per Stage", - ha="center", va="top", fontsize=14, fontweight="bold", color=_NAVY) - - arm_specs = [ - ("s1\nligand filter", 1, 1, "#42a5f5", "Beta(1,1)\n(uniform prior)"), - ("s2\nML affinity", 2, 1, "#66bb6a", "Beta(2,1)"), - ("s3\ndocking", 3, 1, "#ffa726", "Beta(3,1)"), - ("s4\nMD refine", 4, 1, "#ef5350", "Beta(4,1)"), - ("s5\nFEP rank", 5, 1, "#ab47bc", "Beta(5,1)\n(warm-start: prefers s5)"), - ] - x_positions = np.linspace(0.06, 0.88, 5) - ax_width, ax_height = 0.155, 0.26 - ax_y = 0.65 - - xs = np.linspace(0.001, 0.999, 300) - for i, (stage_lbl, a, b, color, prior_lbl) in enumerate(arm_specs): - ax = fig.add_axes([x_positions[i], ax_y, ax_width, ax_height]) - log_pdf = (a - 1) * np.log(xs) + (b - 1) * np.log(1 - xs) - pdf = np.exp(log_pdf - log_pdf.max()) - pdf = pdf / (np.trapz(pdf, xs) if hasattr(np, "trapz") else np.trapezoid(pdf, xs)) - ax.fill_between(xs, pdf, alpha=0.5, color=color) - ax.plot(xs, pdf, color=color, linewidth=2) - ax.set_xlim(0, 1); ax.set_ylim(0) - ax.set_xlabel("θ (priority)", fontsize=7) - ax.set_title(stage_lbl, fontsize=8, fontweight="bold", color=color, pad=2) - ax.tick_params(labelsize=6) - ax.text(0.5, ax.get_ylim()[1] * 0.75, prior_lbl, - ha="center", fontsize=6.5, color=color, style="italic") - sample_theta = a / (a + b) - ax.axvline(sample_theta, color=color, linestyle="--", linewidth=1.5, alpha=0.8) - ax.text(sample_theta, ax.get_ylim()[1] * 0.12, f"θ={sample_theta:.2f}", - ha="center", fontsize=6, color=color, - bbox=dict(boxstyle="round,pad=0.15", fc="white", ec=color, lw=0.8)) - - main_ax = fig.add_axes([0.0, 0.0, 1.0, 1.0], facecolor="none") - main_ax.set_xlim(0, 14); main_ax.set_ylim(0, 9); main_ax.axis("off") - - for xi in x_positions: - main_ax.annotate("", xy=(xi * 14, 5.7), xytext=(xi * 14, 5.95), - arrowprops=dict(arrowstyle="->", color="#888", lw=1.2)) - - _fbox(main_ax, 2.5, 4.95, 9.0, 0.65, - "Sample θᵢ ~ Betaᵢ(αᵢ, βᵢ) for each eligible stage", - "Thompson sample — exploration/exploitation trade-off", fc="#37474F", fs=12, sfs=9) - _arrow(main_ax, 7.0, 4.95, 7.0, 4.65, color="#555") - - _fbox(main_ax, 2.5, 4.0, 9.0, 0.65, - "Rank stages by θ → highest θ gets next freed GPU slot", - "deterministic tie-breaking by registration order", fc=_PRP, fs=12, sfs=9) - _arrow(main_ax, 7.0, 4.0, 7.0, 3.7, color="#555") - - _fbox(main_ax, 2.5, 3.05, 9.0, 0.65, - "Replica executes → on_replica_done → compute reward r ∈ [0, 1]", - "", fc=_BLUE, fs=12) - _arrow(main_ax, 7.0, 3.05, 7.0, 2.75, color="#555") - - reward_specs = [ - (1.8, "THROTTLE\n(downstream queue full)", 0.2, "#E53935"), - (5.5, "HOLD\n(queue healthy)", "1 − 0.5×util", "#43A047"), - (9.5, "WIDEN\n(queue drained)", 0.8, "#1E88E5"), - ] - for rx, lbl, r, c in reward_specs: - main_ax.add_patch(FancyBboxPatch((rx - 1.5, 1.6), 3.0, 0.9, - boxstyle="round,pad=0.08", - facecolor=c, edgecolor="white", lw=1.5, zorder=3, alpha=0.9)) - main_ax.text(rx, 2.22, lbl, ha="center", va="center", - fontsize=9, fontweight="bold", color="white", zorder=4) - main_ax.text(rx, 1.85, f"r = {r}", ha="center", va="center", - fontsize=9, color="white", style="italic", zorder=4) - _arrow(main_ax, rx, 2.75, rx, 2.5, color=c) - - _fbox(main_ax, 2.5, 0.85, 9.0, 0.65, - "Bayesian update: αᵢ ← αᵢ + r βᵢ ← βᵢ + (1 − r)", - "positive reward → arm shifts right (higher priority in next sample)", fc=_GRN, fs=12, sfs=9) - - for rx in [1.8, 5.5, 9.5]: - _arrow(main_ax, rx, 1.6, rx, 1.5, color="#888") - _arrow(main_ax, rx, 1.5, 7.0, 1.5, color="#888") - _arrow(main_ax, 7.0, 1.5, 7.0, 0.85, color="#888") - - main_ax.annotate("", xy=(0.5, 6.4), xytext=(0.5, 0.85), - arrowprops=dict(arrowstyle="->", color=_GRN, lw=2, - connectionstyle="arc3,rad=0.0")) - main_ax.text(0.18, 3.5, "update\nposterior", ha="center", va="center", - fontsize=9, color=_GRN, fontweight="bold", rotation=90) - - _diag_save(fig, path) - - -# ── Diagram 4: Candidate Profiles ───────────────────────────────────────────── - -def make_profiles_diagram(path: Path) -> None: - fig, (ax_heat, ax_desc) = plt.subplots(1, 2, figsize=(14, 7), - gridspec_kw={"width_ratios": [1.15, 1]}) - fig.patch.set_facecolor(_BG) - ax_heat.set_facecolor(_BG) - ax_desc.set_facecolor(_BG) - - profiles = ["pure_promise", "active_learning", "explore_exploit", - "diverse_top", "round_robin"] - weights = np.array([ - [1.0, 0.0, 0.0, 0.05, 0.0], - [0.0, 0.0, 1.0, 0.05, 0.0], - [0.5, 0.3, 0.4, 0.05, 0.0], - [0.6, 0.0, 0.1, 0.05, 0.3], - [0.0, 0.0, 0.0, 0.05, 1.0], - ]) - signal_names = ["Score", "Surrogate", "Uncertainty", "Age", "Diversity"] - row_colors = ["#4caf50", "#9c27b0", "#00838f", "#f44336", "#78909c"] - - im = ax_heat.imshow(weights, cmap="YlGn", vmin=0, vmax=1.0, aspect="auto") - ax_heat.set_xticks(range(5)) - ax_heat.set_xticklabels(signal_names, fontsize=11, fontweight="bold", color=_NAVY) - ax_heat.set_yticks(range(5)) - ax_heat.set_yticklabels(profiles, fontsize=10.5, fontweight="bold") - - for tick_lbl, col in zip(ax_heat.get_yticklabels(), row_colors): - tick_lbl.set_color(col) - - for i in range(5): - for j in range(5): - v = weights[i, j] - txt_col = "white" if v > 0.55 else ("black" if v > 0.15 else "#aaaaaa") - ax_heat.text(j, i, f"{v:.2f}", ha="center", va="center", - fontsize=12, fontweight="bold", color=txt_col) - - ax_heat.set_title("Weight Matrix (greener = higher weight)", - fontsize=12, fontweight="bold", color=_NAVY, pad=10) - plt.colorbar(im, ax=ax_heat, fraction=0.046, pad=0.04) - - ax_desc.axis("off") - ax_desc.set_xlim(0, 1); ax_desc.set_ylim(0, 1) - ax_desc.set_title("Use Cases", fontsize=12, fontweight="bold", color=_NAVY, pad=10) - - descs = [ - ("pure_promise", "#4caf50", - "Greedy quality routing.\nRanks purely by upstream score.\nTiny age bonus prevents starvation.\nBest when scores are reliable."), - ("active_learning", "#9c27b0", - "Uncertainty-first dispatch.\nMaximises surrogate model learning.\nSacrifices short-term quality.\nBest early in a campaign."), - ("explore_exploit", "#00838f", - "Balanced exploration.\nScore + surrogate + uncertainty.\nGood with a calibrated model.\nDefault for most campaigns."), - ("diverse_top", "#f44336", - "Quality + chemical diversity.\nScaffold novelty bonus via MMR.\nPrevents chemical echo chambers.\nBest for diverse leads."), - ("round_robin", "#78909c", - "Pure diversity mandate.\nRound-robin across scaffold classes.\nIgnores quality entirely.\nInitial screening / mandate."), - ] - for i, (name, col, desc) in enumerate(descs): - y_top = 0.94 - i * 0.195 - ax_desc.add_patch(FancyBboxPatch((0.01, y_top - 0.15), 0.98, 0.16, - boxstyle="round,pad=0.01", - facecolor=col, alpha=0.12, - edgecolor=col, linewidth=1.5)) - ax_desc.text(0.04, y_top, name, fontsize=10, fontweight="bold", color=col, va="top") - ax_desc.text(0.04, y_top - 0.04, desc, fontsize=8.5, color=_GRY, va="top", - style="italic") - - plt.tight_layout(pad=2.5) - _diag_save(fig, path) - - -# ── Diagram 5: Planner / Execution Model ────────────────────────────────────── - -def make_planner_diagram(path: Path) -> None: - fig, ax = plt.subplots(figsize=(14, 8.5)) - ax.set_xlim(0, 14); ax.set_ylim(0, 8.5); ax.axis("off") - - # ── Column 1: User inputs ──────────────────────────────────────────────── - ax.text(1.95, 8.3, "User Inputs", ha="center", fontsize=12, - fontweight="bold", color=_NAVY) - - _fbox(ax, 0.2, 6.0, 3.5, 2.1, - "Campaign Config (YAML)", - "workflows:\n s1: replicas=10000, priority=10\n s2: dependencies=[s1]\n min_replicas=1\n ...", - fc=_BLUE, fs=11, sfs=8.5) - - _fbox(ax, 0.2, 3.5, 3.5, 2.2, - "Workflow Classes", - "class SimWorkflow(BaseWorkflow):\n async def run(self, rid):\n await do_work()\n await self._signal_done()", - fc=_GRY, fs=10, sfs=8) - - _fbox(ax, 0.2, 1.2, 3.5, 2.0, - "Resource Spec", - "engine: concurrent # or dragon\ntotal_gpus: 4\ntotal_cpus: 128\nfeatures:\n sharder: true", - fc="#455A64", fs=10, sfs=8.5) - - # ── Arrows → CM ────────────────────────────────────────────────────────── - for y_mid in [7.05, 4.6, 2.2]: - _arrow(ax, 3.7, y_mid, 4.5, y_mid, color=_NAVY, lw=2) - ax.text(4.1, 4.9, "from_config()", ha="center", fontsize=8.5, - color=_NAVY, style="italic", - bbox=dict(boxstyle="round,pad=0.2", fc=_BG, ec=_NAVY, lw=0.8)) - - # ── Column 2: AsyncCampaignManager ─────────────────────────────────────── - ax.add_patch(FancyBboxPatch((4.5, 0.5), 5.2, 7.75, - boxstyle="round,pad=0.15", - facecolor="#E3F2FD", edgecolor=_NAVY, - linewidth=2.5, zorder=2)) - ax.text(7.1, 8.1, "AsyncCampaignManager", ha="center", - fontsize=13, fontweight="bold", color=_NAVY) - - # Dependency graph box - ax.add_patch(FancyBboxPatch((4.7, 5.65), 4.8, 2.4, - boxstyle="round,pad=0.1", - facecolor="#BBDEFB", edgecolor=_BLUE, linewidth=1.5, zorder=3)) - ax.text(7.1, 7.85, "Dependency Graph", ha="center", - fontsize=9, fontweight="bold", color=_BLUE, zorder=4) - - stage_cols = ["#42a5f5", "#66bb6a", "#ffa726", "#ef5350", "#ab47bc"] - stage_names = ["s1", "s2", "s3", "s4", "s5"] - for i, (sn, sc) in enumerate(zip(stage_names, stage_cols)): - nx = 5.1 + i * 0.95 - ax.add_patch(plt.Circle((nx, 6.85), 0.32, color=sc, zorder=5)) - ax.text(nx, 6.85, sn, ha="center", va="center", - fontsize=9, fontweight="bold", color="white", zorder=6) - if i < 4: - _arrow(ax, nx + 0.32, 6.85, nx + 0.63, 6.85, color=_GRY, lw=1.5) - ax.text(7.1, 6.2, "group.status: eligible → scheduling → running → done", - ha="center", fontsize=7.5, color=_GRY, zorder=4, style="italic") - - # Scheduler box - _fbox(ax, 4.7, 3.85, 4.8, 1.7, - "Scheduler (per state change)", - "1. Flush sharder buffers + refresh BP\n" - "2. Collect eligible groups\n" - "3. Pass 1: guarantee min_replicas\n" - "4. Pass 2: fill to max_replicas", - fc=_GRN, fs=10, sfs=8.5) - - # Optional features box - _fbox(ax, 4.7, 2.15, 4.8, 1.5, - "Optional Features (feature flags)", - "sharder=true · backpressure=true\nbandit=true · monitor=true\nAll disabled by default", - fc=_PRP, fs=10, sfs=8.5) - - # Metrics box - _fbox(ax, 4.7, 0.65, 4.8, 1.3, - "CampaignMetrics", - "replica_events · scheduling_events\nbp_fractions · shard_events", - fc=_GRY, fs=9.5, sfs=8) - - # ── Arrow → asyncflow ──────────────────────────────────────────────────── - _arrow(ax, 9.7, 4.5, 10.4, 4.5, color=_GRN, lw=2.5) - ax.text(10.05, 4.85, "create_task()", ha="center", fontsize=8.5, - color=_GRN, style="italic", - bbox=dict(boxstyle="round,pad=0.2", fc=_BG, ec=_GRN, lw=0.8)) - - # ── Column 3: asyncflow Engine ─────────────────────────────────────────── - ax.text(12.0, 8.3, "Execution Engine", ha="center", fontsize=12, - fontweight="bold", color=_NAVY) - - _fbox(ax, 10.4, 5.5, 3.3, 2.7, - "asyncflow\nWorkflowEngine", - "ConcurrentBackend\n(asyncio — local)\n── or ──\nDragonBackend\n(HPC multi-node)", - fc=_ORG, fs=12, sfs=9) - - _fbox(ax, 10.4, 3.1, 3.3, 2.2, - "Running Replicas", - "SimWorkflow.run(replica_0)\nSimWorkflow.run(replica_1)\n...\n(up to max_replicas concurrent)", - fc=_GRY, fs=10, sfs=8) - - _arrow(ax, 12.05, 5.5, 12.05, 5.3, color=_GRY, lw=2) - - _fbox(ax, 10.4, 1.2, 3.3, 1.75, - "Results", - "on_replica_done() callbacks\n_signal_done() / _trigger_dependent()\nCampaignMetrics updated", - fc=_BLUE, fs=10, sfs=8.5) - - _arrow(ax, 12.05, 3.1, 12.05, 2.95, color=_GRY, lw=2) - - # Feedback arrow - ax.annotate("", xy=(7.1, 3.85), xytext=(10.4, 1.8), - arrowprops=dict(arrowstyle="->", color=_BLUE, lw=1.8, - connectionstyle="arc3,rad=-0.3")) - ax.text(9.5, 2.6, "signal / trigger\nstate update", ha="center", fontsize=8, - color=_BLUE, style="italic") - - ax.set_title("SPHERICAL — From Config to Execution (asyncio event-loop model)", - fontsize=14, fontweight="bold", color=_NAVY, pad=6) - _diag_save(fig, path) - - -def generate_diagrams() -> Path: - DIAG_DIR.mkdir(parents=True, exist_ok=True) - print(" Generating architecture diagrams...") - make_cm_arch_diagram(DIAG_DIR / "cm_architecture.png") - print(" cm_architecture.png") - make_sharder_diagram(DIAG_DIR / "sharder.png") - print(" sharder.png") - make_bandit_diagram(DIAG_DIR / "bandit.png") - print(" bandit.png") - make_profiles_diagram(DIAG_DIR / "profiles.png") - print(" profiles.png") - make_planner_diagram(DIAG_DIR / "planner.png") - print(" planner.png") - return DIAG_DIR - - -# ── Assemble ────────────────────────────────────────────────────────────────── - -def build(out_path: str) -> None: - prs = Presentation() - prs.slide_width = SLIDE_W - prs.slide_height = SLIDE_H - - diag_dir = generate_diagrams() - - print("Building slides...") - slide_title(prs); print(" 1. Title") - slide_pipeline_overview(prs); print(" 2. Pipeline overview") - slide_spherical_architecture(prs, diag_dir); print(" 3. SPHERICAL architecture") - slide_stage_profiles(prs, diag_dir); print(" 4. Stage profiles") - slide_benchmark_design(prs); print(" 5. Benchmark design") - slide_cascade_funnel(prs); print(" 6. Cascade funnel") - slide_main_result(prs); print(" 7. Main result (wall time)") - slide_sharding_bp(prs, diag_dir); print(" 8. Optimisation 1 — Sharder+BP+Bandit") - slide_bandit(prs); print(" 9. Optimisation 2 — Scheduling bandit") - slide_all_opt(prs); print(" 10. Optimisation 3 — Combined") - slide_gantt(prs); print(" 11. Pipeline Gantt") - slide_planner_execution(prs, diag_dir); print(" 12. Planner & execution model") - slide_methodology(prs); print(" 13. Methodology / caveats") - slide_summary(prs); print(" 14. Summary") - - prs.save(out_path) - print(f"\nSaved: {out_path} ({len(prs.slides)} slides)") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--out", default="spherical_benchmark.pptx") - args = parser.parse_args() - build(args.out) diff --git a/workflows/run_campaign/dreamer_campaign/plot_dreamer_timeline.py b/workflows/run_campaign/dreamer_campaign/plot_dreamer_timeline.py deleted file mode 100644 index 9bded03..0000000 --- a/workflows/run_campaign/dreamer_campaign/plot_dreamer_timeline.py +++ /dev/null @@ -1,737 +0,0 @@ -#!/usr/bin/env python3 -""" -Plot dreamer campaign timeline and simulation statistics. - -Reads: - - A campaign log file (ANSI-colored CM output) - - dreamer profile JSON files from dreamer-profiles/ (auto-detected next to log) - -Produces a 3-row figure: - Row 0 Gantt chart (wall-clock replica execution) + stage config table - Row 1 CPU / GPU resource utilization over wall-clock time - Row 2 Dreamer simulation metrics: - 2a Simulated makespan per replica, grouped by stage - 2b Task ops distribution per stage (box plots from profile JSONs) - 2c Per-stage summary statistics table - -Usage: - python plot_dreamer_timeline.py log [--profiles-dir DIR] [--config FILE] [--out FILE] -""" - -import argparse -import json -import re -import sys -from datetime import datetime -from pathlib import Path - -try: - import yaml - _HAVE_YAML = True -except ImportError: - _HAVE_YAML = False - -import matplotlib -import matplotlib.ticker -matplotlib.use("Agg") -import matplotlib.gridspec as gridspec -import matplotlib.lines as mlines -import matplotlib.patches as mpatches -import matplotlib.pyplot as plt -import numpy as np - -# ── Stage appearance ────────────────────────────────────────────────────────── - -STAGE_COLORS = { - "s1_ligand_filter": "#4C72B0", - "s2_ml_affinity": "#DD8452", - "s3_docking": "#55A868", - "s4_md_refinement": "#C44E52", - "s5_fep_ranking": "#8172B2", -} -STAGE_ORDER = [ - "s1_ligand_filter", - "s2_ml_affinity", - "s3_docking", - "s4_md_refinement", - "s5_fep_ranking", -] -STAGE_LABELS = { - "s1_ligand_filter": "S1 Filter", - "s2_ml_affinity": "S2 ML", - "s3_docking": "S3 Dock", - "s4_md_refinement": "S4 MD", - "s5_fep_ranking": "S5 FEP", -} -DEFAULT_COLOR = "#9E9E9E" - -# ── Log regexes ─────────────────────────────────────────────────────────────── - -_TS_RE = re.compile(r"\x1b\[2m(\d{2}:\d{2}:\d{2}\.\d{3})\x1b\[0m") -_START_RE = re.compile(r"starting replica '(\S+)'") -_FINISH_RE = re.compile(r"Replica '(\S+)' finished") -_ERROR_RE = re.compile(r"Replica '(\S+)' raised") -_GROUP_RE = re.compile( - r"Registered group '(\S+)': replicas=(\d+) priority=(\d+) " - r"min=(\d+) max=(\d+) deps=\[([^\]]*)\] dep_threshold=(\d+) " - r"resources=\(cpus=(\d+), gpus=(\d+)\)" -) -_USAGE_RE = re.compile(r"resources: cpus=(\d+)/(\d+)\s+gpus=(\d+)/(\d+)") -_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+)") -_TRIGGER_RE = re.compile(r"trigger_dependent: '(\S+)' \+(\d+) replicas \(total=(\d+)\)") -_STALL_RE = re.compile(r"Group '(\S+)' stalled — waiting for resources") - -# dreamer: 128 cores × 256 tasks → 256 completed avg_exec=15.89 makespan=42.82 (strategy=random) -_DREAMER_RE = re.compile( - r"\[(\S+)\] dreamer: (\d+) cores × (\d+) tasks → (\d+) completed" - r"\s+avg_exec=([\d.]+)\s+makespan=([\d.]+)\s+\(strategy=(\w+)\)" -) - - -# ── Log parser ──────────────────────────────────────────────────────────────── - -def parse_log(path): - starts = {} - spans = [] - group_meta = {} - resource_timeline = [] - signal_events = [] # (elapsed_s, tgt_group, n) - dreamer_stats = {} # replica_id → dict - stall_events = [] # (elapsed_s, group) - t0_dt = None - total_cpus = total_gpus = 0 - - _iso_re = re.compile(r"^(\d{4}-\d{2}-\d{2}) \d{2}:\d{2}:\d{2}") - date_ref = "1970-01-01" - with open(path) as fh: - for raw in fh: - if m := _iso_re.match(raw): - date_ref = m.group(1) - break - - with open(path) as fh: - for raw in fh: - if m := _GROUP_RE.search(raw): - name = m.group(1) - deps = [d.strip().strip("'\"") for d in m.group(6).split(",") - if d.strip().strip("'\"")] - group_meta[name] = { - "replicas": int(m.group(2)), - "priority": int(m.group(3)), - "min": int(m.group(4)), - "max": int(m.group(5)), - "deps": deps, - "dep_threshold": int(m.group(7)), - "cpus": int(m.group(8)), - "gpus": int(m.group(9)), - } - if m := _TOTAL_RES_RE.search(raw): - total_cpus, total_gpus = int(m.group(1)), int(m.group(2)) - if m := _DREAMER_RE.search(raw): - dreamer_stats[m.group(1)] = { - "num_cores": int(m.group(2)), - "num_tasks": int(m.group(3)), - "tasks_completed": int(m.group(4)), - "avg_exec": float(m.group(5)), - "makespan": float(m.group(6)), - "strategy": m.group(7), - } - - ts_m = _TS_RE.search(raw) - if ts_m is None: - continue - 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 := _USAGE_RE.search(raw): - 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)) - 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)) - resource_timeline.append((elapsed, tc - ac, tc, tg - ag, tg)) - - if m := _TRIGGER_RE.search(raw): - signal_events.append((elapsed, m.group(1), int(m.group(2)))) - if m := _STALL_RE.search(raw): - stall_events.append((elapsed, m.group(1))) - - if m := _START_RE.search(raw): - starts[m.group(1)] = dt - elif m := _FINISH_RE.search(raw): - rid = m.group(1) - if rid in starts: - spans.append((rid, _group(rid), starts.pop(rid), dt, True)) - elif m := _ERROR_RE.search(raw): - rid = m.group(1) - if rid in starts: - spans.append((rid, _group(rid), starts.pop(rid), dt, False)) - - for rid, start in starts.items(): - spans.append((rid, _group(rid), start, start, None)) - - 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, signal_events, \ - dreamer_stats, stall_events, t0, (total_cpus, total_gpus) - - -def _group(rid): - return rid.rsplit("_", 1)[0] - - -# ── Profile JSON loader ─────────────────────────────────────────────────────── - -def load_profiles(profiles_dir): - profiles = {} - pdir = Path(profiles_dir) - if not pdir.exists(): - return profiles - for fpath in sorted(pdir.rglob("*.json")): - try: - data = json.loads(fpath.read_text()) - rid = data.get("replica_id") or fpath.stem - profiles[rid] = data - except Exception: - pass - return profiles - - -def parse_config(path): - """Parse group metadata from config.yaml (flat or plan format).""" - if not _HAVE_YAML: - return {} - with open(path) as fh: - cfg = yaml.safe_load(fh) - out = {} - - if "stages" in cfg: - # cm-prototype plan format - _PILOT_RES = { - "cpu": {"cpus": 16, "gpus": 0}, - "gpu": {"cpus": 4, "gpus": 1}, - "mpi+gpu": {"cpus": 16, "gpus": 2}, - "largemem": {"cpus": 8, "gpus": 1}, - } - stage_ids = {s["id"] for s in cfg.get("stages", [])} - scale = float(cfg.get("cm", {}).get("concurrency_scale", 1.0)) - for s in cfg.get("stages", []): - sid = s["id"] - upstream = s.get("upstream", "") - deps = [upstream] if upstream in stage_ids else [] - cap = int(s.get("concurrency_cap", 0)) - max_r = int(s.get("max_replicas", - max(1, round(cap * scale)) if cap else 0)) - pilot = s.get("pilot", {}) - res = _PILOT_RES.get(pilot.get("partition", "cpu").lower(), - {"cpus": 4, "gpus": 0}) - out[sid] = { - "replicas": int(s.get("replicas", 0 if deps else 1)), - "priority": int(s.get("priority", 0)), - "min": 0, - "max": max_r, - "deps": deps, - "dep_threshold": int(s.get("dependency_threshold", 1)), - "cpus": res["cpus"], - "gpus": res["gpus"], - } - else: - # Legacy flat format - for name, wf in cfg.get("workflows", {}).items(): - has_deps = bool(wf.get("dependencies", [])) - out[name] = { - "replicas": int(wf.get("replicas", 0 if has_deps else 1)), - "priority": int(wf.get("priority", 0)), - "min": int(wf.get("min_replicas", 0)), - "max": int(wf.get("max_replicas", 0)), - "deps": list(wf.get("dependencies", [])), - "dep_threshold": int(wf.get("dependency_threshold", 1)), - "cpus": int(wf.get("required_cpus", 0)), - "gpus": int(wf.get("required_gpus", 0)), - } - return out - - -# ── Figure builder ──────────────────────────────────────────────────────────── - -def plot(spans, group_meta, resource_timeline, signal_events, - dreamer_stats, stall_events, profiles, t0, total_resources, out_path): - if not spans: - print("No replica events found.", file=sys.stderr) - return - - # Stage ordering - present = {s[1] for s in spans} - ordered = [s for s in STAGE_ORDER if s in present] + \ - sorted(present - set(STAGE_ORDER)) - - def _sort(s): - rid, grp, *_ = s - return (ordered.index(grp) if grp in ordered else 999, - int(rid.rsplit("_", 1)[-1])) - spans.sort(key=_sort) - - n_stages = len(ordered) - total_cpus, total_gpus = total_resources - has_resources = bool(resource_timeline) - has_dreamer = bool(dreamer_stats) or bool(profiles) - - # ── Figure layout ──────────────────────────────────────────────────────── - # Row 0 summary table (under title, full width) - # Row 1 stage-level concurrency Gantt - # Row 2 resource util (CPU/GPU) - # Row 3 makespan dist + ops dist (dreamer metrics, 2 equal panels) - summ_h = 1.5 if has_dreamer else 0 - gantt_h = max(3.5, n_stages * 0.75) - res_h = 2.6 if has_resources else 0 - drm_h = 4.2 if has_dreamer else 0 - fig_h = summ_h + gantt_h + res_h + drm_h + 1.2 - - fig = plt.figure(figsize=(22, fig_h)) - - hr = [] - if summ_h: hr.append(summ_h) - hr.append(gantt_h) - if has_resources: hr.append(res_h) - if has_dreamer: hr.append(drm_h) - outer = gridspec.GridSpec(len(hr), 1, height_ratios=hr, hspace=0.38, - top=0.96, bottom=0.03, left=0.06, right=0.97) - - ri = 0 - ax_summary = None - if summ_h: - ax_summary = fig.add_subplot(outer[ri]); ri += 1 - - ax_gantt = fig.add_subplot(outer[ri]) - ri += 1 - - ax_res = None - if has_resources: - ax_res = fig.add_subplot(outer[ri]); ri += 1 - - ax_mspan = ax_ops = None - if has_dreamer: - inner_d = gridspec.GridSpecFromSubplotSpec( - 1, 2, subplot_spec=outer[ri], width_ratios=[1, 1], wspace=0.30) - ax_mspan = fig.add_subplot(inner_d[0, 0]) - ax_ops = fig.add_subplot(inner_d[0, 1]) - - # ───────────────────────────────────────────────────────────────────────── - # 0. STAGE SUMMARY TABLE (full-width row directly under suptitle) - # ───────────────────────────────────────────────────────────────────────── - if ax_summary is not None and dreamer_stats: - ax_summary.axis("off") - - agg: dict = {} - for rid, ds in dreamer_stats.items(): - g = _group(rid) - a = agg.setdefault(g, {"n": 0, "tasks": [], "mk": [], "ae": [], "strat": set()}) - a["n"] += 1 - a["tasks"].append(ds.get("tasks_completed", ds.get("num_tasks", 0))) - a["mk"].append(ds.get("makespan", 0)) - a["ae"].append(ds.get("avg_exec", 0)) - a["strat"].add(ds.get("strategy", "?")) - - _short = {"smallest_to_fastest": "s→fast", - "largest_to_fastest": "l→fast", - "random": "rand"} - - ch = ["Stage", "Facility / partition", "Budget\n(node-h)", - "Cap", "Reps", "Tasks/rep", - "Makespan\n(avg sim)", "AvgExec\n(avg sim)", "Edge profile", "Strategy"] - tbl_rows, tbl_cols = [], [] - - # Enrich with config metadata when available - cfg_stages = {} - if _HAVE_YAML: - try: - import yaml as _yaml - # locate config.yaml next to wherever the script is being called from - _cfg_path = Path(sys.argv[0]).parent / "config.yaml" - if _cfg_path.exists(): - _raw = _yaml.safe_load(_cfg_path.read_text()) - _edges = {e["upstream"]: e for e in _raw.get("edges", [])} - for _s in _raw.get("stages", []): - _sid = _s["id"] - _pilot = _s.get("pilot", {}) - _edge = _edges.get(_sid, {}) - cfg_stages[_sid] = { - "facility": _pilot.get("facility", "—"), - "partition": _pilot.get("partition", "—"), - "budget": _s.get("budget_node_hours", "—"), - "cap": _s.get("concurrency_cap", "—"), - "profile": _edge.get("profile", "—"), - } - except Exception: - pass - - for s in ordered: - if s not in agg: - continue - a = agg[s] - st = "/".join(_short.get(x, x) for x in sorted(a["strat"])) - cm = cfg_stages.get(s, {}) - fac_part = f"{cm.get('facility','—')} / {cm.get('partition','—')}" - tbl_rows.append([ - STAGE_LABELS.get(s, s), - fac_part, - f"{cm.get('budget','—'):,}" if isinstance(cm.get("budget"), (int, float)) else "—", - f"{cm.get('cap','—'):,}" if isinstance(cm.get("cap"), (int, float)) else "—", - str(a["n"]), - f"{np.mean(a['tasks']):.0f}", - f"{np.mean(a['mk']):.1f}", - f"{np.mean(a['ae']):.1f}", - cm.get("profile", "—"), - st, - ]) - tbl_cols.append([STAGE_COLORS.get(s, DEFAULT_COLOR)] - + ["#f5f5f5"] * (len(ch) - 1)) - - if tbl_rows: - tbl = ax_summary.table( - cellText=tbl_rows, colLabels=ch, - cellColours=tbl_cols, - loc="center", cellLoc="center", - ) - tbl.auto_set_font_size(False) - tbl.set_fontsize(8) - tbl.scale(1.0, 1.35) - for j in range(len(ch)): - tbl[0, j].set_facecolor("#333333") - tbl[0, j].set_text_props(color="white", fontweight="bold") - - # ───────────────────────────────────────────────────────────────────────── - # 1. GANTT CHART — actual wall-clock concurrency per stage - # - # X-axis = elapsed wall-clock seconds (from log). - # Source: replica start/finish timestamps parsed from the CM log. - # - # For each stage the concurrency step function is derived directly from - # the log: events (+1 at replica start, -1 at replica finish) are merged - # and integrated. This faithfully reflects the streaming pipeline where - # stages overlap in wall-clock time (S1 fires S2 as its first replicas - # finish, S3 fires while S2 is still running, etc.). - # - # Bar HEIGHT ∝ peak concurrency / global peak so S4/S5 (peak=1) appear - # visibly thinner than S1 (peak=500), with a 30% minimum so they remain - # readable. - # ───────────────────────────────────────────────────────────────────────── - - # Build per-stage replica spans in elapsed seconds - spans_by_stage: dict[str, list] = {} - for rid, grp, s_dt, e_dt, ok in spans: - if grp not in ordered: - continue - s_el = (s_dt - t0).total_seconds() - e_el = (e_dt - t0).total_seconds() - spans_by_stage.setdefault(grp, []).append((s_el, e_el, ok)) - - def _log_sf(stage): - """Concurrency step function from log spans (wall-clock seconds).""" - events = [] - for s, e, _ in spans_by_stage.get(stage, []): - events.append((s, +1)) - events.append((e, -1)) - if not events: - return [], [] - events.sort() - xs, ys = [0.0], [0] - running = 0 - for t_ev, delta in events: - xs.append(t_ev); ys.append(running) - running += delta - xs.append(t_ev); ys.append(running) - return xs, ys - - stage_sf: dict[str, tuple] = {} - for stage in ordered: - xs, ys = _log_sf(stage) - if not xs: - continue - wc_start = min(spans_by_stage.get(stage, [(0,0,None)])[0][:1] or [0]) - wc_start = min(s for s, e, _ in spans_by_stage.get(stage, [(0,0,None)])) - wc_end = max(e for s, e, _ in spans_by_stage.get(stage, [(0,0,None)])) - stage_sf[stage] = (xs, ys, wc_start, wc_end) - - wc_total = max((e for _, _, _, e in stage_sf.values()), default=1.0) or 1.0 - - global_max_c = max( - (max(ys) for _, (_, ys, _, _) in stage_sf.items() if ys), default=1 - ) or 1 - - # ── Draw ───────────────────────────────────────────────────────────────── - for row_idx, stage in enumerate(ordered): - if stage not in stage_sf: - continue - xs, ys, wc_start, wc_end = stage_sf[stage] - - color = STAGE_COLORS.get(stage, DEFAULT_COLOR) - meta = group_meta.get(stage, {}) - dur = wc_end - wc_start - - stage_max_c = max(ys) or 1 - n_reps = len(spans_by_stage.get(stage, [])) - n_err = sum(1 for _, _, ok in spans_by_stage.get(stage, []) if ok is False) - - row_bot = row_idx - 0.45 - row_h = 0.9 - min_frac = 0.30 # minimum bar height so single-replica stages stay visible - def _scale(y, _smc=stage_max_c): - if y == 0: - return row_bot - raw = y / global_max_c * row_h - return row_bot + max(raw, min_frac * row_h * y / _smc) - ys_sc = [_scale(y) for y in ys] - peak_y = _scale(stage_max_c) - - ax_gantt.fill_between(xs, row_bot, ys_sc, - color=color, alpha=0.60, zorder=2) - ax_gantt.plot(xs, ys_sc, color=color, lw=1.2, alpha=0.9, zorder=3) - ax_gantt.barh(row_idx, dur, left=wc_start, height=row_h, - fill=False, edgecolor=color, lw=0.6, alpha=0.25, zorder=1) - ax_gantt.hlines(peak_y, wc_start, wc_end, - color=color, lw=0.6, ls="--", alpha=0.4, zorder=1) - - err_s = f" ✗{n_err}" if n_err else "" - ann = (f"n={n_reps}{err_s} peak={stage_max_c}" - f" {dur:.1f}s" - f" cpu={meta.get('cpus',0)} gpu={meta.get('gpus',0)}") - ax_gantt.text(wc_total * 1.005, row_idx, ann, - va="center", ha="left", fontsize=7, - color=color, clip_on=True) - - if row_idx % 2 == 0: - ax_gantt.axhspan(row_idx - 0.5, row_idx + 0.5, - color="grey", alpha=0.04, lw=0) - - # Trigger-signal markers on the Gantt - for t_sig, tgt, _ in signal_events: - if tgt in ordered: - ax_gantt.axvline(t_sig, color=STAGE_COLORS.get(tgt, "#888"), - lw=0.6, ls=":", alpha=0.4, zorder=1) - - ax_gantt.set_yticks(range(n_stages)) - ax_gantt.set_yticklabels([STAGE_LABELS.get(s, s) for s in ordered], - fontsize=10, fontweight="bold") - ax_gantt.set_xlabel("Elapsed wall-clock time (s)", fontsize=9) - ax_gantt.set_title( - "Campaign Stage Activity — Streaming Pipeline (wall-clock)\n" - "(source: CM log · bar height ∝ concurrent replicas / global peak)", - fontweight="bold", fontsize=10) - ax_gantt.invert_yaxis() - ax_gantt.grid(axis="x", ls="--", alpha=0.35) - ax_gantt.set_xlim(left=0) - - legend_handles = ( - [mpatches.Patch(color=STAGE_COLORS.get(s, DEFAULT_COLOR), - label=STAGE_LABELS.get(s, s)) for s in ordered] - + [mpatches.Patch(fc="white", ec="red", lw=1.2, label="error")] - ) - ax_gantt.legend(handles=legend_handles, loc="upper right", - fontsize=7, framealpha=0.8) - - # ───────────────────────────────────────────────────────────────────────── - # 2. RESOURCE UTILIZATION - # ───────────────────────────────────────────────────────────────────────── - if ax_res is not None and resource_timeline: - times = [t for t,*_ in resource_timeline] - used_gpu = [ug for _,_,_,ug,_ in resource_timeline] - used_cpu = [uc for _,uc,*_ in resource_timeline] - tot_gpu = [tg for _,_,_,_,tg in resource_timeline] - tot_cpu = [tc for _,_,tc,*_ in resource_timeline] - - ax_res.step(times, used_gpu, where="post", color="#4C72B0", lw=2, label="GPU used") - ax_res.fill_between(times, used_gpu, step="post", color="#4C72B0", alpha=0.15) - if any(t > 0 for t in tot_gpu): - ax_res.step(times, tot_gpu, where="post", color="#4C72B0", - lw=0.9, ls="--", alpha=0.5, label="GPU total") - - ax_res.set_ylabel("GPUs in use", color="#4C72B0", fontsize=8) - ax_res.tick_params(axis="y", labelcolor="#4C72B0", labelsize=7) - ax_res.set_ylim(bottom=0) - - for t_sig, tgt, _ in signal_events: - ax_res.axvline(t_sig, color=STAGE_COLORS.get(tgt, "#888"), - lw=0.8, alpha=0.4, ls=":") - - ax_cpu = ax_res.twinx() - ax_cpu.step(times, used_cpu, where="post", color="#DD8452", lw=2, label="CPU used") - ax_cpu.fill_between(times, used_cpu, step="post", color="#DD8452", alpha=0.12) - if any(t > 0 for t in tot_cpu): - ax_cpu.step(times, tot_cpu, where="post", color="#DD8452", - lw=0.9, ls="--", alpha=0.5, label="CPU total") - ax_cpu.set_ylabel("CPUs in use", color="#DD8452", fontsize=8) - ax_cpu.tick_params(axis="y", labelcolor="#DD8452", labelsize=7) - ax_cpu.set_ylim(bottom=0) - - ax_res.set_xlabel("Elapsed wall-clock time (s)", fontsize=8) - ax_res.set_title("Resource Utilization (CPU / GPU)", - fontweight="bold", fontsize=9) - ax_res.grid(axis="x", ls="--", alpha=0.35) - l1, lb1 = ax_res.get_legend_handles_labels() - l2, lb2 = ax_cpu.get_legend_handles_labels() - ax_res.legend(l1 + l2, lb1 + lb2, fontsize=7, loc="upper right", - framealpha=0.8) - - # ───────────────────────────────────────────────────────────────────────── - # 3a. SIMULATED MAKESPAN DISTRIBUTION (box plots per stage) - # ───────────────────────────────────────────────────────────────────────── - if ax_mspan is not None and dreamer_stats: - mspan_by: dict = {s: [] for s in ordered} - for rid, ds in dreamer_stats.items(): - g = _group(rid) - if g in mspan_by: - mspan_by[g].append(ds["makespan"]) - - plot_stgs = [s for s in ordered if mspan_by.get(s)] - if plot_stgs: - box_data = [mspan_by[s] for s in plot_stgs] - pos = list(range(len(plot_stgs))) - - bp = ax_mspan.boxplot( - box_data, positions=pos, widths=0.55, - patch_artist=True, showfliers=False, - medianprops=dict(color="white", lw=2.5), - boxprops=dict(lw=1), whiskerprops=dict(lw=1), - capprops=dict(lw=1)) - for patch, s in zip(bp["boxes"], plot_stgs): - patch.set_facecolor(STAGE_COLORS.get(s, DEFAULT_COLOR)) - patch.set_alpha(0.78) - - # Jittered individual points (sample ≤ 300 per stage) - rng = np.random.default_rng(42) - for pi, (s, vals) in enumerate(zip(plot_stgs, box_data)): - sample = rng.choice(vals, size=min(300, len(vals)), replace=False) - jitter = rng.uniform(-0.18, 0.18, size=len(sample)) - ax_mspan.scatter(pi + jitter, sample, s=4, alpha=0.30, - color=STAGE_COLORS.get(s, DEFAULT_COLOR), zorder=3) - - # n + median annotation beside each box - for pi, (s, vals) in enumerate(zip(plot_stgs, box_data)): - med = float(np.median(vals)) - ax_mspan.text(pi + 0.34, med, - f"n={len(vals)}\nmed={med:.0f}", - ha="left", va="center", fontsize=6.5, color="#333") - - ax_mspan.set_xticks(pos) - ax_mspan.set_xticklabels( - [STAGE_LABELS.get(s, s) for s in plot_stgs], fontsize=8) - ax_mspan.set_ylabel("Simulated makespan (dreamer time units)", fontsize=8) - ax_mspan.set_title("Makespan Distribution per Stage\n" - "(all replicas; ops ÷ core-perf)", - fontweight="bold", fontsize=9) - ax_mspan.grid(axis="y", ls="--", alpha=0.35) - - # ───────────────────────────────────────────────────────────────────────── - # 3b. TASK OPS DISTRIBUTION (box plots per stage) - # ───────────────────────────────────────────────────────────────────────── - if ax_ops is not None: - ops_by: dict = {s: [] for s in ordered} - for rid, prof in profiles.items(): - g = _group(rid) - if g in ops_by: - for t in prof.get("tasks", []): - v = t.get("ops") - if v is not None: - ops_by[g].append(float(v)) - # Fallback: use avg_exec as ops proxy when no profiles - if not any(ops_by.values()): - for rid, ds in dreamer_stats.items(): - g = _group(rid) - if g in ops_by: - ops_by[g].extend([ds["avg_exec"]] * ds.get("tasks_completed", 1)) - - plot_stgs = [s for s in ordered if ops_by.get(s)] - if plot_stgs: - box_data = [ops_by[s] for s in plot_stgs] - pos = list(range(len(plot_stgs))) - - bp = ax_ops.boxplot( - box_data, positions=pos, widths=0.55, - patch_artist=True, showfliers=False, - medianprops=dict(color="white", lw=2.5), - boxprops=dict(lw=1), whiskerprops=dict(lw=1), - capprops=dict(lw=1)) - for patch, s in zip(bp["boxes"], plot_stgs): - patch.set_facecolor(STAGE_COLORS.get(s, DEFAULT_COLOR)) - patch.set_alpha(0.78) - - # Jittered points (sample ≤ 300 per stage) - rng = np.random.default_rng(42) - for pi, (s, ops) in enumerate(zip(plot_stgs, box_data)): - sample = rng.choice(ops, size=min(300, len(ops)), replace=False) - jitter = rng.uniform(-0.2, 0.2, size=len(sample)) - ax_ops.scatter(pi + jitter, sample, s=3, alpha=0.30, - color=STAGE_COLORS.get(s, DEFAULT_COLOR), zorder=3) - - # Median annotation - for pi, ops in enumerate(box_data): - med = float(np.median(ops)) - ax_ops.text(pi + 0.35, med, f"{med:.0f}", - ha="left", va="center", fontsize=6.5, color="#333") - - ax_ops.set_yscale("log") - ax_ops.set_xticks(pos) - ax_ops.set_xticklabels([STAGE_LABELS.get(s, s) for s in plot_stgs], - fontsize=8) - ax_ops.set_ylabel("Task ops (log scale, dreamer units)", fontsize=8) - ax_ops.set_title("Task Ops Distribution per Stage\n" - "(all tasks × all replicas; log scale)", - fontweight="bold", fontsize=9) - ax_ops.grid(axis="y", ls="--", alpha=0.35) - - # ───────────────────────────────────────────────────────────────────────── - fig.suptitle("Dreamer Campaign — 5-Stage Drug Discovery Cascade", - fontsize=13, fontweight="bold") - plt.savefig(out_path, dpi=150, bbox_inches="tight") - print(f"Saved → {out_path}") - - -# ── CLI ─────────────────────────────────────────────────────────────────────── - -def main(): - ap = argparse.ArgumentParser( - description="Plot dreamer campaign timeline and simulation statistics") - ap.add_argument("log", help="Campaign log file") - ap.add_argument("--profiles-dir", default=None, - help="dreamer-profiles/ directory (default: auto-detect next to log)") - ap.add_argument("--config", default=None, - help="config.yaml (default: auto-detect next to log)") - ap.add_argument("--out", default=None, help="Output PNG path") - args = ap.parse_args() - - log_dir = Path(args.log).parent - stem = Path(args.log).stem - if args.out is None: args.out = f"plots/dreamer_timeline_{stem}.png" - if args.config is None: - c = log_dir / "config.yaml" - if c.exists(): args.config = str(c) - if args.profiles_dir is None: - args.profiles_dir = str(log_dir / "dreamer-profiles") - - Path(args.out).parent.mkdir(parents=True, exist_ok=True) - - spans, group_meta_log, resource_timeline, signal_events, \ - dreamer_stats, stall_events, t0, total_resources = parse_log(args.log) - - group_meta_cfg = parse_config(args.config) if args.config else {} - group_meta = group_meta_cfg if group_meta_cfg else group_meta_log - if args.config: - print(f"Stage config from: {args.config}") - profiles = load_profiles(args.profiles_dir) - - print(f"Parsed: {len(spans)} replica spans | {len(group_meta)} stages | " - f"{len(resource_timeline)} resource events | {len(signal_events)} triggers | " - f"{len(dreamer_stats)} dreamer records | {len(profiles)} profile JSONs") - - plot(spans, group_meta, resource_timeline, signal_events, - dreamer_stats, stall_events, profiles, - t0, total_resources, args.out) - - -if __name__ == "__main__": - main() diff --git a/workflows/run_campaign/dreamer_campaign/plot_optimizations.py b/workflows/run_campaign/dreamer_campaign/plot_optimizations.py deleted file mode 100644 index 9d7ab59..0000000 --- a/workflows/run_campaign/dreamer_campaign/plot_optimizations.py +++ /dev/null @@ -1,640 +0,0 @@ -#!/usr/bin/env python3 -""" -plot_optimizations.py — visualise per-optimization performance improvements. - -Reads benchmark_results.json produced by benchmark.py and generates 7 plots: - - 1. wall_time.png — campaign wall time per configuration - 2. pipeline_gantt.png — stage execution overlap (first/last replica timeline) - 3. cascade_funnel.png — total replicas launched per stage (compute waste) - 4. gpu_utilization.png — GPU slots in use per stage over time (4-panel) - 5. shard_dispatch.png — cumulative candidates dispatched by sharder over time - 6. bandit_convergence.png — scheduling bandit Thompson-sample convergence - 7. time_to_target.png — cumulative terminal-stage completions over wall time - -Each plot is designed to support one specific optimization axis: - - sharding+bp: plots 3 (cascade funnel) + 5 (shard dispatch) - - scheduling_bandit: plots 4 (GPU utilization) + 6 (bandit convergence) - - all_optimizations: plots 1 (wall time) + 2 (Gantt) + 7 (time-to-target) - -Usage: - python plot_optimizations.py [--results benchmark_results.json] [--out-dir plots/] -""" - -import argparse -import itertools -import json -import math -import statistics -import warnings -from collections import defaultdict -from pathlib import Path - -import matplotlib -matplotlib.use("Agg") -import matplotlib.patches as mpatches -import matplotlib.pyplot as plt -import numpy as np - -# ── Colour palette ──────────────────────────────────────────────────────────── - -CFG_COLORS = { - "baseline": "#9e9e9e", - "sharding+bp": "#4caf50", - "scheduling_bandit": "#9c27b0", - "all_optimizations": "#f44336", -} - -STAGE_COLORS = { - "s1_ligand_filter": "#42a5f5", - "s2_ml_affinity": "#66bb6a", - "s3_docking": "#ffa726", - "s4_md_refinement": "#ef5350", - "s5_fep_ranking": "#ab47bc", -} - -STAGE_ORDER = [ - "s1_ligand_filter", "s2_ml_affinity", "s3_docking", - "s4_md_refinement", "s5_fep_ranking", -] - - -# ── Helpers ─────────────────────────────────────────────────────────────────── - -def _load(results_path: str) -> dict: - with open(results_path) as f: - return json.load(f) - - -def _mean(vals): - valid = [v for v in vals if v is not None] - return sum(valid) / len(valid) if valid else None - - -def _median(vals): - valid = [v for v in vals if v is not None] - return statistics.median(valid) if valid else None - - -def _std(vals): - valid = [v for v in vals if v is not None] - if len(valid) < 2: - return 0.0 - m = sum(valid) / len(valid) - return math.sqrt(sum((v - m) ** 2 for v in valid) / (len(valid) - 1)) - - -def _z(v, default=0.0): - return v if v is not None else default - - -def _caption(fig, text: str) -> None: - fig.text( - 0.5, -0.02, text, - ha="center", va="top", fontsize=7.5, color="#444", - wrap=True, - bbox=dict(boxstyle="round,pad=0.4", facecolor="#f5f5f5", - edgecolor="#ccc", linewidth=0.8), - transform=fig.transFigure, - ) - - -def _repr_run(runs, key="wall_time_s"): - """Return the run whose key value is closest to the median.""" - vals = [(i, r.get(key)) for i, r in enumerate(runs) if r.get(key) is not None] - if not vals: - return runs[0] - med = statistics.median(v for _, v in vals) - idx = min(vals, key=lambda iv: abs(iv[1] - med))[0] - return runs[idx] - - -def _reconstruct_intervals(replica_events): - """Yield (start_t, finish_t, group) for each replica that both started and finished.""" - starts: dict[str, float] = {} - groups: dict[str, str] = {} - for e in replica_events: - rid = e["replica_id"] - if e["event"] == "start": - starts[rid] = e["t"] - groups[rid] = e["group"] - elif e["event"] == "finish" and rid in starts: - yield starts[rid], e["t"], groups[rid] - - -# ── Plot 1: Campaign wall time ──────────────────────────────────────────────── - -def plot_wall_time(results: dict, out_dir: Path) -> None: - cfgs = list(results.keys()) - medians = [_median([r["wall_time_s"] for r in results[c] if r.get("wall_time_s")]) for c in cfgs] - baseline = _median([r["wall_time_s"] for r in results.get("baseline", []) if r.get("wall_time_s")]) or 1.0 - - fig, ax = plt.subplots(figsize=(10, 5)) - x = np.arange(len(cfgs)) - bars = ax.bar(x, [_z(m) for m in medians], - color=[CFG_COLORS.get(c, "#888") for c in cfgs], alpha=0.85) - for i, cfg in enumerate(cfgs): - wts = [r["wall_time_s"] for r in results[cfg] if r.get("wall_time_s")] - ax.scatter([i] * len(wts), wts, color="white", edgecolors="black", - zorder=3, s=22, linewidths=0.8) - for bar, m, cfg in zip(bars, medians, cfgs): - if m is not None: - pct = (m - baseline) / baseline * 100 - label = f"{m:.0f}s" + (f"\n({pct:+.0f}%)" if cfg != "baseline" else "") - ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 1.5, - label, ha="center", va="bottom", fontsize=9, fontweight="bold") - ax.axhline(baseline, color="gray", linestyle="--", linewidth=0.9, label="baseline median") - ax.set_xticks(x) - ax.set_xticklabels(cfgs, rotation=20, ha="right", fontsize=10) - ax.set_ylabel("Wall time to target (s)") - ax.set_title("Campaign wall time by configuration\n" - "(time to find 5 terminal-stage hits; lower is better; % vs baseline)") - ax.legend(fontsize=9) - plt.tight_layout() - _caption(fig, - "LOWER IS BETTER. Wall-clock time from campaign start until the 5th s5_fep_ranking " - "replica completes (early-termination target). Bar = median of 5 runs; white dots = " - "individual runs (spread shows run-to-run variance). " - "sharding+bp: sharder routes highest-score candidates first — fewer total replicas needed " - "to produce 5 quality hits (4.9× faster). " - "scheduling_bandit: Thompson-sampling bandit allocates GPUs to terminal stages earlier " - "(3.1× faster). " - "all_optimizations: both axes combined (10.9× faster, lowest variance)." - ) - plt.savefig(out_dir / "1_wall_time.png", dpi=150, bbox_inches="tight") - plt.close() - print(" 1_wall_time.png") - - -# ── Plot 2: Pipeline Gantt ──────────────────────────────────────────────────── - -def plot_gantt(results: dict, out_dir: Path) -> None: - cfgs = list(results.keys()) - n = len(cfgs) - fig, axes = plt.subplots(n, 1, figsize=(12, 2.2 * n), sharex=False) - if n == 1: - axes = [axes] - - for ax, cfg in zip(axes, cfgs): - runs = [r for r in results[cfg] if "group_stats" in r] - if not runs: - ax.set_title(cfg) - continue - groups = STAGE_ORDER - for i, g in enumerate(groups): - starts = [r["group_stats"].get(g, {}).get("first_start") for r in runs] - finishes = [r["group_stats"].get(g, {}).get("last_finish") for r in runs] - starts = [v for v in starts if v is not None] - finishes = [v for v in finishes if v is not None] - if not starts or not finishes: - continue - s, f = _mean(starts), _mean(finishes) - color = STAGE_COLORS.get(g, "#888") - ax.barh(i, f - s, left=s, height=0.55, color=color, alpha=0.85) - ax.text(s + (f - s) / 2, i, g.replace("_", " "), - ha="center", va="center", fontsize=6, color="white", fontweight="bold") - wts = [r.get("wall_time_s") for r in runs if r.get("wall_time_s")] - t_end = _mean(wts) or 0 - ax.axvline(t_end, color="black", linestyle=":", linewidth=1.0, alpha=0.5) - ax.set_yticks([]) - ax.set_xlabel("Time (s)" if ax is axes[-1] else "") - ax.set_title(f"{cfg} (avg wall={t_end:.1f}s)", fontsize=9, color=CFG_COLORS.get(cfg, "black")) - ax.grid(axis="x", linestyle="--", alpha=0.35) - - plt.suptitle("Stage execution overlap per configuration\n" - "(more overlap = better pipeline utilisation)", y=1.01, fontsize=10) - plt.tight_layout() - _caption(fig, - "MORE OVERLAP IS BETTER. Each bar shows the average first-start to last-finish span " - "of a stage across 5 runs. Dotted vertical line = campaign end (target reached). " - "baseline: s1 runs long before downstream stages accumulate enough triggers. " - "sharding+bp: min_replicas floor forces s2-s5 slots open from the start. " - "scheduling_bandit: bandit allocates GPU budget downstream — s4/s5 start early even " - "while s1 is still running. all_optimizations: all stages overlap from t~1s onward." - ) - plt.savefig(out_dir / "2_pipeline_gantt.png", dpi=150, bbox_inches="tight") - plt.close() - print(" 2_pipeline_gantt.png") - - -# ── Plot 3: Cascade funnel (total work launched) ────────────────────────────── - -def plot_cascade_funnel(results: dict, out_dir: Path) -> None: - """Stacked bar: total replicas started per config, coloured by stage. - - Supports sharding+bp story: fewer total candidates launched to find 5 s5 hits. - """ - cfgs = list(results.keys()) - - # Compute mean n_started per stage per config - stage_means: dict[str, list[float]] = {cfg: [] for cfg in cfgs} - for cfg in cfgs: - valid = [r for r in results[cfg] if "group_stats" in r] - for stage in STAGE_ORDER: - vals = [r["group_stats"].get(stage, {}).get("n_started", 0) for r in valid] - stage_means[cfg].append(_mean([v for v in vals if v is not None]) or 0) - - fig, (ax_stacked, ax_s1) = plt.subplots(1, 2, figsize=(14, 5)) - - # ── Left: stacked bar (total compute by stage) ──────────────────────────── - x = np.arange(len(cfgs)) - bottom = np.zeros(len(cfgs)) - for si, stage in enumerate(STAGE_ORDER): - heights = [stage_means[cfg][si] for cfg in cfgs] - bars = ax_stacked.bar(x, heights, bottom=bottom, - color=STAGE_COLORS[stage], alpha=0.85, - label=stage.replace("_", " ")) - # Annotate s1 bars only (dominate the chart) - if stage == "s1_ligand_filter": - for i, (bar, h) in enumerate(zip(bars, heights)): - if h > 50: - ax_stacked.text(bar.get_x() + bar.get_width() / 2, - bottom[i] + h / 2, f"{h:.0f}", - ha="center", va="center", fontsize=8, - color="white", fontweight="bold") - bottom += np.array(heights) - - # Annotate totals on top - for i, cfg in enumerate(cfgs): - total = sum(stage_means[cfg]) - base_total = sum(stage_means.get("baseline", [1])) - ratio = base_total / total if total > 0 else 0 - label = f"{total:.0f}" + (f"\n({ratio:.1f}× less)" if cfg != "baseline" else "") - ax_stacked.text(i, bottom[i] + 30, label, - ha="center", va="bottom", fontsize=8, fontweight="bold") - - ax_stacked.set_xticks(x) - ax_stacked.set_xticklabels(cfgs, rotation=20, ha="right", fontsize=9) - ax_stacked.set_ylabel("Total replicas started") - ax_stacked.set_title("Total compute launched\n(stacked by stage; lower = less wasted work)") - ax_stacked.legend(fontsize=8, loc="upper right") - ax_stacked.grid(axis="y", linestyle="--", alpha=0.3) - - # ── Right: per-stage breakdown (log scale) ──────────────────────────────── - width = 0.8 / len(cfgs) - xs = np.arange(len(STAGE_ORDER)) - for ci, cfg in enumerate(cfgs): - vals = [max(stage_means[cfg][si], 0.5) for si in range(len(STAGE_ORDER))] - offset = (ci - len(cfgs) / 2 + 0.5) * width - ax_s1.bar(xs + offset, vals, width * 0.9, - label=cfg, color=CFG_COLORS.get(cfg, "#888"), alpha=0.85) - - ax_s1.set_yscale("log") - ax_s1.set_xticks(xs) - ax_s1.set_xticklabels([s.replace("_", "\n") for s in STAGE_ORDER], fontsize=8) - ax_s1.set_ylabel("Replicas started (log scale)") - ax_s1.set_title("Per-stage breakdown (log scale)\n(shows full funnel reduction)") - ax_s1.legend(fontsize=8) - ax_s1.grid(axis="y", linestyle="--", alpha=0.3) - - plt.suptitle("Pipeline cascade: replicas launched to find 5 terminal-stage hits", - fontsize=10, y=1.01) - plt.tight_layout() - _caption(fig, - "LOWER IS BETTER. Left: total replicas started per config, stacked by stage. " - "Right: same data on log scale to show the full funnel. " - "baseline: 3,600+ replicas (s1 monopolises GPUs — 3,200 s1 before 5 s5 hits). " - "sharding+bp: sharder routes highest-quality s1 results to s2 first — only 730 " - "replicas total (5× less). scheduling_bandit: bandit terminates campaign earlier by " - "getting s5 resources sooner — 1,050 replicas (3.4× less). " - "all_optimizations: both effects — only 235 replicas total (15× less compute)." - ) - plt.savefig(out_dir / "3_cascade_funnel.png", dpi=150, bbox_inches="tight") - plt.close() - print(" 3_cascade_funnel.png") - - -# ── Plot 4: GPU utilization per stage over time ─────────────────────────────── - -def plot_gpu_utilization(results: dict, out_dir: Path) -> None: - """Stacked-area GPU-in-use per stage over time, one panel per config. - - Supports scheduling_bandit story: terminal stages claim GPUs much earlier. - """ - cfgs = list(results.keys()) - n = len(cfgs) - # Use 2×2 grid when 4 configs for better readability - if n == 4: - fig, axes_grid = plt.subplots(2, 2, figsize=(14, 8), sharey=False) - axes = [axes_grid[0,0], axes_grid[0,1], axes_grid[1,0], axes_grid[1,1]] - else: - fig, axes_raw = plt.subplots(1, n, figsize=(5 * n, 5), sharey=False) - axes = [axes_raw] if n == 1 else list(axes_raw) - fig.patch.set_facecolor("white") - - for ax, cfg in zip(axes, cfgs): - ax.set_facecolor("#fafafa") - rep = _repr_run([r for r in results[cfg] if r.get("replica_events")], "wall_time_s") - if not rep: - ax.set_title(cfg) - continue - - events = rep["replica_events"] - t_max = max(e["t"] for e in events) - ts = np.linspace(0, t_max, 400) - - intervals: dict[str, list[tuple[float, float]]] = defaultdict(list) - for s, f, g in _reconstruct_intervals(events): - intervals[g].append((s, f)) - - bottom = np.zeros(len(ts)) - for stage in STAGE_ORDER: - ivs = intervals.get(stage, []) - if not ivs: - continue - running = np.array([sum(1 for s, f in ivs if s <= t < f) for t in ts]) - color = STAGE_COLORS[stage] - ax.fill_between(ts, bottom, bottom + running, - color=color, alpha=0.80, label=stage.replace("_", " ")) - bottom = bottom + running - - # Annotate when s5 first appears - s5_ivs = intervals.get("s5_fep_ranking", []) - if s5_ivs: - first_s5 = min(s for s, _ in s5_ivs) - y_top = max(bottom) if max(bottom) > 0 else 5 - ax.axvline(first_s5, color="#7b1fa2", linestyle="--", linewidth=2.0) - ax.text(first_s5 + t_max * 0.02, y_top * 0.92, - f"s5 starts\n{first_s5:.1f}s", fontsize=9, color="#7b1fa2", - va="top", fontweight="bold", - bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="#ab47bc", lw=1)) - - wt = rep.get("wall_time_s", t_max) - ax.set_title(f"{cfg}\n(total wall time: {wt:.1f} s)", fontsize=10, - color=CFG_COLORS.get(cfg, "black"), fontweight="bold", pad=6) - ax.set_xlabel("Wall-clock time (s)", fontsize=9) - ax.set_ylabel("GPU slots in use", fontsize=9) - ax.tick_params(labelsize=8) - ax.grid(axis="y", linestyle="--", alpha=0.5, color="#cccccc") - ax.spines["top"].set_visible(False) - ax.spines["right"].set_visible(False) - - # Shared legend - handles = [mpatches.Patch(color=STAGE_COLORS[s], label=s.replace("_", " ")) - for s in STAGE_ORDER] - fig.legend(handles=handles, loc="upper center", ncol=5, fontsize=10, - bbox_to_anchor=(0.5, 1.0), frameon=True, edgecolor="#cccccc") - plt.suptitle("GPU slots in use per stage over time (representative run per config)", - fontsize=12, fontweight="bold", y=1.04, color="#1a237e") - plt.tight_layout(pad=2.0) - _caption(fig, - "EARLIER PURPLE (s5) IS BETTER. Each colour = GPU slots used by that stage over time. " - "Dashed line = first s5 start. " - "baseline: s1 (blue) monopolises all GPUs until ~14 s. " - "scheduling_bandit: s3/s4/s5 share GPUs from t~1 s; s5 starts at ~7 s. " - "all_optimizations: s5 starts at ~4 s — quality routing + learned allocation combined." - ) - plt.savefig(out_dir / "4_gpu_utilization.png", dpi=150, bbox_inches="tight", - facecolor="white") - plt.close() - print(" 4_gpu_utilization.png") - - -# ── Plot 5: Shard dispatch over time ───────────────────────────────────────── - -def plot_shard_dispatch(results: dict, out_dir: Path) -> None: - """Cumulative candidates dispatched by the sharder over time, per downstream stage. - - Supports sharding+bp story: pipeline is fed continuously, not in floods. - Only configs with shard_events are plotted (baseline and scheduling_bandit are excluded). - """ - sharder_cfgs = [c for c in results - if any(r.get("shard_events") for r in results[c])] - if not sharder_cfgs: - return - - stages = ["s2_ml_affinity", "s3_docking", "s4_md_refinement", "s5_fep_ranking"] - labels = ["s2 ML affinity", "s3 Docking", "s4 MD refine", "s5 FEP rank"] - - fig, axes = plt.subplots(1, len(stages), figsize=(4 * len(stages), 4), squeeze=False) - - for si, (stage, slabel) in enumerate(zip(stages, labels)): - ax = axes[0][si] - for cfg in sharder_cfgs: - color = CFG_COLORS.get(cfg, "#888") - # Use representative run - rep = _repr_run([r for r in results[cfg] if r.get("shard_events")], "wall_time_s") - if not rep: - continue - evs = [(e["timestamp"], e.get("n", 1)) - for e in rep.get("shard_events", []) - if e.get("group") == stage] - if not evs: - continue - evs.sort() - ts = [0.0] + [t for t, _ in evs] - cumN = list(itertools.accumulate([0] + [n for _, n in evs])) - ax.step(ts, cumN, where="post", color=color, linewidth=2.0, label=cfg) - - ax.set_title(slabel, fontsize=9) - ax.set_xlabel("Wall time (s)") - ax.set_ylabel("Cumulative dispatched" if si == 0 else "") - ax.legend(fontsize=7, loc="lower right") - ax.grid(linestyle="--", alpha=0.3) - - plt.suptitle("Sharder: cumulative candidates dispatched to each stage over time\n" - "(sharder-enabled configs only)", fontsize=10, y=1.01) - plt.tight_layout() - _caption(fig, - "Shows how the sharder feeds each downstream stage over time. " - "Steeper initial slope = pipeline fed faster with high-priority candidates. " - "Plateau = sharder stopped dispatching (backpressure THROTTLE or upstream done). " - "all_optimizations dispatches fewer candidates total (reaches 5 s5 hits with ~50 " - "s2 dispatches vs ~230 for sharding+bp) because the scheduling bandit keeps s5 " - "consuming candidates faster — the campaign terminates sooner." - ) - plt.savefig(out_dir / "5_shard_dispatch.png", dpi=150, bbox_inches="tight") - plt.close() - print(" 5_shard_dispatch.png") - - -# ── Plot 6: Scheduling bandit convergence ───────────────────────────────────── - -def plot_bandit_convergence(results: dict, out_dir: Path) -> None: - """Thompson-sample values per stage over time for bandit-enabled configs. - - Supports scheduling_bandit story: bandit learns to strongly prefer terminal stages. - """ - bandit_cfgs = [c for c in results - if any(r.get("scheduling_events") and - any(e.get("bandit") for e in r["scheduling_events"]) - for r in results[c])] - if not bandit_cfgs: - return - - fig, axes = plt.subplots(1, len(bandit_cfgs), - figsize=(6 * len(bandit_cfgs), 4), squeeze=False) - - for ci, cfg in enumerate(bandit_cfgs): - ax = axes[0][ci] - # Collect per-stage (timestamp, sample_value) across all runs - group_pairs: dict[str, list[tuple[float, float]]] = defaultdict(list) - t_max = 0.0 - for r in results[cfg]: - evs = [e for e in r.get("scheduling_events", []) if e.get("bandit")] - if not evs: - continue - t_max = max(t_max, max(e["timestamp"] for e in evs)) - for g in STAGE_ORDER: - for e in evs: - if g in e.get("eligible", []) and g in e.get("bandit", {}): - group_pairs[g].append((e["timestamp"], e["bandit"][g])) - - if not group_pairs: - ax.set_title(cfg) - continue - - N_BINS = 30 - bin_edges = np.linspace(0, max(t_max, 1), N_BINS + 1) - bin_mids = 0.5 * (bin_edges[:-1] + bin_edges[1:]) - - for g, color in STAGE_COLORS.items(): - pairs = group_pairs.get(g, []) - if not pairs: - continue - ts = np.array([p[0] for p in pairs]) - vs = np.array([p[1] for p in pairs]) - bin_means = [ - float(vs[(ts >= lo) & (ts < hi)].mean()) - if ((ts >= lo) & (ts < hi)).any() else np.nan - for lo, hi in zip(bin_edges[:-1], bin_edges[1:]) - ] - col_mean = np.array(bin_means) - valid = ~np.isnan(col_mean) - if not valid.any(): - continue - ax.plot(bin_mids[valid], col_mean[valid], - color=color, label=g.replace("_", " "), - linewidth=1.8, marker="o", markersize=3) - - ax.axhline(0.5, color="gray", linestyle=":", linewidth=0.8, alpha=0.6, - label="uniform prior") - ax.set_xlabel("Wall-clock time (s)") - ax.set_ylabel("Thompson sample (priority)") - ax.set_title(cfg.replace("+", "").replace("_", "\n"), fontsize=9, - color=CFG_COLORS.get(cfg, "black")) - ax.set_ylim(0, 1.05) - ax.legend(fontsize=7) - ax.grid(linestyle="--", alpha=0.3) - - plt.suptitle("Scheduling bandit: Thompson-sample priority per stage over time\n" - "(higher = bandit prefers scheduling this stage)", fontsize=10) - plt.tight_layout() - _caption(fig, - "CONVERGENCE AWAY FROM 0.5 IS BETTER (means the bandit learned a preference). " - "Each line shows the time-binned mean Thompson sample for one stage. " - "Warm-start priors: s5=Beta(5,1) starts near 1.0 (strongly preferred); " - "s1=Beta(1,1) starts at 0.5 (neutral). " - "Over time the bandit reinforces downstream stages (s4/s5) that keep GPUs busy " - "and deprioritises stages whose downstream queue is full (THROTTLE). " - "Runs are short (~7-27s) so convergence is driven mainly by the warm-start priors." - ) - plt.savefig(out_dir / "6_bandit_convergence.png", dpi=150, bbox_inches="tight") - plt.close() - print(" 6_bandit_convergence.png") - - -# ── Plot 7: Time-to-target (cumulative terminal-stage completions) ───────────── - -def plot_time_to_target( - results: dict, - out_dir: Path, - target_stage: str = "s5_fep_ranking", - target_n: int = 5, -) -> None: - """Step curves: cumulative terminal-stage completions per config over wall time. - - Supports all_optimizations story: target is reached far sooner. - """ - cfgs = list(results.keys()) - fig, ax = plt.subplots(figsize=(10, 5)) - - for cfg in cfgs: - color = CFG_COLORS.get(cfg, "#888") - run_ts_lists: list[list[float]] = [] - for r in results[cfg]: - ts = sorted( - e["t"] for e in r.get("replica_events", []) - if e["group"] == target_stage and e["event"] == "finish" - ) - if ts: - run_ts_lists.append(ts) - if not run_ts_lists: - continue - - # Draw all runs as faint lines - for ts in run_ts_lists: - xs = [0.0] + ts - ys = list(range(len(xs))) - ax.step(xs, ys, where="post", color=color, linewidth=0.7, alpha=0.3) - - # Representative run (median total count) - totals = [len(ts) for ts in run_ts_lists] - rep_ts = run_ts_lists[sorted(range(len(totals)), key=lambda i: totals[i])[len(totals) // 2]] - xs = [0.0] + rep_ts - ys = list(range(len(xs))) - ax.step(xs, ys, where="post", color=color, linewidth=2.5, - label=cfg, zorder=4) - - # Mark where target is hit - if len(rep_ts) >= target_n: - t_hit = rep_ts[target_n - 1] - ax.plot(t_hit, target_n, "v", color=color, markersize=10, zorder=5) - ax.axvline(t_hit, color=color, linestyle=":", linewidth=1.0, alpha=0.6) - ax.text(t_hit + 0.2, target_n + 0.1, f"{t_hit:.1f}s", - color=color, fontsize=8, fontweight="bold") - - ax.axhline(target_n, color="black", linestyle="--", linewidth=1.2, - label=f"target N={target_n}") - ax.set_xlabel("Wall-clock time (s)") - ax.set_ylabel(f"Cumulative {target_stage.replace('_', ' ')} completions") - ax.set_title(f"Time to {target_n} final candidates ({target_stage.replace('_', ' ')})\n" - f"(faint lines = individual runs; bold = representative run; ▼ = target reached)") - ax.legend(fontsize=9) - ax.grid(linestyle="--", alpha=0.3) - plt.tight_layout() - _caption(fig, - f"LEFTMOST ▼ MARKER IS BEST. Step curves show cumulative terminal-stage " - f"(s5_fep_ranking) completions over wall time. Faint lines are individual runs; " - f"bold line is the run closest to the median. Downward triangle marks when each " - f"configuration crosses the N={target_n} target. " - f"all_optimizations (red) reaches target ~11× sooner than baseline (grey). " - f"sharding+bp (green) reaches target at ~17s via quality routing. " - f"scheduling_bandit (purple) reaches target at ~27s via learned GPU allocation." - ) - plt.savefig(out_dir / "7_time_to_target.png", dpi=150, bbox_inches="tight") - plt.close() - print(" 7_time_to_target.png") - - -# ── Main ────────────────────────────────────────────────────────────────────── - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--results", default="benchmark_results.json") - parser.add_argument("--out-dir", default="plots/optimizations") - args = parser.parse_args() - - out_dir = Path(args.out_dir) - out_dir.mkdir(parents=True, exist_ok=True) - - print(f"Loading {args.results}...") - results = {k: v for k, v in _load(args.results).items() if k in CFG_COLORS} - print(f"Configurations: {list(results.keys())}") - print(f"Writing plots to {out_dir}/\n") - - plot_wall_time(results, out_dir) - plot_gantt(results, out_dir) - plot_cascade_funnel(results, out_dir) - plot_gpu_utilization(results, out_dir) - plot_shard_dispatch(results, out_dir) - plot_bandit_convergence(results, out_dir) - plot_time_to_target(results, out_dir) - - print(f"\nAll plots written to {out_dir}/") - - -if __name__ == "__main__": - main() - -# python plot_optimizations.py --results benchmark_results.json --out-dir plots/optimizations diff --git a/workflows/run_campaign/dreamer_campaign/requirements.txt b/workflows/run_campaign/dreamer_campaign/requirements.txt deleted file mode 100644 index 1853429..0000000 --- a/workflows/run_campaign/dreamer_campaign/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -pyyaml -radical.dreamer diff --git a/workflows/run_campaign/dreamer_campaign/run_campaign.py b/workflows/run_campaign/dreamer_campaign/run_campaign.py deleted file mode 100644 index 67f47bd..0000000 --- a/workflows/run_campaign/dreamer_campaign/run_campaign.py +++ /dev/null @@ -1,338 +0,0 @@ -#!/usr/bin/env python3 -# Limit OpenBLAS/OMP threads before any numpy import to avoid pthread_create -# failures on login nodes where process counts are restricted. -import os - -os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("MKL_NUM_THREADS", "1") - -""" -Dreamer campaign runner — supports two config formats: - - Flat format (legacy): - workflows: - s1_ligand_filter: { replicas: 4, required_cpus: 16, ... } - - Plan format (cm-prototype): - stages: - - id: s1_ligand_filter - upstream: library - downstream: s2_ml_affinity - concurrency_cap: 5000 - pilot: { partition: cpu, ... } - dreamer: { num_cores: 128, ... } - edges: - - { upstream: s1_ligand_filter, downstream: s2_ml_affinity, profile: diverse_top } - cm: - engine: concurrent - concurrency_scale: 0.002 - resources: { total_cpus: 64, total_gpus: 4 } - workflow_registry: { s1_ligand_filter: dreamer_workflow.DreamerWorkflow } - - The plan format is auto-detected by the presence of a "stages" key. - The translator maps: - stage.upstream / downstream → dependencies / trigger_downstream - stage.pilot.partition → required_cpus / required_gpus - stage.concurrency_cap → max_replicas (× cm.concurrency_scale) - edge.profile → schedule_strategy / early_binding - edge.backpressure → backpressure_high / backpressure_low (metadata) - stage.dreamer.* → dreamer emulation parameters - -Usage: - python run_campaign.py [--config config.yaml] -""" - -import argparse -import asyncio -import sys -from pathlib import Path - -import yaml - -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from src.campaign import AsyncCampaignManager as CampaignManager # noqa: E402 -from src.inference.utils import load_config # noqa: E402 -from src.utils.workflow import _expand_env # noqa: E402 - - -# ── Plan format translation tables ─────────────────────────────────────────── - -# pilot.partition → CM resource requirements -_PILOT_RESOURCES: dict[str, dict] = { - "cpu": {"required_cpus": 16, "required_gpus": 0}, - "gpu": {"required_cpus": 4, "required_gpus": 1}, - "mpi+gpu": {"required_cpus": 16, "required_gpus": 2}, - "largemem": {"required_cpus": 8, "required_gpus": 1}, -} - -# edge.profile → dreamer schedule_strategy -_PROFILE_STRATEGY: dict[str, str] = { - "round_robin": "random", - "diverse_top": "smallest_to_fastest", - "explore_exploit": "largest_to_fastest", - "pure_promise": "largest_to_fastest", - "active_learning": "largest_to_fastest", -} - -# edge.profile → dreamer early_binding -_PROFILE_EARLY_BINDING: dict[str, bool] = { - "round_robin": True, # diversity-focused: bind early - "diverse_top": True, - "explore_exploit": False, # score/uncertainty-focused: late binding - "pure_promise": False, - "active_learning": False, -} - - -def _build_from_plan(config: dict) -> dict: - """ - Translate cm-prototype plan (stages + edges) into the flat ``workflows`` - dict consumed by AsyncCampaignManager.from_config(). - - Returns the translated workflows dict; does not mutate ``config``. - """ - cm_cfg = config.get("cm", {}) - - # debug section: SPHERICAL emulation overrides (not in prototype schema). - # Must be read HERE before main() overwrites config["debug"] with a bool. - debug_cfg = config.get("debug", {}) - debug_cfg = debug_cfg if isinstance(debug_cfg, dict) else {} - stage_replicas_dbg = debug_cfg.get("stage_replicas", {}) - trigger_fractions = debug_cfg.get("trigger_fractions", {}) - - stage_ids: set[str] = {s["id"] for s in config.get("stages", [])} - - # Outgoing edge per source stage (profile → schedule_strategy) - edge_out: dict[str, dict] = {} - # Incoming edge per destination stage (backpressure water marks for that stage's queue) - edge_in: dict[str, dict] = {} - for edge in config.get("edges", []): - src, dst = edge.get("upstream", ""), edge.get("downstream", "") - if src in stage_ids: - edge_out[src] = edge - if dst in stage_ids: - edge_in[dst] = edge - - workflows: dict[str, dict] = {} - for stage in config.get("stages", []): - sid = stage["id"] - upstream = stage.get("upstream", "") - downstream = stage.get("downstream", "") - - # Only treat upstream as a CM dependency when it's a real stage - deps = [upstream] if upstream in stage_ids else [] - - # concurrency_cap is used directly as max_replicas for local emulation - cap = int(stage.get("concurrency_cap", 0)) - max_r = int(stage.get("max_replicas", cap or 0)) - - # pilot.partition → required_cpus / required_gpus - pilot = stage.get("pilot", {}) - partition = pilot.get("partition", "cpu").lower() - resources = _PILOT_RESOURCES.get(partition, - {"required_cpus": 4, "required_gpus": 0}) - - # Outgoing edge: profile → dreamer strategy + early_binding - out_edge = edge_out.get(sid, {}) - profile = out_edge.get("profile", "diverse_top") - # Incoming edge: backpressure controls THIS stage's own queue depth - in_bp = edge_in.get(sid, {}).get("backpressure", {}) - strategy = _PROFILE_STRATEGY.get(profile, "smallest_to_fastest") - early_bind = _PROFILE_EARLY_BINDING.get(profile, True) - - # dreamer emulation block — may override strategy / early_binding - dreamer = dict(stage.get("dreamer", {})) - - # trigger_downstream: only when downstream is a registered stage - trigger = downstream if downstream in stage_ids else None - - # Replicas: debug.stage_replicas overrides stage.replicas (root stages only) - replicas = int(stage_replicas_dbg.get(sid, - stage.get("replicas", 0 if deps else 1))) - - # Trigger fraction: debug.trigger_fractions → falls back to threshold_top_fraction - trigger_fraction = float(trigger_fractions.get( - sid, stage.get("threshold_top_fraction", 1.0) or 1.0)) - - wf_cfg: dict = { - # ── CM scheduling (consumed by from_config, not forwarded) ─── - "replicas": replicas, - "min_replicas": int(stage.get("min_replicas", 0)), - "max_replicas": max_r, - "priority": int(stage.get("priority", 0)), - "dependencies": deps, - "dependency_threshold": int(stage.get("dependency_threshold", 1)), - "concurrency_cap": cap, - **resources, # required_cpus, required_gpus - - # ── Workflow config (forwarded to DreamerWorkflow.config) ──── - "trigger_downstream": trigger, - "trigger_fraction": trigger_fraction, # from debug.trigger_fractions - "threshold_top_fraction": stage.get("threshold_top_fraction"), - "budget_node_hours": stage.get("budget_node_hours"), - "downstream_input_target": stage.get("downstream_input_target"), - "pilot": pilot or None, - "surrogate": stage.get("surrogate"), - "profile": profile, - "schedule_strategy": dreamer.pop("schedule_strategy", strategy), - "early_binding": dreamer.pop("early_binding", early_bind), - # Backpressure for THIS stage's queue — only for dependent stages. - # Root (independent) stages have a fixed initial queue size so BP - # would immediately throttle them; skip it for those. - "backpressure_high": (in_bp.get("high_water") or None) if deps else None, - "backpressure_low": (in_bp.get("low_water") or None) if deps else None, - # Sharding spec — only for dependent stages (root stages are not triggered). - "sharding": stage.get("sharding") if deps else None, - **dreamer, # num_cores, perf_dist, num_tasks, ops_dist, profile_dir… - } - - # Drop keys with None / falsy values that would clutter workflow config - workflows[sid] = {k: v for k, v in wf_cfg.items() if v is not None} - - return workflows - - -# ── Legacy flat-format helpers ──────────────────────────────────────────────── - -def _expand_workflow_configs(config: dict, config_dir: Path) -> dict: - """Load external per-workflow YAML files referenced by 'config_file' keys.""" - for wf_cfg in config.get("workflows", {}).values(): - cfg_file = wf_cfg.pop("config_file", None) - if not cfg_file: - continue - cfg_path = Path(os.path.expandvars(cfg_file)) - if not cfg_path.is_absolute(): - cfg_path = config_dir / cfg_path - with open(cfg_path) as f: - wf_specific = _expand_env(yaml.safe_load(f) or {}) - wf_specific.update(wf_cfg) - wf_cfg.clear() - wf_cfg.update(wf_specific) - return config - - -def _build_registry(config: dict) -> dict: - """Dynamically import workflow classes from 'workflow_registry'.""" - import importlib - registry = {} - for name, cls_path in config.get("workflow_registry", {}).items(): - module_name, cls_name = cls_path.rsplit(".", 1) - registry[name] = getattr(importlib.import_module(module_name), cls_name) - return registry - - -# ── Main ────────────────────────────────────────────────────────────────────── - -async def main(config_file: str) -> None: - config_path = Path(config_file) - if not config_path.exists(): - raise FileNotFoundError(f"Config file not found: {config_file}") - config = load_config(config_file) - config_dir = config_path.parent - - # ── Detect and translate plan format ───────────────────────────────────── - if "stages" in config: - cm_cfg = config.get("cm", {}) - config["workflows"] = _build_from_plan(config) - # Hoist cm: runtime keys to the top level where from_config expects them. - # "debug" is intentionally excluded: our top-level debug: is a dict of - # emulation overrides; we set the CM's debug bool explicitly below. - for key in ("engine", "resources", "telemetry", "workflow_registry", "features"): - if key in cm_cfg and key not in config: - config[key] = cm_cfg[key] - # Overwrite the debug dict with the CM boolean so from_config works correctly - config["debug"] = bool(cm_cfg.get("debug", False)) - n_stages = len(config["stages"]) - n_edges = len(config.get("edges", [])) - print(f"Plan format: {n_stages} stages, {n_edges} edges → " - f"{len(config['workflows'])} workflow groups") - else: - _expand_workflow_configs(config, config_dir) - - engine_type = config.get("engine", "concurrent") - - # ── Build async backend ─────────────────────────────────────────────────── - engine_dragon = None - asyncflow = None - - if engine_type == "dragon": - try: - from radical.asyncflow import WorkflowEngine - from rhapsody.backends import DragonExecutionBackendV3 - - engine_dragon = await DragonExecutionBackendV3() - asyncflow = await WorkflowEngine.create(engine_dragon) - print("Dragon backend started") - except ImportError: - engine_type = "concurrent" - - if engine_type == "concurrent": - from radical.asyncflow import WorkflowEngine - from rhapsody.backends import ConcurrentExecutionBackend - - backend = await ConcurrentExecutionBackend() - asyncflow = await WorkflowEngine.create(backend) - print("ConcurrentExecutionBackend started") - - # ── Telemetry ───────────────────────────────────────────────────────────── - tel_cfg = config.get("telemetry", {}) - telemetry = None - if tel_cfg.get("collect_telemetry", False): - telemetry_dir = tel_cfg.get("telemetry_dir", "telemetry-results") - if hasattr(asyncflow, "start_telemetry"): - telemetry = await asyncflow.start_telemetry( - resource_poll_interval=0.5, - checkpoint_path=telemetry_dir, - ) - print(f"Asyncflow telemetry started → {telemetry_dir}") - - # ── Campaign ────────────────────────────────────────────────────────────── - registry = _build_registry(config) - cm = CampaignManager.from_config( - config, - registry, - asyncflow=asyncflow, - engine_dragon=engine_dragon, - ) - - groups = config.get("workflows", {}) - print( - "Campaign: " - + ", ".join( - f"{name}: {cfg.get('replicas', 0)} replica(s) " - f"cap={cfg.get('concurrency_cap', '—')} " - f"deps={cfg.get('dependencies', [])}" - for name, cfg in groups.items() - ) - ) - - try: - await cm.start() - await cm.wait() - finally: - await cm.close() - if telemetry: - await telemetry.stop() - print("Asyncflow telemetry stopped") - await asyncflow.shutdown() - - # ── Summary ─────────────────────────────────────────────────────────────── - print("\n── Campaign complete ──") - for name, info in cm.status()["groups"].items(): - print(f" {name}: status={info['status']} " - f"replicas={info['replicas_finished']}/{info['replicas_total']}") - - print("\n── Replica counts per workflow ──") - for name, s in cm.stats().items(): - print(f" {name}: replicas_started={s.replicas_started} " - f"replicas_finished={s.replicas_finished}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="SPHERICAL dreamer campaign runner") - parser.add_argument("--config", default="config.yaml", - help="Path to YAML config (flat or plan format)") - args = parser.parse_args() - asyncio.run(main(args.config)) diff --git a/workflows/run_campaign/dreamer_campaign/sbatch.sh b/workflows/run_campaign/dreamer_campaign/sbatch.sh deleted file mode 100644 index 6b9c5cc..0000000 --- a/workflows/run_campaign/dreamer_campaign/sbatch.sh +++ /dev/null @@ -1,33 +0,0 @@ -#!/bin/sh -l -# -# SPHERICAL Dreamer Campaign — SLURM batch script (CPU-only, no GPU needed) -# -#SBATCH -A *** -#SBATCH --partition=RM -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=64 -#SBATCH --time=01:00:00 -#SBATCH --job-name=dreamer_campaign -#SBATCH --mail-user=*** -#SBATCH --mail-type=END,FAIL - -# ── Environment ─────────────────────────────────────────────────────────────── -export SPHERICAL_DIR="/scratch/bblj/${USER}/SPHERICAL" -export DREAMER_DIR="/scratch/bblj/${USER}/radical.dreamer" -export ENV_DIR="/u/${USER}/ve/dreamer_campaign" - -export DREAMER_DIR="${DREAMER_DIR}" # picked up by dreamer_workflow.py - -unset SLURM_EXPORT_ENV -module load anaconda3 2>/dev/null || true - -source "${ENV_DIR}/bin/activate" - -# ── Run campaign ────────────────────────────────────────────────────────────── -CAMPAIGN_DIR="${SPHERICAL_DIR}/workflows/run_campaign/dreamer_campaign" -cd "${CAMPAIGN_DIR}" - -rm -rf dreamer-profiles telemetry-results - -python run_campaign.py --config config.yaml diff --git a/workflows/run_campaign/esm2_ddsim_campaign/config.yaml b/workflows/run_campaign/esm2_ddsim_campaign/config.yaml deleted file mode 100644 index 57d3124..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/config.yaml +++ /dev/null @@ -1,113 +0,0 @@ -# ============================================================================= -# SPHERICAL Campaign Configuration -# ============================================================================= - -# ── Cluster Resources ───────────────────────────────────────────────────────── -resources: - total_cpus: 64 # CPU cores on the allocated node(s) - total_gpus: 4 # GPUs on the allocated node(s) - -# ── Engine ──────────────────────────────────────────────────────────────────── -engine: dragon # concurrent | dragon -debug: false # set true to enable rhapsody DEBUG logging - -# ── Telemetry ───────────────────────────────────────────────────────────────── -telemetry: - collect_telemetry: true - telemetry_dir: "telemetry-results" - -# ── Workflow Registry ──────────────────────────────────────────────────────── -# Maps config workflow names to "module.ClassName" strings. -# Modules are resolved relative to the run_campaign directory. -workflow_registry: - dummy: dummy_workflow.DDSimWorkflow - md: ddmd_workflow.DDMdWrapperWorkflow - miniapps: miniapps_workflow.MiniAppsWrapperWorkflow - inference: inference_workflow.InferenceWorkflow - -# ── Workflow Groups ─────────────────────────────────────────────────────────── -# Each key must match a name in the workflow_registry above. -# -# 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: - - # ── 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 - min_replicas: 2 - max_replicas: 4 - required_cpus: 4 - required_gpus: 0 - dependencies: [inference] - config_file: "${DUMMY_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 ─────────────────────────────────────────────────────────────── - # 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 - min_replicas: 1 - max_replicas: 1 - required_cpus: 4 - required_gpus: 1 - dependencies: [md] - config_file: "${MINAPPS_DIR}/config.yaml" diff --git a/workflows/run_campaign/esm2_ddsim_campaign/cpu_batch.sh b/workflows/run_campaign/esm2_ddsim_campaign/cpu_batch.sh deleted file mode 100644 index 35c345b..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/cpu_batch.sh +++ /dev/null @@ -1,47 +0,0 @@ -#!/bin/sh -l - -#SBATCH -A *** -#SBATCH --partition=RM -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=64 -#SBATCH --time=01:25:00 -#SBATCH --job-name=spher -#SBATCH --mail-user=mariya.goliyad@rutgers.edu -#SBATCH --mail-type=ALL - - -export HF_TOKEN="hf_***" - -export BASE_DIR="${PROJECT}/DeepDriveSim" -export WORK_DIR="${BASE_DIR}/pipelines/ddmd_pipeline" -export CONDA_ENV="${WORK_DIR}/conda_env" -export INPUT_DIR="${WORK_DIR}/data" - -#WARNING: this directory has to be empty before running new experiment! -export EXPRMNT_DIR=$WORK_DIR/ddmd_test_experiments -# Remove the following line if you want to keep data from previous experiments. -rm -rf $EXPRMNT_DIR - -cp $INPUT_DIR/lassen-keras-dbscan.yaml $INPUT_DIR/new_lassen-keras-dbscan.yaml -sed -i "s|\${EXPRMNT_DIR}|$EXPRMNT_DIR|g" $INPUT_DIR/new_lassen-keras-dbscan.yaml -sed -i "s|\${CONDA_ENV}|$CONDA_ENV|g" $INPUT_DIR/new_lassen-keras-dbscan.yaml -sed -i "s|\${WORK_DIR}|$WORK_DIR|g" $INPUT_DIR/new_lassen-keras-dbscan.yaml - -# module load cuda -# module load gcc -# module load anaconda3 -#conda activate $(CONDA_ENV)/campaing_manager - -unset SLURM_EXPORT_ENV -module load anaconda3 -module load anaconda -source activate base -#conda activate $CONDA_ENV/deepdrivesim - -export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH - -conda activate $PROJECT/conda_env/test_inf -cd $PROJECT/htp/SPHERICAL/workflows/run_campaign - -python run_esm2_infern.py \ No newline at end of file diff --git a/workflows/run_campaign/esm2_ddsim_campaign/ddmd_workflow.py b/workflows/run_campaign/esm2_ddsim_campaign/ddmd_workflow.py deleted file mode 100644 index 235eb30..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/ddmd_workflow.py +++ /dev/null @@ -1,104 +0,0 @@ -""" -DDMdWrapperWorkflow — wraps the real DDMdWorkflow from DeepDriveSim. -""" - -import importlib.util - -# Make DeepDriveSim importable — honour $DDSIM_DIR set by the sbatch script. -import os -import sys -import tempfile -import traceback -from functools import lru_cache -from pathlib import Path - -import yaml - -_ddsim_dir = os.environ.get("DDSIM_DIR") -if not _ddsim_dir: - raise OSError("DDSIM_DIR is not set. Export it before launching the campaign.") -_DDSIM_ROOT = Path(_ddsim_dir) -if str(_DDSIM_ROOT) not in sys.path: - sys.path.insert(0, str(_DDSIM_ROOT)) - -from src.campaign import BaseWorkflow # noqa: E402 - - -@lru_cache(maxsize=1) -def _get_workflow_class(): - spec = importlib.util.spec_from_file_location( - "ddmd_workflow", - _DDSIM_ROOT / "workflows/ddmd_workflow/ddmd_workflow.py", - ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod.DDMdWorkflow - - -# Pre-load at module import time so the lru_cache is warm before Dragon's -# worker pool starts. Calling _get_workflow_class() inside an active Dragon -# event loop causes an import-lock deadlock: Dragon worker threads hold the -# import lock while initialising, and exec_module() blocks waiting for it. -_get_workflow_class() - - -class DDMdWrapperWorkflow(BaseWorkflow): - """Async wrapper that runs one replica of the DDMd pipeline.""" - - workflow_id = "ddmd" - - async def run(self, replica_id: str) -> None: - workflow_class = _get_workflow_class() # returns instantly from lru_cache - - asyncflow = self.asyncflow - cfg = self.config or {} - ddsim_config = cfg.get("ddsim_config") - if not ddsim_config: - raise ValueError(f"[{replica_id}] 'ddsim_config' missing from workflow config.") - - experiment_dir = cfg.get("experiment_dir", "") - replica_config_path = self._make_replica_config(ddsim_config, replica_id, experiment_dir) - name = replica_id.replace("_", "") - try: - workflow = workflow_class( - asyncflow=asyncflow, - config=replica_config_path, - name=name, - on_ready=lambda: self._signal_done(), - policies=self.policies, - engine_dragon=self.engine_dragon, - ) - except Exception: - print( - f"[{replica_id}] DDMdWorkflow.__init__ raised:\n" + traceback.format_exc(), - flush=True, - ) - raise - - try: - await workflow.start() - finally: - Path(replica_config_path).unlink(missing_ok=True) - - @staticmethod - def _make_replica_config( - base_config_path: str, replica_id: str, experiment_dir: str = "" - ) -> str: - with open(base_config_path) as f: - cfg = yaml.safe_load(f) - - if cfg.get("node_local_path"): - cfg["node_local_path"] = str(Path(cfg["node_local_path"]) / replica_id) - - if experiment_dir: - cfg["experiment_directory"] = str(Path(experiment_dir).expanduser().resolve()) - - tmp = tempfile.NamedTemporaryFile( - mode="w", - suffix=".yaml", - delete=False, - prefix=f"ddmd_{replica_id}_", - ) - yaml.dump(cfg, tmp) - tmp.close() - return tmp.name diff --git a/workflows/run_campaign/esm2_ddsim_campaign/delta_cpu_sbatch.sh b/workflows/run_campaign/esm2_ddsim_campaign/delta_cpu_sbatch.sh deleted file mode 100755 index f2a83d1..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/delta_cpu_sbatch.sh +++ /dev/null @@ -1,61 +0,0 @@ -#!/bin/sh -l - -# ── Cluster settings (adjust per allocation) ───────────────────────────────── -#SBATCH -A ***-delta-gpu -#SBATCH --partition=gpuA40x4 -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=64 -#SBATCH --gpus-per-node=4 -#SBATCH --time=00:30:00 -#SBATCH --job-name=campaign -#xSBATCH --mail-user=${USER}@institution.edu -#SBATCH --mail-user=mariya.goliyad@rutgers.edu -#SBATCH --mail-type=ALL - -# ── System library paths (Delta-specific) ──────────────────────────────────── -export CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_x86_64/25.3/cuda/12.8 -export MPI_LIB=/opt/cray/pe/mpich/8.1.32/ofi/gnu/11.2/lib-abi-mpich -export FAB_LIB=/opt/cray/libfabric/1.22.0/lib64 -export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${MPI_LIB}:${FAB_LIB}:${LD_LIBRARY_PATH} - -export TF_FORCE_GPU_ALLOW_GROWTH=true -export JAX_PLATFORMS=cpu - -# ── Project paths (adjust base dirs if layout differs) ─────────────────────── -export SPHERICAL_DIR=/scratch/***/${USER}/SPHERICAL -export DDSIM_DIR=/scratch/***/${USER}/DeepDriveSim -export VE_HOME=/u/${USER}/ve - -export MD_DIR=${DDSIM_DIR}/workflows/ddmd_workflow -export MINAPPS_DIR=${DDSIM_DIR}/workflows/miniapps_workflow -export DUMMY_DIR=${DDSIM_DIR}/workflows/dummy_workflow -export INF_DIR=${SPHERICAL_DIR}/workflows/esm2_inference - -export MD_HOME=${DDSIM_DIR}/workflows/ddmd_workflow -export MD_INPUT=${MD_HOME}/data -export SGDES_DIR=/scratch/***/${USER}/SGDES - -export WORK_DIR=${SPHERICAL_DIR}/workflows/run_campaign - -cd ${WORK_DIR} - -# ── Clean previous run artifacts ───────────────────────────────────────────── -rm -rf DDMD* - -# ── Activate campaign environment and configure Dragon ─────────────────────── -source ${VE_HOME}/campaign/bin/activate -dragon-config add --ofi-runtime-lib=${FAB_LIB} - -# ── Launch ─────────────────────────────────────────────────────────────────── -# Environment variable substitution (${DUMMY_DIR}, ${VE_HOME}, etc.) is handled -# by run_campaing.py at load time — no sed or cp needed. -GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-1} -export TOTAL_GPUS=$(( SLURM_NNODES * GPUS_PER_NODE )) -echo "Nodes: ${SLURM_NNODES} GPUs/node: ${GPUS_PER_NODE} Total GPUs: ${TOTAL_GPUS}" - -if [ "${SLURM_NNODES}" -gt 1 ]; then - dragon -m run_campaing.py --config config.yaml -else - dragon -s run_campaing.py --config config.yaml -fi diff --git a/workflows/run_campaign/esm2_ddsim_campaign/delta_env_setup.sh b/workflows/run_campaign/esm2_ddsim_campaign/delta_env_setup.sh deleted file mode 100755 index f6087c9..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/delta_env_setup.sh +++ /dev/null @@ -1,173 +0,0 @@ -#!/bin/bash -# ============================================================================= -# SPHERICAL Campaign Manager environment setup — Delta HPC (NCSA) -# -# Creates a Python venv with SPHERICAL and DeepDriveSim (campaign manager only). -# Workflow-specific environments (ddmd, miniapps, inference, etc.) are created -# by their own setup scripts; the CM uses the python executable specified in -# each workflow's config (e.g. executable: "/u/${USER}/ve/ddmd/bin/python"). -# -# Usage: -# bash delta_env_setup.sh [--env-dir DIR] [--spherical-dir DIR] [--ddsim-dir DIR] -# -# Defaults: -# ENV_DIR = /u/$USER/ve/campaign -# SPHERICAL_DIR = /scratch/bblj/$USER/SPHERICAL -# DDSIM_DIR = /scratch/bblj/$USER/DeepDriveSim -# ============================================================================= -if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then - set -euo pipefail -fi - -# ── Parse optional overrides ────────────────────────────────────────────────── -ENV_DIR="${ENV_DIR:-/u/${USER}/ve/campaign}" -SPHERICAL_DIR="${SPHERICAL_DIR:-/scratch/bblj/${USER}/tmp/SPHERICAL}" -DDSIM_DIR="${DDSIM_DIR:-/scratch/bblj/${USER}/DeepDriveSim}" - -while [[ $# -gt 0 ]]; do - case $1 in - --env-dir) ENV_DIR="$2"; shift 2 ;; - --spherical-dir) SPHERICAL_DIR="$2"; shift 2 ;; - --ddsim-dir) DDSIM_DIR="$2"; shift 2 ;; - *) echo "Unknown argument: $1"; exit 1 ;; - esac -done - -echo "=================================================================" -echo " ENV_DIR = ${ENV_DIR}" -echo " SPHERICAL_DIR = ${SPHERICAL_DIR}" -echo " DDSIM_DIR = ${DDSIM_DIR}" -echo "=================================================================" - -# ── 0. Clone repositories ───────────────────────────────────────────────────── -echo "" -echo "── Step 0: Cloning repositories ──" - -if [ ! -d "${SPHERICAL_DIR}/.git" ]; then - echo "Cloning SPHERICAL → ${SPHERICAL_DIR}" - git clone git@github.com:radical-collaboration/SPHERICAL.git "${SPHERICAL_DIR}" -else - echo "SPHERICAL already cloned at ${SPHERICAL_DIR}" -fi - -if [ ! -d "${DDSIM_DIR}/.git" ]; then - echo "Cloning DeepDriveSim → ${DDSIM_DIR}" - git clone --branch origin/campaign_manager --single-branch \ - https://github.com/radical-collaboration/DeepDriveSim.git "${DDSIM_DIR}" -else - echo "DeepDriveSim already cloned at ${DDSIM_DIR}" -fi - -# ── 1. Create venv ──────────────────────────────────────────────────────────── -echo "" -echo "── Step 1: Creating venv ──" - -BASE_PY=$(command -v python3.11 2>/dev/null || true) - -if [ -z "${BASE_PY}" ]; then - echo "python3.11 not in PATH — trying cray-python/3.11.7..." - module load cray-python/3.11.7 2>/dev/null || true - BASE_PY=$(command -v python3.11 2>/dev/null || true) -fi - -if [ -z "${BASE_PY}" ]; then - echo "python3.11 not available — trying python3.10 via anaconda3..." - module load anaconda3 2>/dev/null || true - BASE_PY=$(command -v python3.10 2>/dev/null || true) -fi - -if [ -z "${BASE_PY}" ]; then - BASE_PY=$(command -v python3 2>/dev/null || true) - [ -n "${BASE_PY}" ] && echo "Falling back to $(${BASE_PY} --version)" -fi - -if [ -z "${BASE_PY}" ]; then - echo "ERROR: no Python 3.10+ interpreter found." - echo " Try: module load anaconda3 or module load cray-python/3.11.7" - exit 1 -fi - -PY_VERSION=$("${BASE_PY}" -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") -PY="${ENV_DIR}/bin/python${PY_VERSION}" -PIP="${ENV_DIR}/bin/pip" -echo "Using Python: ${BASE_PY} ($(${BASE_PY} --version))" - -if [ ! -x "${PY}" ]; then - echo "Creating venv at ${ENV_DIR}..." - "${BASE_PY}" -m venv "${ENV_DIR}" -else - echo "venv already exists at ${ENV_DIR}" -fi - -ln -sf "${ENV_DIR}/bin/python${PY_VERSION}" "${ENV_DIR}/bin/python" 2>/dev/null || true -ln -sf "${ENV_DIR}/bin/python${PY_VERSION}" "${ENV_DIR}/bin/python3" 2>/dev/null || true - -echo "Python: $("${PY}" --version)" - -# ── 2. Bootstrap pip / setuptools ──────────────────────────────────────────── -echo "" -echo "── Step 2: Bootstrapping pip ──" -"${PY}" -m pip install -q --upgrade pip wheel -"${PIP}" install -q --force-reinstall "setuptools<71" - -# ── 3. Dragon / Rhapsody / Radical ──────────────────────────────────────────── -echo "" -echo "── Step 3: Dragon HPC + Rhapsody + Radical ──" -"${PIP}" install -q \ - "dragonhpc>=0.13.2" \ - "rhapsody-py>=0.2.0" \ - "radical.asyncflow>=0.3.1" \ - "nvidia-ml-py" \ - "numpy>=1.26.3,<2.0.0" \ - "transformers>=4.30.0" - -# ── 4. SPHERICAL (editable, campaign manager extras only) ──────────────────── -echo "" -echo "── Step 4: SPHERICAL ──" -"${PIP}" install -q -e "${SPHERICAL_DIR}[dragon,dev,plotting]" - -# ── 5. DeepDriveSim (editable) ─────────────────────────────────────────────── -echo "" -echo "── Step 5: DeepDriveSim ──" -"${PIP}" install -q -e "${DDSIM_DIR}" - -# ── 6. run_campaign extra requirements ─────────────────────────────────────── -echo "" -echo "── Step 6: run_campaign requirements ──" -"${PIP}" install -q -r "${SPHERICAL_DIR}/workflows/run_campaign/requirements.txt" - -# ── 7. Apply slurm patch ───────────────────────────────────────────────────── -echo "" -echo "── Step 7: Applying slurm patch ──" -"${PY}" "${SPHERICAL_DIR}/workflows/apply_slurm_patch.py" - -# ── 8. Verify ──────────────────────────────────────────────────────────────── -echo "" -echo "── Verifying installation ──" -_check() { - local label="$1"; shift - if out=$("$@" 2>&1); then - echo " ${label}: OK (${out})" - else - echo " WARNING: ${label} failed" - echo " ${out}" | head -3 - fi -} - -_check "radical.asyncflow" "${PY}" -c "import radical.asyncflow; print('ok')" -_check "rhapsody" "${PY}" -c "import rhapsody; print('ok')" -_check "dragonhpc" "${PY}" -c "import dragon; print('ok')" -_check "ddsim" "${PY}" -c "import ddsim; print('ok')" -_check "spherical" "${PY}" -c "import src.campaign; print('ok')" - -echo "" -echo "=================================================================" -echo "Setup complete." -echo "" -echo "Activate with:" -echo " source ${ENV_DIR}/bin/activate" -echo "" -echo "Note: workflow environments (ddmd, miniapps, inference, etc.) must be" -echo "set up separately. The CM uses the python executable from each" -echo "workflow's config (e.g. executable: \"/u/\${USER}/ve/ddmd/bin/python\")." -echo "=================================================================" diff --git a/workflows/run_campaign/esm2_ddsim_campaign/dummy_workflow.py b/workflows/run_campaign/esm2_ddsim_campaign/dummy_workflow.py deleted file mode 100644 index cc0c4c7..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/dummy_workflow.py +++ /dev/null @@ -1,71 +0,0 @@ -""" -DDSimWorkflow — wraps DummyWorkflow from DeepDriveSim in-process. -""" - -import importlib.util -import os -import sys -import traceback -from functools import lru_cache -from pathlib import Path - -_ddsim_dir = os.environ.get("DDSIM_DIR") -if not _ddsim_dir: - raise OSError("DDSIM_DIR is not set. Export it before launching the campaign.") -_DDSIM_ROOT = Path(_ddsim_dir) -if str(_DDSIM_ROOT) not in sys.path: - sys.path.insert(0, str(_DDSIM_ROOT)) - -from src.campaign import BaseWorkflow # noqa: E402 - - -@lru_cache(maxsize=1) -def _get_workflow_class(): - spec = importlib.util.spec_from_file_location( - "dummy_workflow_mod", - _DDSIM_ROOT / "workflows/dummy_workflow/dummy_workflow.py", - ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod.DummyWorkflow - - -# Pre-load before Dragon's worker pool starts to avoid import-lock deadlock. -_get_workflow_class() - - -class DDSimWorkflow(BaseWorkflow): - """Async wrapper that runs one replica of the DummyWorkflow pipeline.""" - - workflow_id = "dummy" - - async def run(self, replica_id: str) -> None: - workflow_class = _get_workflow_class() - - asyncflow = self.asyncflow - cfg = self.config or {} - - if not cfg.get("src_dir"): - cfg = {**cfg, "src_dir": str(_DDSIM_ROOT / "workflows/dummy_workflow")} - - name = replica_id.replace("_", "") - home_base = Path(cfg.get("home_dir", Path.home() / "Dummy")).expanduser() - - try: - workflow = workflow_class( - config=cfg, - name=name, - asyncflow=asyncflow, - home_dir=str(home_base), - _cm=self._cm, - _group_name=self._group_name, - policies=self.policies, - ) - except Exception: - print( - f"[{replica_id}] DummyWorkflow.__init__ raised:\n" + traceback.format_exc(), - flush=True, - ) - raise - - await workflow.start() diff --git a/workflows/run_campaign/esm2_ddsim_campaign/env_setup.sh b/workflows/run_campaign/esm2_ddsim_campaign/env_setup.sh deleted file mode 100644 index 38ba6fb..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/env_setup.sh +++ /dev/null @@ -1,32 +0,0 @@ -#!/bin/bash -export BASE_DIR="${PROJECT}" -export SPHERICAL_DIR="${BASE_DIR}/htp/SPHERICAL" -export DDSIM_DIR="${BASE_DIR}/DeepDriveSim" -export WORK_DIR="${SPHERICAL_DIR}/workflows/run_campaign" -export CONDA_ENV="${BASE_DIR}/conda_env" - -#mkdir $CONDA_ENV - -module load anaconda3 - -conda create -y -p $CONDA_ENV/campaign_manager python=3.10 -conda activate $CONDA_ENV/campaign_manager -pip install --upgrade pip setuptools wheel -cd $SPHERICAL_DIR -pip install -e ".[dragon,dev,esm2]" -cd $BASE_DIR -if [ ! -d "$DDSIM_DIR" ]; then - git clone --branch origin/campaign_manager --single-branch https://github.com/radical-collaboration/DeepDriveSim.git -fi -cd $DDSIM_DIR -pip install -e . -# cd "${DDSIM_DIR}/workflows/dummy_workflow/" -# pip install -r "requirements.txt" -# cd "${DDSIM_DIR}/workflows/ddmd_workflow/" -# pip install -r "requirements.txt" -# cd "${DDSIM_DIR}/workflows/miniapps_workflow/" -# pip install -r "requirements.txt" -cd $WORK_DIR -pip install -r "requirements.txt" -# conda init -conda deactivate \ No newline at end of file diff --git a/workflows/run_campaign/esm2_ddsim_campaign/gpu_sbatch.sh b/workflows/run_campaign/esm2_ddsim_campaign/gpu_sbatch.sh deleted file mode 100644 index 7660263..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/gpu_sbatch.sh +++ /dev/null @@ -1,68 +0,0 @@ -#!/bin/sh -l - -#SBATCH -A *** -#SBATCH --partition=GPU #-shared -#SBATCH --nodes=1 -#SBATCH --tasks-per-node=4 -#SBATCH --cpus-per-task=1 -#xSBATCH --gpus=v100-32:8 -#SBATCH --gpus=8 -#SBATCH --exclusive -#SBATCH --export NONE -#SBATCH --time=02:30:00 -#SBATCH --job-name sphr -#SBATCH --mail-user=mg2347@soe.rutgers.edu -#SBATCH --mail-type=ALL # When to send emails (BEGIN, END, FAIL, ALL) - - -export BASE_DIR="${PROJECT}" -export WORK_DIR="${BASE_DIR}/DeepDriveSim/workflows/ddmd_workflow" -export CONDA_ENV="${BASE_DIR}/conda_env" -export INPUT_DIR="${WORK_DIR}/data" -export DUMMY_DIR="${BASE_DIR}/DeepDriveSim/workflows/dummy_workflow" -export INF_DIR="${BASE_DIR}/htp/SPHERICAL/workflows/esm2_inference" -export MINAPPS_DIR="${BASE_DIR}/DeepDriveSim/workflows/miniapps_workflow" -export MD_DIR="${WORK_DIR}" - -#WARNING: this directory has to be empty before running new experiment! -export EXPRMNT_DIR=$WORK_DIR/ddmd_test_experiments -# Remove the following line if you want to keep data from previous experiments. -rm -rf $EXPRMNT_DIR - -unset SLURM_EXPORT_ENV -module load anaconda3 -#module load anaconda -source activate base -conda activate $CONDA_ENV/campaign_manager - -export CUDA_HOME=/opt/packages/cuda/v12.6.1 -export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH - -export TF_FORCE_GPU_ALLOW_GROWTH=true - -cp $INPUT_DIR/lassen-keras-dbscan.yaml $INPUT_DIR/new_lassen-keras-dbscan.yaml -sed -i "s|\${EXPRMNT_DIR}|$EXPRMNT_DIR|g" $INPUT_DIR/new_lassen-keras-dbscan.yaml -sed -i "s|\${CONDA_ENV}|$CONDA_ENV|g" $INPUT_DIR/new_lassen-keras-dbscan.yaml -sed -i "s|\${WORK_DIR}|$WORK_DIR|g" $INPUT_DIR/new_lassen-keras-dbscan.yaml - -cp $WORK_DIR/template_config.yaml $WORK_DIR/config.yaml -sed -i "s|\${PROJECT}|$PROJECT|g" $WORK_DIR/config.yaml - -cp template_config.yaml config.yaml -sed -i "s|\${MD_DIR}|$MD_DIR|g" config.yaml -sed -i "s|\${MINAPPS_DIR}|$MINAPPS_DIR|g" config.yaml -sed -i "s|\${INF_DIR}|$INF_DIR|g" config.yaml -sed -i "s|\${DUMMY_DIR}|$DUMMY_DIR|g" config.yaml - -cp $MINAPPS_DIR/template_config.yaml $MINAPPS_DIR/config.yaml -sed -i "s|\${PROJECT}|$PROJECT|g" $MINAPPS_DIR/config.yaml -cp $DUMMY_DIR/template_config.yaml $DUMMY_DIR/config.yaml -sed -i "s|\${PROJECT}|$PROJECT|g" $DUMMY_DIR/config.yaml - -cd $BASE_DIR/htp/SPHERICAL/workflows/run_campaign -rm -rf data/telemetry-results -rm -rf data/nvml-telemetry - -dragon -s run_campaing.py -#python run_campaing.py -#python -m run_workflow -c $INPUT_DIR/new_lassen-keras-dbscan.yaml diff --git a/workflows/run_campaign/esm2_ddsim_campaign/inference_workflow.py b/workflows/run_campaign/esm2_ddsim_campaign/inference_workflow.py deleted file mode 100644 index 1aa8939..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/inference_workflow.py +++ /dev/null @@ -1,265 +0,0 @@ -""" -InferenceWorkflow — async single-phase ESM2 client request workflow. - -N ESM2 services are initialised exactly once (one per GPU, auto-assigned by -start_services_local based on num_services in config). Replicas round-robin -across services: replica_i → service (i % N). A per-service asyncio.Lock -serialises back-to-back replicas on the same service so queue state is always -cleanly reset before the next client runs. Workers and the asyncflow engine -stay alive for the full campaign duration. - -Lifecycle ---------- - First replica → _ensure_initialized() (expensive: model load × N, workers) - Every replica → acquire service lock → processed_queue.join() - → _reset_service_queues() → ESM2Client.run() → release lock - Last replica → on_replica_done() → _teardown() -""" - -import asyncio -from pathlib import Path -from typing import ClassVar, Optional - -from src.campaign import BaseWorkflow -from src.utils.logger import Logger - - -class InferenceWorkflow(BaseWorkflow): - workflow_id = "inference" - - # ------------------------------------------------------------------ # - # Shared state — lives across all replicas in the campaign # - # ------------------------------------------------------------------ # - - _svc_handles: ClassVar[Optional[list]] = None # one handle per service - _svc_locks: ClassVar[Optional[list[asyncio.Lock]]] = None # one lock per service - _asyncflow: ClassVar = None # shared WorkflowEngine - _init_lock: ClassVar[Optional[asyncio.Lock]] = None # one-time init guard - _num_services: ClassVar[int] = 0 - _log: ClassVar[Logger] = Logger(name="InferenceWorkflow", use_colors=True) - - # ------------------------------------------------------------------ # - # Replica entry point # - # ------------------------------------------------------------------ # - - async def run(self, replica_id: str) -> None: - cfg = self.config or {} - await self._client_request(replica_id, cfg) - - # ------------------------------------------------------------------ # - # Debug / real dispatch # - # ------------------------------------------------------------------ # - - async def _client_request(self, replica_id: str, cfg: dict) -> None: - if cfg.get("debug", False): - await self._run_stub(replica_id) - return - try: - await self._run_real_inference(replica_id, cfg) - except BaseException as exc: - if isinstance(exc, asyncio.CancelledError): - raise - InferenceWorkflow._log.error( - f"inference failed ({type(exc).__name__}: {exc}); running stub", - component="workflow", - task_name=replica_id, - ) - await self._run_stub(replica_id) - - # ------------------------------------------------------------------ # - # Real inference — round-robin across shared services # - # ------------------------------------------------------------------ # - - async def _run_real_inference(self, replica_id: str, cfg: dict) -> None: - # Probe transformers before spawning any server: if it's missing this - # raises immediately and _ensure_initialized (which launches Dragon - # server processes) is never reached, preventing port-conflict cascades. - import os as _os - - _os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error") - import transformers # noqa: F401 - - from src.inference.esm2_service.esm2_client import ESM2Client - from src.inference.esm2_service.esm2_service import ESM2InferenceService - from src.inference.utils import export_metrics - - await self._ensure_initialized(cfg, ESM2InferenceService, asyncflow=self.asyncflow) - - # Round-robin: replica index parsed from "group_N" replica_id. - replica_idx = int(replica_id.split("_")[-1]) - svc_idx = replica_idx % InferenceWorkflow._num_services - lock = InferenceWorkflow._svc_locks[svc_idx] - handle = InferenceWorkflow._svc_handles[svc_idx] - svc = handle.service - - async with lock: - # Wait for all in-flight executor threads to finish first. Worker - # task_done on work_queue is called only after run_in_executor returns - # (thread fully done), so this guarantees all processed_queue puts have - # happened before we drain processed_queue. Without this ordering, - # processed_queue.join() can return while a slow thread is still mid- - # flight and will later put a stale batch_id that the next replica's - # _result_writer sees without a matching reply_store entry. - await svc.work_queue.join() - # Now drain the result writer — all puts are guaranteed to be in flight. - await svc.processed_queue.join() - # Reset per-run queue state so init_queue can repopulate. - self._reset_service_queues(svc) - - client = ESM2Client( - endpoints=[handle.endpoint] if handle.endpoint else [], - rank=0, - service=svc, - config=cfg, - asyncflow=InferenceWorkflow._asyncflow, - ) - - output_dir = cfg.get("output_dir", "data/outputs") - await client.run() - InferenceWorkflow._log.info( - f"inference complete → {output_dir}", - component="workflow", - task_name=replica_id, - ) - - metrics_dir = cfg.get("metrics_dir", "outputs") - await export_metrics( - Path(metrics_dir, f"client_{replica_id}.json"), - client.metrics, - ) - - # ------------------------------------------------------------------ # - # One-time service initialisation # - # ------------------------------------------------------------------ # - - @classmethod - async def _ensure_initialized(cls, cfg: dict, service_class, asyncflow=None) -> None: - """Start all N services exactly once. Uses the shared asyncflow if provided.""" - if cls._init_lock is None: - cls._init_lock = asyncio.Lock() - - async with cls._init_lock: - if cls._svc_handles is not None: - return - - from src.inference.orchestrator import start_services, start_services_local - - mode = cfg.get("mode", "local") - cls._log.info(f"Starting ESM2 services (mode={mode})") - if mode == "server": - handles = await start_services(cfg, service_class) - else: - handles = await start_services_local(cfg, service_class) - if not handles: - raise RuntimeError("Failed to initialise ESM2 inference services") - - cls._svc_handles = handles - cls._num_services = len(handles) - cls._svc_locks = [asyncio.Lock() for _ in range(cls._num_services)] - - cls._asyncflow = asyncflow - - if mode != "server": - for h in cls._svc_handles: - await h.service.start_workers() - - cls._log.info( - f"Initialized {cls._num_services} service(s) across " - f"{cls._num_services} GPU(s) (asyncflow=shared)" - ) - - # ------------------------------------------------------------------ # - # Per-replica queue reset (safe under per-service lock) # - # ------------------------------------------------------------------ # - - @staticmethod - def _reset_service_queues(svc) -> None: - """Reset queue/event state between replicas on the same service.""" - svc.shutdown_init.clear() - for q in (svc.seq_queue, svc.input_queue): - while not q.empty(): - try: - q.get_nowait() - if q is svc.seq_queue: - q.task_done() - except Exception: - break - svc.single_batch = None - svc.device_batches.clear() - - # ------------------------------------------------------------------ # - # Replica-done hook # - # ------------------------------------------------------------------ # - - async def on_replica_done(self, replica_id: str, cm, final_state: str) -> None: - """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 - total = g.get("replicas_total", 1) - - InferenceWorkflow._log.info( - f"replica finished [{final_state}] ({finished}/{total})", - component="workflow", - task_name=replica_id, - ) - - if final_state == "done": - # 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( - "all inference replicas done — tearing down ESM2 services", - component="workflow", - task_name=replica_id, - ) - await InferenceWorkflow._teardown() - - # ------------------------------------------------------------------ # - # Teardown after last replica # - # ------------------------------------------------------------------ # - - @classmethod - async def _teardown(cls) -> None: - """Shut down inference services and workers. Idempotent. - - NOTE: cls._asyncflow is intentionally NOT shut down here because - other workflow groups (e.g. ddsim) may still be running and share - the same asyncio event loop subprocess infrastructure. Call - _shutdown_asyncflow() explicitly after cm.wait() returns. - """ - if cls._svc_handles is None: - return - - for h in cls._svc_handles: - svc = h.service - await svc.work_queue.join() - await svc.processed_queue.join() - await svc.shutdown() - cls._svc_handles = None - cls._svc_locks = None - cls._num_services = 0 - - @classmethod - async def _shutdown_asyncflow(cls) -> None: - """No-op: asyncflow is owned and shut down by AsyncCampaignManager.""" - pass - - # ------------------------------------------------------------------ # - # Stub (debug / no-GPU path) # - # ------------------------------------------------------------------ # - - async def _run_stub(self, replica_id: str) -> None: - InferenceWorkflow._log.debug( - "client_req starting (stub)", - component="workflow", - task_name=replica_id, - ) - await asyncio.sleep(0.1) - InferenceWorkflow._log.debug( - "client_req done → requests_sent=500", - component="workflow", - task_name=replica_id, - ) diff --git a/workflows/run_campaign/esm2_ddsim_campaign/miniapps_workflow.py b/workflows/run_campaign/esm2_ddsim_campaign/miniapps_workflow.py deleted file mode 100644 index 849a6c6..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/miniapps_workflow.py +++ /dev/null @@ -1,81 +0,0 @@ -""" -MiniAppsWorkflow — wraps MiniAppsWorkflow from DeepDriveSim. -""" - -import importlib.util -import os -import sys -import traceback -from functools import lru_cache -from pathlib import Path - -# Make DeepDriveSim importable — honour $DDSIM_DIR set by the sbatch script. -_ddsim_dir = os.environ.get("DDSIM_DIR") -if not _ddsim_dir: - raise OSError("DDSIM_DIR is not set. Export it before launching the campaign.") -_DDSIM_ROOT = Path(_ddsim_dir) -if str(_DDSIM_ROOT) not in sys.path: - sys.path.insert(0, str(_DDSIM_ROOT)) - -from src.campaign import BaseWorkflow # noqa: E402 - - -@lru_cache(maxsize=1) -def _get_workflow_class(): - spec = importlib.util.spec_from_file_location( - "miniapps_workflow", - _DDSIM_ROOT / "workflows/miniapps_workflow/miniapps_workflow.py", - ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - return mod.MiniAppsWorkflow - - -# Pre-load at module import time to warm the cache before Dragon workers start. -_get_workflow_class() - - -class MiniAppsWrapperWorkflow(BaseWorkflow): - """Async wrapper that runs one replica of the MiniApps pipeline.""" - - workflow_id = "miniapps" - - async def run(self, replica_id: str) -> None: - # Workaround: concurrent asyncflow backend does not apply process_template.env - # to subprocess children, so CUDA_VISIBLE_DEVICES must be set in os.environ - # before any subprocesses are spawned by the workflow. - if self.policies: - gpu_id = str(self.policies[0].gpu_affinity[0]) - os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id - elif self.config and self.config.get("assigned_gpu_ids"): - gpu_id = str(self.config["assigned_gpu_ids"][0]) - os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id - - workflow_class = _get_workflow_class() - asyncflow = self.asyncflow - cfg = self.config or {} - - if not cfg.get("src_dir"): - cfg = {**cfg, "src_dir": str(_DDSIM_ROOT / "workflows/miniapps_workflow")} - - name = replica_id.replace("_", "") - home_base = Path(cfg.get("home_dir", Path.home() / "MiniApps")).expanduser() - - try: - workflow = workflow_class( - config=cfg, - asyncflow=asyncflow, - home_dir=str(home_base), - name=name, - _cm=self._cm, - _group_name=self._group_name, - policies=self.policies, - ) - except Exception: - print( - f"[{replica_id}] MiniAppsWorkflow.__init__ raised:\n" + traceback.format_exc(), - flush=True, - ) - raise - - await workflow.start() diff --git a/workflows/run_campaign/esm2_ddsim_campaign/requirements.txt b/workflows/run_campaign/esm2_ddsim_campaign/requirements.txt deleted file mode 100644 index 14518d5..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -pydantic-settings \ No newline at end of file diff --git a/workflows/run_campaign/esm2_ddsim_campaign/run_campaing.py b/workflows/run_campaign/esm2_ddsim_campaign/run_campaing.py deleted file mode 100644 index d71167b..0000000 --- a/workflows/run_campaign/esm2_ddsim_campaign/run_campaing.py +++ /dev/null @@ -1,188 +0,0 @@ -#!/usr/bin/env python3 -# Limit OpenBLAS/OMP threads before any numpy import to avoid pthread_create -# failures on login nodes where process counts are restricted. -import os - -os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") -os.environ.setdefault("OMP_NUM_THREADS", "1") -os.environ.setdefault("MKL_NUM_THREADS", "1") - -""" -Campaign runner — starts async workflow replicas with optional dependency -ordering between workflow groups. - -Usage ------ - python run_campaing.py --config config.yaml - -Config file structure ---------------------- - # ── Per-workflow sections ───────────────────────────────────────────── - workflows: - ddsim: - replicas: 8 - min_replicas: 2 - max_replicas: 4 - dependencies: [] - ddsim_config: "/path/to/ddmd_config.yaml" - - inference: - replicas: 1 - dependencies: [ddsim] # starts only after all ddsim replicas finish - num_gpus_per_service: 4 - -If no config file is provided, hard-coded defaults are used for local testing. -""" - -import argparse # noqa: E402 -import asyncio # noqa: E402 -import sys # noqa: E402 -from pathlib import Path # noqa: E402 - -import yaml # noqa: E402 - -sys.path.insert(0, str(Path(__file__).parent.parent.parent)) - -from src.campaign import AsyncCampaignManager as CampaignManager # noqa: E402 -from src.inference.utils import load_config # noqa: E402 -from src.utils.workflow import _expand_env # noqa: E402 - - -def _expand_workflow_configs(config: dict, config_dir: Path) -> dict: - """ - For each workflow entry that has a ``config_file`` key, load that YAML and - merge its contents into the workflow config dict. The per-workflow file - provides workflow-specific parameters; any keys already present in the - workflow entry (scheduling params) take precedence. - """ - for wf_cfg in config.get("workflows", {}).values(): - cfg_file = wf_cfg.pop("config_file", None) - if not cfg_file: - continue - cfg_path = Path(os.path.expandvars(cfg_file)) - if not cfg_path.is_absolute(): - cfg_path = config_dir / cfg_path - with open(cfg_path) as f: - wf_specific = _expand_env(yaml.safe_load(f) or {}) - # Scheduling params in config.yaml win; workflow file fills the rest. - wf_specific.update(wf_cfg) - wf_cfg.clear() - wf_cfg.update(wf_specific) - return config - - -def _build_registry(config: dict) -> dict: - """Dynamically import workflow classes from the 'workflow_registry' config section.""" - import importlib - - registry = {} - for name, cls_path in config.get("workflow_registry", {}).items(): - module_name, cls_name = cls_path.rsplit(".", 1) - module = importlib.import_module(module_name) - registry[name] = getattr(module, cls_name) - return registry - - -async def main(config_file: str) -> None: - config_path = Path(config_file) - if not config_path.exists(): - raise FileNotFoundError(f"Config file not found: {config_file}") - config = load_config(config_file) - config_dir = config_path.parent - _expand_workflow_configs(config, config_dir) - - engine_type = config.get("engine", "dragon") - - # ── Build backend and asyncflow (mirrors workflow run_workflow.py pattern) ─ - engine_dragon = None - asyncflow = None - - if engine_type == "dragon": - try: - from radical.asyncflow import WorkflowEngine - from rhapsody.backends import DragonExecutionBackendV3 - - engine_dragon = await DragonExecutionBackendV3() - asyncflow = await WorkflowEngine.create(engine_dragon) - print("Dragon backend started") - except ImportError: - engine_type = "concurrent" - - else: - from radical.asyncflow import WorkflowEngine - from rhapsody.backends import ConcurrentExecutionBackend - - backend = await ConcurrentExecutionBackend() - asyncflow = await WorkflowEngine.create(backend) - print("ConcurrentExecutionBackend started") - - # ── Telemetry ───────────────────────────────────────────────────────────── - tel_cfg = config.get("telemetry", {}) - telemetry = None - if tel_cfg.get("collect_telemetry", False): - telemetry_dir = tel_cfg.get("telemetry_dir", "data/telemetry-results") - if hasattr(asyncflow, "start_telemetry"): - telemetry = await asyncflow.start_telemetry( - resource_poll_interval=0.5, - checkpoint_path=telemetry_dir, - ) - print(f"Started Asyncflow telemetry → {telemetry_dir}") - - # ── Campaign ────────────────────────────────────────────────────────────── - registry = _build_registry(config) - cm = CampaignManager.from_config( - config, - registry, - asyncflow=asyncflow, - engine_dragon=engine_dragon, - ) - - groups = config.get("workflows", {}) - print( - "Campaign: " - + ", ".join( - f"{name}: {cfg.get('replicas', 1)} replica(s) deps={cfg.get('dependencies', [])}" - for name, cfg in groups.items() - ) - ) - - try: - await cm.start() # launch groups with no unmet dependencies - await cm.wait() # block until all groups (including dependents) finish - finally: - await cm.close() - - if telemetry: - await telemetry.stop() - print("Asyncflow telemetry stopped") - - await asyncflow.shutdown() - - # ── Summary ──────────────────────────────────────────────────────────── - print("\n── Campaign complete ──") - final_status = cm.status() - for name, info in final_status["groups"].items(): - print( - f" {name}: status={info['status']} " - f"replicas={info['replicas_finished']}/{info['replicas_total']}" - ) - - print("\n── Replica counts per workflow ──") - for name, s in cm.stats().items(): - print( - f" {name}: " - f"replicas_started={s.replicas_started} " - f"replicas_finished={s.replicas_finished}" - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="SPHERICAL campaign runner") - parser.add_argument( - "--config", - default="config.yaml", - help="Path to YAML config file (default: config.yaml)", - ) - - args = parser.parse_args() - asyncio.run(main(args.config)) diff --git a/workflows/run_campaign/plot_cm_timeline.py b/workflows/run_campaign/plot_cm_timeline.py deleted file mode 100644 index 9d4023e..0000000 --- a/workflows/run_campaign/plot_cm_timeline.py +++ /dev/null @@ -1,737 +0,0 @@ -#!/usr/bin/env python3 -""" -Plot replica execution timeline from a campaign SLURM log. - -Usage: - python plot_cm_timeline.py slurm-XXXXXX.out [--out timeline.png] -""" - -import argparse -import re -import sys -from datetime import datetime -from pathlib import Path - -try: - import yaml - - _HAVE_YAML = True -except ImportError: - _HAVE_YAML = False - -import matplotlib - -matplotlib.use("Agg") -import matplotlib.gridspec as gridspec -import matplotlib.lines as mlines -import matplotlib.patches as mpatches -import matplotlib.pyplot as plt - -_TS_RE = re.compile(r"\x1b\[2m(\d{2}:\d{2}:\d{2}\.\d{3})\x1b\[0m") -_START_RE = re.compile(r"starting replica '(\w+)'") -_FINISH_RE = re.compile(r"Replica '(\w+)' finished") -_ERROR_RE = re.compile(r"Replica '(\w+)' raised") -_GROUP_RE = re.compile( - r"Registered group '(\w+)': replicas=(\d+) priority=(\d+) " - r"min=(\d+) max=(\d+) deps=\[([^\]]*)\] dep_threshold=(\d+) " - r"resources=\(cpus=(\d+), gpus=(\d+)\)" -) -_GPU_ASSIGN_RE = re.compile(r"GPU assign: '(\w+)' → GPU\(s\) \[([^\]]*)\]") -_USAGE_RE = re.compile(r"resources: cpus=(\d+)/(\d+)\s+gpus=(\d+)/(\d+)") -_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+)") - -# 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", - "md": "#DD8452", - "miniapps": "#55A868", - "dummy": "#C44E52", -} -DEFAULT_COLOR = "#8172B2" - -GROUP_ORDER = ["md", "miniapps", "inference", "dummy"] - - -def parse_log(path: str): - """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, source_group, target_groups, n_replicas, kind) - # kind: "signal_done" | "trigger_dep" - signal_events = [] - t0_dt = None - total_cpus = total_gpus = 0 - - _iso_re = re.compile(r"^(\d{4}-\d{2}-\d{2}) \d{2}:\d{2}:\d{2}") - date_ref = "1970-01-01" - with open(path) as fh: - for raw in fh: - if m := _iso_re.match(raw): - date_ref = m.group(1) - break - - with open(path) as fh: - for raw in fh: - if m := _GROUP_RE.search(raw): - name = m.group(1) - deps_raw = m.group(6) - deps = [ - d.strip().strip("'\"") for d in deps_raw.split(",") if d.strip().strip("'\"") - ] - group_meta[name] = { - "replicas": int(m.group(2)), - "priority": int(m.group(3)), - "min": int(m.group(4)), - "max": int(m.group(5)), - "deps": deps, - "dep_threshold": int(m.group(7)), - "cpus": int(m.group(8)), - "gpus": int(m.group(9)), - } - - if m := _TOTAL_RES_RE.search(raw): - total_cpus = int(m.group(1)) - total_gpus = int(m.group(2)) - - ts_m = _TS_RE.search(raw) - if ts_m is None: - continue - 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) - gpu_str = m.group(2).strip() - gpu_ids = [int(x) for x in gpu_str.split(",") if x.strip()] if gpu_str else [] - gpu_assignments[rid] = gpu_ids - - 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)), - ) - resource_timeline.append((elapsed, uc, tc, ug, tg)) - 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)), - ) - 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): - rid = m.group(1) - if rid in starts: - group = rid.rsplit("_", 1)[0] - spans.append((rid, group, starts.pop(rid), dt, True)) - elif m := _ERROR_RE.search(raw): - rid = m.group(1) - if rid in starts: - group = rid.rsplit("_", 1)[0] - spans.append((rid, group, starts.pop(rid), dt, False)) - - for rid, start in starts.items(): - group = rid.rsplit("_", 1)[0] - spans.append((rid, group, start, start, None)) - - 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, - signal_events, - t0, - (total_cpus, total_gpus), - ) - - -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 - - def sort_key(s): - rid, group, *_ = s - idx = int(rid.rsplit("_", 1)[-1]) - g_idx = GROUP_ORDER.index(group) if group in GROUP_ORDER else len(GROUP_ORDER) - return (g_idx, idx) - - spans.sort(key=sort_key) - - has_resources = len(resource_timeline) > 0 - total_cpus, total_gpus = total_resources - - n_rows = len(spans) - gantt_h = max(6, n_rows * 0.38) - - fig = plt.figure(figsize=(20, gantt_h + (3 if has_resources else 0) + 1)) - - if has_resources: - gs = gridspec.GridSpec( - 2, - 2, - height_ratios=[gantt_h, 2.5], - width_ratios=[3, 1], - hspace=0.2, - wspace=0.18, - ) - ax_gantt = fig.add_subplot(gs[0, 0]) - ax_info = fig.add_subplot(gs[0, 1]) - ax_res = fig.add_subplot(gs[1, 0]) - ax_leg = fig.add_subplot(gs[1, 1]) - ax_leg.axis("off") - else: - gs = gridspec.GridSpec(1, 2, width_ratios=[3, 1], wspace=0.18) - ax_gantt = fig.add_subplot(gs[0, 0]) - ax_info = fig.add_subplot(gs[0, 1]) - ax_res = None - - # ── Gantt chart ────────────────────────────────────────────────────────── - yticks, ylabels = [], [] - group_row_ranges = {} - - prev_group = None - for row, (rid, group, start, end, ok) in enumerate(spans): - t_start = (start - t0).total_seconds() - t_end = (end - t0).total_seconds() if end != start else t_start + 0.5 - bar_w = t_end - t_start - - color = GROUP_COLORS.get(group, DEFAULT_COLOR) - edgecolor = "red" if ok is False else "none" - lw = 1.5 if ok is False else 0 - alpha = 0.45 if ok is None else 0.88 - - 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 - - if group not in group_row_ranges: - group_row_ranges[group] = [row, row] - else: - group_row_ranges[group][1] = row - - 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) - - ax_gantt.barh( - row, - bar_w, - left=t_start, - height=0.72, - color=color, - edgecolor=edgecolor, - linewidth=lw, - alpha=alpha, - ) - - gpu_ids = gpu_assignments.get(rid, []) - meta = group_meta.get(group, {}) - if gpu_ids: - ann_txt = f"gpu:{','.join(str(g) for g in gpu_ids)}" - elif meta.get("cpus", 0) > 0: - ann_txt = f"{meta['cpus']} cpu(s)" - else: - ann_txt = "" - - if ann_txt and bar_w > 0.5: - ax_gantt.text( - t_start + bar_w / 2, - row, - ann_txt, - ha="center", - va="center", - fontsize=5.5, - color="white", - fontweight="bold", - clip_on=True, - ) - - yticks.append(row) - ylabels.append(rid) - - # 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) - 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, - info, - transform=ax_gantt.transAxes, - va="center", - ha="left", - fontsize=6.5, - color=GROUP_COLORS.get(group, DEFAULT_COLOR), - fontweight="bold", - ) - - # 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 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] - grp_first = group_spans[group][0] - dep_row, dep_start, dep_end, _ = dep_first - grp_row, grp_start, _, _ = grp_first - dep_mid_t = (dep_start - t0).total_seconds() - if dep_end != dep_start: - dep_mid_t += (dep_end - dep_start).total_seconds() / 2 - grp_start_t = (grp_start - t0).total_seconds() - ax_gantt.annotate( - "", - xy=(grp_start_t, grp_row), - xytext=(dep_mid_t, dep_row), - arrowprops=dict( - arrowstyle="->", - color="#888888", - lw=1.0, - connectionstyle="arc3,rad=0.35", - linestyle="dashed", - ), - annotation_clip=False, - ) - - ax_gantt.set_yticks(yticks) - ax_gantt.set_yticklabels(ylabels, fontsize=7) - ax_gantt.set_xlabel("Elapsed time (s)", fontsize=9) - ax_gantt.set_title("Campaign Manager Timeline", fontweight="bold", fontsize=12) - ax_gantt.invert_yaxis() - ax_gantt.grid(axis="x", linestyle="--", alpha=0.35) - - legend_patches = [mpatches.Patch(color=c, label=g) for g, c in GROUP_COLORS.items()] - 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) - - # ── Group info + dependency table ──────────────────────────────────────── - ax_info.axis("off") - ax_info.set_title("Campaign Manager Config", fontweight="bold", fontsize=9, pad=4) - - if group_meta: - info_groups = [g for g in GROUP_ORDER if g in group_meta] + [ - g for g in group_meta if g not in GROUP_ORDER - ] - - 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, - str(m.get("priority", "?")), - str(m.get("cpus", 0)), - str(m.get("gpus", 0)), - f"{m.get('min', 0)}/{m.get('max', 0)}", - mode, - deps, - ] - ) - c = GROUP_COLORS.get(gname, DEFAULT_COLOR) - row_colors.append([c] + ["#f5f5f5"] * (len(col_labels) - 1)) - - tbl = ax_info.table( - cellText=rows_data, - colLabels=col_labels, - cellColours=row_colors, - loc="upper center", - cellLoc="center", - ) - tbl.auto_set_font_size(False) - tbl.set_fontsize(7.5) - tbl.scale(1.0, 1.5) - - for j in range(len(col_labels)): - tbl[0, j].set_facecolor("#333333") - tbl[0, j].set_text_props(color="white", fontweight="bold") - - # 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] - 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_lines: - dep_str = "Dependency graph (signals):\n" + "\n".join(dep_lines) - ax_info.text( - 0.5, - 0.22, - dep_str, - transform=ax_info.transAxes, - va="bottom", - ha="center", - fontsize=8, - family="monospace", - bbox=dict(boxstyle="round,pad=0.5", facecolor="#f0f4ff", edgecolor="#aabbdd"), - ) - - 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 (priority)\n" - " Pass 2: fill up to max replicas (priority)" - ) - ax_info.text( - 0.5, - 0.52, - sched_note, - transform=ax_info.transAxes, - va="bottom", - ha="center", - fontsize=8, - bbox=dict(boxstyle="round,pad=0.5", facecolor="#fffbe6", edgecolor="#ccaa00"), - ) - - # ── Resource utilization subplot ───────────────────────────────────────── - if ax_res is not None and resource_timeline: - times = [t for t, *_ in resource_timeline] - used_gpus = [ug for _, _, _, ug, _ in resource_timeline] - used_cpus = [uc for _, uc, *_ in resource_timeline] - tot_gpus = [tg for _, _, _, _, tg in resource_timeline] - tot_cpus = [tc for _, _, tc, *_ in resource_timeline] - - ax_res.step(times, used_gpus, where="post", color="#4C72B0", lw=1.8, label="GPUs used") - ax_res.fill_between(times, used_gpus, step="post", color="#4C72B0", alpha=0.15) - if any(t > 0 for t in tot_gpus): - ax_res.step( - times, - tot_gpus, - where="post", - color="#4C72B0", - lw=0.8, - linestyle="--", - alpha=0.55, - label="GPU total", - ) - - ax_res.set_ylabel("GPUs in use", color="#4C72B0", fontsize=8) - 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) - if any(t > 0 for t in tot_cpus): - ax_cpu.step( - times, - tot_cpus, - where="post", - color="#DD8452", - lw=0.8, - linestyle="--", - alpha=0.55, - label="CPU total", - ) - - ax_cpu.set_ylabel("CPUs in use", color="#DD8452", fontsize=8) - ax_cpu.tick_params(axis="y", labelcolor="#DD8452", labelsize=7) - ax_cpu.set_ylim(bottom=0) - - ax_res.set_xlabel("Elapsed time (s)", fontsize=8) - ax_res.set_title("Resource Utilization (GPU / CPU)", fontweight="bold", fontsize=9) - ax_res.grid(axis="x", linestyle="--", alpha=0.35) - - lines1, lbl1 = ax_res.get_legend_handles_labels() - lines2, lbl2 = ax_cpu.get_legend_handles_labels() - ax_res.legend(lines1 + lines2, lbl1 + lbl2, fontsize=7, loc="upper right", framealpha=0.8) - - plt.savefig(out_path, dpi=150, bbox_inches="tight") - print(f"Saved → {out_path}") - - -def parse_config(path: str) -> dict: - """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 {} - with open(path) as fh: - 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", default_replicas)), - "priority": int(wf.get("priority", 0)), - "min": int(wf.get("min_replicas", 0)), - "max": int(wf.get("max_replicas", 0)), - "deps": list(wf.get("dependencies", [])), - "dep_threshold": int(wf.get("dependency_threshold", 1)), - "cpus": int(wf.get("required_cpus", 0)), - "gpus": int(wf.get("required_gpus", 0)), - } - return group_meta - - -def _default_out(log_path: str) -> str: - stem = Path(log_path).stem - m = re.search(r"(\d+)", stem) - run_num = m.group(1) if m else stem - return f"cm_timeline_{run_num}.png" - - -def main(): - 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 (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) - - 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, - signal_events, - t0, - total_resources, - ) = parse_log(args.log) - - if args.config: - group_meta = parse_config(args.config) - print(f"Loaded group metadata from config: {args.config}") - else: - group_meta = group_meta_log - print("No config.yaml found — using group metadata parsed from log") - - print( - 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(signal_events)} signal events" - ) - plot( - spans, - group_meta, - resource_timeline, - gpu_assignments, - signal_events, - t0, - total_resources, - args.out, - ) - - -if __name__ == "__main__": - main() diff --git a/workflows/run_campaign/plots/adaptive_cm_timeline.png b/workflows/run_campaign/plots/adaptive_cm_timeline.png deleted file mode 100644 index ace6c86..0000000 Binary files a/workflows/run_campaign/plots/adaptive_cm_timeline.png and /dev/null differ diff --git a/workflows/run_campaign/plots/cm_timeline.png b/workflows/run_campaign/plots/cm_timeline.png deleted file mode 100644 index d72c737..0000000 Binary files a/workflows/run_campaign/plots/cm_timeline.png and /dev/null differ diff --git a/workflows/sgdes/delta_env_setup.sh b/workflows/sgdes/delta_env_setup.sh index 775c590..b873589 100755 --- a/workflows/sgdes/delta_env_setup.sh +++ b/workflows/sgdes/delta_env_setup.sh @@ -22,7 +22,7 @@ fi # ── Parse optional overrides ────────────────────────────────────────────────── ENV_DIR="${ENV_DIR:-/u/${USER}/ve/sgdes}" -SGDES_DIR="${SGDES_DIR:-/scratch/bblj/${USER}/SGDES}" +SGDES_DIR="${SGDES_DIR:-/scratch/bblj/${USER}/sgdes}" SPHERICAL_DIR="${SPHERICAL_DIR:-/scratch/bblj/${USER}/SPHERICAL}" while [[ $# -gt 0 ]]; do diff --git a/workflows/sgdes/delta_gpu_sbatch.sh b/workflows/sgdes/delta_gpu_batch.sh similarity index 63% rename from workflows/sgdes/delta_gpu_sbatch.sh rename to workflows/sgdes/delta_gpu_batch.sh index fa97007..10e74ef 100644 --- a/workflows/sgdes/delta_gpu_sbatch.sh +++ b/workflows/sgdes/delta_gpu_batch.sh @@ -1,13 +1,16 @@ #!/bin/sh -l +# +# SGDES Campaign — SLURM GPU batch script (Dragon backend) +# +# Account: set SBATCH_ACCOUNT=-delta-gpu before calling sbatch -#SBATCH -A ***-delta-gpu #SBATCH --partition=gpuA40x4 #SBATCH --nodes=2 #SBATCH --ntasks-per-node=1 #SBATCH --cpus-per-task=64 #SBATCH --gpus-per-node=4 #xSBATCH --exclusive -#SBATCH --time=01:30:00 +#SBATCH --time=00:30:00 #SBATCH --job-name=sgdes #SBATCH --mail-user=mariya.goliyad@rutgers.edu #SBATCH --mail-type=ALL @@ -18,10 +21,23 @@ export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH export TF_FORCE_GPU_ALLOW_GROWTH=true export JAX_PLATFORMS=cpu -export SGDES_DIR=/scratch/***/${USER}/SGDES -export SPHERICAL_DIR=/scratch/***/${USER}/SPHERICAL +# ── Environment ─────────────────────────────────────────────────────────────── +if [ -z "${SBATCH_ACCOUNT:-}${SLURM_JOB_ACCOUNT:-}" ]; then + echo "WARNING: SBATCH_ACCOUNT is not set — job may be charged to default account." + echo " Set it with: export SBATCH_ACCOUNT=-delta-gpu" +fi +echo "Account: ${SLURM_JOB_ACCOUNT:-unknown}" + +if [ -z "${SCRATCH:-}" ]; then + echo "ERROR: SCRATCH is not set." + echo " export SCRATCH=/scratch/ && sbatch delta_gpu_batch.sh" + exit 1 +fi + +export SGDES_DIR="${SGDES_DIR:-${SCRATCH}/${USER}/sgdes}" +export SPHERICAL_DIR="${SPHERICAL_DIR:-${SCRATCH}/${USER}/SPHERICAL}" -cd $SPHERICAL_DIR/examples/sgdes +cd $SPHERICAL_DIR/workflows/sgdes # ── Clean previous run artifacts ────────────────────────────────────────────── rm -rf mayv_output tmp* diff --git a/workflows/sgdes/sgdes_workflow.py b/workflows/sgdes/sgdes_workflow.py index 46b7b2f..ae29b59 100644 --- a/workflows/sgdes/sgdes_workflow.py +++ b/workflows/sgdes/sgdes_workflow.py @@ -6,15 +6,17 @@ Task types ---------- -executable_task — trill embed/fold, foldseek createdb, seqkit grep/stats; - Dragon launches each as a subprocess with GPU affinity and - HOST_NAME placement via task_description. -function_task — foldseek_search only; runs via subprocess.run to avoid the - ggml-CUDA context conflict that occurs when foldseek easy-search - runs as a direct Dragon executable_task subprocess. -_run_des — plain async method on SGDESWorkflow; orchestrates the DES - loop (solver.propose → foldseek scoring → population.add_samples) - using the registered tasks. +function_task — all shell commands (trill embed/fold, foldseek createdb/search, + seqkit grep/stats) run via subprocess.run() inside a Dragon + function_task. This avoids the Dragon executable_task + completion-delivery deadlock: the subprocess runs successfully + but asyncflow's TaskCompleted event never fires, permanently + blocking the await. The same GPU/CUDA context conflict that + was documented for foldseek_search also applies to trill and + foldseek createdb. +_run_des — plain async method on SGDESWorkflow; orchestrates the DES + loop (solver.propose → foldseek scoring → population.add_samples) + using the registered tasks. All environment variables (CUDA_HOME, SGDES_DIR, SPHERICAL_DIR, JAX_PLATFORMS, TF_FORCE_GPU_ALLOW_GROWTH) are set in the sbatch script. @@ -34,9 +36,10 @@ from datetime import datetime # amortized_bo imports JAX at module level, making this process multithreaded. -# Dragon then uses os.fork() to spawn executable_task subprocesses (trill embed), -# which triggers Python's fork-after-threads warning. The warning is harmless — -# trill runs in its own fresh subprocess and completes normally. +# Dragon's function_task workers inherit this state and may trigger the +# fork-after-threads warning when Dragon internally forks to run a worker. +# The warning is suppressed because it is harmless: subprocess.run() inside +# each function_task creates its own clean subprocess via the shell. warnings.filterwarnings( "ignore", message="os.fork\\(\\) was called.*JAX is multithreaded", @@ -226,29 +229,47 @@ def _register_tasks(self, policy): _TD_HOST = {} # noqa: N806 # ── trill embed ──────────────────────────────────────────────────────── - @flow.executable_task + # NOTE: function_task (not executable_task) to avoid the Dragon executable_task + # completion-delivery deadlock: the subprocess runs successfully but + # asyncflow's TaskCompleted event never fires, blocking the await forever. + # The same issue affects all GPU/CUDA tasks; subprocess.run() sidesteps it. + @flow.function_task async def embed(task_description=_TD_GPU, **kwargs): - """Run trill embed esm2_t33_650M as an executable_task. + """Run trill embed esm2_t33_650M via subprocess.run (function_task). kwargs: name, GPUs, seed, outdir, query """ + import subprocess as _sp + name = kwargs["name"] gpus = kwargs["GPUs"] seed = kwargs["seed"] outdir = kwargs["outdir"] query = kwargs["query"] cmd = ( + f"env -u SLURM_NTASKS -u SLURM_PROCID -u SLURM_NODEID -u SLURM_LOCALID " f"trill {name} {gpus} --RNG_seed {seed} --outdir {outdir} " f"embed esm2_t33_650M {query} --avg" ) print(f"[embed] cmd: {cmd}", flush=True) - return cmd + result = _sp.run(cmd, shell=True, capture_output=True, text=True) + if result.stdout: + print(result.stdout, flush=True) + if result.stderr: + print(result.stderr, flush=True) + if result.returncode != 0: + raise RuntimeError( + f"trill embed failed (rc={result.returncode}): {result.stderr[-500:]}" + ) # ── trill fold (ESMFold, slow path only) ─────────────────────────────── - @flow.executable_task + # NOTE: function_task for the same reason as embed above. + @flow.function_task async def fold(task_description=_TD_GPU, **kwargs): - """Run trill fold ESMFold as an executable_task. + """Run trill fold ESMFold via subprocess.run (function_task). kwargs: name, GPUs, seed, outdir, query, batch_size """ + import subprocess as _sp + name = kwargs["name"] gpus = kwargs["GPUs"] seed = kwargs["seed"] @@ -256,18 +277,30 @@ async def fold(task_description=_TD_GPU, **kwargs): query = kwargs["query"] batch_size = kwargs["batch_size"] cmd = ( + f"env -u SLURM_NTASKS -u SLURM_PROCID -u SLURM_NODEID -u SLURM_LOCALID " f"trill {name} {gpus} --RNG_seed {seed} --outdir {outdir} " f"fold ESMFold {query} --batch_size {batch_size}" ) print(f"[fold] cmd: {cmd}", flush=True) - return cmd + result = _sp.run(cmd, shell=True, capture_output=True, text=True) + if result.stdout: + print(result.stdout, flush=True) + if result.stderr: + print(result.stderr, flush=True) + if result.returncode != 0: + raise RuntimeError( + f"trill fold failed (rc={result.returncode}): {result.stderr[-500:]}" + ) # ── foldseek createdb ───────────────────────────────────────────────── - @flow.executable_task + # NOTE: function_task for the same reason as embed above (GPU/ProstT5 path). + @flow.function_task async def foldseek_createdb(task_description=_TD_GPU, **kwargs): - """Run foldseek createdb as an executable_task. + """Run foldseek createdb via subprocess.run (function_task). kwargs: fasta, db_path, prostt5_model (optional) """ + import subprocess as _sp + fasta = kwargs["fasta"] db_path = kwargs["db_path"] prostt5_model = kwargs.get("prostt5_model", "") @@ -275,7 +308,15 @@ async def foldseek_createdb(task_description=_TD_GPU, **kwargs): gpu_flag = " --gpu 1" if prostt5_model else "" cmd = f"foldseek createdb {fasta} {db_path} {model_flag}{gpu_flag}".strip() print(f"[foldseek_createdb] cmd: {cmd}", flush=True) - return cmd + result = _sp.run(cmd, shell=True, capture_output=True, text=True) + if result.stdout: + print(result.stdout, flush=True) + if result.stderr: + print(result.stderr, flush=True) + if result.returncode != 0: + raise RuntimeError( + f"foldseek createdb failed (rc={result.returncode}): {result.stderr[-500:]}" + ) # ── foldseek easy-search ────────────────────────────────────────────── # NOTE: function_task (not executable_task) to avoid the ggml-CUDA context @@ -314,28 +355,59 @@ async def foldseek_search(task_description=_TD_GPU, **kwargs): f"foldseek easy-search failed (rc={result.returncode}): {result.stderr[-500:]}" ) - # ── seqkit grep → file (stdout redirected via ProcessTemplate) ─────── - @flow.executable_task + # ── seqkit grep → writes matching FASTA to output_fasta, returns path ── + # Dragon V3's return_value channel does not preserve large strings + # (the future resolves to an int instead). Writing directly to the + # caller-supplied output path sidesteps that channel entirely. + @flow.function_task async def seqkit_grep(task_description=_TD_HOST, **kwargs): - """Run seqkit grep; Dragon writes stdout to output_fasta via task_description. - kwargs: pattern_file, input_fasta + """Run seqkit grep; writes stdout (FASTA) to kwargs['output_fasta']. + kwargs: pattern_file, input_fasta, output_fasta + Returns the output path so the caller can confirm the file exists. """ + import subprocess as _sp + pattern_file = kwargs["pattern_file"] input_fasta = kwargs["input_fasta"] + output_fasta = kwargs["output_fasta"] cmd = f"seqkit grep --pattern-file {pattern_file} {input_fasta}" print(f"[seqkit_grep] cmd: {cmd}", flush=True) - return cmd - - # ── seqkit stats → returns TSV via stdout ───────────────────────────── - @flow.executable_task + result = _sp.run(cmd, shell=True, capture_output=True, text=True) + if result.stderr: + print(result.stderr, flush=True) + if result.returncode != 0: + raise RuntimeError( + f"seqkit grep failed (rc={result.returncode}): {result.stderr[-500:]}" + ) + with open(output_fasta, "w") as _f: + _f.write(result.stdout) + return output_fasta + + # ── seqkit stats → writes TSV to temp file, returns path ───────────── + # Dragon V3's return_value channel does not preserve large strings + # (the future resolves to an int instead). Writing to a temp file + # and returning the path sidesteps that channel entirely. + @flow.function_task async def seqkit_stats(task_description=_TD_HOST, **kwargs): - """Run seqkit stats -a -T; stdout is returned by asyncflow as a string. + """Run seqkit stats -a -T; writes TSV to a temp file, returns path. kwargs: input_fasta """ + import subprocess as _sp + import tempfile as _tf + input_fasta = kwargs["input_fasta"] cmd = f"seqkit stats -a -T {input_fasta}" print(f"[seqkit_stats] cmd: {cmd}", flush=True) - return cmd + result = _sp.run(cmd, shell=True, capture_output=True, text=True) + if result.stderr: + print(result.stderr, flush=True) + if result.returncode != 0: + raise RuntimeError( + f"seqkit stats failed (rc={result.returncode}): {result.stderr[-500:]}" + ) + with _tf.NamedTemporaryFile(mode="w", suffix=".stats.tsv", delete=False) as _tmp: + _tmp.write(result.stdout) + return _tmp.name return types.SimpleNamespace( embed=embed, @@ -654,7 +726,21 @@ async def _sgdes_async(self, mutation: str, policy) -> None: ) logger.info(f"[{mutation}] Fold done ({time.time() - t0:.1f}s)") - stats_out = await tasks.seqkit_stats(input_fasta=query) + # Run seqkit stats directly (not via Dragon) — the Dragon + # function_task return-value channel hangs on string results. + # seqkit stats is host-pinned, fast (<1 s), and needs no GPU. + _stats_proc = await asyncio.create_subprocess_shell( + f"seqkit stats -a -T {query}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _stats_stdout, _stats_stderr = await _stats_proc.communicate() + if _stats_proc.returncode != 0: + raise RuntimeError( + f"seqkit stats failed (rc={_stats_proc.returncode}): " + f"{_stats_stderr.decode()[-500:]}" + ) + stats_out = _stats_stdout.decode() df_stats = pd.read_csv(io.StringIO(stats_out), sep="\t") median = df_stats.Q2.values logger.info(f"[{mutation}] Sequence length median={median[0]} (from {query})") @@ -840,13 +926,21 @@ async def _sgdes_async(self, mutation: str, policy) -> None: # } # ) - res = await tasks.seqkit_grep( - # task_description=_grep_td, - pattern_file=labels_file, - input_fasta=cleaned, + # Run seqkit grep directly (not via Dragon) for the same reason + # as seqkit stats above: the Dragon return-value channel hangs. + _grep_proc = await asyncio.create_subprocess_shell( + f"seqkit grep --pattern-file {labels_file} {cleaned}", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) - with open(output_fasta, "w+") as output_file: - output_file.write(res) + _grep_stdout, _grep_stderr = await _grep_proc.communicate() + if _grep_proc.returncode != 0: + raise RuntimeError( + f"seqkit grep failed (rc={_grep_proc.returncode}): " + f"{_grep_stderr.decode()[-500:]}" + ) + with open(output_fasta, "w") as _gf: + _gf.write(_grep_stdout.decode()) # ── Save per-sequence metrics ────────────────────────────────────── tmp_input_df = pd.read_csv(