diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index bf1bf8c..d5f6b12 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -12,7 +12,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11"] + python-version: ["3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v4 diff --git a/README.md b/README.md index 9622ddb..e12e6e0 100644 --- a/README.md +++ b/README.md @@ -4,15 +4,13 @@ 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 - **Dragon/Asyncflow Integration** — HPC runtime for distributed multi-node execution via DragonHPC - **Automatic Device Detection** — CUDA GPUs if available, CPU fallback - **YAML Config with Env-Var Expansion** — `${VAR}` references in config files are resolved at load time -- **Telemetry & Visualization** — asyncflow native JSONL telemetry with workflow dashboard plots; campaign replica timeline and resource utilization charts from SLURM logs +- **Telemetry & Visualization** — asyncflow native JSONL telemetry with workflow dashboard plots --- @@ -35,15 +33,6 @@ spherical/ │ ├── esm2_inference/ # Standalone ESM2 inference runner │ │ ├── run_esm2_infern.py │ │ └── 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 @@ -106,53 +95,6 @@ engine: concurrent # "concurrent" or "dragon" service_python: "${VE_HOME}/esm2/bin/python" # resolved at load time ``` -### Multi-workflow Campaign - -Two campaigns ship under `workflows/run_campaign/`: - -```bash -# 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: - -```yaml -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 - concurrency_floor: 2 - concurrency_cap: 4 - priority: 5 - required_cpus: 20 - dependencies: [] - - inference: - replicas: 16 - concurrency_floor: 1 - concurrency_cap: 4 - priority: 10 - required_cpus: 32 - required_gpus: 1 - dependencies: [ddsim] -``` - ### SGDES Protein Engineering ```bash @@ -230,134 +172,6 @@ bash workflows/plot_telemetry.sh \ --out-dir plots/sgdes ``` -### Campaign Manager replica timeline - -`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 slurm-.out \ - [--config workflows/run_campaign/esm2_ddsim_campaign/config.yaml] \ - [--out timeline.png] -``` - -| Output element | Description | -|----------------|-------------| -| Gantt chart | One bar per replica, coloured by workflow group; red border = error; dependency arrows show signal-ready flow | -| Resource panel | Step plot of GPU and CPU slots in use over elapsed time (from scheduler log lines) | -| Config table | Replicas, priority, CPU/GPU requirements, min/max, and dependency graph per group | - -`config.yaml` is auto-detected when it sits next to the log file. If found, -group metadata is taken from the config (authoritative); otherwise it is parsed -from the log lines. - -**Example**: -```bash -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/workflows/sgdes/apply_slurm_patch.py b/workflows/sgdes/apply_slurm_patch.py deleted file mode 100644 index b719c4c..0000000 --- a/workflows/sgdes/apply_slurm_patch.py +++ /dev/null @@ -1,113 +0,0 @@ -""" -Patch Dragon's slurm.py safely: - -- Replaces --ntasks with --ntasks-per-node=1 -- Adds --overlap -- Backs up original slurm.py -- Deletes __pycache__ to force recompilation -""" - -import os -import shutil - - -def patch_slurm_file(): - try: - import dragon.launcher.wlm.slurm as slurm_module - - slurm_path = slurm_module.__file__ - except ImportError: - print("[ERROR] Could not find Dragon slurm module.") - return - - # Backup original - backup_path = slurm_path + ".bak" - if not os.path.exists(backup_path): - shutil.copy(slurm_path, backup_path) - print(f"[INFO] Backup created at {backup_path}") - else: - print(f"[INFO] Backup already exists at {backup_path}") - - # Read original file - with open(slurm_path) as f: - lines = f.readlines() - - new_lines = [] - inside_get_wlm = False - - for line in lines: - stripped = line.strip() - - # Patch SRUN_COMMAND_LINE inside class - if stripped.startswith("SRUN_COMMAND_LINE"): - indent = line[: line.find("SRUN_COMMAND_LINE")] - new_lines.append( - indent - + "SRUN_COMMAND_LINE = (\n" - + indent - + ' "srun "\n' - + indent - + ' "--nodes={nnodes} "\n' - + indent - + ' "--ntasks-per-node=1 "\n' - + indent - + ' "--cpu_bind=none "\n' - + indent - + ' "--overlap "\n' - + indent - + ' "-u -l -W 0"\n' - + indent - + ")\n" - ) - continue - - # Patch _get_wlm_launch_be_args method - if stripped.startswith("def _get_wlm_launch_be_args"): - inside_get_wlm = True - indent = line[: line.find("def")] - new_lines.append(line) # keep method definition line - # insert new method body - method_indent = indent + " " * 4 - new_lines.append(method_indent + "slurm_launch_be_args = [\n") - new_lines.append(method_indent + ' "srun",\n') - new_lines.append(method_indent + " f\"--nodes={args_map['nnodes']}\",\n") - new_lines.append(method_indent + ' "--ntasks-per-node=1",\n') - new_lines.append(method_indent + ' "--cpu_bind=none",\n') - new_lines.append(method_indent + ' "--overlap",\n') - new_lines.append(method_indent + "]\n") - new_lines.append(method_indent + "return slurm_launch_be_args + launch_args\n") - continue - - # Skip original method body lines - if inside_get_wlm: - if ( - stripped == "" - or stripped.startswith("return") - or stripped.startswith("slurm_launch_be_args") - ): - continue - if stripped == "": - inside_get_wlm = False - continue - continue - - # Keep other lines - new_lines.append(line) - - # Write patched file - with open(slurm_path, "w") as f: - f.writelines(new_lines) - - print(f"[SUCCESS] Dragon slurm.py patched at {slurm_path}") - - # Delete __pycache__ folder to force recompilation - pycache_dir = os.path.join(os.path.dirname(slurm_path), "__pycache__") - if os.path.exists(pycache_dir): - shutil.rmtree(pycache_dir) - print(f"[INFO] Removed __pycache__ at {pycache_dir}") - else: - print(f"[INFO] No __pycache__ found at {pycache_dir}") - - -if __name__ == "__main__": - patch_slurm_file() diff --git a/workflows/sgdes/delta_gpu_batch.sh b/workflows/sgdes/delta_gpu_batch.sh index 10e74ef..8a72bec 100644 --- a/workflows/sgdes/delta_gpu_batch.sh +++ b/workflows/sgdes/delta_gpu_batch.sh @@ -5,10 +5,10 @@ # Account: set SBATCH_ACCOUNT=-delta-gpu before calling sbatch #SBATCH --partition=gpuA40x4 -#SBATCH --nodes=2 +#SBATCH --nodes=1 #SBATCH --ntasks-per-node=1 #SBATCH --cpus-per-task=64 -#SBATCH --gpus-per-node=4 +#SBATCH --gpus-per-node=1 #xSBATCH --exclusive #SBATCH --time=00:30:00 #SBATCH --job-name=sgdes