diff --git a/CLAUDE.md b/CLAUDE.md index 9c66af3c..0b12b53e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -101,10 +101,12 @@ export TORC_API_URL="http://localhost:8080/torc-service/v1" # Quick workflow execution (convenience commands) ./target/release/torc run examples/sample_workflow.yaml # Create and run locally -./target/release/torc submit examples/sample_workflow.yaml # Create and submit to scheduler +./target/release/torc submit examples/sample_workflow.yaml # Submit (requires scheduler actions) +./target/release/torc submit-slurm --account myproject examples/sample_workflow.yaml # Auto-generate Slurm schedulers # Or use explicit workflow management ./target/release/torc workflows create examples/sample_workflow.yaml +./target/release/torc workflows create-slurm --account myproject examples/sample_workflow.yaml # With Slurm schedulers ./target/release/torc workflows submit # Submit to scheduler ./target/release/torc workflows run # Run locally @@ -287,12 +289,15 @@ List endpoints support `offset` and `limit` query parameters: 5. Launch interactive UI: `torc tui` ### Submitting a Workflow to Scheduler -**Quick method:** -- `torc submit ` - Create from spec and submit to scheduler (requires on_workflow_start/schedule_nodes action) +**Quick method (Slurm with auto-generated schedulers):** +- `torc submit-slurm --account ` - Auto-generate Slurm schedulers, create workflow, and submit + +**Quick method (pre-configured schedulers):** +- `torc submit ` - Create from spec and submit (requires on_workflow_start/schedule_nodes action in spec) - `torc submit ` - Submit existing workflow to scheduler **Explicit method:** -1. Create workflow: `torc workflows create ` +1. Create workflow: `torc workflows create ` or `torc workflows create-slurm --account ` 2. Submit workflow: `torc workflows submit ` ### Debugging @@ -337,10 +342,12 @@ sqlite3 server/db/sqlite/dev.db **Quick Workflow Execution** (convenience commands): - `torc run ` - Create from spec and run locally, or run existing workflow -- `torc submit ` - Create from spec and submit to scheduler, or submit existing workflow +- `torc submit ` - Submit workflow to scheduler (requires pre-configured scheduler actions) +- `torc submit-slurm --account ` - Auto-generate Slurm schedulers and submit **Workflow Management**: - `torc workflows create ` - Create workflow from specification +- `torc workflows create-slurm --account ` - Create workflow with auto-generated Slurm schedulers - `torc workflows new` - Create empty workflow interactively - `torc workflows list` - List all workflows - `torc workflows submit ` - Submit workflow to scheduler (requires on_workflow_start/schedule_nodes action) @@ -349,6 +356,14 @@ sqlite3 server/db/sqlite/dev.db - `torc workflows status ` - Check workflow status - `torc workflows cancel ` - Cancel workflow +**Slurm/HPC Commands**: +- `torc slurm generate --account ` - Generate Slurm schedulers for a workflow +- `torc hpc list` - List available HPC profiles +- `torc hpc detect` - Detect current HPC system +- `torc hpc show ` - Show profile details +- `torc hpc partitions ` - List partitions for a profile +- `torc hpc match --cpus N --memory SIZE` - Find matching partitions + **Job Management**: - `torc jobs list ` - List jobs for workflow - `torc jobs get ` - Get job details @@ -357,6 +372,7 @@ sqlite3 server/db/sqlite/dev.db **Execution**: - `torc run ` - Run workflow locally (top-level command) - `torc submit ` - Submit workflow to scheduler (top-level command) +- `torc submit-slurm --account ` - Submit with auto-generated Slurm schedulers - `torc tui` - Interactive terminal UI **Utilities**: diff --git a/docs/book.toml b/docs/book.toml index e06615d0..0f87dc29 100644 --- a/docs/book.toml +++ b/docs/book.toml @@ -12,6 +12,7 @@ create-missing = true [output.html] default-theme = "light" preferred-dark-theme = "navy" +favicon = "images/torc_icon_256.png" git-repository-url = "https://github.com/NREL/torc" edit-url-template = "https://github.com/NREL/torc/edit/main/docs/{path}" site-url = "/torc/" diff --git a/docs/src/SUMMARY.md b/docs/src/SUMMARY.md index 1d566b9e..1477c8ce 100644 --- a/docs/src/SUMMARY.md +++ b/docs/src/SUMMARY.md @@ -7,6 +7,8 @@ - [Getting Started](./getting-started.md) - [Installation](./installation.md) - [Quick Start](./quick-start.md) + - [Quick Start (Local)](./quick-start-local.md) + - [Quick Start (HPC)](./quick-start-hpc.md) # Understanding Torc @@ -21,6 +23,7 @@ - [Dependency Resolution](./explanation/dependencies.md) - [Parallelization Strategies](./explanation/parallelization.md) - [Workflow Actions](./explanation/workflow-actions.md) + - [Slurm Workflows](./explanation/slurm-workflows.md) - [Design](./explanation/design/README.md) - [Server API Handler](./explanation/design/server.md) - [Central Database](./explanation/design/database.md) @@ -32,10 +35,12 @@ - [Configuration Files](./how-to/configuration-files.md) - [Creating Workflows](./how-to/creating-workflows.md) - [Working with Slurm](./how-to/slurm.md) + - [HPC Profiles](./how-to/hpc-profiles.md) - [Job Checkpointing](./how-to/checkpointing.md) - [Resource Monitoring](./how-to/resource-monitoring.md) - [Terminal UI (TUI)](./how-to/tui.md) - [Web Dashboard](./how-to/dashboard.md) + - [Visualizing Workflow Structure](./how-to/visualizing-workflows.md) - [Debugging Workflows](./how-to/debugging.md) - [Debugging Slurm Workflows](./how-to/debugging-slurm.md) - [Authentication](./how-to/authentication.md) @@ -49,6 +54,7 @@ - [Workflow Specification Formats](./reference/workflow-formats.md) - [Job Parameterization](./reference/parameterization.md) - [Resource Requirements](./reference/resources.md) + - [HPC Profiles](./reference/hpc-profiles.md) - [Resource Monitoring](./reference/resource-monitoring.md) - [OpenAPI Specification](./reference/openapi.md) - [Configuration](./reference/configuration.md) @@ -68,6 +74,7 @@ - [Multi-Stage Workflows with Barriers](./tutorials/multi-stage-barrier.md) - [Map Python functions across workers](./tutorials/map_python_function_across_workers.md) - [Filtering CLI Output with Nushell](./tutorials/filtering-with-nushell.md) + - [Custom HPC Profile](./tutorials/custom-hpc-profile.md) --- diff --git a/docs/src/explanation/slurm-workflows.md b/docs/src/explanation/slurm-workflows.md new file mode 100644 index 00000000..7d206ed7 --- /dev/null +++ b/docs/src/explanation/slurm-workflows.md @@ -0,0 +1,433 @@ +# Slurm Workflows + +This document explains how Torc simplifies running workflows on Slurm-based HPC systems. The key insight is that **you don't need to understand Slurm schedulers or workflow actions** to run workflows on HPC systems—Torc handles this automatically. + +## The Simple Approach + +Running a workflow on Slurm requires just two things: + +1. **Define your jobs with resource requirements** +2. **Submit with `submit-slurm`** + +That's it. Torc will analyze your workflow, generate appropriate Slurm configurations, and submit everything for execution. + +> **⚠️ Important:** The `submit-slurm` command uses heuristics to auto-generate Slurm schedulers and workflow actions. For complex workflows with unusual dependency patterns, the generated configuration may not be optimal and could result in suboptimal allocation timing. **Always preview the configuration first** using `torc slurm generate` (see [Previewing Generated Configuration](#previewing-generated-configuration)) before submitting production workflows. + +### Example Workflow + +Here's a complete workflow specification that runs on Slurm: + +```yaml +name: data_analysis_pipeline +description: Analyze experimental data with preprocessing, training, and evaluation + +resource_requirements: + - name: light + num_cpus: 4 + memory: 8g + runtime: PT30M + + - name: compute + num_cpus: 32 + memory: 64g + runtime: PT2H + + - name: gpu + num_cpus: 16 + num_gpus: 2 + memory: 128g + runtime: PT4H + +jobs: + - name: preprocess + command: python preprocess.py --input data/ --output processed/ + resource_requirements: light + + - name: train_model + command: python train.py --data processed/ --output model/ + resource_requirements: gpu + depends_on: [preprocess] + + - name: evaluate + command: python evaluate.py --model model/ --output results/ + resource_requirements: compute + depends_on: [train_model] + + - name: generate_report + command: python report.py --results results/ + resource_requirements: light + depends_on: [evaluate] +``` + +### Submitting the Workflow + +```bash +torc submit-slurm --account myproject workflow.yaml +``` + +Torc will: +1. Detect which HPC system you're on (e.g., NREL Kestrel) +2. Match each job's requirements to appropriate partitions +3. Generate Slurm scheduler configurations +4. Create workflow actions that stage resource allocation based on dependencies +5. Submit the workflow for execution + +## How It Works + +When you use `submit-slurm`, Torc performs intelligent analysis of your workflow: + +### 1. Per-Job Scheduler Generation + +Each job gets its own Slurm scheduler configuration based on its resource requirements. This means: +- Jobs are matched to the most appropriate partition +- Memory, CPU, and GPU requirements are correctly specified +- Walltime is set to the partition's maximum (explained below) + +### 2. Staged Resource Allocation + +Torc analyzes job dependencies and creates **staged workflow actions**: + +- **Jobs without dependencies** trigger `on_workflow_start` — resources are allocated immediately +- **Jobs with dependencies** trigger `on_jobs_ready` — resources are allocated only when the job becomes ready to run + +This prevents wasting allocation time on resources that aren't needed yet. For example, in the workflow above: +- `preprocess` resources are allocated at workflow start +- `train_model` resources are allocated when `preprocess` completes +- `evaluate` resources are allocated when `train_model` completes +- `generate_report` resources are allocated when `evaluate` completes + +### 3. Conservative Walltime + +Torc sets the walltime to the **partition's maximum** rather than your job's estimated runtime. This provides: +- Headroom for jobs that run slightly longer than expected +- No additional cost since Torc workers exit when work completes +- Protection against job termination due to tight time limits + +For example, if your job requests 3 hours and matches the "short" partition (4 hours max), the allocation will request 4 hours. + +### 4. HPC Profile Knowledge + +Torc includes built-in knowledge of HPC systems like NREL Kestrel, including: +- Available partitions and their resource limits +- GPU configurations +- Memory and CPU specifications +- Special requirements (e.g., minimum node counts for high-bandwidth partitions) + +> **Using an unsupported HPC?** Please [request built-in support](https://github.com/NREL/torc/issues) so everyone benefits. You can also [create a custom profile](../tutorials/custom-hpc-profile.md) for immediate use. + +## Resource Requirements Specification + +Resource requirements are the key to the simplified workflow. Define them once and reference them from jobs: + +```yaml +resource_requirements: + - name: small + num_cpus: 4 + num_gpus: 0 + num_nodes: 1 + memory: 8g + runtime: PT1H + + - name: gpu_training + num_cpus: 32 + num_gpus: 4 + num_nodes: 1 + memory: 256g + runtime: PT8H +``` + +### Fields + +| Field | Description | Example | +|-------|-------------|---------| +| `name` | Reference name for jobs | `"compute"` | +| `num_cpus` | CPU cores required | `32` | +| `num_gpus` | GPUs required (0 if none) | `2` | +| `num_nodes` | Nodes required | `1` | +| `memory` | Memory with unit suffix | `"64g"`, `"512m"` | +| `runtime` | ISO8601 duration | `"PT2H"`, `"PT30M"` | + +### Runtime Format + +Use ISO8601 duration format: +- `PT30M` — 30 minutes +- `PT2H` — 2 hours +- `PT1H30M` — 1 hour 30 minutes +- `P1D` — 1 day +- `P2DT4H` — 2 days 4 hours + +## Job Dependencies + +Define dependencies explicitly or implicitly through file/data relationships: + +### Explicit Dependencies + +```yaml +jobs: + - name: step1 + command: ./step1.sh + resource_requirements: small + + - name: step2 + command: ./step2.sh + resource_requirements: small + depends_on: [step1] + + - name: step3 + command: ./step3.sh + resource_requirements: small + depends_on: [step1, step2] # Waits for both +``` + +### Implicit Dependencies (via Files) + +```yaml +files: + - name: raw_data + path: /data/raw.csv + - name: processed_data + path: /data/processed.csv + +jobs: + - name: process + command: python process.py + input_files: [raw_data] + output_files: [processed_data] + resource_requirements: compute + + - name: analyze + command: python analyze.py + input_files: [processed_data] # Creates implicit dependency on 'process' + resource_requirements: compute +``` + +## Previewing Generated Configuration + +> **Recommended Practice:** Always preview the generated configuration before submitting to Slurm, especially for complex workflows. This allows you to verify that schedulers and actions are appropriate for your workflow structure. + +### Viewing the Execution Plan + +Before generating schedulers, visualize how your workflow will execute in stages: + +```bash +torc workflows execution-plan workflow.yaml +``` + +This shows the execution stages, which jobs run at each stage, and (if schedulers are defined) when Slurm allocations are requested. See [Visualizing Workflow Structure](../how-to/visualizing-workflows.md) for detailed examples. + +### Generating Slurm Configuration + +Preview what Torc will generate: + +```bash +torc slurm generate --account myproject --profile kestrel workflow.yaml +``` + +This outputs the complete workflow with generated schedulers and actions: + +```yaml +name: data_analysis_pipeline +# ... original content ... + +jobs: + - name: preprocess + command: python preprocess.py --input data/ --output processed/ + resource_requirements: light + scheduler: preprocess_scheduler + + # ... more jobs ... + +slurm_schedulers: + - name: preprocess_scheduler + account: myproject + mem: 8g + nodes: 1 + walltime: "04:00:00" + + - name: train_model_scheduler + account: myproject + mem: 128g + nodes: 1 + gres: "gpu:2" + walltime: "04:00:00" + + # ... more schedulers ... + +actions: + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: preprocess_scheduler + scheduler_type: slurm + num_allocations: 1 + + - trigger_type: on_jobs_ready + action_type: schedule_nodes + jobs: [train_model] + scheduler: train_model_scheduler + scheduler_type: slurm + num_allocations: 1 + + # ... more actions ... +``` + +Save the output to inspect or modify before submission: + +```bash +torc slurm generate --account myproject workflow.yaml -o workflow_with_schedulers.yaml +``` + +## Torc Server Considerations + +The Torc server must be accessible to compute nodes. Options include: + +1. **Shared server** (Recommended): A team member allocates a dedicated server in the HPC environment +2. **Login node**: Suitable for small workflows with few, long-running jobs + +For large workflows with many short jobs, a dedicated server prevents overloading login nodes. + +## Best Practices + +### 1. Focus on Resource Requirements + +Spend time accurately defining resource requirements. Torc handles the rest: + +```yaml +resource_requirements: + # Be specific about what each job type needs + - name: io_heavy + num_cpus: 4 + memory: 32g # High memory for data loading + runtime: PT1H + + - name: compute_heavy + num_cpus: 64 + memory: 16g # Less memory, more CPU + runtime: PT4H +``` + +### 2. Use Meaningful Names + +Name resource requirements by their purpose, not by partition: + +```yaml +# Good - describes the workload +resource_requirements: + - name: data_preprocessing + - name: model_training + - name: inference + +# Avoid - ties you to specific infrastructure +resource_requirements: + - name: short_partition + - name: gpu_h100 +``` + +### 3. Group Similar Jobs + +Jobs with similar requirements can share resource requirement definitions: + +```yaml +resource_requirements: + - name: quick_task + num_cpus: 2 + memory: 4g + runtime: PT15M + +jobs: + - name: validate_input + command: ./validate.sh + resource_requirements: quick_task + + - name: check_output + command: ./check.sh + resource_requirements: quick_task + depends_on: [main_process] +``` + +### 4. Test Locally First + +Validate your workflow logic locally before submitting to HPC: + +```bash +# Run locally (without Slurm) +torc run workflow.yaml + +# Then submit to HPC +torc submit-slurm --account myproject workflow.yaml +``` + +## Limitations and Caveats + +The auto-generation in `submit-slurm` uses heuristics that work well for common workflow patterns but may not be optimal for all cases: + +### When Auto-Generation Works Well + +- **Linear pipelines**: A → B → C → D +- **Fan-out patterns**: One job unblocks many (e.g., preprocess → 100 work jobs) +- **Fan-in patterns**: Many jobs unblock one (e.g., 100 work jobs → postprocess) +- **Simple DAGs**: Clear dependency structures with distinct resource tiers + +### When to Use Manual Configuration + +Consider using `torc slurm generate` to preview and manually adjust, or define schedulers manually, when: + +- **Complex dependency graphs**: Multiple interleaved dependency patterns +- **Shared schedulers**: You want multiple jobs to share the same Slurm allocation +- **Custom timing**: Specific requirements for when allocations should be requested +- **Resource optimization**: Fine-tuning to minimize allocation waste +- **Multi-node jobs**: Jobs requiring coordination across multiple nodes + +### What Could Go Wrong + +Without previewing, auto-generation might: + +1. **Request allocations too early**: Wasting queue time waiting for dependencies +2. **Request allocations too late**: Adding latency to job startup +3. **Create suboptimal scheduler groupings**: Not sharing allocations when beneficial +4. **Miss optimization opportunities**: Not recognizing patterns that could share resources + +**Best Practice**: For production workflows, always run `torc slurm generate` first, review the output, and submit the reviewed configuration with `torc submit`. + +## Advanced: Manual Scheduler Configuration + +For advanced users who need fine-grained control, you can define schedulers and actions manually. See [Working with Slurm](../how-to/slurm.md) for details. + +Common reasons for manual configuration: +- Non-standard partition requirements +- Custom Slurm directives (e.g., `--constraint`) +- Multi-node jobs with specific topology requirements +- Reusing allocations across multiple jobs for efficiency + +## Troubleshooting + +### "No partition found for job" + +Your resource requirements exceed what's available. Check: +- Memory doesn't exceed partition limits +- Runtime doesn't exceed partition walltime +- GPU count is available on GPU partitions + +Use `torc hpc partitions ` to see available resources. + +### Jobs Not Starting + +Ensure the Torc server is accessible from compute nodes: +```bash +# From a compute node +curl $TORC_API_URL/health +``` + +### Wrong Partition Selected + +Use `torc hpc match` to see which partitions match your requirements: +```bash +torc hpc match kestrel --cpus 32 --memory 64g --walltime 2h --gpus 2 +``` + +## See Also + +- [Visualizing Workflow Structure](../how-to/visualizing-workflows.md) — Execution plans and DAG visualization +- [HPC Profiles](../how-to/hpc-profiles.md) — Detailed HPC profile usage +- [Working with Slurm](../how-to/slurm.md) — Advanced Slurm configuration +- [Resource Requirements Reference](../reference/resources.md) — Complete specification +- [Workflow Actions](./workflow-actions.md) — Understanding actions diff --git a/docs/src/how-to/hpc-profiles.md b/docs/src/how-to/hpc-profiles.md new file mode 100644 index 00000000..1261c7de --- /dev/null +++ b/docs/src/how-to/hpc-profiles.md @@ -0,0 +1,157 @@ +# Working with HPC Profiles + +HPC (High-Performance Computing) profiles provide pre-configured knowledge about specific HPC systems, including their partitions, resource limits, and optimal settings. Torc uses this information to automatically match job requirements to appropriate partitions. + +## Overview + +HPC profiles contain: +- **Partition definitions**: Available queues with their resource limits (CPUs, memory, walltime, GPUs) +- **Detection rules**: How to identify when you're on a specific HPC system +- **Default settings**: Account names and other system-specific defaults + +Built-in profiles are available for systems like NREL's Kestrel. You can also define custom profiles for private clusters. + +## Listing Available Profiles + +View all known HPC profiles: + +```bash +torc hpc list +``` + +Example output: +``` +Known HPC profiles: + +╭─────────┬──────────────┬────────────┬──────────╮ +│ Name │ Display Name │ Partitions │ Detected │ +├─────────┼──────────────┼────────────┼──────────┤ +│ kestrel │ NREL Kestrel │ 15 │ ✓ │ +╰─────────┴──────────────┴────────────┴──────────╯ +``` + +The "Detected" column shows if Torc recognizes you're currently on that system. + +## Detecting the Current System + +Torc can automatically detect which HPC system you're on: + +```bash +torc hpc detect +``` + +Detection works through environment variables. For example, NREL Kestrel is detected when `NREL_CLUSTER=kestrel` is set. + +## Viewing Profile Details + +See detailed information about a specific profile: + +```bash +torc hpc show kestrel +``` + +This displays: +- Profile name and description +- Detection method +- Default account (if configured) +- Number of partitions + +## Viewing Available Partitions + +List all partitions for a profile: + +```bash +torc hpc partitions kestrel +``` + +Example output: +``` +Partitions for kestrel: + +╭──────────┬─────────────┬───────────┬─────────────────┬─────────────────╮ +│ Name │ CPUs/Node │ Mem/Node │ Max Walltime │ GPUs │ +├──────────┼─────────────┼───────────┼─────────────────┼─────────────────┤ +│ debug │ 104 │ 240 GB │ 1h │ - │ +│ short │ 104 │ 240 GB │ 4h │ - │ +│ standard │ 104 │ 240 GB │ 48h │ - │ +│ gpu-h100 │ 2 │ 240 GB │ 48h │ 4 (H100) │ +│ ... │ ... │ ... │ ... │ ... │ +╰──────────┴─────────────┴───────────┴─────────────────┴─────────────────╯ +``` + +## Finding Matching Partitions + +Find partitions that can satisfy specific resource requirements: + +```bash +torc hpc match kestrel --cpus 32 --memory 64g --walltime 2h +``` + +Options: +- `--cpus `: Required CPU cores +- `--memory `: Required memory (e.g., `64g`, `512m`) +- `--walltime `: Required walltime (e.g., `2h`, `4:00:00`) +- `--gpus `: Required GPUs (optional) + +This is useful for understanding which partitions your jobs will be assigned to. + +## Custom HPC Profiles + +If your HPC system doesn't have a built-in profile, you have two options: + +> **Request Built-in Support** (Recommended) +> +> If your HPC is widely used, please [open an issue](https://github.com/NREL/torc/issues) requesting built-in support. Include: +> - Your HPC system name and organization +> - Partition names with resource limits (CPUs, memory, walltime, GPUs) +> - Detection method (environment variable or hostname pattern) +> +> Built-in profiles benefit everyone using that system and are maintained by the Torc team. + +If you need to use your HPC immediately or have a private cluster, you can define a custom profile in your configuration file. See the [Custom HPC Profile Tutorial](../tutorials/custom-hpc-profile.md) for a complete walkthrough. + +### Quick Example + +Define custom profiles in your configuration file: + +```toml +# ~/.config/torc/config.toml + +[client.hpc.custom_profiles.mycluster] +display_name = "My Research Cluster" +description = "Internal research HPC system" +detect_env_var = "MY_CLUSTER=research" +default_account = "default_project" + +[[client.hpc.custom_profiles.mycluster.partitions]] +name = "compute" +cpus_per_node = 64 +memory_mb = 256000 +max_walltime_secs = 172800 +shared = false + +[[client.hpc.custom_profiles.mycluster.partitions]] +name = "gpu" +cpus_per_node = 32 +memory_mb = 128000 +max_walltime_secs = 86400 +gpus_per_node = 4 +gpu_type = "A100" +shared = false +``` + +See [Configuration Reference](../reference/configuration.md) for full configuration options. + +## Using Profiles with Slurm Workflows + +HPC profiles are used by Slurm-related commands to automatically generate scheduler configurations. See [Working with Slurm](./slurm.md) for details on: +- `torc submit-slurm` - Submit workflows with auto-generated schedulers +- `torc workflows create-slurm` - Create workflows with auto-generated schedulers + +## See Also + +- [Working with Slurm](./slurm.md) +- [Custom HPC Profile Tutorial](../tutorials/custom-hpc-profile.md) +- [HPC Profiles Reference](../reference/hpc-profiles.md) +- [Configuration Reference](../reference/configuration.md) +- [Resource Requirements Reference](../reference/resources.md) diff --git a/docs/src/how-to/slurm.md b/docs/src/how-to/slurm.md index 935a00f2..b19bf29b 100644 --- a/docs/src/how-to/slurm.md +++ b/docs/src/how-to/slurm.md @@ -1,123 +1,305 @@ # Working with Slurm -## Torc server -- **External server**: A single member of your team allocates a shared server in - the HPC environment to host a Torc server. This is recommended if your - operations team provides this capability. The Torc server must be running on a - network accessible to compute nodes. -- **Login node**: By default, the server runs single-threaded. If you have small - job counts and each job runs a long time, the overhead on the login node will - be minimal. However, if you allocate hundreds of nodes with many thousands of - short jobs, the Torc server process may exceed allowed resource limits. Check - with your operations team if you have doubts. +This guide covers advanced Slurm configuration for users who need fine-grained control over their HPC workflows. +> **For most users**: See [Slurm Workflows](../explanation/slurm-workflows.md) for the recommended approach using `torc submit-slurm`. You don't need to manually configure schedulers or actions—Torc handles this automatically. -## Basic Slurm Configuration +## When to Use Manual Configuration -Define a Slurm scheduler in your workflow spec that matches your jobs' resource requirements: +Manual Slurm configuration is useful when you need: + +- Custom Slurm directives (e.g., `--constraint`, `--exclusive`) +- Multi-node jobs with specific topology requirements +- Shared allocations across multiple jobs for efficiency +- Non-standard partition configurations +- Fine-tuned control over allocation timing + +## Torc Server Requirements + +The Torc server must be accessible from compute nodes: + +- **External server** (Recommended): A team member allocates a shared server in the HPC environment. This is recommended if your operations team provides this capability. +- **Login node**: Suitable for small workflows. The server runs single-threaded by default. If you have many thousands of short jobs, check with your operations team about resource limits. + +## Manual Scheduler Configuration + +### Defining Slurm Schedulers + +Define schedulers in your workflow specification: ```yaml slurm_schedulers: - name: standard account: my_project nodes: 1 - walltime: 12:00:00 + walltime: "12:00:00" + partition: compute + mem: 64G + + - name: gpu_nodes + account: my_project + nodes: 1 + walltime: "08:00:00" + partition: gpu + gres: "gpu:4" + mem: 256G ``` -Then define an action to schedule the node on workflow start: +### Scheduler Fields + +| Field | Description | Required | +|-------|-------------|----------| +| `name` | Scheduler identifier | Yes | +| `account` | Slurm account/allocation | Yes | +| `nodes` | Number of nodes | Yes | +| `walltime` | Time limit (HH:MM:SS or D-HH:MM:SS) | Yes | +| `partition` | Slurm partition | No | +| `mem` | Memory per node | No | +| `gres` | Generic resources (e.g., GPUs) | No | +| `qos` | Quality of Service | No | +| `ntasks_per_node` | Tasks per node | No | +| `tmp` | Temporary disk space | No | +| `extra` | Additional sbatch arguments | No | + +### Defining Workflow Actions + +Actions trigger scheduler allocations: + ```yaml - - trigger_type: "on_workflow_start" - action_type: "schedule_nodes" - scheduler: "standard" - scheduler_type: "slurm" +actions: + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: standard + scheduler_type: slurm num_allocations: 1 -``` -Start the workflow with the workflow specification or create it and then pass the ID. -```bash -torc submit + - trigger_type: on_jobs_ready + action_type: schedule_nodes + jobs: [train_model] + scheduler: gpu_nodes + scheduler_type: slurm + num_allocations: 2 ``` -When Slurm allocates the node, a Torc worker will start pulling appropriate jobs from the database. +### Action Trigger Types + +| Trigger | Description | +|---------|-------------| +| `on_workflow_start` | Fires when workflow is submitted | +| `on_jobs_ready` | Fires when specified jobs become ready | +| `on_jobs_complete` | Fires when specified jobs complete | +| `on_workflow_complete` | Fires when all jobs complete | + +### Assigning Jobs to Schedulers + +Reference schedulers in job definitions: + +```yaml +jobs: + - name: preprocess + command: ./preprocess.sh + scheduler: standard -## Scheduling Compute Nodes + - name: train + command: python train.py + scheduler: gpu_nodes + depends_on: [preprocess] +``` -Three main approaches for running Torc workers on Slurm: +## Scheduling Strategies -### Approach 1: Many Single-Node Allocations +### Strategy 1: Many Single-Node Allocations Submit multiple Slurm jobs, each with its own Torc worker: ```yaml slurm_schedulers: - - name: "work_scheduler" - account: "my_account" + - name: work_scheduler + account: my_account nodes: 1 walltime: "04:00:00" -``` -```bash -torc slurm schedule-nodes -n 10 $WORKFLOW_ID +actions: + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: work_scheduler + scheduler_type: slurm + num_allocations: 10 ``` **When to use:** - Jobs have diverse resource requirements - Want independent time limits per job -- Need fine-grained control over resource allocation - Cluster has low queue wait times **Benefits:** -- Maximum flexibility in job scheduling -- Each job gets independent time limit -- Better resource matching (no wasted nodes) -- Fault isolation (one job failure doesn't affect others) +- Maximum scheduling flexibility +- Independent time limits per allocation +- Fault isolation **Drawbacks:** -- More Slurm queue time (multiple jobs to schedule) -- Higher Slurm scheduler overhead -- Potential for queue backlog on busy clusters +- More Slurm queue overhead +- Multiple jobs to schedule -### Approach 2: Multi-node Allocation, One Worker Per Node +### Strategy 2: Multi-Node Allocation, One Worker Per Node -Launch multiple workers, one per node: +Launch multiple workers within a single allocation: ```yaml slurm_schedulers: - - name: "work_scheduler" - account: "my_account" + - name: work_scheduler + account: my_account nodes: 10 walltime: "04:00:00" -``` -```bash -torc slurm schedule-nodes -n 1 $WORKFLOW_ID --start-one-worker-per-node +actions: + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: work_scheduler + scheduler_type: slurm + num_allocations: 1 + start_one_worker_per_node: true ``` +**When to use:** +- Many jobs with similar requirements +- Want faster queue scheduling (larger jobs often prioritized) + **Benefits:** -- Slurm typically prioritizes jobs with higher node counts -- Efficient for workflows with many jobs having similar resource requirements +- Single queue wait +- Often prioritized by Slurm scheduler **Drawbacks:** -- All Torc jobs share the same Slurm allocation time limit -- Wasted resources if some jobs run much longer than others +- Shared time limit for all workers +- Less flexibility -### Approach 3: One Worker Per Slurm Allocation +### Strategy 3: Single Worker Per Allocation -Submit a single Slurm job that runs the Torc worker on the head node. +One Torc worker handles all nodes: ```yaml slurm_schedulers: - - name: "work_scheduler" - account: "my_account" + - name: work_scheduler + account: my_account nodes: 10 walltime: "04:00:00" + +actions: + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: work_scheduler + scheduler_type: slurm + num_allocations: 1 ``` + +**When to use:** +- Your application manages node coordination +- Need full control over compute resources + +## Staged Allocations + +For pipelines with distinct phases, stage allocations to avoid wasted resources: + +```yaml +slurm_schedulers: + - name: preprocess_sched + account: my_project + nodes: 2 + walltime: "01:00:00" + + - name: compute_sched + account: my_project + nodes: 20 + walltime: "08:00:00" + + - name: postprocess_sched + account: my_project + nodes: 1 + walltime: "00:30:00" + +actions: + # Preprocessing starts immediately + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: preprocess_sched + scheduler_type: slurm + num_allocations: 1 + + # Compute nodes allocated when compute jobs are ready + - trigger_type: on_jobs_ready + action_type: schedule_nodes + jobs: [compute_step] + scheduler: compute_sched + scheduler_type: slurm + num_allocations: 1 + start_one_worker_per_node: true + + # Postprocessing allocated when those jobs are ready + - trigger_type: on_jobs_ready + action_type: schedule_nodes + jobs: [postprocess] + scheduler: postprocess_sched + scheduler_type: slurm + num_allocations: 1 +``` + +> **Note**: The `torc submit-slurm` command handles this automatically by analyzing job dependencies. + +## Custom Slurm Directives + +Use the `extra` field for additional sbatch arguments: + +```yaml +slurm_schedulers: + - name: exclusive_nodes + account: my_project + nodes: 4 + walltime: "04:00:00" + extra: "--exclusive --constraint=skylake" +``` + +## Submitting Workflows + +### With Manual Configuration + ```bash -torc slurm schedule-nodes -n 1 $WORKFLOW_ID +# Submit workflow with pre-defined schedulers and actions +torc submit workflow.yaml ``` -**Benefits:** -- Your job has full control over all compute nodes in the allocation. +### Scheduling Additional Nodes -**Drawbacks:** -- Complexity: your code must find all compute nodes and coordinate worker startup. +Add more allocations to a running workflow: + +```bash +torc slurm schedule-nodes -n 5 $WORKFLOW_ID +``` + +## Debugging + +### Check Slurm Job Status + +```bash +squeue -u $USER +``` + +### View Torc Worker Logs + +Workers log to the Slurm output file. Check: +```bash +cat slurm-.out +``` + +### Verify Server Connectivity + +From a compute node: +```bash +curl $TORC_API_URL/health +``` + +## See Also + +- [Slurm Workflows](../explanation/slurm-workflows.md) — Simplified workflow approach +- [HPC Profiles](./hpc-profiles.md) — Automatic partition matching +- [Workflow Actions](../explanation/workflow-actions.md) — Action system details +- [Debugging Slurm Workflows](./debugging-slurm.md) — Troubleshooting guide diff --git a/docs/src/how-to/visualizing-workflows.md b/docs/src/how-to/visualizing-workflows.md new file mode 100644 index 00000000..d2bc7b28 --- /dev/null +++ b/docs/src/how-to/visualizing-workflows.md @@ -0,0 +1,171 @@ +# Visualizing Workflow Structure + +Understanding how your workflow will execute—which jobs run in parallel, how dependencies create stages, and when Slurm allocations are requested—is essential for debugging and optimization. Torc provides several tools for visualizing workflow structure. + +## Execution Plan Command + +The `torc workflows execution-plan` command analyzes a workflow and displays its execution stages, showing how jobs are grouped and when schedulers allocate resources. + +### Basic Usage + +```bash +# From a specification file +torc workflows execution-plan workflow.yaml + +# From an existing workflow +torc workflows execution-plan +``` + +### Example Output + +For a workflow with two independent processing pipelines that merge at the end: + +``` +Workflow: two_subgraph_pipeline +Total Jobs: 15 + +▶ Stage 1: Workflow Start + Scheduler Allocations: + • prep_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • prep_a + • prep_b + +→ Stage 2: When jobs 'prep_a', 'prep_b' complete + Scheduler Allocations: + • work_a_sched (slurm) - 1 allocation(s) + • work_b_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • work_a_{1..5} + • work_b_{1..5} + +→ Stage 3: When 10 jobs complete + Scheduler Allocations: + • post_a_sched (slurm) - 1 allocation(s) + • post_b_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • post_a + • post_b + +→ Stage 4: When jobs 'post_a', 'post_b' complete + Scheduler Allocations: + • final_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • final + +Total Stages: 4 +``` + +### What the Execution Plan Shows + +1. **Stages**: Groups of jobs that become ready at the same time based on dependency resolution +2. **Scheduler Allocations**: Which Slurm schedulers request resources at each stage (for workflows with Slurm configuration) +3. **Jobs Becoming Ready**: Which jobs transition to "ready" status at each stage +4. **Subgraphs**: Independent branches of the workflow that can execute in parallel + +### Workflows Without Slurm Schedulers + +For workflows without pre-defined Slurm schedulers, the execution plan shows the job stages without scheduler information: + +```bash +torc workflows execution-plan workflow_no_slurm.yaml +``` + +``` +Workflow: my_pipeline +Total Jobs: 10 + +▶ Stage 1: Workflow Start + Jobs Becoming Ready: + • preprocess + +→ Stage 2: When job 'preprocess' completes + Jobs Becoming Ready: + • work_{1..5} + +→ Stage 3: When 5 jobs complete + Jobs Becoming Ready: + • postprocess + +Total Stages: 3 +``` + +This helps you understand the workflow topology before adding Slurm configuration with `torc slurm generate`. + +### Use Cases + +- **Validate workflow structure**: Ensure dependencies create the expected execution order +- **Identify parallelism**: See which jobs can run concurrently +- **Debug slow workflows**: Find stages that serialize unnecessarily +- **Plan Slurm allocations**: Understand when resources will be requested +- **Verify auto-generated schedulers**: Check that `torc slurm generate` created appropriate staging + +## DAG Visualization in the Dashboard + +The [web dashboard](./dashboard.md) provides interactive DAG (Directed Acyclic Graph) visualization. + +### Viewing the DAG + +1. Navigate to the **Details** tab +2. Select a workflow +3. Click **View DAG** in the Visualization section + +### DAG Types + +The dashboard supports three DAG visualization types: + +| Type | Description | +|------|-------------| +| **Job Dependencies** | Shows explicit and implicit dependencies between jobs | +| **Job-File Relations** | Shows how jobs connect through input/output files | +| **Job-UserData Relations** | Shows how jobs connect through user data | + +### DAG Features + +- **Color-coded nodes**: Jobs are colored by status (ready, running, completed, failed, etc.) +- **Interactive**: Zoom, pan, and click nodes for details +- **Layout**: Automatic hierarchical layout using Dagre algorithm +- **Legend**: Status color reference + +## TUI DAG View + +The terminal UI (`torc tui`) also includes DAG visualization: + +1. Select a workflow +2. Press `d` to toggle the DAG view +3. Use arrow keys to navigate + +## Comparing Visualization Tools + +| Tool | Best For | +|------|----------| +| `execution-plan` | Understanding execution stages, Slurm allocation timing | +| Dashboard DAG | Interactive exploration, status monitoring | +| TUI DAG | Quick terminal-based visualization | + +## Example: Analyzing a Complex Workflow + +Consider a workflow with preprocessing, parallel work, and aggregation: + +```bash +# First, view the execution plan +torc workflows execution-plan examples/subgraphs/subgraphs_workflow.yaml + +# If no schedulers, generate them +torc slurm generate --account myproject examples/subgraphs/subgraphs_workflow_no_slurm.yaml + +# View the plan again to see scheduler allocations +torc workflows execution-plan examples/subgraphs/subgraphs_workflow.yaml +``` + +The execution plan helps you verify that: +- Independent subgraphs are correctly identified +- Stages align with your expected execution order +- Slurm allocations are timed appropriately + +## See Also + +- [Web Dashboard](./dashboard.md) — Full dashboard documentation +- [Slurm Workflows](../explanation/slurm-workflows.md) — Understanding Slurm integration +- [Workflow Actions](../explanation/workflow-actions.md) — How actions trigger scheduler allocations +- [Subgraphs Example](https://github.com/NREL/torc/tree/main/examples/subgraphs) — Complete example with multiple subgraphs diff --git a/docs/src/quick-start-hpc.md b/docs/src/quick-start-hpc.md new file mode 100644 index 00000000..c9aacc4d --- /dev/null +++ b/docs/src/quick-start-hpc.md @@ -0,0 +1,181 @@ +# Quick Start (HPC) + +This guide walks you through running your first Torc workflow on an HPC cluster with Slurm. +Jobs are submitted to Slurm and run on compute nodes. + +For local execution (testing, development, or non-HPC environments), see [Quick Start (Local)](./quick-start-local.md). + +## Prerequisites + +- Access to an HPC cluster with Slurm +- A Slurm account/allocation for submitting jobs +- Torc installed (see [Installation](./installation.md)) + +## Start the Server + +On the login node, start a Torc server with a local database: + +```console +torc-server run --database torc.db --completion-check-interval-secs 5 +``` + +> **Note:** For larger deployments, your team may provide a shared Torc server. +> In that case, skip this step and set `TORC_API_URL` to the shared server address. + +## Check Your HPC Profile + +Torc includes built-in profiles for common HPC systems. Check if your system is detected: + +```console +torc hpc detect +``` + +If detected, you'll see your HPC system name. To see available partitions: + +```console +torc hpc partitions +``` + +> **Note:** If your HPC system isn't detected, see [Custom HPC Profile](./tutorials/custom-hpc-profile.md) +> or [request built-in support](https://github.com/NREL/torc/issues). + +## Create a Workflow with Resource Requirements + +Save this as `workflow.yaml`: + +```yaml +name: hpc_hello_world +description: Simple HPC workflow + +resource_requirements: + - name: small + num_cpus: 4 + memory: 8g + runtime: PT30M + +jobs: + - name: job1 + command: echo "Hello from compute node!" && hostname + resource_requirements: small + + - name: job2 + command: echo "Hello again!" && hostname + resource_requirements: small + depends_on: [job1] +``` + +Key differences from local workflows: +- **resource_requirements**: Define CPU, memory, and runtime needs +- Jobs reference these requirements by name +- Torc matches requirements to appropriate Slurm partitions + +## Submit the Workflow + +Submit with your Slurm account: + +```console +torc submit-slurm --account workflow.yaml +``` + +Torc will: +1. Detect your HPC system +2. Match job requirements to appropriate partitions +3. Generate Slurm scheduler configurations +4. Create and submit the workflow + +## Monitor Progress + +Check workflow status: + +```console +torc workflows list +torc jobs list +``` + +Or use the interactive TUI: + +```console +torc tui +``` + +Check Slurm queue: + +```console +squeue -u $USER +``` + +## View Results + +Once jobs complete: + +```console +torc results list +``` + +Job output is stored in the `output/` directory by default. + +## Example: Multi-Stage Pipeline + +A more realistic workflow with different resource requirements per stage: + +```yaml +name: analysis_pipeline +description: Data processing pipeline + +resource_requirements: + - name: light + num_cpus: 4 + memory: 8g + runtime: PT30M + + - name: compute + num_cpus: 32 + memory: 64g + runtime: PT2H + + - name: gpu + num_cpus: 8 + num_gpus: 1 + memory: 32g + runtime: PT1H + +jobs: + - name: preprocess + command: python preprocess.py + resource_requirements: light + + - name: train + command: python train.py + resource_requirements: gpu + depends_on: [preprocess] + + - name: evaluate + command: python evaluate.py + resource_requirements: compute + depends_on: [train] +``` + +Torc stages resource allocation based on dependencies: +- `preprocess` resources are allocated at workflow start +- `train` resources are allocated when `preprocess` completes +- `evaluate` resources are allocated when `train` completes + +This prevents wasting allocation time on resources that aren't needed yet. + +## Preview Before Submitting + +For production workflows, preview the generated Slurm configuration first: + +```console +torc slurm generate --account workflow.yaml +``` + +This shows what schedulers and actions Torc will create without submitting anything. + +## Next Steps + +- [Slurm Workflows](./explanation/slurm-workflows.md) — How Torc manages Slurm +- [Resource Requirements](./reference/resources.md) — All resource options +- [HPC Profiles](./how-to/hpc-profiles.md) — Managing HPC configurations +- [Working with Slurm](./how-to/slurm.md) — Advanced Slurm configuration +- [Debugging Slurm Workflows](./how-to/debugging-slurm.md) — Troubleshooting diff --git a/docs/src/quick-start-local.md b/docs/src/quick-start-local.md new file mode 100644 index 00000000..5a053db2 --- /dev/null +++ b/docs/src/quick-start-local.md @@ -0,0 +1,115 @@ +# Quick Start (Local) + +This guide walks you through creating and running your first Torc workflow with local execution. +Jobs run directly on the current machine, making this ideal for testing, development, or +non-HPC environments. + +For running workflows on HPC clusters with Slurm, see [Quick Start (HPC)](./quick-start-hpc.md). + +## Start the Server + +Start a Torc server with a local database. Setting `--completion-check-interval-secs` ensures +job completions are processed quickly (use this for personal servers, not shared deployments). + +```console +torc-server run --database torc.db --completion-check-interval-secs 5 +``` + +## Test the Connection + +In a new terminal, verify the client can connect: + +```console +torc workflows list +``` + +## Create a Workflow + +Save this as `workflow.yaml`: + +```yaml +name: hello_world +description: Simple hello world workflow + +jobs: + - name: job 1 + command: echo "Hello from torc!" + - name: job 2 + command: echo "Hello again from torc!" +``` + +> **Note:** Torc also accepts `.json5` and `.kdl` workflow specifications. +> See [Workflow Specification Formats](./reference/workflow-formats.md) for details. + +## Run the Workflow + +Run jobs locally with a short poll interval for demo purposes: + +```console +torc run workflow.yaml --poll-interval 1 +``` + +This creates the workflow, initializes it, and runs all jobs on the current machine. + +## View Results + +```console +torc results list +``` + +Or use the TUI for an interactive view: + +```console +torc tui +``` + +## Example: Diamond Workflow + +A workflow with fan-out and fan-in dependencies: + +```yaml +name: diamond_workflow +description: Example workflow with implicit dependencies + +jobs: + - name: preprocess + command: "bash tests/scripts/preprocess.sh -i ${files.input.f1} -o ${files.output.f2} -o ${files.output.f3}" + + - name: work1 + command: "bash tests/scripts/work.sh -i ${files.input.f2} -o ${files.output.f4}" + + - name: work2 + command: "bash tests/scripts/work.sh -i ${files.input.f3} -o ${files.output.f5}" + + - name: postprocess + command: "bash tests/scripts/postprocess.sh -i ${files.input.f4} -i ${files.input.f5} -o ${files.output.f6}" + +files: + - name: f1 + path: f1.json + - name: f2 + path: f2.json + - name: f3 + path: f3.json + - name: f4 + path: f4.json + - name: f5 + path: f5.json + - name: f6 + path: f6.json +``` + +Dependencies are automatically inferred from file inputs/outputs: +- `work1` and `work2` wait for `preprocess` (depend on its output files) +- `postprocess` waits for both `work1` and `work2` to complete + +## More Examples + +The [examples directory](https://github.com/NREL/torc/tree/main/examples) contains many more +workflow examples in YAML, JSON5, and KDL formats. + +## Next Steps + +- [Quick Start (HPC)](./quick-start-hpc.md) - Run workflows on Slurm clusters +- [Creating Workflows](./how-to/creating-workflows.md) - Detailed workflow creation guide +- [Terminal UI](./how-to/tui.md) - Interactive workflow monitoring diff --git a/docs/src/quick-start.md b/docs/src/quick-start.md index 0995190e..46dab477 100644 --- a/docs/src/quick-start.md +++ b/docs/src/quick-start.md @@ -1,109 +1,21 @@ # Quick Start -This guide will walk you through creating and running your first Torc workflows. These examples -rely on local execution. Refer to the Tutorials for HPC examples. +Choose the guide that matches your environment: -## Start the server -This will create a torc database in the current directory if it doesn't exist. -Setting `--completion-check-interval-secs` (or `-c`) will ensure that the server processes job completions quickly. -This should not be set on a shared server. +## [Quick Start (HPC)](./quick-start-hpc.md) -```console -torc-server run --database torc.db --completion-check-interval-secs 5 -``` +**For HPC clusters with Slurm** — Run workflows on compute nodes via Slurm. -## Test the client connection -```console -torc workflows list -``` +- Start server on login node +- Define jobs with resource requirements (CPU, memory, runtime) +- Submit with `torc submit-slurm --account workflow.yaml` +- Jobs run on compute nodes -## Create a Workflow Specification +## [Quick Start (Local)](./quick-start-local.md) -Torc supports YAML, JSON5, and KDL formats. Save as `workflow.yaml`: +**For local execution** — Run workflows on the current machine. -```yaml -name: hello_world -description: Simple hello world workflow - -jobs: - - name: job 1 - command: echo "Hello from torc!" - - name: job 2 - command: echo "Hello again from torc!" -``` - -> **Note:** Torc also accepts `.json5` and `.kdl` workflow specifications. See the [Workflow Specification Formats](./reference/workflow-formats.md) reference for details on all supported formats. - -## Run Jobs -Run the jobs on the current computer. Use a short poll interval for demo purposes. -This will automatically initialize the jobs if you skipped that step. - -```console -torc run --poll-interval 1 -``` - -## View Results - -```console -# View job results -torc results list -``` - -# Or use the TUI to view the results in table. -```console -torc-tui -``` - -## Example: Diamond Workflow - -A workflow with fan-out and fan-in dependencies. You can find this example in the repository: -- [diamond_workflow.yaml](https://github.com/NREL/torc/blob/main/examples/yaml/diamond_workflow.yaml) -- [diamond_workflow.json5](https://github.com/NREL/torc/blob/main/examples/json/diamond_workflow.json5) -- [diamond_workflow.kdl](https://github.com/NREL/torc/blob/main/examples/kdl/diamond_workflow.kdl) - -```yaml -name: diamond_workflow -description: Example workflow with implicit dependencies - -jobs: - - name: preprocess - command: "bash tests/scripts/preprocess.sh -i ${files.input.f1} -o ${files.output.f2} -o ${files.output.f3}" - - - name: work1 - command: "bash tests/scripts/work.sh -i ${files.input.f2} -o ${files.output.f4}" - - - name: work2 - command: "bash tests/scripts/work.sh -i ${files.input.f3} -o ${files.output.f5}" - - - name: postprocess - command: "bash tests/scripts/postprocess.sh -i ${files.input.f4} -i ${files.input.f5} -o ${files.output.f6}" - -# File definitions - representing the data files in the workflow -files: - - name: f1 - path: f1.json - - name: f2 - path: f2.json - - name: f3 - path: f3.json - - name: f4 - path: f4.json - - name: f5 - path: f5.json - - name: f6 - path: f6.json -``` - -Dependencies are automatically inferred: -- `work1` and `work2` wait for `preprocess` (depend on its output files) -- `postprocess` waits for both `work1` and `work2` to complete - -## More Examples - -The [examples directory](https://github.com/NREL/torc/tree/main/examples) contains many more workflow examples in all supported formats: -- Simple workflows and resource monitoring -- Workflow actions for automation -- Slurm integration examples -- Parameterized workflows - -Browse [examples/yaml](https://github.com/NREL/torc/tree/main/examples/yaml), [examples/json](https://github.com/NREL/torc/tree/main/examples/json), or [examples/kdl](https://github.com/NREL/torc/tree/main/examples/kdl) to explore the full collection. +- Ideal for testing, development, or non-HPC environments +- Start server locally +- Run with `torc run workflow.yaml` +- Jobs run on the current machine diff --git a/docs/src/reference/configuration.md b/docs/src/reference/configuration.md index 9f5cbd20..3ae9fd41 100644 --- a/docs/src/reference/configuration.md +++ b/docs/src/reference/configuration.md @@ -70,6 +70,81 @@ memory_gb = 32.0 num_gpus = 1 ``` +### `[client.hpc]` Section + +Settings for HPC profile system (used by `torc hpc` and `torc slurm` commands). + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `profile_overrides` | table | `{}` | Override settings for built-in HPC profiles | +| `custom_profiles` | table | `{}` | Define custom HPC profiles | + +### `[client.hpc.profile_overrides.]` Section + +Override settings for built-in profiles (e.g., `kestrel`). + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `default_account` | string | (none) | Default Slurm account for this profile | + +### `[client.hpc.custom_profiles.]` Section + +Define a custom HPC profile. + +| Option | Type | Required | Description | +|--------|------|----------|-------------| +| `display_name` | string | No | Human-readable name | +| `description` | string | No | Profile description | +| `detect_env_var` | string | No | Environment variable for detection (`NAME=value`) | +| `detect_hostname` | string | No | Regex pattern for hostname detection | +| `default_account` | string | No | Default Slurm account | +| `partitions` | array | Yes | List of partition configurations | + +### `[[client.hpc.custom_profiles..partitions]]` Section + +Define partitions for a custom profile. + +| Option | Type | Required | Description | +|--------|------|----------|-------------| +| `name` | string | Yes | Partition name | +| `cpus_per_node` | int | Yes | CPU cores per node | +| `memory_mb` | int | Yes | Memory per node in MB | +| `max_walltime_secs` | int | Yes | Maximum walltime in seconds | +| `gpus_per_node` | int | No | GPUs per node | +| `gpu_type` | string | No | GPU model (e.g., "H100") | +| `shared` | bool | No | Whether partition supports shared jobs | +| `min_nodes` | int | No | Minimum required nodes | +| `requires_explicit_request` | bool | No | Must be explicitly requested | + +### HPC Example + +```toml +[client.hpc.profile_overrides.kestrel] +default_account = "my_default_account" + +[client.hpc.custom_profiles.mycluster] +display_name = "My Research Cluster" +description = "Internal research HPC system" +detect_env_var = "MY_CLUSTER=research" +default_account = "default_project" + +[[client.hpc.custom_profiles.mycluster.partitions]] +name = "compute" +cpus_per_node = 64 +memory_mb = 256000 +max_walltime_secs = 172800 +shared = false + +[[client.hpc.custom_profiles.mycluster.partitions]] +name = "gpu" +cpus_per_node = 32 +memory_mb = 128000 +max_walltime_secs = 86400 +gpus_per_node = 4 +gpu_type = "A100" +shared = false +``` + ## Server Configuration Settings for `torc-server`. diff --git a/docs/src/reference/hpc-profiles.md b/docs/src/reference/hpc-profiles.md new file mode 100644 index 00000000..49e31b46 --- /dev/null +++ b/docs/src/reference/hpc-profiles.md @@ -0,0 +1,351 @@ +# HPC Profiles Reference + +Complete reference for HPC profile system and CLI commands. + +## Overview + +HPC profiles contain pre-configured knowledge about High-Performance Computing systems, enabling automatic Slurm scheduler generation based on job resource requirements. + +## CLI Commands + +### `torc hpc list` + +List all available HPC profiles. + +```bash +torc hpc list [OPTIONS] +``` + +**Options:** +| Option | Description | +|--------|-------------| +| `-f, --format ` | Output format: `table` or `json` | + +**Output columns:** +- **Name**: Profile identifier used in commands +- **Display Name**: Human-readable name +- **Partitions**: Number of configured partitions +- **Detected**: Whether current system matches this profile + +--- + +### `torc hpc detect` + +Detect the current HPC system. + +```bash +torc hpc detect [OPTIONS] +``` + +**Options:** +| Option | Description | +|--------|-------------| +| `-f, --format ` | Output format: `table` or `json` | + +Returns the detected profile name, or indicates no match. + +--- + +### `torc hpc show` + +Display detailed information about an HPC profile. + +```bash +torc hpc show [OPTIONS] +``` + +**Arguments:** +| Argument | Description | +|----------|-------------| +| `` | Profile name (e.g., `kestrel`) | + +**Options:** +| Option | Description | +|--------|-------------| +| `-f, --format ` | Output format: `table` or `json` | + +--- + +### `torc hpc partitions` + +List partitions for an HPC profile. + +```bash +torc hpc partitions [OPTIONS] +``` + +**Arguments:** +| Argument | Description | +|----------|-------------| +| `` | Profile name (e.g., `kestrel`) | + +**Options:** +| Option | Description | +|--------|-------------| +| `-f, --format ` | Output format: `table` or `json` | + +**Output columns:** +- **Name**: Partition name +- **CPUs/Node**: CPU cores per node +- **Mem/Node**: Memory per node +- **Max Walltime**: Maximum job duration +- **GPUs**: GPU count and type (if applicable) +- **Shared**: Whether partition supports shared jobs +- **Notes**: Special requirements or features + +--- + +### `torc hpc match` + +Find partitions matching resource requirements. + +```bash +torc hpc match [OPTIONS] +``` + +**Arguments:** +| Argument | Description | +|----------|-------------| +| `` | Profile name (e.g., `kestrel`) | + +**Options:** +| Option | Description | +|--------|-------------| +| `--cpus ` | Required CPU cores | +| `--memory ` | Required memory (e.g., `64g`, `512m`) | +| `--walltime ` | Required walltime (e.g., `2h`, `4:00:00`) | +| `--gpus ` | Required GPUs | +| `-f, --format ` | Output format: `table` or `json` | + +**Memory format:** `` where unit is `k`, `m`, `g`, or `t` (case-insensitive). + +**Walltime formats:** +- `HH:MM:SS` (e.g., `04:00:00`) +- `h` (e.g., `4h`) +- `m` (e.g., `30m`) +- `s` (e.g., `3600s`) + +--- + +### `torc slurm generate` + +Generate Slurm schedulers for a workflow based on job resource requirements. + +```bash +torc slurm generate [OPTIONS] --account +``` + +**Arguments:** +| Argument | Description | +|----------|-------------| +| `` | Path to workflow specification file (YAML, JSON, or JSON5) | + +**Options:** +| Option | Description | +|--------|-------------| +| `--account ` | Slurm account to use (required) | +| `--profile ` | HPC profile to use (auto-detected if not specified) | +| `-o, --output ` | Output file path (prints to stdout if not specified) | +| `--no-actions` | Don't add workflow actions for scheduling nodes | +| `--force` | Overwrite existing schedulers in the workflow | + +**Generated artifacts:** + +1. **Slurm schedulers**: One for each unique resource requirement +2. **Job scheduler assignments**: Each job linked to appropriate scheduler +3. **Workflow actions**: `on_workflow_start`/`schedule_nodes` actions (unless `--no-actions`) + +**Scheduler naming:** `_scheduler` + +--- + +## Built-in Profiles + +### NREL Kestrel + +**Profile name:** `kestrel` + +**Detection:** Environment variable `NREL_CLUSTER=kestrel` + +**Partitions:** + +| Partition | CPUs | Memory | Max Walltime | GPUs | Notes | +|-----------|------|--------|--------------|------|-------| +| `debug` | 104 | 240 GB | 1h | - | Quick testing | +| `short` | 104 | 240 GB | 4h | - | Short jobs | +| `standard` | 104 | 240 GB | 48h | - | General workloads | +| `long` | 104 | 240 GB | 240h | - | Extended jobs | +| `medmem` | 104 | 480 GB | 48h | - | Medium memory | +| `bigmem` | 104 | 2048 GB | 48h | - | High memory | +| `shared` | 104 | 240 GB | 48h | - | Shared node access | +| `hbw` | 104 | 240 GB | 48h | - | High-bandwidth memory, min 10 nodes | +| `nvme` | 104 | 240 GB | 48h | - | NVMe local storage | +| `gpu-h100` | 2 | 240 GB | 48h | 4x H100 | GPU compute | + +**Node specifications:** +- **Standard nodes**: 104 cores (2x Intel Xeon Sapphire Rapids), 240 GB RAM +- **GPU nodes**: 4x NVIDIA H100 80GB HBM3, 128 cores, 2 TB RAM + +--- + +## Configuration + +### Custom Profiles + +> **Don't see your HPC?** Please [request built-in support](https://github.com/NREL/torc/issues) so everyone benefits. See the [Custom HPC Profile Tutorial](../tutorials/custom-hpc-profile.md) for creating a profile while you wait. + +Define custom profiles in your Torc configuration file: + +```toml +# ~/.config/torc/config.toml + +[client.hpc.custom_profiles.mycluster] +display_name = "My Cluster" +description = "Description of the cluster" +detect_env_var = "CLUSTER_NAME=mycluster" +detect_hostname = ".*\\.mycluster\\.org" +default_account = "myproject" + +[[client.hpc.custom_profiles.mycluster.partitions]] +name = "compute" +cpus_per_node = 64 +memory_mb = 256000 +max_walltime_secs = 172800 +shared = false + +[[client.hpc.custom_profiles.mycluster.partitions]] +name = "gpu" +cpus_per_node = 32 +memory_mb = 128000 +max_walltime_secs = 86400 +gpus_per_node = 4 +gpu_type = "A100" +shared = false +``` + +### Profile Override + +Override settings for built-in profiles: + +```toml +[client.hpc.profile_overrides.kestrel] +default_account = "my_default_account" +``` + +### Configuration Options + +**`[client.hpc]` Section:** + +| Option | Type | Description | +|--------|------|-------------| +| `profile_overrides` | table | Override settings for built-in profiles | +| `custom_profiles` | table | Define custom HPC profiles | + +**Profile override options:** + +| Option | Type | Description | +|--------|------|-------------| +| `default_account` | string | Default Slurm account for this profile | + +**Custom profile options:** + +| Option | Type | Required | Description | +|--------|------|----------|-------------| +| `display_name` | string | No | Human-readable name | +| `description` | string | No | Profile description | +| `detect_env_var` | string | No | Environment variable for detection (`NAME=value`) | +| `detect_hostname` | string | No | Regex pattern for hostname detection | +| `default_account` | string | No | Default Slurm account | +| `partitions` | array | Yes | List of partition configurations | + +**Partition options:** + +| Option | Type | Required | Description | +|--------|------|----------|-------------| +| `name` | string | Yes | Partition name | +| `cpus_per_node` | int | Yes | CPU cores per node | +| `memory_mb` | int | Yes | Memory per node in MB | +| `max_walltime_secs` | int | Yes | Maximum walltime in seconds | +| `gpus_per_node` | int | No | GPUs per node | +| `gpu_type` | string | No | GPU model (e.g., "H100") | +| `shared` | bool | No | Whether partition supports shared jobs | +| `min_nodes` | int | No | Minimum required nodes | +| `requires_explicit_request` | bool | No | Must be explicitly requested | + +--- + +## Resource Matching Algorithm + +When generating schedulers, Torc uses this algorithm to match resource requirements to partitions: + +1. **Filter by resources**: Partitions must satisfy: + - CPUs >= required CPUs + - Memory >= required memory + - GPUs >= required GPUs (if specified) + - Max walltime >= required runtime + +2. **Exclude debug partitions**: Unless no other partition matches + +3. **Prefer best fit**: + - Partitions that exactly match resource needs + - Non-shared partitions over shared + - Shorter max walltime over longer + +4. **Handle special requirements**: + - GPU jobs only match GPU partitions + - Respect `requires_explicit_request` flag + - Honor `min_nodes` constraints + +--- + +## Generated Scheduler Format + +Example generated Slurm scheduler: + +```yaml +slurm_schedulers: + - name: medium_scheduler + account: myproject + nodes: 1 + mem: 64g + walltime: 04:00:00 + gres: null + partition: null # Let Slurm choose based on resources +``` + +Corresponding workflow action: + +```yaml +actions: + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: medium_scheduler + scheduler_type: slurm + num_allocations: 1 +``` + +--- + +## Runtime Format Parsing + +Resource requirements use ISO 8601 duration format for runtime: + +| Format | Example | Meaning | +|--------|---------|---------| +| `PTnH` | `PT4H` | 4 hours | +| `PTnM` | `PT30M` | 30 minutes | +| `PTnS` | `PT3600S` | 3600 seconds | +| `PTnHnM` | `PT2H30M` | 2 hours 30 minutes | +| `PnDTnH` | `P1DT12H` | 1 day 12 hours | + +Generated walltime uses `HH:MM:SS` format (e.g., `04:00:00`). + +--- + +## See Also + +- [Working with HPC Profiles](../how-to/hpc-profiles.md) +- [Custom HPC Profile Tutorial](../tutorials/custom-hpc-profile.md) +- [Working with Slurm](../how-to/slurm.md) +- [Resource Requirements](./resources.md) +- [Configuration Reference](./configuration.md) diff --git a/docs/src/tutorials/README.md b/docs/src/tutorials/README.md index f2811ccd..4d835900 100644 --- a/docs/src/tutorials/README.md +++ b/docs/src/tutorials/README.md @@ -15,6 +15,7 @@ This section contains learning-oriented lessons to help you get started with Tor 9. [Multi-Stage Workflows with Barriers](./multi-stage-barrier.md) - Scale to thousands of jobs efficiently 10. [Map Python Functions](./map_python_function_across_workers.md) - Distribute Python functions across workers 11. [Filtering CLI Output with Nushell](./filtering-with-nushell.md) - Filter jobs, results, and user data with readable queries +12. [Custom HPC Profile](./custom-hpc-profile.md) - Create an HPC profile for unsupported clusters Start with the Configuration Files tutorial to set up your environment, then try the Dashboard Deployment tutorial if you want to use the web interface. diff --git a/docs/src/tutorials/custom-hpc-profile.md b/docs/src/tutorials/custom-hpc-profile.md new file mode 100644 index 00000000..2c585256 --- /dev/null +++ b/docs/src/tutorials/custom-hpc-profile.md @@ -0,0 +1,326 @@ +# Creating a Custom HPC Profile + +This tutorial walks you through creating a custom HPC profile for a cluster that Torc doesn't have built-in support for. + +## Before You Start + +> **Request Built-in Support First!** +> +> If your HPC system is widely used, consider requesting that Torc developers add it as a built-in profile. This benefits everyone using that system. +> +> Open an issue at [github.com/NREL/torc/issues](https://github.com/NREL/torc/issues) with: +> - Your HPC system name and organization +> - Partition names and their resource limits (CPUs, memory, walltime, GPUs) +> - How to detect the system (environment variable or hostname pattern) +> - Any special requirements (minimum nodes, exclusive partitions, etc.) +> +> Built-in profiles are maintained by the Torc team and stay up-to-date as systems change. + +## When to Create a Custom Profile + +Create a custom profile when: +- Your HPC isn't supported and you need to use it immediately +- You have a private or internal cluster +- You want to test profile configurations before submitting upstream + +## Step 1: Gather Partition Information + +First, collect information about your HPC's partitions. On most Slurm systems: + +```bash +# List all partitions +sinfo -s + +# Get detailed partition info +sinfo -o "%P %c %m %l %G" +``` + +For this tutorial, let's say your cluster "ResearchCluster" has these partitions: + +| Partition | CPUs/Node | Memory | Max Walltime | GPUs | +|-----------|-----------|--------|--------------|------| +| `batch` | 48 | 192 GB | 72 hours | - | +| `short` | 48 | 192 GB | 4 hours | - | +| `gpu` | 32 | 256 GB | 48 hours | 4x A100 | +| `himem` | 48 | 1024 GB | 48 hours | - | + +## Step 2: Identify Detection Method + +Determine how Torc can detect when you're on this system. Common methods: + +**Environment variable** (most common): +```bash +echo $CLUSTER_NAME # e.g., "research" +echo $SLURM_CLUSTER # e.g., "researchcluster" +``` + +**Hostname pattern**: +```bash +hostname # e.g., "login01.research.edu" +``` + +For this tutorial, we'll use the environment variable `CLUSTER_NAME=research`. + +## Step 3: Create the Configuration File + +Create or edit your Torc configuration file: + +```bash +# Linux +mkdir -p ~/.config/torc +nano ~/.config/torc/config.toml + +# macOS +mkdir -p ~/Library/Application\ Support/torc +nano ~/Library/Application\ Support/torc/config.toml +``` + +Add your custom profile: + +```toml +# Custom HPC Profile for ResearchCluster +[client.hpc.custom_profiles.research] +display_name = "Research Cluster" +description = "University Research HPC System" +detect_env_var = "CLUSTER_NAME=research" +default_account = "my_project" + +# Batch partition - general purpose +[[client.hpc.custom_profiles.research.partitions]] +name = "batch" +cpus_per_node = 48 +memory_mb = 192000 # 192 GB in MB +max_walltime_secs = 259200 # 72 hours in seconds +shared = false + +# Short partition - quick jobs +[[client.hpc.custom_profiles.research.partitions]] +name = "short" +cpus_per_node = 48 +memory_mb = 192000 +max_walltime_secs = 14400 # 4 hours +shared = true # Allows sharing nodes + +# GPU partition +[[client.hpc.custom_profiles.research.partitions]] +name = "gpu" +cpus_per_node = 32 +memory_mb = 256000 # 256 GB +max_walltime_secs = 172800 # 48 hours +gpus_per_node = 4 +gpu_type = "A100" +shared = false + +# High memory partition +[[client.hpc.custom_profiles.research.partitions]] +name = "himem" +cpus_per_node = 48 +memory_mb = 1048576 # 1024 GB (1 TB) +max_walltime_secs = 172800 # 48 hours +shared = false +``` + +## Step 4: Verify the Profile + +Check that Torc recognizes your profile: + +```bash +# List all profiles +torc hpc list +``` + +You should see your custom profile: + +``` +Known HPC profiles: + +╭──────────┬──────────────────┬────────────┬──────────╮ +│ Name │ Display Name │ Partitions │ Detected │ +├──────────┼──────────────────┼────────────┼──────────┤ +│ kestrel │ NREL Kestrel │ 15 │ │ +│ research │ Research Cluster │ 4 │ ✓ │ +╰──────────┴──────────────────┴────────────┴──────────╯ +``` + +View the partitions: + +```bash +torc hpc partitions research +``` + +``` +Partitions for research: + +╭─────────┬───────────┬───────────┬─────────────┬──────────╮ +│ Name │ CPUs/Node │ Mem/Node │ Max Walltime│ GPUs │ +├─────────┼───────────┼───────────┼─────────────┼──────────┤ +│ batch │ 48 │ 192 GB │ 72h │ - │ +│ short │ 48 │ 192 GB │ 4h │ - │ +│ gpu │ 32 │ 256 GB │ 48h │ 4 (A100) │ +│ himem │ 48 │ 1024 GB │ 48h │ - │ +╰─────────┴───────────┴───────────┴─────────────┴──────────╯ +``` + +## Step 5: Test Partition Matching + +Verify that Torc correctly matches resource requirements to partitions: + +```bash +# Should match 'short' partition +torc hpc match research --cpus 8 --memory 16g --walltime 2h + +# Should match 'gpu' partition +torc hpc match research --cpus 16 --memory 64g --walltime 8h --gpus 2 + +# Should match 'himem' partition +torc hpc match research --cpus 24 --memory 512g --walltime 24h +``` + +## Step 6: Test Scheduler Generation + +Create a test workflow to verify scheduler generation: + +```yaml +# test_workflow.yaml +name: profile_test +description: Test custom HPC profile + +resource_requirements: + - name: standard + num_cpus: 16 + memory: 64g + runtime: PT2H + + - name: gpu_compute + num_cpus: 16 + num_gpus: 2 + memory: 128g + runtime: PT8H + +jobs: + - name: preprocess + command: echo "preprocessing" + resource_requirements: standard + + - name: train + command: echo "training" + resource_requirements: gpu_compute + depends_on: [preprocess] +``` + +Generate schedulers: + +```bash +torc slurm generate --account my_project --profile research test_workflow.yaml +``` + +You should see the generated workflow with appropriate schedulers for each partition. + +## Step 7: Use Your Profile + +Now you can submit workflows using your custom profile: + +```bash +# Auto-detect the profile (if on the cluster) +torc submit-slurm --account my_project workflow.yaml + +# Or explicitly specify the profile +torc submit-slurm --account my_project --hpc-profile research workflow.yaml +``` + +## Advanced Configuration + +### Hostname-Based Detection + +If your cluster doesn't set a unique environment variable, use hostname detection: + +```toml +[client.hpc.custom_profiles.research] +display_name = "Research Cluster" +detect_hostname = ".*\\.research\\.edu" # Regex pattern +``` + +### Minimum Node Requirements + +Some partitions require a minimum number of nodes: + +```toml +[[client.hpc.custom_profiles.research.partitions]] +name = "large_scale" +cpus_per_node = 128 +memory_mb = 512000 +max_walltime_secs = 172800 +min_nodes = 16 # Must request at least 16 nodes +``` + +### Explicit Request Partitions + +Some partitions shouldn't be auto-selected: + +```toml +[[client.hpc.custom_profiles.research.partitions]] +name = "priority" +cpus_per_node = 48 +memory_mb = 192000 +max_walltime_secs = 86400 +requires_explicit_request = true # Only used when explicitly requested +``` + +## Troubleshooting + +### Profile Not Detected + +If `torc hpc detect` doesn't find your profile: + +1. Check the environment variable or hostname: + ```bash + echo $CLUSTER_NAME + hostname + ``` + +2. Verify the detection pattern in your config matches exactly + +3. Test with explicit profile specification: + ```bash + torc hpc show research + ``` + +### No Partition Found for Job + +If `torc slurm generate` can't find a matching partition: + +1. Check if any partition satisfies all requirements: + ```bash + torc hpc match research --cpus 32 --memory 128g --walltime 8h + ``` + +2. Verify memory is specified in MB in the config (not GB) + +3. Verify walltime is in seconds (not hours) + +### Configuration File Location + +Torc looks for config files in these locations: + +- **Linux**: `~/.config/torc/config.toml` +- **macOS**: `~/Library/Application Support/torc/config.toml` +- **Windows**: `%APPDATA%\torc\config.toml` + +You can also use the `TORC_CONFIG` environment variable to specify a custom path. + +## Contributing Your Profile + +If your HPC is used by others, please contribute it upstream: + +1. Fork the [Torc repository](https://github.com/NREL/torc) +2. Add your profile to `src/client/hpc_profiles.rs` +3. Add tests for your profile +4. Submit a pull request + +Or simply [open an issue](https://github.com/NREL/torc/issues) with your partition information and we'll add it for you. + +## See Also + +- [Working with HPC Profiles](../how-to/hpc-profiles.md) - General HPC profile usage +- [HPC Profiles Reference](../reference/hpc-profiles.md) - Complete configuration options +- [Slurm Workflows](../explanation/slurm-workflows.md) - Simplified Slurm approach diff --git a/examples/kdl/hundred_jobs_parameterized.kdl b/examples/kdl/hundred_jobs_parameterized.kdl index 64326b34..f6d62943 100644 --- a/examples/kdl/hundred_jobs_parameterized.kdl +++ b/examples/kdl/hundred_jobs_parameterized.kdl @@ -6,20 +6,34 @@ name "hundred_jobs_parameterized" user "testuser" description "A workflow with 100 independent jobs using parameterization" -// Resource requirements - minimal resources for simple echo commands -resource_requirements "minimal" { +resource_requirements "work" { + num_cpus 50 + num_gpus 0 + num_nodes 1 + memory "20g" + runtime "P0DT1H" +} + +resource_requirements "small" { num_cpus 1 num_gpus 0 num_nodes 1 memory "1g" - runtime "P0DT1M" + runtime "P0DT1H" } + // Job definitions - ONE job specification expands to 100 jobs job "job_{i:03d}" { command "echo hello {i}" - resource_requirements "minimal" + resource_requirements "work" parameters { i "1:100" } } + +job "postprocess" { + command "echo postprocess" + resource_requirements "small" + depends_on_regexes "job_\\d+" +} diff --git a/examples/kdl/mixed_workload.kdl b/examples/kdl/mixed_workload.kdl new file mode 100644 index 00000000..f752ec36 --- /dev/null +++ b/examples/kdl/mixed_workload.kdl @@ -0,0 +1,53 @@ +// Simple Workflow with 100 Independent Jobs (Parameterized Version) +// This demonstrates how parameterization makes it easy to create many parallel jobs +// Compare this file to hundred_jobs_workflow.kdl - same result, 90% less code! + +name "mixed_workload" +description "A workflow with mixed job workloads" + +resource_requirements "short_work" { + num_cpus 50 + num_gpus 0 + num_nodes 1 + memory "20g" + runtime "P0DT1H" +} + +resource_requirements "long_work" { + num_cpus 10 + num_gpus 0 + num_nodes 1 + memory "50g" + runtime "P0DT48H" +} + +resource_requirements "small" { + num_cpus 1 + num_gpus 0 + num_nodes 1 + memory "1g" + runtime "P0DT1H" +} + + +job "short_job_{i:03d}" { + command "echo hello {i}" + resource_requirements "short_work" + parameters { + i "1:100" + } +} + +job "long_job_{i:03d}" { + command "echo hello {i}" + resource_requirements "long_work" + parameters { + i "1:20" + } +} + +job "postprocess" { + command "echo postprocess" + resource_requirements "small" + depends_on_regexes "short_job_\\d+" "long_job_\\d+" +} diff --git a/examples/kdl/resource_monitoring_demo.kdl b/examples/kdl/resource_monitoring_demo.kdl index 5a594ae0..adce4abe 100644 --- a/examples/kdl/resource_monitoring_demo.kdl +++ b/examples/kdl/resource_monitoring_demo.kdl @@ -14,26 +14,42 @@ resource_monitor { } // Resource requirements -resource_requirements "medium" { - num_cpus 2 +resource_requirements "cpu" { + num_cpus 3 num_gpus 0 num_nodes 1 - memory "2g" - runtime "PT5M" + memory "1g" + runtime "PT15M" +} + +resource_requirements "memory" { + num_cpus 1 + num_gpus 0 + num_nodes 1 + memory "5g" + runtime "PT10M" +} + +resource_requirements "mixed" { + num_cpus 3 + num_gpus 0 + num_nodes 1 + memory "5g" + runtime "PT15M" } // Three independent jobs that run concurrently job "cpu_heavy_job" { command "python3 examples/scripts/cpu_intensive.py" - resource_requirements "medium" + resource_requirements "cpu" } job "memory_heavy_job" { command "python3 examples/scripts/memory_intensive.py" - resource_requirements "medium" + resource_requirements "memory" } job "mixed_workload_job" { command "python3 examples/scripts/mixed_workload.py" - resource_requirements "medium" + resource_requirements "mixed" } diff --git a/examples/kdl/simple_slurm.kdl b/examples/kdl/simple_slurm.kdl new file mode 100644 index 00000000..0cd55112 --- /dev/null +++ b/examples/kdl/simple_slurm.kdl @@ -0,0 +1,41 @@ +// Simple Workflow with Slurm Actions in KDL format +// Demonstrates scheduler actions for different workflow stages + +name "Simple Workflow with Actions" +user "demo" +description "Demonstrates scheduler actions for different workflow stages" + +// Resource requirements +resource_requirements "small" { + num_cpus 2 + num_gpus 0 + num_nodes 1 + memory "4g" + runtime "PT30M" +} + +resource_requirements "medium" { + num_cpus 8 + num_gpus 0 + num_nodes 1 + memory "16g" + runtime "PT2H" +} + +// Jobs +job "setup" { + command "echo 'Setting up...' && sleep 2" + resource_requirements "small" +} + +job "process" { + command "echo 'Processing data...' && sleep 5" + depends_on_job "setup" + resource_requirements "medium" +} + +job "finalize" { + command "echo 'Finalizing...' && sleep 2" + depends_on_job "process" + resource_requirements "small" +} diff --git a/examples/kdl/workflow_actions_ml_training.kdl b/examples/kdl/workflow_actions_ml_training.kdl index 602eea7c..0a0a0d37 100644 --- a/examples/kdl/workflow_actions_ml_training.kdl +++ b/examples/kdl/workflow_actions_ml_training.kdl @@ -90,7 +90,7 @@ action { action { trigger_type "on_jobs_ready" action_type "schedule_nodes" - job_name_regex "train_model_.*" + job_name_regexes "train_model_.*" scheduler "gpu_cluster" scheduler_type "slurm" num_allocations 2 @@ -102,7 +102,7 @@ action { action { trigger_type "on_jobs_complete" action_type "run_commands" - job_name_regex "train_model_.*" + job_name_regexes "train_model_.*" command "echo 'All training jobs completed. Archiving models...'" command "tar -czf results/trained_models.tar.gz models/" command "echo 'Generating training summary...'" diff --git a/examples/subgraphs/README.md b/examples/subgraphs/README.md new file mode 100644 index 00000000..336f6406 --- /dev/null +++ b/examples/subgraphs/README.md @@ -0,0 +1,190 @@ +# Sub-graph Workflow Example + +This directory contains example workflow specifications demonstrating **2 independent sub-graphs** with **4 execution stages**. The same workflow is provided in all supported formats: + +- `subgraphs_workflow.json5` - JSON5 format (with comments) +- `subgraphs_workflow.yaml` - YAML format +- `subgraphs_workflow.kdl` - KDL format + +## Workflow Structure + +``` + ┌─────────────────────────────────────────┐ + │ STAGE 1 (prep) │ + │ 1 Slurm node (shared) │ + │ │ + │ ┌─────────┐ ┌─────────┐ │ + input_a.txt ───>│ │ prep_a │ │ prep_b │<───── input_b.txt + │ └────┬────┘ └────┬────┘ │ + └────────┼─────────────────┼──────────────┘ + │ │ + prep_a_out.txt prep_b_out.txt + │ │ + ┌───────────────────┴──┐ ┌──┴───────────────────┐ + │ │ │ │ + ▼ │ │ ▼ +┌────────────────────┐ │ │ ┌────────────────────┐ +│ STAGE 2 (work_a) │ │ │ │ STAGE 2 (work_b) │ +│ 3 Slurm nodes │ │ │ │ 2 GPU nodes │ +│ │ │ │ │ │ +│ ┌──────────────┐ │ │ │ │ ┌──────────────┐ │ +│ │ work_a_1..5 │ │ │ │ │ │ work_b_1..5 │ │ +│ │ (5 jobs) │ │ │ │ │ │ (5 jobs) │ │ +│ └──────┬───────┘ │ │ │ │ └──────┬───────┘ │ +└─────────┼──────────┘ │ │ └─────────┼──────────┘ + │ │ │ │ + work_a_*_out.txt │ │ work_b_*_out.txt + │ │ │ │ + ▼ │ │ ▼ +┌────────────────────┐ │ │ ┌────────────────────┐ +│ STAGE 3 (post_a) │ │ │ │ STAGE 3 (post_b) │ +│ 1 Slurm node │ │ │ │ 1 Slurm node │ +│ │ │ │ │ │ +│ ┌────────┐ │ │ │ │ ┌────────┐ │ +│ │ post_a │ │ │ │ │ │ post_b │ │ +│ └───┬────┘ │ │ │ │ └───┬────┘ │ +└────────┼───────────┘ │ │ └──────────┼─────────┘ + │ │ │ │ + post_a_out.txt │ │ post_b_out.txt + │ │ │ │ + └──────────────────────┼───────────┼─────────────────────┘ + │ │ + ▼ ▼ + ┌─────────────────────────────────────────┐ + │ STAGE 4 (final) │ + │ 1 Slurm node │ + │ │ + │ ┌───────┐ │ + │ │ final │ │ + │ └───┬───┘ │ + └──────────────────┼──────────────────────┘ + │ + final_out.txt +``` + +## Key Features Demonstrated + +### 1. Implicit File Dependencies +All job dependencies are defined through `input_files` and `output_files` rather than explicit `depends_on`: +- `prep_a` reads `input_a.txt` and writes `prep_a_out.txt` +- `work_a_*` jobs read `prep_a_out.txt` (implicit dependency on `prep_a`) +- `post_a` reads all `work_a_*_out.txt` files (implicit dependency on all work_a jobs) + +### 2. Two Independent Sub-graphs +The workflow splits into two parallel processing pipelines: +- **Sub-graph A**: `prep_a` → `work_a_1..5` → `post_a` (CPU-intensive) +- **Sub-graph B**: `prep_b` → `work_b_1..5` → `post_b` (GPU-accelerated) + +These sub-graphs run completely independently until they merge at `final`. + +### 3. Parameterized Jobs +Work jobs use parameters to generate multiple instances: +```yaml +name: "work_a_{i}" +parameters: + i: "1:5" # Creates work_a_1, work_a_2, ..., work_a_5 +``` + +### 4. Parameterized Files +Output files also use parameters: +```yaml +name: "work_a_{i}_out" +path: "output/work_a_{i}.txt" +parameters: + i: "1:5" # Creates work_a_1_out, work_a_2_out, ..., work_a_5_out +``` + +### 5. Different Resource Requirements +Jobs have different resource profiles: +- `small`: 1 CPU, 2GB RAM (prep jobs) +- `work_large`: 8 CPUs, 32GB RAM (CPU work jobs) +- `work_gpu`: 4 CPUs, 16GB RAM, 1 GPU (GPU work jobs) +- `medium`: 2 CPUs, 8GB RAM (post jobs) +- `large`: 4 CPUs, 16GB RAM (final job) + +### 6. Stage-aware Slurm Scheduling +Each stage has its own scheduler action: +- **Stage 1**: `on_workflow_start` triggers `prep_sched` (1 node for both prep jobs) +- **Stage 2**: `on_jobs_ready` triggers `work_a_sched` (3 nodes) AND `work_b_sched` (2 GPU nodes) simultaneously +- **Stage 3**: `on_jobs_ready` triggers separate schedulers for `post_a` and `post_b` +- **Stage 4**: `on_jobs_ready` triggers `final_sched` (1 node) + +## Running the Example + +### With pre-defined schedulers + +View the execution plan: +```bash +torc workflows execution-plan examples/subgraphs/subgraphs_workflow.yaml +``` + +Output: +``` +Workflow: two_subgraph_pipeline +Total Jobs: 15 + +▶ Stage 1: Workflow Start + Scheduler Allocations: + • prep_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • prep_a + • prep_b + +→ Stage 2: When jobs 'prep_a', 'prep_b' complete + Scheduler Allocations: + • work_a_sched (slurm) - 1 allocation(s) + • work_b_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • work_a_{1..5} + • work_b_{1..5} + +→ Stage 3: When 10 jobs complete + Scheduler Allocations: + • post_a_sched (slurm) - 1 allocation(s) + • post_b_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • post_a + • post_b + +→ Stage 4: When jobs 'post_a', 'post_b' complete + Scheduler Allocations: + • final_sched (slurm) - 1 allocation(s) + Jobs Becoming Ready: + • final + +Total Stages: 4 +``` + +### Auto-generating Slurm schedulers + +The `*_no_slurm.*` files contain the same workflow without Slurm schedulers or actions. +Use `torc slurm generate` to auto-generate them: + +```bash +torc slurm generate --account myproject --profile kestrel examples/subgraphs/subgraphs_workflow_no_slurm.yaml +``` + +This will: +1. Expand parameterized jobs and files +2. Analyze the workflow graph for dependencies +3. Group jobs by (resource_requirements, has_dependencies) +4. Generate appropriate schedulers and `on_workflow_start`/`on_jobs_ready` actions + +Output shows 5 schedulers created: +- `small_scheduler` (prep jobs, `on_workflow_start`) +- `work_large_deferred_scheduler` (work_a jobs, `on_jobs_ready`) +- `work_gpu_deferred_scheduler` (work_b jobs, `on_jobs_ready`) +- `medium_deferred_scheduler` (post jobs, `on_jobs_ready`) +- `large_deferred_scheduler` (final job, `on_jobs_ready`) + +## Total Resources + +| Stage | Scheduler | Nodes | Partition | Purpose | +|-------|-----------|-------|-----------|---------| +| 1 | prep_sched | 1 | standard | Run both prep jobs | +| 2 | work_a_sched | 3 | standard | Run 5 CPU work jobs | +| 2 | work_b_sched | 2 | gpu | Run 5 GPU work jobs | +| 3 | post_a_sched | 1 | standard | Post-process sub-graph A | +| 3 | post_b_sched | 1 | standard | Post-process sub-graph B | +| 4 | final_sched | 1 | standard | Aggregate results | +| **Total** | | **9** | | | diff --git a/examples/subgraphs/subgraphs_workflow.json5 b/examples/subgraphs/subgraphs_workflow.json5 new file mode 100644 index 00000000..2e803210 --- /dev/null +++ b/examples/subgraphs/subgraphs_workflow.json5 @@ -0,0 +1,213 @@ +{ + name: "two_subgraph_pipeline", + description: "Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies", + + // ============================================================================ + // FILES - Define all input/output files. Dependencies are implicit based on + // which jobs read/write these files. + // ============================================================================ + files: [ + // Stage 1 inputs (must exist before workflow starts) + {name: "input_a", path: "data/input_a.txt", must_exist: true}, + {name: "input_b", path: "data/input_b.txt", must_exist: true}, + + // Stage 1 outputs -> Stage 2 inputs + {name: "prep_a_out", path: "output/prep_a.txt"}, + {name: "prep_b_out", path: "output/prep_b.txt"}, + + // Stage 2 outputs -> Stage 3 inputs (parameterized) + {name: "work_a_{i}_out", path: "output/work_a_{i}.txt", parameters: {i: "1:5"}}, + {name: "work_b_{i}_out", path: "output/work_b_{i}.txt", parameters: {i: "1:5"}}, + + // Stage 3 outputs -> Stage 4 inputs + {name: "post_a_out", path: "output/post_a.txt"}, + {name: "post_b_out", path: "output/post_b.txt"}, + + // Stage 4 output (final result) + {name: "final_out", path: "output/final.txt"}, + ], + + // ============================================================================ + // JOBS - Organized by stage. Dependencies are implicit via input_files. + // ============================================================================ + jobs: [ + // --- Stage 1: Preprocessing (both run on same node) --- + { + name: "prep_a", + command: "./scripts/prep.sh a", + input_files: ["input_a"], + output_files: ["prep_a_out"], + resource_requirements: "small", + }, + { + name: "prep_b", + command: "./scripts/prep.sh b", + input_files: ["input_b"], + output_files: ["prep_b_out"], + resource_requirements: "small", + }, + + // --- Stage 2: Work (two independent sub-graphs) --- + // Sub-graph A: CPU-intensive work + { + name: "work_a_{i}", + command: "./scripts/work.sh a {i}", + input_files: ["prep_a_out"], + output_files: ["work_a_{i}_out"], + resource_requirements: "work_large", + parameters: {i: "1:5"}, + }, + // Sub-graph B: GPU-accelerated work + { + name: "work_b_{i}", + command: "./scripts/work.sh b {i}", + input_files: ["prep_b_out"], + output_files: ["work_b_{i}_out"], + resource_requirements: "work_gpu", + parameters: {i: "1:5"}, + }, + + // --- Stage 3: Post-processing (one per sub-graph) --- + { + name: "post_a", + command: "./scripts/post.sh a", + input_files: ["work_a_1_out", "work_a_2_out", "work_a_3_out", "work_a_4_out", "work_a_5_out"], + output_files: ["post_a_out"], + resource_requirements: "medium", + }, + { + name: "post_b", + command: "./scripts/post.sh b", + input_files: ["work_b_1_out", "work_b_2_out", "work_b_3_out", "work_b_4_out", "work_b_5_out"], + output_files: ["post_b_out"], + resource_requirements: "medium", + }, + + // --- Stage 4: Final aggregation --- + { + name: "final", + command: "./scripts/aggregate.sh", + input_files: ["post_a_out", "post_b_out"], + output_files: ["final_out"], + resource_requirements: "large", + }, + ], + + // ============================================================================ + // RESOURCE REQUIREMENTS - Different profiles for different job types + // ============================================================================ + resource_requirements: [ + {name: "small", num_cpus: 1, memory: "2g", runtime: "PT30M"}, + {name: "work_large", num_cpus: 8, memory: "32g", runtime: "PT2H"}, + {name: "work_gpu", num_cpus: 4, memory: "16g", num_gpus: 1, runtime: "PT4H"}, + {name: "medium", num_cpus: 2, memory: "8g", runtime: "PT1H"}, + {name: "large", num_cpus: 4, memory: "16g", runtime: "PT2H"}, + ], + + // ============================================================================ + // SLURM SCHEDULERS - Each stage/sub-graph gets its own scheduler + // ============================================================================ + slurm_schedulers: [ + // Stage 1: Both prep jobs share one node + { + name: "prep_sched", + account: "myproject", + partition: "standard", + num_nodes: 1, + walltime: "01:00:00", + }, + // Stage 2: Sub-graph A gets 3 CPU nodes + { + name: "work_a_sched", + account: "myproject", + partition: "standard", + num_nodes: 3, + walltime: "04:00:00", + }, + // Stage 2: Sub-graph B gets 2 GPU nodes + { + name: "work_b_sched", + account: "myproject", + partition: "gpu", + num_nodes: 2, + walltime: "06:00:00", + extra_args: "--gres=gpu:1", + }, + // Stage 3: Each post job gets its own node + { + name: "post_a_sched", + account: "myproject", + partition: "standard", + num_nodes: 1, + walltime: "02:00:00", + }, + { + name: "post_b_sched", + account: "myproject", + partition: "standard", + num_nodes: 1, + walltime: "02:00:00", + }, + // Stage 4: Final job gets its own node + { + name: "final_sched", + account: "myproject", + partition: "standard", + num_nodes: 1, + walltime: "03:00:00", + }, + ], + + // ============================================================================ + // ACTIONS - Schedule compute nodes when jobs become ready + // ============================================================================ + actions: [ + // Stage 1: Triggered at workflow start + { + trigger_type: "on_workflow_start", + action_type: "schedule_nodes", + scheduler: "prep_sched", + scheduler_type: "slurm", + jobs: ["prep_a", "prep_b"], + }, + // Stage 2: Triggered when work jobs become ready (after prep completes) + // These trigger simultaneously since sub-graphs are independent + { + trigger_type: "on_jobs_ready", + action_type: "schedule_nodes", + scheduler: "work_a_sched", + scheduler_type: "slurm", + job_name_regexes: ["^work_a_\\d+$"], + }, + { + trigger_type: "on_jobs_ready", + action_type: "schedule_nodes", + scheduler: "work_b_sched", + scheduler_type: "slurm", + job_name_regexes: ["^work_b_\\d+$"], + }, + // Stage 3: Triggered when post jobs become ready (after all work completes) + { + trigger_type: "on_jobs_ready", + action_type: "schedule_nodes", + scheduler: "post_a_sched", + scheduler_type: "slurm", + jobs: ["post_a"], + }, + { + trigger_type: "on_jobs_ready", + action_type: "schedule_nodes", + scheduler: "post_b_sched", + scheduler_type: "slurm", + jobs: ["post_b"], + }, + // Stage 4: Triggered when final job becomes ready + { + trigger_type: "on_jobs_ready", + action_type: "schedule_nodes", + scheduler: "final_sched", + scheduler_type: "slurm", + jobs: ["final"], + }, + ], +} diff --git a/examples/subgraphs/subgraphs_workflow.kdl b/examples/subgraphs/subgraphs_workflow.kdl new file mode 100644 index 00000000..6fe26d64 --- /dev/null +++ b/examples/subgraphs/subgraphs_workflow.kdl @@ -0,0 +1,262 @@ +// Two Sub-graph Pipeline +// Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies +// +// Structure: +// Stage 1: prep_a, prep_b (run on 1 shared node) +// Stage 2: work_a_1..5, work_b_1..5 (2 independent sub-graphs, different schedulers) +// Stage 3: post_a, post_b (each on its own node) +// Stage 4: final (aggregates both sub-graphs) + +name "two_subgraph_pipeline" +description "Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies" + +// ========================================================================== +// FILES - All dependencies are implicit based on input_files/output_files +// ========================================================================== + +// Stage 1 inputs (must exist before workflow starts) +file "input_a" path="input_a.txt" must_exist=#true +file "input_b" path="input_b.txt" must_exist=#true + +// Stage 1 outputs -> Stage 2 inputs +file "prep_a_out" path="output/prep_a.txt" +file "prep_b_out" path="output/prep_b.txt" + +// Stage 2 outputs -> Stage 3 inputs (parameterized) +file "work_a_{i}_out" path="output/work_a_{i}.txt" { + parameters { + i "1:5" + } +} +file "work_b_{i}_out" path="output/work_b_{i}.txt" { + parameters { + i "1:5" + } +} + +// Stage 3 outputs -> Stage 4 inputs +file "post_a_out" path="output/post_a.txt" +file "post_b_out" path="output/post_b.txt" + +// Stage 4 output (final result) +file "final_out" path="output/final.txt" + +// ========================================================================== +// RESOURCE REQUIREMENTS +// ========================================================================== + +resource_requirements "small" { + num_cpus 1 + memory "2g" + runtime "PT30M" +} + +resource_requirements "work_large" { + num_cpus 8 + memory "32g" + runtime "PT2H" +} + +resource_requirements "work_gpu" { + num_cpus 4 + memory "16g" + num_gpus 1 + runtime "PT4H" +} + +resource_requirements "medium" { + num_cpus 2 + memory "8g" + runtime "PT1H" +} + +resource_requirements "large" { + num_cpus 4 + memory "16g" + runtime "PT2H" +} + +// ========================================================================== +// SLURM SCHEDULERS - Each stage/sub-graph gets its own scheduler +// ========================================================================== + +// Stage 1: Both prep jobs share one node +slurm_scheduler "prep_sched" { + account "myproject" + partition "standard" + nodes 1 + walltime "01:00:00" +} + +// Stage 2: Sub-graph A gets 3 CPU nodes +slurm_scheduler "work_a_sched" { + account "myproject" + partition "standard" + nodes 3 + walltime "04:00:00" +} + +// Stage 2: Sub-graph B gets 2 GPU nodes +slurm_scheduler "work_b_sched" { + account "myproject" + partition "gpu" + nodes 2 + walltime "06:00:00" + extra "--gres=gpu:1" +} + +// Stage 3: Each post job gets its own node +slurm_scheduler "post_a_sched" { + account "myproject" + partition "standard" + nodes 1 + walltime "02:00:00" +} + +slurm_scheduler "post_b_sched" { + account "myproject" + partition "standard" + nodes 1 + walltime "02:00:00" +} + +// Stage 4: Final job gets its own node +slurm_scheduler "final_sched" { + account "myproject" + partition "standard" + nodes 1 + walltime "03:00:00" +} + +// ========================================================================== +// JOBS - Organized by stage +// ========================================================================== + +// --- Stage 1: Preprocessing --- +job "prep_a" { + command "./scripts/prep.sh a" + input_file "input_a" + output_file "prep_a_out" + resource_requirements "small" +} + +job "prep_b" { + command "./scripts/prep.sh b" + input_file "input_b" + output_file "prep_b_out" + resource_requirements "small" +} + +// --- Stage 2: Work (two independent sub-graphs) --- + +// Sub-graph A: CPU-intensive work +job "work_a_{i}" { + command "./scripts/work.sh a {i}" + input_file "prep_a_out" + output_file "work_a_{i}_out" + resource_requirements "work_large" + parameters { + i "1:5" + } +} + +// Sub-graph B: GPU-accelerated work +job "work_b_{i}" { + command "./scripts/work.sh b {i}" + input_file "prep_b_out" + output_file "work_b_{i}_out" + resource_requirements "work_gpu" + parameters { + i "1:5" + } +} + +// --- Stage 3: Post-processing --- +job "post_a" { + command "./scripts/post.sh a" + input_file "work_a_1_out" + input_file "work_a_2_out" + input_file "work_a_3_out" + input_file "work_a_4_out" + input_file "work_a_5_out" + output_file "post_a_out" + resource_requirements "medium" +} + +job "post_b" { + command "./scripts/post.sh b" + input_file "work_b_1_out" + input_file "work_b_2_out" + input_file "work_b_3_out" + input_file "work_b_4_out" + input_file "work_b_5_out" + output_file "post_b_out" + resource_requirements "medium" +} + +// --- Stage 4: Final aggregation --- +job "final" { + command "./scripts/aggregate.sh" + input_file "post_a_out" + input_file "post_b_out" + output_file "final_out" + resource_requirements "large" +} + +// ========================================================================== +// ACTIONS - Schedule compute nodes when jobs become ready +// ========================================================================== + +// Stage 1: Triggered at workflow start +action { + trigger_type "on_workflow_start" + action_type "schedule_nodes" + scheduler "prep_sched" + scheduler_type "slurm" + job "prep_a" + job "prep_b" +} + +// Stage 2: Triggered when work jobs become ready +// These trigger simultaneously since sub-graphs are independent +action { + trigger_type "on_jobs_ready" + action_type "schedule_nodes" + scheduler "work_a_sched" + scheduler_type "slurm" + job_name_regexes "^work_a_\\d+$" +} + +action { + trigger_type "on_jobs_ready" + action_type "schedule_nodes" + scheduler "work_b_sched" + scheduler_type "slurm" + job_name_regexes "^work_b_\\d+$" +} + +// Stage 3: Triggered when post jobs become ready +action { + trigger_type "on_jobs_ready" + action_type "schedule_nodes" + scheduler "post_a_sched" + scheduler_type "slurm" + job "post_a" +} + +action { + trigger_type "on_jobs_ready" + action_type "schedule_nodes" + scheduler "post_b_sched" + scheduler_type "slurm" + job "post_b" +} + +// Stage 4: Triggered when final job becomes ready +action { + trigger_type "on_jobs_ready" + action_type "schedule_nodes" + scheduler "final_sched" + scheduler_type "slurm" + job "final" +} diff --git a/examples/subgraphs/subgraphs_workflow.yaml b/examples/subgraphs/subgraphs_workflow.yaml new file mode 100644 index 00000000..336c5602 --- /dev/null +++ b/examples/subgraphs/subgraphs_workflow.yaml @@ -0,0 +1,240 @@ +# Two Sub-graph Pipeline +# Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies +# +# Structure: +# Stage 1: prep_a, prep_b (run on 1 shared node) +# Stage 2: work_a_1..5, work_b_1..5 (2 independent sub-graphs, different schedulers) +# Stage 3: post_a, post_b (each on its own node) +# Stage 4: final (aggregates both sub-graphs) + +name: two_subgraph_pipeline +description: Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies + +# ============================================================================== +# FILES - All dependencies are implicit based on input_files/output_files +# ============================================================================== +files: + # Stage 1 inputs (must exist before workflow starts) + - name: input_a + path: input_a.txt + must_exist: true + + - name: input_b + path: input_b.txt + must_exist: true + + # Stage 1 outputs -> Stage 2 inputs + - name: prep_a_out + path: output/prep_a.txt + + - name: prep_b_out + path: output/prep_b.txt + + # Stage 2 outputs -> Stage 3 inputs (parameterized) + - name: "work_a_{i}_out" + path: "output/work_a_{i}.txt" + parameters: + i: "1:5" + + - name: "work_b_{i}_out" + path: "output/work_b_{i}.txt" + parameters: + i: "1:5" + + # Stage 3 outputs -> Stage 4 inputs + - name: post_a_out + path: output/post_a.txt + + - name: post_b_out + path: output/post_b.txt + + # Stage 4 output (final result) + - name: final_out + path: output/final.txt + +# ============================================================================== +# JOBS - Organized by stage +# ============================================================================== +jobs: + # --- Stage 1: Preprocessing --- + - name: prep_a + command: ./scripts/prep.sh a + input_files: [input_a] + output_files: [prep_a_out] + resource_requirements: small + + - name: prep_b + command: ./scripts/prep.sh b + input_files: [input_b] + output_files: [prep_b_out] + resource_requirements: small + + # --- Stage 2: Work (two independent sub-graphs) --- + # Sub-graph A: CPU-intensive work + - name: "work_a_{i}" + command: "./scripts/work.sh a {i}" + input_files: [prep_a_out] + output_files: ["work_a_{i}_out"] + resource_requirements: work_large + parameters: + i: "1:5" + + # Sub-graph B: GPU-accelerated work + - name: "work_b_{i}" + command: "./scripts/work.sh b {i}" + input_files: [prep_b_out] + output_files: ["work_b_{i}_out"] + resource_requirements: work_gpu + parameters: + i: "1:5" + + # --- Stage 3: Post-processing --- + - name: post_a + command: ./scripts/post.sh a + input_files: + - work_a_1_out + - work_a_2_out + - work_a_3_out + - work_a_4_out + - work_a_5_out + output_files: [post_a_out] + resource_requirements: medium + + - name: post_b + command: ./scripts/post.sh b + input_files: + - work_b_1_out + - work_b_2_out + - work_b_3_out + - work_b_4_out + - work_b_5_out + output_files: [post_b_out] + resource_requirements: medium + + # --- Stage 4: Final aggregation --- + - name: final + command: ./scripts/aggregate.sh + input_files: [post_a_out, post_b_out] + output_files: [final_out] + resource_requirements: large + +# ============================================================================== +# RESOURCE REQUIREMENTS +# ============================================================================== +resource_requirements: + - name: small + num_cpus: 1 + memory: 2g + runtime: PT30M + + - name: work_large + num_cpus: 8 + memory: 32g + runtime: PT2H + + - name: work_gpu + num_cpus: 4 + memory: 16g + num_gpus: 1 + runtime: PT4H + + - name: medium + num_cpus: 2 + memory: 8g + runtime: PT1H + + - name: large + num_cpus: 4 + memory: 16g + runtime: PT2H + +# ============================================================================== +# SLURM SCHEDULERS - Each stage/sub-graph gets its own scheduler +# ============================================================================== +slurm_schedulers: + # Stage 1: Both prep jobs share one node + - name: prep_sched + account: myproject + partition: standard + num_nodes: 1 + walltime: "01:00:00" + + # Stage 2: Sub-graph A gets 3 CPU nodes + - name: work_a_sched + account: myproject + partition: standard + num_nodes: 3 + walltime: "04:00:00" + + # Stage 2: Sub-graph B gets 2 GPU nodes + - name: work_b_sched + account: myproject + partition: gpu + num_nodes: 2 + walltime: "06:00:00" + extra_args: "--gres=gpu:1" + + # Stage 3: Each post job gets its own node + - name: post_a_sched + account: myproject + partition: standard + num_nodes: 1 + walltime: "02:00:00" + + - name: post_b_sched + account: myproject + partition: standard + num_nodes: 1 + walltime: "02:00:00" + + # Stage 4: Final job gets its own node + - name: final_sched + account: myproject + partition: standard + num_nodes: 1 + walltime: "03:00:00" + +# ============================================================================== +# ACTIONS - Schedule compute nodes when jobs become ready +# ============================================================================== +actions: + # Stage 1: Triggered at workflow start + - trigger_type: on_workflow_start + action_type: schedule_nodes + scheduler: prep_sched + scheduler_type: slurm + jobs: [prep_a, prep_b] + + # Stage 2: Triggered when work jobs become ready + # These trigger simultaneously since sub-graphs are independent + - trigger_type: on_jobs_ready + action_type: schedule_nodes + scheduler: work_a_sched + scheduler_type: slurm + job_name_regexes: ["^work_a_\\d+$"] + + - trigger_type: on_jobs_ready + action_type: schedule_nodes + scheduler: work_b_sched + scheduler_type: slurm + job_name_regexes: ["^work_b_\\d+$"] + + # Stage 3: Triggered when post jobs become ready + - trigger_type: on_jobs_ready + action_type: schedule_nodes + scheduler: post_a_sched + scheduler_type: slurm + jobs: [post_a] + + - trigger_type: on_jobs_ready + action_type: schedule_nodes + scheduler: post_b_sched + scheduler_type: slurm + jobs: [post_b] + + # Stage 4: Triggered when final job becomes ready + - trigger_type: on_jobs_ready + action_type: schedule_nodes + scheduler: final_sched + scheduler_type: slurm + jobs: [final] diff --git a/examples/subgraphs/subgraphs_workflow_no_slurm.json5 b/examples/subgraphs/subgraphs_workflow_no_slurm.json5 new file mode 100644 index 00000000..ae4259db --- /dev/null +++ b/examples/subgraphs/subgraphs_workflow_no_slurm.json5 @@ -0,0 +1,97 @@ +{ + // Two Sub-graph Pipeline (without Slurm schedulers/actions) + // Run `torc slurm generate --account subgraphs_workflow_no_slurm.json5` to auto-generate + + name: "two_subgraph_pipeline", + description: "Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies", + + files: [ + // Stage 1 inputs + {name: "input_a", path: "input_a.txt", must_exist: true}, + {name: "input_b", path: "input_b.txt", must_exist: true}, + + // Stage 1 outputs + {name: "prep_a_out", path: "output/prep_a.txt"}, + {name: "prep_b_out", path: "output/prep_b.txt"}, + + // Stage 2 outputs (parameterized) + {name: "work_a_{i}_out", path: "output/work_a_{i}.txt", parameters: {i: "1:5"}}, + {name: "work_b_{i}_out", path: "output/work_b_{i}.txt", parameters: {i: "1:5"}}, + + // Stage 3 outputs + {name: "post_a_out", path: "output/post_a.txt"}, + {name: "post_b_out", path: "output/post_b.txt"}, + + // Stage 4 output + {name: "final_out", path: "output/final.txt"}, + ], + + jobs: [ + // Stage 1: Preprocessing + { + name: "prep_a", + command: "./scripts/prep.sh a", + input_files: ["input_a"], + output_files: ["prep_a_out"], + resource_requirements: "small", + }, + { + name: "prep_b", + command: "./scripts/prep.sh b", + input_files: ["input_b"], + output_files: ["prep_b_out"], + resource_requirements: "small", + }, + + // Stage 2: Work (two independent sub-graphs) + { + name: "work_a_{i}", + command: "./scripts/work.sh a {i}", + input_files: ["prep_a_out"], + output_files: ["work_a_{i}_out"], + resource_requirements: "work_large", + parameters: {i: "1:5"}, + }, + { + name: "work_b_{i}", + command: "./scripts/work.sh b {i}", + input_files: ["prep_b_out"], + output_files: ["work_b_{i}_out"], + resource_requirements: "work_gpu", + parameters: {i: "1:5"}, + }, + + // Stage 3: Post-processing + { + name: "post_a", + command: "./scripts/post.sh a", + input_files: ["work_a_1_out", "work_a_2_out", "work_a_3_out", "work_a_4_out", "work_a_5_out"], + output_files: ["post_a_out"], + resource_requirements: "medium", + }, + { + name: "post_b", + command: "./scripts/post.sh b", + input_files: ["work_b_1_out", "work_b_2_out", "work_b_3_out", "work_b_4_out", "work_b_5_out"], + output_files: ["post_b_out"], + resource_requirements: "medium", + }, + + // Stage 4: Final aggregation + { + name: "final", + command: "./scripts/aggregate.sh", + input_files: ["post_a_out", "post_b_out"], + output_files: ["final_out"], + resource_requirements: "large", + }, + ], + + resource_requirements: [ + {name: "small", num_cpus: 1, memory: "2g", runtime: "PT30M"}, + {name: "work_large", num_cpus: 8, memory: "32g", runtime: "PT2H"}, + {name: "work_gpu", num_cpus: 4, memory: "16g", num_gpus: 1, runtime: "PT4H"}, + {name: "medium", num_cpus: 2, memory: "8g", runtime: "PT1H"}, + {name: "large", num_cpus: 4, memory: "16g", runtime: "PT2H"}, + ], +} diff --git a/examples/subgraphs/subgraphs_workflow_no_slurm.kdl b/examples/subgraphs/subgraphs_workflow_no_slurm.kdl new file mode 100644 index 00000000..9bb46ad0 --- /dev/null +++ b/examples/subgraphs/subgraphs_workflow_no_slurm.kdl @@ -0,0 +1,132 @@ +// Two Sub-graph Pipeline (without Slurm schedulers/actions) +// Run `torc slurm generate --account subgraphs_workflow_no_slurm.kdl` to auto-generate + +name "two_subgraph_pipeline" +description "Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies" + +// Stage 1 inputs +file "input_a" path="input_a.txt" must_exist=#true +file "input_b" path="input_b.txt" must_exist=#true + +// Stage 1 outputs +file "prep_a_out" path="output/prep_a.txt" +file "prep_b_out" path="output/prep_b.txt" + +// Stage 2 outputs (parameterized) +file "work_a_{i}_out" path="output/work_a_{i}.txt" { + parameters { + i "1:5" + } +} +file "work_b_{i}_out" path="output/work_b_{i}.txt" { + parameters { + i "1:5" + } +} + +// Stage 3 outputs +file "post_a_out" path="output/post_a.txt" +file "post_b_out" path="output/post_b.txt" + +// Stage 4 output +file "final_out" path="output/final.txt" + +// Resource requirements +resource_requirements "small" { + num_cpus 1 + memory "2g" + runtime "PT30M" +} + +resource_requirements "work_large" { + num_cpus 8 + memory "32g" + runtime "PT2H" +} + +resource_requirements "work_gpu" { + num_cpus 4 + memory "16g" + num_gpus 1 + runtime "PT4H" +} + +resource_requirements "medium" { + num_cpus 2 + memory "8g" + runtime "PT1H" +} + +resource_requirements "large" { + num_cpus 4 + memory "16g" + runtime "PT2H" +} + +// Stage 1: Preprocessing +job "prep_a" { + command "./scripts/prep.sh a" + input_file "input_a" + output_file "prep_a_out" + resource_requirements "small" +} + +job "prep_b" { + command "./scripts/prep.sh b" + input_file "input_b" + output_file "prep_b_out" + resource_requirements "small" +} + +// Stage 2: Work (two independent sub-graphs) +job "work_a_{i}" { + command "./scripts/work.sh a {i}" + input_file "prep_a_out" + output_file "work_a_{i}_out" + resource_requirements "work_large" + parameters { + i "1:5" + } +} + +job "work_b_{i}" { + command "./scripts/work.sh b {i}" + input_file "prep_b_out" + output_file "work_b_{i}_out" + resource_requirements "work_gpu" + parameters { + i "1:5" + } +} + +// Stage 3: Post-processing +job "post_a" { + command "./scripts/post.sh a" + input_file "work_a_1_out" + input_file "work_a_2_out" + input_file "work_a_3_out" + input_file "work_a_4_out" + input_file "work_a_5_out" + output_file "post_a_out" + resource_requirements "medium" +} + +job "post_b" { + command "./scripts/post.sh b" + input_file "work_b_1_out" + input_file "work_b_2_out" + input_file "work_b_3_out" + input_file "work_b_4_out" + input_file "work_b_5_out" + output_file "post_b_out" + resource_requirements "medium" +} + +// Stage 4: Final aggregation +job "final" { + command "./scripts/aggregate.sh" + input_file "post_a_out" + input_file "post_b_out" + output_file "final_out" + resource_requirements "large" +} diff --git a/examples/subgraphs/subgraphs_workflow_no_slurm.yaml b/examples/subgraphs/subgraphs_workflow_no_slurm.yaml new file mode 100644 index 00000000..83e6ae2f --- /dev/null +++ b/examples/subgraphs/subgraphs_workflow_no_slurm.yaml @@ -0,0 +1,132 @@ +# Two Sub-graph Pipeline (without Slurm schedulers/actions) +# Run `torc slurm generate --account subgraphs_workflow_no_slurm.yaml` to auto-generate + +name: two_subgraph_pipeline +description: Demonstrates 2 independent sub-graphs with 4 stages, implicit file dependencies + +files: + # Stage 1 inputs + - name: input_a + path: input_a.txt + must_exist: true + + - name: input_b + path: input_b.txt + must_exist: true + + # Stage 1 outputs + - name: prep_a_out + path: output/prep_a.txt + + - name: prep_b_out + path: output/prep_b.txt + + # Stage 2 outputs (parameterized) + - name: "work_a_{i}_out" + path: "output/work_a_{i}.txt" + parameters: + i: "1:5" + + - name: "work_b_{i}_out" + path: "output/work_b_{i}.txt" + parameters: + i: "1:5" + + # Stage 3 outputs + - name: post_a_out + path: output/post_a.txt + + - name: post_b_out + path: output/post_b.txt + + # Stage 4 output + - name: final_out + path: output/final.txt + +jobs: + # Stage 1: Preprocessing + - name: prep_a + command: ./scripts/prep.sh a + input_files: [input_a] + output_files: [prep_a_out] + resource_requirements: small + + - name: prep_b + command: ./scripts/prep.sh b + input_files: [input_b] + output_files: [prep_b_out] + resource_requirements: small + + # Stage 2: Work (two independent sub-graphs) + - name: "work_a_{i}" + command: "./scripts/work.sh a {i}" + input_files: [prep_a_out] + output_files: ["work_a_{i}_out"] + resource_requirements: work_large + parameters: + i: "1:5" + + - name: "work_b_{i}" + command: "./scripts/work.sh b {i}" + input_files: [prep_b_out] + output_files: ["work_b_{i}_out"] + resource_requirements: work_gpu + parameters: + i: "1:5" + + # Stage 3: Post-processing + - name: post_a + command: ./scripts/post.sh a + input_files: + - work_a_1_out + - work_a_2_out + - work_a_3_out + - work_a_4_out + - work_a_5_out + output_files: [post_a_out] + resource_requirements: medium + + - name: post_b + command: ./scripts/post.sh b + input_files: + - work_b_1_out + - work_b_2_out + - work_b_3_out + - work_b_4_out + - work_b_5_out + output_files: [post_b_out] + resource_requirements: medium + + # Stage 4: Final aggregation + - name: final + command: ./scripts/aggregate.sh + input_files: [post_a_out, post_b_out] + output_files: [final_out] + resource_requirements: large + +resource_requirements: + - name: small + num_cpus: 1 + memory: 2g + runtime: PT30M + + - name: work_large + num_cpus: 8 + memory: 32g + runtime: PT2H + + - name: work_gpu + num_cpus: 4 + memory: 16g + num_gpus: 1 + runtime: PT4H + + - name: medium + num_cpus: 2 + memory: 8g + runtime: PT1H + + - name: large + num_cpus: 4 + memory: 16g + runtime: PT2H diff --git a/examples/yaml/resource_monitoring_demo.yaml b/examples/yaml/resource_monitoring_demo.yaml index 658bad67..362da3d9 100644 --- a/examples/yaml/resource_monitoring_demo.yaml +++ b/examples/yaml/resource_monitoring_demo.yaml @@ -24,15 +24,15 @@ resource_monitor: jobs: - name: cpu_heavy_job command: "python3 examples/scripts/cpu_intensive.py" - resource_requirements: medium + resource_requirements: cpu - name: memory_heavy_job command: "python3 examples/scripts/memory_intensive.py" - resource_requirements: medium + resource_requirements: memory - name: mixed_workload_job command: "python3 examples/scripts/mixed_workload.py" - resource_requirements: medium + resource_requirements: mixed # No files needed for this demo files: [] @@ -40,14 +40,22 @@ files: [] # No user data needed user_data: [] -# Resource requirements resource_requirements: - - name: medium + - name: cpu + num_cpus: 3 + num_gpus: 0 + num_nodes: 1 + memory: 1g + runtime: PT15M # ISO 8601 duration: 15 minutes + - name: memory num_cpus: 2 num_gpus: 0 num_nodes: 1 - memory: 2g - runtime: PT5M # ISO 8601 duration: 5 minutes - -# No Slurm schedulers (runs locally) -slurm_schedulers: [] + memory: 3g + runtime: PT10M + - name: mixed + num_cpus: 3 + num_gpus: 0 + num_nodes: 1 + memory: 3g + runtime: PT15M diff --git a/src/cli.rs b/src/cli.rs index 2f7eefa4..562daa7d 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -10,6 +10,7 @@ use crate::client::commands::compute_nodes::ComputeNodeCommands; use crate::client::commands::config::ConfigCommands; use crate::client::commands::events::EventCommands; use crate::client::commands::files::FileCommands; +use crate::client::commands::hpc::HpcCommands; use crate::client::commands::job_dependencies::JobDependencyCommands; use crate::client::commands::jobs::JobCommands; use crate::client::commands::reports::ReportCommands; @@ -83,7 +84,8 @@ pub enum Commands { }, /// Submit a workflow to scheduler (create from spec file or submit existing workflow by ID) /// - /// Requires workflow to have an on_workflow_start action with schedule_nodes + /// Requires workflow to have an on_workflow_start action with schedule_nodes. + /// For Slurm workflows without pre-configured schedulers, use `submit-slurm` instead. Submit { /// Path to workflow spec file (JSON/JSON5/YAML) or workflow ID #[arg()] @@ -95,6 +97,51 @@ pub enum Commands { #[arg(long, default_value = "false")] skip_checks: bool, }, + /// Submit a workflow to Slurm with auto-generated schedulers + /// + /// Automatically generates Slurm schedulers based on job resource requirements + /// and HPC profile. + /// + /// WARNING: This command uses heuristics to generate schedulers and workflow + /// actions. For complex workflows with unusual dependency patterns, the + /// generated configuration may not be optimal and could waste allocation time. + /// + /// RECOMMENDED: Preview the generated configuration first with: + /// + /// torc slurm generate --account workflow.yaml + /// + /// Review the schedulers and actions to ensure they are appropriate for your + /// workflow before submitting. You can save the output and submit manually: + /// + /// torc slurm generate --account -o workflow_with_schedulers.yaml workflow.yaml + /// torc submit workflow_with_schedulers.yaml + #[command(name = "submit-slurm")] + SubmitSlurm { + /// Path to workflow spec file (JSON/JSON5/YAML/KDL) + #[arg()] + workflow_spec: String, + /// Slurm account to use for allocations + #[arg(long)] + account: String, + /// HPC profile to use (auto-detected if not specified) + #[arg(long)] + hpc_profile: Option, + /// Bundle all nodes into a single Slurm allocation per scheduler + /// + /// By default, creates one Slurm allocation per node (N×1 mode), which allows + /// jobs to start as nodes become available and provides better fault tolerance. + /// + /// With this flag, creates one large allocation with all nodes (1×N mode), + /// which requires all nodes to be available simultaneously but uses a single sbatch. + #[arg(long)] + single_allocation: bool, + /// Ignore missing data (defaults to false) + #[arg(short, long, default_value = "false")] + ignore_missing_data: bool, + /// Skip validation checks (e.g., scheduler node requirements). Use with caution. + #[arg(long, default_value = "false")] + skip_checks: bool, + }, /// Workflow management commands Workflows { #[command(subcommand)] @@ -150,6 +197,11 @@ pub enum Commands { #[command(subcommand)] command: ScheduledComputeNodeCommands, }, + /// HPC system profiles and partition information + Hpc { + #[command(subcommand)] + command: HpcCommands, + }, /// Generate reports and analytics Reports { #[command(subcommand)] diff --git a/src/client/commands/events.rs b/src/client/commands/events.rs index 64d91835..854aa75c 100644 --- a/src/client/commands/events.rs +++ b/src/client/commands/events.rs @@ -1,3 +1,6 @@ +use std::thread; +use std::time::{Duration as StdDuration, Instant}; + use crate::client::apis::configuration::Configuration; use crate::client::apis::default_api; use crate::client::commands::get_env_user_name; @@ -313,9 +316,6 @@ fn handle_monitor_events( category: &Option, format: &str, ) { - use std::thread; - use std::time::{Duration as StdDuration, Instant}; - // Get the latest event timestamp to start monitoring from (in milliseconds since epoch) // Use list_events with limit=1 and reverse_sort=true to get the newest event let mut last_timestamp_ms: i64 = match default_api::list_events( diff --git a/src/client/commands/hpc.rs b/src/client/commands/hpc.rs new file mode 100644 index 00000000..51c025d0 --- /dev/null +++ b/src/client/commands/hpc.rs @@ -0,0 +1,589 @@ +//! HPC system profile commands +//! +//! Commands for listing, detecting, and querying HPC system profiles. + +use clap::Subcommand; +use serde_json; +use tabled::Tabled; + +use crate::client::commands::slurm::{parse_memory_mb, parse_walltime_secs}; +use crate::client::hpc::{HpcDetection, HpcPartition, HpcProfile, HpcProfileRegistry}; +use crate::config::{ClientHpcConfig, HpcPartitionConfig, HpcProfileConfig, TorcConfig}; + +use super::table_format::display_table_with_count; + +/// Create an HPC profile registry with built-in profiles and user-defined profiles from config +/// +/// This is a public version for use by other modules (e.g., main.rs for submit command) +pub fn create_registry_with_config_public(hpc_config: &ClientHpcConfig) -> HpcProfileRegistry { + create_registry_with_config(hpc_config) +} + +/// Create an HPC profile registry with built-in profiles and user-defined profiles from config +fn create_registry_with_config(hpc_config: &ClientHpcConfig) -> HpcProfileRegistry { + let mut registry = HpcProfileRegistry::with_builtin_profiles(); + + // Apply profile overrides to built-in profiles + // Note: This modifies the default_account for profiles + // In a real implementation, we'd need mutable access to profiles + + // Add custom profiles from config + for (name, profile_config) in &hpc_config.custom_profiles { + let profile = config_to_profile(name, profile_config); + registry.register(profile); + } + + registry +} + +/// Convert a config profile to an HpcProfile +fn config_to_profile(name: &str, config: &HpcProfileConfig) -> HpcProfile { + let mut detection = Vec::new(); + + // Parse detect_env_var (format: "NAME=value") + if let Some(env_var) = &config.detect_env_var { + if let Some((var_name, var_value)) = env_var.split_once('=') { + detection.push(HpcDetection::EnvVar { + name: var_name.to_string(), + value: var_value.to_string(), + }); + } + } + + // Parse hostname pattern + if let Some(pattern) = &config.detect_hostname { + detection.push(HpcDetection::HostnamePattern { + pattern: pattern.clone(), + }); + } + + // Convert partitions + let partitions: Vec = config + .partitions + .iter() + .map(|p| config_to_partition(p)) + .collect(); + + HpcProfile { + name: name.to_string(), + display_name: config.display_name.clone(), + description: config.description.clone(), + detection, + default_account: config.default_account.clone(), + partitions, + charge_factor_cpu: config.charge_factor_cpu, + charge_factor_gpu: config.charge_factor_gpu, + metadata: std::collections::HashMap::new(), + } +} + +/// Convert a config partition to an HpcPartition +fn config_to_partition(config: &HpcPartitionConfig) -> HpcPartition { + HpcPartition { + name: config.name.clone(), + description: config.description.clone(), + cpus_per_node: config.cpus_per_node, + memory_mb: config.memory_mb, + max_walltime_secs: config.max_walltime_secs, + max_nodes: None, + max_nodes_per_user: None, + min_nodes: None, + gpus_per_node: config.gpus_per_node, + gpu_type: config.gpu_type.clone(), + gpu_memory_gb: config.gpu_memory_gb, + local_disk_gb: None, + shared: config.shared, + requires_explicit_request: config.requires_explicit_request, + default_qos: None, + features: vec![], + } +} + +#[derive(Subcommand)] +pub enum HpcCommands { + /// List known HPC system profiles + List, + + /// Detect the current HPC system + Detect, + + /// Show details of an HPC profile + Show { + /// Profile name (e.g., "kestrel") + #[arg()] + name: String, + }, + + /// Show partitions for an HPC profile + Partitions { + /// Profile name (e.g., "kestrel"). If not specified, tries to detect current system. + #[arg()] + name: Option, + + /// Filter to GPU partitions only + #[arg(long)] + gpu: bool, + + /// Filter to CPU-only partitions + #[arg(long)] + cpu: bool, + + /// Filter to shared partitions + #[arg(long)] + shared: bool, + }, + + /// Find partitions matching resource requirements + Match { + /// Number of CPUs required + #[arg(long, default_value = "1")] + cpus: u32, + + /// Memory required (e.g., "100g", "512m", or MB as number) + #[arg(long, default_value = "1g")] + memory: String, + + /// Wall time required (e.g., "4:00:00", "2-00:00:00") + #[arg(long, default_value = "1:00:00")] + walltime: String, + + /// Number of GPUs required + #[arg(long)] + gpus: Option, + + /// Profile name (if not specified, tries to detect current system) + #[arg(long)] + profile: Option, + }, +} + +#[derive(Tabled)] +struct ProfileListRow { + #[tabled(rename = "Name")] + name: String, + #[tabled(rename = "Display Name")] + display_name: String, + #[tabled(rename = "Partitions")] + partition_count: usize, + #[tabled(rename = "Detected")] + detected: String, +} + +#[derive(Tabled)] +struct PartitionRow { + #[tabled(rename = "Partition")] + name: String, + #[tabled(rename = "CPUs")] + cpus: u32, + #[tabled(rename = "Memory")] + memory: String, + #[tabled(rename = "Max Time")] + max_time: String, + #[tabled(rename = "GPUs")] + gpus: String, + #[tabled(rename = "Shared")] + shared: String, + #[tabled(rename = "Explicit")] + explicit: String, +} + +#[derive(Tabled)] +struct MatchRow { + #[tabled(rename = "Partition")] + name: String, + #[tabled(rename = "CPUs")] + cpus: u32, + #[tabled(rename = "Memory")] + memory: String, + #[tabled(rename = "Max Time")] + max_time: String, + #[tabled(rename = "GPUs")] + gpus: String, + #[tabled(rename = "Notes")] + notes: String, +} + +pub fn handle_hpc_commands(command: &HpcCommands, format: &str) { + // Load config to get user-defined profiles + let config = TorcConfig::load().unwrap_or_default(); + let registry = create_registry_with_config(&config.client.hpc); + + match command { + HpcCommands::List => { + let rows: Vec = registry + .profiles() + .iter() + .map(|p| ProfileListRow { + name: p.name.clone(), + display_name: p.display_name.clone(), + partition_count: p.partitions.len(), + detected: if p.detect() { "Yes" } else { "-" }.to_string(), + }) + .collect(); + + if format == "json" { + let json_output: Vec<_> = registry + .profiles() + .iter() + .map(|p| { + serde_json::json!({ + "name": p.name, + "display_name": p.display_name, + "description": p.description, + "partition_count": p.partitions.len(), + "detected": p.detect(), + }) + }) + .collect(); + println!("{}", serde_json::to_string_pretty(&json_output).unwrap()); + } else { + println!("Known HPC profiles:\n"); + display_table_with_count(&rows, "profiles"); + } + } + + HpcCommands::Detect => { + if let Some(profile) = registry.detect() { + if format == "json" { + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "detected": true, + "name": profile.name, + "display_name": profile.display_name, + "description": profile.description, + })) + .unwrap() + ); + } else { + println!( + "Detected HPC system: {} ({})", + profile.display_name, profile.name + ); + if !profile.description.is_empty() { + println!(" {}", profile.description); + } + // Show what triggered detection + for detection in &profile.detection { + if detection.matches() { + match detection { + crate::client::hpc::HpcDetection::EnvVar { name, value } => { + println!(" Detection: {}={}", name, value); + } + crate::client::hpc::HpcDetection::HostnamePattern { pattern } => { + println!(" Detection: hostname matches {}", pattern); + } + crate::client::hpc::HpcDetection::FileExists { path } => { + println!(" Detection: file exists {}", path); + } + } + } + } + } + } else { + if format == "json" { + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "detected": false, + "message": "No known HPC system detected", + })) + .unwrap() + ); + } else { + println!("No known HPC system detected."); + println!("\nKnown systems:"); + for profile in registry.profiles() { + println!(" - {} ({})", profile.display_name, profile.name); + } + } + } + } + + HpcCommands::Show { name } => { + if let Some(profile) = registry.get(name) { + if format == "json" { + println!("{}", serde_json::to_string_pretty(&profile).unwrap()); + } else { + println!("HPC Profile: {} ({})", profile.display_name, profile.name); + println!(); + if !profile.description.is_empty() { + println!("Description: {}", profile.description); + } + println!("Detected: {}", if profile.detect() { "Yes" } else { "No" }); + println!(); + + // Show detection methods + println!("Detection methods:"); + for detection in &profile.detection { + match detection { + crate::client::hpc::HpcDetection::EnvVar { name, value } => { + println!(" - Environment variable: {}={}", name, value); + } + crate::client::hpc::HpcDetection::HostnamePattern { pattern } => { + println!(" - Hostname pattern: {}", pattern); + } + crate::client::hpc::HpcDetection::FileExists { path } => { + println!(" - File exists: {}", path); + } + } + } + println!(); + + // Show charge factors + println!("Charge factors:"); + println!(" CPU: {} AU per node-hour", profile.charge_factor_cpu); + println!(" GPU: {} AU per node-hour", profile.charge_factor_gpu); + println!(); + + // Show partition summary + let cpu_count = profile.cpu_partitions().len(); + let gpu_count = profile.gpu_partitions().len(); + println!( + "Partitions: {} total ({} CPU, {} GPU)", + profile.partitions.len(), + cpu_count, + gpu_count + ); + println!(); + println!( + "Use 'torc hpc partitions {}' to see partition details.", + name + ); + + // Show metadata + if !profile.metadata.is_empty() { + println!(); + println!("Additional information:"); + for (key, value) in &profile.metadata { + println!(" {}: {}", key, value); + } + } + } + } else { + eprintln!("Unknown HPC profile: {}", name); + eprintln!("\nKnown profiles:"); + for p in registry.profiles() { + eprintln!(" - {}", p.name); + } + std::process::exit(1); + } + } + + HpcCommands::Partitions { + name, + gpu, + cpu, + shared, + } => { + let profile = if let Some(n) = name { + registry.get(n) + } else { + registry.detect() + }; + + let profile = match profile { + Some(p) => p, + None => { + if name.is_some() { + eprintln!("Unknown HPC profile: {}", name.as_ref().unwrap()); + } else { + eprintln!("No HPC profile specified and no system detected."); + eprintln!("Use 'torc hpc partitions ' or run on an HPC system."); + } + std::process::exit(1); + } + }; + + let mut partitions: Vec<_> = profile.partitions.iter().collect(); + + // Apply filters + if *gpu { + partitions.retain(|p| p.gpus_per_node.is_some()); + } + if *cpu { + partitions.retain(|p| p.gpus_per_node.is_none()); + } + if *shared { + partitions.retain(|p| p.shared); + } + + if format == "json" { + println!("{}", serde_json::to_string_pretty(&partitions).unwrap()); + } else { + println!( + "Partitions for {} ({}):\n", + profile.display_name, profile.name + ); + + let rows: Vec = partitions + .iter() + .map(|p| PartitionRow { + name: p.name.clone(), + cpus: p.cpus_per_node, + memory: format!("{:.0}G", p.memory_gb()), + max_time: p.max_walltime_str(), + gpus: p + .gpus_per_node + .map(|g| format!("{}x {}", g, p.gpu_type.as_deref().unwrap_or("GPU"))) + .unwrap_or_else(|| "-".to_string()), + shared: if p.shared { "Yes" } else { "-" }.to_string(), + explicit: if p.requires_explicit_request { + "Yes" + } else { + "-" + } + .to_string(), + }) + .collect(); + + display_table_with_count(&rows, "partitions"); + + println!(); + println!("Notes:"); + println!(" - 'Explicit' means partition must be requested with -p/--partition"); + println!(" - Use 'torc hpc match' to find partitions for specific requirements"); + } + } + + HpcCommands::Match { + cpus, + memory, + walltime, + gpus, + profile: profile_name, + } => { + let profile = if let Some(n) = profile_name { + registry.get(n) + } else { + registry.detect() + }; + + let profile = match profile { + Some(p) => p, + None => { + if profile_name.is_some() { + eprintln!("Unknown HPC profile: {}", profile_name.as_ref().unwrap()); + } else { + eprintln!("No HPC profile specified and no system detected."); + eprintln!("Use --profile or run on an HPC system."); + } + std::process::exit(1); + } + }; + + // Parse memory + let memory_mb = match parse_memory_mb(memory) { + Ok(m) => m, + Err(e) => { + eprintln!("Invalid memory format: {}", e); + std::process::exit(1); + } + }; + + // Parse walltime + let walltime_secs = match parse_walltime_secs(walltime) { + Ok(s) => s, + Err(e) => { + eprintln!("Invalid walltime format: {}", e); + std::process::exit(1); + } + }; + + let matching = profile.find_matching_partitions(*cpus, memory_mb, walltime_secs, *gpus); + + if format == "json" { + let output = serde_json::json!({ + "profile": profile.name, + "requirements": { + "cpus": cpus, + "memory_mb": memory_mb, + "walltime_secs": walltime_secs, + "gpus": gpus, + }, + "matching_partitions": matching, + "best_partition": profile.find_best_partition(*cpus, memory_mb, walltime_secs, *gpus), + }); + println!("{}", serde_json::to_string_pretty(&output).unwrap()); + } else { + println!("Requirements:"); + println!(" CPUs: {}", cpus); + println!(" Memory: {} ({} MB)", memory, memory_mb); + println!(" Walltime: {} ({} seconds)", walltime, walltime_secs); + if let Some(g) = gpus { + println!(" GPUs: {}", g); + } + println!(); + + if matching.is_empty() { + println!( + "No partitions match these requirements on {}.", + profile.display_name + ); + println!(); + println!("Possible reasons:"); + println!(" - Requested CPUs exceed maximum per node"); + println!(" - Requested memory exceeds maximum per node"); + println!(" - Requested walltime exceeds partition limits"); + if gpus.map(|g| g > 0).unwrap_or(false) { + println!(" - No GPU partitions with enough GPUs available"); + } + } else { + println!( + "Matching partitions on {} ({}):\n", + profile.display_name, profile.name + ); + + let best = profile.find_best_partition(*cpus, memory_mb, walltime_secs, *gpus); + + let rows: Vec = matching + .iter() + .map(|p| { + let mut notes = Vec::new(); + if best.map(|b| b.name == p.name).unwrap_or(false) { + notes.push("Best match".to_string()); + } + if p.shared { + notes.push("Shared".to_string()); + } + if p.requires_explicit_request { + notes.push("Requires -p".to_string()); + } + if let Some(min) = p.min_nodes { + notes.push(format!("Min {} nodes", min)); + } + + MatchRow { + name: p.name.clone(), + cpus: p.cpus_per_node, + memory: format!("{:.0}G", p.memory_gb()), + max_time: p.max_walltime_str(), + gpus: p + .gpus_per_node + .map(|g| g.to_string()) + .unwrap_or_else(|| "-".to_string()), + notes: if notes.is_empty() { + "-".to_string() + } else { + notes.join(", ") + }, + } + }) + .collect(); + + display_table_with_count(&rows, "matching partitions"); + + if let Some(best) = best { + println!(); + println!("Recommended: {} partition", best.name); + if best.requires_explicit_request { + println!(" Use: --partition={}", best.name); + } else { + println!(" (Auto-routed based on requirements)"); + } + } + } + } + } + } +} diff --git a/src/client/commands/jobs.rs b/src/client/commands/jobs.rs index 4f6ae8c6..22f42b79 100644 --- a/src/client/commands/jobs.rs +++ b/src/client/commands/jobs.rs @@ -1,3 +1,8 @@ +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::io::{self, Write}; +use std::path::Path; + use crate::client::apis::configuration::Configuration; use crate::client::apis::default_api; use crate::client::commands::get_env_user_name; @@ -598,7 +603,6 @@ pub fn handle_job_commands(config: &Configuration, command: &JobCommands, format job_count, selected_workflow_id ); print!("Are you sure? (y/N): "); - use std::io::{self, Write}; io::stdout().flush().unwrap(); let mut input = String::new(); @@ -697,8 +701,6 @@ pub fn handle_job_commands(config: &Configuration, command: &JobCommands, format workflow_id, job_id, } => { - use std::collections::HashMap; - // Get jobs - either a single job or all jobs for a workflow let jobs: Vec = if let Some(jid) = job_id { // Get single job @@ -842,9 +844,6 @@ pub fn create_jobs_from_file( runtime_per_job: &str, _format: &str, ) -> Result> { - use std::fs; - use std::path::Path; - // Read the file let file_path = Path::new(file_path); if !file_path.exists() { @@ -937,9 +936,7 @@ pub fn get_current_job_count( pub fn get_existing_job_names( config: &Configuration, workflow_id: i64, -) -> Result, Box> { - use std::collections::HashSet; - +) -> Result, Box> { let mut names = HashSet::new(); let mut offset = 0; const PAGE_SIZE: i64 = 1000; diff --git a/src/client/commands/mod.rs b/src/client/commands/mod.rs index cd9fa24f..26d90cce 100644 --- a/src/client/commands/mod.rs +++ b/src/client/commands/mod.rs @@ -2,6 +2,7 @@ pub mod compute_nodes; pub mod config; pub mod events; pub mod files; +pub mod hpc; pub mod job_dependencies; pub mod jobs; pub mod pagination; @@ -14,9 +15,11 @@ pub mod table_format; pub mod user_data; pub mod workflows; +use std::env; +use std::io::{self, Write}; + use crate::client::apis::configuration::Configuration; use crate::client::apis::default_api; -use std::env; /// Helper function to prompt user to select a workflow when workflow_id is not provided pub fn select_workflow_interactively( @@ -65,7 +68,6 @@ pub fn select_workflow_interactively( } println!("\nEnter workflow ID: "); - use std::io::{self, Write}; io::stdout().flush().unwrap(); let mut input = String::new(); match io::stdin().read_line(&mut input) { diff --git a/src/client/commands/slurm.rs b/src/client/commands/slurm.rs index 57436b70..9708d8ab 100644 --- a/src/client/commands/slurm.rs +++ b/src/client/commands/slurm.rs @@ -13,13 +13,20 @@ use std::process::Command; use crate::client::apis::configuration::Configuration; use crate::client::apis::default_api; use crate::client::commands::get_env_user_name; +use crate::client::commands::hpc::create_registry_with_config_public; use crate::client::commands::{ print_error, select_workflow_interactively, table_format::display_table_with_count, }; +use crate::client::hpc::HpcProfile; use crate::client::hpc::hpc_interface::HpcInterface; +use crate::client::workflow_graph::WorkflowGraph; use crate::client::workflow_manager::WorkflowManager; +use crate::client::workflow_spec::{ + ResourceRequirementsSpec, SlurmSchedulerSpec, WorkflowActionSpec, WorkflowSpec, +}; use crate::config::TorcConfig; use crate::models; +use crate::time_utils::duration_string_to_seconds; use tabled::Tabled; #[derive(Tabled)] @@ -274,6 +281,399 @@ pub enum SlurmCommands { #[arg(long, default_value = "false")] save_json: bool, }, + /// Generate Slurm schedulers for a workflow based on job resource requirements + Generate { + /// Path to workflow specification file (YAML, JSON, JSON5, or KDL) + #[arg()] + workflow_file: PathBuf, + + /// Slurm account to use + #[arg(long)] + account: String, + + /// HPC profile to use (if not specified, tries to detect current system) + #[arg(long)] + profile: Option, + + /// Output file path (if not specified, prints to stdout) + #[arg(short, long)] + output: Option, + + /// Bundle all nodes into a single Slurm allocation per scheduler + /// + /// By default, creates one Slurm allocation per node (N×1 mode), which allows + /// jobs to start as nodes become available and provides better fault tolerance. + /// + /// With this flag, creates one large allocation with all nodes (1×N mode), + /// which requires all nodes to be available simultaneously but uses a single sbatch. + #[arg(long)] + single_allocation: bool, + + /// Don't add workflow actions for scheduling nodes + #[arg(long)] + no_actions: bool, + + /// Force overwrite of existing schedulers in the workflow + #[arg(long)] + force: bool, + }, +} + +/// Convert seconds to Slurm walltime format (HH:MM:SS or D-HH:MM:SS) +pub fn secs_to_walltime(secs: u64) -> String { + let hours = secs / 3600; + let mins = (secs % 3600) / 60; + let s = secs % 60; + + if hours >= 24 { + let days = hours / 24; + let h = hours % 24; + format!("{}-{:02}:{:02}:{:02}", days, h, mins, s) + } else { + format!("{:02}:{:02}:{:02}", hours, mins, s) + } +} + +/// Generate a scheduler name from a resource requirements name +fn scheduler_name_from_rr(rr_name: &str) -> String { + format!("{}_scheduler", rr_name) +} + +/// Generate Slurm schedulers for a workflow spec based on resource requirements +/// +/// This creates one scheduler per unique resource requirement (not per job). +/// All jobs with the same resource requirements share a scheduler. +/// Actions are generated based on whether any job using that scheduler has dependencies. +/// +/// # Arguments +/// * `spec` - Workflow specification to modify +/// * `profile` - HPC profile with partition information +/// * `account` - Slurm account to use +/// * `single_allocation` - If true, create 1 allocation with N nodes (1×N mode). +/// If false (default), create N allocations with 1 node each (N×1 mode). +/// * `add_actions` - Whether to add workflow actions for scheduling +/// * `force` - Force overwrite of existing schedulers +pub fn generate_schedulers_for_workflow( + spec: &mut WorkflowSpec, + profile: &HpcProfile, + account: &str, + single_allocation: bool, + add_actions: bool, + force: bool, +) -> Result { + // Check if workflow already has schedulers or actions + if !force { + if spec.slurm_schedulers.is_some() && !spec.slurm_schedulers.as_ref().unwrap().is_empty() { + return Err( + "Workflow already has slurm_schedulers defined. Use a workflow spec without schedulers (e.g., *_no_slurm.* variants).".to_string(), + ); + } + if spec.actions.is_some() && !spec.actions.as_ref().unwrap().is_empty() { + return Err( + "Workflow already has actions defined. Use a workflow spec without actions (e.g., *_no_slurm.* variants).".to_string(), + ); + } + } + + // Expand parameters before building the graph to properly detect file-based dependencies + spec.expand_parameters() + .map_err(|e| format!("Failed to expand parameters: {}", e))?; + + // Build a map of resource requirements by name + let rr_map: std::collections::HashMap = spec + .resource_requirements + .as_ref() + .map(|rrs| rrs.iter().map(|rr| (rr.name.clone(), rr)).collect()) + .unwrap_or_default(); + + if rr_map.is_empty() { + return Err( + "Workflow has no resource_requirements defined. Cannot generate schedulers." + .to_string(), + ); + } + + // Build workflow graph for dependency analysis and job grouping + let graph = WorkflowGraph::from_spec(spec) + .map_err(|e| format!("Failed to build workflow graph: {}", e))?; + + let mut schedulers: Vec = Vec::new(); + let mut actions: Vec = Vec::new(); + let mut warnings: Vec = Vec::new(); + + // Get scheduler groups from the graph + // Groups jobs by (resource_requirements, has_dependencies) + let scheduler_groups = graph.scheduler_groups(); + + // Check for jobs without resource requirements + for job in &spec.jobs { + if job.resource_requirements.is_none() { + warnings.push(format!( + "Job '{}' has no resource_requirements, skipping scheduler generation", + job.name + )); + } + } + + // Create schedulers and actions for each group + for group in &scheduler_groups { + let rr_name = &group.resource_requirements; + let has_deps = group.has_dependencies; + let rr = match rr_map.get(rr_name) { + Some(rr) => *rr, + None => { + return Err(format!("Resource requirements '{}' not found", rr_name)); + } + }; + + // Parse resource requirements + let memory_mb = parse_memory_mb(&rr.memory)?; + let runtime_secs = duration_string_to_seconds(&rr.runtime)? as u64; + let gpus = if rr.num_gpus > 0 { + Some(rr.num_gpus as u32) + } else { + None + }; + + // Find best partition + let partition = match profile.find_best_partition( + rr.num_cpus as u32, + memory_mb, + runtime_secs, + gpus, + ) { + Some(p) => p, + None => { + warnings.push(format!( + "No partition found for resource requirements '{}' (CPUs: {}, Memory: {}, Runtime: {}, GPUs: {:?})", + rr_name, rr.num_cpus, rr.memory, rr.runtime, gpus + )); + continue; + } + }; + + // Scheduler naming: jobs with deps get "_deferred" suffix + let scheduler_name = if has_deps { + format!("{}_deferred_scheduler", rr_name) + } else { + scheduler_name_from_rr(rr_name) + }; + + // Calculate total nodes needed based on job count and partition capacity + // Jobs per node is the minimum of what fits by CPU, memory, and GPU constraints + let jobs_per_node_by_cpu = partition.cpus_per_node / rr.num_cpus as u32; + let jobs_per_node_by_mem = (partition.memory_mb / memory_mb) as u32; + let jobs_per_node_by_gpu = match (gpus, partition.gpus_per_node) { + (Some(job_gpus), Some(node_gpus)) if job_gpus > 0 => node_gpus / job_gpus, + _ => u32::MAX, // No GPU constraint + }; + let jobs_per_node = std::cmp::max( + 1, + std::cmp::min( + jobs_per_node_by_cpu, + std::cmp::min(jobs_per_node_by_mem, jobs_per_node_by_gpu), + ), + ); + // Total nodes needed to run all jobs concurrently (respecting num_nodes per job) + let nodes_per_job = rr.num_nodes as u32; + let total_nodes_needed = + ((group.job_count as u32 + jobs_per_node - 1) / jobs_per_node) * nodes_per_job; + let total_nodes_needed = std::cmp::max(1, total_nodes_needed) as i64; + + // Allocation strategy: + // - N×1 mode (default): N separate 1-node allocations, jobs start as nodes become available + // - 1×N mode (--single-allocation): 1 allocation with N nodes, all nodes must be available + let (nodes_per_alloc, effective_num_allocations) = if single_allocation { + // 1×N mode: single allocation with all nodes + (total_nodes_needed, 1i64) + } else { + // N×1 mode: many single-node allocations + (1i64, total_nodes_needed) + }; + + // Create scheduler for this group + // Use the partition's max walltime for headroom + let mut scheduler = SlurmSchedulerSpec { + name: Some(scheduler_name.clone()), + account: account.to_string(), + partition: if partition.requires_explicit_request { + Some(partition.name.clone()) + } else { + None // Auto-routed + }, + mem: Some(rr.memory.clone()), + walltime: secs_to_walltime(partition.max_walltime_secs), + nodes: nodes_per_alloc, + gres: None, + ntasks_per_node: None, + qos: partition.default_qos.clone(), + tmp: None, + extra: None, + }; + + // Add GPU gres if needed + if let Some(gpu_count) = gpus { + scheduler.gres = Some(format!("gpu:{}", gpu_count)); + } + + schedulers.push(scheduler); + + // Create action for this scheduler + if add_actions { + // If requesting multiple nodes, start one worker per node + let start_one_worker_per_node = if nodes_per_alloc > 1 { + Some(true) + } else { + None + }; + + let (trigger_type, job_name_regexes) = if has_deps { + // Jobs with dependencies: trigger on_jobs_ready when they become unblocked + // This fires when the first job in the group becomes ready + ("on_jobs_ready", Some(group.job_name_patterns.clone())) + } else { + // Jobs without dependencies: trigger on_workflow_start + ("on_workflow_start", None) + }; + + let action = WorkflowActionSpec { + trigger_type: trigger_type.to_string(), + action_type: "schedule_nodes".to_string(), + jobs: None, + job_name_regexes, + commands: None, + scheduler: Some(scheduler_name.clone()), + scheduler_type: Some("slurm".to_string()), + num_allocations: Some(effective_num_allocations), + start_one_worker_per_node, + max_parallel_jobs: None, + persistent: None, + }; + actions.push(action); + } + } + + if schedulers.is_empty() { + return Err("No schedulers could be generated. Check resource requirements.".to_string()); + } + + // Update jobs to reference their scheduler based on (resource_requirement, has_dependencies) + for job in &mut spec.jobs { + if let Some(rr_name) = &job.resource_requirements { + let has_deps = graph.has_dependencies(&job.name); + let scheduler_name = if has_deps { + format!("{}_deferred_scheduler", rr_name) + } else { + scheduler_name_from_rr(rr_name) + }; + job.scheduler = Some(scheduler_name); + } + } + + // Update workflow with schedulers and actions + spec.slurm_schedulers = Some(schedulers); + if add_actions && !actions.is_empty() { + let mut existing_actions = spec.actions.take().unwrap_or_default(); + existing_actions.extend(actions.clone()); + spec.actions = Some(existing_actions); + } + + Ok(GenerateResult { + scheduler_count: spec.slurm_schedulers.as_ref().map(|s| s.len()).unwrap_or(0), + action_count: actions.len(), + warnings, + }) +} + +/// Result of generating schedulers +pub struct GenerateResult { + pub scheduler_count: usize, + pub action_count: usize, + pub warnings: Vec, +} + +/// Parse memory string like "100g", "512m", "1024" (MB) into MB +pub fn parse_memory_mb(s: &str) -> Result { + let s = s.trim().to_lowercase(); + if s.is_empty() { + return Err("Empty memory string".to_string()); + } + + // Check for suffix + if let Some(num_str) = s.strip_suffix('g') { + let num: f64 = num_str + .parse() + .map_err(|_| format!("Invalid number: {}", num_str))?; + Ok((num * 1024.0) as u64) + } else if let Some(num_str) = s.strip_suffix('m') { + let num: u64 = num_str + .parse() + .map_err(|_| format!("Invalid number: {}", num_str))?; + Ok(num) + } else if let Some(num_str) = s.strip_suffix('k') { + let num: f64 = num_str + .parse() + .map_err(|_| format!("Invalid number: {}", num_str))?; + Ok((num / 1024.0) as u64) + } else { + // Assume MB + s.parse() + .map_err(|_| format!("Invalid memory value: {}", s)) + } +} + +/// Parse walltime string like "4:00:00", "2-00:00:00" into seconds +pub fn parse_walltime_secs(s: &str) -> Result { + let s = s.trim(); + + // Check for day format: D-HH:MM:SS + if let Some((days_str, rest)) = s.split_once('-') { + let days: u64 = days_str + .parse() + .map_err(|_| format!("Invalid days: {}", days_str))?; + let hms_secs = parse_hms(rest)?; + return Ok(days * 24 * 3600 + hms_secs); + } + + // Check for hours format: H, HH:MM, or HH:MM:SS + parse_hms(s) +} + +fn parse_hms(s: &str) -> Result { + let parts: Vec<&str> = s.split(':').collect(); + match parts.len() { + 1 => { + // Just hours + let hours: u64 = parts[0] + .parse() + .map_err(|_| format!("Invalid hours: {}", parts[0]))?; + Ok(hours * 3600) + } + 2 => { + // HH:MM + let hours: u64 = parts[0] + .parse() + .map_err(|_| format!("Invalid hours: {}", parts[0]))?; + let mins: u64 = parts[1] + .parse() + .map_err(|_| format!("Invalid minutes: {}", parts[1]))?; + Ok(hours * 3600 + mins * 60) + } + 3 => { + // HH:MM:SS + let hours: u64 = parts[0] + .parse() + .map_err(|_| format!("Invalid hours: {}", parts[0]))?; + let mins: u64 = parts[1] + .parse() + .map_err(|_| format!("Invalid minutes: {}", parts[1]))?; + let secs: u64 = parts[2] + .parse() + .map_err(|_| format!("Invalid seconds: {}", parts[2]))?; + Ok(hours * 3600 + mins * 60 + secs) + } + _ => Err(format!("Invalid time format: {}", s)), + } } pub fn handle_slurm_commands(config: &Configuration, command: &SlurmCommands, format: &str) { @@ -660,6 +1060,26 @@ pub fn handle_slurm_commands(config: &Configuration, command: &SlurmCommands, fo }); run_sacct_for_workflow(config, wf_id, output_dir, *save_json, format); } + SlurmCommands::Generate { + workflow_file, + account, + profile: profile_name, + output, + single_allocation, + no_actions, + force, + } => { + handle_generate( + workflow_file, + account, + profile_name.as_deref(), + output.as_ref(), + *single_allocation, + *no_actions, + *force, + format, + ); + } } } @@ -1897,3 +2317,128 @@ pub fn run_sacct_for_workflow( } } } + +/// Handle the generate command - generates Slurm schedulers for a workflow +fn handle_generate( + workflow_file: &PathBuf, + account: &str, + profile_name: Option<&str>, + output: Option<&PathBuf>, + single_allocation: bool, + no_actions: bool, + force: bool, + format: &str, +) { + // Load HPC config and registry + let torc_config = TorcConfig::load().unwrap_or_default(); + let registry = create_registry_with_config_public(&torc_config.client.hpc); + + // Get the HPC profile + let profile = if let Some(n) = profile_name { + registry.get(n) + } else { + registry.detect() + }; + + let profile = match profile { + Some(p) => p, + None => { + if profile_name.is_some() { + eprintln!("Unknown HPC profile: {}", profile_name.unwrap()); + } else { + eprintln!("No HPC profile specified and no system detected."); + eprintln!("Use --profile or run on an HPC system."); + } + std::process::exit(1); + } + }; + + // Parse the workflow spec (supports YAML, JSON, JSON5, and KDL) + let mut spec: WorkflowSpec = match WorkflowSpec::from_spec_file(workflow_file) { + Ok(s) => s, + Err(e) => { + eprintln!("Failed to parse workflow file: {}", e); + std::process::exit(1); + } + }; + + // Generate schedulers + let result = match generate_schedulers_for_workflow( + &mut spec, + profile, + account, + single_allocation, + !no_actions, + force, + ) { + Ok(r) => r, + Err(e) => { + eprintln!("Error: {}", e); + std::process::exit(1); + } + }; + + // Output the result in the same format as the input + let output_content = match workflow_file.extension().and_then(|e| e.to_str()) { + Some("json") => serde_json::to_string_pretty(&spec).unwrap(), + Some("json5") => serde_json::to_string_pretty(&spec).unwrap(), // Output as JSON + Some("kdl") => spec.to_kdl_str(), + _ => serde_yaml::to_string(&spec).unwrap(), // YAML, YML, or unknown + }; + + if let Some(output_path) = output { + match std::fs::write(output_path, &output_content) { + Ok(_) => { + if format != "json" { + println!("Generated workflow written to: {}", output_path.display()); + println!(); + println!("Summary:"); + println!(" Schedulers generated: {}", result.scheduler_count); + println!(" Actions added: {}", result.action_count); + println!( + " Profile used: {} ({})", + profile.display_name, profile.name + ); + + if !result.warnings.is_empty() { + println!(); + println!("Warnings:"); + for warning in &result.warnings { + println!(" - {}", warning); + } + } + } + } + Err(e) => { + eprintln!("Failed to write output file: {}", e); + std::process::exit(1); + } + } + } else { + // Print to stdout + if format == "json" { + println!("{}", serde_json::to_string_pretty(&spec).unwrap()); + } else { + println!("{}", output_content); + + // Print summary to stderr so it doesn't mix with the workflow output + // Use // for KDL-compatible comments + eprintln!(); + eprintln!("// Summary:"); + eprintln!("// Schedulers generated: {}", result.scheduler_count); + eprintln!("// Actions added: {}", result.action_count); + eprintln!( + "// Profile used: {} ({})", + profile.display_name, profile.name + ); + + if !result.warnings.is_empty() { + eprintln!("//"); + eprintln!("// Warnings:"); + for warning in &result.warnings { + eprintln!("// - {}", warning); + } + } + } + } +} diff --git a/src/client/commands/workflows.rs b/src/client/commands/workflows.rs index 2c7c787c..a3798cb4 100644 --- a/src/client/commands/workflows.rs +++ b/src/client/commands/workflows.rs @@ -1,7 +1,11 @@ +use std::io::{self, Write}; + use clap::Subcommand; use crate::client::apis::configuration::Configuration; use crate::client::apis::default_api; +use crate::client::commands::hpc::create_registry_with_config_public; +use crate::client::commands::slurm::generate_schedulers_for_workflow; use crate::client::commands::{ get_env_user_name, get_user_name, pagination, print_error, select_workflow_interactively, table_format::display_table_with_count, @@ -60,7 +64,7 @@ struct WorkflowActionTableRow { #[derive(Subcommand)] pub enum WorkflowCommands { - /// Create a workflow from a specification file (supports JSON, JSON5, and YAML formats) + /// Create a workflow from a specification file (supports JSON, JSON5, YAML, and KDL formats) Create { /// Path to specification file containing WorkflowSpec /// @@ -68,6 +72,7 @@ pub enum WorkflowCommands { /// - JSON (.json): Standard JSON format /// - JSON5 (.json5): JSON with comments and trailing commas /// - YAML (.yaml, .yml): Human-readable YAML format + /// - KDL (.kdl): KDL document format /// /// Format is auto-detected from file extension, with fallback parsing attempted #[arg()] @@ -86,6 +91,43 @@ pub enum WorkflowCommands { #[arg(long)] dry_run: bool, }, + /// Create a workflow with auto-generated Slurm schedulers + /// + /// Automatically generates Slurm schedulers based on job resource requirements + /// and HPC profile. For Slurm workflows without pre-configured schedulers. + #[command(name = "create-slurm")] + CreateSlurm { + /// Path to specification file containing WorkflowSpec + #[arg()] + file: String, + /// Slurm account to use for allocations + #[arg(long)] + account: String, + /// HPC profile to use (auto-detected if not specified) + #[arg(long)] + hpc_profile: Option, + /// Bundle all nodes into a single Slurm allocation per scheduler + /// + /// By default, creates one Slurm allocation per node (N×1 mode), which allows + /// jobs to start as nodes become available and provides better fault tolerance. + /// + /// With this flag, creates one large allocation with all nodes (1×N mode), + /// which requires all nodes to be available simultaneously but uses a single sbatch. + #[arg(long)] + single_allocation: bool, + /// User that owns the workflow (defaults to USER environment variable) + #[arg(short, long, env = "USER")] + user: String, + /// Disable resource monitoring (default: enabled with summary granularity and 5s sample rate) + #[arg(long, default_value = "false")] + no_resource_monitoring: bool, + /// Skip validation checks (e.g., scheduler node requirements). Use with caution. + #[arg(long, default_value = "false")] + skip_checks: bool, + /// Validate the workflow specification without creating it (dry-run mode) + #[arg(long)] + dry_run: bool, + }, /// Create a new empty workflow New { /// Name of the workflow @@ -307,12 +349,13 @@ fn show_execution_plan_from_spec(file_path: &str, format: &str) { match crate::client::execution_plan::ExecutionPlan::from_spec(&spec) { Ok(plan) => { if format == "json" { - // For JSON output, create a structured representation - let stages_json: Vec = plan.stages.iter().map(|stage| { + // For JSON output, use the new DAG-based event structure + let events_json: Vec = plan.events.values().map(|event| { serde_json::json!({ - "stage_number": stage.stage_number + 1, // Display as 1-based - "trigger": stage.trigger_description, - "scheduler_allocations": stage.scheduler_allocations.iter().map(|alloc| { + "id": event.id, + "trigger": event.trigger, + "trigger_description": event.trigger_description, + "scheduler_allocations": event.scheduler_allocations.iter().map(|alloc| { serde_json::json!({ "scheduler": alloc.scheduler, "scheduler_type": alloc.scheduler_type, @@ -320,7 +363,9 @@ fn show_execution_plan_from_spec(file_path: &str, format: &str) { "job_names": alloc.jobs, }) }).collect::>(), - "jobs_becoming_ready": stage.jobs_becoming_ready, + "jobs_becoming_ready": event.jobs_becoming_ready, + "depends_on_events": event.depends_on_events, + "unlocks_events": event.unlocks_events, }) }).collect(); @@ -328,9 +373,11 @@ fn show_execution_plan_from_spec(file_path: &str, format: &str) { "status": "success", "source": "spec_file", "workflow_name": spec.name, - "total_stages": plan.stages.len(), + "total_events": plan.events.len(), "total_jobs": spec.jobs.len(), - "stages": stages_json, + "root_events": plan.root_events, + "leaf_events": plan.leaf_events, + "events": events_json, }); match serde_json::to_string_pretty(&output) { @@ -398,17 +445,50 @@ fn show_execution_plan_from_database(config: &Configuration, workflow_id: i64, f } }; + // Fetch slurm schedulers for this workflow + let slurm_schedulers = match default_api::list_slurm_schedulers( + config, + workflow_id, + None, // offset + None, // limit + None, // sort_by + None, // reverse_sort + None, // name + None, // account + None, // gres + None, // mem + None, // nodes + None, // partition + None, // qos + None, // tmp + None, // walltime + ) { + Ok(response) => response.items.unwrap_or_default(), + Err(e) => { + eprintln!( + "Warning: Could not fetch slurm schedulers for workflow {}: {}", + workflow_id, e + ); + vec![] + } + }; + // Build execution plan from database models match crate::client::execution_plan::ExecutionPlan::from_database_models( - &workflow, &jobs, &actions, + &workflow, + &jobs, + &actions, + &slurm_schedulers, ) { Ok(plan) => { if format == "json" { - let stages_json: Vec = plan.stages.iter().map(|stage| { + // For JSON output, use the new DAG-based event structure + let events_json: Vec = plan.events.values().map(|event| { serde_json::json!({ - "stage_number": stage.stage_number + 1, - "trigger": stage.trigger_description, - "scheduler_allocations": stage.scheduler_allocations.iter().map(|alloc| { + "id": event.id, + "trigger": event.trigger, + "trigger_description": event.trigger_description, + "scheduler_allocations": event.scheduler_allocations.iter().map(|alloc| { serde_json::json!({ "scheduler": alloc.scheduler, "scheduler_type": alloc.scheduler_type, @@ -416,7 +496,9 @@ fn show_execution_plan_from_database(config: &Configuration, workflow_id: i64, f "job_names": alloc.jobs, }) }).collect::>(), - "jobs_becoming_ready": stage.jobs_becoming_ready, + "jobs_becoming_ready": event.jobs_becoming_ready, + "depends_on_events": event.depends_on_events, + "unlocks_events": event.unlocks_events, }) }).collect(); @@ -425,9 +507,11 @@ fn show_execution_plan_from_database(config: &Configuration, workflow_id: i64, f "source": "database", "workflow_id": workflow_id, "workflow_name": workflow.name, - "total_stages": plan.stages.len(), + "total_events": plan.events.len(), "total_jobs": jobs.len(), - "stages": stages_json, + "root_events": plan.root_events, + "leaf_events": plan.leaf_events, + "events": events_json, }); match serde_json::to_string_pretty(&output) { @@ -724,8 +808,6 @@ fn handle_reset_status( eprintln!("Force mode is enabled (will ignore running/pending jobs check)."); } print!("\nDo you want to continue? (y/N): "); - - use std::io::{self, Write}; io::stdout().flush().unwrap(); let mut input = String::new(); @@ -1164,8 +1246,6 @@ fn handle_initialize( println!("\nWarning: This workflow has already been initialized."); println!("Some jobs already have initialized status."); print!("\nDo you want to continue? (y/N): "); - - use std::io::{self, Write}; io::stdout().flush().unwrap(); let mut input = String::new(); @@ -1522,8 +1602,6 @@ fn handle_delete(config: &Configuration, ids: &[i64], no_prompts: bool, force: b println!(" - All job dependencies and relationships"); println!("\nThis action cannot be undone."); print!("\nAre you sure you want to delete this workflow? (y/N): "); - - use std::io::{self, Write}; io::stdout().flush().unwrap(); let mut input = String::new(); @@ -2046,6 +2124,170 @@ fn handle_create( } } +fn handle_create_slurm( + config: &Configuration, + file: &str, + account: &str, + hpc_profile: Option<&str>, + single_allocation: bool, + user: &str, + no_resource_monitoring: bool, + skip_checks: bool, + dry_run: bool, + format: &str, +) { + // Handle dry-run mode first + if dry_run { + let result = WorkflowSpec::validate_spec(file); + if format == "json" { + match serde_json::to_string_pretty(&result) { + Ok(json) => println!("{}", json), + Err(e) => { + eprintln!("Error serializing validation result: {}", e); + std::process::exit(1); + } + } + } else { + println!("Workflow Validation Results (with Slurm scheduler generation)"); + println!("=============================================================="); + println!(); + println!("Note: Dry-run validates the spec before scheduler generation."); + println!("Use 'torc slurm generate' to preview generated schedulers."); + println!(); + + let summary = &result.summary; + println!("Workflow: {}", summary.workflow_name); + println!("Jobs: {}", summary.job_count); + println!( + "Resource requirements: {}", + summary.resource_requirements_count + ); + println!(); + + if !result.errors.is_empty() { + eprintln!("Errors ({}):", result.errors.len()); + for error in &result.errors { + eprintln!(" - {}", error); + } + } + + if result.valid { + println!("Validation: PASSED"); + } else { + eprintln!("Validation: FAILED"); + } + } + + if !result.valid { + std::process::exit(1); + } + return; + } + + // Load HPC config and registry + let torc_config = TorcConfig::load().unwrap_or_default(); + let registry = create_registry_with_config_public(&torc_config.client.hpc); + + // Get the HPC profile + let profile = if let Some(name) = hpc_profile { + registry.get(name) + } else { + registry.detect() + }; + + let profile = match profile { + Some(p) => p, + None => { + if hpc_profile.is_some() { + eprintln!("Unknown HPC profile: {}", hpc_profile.unwrap()); + } else { + eprintln!("No HPC profile specified and no system detected."); + eprintln!("Use --hpc-profile to specify a profile."); + } + std::process::exit(1); + } + }; + + // Parse the workflow spec + let mut spec = match WorkflowSpec::from_spec_file(file) { + Ok(s) => s, + Err(e) => { + eprintln!("Failed to parse workflow file: {}", e); + std::process::exit(1); + } + }; + + // Generate schedulers + // Don't allow force=true - if schedulers already exist, user should use the _no_slurm variant + match generate_schedulers_for_workflow( + &mut spec, + profile, + account, + single_allocation, + true, + false, + ) { + Ok(result) => { + if format != "json" { + eprintln!( + "Auto-generated {} scheduler(s) and {} action(s) using {} profile", + result.scheduler_count, result.action_count, profile.name + ); + for warning in &result.warnings { + eprintln!(" Warning: {}", warning); + } + } + } + Err(e) => { + eprintln!("Error generating schedulers: {}", e); + std::process::exit(1); + } + } + + // Write modified spec to temp file + let temp_dir = std::env::temp_dir(); + let temp_file = temp_dir.join(format!("torc_workflow_{}.yaml", std::process::id())); + match std::fs::write(&temp_file, serde_yaml::to_string(&spec).unwrap()) { + Ok(_) => {} + Err(e) => { + eprintln!("Failed to write temporary workflow file: {}", e); + std::process::exit(1); + } + } + + // Create workflow from modified spec + match WorkflowSpec::create_workflow_from_spec( + config, + &temp_file, + user, + !no_resource_monitoring, + skip_checks, + ) { + Ok(workflow_id) => { + if format == "json" { + let json_output = serde_json::json!({ + "workflow_id": workflow_id, + "status": "success", + "message": format!("Workflow created successfully with ID: {}", workflow_id) + }); + match serde_json::to_string_pretty(&json_output) { + Ok(json) => println!("{}", json), + Err(e) => { + eprintln!("Error serializing JSON: {}", e); + std::process::exit(1); + } + } + } else { + println!("Created workflow {}", workflow_id); + } + } + Err(e) => { + eprintln!("Error creating workflow from spec: {}", e); + std::process::exit(1); + } + } +} + pub fn handle_workflow_commands(config: &Configuration, command: &WorkflowCommands, format: &str) { match command { WorkflowCommands::Create { @@ -2065,6 +2307,29 @@ pub fn handle_workflow_commands(config: &Configuration, command: &WorkflowComman format, ); } + WorkflowCommands::CreateSlurm { + file, + account, + hpc_profile, + single_allocation, + user, + no_resource_monitoring, + skip_checks, + dry_run, + } => { + handle_create_slurm( + config, + file, + account, + hpc_profile.as_deref(), + *single_allocation, + user, + *no_resource_monitoring, + *skip_checks, + *dry_run, + format, + ); + } WorkflowCommands::New { name, description, diff --git a/src/client/execution_plan.rs b/src/client/execution_plan.rs index 8f5f890a..9676ed57 100644 --- a/src/client/execution_plan.rs +++ b/src/client/execution_plan.rs @@ -1,10 +1,25 @@ +//! Execution Plan - Visualization of workflow execution as a DAG of events +//! +//! This module provides a high-level view of how a workflow will execute, +//! showing which jobs become ready at each event and what scheduler +//! allocations are triggered. +//! +//! The execution plan is represented as a Directed Acyclic Graph (DAG) where: +//! - Each node is an "event" (workflow start or a set of jobs completing) +//! - Edges represent dependencies between events +//! - Independent sub-workflows have separate event chains +//! +//! Built on top of WorkflowGraph for consistent dependency analysis. + +use crate::client::workflow_graph::WorkflowGraph; use crate::client::workflow_spec::{WorkflowActionSpec, WorkflowSpec}; use crate::models::{JobModel, WorkflowActionModel, WorkflowModel}; use regex::Regex; +use serde::Serialize; use std::collections::{HashMap, HashSet}; /// Represents a scheduler allocation in the execution plan -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Serialize)] pub struct SchedulerAllocation { pub scheduler: String, pub scheduler_type: String, @@ -12,387 +27,666 @@ pub struct SchedulerAllocation { pub jobs: Vec, } -/// Represents a stage in the workflow execution plan -#[derive(Debug, Clone)] -pub struct ExecutionStage { - pub stage_number: usize, +/// What triggers an execution event +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type", content = "data")] +pub enum EventTrigger { + /// Workflow starts + WorkflowStart, + /// Specific jobs complete + JobsComplete { jobs: Vec }, +} + +/// An event in the execution plan DAG +#[derive(Debug, Clone, Serialize)] +pub struct ExecutionEvent { + /// Unique identifier for this event + pub id: String, + /// What triggers this event + pub trigger: EventTrigger, + /// Human-readable description of the trigger pub trigger_description: String, + /// Scheduler allocations triggered by this event pub scheduler_allocations: Vec, + /// Jobs that become ready when this event fires pub jobs_becoming_ready: Vec, + /// Event IDs that must complete before this event can fire + pub depends_on_events: Vec, + /// Event IDs that depend on this event + pub unlocks_events: Vec, } -/// Represents the complete execution plan for a workflow -#[derive(Debug)] +/// Represents the complete execution plan for a workflow as a DAG +#[derive(Debug, Serialize)] pub struct ExecutionPlan { - pub stages: Vec, + /// All events in the plan, keyed by event ID + pub events: HashMap, + /// Event IDs that have no dependencies (entry points) + pub root_events: Vec, + /// Event IDs that nothing depends on (exit points) + pub leaf_events: Vec, + /// The underlying workflow graph (if built from spec) + #[serde(skip)] + pub graph: Option, +} + +// Legacy stage-based interface for backwards compatibility +/// Represents a stage in the workflow execution plan (legacy format) +#[derive(Debug, Clone, Serialize)] +pub struct ExecutionStage { + pub stage_number: usize, + pub trigger_description: String, + pub scheduler_allocations: Vec, + pub jobs_becoming_ready: Vec, } impl ExecutionPlan { /// Build an execution plan from a workflow specification pub fn from_spec(spec: &WorkflowSpec) -> Result> { - let mut stages = Vec::new(); - - // Build dependency graph (job_name -> jobs it's blocked by) - let dependency_graph = build_dependency_graph(spec)?; + let graph = WorkflowGraph::from_spec(spec)?; + let mut events: HashMap = HashMap::new(); - // Build reverse dependency graph (job_name -> jobs that depend on it) - let reverse_deps = build_reverse_dependencies(&dependency_graph); + // Event 0: Workflow start - root jobs become ready + let root_jobs: Vec = graph.roots().iter().map(|s| s.to_string()).collect(); - // Topologically sort jobs into levels - let job_levels = topological_sort(spec, &dependency_graph)?; + // Find on_workflow_start scheduler allocations + let mut start_allocations = Vec::new(); + if let Some(ref actions) = spec.actions { + for action in actions { + if action.trigger_type == "on_workflow_start" { + if let Some(alloc) = build_scheduler_allocation(spec, action)? { + start_allocations.push(alloc); + } + } + } + } - // Stage 0: on_workflow_start actions - let workflow_start_stage = build_workflow_start_stage(spec)?; - if !workflow_start_stage.scheduler_allocations.is_empty() - || !workflow_start_stage.jobs_becoming_ready.is_empty() - { - stages.push(workflow_start_stage); + let start_event = ExecutionEvent { + id: "start".to_string(), + trigger: EventTrigger::WorkflowStart, + trigger_description: "Workflow Start".to_string(), + scheduler_allocations: start_allocations, + jobs_becoming_ready: root_jobs.clone(), + depends_on_events: vec![], + unlocks_events: vec![], // Will be filled in later + }; + events.insert("start".to_string(), start_event); + + // Build events for each distinct set of jobs that complete together + // Key insight: jobs at the same level that share the same dependents form an event + // But we need to track events by the jobs they unlock, not by level + + // Build a map: job -> event_id for the event where this job completes + let mut job_completion_event: HashMap = HashMap::new(); + + // Root jobs complete as part of start event + for job in &root_jobs { + job_completion_event.insert(job.clone(), "start".to_string()); } - // Build stages for each job level - for (level_idx, level_jobs) in job_levels.iter().enumerate() { - if level_jobs.is_empty() { - continue; + // Build events by analyzing what jobs become ready when other jobs complete + // We iterate through jobs by dependency depth to ensure proper ordering + let mut processed_jobs: HashSet = root_jobs.iter().cloned().collect(); + + // Keep processing until all jobs have been assigned to events + loop { + // Find jobs that become ready based on currently processed jobs + let newly_ready = graph.jobs_unblocked_by(&processed_jobs); + if newly_ready.is_empty() { + break; } - let stage_number = stages.len(); - let stage = build_job_completion_stage( - spec, - stage_number, - level_idx, - level_jobs, - &job_levels, - &reverse_deps, - )?; - - // Only add the stage if it has actions or jobs becoming ready - if !stage.scheduler_allocations.is_empty() || !stage.jobs_becoming_ready.is_empty() { - stages.push(stage); + // Group newly ready jobs by their dependencies (the jobs they're waiting on) + // Jobs with the same dependencies should be in the same event + let mut jobs_by_deps: HashMap, Vec> = HashMap::new(); + + for job in &newly_ready { + if let Some(deps) = graph.dependencies_of(job) { + let mut dep_list: Vec = deps.iter().cloned().collect(); + dep_list.sort(); + jobs_by_deps.entry(dep_list).or_default().push(job.clone()); + } } + + // Create an event for each unique dependency set + for (deps, jobs_becoming_ready) in jobs_by_deps { + // Generate a descriptive event ID based on the triggering jobs + let event_id = generate_event_id(&deps); + + // Find which events the dependencies belong to + let mut depends_on_event_ids: HashSet = HashSet::new(); + for dep in &deps { + if let Some(evt_id) = job_completion_event.get(dep) { + depends_on_event_ids.insert(evt_id.clone()); + } + } + let depends_on_events: Vec = depends_on_event_ids.into_iter().collect(); + + // Find matching scheduler actions + let mut scheduler_allocations = Vec::new(); + if let Some(ref actions) = spec.actions { + let matching = graph.matching_actions(&jobs_becoming_ready, actions); + for action in matching { + if let Some(alloc) = build_scheduler_allocation(spec, action)? { + scheduler_allocations.push(alloc); + } + } + } + + let trigger_description = build_trigger_description(&deps); + + let event = ExecutionEvent { + id: event_id.clone(), + trigger: EventTrigger::JobsComplete { jobs: deps.clone() }, + trigger_description, + scheduler_allocations, + jobs_becoming_ready: jobs_becoming_ready.clone(), + depends_on_events, + unlocks_events: vec![], // Will be filled in below + }; + events.insert(event_id.clone(), event); + + // Mark these jobs as completing in this event + for job in &jobs_becoming_ready { + job_completion_event.insert(job.clone(), event_id.clone()); + } + } + + // Add newly ready jobs to processed set + processed_jobs.extend(newly_ready); } - Ok(ExecutionPlan { stages }) + // Build unlocks_events by reversing the depends_on_events relationships + let event_deps: Vec<(String, Vec)> = events + .values() + .map(|e| (e.id.clone(), e.depends_on_events.clone())) + .collect(); + + for (event_id, deps) in event_deps { + for dep_event_id in deps { + if let Some(dep_event) = events.get_mut(&dep_event_id) { + if !dep_event.unlocks_events.contains(&event_id) { + dep_event.unlocks_events.push(event_id.clone()); + } + } + } + } + + // Identify root and leaf events + let root_events: Vec = events + .values() + .filter(|e| e.depends_on_events.is_empty()) + .map(|e| e.id.clone()) + .collect(); + + let leaf_events: Vec = events + .values() + .filter(|e| e.unlocks_events.is_empty()) + .map(|e| e.id.clone()) + .collect(); + + Ok(ExecutionPlan { + events, + root_events, + leaf_events, + graph: Some(graph), + }) } - /// Build an execution plan from database models (workflow, jobs, actions) + /// Build an execution plan from database models (workflow, jobs, actions, slurm_schedulers) pub fn from_database_models( _workflow: &WorkflowModel, jobs: &[JobModel], actions: &[WorkflowActionModel], + slurm_schedulers: &[crate::models::SlurmSchedulerModel], ) -> Result> { - let mut stages = Vec::new(); - - // Build dependency graph from jobs (job_id -> Vec) - let mut job_id_to_name = HashMap::new(); - let mut job_name_to_id = HashMap::new(); - let mut dependency_graph_by_id: HashMap> = HashMap::new(); + // Build scheduler ID to name map + let scheduler_id_to_name: HashMap = slurm_schedulers + .iter() + .filter_map(|s| match (s.id, &s.name) { + (Some(id), Some(name)) => Some((id, name.clone())), + _ => None, + }) + .collect(); + let mut events: HashMap = HashMap::new(); + + // Build dependency graph from jobs (job_name -> Vec) + let mut job_id_to_name: HashMap = HashMap::new(); + let mut job_name_to_id: HashMap = HashMap::new(); + let mut dependency_graph_by_name: HashMap> = HashMap::new(); for job in jobs { let job_id = job.id.ok_or("Job missing ID")?; job_id_to_name.insert(job_id, job.name.clone()); job_name_to_id.insert(job.name.clone(), job_id); - - let deps = job.depends_on_job_ids.clone().unwrap_or_default(); - dependency_graph_by_id.insert(job_id, deps); } - // Convert to name-based graph for topological sort - let mut dependency_graph_by_name: HashMap> = HashMap::new(); + // Try to use depends_on_job_ids first + let mut has_any_deps = false; for job in jobs { - let job_id = job.id.unwrap(); - let job_name = &job.name; - let dep_ids = dependency_graph_by_id.get(&job_id).unwrap(); - let dep_names: Vec = dep_ids + let dep_ids = job.depends_on_job_ids.clone().unwrap_or_default(); + if !dep_ids.is_empty() { + has_any_deps = true; + } + let dep_names: HashSet = dep_ids .iter() .filter_map(|id| job_id_to_name.get(id).cloned()) .collect(); - dependency_graph_by_name.insert(job_name.clone(), dep_names); + dependency_graph_by_name.insert(job.name.clone(), dep_names); } - // Topologically sort jobs into levels - let job_levels = topological_sort_by_names(jobs, &dependency_graph_by_name)?; - - // Stage 0: on_workflow_start actions - let workflow_start_stage = build_workflow_start_stage_from_db( - jobs, - actions, - &dependency_graph_by_name, - &job_name_to_id, - )?; - if !workflow_start_stage.scheduler_allocations.is_empty() - || !workflow_start_stage.jobs_becoming_ready.is_empty() - { - stages.push(workflow_start_stage); - } + // If no explicit dependencies found, compute from file relationships + // This handles cases where workflow hasn't been initialized yet + if !has_any_deps { + // Build file_id -> producing job_id map + let mut file_producers: HashMap = HashMap::new(); + for job in jobs { + if let Some(output_ids) = &job.output_file_ids { + let job_id = job.id.ok_or("Job missing ID")?; + for file_id in output_ids { + file_producers.insert(*file_id, job_id); + } + } + } - // Build stages for each job level - for (level_idx, level_jobs) in job_levels.iter().enumerate() { - if level_jobs.is_empty() { - continue; + // Compute dependencies from input files + for job in jobs { + let job_id = job.id.ok_or("Job missing ID")?; + let mut deps: HashSet = HashSet::new(); + + if let Some(input_ids) = &job.input_file_ids { + for file_id in input_ids { + if let Some(producer_id) = file_producers.get(file_id) { + if *producer_id != job_id { + if let Some(producer_name) = job_id_to_name.get(producer_id) { + deps.insert(producer_name.clone()); + } + } + } + } + } + + dependency_graph_by_name.insert(job.name.clone(), deps); } + } - let stage_number = stages.len(); - let stage = build_job_completion_stage_from_db( - jobs, - actions, - stage_number, - level_idx, - level_jobs, - &job_levels, - &job_name_to_id, - )?; - - // Only add the stage if it has actions or jobs becoming ready - if !stage.scheduler_allocations.is_empty() || !stage.jobs_becoming_ready.is_empty() { - stages.push(stage); + // Find root jobs (no dependencies) + let root_jobs: Vec = jobs + .iter() + .filter(|j| { + dependency_graph_by_name + .get(&j.name) + .map(|deps| deps.is_empty()) + .unwrap_or(true) + }) + .map(|j| j.name.clone()) + .collect(); + + // Find on_workflow_start scheduler allocations + let mut start_allocations = Vec::new(); + for action in actions { + if action.trigger_type == "on_workflow_start" { + if let Some(alloc) = + build_scheduler_allocation_from_db_action(action, jobs, &scheduler_id_to_name)? + { + start_allocations.push(alloc); + } } } - Ok(ExecutionPlan { stages }) - } + let start_event = ExecutionEvent { + id: "start".to_string(), + trigger: EventTrigger::WorkflowStart, + trigger_description: "Workflow Start".to_string(), + scheduler_allocations: start_allocations, + jobs_becoming_ready: root_jobs.clone(), + depends_on_events: vec![], + unlocks_events: vec![], + }; + events.insert("start".to_string(), start_event); + + // Build a map: job -> event_id for the event where this job completes + let mut job_completion_event: HashMap = HashMap::new(); + for job in &root_jobs { + job_completion_event.insert(job.clone(), "start".to_string()); + } - /// Display the execution plan in a human-readable format - pub fn display(&self) { - println!("\n{}", "=".repeat(80)); - println!("Workflow Execution Plan"); - println!("{}", "=".repeat(80)); + // Process jobs level by level, similar to from_spec() + let mut processed_jobs: HashSet = root_jobs.iter().cloned().collect(); - for stage in &self.stages { - println!( - "\n{} Stage {}: {}", - if stage.stage_number == 0 { - "▶" - } else { - "→" - }, - stage.stage_number + 1, // Display as 1-based - stage.trigger_description - ); - println!("{}", "-".repeat(80)); - - if !stage.scheduler_allocations.is_empty() { - println!("\n Scheduler Allocations:"); - for alloc in &stage.scheduler_allocations { - println!( - " • {} ({}) - {} allocation(s)", - alloc.scheduler, alloc.scheduler_type, alloc.num_allocations - ); - let compact = compact_job_list(&alloc.jobs); - println!(" For jobs: {}", compact.join(", ")); + loop { + // Find jobs that become ready based on currently processed jobs + let mut newly_ready = Vec::new(); + for job in jobs { + if processed_jobs.contains(&job.name) { + continue; } - } - - if !stage.jobs_becoming_ready.is_empty() { - println!("\n Jobs Becoming Ready:"); - let compact = compact_job_list(&stage.jobs_becoming_ready); - for item in compact { - println!(" • {}", item); + if let Some(deps) = dependency_graph_by_name.get(&job.name) { + if !deps.is_empty() && deps.iter().all(|d| processed_jobs.contains(d)) { + newly_ready.push(job.name.clone()); + } } } - if stage.scheduler_allocations.is_empty() && stage.jobs_becoming_ready.is_empty() { - println!(" (No actions or jobs in this stage)"); + if newly_ready.is_empty() { + break; } - } - println!("\n{}", "=".repeat(80)); - println!("Total Stages: {}", self.stages.len()); - println!("{}\n", "=".repeat(80)); - } -} - -/// Build dependency graph: job_name -> vec of job names it depends on -fn build_dependency_graph( - spec: &WorkflowSpec, -) -> Result>, Box> { - let mut deps = HashMap::new(); - - for job in &spec.jobs { - let mut job_deps = Vec::new(); + // Group newly ready jobs by their dependencies + let mut jobs_by_deps: HashMap, Vec> = HashMap::new(); + + for job_name in &newly_ready { + if let Some(deps) = dependency_graph_by_name.get(job_name) { + let mut dep_list: Vec = deps.iter().cloned().collect(); + dep_list.sort(); + jobs_by_deps + .entry(dep_list) + .or_default() + .push(job_name.clone()); + } + } - // Add explicit dependencies from depends_on - if let Some(ref names) = job.depends_on { - job_deps.extend(names.clone()); - } + // Create an event for each unique dependency set + for (deps, jobs_becoming_ready) in jobs_by_deps { + // Generate a descriptive event ID based on the triggering jobs + let event_id = generate_event_id(&deps); - // Add dependencies from depends_on_regexes - if let Some(ref regexes) = job.depends_on_regexes { - for regex_str in regexes { - let re = Regex::new(regex_str)?; - for other_job in &spec.jobs { - if re.is_match(&other_job.name) && !job_deps.contains(&other_job.name) { - job_deps.push(other_job.name.clone()); + // Find which events the dependencies belong to + let mut depends_on_event_ids: HashSet = HashSet::new(); + for dep in &deps { + if let Some(evt_id) = job_completion_event.get(dep) { + depends_on_event_ids.insert(evt_id.clone()); } } - } - } - - // Add implicit dependencies from input files - if let Some(ref input_files) = job.input_files { - for input_file in input_files { - // Find jobs that produce this file - for other_job in &spec.jobs { - if let Some(ref output_files) = other_job.output_files { - if output_files.contains(input_file) && !job_deps.contains(&other_job.name) - { - job_deps.push(other_job.name.clone()); + let depends_on_events: Vec = depends_on_event_ids.into_iter().collect(); + + // Find matching scheduler actions + let mut scheduler_allocations = Vec::new(); + for action in actions { + if action.trigger_type == "on_jobs_ready" { + let empty_vec = vec![]; + let action_job_ids = action.job_ids.as_ref().unwrap_or(&empty_vec); + + let matches_any = action_job_ids.iter().any(|action_job_id| { + jobs_becoming_ready + .iter() + .any(|job_name| job_name_to_id.get(job_name) == Some(action_job_id)) + }); + + if matches_any { + if let Some(alloc) = build_scheduler_allocation_from_db_action( + action, + jobs, + &scheduler_id_to_name, + )? { + scheduler_allocations.push(alloc); + } } } } + + let trigger_description = build_trigger_description(&deps); + + let event = ExecutionEvent { + id: event_id.clone(), + trigger: EventTrigger::JobsComplete { jobs: deps.clone() }, + trigger_description, + scheduler_allocations, + jobs_becoming_ready: jobs_becoming_ready.clone(), + depends_on_events, + unlocks_events: vec![], + }; + events.insert(event_id.clone(), event); + + // Mark these jobs as completing in this event + for job in &jobs_becoming_ready { + job_completion_event.insert(job.clone(), event_id.clone()); + } } + + // Add newly ready jobs to processed set + processed_jobs.extend(newly_ready); } - // Add implicit dependencies from input user data - if let Some(ref input_data) = job.input_user_data { - for input_datum in input_data { - // Find jobs that produce this data - for other_job in &spec.jobs { - if let Some(ref output_user_data) = other_job.output_user_data { - if output_user_data.contains(input_datum) - && !job_deps.contains(&other_job.name) - { - job_deps.push(other_job.name.clone()); - } + // Build unlocks_events by reversing the depends_on_events relationships + let event_deps: Vec<(String, Vec)> = events + .values() + .map(|e| (e.id.clone(), e.depends_on_events.clone())) + .collect(); + + for (event_id, deps) in event_deps { + for dep_event_id in deps { + if let Some(dep_event) = events.get_mut(&dep_event_id) { + if !dep_event.unlocks_events.contains(&event_id) { + dep_event.unlocks_events.push(event_id.clone()); } } } } - deps.insert(job.name.clone(), job_deps); + // Identify root and leaf events + let root_events: Vec = events + .values() + .filter(|e| e.depends_on_events.is_empty()) + .map(|e| e.id.clone()) + .collect(); + + let leaf_events: Vec = events + .values() + .filter(|e| e.unlocks_events.is_empty()) + .map(|e| e.id.clone()) + .collect(); + + Ok(ExecutionPlan { + events, + root_events, + leaf_events, + graph: None, + }) } - Ok(deps) -} - -/// Build reverse dependency graph: job_name -> vec of jobs that depend on it -fn build_reverse_dependencies(deps: &HashMap>) -> HashMap> { - let mut reverse = HashMap::new(); + /// Display the execution plan in a human-readable format + pub fn display(&self) { + println!("\n{}", "=".repeat(80)); + println!("Workflow Execution Plan (DAG)"); + println!("{}", "=".repeat(80)); - for (job, dependencies) in deps { - for dep in dependencies { - reverse - .entry(dep.clone()) - .or_insert_with(Vec::new) - .push(job.clone()); + // Show connected components if we have a graph + if let Some(ref graph) = self.graph { + let mut graph_clone = graph.clone(); + let components = graph_clone.connected_components(); + if components.len() > 1 { + println!( + "\nNote: Workflow has {} independent sub-workflows that can run in parallel", + components.len() + ); + } } - } - reverse -} + println!("\nEvents: {} total", self.events.len()); + println!("Root events: {}", self.root_events.join(", ")); + println!("Leaf events: {}", self.leaf_events.join(", ")); -/// Topologically sort jobs into levels based on dependencies -fn topological_sort( - spec: &WorkflowSpec, - deps: &HashMap>, -) -> Result>, Box> { - let mut levels = Vec::new(); - let mut remaining: HashSet = spec.jobs.iter().map(|j| j.name.clone()).collect(); - let mut processed = HashSet::new(); - - while !remaining.is_empty() { - let mut current_level = Vec::new(); - - // Find all jobs whose dependencies are satisfied - for job in &spec.jobs { - if remaining.contains(&job.name) { - let job_deps = deps.get(&job.name).unwrap(); - if job_deps.iter().all(|d| processed.contains(d)) { - current_level.push(job.name.clone()); + // Display events in topological order (BFS from roots) + let mut displayed = HashSet::new(); + let mut queue: Vec = self.root_events.clone(); + + while !queue.is_empty() { + // Find events that can be displayed (all dependencies displayed) + let mut next_queue = Vec::new(); + let mut events_to_display = Vec::new(); + + for event_id in &queue { + if displayed.contains(event_id) { + continue; + } + + if let Some(event) = self.events.get(event_id) { + if event + .depends_on_events + .iter() + .all(|d| displayed.contains(d)) + { + events_to_display.push(event_id.clone()); + } else { + next_queue.push(event_id.clone()); + } } } - } - if current_level.is_empty() { - return Err("Circular dependency detected in job graph".into()); - } + // Display these events + for event_id in events_to_display { + if let Some(event) = self.events.get(&event_id) { + let symbol = match &event.trigger { + EventTrigger::WorkflowStart => "▶", + EventTrigger::JobsComplete { .. } => "→", + }; + + println!("\n{} {}", symbol, event.trigger_description); + println!("{}", "-".repeat(80)); + + // Show flow information + if !event.depends_on_events.is_empty() { + let dep_descriptions: Vec = event + .depends_on_events + .iter() + .filter_map(|id| self.events.get(id)) + .map(|e| format!("'{}'", e.trigger_description)) + .collect(); + println!(" After: {}", dep_descriptions.join(", ")); + } - // Mark these jobs as processed - for job in ¤t_level { - remaining.remove(job); - processed.insert(job.clone()); - } + if !event.unlocks_events.is_empty() { + let unlock_descriptions: Vec = event + .unlocks_events + .iter() + .filter_map(|id| self.events.get(id)) + .map(|e| format!("'{}'", e.trigger_description)) + .collect(); + println!(" Then: {}", unlock_descriptions.join(", ")); + } - levels.push(current_level); - } + if !event.scheduler_allocations.is_empty() { + println!("\n Scheduler Allocations:"); + for alloc in &event.scheduler_allocations { + println!( + " • {} ({}) - {} allocation(s)", + alloc.scheduler, alloc.scheduler_type, alloc.num_allocations + ); + let compact = compact_job_list(&alloc.jobs); + println!(" For jobs: {}", compact.join(", ")); + } + } - Ok(levels) -} + if !event.jobs_becoming_ready.is_empty() { + println!("\n Jobs Becoming Ready:"); + let compact = compact_job_list(&event.jobs_becoming_ready); + for item in compact { + println!(" • {}", item); + } + } -/// Build stage for workflow start actions -fn build_workflow_start_stage( - spec: &WorkflowSpec, -) -> Result> { - let mut scheduler_allocations = Vec::new(); + displayed.insert(event_id.clone()); - // Find on_workflow_start actions - if let Some(ref actions) = spec.actions { - for action in actions { - if action.trigger_type == "on_workflow_start" { - if let Some(alloc) = build_scheduler_allocation(spec, action)? { - scheduler_allocations.push(alloc); + // Add events this one unlocks to the next queue + for unlocked in &event.unlocks_events { + if !displayed.contains(unlocked) && !next_queue.contains(unlocked) { + next_queue.push(unlocked.clone()); + } + } } } + + queue = next_queue; } - } - // Build dependency graph to find jobs with no dependencies - let dependency_graph = build_dependency_graph(spec)?; + println!("\n{}", "=".repeat(80)); + println!("Total Events: {}", self.events.len()); + println!("{}\n", "=".repeat(80)); + } - // Jobs with empty dependency lists become ready at start - let mut jobs_becoming_ready = Vec::new(); - for job in &spec.jobs { - if let Some(deps) = dependency_graph.get(&job.name) { - if deps.is_empty() { - jobs_becoming_ready.push(job.name.clone()); - } - } + /// Get the underlying workflow graph (if available) + pub fn workflow_graph(&self) -> Option<&WorkflowGraph> { + self.graph.as_ref() } - Ok(ExecutionStage { - stage_number: 0, - trigger_description: "Workflow Start".to_string(), - scheduler_allocations, - jobs_becoming_ready, - }) -} + /// Convert to legacy stage format for backwards compatibility + /// Note: This flattens the DAG into a linear sequence, which may not + /// accurately represent parallel subgraphs + pub fn to_stages(&self) -> Vec { + let mut stages = Vec::new(); + let mut displayed = HashSet::new(); + let mut queue: Vec = self.root_events.clone(); + let mut stage_number = 0; -/// Build stage for when a level of jobs completes -fn build_job_completion_stage( - spec: &WorkflowSpec, - stage_number: usize, - level_idx: usize, - level_jobs: &[String], - all_levels: &[Vec], - _reverse_deps: &HashMap>, -) -> Result> { - let mut scheduler_allocations = Vec::new(); - - // Find jobs that become ready after this level completes - let jobs_becoming_ready: Vec = if level_idx + 1 < all_levels.len() { - all_levels[level_idx + 1].clone() - } else { - Vec::new() - }; + while !queue.is_empty() { + let mut next_queue = Vec::new(); + let mut events_at_level = Vec::new(); - // Find on_jobs_ready actions that match the jobs becoming ready - if let Some(ref actions) = spec.actions { - for action in actions { - if action.trigger_type == "on_jobs_ready" { - // Check if this action matches any of the jobs becoming ready - let action_jobs = get_matching_jobs(spec, action)?; - let matches_any = action_jobs.iter().any(|j| jobs_becoming_ready.contains(j)); + for event_id in &queue { + if displayed.contains(event_id) { + continue; + } - if matches_any { - if let Some(alloc) = build_scheduler_allocation(spec, action)? { - scheduler_allocations.push(alloc); + if let Some(event) = self.events.get(event_id) { + if event + .depends_on_events + .iter() + .all(|d| displayed.contains(d)) + { + events_at_level.push(event_id.clone()); + } else { + next_queue.push(event_id.clone()); } } } + + // Merge events at the same level into a single stage + if !events_at_level.is_empty() { + let mut merged_allocations = Vec::new(); + let mut merged_jobs = Vec::new(); + let mut trigger_parts = Vec::new(); + + for event_id in &events_at_level { + if let Some(event) = self.events.get(event_id) { + merged_allocations.extend(event.scheduler_allocations.clone()); + merged_jobs.extend(event.jobs_becoming_ready.clone()); + trigger_parts.push(event.trigger_description.clone()); + displayed.insert(event_id.clone()); + + for unlocked in &event.unlocks_events { + if !displayed.contains(unlocked) && !next_queue.contains(unlocked) { + next_queue.push(unlocked.clone()); + } + } + } + } + + let trigger_description = if trigger_parts.len() == 1 { + trigger_parts[0].clone() + } else { + trigger_parts.join(" AND ") + }; + + stages.push(ExecutionStage { + stage_number, + trigger_description, + scheduler_allocations: merged_allocations, + jobs_becoming_ready: merged_jobs, + }); + stage_number += 1; + } + + queue = next_queue; } + + stages } +} - // Build trigger description - let trigger_description = if level_jobs.len() == 1 { +/// Build trigger description for a set of completing jobs +fn build_trigger_description(level_jobs: &[String]) -> String { + if level_jobs.len() == 1 { format!("When job '{}' completes", level_jobs[0]) } else if level_jobs.len() <= 3 { format!( @@ -414,14 +708,49 @@ fn build_job_completion_stage( .collect::>() .join(", ") ) - }; + } +} + +/// Generate a descriptive event ID based on the jobs that trigger it +fn generate_event_id(trigger_jobs: &[String]) -> String { + if trigger_jobs.is_empty() { + return "start".to_string(); + } + + // For a single job, use "after_" + if trigger_jobs.len() == 1 { + return format!("after_{}", trigger_jobs[0]); + } - Ok(ExecutionStage { - stage_number, - trigger_description, - scheduler_allocations, - jobs_becoming_ready, - }) + // For multiple jobs, try to find a common prefix + let first = &trigger_jobs[0]; + let mut common_prefix_len = first.len(); + + for job in trigger_jobs.iter().skip(1) { + let matching = first + .chars() + .zip(job.chars()) + .take_while(|(a, b)| a == b) + .count(); + common_prefix_len = common_prefix_len.min(matching); + } + + // Use common prefix if it's meaningful (at least 2 chars, not just "work" or similar) + if common_prefix_len >= 3 { + let prefix = &first[..common_prefix_len]; + // Clean up trailing underscores or numbers + let clean_prefix = prefix.trim_end_matches(|c: char| c == '_' || c.is_ascii_digit()); + if clean_prefix.len() >= 3 { + return format!("after_{}_jobs", clean_prefix); + } + } + + // Fall back to listing first job and count + format!( + "after_{}_and_{}_more", + trigger_jobs[0], + trigger_jobs.len() - 1 + ) } /// Get job names that match an action's job_names or job_name_regexes @@ -484,157 +813,11 @@ fn build_scheduler_allocation( })) } -/// Topologically sort jobs by name -fn topological_sort_by_names( - jobs: &[JobModel], - deps: &HashMap>, -) -> Result>, Box> { - let mut levels = Vec::new(); - let mut remaining: HashSet = jobs.iter().map(|j| j.name.clone()).collect(); - let mut processed = HashSet::new(); - - while !remaining.is_empty() { - let mut current_level = Vec::new(); - - // Find all jobs whose dependencies are satisfied - for job in jobs { - if remaining.contains(&job.name) { - let job_deps = deps.get(&job.name).unwrap(); - if job_deps.iter().all(|d| processed.contains(d)) { - current_level.push(job.name.clone()); - } - } - } - - if current_level.is_empty() { - return Err("Circular dependency detected in job graph".into()); - } - - // Mark these jobs as processed - for job in ¤t_level { - remaining.remove(job); - processed.insert(job.clone()); - } - - levels.push(current_level); - } - - Ok(levels) -} - -/// Build stage for workflow start from database models -fn build_workflow_start_stage_from_db( - jobs: &[JobModel], - actions: &[WorkflowActionModel], - dependency_graph: &HashMap>, - _job_name_to_id: &HashMap, -) -> Result> { - let mut scheduler_allocations = Vec::new(); - - // Find on_workflow_start actions - for action in actions { - if action.trigger_type == "on_workflow_start" { - if let Some(alloc) = build_scheduler_allocation_from_db_action(action, jobs)? { - scheduler_allocations.push(alloc); - } - } - } - - // Jobs with empty dependency lists become ready at start - let mut jobs_becoming_ready = Vec::new(); - for job in jobs { - if let Some(deps) = dependency_graph.get(&job.name) { - if deps.is_empty() { - jobs_becoming_ready.push(job.name.clone()); - } - } - } - - Ok(ExecutionStage { - stage_number: 0, - trigger_description: "Workflow Start".to_string(), - scheduler_allocations, - jobs_becoming_ready, - }) -} - -/// Build stage for when a level of jobs completes from database models -fn build_job_completion_stage_from_db( - jobs: &[JobModel], - actions: &[WorkflowActionModel], - stage_number: usize, - level_idx: usize, - level_jobs: &[String], - all_levels: &[Vec], - job_name_to_id: &HashMap, -) -> Result> { - let mut scheduler_allocations = Vec::new(); - - // Find jobs that become ready after this level completes - let jobs_becoming_ready: Vec = if level_idx + 1 < all_levels.len() { - all_levels[level_idx + 1].clone() - } else { - Vec::new() - }; - - // Find on_jobs_ready actions that match the jobs becoming ready - for action in actions { - if action.trigger_type == "on_jobs_ready" { - // Check if this action matches any of the jobs becoming ready - let empty_vec = vec![]; - let action_job_ids = action.job_ids.as_ref().unwrap_or(&empty_vec); - - let matches_any = action_job_ids.iter().any(|action_job_id| { - jobs_becoming_ready - .iter() - .any(|job_name| job_name_to_id.get(job_name) == Some(action_job_id)) - }); - - if matches_any { - if let Some(alloc) = build_scheduler_allocation_from_db_action(action, jobs)? { - scheduler_allocations.push(alloc); - } - } - } - } - - // Build trigger description - let trigger_description = if level_jobs.len() == 1 { - format!("When job '{}' completes", level_jobs[0]) - } else if level_jobs.len() <= 3 { - format!( - "When jobs {} complete", - level_jobs - .iter() - .map(|j| format!("'{}'", j)) - .collect::>() - .join(", ") - ) - } else { - format!( - "When {} jobs complete ({}...)", - level_jobs.len(), - level_jobs - .iter() - .take(2) - .map(|j| format!("'{}'", j)) - .collect::>() - .join(", ") - ) - }; - - Ok(ExecutionStage { - stage_number, - trigger_description, - scheduler_allocations, - jobs_becoming_ready, - }) -} - /// Build a scheduler allocation from a database action model fn build_scheduler_allocation_from_db_action( action: &WorkflowActionModel, workflow_jobs: &[JobModel], + scheduler_id_to_name: &HashMap, ) -> Result, Box> { if action.action_type != "schedule_nodes" { return Ok(None); @@ -660,9 +843,15 @@ fn build_scheduler_allocation_from_db_action( } } - // For scheduler, we'd need to look it up from the database - // For now, use a placeholder based on the scheduler_type - let scheduler = format!("{}_scheduler", scheduler_type); + // Look up the scheduler name from the ID + let scheduler = if let Some(scheduler_id) = config["scheduler_id"].as_i64() { + scheduler_id_to_name + .get(&scheduler_id) + .cloned() + .unwrap_or_else(|| format!("{}_scheduler", scheduler_type)) + } else { + format!("{}_scheduler", scheduler_type) + }; Ok(Some(SchedulerAllocation { scheduler, diff --git a/src/client/hpc/kestrel.rs b/src/client/hpc/kestrel.rs new file mode 100644 index 00000000..e7c6c66f --- /dev/null +++ b/src/client/hpc/kestrel.rs @@ -0,0 +1,327 @@ +//! NREL Kestrel HPC profile +//! +//! Kestrel is NREL's flagship HPC system featuring: +//! - 2,240 standard CPU nodes (104 cores, 240GB RAM each) +//! - 156 GPU nodes with 4x NVIDIA H100 GPUs (80GB each) +//! - Various specialized partitions for different workload types +//! +//! Detection: Environment variable NREL_CLUSTER=kestrel + +use super::profiles::{HpcDetection, HpcPartition, HpcProfile}; + +/// Create the Kestrel HPC profile +pub fn kestrel_profile() -> HpcProfile { + HpcProfile { + name: "kestrel".to_string(), + display_name: "NREL Kestrel".to_string(), + description: "NREL's flagship HPC system with CPU and GPU nodes".to_string(), + detection: vec![HpcDetection::EnvVar { + name: "NREL_CLUSTER".to_string(), + value: "kestrel".to_string(), + }], + default_account: None, + partitions: kestrel_partitions(), + charge_factor_cpu: 10.0, + charge_factor_gpu: 100.0, + metadata: [ + ( + "documentation".to_string(), + "https://nrel.github.io/HPC/Documentation/Systems/Kestrel/Running/".to_string(), + ), + ("support_email".to_string(), "HPC-Help@nrel.gov".to_string()), + ] + .into_iter() + .collect(), + } +} + +fn kestrel_partitions() -> Vec { + vec![ + // Debug partition + HpcPartition { + name: "debug".to_string(), + description: "Nodes dedicated to developing and troubleshooting jobs".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 3600, // 1 hour + max_nodes: Some(2), + max_nodes_per_user: Some(2), + min_nodes: None, + gpus_per_node: Some(2), // Max 2 GPUs per user in debug + gpu_type: Some("h100".to_string()), + gpu_memory_gb: Some(80), + local_disk_gb: None, + shared: true, + requires_explicit_request: true, + default_qos: None, + features: vec!["debug".to_string()], + }, + // Short partition (<=4 hours) + HpcPartition { + name: "short".to_string(), + description: "Nodes that prefer jobs with walltimes <= 4 hours".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, // ~240G usable (984256M total but we use practical limit) + max_walltime_secs: 4 * 3600, // 4 hours + max_nodes: Some(2240), + max_nodes_per_user: None, + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: false, // Auto-routed based on walltime + default_qos: None, + features: vec![], + }, + // Standard partition (<=2 days) + HpcPartition { + name: "standard".to_string(), + description: "Nodes that prefer jobs with walltimes <= 2 days".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 2 * 24 * 3600, // 2 days + max_nodes: Some(2240), + max_nodes_per_user: Some(1050), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: false, + default_qos: None, + features: vec![], + }, + // Long partition (>2 days, up to 10 days) + HpcPartition { + name: "long".to_string(), + description: "Nodes that prefer jobs with walltimes > 2 days (max 10 days)".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 10 * 24 * 3600, // 10 days + max_nodes: Some(430), + max_nodes_per_user: Some(215), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: false, + default_qos: None, + features: vec![], + }, + // Medium memory partition (1TB RAM) + HpcPartition { + name: "medmem".to_string(), + description: "Nodes with 1TB of RAM".to_string(), + cpus_per_node: 104, + memory_mb: 1_000_000, // ~1TB + max_walltime_secs: 10 * 24 * 3600, // 10 days + max_nodes: Some(64), + max_nodes_per_user: Some(32), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: false, // Auto-routed based on memory + default_qos: None, + features: vec!["highmem".to_string()], + }, + // Big memory partition (2TB RAM, short walltime) + HpcPartition { + name: "bigmem".to_string(), + description: "Nodes with 2TB RAM and 5.6TB NVMe local disk".to_string(), + cpus_per_node: 104, + memory_mb: 2_000_000, // ~2TB + max_walltime_secs: 2 * 24 * 3600, // 2 days + max_nodes: Some(10), + max_nodes_per_user: Some(4), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: Some(5600), + shared: false, + requires_explicit_request: false, // Auto-routed based on memory + default_qos: None, + features: vec!["bigmem".to_string(), "nvme".to_string()], + }, + // Big memory long partition + HpcPartition { + name: "bigmeml".to_string(), + description: "Bigmem nodes for jobs > 2 days (max 10 days)".to_string(), + cpus_per_node: 104, + memory_mb: 2_000_000, + max_walltime_secs: 10 * 24 * 3600, // 10 days + max_nodes: Some(4), + max_nodes_per_user: Some(2), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: Some(5600), + shared: false, + requires_explicit_request: false, + default_qos: None, + features: vec!["bigmem".to_string(), "nvme".to_string()], + }, + // High bandwidth partition (dual NIC) + HpcPartition { + name: "hbw".to_string(), + description: "CPU nodes with dual network interface cards for multi-node jobs" + .to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 2 * 24 * 3600, // 2 days + max_nodes: Some(512), + max_nodes_per_user: Some(256), + min_nodes: Some(2), // Minimum 2 nodes required + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: true, // Must specify -p hbw + default_qos: None, + features: vec!["dual-nic".to_string()], + }, + // High bandwidth long partition + HpcPartition { + name: "hbwl".to_string(), + description: "HBW nodes for jobs > 2 days (max 10 days)".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 10 * 24 * 3600, // 10 days + max_nodes: Some(128), + max_nodes_per_user: Some(64), + min_nodes: Some(2), + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: true, + default_qos: None, + features: vec!["dual-nic".to_string()], + }, + // NVMe partition + HpcPartition { + name: "nvme".to_string(), + description: "CPU nodes with 1.7TB NVMe local drives".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 2 * 24 * 3600, // 2 days + max_nodes: Some(256), + max_nodes_per_user: Some(128), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: Some(1700), + shared: false, + requires_explicit_request: true, // Must specify -p nvme + default_qos: None, + features: vec!["nvme".to_string()], + }, + // Shared partition + HpcPartition { + name: "shared".to_string(), + description: "Nodes that can be shared by multiple users and jobs".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 2 * 24 * 3600, // 2 days + max_nodes: Some(128), + max_nodes_per_user: Some(64), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: true, + requires_explicit_request: true, // Must specify -p shared + default_qos: None, + features: vec!["shared".to_string()], + }, + // Shared long partition + HpcPartition { + name: "sharedl".to_string(), + description: "Shared nodes for jobs > 2 days".to_string(), + cpus_per_node: 104, + memory_mb: 240_000, + max_walltime_secs: 10 * 24 * 3600, // Docs say 2 days but listing says 10 days pattern + max_nodes: Some(32), + max_nodes_per_user: Some(16), + min_nodes: None, + gpus_per_node: None, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: true, + requires_explicit_request: true, + default_qos: None, + features: vec!["shared".to_string()], + }, + // GPU H100 partition (short walltime, <= 4 hours preferred) + HpcPartition { + name: "gpu-h100s".to_string(), + description: "GPU nodes preferring jobs <= 4 hours".to_string(), + cpus_per_node: 128, + memory_mb: 360_000, // ~384G base, some have more + max_walltime_secs: 4 * 3600, // 4 hours + max_nodes: Some(156), + max_nodes_per_user: None, + min_nodes: None, + gpus_per_node: Some(4), + gpu_type: Some("h100".to_string()), + gpu_memory_gb: Some(80), + local_disk_gb: Some(3400), // 3.4TB + shared: true, // GPU nodes are always shared + requires_explicit_request: false, + default_qos: None, + features: vec!["gpu".to_string(), "h100".to_string()], + }, + // GPU H100 partition (standard, <= 2 days) + HpcPartition { + name: "gpu-h100".to_string(), + description: "GPU nodes with 4x NVIDIA H100 SXM 80GB".to_string(), + cpus_per_node: 128, + memory_mb: 360_000, + max_walltime_secs: 2 * 24 * 3600, // 2 days + max_nodes: Some(156), + max_nodes_per_user: None, + min_nodes: None, + gpus_per_node: Some(4), + gpu_type: Some("h100".to_string()), + gpu_memory_gb: Some(80), + local_disk_gb: Some(3400), + shared: true, + requires_explicit_request: false, + default_qos: None, + features: vec!["gpu".to_string(), "h100".to_string()], + }, + // GPU H100 long partition (> 2 days) + HpcPartition { + name: "gpu-h100l".to_string(), + description: "GPU nodes for jobs > 2 days".to_string(), + cpus_per_node: 128, + memory_mb: 360_000, + max_walltime_secs: 10 * 24 * 3600, // 10 days (assumed from pattern) + max_nodes: Some(39), + max_nodes_per_user: None, + min_nodes: None, + gpus_per_node: Some(4), + gpu_type: Some("h100".to_string()), + gpu_memory_gb: Some(80), + local_disk_gb: Some(3400), + shared: true, + requires_explicit_request: false, + default_qos: None, + features: vec!["gpu".to_string(), "h100".to_string()], + }, + ] +} diff --git a/src/client/hpc/mod.rs b/src/client/hpc/mod.rs index 6edde902..36e9c5ed 100644 --- a/src/client/hpc/mod.rs +++ b/src/client/hpc/mod.rs @@ -3,15 +3,21 @@ //! This module provides abstractions for working with HPC schedulers like Slurm. //! It includes traits for HPC interfaces and concrete implementations for different //! scheduler types. +//! +//! It also provides HPC system profiles for known HPC systems (like NREL Kestrel) +//! that include partition configurations, resource limits, and auto-detection. pub mod common; pub mod hpc_interface; pub mod hpc_manager; +pub mod kestrel; +pub mod profiles; pub mod slurm_interface; pub use common::{HpcJobInfo, HpcJobStats, HpcJobStatus, HpcType}; pub use hpc_interface::HpcInterface; pub use hpc_manager::HpcManager; +pub use profiles::{HpcDetection, HpcPartition, HpcProfile, HpcProfileRegistry}; pub use slurm_interface::SlurmInterface; use anyhow::Result; diff --git a/src/client/hpc/profiles.rs b/src/client/hpc/profiles.rs new file mode 100644 index 00000000..bf8daf7e --- /dev/null +++ b/src/client/hpc/profiles.rs @@ -0,0 +1,383 @@ +//! HPC profile definitions for known HPC systems +//! +//! This module provides data structures for defining HPC system profiles, +//! including partition configurations, resource limits, and auto-detection. + +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::env; +use std::time::Duration; + +/// How to detect if we're running on a particular HPC system +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum HpcDetection { + /// Detect by environment variable value + EnvVar { + /// Environment variable name + name: String, + /// Expected value + value: String, + }, + /// Detect by hostname pattern (regex) + HostnamePattern { + /// Regex pattern to match hostname + pattern: String, + }, + /// Detect by existence of a file + FileExists { + /// Path to check + path: String, + }, +} + +impl HpcDetection { + /// Check if this detection method matches the current environment + pub fn matches(&self) -> bool { + match self { + HpcDetection::EnvVar { name, value } => { + env::var(name).map(|v| v == *value).unwrap_or(false) + } + HpcDetection::HostnamePattern { pattern } => { + if let Ok(hostname) = hostname::get() { + if let Some(hostname_str) = hostname.to_str() { + if let Ok(re) = regex::Regex::new(pattern) { + return re.is_match(hostname_str); + } + } + } + false + } + HpcDetection::FileExists { path } => std::path::Path::new(path).exists(), + } + } +} + +/// A partition on an HPC system +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HpcPartition { + /// Partition name (as used with --partition) + pub name: String, + + /// Human-readable description + #[serde(default)] + pub description: String, + + /// Number of CPUs per node + pub cpus_per_node: u32, + + /// Memory per node in MB + pub memory_mb: u64, + + /// Maximum wall time in seconds + pub max_walltime_secs: u64, + + /// Total nodes in partition (optional) + #[serde(default)] + pub max_nodes: Option, + + /// Maximum nodes per user (optional) + #[serde(default)] + pub max_nodes_per_user: Option, + + /// Minimum nodes per job (e.g., for high-bandwidth partitions) + #[serde(default)] + pub min_nodes: Option, + + /// Number of GPUs per node (if any) + #[serde(default)] + pub gpus_per_node: Option, + + /// GPU type (e.g., "h100", "a100") + #[serde(default)] + pub gpu_type: Option, + + /// GPU memory in GB per GPU + #[serde(default)] + pub gpu_memory_gb: Option, + + /// Local disk storage in GB (if any) + #[serde(default)] + pub local_disk_gb: Option, + + /// Whether the partition supports shared node access + #[serde(default)] + pub shared: bool, + + /// Whether partition must be explicitly requested (vs auto-routed) + #[serde(default)] + pub requires_explicit_request: bool, + + /// Default QOS for this partition + #[serde(default)] + pub default_qos: Option, + + /// Additional constraints or features + #[serde(default)] + pub features: Vec, +} + +impl HpcPartition { + /// Get the maximum wall time as a Duration + pub fn max_walltime(&self) -> Duration { + Duration::from_secs(self.max_walltime_secs) + } + + /// Format wall time as HH:MM:SS string + pub fn max_walltime_str(&self) -> String { + let secs = self.max_walltime_secs; + let hours = secs / 3600; + let mins = (secs % 3600) / 60; + let s = secs % 60; + + if hours >= 24 { + let days = hours / 24; + let h = hours % 24; + format!("{}-{:02}:{:02}:{:02}", days, h, mins, s) + } else { + format!("{:02}:{:02}:{:02}", hours, mins, s) + } + } + + /// Get memory in GB + pub fn memory_gb(&self) -> f64 { + self.memory_mb as f64 / 1024.0 + } + + /// Check if this partition can satisfy the given requirements + pub fn can_satisfy( + &self, + cpus: u32, + memory_mb: u64, + walltime_secs: u64, + gpus: Option, + ) -> bool { + // Check CPU + if cpus > self.cpus_per_node { + return false; + } + + // Check memory + if memory_mb > self.memory_mb { + return false; + } + + // Check wall time + if walltime_secs > self.max_walltime_secs { + return false; + } + + // Check GPUs if requested + if let Some(requested_gpus) = gpus { + if requested_gpus > 0 { + match self.gpus_per_node { + Some(available) if requested_gpus <= available => {} + _ => return false, + } + } + } + + true + } +} + +/// An HPC system profile +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HpcProfile { + /// System identifier (e.g., "kestrel", "perlmutter") + pub name: String, + + /// Human-readable display name + pub display_name: String, + + /// Optional description + #[serde(default)] + pub description: String, + + /// Detection methods (any match triggers detection) + pub detection: Vec, + + /// Default account (can be overridden in config) + #[serde(default)] + pub default_account: Option, + + /// Available partitions + pub partitions: Vec, + + /// Charge factor for CPU jobs (AU per node-hour) + #[serde(default = "default_charge_factor")] + pub charge_factor_cpu: f64, + + /// Charge factor for GPU jobs (AU per node-hour) + #[serde(default = "default_charge_factor_gpu")] + pub charge_factor_gpu: f64, + + /// Additional metadata + #[serde(default)] + pub metadata: HashMap, +} + +fn default_charge_factor() -> f64 { + 1.0 +} + +fn default_charge_factor_gpu() -> f64 { + 10.0 +} + +impl HpcProfile { + /// Check if this profile matches the current environment + pub fn detect(&self) -> bool { + self.detection.iter().any(|d| d.matches()) + } + + /// Get a partition by name + pub fn get_partition(&self, name: &str) -> Option<&HpcPartition> { + self.partitions.iter().find(|p| p.name == name) + } + + /// Find partitions that can satisfy the given requirements + pub fn find_matching_partitions( + &self, + cpus: u32, + memory_mb: u64, + walltime_secs: u64, + gpus: Option, + ) -> Vec<&HpcPartition> { + self.partitions + .iter() + .filter(|p| p.can_satisfy(cpus, memory_mb, walltime_secs, gpus)) + .collect() + } + + /// Find the best partition for the given requirements + /// Prefers: GPU partitions if GPUs requested, shared if small job, otherwise standard + /// Avoids: debug partitions (they're for development, not production) + pub fn find_best_partition( + &self, + cpus: u32, + memory_mb: u64, + walltime_secs: u64, + gpus: Option, + ) -> Option<&HpcPartition> { + let matching = self.find_matching_partitions(cpus, memory_mb, walltime_secs, gpus); + + if matching.is_empty() { + return None; + } + + // Filter out debug partitions for automatic selection + let non_debug: Vec<_> = matching + .iter() + .filter(|p| !p.name.to_lowercase().contains("debug")) + .copied() + .collect(); + + // Use non-debug partitions if available, otherwise fall back to all matching + let candidates = if non_debug.is_empty() { + &matching + } else { + &non_debug + }; + + // If GPUs requested, prefer GPU partitions that don't require explicit request + if gpus.map(|g| g > 0).unwrap_or(false) { + // First try auto-routed GPU partitions + if let Some(gpu_partition) = candidates + .iter() + .find(|p| p.gpus_per_node.is_some() && !p.requires_explicit_request) + { + return Some(gpu_partition); + } + // Fall back to any GPU partition + if let Some(gpu_partition) = candidates.iter().find(|p| p.gpus_per_node.is_some()) { + return Some(gpu_partition); + } + } + + // For small jobs, prefer shared partitions that don't require explicit request + let is_small_job = cpus <= 26 && memory_mb <= 60_000; // ~1/4 of standard node + if is_small_job { + // First try auto-routed shared partitions + if let Some(shared_partition) = candidates + .iter() + .find(|p| p.shared && !p.requires_explicit_request) + { + return Some(shared_partition); + } + } + + // Prefer partitions that don't require explicit request (auto-routed) + if let Some(auto_partition) = candidates.iter().find(|p| !p.requires_explicit_request) { + return Some(auto_partition); + } + + // Return first matching from candidates + candidates.first().copied() + } + + /// Get all GPU partitions + pub fn gpu_partitions(&self) -> Vec<&HpcPartition> { + self.partitions + .iter() + .filter(|p| p.gpus_per_node.is_some()) + .collect() + } + + /// Get all CPU-only partitions + pub fn cpu_partitions(&self) -> Vec<&HpcPartition> { + self.partitions + .iter() + .filter(|p| p.gpus_per_node.is_none()) + .collect() + } +} + +/// Registry of known HPC profiles +#[derive(Debug, Clone, Default)] +pub struct HpcProfileRegistry { + profiles: Vec, +} + +impl HpcProfileRegistry { + /// Create a new empty registry + pub fn new() -> Self { + Self { + profiles: Vec::new(), + } + } + + /// Create a registry with all built-in profiles + pub fn with_builtin_profiles() -> Self { + let mut registry = Self::new(); + registry.register(super::kestrel::kestrel_profile()); + registry + } + + /// Register a profile + pub fn register(&mut self, profile: HpcProfile) { + // Remove existing profile with same name + self.profiles.retain(|p| p.name != profile.name); + self.profiles.push(profile); + } + + /// Get all registered profiles + pub fn profiles(&self) -> &[HpcProfile] { + &self.profiles + } + + /// Get a profile by name + pub fn get(&self, name: &str) -> Option<&HpcProfile> { + self.profiles.iter().find(|p| p.name == name) + } + + /// Detect the current HPC system + pub fn detect(&self) -> Option<&HpcProfile> { + self.profiles.iter().find(|p| p.detect()) + } + + /// Get profile names + pub fn names(&self) -> Vec<&str> { + self.profiles.iter().map(|p| p.name.as_str()).collect() + } +} diff --git a/src/client/hpc/slurm_interface.rs b/src/client/hpc/slurm_interface.rs index b6c4f929..3c6c2264 100644 --- a/src/client/hpc/slurm_interface.rs +++ b/src/client/hpc/slurm_interface.rs @@ -13,6 +13,9 @@ use std::thread; use std::time::Duration; use sysinfo::SystemExt; +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; + use super::common::{HpcJobInfo, HpcJobStats, HpcJobStatus}; use super::hpc_interface::HpcInterface; @@ -278,7 +281,6 @@ impl HpcInterface for SlurmInterface { // Make the script executable #[cfg(unix)] { - use std::os::unix::fs::PermissionsExt; let mut perms = fs::metadata(filename)?.permissions(); perms.set_mode(0o755); fs::set_permissions(filename, perms)?; diff --git a/src/client/mod.rs b/src/client/mod.rs index 9d06dbb3..ae53e318 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -18,6 +18,7 @@ pub mod log_paths; pub mod parameter_expansion; pub mod resource_monitor; pub mod utils; +pub mod workflow_graph; pub mod workflow_manager; pub mod workflow_spec; @@ -25,8 +26,8 @@ pub mod workflow_spec; pub use apis::configuration::Configuration; pub use apis::default_api; pub use hpc::{ - HpcInterface, HpcJobInfo, HpcJobStats, HpcJobStatus, HpcManager, HpcType, SlurmInterface, - create_hpc_interface, + HpcDetection, HpcInterface, HpcJobInfo, HpcJobStats, HpcJobStatus, HpcManager, HpcPartition, + HpcProfile, HpcProfileRegistry, HpcType, SlurmInterface, create_hpc_interface, }; pub use job_runner::JobRunner; // JobModel is re-exported from models (which re-exports from crate::models) diff --git a/src/client/parameter_expansion.rs b/src/client/parameter_expansion.rs index ce20830f..d592aa0e 100644 --- a/src/client/parameter_expansion.rs +++ b/src/client/parameter_expansion.rs @@ -54,9 +54,26 @@ impl ParameterValue { /// - Integer ranges: "1:100" (inclusive), "1:100:5" (with step) /// - Float ranges: "0.0:1.0:0.1" /// - Lists: "[1,5,10,50,100]" or "['train','test','validation']" +/// +/// Also tolerates curly braces around values (e.g., "{1:100}" is treated as "1:100") +/// since users sometimes confuse parameter value syntax with template substitution syntax. pub fn parse_parameter_value(value: &str) -> Result, String> { let trimmed = value.trim(); + // Strip curly braces if they wrap the entire value + // This handles the common mistake of using {1:100} instead of 1:100 + // (users confuse parameter values with template substitution syntax like {index}) + let trimmed = if trimmed.starts_with('{') && trimmed.ends_with('}') && !trimmed.contains(',') { + // Only strip if it looks like a wrapped range, not a JSON object + trimmed + .strip_prefix('{') + .and_then(|s| s.strip_suffix('}')) + .unwrap_or(trimmed) + .trim() + } else { + trimmed + }; + // Check for list notation if trimmed.starts_with('[') && trimmed.ends_with(']') { return parse_list(trimmed); @@ -355,6 +372,20 @@ mod tests { assert_eq!(values[4], ParameterValue::Integer(5)); } + #[test] + fn test_parse_integer_range_with_curly_braces() { + // Users sometimes confuse parameter values with template substitution syntax + // e.g., writing {1:1000} instead of 1:1000 + let values = parse_parameter_value("{1:5}").unwrap(); + assert_eq!(values.len(), 5); + assert_eq!(values[0], ParameterValue::Integer(1)); + assert_eq!(values[4], ParameterValue::Integer(5)); + + // Should also work with spaces + let values = parse_parameter_value("{ 1:100 }").unwrap(); + assert_eq!(values.len(), 100); + } + #[test] fn test_parse_integer_range_with_step() { let values = parse_parameter_value("0:10:2").unwrap(); diff --git a/src/client/workflow_graph.rs b/src/client/workflow_graph.rs new file mode 100644 index 00000000..e31dd5f1 --- /dev/null +++ b/src/client/workflow_graph.rs @@ -0,0 +1,817 @@ +//! Workflow Graph - A unified graph-based representation of workflow execution +//! +//! This module provides a directed acyclic graph (DAG) representation of workflow +//! jobs and their dependencies. It supports: +//! +//! - Dependency analysis (topological sorting, levels) +//! - Sub-graph detection (connected components) +//! - Scheduler group generation +//! - Execution plan creation +//! +//! The graph structure enables sophisticated scheduling strategies and visualization. + +use regex::Regex; +use std::collections::{HashMap, HashSet, VecDeque}; + +use crate::client::parameter_expansion::parse_parameter_value; +use crate::client::workflow_spec::{JobSpec, WorkflowActionSpec, WorkflowSpec}; + +/// A node in the workflow graph representing a job (or parameterized job template) +#[derive(Debug, Clone)] +pub struct JobNode { + /// Job name (may contain parameter placeholders like `{index}`) + pub name: String, + /// Resource requirements name + pub resource_requirements: Option, + /// Number of job instances (1 for non-parameterized, N for parameterized) + pub instance_count: usize, + /// Regex pattern matching all instances of this job + pub name_pattern: String, + /// Assigned scheduler name + pub scheduler: Option, + /// Original job spec reference data + pub command: String, +} + +/// Represents a group of jobs that share scheduling characteristics +#[derive(Debug, Clone)] +pub struct SchedulerGroup { + /// Resource requirements name + pub resource_requirements: String, + /// Whether jobs in this group have dependencies + pub has_dependencies: bool, + /// Total job count across all jobs in this group + pub job_count: usize, + /// Job name patterns for matching (regex patterns) + pub job_name_patterns: Vec, + /// Job names in this group + pub job_names: Vec, +} + +/// A connected component (independent sub-workflow) within the graph +#[derive(Debug, Clone)] +pub struct WorkflowComponent { + /// Job names in this component + pub jobs: HashSet, + /// Root jobs (no dependencies within the component) + pub roots: Vec, + /// Leaf jobs (nothing depends on them within the component) + pub leaves: Vec, +} + +/// The main workflow graph structure +#[derive(Debug, Clone)] +pub struct WorkflowGraph { + /// Job nodes indexed by name + nodes: HashMap, + /// Forward edges: job → jobs it depends on (blockers) + depends_on: HashMap>, + /// Reverse edges: job → jobs that depend on it (dependents) + depended_by: HashMap>, + /// Cached topological levels (lazily computed) + levels: Option>>, + /// Cached connected components (lazily computed) + components: Option>, +} + +impl WorkflowGraph { + /// Create a new empty graph + pub fn new() -> Self { + Self { + nodes: HashMap::new(), + depends_on: HashMap::new(), + depended_by: HashMap::new(), + levels: None, + components: None, + } + } + + /// Build a workflow graph from a workflow specification + pub fn from_spec(spec: &WorkflowSpec) -> Result> { + let mut graph = Self::new(); + + // First pass: add all job nodes + for job in &spec.jobs { + let instance_count = count_job_instances(job); + let name_pattern = + build_job_name_pattern(&job.name, job.parameters.is_some(), instance_count); + + let node = JobNode { + name: job.name.clone(), + resource_requirements: job.resource_requirements.clone(), + instance_count, + name_pattern, + scheduler: job.scheduler.clone(), + command: job.command.clone(), + }; + + graph.nodes.insert(job.name.clone(), node); + graph.depends_on.insert(job.name.clone(), HashSet::new()); + graph.depended_by.insert(job.name.clone(), HashSet::new()); + } + + // Second pass: build dependency edges + for job in &spec.jobs { + let mut dependencies = HashSet::new(); + + // Explicit dependencies from depends_on + if let Some(ref deps) = job.depends_on { + for dep in deps { + if graph.nodes.contains_key(dep) { + dependencies.insert(dep.clone()); + } + } + } + + // Dependencies from depends_on_regexes + if let Some(ref regexes) = job.depends_on_regexes { + for regex_str in regexes { + let re = Regex::new(regex_str)?; + for other_job in &spec.jobs { + if re.is_match(&other_job.name) && other_job.name != job.name { + dependencies.insert(other_job.name.clone()); + } + } + } + } + + // Implicit dependencies from input files + if let Some(ref input_files) = job.input_files { + for input_file in input_files { + for other_job in &spec.jobs { + if let Some(ref output_files) = other_job.output_files { + if output_files.contains(input_file) && other_job.name != job.name { + dependencies.insert(other_job.name.clone()); + } + } + } + } + } + + // Implicit dependencies from input user data + if let Some(ref input_data) = job.input_user_data { + for input_datum in input_data { + for other_job in &spec.jobs { + if let Some(ref output_data) = other_job.output_user_data { + if output_data.contains(input_datum) && other_job.name != job.name { + dependencies.insert(other_job.name.clone()); + } + } + } + } + } + + // Add edges + for dep in &dependencies { + graph + .depends_on + .get_mut(&job.name) + .unwrap() + .insert(dep.clone()); + graph + .depended_by + .get_mut(dep) + .unwrap() + .insert(job.name.clone()); + } + } + + Ok(graph) + } + + /// Get all job names in the graph + pub fn job_names(&self) -> impl Iterator { + self.nodes.keys() + } + + /// Get a job node by name + pub fn get_job(&self, name: &str) -> Option<&JobNode> { + self.nodes.get(name) + } + + /// Get the number of jobs in the graph + pub fn job_count(&self) -> usize { + self.nodes.len() + } + + /// Get total instance count (accounting for parameterized jobs) + pub fn total_instance_count(&self) -> usize { + self.nodes.values().map(|n| n.instance_count).sum() + } + + /// Check if a job has any dependencies + pub fn has_dependencies(&self, job: &str) -> bool { + self.depends_on + .get(job) + .map(|deps| !deps.is_empty()) + .unwrap_or(false) + } + + /// Get the jobs that a job depends on (its blockers) + pub fn dependencies_of(&self, job: &str) -> Option<&HashSet> { + self.depends_on.get(job) + } + + /// Get the jobs that depend on a job (its dependents) + pub fn dependents_of(&self, job: &str) -> Option<&HashSet> { + self.depended_by.get(job) + } + + /// Get root jobs (jobs with no dependencies) + pub fn roots(&self) -> Vec<&str> { + self.nodes + .keys() + .filter(|name| { + self.depends_on + .get(*name) + .map(|deps| deps.is_empty()) + .unwrap_or(true) + }) + .map(|s| s.as_str()) + .collect() + } + + /// Get leaf jobs (jobs that nothing depends on) + pub fn leaves(&self) -> Vec<&str> { + self.nodes + .keys() + .filter(|name| { + self.depended_by + .get(*name) + .map(|deps| deps.is_empty()) + .unwrap_or(true) + }) + .map(|s| s.as_str()) + .collect() + } + + /// Compute topological levels (jobs grouped by dependency depth) + /// + /// Level 0 contains jobs with no dependencies. + /// Level N contains jobs whose dependencies are all in levels < N. + pub fn topological_levels(&mut self) -> Result<&Vec>, Box> { + if self.levels.is_some() { + return Ok(self.levels.as_ref().unwrap()); + } + + let mut levels = Vec::new(); + let mut remaining: HashSet = self.nodes.keys().cloned().collect(); + let mut processed = HashSet::new(); + + while !remaining.is_empty() { + let mut current_level = Vec::new(); + + // Find all jobs whose dependencies are satisfied + for name in &remaining { + let deps = self.depends_on.get(name).unwrap(); + if deps.iter().all(|d| processed.contains(d)) { + current_level.push(name.clone()); + } + } + + if current_level.is_empty() { + return Err("Circular dependency detected in workflow graph".into()); + } + + // Mark these jobs as processed + for job in ¤t_level { + remaining.remove(job); + processed.insert(job.clone()); + } + + levels.push(current_level); + } + + self.levels = Some(levels); + Ok(self.levels.as_ref().unwrap()) + } + + /// Find connected components (independent sub-workflows) + /// + /// Each component can be scheduled independently of others. + pub fn connected_components(&mut self) -> &Vec { + if self.components.is_some() { + return self.components.as_ref().unwrap(); + } + + let mut components = Vec::new(); + let mut visited: HashSet = HashSet::new(); + + for start_job in self.nodes.keys() { + if visited.contains(start_job) { + continue; + } + + // BFS to find all connected jobs (treating graph as undirected) + let mut component_jobs = HashSet::new(); + let mut queue = VecDeque::new(); + queue.push_back(start_job.clone()); + + while let Some(job) = queue.pop_front() { + if visited.contains(&job) { + continue; + } + visited.insert(job.clone()); + component_jobs.insert(job.clone()); + + // Add dependencies (forward edges) + if let Some(deps) = self.depends_on.get(&job) { + for dep in deps { + if !visited.contains(dep) { + queue.push_back(dep.clone()); + } + } + } + + // Add dependents (reverse edges) + if let Some(dependents) = self.depended_by.get(&job) { + for dependent in dependents { + if !visited.contains(dependent) { + queue.push_back(dependent.clone()); + } + } + } + } + + // Find roots and leaves within this component + let roots: Vec = component_jobs + .iter() + .filter(|name| { + self.depends_on + .get(*name) + .map(|deps| deps.iter().all(|d| !component_jobs.contains(d))) + .unwrap_or(true) + }) + .cloned() + .collect(); + + let leaves: Vec = component_jobs + .iter() + .filter(|name| { + self.depended_by + .get(*name) + .map(|deps| deps.iter().all(|d| !component_jobs.contains(d))) + .unwrap_or(true) + }) + .cloned() + .collect(); + + components.push(WorkflowComponent { + jobs: component_jobs, + roots, + leaves, + }); + } + + self.components = Some(components); + self.components.as_ref().unwrap() + } + + /// Extract a sub-graph containing only the specified jobs + pub fn subgraph(&self, job_names: &HashSet) -> Self { + let mut subgraph = Self::new(); + + // Copy relevant nodes + for name in job_names { + if let Some(node) = self.nodes.get(name) { + subgraph.nodes.insert(name.clone(), node.clone()); + subgraph.depends_on.insert(name.clone(), HashSet::new()); + subgraph.depended_by.insert(name.clone(), HashSet::new()); + } + } + + // Copy relevant edges (only those within the subgraph) + for name in job_names { + if let Some(deps) = self.depends_on.get(name) { + for dep in deps { + if job_names.contains(dep) { + subgraph + .depends_on + .get_mut(name) + .unwrap() + .insert(dep.clone()); + subgraph + .depended_by + .get_mut(dep) + .unwrap() + .insert(name.clone()); + } + } + } + } + + subgraph + } + + /// Generate scheduler groups based on (resource_requirements, has_dependencies) + /// + /// Jobs are grouped by their resource requirements and dependency status. + /// This is used for scheduler generation. + pub fn scheduler_groups(&self) -> Vec { + // Group by (resource_requirements, has_dependencies) + let mut groups: HashMap<(String, bool), SchedulerGroup> = HashMap::new(); + + for (name, node) in &self.nodes { + let rr_name = match &node.resource_requirements { + Some(rr) => rr.clone(), + None => continue, // Skip jobs without resource requirements + }; + + let has_deps = self.has_dependencies(name); + let key = (rr_name.clone(), has_deps); + + let group = groups.entry(key).or_insert_with(|| SchedulerGroup { + resource_requirements: rr_name, + has_dependencies: has_deps, + job_count: 0, + job_name_patterns: Vec::new(), + job_names: Vec::new(), + }); + + group.job_count += node.instance_count; + group.job_name_patterns.push(node.name_pattern.clone()); + group.job_names.push(name.clone()); + } + + groups.into_values().collect() + } + + /// Generate scheduler groups per connected component + /// + /// Returns a map of component index to scheduler groups for that component. + pub fn scheduler_groups_by_component( + &mut self, + ) -> Vec<(WorkflowComponent, Vec)> { + let components = self.connected_components().clone(); + let mut result = Vec::new(); + + for component in components { + let subgraph = self.subgraph(&component.jobs); + let groups = subgraph.scheduler_groups(); + result.push((component, groups)); + } + + result + } + + /// Find the critical path (longest path through the graph) + /// + /// Returns job names along the critical path and the total instance count. + pub fn critical_path(&mut self) -> Result<(Vec, usize), Box> { + // Use dynamic programming on topological order + let levels = self.topological_levels()?.clone(); + + // dist[job] = (max distance to reach this job, predecessor) + let mut dist: HashMap)> = HashMap::new(); + + for name in self.nodes.keys() { + dist.insert(name.clone(), (0, None)); + } + + // Process in topological order + for level in &levels { + for job in level { + let node = self.nodes.get(job).unwrap(); + let job_weight = node.instance_count; + + if let Some(dependents) = self.depended_by.get(job) { + for dependent in dependents { + let current_dist = dist.get(job).unwrap().0 + job_weight; + let dependent_dist = dist.get(dependent).unwrap().0; + + if current_dist > dependent_dist { + dist.insert(dependent.clone(), (current_dist, Some(job.clone()))); + } + } + } + } + } + + // Find the job with maximum distance (end of critical path) + let (end_job, (_max_dist, _)) = dist + .iter() + .max_by_key(|(_, (d, _))| d) + .ok_or("Empty graph")?; + + // Backtrack to find the path + let mut path = vec![end_job.clone()]; + let mut current = end_job.clone(); + + while let Some((_, Some(prev))) = dist.get(¤t) { + path.push(prev.clone()); + current = prev.clone(); + } + + path.reverse(); + + // Calculate total instance count along critical path + let total: usize = path + .iter() + .filter_map(|name| self.nodes.get(name)) + .map(|n| n.instance_count) + .sum(); + + Ok((path, total)) + } + + /// Get jobs that become ready when a set of jobs complete + pub fn jobs_unblocked_by(&self, completed_jobs: &HashSet) -> Vec { + let mut unblocked = Vec::new(); + + for (name, deps) in &self.depends_on { + if completed_jobs.contains(name) { + continue; // Already completed + } + + // Check if all dependencies are in completed_jobs + if deps.iter().all(|d| completed_jobs.contains(d)) && !deps.is_empty() { + unblocked.push(name.clone()); + } + } + + unblocked + } + + /// Find actions that should trigger when specific jobs become ready + pub fn matching_actions<'a>( + &self, + jobs_becoming_ready: &[String], + actions: &'a [WorkflowActionSpec], + ) -> Vec<&'a WorkflowActionSpec> { + let mut matching = Vec::new(); + + for action in actions { + if action.trigger_type != "on_jobs_ready" { + continue; + } + + // Check job_name_regexes + if let Some(ref regexes) = action.job_name_regexes { + for regex_str in regexes { + if let Ok(re) = Regex::new(regex_str) { + if jobs_becoming_ready.iter().any(|j| re.is_match(j)) { + matching.push(action); + break; + } + } + } + } + + // Check exact job names + if let Some(ref job_names) = action.jobs { + if jobs_becoming_ready.iter().any(|j| job_names.contains(j)) { + if !matching.contains(&action) { + matching.push(action); + } + } + } + } + + matching + } +} + +impl Default for WorkflowGraph { + fn default() -> Self { + Self::new() + } +} + +/// Count job instances for parameterized jobs +fn count_job_instances(job: &JobSpec) -> usize { + if let Some(params) = &job.parameters { + let mut count = 1usize; + for value in params.values() { + let param_count = parse_parameter_value(value).map(|pv| pv.len()).unwrap_or(1); + count *= param_count; + } + count + } else { + 1 + } +} + +/// Build a regex pattern for matching job instances +fn build_job_name_pattern(name: &str, is_parameterized: bool, instance_count: usize) -> String { + if is_parameterized && instance_count > 1 { + // Convert parameterized job name like "work_{index}" to regex "^work_.*$" + let param_regex = Regex::new(r"\{[^}]+\}").unwrap(); + let pattern = param_regex.replace_all(name, ".*").to_string(); + // Simplify consecutive .* patterns + let pattern = pattern.replace(".*.*", ".*"); + format!("^{}$", pattern) + } else { + // Exact match for non-parameterized jobs + format!("^{}$", regex::escape(name)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + fn create_test_spec() -> WorkflowSpec { + WorkflowSpec { + name: "test_workflow".to_string(), + jobs: vec![ + JobSpec { + name: "preprocess".to_string(), + command: "preprocess.sh".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }, + JobSpec { + name: "work_{i}".to_string(), + command: "work.sh".to_string(), + resource_requirements: Some("medium".to_string()), + depends_on: Some(vec!["preprocess".to_string()]), + parameters: Some({ + let mut m = HashMap::new(); + m.insert("i".to_string(), "1:10".to_string()); + m + }), + ..Default::default() + }, + JobSpec { + name: "postprocess".to_string(), + command: "postprocess.sh".to_string(), + resource_requirements: Some("small".to_string()), + depends_on_regexes: Some(vec!["^work_.*$".to_string()]), + ..Default::default() + }, + ], + ..Default::default() + } + } + + #[test] + fn test_graph_construction() { + let spec = create_test_spec(); + let graph = WorkflowGraph::from_spec(&spec).unwrap(); + + assert_eq!(graph.job_count(), 3); + assert_eq!(graph.total_instance_count(), 12); // 1 + 10 + 1 + } + + #[test] + fn test_roots_and_leaves() { + let spec = create_test_spec(); + let graph = WorkflowGraph::from_spec(&spec).unwrap(); + + let roots = graph.roots(); + assert_eq!(roots.len(), 1); + assert!(roots.contains(&"preprocess")); + + let leaves = graph.leaves(); + assert_eq!(leaves.len(), 1); + assert!(leaves.contains(&"postprocess")); + } + + #[test] + fn test_dependencies() { + let spec = create_test_spec(); + let graph = WorkflowGraph::from_spec(&spec).unwrap(); + + assert!(!graph.has_dependencies("preprocess")); + assert!(graph.has_dependencies("work_{i}")); + assert!(graph.has_dependencies("postprocess")); + + let work_deps = graph.dependencies_of("work_{i}").unwrap(); + assert!(work_deps.contains("preprocess")); + + let post_deps = graph.dependencies_of("postprocess").unwrap(); + assert!(post_deps.contains("work_{i}")); + } + + #[test] + fn test_topological_levels() { + let spec = create_test_spec(); + let mut graph = WorkflowGraph::from_spec(&spec).unwrap(); + + let levels = graph.topological_levels().unwrap(); + assert_eq!(levels.len(), 3); + assert!(levels[0].contains(&"preprocess".to_string())); + assert!(levels[1].contains(&"work_{i}".to_string())); + assert!(levels[2].contains(&"postprocess".to_string())); + } + + #[test] + fn test_connected_components_single() { + let spec = create_test_spec(); + let mut graph = WorkflowGraph::from_spec(&spec).unwrap(); + + let components = graph.connected_components(); + assert_eq!(components.len(), 1); + assert_eq!(components[0].jobs.len(), 3); + } + + #[test] + fn test_connected_components_multiple() { + // Create spec with two independent pipelines + let spec = WorkflowSpec { + name: "multi_pipeline".to_string(), + jobs: vec![ + // Pipeline A + JobSpec { + name: "a_start".to_string(), + command: "a.sh".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }, + JobSpec { + name: "a_end".to_string(), + command: "a.sh".to_string(), + resource_requirements: Some("small".to_string()), + depends_on: Some(vec!["a_start".to_string()]), + ..Default::default() + }, + // Pipeline B (independent) + JobSpec { + name: "b_start".to_string(), + command: "b.sh".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }, + JobSpec { + name: "b_end".to_string(), + command: "b.sh".to_string(), + resource_requirements: Some("small".to_string()), + depends_on: Some(vec!["b_start".to_string()]), + ..Default::default() + }, + ], + ..Default::default() + }; + + let mut graph = WorkflowGraph::from_spec(&spec).unwrap(); + let components = graph.connected_components(); + + assert_eq!(components.len(), 2); + + // Each component should have 2 jobs + for component in components { + assert_eq!(component.jobs.len(), 2); + assert_eq!(component.roots.len(), 1); + assert_eq!(component.leaves.len(), 1); + } + } + + #[test] + fn test_scheduler_groups() { + let spec = create_test_spec(); + let graph = WorkflowGraph::from_spec(&spec).unwrap(); + + let groups = graph.scheduler_groups(); + + // Should have 4 groups: + // (small, no_deps), (medium, has_deps), (small, has_deps) + // Wait, preprocess is small/no_deps, work is medium/has_deps, postprocess is small/has_deps + assert_eq!(groups.len(), 3); + + // Find the work group + let work_group = groups + .iter() + .find(|g| g.resource_requirements == "medium") + .unwrap(); + assert_eq!(work_group.job_count, 10); // Parameterized + assert!(work_group.has_dependencies); + } + + #[test] + fn test_subgraph() { + let spec = create_test_spec(); + let graph = WorkflowGraph::from_spec(&spec).unwrap(); + + let mut subset = HashSet::new(); + subset.insert("preprocess".to_string()); + subset.insert("work_{i}".to_string()); + + let subgraph = graph.subgraph(&subset); + assert_eq!(subgraph.job_count(), 2); + + // work_{i} should still depend on preprocess + assert!(subgraph.has_dependencies("work_{i}")); + assert!(!subgraph.has_dependencies("preprocess")); + } + + #[test] + fn test_jobs_unblocked_by() { + let spec = create_test_spec(); + let graph = WorkflowGraph::from_spec(&spec).unwrap(); + + // When preprocess completes, work_{i} becomes ready + let mut completed = HashSet::new(); + completed.insert("preprocess".to_string()); + + let unblocked = graph.jobs_unblocked_by(&completed); + assert_eq!(unblocked.len(), 1); + assert!(unblocked.contains(&"work_{i}".to_string())); + + // When work_{i} also completes, postprocess becomes ready + completed.insert("work_{i}".to_string()); + let unblocked = graph.jobs_unblocked_by(&completed); + assert_eq!(unblocked.len(), 1); + assert!(unblocked.contains(&"postprocess".to_string())); + } +} diff --git a/src/client/workflow_spec.rs b/src/client/workflow_spec.rs index 5c290b17..de3d2a55 100644 --- a/src/client/workflow_spec.rs +++ b/src/client/workflow_spec.rs @@ -2,12 +2,13 @@ use crate::client::apis::{configuration::Configuration, default_api}; use crate::client::parameter_expansion::{ ParameterValue, cartesian_product, parse_parameter_value, substitute_parameters, zip_parameters, }; +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::path::Path; + use crate::models; use regex::Regex; use serde::{Deserialize, Serialize}; -use std::collections::HashMap; -use std::fs; -use std::path::Path; /// Result of validating a workflow specification (dry-run) #[derive(Clone, Debug, Serialize, Deserialize)] @@ -552,6 +553,12 @@ impl WorkflowSpec { } } + /// Deserialize a WorkflowSpec from a serde_json::Value + /// This is the common conversion point for all file formats + pub fn from_json_value(value: serde_json::Value) -> Result> { + Ok(serde_json::from_value(value)?) + } + /// Expand all parameterized jobs and files in this workflow spec /// This modifies the spec in-place, replacing parameterized specs with their expanded versions /// @@ -851,8 +858,6 @@ impl WorkflowSpec { /// # Returns /// A `ValidationResult` containing validation status and summary pub fn validate_spec>(path: P) -> ValidationResult { - use std::collections::{HashMap, HashSet}; - let mut errors = Vec::new(); let mut warnings = Vec::new(); @@ -1870,8 +1875,6 @@ impl WorkflowSpec { jobs: &'a [JobSpec], dependencies: &HashMap>, ) -> Result>, Box> { - use std::collections::HashSet; - let mut levels = Vec::new(); let mut remaining: HashSet = jobs.iter().map(|j| j.name.clone()).collect(); let mut processed = HashSet::new(); @@ -2141,24 +2144,35 @@ impl WorkflowSpec { Ok((job_name_to_id, created_jobs)) } - /// Parse parameters from a KDL node's children - /// Expects a structure like: - /// ```kdl - /// parameters { - /// i "1:100" - /// lr "[0.001,0.01,0.1]" - /// } - /// ``` - /// Returns a HashMap of parameter name -> value string + /// Convert a byte offset to (line, column) for error reporting + #[cfg(feature = "client")] + fn offset_to_line_col(content: &str, offset: usize) -> (usize, usize) { + let mut line = 1; + let mut col = 1; + for (i, ch) in content.char_indices() { + if i >= offset { + break; + } + if ch == '\n' { + line += 1; + col = 1; + } else { + col += 1; + } + } + (line, col) + } + + /// Convert a KDL parameters block to a JSON object #[cfg(feature = "client")] - fn parse_kdl_parameters( + fn kdl_parameters_to_json( node: &KdlNode, - ) -> Result>, Box> { + ) -> Result, Box> { let Some(children) = node.children() else { return Ok(None); }; - let mut params = HashMap::new(); + let mut params = serde_json::Map::new(); for child in children.nodes() { let param_name = child.name().value().to_string(); let param_value = child @@ -2167,153 +2181,19 @@ impl WorkflowSpec { .and_then(|e| e.value().as_string()) .ok_or_else(|| format!("Parameter '{}' must have a string value", param_name))? .to_string(); - params.insert(param_name, param_value); + params.insert(param_name, serde_json::Value::String(param_value)); } if params.is_empty() { Ok(None) } else { - Ok(Some(params)) + Ok(Some(serde_json::Value::Object(params))) } } - /// Parse a WorkflowSpec from a KDL string + /// Convert a KDL job node to a JSON object #[cfg(feature = "client")] - fn from_kdl_str(content: &str) -> Result> { - let doc: KdlDocument = content - .parse() - .map_err(|e| format!("Failed to parse KDL document: {}", e))?; - - let mut spec = WorkflowSpec::default(); - let mut jobs = Vec::new(); - let mut files = Vec::new(); - let mut user_data = Vec::new(); - let mut resource_requirements = Vec::new(); - let mut slurm_schedulers = Vec::new(); - let mut actions = Vec::new(); - - for node in doc.nodes() { - match node.name().value() { - "name" => { - spec.name = node - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("name must have a string value")? - .to_string(); - } - "user" => { - spec.user = Some( - node.entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("user must have a string value")? - .to_string(), - ); - } - "description" => { - spec.description = Some( - node.entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("description must have a string value")? - .to_string(), - ); - } - "compute_node_expiration_buffer_seconds" => { - spec.compute_node_expiration_buffer_seconds = node - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()); - } - "compute_node_wait_for_new_jobs_seconds" => { - spec.compute_node_wait_for_new_jobs_seconds = node - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()); - } - "compute_node_ignore_workflow_completion" => { - spec.compute_node_ignore_workflow_completion = - node.entries().first().and_then(|e| e.value().as_bool()); - } - "compute_node_wait_for_healthy_database_minutes" => { - spec.compute_node_wait_for_healthy_database_minutes = node - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()); - } - "jobs_sort_method" => { - if let Some(value_str) = - node.entries().first().and_then(|e| e.value().as_string()) - { - spec.jobs_sort_method = match value_str { - "gpus_runtime_memory" => { - Some(models::ClaimJobsSortMethod::GpusRuntimeMemory) - } - "gpus_memory_runtime" => { - Some(models::ClaimJobsSortMethod::GpusMemoryRuntime) - } - "none" => Some(models::ClaimJobsSortMethod::None), - _ => { - return Err( - format!("Invalid jobs_sort_method: {}", value_str).into() - ); - } - }; - } - } - "parameters" => { - spec.parameters = Self::parse_kdl_parameters(node)?; - } - "job" => { - let job_spec = Self::parse_kdl_job(node)?; - jobs.push(job_spec); - } - "file" => { - let file_spec = Self::parse_kdl_file(node)?; - files.push(file_spec); - } - "user_data" => { - let user_data_spec = Self::parse_kdl_user_data(node)?; - user_data.push(user_data_spec); - } - "resource_requirements" => { - let resource_req_spec = Self::parse_kdl_resource_requirements(node)?; - resource_requirements.push(resource_req_spec); - } - "slurm_scheduler" => { - let scheduler_spec = Self::parse_kdl_slurm_scheduler(node)?; - slurm_schedulers.push(scheduler_spec); - } - "action" => { - let action_spec = Self::parse_kdl_action(node)?; - actions.push(action_spec); - } - "resource_monitor" => { - let monitor_config = Self::parse_kdl_resource_monitor(node)?; - spec.resource_monitor = Some(monitor_config); - } - _ => { - // Ignore unknown nodes - } - } - } - - spec.jobs = jobs; - spec.files = Some(files).filter(|v| !v.is_empty()); - spec.user_data = Some(user_data).filter(|v| !v.is_empty()); - spec.resource_requirements = Some(resource_requirements).filter(|v| !v.is_empty()); - spec.slurm_schedulers = Some(slurm_schedulers).filter(|v| !v.is_empty()); - spec.actions = Some(actions).filter(|v| !v.is_empty()); - - Ok(spec) - } - - #[cfg(feature = "client")] - fn parse_kdl_job(node: &KdlNode) -> Result> { + fn kdl_job_to_json(node: &KdlNode) -> Result> { let name = node .entries() .first() @@ -2321,140 +2201,137 @@ impl WorkflowSpec { .ok_or("job must have a name")? .to_string(); - let mut job_spec = JobSpec { - name, - ..Default::default() - }; + let mut obj = serde_json::Map::new(); + obj.insert("name".to_string(), serde_json::Value::String(name)); + + // Collect array fields + let mut depends_on: Vec = Vec::new(); + let mut depends_on_regexes: Vec = Vec::new(); + let mut input_files: Vec = Vec::new(); + let mut output_files: Vec = Vec::new(); + let mut input_user_data: Vec = Vec::new(); + let mut output_user_data: Vec = Vec::new(); if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { "command" => { - job_spec.command = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("command must have a string value")? - .to_string(); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "command".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "invocation_script" => { - job_spec.invocation_script = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "invocation_script".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "cancel_on_blocking_job_failure" => { - job_spec.cancel_on_blocking_job_failure = - child.entries().first().and_then(|e| e.value().as_bool()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert( + "cancel_on_blocking_job_failure".to_string(), + serde_json::Value::Bool(v), + ); + } } "supports_termination" => { - job_spec.supports_termination = - child.entries().first().and_then(|e| e.value().as_bool()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert( + "supports_termination".to_string(), + serde_json::Value::Bool(v), + ); + } } "resource_requirements" => { - job_spec.resource_requirements = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "resource_requirements".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } - "depends_on_job" => { - if job_spec.depends_on.is_none() { - job_spec.depends_on = Some(Vec::new()); + "depends_on" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + depends_on.push(serde_json::Value::String(v.to_string())); } - if let Some(job_name) = - child.entries().first().and_then(|e| e.value().as_string()) + } + "depends_on_regexes" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - job_spec - .depends_on - .as_mut() - .unwrap() - .push(job_name.to_string()); + depends_on_regexes.push(serde_json::Value::String(v.to_string())); } } "input_file" => { - if job_spec.input_files.is_none() { - job_spec.input_files = Some(Vec::new()); - } - if let Some(file_name) = - child.entries().first().and_then(|e| e.value().as_string()) + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - job_spec - .input_files - .as_mut() - .unwrap() - .push(file_name.to_string()); + input_files.push(serde_json::Value::String(v.to_string())); } } "output_file" => { - if job_spec.output_files.is_none() { - job_spec.output_files = Some(Vec::new()); - } - if let Some(file_name) = - child.entries().first().and_then(|e| e.value().as_string()) + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - job_spec - .output_files - .as_mut() - .unwrap() - .push(file_name.to_string()); + output_files.push(serde_json::Value::String(v.to_string())); } } "input_user_data" => { - if job_spec.input_user_data.is_none() { - job_spec.input_user_data = Some(Vec::new()); - } - if let Some(data_name) = - child.entries().first().and_then(|e| e.value().as_string()) + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - job_spec - .input_user_data - .as_mut() - .unwrap() - .push(data_name.to_string()); + input_user_data.push(serde_json::Value::String(v.to_string())); } } "output_user_data" => { - if job_spec.output_user_data.is_none() { - job_spec.output_user_data = Some(Vec::new()); - } - if let Some(data_name) = - child.entries().first().and_then(|e| e.value().as_string()) + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - job_spec - .output_user_data - .as_mut() - .unwrap() - .push(data_name.to_string()); + output_user_data.push(serde_json::Value::String(v.to_string())); } } "scheduler" => { - job_spec.scheduler = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "scheduler".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "parameters" => { - job_spec.parameters = Self::parse_kdl_parameters(child)?; + if let Some(params) = Self::kdl_parameters_to_json(child)? { + obj.insert("parameters".to_string(), params); + } } "parameter_mode" => { - job_spec.parameter_mode = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "parameter_mode".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "use_parameters" => { - // Parse use_parameters as multiple string arguments: use_parameters "lr" "batch_size" - let param_names: Vec = child + let param_names: Vec = child .entries() .iter() - .filter_map(|e| e.value().as_string().map(|s| s.to_string())) + .filter_map(|e| { + e.value() + .as_string() + .map(|s| serde_json::Value::String(s.to_string())) + }) .collect(); if !param_names.is_empty() { - job_spec.use_parameters = Some(param_names); + obj.insert( + "use_parameters".to_string(), + serde_json::Value::Array(param_names), + ); } } _ => {} @@ -2462,11 +2339,50 @@ impl WorkflowSpec { } } - Ok(job_spec) + // Add collected arrays if non-empty + if !depends_on.is_empty() { + obj.insert( + "depends_on".to_string(), + serde_json::Value::Array(depends_on), + ); + } + if !depends_on_regexes.is_empty() { + obj.insert( + "depends_on_regexes".to_string(), + serde_json::Value::Array(depends_on_regexes), + ); + } + if !input_files.is_empty() { + obj.insert( + "input_files".to_string(), + serde_json::Value::Array(input_files), + ); + } + if !output_files.is_empty() { + obj.insert( + "output_files".to_string(), + serde_json::Value::Array(output_files), + ); + } + if !input_user_data.is_empty() { + obj.insert( + "input_user_data".to_string(), + serde_json::Value::Array(input_user_data), + ); + } + if !output_user_data.is_empty() { + obj.insert( + "output_user_data".to_string(), + serde_json::Value::Array(output_user_data), + ); + } + + Ok(serde_json::Value::Object(obj)) } + /// Convert a KDL file node to a JSON object #[cfg(feature = "client")] - fn parse_kdl_file(node: &KdlNode) -> Result> { + fn kdl_file_to_json(node: &KdlNode) -> Result> { let name = node .entries() .first() @@ -2474,47 +2390,59 @@ impl WorkflowSpec { .ok_or("file must have a name")? .to_string(); - // Path can be specified as a property (file "name" path="/path") - // or as a child node for parameterized files - let mut path = node - .get("path") - .and_then(|e| e.as_string()) - .map(|s| s.to_string()); + let mut obj = serde_json::Map::new(); + obj.insert("name".to_string(), serde_json::Value::String(name)); - let mut parameters = None; - let mut parameter_mode = None; - let mut use_parameters = None; + // Path can be specified as a property (file "name" path="/path") + if let Some(path) = node.get("path").and_then(|e| e.as_string()) { + obj.insert( + "path".to_string(), + serde_json::Value::String(path.to_string()), + ); + } - // Check for child nodes (path, parameters, parameter_mode, use_parameters) + // Check for child nodes if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { "path" => { - path = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "path".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "parameters" => { - parameters = Self::parse_kdl_parameters(child)?; + if let Some(params) = Self::kdl_parameters_to_json(child)? { + obj.insert("parameters".to_string(), params); + } } "parameter_mode" => { - parameter_mode = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "parameter_mode".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "use_parameters" => { - // Parse use_parameters as multiple string arguments: use_parameters "lr" "batch_size" - let param_names: Vec = child + let param_names: Vec = child .entries() .iter() - .filter_map(|e| e.value().as_string().map(|s| s.to_string())) + .filter_map(|e| { + e.value() + .as_string() + .map(|s| serde_json::Value::String(s.to_string())) + }) .collect(); if !param_names.is_empty() { - use_parameters = Some(param_names); + obj.insert( + "use_parameters".to_string(), + serde_json::Value::Array(param_names), + ); } } _ => {} @@ -2522,33 +2450,38 @@ impl WorkflowSpec { } } - let path = path.ok_or("file must have a path property")?; + // Validate required path field + if !obj.contains_key("path") { + return Err("file must have a path property".into()); + } - Ok(FileSpec { - name, - path, - parameters, - parameter_mode, - use_parameters, - }) + Ok(serde_json::Value::Object(obj)) } + /// Convert a KDL user_data node to a JSON object #[cfg(feature = "client")] - fn parse_kdl_user_data(node: &KdlNode) -> Result> { - let name = node - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + fn kdl_user_data_to_json( + node: &KdlNode, + ) -> Result> { + let mut obj = serde_json::Map::new(); + + // Name is optional + if let Some(name) = node.entries().first().and_then(|e| e.value().as_string()) { + obj.insert( + "name".to_string(), + serde_json::Value::String(name.to_string()), + ); + } - let mut is_ephemeral = None; let mut data_str: Option<&str> = None; if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { "is_ephemeral" => { - is_ephemeral = child.entries().first().and_then(|e| e.value().as_bool()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert("is_ephemeral".to_string(), serde_json::Value::Bool(v)); + } } "data" => { data_str = child.entries().first().and_then(|e| e.value().as_string()); @@ -2558,20 +2491,19 @@ impl WorkflowSpec { } } + // Parse data string as JSON let data_str = data_str.ok_or("user_data must have a data property")?; let data: serde_json::Value = serde_json::from_str(data_str)?; + obj.insert("data".to_string(), data); - Ok(UserDataSpec { - is_ephemeral, - name, - data: Some(data), - }) + Ok(serde_json::Value::Object(obj)) } + /// Convert a KDL resource_requirements node to a JSON object #[cfg(feature = "client")] - fn parse_kdl_resource_requirements( + fn kdl_resource_requirements_to_json( node: &KdlNode, - ) -> Result> { + ) -> Result> { let name = node .entries() .first() @@ -2579,359 +2511,919 @@ impl WorkflowSpec { .ok_or("resource_requirements must have a name")? .to_string(); - let mut spec = ResourceRequirementsSpec { - name, - num_cpus: 0, - num_gpus: 0, - num_nodes: 0, - memory: String::new(), - runtime: String::new(), - }; + let mut obj = serde_json::Map::new(); + obj.insert("name".to_string(), serde_json::Value::String(name)); if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { "num_cpus" => { - spec.num_cpus = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()) - .ok_or("num_cpus must have a valid integer value")?; + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "num_cpus".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } } "num_gpus" => { - spec.num_gpus = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()) - .ok_or("num_gpus must have a valid integer value")?; + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "num_gpus".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } } "num_nodes" => { - spec.num_nodes = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()) - .ok_or("num_nodes must have a valid integer value")?; + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "num_nodes".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } } "memory" => { - spec.memory = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("memory must have a string value")? - .to_string(); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "memory".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "runtime" => { - spec.runtime = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("runtime must have a string value")? - .to_string(); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "runtime".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } _ => {} } } } - Ok(spec) + Ok(serde_json::Value::Object(obj)) } + /// Convert a KDL slurm_scheduler node to a JSON object #[cfg(feature = "client")] - fn parse_kdl_slurm_scheduler( + fn kdl_slurm_scheduler_to_json( node: &KdlNode, - ) -> Result> { - let name = node - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + ) -> Result> { + let mut obj = serde_json::Map::new(); - let mut spec = SlurmSchedulerSpec { - name, - account: String::new(), - gres: None, - mem: None, - nodes: 0, - ntasks_per_node: None, - partition: None, - qos: None, - tmp: None, - walltime: String::new(), - extra: None, - }; + // Name is optional + if let Some(name) = node.entries().first().and_then(|e| e.value().as_string()) { + obj.insert( + "name".to_string(), + serde_json::Value::String(name.to_string()), + ); + } if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { "account" => { - spec.account = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("account must have a string value")? - .to_string(); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "account".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "gres" => { - spec.gres = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "gres".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "mem" => { - spec.mem = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert("mem".to_string(), serde_json::Value::String(v.to_string())); + } } "nodes" => { - spec.nodes = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()) - .ok_or("nodes must have a valid integer value")?; + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "nodes".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } } "ntasks_per_node" => { - spec.ntasks_per_node = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()); + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "ntasks_per_node".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } } "partition" => { - spec.partition = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "partition".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "qos" => { - spec.qos = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert("qos".to_string(), serde_json::Value::String(v.to_string())); + } } "tmp" => { - spec.tmp = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert("tmp".to_string(), serde_json::Value::String(v.to_string())); + } } "walltime" => { - spec.walltime = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("walltime must have a string value")? - .to_string(); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "walltime".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } "extra" => { - spec.extra = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "extra".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } _ => {} } } } - // Validate required fields - if spec.walltime.is_empty() { - return Err("walltime is required for Slurm scheduler".into()); - } - - Ok(spec) + Ok(serde_json::Value::Object(obj)) } + /// Convert a KDL action node to a JSON object #[cfg(feature = "client")] - fn parse_kdl_resource_monitor( - node: &KdlNode, - ) -> Result> - { - let mut config = crate::client::resource_monitor::ResourceMonitorConfig::default(); + fn kdl_action_to_json(node: &KdlNode) -> Result> { + let mut obj = serde_json::Map::new(); + + // Collect array fields + let mut job_names: Vec = Vec::new(); + let mut job_name_regexes: Vec = Vec::new(); + let mut commands: Vec = Vec::new(); if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { - "enabled" => { - config.enabled = child - .entries() - .first() - .and_then(|e| e.value().as_bool()) - .ok_or("enabled must have a boolean value")?; + "trigger_type" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "trigger_type".to_string(), + serde_json::Value::String(v.to_string()), + ); + } } - "granularity" => { - if let Some(value_str) = - child.entries().first().and_then(|e| e.value().as_string()) + "action_type" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - config.granularity = match value_str { - "summary" => { - crate::client::resource_monitor::MonitorGranularity::Summary - } - "time_series" => { - crate::client::resource_monitor::MonitorGranularity::TimeSeries - } - _ => { - return Err( - format!("Invalid granularity: {}", value_str).into() - ); - } - }; + obj.insert( + "action_type".to_string(), + serde_json::Value::String(v.to_string()), + ); } } - "sample_interval_seconds" => { - config.sample_interval_seconds = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()) - .ok_or("sample_interval_seconds must have a valid integer value")?; + "job" => { + // Collect individual job entries: job "prep_a" / job "prep_b" + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + job_names.push(serde_json::Value::String(v.to_string())); + } } - "generate_plots" => { - config.generate_plots = child - .entries() - .first() - .and_then(|e| e.value().as_bool()) - .ok_or("generate_plots must have a boolean value")?; + "jobs" => { + // Parse jobs as multiple string arguments: jobs "job1" "job2" "job3" + for e in child.entries().iter() { + if let Some(s) = e.value().as_string() { + job_names.push(serde_json::Value::String(s.to_string())); + } + } + } + "job_name_regexes" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + job_name_regexes.push(serde_json::Value::String(v.to_string())); + } + } + "command" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + commands.push(serde_json::Value::String(v.to_string())); + } + } + "scheduler" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "scheduler".to_string(), + serde_json::Value::String(v.to_string()), + ); + } + } + "scheduler_type" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) + { + obj.insert( + "scheduler_type".to_string(), + serde_json::Value::String(v.to_string()), + ); + } + } + "num_allocations" => { + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "num_allocations".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } + } + "start_one_worker_per_node" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert( + "start_one_worker_per_node".to_string(), + serde_json::Value::Bool(v), + ); + } + } + "max_parallel_jobs" => { + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) + { + obj.insert( + "max_parallel_jobs".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } + } + "persistent" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert("persistent".to_string(), serde_json::Value::Bool(v)); + } } _ => {} } } } - Ok(config) + // Add collected arrays if non-empty + if !job_names.is_empty() { + obj.insert("jobs".to_string(), serde_json::Value::Array(job_names)); + } + if !job_name_regexes.is_empty() { + obj.insert( + "job_name_regexes".to_string(), + serde_json::Value::Array(job_name_regexes), + ); + } + if !commands.is_empty() { + obj.insert("commands".to_string(), serde_json::Value::Array(commands)); + } + + Ok(serde_json::Value::Object(obj)) } + /// Convert a KDL resource_monitor node to a JSON object #[cfg(feature = "client")] - fn parse_kdl_action(node: &KdlNode) -> Result> { - let mut spec = WorkflowActionSpec { - trigger_type: String::new(), - action_type: String::new(), - jobs: None, - job_name_regexes: None, - commands: None, - scheduler: None, - scheduler_type: None, - num_allocations: None, - start_one_worker_per_node: None, - max_parallel_jobs: None, - persistent: None, - }; + fn kdl_resource_monitor_to_json( + node: &KdlNode, + ) -> Result> { + let mut obj = serde_json::Map::new(); if let Some(children) = node.children() { for child in children.nodes() { match child.name().value() { - "trigger_type" => { - spec.trigger_type = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("trigger_type must have a string value")? - .to_string(); - } - "action_type" => { - spec.action_type = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .ok_or("action_type must have a string value")? - .to_string(); - } - "jobs" => { - // Parse jobs as multiple string arguments: jobs "job1" "job2" "job3" - let job_names: Vec = child - .entries() - .iter() - .filter_map(|e| e.value().as_string().map(|s| s.to_string())) - .collect(); - if !job_names.is_empty() { - spec.jobs = Some(job_names); + "enabled" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert("enabled".to_string(), serde_json::Value::Bool(v)); } } - "job_name_regexes" => { - if spec.job_name_regexes.is_none() { - spec.job_name_regexes = Some(Vec::new()); - } - if let Some(regex) = - child.entries().first().and_then(|e| e.value().as_string()) + "granularity" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_string()) { - spec.job_name_regexes - .as_mut() - .unwrap() - .push(regex.to_string()); + obj.insert( + "granularity".to_string(), + serde_json::Value::String(v.to_string()), + ); } } - "command" => { - if spec.commands.is_none() { - spec.commands = Some(Vec::new()); - } - if let Some(command) = - child.entries().first().and_then(|e| e.value().as_string()) + "sample_interval_seconds" => { + if let Some(v) = + child.entries().first().and_then(|e| e.value().as_integer()) { - spec.commands.as_mut().unwrap().push(command.to_string()); + obj.insert( + "sample_interval_seconds".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); } } - "scheduler" => { - spec.scheduler = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + "generate_plots" => { + if let Some(v) = child.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert("generate_plots".to_string(), serde_json::Value::Bool(v)); + } } - "scheduler_type" => { - spec.scheduler_type = child - .entries() - .first() - .and_then(|e| e.value().as_string()) - .map(|s| s.to_string()); + _ => {} + } + } + } + + Ok(serde_json::Value::Object(obj)) + } + + /// Convert a KDL document string to a serde_json::Value + /// This is the intermediate representation used by all file formats + #[cfg(feature = "client")] + fn kdl_to_json_value(content: &str) -> Result> { + let doc: KdlDocument = content.parse().map_err(|e: kdl::KdlError| { + // Extract detailed diagnostic information from KDL parse errors + let mut error_msg = String::from("Failed to parse KDL document:\n"); + for diag in e.diagnostics.iter() { + let offset = diag.span.offset(); + let (line, col) = Self::offset_to_line_col(content, offset); + + if let Some(msg) = &diag.message { + error_msg.push_str(&format!(" Line {}, column {}: {}", line, col, msg)); + } else { + error_msg.push_str(&format!(" Line {}, column {}: syntax error", line, col)); + } + if let Some(label) = &diag.label { + error_msg.push_str(&format!(" ({})", label)); + } + error_msg.push('\n'); + if let Some(help) = &diag.help { + error_msg.push_str(&format!(" Help: {}\n", help)); + } + } + // Show the problematic line if we can + if let Some(first_diag) = e.diagnostics.first() { + let offset = first_diag.span.offset(); + let (line_num, col) = Self::offset_to_line_col(content, offset); + if let Some(line_content) = content.lines().nth(line_num.saturating_sub(1)) { + error_msg.push_str(&format!("\n {} | {}\n", line_num, line_content)); + error_msg.push_str(&format!( + " {} | {}^\n", + " ".repeat(line_num.to_string().len()), + " ".repeat(col.saturating_sub(1)) + )); + } + } + error_msg + })?; + + let mut obj = serde_json::Map::new(); + let mut jobs: Vec = Vec::new(); + let mut files: Vec = Vec::new(); + let mut user_data: Vec = Vec::new(); + let mut resource_requirements: Vec = Vec::new(); + let mut slurm_schedulers: Vec = Vec::new(); + let mut actions: Vec = Vec::new(); + + for node in doc.nodes() { + match node.name().value() { + "name" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_string()) { + obj.insert("name".to_string(), serde_json::Value::String(v.to_string())); } - "num_allocations" => { - spec.num_allocations = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()); + } + "user" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_string()) { + obj.insert("user".to_string(), serde_json::Value::String(v.to_string())); } - "start_one_worker_per_node" => { - spec.start_one_worker_per_node = - child.entries().first().and_then(|e| e.value().as_bool()); + } + "description" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_string()) { + obj.insert( + "description".to_string(), + serde_json::Value::String(v.to_string()), + ); } - "max_parallel_jobs" => { - spec.max_parallel_jobs = child - .entries() - .first() - .and_then(|e| e.value().as_integer()) - .and_then(|i| i.try_into().ok()); + } + "compute_node_expiration_buffer_seconds" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_integer()) { + obj.insert( + "compute_node_expiration_buffer_seconds".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); } - "persistent" => { - spec.persistent = child.entries().first().and_then(|e| e.value().as_bool()); + } + "compute_node_wait_for_new_jobs_seconds" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_integer()) { + obj.insert( + "compute_node_wait_for_new_jobs_seconds".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); } - _ => {} + } + "compute_node_ignore_workflow_completion" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_bool()) { + obj.insert( + "compute_node_ignore_workflow_completion".to_string(), + serde_json::Value::Bool(v), + ); + } + } + "compute_node_wait_for_healthy_database_minutes" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_integer()) { + obj.insert( + "compute_node_wait_for_healthy_database_minutes".to_string(), + serde_json::Value::Number(serde_json::Number::from(v as i64)), + ); + } + } + "jobs_sort_method" => { + if let Some(v) = node.entries().first().and_then(|e| e.value().as_string()) { + obj.insert( + "jobs_sort_method".to_string(), + serde_json::Value::String(v.to_string()), + ); + } + } + "parameters" => { + if let Some(params) = Self::kdl_parameters_to_json(node)? { + obj.insert("parameters".to_string(), params); + } + } + "job" => { + jobs.push(Self::kdl_job_to_json(node)?); + } + "file" => { + files.push(Self::kdl_file_to_json(node)?); + } + "user_data" => { + user_data.push(Self::kdl_user_data_to_json(node)?); + } + "resource_requirements" => { + resource_requirements.push(Self::kdl_resource_requirements_to_json(node)?); + } + "slurm_scheduler" => { + slurm_schedulers.push(Self::kdl_slurm_scheduler_to_json(node)?); + } + "action" => { + actions.push(Self::kdl_action_to_json(node)?); + } + "resource_monitor" => { + obj.insert( + "resource_monitor".to_string(), + Self::kdl_resource_monitor_to_json(node)?, + ); + } + _ => { + // Ignore unknown nodes } } } - // Validate required fields - if spec.trigger_type.is_empty() { - return Err("trigger_type is required for action".into()); + // Add collected arrays - jobs is required (can be empty), others are optional + obj.insert("jobs".to_string(), serde_json::Value::Array(jobs)); + if !files.is_empty() { + obj.insert("files".to_string(), serde_json::Value::Array(files)); + } + if !user_data.is_empty() { + obj.insert("user_data".to_string(), serde_json::Value::Array(user_data)); + } + if !resource_requirements.is_empty() { + obj.insert( + "resource_requirements".to_string(), + serde_json::Value::Array(resource_requirements), + ); + } + if !slurm_schedulers.is_empty() { + obj.insert( + "slurm_schedulers".to_string(), + serde_json::Value::Array(slurm_schedulers), + ); } - if spec.action_type.is_empty() { - return Err("action_type is required for action".into()); + if !actions.is_empty() { + obj.insert("actions".to_string(), serde_json::Value::Array(actions)); } - Ok(spec) + Ok(serde_json::Value::Object(obj)) + } + + /// Serialize WorkflowSpec to KDL format + #[cfg(feature = "client")] + pub fn to_kdl_str(&self) -> String { + let mut lines = Vec::new(); + + // Helper to escape strings for KDL + fn kdl_escape(s: &str) -> String { + // Use raw strings for multi-line or strings with special chars + if s.contains('\n') || s.contains('"') || s.contains('\\') { + // Count the number of # needed for raw string + let mut hashes = 0; + loop { + let delimiter: String = std::iter::repeat('#').take(hashes).collect(); + if !s.contains(&format!("\"{}", delimiter)) { + break; + } + hashes += 1; + } + let delimiter: String = std::iter::repeat('#').take(hashes).collect(); + // KDL raw string format: r#"..."# where # count can vary + format!("r{}\"{}\"{}", delimiter, s, delimiter) + } else { + format!("\"{}\"", s) + } + } + + // Top-level fields + lines.push(format!("name {}", kdl_escape(&self.name))); + if let Some(ref user) = self.user { + lines.push(format!("user {}", kdl_escape(user))); + } + if let Some(ref desc) = self.description { + lines.push(format!("description {}", kdl_escape(desc))); + } + if let Some(val) = self.compute_node_expiration_buffer_seconds { + lines.push(format!("compute_node_expiration_buffer_seconds {}", val)); + } + if let Some(val) = self.compute_node_wait_for_new_jobs_seconds { + lines.push(format!("compute_node_wait_for_new_jobs_seconds {}", val)); + } + if let Some(val) = self.compute_node_ignore_workflow_completion { + lines.push(format!( + "compute_node_ignore_workflow_completion {}", + if val { "#true" } else { "#false" } + )); + } + if let Some(val) = self.compute_node_wait_for_healthy_database_minutes { + lines.push(format!( + "compute_node_wait_for_healthy_database_minutes {}", + val + )); + } + if let Some(ref method) = self.jobs_sort_method { + let method_str = match method { + models::ClaimJobsSortMethod::GpusRuntimeMemory => "gpus_runtime_memory", + models::ClaimJobsSortMethod::GpusMemoryRuntime => "gpus_memory_runtime", + models::ClaimJobsSortMethod::None => "none", + }; + lines.push(format!("jobs_sort_method \"{}\"", method_str)); + } + + // Parameters + if let Some(ref params) = self.parameters { + if !params.is_empty() { + lines.push("parameters {".to_string()); + for (key, value) in params { + lines.push(format!(" {} {}", key, kdl_escape(value))); + } + lines.push("}".to_string()); + } + } + + lines.push(String::new()); // Empty line for readability + + // Files + if let Some(ref files) = self.files { + for file in files { + Self::file_spec_to_kdl(&mut lines, file, &kdl_escape); + } + if !files.is_empty() { + lines.push(String::new()); + } + } + + // User data + if let Some(ref user_data) = self.user_data { + for ud in user_data { + Self::user_data_spec_to_kdl(&mut lines, ud, &kdl_escape); + } + if !user_data.is_empty() { + lines.push(String::new()); + } + } + + // Resource requirements + if let Some(ref reqs) = self.resource_requirements { + for req in reqs { + Self::resource_requirements_spec_to_kdl(&mut lines, req, &kdl_escape); + } + if !reqs.is_empty() { + lines.push(String::new()); + } + } + + // Resource monitor + if let Some(ref monitor) = self.resource_monitor { + lines.push("resource_monitor {".to_string()); + lines.push(format!( + " enabled {}", + if monitor.enabled { "#true" } else { "#false" } + )); + let granularity = match monitor.granularity { + crate::client::resource_monitor::MonitorGranularity::Summary => "summary", + crate::client::resource_monitor::MonitorGranularity::TimeSeries => "time_series", + }; + lines.push(format!(" granularity \"{}\"", granularity)); + lines.push(format!( + " sample_interval_seconds {}", + monitor.sample_interval_seconds + )); + lines.push(format!( + " generate_plots {}", + if monitor.generate_plots { + "#true" + } else { + "#false" + } + )); + lines.push("}".to_string()); + lines.push(String::new()); + } + + // Jobs + for job in &self.jobs { + Self::job_spec_to_kdl(&mut lines, job, &kdl_escape); + } + if !self.jobs.is_empty() { + lines.push(String::new()); + } + + // Slurm schedulers (placed after jobs since they may be auto-generated) + if let Some(ref schedulers) = self.slurm_schedulers { + for sched in schedulers { + Self::slurm_scheduler_spec_to_kdl(&mut lines, sched, &kdl_escape); + } + if !schedulers.is_empty() { + lines.push(String::new()); + } + } + + // Actions (placed last since they may be auto-generated) + if let Some(ref actions) = self.actions { + for action in actions { + Self::action_spec_to_kdl(&mut lines, action, &kdl_escape); + } + } + + lines.join("\n") + } + + #[cfg(feature = "client")] + fn file_spec_to_kdl(lines: &mut Vec, file: &FileSpec, escape: &dyn Fn(&str) -> String) { + let has_params = file + .parameters + .as_ref() + .map(|p| !p.is_empty()) + .unwrap_or(false); + let has_mode = file.parameter_mode.is_some(); + let has_use_params = file.use_parameters.is_some(); + + if !has_params && !has_mode && !has_use_params { + // Simple form: file "name" path="value" + lines.push(format!( + "file {} path={}", + escape(&file.name), + escape(&file.path) + )); + } else { + lines.push(format!("file {} {{", escape(&file.name))); + lines.push(format!(" path {}", escape(&file.path))); + if let Some(ref params) = file.parameters { + if !params.is_empty() { + lines.push(" parameters {".to_string()); + for (key, value) in params { + lines.push(format!(" {} {}", key, escape(value))); + } + lines.push(" }".to_string()); + } + } + if let Some(ref mode) = file.parameter_mode { + lines.push(format!(" parameter_mode {}", escape(mode))); + } + if let Some(ref use_params) = file.use_parameters { + for param in use_params { + lines.push(format!(" use_parameter {}", escape(param))); + } + } + lines.push("}".to_string()); + } + } + + #[cfg(feature = "client")] + fn user_data_spec_to_kdl( + lines: &mut Vec, + ud: &UserDataSpec, + escape: &dyn Fn(&str) -> String, + ) { + let name = ud.name.as_deref().unwrap_or("unnamed"); + lines.push(format!("user_data {} {{", escape(name))); + if ud.is_ephemeral.unwrap_or(false) { + lines.push(" is_ephemeral #true".to_string()); + } + if let Some(ref data) = ud.data { + // Serialize JSON value to string + let data_str = serde_json::to_string(data).unwrap_or_default(); + lines.push(format!(" data {}", escape(&data_str))); + } + lines.push("}".to_string()); + } + + #[cfg(feature = "client")] + fn resource_requirements_spec_to_kdl( + lines: &mut Vec, + req: &ResourceRequirementsSpec, + escape: &dyn Fn(&str) -> String, + ) { + lines.push(format!("resource_requirements {} {{", escape(&req.name))); + lines.push(format!(" num_cpus {}", req.num_cpus)); + lines.push(format!(" num_gpus {}", req.num_gpus)); + lines.push(format!(" num_nodes {}", req.num_nodes)); + lines.push(format!(" memory {}", escape(&req.memory))); + lines.push(format!(" runtime {}", escape(&req.runtime))); + lines.push("}".to_string()); + } + + #[cfg(feature = "client")] + fn slurm_scheduler_spec_to_kdl( + lines: &mut Vec, + sched: &SlurmSchedulerSpec, + escape: &dyn Fn(&str) -> String, + ) { + if let Some(ref name) = sched.name { + lines.push(format!("slurm_scheduler {} {{", escape(name))); + } else { + lines.push("slurm_scheduler {".to_string()); + } + lines.push(format!(" account {}", escape(&sched.account))); + if let Some(ref gres) = sched.gres { + lines.push(format!(" gres {}", escape(gres))); + } + if let Some(ref mem) = sched.mem { + lines.push(format!(" mem {}", escape(mem))); + } + lines.push(format!(" nodes {}", sched.nodes)); + if let Some(ntasks) = sched.ntasks_per_node { + lines.push(format!(" ntasks_per_node {}", ntasks)); + } + if let Some(ref partition) = sched.partition { + lines.push(format!(" partition {}", escape(partition))); + } + if let Some(ref qos) = sched.qos { + lines.push(format!(" qos {}", escape(qos))); + } + if let Some(ref tmp) = sched.tmp { + lines.push(format!(" tmp {}", escape(tmp))); + } + lines.push(format!(" walltime {}", escape(&sched.walltime))); + if let Some(ref extra) = sched.extra { + lines.push(format!(" extra {}", escape(extra))); + } + lines.push("}".to_string()); + } + + #[cfg(feature = "client")] + fn action_spec_to_kdl( + lines: &mut Vec, + action: &WorkflowActionSpec, + escape: &dyn Fn(&str) -> String, + ) { + lines.push("action {".to_string()); + lines.push(format!(" trigger_type {}", escape(&action.trigger_type))); + lines.push(format!(" action_type {}", escape(&action.action_type))); + if let Some(ref jobs) = action.jobs { + for job in jobs { + lines.push(format!(" job {}", escape(job))); + } + } + if let Some(ref regexes) = action.job_name_regexes { + for regex in regexes { + lines.push(format!(" job_name_regexes {}", escape(regex))); + } + } + if let Some(ref commands) = action.commands { + for cmd in commands { + lines.push(format!(" command {}", escape(cmd))); + } + } + if let Some(ref scheduler) = action.scheduler { + lines.push(format!(" scheduler {}", escape(scheduler))); + } + if let Some(ref scheduler_type) = action.scheduler_type { + lines.push(format!(" scheduler_type {}", escape(scheduler_type))); + } + if let Some(count) = action.num_allocations { + lines.push(format!(" num_allocations {}", count)); + } + if let Some(val) = action.start_one_worker_per_node { + lines.push(format!( + " start_one_worker_per_node {}", + if val { "#true" } else { "#false" } + )); + } + if let Some(max) = action.max_parallel_jobs { + lines.push(format!(" max_parallel_jobs {}", max)); + } + if let Some(val) = action.persistent { + lines.push(format!( + " persistent {}", + if val { "#true" } else { "#false" } + )); + } + lines.push("}".to_string()); + } + + #[cfg(feature = "client")] + fn job_spec_to_kdl(lines: &mut Vec, job: &JobSpec, escape: &dyn Fn(&str) -> String) { + lines.push(format!("job {} {{", escape(&job.name))); + lines.push(format!(" command {}", escape(&job.command))); + if let Some(ref script) = job.invocation_script { + lines.push(format!(" invocation_script {}", escape(script))); + } + if let Some(val) = job.cancel_on_blocking_job_failure { + lines.push(format!( + " cancel_on_blocking_job_failure {}", + if val { "#true" } else { "#false" } + )); + } + if let Some(val) = job.supports_termination { + lines.push(format!( + " supports_termination {}", + if val { "#true" } else { "#false" } + )); + } + if let Some(ref req) = job.resource_requirements { + lines.push(format!(" resource_requirements {}", escape(req))); + } + if let Some(ref deps) = job.depends_on { + for dep in deps { + lines.push(format!(" depends_on {}", escape(dep))); + } + } + if let Some(ref regexes) = job.depends_on_regexes { + for regex in regexes { + lines.push(format!(" depends_on_regexes {}", escape(regex))); + } + } + if let Some(ref files) = job.input_files { + for file in files { + lines.push(format!(" input_file {}", escape(file))); + } + } + if let Some(ref files) = job.output_files { + for file in files { + lines.push(format!(" output_file {}", escape(file))); + } + } + if let Some(ref ud) = job.input_user_data { + for name in ud { + lines.push(format!(" input_user_data {}", escape(name))); + } + } + if let Some(ref ud) = job.output_user_data { + for name in ud { + lines.push(format!(" output_user_data {}", escape(name))); + } + } + if let Some(ref sched) = job.scheduler { + lines.push(format!(" scheduler {}", escape(sched))); + } + if let Some(ref params) = job.parameters { + if !params.is_empty() { + lines.push(" parameters {".to_string()); + for (key, value) in params { + lines.push(format!(" {} {}", key, escape(value))); + } + lines.push(" }".to_string()); + } + } + lines.push("}".to_string()); } - /// Deserialize a WorkflowSpec from a specification file (JSON, JSON5, or YAML) + /// Deserialize a WorkflowSpec from a specification file (JSON, JSON5, YAML, or KDL) + /// All formats are first converted to serde_json::Value, then to WorkflowSpec, + /// ensuring consistent behavior across all file formats. pub fn from_spec_file>( path: P, ) -> Result> { @@ -2944,24 +3436,26 @@ impl WorkflowSpec { .and_then(|ext| ext.to_str()) .unwrap_or(""); - let workflow_spec: WorkflowSpec = match extension.to_lowercase().as_str() { + // Parse to JSON Value first, then convert to WorkflowSpec + // This ensures consistent behavior across all formats + let json_value: serde_json::Value = match extension.to_lowercase().as_str() { "json" => serde_json::from_str(&file_content)?, "json5" => json5::from_str(&file_content)?, "yaml" | "yml" => serde_yaml::from_str(&file_content)?, #[cfg(feature = "client")] - "kdl" => Self::from_kdl_str(&file_content)?, + "kdl" => Self::kdl_to_json_value(&file_content)?, _ => { // Try to parse as JSON first, then JSON5, then YAML, then KDL - if let Ok(spec) = serde_json::from_str::(&file_content) { - spec - } else if let Ok(spec) = json5::from_str::(&file_content) { - spec - } else if let Ok(spec) = serde_yaml::from_str::(&file_content) { - spec + if let Ok(value) = serde_json::from_str::(&file_content) { + value + } else if let Ok(value) = json5::from_str::(&file_content) { + value + } else if let Ok(value) = serde_yaml::from_str::(&file_content) { + value } else { #[cfg(feature = "client")] { - Self::from_kdl_str(&file_content)? + Self::kdl_to_json_value(&file_content)? } #[cfg(not(feature = "client"))] { @@ -2971,11 +3465,13 @@ impl WorkflowSpec { } }; - Ok(workflow_spec) + Self::from_json_value(json_value) } /// Deserialize a WorkflowSpec from string content with a specified format /// Useful for testing or when content is already loaded + /// All formats are first converted to serde_json::Value, then to WorkflowSpec, + /// ensuring consistent behavior across all file formats. /// /// # Arguments /// * `content` - The workflow spec content as a string @@ -2984,18 +3480,19 @@ impl WorkflowSpec { content: &str, format: &str, ) -> Result> { - let workflow_spec: WorkflowSpec = match format.to_lowercase().as_str() { + // Parse to JSON Value first, then convert to WorkflowSpec + let json_value: serde_json::Value = match format.to_lowercase().as_str() { "json" => serde_json::from_str(content)?, "json5" => json5::from_str(content)?, "yaml" | "yml" => serde_yaml::from_str(content)?, #[cfg(feature = "client")] - "kdl" => Self::from_kdl_str(content)?, + "kdl" => Self::kdl_to_json_value(content)?, #[cfg(not(feature = "client"))] "kdl" => return Err("KDL format requires 'client' feature".into()), _ => return Err(format!("Unknown format: {}", format).into()), }; - Ok(workflow_spec) + Self::from_json_value(json_value) } /// Perform variable substitution on job commands and invocation scripts diff --git a/src/config/client.rs b/src/config/client.rs index e0aa8290..46b79be7 100644 --- a/src/config/client.rs +++ b/src/config/client.rs @@ -1,6 +1,7 @@ //! Client configuration for the torc CLI use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::path::PathBuf; /// Configuration for the torc CLI client @@ -24,6 +25,9 @@ pub struct ClientConfig { /// Slurm scheduler configuration pub slurm: ClientSlurmConfig, + + /// HPC profile configuration + pub hpc: ClientHpcConfig, } impl Default for ClientConfig { @@ -35,6 +39,7 @@ impl Default for ClientConfig { log_level: "info".to_string(), run: ClientRunConfig::default(), slurm: ClientSlurmConfig::default(), + hpc: ClientHpcConfig::default(), } } } @@ -99,6 +104,112 @@ impl Default for ClientSlurmConfig { } } +/// Configuration for HPC profiles +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct ClientHpcConfig { + /// Default account to use for HPC jobs + pub default_account: Option, + + /// Profile overrides - allows customizing built-in profiles + /// Key is the profile name (e.g., "kestrel") + pub profile_overrides: HashMap, + + /// Custom profiles defined by the user + pub custom_profiles: HashMap, +} + +/// Override settings for a built-in HPC profile +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct HpcProfileOverride { + /// Override the default account for this profile + pub default_account: Option, +} + +/// Configuration for a custom HPC profile +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HpcProfileConfig { + /// Display name for the profile + pub display_name: String, + + /// Description of the HPC system + #[serde(default)] + pub description: String, + + /// Detection via environment variable (name=value) + #[serde(default)] + pub detect_env_var: Option, + + /// Detection via hostname pattern (regex) + #[serde(default)] + pub detect_hostname: Option, + + /// Default account for this profile + #[serde(default)] + pub default_account: Option, + + /// Charge factor for CPU jobs + #[serde(default = "default_charge_factor")] + pub charge_factor_cpu: f64, + + /// Charge factor for GPU jobs + #[serde(default = "default_charge_factor_gpu")] + pub charge_factor_gpu: f64, + + /// Partition configurations + #[serde(default)] + pub partitions: Vec, +} + +fn default_charge_factor() -> f64 { + 1.0 +} + +fn default_charge_factor_gpu() -> f64 { + 10.0 +} + +/// Configuration for an HPC partition +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HpcPartitionConfig { + /// Partition name + pub name: String, + + /// Description + #[serde(default)] + pub description: String, + + /// CPUs per node + pub cpus_per_node: u32, + + /// Memory per node in MB + pub memory_mb: u64, + + /// Maximum wall time in seconds + pub max_walltime_secs: u64, + + /// GPUs per node (if any) + #[serde(default)] + pub gpus_per_node: Option, + + /// GPU type (e.g., "h100", "a100") + #[serde(default)] + pub gpu_type: Option, + + /// GPU memory in GB + #[serde(default)] + pub gpu_memory_gb: Option, + + /// Whether the partition supports shared access + #[serde(default)] + pub shared: bool, + + /// Whether partition must be explicitly requested + #[serde(default)] + pub requires_explicit_request: bool, +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/config/loader.rs b/src/config/loader.rs index 97ec1fc4..4cc98841 100644 --- a/src/config/loader.rs +++ b/src/config/loader.rs @@ -216,6 +216,30 @@ poll_interval = 60 # Keep submission scripts after job submission (useful for debugging) keep_submission_scripts = false +[client.hpc] +# Default account to use for HPC jobs (applies to all profiles) +# default_account = "my_project" + +# Override settings for built-in profiles +# [client.hpc.profile_overrides.kestrel] +# default_account = "my_kestrel_account" + +# Define custom HPC profiles +# [[client.hpc.custom_profiles]] +# name = "my_cluster" +# display_name = "My Custom Cluster" +# description = "Our department's HPC cluster" +# detect_env_var = "MY_CLUSTER=prod" +# default_account = "dept_account" +# charge_factor_cpu = 1.0 +# charge_factor_gpu = 10.0 +# +# [[client.hpc.custom_profiles.my_cluster.partitions]] +# name = "compute" +# cpus_per_node = 64 +# memory_mb = 256000 +# max_walltime_secs = 172800 # 2 days + [server] # Hostname/IP to bind to url = "localhost" diff --git a/src/config/mod.rs b/src/config/mod.rs index 1cc49bf8..618bda2b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -28,7 +28,10 @@ mod dash; mod loader; mod server; -pub use client::{ClientConfig, ClientRunConfig, ClientSlurmConfig}; +pub use client::{ + ClientConfig, ClientHpcConfig, ClientRunConfig, ClientSlurmConfig, HpcPartitionConfig, + HpcProfileConfig, HpcProfileOverride, +}; pub use dash::DashConfig; pub use loader::{ConfigPaths, TorcConfig}; pub use server::{ServerConfig, ServerLoggingConfig}; diff --git a/src/main.rs b/src/main.rs index 3bf3085c..4823e573 100644 --- a/src/main.rs +++ b/src/main.rs @@ -9,6 +9,7 @@ use torc::client::commands::compute_nodes::handle_compute_node_commands; use torc::client::commands::config::handle_config_commands; use torc::client::commands::events::handle_event_commands; use torc::client::commands::files::handle_file_commands; +use torc::client::commands::hpc::handle_hpc_commands; use torc::client::commands::job_dependencies::handle_job_dependency_commands; use torc::client::commands::jobs::handle_job_commands; use torc::client::commands::reports::handle_report_commands; @@ -202,13 +203,20 @@ fn main() { eprintln!( "The spec does not define an on_workflow_start action with schedule_nodes." ); - eprintln!("To submit to a scheduler, add a workflow action like:"); + eprintln!("To submit to Slurm, either:"); eprintln!(); - eprintln!(" actions:"); - eprintln!(" - trigger_type: on_workflow_start"); - eprintln!(" action_type: schedule_nodes"); - eprintln!(" scheduler_type: slurm"); - eprintln!(" scheduler: \"my-cluster\""); + eprintln!(" 1. Use 'torc submit-slurm' to auto-generate schedulers:"); + eprintln!( + " torc submit-slurm --account {}", + workflow_spec_or_id + ); + eprintln!(); + eprintln!(" 2. Add a workflow action manually:"); + eprintln!(" actions:"); + eprintln!(" - trigger_type: on_workflow_start"); + eprintln!(" action_type: schedule_nodes"); + eprintln!(" scheduler_type: slurm"); + eprintln!(" scheduler: \"my-scheduler\""); eprintln!(); eprintln!("Or run locally instead:"); eprintln!(" torc run {}", workflow_spec_or_id); @@ -219,6 +227,7 @@ fn main() { let user = std::env::var("USER") .or_else(|_| std::env::var("USERNAME")) .unwrap_or_else(|_| "unknown".to_string()); + match WorkflowSpec::create_workflow_from_spec( &config, workflow_spec_or_id, @@ -302,6 +311,136 @@ fn main() { } } } + Commands::SubmitSlurm { + workflow_spec, + account, + hpc_profile, + single_allocation, + ignore_missing_data, + skip_checks, + } => { + use torc::client::commands::slurm::generate_schedulers_for_workflow; + + // Load the workflow spec + let mut spec = match WorkflowSpec::from_spec_file(workflow_spec) { + Ok(spec) => spec, + Err(e) => { + eprintln!("Error loading workflow spec: {}", e); + std::process::exit(1); + } + }; + + // Get HPC profile + let torc_config = TorcConfig::load().unwrap_or_default(); + let registry = torc::client::commands::hpc::create_registry_with_config_public( + &torc_config.client.hpc, + ); + + let profile = if let Some(name) = hpc_profile { + registry.get(name) + } else { + registry.detect() + }; + + let profile = match profile { + Some(p) => p, + None => { + if hpc_profile.is_some() { + eprintln!("Unknown HPC profile: {}", hpc_profile.as_ref().unwrap()); + } else { + eprintln!("No HPC profile specified and no system detected."); + eprintln!("Use --hpc-profile to specify a profile."); + } + std::process::exit(1); + } + }; + + // Generate schedulers + match generate_schedulers_for_workflow( + &mut spec, + profile, + account, + *single_allocation, + true, + true, + ) { + Ok(result) => { + eprintln!( + "Auto-generated {} scheduler(s) and {} action(s) using {} profile", + result.scheduler_count, result.action_count, profile.name + ); + for warning in &result.warnings { + eprintln!(" Warning: {}", warning); + } + } + Err(e) => { + eprintln!("Error generating schedulers: {}", e); + std::process::exit(1); + } + } + + // Print warning about auto-generated configuration + eprintln!(); + eprintln!("WARNING: Schedulers and actions were auto-generated using heuristics."); + eprintln!(" For complex workflows, this may not be optimal."); + eprintln!(); + eprintln!("TIP: To preview and validate the configuration before submitting, use:"); + eprintln!( + " torc slurm generate --account {} {}", + account, workflow_spec + ); + eprintln!(); + + // Write modified spec to temp file + let temp_dir = std::env::temp_dir(); + let temp_file = + temp_dir.join(format!("torc_submit_workflow_{}.yaml", std::process::id())); + std::fs::write(&temp_file, serde_yaml::to_string(&spec).unwrap()) + .expect("Failed to write temporary workflow file"); + + // Create workflow from spec + let user = std::env::var("USER") + .or_else(|_| std::env::var("USERNAME")) + .unwrap_or_else(|_| "unknown".to_string()); + + let workflow_id = match WorkflowSpec::create_workflow_from_spec( + &config, + &temp_file, + &user, + true, + *skip_checks, + ) { + Ok(id) => { + println!("Created workflow {}", id); + id + } + Err(e) => { + eprintln!("Error creating workflow from spec: {}", e); + std::process::exit(1); + } + }; + + // Submit the workflow + match default_api::get_workflow(&config, workflow_id) { + Ok(workflow) => { + let workflow_manager = + WorkflowManager::new(config.clone(), torc_config, workflow); + match workflow_manager.start(*ignore_missing_data) { + Ok(()) => { + println!("Successfully submitted workflow {}", workflow_id); + } + Err(e) => { + eprintln!("Error submitting workflow {}: {}", workflow_id, e); + std::process::exit(1); + } + } + } + Err(e) => { + eprintln!("Error getting workflow {}: {}", workflow_id, e); + std::process::exit(1); + } + } + } Commands::Workflows { command } => { handle_workflow_commands(&config, command, &format); } @@ -335,6 +474,9 @@ fn main() { Commands::ScheduledComputeNodes { command } => { handle_scheduled_compute_node_commands(&config, command, &format); } + Commands::Hpc { command } => { + handle_hpc_commands(command, &format); + } Commands::Reports { command } => { handle_report_commands(&config, command, &format); } diff --git a/src/server/api/workflow_actions.rs b/src/server/api/workflow_actions.rs index b8ec91c0..7faaec8c 100644 --- a/src/server/api/workflow_actions.rs +++ b/src/server/api/workflow_actions.rs @@ -11,6 +11,7 @@ use crate::server::api_types::{ }; use crate::models; +use crate::models::JobStatus; use super::{ApiContext, database_error}; @@ -724,8 +725,6 @@ impl WorkflowActionsApiImpl { job_ids: &[i64], trigger_type: &str, ) -> Result { - use crate::models::JobStatus; - let mut count = 0i64; for job_id in job_ids { diff --git a/src/tui/api.rs b/src/tui/api.rs index 0a3809fa..534b0397 100644 --- a/src/tui/api.rs +++ b/src/tui/api.rs @@ -1,9 +1,10 @@ use crate::client::apis::configuration::Configuration; use crate::client::apis::default_api; use crate::client::config::TorcConfig; +use crate::client::workflow_spec::WorkflowSpec; use crate::models::{ - EventModel, FileModel, JobDependencyModel, JobModel, ResultModel, ScheduledComputeNodesModel, - WorkflowModel, + EventModel, FileModel, JobDependencyModel, JobModel, JobStatus, ResultModel, + ScheduledComputeNodesModel, WorkflowModel, }; use anyhow::{Context, Result}; @@ -225,8 +226,6 @@ impl TorcClient { } pub fn cancel_job(&self, job_id: i64) -> Result<()> { - use crate::models::JobStatus; - // Get the existing job, update status, and PUT back let mut job = self.get_job(job_id)?; job.status = Some(JobStatus::Canceled); @@ -237,8 +236,6 @@ impl TorcClient { } pub fn terminate_job(&self, job_id: i64) -> Result<()> { - use crate::models::JobStatus; - let mut job = self.get_job(job_id)?; job.status = Some(JobStatus::Terminated); @@ -248,8 +245,6 @@ impl TorcClient { } pub fn retry_job(&self, job_id: i64) -> Result<()> { - use crate::models::JobStatus; - let mut job = self.get_job(job_id)?; job.status = Some(JobStatus::Ready); @@ -271,8 +266,6 @@ impl TorcClient { } pub fn create_workflow_from_file(&self, path: &str) -> Result { - use crate::client::workflow_spec::WorkflowSpec; - // Get current user let user = std::env::var("USER") .or_else(|_| std::env::var("USERNAME")) diff --git a/src/tui/app.rs b/src/tui/app.rs index 47bbdd42..a0a48052 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -1,8 +1,16 @@ +use std::collections::HashMap; +use std::path::PathBuf; + +use anyhow::Result; +use petgraph::graph::NodeIndex; +use ratatui::widgets::TableState; + +use crate::client::log_paths::{ + get_job_stderr_path, get_job_stdout_path, get_slurm_stderr_path, get_slurm_stdout_path, +}; use crate::models::{ EventModel, FileModel, JobModel, ResultModel, ScheduledComputeNodesModel, WorkflowModel, }; -use anyhow::Result; -use ratatui::widgets::TableState; use super::api::TorcClient; use super::components::{ @@ -772,9 +780,6 @@ impl App { } pub fn build_dag_from_jobs(&mut self) { - use petgraph::graph::NodeIndex; - use std::collections::HashMap; - let mut dag = DagLayout::new(); let mut job_id_to_node: HashMap = HashMap::new(); @@ -1549,9 +1554,6 @@ impl App { } fn load_job_logs(&self, viewer: &mut LogViewer) -> Result<()> { - use crate::client::log_paths::{get_job_stderr_path, get_job_stdout_path}; - use std::path::PathBuf; - // Try to find log files based on job results if let Some(workflow_id) = self.selected_workflow_id { let results = self.client.list_results(workflow_id)?; @@ -1644,9 +1646,6 @@ impl App { } fn load_slurm_logs(&self, viewer: &mut LogViewer, scheduler_id: &str) -> Result<()> { - use crate::client::log_paths::{get_slurm_stderr_path, get_slurm_stdout_path}; - use std::path::PathBuf; - // Default output directory is "output" in the current working directory let output_dir = PathBuf::from("output"); diff --git a/src/tui/components.rs b/src/tui/components.rs index 777ae00c..47c725b1 100644 --- a/src/tui/components.rs +++ b/src/tui/components.rs @@ -1,5 +1,8 @@ //! Reusable UI components for the TUI +use std::fs; +use std::path::Path; + use ratatui::{ Frame, layout::{Alignment, Constraint, Direction, Layout, Rect}, @@ -790,9 +793,6 @@ impl FileViewer { } pub fn load_content(&mut self) -> Result<(), String> { - use std::fs; - use std::path::Path; - let path = Path::new(&self.file_path); if !path.exists() { diff --git a/src/tui/dag.rs b/src/tui/dag.rs index 9a41eb83..f586159a 100644 --- a/src/tui/dag.rs +++ b/src/tui/dag.rs @@ -73,6 +73,40 @@ impl DagLayout { layers[max_predecessor_layer].push(node); } + // Sort nodes within each layer to group related subgraphs together + // Group by their parent nodes to keep subgraphs visually connected + for layer in &mut layers { + layer.sort_by(|a, b| { + // Get predecessor indices as sort keys + let a_preds: Vec = self + .graph + .edges_directed(*a, petgraph::Direction::Incoming) + .map(|e| e.source().index()) + .collect(); + let b_preds: Vec = self + .graph + .edges_directed(*b, petgraph::Direction::Incoming) + .map(|e| e.source().index()) + .collect(); + + // Sort by first predecessor, then by job name for consistency + match (a_preds.first(), b_preds.first()) { + (Some(a_pred), Some(b_pred)) => a_pred.cmp(b_pred).then_with(|| { + let a_name = &self.graph[*a].name; + let b_name = &self.graph[*b].name; + a_name.cmp(b_name) + }), + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + (None, None) => { + let a_name = &self.graph[*a].name; + let b_name = &self.graph[*b].name; + a_name.cmp(b_name) + } + } + }); + } + layers } diff --git a/src/tui/mod.rs b/src/tui/mod.rs index ff2ac593..f92104de 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -1,3 +1,5 @@ +use std::io; + use anyhow::Result; use crossterm::{ event::{self, DisableMouseCapture, EnableMouseCapture, Event, KeyCode, KeyEventKind}, @@ -5,7 +7,8 @@ use crossterm::{ terminal::{EnterAlternateScreen, LeaveAlternateScreen, disable_raw_mode, enable_raw_mode}, }; use ratatui::{Terminal, backend::CrosstermBackend}; -use std::io; + +use crate::client::apis::default_api; mod api; mod app; @@ -18,8 +21,6 @@ use components::StatusMessage; /// Check if the Torc server is reachable by calling the ping endpoint fn check_server_connection(base_url: &str) -> bool { - use crate::client::apis::default_api; - let config = crate::client::apis::configuration::Configuration { base_path: base_url.to_string(), ..Default::default() diff --git a/src/tui/ui.rs b/src/tui/ui.rs index f2941c63..61c03991 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -1,4 +1,7 @@ +use std::collections::HashMap; + use chrono::{DateTime, Local, Utc}; +use petgraph::visit::{EdgeRef, Topo}; use ratatui::{ Frame, layout::{Constraint, Direction, Layout, Rect}, @@ -936,8 +939,30 @@ fn draw_dag(f: &mut Frame, area: Rect, app: &App) { )])); } + // Group jobs in this layer by their predecessors to show subgraphs + let mut current_pred_group: Option = None; + // Display all jobs in this layer for &node_idx in layer { + // Check if this job belongs to a different subgraph + let first_pred: Option = dag + .graph + .edges_directed(node_idx, petgraph::Direction::Incoming) + .next() + .map(|e| e.source().index()); + + // Add a separator between different subgraph groups within the same layer + if layer.len() > 1 + && first_pred != current_pred_group + && current_pred_group.is_some() + { + lines.push(Line::from(vec![Span::styled( + " ─ ─ ─", + Style::default().fg(Color::DarkGray), + )])); + } + current_pred_group = first_pred; + let node_data = &dag.graph[node_idx]; // Determine color based on status @@ -1014,9 +1039,6 @@ fn draw_dag(f: &mut Frame, area: Rect, app: &App) { fn dag_compute_layers( graph: &petgraph::Graph, ) -> Vec> { - use petgraph::visit::{EdgeRef, Topo}; - use std::collections::HashMap; - let mut layers: Vec> = Vec::new(); let mut node_layer: HashMap = HashMap::new(); @@ -1040,5 +1062,37 @@ fn dag_compute_layers( layers[max_predecessor_layer].push(node); } + // Sort nodes within each layer to group related subgraphs together + // Group by their parent nodes to keep subgraphs visually connected + for layer in &mut layers { + layer.sort_by(|a, b| { + // Get predecessor indices as sort keys + let a_preds: Vec = graph + .edges_directed(*a, petgraph::Direction::Incoming) + .map(|e| e.source().index()) + .collect(); + let b_preds: Vec = graph + .edges_directed(*b, petgraph::Direction::Incoming) + .map(|e| e.source().index()) + .collect(); + + // Sort by first predecessor, then by job name for consistency + match (a_preds.first(), b_preds.first()) { + (Some(a_pred), Some(b_pred)) => a_pred.cmp(b_pred).then_with(|| { + let a_name = &graph[*a].name; + let b_name = &graph[*b].name; + a_name.cmp(b_name) + }), + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + (None, None) => { + let a_name = &graph[*a].name; + let b_name = &graph[*b].name; + a_name.cmp(b_name) + } + } + }); + } + layers } diff --git a/tests/test_claim_jobs_based_on_resources.rs b/tests/test_claim_jobs_based_on_resources.rs index 095799cc..1dcd6e07 100644 --- a/tests/test_claim_jobs_based_on_resources.rs +++ b/tests/test_claim_jobs_based_on_resources.rs @@ -1,5 +1,9 @@ mod common; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use std::thread; + use common::{ ServerProcess, create_custom_resources_workflow, create_dependency_chain_workflow, create_diverse_jobs_workflow, create_gpu_workflow, create_high_cpu_workflow, @@ -1623,10 +1627,6 @@ fn test_prepare_jobs_different_sort_methods_different_orders(start_server: &Serv /// ensuring the server's database locking mechanism correctly prevents race conditions #[rstest] fn test_prepare_jobs_concurrent_allocation(start_server: &ServerProcess) { - use std::collections::{HashMap, HashSet}; - use std::sync::{Arc, Mutex}; - use std::thread; - let config = &start_server.config; let num_threads = thread::available_parallelism() diff --git a/tests/test_hpc.rs b/tests/test_hpc.rs new file mode 100644 index 00000000..6737577d --- /dev/null +++ b/tests/test_hpc.rs @@ -0,0 +1,1161 @@ +//! Tests for HPC profile system and scheduler generation + +use rstest::rstest; +use std::collections::HashMap; +use torc::client::commands::slurm::{ + generate_schedulers_for_workflow, parse_memory_mb, parse_walltime_secs, secs_to_walltime, +}; +use torc::client::hpc::kestrel::kestrel_profile; +use torc::client::hpc::{HpcDetection, HpcPartition, HpcProfile, HpcProfileRegistry}; +use torc::client::workflow_spec::{JobSpec, ResourceRequirementsSpec, WorkflowSpec}; +use torc::time_utils::duration_string_to_seconds; + +// ============== Utility Function Tests ============== + +#[rstest] +fn test_parse_memory_mb() { + assert_eq!(parse_memory_mb("100g").unwrap(), 102400); + assert_eq!(parse_memory_mb("1G").unwrap(), 1024); + assert_eq!(parse_memory_mb("512m").unwrap(), 512); + assert_eq!(parse_memory_mb("512M").unwrap(), 512); + assert_eq!(parse_memory_mb("1024").unwrap(), 1024); + assert_eq!(parse_memory_mb("1024k").unwrap(), 1); +} + +#[rstest] +fn test_parse_walltime_secs() { + assert_eq!(parse_walltime_secs("1:00:00").unwrap(), 3600); + assert_eq!(parse_walltime_secs("4:00:00").unwrap(), 14400); + assert_eq!(parse_walltime_secs("1-00:00:00").unwrap(), 86400); + assert_eq!(parse_walltime_secs("2-00:00:00").unwrap(), 172800); + assert_eq!(parse_walltime_secs("10-00:00:00").unwrap(), 864000); + assert_eq!(parse_walltime_secs("0:30:00").unwrap(), 1800); +} + +#[rstest] +fn test_duration_string_to_seconds() { + // Test ISO 8601 duration parsing using the consolidated function from time_utils + assert_eq!(duration_string_to_seconds("PT1H").unwrap(), 3600); + assert_eq!(duration_string_to_seconds("PT30M").unwrap(), 1800); + assert_eq!(duration_string_to_seconds("PT1H30M").unwrap(), 5400); + assert_eq!(duration_string_to_seconds("P1D").unwrap(), 86400); + assert_eq!(duration_string_to_seconds("P1DT2H").unwrap(), 93600); + assert_eq!(duration_string_to_seconds("P0DT1M").unwrap(), 60); + assert_eq!(duration_string_to_seconds("PT4H").unwrap(), 14400); +} + +#[rstest] +fn test_secs_to_walltime() { + assert_eq!(secs_to_walltime(3600), "01:00:00"); + assert_eq!(secs_to_walltime(14400), "04:00:00"); + assert_eq!(secs_to_walltime(86400), "1-00:00:00"); + assert_eq!(secs_to_walltime(172800), "2-00:00:00"); + assert_eq!(secs_to_walltime(93600), "1-02:00:00"); // 1 day 2 hours +} + +// ============== Profile System Tests ============== + +fn create_test_partition( + name: &str, + cpus: u32, + memory_mb: u64, + walltime_secs: u64, + gpus: Option, +) -> HpcPartition { + HpcPartition { + name: name.to_string(), + description: String::new(), + cpus_per_node: cpus, + memory_mb, + max_walltime_secs: walltime_secs, + max_nodes: None, + max_nodes_per_user: None, + min_nodes: None, + gpus_per_node: gpus, + gpu_type: None, + gpu_memory_gb: None, + local_disk_gb: None, + shared: false, + requires_explicit_request: false, + default_qos: None, + features: vec![], + } +} + +fn create_test_profile(name: &str, partitions: Vec) -> HpcProfile { + HpcProfile { + name: name.to_string(), + display_name: format!("Test {}", name), + description: String::new(), + detection: vec![], + default_account: None, + partitions, + charge_factor_cpu: 1.0, + charge_factor_gpu: 10.0, + metadata: HashMap::new(), + } +} + +#[rstest] +fn test_partition_can_satisfy_basic() { + let partition = create_test_partition("standard", 104, 245760, 172800, None); + + // Should satisfy small request + assert!(partition.can_satisfy(4, 8192, 3600, None)); + // Should satisfy request up to limits + assert!(partition.can_satisfy(104, 245760, 172800, None)); + // Should fail if CPUs exceed + assert!(!partition.can_satisfy(105, 8192, 3600, None)); + // Should fail if memory exceeds + assert!(!partition.can_satisfy(4, 300000, 3600, None)); + // Should fail if walltime exceeds + assert!(!partition.can_satisfy(4, 8192, 200000, None)); +} + +#[rstest] +fn test_partition_can_satisfy_gpu() { + let partition = create_test_partition("gpu-h100", 128, 2097152, 172800, Some(4)); + + // Should satisfy GPU request within limits + assert!(partition.can_satisfy(64, 200000, 3600, Some(2))); + // Should fail if GPUs exceed + assert!(!partition.can_satisfy(64, 200000, 3600, Some(5))); + + // Non-GPU partition should not satisfy GPU requests + let cpu_partition = create_test_partition("standard", 104, 245760, 172800, None); + assert!(!cpu_partition.can_satisfy(4, 8192, 3600, Some(1))); +} + +#[rstest] +fn test_env_var_detection() { + let profile = HpcProfile { + name: "test".to_string(), + display_name: "Test Profile".to_string(), + description: "Test".to_string(), + detection: vec![HpcDetection::EnvVar { + name: "TEST_CLUSTER".to_string(), + value: "test".to_string(), + }], + default_account: None, + partitions: vec![], + charge_factor_cpu: 1.0, + charge_factor_gpu: 10.0, + metadata: HashMap::new(), + }; + + // Detection should work when env var matches + // SAFETY: Tests run serially and we restore the var + unsafe { + std::env::set_var("TEST_CLUSTER", "test"); + } + assert!(profile.detect()); + + // Detection should fail when env var doesn't match + unsafe { + std::env::set_var("TEST_CLUSTER", "other"); + } + assert!(!profile.detect()); + + unsafe { + std::env::remove_var("TEST_CLUSTER"); + } +} + +#[rstest] +fn test_profile_registry() { + let mut registry = HpcProfileRegistry::new(); + + let profile = create_test_profile( + "test", + vec![create_test_partition("standard", 64, 128000, 86400, None)], + ); + + registry.register(profile); + + assert!(registry.get("test").is_some()); + assert!(registry.get("nonexistent").is_none()); +} + +#[rstest] +fn test_walltime_format() { + let partition = create_test_partition("test", 64, 128000, 90061, None); // 25h 1m 1s + + let formatted = partition.max_walltime_str(); + assert!(formatted.contains("25") || formatted.contains("1-01")); +} + +// ============== Kestrel Profile Tests ============== + +#[rstest] +fn test_kestrel_profile_basics() { + let profile = kestrel_profile(); + assert_eq!(profile.name, "kestrel"); + assert_eq!(profile.display_name, "NREL Kestrel"); + assert!(!profile.partitions.is_empty()); +} + +#[rstest] +fn test_kestrel_has_expected_partitions() { + let profile = kestrel_profile(); + let partition_names: Vec<&str> = profile.partitions.iter().map(|p| p.name.as_str()).collect(); + + // Check for key partitions + assert!(partition_names.contains(&"debug")); + assert!(partition_names.contains(&"short")); + assert!(partition_names.contains(&"standard")); + assert!(partition_names.contains(&"gpu-h100")); +} + +#[rstest] +fn test_kestrel_standard_partition() { + let profile = kestrel_profile(); + let standard = profile + .get_partition("standard") + .expect("Standard partition not found"); + + assert_eq!(standard.cpus_per_node, 104); + assert_eq!(standard.memory_mb, 240_000); + assert_eq!(standard.max_walltime_secs, 172800); // 48 hours + assert!(standard.gpus_per_node.is_none()); +} + +#[rstest] +fn test_kestrel_gpu_partition() { + let profile = kestrel_profile(); + let gpu = profile + .get_partition("gpu-h100") + .expect("GPU partition not found"); + + assert_eq!(gpu.gpus_per_node, Some(4)); + assert!(gpu.gpu_type.is_some()); +} + +#[rstest] +fn test_kestrel_find_matching_partitions() { + let profile = kestrel_profile(); + + // Small CPU job should match multiple partitions + let matches = profile.find_matching_partitions(4, 8192, 3600, None); + assert!(!matches.is_empty()); + + // GPU job should only match GPU partitions + let gpu_matches = profile.find_matching_partitions(64, 200000, 3600, Some(2)); + assert!(!gpu_matches.is_empty()); + for partition in &gpu_matches { + assert!(partition.gpus_per_node.is_some()); + } +} + +#[rstest] +fn test_kestrel_hbw_requires_min_nodes() { + let profile = kestrel_profile(); + let hbw = profile + .get_partition("hbw") + .expect("HBW partition not found"); + + assert!(hbw.min_nodes.is_some()); +} + +// ============== Scheduler Generation Tests ============== + +#[rstest] +fn test_generate_schedulers_basic() { + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: Some("Test workflow".to_string()), + jobs: vec![ + JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }, + JobSpec { + name: "job2".to_string(), + command: "echo world".to_string(), + resource_requirements: Some("medium".to_string()), + depends_on: Some(vec!["job1".to_string()]), + ..Default::default() + }, + ], + resource_requirements: Some(vec![ + ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT1H".to_string(), + }, + ResourceRequirementsSpec { + name: "medium".to_string(), + num_cpus: 32, + num_gpus: 0, + num_nodes: 1, + memory: "64g".to_string(), + runtime: "PT4H".to_string(), + }, + ]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + // Should generate 2 schedulers: + // - small_scheduler for job1 (no dependencies, on_workflow_start) + // - medium_deferred_scheduler for job2 (has dependencies, on_jobs_ready) + // Schedulers are grouped by (resource_requirements, has_dependencies) + assert_eq!(result.scheduler_count, 2); + assert_eq!(result.action_count, 2); + + // Check that slurm_schedulers were added + assert!(spec.slurm_schedulers.is_some()); + let schedulers = spec.slurm_schedulers.as_ref().unwrap(); + assert_eq!(schedulers.len(), 2); + + // Check scheduler names - grouped by (resource_requirement, has_deps) + let scheduler_names: Vec<&str> = schedulers + .iter() + .filter_map(|s| s.name.as_deref()) + .collect(); + assert!(scheduler_names.contains(&"small_scheduler")); + assert!(scheduler_names.contains(&"medium_deferred_scheduler")); + + // Check that jobs were assigned to correct schedulers + // job1 (no deps) → small_scheduler + // job2 (has deps) → medium_deferred_scheduler + assert_eq!(spec.jobs[0].scheduler.as_ref().unwrap(), "small_scheduler"); + assert_eq!( + spec.jobs[1].scheduler.as_ref().unwrap(), + "medium_deferred_scheduler" + ); + + // Check that workflow actions were added + assert!(spec.actions.is_some()); + let actions = spec.actions.as_ref().unwrap(); + assert_eq!(actions.len(), 2); + + // Jobs without dependencies use on_workflow_start + let small_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("small_scheduler")) + .unwrap(); + assert_eq!(small_action.trigger_type, "on_workflow_start"); + assert_eq!(small_action.action_type, "schedule_nodes"); + + // Jobs with dependencies use on_jobs_ready + let medium_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("medium_deferred_scheduler")) + .unwrap(); + assert_eq!(medium_action.trigger_type, "on_jobs_ready"); + assert_eq!(medium_action.action_type, "schedule_nodes"); +} + +#[rstest] +fn test_generate_schedulers_with_gpus() { + let mut spec = WorkflowSpec { + name: "gpu_workflow".to_string(), + description: Some("GPU workflow".to_string()), + jobs: vec![JobSpec { + name: "gpu_job".to_string(), + command: "python train.py".to_string(), + resource_requirements: Some("gpu_heavy".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "gpu_heavy".to_string(), + num_cpus: 64, + num_gpus: 2, + num_nodes: 1, + memory: "200g".to_string(), + runtime: "PT8H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + assert_eq!(result.scheduler_count, 1); + + let schedulers = spec.slurm_schedulers.as_ref().unwrap(); + assert_eq!(schedulers.len(), 1); + + let gpu_scheduler = &schedulers[0]; + // Per-resource-requirement scheduler naming: rr_name + "_scheduler" + assert_eq!(gpu_scheduler.name.as_deref(), Some("gpu_heavy_scheduler")); + assert_eq!(gpu_scheduler.account, "testaccount"); + // GPU scheduler should have gres set + assert!(gpu_scheduler.gres.is_some()); + assert!(gpu_scheduler.gres.as_ref().unwrap().contains("gpu")); +} + +#[rstest] +fn test_generate_schedulers_no_actions() { + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + // Pass add_actions = false + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, false, false) + .unwrap(); + + assert_eq!(result.scheduler_count, 1); + assert_eq!(result.action_count, 0); + + // Schedulers should be added + assert!(spec.slurm_schedulers.is_some()); + + // But no actions + assert!(spec.actions.is_none() || spec.actions.as_ref().unwrap().is_empty()); +} + +#[rstest] +fn test_generate_schedulers_shared_by_jobs() { + // Jobs with the same resource requirements share a scheduler + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![ + JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }, + JobSpec { + name: "job2".to_string(), + command: "echo world".to_string(), + resource_requirements: Some("small".to_string()), // Same requirements + ..Default::default() + }, + ], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + // Only one scheduler since both jobs use the same resource requirements + assert_eq!(result.scheduler_count, 1); + + let schedulers = spec.slurm_schedulers.as_ref().unwrap(); + assert_eq!(schedulers.len(), 1); + + // Both jobs should share the same scheduler + assert_eq!(spec.jobs[0].scheduler.as_ref().unwrap(), "small_scheduler"); + assert_eq!(spec.jobs[1].scheduler.as_ref().unwrap(), "small_scheduler"); +} + +#[rstest] +fn test_generate_schedulers_errors_no_resource_requirements() { + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("nonexistent".to_string()), + ..Default::default() + }], + resource_requirements: None, // No resource requirements defined + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false); + + // Should return an error when no resource requirements are defined + match result { + Err(e) => assert!(e.contains("resource_requirements")), + Ok(_) => panic!("Expected error but got Ok"), + } +} + +#[rstest] +fn test_generate_schedulers_existing_schedulers_no_force() { + use torc::client::workflow_spec::SlurmSchedulerSpec; + + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + scheduler: Some("existing_scheduler".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT1H".to_string(), + }]), + slurm_schedulers: Some(vec![SlurmSchedulerSpec { + name: Some("existing_scheduler".to_string()), + account: "test".to_string(), + nodes: 1, + walltime: "01:00:00".to_string(), + gres: None, + mem: None, + ntasks_per_node: None, + partition: None, + qos: None, + tmp: None, + extra: None, + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + // force = false should return error when slurm_schedulers already exists + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false); + + match result { + Err(e) => assert!(e.contains("already has slurm_schedulers")), + Ok(_) => panic!("Expected error but got Ok"), + } +} + +#[rstest] +fn test_generate_schedulers_existing_schedulers_with_force() { + use torc::client::workflow_spec::SlurmSchedulerSpec; + + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + scheduler: Some("existing_scheduler".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT1H".to_string(), + }]), + slurm_schedulers: Some(vec![SlurmSchedulerSpec { + name: Some("existing_scheduler".to_string()), + account: "test".to_string(), + nodes: 1, + walltime: "01:00:00".to_string(), + gres: None, + mem: None, + ntasks_per_node: None, + partition: None, + qos: None, + tmp: None, + extra: None, + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + // force = true should succeed even when slurm_schedulers already exists + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, true) + .unwrap(); + + // Job should be reassigned to scheduler based on resource requirement name + assert_eq!(spec.jobs[0].scheduler.as_ref().unwrap(), "small_scheduler"); + + // New scheduler should be generated + assert_eq!(result.scheduler_count, 1); +} + +#[rstest] +fn test_generate_schedulers_sets_correct_account() { + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let _result = generate_schedulers_for_workflow( + &mut spec, + &profile, + "my_project_account", + false, + true, + false, + ) + .unwrap(); + + let scheduler = &spec.slurm_schedulers.as_ref().unwrap()[0]; + assert_eq!(scheduler.account, "my_project_account"); +} + +#[rstest] +fn test_generate_schedulers_sets_walltime() { + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("long_job".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "long_job".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT12H".to_string(), // 12 hours - matches standard partition + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let _result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + let scheduler = &spec.slurm_schedulers.as_ref().unwrap()[0]; + // Walltime should be set to the partition's max (2 days for standard), not the job's runtime. + // This provides headroom for jobs that run slightly longer than expected. + assert_eq!(scheduler.walltime, "2-00:00:00"); +} + +#[rstest] +fn test_generate_schedulers_sets_memory() { + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + description: None, + jobs: vec![JobSpec { + name: "job1".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("mem_job".to_string()), + ..Default::default() + }], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "mem_job".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "128g".to_string(), + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let _result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + let scheduler = &spec.slurm_schedulers.as_ref().unwrap()[0]; + // Memory should be set + assert_eq!(scheduler.mem.as_deref(), Some("128g")); +} + +#[rstest] +fn test_generate_schedulers_per_resource_requirement() { + // Schedulers are created per (resource_requirement, has_dependencies) + // Jobs with same resource requirements but different dependency status get separate schedulers + let mut spec = WorkflowSpec { + name: "staged_workflow".to_string(), + description: None, + jobs: vec![ + JobSpec { + name: "setup".to_string(), + command: "echo setup".to_string(), + resource_requirements: Some("small".to_string()), + depends_on: None, // No dependencies + ..Default::default() + }, + JobSpec { + name: "process".to_string(), + command: "echo process".to_string(), + resource_requirements: Some("medium".to_string()), + depends_on: Some(vec!["setup".to_string()]), // Depends on setup + ..Default::default() + }, + JobSpec { + name: "finalize".to_string(), + command: "echo finalize".to_string(), + resource_requirements: Some("small".to_string()), // Same as setup + depends_on: Some(vec!["process".to_string()]), // Depends on process + ..Default::default() + }, + ], + resource_requirements: Some(vec![ + ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 2, + num_gpus: 0, + num_nodes: 1, + memory: "4g".to_string(), + runtime: "PT30M".to_string(), + }, + ResourceRequirementsSpec { + name: "medium".to_string(), + num_cpus: 8, + num_gpus: 0, + num_nodes: 1, + memory: "16g".to_string(), + runtime: "PT2H".to_string(), + }, + ]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + // 3 schedulers: + // - small_scheduler for setup (no deps, on_workflow_start) + // - medium_deferred_scheduler for process (has deps, on_jobs_ready) + // - small_deferred_scheduler for finalize (has deps, on_jobs_ready) + // Stage-aware scheduling launches nodes when jobs become ready + assert_eq!(result.scheduler_count, 3); + assert_eq!(result.action_count, 3); + + let actions = spec.actions.as_ref().unwrap(); + assert_eq!(actions.len(), 3); + + // Jobs should be assigned to schedulers based on (resource_requirement, has_deps) + assert_eq!(spec.jobs[0].scheduler.as_deref(), Some("small_scheduler")); // setup (no deps) + assert_eq!( + spec.jobs[1].scheduler.as_deref(), + Some("medium_deferred_scheduler") + ); // process (has deps) + assert_eq!( + spec.jobs[2].scheduler.as_deref(), + Some("small_deferred_scheduler") + ); // finalize (has deps) + + // Jobs without dependencies use on_workflow_start + let small_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("small_scheduler")) + .unwrap(); + assert_eq!(small_action.trigger_type, "on_workflow_start"); + + // Jobs with dependencies use on_jobs_ready + let medium_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("medium_deferred_scheduler")) + .unwrap(); + assert_eq!(medium_action.trigger_type, "on_jobs_ready"); + + let finalize_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("small_deferred_scheduler")) + .unwrap(); + assert_eq!(finalize_action.trigger_type, "on_jobs_ready"); +} + +/// Test that num_allocations is auto-calculated based on job count and partition capacity +#[test] +fn test_generate_schedulers_auto_calculates_allocations() { + use torc::client::workflow_spec::{JobSpec, ResourceRequirementsSpec, WorkflowSpec}; + + // Create a workflow with 10 jobs, each requiring 26 CPUs + // On Kestrel (104 CPUs/node), 4 jobs fit per node + // So we need 10/4 = 3 nodes (rounded up) + let jobs: Vec = (0..10) + .map(|i| JobSpec { + name: format!("job_{:03}", i), + command: "echo hello".to_string(), + resource_requirements: Some("compute".to_string()), + ..Default::default() + }) + .collect(); + + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + user: Some("testuser".to_string()), + jobs, + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "compute".to_string(), + num_cpus: 26, // 104 / 26 = 4 jobs per node + num_gpus: 0, + num_nodes: 1, + memory: "10g".to_string(), + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + + // Pass None for num_allocations to trigger auto-calculation + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + assert_eq!(result.scheduler_count, 1); + assert_eq!(result.action_count, 1); + + let actions = spec.actions.as_ref().unwrap(); + let action = &actions[0]; + + // 10 jobs, 4 jobs per node (104 CPUs / 26 CPUs) = 3 nodes needed (rounded up) + // With 1 node per allocation = 3 allocations + assert_eq!(action.num_allocations, Some(3)); +} + +/// Test auto-calculation with parameterized jobs +#[test] +fn test_generate_schedulers_auto_calculates_with_parameters() { + // One parameterized job that expands to 100 jobs + let mut parameters = HashMap::new(); + parameters.insert("i".to_string(), "1:100".to_string()); + + let jobs = vec![JobSpec { + name: "job_{i:03d}".to_string(), + command: "echo hello".to_string(), + resource_requirements: Some("small".to_string()), + parameters: Some(parameters), + ..Default::default() + }]; + + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + user: Some("testuser".to_string()), + jobs, + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 52, // 104 / 52 = 2 jobs per node + num_gpus: 0, + num_nodes: 1, + memory: "10g".to_string(), + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + + // Pass None for num_allocations to trigger auto-calculation + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + assert_eq!(result.scheduler_count, 1); + assert_eq!(result.action_count, 1); + + let actions = spec.actions.as_ref().unwrap(); + let action = &actions[0]; + + // 100 jobs (from parameterized expansion), 2 jobs per node (104 CPUs / 52 CPUs) = 50 nodes needed + // With 1 node per allocation = 50 allocations + assert_eq!(action.num_allocations, Some(50)); +} + +/// Test stage-aware scheduling: jobs with and without dependencies get separate schedulers. +/// This enables launching compute nodes only when jobs become ready. +#[test] +fn test_generate_schedulers_stage_aware_for_dependent_jobs() { + // job1: no dependencies → scheduled at on_workflow_start + // job2: depends on job1 → scheduled at on_jobs_ready when job1 completes + // Both use the same resource requirements but get separate schedulers + let jobs = vec![ + JobSpec { + name: "job1".to_string(), + command: "echo job1".to_string(), + resource_requirements: Some("small".to_string()), + ..Default::default() + }, + JobSpec { + name: "job2".to_string(), + command: "echo job2".to_string(), + resource_requirements: Some("small".to_string()), + depends_on: Some(vec!["job1".to_string()]), + ..Default::default() + }, + ]; + + let mut spec = WorkflowSpec { + name: "test_workflow".to_string(), + user: Some("testuser".to_string()), + jobs, + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "small".to_string(), + num_cpus: 4, + num_gpus: 0, + num_nodes: 1, + memory: "8g".to_string(), + runtime: "PT30M".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + // Should generate 2 schedulers for stage-aware scheduling: + // - small_scheduler (on_workflow_start) for job1 + // - small_deferred_scheduler (on_jobs_ready) for job2 + assert_eq!(result.scheduler_count, 2); + assert_eq!(result.action_count, 2); + + let schedulers = spec.slurm_schedulers.as_ref().unwrap(); + assert_eq!(schedulers.len(), 2); + + // Jobs are assigned to different schedulers based on dependency status + assert_eq!(spec.jobs[0].scheduler, Some("small_scheduler".to_string())); // no deps + assert_eq!( + spec.jobs[1].scheduler, + Some("small_deferred_scheduler".to_string()) + ); // has deps + + // Verify trigger types + let actions = spec.actions.as_ref().unwrap(); + assert_eq!(actions.len(), 2); + + let job1_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("small_scheduler")) + .unwrap(); + assert_eq!(job1_action.trigger_type, "on_workflow_start"); + + let job2_action = actions + .iter() + .find(|a| a.scheduler.as_deref() == Some("small_deferred_scheduler")) + .unwrap(); + assert_eq!(job2_action.trigger_type, "on_jobs_ready"); +} + +/// Test that jobs-per-node calculation considers memory, not just CPUs. +/// When memory is the limiting factor, we should allocate more nodes. +#[rstest] +fn test_generate_schedulers_memory_constrained_allocation() { + // Create 10 jobs that are memory-heavy: 8 CPUs, 120GB each + // On Kestrel standard nodes (104 CPUs, 240GB): + // - CPU-based: 104/8 = 13 jobs per node + // - Memory-based: 240,000MB / 122,880MB = ~1.95 = 1 job per node + // Memory should be the limiting factor, so we need 10 nodes for 10 jobs + let jobs: Vec = (0..10) + .map(|i| JobSpec { + name: format!("memory_job_{}", i), + command: "echo heavy".to_string(), + resource_requirements: Some("memory_heavy".to_string()), + ..Default::default() + }) + .collect(); + + let mut spec = WorkflowSpec { + name: "memory_test".to_string(), + description: Some("Test memory-constrained allocation".to_string()), + jobs, + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "memory_heavy".to_string(), + num_cpus: 8, // Small CPU requirement + num_gpus: 0, + num_nodes: 1, + memory: "120g".to_string(), // Large memory requirement (120GB = 122,880MB) + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + assert_eq!(result.scheduler_count, 1); + assert_eq!(result.action_count, 1); + + // Check the action's num_allocations + let actions = spec.actions.as_ref().unwrap(); + assert_eq!(actions.len(), 1); + + let action = &actions[0]; + // With memory as limiting factor (1 job per node), we need 10 allocations for 10 jobs + // If only CPU was considered, it would be ceil(10/13) = 1 allocation (wrong!) + assert_eq!( + action.num_allocations, + Some(10), + "Should allocate 10 nodes for 10 memory-heavy jobs (1 job per node due to 120GB memory)" + ); +} + +/// Test mixed constraint: some jobs CPU-limited, some memory-limited +#[rstest] +fn test_generate_schedulers_cpu_vs_memory_constraint() { + let mut spec = WorkflowSpec { + name: "mixed_constraint_test".to_string(), + description: Some("Test CPU vs memory constraints".to_string()), + jobs: vec![ + // 4 CPU-limited jobs: 52 CPUs, 60GB each + // On 104 CPU / 240GB node: 104/52=2 by CPU, 240000/61440=3.9 by memory -> CPU wins (2 per node) + // 4 jobs / 2 per node = 2 allocations + JobSpec { + name: "cpu_job_1".to_string(), + command: "echo cpu".to_string(), + resource_requirements: Some("cpu_heavy".to_string()), + ..Default::default() + }, + JobSpec { + name: "cpu_job_2".to_string(), + command: "echo cpu".to_string(), + resource_requirements: Some("cpu_heavy".to_string()), + ..Default::default() + }, + JobSpec { + name: "cpu_job_3".to_string(), + command: "echo cpu".to_string(), + resource_requirements: Some("cpu_heavy".to_string()), + ..Default::default() + }, + JobSpec { + name: "cpu_job_4".to_string(), + command: "echo cpu".to_string(), + resource_requirements: Some("cpu_heavy".to_string()), + ..Default::default() + }, + ], + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "cpu_heavy".to_string(), + num_cpus: 52, // Half the CPUs + num_gpus: 0, + num_nodes: 1, + memory: "60g".to_string(), // Only 1/4 of memory + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let _result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + let actions = spec.actions.as_ref().unwrap(); + let action = &actions[0]; + + // CPU is limiting: 104/52 = 2 jobs per node + // Memory would allow: 240000/61440 = 3.9 = 3 jobs per node + // min(2, 3) = 2 jobs per node + // 4 jobs / 2 per node = 2 allocations + assert_eq!( + action.num_allocations, + Some(2), + "Should allocate 2 nodes for 4 CPU-heavy jobs (2 jobs per node, CPU-limited)" + ); +} + +/// Test that GPU constraints are considered in jobs-per-node calculation. +/// When GPUs are the limiting factor, we should allocate more nodes. +#[rstest] +fn test_generate_schedulers_gpu_constrained_allocation() { + // Create 8 jobs that need 2 GPUs each + // On Kestrel GPU nodes (128 CPUs, 360GB, 4 GPUs): + // - CPU-based: 128/32 = 4 jobs per node + // - Memory-based: 360,000MB / 92,160MB = 3.9 = 3 jobs per node + // - GPU-based: 4/2 = 2 jobs per node + // GPU should be the limiting factor, so we need 4 nodes for 8 jobs + let jobs: Vec = (0..8) + .map(|i| JobSpec { + name: format!("gpu_job_{}", i), + command: "python train.py".to_string(), + resource_requirements: Some("gpu_training".to_string()), + ..Default::default() + }) + .collect(); + + let mut spec = WorkflowSpec { + name: "gpu_test".to_string(), + description: Some("Test GPU-constrained allocation".to_string()), + jobs, + resource_requirements: Some(vec![ResourceRequirementsSpec { + name: "gpu_training".to_string(), + num_cpus: 32, // 1/4 of node CPUs + num_gpus: 2, // Half the GPUs - this should be limiting + num_nodes: 1, + memory: "90g".to_string(), // ~1/4 of node memory + runtime: "PT1H".to_string(), + }]), + ..Default::default() + }; + + let profile = kestrel_profile(); + let result = + generate_schedulers_for_workflow(&mut spec, &profile, "testaccount", false, true, false) + .unwrap(); + + assert_eq!(result.scheduler_count, 1); + assert_eq!(result.action_count, 1); + + let actions = spec.actions.as_ref().unwrap(); + let action = &actions[0]; + + // GPU is limiting: 4/2 = 2 jobs per node + // CPU would allow: 128/32 = 4 jobs per node + // Memory would allow: 360000/92160 = 3.9 = 3 jobs per node + // min(4, 3, 2) = 2 jobs per node + // 8 jobs / 2 per node = 4 allocations + assert_eq!( + action.num_allocations, + Some(4), + "Should allocate 4 nodes for 8 GPU jobs (2 jobs per node due to GPU constraint)" + ); +} diff --git a/tests/test_parameterization.rs b/tests/test_parameterization.rs index b514fb5b..69714701 100644 --- a/tests/test_parameterization.rs +++ b/tests/test_parameterization.rs @@ -127,17 +127,19 @@ fn test_kdl_example_file_hundred_jobs() { let mut spec = WorkflowSpec::from_spec_file(&path).expect("Failed to parse KDL example file"); assert_eq!(spec.name, "hundred_jobs_parameterized"); - assert_eq!(spec.jobs.len(), 1); + // 2 jobs before expansion: parameterized job template + postprocess + assert_eq!(spec.jobs.len(), 2); assert!(spec.jobs[0].parameters.is_some()); // Expand parameters spec.expand_parameters() .expect("Failed to expand parameters"); - // Should have 100 jobs after expansion - assert_eq!(spec.jobs.len(), 100); + // Should have 101 jobs after expansion: 100 parameterized + 1 postprocess + assert_eq!(spec.jobs.len(), 101); assert_eq!(spec.jobs[0].name, "job_001"); assert_eq!(spec.jobs[99].name, "job_100"); + assert_eq!(spec.jobs[100].name, "postprocess"); } #[rstest] diff --git a/tests/test_workflow_actions.rs b/tests/test_workflow_actions.rs index 309fb72f..97505ef4 100644 --- a/tests/test_workflow_actions.rs +++ b/tests/test_workflow_actions.rs @@ -1,9 +1,14 @@ mod common; +use std::thread; +use std::time::Duration; + use common::{ServerProcess, create_test_workflow, start_server}; use rstest::rstest; use serde_json::json; use torc::client::default_api; +use torc::client::workflow_manager::WorkflowManager; +use torc::config::TorcConfig; use torc::models::JobModel; /// Helper function to create a test job @@ -516,11 +521,6 @@ fn test_action_status_lifecycle(start_server: &ServerProcess) { /// - Expected: The workflow action should trigger again when postprocess_job becomes ready #[rstest] fn test_action_executed_flag_reset_on_reinitialize(start_server: &ServerProcess) { - use std::thread; - use std::time::Duration; - use torc::client::workflow_manager::WorkflowManager; - use torc::config::TorcConfig; - let config = &start_server.config; let workflow = create_test_workflow(config, "action_reinit_test_workflow"); let workflow_id = workflow.id.unwrap(); diff --git a/tests/test_workflow_manager.rs b/tests/test_workflow_manager.rs index 857a5a87..d1ac5725 100644 --- a/tests/test_workflow_manager.rs +++ b/tests/test_workflow_manager.rs @@ -6,8 +6,8 @@ mod common; use common::{ - ServerProcess, create_test_compute_node, create_test_file, create_test_user_data, - create_test_workflow_advanced, start_server, + ServerProcess, create_diamond_workflow, create_test_compute_node, create_test_file, + create_test_user_data, create_test_workflow_advanced, start_server, }; use rstest::rstest; use std::fs; @@ -2407,11 +2407,6 @@ fn test_user_data_dependency_chain(start_server: &ServerProcess) { /// 5. Verify work1 is set to Ready (not Blocked), since preprocess is already Completed #[rstest] fn test_reinitialize_with_file_change_depends_on_complete_job(start_server: &ServerProcess) { - use common::create_diamond_workflow; - use std::fs; - use std::thread; - use std::time::Duration; - let config = start_server.config.clone(); // Create a temporary directory for workflow files diff --git a/tests/test_workflow_spec.rs b/tests/test_workflow_spec.rs index fee3ab84..ec425d57 100644 --- a/tests/test_workflow_spec.rs +++ b/tests/test_workflow_spec.rs @@ -1,9 +1,11 @@ mod common; +use std::fs; +use std::path::PathBuf; + use common::{ServerProcess, start_server}; use rstest::rstest; use serde_json; -use std::fs; use tempfile::NamedTempFile; use torc::client::default_api; use torc::client::workflow_spec::{ @@ -1797,8 +1799,6 @@ fn test_create_workflow_with_mixed_exact_and_regex_dependencies(start_server: &S #[rstest] fn test_create_workflows_from_all_example_files(start_server: &ServerProcess) { - use std::path::PathBuf; - // Define the subdirectories to check let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples"); let subdirs = vec!["yaml", "json", "kdl"]; @@ -3061,3 +3061,876 @@ fn test_validate_spec_no_warning_with_scheduler_assignments() { result.warnings ); } + +// ============================================================================= +// Subgraph Workflow Tests +// ============================================================================= + +/// Test that the subgraph workflow examples parse correctly in all formats +#[test] +fn test_subgraph_workflow_parses_in_all_formats() { + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + + // Test each format with slurm schedulers + let formats = vec![ + "subgraphs_workflow.json", + "subgraphs_workflow.json5", + "subgraphs_workflow.yaml", + "subgraphs_workflow.kdl", + ]; + + for format in formats { + let spec_file = examples_dir.join(format); + if !spec_file.exists() { + eprintln!("Skipping {} (file not found)", format); + continue; + } + + let mut spec = WorkflowSpec::from_spec_file(&spec_file) + .unwrap_or_else(|e| panic!("Failed to parse {}: {}", format, e)); + + assert_eq!( + spec.name, "two_subgraph_pipeline", + "Workflow name mismatch for {}", + format + ); + + // Expand parameters to get the full job list + spec.expand_parameters() + .unwrap_or_else(|e| panic!("Failed to expand parameters for {}: {}", format, e)); + + assert_eq!(spec.jobs.len(), 15, "Expected 15 jobs for {}", format); + assert!( + spec.slurm_schedulers.is_some(), + "Expected slurm_schedulers for {}", + format + ); + assert!(spec.actions.is_some(), "Expected actions for {}", format); + + eprintln!( + "✓ {} parses correctly with {} jobs", + format, + spec.jobs.len() + ); + } +} + +/// Test that the no_slurm versions parse correctly +#[test] +fn test_subgraph_workflow_no_slurm_parses_in_all_formats() { + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + + let formats = vec![ + "subgraphs_workflow_no_slurm.json", + "subgraphs_workflow_no_slurm.json5", + "subgraphs_workflow_no_slurm.yaml", + "subgraphs_workflow_no_slurm.kdl", + ]; + + for format in formats { + let spec_file = examples_dir.join(format); + if !spec_file.exists() { + eprintln!("Skipping {} (file not found)", format); + continue; + } + + let mut spec = WorkflowSpec::from_spec_file(&spec_file) + .unwrap_or_else(|e| panic!("Failed to parse {}: {}", format, e)); + + assert_eq!( + spec.name, "two_subgraph_pipeline", + "Workflow name mismatch for {}", + format + ); + + // Expand parameters to get the full job list + spec.expand_parameters() + .unwrap_or_else(|e| panic!("Failed to expand parameters for {}: {}", format, e)); + + assert_eq!(spec.jobs.len(), 15, "Expected 15 jobs for {}", format); + assert!( + spec.slurm_schedulers.is_none(), + "Expected no slurm_schedulers for {}", + format + ); + assert!(spec.actions.is_none(), "Expected no actions for {}", format); + + eprintln!( + "✓ {} parses correctly with {} jobs (no slurm)", + format, + spec.jobs.len() + ); + } +} + +/// Test that execution plans have 4 stages for both slurm and no_slurm versions +#[test] +fn test_subgraph_workflow_execution_plan_has_4_stages() { + use torc::client::execution_plan::ExecutionPlan; + + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + + // Test with slurm schedulers + let slurm_spec_file = examples_dir.join("subgraphs_workflow.yaml"); + if slurm_spec_file.exists() { + let mut spec = + WorkflowSpec::from_spec_file(&slurm_spec_file).expect("Failed to parse slurm workflow"); + spec.expand_parameters() + .expect("Failed to expand parameters"); + + let plan = ExecutionPlan::from_spec(&spec).expect("Failed to build execution plan"); + + // With the DAG structure, we have 6 events for the subgraph workflow: + // 1. start event (prep_a, prep_b) + // 2. prep_a completes -> work_a_1..5 + // 3. prep_b completes -> work_b_1..5 + // 4. work_a_* complete -> post_a + // 5. work_b_* complete -> post_b + // 6. post_a, post_b complete -> final + assert_eq!( + plan.events.len(), + 6, + "Expected 6 events for slurm workflow (DAG structure), got {}", + plan.events.len() + ); + + // Verify there's exactly one root event (start) + assert_eq!(plan.root_events.len(), 1, "Should have 1 root event"); + + // Verify the start event has workflow start trigger + let start_event = plan.events.get(&plan.root_events[0]).unwrap(); + assert!( + start_event.trigger_description.contains("Workflow Start"), + "Root event should be workflow start" + ); + + eprintln!( + "✓ Slurm workflow has {} events (DAG structure)", + plan.events.len() + ); + } + + // Test without slurm schedulers + let no_slurm_spec_file = examples_dir.join("subgraphs_workflow_no_slurm.yaml"); + if no_slurm_spec_file.exists() { + let mut spec = WorkflowSpec::from_spec_file(&no_slurm_spec_file) + .expect("Failed to parse no_slurm workflow"); + spec.expand_parameters() + .expect("Failed to expand parameters"); + + let plan = ExecutionPlan::from_spec(&spec).expect("Failed to build execution plan"); + + assert_eq!( + plan.events.len(), + 6, + "Expected 6 events for no_slurm workflow (DAG structure), got {}", + plan.events.len() + ); + + eprintln!( + "✓ No-slurm workflow has {} events (DAG structure)", + plan.events.len() + ); + } +} + +/// Test that slurm and no_slurm versions produce the same execution plan structure +#[test] +fn test_subgraph_workflow_slurm_and_no_slurm_have_same_events() { + use torc::client::execution_plan::ExecutionPlan; + + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + + let slurm_spec_file = examples_dir.join("subgraphs_workflow.yaml"); + let no_slurm_spec_file = examples_dir.join("subgraphs_workflow_no_slurm.yaml"); + + if !slurm_spec_file.exists() || !no_slurm_spec_file.exists() { + eprintln!("Skipping test - example files not found"); + return; + } + + // Parse and expand both specs + let mut slurm_spec = + WorkflowSpec::from_spec_file(&slurm_spec_file).expect("Failed to parse slurm spec"); + slurm_spec + .expand_parameters() + .expect("Failed to expand slurm parameters"); + + let mut no_slurm_spec = + WorkflowSpec::from_spec_file(&no_slurm_spec_file).expect("Failed to parse no_slurm spec"); + no_slurm_spec + .expand_parameters() + .expect("Failed to expand no_slurm parameters"); + + // Build execution plans + let slurm_plan = + ExecutionPlan::from_spec(&slurm_spec).expect("Failed to build slurm execution plan"); + let no_slurm_plan = + ExecutionPlan::from_spec(&no_slurm_spec).expect("Failed to build no_slurm execution plan"); + + // Verify same number of events + assert_eq!( + slurm_plan.events.len(), + no_slurm_plan.events.len(), + "Slurm and no_slurm workflows should have the same number of events" + ); + + // Collect all jobs becoming ready from both plans + let mut slurm_all_jobs: Vec = slurm_plan + .events + .values() + .flat_map(|e| e.jobs_becoming_ready.clone()) + .collect(); + slurm_all_jobs.sort(); + + let mut no_slurm_all_jobs: Vec = no_slurm_plan + .events + .values() + .flat_map(|e| e.jobs_becoming_ready.clone()) + .collect(); + no_slurm_all_jobs.sort(); + + assert_eq!( + slurm_all_jobs, no_slurm_all_jobs, + "Both plans should make the same jobs ready" + ); + + eprintln!( + "✓ Both versions have {} events with {} total jobs", + slurm_plan.events.len(), + slurm_all_jobs.len() + ); +} + +/// Test that all format pairs (slurm vs no_slurm) produce identical execution plans +#[test] +fn test_subgraph_workflow_all_formats_produce_same_execution_plan() { + use torc::client::execution_plan::ExecutionPlan; + + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + + // Compare JSON vs YAML (slurm versions) + let json_file = examples_dir.join("subgraphs_workflow.json"); + let yaml_file = examples_dir.join("subgraphs_workflow.yaml"); + + if json_file.exists() && yaml_file.exists() { + let mut json_spec = WorkflowSpec::from_spec_file(&json_file).expect("Failed to parse JSON"); + json_spec + .expand_parameters() + .expect("Failed to expand JSON parameters"); + + let mut yaml_spec = WorkflowSpec::from_spec_file(&yaml_file).expect("Failed to parse YAML"); + yaml_spec + .expand_parameters() + .expect("Failed to expand YAML parameters"); + + let json_plan = ExecutionPlan::from_spec(&json_spec).expect("Failed to build JSON plan"); + let yaml_plan = ExecutionPlan::from_spec(&yaml_spec).expect("Failed to build YAML plan"); + + assert_eq!( + json_plan.events.len(), + yaml_plan.events.len(), + "JSON and YAML should have same number of events" + ); + + // Verify total job counts match + let json_job_count: usize = json_plan + .events + .values() + .map(|e| e.jobs_becoming_ready.len()) + .sum(); + let yaml_job_count: usize = yaml_plan + .events + .values() + .map(|e| e.jobs_becoming_ready.len()) + .sum(); + + assert_eq!( + json_job_count, yaml_job_count, + "Total job counts should match between JSON and YAML" + ); + + eprintln!("✓ JSON and YAML produce identical execution plans"); + } + + // Compare no_slurm versions + let json_no_slurm = examples_dir.join("subgraphs_workflow_no_slurm.json"); + let yaml_no_slurm = examples_dir.join("subgraphs_workflow_no_slurm.yaml"); + + if json_no_slurm.exists() && yaml_no_slurm.exists() { + let mut json_spec = + WorkflowSpec::from_spec_file(&json_no_slurm).expect("Failed to parse JSON no_slurm"); + json_spec + .expand_parameters() + .expect("Failed to expand parameters"); + + let mut yaml_spec = + WorkflowSpec::from_spec_file(&yaml_no_slurm).expect("Failed to parse YAML no_slurm"); + yaml_spec + .expand_parameters() + .expect("Failed to expand parameters"); + + let json_plan = ExecutionPlan::from_spec(&json_spec).expect("Failed to build plan"); + let yaml_plan = ExecutionPlan::from_spec(&yaml_spec).expect("Failed to build plan"); + + assert_eq!( + json_plan.events.len(), + yaml_plan.events.len(), + "No_slurm JSON and YAML should have same number of events" + ); + + eprintln!("✓ No_slurm JSON and YAML produce identical execution plans"); + } +} + +/// Test subgraph workflow job structure +#[test] +fn test_subgraph_workflow_job_structure() { + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + let spec_file = examples_dir.join("subgraphs_workflow_no_slurm.yaml"); + + if !spec_file.exists() { + eprintln!("Skipping test - example file not found"); + return; + } + + let mut spec = WorkflowSpec::from_spec_file(&spec_file).expect("Failed to parse spec"); + spec.expand_parameters() + .expect("Failed to expand parameters"); + + // Verify job counts after expansion + // prep_a, prep_b = 2 + // work_a_1..5, work_b_1..5 = 10 + // post_a, post_b = 2 + // final = 1 + // Total = 15 + assert_eq!(spec.jobs.len(), 15, "Expected 15 jobs after expansion"); + + // Verify prep jobs have no dependencies + let prep_a = spec.jobs.iter().find(|j| j.name == "prep_a"); + assert!(prep_a.is_some(), "prep_a job not found"); + assert!( + prep_a.unwrap().depends_on.is_none() + || prep_a.unwrap().depends_on.as_ref().unwrap().is_empty(), + "prep_a should have no explicit dependencies" + ); + + // Verify work jobs have input_files + let work_a_1 = spec.jobs.iter().find(|j| j.name == "work_a_1"); + assert!(work_a_1.is_some(), "work_a_1 job not found"); + assert!( + work_a_1.unwrap().input_files.is_some(), + "work_a_1 should have input_files" + ); + + // Verify post jobs have input_files from work jobs + let post_a = spec.jobs.iter().find(|j| j.name == "post_a"); + assert!(post_a.is_some(), "post_a job not found"); + let post_a_inputs = post_a.unwrap().input_files.as_ref().unwrap(); + assert_eq!( + post_a_inputs.len(), + 5, + "post_a should have 5 input files from work_a jobs" + ); + + // Verify final job has input_files from both post jobs + let final_job = spec.jobs.iter().find(|j| j.name == "final"); + assert!(final_job.is_some(), "final job not found"); + let final_inputs = final_job.unwrap().input_files.as_ref().unwrap(); + assert_eq!( + final_inputs.len(), + 2, + "final should have 2 input files (post_a_out, post_b_out)" + ); + + eprintln!("✓ Job structure verified: 15 jobs with correct dependencies"); +} + +/// Test that subgraph workflow resource requirements are preserved +#[test] +fn test_subgraph_workflow_resource_requirements() { + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + let spec_file = examples_dir.join("subgraphs_workflow_no_slurm.yaml"); + + if !spec_file.exists() { + eprintln!("Skipping test - example file not found"); + return; + } + + let spec = WorkflowSpec::from_spec_file(&spec_file).expect("Failed to parse spec"); + + let resource_reqs = spec + .resource_requirements + .as_ref() + .expect("Missing resource_requirements"); + assert_eq!( + resource_reqs.len(), + 5, + "Expected 5 resource requirement definitions" + ); + + // Verify specific resource requirements + let small = resource_reqs.iter().find(|r| r.name == "small"); + assert!(small.is_some(), "small resource requirement not found"); + assert_eq!(small.unwrap().num_cpus, 1); + + let work_large = resource_reqs.iter().find(|r| r.name == "work_large"); + assert!( + work_large.is_some(), + "work_large resource requirement not found" + ); + assert_eq!(work_large.unwrap().num_cpus, 8); + assert_eq!(work_large.unwrap().memory, "32g"); + + let work_gpu = resource_reqs.iter().find(|r| r.name == "work_gpu"); + assert!( + work_gpu.is_some(), + "work_gpu resource requirement not found" + ); + assert_eq!(work_gpu.unwrap().num_gpus, 1); + + eprintln!("✓ Resource requirements verified"); +} + +/// Integration test: create subgraph workflows on server +#[rstest] +fn test_create_subgraph_workflows_from_examples(start_server: &ServerProcess) { + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + + // Test both slurm and no_slurm YAML versions + let test_files = vec![ + ("subgraphs_workflow.yaml", true), // has schedulers + ("subgraphs_workflow_no_slurm.yaml", false), // no schedulers + ]; + + for (filename, has_schedulers) in test_files { + let spec_file = examples_dir.join(filename); + if !spec_file.exists() { + eprintln!("Skipping {} (file not found)", filename); + continue; + } + + let workflow_id = WorkflowSpec::create_workflow_from_spec( + &start_server.config, + &spec_file, + "test_user", + false, + true, // skip_checks - we don't have a real Slurm environment + ) + .unwrap_or_else(|e| panic!("Failed to create workflow from {}: {}", filename, e)); + + assert!(workflow_id > 0, "Invalid workflow ID for {}", filename); + + // Verify the workflow was created + let workflow = default_api::get_workflow(&start_server.config, workflow_id) + .expect("Failed to get workflow"); + assert_eq!(workflow.name, "two_subgraph_pipeline"); + + // Verify job count + let jobs = default_api::list_jobs( + &start_server.config, + workflow_id, + None, + None, + None, + None, + None, + None, + None, + None, + ) + .expect("Failed to list jobs"); + + let job_count = jobs.items.as_ref().map(|j| j.len()).unwrap_or(0); + assert_eq!( + job_count, 15, + "Expected 15 jobs for {}, got {}", + filename, job_count + ); + + // Verify schedulers if present + if has_schedulers { + let response = default_api::list_slurm_schedulers( + &start_server.config, + workflow_id, + Some(0), // offset + Some(50), // limit + None, // sort_by + None, // reverse_sort + None, // name filter + None, // account filter + None, // gres filter + None, // mem filter + None, // nodes filter + None, // partition filter + None, // qos filter + None, // tmp filter + None, // walltime filter + ) + .expect("Failed to list schedulers"); + let sched_count = response.items.unwrap_or_default().len(); + assert!( + sched_count > 0, + "Expected schedulers for {}, got {}", + filename, + sched_count + ); + eprintln!( + "✓ {} created with {} jobs and {} schedulers", + filename, job_count, sched_count + ); + } else { + eprintln!( + "✓ {} created with {} jobs (no schedulers)", + filename, job_count + ); + } + + // Clean up + let _ = default_api::delete_workflow(&start_server.config, workflow_id, None); + } +} + +/// Test that generate_schedulers_for_workflow assigns correct trigger types +/// Jobs without dependencies get on_workflow_start, jobs with dependencies get on_jobs_ready +#[test] +fn test_subgraph_workflow_generated_actions_have_correct_triggers() { + use torc::client::commands::slurm::generate_schedulers_for_workflow; + use torc::client::hpc::kestrel::kestrel_profile; + + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + let no_slurm_spec_file = examples_dir.join("subgraphs_workflow_no_slurm.yaml"); + + if !no_slurm_spec_file.exists() { + eprintln!("Skipping test - example file not found"); + return; + } + + // Parse the no_slurm spec + let mut spec = + WorkflowSpec::from_spec_file(&no_slurm_spec_file).expect("Failed to parse no_slurm spec"); + + // Generate schedulers + let profile = kestrel_profile(); + let result = generate_schedulers_for_workflow( + &mut spec, + &profile, + "testaccount", + false, // not single allocation + true, // add actions + false, // don't force + ) + .expect("Failed to generate schedulers"); + + eprintln!( + "Generated {} schedulers and {} actions", + result.scheduler_count, result.action_count + ); + + let actions = spec + .actions + .as_ref() + .expect("Should have generated actions"); + + // Build map of scheduler -> trigger_type + let scheduler_triggers: std::collections::HashMap = actions + .iter() + .filter_map(|a| a.scheduler.clone().map(|s| (s, a.trigger_type.clone()))) + .collect(); + + // Verify each job has the correct trigger type based on dependencies + // prep_a, prep_b: no dependencies -> on_workflow_start + // work_*: depend on prep_* outputs -> on_jobs_ready + // post_*: depend on work_* outputs -> on_jobs_ready + // final: depends on post_* outputs -> on_jobs_ready + for job in &spec.jobs { + let sched = job + .scheduler + .as_ref() + .expect(&format!("Job {} should have scheduler assigned", job.name)); + let trigger = scheduler_triggers + .get(sched) + .expect(&format!("Scheduler {} should have action", sched)); + + let expected_trigger = if job.name == "prep_a" || job.name == "prep_b" { + "on_workflow_start" + } else { + "on_jobs_ready" + }; + + assert_eq!( + trigger, expected_trigger, + "Job {} (scheduler {}) should have trigger {}, got {}", + job.name, sched, expected_trigger, trigger + ); + } + + eprintln!("✓ All jobs have correct trigger types"); +} + +/// Test that execution plan from database correctly computes dependencies from file relationships +#[test] +fn test_subgraph_workflow_execution_plan_from_database() { + use torc::client::execution_plan::ExecutionPlan; + + let start_server = common::start_server(); + + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + let spec_file = examples_dir.join("subgraphs_workflow.yaml"); + + if !spec_file.exists() { + eprintln!("Skipping test - example file not found"); + return; + } + + // Create workflow on server + let workflow_id = WorkflowSpec::create_workflow_from_spec( + &start_server.config, + &spec_file, + "test_user", + false, + true, // skip_checks + ) + .expect("Failed to create workflow"); + + // Fetch workflow, jobs (with relationships), and actions from server + let workflow = default_api::get_workflow(&start_server.config, workflow_id) + .expect("Failed to get workflow"); + + let jobs = default_api::list_jobs( + &start_server.config, + workflow_id, + None, + None, + None, + None, + Some(10000), + None, + None, + Some(true), // include_relationships - this is key! + ) + .expect("Failed to list jobs") + .items + .unwrap_or_default(); + + let actions = default_api::get_workflow_actions(&start_server.config, workflow_id) + .expect("Failed to get actions"); + + let slurm_schedulers = default_api::list_slurm_schedulers( + &start_server.config, + workflow_id, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + .map(|r| r.items.unwrap_or_default()) + .unwrap_or_default(); + + // Build execution plan from database models + let plan = ExecutionPlan::from_database_models(&workflow, &jobs, &actions, &slurm_schedulers) + .expect("Failed to build execution plan from database"); + + // With the DAG structure, we have 6 events: + // 1. start event (prep_a, prep_b) + // 2. prep_a completes -> work_a_1..5 + // 3. prep_b completes -> work_b_1..5 + // 4. work_a_* complete -> post_a + // 5. work_b_* complete -> post_b + // 6. post_a, post_b complete -> final + assert_eq!( + plan.events.len(), + 6, + "Expected 6 events from database (DAG structure), got {}", + plan.events.len() + ); + + // Find the start event + let start_event = plan + .events + .get(&plan.root_events[0]) + .expect("Start event not found"); + + // Verify start event has 2 jobs (prep_a, prep_b) + assert_eq!( + start_event.jobs_becoming_ready.len(), + 2, + "Start event should have 2 jobs, got {} - {:?}", + start_event.jobs_becoming_ready.len(), + start_event.jobs_becoming_ready + ); + assert!( + start_event + .jobs_becoming_ready + .contains(&"prep_a".to_string()), + "Start event should contain prep_a" + ); + assert!( + start_event + .jobs_becoming_ready + .contains(&"prep_b".to_string()), + "Start event should contain prep_b" + ); + + // Collect all jobs becoming ready across all events + let all_jobs: Vec = plan + .events + .values() + .flat_map(|e| e.jobs_becoming_ready.clone()) + .collect(); + + // Should have 15 total jobs + assert_eq!( + all_jobs.len(), + 15, + "Total jobs across all events should be 15, got {}", + all_jobs.len() + ); + + // Verify final job is in a leaf event + let leaf_event = plan + .events + .get(&plan.leaf_events[0]) + .expect("Leaf event not found"); + assert!( + leaf_event + .jobs_becoming_ready + .contains(&"final".to_string()), + "Leaf event should contain final job" + ); + + // Clean up + let _ = default_api::delete_workflow(&start_server.config, workflow_id, None); + + eprintln!("✓ Execution plan from database has correct 6 events (DAG structure)"); +} + +/// Test that execution plan from spec matches execution plan from database +#[test] +fn test_subgraph_workflow_execution_plan_spec_vs_database() { + use torc::client::execution_plan::ExecutionPlan; + + let start_server = common::start_server(); + + let examples_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("examples/subgraphs"); + let spec_file = examples_dir.join("subgraphs_workflow.yaml"); + + if !spec_file.exists() { + eprintln!("Skipping test - example file not found"); + return; + } + + // Build execution plan directly from spec + let mut spec = WorkflowSpec::from_spec_file(&spec_file).expect("Failed to parse spec"); + spec.expand_parameters() + .expect("Failed to expand parameters"); + let spec_plan = ExecutionPlan::from_spec(&spec).expect("Failed to build plan from spec"); + + // Create workflow on server + let workflow_id = WorkflowSpec::create_workflow_from_spec( + &start_server.config, + &spec_file, + "test_user", + false, + true, // skip_checks + ) + .expect("Failed to create workflow"); + + // Fetch workflow, jobs (with relationships), and actions from server + let workflow = default_api::get_workflow(&start_server.config, workflow_id) + .expect("Failed to get workflow"); + + let jobs = default_api::list_jobs( + &start_server.config, + workflow_id, + None, + None, + None, + None, + Some(10000), + None, + None, + Some(true), // include_relationships + ) + .expect("Failed to list jobs") + .items + .unwrap_or_default(); + + let actions = default_api::get_workflow_actions(&start_server.config, workflow_id) + .expect("Failed to get actions"); + + let slurm_schedulers = default_api::list_slurm_schedulers( + &start_server.config, + workflow_id, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + .map(|r| r.items.unwrap_or_default()) + .unwrap_or_default(); + + // Build execution plan from database models + let db_plan = + ExecutionPlan::from_database_models(&workflow, &jobs, &actions, &slurm_schedulers) + .expect("Failed to build plan from database"); + + // Compare event counts + assert_eq!( + spec_plan.events.len(), + db_plan.events.len(), + "Spec and database should have same number of events" + ); + + // Collect all jobs from both plans and compare + let mut spec_all_jobs: Vec = spec_plan + .events + .values() + .flat_map(|e| e.jobs_becoming_ready.clone()) + .collect(); + spec_all_jobs.sort(); + + let mut db_all_jobs: Vec = db_plan + .events + .values() + .flat_map(|e| e.jobs_becoming_ready.clone()) + .collect(); + db_all_jobs.sort(); + + assert_eq!( + spec_all_jobs, + db_all_jobs, + "Spec and database should have same total jobs.\nSpec: {:?}\nDB: {:?}", + spec_all_jobs.len(), + db_all_jobs.len() + ); + + eprintln!( + "✓ Both plans have {} events with {} total jobs", + spec_plan.events.len(), + spec_all_jobs.len() + ); + + // Clean up + let _ = default_api::delete_workflow(&start_server.config, workflow_id, None); + + eprintln!("✓ Execution plan from spec matches execution plan from database"); +} diff --git a/torc-dash/src/main.rs b/torc-dash/src/main.rs index f540f852..115017d7 100644 --- a/torc-dash/src/main.rs +++ b/torc-dash/src/main.rs @@ -20,6 +20,8 @@ use axum::{ use clap::Parser; use rust_embed::Embed; use serde::{Deserialize, Serialize}; +use std::path::Path as FsPath; +use std::process::Command as StdCommand; use std::sync::Arc; use tokio::io::{AsyncBufReadExt, BufReader}; use tokio::process::Command; @@ -360,6 +362,9 @@ async fn main() -> Result<()> { post(cli_slurm_parse_logs_handler), ) .route("/api/cli/slurm-sacct", post(cli_slurm_sacct_handler)) + // Slurm scheduler generation endpoints + .route("/api/cli/slurm-generate", post(cli_slurm_generate_handler)) + .route("/api/cli/hpc-profiles", get(cli_hpc_profiles_handler)) // Server management endpoints .route("/api/server/start", post(server_start_handler)) .route("/api/server/stop", post(server_stop_handler)) @@ -736,24 +741,71 @@ async fn cli_reset_status_handler( Json(result) } -/// Get execution plan for a workflow (includes scheduler allocations per stage) +#[derive(Serialize)] +struct ExecutionPlanResponse { + success: bool, + /// Parsed execution plan data + data: Option, + error: Option, +} + +/// Get the execution plan for a workflow async fn cli_execution_plan_handler( State(state): State>, Json(req): Json, ) -> impl IntoResponse { - let result = run_torc_command( - &state.torc_bin, - &[ - "-f", - "json", - "workflows", - "execution-plan", - &req.workflow_id, - ], - &state.api_url, - ) - .await; - Json(result) + let args = vec![ + "-f", + "json", + "workflows", + "execution-plan", + &req.workflow_id, + ]; + + info!("Running: {} {}", state.torc_bin, args.join(" ")); + + let output = Command::new(&state.torc_bin) + .args(&args) + .env("TORC_API_URL", &state.api_url) + .output() + .await; + + match output { + Ok(output) => { + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + + if !output.status.success() { + return Json(ExecutionPlanResponse { + success: false, + data: None, + error: Some(format!("Command failed: {}", stderr)), + }); + } + + // Parse the JSON output + match serde_json::from_str::(&stdout) { + Ok(data) => Json(ExecutionPlanResponse { + success: true, + data: Some(data), + error: None, + }), + Err(e) => Json(ExecutionPlanResponse { + success: false, + data: None, + error: Some(format!( + "Failed to parse JSON output: {}. Output: {}", + e, stdout + )), + }), + } + } + Err(e) => Json(ExecutionPlanResponse { + success: false, + data: None, + error: Some(format!("Failed to execute command: {}", e)), + }), + } } /// Streaming workflow run handler using Server-Sent Events @@ -975,9 +1027,7 @@ async fn cli_run_stream_handler( /// Read file contents from filesystem async fn cli_read_file_handler(Json(req): Json) -> impl IntoResponse { - use std::path::Path; - - let path = Path::new(&req.path); + let path = FsPath::new(&req.path); let exists = path.exists(); if !exists { @@ -1215,9 +1265,7 @@ async fn cli_plot_resources_handler( async fn cli_list_resource_dbs_handler( Json(req): Json, ) -> impl IntoResponse { - use std::path::Path; - - let base_path = Path::new(&req.base_dir); + let base_path = FsPath::new(&req.base_dir); if !base_path.exists() { return Json(ListResourceDbsResponse { @@ -1462,6 +1510,259 @@ async fn cli_slurm_sacct_handler( } } +// ============== Slurm Scheduler Generation Handlers ============== + +#[derive(Deserialize)] +struct SlurmGenerateRequest { + /// Workflow specification (as JSON) + spec: serde_json::Value, + /// Slurm account name + account: String, + /// HPC profile name (optional, auto-detected if not provided) + #[serde(default)] + profile: Option, +} + +#[derive(Serialize)] +struct SlurmGenerateResponse { + success: bool, + /// Generated slurm_schedulers array + schedulers: Option>, + /// Generated actions array + actions: Option>, + /// Profile that was used + profile_used: Option, + error: Option, +} + +#[derive(Serialize)] +struct HpcProfileInfo { + name: String, + display_name: String, + description: String, + is_detected: bool, +} + +#[derive(Serialize)] +struct HpcProfilesResponse { + success: bool, + profiles: Vec, + detected_profile: Option, + error: Option, +} + +/// Generate Slurm schedulers from a workflow specification +async fn cli_slurm_generate_handler( + State(state): State>, + Json(req): Json, +) -> impl IntoResponse { + // Write the spec to a temp file + let unique_id = uuid::Uuid::new_v4(); + let temp_path = format!("/tmp/torc_spec_{}.json", unique_id); + + let spec_str = match serde_json::to_string_pretty(&req.spec) { + Ok(s) => s, + Err(e) => { + return Json(SlurmGenerateResponse { + success: false, + schedulers: None, + actions: None, + profile_used: None, + error: Some(format!("Failed to serialize spec: {}", e)), + }); + } + }; + + if let Err(e) = tokio::fs::write(&temp_path, &spec_str).await { + return Json(SlurmGenerateResponse { + success: false, + schedulers: None, + actions: None, + profile_used: None, + error: Some(format!("Failed to write temp file: {}", e)), + }); + } + + // Build command arguments + let mut args = vec![ + "-f".to_string(), + "json".to_string(), + "slurm".to_string(), + "generate".to_string(), + "--account".to_string(), + req.account.clone(), + ]; + + if let Some(ref profile) = req.profile { + args.push("--profile".to_string()); + args.push(profile.clone()); + } + + args.push(temp_path.clone()); + + let args_refs: Vec<&str> = args.iter().map(|s| s.as_str()).collect(); + info!("Running: {} {}", state.torc_bin, args_refs.join(" ")); + + let output = Command::new(&state.torc_bin) + .args(&args_refs) + .env("TORC_API_URL", &state.api_url) + .output() + .await; + + // Clean up temp file + let _ = tokio::fs::remove_file(&temp_path).await; + + match output { + Ok(output) => { + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + + if !output.status.success() { + return Json(SlurmGenerateResponse { + success: false, + schedulers: None, + actions: None, + profile_used: None, + error: Some(format!("Command failed: {}", stderr)), + }); + } + + // Parse the JSON output - it should contain the full spec with schedulers and actions + match serde_json::from_str::(&stdout) { + Ok(data) => { + let schedulers = data + .get("slurm_schedulers") + .and_then(|v| v.as_array()) + .cloned(); + let actions = data.get("actions").and_then(|v| v.as_array()).cloned(); + let profile_used = data + .get("profile_used") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + + Json(SlurmGenerateResponse { + success: true, + schedulers, + actions, + profile_used, + error: None, + }) + } + Err(e) => Json(SlurmGenerateResponse { + success: false, + schedulers: None, + actions: None, + profile_used: None, + error: Some(format!( + "Failed to parse JSON output: {}. Output: {}", + e, stdout + )), + }), + } + } + Err(e) => Json(SlurmGenerateResponse { + success: false, + schedulers: None, + actions: None, + profile_used: None, + error: Some(format!("Failed to execute command: {}", e)), + }), + } +} + +/// List available HPC profiles and detect current profile +async fn cli_hpc_profiles_handler(State(state): State>) -> impl IntoResponse { + // Run: torc hpc list -f json + let args = vec!["-f", "json", "hpc", "list"]; + + info!("Running: {} {}", state.torc_bin, args.join(" ")); + + let output = Command::new(&state.torc_bin) + .args(&args) + .env("TORC_API_URL", &state.api_url) + .output() + .await; + + match output { + Ok(output) => { + let stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + + if !output.status.success() { + return Json(HpcProfilesResponse { + success: false, + profiles: vec![], + detected_profile: None, + error: Some(format!("Command failed: {}", stderr)), + }); + } + + // Parse the JSON output - it's an array of profiles directly + match serde_json::from_str::>(&stdout) { + Ok(items) => { + let mut profiles = Vec::new(); + let mut detected_profile = None; + + for item in items { + let name = item + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let display_name = item + .get("display_name") + .and_then(|v| v.as_str()) + .unwrap_or(&name) + .to_string(); + let description = item + .get("description") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let is_detected = item + .get("detected") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if is_detected { + detected_profile = Some(name.clone()); + } + + profiles.push(HpcProfileInfo { + name, + display_name, + description, + is_detected, + }); + } + + Json(HpcProfilesResponse { + success: true, + profiles, + detected_profile, + error: None, + }) + } + Err(e) => Json(HpcProfilesResponse { + success: false, + profiles: vec![], + detected_profile: None, + error: Some(format!( + "Failed to parse JSON output: {}. Output: {}", + e, stdout + )), + }), + } + } + Err(e) => Json(HpcProfilesResponse { + success: false, + profiles: vec![], + detected_profile: None, + error: Some(format!("Failed to execute command: {}", e)), + }), + } +} + // ============== Server Management Handlers ============== #[derive(Deserialize)] @@ -1646,7 +1947,6 @@ async fn server_stop_handler(State(state): State>) -> impl IntoRes // Try to kill the process #[cfg(unix)] { - use std::process::Command as StdCommand; let result = StdCommand::new("kill").arg(pid.to_string()).status(); match result { @@ -1683,7 +1983,6 @@ async fn server_stop_handler(State(state): State>) -> impl IntoRes #[cfg(not(unix))] { // On Windows, try taskkill - use std::process::Command as StdCommand; let result = StdCommand::new("taskkill") .args(["/PID", &pid.to_string(), "/F"]) .status(); @@ -1722,7 +2021,6 @@ async fn server_status_handler(State(state): State>) -> impl IntoR if let Some(pid) = managed.pid { #[cfg(unix)] { - use std::process::Command as StdCommand; // Check if process exists by sending signal 0 if let Ok(status) = StdCommand::new("kill") .args(["-0", &pid.to_string()]) @@ -1735,7 +2033,6 @@ async fn server_status_handler(State(state): State>) -> impl IntoR #[cfg(not(unix))] { // On Windows, check with tasklist - use std::process::Command as StdCommand; if let Ok(output) = StdCommand::new("tasklist") .args(["/FI", &format!("PID eq {}", pid)]) .output() diff --git a/torc-dash/static/css/style.css b/torc-dash/static/css/style.css index 71a221ca..99bce5f4 100644 --- a/torc-dash/static/css/style.css +++ b/torc-dash/static/css/style.css @@ -1100,6 +1100,51 @@ textarea.text-input { font-size: 13px; } +/* Execution plan event types */ +.plan-stage.event-root { + border-left-color: var(--success-color); +} + +.plan-stage.event-root .plan-stage-number { + background: var(--success-color); +} + +.plan-stage.event-leaf { + border-left-color: var(--warning-color); +} + +.plan-stage.event-leaf .plan-stage-number { + background: var(--warning-color); +} + +.plan-section { + margin-bottom: 12px; +} + +.plan-section:last-child { + margin-bottom: 0; +} + +.plan-flow { + display: flex; + align-items: center; + gap: 8px; + padding-top: 8px; + border-top: 1px dashed var(--border-color); + margin-top: 8px; +} + +.plan-flow .flow-arrow { + color: var(--primary-color); + font-size: 16px; +} + +.plan-flow .flow-text { + font-size: 12px; + color: var(--text-secondary); + font-style: italic; +} + /* Debugging table styling */ .debug-table-row { cursor: pointer; @@ -1839,6 +1884,56 @@ textarea.text-input { margin-bottom: 16px; } +/* Auto-generate section in Schedulers step */ +.wizard-auto-generate-section { + background: var(--bg-tertiary); + border: 1px solid var(--border-color); + border-radius: var(--border-radius); + padding: 16px; + margin-bottom: 16px; +} + +.wizard-auto-generate-section h5 { + margin: 0 0 8px 0; + font-size: 14px; + color: var(--text-primary); +} + +.wizard-auto-generate-form { + display: flex; + flex-wrap: wrap; + gap: 12px; + align-items: flex-end; +} + +.wizard-auto-generate-form .form-group { + flex: 1; + min-width: 180px; + margin-bottom: 0; +} + +.wizard-auto-generate-form .form-group:last-child { + flex: 0 0 auto; +} + +.wizard-section-divider { + border: none; + border-top: 1px solid var(--border-color); + margin: 16px 0; +} + +.wizard-manual-schedulers-header h5 { + margin: 0 0 8px 0; + font-size: 14px; + color: var(--text-secondary); +} + +.wizard-warning { + color: var(--warning-color, #ffc107); + font-size: 13px; + margin-top: 8px; +} + .wizard-schedulers-list { display: flex; flex-direction: column; diff --git a/torc-dash/static/index.html b/torc-dash/static/index.html index 80ec7688..2b558d82 100644 --- a/torc-dash/static/index.html +++ b/torc-dash/static/index.html @@ -576,11 +576,13 @@

Jobs

+ +
+
-

Slurm Schedulers

+

Manual Schedulers

-

Schedulers define Slurm job allocation settings. Jobs can reference schedulers by name to run on HPC clusters.

diff --git a/torc-dash/static/js/api.js b/torc-dash/static/js/api.js index 7c6820b5..fb140c2d 100644 --- a/torc-dash/static/js/api.js +++ b/torc-dash/static/js/api.js @@ -334,11 +334,11 @@ class TorcAPI { } /** - * Get execution plan for a workflow (includes scheduler allocations per stage) + * Get execution plan for a workflow * @param {string} workflowId - Workflow ID - * @returns {object} CLI response with JSON execution plan in stdout + * @returns {object} Response with execution plan data */ - async cliExecutionPlan(workflowId) { + async getExecutionPlan(workflowId) { return this.cliRequest('/api/cli/execution-plan', { workflow_id: workflowId }); } @@ -388,6 +388,37 @@ class TorcAPI { }); } + // ==================== HPC Profiles & Slurm Generation ==================== + + /** + * Get available HPC profiles and detect current system + * @returns {object} Response with profiles array and detected_profile + */ + async getHpcProfiles() { + try { + const response = await fetch('/api/cli/hpc-profiles'); + return await response.json(); + } catch (error) { + console.error('HPC profiles error:', error); + return { success: false, profiles: [], error: error.message }; + } + } + + /** + * Generate Slurm schedulers from a workflow specification + * @param {object} spec - Workflow specification object + * @param {string} account - Slurm account name + * @param {string} [profile] - HPC profile name (auto-detected if not provided) + * @returns {object} Response with schedulers and actions arrays + */ + async generateSlurmSchedulers(spec, account, profile = null) { + const body = { spec, account }; + if (profile) { + body.profile = profile; + } + return this.cliRequest('/api/cli/slurm-generate', body); + } + // ==================== Server Management ==================== /** diff --git a/torc-dash/static/js/app-modals.js b/torc-dash/static/js/app-modals.js index eda2ce45..f7d13c06 100644 --- a/torc-dash/static/js/app-modals.js +++ b/torc-dash/static/js/app-modals.js @@ -327,22 +327,15 @@ Object.assign(TorcDashboard.prototype, { content.innerHTML = '
Loading execution plan...
'; try { - // Get execution plan from CLI (includes scheduler allocations) - const result = await api.cliExecutionPlan(workflowId); + // Get execution plan from the CLI + const response = await api.getExecutionPlan(workflowId); - if (!result.success) { - content.innerHTML = `
Error: ${this.escapeHtml(result.stderr || result.stdout || 'Unknown error')}
`; + if (!response.success) { + content.innerHTML = `
Error: ${response.error || 'Unknown error'}
`; return; } - // Parse the JSON output - let plan; - try { - plan = JSON.parse(result.stdout); - } catch (parseError) { - content.innerHTML = `
Invalid JSON response from server: ${this.escapeHtml(parseError.message)}
`; - return; - } + const plan = response.data; content.innerHTML = this.renderExecutionPlan(plan); } catch (error) { content.innerHTML = `
Error loading execution plan: ${this.escapeHtml(error.message)}
`; @@ -350,55 +343,137 @@ Object.assign(TorcDashboard.prototype, { }, renderExecutionPlan(plan) { - if (!plan.stages || plan.stages.length === 0) { - return '
No execution stages computed
'; + if (!plan || !plan.events || plan.events.length === 0) { + return '
No execution events computed
'; + } + + const events = plan.events; + const rootEvents = plan.root_events || []; + const leafEvents = plan.leaf_events || []; + + // Build event map for quick lookup + const eventMap = {}; + events.forEach(e => eventMap[e.id] = e); + + // Count total jobs + let totalJobs = 0; + events.forEach(e => totalJobs += (e.jobs_becoming_ready || []).length); + + // Render events in a topological order using BFS from roots + const visited = new Set(); + const orderedEvents = []; + const queue = [...rootEvents]; + + while (queue.length > 0) { + const eventId = queue.shift(); + if (visited.has(eventId)) continue; + visited.add(eventId); + + const event = eventMap[eventId]; + if (event) { + orderedEvents.push(event); + // Add unlocked events to queue + (event.unlocks_events || []).forEach(next => { + if (!visited.has(next)) { + queue.push(next); + } + }); + } } + // Also add any events not reachable from roots (shouldn't happen, but just in case) + events.forEach(e => { + if (!visited.has(e.id)) { + orderedEvents.push(e); + } + }); + return `
- Workflow: ${this.escapeHtml(plan.workflow_name || 'Unnamed')} | - Total Stages: ${plan.total_stages} | - Total Jobs: ${plan.total_jobs} + Total Events: ${events.length} | + Total Jobs: ${totalJobs}
- ${plan.stages.map(stage => ` -
-
-
${stage.stage_number}
-
${this.escapeHtml(stage.trigger)}
-
-
-
Jobs Becoming Ready (${stage.jobs_becoming_ready.length})
+
+ ${orderedEvents.map((event, idx) => this.renderExecutionEvent(event, idx, rootEvents, leafEvents)).join('')} +
+ `; + }, + + renderExecutionEvent(event, index, rootEvents, leafEvents) { + const isRoot = rootEvents.includes(event.id); + const isLeaf = leafEvents.includes(event.id); + const jobs = event.jobs_becoming_ready || []; + const schedulers = event.scheduler_allocations || []; + const unlocks = event.unlocks_events || []; + + // Determine event type icon and style + let eventIcon = '→'; + let eventClass = ''; + if (isRoot) { + eventIcon = '▶'; + eventClass = 'event-root'; + } else if (isLeaf) { + eventIcon = '◆'; + eventClass = 'event-leaf'; + } + + // Format trigger description + const triggerDesc = event.trigger_description || this.formatEventTrigger(event.trigger); + + return ` +
+
+
${eventIcon}
+
${this.escapeHtml(triggerDesc)}
+
+
+ ${schedulers.length > 0 ? ` +
+
Scheduler Allocations
+
    + ${schedulers.map(s => ` +
  • + ${this.escapeHtml(s.scheduler)} + (${this.escapeHtml(s.scheduler_type)}) + - ${s.num_allocations} allocation${s.num_allocations !== 1 ? 's' : ''} +
    • + `).join('')} +
+
+ ` : ''} +
+
Jobs Becoming Ready (${jobs.length})
    - ${stage.jobs_becoming_ready.slice(0, 10).map(jobName => ` -
  • ${this.escapeHtml(jobName)}
    • + ${jobs.slice(0, 10).map(job => ` +
  • ${this.escapeHtml(job)}
    • `).join('')} - ${stage.jobs_becoming_ready.length > 10 ? `
  • ... and ${stage.jobs_becoming_ready.length - 10} more
    • ` : ''} + ${jobs.length > 10 ? `
  • ... and ${jobs.length - 10} more
    • ` : ''}
- ${this.renderSchedulerAllocations(stage.scheduler_allocations)}
+ ${unlocks.length > 0 ? ` +
+ ↓ + unlocks ${unlocks.length} event${unlocks.length !== 1 ? 's' : ''} +
+ ` : ''}
- `).join('')} +
`; }, - renderSchedulerAllocations(allocations) { - if (!allocations || allocations.length === 0) { - return ''; + formatEventTrigger(trigger) { + if (!trigger) return 'Unknown trigger'; + + if (trigger.type === 'WorkflowStart') { + return 'Workflow Start'; + } else if (trigger.type === 'JobsComplete') { + const jobs = trigger.data?.jobs || []; + if (jobs.length === 0) return 'Jobs Complete'; + if (jobs.length === 1) return `When job '${jobs[0]}' completes`; + if (jobs.length <= 3) return `When jobs complete: ${jobs.map(j => `'${j}'`).join(', ')}`; + return `When ${jobs.length} jobs complete ('${jobs[0]}', '${jobs[1]}'...)`; } - return ` -
-
Scheduler Allocations
- ${allocations.map(alloc => ` -
-
Scheduler: ${this.escapeHtml(alloc.scheduler)} (${this.escapeHtml(alloc.scheduler_type)})
-
Allocations: ${alloc.num_allocations}
- ${alloc.job_names && alloc.job_names.length > 0 ? ` -
Jobs: ${alloc.job_names.slice(0, 5).map(n => this.escapeHtml(n)).join(', ')}${alloc.job_names.length > 5 ? ` ... and ${alloc.job_names.length - 5} more` : ''}
- ` : ''} -
- `).join('')} -
- `; + return 'Unknown trigger'; }, }); diff --git a/torc-dash/static/js/app-wizard.js b/torc-dash/static/js/app-wizard.js index 22b49959..8b1b9dfa 100644 --- a/torc-dash/static/js/app-wizard.js +++ b/torc-dash/static/js/app-wizard.js @@ -22,6 +22,12 @@ Object.assign(TorcDashboard.prototype, { }; this.wizardParallelizationStrategy = 'resource_aware'; + // HPC profile state for auto-generation + this.wizardHpcProfiles = []; + this.wizardDetectedProfile = null; + this.wizardSelectedProfile = null; + this.wizardSlurmAccount = ''; + this.resourcePresets = { 'small': { name: 'Small', num_cpus: 1, memory: '1g', num_gpus: 0 }, 'medium': { name: 'Medium', num_cpus: 4, memory: '10g', num_gpus: 0 }, @@ -71,6 +77,10 @@ Object.assign(TorcDashboard.prototype, { this.wizardResourceMonitor = { enabled: true, granularity: 'summary', sample_interval_seconds: 5 }; this.wizardParallelizationStrategy = 'resource_aware'; + // Reset HPC profile state but keep the profiles list and detected profile + this.wizardSelectedProfile = this.wizardDetectedProfile; + this.wizardSlurmAccount = ''; + const nameInput = document.getElementById('wizard-name'); const descInput = document.getElementById('wizard-description'); if (nameInput) nameInput.value = ''; @@ -115,7 +125,13 @@ Object.assign(TorcDashboard.prototype, { }); if (step === 2) this.wizardRenderJobs(); - else if (step === 3) this.wizardRenderSchedulers(); + else if (step === 3) { + // Load HPC profiles if not already loaded + if (this.wizardHpcProfiles.length === 0) { + this.wizardLoadHpcProfiles(); + } + this.wizardRenderSchedulers(); + } else if (step === 4) this.wizardRenderActions(); document.getElementById('wizard-prev').disabled = step === 1; @@ -420,6 +436,12 @@ Object.assign(TorcDashboard.prototype, { const container = document.getElementById('wizard-schedulers-list'); if (!container) return; + // Render the auto-generate section at the top + const autoGenContainer = document.getElementById('wizard-auto-generate-container'); + if (autoGenContainer) { + autoGenContainer.innerHTML = this.wizardRenderAutoGenerateUI(); + } + const expandedSchedulerIds = []; container.querySelectorAll('.wizard-scheduler-card.expanded').forEach(card => { const schedulerId = parseInt(card.dataset.schedulerId); @@ -427,7 +449,7 @@ Object.assign(TorcDashboard.prototype, { }); if (this.wizardSchedulers.length === 0) { - container.innerHTML = `

No schedulers defined

Click "+ Add Scheduler" to define a Slurm scheduler.

Schedulers are optional.

`; + container.innerHTML = `

No schedulers defined yet

Use "Auto-Generate Schedulers" above or click "+ Add Scheduler" below.

`; return; } @@ -761,4 +783,190 @@ Object.assign(TorcDashboard.prototype, { this.showToast('Error creating workflow: ' + error.message, 'error'); } }, + + // ==================== HPC Profile & Auto-Generation ==================== + + /** + * Load available HPC profiles from the server + */ + async wizardLoadHpcProfiles() { + try { + const result = await api.getHpcProfiles(); + if (result.success) { + this.wizardHpcProfiles = result.profiles || []; + this.wizardDetectedProfile = result.detected_profile; + // Auto-select detected profile + if (this.wizardDetectedProfile && !this.wizardSelectedProfile) { + this.wizardSelectedProfile = this.wizardDetectedProfile; + } + // Update UI if we're on the schedulers step + if (this.wizardStep === 3) { + this.wizardRenderSchedulers(); + } + } else { + console.warn('Failed to load HPC profiles:', result.error); + } + } catch (error) { + console.error('Error loading HPC profiles:', error); + } + }, + + /** + * Handle HPC profile selection change + */ + wizardSelectHpcProfile(profileName) { + this.wizardSelectedProfile = profileName || null; + }, + + /** + * Handle Slurm account input change + */ + wizardSetSlurmAccount(account) { + this.wizardSlurmAccount = account; + }, + + /** + * Auto-generate Slurm schedulers based on job resource requirements + */ + async wizardAutoGenerateSchedulers() { + const account = this.wizardSlurmAccount?.trim(); + if (!account) { + this.showToast('Please enter a Slurm account name', 'warning'); + return; + } + + if (!this.wizardSelectedProfile) { + this.showToast('Please select an HPC profile', 'warning'); + return; + } + + if (this.wizardJobs.length === 0) { + this.showToast('No jobs defined. Add jobs in step 2 first.', 'warning'); + return; + } + + // Build a partial spec with just jobs and resource_requirements + const spec = this.wizardGenerateSpec(); + + // Show loading state + const generateBtn = document.getElementById('wizard-auto-generate-btn'); + if (generateBtn) { + generateBtn.disabled = true; + generateBtn.textContent = 'Generating...'; + } + + try { + const result = await api.generateSlurmSchedulers(spec, account, this.wizardSelectedProfile); + + if (result.success && result.schedulers) { + // Apply the generated schedulers + this.wizardApplyGeneratedSchedulers(result.schedulers, result.actions); + this.showToast(`Generated ${result.schedulers.length} scheduler(s) and ${result.actions?.length || 0} action(s)`, 'success'); + } else { + this.showToast('Error: ' + (result.error || 'Failed to generate schedulers'), 'error'); + } + } catch (error) { + this.showToast('Error generating schedulers: ' + error.message, 'error'); + } finally { + if (generateBtn) { + generateBtn.disabled = false; + generateBtn.textContent = 'Auto-Generate Schedulers'; + } + } + }, + + /** + * Apply generated schedulers and actions to the wizard state + */ + wizardApplyGeneratedSchedulers(schedulers, actions) { + // Clear existing schedulers and actions + this.wizardSchedulers = []; + this.wizardActions = []; + this.wizardSchedulerIdCounter = 0; + this.wizardActionIdCounter = 0; + + // Add generated schedulers + for (const s of schedulers) { + const schedulerId = ++this.wizardSchedulerIdCounter; + this.wizardSchedulers.push({ + id: schedulerId, + name: s.name || '', + account: s.account || this.wizardSlurmAccount, + nodes: s.nodes || 1, + walltime: s.walltime || '01:00:00', + partition: s.partition || '', + qos: s.qos || '', + gres: s.gres || '', + mem: s.mem || '', + tmp: s.tmp || '', + extra: s.extra || '' + }); + } + + // Add generated actions + if (actions) { + for (const a of actions) { + const actionId = ++this.wizardActionIdCounter; + this.wizardActions.push({ + id: actionId, + trigger_type: a.trigger_type || 'on_workflow_start', + scheduler: a.scheduler || '', + jobs: a.jobs || [], + num_allocations: a.num_allocations || 1, + max_parallel_jobs: a.max_parallel_jobs || 10, + start_one_worker_per_node: a.start_one_worker_per_node || false + }); + } + } + + // Re-render the schedulers view + this.wizardRenderSchedulers(); + }, + + /** + * Render the auto-generate UI for step 3 + */ + wizardRenderAutoGenerateUI() { + const hasProfiles = this.wizardHpcProfiles.length > 0; + const profileOptions = this.wizardHpcProfiles.map(p => { + const isDetected = p.is_detected ? ' (detected)' : ''; + const isSelected = p.name === this.wizardSelectedProfile; + return ``; + }).join(''); + + return ` +
+
Auto-Generate from HPC Profile
+

Automatically generate Slurm schedulers based on job resource requirements.

+
+
+ + +
+
+ + +
+
+ +
+
+ ${!hasProfiles ? '

No HPC profiles available. You can still add schedulers manually below.

' : ''} +
+
+
+
Schedulers
+

Or manually configure Slurm schedulers below.

+
+ `; + }, }); diff --git a/torc-dash/static/js/dag.js b/torc-dash/static/js/dag.js index b9c556f9..8ee62574 100644 --- a/torc-dash/static/js/dag.js +++ b/torc-dash/static/js/dag.js @@ -237,26 +237,30 @@ class DAGVisualizer { } // Create edges from relationships + // API returns producer_job_id (job that outputs the file) and consumer_job_id (job that inputs the file) const edges = []; if (relationships) { for (const rel of relationships) { - if (rel.relationship_type === 'input') { - // File -> Job (input) + // Producer job -> File (output edge) + if (rel.producer_job_id) { edges.push({ data: { - source: `file-${rel.file_id}`, - target: `job-${rel.job_id}`, + id: `edge-job-${rel.producer_job_id}-file-${rel.file_id}`, + source: `job-${rel.producer_job_id}`, + target: `file-${rel.file_id}`, }, - classes: 'input', + classes: 'output', }); - } else if (rel.relationship_type === 'output') { - // Job -> File (output) + } + // File -> Consumer job (input edge) + if (rel.consumer_job_id) { edges.push({ data: { - source: `job-${rel.job_id}`, - target: `file-${rel.file_id}`, + id: `edge-file-${rel.file_id}-job-${rel.consumer_job_id}`, + source: `file-${rel.file_id}`, + target: `job-${rel.consumer_job_id}`, }, - classes: 'output', + classes: 'input', }); } } @@ -311,24 +315,30 @@ class DAGVisualizer { } // Create edges from relationships + // API returns producer_job_id (job that outputs user data) and consumer_job_id (job that inputs user data) const edges = []; if (relationships) { for (const rel of relationships) { - if (rel.relationship_type === 'input') { + // Producer job -> UserData (output edge) + if (rel.producer_job_id) { edges.push({ data: { - source: `ud-${rel.user_data_id}`, - target: `job-${rel.job_id}`, + id: `edge-job-${rel.producer_job_id}-ud-${rel.user_data_id}`, + source: `job-${rel.producer_job_id}`, + target: `ud-${rel.user_data_id}`, }, - classes: 'input', + classes: 'output', }); - } else if (rel.relationship_type === 'output') { + } + // UserData -> Consumer job (input edge) + if (rel.consumer_job_id) { edges.push({ data: { - source: `job-${rel.job_id}`, - target: `ud-${rel.user_data_id}`, + id: `edge-ud-${rel.user_data_id}-job-${rel.consumer_job_id}`, + source: `ud-${rel.user_data_id}`, + target: `job-${rel.consumer_job_id}`, }, - classes: 'output', + classes: 'input', }); } }