# XGFabric Workflow Orchestrator - Implementation Plan

## Overview

Implement a Python client application (`examples/xgfabric.py`) that orchestrates the CFDaAI (CFD + ML) workflow across two HPC clusters (UCSB and Perlmutter) using the RADICAL Edge framework. The client fetches sensor data, starts work on UCSB immediately, submits a job to Perlmutter to spawn a sub-edge, and migrates work to Perlmutter once the allocation starts.

## Architecture

```
Client (local)
    │
    ├── BridgeClient ──► Bridge
    │                      │
    │                      ├── "ucsb" edge (always online, SSH-based)
    │                      │      └── rhapsody, staging plugins
    │                      │
    │                      ├── "perlmutter" edge (login node, always online)
    │                      │      └── psij, queue_info plugins
    │                      │
    │                      └── "perlmutter.1" edge (compute node, spawned via job)
    │                             └── rhapsody, staging plugins
    │
    └── Local workspace: /tmp/xgfabric_workspace/
```

## Execution Flow

```
1. Initialize
   └── Connect to bridge, verify "ucsb" and "perlmutter" edges exist

2. Acquire sensor data (local)
   └── Call CSPOT (senspot-get) locally → sensor_out.csv

3. Submit Perlmutter allocation job (async)
   └── psij.submit_job() on "perlmutter" edge
       └── Job runs: radical-edge-service.py --name perlmutter.1

4. Stage sensor data to UCSB
   └── staging.put() to ~/xgfabric/intheloop/data/

5. Run CFD simulations on UCSB (batched)
   └── FOR each batch of N simulations:
       ├── rhapsody.submit_tasks() → runme.sh per simulation
       ├── rhapsody.wait_tasks()
       ├── Check if perlmutter.1 is online
       └── If online AND remaining work benefits from GPU: break

6. Decision point - where to continue?
   ├── IF perlmutter.1 online:
   │   ├── Migrate sim results (ucsb → local → perlmutter.1)
   │   └── Set active_cluster = perlmutter.1
   └── ELSE:
       └── Set active_cluster = ucsb (fallback, warn about slow GPU training)

7. Run ML training on active_cluster
   └── rhapsody.submit_tasks() for PCR, PINN, FNO
   └── rhapsody.wait_tasks()

8. Run evaluation on active_cluster
   └── rhapsody tasks for metrics computation

9. Cleanup
   └── Download final results to local
   └── Close sessions
   └── Optionally cancel Perlmutter job
```

## File Structure

```
examples/xgfabric.py    # Single-file implementation (~400 lines)
```

## Implementation Details

### 1. Configuration (dataclass at top of file)

```python
@dataclass
class Config:
    bridge_url: str = "https://localhost:8000"
    bridge_cert: str | None = None

    ucsb_edge: str = "ucsb"
    perlmutter_edge: str = "perlmutter"
    perlmutter_child: str = "perlmutter.1"

    workflow_path: str = "~/xgfabric/intheloop"
    local_workspace: str = "/tmp/xgfabric_workspace"

    # CSPOT
    cspot_woof_url: str = "woof://128.111.45.61/davisstations/daviscupsout"
    cspot_limit: int = 72

    # Perlmutter job
    pm_queue: str = "regular"
    pm_account: str = "m5290"
    pm_duration: int = 3600
    pm_nodes: int = 1

    # Workflow
    num_simulations: int = 16
    train_models: list = field(default_factory=lambda: ["pcr", "pinn", "fno"])
    batch_size: int = 4
```

### 2. Main Orchestrator Class

```python
class XGFabricOrchestrator:
    def __init__(self, config: Config):
        self.cfg = config
        self.bc: BridgeClient = None
        self.pm_job_id: str = None
        self.active_cluster: str = None  # "ucsb" or "perlmutter.1"

    def run(self) -> bool:
        try:
            self._connect()
            self._verify_edges()

            # Phase 1: Local data acquisition
            sensor_csv = self._acquire_sensor_data()

            # Phase 2: Submit perlmutter allocation (async)
            self._submit_perlmutter_job()

            # Phase 3: Stage data to UCSB and run simulations (batched)
            self._stage_to_ucsb(sensor_csv)
            sim_results = self._run_ucsb_simulations_batched()

            # Phase 4: Decide where to continue
            if self._is_perlmutter1_online():
                print("perlmutter.1 online - migrating data")
                self._migrate_to_perlmutter1(sim_results)
                self.active_cluster = "perlmutter.1"
            else:
                print("perlmutter.1 not available - continuing on UCSB (slow)")
                self.active_cluster = "ucsb"

            # Phase 5: Training and evaluation on active cluster
            self._run_training()
            self._run_evaluation()

            return True
        except Exception as e:
            print(f"Workflow failed: {e}")
            return False
        finally:
            self._cleanup()
```

### 3. Key Methods

**`_verify_edges()`**
- `bc.list_edges()` → check "ucsb" and "perlmutter" are present
- Raise `RuntimeError` if missing

**`_acquire_sensor_data()`**
- Adapt logic from `load_historical_data.py`
- Find `senspot-get` binary, invoke subprocess
- Parse output, write to local `sensor_out.csv`
- Return path to CSV

**`_submit_perlmutter_job()`**
- Get EdgeClient for "perlmutter"
- Get psij plugin
- Build job spec (from `example_edge_spawn.py` pattern):
  ```python
  job_spec = {
      "executable": "radical-edge-service.py",
      "arguments": ["--url", bridge_url, "--name", "perlmutter.1"],
      "attributes": {
          "queue_name": self.cfg.pm_queue,
          "account": self.cfg.pm_account,
          "duration": str(self.cfg.pm_duration),
          "node_count": self.cfg.pm_nodes,
      }
  }
  psij.submit_job(job_spec, "slurm")
  ```

**`_stage_to_ucsb(sensor_csv)`**
- Get staging plugin for UCSB
- `staging.put(sensor_csv, "~/xgfabric/intheloop/data/sensor_out.csv")`

**`_run_ucsb_simulations_batched()`**
- Get rhapsody plugin for UCSB
- Build simulation tasks (call runme.sh with wind parameters)
- Submit in batches of `batch_size` (default 4):
  ```python
  for batch in chunked(tasks, batch_size):
      submitted = rhapsody.submit_tasks(batch)
      rhapsody.wait_tasks([t['uid'] for t in submitted])
      if self._is_perlmutter1_online():
          break  # Migrate opportunity
  ```
- Return list of completed result file paths

**`_is_perlmutter1_online()`**
- Non-blocking check: `"perlmutter.1" in bc.list_edges()`
- Return True/False immediately

**`_migrate_to_perlmutter1(sim_results)`**
- Get staging plugins for both edges
- For each result file:
  - `ucsb_staging.get(remote, local)`
  - `pm1_staging.put(local, remote)`

**`_run_training()`**
- Get rhapsody plugin for active_cluster
- Submit training tasks for PCR, PINN, FNO
- Wait for completion

**`_cleanup()`**
- Close all plugin sessions
- Close BridgeClient
- Optionally cancel perlmutter job

### 4. Rhapsody Task Definitions

**Simulation task:**
```python
{
    "executable": f"{workflow_path}/simulation/runme.sh",
    "arguments": [str(wind_speed), str(wind_dir), str(sim_id)],
    "environment": {"WORK_DIR": f"{workflow_path}"}
}
```

**Training tasks:**
```python
# PCR (CPU)
{
    "executable": f"{workflow_path}/training/pcr/train_pcr.sh",
    "arguments": [data_dir, output_dir]
}

# PINN (GPU)
{
    "executable": "python3",
    "arguments": [f"{workflow_path}/training/pinn/train_pinn.py",
                  "--data", data_dir, "--output", output_dir]
}

# FNO (GPU)
{
    "executable": "python3",
    "arguments": [f"{workflow_path}/training/fno/train_fno.py",
                  "--data", data_dir, "--output", output_dir]
}
```

### 5. Data Staging Strategy

Three-way copy pattern (UCSB → Local → Perlmutter.1):

```python
def _migrate_to_perlmutter1(self, files: list[str]):
    ucsb_staging = self._get_ucsb_staging()
    pm1_staging = self._get_pm1_staging()

    local_dir = Path(self.cfg.local_workspace) / "staging"
    local_dir.mkdir(parents=True, exist_ok=True)

    for remote_path in files:
        filename = Path(remote_path).name
        local_path = str(local_dir / filename)

        # Download from UCSB
        ucsb_staging.get(remote_path, local_path)

        # Upload to Perlmutter.1
        pm1_target = f"{self.cfg.workflow_path}/input_data/{filename}"
        pm1_staging.put(local_path, pm1_target)
```

### 6. Error Handling

- Wrap main flow in try/except/finally
- Always call `_cleanup()` in finally block
- Log errors with context
- Return False on failure for CLI exit code

### 7. CLI Entry Point

```python
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--bridge-url", default=os.environ.get("RADICAL_BRIDGE_URL"))
    parser.add_argument("--bridge-cert", default=os.environ.get("RADICAL_BRIDGE_CERT"))
    parser.add_argument("--dry-run", action="store_true")
    # ... more args

    args = parser.parse_args()
    config = Config(bridge_url=args.bridge_url, ...)

    orchestrator = XGFabricOrchestrator(config)
    success = orchestrator.run()
    sys.exit(0 if success else 1)
```

## Files to Create/Modify

| File | Action | Description |
|------|--------|-------------|
| `examples/xgfabric.py` | Create | Main orchestrator script |

## Reference Files

- `examples/example_edge_spawn.py:50-59` - Pattern for waiting for edge to come online
- `examples/example_edge_spawn.py:76-88` - Job spec for spawning sub-edge
- `src/radical/edge/plugin_staging.py` - Staging plugin API
- `tmp/intheloop/data/load_historical_data.py` - CSPOT data fetching logic
- `tmp/intheloop/simulation/runme.sh` - Simulation runner script

## Verification

1. **Unit test**: Mock BridgeClient and verify orchestration logic
2. **Integration test**:
   - Start bridge locally
   - Start mock edges (or real UCSB/Perlmutter)
   - Run orchestrator
   - Verify data staged correctly
   - Verify tasks submitted and completed
3. **Manual test**:
   ```bash
   # Terminal 1: Bridge
   ./bin/radical-edge-bridge.py

   # Terminal 2: Mock edges (or real)
   # ... start ucsb and perlmutter edges

   # Terminal 3: Run orchestrator
   python examples/xgfabric.py --bridge-url https://localhost:8000
   ```

## Fallback Behavior (matches current shell workflow)

The current shell workflow on UCSB:
- **Simulations**: SSH distributed across UCSB machines (parallel)
- **PCR Training**: SSH distributed across UCSB machines (CPU-based)
- **PINN/FNO Training**: Runs locally with warning "No GPU detected - may be slow"

The script doesn't fail without GPU - it warns and proceeds. So the orchestrator should:
1. If perlmutter.1 comes online: migrate and use GPU for PINN/FNO
2. If perlmutter.1 doesn't come online: complete everything on UCSB (slow for PINN/FNO)

## Batched Execution with Migration Checks

Simulations run in batches with periodic checks for perlmutter.1:
- Submit batch of N simulations (e.g., 4 at a time)
- After each batch completes, check if perlmutter.1 is online
- If online: migrate intermediate data and continue on Perlmutter
- If not: continue with next batch on UCSB
