diff --git a/lib/python/examples/MIGRATING_TO_LAUNCHER.md b/lib/python/examples/MIGRATING_TO_LAUNCHER.md index f52d1380b..04d33083b 100644 --- a/lib/python/examples/MIGRATING_TO_LAUNCHER.md +++ b/lib/python/examples/MIGRATING_TO_LAUNCHER.md @@ -8,7 +8,7 @@ implementation; mirror it. - **How to *run*** the launcher and the experiment-YAML shape: see [`examples/README.md`](README.md). - **Why** we moved off thousands of JSON files: see - [`async_cifar10/MIGRATION_PLAN.md`](async_cifar10/MIGRATION_PLAN.md). + [`async_cifar10/MIGRATION_PLAN.md`](async_cifar10/MIGRATION_PLAN.md) (historical). - This doc: **what to change** (trainer + aggregator + metadata) to migrate a new example, and what legacy to delete. @@ -35,22 +35,47 @@ so no example stores per-trainer JSON anymore. `async_cifar10/metadata` is a symlink to `../_metadata`, so an example points at the shared bundle by symlinking (or via `experiment.metadata.dir`). +### What the launcher injects automatically per trainer + +`flame/launch/spawner.py:ConfigGenerator.generate_trainer_config()` merges +these layers in order (last wins for leaf values): + +1. `configs/trainer_base.yaml` — static per-example template +2. `baseline.trainer` dict from `baselines.yaml` — selector/optimizer/notify defaults +3. Per-trainer metadata from `_metadata/`: + - `taskid` (from registry `task_id`) + - `training_delay_s` (from registry, 4–18 s range for the 300-trainer pool) + - `trainer_indices_list` (from dataset splits; for non-CIFAR examples see §6) + - **All five availability traces pre-injected**: `avl_events_syn_0`, `avl_events_syn_20`, + `avl_events_syn_50`, `avl_events_mobiperf_2st`, `avl_events_mobiperf_3st_50`, + `avl_events_mobiperf_3st_75` — trainer selects at runtime via + `client_notify.trace` +4. `experiment.trainer.config_overrides` — experiment-level deep-merge +5. Dotted-key overrides: `job.id`, `job.name`, `hyperparameters.*` + +All five traces are pre-injected so the trainer never has to be re-spawned to +switch availability modes within a batch. + ### Adding a new dataset's data -1. **Dataset splits.** Generate `dataset_splits/_alpha_n.yaml` with - a `trainer_data_splits` map (`trainer_001 … trainer_NNN` → index lists). - Reuse the existing split-generation script; do **not** hand-write splits. -2. **Trainer population.** If the new example needs a different `N` or device +1. **Index-style splits** (e.g., CIFAR-10). Generate + `dataset_splits/_alpha_n.yaml` with a `trainer_data_splits` + map (`trainer_001 … trainer_NNN` → list of sample indices). Reuse the + existing split-generation script; do **not** hand-write splits. +2. **Path-style datasets** (e.g., NLP datasets from H5 partitions). The shared + `dataset_splits` YAML cannot hold per-trainer index lists; instead, store + `data_file_path`, `partition_file_path`, and `client_idx` directly in + `trainer_base.yaml` as placeholders, and inject them via + `experiment.trainer.config_overrides.hyperparameters`. See §6 for the + fwdllm-specific guidance. +3. **Trainer population.** If the new example needs a different `N` or device timing, extend `trainer_registry.yaml` (or add a parallel registry and point `experiment.metadata.registry` at it). The registry is dataset-agnostic — reuse it when the device population/timing is the same. -3. **Availability traces.** Reuse `synthetic_traces.yaml` / `mobiperf_traces.yaml` +4. **Availability traces.** Reuse `synthetic_traces.yaml` / `mobiperf_traces.yaml` as-is — availability is independent of the dataset. Only add traces if you need new patterns. -So data, availability, and timing are **shared**; usually only a new -`dataset_splits/*.yaml` is required per dataset. - --- ## 2. Aggregator-side changes @@ -61,31 +86,120 @@ So data, availability, and timing are **shared**; usually only a new - `main_fedavg_agg.py` → `flame.mode.horizontal.top_aggregator` (base syncfl) Delete legacy `main*.py` that don't map to a stack you run. + 2. **Config intake via `load_config_from_argv()`** (`flame/launch/cli.py`). The launcher spawns aggregators with `--config-json`; manual runs can use `--config `. Do not parse a positional config path. + + ```python + from flame.launch.cli import load_config_from_argv + config = load_config_from_argv() + ``` + 3. **Gate wandb behind `--log_to_wandb`.** No module-level `wandb.init()` (it runs on import and blocks non-wandb runs). Mirror `main_fedavg_agg.py`'s `initialize_wandb()` + `self.log_to_wandb` pattern. + 4. **Availability comes from `_metadata`**, not per-trainer JSON. For oracular stacks, `read_trainer_unavailability(trace)` reads `_metadata/{trainer_registry.yaml, availability_traces/*}`. Use `track_trainer_avail.get('trace', '')` — the key is absent when tracking is disabled. +5. **Telemetry**: the launcher sets `FLAME_TELEMETRY_DIR` before spawning; the + aggregator should write structured JSONL events there. Typical events: + `selection` (selected trainer IDs, scores, round) and `aggregation` + (aggregation time, staleness stats). See §5 for the full telemetry contract. + +--- + ## 3. Trainer-side changes The trainer is spawned with `--config-json`; everything per-trainer arrives in `hyperparameters`. The trainer must read (not hardcode): - `trainer_indices_list` — dataset sample indices (from `dataset_splits`). + For path-style datasets, use `data_file_path` + `partition_file_path` + + `client_idx` instead (§6). - `training_delay_s` — per-trainer delay (from registry). -- `avl_events_` — availability events for the active trace. +- `training_delay_enabled` — bool gate on delay (from config_overrides or + trainer_base.yaml default). +- `avl_events_` — all five trace types pre-injected; trainer reads + `client_notify.trace` to pick the active one at runtime. - `client_notify.{enabled,trace}` — availability-aware notify; `trace` is always - present (the launcher defaults it to `syn_0`). + present (launcher defaults it to `syn_0`). + +Trainer entrypoint uses `load_config_from_argv()`: + +```python +from flame.launch.cli import load_config_from_argv +config = load_config_from_argv() +``` + +One trainer `main.py` is typically enough across stacks. + +### Time mode (real vs simulated) + +The launcher passes `--time_mode real|simulated` as a CLI-only arg (not in config +JSON). The trainer reads it from `sys.argv` alongside the config. Two modes: + +- **real**: trainer sleeps `training_delay_s` at true wall-clock pace. Availability + transitions happen at real-time intervals. Use for wall-clock benchmarks. +- **simulated**: no sleeps; the trainer reports its completion time back to the + aggregator so the aggregator can order updates by a *virtual clock*. Availability + states are evaluated against that virtual clock. Use for fast iteration. + +Set `time_mode` in the experiment YAML under `trainer.time_mode`. The spawner +passes it as `--time_mode` to each trainer subprocess. + +### Data streaming (optional) + +By default a trainer loads its **entire** partition at init and reuses it on every +selection. Real devices instead generate data continuously, so a client selected +early should train on *less* data than one selected late. `async_cifar10/main.py` +is the reference implementation — mirror it when the example wants realistic data +growth. -Trainer entrypoint also uses `load_config_from_argv()`. One trainer `main.py` is -typically enough across stacks. +- **Config** (on `hyperparameters`, disabled by default in `trainer_base.yaml`): + ```yaml + data_streaming: + enabled: "False" + full_data_available_after_s: 0 # sim-seconds until 100% data is visible + ``` + Enable per experiment under `trainer.config_overrides.hyperparameters.data_streaming`. +- **Algorithm**: at `load_data()` retain the full pool and a one-time shuffle seeded + by `trainer_id`. A helper exposes a growing prefix: + `visible = floor(min(1, sim_elapsed / X) * total)`, floored at 1 sample; once + `sim_elapsed >= X` the full pool stays visible. `sim_elapsed` is driven by the + sim clock (real-mode: wall-clock elapsed; simulated-mode: aggregator-stamped + virtual time). +- **Where**: the loader is rebuilt at the top of `train()` and `evaluate()` so + the visible subset (and the `dataset_size` reported to the aggregator) reflects + the current time. All a no-op when `enabled: "False"`. + +### Utility counterfactual (optional telemetry) + +Emit a comparison of streamed-prefix statistical utility vs full-pool utility to +measure the data-availability impact on selection bias. + +```yaml +util_counterfactual: + enabled: "True" + every_n_rounds: 1 + sample_size: 256 +``` + +Enable under `trainer.config_overrides.hyperparameters.util_counterfactual`. +The trainer emits `util_disparity` JSONL events; the post-run analyzer produces +streamed-vs-full disparity plots. + +### Memory profiler (optional) + +`async_cifar10/trainer/pytorch/memory_profiler.py` provides a `MemoryProfiler` +class that logs RSS, Python object counts, and torch tensor sizes (CPU + GPU) +every N rounds. Import and instantiate it in the trainer `__init__`; call +`log_memory_before_round()` at the top of the training loop. Useful for catching +leaks in long runs. --- @@ -114,22 +228,288 @@ everything else on a sync stack (`flame/launch/runner.py:_validate_stack`). --- -## 5. Migration checklist for a new example +## 5. Telemetry and post-run analysis + +Telemetry is **auto-enabled** for any run launched via the YAML runner: the runner +sets `FLAME_TELEMETRY_DIR=/telemetry/` before spawning processes; any +child process that imports `flame.telemetry` (or checks the env var) will write +to that directory. No broker or extra setup required. + +### JSONL event contract + +Each process writes one file: `trainer_.jsonl` or `aggregator.jsonl`. Each +line is a JSON object with at least `{"event": "", "ts": , ...}`. + +**Trainer events** (implement these for useful plots): + +| Event | Required fields | Purpose | +|-------|-----------------|---------| +| `trainer_round` | `round`, `loss`, `accuracy`, `train_time_s`, `wait_time_s` | Per-round training stats | +| `avail_change` | `round`, `old_state`, `new_state` | State-machine transitions | +| `util_disparity` | `round`, `streamed_util`, `full_util` | Data-streaming disparity | + +**Aggregator events**: + +| Event | Required fields | Purpose | +|-------|-----------------|---------| +| `selection` | `round`, `selected_ids`, `scores` | Which trainers were chosen | +| `aggregation` | `round`, `agg_time_s`, `staleness_stats` | Aggregation timing | + +### Post-run analysis scripts + +The runner calls analysis best-effort after experiment completion. Two scripts are +available in `async_cifar10/scripts/`: + +- **`analyze_send_recv_lag.py`**: parses `[SEND_RECV_LAG]` log entries from the + aggregator log; reports per-trainer median/p95/max lag. Run manually: + ```bash + python scripts/analyze_send_recv_lag.py /*_aggregator.log [--warn-threshold-s 5.0] + ``` +- **`compare_parity.py`**: compares two experiment runs (e.g., real vs simulated + time_mode) for trajectory parity. + +For full telemetry plot generation, point a post-processing script at the +`telemetry/` directory. The JSONL files are append-only and line-buffered, so +`tail -f` works during a live run. + +### Output directory structure + +``` +experiments/run_YYYYMMDD_HHMMSS_/ + aggregator_config.json # merged aggregator config used + execution_config.yaml # compact metadata refs + spawn commands + snapshot.yaml # git info + metadata SHA256 checksums + YYYYMMDD_HHMMSS_*_aggregator.log # aggregator stdout/stderr + YYYYMMDD_HHMMSS_*_trainers.log # all trainers combined (line-buffered) + YYYYMMDD_HHMMSS_*_resources.log # RAM/GPU monitoring (if enabled) + telemetry/ + aggregator.jsonl + trainer_1.jsonl … trainer_N.jsonl + plots/ # generated by post-run analysis (if run) +``` + +--- + +## 6. fwdllm-specific migration notes + +fwdllm differs from async_cifar10 in two key areas: (a) its dataset is stored as +H5 partitions on disk (not index lists injected by the spawner), and (b) its +existing trainer/aggregator already use `flame.config.Config` and have partial +availability support. This makes the port narrower than a greenfield migration. + +### What's already compatible + +- Trainer (`trainer/fl_main.py`) and aggregator (`aggregator/fl_main.py`) both + use `flame.config.Config`. Only the config-loading call needs to change. +- Availability notification thread and `client_notify` struct are already present. +- Oracular tracking (`track_trainer_avail`) is already implemented in + `FedSGDAggregator.py`. + +### What must change + +**1. Config intake** (trainer and aggregator) + +Replace: +```python +parser.add_argument("--config", required=True) +config = Config(args.config) +``` +with: +```python +from flame.launch.cli import load_config_from_argv +config = load_config_from_argv() +``` +Also add `--time_mode` as a side-channel CLI arg (not in config JSON) if you want +simulated mode support: +```python +import argparse, sys +_p = argparse.ArgumentParser(add_help=False) +_p.add_argument("--time_mode", default="real") +_known, _ = _p.parse_known_args() +time_mode = _known.time_mode +``` + +**2. trainer_base.yaml for fwdllm** + +Create `fwdllm/configs/trainer_base.yaml`. Unlike async_cifar10, the dataset +fields cannot come from `_metadata/dataset_splits/` (H5 paths are not index +lists). Put the dataset fields in the base template with placeholder comments, and +let experiments override them via `config_overrides.hyperparameters`: + +```yaml +# fwdllm/configs/trainer_base.yaml (minimal skeleton) +backend: mqtt +brokers: + - host: localhost + port: 1883 +channels: + # ... same channel block as current trainer JSON configs +hyperparameters: + # Dataset — override per experiment or per baseline + dataset: agnews # override in experiment YAML + data_file_path: PLACEHOLDER # inject via config_overrides + partition_file_path: PLACEHOLDER # inject via config_overrides + partition_method: niid_label_clients=100_alpha=1 + # Model + model_type: distilbert + model_name: distilbert-base-uncased + max_seq_length: 64 + peft_method: adapter + # FL + fl_algorithm: FedFwd + epochs: 1 + comm_round: 3000 + learning_rate: 0.01 + server_lr: 0.1 + train_batch_size: 32 + eval_batch_size: 32 + forward_mode: "True" + use_adapter: "True" + freeze_layers: "True" + fp16: "True" + # Per-trainer fields — injected by spawner + client_idx: 0 # injected: (trainer_id - 1) % 100 + training_delay_s: 4.0 # injected from registry + training_delay_enabled: "False" + training_delay_factor: "10" + # Availability + wait_until_next_avl: "True" + client_notify: + enabled: "False" + trace: syn_0 + # avl_events_* — injected by spawner for all five trace types +``` + +**3. Per-trainer `client_idx` injection** + +fwdllm uses `client_idx` (0–99) to select the trainer's data partition from the +H5 file. The spawner injects `taskid` and `training_delay_s` automatically from +the registry. For `client_idx`, add it to `experiment.trainer.config_overrides` +or implement a custom injection hook in a thin wrapper around `TrainerSpawner`. + +Simplest approach: add `client_idx` as a dotted-key override per trainer. In the +experiment YAML, set a formula under `config_overrides.hyperparameters.client_idx` +that the spawner evaluates to `(trainer_id - 1) % num_clients`. If the spawner +doesn't support per-trainer formula evaluation, inject it via a small wrapper +script that calls the spawner with per-trainer overrides. + +**4. Baselines in `_metadata/baselines.yaml`** + +Add fwdllm baselines. Minimal entries: + +```yaml +fedfwd_async_oort: + description: FedFwd + async_oort selector + fedbuff optimizer (asyncfl stack) + example: + aggregator_main: aggregator/main_asyncfl_agg.py + aggregator: + selector: + sort: async_oort + kwargs: + evalGoalFactor: 1.0 + optimizer: + sort: fedbuff + kwargs: {} + hyperparameters: + trackTrainerAvail: + enabled: "False" + type: "NA" + trainer: + hyperparameters: + client_notify: + enabled: "True" + fl_algorithm: FedFwd + forward_mode: "True" + +fedfwd_oort_oracular: + description: FedFwd + oort selector + fedavg optimizer + oracular tracking (syncfl stack) + example: + aggregator_main: aggregator/main_oort_sync_agg.py + aggregator: + selector: + sort: oort + kwargs: {} + optimizer: + sort: fedavg + kwargs: {} + hyperparameters: + trackTrainerAvail: + enabled: "True" + type: oracular + trainer: + hyperparameters: + client_notify: + enabled: "False" + fl_algorithm: FedFwd +``` + +**5. Availability traces** — reuse `synthetic_traces.yaml` and +`mobiperf_traces.yaml` as-is. fwdllm currently embeds trace arrays inline in +trainer JSONs; the launcher will inject them from `_metadata/` instead. No new +trace files needed for the standard syn_0/syn_20/syn_50/mobiperf variants. + +**6. Aggregator entrypoints** + +fwdllm currently has a single `aggregator/fl_main.py`. Split into per-stack +files matching async_cifar10's pattern, then register each in `baselines.yaml`: + +| Stack | File | Baseline | +|-------|------|---------| +| asyncfl | `aggregator/main_asyncfl_agg.py` | fedfwd_async_oort | +| syncfl (oort) | `aggregator/main_oort_sync_agg.py` | fedfwd_oort_oracular | + +Each file only changes the `TopAggregator` import; all FL logic stays in +`FedSGDAggregator.py`. + +**7. Wandb gating** + +fwdllm aggregator likely calls `wandb.init()` unconditionally. Gate it behind +`--log_to_wandb` following `main_fedavg_agg.py`'s `initialize_wandb()` pattern. +The launcher sets `log_to_wandb: false` by default; override in the experiment +YAML if needed. + +**8. Telemetry** + +fwdllm does not yet emit structured JSONL telemetry. Add `trainer_round` events +at minimum to get loss/accuracy curves from the post-run analyzer. The env var +`FLAME_TELEMETRY_DIR` is set by the launcher; check it and open a JSONL file: + +```python +import json, os, pathlib +_tel_dir = os.environ.get("FLAME_TELEMETRY_DIR") +_tel_file = open(pathlib.Path(_tel_dir) / f"trainer_{trainer_id}.jsonl", "a") if _tel_dir else None + +def _emit(event, **fields): + if _tel_file: + _tel_file.write(json.dumps({"event": event, **fields}) + "\n") + _tel_file.flush() +``` + +--- + +## 7. Migration checklist for a new example 1. Symlink `/metadata → ../_metadata` (or set `experiment.metadata.dir`). 2. Add `dataset_splits/_alpha_n.yaml`; reuse registry + traces. + For path-style datasets (fwdllm), skip this step and inject paths via + `config_overrides` instead (§6). 3. Add `/configs/trainer_base.yaml` (static per-example trainer template; per-trainer fields are injected by the launcher). 4. Provide one aggregator entrypoint per stack you run (§2); delete the rest. -5. Make trainer + aggregator entrypoints intake `--config-json` and read all - per-trainer values from `hyperparameters` (§2–§3). +5. Make trainer + aggregator entrypoints use `load_config_from_argv()` and read + all per-trainer values from `hyperparameters` (§2–§3). 6. Add the example's baselines to `baselines.yaml` with `example.aggregator_main`. 7. Add `expt_scripts/_n10_*_smoke.yaml` experiment YAMLs. 8. Validate without spawning: load the YAML, resolve the baseline, build the - aggregator config, and run `_validate_stack` (see the dry-run snippet in the - commit history / `runner` API). Then run a 10-trainer smoke test. -9. Delete the example's legacy JSON config dirs and mark its shell scripts - deprecated (§Legacy decommission). + aggregator config, and run `_validate_stack`. Then run a 10-trainer smoke test. +9. (Optional) Add `trainer_round` / `avail_change` JSONL telemetry events to + enable post-run analysis plots (§5). +10. (Optional) Port the `data_streaming` block into `trainer_base.yaml` + the + trainer `main.py` if the example wants streamed data growth (§3 Data streaming). +11. (Optional) Add `util_counterfactual` support if comparing streamed vs full + utility matters for this example (§3 Utility counterfactual). +12. Delete the example's legacy JSON config dirs and mark its shell scripts + deprecated (§Legacy decommission). --- @@ -155,7 +535,7 @@ Keep (still current): | Example | Status | |---------|--------| -| `async_cifar10` | Migrated (reference). 5 baselines: felix, fedbuff, fedavg, oort, refl. | +| `async_cifar10` | Migrated (reference). 5 baselines: felix, fedbuff, fedavg, oort, refl. Includes telemetry, streaming, time_mode, memory profiler. | | `feddance_cifar10` | Has launcher YAMLs; FedDance baseline blockers tracked in `async_cifar10/FEDDANCE_TODO.md`. | | `async_google_speech` | **TODO** — migrate per this guide (needs google-speech dataset splits + per-stack aggregator entrypoints). | -| `fwdllm` | **TODO** — migrate per this guide. | +| `fwdllm` | **TODO** — see §6 for fwdllm-specific steps. Main blockers: (1) config intake switch to `load_config_from_argv()`; (2) per-trainer `client_idx` injection; (3) split aggregator into per-stack entrypoints; (4) add baselines to `baselines.yaml`; (5) gate wandb; (6) add JSONL telemetry. Dataset splits are H5-path-based, not index-list-based — use `config_overrides` for data paths rather than `_metadata/dataset_splits/`. | diff --git a/lib/python/examples/async_cifar10/MIGRATION_PLAN.md b/lib/python/examples/async_cifar10/MIGRATION_PLAN.md deleted file mode 100644 index d8c7e70a9..000000000 --- a/lib/python/examples/async_cifar10/MIGRATION_PLAN.md +++ /dev/null @@ -1,1930 +0,0 @@ -# Async CIFAR-10 Configuration Refactoring Plan - -## Executive Summary - -**Goal**: Migrate from static config files to programmatic configuration for scalable, maintainable federated learning experiments. - -**Current Problem**: 5,924 trainer config files across 20+ directories with massive redundancy, making it nearly impossible to: -- Compare configurations across experiments -- Update shared parameters (e.g., availability traces) -- Maintain consistency across dataset splits -- Scale to new datasets/examples - -**Solution**: Centralized metadata with programmatic trainer spawning. - ---- - -## Current Architecture Analysis - -### File Structure Issues -``` -trainer/ -├── config_dir0.1_num300_traceFail_6d_3state_oort/ -│ ├── trainer_1.json (312 lines) -│ ├── trainer_2.json (312 lines) -│ └── ... (300 files × 312 lines = 93,600 lines!) -├── config_dir100_num300_traceFail_6d_3state_oort/ -│ └── ... (another 300 files) -└── ... (20+ directories) -``` - -### Redundancy Analysis - -**Identical Across All Trainers (70% of config):** -- `backend`, `brokers`, `groupAssociation` -- `channels` definition (47 lines) -- `dataset`, `dependencies` -- `hyperparameters`: `batchSize`, `learningRate`, `rounds`, `epochs` -- `use_oort_loss_fn`, `heartbeats`, `client_notify` -- `baseModel`, `job`, `registry`, `selector`, `optimizer` -- `maxRunTime`, `realm`, `role` - -**Trainer-Specific Properties (Intrinsic, 15%):** -- `taskid` (trainer ID - deterministic hash) -- `training_delay_s` (4.0, 8.0, 12.0, 16.0s - speed classes) -- Availability traces (mobiperf: per-trainer real-world patterns) - -**Dataset/Experiment-Specific (15%):** -- `trainer_indices_list` (local dataset indices - varies by alpha) -- Synthetic availability traces (syn_0, syn_20, syn_50 - uniform patterns) - ---- - -## Proposed Architecture - -### Directory Structure -``` -examples/ -├── metadata/ # NEW: Centralized metadata -│ ├── trainer_registry.yaml # Static trainer properties -│ ├── dataset_splits/ # Dataset partitions -│ │ ├── cifar10_alpha0.1_n300.yaml -│ │ ├── cifar10_alpha1.0_n300.yaml -│ │ └── cifar10_alpha100_n300.yaml -│ └── availability_traces/ # Trace definitions -│ ├── mobiperf_traces.yaml # Real-world per-trainer -│ └── synthetic_traces.yaml # Uniform synthetic -│ -├── async_cifar10/ -│ ├── configs/ # NEW: Template configs only -│ │ ├── aggregator_base.yaml # Base agg config -│ │ └── trainer_base.yaml # Base trainer template -│ ├── launch/ # NEW: Experiment launchers -│ │ ├── spawner.py # Programmatic trainer spawning -│ │ ├── run_experiment.py # High-level experiment runner -│ │ └── experiment_configs.yaml # Experiment definitions -│ ├── aggregator/ -│ │ └── pytorch/ -│ │ └── main_oort_agg.py # (minimal changes) -│ └── trainer/ -│ └── pytorch/ -│ └── main.py # (refactored to use metadata) -│ -└── README.md # Updated documentation -``` - -### Key Files Design - -#### 1. `metadata/trainer_registry.yaml` -```yaml -# Global trainer registry - 300 trainers with intrinsic properties -trainers: - trainer_001: - task_id: "505f9fc483cf4df68a2409257b5fad7d3c580370" - training_delay_s: 4.0 - speed_class: "fast" - mobiperf_trace_id: "device_001" - - trainer_002: - task_id: "505f9fc483cf4df68a2409257b5fad7d3c580371" - training_delay_s: 16.0 - speed_class: "slow" - mobiperf_trace_id: "device_002" - - # ... 298 more trainers - -# Speed distribution (example) -speed_distribution: - fast: [1-50] # 4s delay - medium: [51-150] # 8s delay - slow: [151-300] # 12-16s delay -``` - -#### 2. `metadata/dataset_splits/cifar10_alpha0.1_n300.yaml` -```yaml -# CIFAR-10 split with Dirichlet α=0.1, 300 clients -dataset_name: "cifar10" -num_trainers: 300 -dirichlet_alpha: 0.1 -total_samples: 50000 - -trainer_data_splits: - trainer_001: [8630, 6636, 35954, ...] # 217 indices - trainer_002: [38308, 4308, 7268, ...] # 13 indices - trainer_003: [...] - # ... 300 entries -``` - -#### 3. `metadata/availability_traces/mobiperf_traces.yaml` -```yaml -# Real-world MobiPerf availability traces per device -traces: - device_001: - name: "mobiperf_device_001" - states_2st: [(0, 'AVL_TRAIN'), (250, 'UN_AVL'), ...] - states_3st_50: [(0, 'AVL_TRAIN'), (250, 'UN_AVL'), (1053, 'AVL_TRAIN'), (7162, 'AVL_EVAL'), ...] - states_3st_75: [(0, 'AVL_TRAIN'), (300, 'UN_AVL'), ...] - - device_002: - name: "mobiperf_device_002" - states_2st: [(0, 'AVL_TRAIN'), (300, 'UN_AVL'), ...] - states_3st_50: [(0, 'AVL_TRAIN'), (300, 'UN_AVL'), ...] - states_3st_75: [(0, 'AVL_TRAIN'), (300, 'UN_AVL'), ...] -``` - -#### 4. `metadata/availability_traces/synthetic_traces.yaml` -```yaml -# Synthetic availability traces (uniform across all trainers) -traces: - syn_0: - description: "Always available" - pattern: [(0, 'AVL_TRAIN')] - - syn_20: - description: "20% unavailability" - pattern: [(0, 'AVL_TRAIN'), (600, 'UN_AVL'), (1200, 'AVL_TRAIN'), ...] - - syn_50: - description: "50% unavailability" - pattern: [(0, 'AVL_TRAIN'), (4800, 'UN_AVL'), (5400, 'AVL_TRAIN'), ...] -``` - -#### 5. `async_cifar10/launch/spawner.py` -```python -"""Programmatic trainer spawner - replaces exec_300_trainers_2state.sh""" -import yaml -import subprocess -from pathlib import Path -from typing import Dict, List - -class TrainerSpawner: - def __init__(self, metadata_dir: Path): - self.metadata_dir = metadata_dir - self.trainer_registry = self._load_yaml('trainer_registry.yaml') - self.mobiperf_traces = self._load_yaml('availability_traces/mobiperf_traces.yaml') - self.synthetic_traces = self._load_yaml('availability_traces/synthetic_traces.yaml') - - def load_dataset_split(self, dataset: str, alpha: float, num_trainers: int): - """Load dataset split configuration.""" - split_file = f"dataset_splits/{dataset}_alpha{alpha}_n{num_trainers}.yaml" - return self._load_yaml(split_file) - - def generate_trainer_config(self, - trainer_id: int, - dataset_split: Dict, - trace_type: str, - trace_name: str, - base_config: Dict) -> Dict: - """Generate runtime trainer config from metadata.""" - trainer_key = f"trainer_{trainer_id:03d}" - trainer_meta = self.trainer_registry['trainers'][trainer_key] - - # Build config dict (no file needed) - config = base_config.copy() - config['taskid'] = trainer_meta['task_id'] - config['hyperparameters']['training_delay_s'] = trainer_meta['training_delay_s'] - config['hyperparameters']['trainer_indices_list'] = dataset_split['trainer_data_splits'][trainer_key] - - # Set availability trace - if trace_type == 'synthetic': - config['hyperparameters'][f'avl_events_{trace_name}'] = \ - self.synthetic_traces['traces'][trace_name]['pattern'] - elif trace_type == 'mobiperf': - trace_id = trainer_meta['mobiperf_trace_id'] - config['hyperparameters']['avl_events_mobiperf_2st'] = \ - self.mobiperf_traces['traces'][trace_id]['states_2st'] - # ... handle 3st variants - - return config - - def spawn_trainers(self, - num_trainers: int, - dataset_split: Dict, - trace_type: str, - trace_name: str, - base_config: Dict, - num_gpus: int = 8, - delay_between_trainers: float = 1.0): - """Spawn trainer processes with generated configs.""" - processes = [] - - for trainer_id in range(1, num_trainers + 1): - # Generate config in-memory - config = self.generate_trainer_config( - trainer_id, dataset_split, trace_type, trace_name, base_config - ) - - # Assign GPU - gpu_id = trainer_id % num_gpus - - # Spawn trainer with config passed as JSON string or temp file - cmd = [ - 'python', '../pytorch/main.py', - '--config-dict', json.dumps(config) # OR write temp file - ] - - env = os.environ.copy() - env['CUDA_VISIBLE_DEVICES'] = str(gpu_id) - - proc = subprocess.Popen(cmd, env=env) - processes.append((trainer_id, proc)) - - time.sleep(delay_between_trainers) - - return processes -``` - -#### 6. `async_cifar10/launch/run_experiment.py` -```python -"""High-level experiment runner.""" -import argparse -from spawner import TrainerSpawner -from pathlib import Path - -def run_experiment(exp_config: Dict): - """Run a complete FL experiment.""" - # Load metadata - metadata_dir = Path(__file__).parent.parent.parent / 'metadata' - spawner = TrainerSpawner(metadata_dir) - - # Load dataset split - dataset_split = spawner.load_dataset_split( - dataset=exp_config['dataset'], - alpha=exp_config['dirichlet_alpha'], - num_trainers=exp_config['num_trainers'] - ) - - # Start aggregator - agg_proc = start_aggregator(exp_config['aggregator_config']) - time.sleep(15) # Wait for agg ready - - # Spawn trainers - trainer_procs = spawner.spawn_trainers( - num_trainers=exp_config['num_trainers'], - dataset_split=dataset_split, - trace_type=exp_config['trace_type'], - trace_name=exp_config['trace_name'], - base_config=load_base_config(), - num_gpus=exp_config['num_gpus'] - ) - - # Monitor and manage experiment - monitor_experiment(agg_proc, trainer_procs, exp_config) - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('--experiment', required=True, - help='Experiment name from experiment_configs.yaml') - args = parser.parse_args() - - # Load experiment definition - exp_config = load_experiment_config(args.experiment) - run_experiment(exp_config) -``` - -#### 7. `async_cifar10/launch/experiment_configs.yaml` -```yaml -# High-level experiment definitions (replaces multiple script files) -experiments: - oort_n300_oracular_syn0: - description: "Oort selector, 300 trainers, oracular tracking, always available" - dataset: "cifar10" - dirichlet_alpha: 0.1 - num_trainers: 300 - trace_type: "synthetic" - trace_name: "syn_0" - selector: "oort" - aggregator_concurrency: 13 - aggregation_goal: 10 - tracking_mode: "oracular" - num_gpus: 8 - aggregator_config: "configs/aggregator_base.yaml" - - oort_n300_oracular_syn20: - description: "Oort selector, 300 trainers, oracular tracking, 20% unavailable" - # Similar structure... - trace_name: "syn_20" - - oort_n300_oracular_mobiperf: - description: "Oort selector, 300 trainers, oracular tracking, real-world traces" - trace_type: "mobiperf" - trace_name: "mobiperf_2st" -``` - ---- - -## Migration Strategy - -### Phase 1: Automated Metadata Extraction & Verification (Week 1) -**Goal**: Programmatically extract static properties into centralized metadata with full validation - -**Critical Principle**: 100% automated extraction with verification - NO MANUAL DATA ENTRY - ---- - -#### Step 1.1: Setup Extraction Infrastructure (Day 1) - -**Create**: `scripts/extract_metadata.py` - -```python -""" -Automated metadata extraction from existing trainer configs. -Handles JSON parsing, deduplication, and YAML generation. -""" -import json -import yaml -from pathlib import Path -from collections import defaultdict -from typing import Dict, List, Set -import hashlib - -class MetadataExtractor: - def __init__(self, trainer_config_dirs: List[Path], output_dir: Path): - self.config_dirs = trainer_config_dirs - self.output_dir = output_dir - self.trainer_registry = {} - self.dataset_splits = defaultdict(dict) - self.mobiperf_traces = {} - self.synthetic_traces = {} - - def extract_all(self): - """Run full extraction pipeline.""" - print("Phase 1.1: Scanning trainer configs...") - self._scan_all_configs() - - print("Phase 1.2: Building trainer registry...") - self._build_trainer_registry() - - print("Phase 1.3: Extracting dataset splits...") - self._extract_dataset_splits() - - print("Phase 1.4: Extracting availability traces...") - self._extract_traces() - - print("Phase 1.5: Writing YAML files...") - self._write_yaml_files() - - print("✓ Extraction complete!") - - def _scan_all_configs(self): - """Scan all config directories and load JSON files.""" - self.all_configs = {} - for config_dir in self.config_dirs: - dir_name = config_dir.name - self.all_configs[dir_name] = {} - - # Load all trainer_N.json files - for config_file in sorted(config_dir.glob("trainer_*.json")): - if config_file.name == "trainer_0_test.json": - continue # Skip test files - - trainer_num = int(config_file.stem.split('_')[1]) - with open(config_file) as f: - self.all_configs[dir_name][trainer_num] = json.load(f) - - print(f" Loaded configs from {len(self.all_configs)} directories") - for dir_name, configs in self.all_configs.items(): - print(f" {dir_name}: {len(configs)} trainers") - - def _build_trainer_registry(self): - """Extract trainer-specific intrinsic properties.""" - # Use first config directory as reference (properties should be consistent) - reference_dir = list(self.all_configs.keys())[0] - reference_configs = self.all_configs[reference_dir] - - for trainer_num, config in sorted(reference_configs.items()): - trainer_key = f"trainer_{trainer_num:03d}" - - self.trainer_registry[trainer_key] = { - 'trainer_id': trainer_num, - 'task_id': config['taskid'], - 'training_delay_s': float(config['hyperparameters']['training_delay_s']), - } - - # Infer speed class from delay - delay = float(config['hyperparameters']['training_delay_s']) - if delay <= 4.0: - speed_class = 'fast' - elif delay <= 8.0: - speed_class = 'medium' - elif delay <= 12.0: - speed_class = 'slow' - else: - speed_class = 'very_slow' - - self.trainer_registry[trainer_key]['speed_class'] = speed_class - - print(f" Built registry for {len(self.trainer_registry)} trainers") - - def _extract_dataset_splits(self): - """Extract dataset index assignments per trainer per alpha.""" - for dir_name, configs in self.all_configs.items(): - # Parse directory name: config_dir_num_ - # e.g., config_dir0.1_num300_traceFail_6d_3state_oort - if not dir_name.startswith('config_dir'): - continue - - parts = dir_name.split('_') - alpha_str = parts[0].replace('config_dir', '') - - try: - alpha = float(alpha_str) - except ValueError: - print(f" Warning: Could not parse alpha from {dir_name}") - continue - - split_key = f"cifar10_alpha{alpha}_n300" - - for trainer_num, config in configs.items(): - trainer_key = f"trainer_{trainer_num:03d}" - indices = config['hyperparameters']['trainer_indices_list'] - - if split_key not in self.dataset_splits: - self.dataset_splits[split_key] = { - 'dataset_name': 'cifar10', - 'num_trainers': 300, - 'dirichlet_alpha': alpha, - 'total_samples': 50000, - 'trainer_data_splits': {} - } - - self.dataset_splits[split_key]['trainer_data_splits'][trainer_key] = indices - - print(f" Extracted {len(self.dataset_splits)} dataset splits") - for split_name, split_data in self.dataset_splits.items(): - num_trainers = len(split_data['trainer_data_splits']) - print(f" {split_name}: {num_trainers} trainers") - - def _extract_traces(self): - """Extract availability traces (both mobiperf and synthetic).""" - # Use first config directory as reference - reference_dir = list(self.all_configs.keys())[0] - reference_configs = self.all_configs[reference_dir] - - # Extract synthetic traces (uniform across all trainers) - first_config = reference_configs[1] - hp = first_config['hyperparameters'] - - self.synthetic_traces = { - 'description': 'Synthetic availability traces - uniform across all trainers', - 'traces': { - 'syn_0': { - 'description': 'Always available (0% unavailability)', - 'pattern': eval(hp['avl_events_syn_0']) - }, - 'syn_20': { - 'description': '20% unavailability', - 'pattern': eval(hp['avl_events_syn_20']) - }, - 'syn_50': { - 'description': '50% unavailability', - 'pattern': eval(hp['avl_events_syn_50']) - } - } - } - - # Extract mobiperf traces (per-trainer) - self.mobiperf_traces = { - 'description': 'Real-world MobiPerf availability traces - per device', - 'traces': {} - } - - for trainer_num, config in reference_configs.items(): - trainer_key = f"trainer_{trainer_num:03d}" - device_id = f"device_{trainer_num:03d}" - hp = config['hyperparameters'] - - self.mobiperf_traces['traces'][device_id] = { - 'trainer': trainer_key, - 'states_2st': eval(hp['avl_events_mobiperf_2st']), - 'states_3st_50': eval(hp['avl_events_mobiperf_3st_50']), - 'states_3st_75': eval(hp['avl_events_mobiperf_3st_75']), - } - - # Link device to trainer in registry - self.trainer_registry[trainer_key]['mobiperf_device_id'] = device_id - - print(f" Extracted {len(self.synthetic_traces['traces'])} synthetic traces") - print(f" Extracted {len(self.mobiperf_traces['traces'])} mobiperf device traces") - - def _write_yaml_files(self): - """Write all metadata to YAML files.""" - self.output_dir.mkdir(parents=True, exist_ok=True) - - # Write trainer registry - registry_file = self.output_dir / 'trainer_registry.yaml' - with open(registry_file, 'w') as f: - yaml.dump({ - 'description': 'Global trainer registry with intrinsic properties', - 'num_trainers': len(self.trainer_registry), - 'trainers': self.trainer_registry - }, f, default_flow_style=False, sort_keys=False) - print(f" ✓ Wrote {registry_file}") - - # Write dataset splits - splits_dir = self.output_dir / 'dataset_splits' - splits_dir.mkdir(exist_ok=True) - for split_name, split_data in self.dataset_splits.items(): - split_file = splits_dir / f"{split_name}.yaml" - with open(split_file, 'w') as f: - yaml.dump(split_data, f, default_flow_style=False) - print(f" ✓ Wrote {split_file}") - - # Write availability traces - traces_dir = self.output_dir / 'availability_traces' - traces_dir.mkdir(exist_ok=True) - - synthetic_file = traces_dir / 'synthetic_traces.yaml' - with open(synthetic_file, 'w') as f: - yaml.dump(self.synthetic_traces, f, default_flow_style=False) - print(f" ✓ Wrote {synthetic_file}") - - mobiperf_file = traces_dir / 'mobiperf_traces.yaml' - with open(mobiperf_file, 'w') as f: - yaml.dump(self.mobiperf_traces, f, default_flow_style=False) - print(f" ✓ Wrote {mobiperf_file}") - -# Main execution -if __name__ == '__main__': - import sys - - # Define config directories to process - base_dir = Path(__file__).parent.parent / 'trainer' - config_dirs = [ - base_dir / 'config_dir0.1_num300_traceFail_6d_3state_oort', - base_dir / 'config_dir1_num300_traceFail_6d_3state_oort', - base_dir / 'config_dir10_num300_traceFail_6d_3state_oort', - base_dir / 'config_dir100_num300_traceFail_6d_3state_oort', - ] - - output_dir = Path(__file__).parent.parent.parent / 'metadata' - - extractor = MetadataExtractor(config_dirs, output_dir) - extractor.extract_all() - - print("\n" + "="*60) - print("✓ Phase 1.1 Complete: Metadata extraction successful") - print("="*60) -``` - -**Run**: -```bash -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 -python scripts/extract_metadata.py -``` - -**Expected Output**: -- `metadata/trainer_registry.yaml` (300 trainers) -- `metadata/dataset_splits/cifar10_alpha*.yaml` (4 files) -- `metadata/availability_traces/synthetic_traces.yaml` -- `metadata/availability_traces/mobiperf_traces.yaml` - ---- - -#### Step 1.2: Comprehensive Validation (Day 2-3) - -**Create**: `scripts/validate_metadata.py` - -```python -""" -Comprehensive validation: verify YAML metadata matches original JSON configs. -This is the GATE before proceeding to Phase 2. -""" -import json -import yaml -from pathlib import Path -from typing import Dict, List, Tuple -import sys - -class MetadataValidator: - def __init__(self, metadata_dir: Path, trainer_config_dirs: List[Path]): - self.metadata_dir = metadata_dir - self.config_dirs = trainer_config_dirs - self.errors = [] - self.warnings = [] - - def validate_all(self) -> bool: - """Run all validation checks. Returns True if all pass.""" - print("="*70) - print("METADATA VALIDATION SUITE") - print("="*70) - - all_passed = True - - checks = [ - ("Trainer Registry", self._validate_trainer_registry), - ("Dataset Splits", self._validate_dataset_splits), - ("Synthetic Traces", self._validate_synthetic_traces), - ("MobiPerf Traces", self._validate_mobiperf_traces), - ("Cross-Reference Integrity", self._validate_cross_references), - ("Completeness Check", self._validate_completeness), - ] - - for check_name, check_func in checks: - print(f"\n{'─'*70}") - print(f"Running: {check_name}") - print('─'*70) - passed = check_func() - - if passed: - print(f"✓ {check_name}: PASSED") - else: - print(f"✗ {check_name}: FAILED") - all_passed = False - - print("\n" + "="*70) - print("VALIDATION SUMMARY") - print("="*70) - - if all_passed: - print("✓ ALL VALIDATIONS PASSED") - print("\n🎉 Metadata is verified correct. Safe to proceed to Phase 2.") - else: - print(f"✗ VALIDATION FAILED with {len(self.errors)} errors") - print("\nErrors:") - for error in self.errors[:10]: # Show first 10 - print(f" • {error}") - if len(self.errors) > 10: - print(f" ... and {len(self.errors) - 10} more errors") - print("\n⚠ DO NOT PROCEED TO PHASE 2 until all errors are fixed.") - - if self.warnings: - print(f"\nWarnings: {len(self.warnings)}") - for warning in self.warnings[:5]: - print(f" ⚠ {warning}") - - return all_passed - - def _load_metadata(self): - """Load all YAML metadata files.""" - with open(self.metadata_dir / 'trainer_registry.yaml') as f: - self.trainer_registry = yaml.safe_load(f) - - self.dataset_splits = {} - splits_dir = self.metadata_dir / 'dataset_splits' - for split_file in splits_dir.glob('*.yaml'): - with open(split_file) as f: - self.dataset_splits[split_file.stem] = yaml.safe_load(f) - - traces_dir = self.metadata_dir / 'availability_traces' - with open(traces_dir / 'synthetic_traces.yaml') as f: - self.synthetic_traces = yaml.safe_load(f) - with open(traces_dir / 'mobiperf_traces.yaml') as f: - self.mobiperf_traces = yaml.safe_load(f) - - def _load_original_config(self, config_dir: Path, trainer_num: int) -> Dict: - """Load original JSON config for comparison.""" - config_file = config_dir / f"trainer_{trainer_num}.json" - with open(config_file) as f: - return json.load(f) - - def _validate_trainer_registry(self) -> bool: - """Validate trainer registry against original configs.""" - self._load_metadata() - passed = True - - # Check count - expected_count = 300 - actual_count = self.trainer_registry['num_trainers'] - if actual_count != expected_count: - self.errors.append(f"Registry has {actual_count} trainers, expected {expected_count}") - passed = False - - # Validate each trainer against first config directory - reference_dir = self.config_dirs[0] - print(f" Validating against: {reference_dir.name}") - - for trainer_key, trainer_meta in self.trainer_registry['trainers'].items(): - trainer_num = trainer_meta['trainer_id'] - - try: - orig_config = self._load_original_config(reference_dir, trainer_num) - except FileNotFoundError: - self.errors.append(f"{trainer_key}: Original config not found") - passed = False - continue - - # Check task_id - if trainer_meta['task_id'] != orig_config['taskid']: - self.errors.append( - f"{trainer_key}: task_id mismatch - " - f"YAML={trainer_meta['task_id']}, JSON={orig_config['taskid']}" - ) - passed = False - - # Check training_delay_s - orig_delay = float(orig_config['hyperparameters']['training_delay_s']) - if trainer_meta['training_delay_s'] != orig_delay: - self.errors.append( - f"{trainer_key}: training_delay_s mismatch - " - f"YAML={trainer_meta['training_delay_s']}, JSON={orig_delay}" - ) - passed = False - - if passed: - print(f" ✓ All {actual_count} trainers validated") - - return passed - - def _validate_dataset_splits(self) -> bool: - """Validate dataset splits against original configs.""" - passed = True - - for split_name, split_data in self.dataset_splits.items(): - alpha = split_data['dirichlet_alpha'] - print(f" Validating split: {split_name} (alpha={alpha})") - - # Find corresponding config directory - config_dir = None - for cd in self.config_dirs: - if f"config_dir{alpha}_" in cd.name: - config_dir = cd - break - - if not config_dir: - self.errors.append(f"{split_name}: No matching config directory found") - passed = False - continue - - # Validate each trainer's indices - for trainer_key, indices in split_data['trainer_data_splits'].items(): - trainer_num = int(trainer_key.split('_')[1]) - - try: - orig_config = self._load_original_config(config_dir, trainer_num) - orig_indices = orig_config['hyperparameters']['trainer_indices_list'] - - if indices != orig_indices: - self.errors.append( - f"{split_name}/{trainer_key}: Indices mismatch - " - f"lengths YAML={len(indices)}, JSON={len(orig_indices)}" - ) - passed = False - except FileNotFoundError: - self.errors.append(f"{split_name}/{trainer_key}: Original config not found") - passed = False - - if passed: - num_trainers = len(split_data['trainer_data_splits']) - print(f" ✓ {num_trainers} trainers validated") - - return passed - - def _validate_synthetic_traces(self) -> bool: - """Validate synthetic traces.""" - passed = True - - # Check against first trainer config - reference_dir = self.config_dirs[0] - orig_config = self._load_original_config(reference_dir, 1) - hp = orig_config['hyperparameters'] - - for trace_name in ['syn_0', 'syn_20', 'syn_50']: - yaml_pattern = self.synthetic_traces['traces'][trace_name]['pattern'] - json_pattern = eval(hp[f'avl_events_{trace_name}']) - - if yaml_pattern != json_pattern: - self.errors.append(f"Synthetic trace {trace_name} mismatch") - passed = False - else: - print(f" ✓ {trace_name}: {len(yaml_pattern)} events") - - return passed - - def _validate_mobiperf_traces(self) -> bool: - """Validate mobiperf traces for all trainers.""" - passed = True - - reference_dir = self.config_dirs[0] - - for device_id, trace_data in self.mobiperf_traces['traces'].items(): - trainer_key = trace_data['trainer'] - trainer_num = int(trainer_key.split('_')[1]) - - try: - orig_config = self._load_original_config(reference_dir, trainer_num) - hp = orig_config['hyperparameters'] - - # Check all three trace variants - for variant in ['2st', '3st_50', '3st_75']: - yaml_trace = trace_data[f'states_{variant}'] - json_trace = eval(hp[f'avl_events_mobiperf_{variant}']) - - if yaml_trace != json_trace: - self.errors.append( - f"{device_id} ({trainer_key}): mobiperf_{variant} mismatch" - ) - passed = False - except Exception as e: - self.errors.append(f"{device_id}: Validation error - {e}") - passed = False - - if passed: - num_devices = len(self.mobiperf_traces['traces']) - print(f" ✓ {num_devices} device traces validated (3 variants each)") - - return passed - - def _validate_cross_references(self) -> bool: - """Validate that all cross-references are valid.""" - passed = True - - # Check that mobiperf device IDs in registry exist in traces - for trainer_key, trainer_meta in self.trainer_registry['trainers'].items(): - device_id = trainer_meta.get('mobiperf_device_id') - if device_id not in self.mobiperf_traces['traces']: - self.errors.append(f"{trainer_key}: References non-existent device {device_id}") - passed = False - - # Check that dataset splits reference valid trainers - for split_name, split_data in self.dataset_splits.items(): - for trainer_key in split_data['trainer_data_splits'].keys(): - if trainer_key not in self.trainer_registry['trainers']: - self.errors.append(f"{split_name}: References non-existent {trainer_key}") - passed = False - - if passed: - print(" ✓ All cross-references valid") - - return passed - - def _validate_completeness(self) -> bool: - """Check that no data is missing.""" - passed = True - - # Check all 300 trainers present - trainer_ids = set(range(1, 301)) - actual_ids = {meta['trainer_id'] for meta in self.trainer_registry['trainers'].values()} - missing = trainer_ids - actual_ids - - if missing: - self.errors.append(f"Missing trainers: {sorted(missing)}") - passed = False - else: - print(" ✓ All 300 trainers present") - - # Check all dataset splits have 300 trainers - for split_name, split_data in self.dataset_splits.items(): - if len(split_data['trainer_data_splits']) != 300: - self.errors.append( - f"{split_name}: Has {len(split_data['trainer_data_splits'])}/300 trainers" - ) - passed = False - - if passed: - print(f" ✓ All splits complete (4 splits × 300 trainers)") - - return passed - -# Main execution -if __name__ == '__main__': - metadata_dir = Path(__file__).parent.parent.parent / 'metadata' - base_dir = Path(__file__).parent.parent / 'trainer' - - config_dirs = [ - base_dir / 'config_dir0.1_num300_traceFail_6d_3state_oort', - base_dir / 'config_dir1_num300_traceFail_6d_3state_oort', - base_dir / 'config_dir10_num300_traceFail_6d_3state_oort', - base_dir / 'config_dir100_num300_traceFail_6d_3state_oort', - ] - - validator = MetadataValidator(metadata_dir, config_dirs) - passed = validator.validate_all() - - sys.exit(0 if passed else 1) -``` - -**Run**: -```bash -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 -python scripts/validate_metadata.py -``` - -**Gate Criteria**: Script must exit with code 0 (all checks pass) before proceeding to Phase 2. - ---- - -#### Step 1.3: Phase 1 Deliverables Checklist - -**Before proceeding to Phase 2, verify:** - -- [ ] `extract_metadata.py` runs without errors -- [ ] All YAML files generated: - - [ ] `metadata/trainer_registry.yaml` (300 trainers) - - [ ] `metadata/dataset_splits/cifar10_alpha0.1_n300.yaml` - - [ ] `metadata/dataset_splits/cifar10_alpha1_n300.yaml` - - [ ] `metadata/dataset_splits/cifar10_alpha10_n300.yaml` - - [ ] `metadata/dataset_splits/cifar10_alpha100_n300.yaml` - - [ ] `metadata/availability_traces/synthetic_traces.yaml` - - [ ] `metadata/availability_traces/mobiperf_traces.yaml` -- [ ] `validate_metadata.py` exits with code 0 -- [ ] All 6 validation checks pass: - - [ ] Trainer Registry validation - - [ ] Dataset Splits validation - - [ ] Synthetic Traces validation - - [ ] MobiPerf Traces validation - - [ ] Cross-reference integrity - - [ ] Completeness check -- [ ] Manual spot-check: Compare 5 random trainer configs (JSON vs YAML) by hand -- [ ] Git commit with message: "Phase 1 complete: Extracted and validated metadata" - -**Time Estimate**: 3 days -- Day 1: Write and test extraction script -- Day 2: Write and test validation script -- Day 3: Run full validation, fix any issues, final verification - -**Sign-off Required**: All checkboxes above must be ticked before Phase 2 begins. - -### Phase 2: Programmatic Spawning (Week 2) -**Goal**: Implement trainer spawner without changing experiment behavior - -**Prerequisites**: Phase 1 complete with all validation checks passing - -**Tasks** (detailed implementation in separate section after Phase 1 approval): -1. **Implement `launch/spawner.py`** - - Metadata loader - - Config generator (runtime assembly) - - Process spawner with GPU affinity - -2. **Refactor `trainer/pytorch/main.py`** - - Accept config as dict or JSON string (not just file) - - Keep backward compatibility with file-based configs - -3. **Create minimal base configs** - - `configs/trainer_base.yaml` (common properties only) - - `configs/aggregator_base.yaml` - -**Validation**: Run parallel experiments - old script vs new spawner -- Same trainer IDs spawned -- Same dataset splits loaded -- Same availability traces used - -*Note: Detailed Phase 2 implementation will be provided after Phase 1 sign-off.* - -### Phase 3: Experiment Orchestration (Week 3) -**Goal**: Complete end-to-end experiment management with reproducibility - -**Prerequisites**: Phase 2 complete with parallel validation passing - -**Status**: ✅ **COMPLETE** - -**Tasks**: -1. **Implement `launch/run_experiment.py`** ✅ - - ExperimentRunner class for orchestration - - Aggregator lifecycle management (spawn, wait, terminate) - - Trainer spawning with combined log capture - - Process monitoring and graceful cleanup (Ctrl+C handling) - - Batch experiment execution - -2. **Create experiment configuration system** ✅ - - `launch/experiment_config.py` - Config schema (DatasetConfig, AvailabilityConfig, etc.) - - `experiments/configs/test_oort_syn0.yaml` - Single experiment example - - `experiments/configs/oort_n300_all4unavail.yaml` - Batch of 4 experiments - -3. **Implement reproducibility snapshot** ✅ - - `launch/snapshot.py` - ExperimentSnapshot class - - Captures spawn commands (not rendered configs) - - Records git commit, branch, dirty state - - Computes SHA256 checksums for all metadata files - - Copies metadata and aggregator config to experiment directory - -4. **Implement aggregator spawner** ✅ - - `launch/aggregator_spawner.py` - AggregatorSpawner class - - Process management with log capture - - wait_until_ready() with configurable warmup time - - Clean termination - -5. **Enhanced logging system** ✅ - - Descriptive timestamped log filenames: `DD_MM_YY_HH_MM_.log` - - Example: `03_02_26_14_30_oort_n300_oracular_alpha0p1_syn0_aggregator.log` - - Separate logs: aggregator.log + trainers.log (combined) - - Line-buffered output for real-time monitoring - - Context-free filenames (understand experiment from filename alone) - -**Validation**: Run full experiment end-to-end with new system -- Test with 5 trainers: `./scripts/test_phase3.sh` -- Verify logs captured correctly -- Test Ctrl+C cleanup -- Run full 300-trainer batch experiment - -**Deliverables**: -- ✅ `launch/run_experiment.py` (262 lines) -- ✅ `launch/experiment_config.py` (157 lines) -- ✅ `launch/snapshot.py` (162 lines) -- ✅ `launch/aggregator_spawner.py` (131 lines) -- ✅ `launch/spawner.py` (enhanced with log file support) -- ✅ `experiments/configs/` (example configs) -- ✅ `scripts/test_phase3.sh` (test script) -- ✅ `PHASE3_README.md` (comprehensive documentation) - -**Output Structure**: -``` -experiments/ - run_20260203_143000_test_oort_syn0/ - snapshot.yaml # Reproducibility info - 03_02_26_14_30_oort_n10_oracular_alpha0p1_syn0_aggregator.log - 03_02_26_14_30_oort_n10_oracular_alpha0p1_syn0_trainers.log - metadata/ # Copied from metadata/ - asyncfl_oort_agg.json # Copied aggregator config -``` - -### Phase 4: Cleanup & Documentation (Week 4) -**Goal**: Remove legacy files and update documentation - -**Tasks**: -1. **Archive old configs** - - Move `config_dir*` to `trainer/legacy_configs/` (Git history preserved) - - Keep one example config for reference - -2. **Update shell script** - - Simplify to call Python experiment runner - - Remove bash trainer spawning logic - -3. **Documentation** - - Update `README.md` with new workflow - - Add metadata file format documentation - - Create developer guide for adding new experiments - -4. **Validation script** - - `scripts/validate_metadata.py` - check consistency - - Unit tests for spawner - ---- - -## Code Changes Required - -### 1. `trainer/pytorch/main.py` -**Current**: -```python -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--config', required=True) # File path only - args = parser.parse_args() - config = Config(args.config) # Loads from file -``` - -**New**: -```python -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('--config', help='Config file path (legacy)') - parser.add_argument('--config-dict', help='Config as JSON string') - parser.add_argument('--trainer-id', type=int, help='Trainer ID for metadata lookup') - parser.add_argument('--metadata-dir', help='Path to metadata directory') - args = parser.parse_args() - - if args.config_dict: - config = Config(config_dict=json.loads(args.config_dict)) - elif args.trainer_id and args.metadata_dir: - # Load from metadata + generate config - config = load_from_metadata(args.trainer_id, args.metadata_dir) - else: - config = Config(args.config) # Legacy file-based -``` - -### 2. `aggregator/pytorch/main_oort_agg.py` -**Minimal changes**: Already uses single config file - no spawning needed - -### 3. `flame/config.py` -**Add support for dict-based config**: -```python -class Config: - def __init__(self, config_path=None, config_dict=None): - if config_dict: - self._config = config_dict - elif config_path: - self._config = self._load_from_file(config_path) - else: - raise ValueError("Must provide config_path or config_dict") -``` - ---- - -## Benefits - -### Immediate -- **99% reduction** in config file size (5,924 files → ~10 metadata files) -- **Easy comparison**: `diff` two YAML files instead of 300 pairs -- **Single source of truth**: Update trainer delay once, affects all experiments -- **Git-friendly**: Small YAML diffs instead of 300-file changes - -### Long-term -- **Scalability**: Add new dataset (MNIST, Speech) with one split file -- **Reproducibility**: Experiment definition in one YAML file -- **Debugging**: Clear separation of static metadata vs experiment params -- **Testing**: Mock metadata for unit tests -- **Extensibility**: Easy to add new trainer properties (battery level, network speed) - ---- - -## Migration Risks & Mitigation - -### Risk 1: Data Loss During Extraction -**Mitigation**: -- Keep all original configs in Git history -- Validation script to compare old vs new configs -- Parallel run tests (Phase 2) - -### Risk 2: Breaking Existing Experiments -**Mitigation**: -- Maintain backward compatibility (file-based configs) -- Phased rollout with validation at each step -- Feature flag to switch between old/new system - -### Risk 3: Performance Overhead -**Mitigation**: -- Config generation is O(1) per trainer (fast) -- Spawn delay already 1s per trainer (no change) -- Benchmark before/after - ---- - -## Success Criteria - -1. **Functionally equivalent**: New system produces identical experiment results -2. **Reduced complexity**: <1% of original config file count -3. **Maintainability**: Can add ne Gate | -|-------|----------|-------------|------| -| **1. Metadata Extraction** | **3 days** | YAML metadata files + validation | ✓ Validation script passes | -| 2. Programmatic Spawning | 5 days | `spawner.py` + refactored `main.py` | ✓ Parallel validation passes | -| 3. Experiment Runner | 5 days | `run_experiment.py` + experiment configs | ✓ Full experiment succeeds | -| 4. Cleanup & Docs | 5 days | Updated docs, archived legacy, tests | ✓ Review complete | -| **Total** | **18 days** | Production-ready refactored system | | - ---- - -## Execution Plan - -### Immediate Next Steps (START HERE) - -1. **Create extraction script** (`scripts/extract_metadata.py`) - - Copy full implementation from Phase 1 Step 1.1 above - - Run and verify output files generated - -2. **Create validation script** (`scripts/validate_metadata.py`) - - Copy full implementation from Phase 1 Step 1.2 above - - Run and verify all checks pass - -3. **Phase 1 Gate Review** - - Complete Step 1.3 checklist - - Manual spot-check 5 random trainers - - Git commit with Phase 1 sign-off - -4. **Phase 1 Approval Required** - - User review of generated YAML files - - Confirmation all validation checks pass - - Sign-off before Phase 2 detailed plan - -### After Phase 1 Approval - -- Detailed Phase 2 implementation plan will be provided -- Phase 3 implementation plan follows Phase 2 approval -- Iterative approach ensures no wasted effort on incorrect architecture - -## Next Steps - -1. **Review this plan** - Get feedback on architecture -2. **Proof of concept** - Implement Phase 1 for 10 trainers -3. **Validate extraction** - Ensure no data loss -4. **Iterate** - Refine based on learnings -5. **Full rollout** - Complete all phases - ---- - -## Questions for Discussion - -1. **YAML vs JSON**: Prefer YAML for readability, but JSON for Python compatibility? -2. **Metadata location**: `examples/metadata/` vs `lib/python/flame/metadata/`? -3. **Backward compatibility**: Keep file-based configs indefinitely or deprecate after migration? -4. **Testing strategy**: Unit tests, integration tests, or manual validation sufficient? -5. **Config assembly**: Runtime generation vs pre-generated temp files? -6. **Extension to other examples**: Apply same pattern to `async_mnist`, `async_google_speech`? - ---- - -**Document Version**: 1.0 -**Last Updated**: February 2, 2026 -**Author**: Based on codebase analysis and user requirements - ---- - -## MIGRATION COMPLETE - -### Status: ✅ All Phases Complete (February 3, 2026) - -**Phase 1**: ✅ Metadata extraction and validation (COMPLETE) -**Phase 2**: ✅ Programmatic spawning (COMPLETE) -**Phase 3**: ✅ Experiment orchestration (COMPLETE) - -### Results - -- **Config reduction**: 5,924 files → 7 YAML files + 5 Python scripts (99.9% reduction) -- **Lines of code**: 1.8M lines → ~1,200 lines (99.9% reduction) -- **Maintainability**: Single source of truth for all trainer properties -- **Reproducibility**: Full snapshot system with git tracking and checksums - ---- - -## Usage Guide - -### Quick Start - -**Test with 5 trainers:** -```bash -cd lib/python/examples/async_cifar10 -./scripts/test_phase3.sh -``` - -**Run single experiment (10 trainers):** -```bash -python3 launch/run_experiment.py experiments/configs/test_oort_syn0.yaml -``` - -**Run full batch (300 trainers × 4 experiments):** -```bash -python3 launch/run_experiment.py experiments/configs/oort_n300_all4unavail.yaml -``` - -### Experiment Configuration Format - -Create YAML files in `experiments/configs/`: - -```yaml -experiments: - - name: my_experiment - description: "Optional description" - - trainer: - num_trainers: 50 # Number of trainers to spawn - start_id: 1 # Starting trainer ID (1-300) - dataset: - dirichlet_alpha: 0.1 # 0.1, 1.0, 10, or 100 - availability: - mode: syn_0 # syn_0, syn_20, syn_50, mobiperf_2st, etc. - - aggregator: - selector: oort - tracking_mode: oracular - config_template: expt_scripts_2026/configs/oort_n300_oracular_9may25_syn0.json - - execution: - num_gpus: 8 - sleep_between_spawns: 2.0 - aggregator_warmup_time: 5.0 -``` - -**Batch experiments** - just add multiple experiments to the same file: -```yaml -experiments: - - name: exp1 - # config - - name: exp2 - # config -``` - -### Output Structure - -Each experiment creates: -``` -experiments/run_YYYYMMDD_HHMMSS_/ - snapshot.yaml # Reproducibility info - DD_MM_YY_HH_MM__aggregator.log # Aggregator log - DD_MM_YY_HH_MM__trainers.log # Combined trainers log - aggregator_config.json # Copy of aggregator config -``` - -**Log filename format**: `DD_MM_YY_HH_MM__n__alpha__.log` - -Example: `03_02_26_14_30_oort_n300_oracular_alpha0p1_syn0_aggregator.log` - -### Snapshot File - -The `snapshot.yaml` contains everything needed for reproduction: - -```yaml -experiment: # Full experiment config -spawn_commands: # Exact commands used - aggregator: [...] - trainers: [...] -git_info: # Git commit, branch, dirty state - commit: abc123 - branch: dg/simplify_cifar10_expts - dirty: false -metadata_location: # Path to metadata directory -metadata_checksums: # SHA256 hashes of all metadata files - trainer_registry.yaml: sha256:... - dataset_splits/...: sha256:... -``` - -### Monitoring - -**Live tail logs:** -```bash -# Aggregator -tail -f experiments/run_*/03_02_26_*_aggregator.log - -# Trainers -tail -f experiments/run_*/03_02_26_*_trainers.log -``` - -**Search logs:** -```bash -# Find specific trainer -grep "trainer_042" experiments/run_*/03_02_26_*_trainers.log - -# Find errors -grep -i error experiments/run_*/03_02_26_*.log -``` - -**Graceful termination:** -Press `Ctrl+C` to cleanly stop all processes, close log files, and exit. - -### Availability Modes - -**Synthetic** (uniform across all trainers): -- `syn_0`: Always available (0% failure) -- `syn_20`: 20% unavailability -- `syn_50`: 50% unavailability - -**MobiPerf** (real-world, per-trainer): -- `mobiperf_2st`: 2-state traces -- `mobiperf_3st_50`: 3-state with 50% battery threshold -- `mobiperf_3st_75`: 3-state with 75% battery threshold - -### Dirichlet Alpha Values - -From `metadata/dataset_splits/`: -- `0.1`: Highly non-IID (heterogeneous data) -- `1.0`: Moderately non-IID -- `10.0`: Slightly non-IID -- `100.0`: Nearly IID (homogeneous data) - -### Troubleshooting - -**Aggregator fails to start:** -```bash -# Check MQTT broker -sudo systemctl status mosquitto -sudo systemctl start mosquitto - -# Check log -cat experiments/run_*/03_02_26_*_aggregator.log -``` - -**Trainers fail to spawn:** -```bash -# Verify metadata exists -ls metadata/trainer_registry.yaml - -# Verify base config exists -ls configs/trainer_base.yaml - -# Check trainer log -cat experiments/run_*/03_02_26_*_trainers.log -``` - -**Clean up hung processes:** -```bash -# Kill all trainers -pkill -f "trainer/pytorch/main.py" - -# Kill aggregator -pkill -f "aggregator/pytorch/main_oort_agg.py" -``` - ---- - -## Architecture Details - -### Component Overview - -**Phase 1 - Metadata System:** -- `scripts/extract_metadata.py` (293 lines) - Automated extraction -- `scripts/validate_metadata.py` (247 lines) - 6 validation checks -- `metadata/` (7 YAML files) - Single source of truth - -**Phase 2 - Programmatic Spawning:** -- `launch/spawner.py` (379 lines) - MetadataLoader, ConfigGenerator, TrainerSpawner -- `configs/trainer_base.yaml` - Minimal template (75% smaller) -- `trainer/pytorch/main.py` (modified) - Handles JSON configs + list/string traces - -**Phase 3 - Experiment Orchestration:** -- `launch/run_experiment.py` (262 lines) - ExperimentRunner with signal handling -- `launch/experiment_config.py` (157 lines) - Configuration schema -- `launch/snapshot.py` (176 lines) - Reproducibility snapshots -- `launch/aggregator_spawner.py` (131 lines) - Aggregator management -- `experiments/configs/` - Experiment definitions - -### Key Design Decisions - -1. **Metadata location reference (not copy)**: Metadata doesn't change across runs, so snapshots only record the location and checksums instead of duplicating files - -2. **Spawn commands (not configs)**: Snapshots save the actual commands used to spawn processes, not the generated configs. Configs are deterministically generated from metadata. - -3. **Combined trainer log**: All trainer output goes to single log file with line buffering, easier to monitor than 300 separate files - -4. **Context-free log filenames**: Timestamp + experiment parameters in filename means you can understand the experiment without opening files - -5. **Positional aggregator arg**: Aggregator expects positional config path argument, not `--config` flag - -6. **List/string trace handling**: Trainer code handles both formats - lists from JSON configs, strings from file-based configs - -### File Generation Flow - -``` -Experiment Config (YAML) - ↓ -ExperimentRunner - ↓ -MetadataLoader → loads 7 YAML files - ↓ -ConfigGenerator → combines base + metadata - ↓ -TrainerSpawner → spawns processes with JSON config string - ↓ -Trainer main.py → receives JSON, creates temp file for Config class -``` - -### Validation Results - -**Phase 1**: ✅ All 6 validation checks passed -- Trainer registry completeness -- Dataset splits consistency -- Availability traces correctness -- Cross-references valid -- No missing trainers -- 100% coverage - -**Phase 2**: ✅ Config comparison validation passed -- Generated configs match original JSONs exactly -- All 300 trainers validated -- Test with 5 trainers successful - -**Phase 3**: ✅ End-to-end testing -- 5-trainer test: Successful -- Snapshot creation: Verified -- Log capture: Working -- Ctrl+C cleanup: Clean - ---- - -## Migration Timeline - -| Date | Phase | Activity | Status | -|------|-------|----------|--------| -| Feb 1, 2026 | Planning | Architecture design | ✅ | -| Feb 2, 2026 | Phase 1 | Metadata extraction & validation | ✅ | -| Feb 2, 2026 | Phase 2 | Programmatic spawning | ✅ | -| Feb 3, 2026 | Phase 3 | Experiment orchestration | ✅ | -| Feb 3, 2026 | Testing | End-to-end validation | ✅ | - -**Total implementation time**: 3 days (faster than planned 18 days due to iterative approach) - ---- - -## Benefits Achieved - -1. **Maintenance**: 99.9% reduction in config files -2. **Consistency**: Single source of truth eliminates drift -3. **Scalability**: Add new experiments by editing YAML, not creating 300 JSONs -4. **Reproducibility**: Full snapshot system captures exact experiment state -5. **Debugging**: Combined logs and descriptive filenames -6. **Flexibility**: Easy to change parameters, add trainers, modify traces - ---- - -## Future Enhancements - -1. **Snapshot reproduction**: `--from-snapshot` flag to re-run experiments -2. **WandB integration**: Automatic experiment logging -3. **Progress monitoring**: Real-time dashboard -4. **Parallel experiments**: Run multiple experiments simultaneously -5. **Cloud deployment**: AWS/GCP integration -6. **Apply to other examples**: Extend pattern to async_mnist, etc. - ---- - -**Document Version**: 2.0 -**Last Updated**: February 3, 2026 -**Status**: Migration Complete - System in Production - ---- - -## PHASE 3 STATUS UPDATE - February 3, 2026 Evening - -### Current Implementation Status - -**Phase 1**: ✅ COMPLETE - Metadata extraction validated -**Phase 2**: ✅ COMPLETE - Programmatic spawning implemented -**Phase 3**: ⚠️ IMPLEMENTED BUT UNTESTED - Needs validation and refinement - -### Critical Gap Identified - -**User Requirements Review:** -1. ✅ "Programmatic configuration generation at runtime" - ACHIEVED -2. ⚠️ "One YAML file per execution with all arguments for reproducibility" - NEEDS IMPLEMENTATION -3. ⚠️ "Underlying launcher scripts remain the same" - NEEDS CLARIFICATION - -### The Core Issue - -Current snapshot system saves metadata checksums but lacks a clear execution record. - -**User requirement**: One YAML file per execution with all arguments to reproduce experiments. - -### Optimized Solution: Compact Execution Config with Metadata Keys - -**Key insight**: Since metadata is organized with named keys (e.g., `"syn_0"`, `"cifar10_alpha0.1_n300"`), we only need to store the **keys**, not the full data. This keeps execution configs: -- **Readable**: ~1-2KB, not 50-100KB -- **Editable**: Easy to modify for next run -- **Standard**: Uses same labels as launcher scripts and metadata -- **Reproducible**: Git commit + metadata keys = exact reproduction - -Each experiment run generates `execution_config.yaml`: -```yaml -# Compact, readable execution record (~1-2KB) -execution_timestamp: "2026-02-03T14:30:00" - -git_info: - commit: "abc123def456" - branch: "dg/simplify_cifar10_expts" - dirty: false - -# Metadata references by KEY (not full data) -metadata_refs: - trainer_registry: "metadata/trainer_registry.yaml" - dataset_split_key: "cifar10_alpha0.1_n300" # Key in dataset_splits/ - availability_trace_key: "syn_0" # Key in availability_traces/ - -# Aggregator config by PATH (not embedded) -aggregator: - config_file: "expt_scripts_2026/configs/oort_n300_oracular_9may25_syn0.json" - selector: "oort" - tracking_mode: "oracular" - agg_goal: 10 - -# Experiment parameters (what changes between runs) -experiment: - name: "oort_n300_alpha0.1_syn0" - description: "Oort with 300 trainers, alpha=0.1, synthetic 0% unavailability" - - trainer: - num_trainers: 300 - start_id: 1 - trainer_id_range: [1, 300] - dataset: - alpha: 0.1 - split_key: "cifar10_alpha0.1_n300" - availability: - mode: "syn_0" - trace_key: "syn_0" - battery_threshold: 50 - speedup_factor: 1.0 - - execution: - num_gpus: 8 - sleep_between_spawns: 1.0 - aggregator_warmup_time: 10 - -# Exact spawn commands used -spawn_commands: - aggregator: ["python3", "aggregator/pytorch/main_oort_agg.py", "expt_scripts_2026/configs/oort_n300_oracular_9may25_syn0.json"] - trainers: ["python3", "launch/spawner.py", "--alpha", "0.1", "--availability", "syn_0", "--num-trainers", "300", "--start-id", "1", "--num-gpus", "8"] -``` - -**Benefits**: -1. **Compact**: ~1-2KB vs 50-100KB -2. **Human-readable**: Clear experiment parameters -3. **Easy to edit**: Change keys/params for next run -4. **Uses standard labels**: Same as launcher scripts and metadata -5. **Git-based reproducibility**: Checkout commit → load metadata by keys → spawn -6. **No duplication**: Metadata folder is single source of truth - -### Implementation Plan - -#### Task 1: Create `launch/execution_config_generator.py` -**Purpose**: Generate compact execution configs with metadata keys -**Effort**: 1-2 hours - -Key function: -```python -def create_execution_config(exp_config: ExperimentConfig, - aggregator_config_path: Path, - spawn_commands: Dict) -> Dict: - """Generate compact execution config using metadata keys.""" - import subprocess - - return { - 'execution_timestamp': datetime.now().isoformat(), - 'git_info': { - 'commit': subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode().strip(), - 'branch': subprocess.check_output(['git', 'branch', '--show-current']).decode().strip(), - 'dirty': len(subprocess.check_output(['git', 'status', '--porcelain']).decode()) > 0 - }, - 'metadata_refs': { - 'trainer_registry': 'metadata/trainer_registry.yaml', - 'dataset_split_key': f"cifar10_alpha{exp_config.trainer.dataset.dirichlet_alpha}_n300", - 'availability_trace_key': exp_config.trainer.availability.mode, - }, - 'aggregator': { - 'config_file': str(aggregator_config_path), - 'selector': exp_config.aggregator.selector, - 'tracking_mode': exp_config.aggregator.tracking_mode, - 'agg_goal': exp_config.aggregator.agg_goal, - }, - 'experiment': { - 'name': exp_config.name, - 'description': exp_config.description, - 'trainer': { - 'num_trainers': exp_config.trainer.num_trainers, - 'start_id': exp_config.trainer.start_id, - 'trainer_id_range': [exp_config.trainer.start_id, - exp_config.trainer.start_id + exp_config.trainer.num_trainers - 1], - 'dataset': { - 'alpha': exp_config.trainer.dataset.dirichlet_alpha, - 'split_key': f"cifar10_alpha{exp_config.trainer.dataset.dirichlet_alpha}_n300", - }, - 'availability': { - 'mode': exp_config.trainer.availability.mode, - 'trace_key': exp_config.trainer.availability.mode, - }, - 'battery_threshold': exp_config.trainer.battery_threshold, - 'speedup_factor': exp_config.trainer.speedup_factor, - }, - 'execution': { - 'num_gpus': exp_config.execution.num_gpus, - 'sleep_between_spawns': exp_config.execution.sleep_between_spawns, - 'aggregator_warmup_time': exp_config.execution.aggregator_warmup_time, - } - }, - 'spawn_commands': spawn_commands - } -``` - -#### Task 2: Create `launch/reproduce.py` -**Purpose**: Reproduce experiments from compact execution configs -**Effort**: 1-2 hours - -Usage: -```bash -python3 launch/reproduce.py experiments/run_20260203_143000/execution_config.yaml -``` - -Loads metadata using keys, spawns using exact parameters from config: -```python -def reproduce_from_config(exec_config_path: Path): - """Reproduce experiment using metadata keys from execution config.""" - with open(exec_config_path) as f: - config = yaml.safe_load(f) - - # Warn if git state differs - current_commit = get_git_commit() - if current_commit != config['git_info']['commit']: - print(f"⚠️ Git commit mismatch: {current_commit[:7]} vs {config['git_info']['commit'][:7]}") - - # Load metadata using keys - example_dir = Path(__file__).parent.parent - metadata_loader = MetadataLoader(example_dir / 'metadata') - - # Extract experiment params - exp = config['experiment'] - - # Spawn aggregator - agg_config_path = example_dir / config['aggregator']['config_file'] - agg_spawner = AggregatorSpawner(...) - agg_spawner.spawn(agg_config_path) - - # Spawn trainers using metadata keys - config_gen = ConfigGenerator(metadata_loader, ...) - trainer_spawner = TrainerSpawner(config_gen, ...) - trainer_spawner.spawn_all(2-3 hours (much faster with compact format!) -- **Testing & debugging (Task 4)**: 2-4 hours -- **Optional bash wrapper (Task 5)**: 30 minutes -- **Total**: 4.5-7.5 hours of focused work - -### Recommendation - -Focus on **Tasks 1-4** in this order: -1. **Implement execution_config_generator.py** (1-2 hours) - create compact configs with metadata keys -2. **Integrate into run_experiment.py** (30 min) - auto-generate on each run -3. **Test mini experiment** (1-2 hours) - verify execution_config.yaml is readable and complete -4. **Implement reproduce.py** (1-2 hours) - parse keys and spawn processes -5. **Test reproduction workflow** (1 hour) - ensure exact reproducibility -6. **Scale to 300 trainers** (1 hour) - validate full-scale experiments - -**Key advantages of compact format**: -- ✅ Faster implementation (no data embedding logic) -- ✅ Easier to read/edit/debug -- ✅ Standard labels match launcher scripts -- ✅ Single source of truth (metadata folder) -- ✅ Git-based versioning works naturally - -exec_config = create_execution_config( - exp_config, - aggregator_config_path, - spawn_commands={ - 'aggregator': aggregator_cmd, - 'trainers': trainer_cmd - } -) -exec_config_file = self.current_exp_dir / 'execution_config.yaml' -with open(exec_config_file, 'w') as f: - yaml.dump(exec_config, f, default_flow_style=False, sort_keys=False) -print(f" ✓ Saved execution config: {exec_config_file}") -``` - -#### Task 4: Test End-to-End -**Purpose**: Validate full workflow -**Effort**: 2-4 hours - -1. Run mini experiment (5 trainers) -2. Verify execution_config.yaml generated -3. Test reproduction -4. Debug any issues - -#### Task 5: Bash Script Compatibility (Optional) -**Purpose**: Keep existing script interface -**Effort**: 1 hour - -Wrapper: -```bash -#!/bin/bash -# oort_n300_all4unavail.sh - New implementation -python3 launch/run_experiment.py experiments/configs/oort_n300_all4unavail.yaml -``` - -### Testing Roadmap - -**Phase 3a: Mini Test (5 trainers)** -```bash -python3 launch/run_experiment.py experiments/configs/test_phase3_mini.yaml -``` -compact, ~1-2KB, human-readable) -- [ ] Config uses metadata keys (not embedded data) -- [ ] Config includes exact spawn commands -- [ ] Reproduction from config -- [ ] Aggregator spawns -- [ ] Trainers connect -- [ ] Logs captured -- [ ] execution_config.yaml created (self-contained) -- [ ] Reproduction works -- [ ] Ctrl+C cleanup works - -**Phase 3b: Full Test (300 trainers)** -```bash -python3 launch/run_experiment.py experiments/configs/oort_n300_all4unavail.yaml -``` - -### What's Already Working - -1. **Compact execution configs with metadata keys**: Not yet implemented (1-2 hours) -2. **Reproduction tool using metadata keys**: Not yet implemented (1-2 hours) -3. **End-to-end testing**: Not yet validated (2-4 hours) -4. **Bash script wrapper** (optional): Thin wrapper for compatibility (30 min) - -**Design decision made**: Use compact format with metadata keys instead of embedding all dataefinitions - -### What Needs Work - -1. **Self-contained execution configs**: Not yet implemented -2. **Reproduction tool**: Not yet implemented -3. **End-to-end testing**: Not yet validated -4. **Bash script bridge**: Decision needed on interface - -### Estimated Completion Time - -- **Core functionality (Tasks 1-3)**: 4-6 hours -- **Testing & debugging (Task 4)**: 2-4 hours -- **Optional bash wrapper (Task 5)**: 1 hour -- **Total**: 7-11 hours of focused work - -### Recommendation - -Focus on **Tasks 1-4** in this order: -1. Implement execution_snapshot.py (embed all data) -2. Integrate into run_experiment.py (auto-generate on each run) -3. Test mini experiment (5 trainers) -4. Implement reproduce.py (once we see real execution_config.yaml) -5. Test reproduction workflow -6. Scale to 300 trainers - -This approach ensures we build reproducibility correctly before scaling. - ---- - -**Document Status**: Phase 3 gap analysis complete, action plan defined -**Updated**: February 3, 2026, 8:00 PM -**Next Action**: Implement execution_snapshot.py - - ---- - -## DESIGN REFINEMENT - February 3, 2026 (Final Update) - -### Optimized Approach: Compact Execution Configs - -**Key insight from user feedback**: Don't embed all metadata inline. Use **keys** instead. - -#### Why This Is Better - -**Previous approach** (embedded data): -- ❌ 50-100KB files with duplicate data -- ❌ Hard to read and edit -- ❌ Doesn't leverage organized metadata structure - -**Optimized approach** (metadata keys): -- ✅ 1-2KB readable configs -- ✅ Uses same labels as launcher scripts -- ✅ Easy to modify for next run -- ✅ Metadata folder is single source of truth -- ✅ Git tracks both metadata and execution configs together - -#### Execution Config Format - -```yaml -# execution_config.yaml - Compact and readable (~1-2KB) - -execution_timestamp: "2026-02-03T14:30:00" - -git_info: - commit: "abc123def456" - branch: "dg/simplify_cifar10_expts" - dirty: false - -# References to metadata by KEY (not embedded!) -metadata_refs: - trainer_registry: "metadata/trainer_registry.yaml" - dataset_split_key: "cifar10_alpha0.1_n300" # Key to look up in dataset_splits/ - availability_trace_key: "syn_0" # Key to look up in availability_traces/ - -# Aggregator config by PATH (not embedded!) -aggregator: - config_file: "expt_scripts_2026/configs/oort_n300_oracular_9may25_syn0.json" - selector: "oort" - tracking_mode: "oracular" - agg_goal: 10 - -# Experiment parameters -experiment: - name: "oort_n300_alpha0.1_syn0" - trainer: - num_trainers: 300 - start_id: 1 - dataset: {alpha: 0.1, split_key: "cifar10_alpha0.1_n300"} - availability: {mode: "syn_0", trace_key: "syn_0"} - battery_threshold: 50 - speedup_factor: 1.0 - execution: - num_gpus: 8 - sleep_between_spawns: 1.0 - aggregator_warmup_time: 10 - -# Exact commands used -spawn_commands: - aggregator: ["python3", "aggregator/pytorch/main_oort_agg.py", "..."] - trainers: ["python3", "launch/spawner.py", "--alpha", "0.1", ...] -``` - -#### Reproduction Workflow - -1. **Read execution_config.yaml** -2. **Checkout git commit** (if different from current) -3. **Load metadata using keys**: - - `dataset_split_key` → load from `metadata/dataset_splits/` - - `availability_trace_key` → load from `metadata/availability_traces/` -4. **Load aggregator config** using `config_file` path -5. **Spawn processes** using experiment parameters - -#### Why This Optimizes for User Goals - -1. **Programmatic generation**: ✅ Configs generated at runtime from metadata -2. **Reproducibility**: ✅ Git commit + metadata keys = exact reproduction -3. **Minimal changes between runs**: ✅ Edit keys/parameters, not regenerate files -4. **Readable and informative**: ✅ Uses standard labels from launcher scripts -5. **Easy parsing**: ✅ Simple YAML structure for next launch - -#### Implementation Simplicity - -**Compact format is faster to implement**: -- No data embedding logic needed -- No data extraction from embedded configs -- Simple key lookups in metadata folder -- Standard YAML read/write - -**Estimated total time**: 4.5-7.5 hours (vs 7-11 hours for embedded approach) - ---- - -**Migration Plan Status**: Updated with optimized compact execution config approach -**Ready for**: Implementation of execution_config_generator.py -**Updated**: February 3, 2026, 8:30 PM - diff --git a/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md b/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md new file mode 100644 index 000000000..4cfebf95a --- /dev/null +++ b/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md @@ -0,0 +1,199 @@ +# Plan: Trainer/Aggregator Telemetry + Virtual-Clock Speedup (async_cifar10) + +## Context + +Two related needs in the `async_cifar10` FL example. **Order of work: Task 1 (telemetry) first, then Task 2 (speedup) later.** + +1. **Telemetry & visualization.** Today we only see loss/accuracy/round-time. We can't see *how the selector chooses* (utility vs. speed tradeoff, what it deems available/eligible, how it uses current resources) or *how the aggregator aggregates* (staleness, agg-goal progress, participation). The existing `scripts/plotters/` is a manual, multi-step, hardcoded-path pipeline tuned for an old "fwdllm" log schema; it scrapes regexes out of free-text logs and is not run automatically post-run. We want robust, structured telemetry that supports apples-to-apples comparison **across selector implementations** and auto-generated plots after a run. + +2. **Speedup that actually speeds up.** `speedup_factor=2` did not yield ~2x wall-clock. Root cause: slow-trainer dynamics are simulated with a literal `time.sleep()` after the GPU finishes, and several fixed waits/polls ignore `speedup_factor` entirely. We want to **keep the dynamics** (slow trainers delay their updates; selection still trades off utility vs. speed; staleness still accrues) while **eliminating real wall-clock spent sleeping**. The chosen approach is a **virtual (simulated) clock**: trainers do real GPU compute but report a *simulated completion time* instead of sleeping; the aggregator orders/commits updates by simulated time. + +**Scope / reuse principle:** implement for `async_cifar10` first, but place all shared logic in the `flame` library base classes (aggregator / selector / trainer mixins) so other examples inherit it without duplication. The example only wires config + example-specific hooks. + +--- + +# Task 1 — Structured telemetry + plots (do first) + +### Structured emission (library layer, generic) +- New **`flame/telemetry/` module**: a lightweight `TelemetryWriter` that appends **JSONL event records** to a per-run file (one schema, typed events). Inject via base aggregator/selector/trainer so every example emits the same schema → cross-selector comparison is free, no regex scraping. +- **Event types to emit** (most state already exists; we structure it, not recompute): + - *Selector*: per-`select()` — candidates, eligible set, availability composition (counts of `AVL_TRAIN`/`AVL_EVAL`/`UN_AVL` from `PROP_AVL_STATE`), per-trainer utility/speed used, chosen set, explore-vs-exploit split. Hook the base `AbstractSelector` ([selector/__init__.py](../../flame/selector/__init__.py)) + each concrete selector's decision point. + - *Aggregator*: per-round loss/accuracy/round-time (already logged), plus staleness distribution, agg-goal progress, per-trainer participation, in-flight count (async). + - *Trainer*: per-round real GPU time vs. `sim_round_duration`, wait time, availability state, samples visible (streaming). This is also the data that proves the Task 2 speedup works. + - *Streaming utility disparity (new)*: with data streaming enabled, emit at fine resolution the trainer's **statistical utility on the currently-unlocked prefix** (the real, time-T value) **alongside a counterfactual utility computed over the full dataset** as if it were unlocked from the start, plus their ratio and the visible-sample fraction. Reuse the existing utility computation (Oort `I_m` / `reset_stat_utility` / `fetch_statistical_utility`) and the streaming machinery (`_visible_sample_count`, `_rebuild_stream_loader`, [trainer/pytorch/main.py:428](trainer/pytorch/main.py#L428)) — build a full-pool loader for the counterfactual forward pass. **Configurable** (`util_counterfactual: {mode, every_n_rounds, sample_size}`), **default subsample + every-N-rounds**, gated off otherwise, so the extra forward pass is opt-in. This surfaces the hypothesized gap between real-world streamed utility and the full-dataset assumption. +- Keep `wandb` as-is for live dashboards; JSONL is the source of truth for offline/paper plots. + +### Plotting (fresh module, reuse helpers) +- New **`scripts/analysis/analyze_run.py`** (or `flame/telemetry/plots.py`): a single `analyze_run.py ` that reads the JSONL and emits a PNG bundle + a small summary. **Reuse the existing matplotlib helpers** rather than rewriting: CDF/percentile rendering from `scripts/plotters/cdf_plot.py` (`generateCDF`), multi-run band/interpolation logic from `scripts/plotters/comparative_plotter.py`, and stacked-bar (train vs. stall) from `scripts/plotters/stall_stacked_bar_plot.py`. +- **Plots:** loss/accuracy/round-time over time; selector availability composition over rounds (stacked); utility-vs-speed scatter of selected vs. eligible; selection-frequency/fairness per trainer; staleness CDF; agg-goal/in-flight timeline; trainer real-GPU vs. simulated-time stacked bar (the speedup evidence); **streamed-vs-full utility disparity over time** (per-trainer real vs. counterfactual utility + ratio). A `--compare run_dirs...` mode overlays selectors. +- **Auto post-run:** invoke `analyze_run.py` from the example's launch/teardown so plots land in `/plots/` automatically. + +--- + +# Task 2 — Virtual-clock speedup (do later) + +## Implementation status (2026-05-28) + +Implemented and unit-verified (156 tests green incl. `tests/sim` + `tests/mode`): +- `speedup_factor` removed everywhere; replaced by `time_mode: real|simulated` (default `simulated`) plumbed through experiment_config → spawner (`--time_mode`) → trainer, and into the aggregator config (`hyperparameters.time_mode`). YAML `speedup_factor` lines stripped from ~20 configs. +- Contract: `MessageType.SIM_SEND_TS / SIM_COMPLETION_TS / SIM_ROUND_DURATION`, `PROP_SIM_SEND_TS / PROP_SIM_COMPLETION_TS`, and pure `flame/sim/virtual_clock.py` (`VirtualClock`, `SimReorderBuffer`, `sim_ordered_ends`). +- Trainer: `real` sleeps `D`; `simulated` skips sleeps, reports `sim_completion_ts`/`D`, evaluates availability + streaming against sim-time (`_sim_now`), and skips (rather than busy-waits) when unavailable in sim mode. +- Async aggregator: stamps `sim_send_ts` on distribute; `simulated` commits via `_sim_recv_min` (reorder buffer → ascending `sim_completion_ts`, advances `T_v`) and sources `PROP_ROUND_DURATION` from `D`; `real` keeps arrival-order FIFO. Bounded-recv guard retained. + +Remaining / known gaps: +- **Sync aggregator**: init + `sim_send_ts` stamping done, but the `first_k`-by-`sim_completion_ts` selection and sim-sourced round duration are **not** implemented (sync still uses arrival order in sim mode). The primary target (felix) is async; sync (oort/refl/fedavg) sim-ordering is a follow-up. +- **End-to-end two-run verification** (Q1: `real` vs `simulated` per-round parity + speedup) needs a GPU/MQTT run — not yet executed. +- **Availability in sim mode** uses the trainer's own trace evaluated at the stamped sim-time; correlated-trace / oracular-at-`T_v` selection nuances (and the client-notify-vs-oracular question) want runtime validation. +- `examples/fwdllm` still hardcodes `speedup_factor=1.0` (separate example; left as-is). + +## Why virtual-clock is the right call (correctness of reordering) + +Concern: the async aggregator pops updates FIFO **by physical arrival** (`channel.recv_fifo`, [channel.py:430](../../flame/channel.py#L430)) — if we stop sleeping, all trainers finish ~together and arrival order no longer reflects simulated speed, corrupting staleness/ordering. Verified against both loops: + +- **Async** ([asyncfl/top_aggregator.py:176](../../flame/mode/horizontal/asyncfl/top_aggregator.py#L176)): pops **one** update per iteration via `recv_fifo(..., 1)`; staleness = `agg_round − trainer_version`, so commit order *does* matter. +- **Sync** ([syncfl/top_aggregator.py:231](../../flame/mode/horizontal/syncfl/top_aggregator.py#L231)): waits for `first_k` then aggregates via a weighted average (`optimizer.do`), which is **order-independent**; only *which* k complete and the round duration matter. + +**Decoupling constraint (important):** the aggregator must **not** hold a global profile of client runtimes — that is unrealistic FL. Instead it acts only on information **trainers report about themselves**. A trainer's simulated round duration *is* a local quantity (its own `training_delay_s`, eval delay, streaming reveal — all known to the trainer at start). So: + +- **Upfront ETA announce:** when a trainer is selected and begins (real) compute, it immediately sends a cheap control message announcing `sim_completion_ts = sim_send_ts + its own modeled duration`. This is local, observation-based reporting (like Oort/FedScale clients reporting expected speed), not aggregator coupling. Because the announcement is sent at *start* and is tiny, it arrives well before the heavy weight update — so by the time any update lands, the aggregator already knows every in-flight trainer's reported ETA. +- **Async commit rule:** the aggregator orders in-flight ends by reported `sim_completion_ts`, picks the **smallest**, and does a *targeted* blocking `channel.recv(end_id)` ([channel.py:384](../../flame/channel.py#L384)) on exactly that end. It commits it and advances the virtual clock `T_v` to that ETA. Ordering is by reported sim time, **not physical arrival**, so GPU-contention jitter (different trainers sharing a GPU finishing out of sim-order) cannot mis-order updates. +- **Bounded, correct waiting:** the only wait is physical — blocking until the earliest-ETA update actually arrives. If a simulated-faster trainer is physically slow due to contention, the aggregator waits real compute time for it (unavoidable for correctness; this is the "wait a bit" — bounded by one trainer's real compute, **never** by simulated delays). A real-time budget on the targeted recv maps onto the existing `SEND_TIMEOUT_WAIT_S` path so a trainer that announces then goes unavailable doesn't block forever (drop/stale → advance). +- **Sync:** the k committed updates = the k selected trainers with smallest reported `sim_completion_ts`; round duration = the k-th smallest. Aggregation result is unchanged (order-independent); we stop sleeping and set round wall-time from virtual time. + +This is more efficient than scaling sleeps: wall-clock is bounded by *real GPU time only*, not by the modeled delays, so large `training_delay_s` values cost nothing — and it keeps the aggregator decoupled from any client profile. + +## Resolved design questions (from review) + +**Two clean modes, one logical timeline, no `speedup_factor`.** There is a single logical simulated timeline driven by each trainer's modeled round duration `D` (= `training_delay_s` + eval). The two modes only differ in how they *realize* that timeline: +- **`real` (unsimulated)** — how training behaves in the real world. The trainer **sleeps** `D` so a slow device genuinely takes that long; wall-clock **is** the sim-timeline (1:1). Decisions use **actual arrival order** and the **real recv−send delta**. Authentic baseline. +- **`simulated`** — the same dynamics enacted faster: **no sleep**. The GPU does its real (fast) work, the trainer reports `sim_completion_ts = sim_send_ts + D`, and the aggregator advances a **virtual clock** `T_v`. Wall-clock = real GPU only. + +`speedup_factor` is **removed entirely** from both modes (it never worked end-to-end and is unnecessary: `real` runs at true pace, `simulated` is bounded by GPU compute). Sleep exists **only** in `real` mode — it is not reintroduced anywhere else. + +**Q1 — Two-run verification and the equivalence expectation.** Run the same experiment in `real` and `simulated` (fixed seed). Both read availability / data-visibility / speed / ordering off the *same logical times*, so they're expected to produce the **same per-(model-version) selection sets, staleness, and loss/accuracy** — "same per round/model-version, not per wall-clock-second" — with `simulated` wall-clock ≪ `real`. Equivalence is **exact in a deterministic scripted scenario** (well-separated `D`, controlled/mocked GPU time → arrival order == sim-completion order) and **within tolerance in a live smoke** (the only divergence source is real-GPU jitter reordering two trainers whose `D` are very close). Verify both: scripted parity test + smoke per-round parity & speedup. + +**Q2 — Trainer speed / system-utility source.** In `simulated` mode the wall-clock recv−send delta collapses to ~GPU time and misrepresents speed, so `PROP_ROUND_DURATION` is sourced from `sim_completion_ts − sim_send_ts` (= `D`). In `real` mode the real delta already ≈ `D` (the sleep dominates), so it stays authentic. Either way the value used for Oort's system-utility is the modeled duration. (The current `recv−send` path is *already* wrong under the old `speedup_factor>1`; removing speedup + this fix resolve it.) + +**Q3 — Dual timestamps (real + simulated) for lineage.** Keep both on every event record and in the aggregator's `_track_trainer_version_duration_s` (`real_ts` + `sim_ts`, send & recv), so lineage and wall-clock behavior stay reconstructable. Telemetry events gain `sim_ts`/`sim_completion_ts` alongside the existing real `ts`. + +**Q4 — Availability & data-streaming timing without `speedup_factor`.** Both are expressed in **sim-seconds** and evaluated against the logical sim-time: in `real` mode that equals wall-clock-since-start (1:1, no scaling — replaces today's `real_elapsed × speedup_factor` at [main.py:263](trainer/pytorch/main.py#L263)/[main.py:434](trainer/pytorch/main.py#L434)); in `simulated` mode against the virtual clock stamped on each task (`sim_send_ts = T_v`), so the trainer needs no free-running clock. **Smoke-test impact (expected, and fine):** switch the telemetry smoke to `time_mode: simulated`, delete `speedup_factor`, and keep `full_data_available_after_s` / availability horizons in sim-seconds (e.g. a 600s = 10-min ramp in sim-time). + +### Shared library layer (generic, reused by all examples) +- **A new `flame/sim/virtual_clock.py` mixin** holding `T_v`, `advance(ts)`, and a per-in-flight-end map of **trainer-reported** `sim_completion_ts`. Aggregators compose this in; examples inherit. The aggregator stores only what trainers announce — no client profile. +- **Message contract** in [flame/common/constants.py](../../flame/common/constants.py): (a) a lightweight **ETA-announce control message** (`MessageType.SIM_COMPLETION_TS`) the trainer sends at start of compute; (b) the same field echoed on the final weight update for validation. Backward-compatible: absence ⇒ real-time behavior (non-sim examples untouched). +- **New selector property** `PROP_SIM_COMPLETION_TS` in [selector/properties.py](../../flame/selector/properties.py), populated from the trainer's announcement; aggregator records `sim_send_ts` at distribute-time (reuse existing `PROP_ROUND_START_TIME` / `_track_trainer_version_duration_s`). + +### Trainer changes ([trainer/pytorch/main.py](trainer/pytorch/main.py)) +- Compute the modeled round duration `D` (`training_delay_s` + eval) locally. **`real` mode:** sleep `D` after GPU work ([main.py:571](trainer/pytorch/main.py#L571), [main.py:708](trainer/pytorch/main.py#L708)) — keep it. **`simulated` mode:** skip the sleep; instead **announce** `sim_completion_ts = sim_send_ts + D` (where `sim_send_ts` is the sim-time the aggregator stamped on the task) and echo `D`/`sim_completion_ts` on the outgoing update. +- Drop `speedup_factor` from the timing path entirely. Evaluate availability and data-visibility against the logical sim-time: wall-clock-since-start in `real` mode, the task's stamped sim-time in `simulated` mode (replaces `real_elapsed × speedup_factor` in `_visible_sample_count`/`check_and_update_state_avl`). In `simulated` mode the fixed real-time waits ([main.py:485](trainer/pytorch/main.py#L485), [main.py:668](trainer/pytorch/main.py#L668), polling/heartbeat threads, 20s fallback at [main.py:295](trainer/pytorch/main.py#L295)) must not gate progress. + +### Aggregator changes +- **Async** ([asyncfl/top_aggregator.py](../../flame/mode/horizontal/asyncfl/top_aggregator.py)): in `simulated` mode, consume ETA-announce messages into an in-flight ETA map; replace the single `recv_fifo(...,1)` pop with "select in-flight end with min `sim_completion_ts` → targeted `channel.recv(end_id)` → commit → advance `T_v`"; set `PROP_ROUND_DURATION = sim_completion_ts − sim_send_ts` (Q2); stamp each distributed task with `sim_send_ts = T_v`. In `real` mode keep today's arrival-ordered FIFO recv and real recv−send delta. Both keep the `RECV_TIMEOUT_WAIT_S` guard. Staleness keys off commit order (virtual-clock order in `simulated`, arrival order in `real`). +- **Sync** ([syncfl/top_aggregator.py](../../flame/mode/horizontal/syncfl/top_aggregator.py)): `simulated` mode commits the `first_k` smallest-`sim_completion_ts` responders with round duration from virtual time; `real` mode unchanged. Aggregation math is order-independent (weighted average) either way. +- Gate behind `time_mode: real|simulated`. `real` is the authentic, unchanged-decision baseline; `simulated` reproduces it via the virtual clock and is expected to match per-round (Q1). + +### `time_mode` plumbing (new config; replaces `speedup_factor`) +- **Config field:** add `time_mode: str = "simulated"` to `TrainerConfig` in [experiment_config.py](../../flame/launch/experiment_config.py) (default `simulated`; `real` is opt-in for verification). Load it in the same place `speedup_factor` was read (`load_experiment_config`, ~L190). +- **To the trainer:** the trainer reads `time_mode` from argv. Replace the trainer's `--speedup_factor` arg ([trainer/pytorch/main.py](trainer/pytorch/main.py), `main()`) with `--time_mode`; `TrainerSpawner` ([spawner.py](../../flame/launch/spawner.py)) passes `--time_mode ` instead of `--speedup_factor` (drop that field from the spawner too). `runner.py` passes `time_mode=exp.trainer.time_mode`. +- **To the aggregator:** the aggregator also needs `time_mode` (virtual-clock vs arrival path). Thread it into the aggregator config `hyperparameters.time_mode` via the runner's aggregator `config_overrides` (same mechanism as `agg_goal`), read it in `asyncfl/syncfl top_aggregator.__init__`. +- **Snapshot/exec-config:** drop the `speedup_factor` keys from [snapshot.py:92](../../flame/launch/snapshot.py#L92) and [execution_config_generator.py:118](../../flame/launch/execution_config_generator.py#L118); record `time_mode` instead. + +### `speedup_factor` removal inventory (delete all occurrences) +Code (remove field/arg/reads and the `/ speedup_factor` or `* speedup_factor` arithmetic — the timing now uses sim-time per the two modes): +- [experiment_config.py](../../flame/launch/experiment_config.py): `TrainerConfig.speedup_factor` field + the `.get("speedup_factor", ...)` read. +- [spawner.py](../../flame/launch/spawner.py): ctor param `speedup_factor`, `self.speedup_factor`, and the `--speedup_factor` cmd args. +- [runner.py](../../flame/launch/runner.py): `speedup_factor=exp.trainer.speedup_factor` kwarg. +- [snapshot.py](../../flame/launch/snapshot.py), [execution_config_generator.py](../../flame/launch/execution_config_generator.py): the recorded `speedup_factor` keys. +- [trainer/pytorch/main.py](trainer/pytorch/main.py): ctor param, `self.speedup_factor`, the `--speedup_factor` argparse arg, and every `/ speedup_factor` (avail event ts ~L283, training delay sleep ~L718, eval delay sleep ~L855) and `* speedup_factor` (`_visible_sample_count` sim_elapsed ~L461, util-disparity elapsed ~L712). Replace with the two-mode sim-time logic; remove the stale `TODO(DG): revisit speedup_factor coupling` note. +- YAMLs: delete the `speedup_factor: 1.0` lines across `expt_scripts_2026/*.yaml` and `experiments/configs/*.yaml` (≈25 files; a leftover line is harmless — the loader uses per-field `.get()` — but remove for cleanliness). Add `time_mode:` where a non-default is wanted. +- Out of scope but for consistency: `examples/fwdllm/.../FedSgdTrainer.py` hardcodes `speedup_factor=1.0` + an eval-delay divide; leave or clean separately (it's a different example, always 1.0, so behavior is unaffected). + +### Smoke-test YAML update +[felix_n10_alpha100_syn20_telemetry_smoke.yaml](expt_scripts_2026/felix_n10_alpha100_syn20_telemetry_smoke.yaml): remove `speedup_factor: 60.0`; add `time_mode: simulated`; set `data_streaming.full_data_available_after_s` back to **sim-seconds** (e.g. `600` for a 10-min *sim-time* ramp). Availability stays `syn_20`. For the Q1 two-run verification, add a tiny companion run/YAML (small `D`, few rounds) so a `time_mode: real` run finishes quickly. + +### Implementation order (suggested) +1. `flame/sim/virtual_clock.py` mixin + `MessageType.SIM_COMPLETION_TS` + `PROP_SIM_COMPLETION_TS`; add `time_mode` config plumbing; rip out `speedup_factor` (inventory above). +2. Trainer: modeled `D`, announce `sim_completion_ts` (simulated) / sleep `D` (real); sim-time availability + streaming; dual timestamps in telemetry. +3. Async aggregator: virtual-clock ordering + targeted recv + `PROP_ROUND_DURATION` from sim duration (simulated); keep arrival path (real). Then sync. +4. Tests: `tests/sim` (virtual clock), `tests/mode` (scripted real-vs-sim parity + arrival-order independence). +5. Smoke YAML update; run the two-run verification (Q1). + +--- + +## Discovered during Task 1 testing (carry into Task 2) + +Running the telemetry smoke (`felix_n10_alpha100_syn20_telemetry_smoke.yaml`, async_oort/fedbuff) surfaced several real bugs. Some were fixed inline to unblock telemetry validation; the deeper ones are flagged for Task 2. + +**Already fixed (keep, but revisit under the virtual clock):** +- **`speedup_factor` never reached the trainer.** The launcher built the trainer command without `--speedup_factor`/`--battery_threshold`, so the trainer always ran at `1.0` regardless of YAML — almost certainly why "speedup_factor=2 gave no 2× speedup". A stop-gap wiring fix was committed (spawner passes `--speedup_factor`), but **Task 2 supersedes it entirely: `speedup_factor` is removed in favor of the `real`/`simulated` two-mode design** (see "Resolved design questions" + the removal inventory). The `--battery_threshold` wiring stays. +- **Aggregator hung forever on a quiet in-flight trainer.** `recv_fifo` had no timeout; when every in-flight end went silent (all unavailable, or a stale ghost) the async aggregator blocked indefinitely. Band-aided with `recv_fifo(timeout=...)` + `RECV_TIMEOUT_WAIT_S=30` in asyncfl `_aggregate_weights`, plus a ghost filter (`channel.has`) and `selected_ends` cleanup in `async_oort._cleanup_removed_ends`. **The virtual-clock redesign should replace this band-aid** with the targeted, ETA-ordered receive (which has principled per-end completion times and timeouts). +- UTF-8 stdio in the launcher (latin-1 locales crashed on status glyphs) and an async `channel.ends()==None` guard. + +**Open issues for Task 2:** +- **In-flight accounting leaks (`freed=0`).** Telemetry showed a persistent `in_flight=1` ghost every round: an end gets selected, never returns an update, and is never freed from `selected_ends` (channel cleanup logged `in_flight_after=1, freed=0` each round). The proper fix is explicit in-flight↔availability reconciliation: on a client-notify `AVL_TRAIN→UN_AVL` transition (felix) drop the end from in-flight immediately; for non-notify baselines (fedbuff) a real timeout must drop it — and a returning trainer's late update should be explicitly accepted-or-discarded by version/staleness checks. This is core to the async correctness work. +- **Correlated availability is unrealistic.** All trainers share one `syn_20` "pattern" trace, so they go UN_AVL/AVL in lockstep (the whole system stalls together when the pattern is down). Decorrelate with per-trainer phase offsets (or per-trainer traces) so availability is independent — this also removes the global-stall failure mode at low speedup. +- **Model not learning in the smoke (accuracy flat at ~0.10 = random, loss pinned at ln 10).** Even after data fully unlocked, 100 async rounds produced no learning. Likely a hyperparameter/optimizer issue (client LR `0.001`, fedbuff server LR/`use_oort_lr`, delta-weight scaling) compounded by a very small per-trainer partition (~164 samples in this split). Needs a convergence sanity pass (validate felix actually learns on a known-good config) before drawing conclusions from streamed-vs-full utility plots. Telemetry is correct — it is what surfaced this. +- **`np.str_` end-ids leak into participation/selection dicts** (from `np.random.choice` in oort sampling). Harmless today (subclass of `str`) but a smell; normalize to `str` to avoid subtle set/dict-key surprises. + +--- + +## Tests (enhance pytest; cover simulated + real modes) + +Existing selector tests live in `lib/python/tests/selector/` with `lib/python/tests/conftest.py`. Add: +- **`tests/telemetry/test_event_schema.py`** (Task 1): schema/round-trip + that each selector emits the required event fields (parametrized over selectors for cross-comparability). Include a streaming-utility case: with a known prefix vs. full pool, assert the counterfactual utility ≥/≠ the streamed utility as expected and that disparity events are emitted at the configured cadence. +- **`tests/sim/test_virtual_clock.py`** (Task 2): virtual-clock advance/ordering unit tests; assert min-`sim_completion_ts` selection reproduces a hand-computed staleness sequence. +- **`tests/mode/test_async_aggregation_ordering.py`** (Task 2): feed a scripted set of in-flight ends with known sim durations + shuffled physical arrival; assert commit order and per-update staleness are **identical** to a reference computed by virtual time, and **independent of arrival order** (the jitter guarantee). +- **`tests/mode/test_sync_aggregation_equivalence.py`** (Task 2): assert sync aggregate output is identical in `simulated` vs. `real` mode (order-independence guard / regression). +- **Mode parity / regression** (Task 2): parametrize a tiny end-to-end run over `time_mode ∈ {simulated, real}`; assert correctness parity (final weights/accuracy within tolerance) and that `simulated` wall-clock < `real` (performance guard). + +--- + +## Critical files + +| Area | File | +|---|---| +| Async agg loop | `lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py` | +| Sync agg loop | `lib/python/flame/mode/horizontal/syncfl/top_aggregator.py` | +| Channel recv | `lib/python/flame/channel.py` (`recv` L384, `recv_fifo` L430) | +| Selector base/props | `lib/python/flame/selector/__init__.py`, `lib/python/flame/selector/properties.py` | +| Trainer | `lib/python/examples/async_cifar10/trainer/pytorch/main.py` | +| Msg types | `lib/python/flame/common/constants.py` | +| New: telemetry (Task 1) | `lib/python/flame/telemetry/` | +| New: analysis (Task 1) | `scripts/analysis/analyze_run.py` (reuses `scripts/plotters/` helpers) | +| New: virtual clock (Task 2) | `lib/python/flame/sim/virtual_clock.py` | + +--- + +## Verification + +**Task 1** +1. Run produces `/events.jsonl` and `/plots/*.png` automatically; eyeball selector availability-composition and streamed-vs-full utility-disparity plots. +2. `pytest lib/python/tests/telemetry` green. +3. Cross-selector compare: run two selectors (e.g. async_oort vs. feddance), `analyze_run.py --compare` overlays utility/speed/staleness/accuracy. + +**Task 2** (the two-run equivalence is the core check — see Q1) +4. `pytest lib/python/tests/sim lib/python/tests/mode` green: scripted deterministic scenario (fixed durations + availability + seed) yields **identical** commit order, staleness, selection sets, and resulting weights in `real` vs. `simulated` mode, and independent of physical arrival order. +5. **Two smoke runs, same seed** — one `real`, one `simulated`: assert **per-(model-version) selection sets and loss/accuracy match** (within RNG tolerance), and `simulated` wall-clock ≪ `real`. This is the "same selection and learning over time (per round, not per second)" verification. Use a **short / small-`D` / few-round** scenario so the `real` run (true pace, no speedup) finishes in reasonable wall-clock. + +--- + +## Parity smoke results (2026-05-30, pre-pending-commit fix) + +Runs: `run_20260530_130644_felix_n10_parity_REAL` vs `run_20260530_171846_felix_n10_parity_SIMULATED` +Config: `felix_n10_parity_real_vs_sim.yaml` — 10 trainers, `syn_0` availability, 100 rounds, `agg_goal=5`, `c=8`. + +| Check | Result | Notes | +|---|---|---| +| 1. Selection parity (Jaccard) | **WARN** — 0% exact, mean J=0.477 | Unseeded OORT RNG + participation skew | +| 2. Statistical utility (KS) | **FAIL** — max KS=0.806 | Slow trainers (D≥16s) have n_sim≤2 vs n_real≥32 | +| 3. Aggregation sequence | **FAIL** — 1% exact | Directly downstream of participation skew | +| 4. Staleness distribution | **WARN** — real mean=0.99, sim mean=2.14, diff=1.15 | Sim over-accumulates staleness | +| 5. Participation counts | **FAIL** — avg diff=31, max=47 | Slow trainers starved in sim; fast trainers dominate | +| 6. Convergence (acc/loss) | **PASS** — avg acc diff=0.040 | Both curves still flat (~random) at 100 rounds | + +**Root cause:** before the `_sim_pending_commit` fix, slow trainers (D=16–18s) were released back into selection at round end even though their buffer entry was uncommitted. OORT's system-utility then penalised them (high staleness → poor speed score), causing fast trainers (D=4–13s) to monopolise selection. Trainers 0371/0374/0377 (D=16/18/17s) appeared only 2/2/1 times in sim vs 36/32/32 in real. + +**Fix applied:** `_sim_pending_commit` — a trainer with an uncommitted `_sim_buffer` entry stays in `all_selected` (invisible to the selector) until `pop_min` commits that entry in `_sim_recv_min`. This enforces the real-mode invariant (one in-flight update per trainer at a time) in simulation. + +**Expected improvement in next run:** staleness mean should drop to ~1.0 (matching real), participation skew should narrow, and Jaccard should improve toward ≥0.7. Selection will not be bit-for-bit identical because OORT's RNG is unseeded; statistical distributions across trainers should match. +6. Regression: `real`-mode decisions unchanged vs. pre-Task-2 baseline (same seed). Note a convergence prerequisite from Task 1 findings — confirm felix actually learns on a known-good config before reading equivalence into accuracy curves. diff --git a/lib/python/examples/async_cifar10/aggregator/REFL_README.md b/lib/python/examples/async_cifar10/aggregator/REFL_README.md deleted file mode 100644 index caf45637d..000000000 --- a/lib/python/examples/async_cifar10/aggregator/REFL_README.md +++ /dev/null @@ -1,376 +0,0 @@ -# REFL Integration - Usage Guide - -This guide explains how to use the REFL (Resource-Efficient Federated Learning) integration in Flame. - -## Overview - -The REFL integration provides three main components that can be used independently or together: - -1. **Availability-Aware Selection** (`REFLOortSelector`) - Priority-based client selection using availability predictions -2. **Deadline-Based Aggregation** (`REFLFedAvg`) - Filters slow trainers and handles stale updates with various weighting strategies -3. **Availability Tracking** (`REFLAvailabilityTracker`) - Manages client availability traces and predictions - -## Quick Start - -### IMPORTANT: Use YAML Experiment Configs - -The experiment runner expects YAML configuration files that reference JSON aggregator configs. - -**Correct usage:** -```bash -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - -# Test with 5 trainers -python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml - -# Run ablation study -python3 launch/run_experiment.py experiments/configs/refl_ablation_study.yaml - -# Full scale with 300 trainers -python3 launch/run_experiment.py experiments/configs/refl_full_n300.yaml -``` - -**Note:** The YAML files reference JSON aggregator configs in `expt_scripts_2026/configs/refl_config_*.json` - -### Basic REFL Configuration (JSON Aggregator Config) - -The JSON aggregator configs define the Flame components: - -```json -{ - "selector": { - "sort": "refl_oort", - "kwargs": { - "aggr_num": 10, - "avail_priority": 2, - "availability_trace_file": "metadata/availability_traces/synthetic_traces.yaml" - } - }, - "optimizer": { - "sort": "reflfedavg", - "kwargs": { - "deadline": 100, - "stale_factor": -4 - } - } -} -``` - -**See `expt_scripts_2026/configs/refl_config_*.json` for complete examples.** - -## Component Configuration - -### 1. REFLOortSelector - -Priority-based client selection with availability awareness. - -**Required Parameters:** -- `aggr_num`: Number of clients to select per round - -**Optional Parameters:** -- `avail_priority`: Priority mode (default: 0) - - `0`: No priority (standard Oort) - - `1`: Fill mode - prioritize high-priority, fill from others - - `2`: Strict mode - only select high-priority clients -- `availability_trace_file`: Path to availability trace file (required if avail_priority > 0) -- `avail_probability`: Prediction accuracy 0-1 (default: 1.0) -- `blacklist_rounds`: Max selections before blacklisting or -1 to disable (default: -1) -- `blacklist_max_len`: Max fraction of clients to blacklist (default: 0.3) -- `pacer_step`: Rounds between pacer adjustments or -1 to disable (default: 20) -- `pacer_delta`: Percentage adjustment per pacer step (default: 5) - -**Oort UCB Parameters:** -- `exploration_factor`: Initial exploration rate (default: 0.9) -- `exploration_decay`: Decay rate per round (default: 0.98) -- `exploration_min`: Minimum exploration rate (default: 0.2) -- `round_threshold`: Speed filter percentile (default: 30) -- `alpha`: Staleness weight in utility (default: 2) - -**Example:** -```json -{ - "selector": { - "sort": "refl_oort", - "kwargs": { - "aggr_num": 30, - "avail_priority": 2, - "availability_trace_file": "metadata/availability_traces/mobiperf_traces.yaml", - "blacklist_rounds": 50, - "pacer_step": 20, - "pacer_delta": 5, - "exploration_factor": 0.9, - "exploration_decay": 0.98, - "exploration_min": 0.2, - "round_threshold": 30 - } - } -} -``` - -### 2. REFLFedAvg Optimizer - -Deadline-based aggregation with stale update handling. - -**Deadline Parameters:** -- `deadline`: Fixed deadline in seconds, or 0 for moving average (default: 0) -- `target_ratio`: Target percentile for moving average deadline (default: 0.8) -- `initial_deadline`: Initial deadline for moving average (default: 100.0) - -**Staleness Parameters:** -- `stale_update`: Max rounds a stale update can remain cached, or -1 for no limit (default: -1) -- `stale_factor`: Staleness weighting strategy (default: 1) - - `> 1`: Divide by constant factor - - `1`: Equal weight (standard FedAvg) - - `-1`: Average - divide by average staleness - - `-2`: AdaSGD - divide by (staleness + 1) - - `-3`: DynSGD - multiply by exp(-(staleness + 1)) - - `-4`: REFL - hybrid formula with utility -- `stale_beta`: REFL beta parameter, balances staleness vs utility (default: 0.9) -- `scale_coff`: REFL scaling coefficient (default: 10.0) - -**Example:** -```json -{ - "optimizer": { - "sort": "reflfedavg", - "kwargs": { - "deadline": 0, - "stale_update": 5, - "stale_factor": -4, - "stale_beta": 0.9, - "scale_coff": 10.0, - "target_ratio": 0.8, - "initial_deadline": 150.0 - } - } -} -``` - -### 3. Availability Traces - -Availability traces define when trainers are available/unavailable. - -**Format (YAML):** -```yaml -trace_name: - description: "Description of trace" - traces: - "1": {periods: [[start1, end1], [start2, end2]], duration: total_duration} - "2": {periods: [[start1, end1], [start2, end2]], duration: total_duration} -``` - -**Provided Traces:** -- `synthetic_traces.yaml`: - - `syn_0`: Always available (0% dropout) - - `syn_20`: 20% synthetic dropout - - `syn_50`: 50% synthetic dropout -- `mobiperf_traces.yaml`: - - `mobiperf_2st`: Real-world MobiPerf traces (2-state) - - `mobiperf_3st`: Real-world MobiPerf traces (3-state) - -## Configuration Examples - -### Example 1: Full REFL (All Features) - -```json -{ - "selector": { - "sort": "refl_oort", - "kwargs": { - "aggr_num": 10, - "avail_priority": 2, - "blacklist_rounds": 50, - "pacer_step": 20, - "availability_trace_file": "metadata/availability_traces/mobiperf_traces.yaml" - } - }, - "optimizer": { - "sort": "reflfedavg", - "kwargs": { - "deadline": 100, - "stale_update": 5, - "stale_factor": -4, - "stale_beta": 0.9 - } - } -} -``` - -### Example 2: Availability-Aware Selection Only - -```json -{ - "selector": { - "sort": "refl_oort", - "kwargs": { - "aggr_num": 10, - "avail_priority": 2, - "availability_trace_file": "metadata/availability_traces/synthetic_traces.yaml" - } - }, - "optimizer": { - "sort": "fedavg" - } -} -``` - -### Example 3: Staleness Handling Only - -```json -{ - "selector": { - "sort": "random" - }, - "optimizer": { - "sort": "reflfedavg", - "kwargs": { - "deadline": 100, - "stale_update": 5, - "stale_factor": -4 - } - } -} -``` - -### Example 4: Comparing Staleness Strategies - -**Equal Weighting (Baseline):** -```json -{"optimizer": {"sort": "reflfedavg", "kwargs": {"stale_factor": 1}}} -``` - -**AdaSGD:** -```json -{"optimizer": {"sort": "reflfedavg", "kwargs": {"stale_factor": -2}}} -``` - -**DynSGD:** -```json -{"optimizer": {"sort": "reflfedavg", "kwargs": {"stale_factor": -3}}} -``` - -**REFL (Hybrid):** -```json -{"optimizer": {"sort": "reflfedavg", "kwargs": {"stale_factor": -4, "stale_beta": 0.9}}} -``` - -## Ablation Studies - -Pre-configured ablation configs are provided in `aggregator/`: - -1. **`refl_config_ablation_baseline.json`** - No REFL features (baseline) -2. **`refl_config_ablation_avail.json`** - Availability selection only -3. **`refl_config_ablation_staleness.json`** - Staleness handling only -4. **`refl_config_test.json`** - Full REFL (small scale) -5. **`refl_config_n300.json`** - Full REFL (300 trainers) - -## Monitoring - -### Selector Metrics - -The REFLOortSelector logs: -- Number of priority vs. remaining clients per round -- Blacklisted clients count -- Pacer threshold adjustments -- Exploration/exploitation split - -### Optimizer Metrics - -The REFLFedAvg optimizer logs: -- Fast vs. slow trainer split -- Stale updates: cached, applied, discarded -- Moving average deadline (if adaptive) -- Staleness statistics - -Use `get_statistics()` method for programmatic access: -```python -optimizer = aggregator.optimizer -stats = optimizer.get_statistics() -print(f"Stale cached: {stats['stale_cached_count']}") -print(f"Stale applied: {stats['stale_applied_total']}") -print(f"Moving avg deadline: {stats['moving_avg_deadline']:.2f}s") -``` - -## Troubleshooting - -### Issue: No priority clients selected -**Cause:** Availability trace file not loaded or incorrect trace format -**Solution:** Check trace file path and format. Ensure `avail_priority > 0` and `availability_trace_file` is set. - -### Issue: All trainers marked as slow -**Cause:** Deadline too aggressive -**Solution:** Increase `deadline` or use `deadline: 0` for adaptive moving average. - -### Issue: Stale updates never applied -**Cause:** `stale_update` too restrictive or round duration too short -**Solution:** Increase `stale_update` limit or check trainer/round timing. - -### Issue: Blacklist too large -**Cause:** `blacklist_rounds` too low or `blacklist_max_len` too high -**Solution:** Adjust `blacklist_rounds` upward or reduce `blacklist_max_len`. - -## Performance Tuning - -### For High Availability Environments (low dropout): -```json -{ - "avail_priority": 0, - "deadline": 0, - "stale_factor": 1 -} -``` -Use standard FedAvg since REFL overhead isn't justified. - -### For Moderate Availability: -```json -{ - "avail_priority": 1, - "deadline": 0, - "stale_factor": -4 -} -``` -Use priority fill mode and adaptive deadline. - -### For Low Availability (high dropout): -```json -{ - "avail_priority": 2, - "deadline": 0, - "stale_update": 10, - "stale_factor": -4 -} -``` -Strict priority mode, cache stale updates longer. - -## Comparison with Felix/FedBuff - -To compare REFL against Felix (FedBuff + AsyncOort): - -**REFL Config:** -```json -{ - "selector": {"sort": "refl_oort", "kwargs": {"aggr_num": 10, "avail_priority": 2}}, - "optimizer": {"sort": "reflfedavg", "kwargs": {"deadline": 100, "stale_factor": -4}} -} -``` - -**Felix Config:** -```json -{ - "selector": {"sort": "async_oort", "kwargs": {"aggr_num": 10}}, - "optimizer": {"sort": "fedbuff", "kwargs": {"use_oort_lr": true}} -} -``` - -Ensure identical: -- Dataset splits (same Dirichlet alpha) -- Availability traces -- Hyperparameters (batch size, learning rate, etc.) -- Number of trainers and rounds - -## References - -- **REFL Paper:** https://arxiv.org/abs/2111.01108 -- **REFL EuroSys'23:** ACM EuroSys 2023 -- **Integration Plan:** `experiments/REFL_INTEGRATION_PLAN.md` -- **Flame Documentation:** `../../README.md` diff --git a/lib/python/examples/async_cifar10/aggregator/pytorch/main_asyncfl_agg.py b/lib/python/examples/async_cifar10/aggregator/pytorch/main_asyncfl_agg.py index 14c2e3e97..7bdf9bc54 100644 --- a/lib/python/examples/async_cifar10/aggregator/pytorch/main_asyncfl_agg.py +++ b/lib/python/examples/async_cifar10/aggregator/pytorch/main_asyncfl_agg.py @@ -209,7 +209,10 @@ def train(self) -> None: pass def evaluate(self) -> None: - """Evaluate (test) a model.""" + """Evaluate (test) a model every evalEveryNRounds rounds (default 10).""" + eval_every = self.config.hyperparameters.eval_every_n_rounds or 10 + if self._round % eval_every != 0: + return self.model.eval() test_loss = 0 correct = 0 @@ -268,5 +271,12 @@ def check_and_sleep(self) -> None: config = load_config_from_argv() a = PyTorchCifar10Aggregator(config, args.log_to_wandb, args.wandb_run_name) + + # Structured telemetry (no-op unless $FLAME_TELEMETRY_DIR is set by the + # launcher). One JSONL file for the aggregator process. + from flame import telemetry + + telemetry.configure(role="aggregator", end_id=config.job.job_id) + a.compose() a.run() diff --git a/lib/python/examples/async_cifar10/aggregator/pytorch/main_fedavg_agg.py b/lib/python/examples/async_cifar10/aggregator/pytorch/main_fedavg_agg.py index bc69fa085..809688a17 100644 --- a/lib/python/examples/async_cifar10/aggregator/pytorch/main_fedavg_agg.py +++ b/lib/python/examples/async_cifar10/aggregator/pytorch/main_fedavg_agg.py @@ -205,5 +205,10 @@ def check_and_sleep(self) -> None: config = load_config_from_argv() a = PyTorchCifar10Aggregator(config, args.log_to_wandb) + + from flame import telemetry + + telemetry.configure(role="aggregator", end_id=config.job.job_id) + a.compose() a.run() diff --git a/lib/python/examples/async_cifar10/aggregator/pytorch/main_oort_sync_agg.py b/lib/python/examples/async_cifar10/aggregator/pytorch/main_oort_sync_agg.py index c3e797111..5f8a07042 100644 --- a/lib/python/examples/async_cifar10/aggregator/pytorch/main_oort_sync_agg.py +++ b/lib/python/examples/async_cifar10/aggregator/pytorch/main_oort_sync_agg.py @@ -367,5 +367,10 @@ def check_and_sleep(self) -> None: config = load_config_from_argv() a = PyTorchCifar10Aggregator(config, args.log_to_wandb, args.wandb_run_name) + + from flame import telemetry + + telemetry.configure(role="aggregator", end_id=config.job.job_id) + a.compose() a.run() diff --git a/lib/python/examples/async_cifar10/configs/trainer_base.yaml b/lib/python/examples/async_cifar10/configs/trainer_base.yaml index 6b3b38b38..302d5892b 100644 --- a/lib/python/examples/async_cifar10/configs/trainer_base.yaml +++ b/lib/python/examples/async_cifar10/configs/trainer_base.yaml @@ -54,6 +54,12 @@ hyperparameters: enabled: "False" trace: syn_0 wait_until_next_avl: "True" + # Stream data in over time instead of exposing the full partition at + # init. Disabled by default; override per experiment under + # trainer.config_overrides.hyperparameters.data_streaming. + data_streaming: + enabled: "False" + full_data_available_after_s: 0 # sim seconds until 100% data is visible baseModel: name: "" diff --git a/lib/python/examples/async_cifar10/experiments/PHASE4_COMMANDS.md b/lib/python/examples/async_cifar10/experiments/PHASE4_COMMANDS.md deleted file mode 100644 index 3050a3436..000000000 --- a/lib/python/examples/async_cifar10/experiments/PHASE4_COMMANDS.md +++ /dev/null @@ -1,321 +0,0 @@ -# Phase 4: Testing Commands & Setup - -**Date:** February 6, 2026 -**Status:** Ready to Test - ---- - -## What is Phase 4? - -Phase 4 is the **Testing & Validation** phase of the REFL integration. The implementation (Phases 1-3) is complete, and now we need to: - -1. ✅ **Test the implementation** - Verify components work correctly -2. ✅ **Debug issues** - Fix any runtime errors -3. ✅ **Validate results** - Compare against REFL's published results -4. ✅ **Compare with Felix** - Head-to-head performance comparison - ---- - -## What Was Wrong Before? - -### The Issue - -You ran: -```bash -python3 launch/run_experiment.py aggregator/refl_config_test.json -``` - -And got: -``` -Loaded 0 experiments from refl_config_test.json -``` - -### The Problem - -The experiment runner (`launch/run_experiment.py`) expects a **YAML experiment configuration file**, not a JSON aggregator config directly. - -**Two-level configuration system:** -1. **YAML experiment config** - Defines experiment parameters (which trainers, dataset, execution settings) -2. **JSON aggregator config** - Defines Flame components (selector, optimizer, hyperparameters) - -The YAML file **references** the JSON file. - ---- - -## What Was Fixed? - -### 1. Moved JSON Configs to Correct Location - -```bash -# Before: aggregator/refl_config_test.json -# After: expt_scripts_2026/configs/refl_config_test.json -``` - -All REFL JSON configs now in: `expt_scripts_2026/configs/` - -### 2. Created YAML Experiment Configs - -Created three YAML configs in `experiments/configs/`: - -- **refl_test_5trainers.yaml** - Small test (5 trainers) -- **refl_full_n300.yaml** - Full scale (300 trainers) -- **refl_ablation_study.yaml** - Ablation study (4 experiments) - -### 3. Fixed aggr_num Parameter - -Changed `refl_config_test.json`: -- Before: `"aggr_num": 10` -- After: `"aggr_num": 3` (select 3 out of 5 trainers) - -### 4. Updated Documentation - -- Updated [REFL_README.md](../aggregator/REFL_README.md) with correct commands -- Created [PHASE4_TESTING_GUIDE.md](PHASE4_TESTING_GUIDE.md) with comprehensive testing guide - ---- - -## Correct Commands for Phase 4 - -### Step 1: Small Test (Start Here!) - -```bash -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - -# Run small test with 5 trainers -python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml -``` - -**Expected output:** -``` -====================================================================== -EXPERIMENT BATCH RUNNER -====================================================================== - -Loaded 1 experiments from experiments/configs/refl_test_5trainers.yaml - -====================================================================== -[1/1] RUNNING EXPERIMENT: refl_test_5trainers_syn0 -====================================================================== - -[1/6] Setting up experiment directory... - ✓ Created: experiments/run__refl_oort_n5_oracular_alpha0p1_syn0 - -[2/6] Initializing spawners... - ✓ Spawners initialized - -[3/6] Starting aggregator... - Spawning aggregator with config: expt_scripts_2026/configs/refl_config_test.json - ... -``` - -### Step 2: Check for Errors - -Monitor the aggregator log: -```bash -tail -f expt_scripts_2026/agg_logs/*_aggregator.log -``` - -**Look for:** -- ✅ "Using selector: refl_oort" -- ✅ "Using optimizer: reflfedavg" -- ✅ "Selected trainers: [...]" -- ❌ Any Python exceptions or errors - -### Step 3: Run Ablation Study - -After small test passes: -```bash -python3 launch/run_experiment.py experiments/configs/refl_ablation_study.yaml -``` - -This runs **4 experiments sequentially:** -1. Baseline (no REFL) -2. Availability only -3. Staleness only -4. Full REFL - -### Step 4: Full Scale Test - -After ablations pass: -```bash -python3 launch/run_experiment.py experiments/configs/refl_full_n300.yaml -``` - ---- - -## Expected Behavior - -### What Should Happen - -1. **Experiment loads** - YAML parsed successfully -2. **Aggregator starts** - JSON config loaded, REFL components registered -3. **Trainers spawn** - 5 (or N) trainers start training -4. **Selection happens** - REFLOortSelector picks high-priority trainers -5. **Aggregation happens** - REFLFedAvg filters by deadline and applies staleness weighting -6. **Training converges** - Accuracy improves over rounds - -### What to Check - -**In aggregator logs:** -``` -INFO - Initializing selector: refl_oort -INFO - Initializing optimizer: reflfedavg -INFO - Round 1: Selected trainers: [1, 3, 5] -INFO - Round 1: Applied 0 stale updates, cached 0 -INFO - Round 1: Test accuracy: 0.35 -INFO - Round 2: Selected trainers: [2, 4, 5] -INFO - Round 2: Applied 0 stale updates, cached 0 -INFO - Round 2: Test accuracy: 0.42 -... -``` - -**In trainer logs:** -``` -INFO - Trainer 1 starting -INFO - Round 1: Training epoch 1/1 -INFO - Round 1: Sending weights to aggregator -INFO - Round 2: Training epoch 1/1 -... -``` - ---- - -## Common Issues & Fixes - -### Issue 1: ImportError for flame.availability - -**Error:** -``` -ImportError: No module named 'flame.availability' -``` - -**Fix:** -```bash -# Make sure you're in the correct directory -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - -# Check Python path -python3 -c "import sys; print('\n'.join(sys.path))" -``` - -### Issue 2: Selector not found - -**Error:** -``` -KeyError: 'refl_oort' not found in selector registry -``` - -**Fix:** -```bash -# Verify registration -cd /home/dgarg39/flame/lib/python -python3 -c "from flame.selectors import selector_provider; print('refl_oort' in selector_provider._classes)" -# Should print: True - -# If False, check flame/selectors.py registration -``` - -### Issue 3: Availability trace file not found - -**Error:** -``` -FileNotFoundError: metadata/availability_traces/synthetic_traces.yaml -``` - -**Fix:** -```bash -# Check trace file exists -ls metadata/availability_traces/synthetic_traces.yaml - -# Verify path in JSON config is relative to async_cifar10 directory -``` - -### Issue 4: Not enough available trainers - -**Error:** -``` -ValueError: Not enough trainers available for selection -``` - -**Fix:** -- Lower `avail_priority` to 1 or 0 in JSON config -- Or use `syn_0` availability mode (always available) - ---- - -## File Structure - -``` -async_cifar10/ -├── experiments/ -│ ├── configs/ -│ │ ├── refl_test_5trainers.yaml ← YAML experiment configs -│ │ ├── refl_full_n300.yaml -│ │ └── refl_ablation_study.yaml -│ ├── REFL_INTEGRATION_PLAN.md ← Overall plan -│ ├── PHASE4_TESTING_GUIDE.md ← Detailed testing guide -│ └── PHASE4_COMMANDS.md ← This file -│ -├── expt_scripts_2026/ -│ ├── configs/ -│ │ ├── refl_config_test.json ← JSON aggregator configs -│ │ ├── refl_config_n300.json -│ │ ├── refl_config_ablation_baseline.json -│ │ ├── refl_config_ablation_avail.json -│ │ └── refl_config_ablation_staleness.json -│ ├── agg_logs/ ← Aggregator logs -│ └── trainer_logs/ ← Trainer logs -│ -├── aggregator/ -│ └── REFL_README.md ← Configuration reference -│ -└── launch/ - └── run_experiment.py ← Experiment runner -``` - ---- - -## Next Steps - -1. **Run small test:** - ```bash - cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml - ``` - -2. **Check logs for errors:** - ```bash - tail -f expt_scripts_2026/agg_logs/*_aggregator.log - ``` - -3. **Debug any issues** (see "Common Issues & Fixes" above) - -4. **If successful**, proceed to ablation study and full-scale tests - -5. **Document results** and compare with Felix - ---- - -## Quick Reference - -| Task | Command | -|------|---------| -| Small test (5 trainers) | `python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml` | -| Ablation study (4x100 trainers) | `python3 launch/run_experiment.py experiments/configs/refl_ablation_study.yaml` | -| Full scale (300 trainers) | `python3 launch/run_experiment.py experiments/configs/refl_full_n300.yaml` | -| View aggregator logs | `tail -f expt_scripts_2026/agg_logs/*_aggregator.log` | -| View trainer logs | `ls expt_scripts_2026/trainer_logs/` | -| Check experiments | `ls experiments/run_*/` | -| Kill all processes | `pkill -f "flame"` | - ---- - -## Documentation - -- **[REFL_INTEGRATION_PLAN.md](REFL_INTEGRATION_PLAN.md)** - Overall integration plan with checkpoint -- **[PHASE4_TESTING_GUIDE.md](PHASE4_TESTING_GUIDE.md)** - Comprehensive testing guide -- **[../aggregator/REFL_README.md](../aggregator/REFL_README.md)** - Configuration reference - ---- - -**Ready to test! Start with the small test (5 trainers) and let me know what happens.** diff --git a/lib/python/examples/async_cifar10/experiments/PHASE4_TESTING_GUIDE.md b/lib/python/examples/async_cifar10/experiments/PHASE4_TESTING_GUIDE.md deleted file mode 100644 index 0369a02f3..000000000 --- a/lib/python/examples/async_cifar10/experiments/PHASE4_TESTING_GUIDE.md +++ /dev/null @@ -1,502 +0,0 @@ -# Phase 4: REFL Testing & Validation Guide - -**Date:** February 6, 2026 -**Status:** Testing Phase -**Goal:** Validate REFL implementation and compare with Felix - ---- - -## Overview - -Phase 4 focuses on testing, debugging, and validating the REFL implementation completed in Phases 1-3. This includes unit tests, integration tests, REFL validation, and head-to-head comparison with Felix. - ---- - -## Quick Start - -### 1. Small-Scale Test (5 Trainers) - -Start with a minimal test to verify basic functionality: - -```bash -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - -# Run small test -python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml -``` - -**What to check:** -- ✓ Aggregator starts without errors -- ✓ 5 trainers spawn successfully -- ✓ REFLOortSelector is used (check logs) -- ✓ REFLFedAvg optimizer is used -- ✓ Training progresses normally -- ✓ No Python exceptions in logs - -**Expected output:** -``` -====================================================================== -EXPERIMENT BATCH RUNNER -====================================================================== - -Loaded 1 experiments from experiments/configs/refl_test_5trainers.yaml - -====================================================================== -[1/1] RUNNING EXPERIMENT: refl_test_5trainers_syn0 -====================================================================== -... -``` - -### 2. Ablation Study (100 Trainers each) - -Run all four ablation experiments sequentially: - -```bash -python3 launch/run_experiment.py experiments/configs/refl_ablation_study.yaml -``` - -**Configurations tested:** -1. **Baseline:** Standard Oort + FedAvg (no REFL) -2. **Availability:** REFL priority selection only -3. **Staleness:** REFL staleness handling only -4. **Full REFL:** All features enabled - -**What to compare:** -- Convergence speed (rounds to target accuracy) -- Final accuracy -- Training time per round -- Number of stale updates used -- Fairness metrics - -### 3. Full-Scale Test (300 Trainers) - -After initial tests pass, run full-scale REFL: - -```bash -python3 launch/run_experiment.py experiments/configs/refl_full_n300.yaml -``` - -**Requirements:** -- Availability traces for 300 trainers (see "Trace Generation" below) -- Dataset splits for 300 trainers (already exists) -- 8 GPUs recommended - ---- - -## Current Issues to Resolve - -### Issue 1: Aggregator Config Parameters - -The JSON configs in `expt_scripts_2026/configs/refl_config_*.json` may need adjustment: - -**Check these parameters:** -- `aggr_num`: Should match `agg_goal` in YAML config -- Trace file path: verify `availability_trace_file` points to correct location -- Hyperparameters: ensure `rounds`, `batchSize`, `learningRate` are appropriate - -### Issue 2: Availability Trace Generation - -Current traces may not cover all 300 trainers. Need to: - -```bash -# Check current trace coverage -python3 -c "import yaml; print(len(yaml.safe_load(open('metadata/availability_traces/synthetic_traces.yaml'))))" -``` - -If <300, extend traces or create mappings (see "Trace Generation" section below). - -### Issue 3: Top Aggregator Integration - -The `syncfl/top_aggregator.py` needs to provide round timing metadata: - -**Required changes:** -- Track `round_start_time` and `round_end_time` -- Compute `round_duration` -- Pass to optimizer in `TrainResult` objects - -**Location:** `lib/python/flame/mode/horizontal/syncfl/top_aggregator.py` - -### Issue 4: Selector/Optimizer Registration - -Verify the REFL components are properly registered: - -```bash -cd /home/dgarg39/flame/lib/python -python3 -c "from flame.selectors import selector_provider; print('refl_oort' in selector_provider._classes)" -python3 -c "from flame.optimizers import optimizer_provider; print('reflfedavg' in optimizer_provider._classes)" -``` - -Both should print `True`. - ---- - -## Testing Checklist - -### ✓ Unit Tests (Not Yet Started) - -Create tests for core REFL components: - -**Test availability tracker:** -```python -# test_refl_tracker.py -from flame.availability.refl_tracker import REFLAvailabilityTracker - -def test_load_traces(): - tracker = REFLAvailabilityTracker("metadata/availability_traces/synthetic_traces.yaml") - assert tracker.is_available("trainer_1", 0.0) - -def test_priority_computation(): - # Test priority = true_prob * ucb_score - pass - -def test_split_by_priority(): - # Test priority-based sorting and splitting - pass -``` - -**Test staleness weighting:** -```python -# test_reflfedavg.py -from flame.optimizer.reflfedavg import REFLFedAvg - -def test_equal_weighting(): - # stale_factor = 1 - pass - -def test_adasgd_weighting(): - # stale_factor = -2: weight = 1/(staleness+1) - pass - -def test_refl_weighting(): - # stale_factor = -4: hybrid formula - pass -``` - -**Test deadline filtering:** -```python -def test_fixed_deadline(): - # Filter updates > deadline - pass - -def test_moving_average_deadline(): - # Adaptive deadline at target_ratio percentile - pass -``` - -### ✓ Integration Tests (Ready to Run) - -1. **Test with 5 trainers (syn_0):** - ```bash - python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml - ``` - -2. **Test with different availability modes:** - - `syn_0`: Always available - - `syn_20`: 20% unavailability - - `syn_50`: 50% unavailability - - `mobiperf_2st`: Realistic traces - -3. **Test all staleness strategies:** - - Modify `stale_factor` in JSON config: 1, -1, -2, -3, -4 - - Compare convergence and fairness - -### ✓ REFL Validation (Next Step) - -Reproduce REFL's published results: - -**Target metrics (from REFL paper):** -- CIFAR-10, α=0.1, 100 clients -- Convergence to ~70% accuracy -- 2-3x speedup vs FedAvg -- Better fairness (lower variance across clients) - -**How to validate:** -1. Run REFL with their exact hyperparameters -2. Compare convergence curves -3. Check resource efficiency (compute + communication) -4. Measure fairness metrics (client accuracy variance) - ---- - -## Availability Trace Generation - -### Current State - -Traces exist for synthetic scenarios and MobiPerf: -- `metadata/availability_traces/synthetic_traces.yaml`: Synthetic traces -- `metadata/availability_traces/mobiperf_traces.yaml`: Real-world traces from MobiPerf dataset - -### Extending to 300 Trainers - -If traces don't cover 300 trainers, two options: - -**Option 1: Duplicate traces** -```python -# expand_traces.py -import yaml - -with open('metadata/availability_traces/synthetic_traces.yaml') as f: - traces = yaml.safe_load(f) - -# Duplicate to reach 300 -expanded = {} -for i in range(1, 301): - source_id = ((i - 1) % len(traces)) + 1 - expanded[f'trainer_{i}'] = traces[f'trainer_{source_id}'] - -with open('metadata/availability_traces/synthetic_traces_n300.yaml', 'w') as f: - yaml.dump(expanded, f) -``` - -**Option 2: Generate synthetic traces** -```python -# Using REFL's trace generator -import numpy as np - -def generate_2state_trace(num_trainers, duration, mean_avail=0.8): - """Generate 2-state availability traces.""" - traces = {} - for i in range(1, num_trainers + 1): - # Random availability pattern - trace = np.random.rand(duration) < mean_avail - timestamps = list(range(duration)) - traces[f'trainer_{i}'] = { - 'times': timestamps, - 'availability': trace.tolist() - } - return traces -``` - ---- - -## Syncfl Top Aggregator Modifications - -The top aggregator needs to track round timing and pass it to the optimizer. - -**File:** `lib/python/flame/mode/horizontal/syncfl/top_aggregator.py` - -**Required changes:** - -```python -# In the training loop -round_start_time = time.time() - -# ... selection and training ... - -round_end_time = time.time() -round_duration = round_end_time - round_start_time - -# When creating TrainResult objects -for trainer_id, weights in trainer_weights.items(): - result = TrainResult( - trainer_id=trainer_id, - weights=weights, - # ... other fields ... - completion_time=time.time(), - round_duration=round_duration, - staleness=0, # Will be computed by optimizer - ) -``` - ---- - -## Debugging Tips - -### Check Aggregator Logs - -Aggregator logs are saved to `expt_scripts_2026/agg_logs/`: - -```bash -# View most recent aggregator log -tail -f expt_scripts_2026/agg_logs/*_aggregator.log -``` - -**What to look for:** -- "Using selector: refl_oort" -- "Using optimizer: reflfedavg" -- "Selected trainers: [...]" -- "Applied stale updates: X, cached: Y" - -### Check Trainer Logs - -Trainer logs are in `expt_scripts_2026/trainer_logs/`: - -```bash -# Check if trainers are running -ls expt_scripts_2026/trainer_logs/ - -# View specific trainer -tail -f expt_scripts_2026/trainer_logs/trainer_1.log -``` - -### Common Errors - -**ImportError: No module named 'flame.availability'** -- The availability module isn't in Python path -- Ensure you're running from the correct directory - -**KeyError: 'refl_oort' not found** -- Selector not registered -- Check `flame/selectors.py` registration - -**FileNotFoundError: availability trace file** -- Trace file path incorrect in JSON config -- Use relative path from example directory - -**ValueError: Not enough trainers available** -- Priority selection can't find enough available trainers -- Lower `avail_priority` (try 0 or 1) -- Check trace file has correct trainer IDs - ---- - -## Performance Metrics to Track - -### Convergence Metrics -- **Rounds to target accuracy** (e.g., 70% test accuracy) -- **Final accuracy** after fixed rounds (e.g., 100 rounds) -- **Convergence curve** (accuracy vs rounds) - -### Resource Efficiency -- **Training time per round** (lower is better) -- **Communication cost** (bytes transferred) -- **Compute cost** (GPU hours) - -### Fairness Metrics -- **Client accuracy variance** (lower is better) -- **Minimum client accuracy** (higher is better) -- **Jain's fairness index** - -### REFL-Specific Metrics -- **Stale update statistics**: - - Number cached per round - - Number applied per round - - Number discarded (too stale) -- **Priority selection**: - - High priority selections vs low priority - - Blacklist length over time -- **Pacer adjustments**: - - `round_threshold` evolution - - Effect on selection diversity - ---- - -## Next Steps - -### Immediate (Week 1) -1. ✅ Create YAML experiment configs (DONE) -2. ⏳ Run small-scale test (5 trainers) -3. ⏳ Debug any import/registration errors -4. ⏳ Verify selector and optimizer are used - -### Short-term (Week 2) -1. ⏳ Add round timing to top aggregator -2. ⏳ Create unit tests for REFL components -3. ⏳ Run ablation study (100 trainers) -4. ⏳ Extend availability traces to 300 trainers - -### Medium-term (Week 3-4) -1. ⏳ Run full-scale REFL (300 trainers) -2. ⏳ Reproduce REFL's published results -3. ⏳ Compare against Felix on identical workloads -4. ⏳ Document performance differences - ---- - -## Comparison with Felix - -### Experimental Setup - -For a fair comparison, use **identical** configurations: - -**Dataset:** -- CIFAR-10, α=0.1 (same Dirichlet split) -- Same train/test splits for all 300 trainers - -**Availability:** -- Same trace file -- Same trainer-to-trace mappings - -**Hyperparameters:** -- Same batch size, learning rate, epochs per round -- Same number of rounds - -**Selection:** -- Same `agg_goal` (number of trainers per round) -- Compare: Felix vs REFL vs Oort baseline - -### Run Comparison - -```bash -# 1. Baseline: Standard Oort -python3 launch/run_experiment.py experiments/configs/refl_ablation_study.yaml - -# 2. Felix (existing) -python3 launch/run_experiment.py experiments/configs/.yaml - -# 3. Full REFL -python3 launch/run_experiment.py experiments/configs/refl_full_n300.yaml -``` - -### Analysis - -Compare across all three: -- Convergence curves -- Resource efficiency -- Fairness metrics -- Robustness to availability changes - -**Expected outcome:** -- Felix: Best handling of asynchrony -- REFL: Better fairness and resource efficiency -- Oort: Baseline performance - ---- - -## Success Criteria (Updated) - -### Phase 4 Complete When: - -- [ ] Small-scale test runs without errors -- [ ] All REFL components working (selector + optimizer) -- [ ] Ablation study shows independent feature effects -- [ ] Full-scale experiment (300 trainers) runs successfully -- [ ] Convergence matches expected REFL behavior -- [ ] Resource efficiency gains validated -- [ ] Fairness improvements demonstrated -- [ ] Head-to-head comparison with Felix completed -- [ ] Documentation updated with findings - ---- - -## Useful Commands - -```bash -# Navigate to async_cifar10 -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - -# Run experiment -python3 launch/run_experiment.py experiments/configs/.yaml - -# Check experiment results -ls experiments/ - -# View logs -tail -f expt_scripts_2026/agg_logs/*_aggregator.log - -# Kill all flame processes (if needed) -pkill -f "flame" - -# Check available configs -ls experiments/configs/refl*.yaml -ls expt_scripts_2026/configs/refl*.json -``` - ---- - -## Contact & Questions - -For issues or questions about Phase 4 testing, refer to: -- [REFL_INTEGRATION_PLAN.md](REFL_INTEGRATION_PLAN.md) - Overall integration plan -- [aggregator/REFL_README.md](#) - Configuration guide (if it exists in aggregator/) -- REFL paper: "REFL: Resource-Efficient Federated Learning" diff --git a/lib/python/examples/async_cifar10/experiments/QA_REFL_IMPLEMENTATION.md b/lib/python/examples/async_cifar10/experiments/QA_REFL_IMPLEMENTATION.md deleted file mode 100644 index 03de76bfe..000000000 --- a/lib/python/examples/async_cifar10/experiments/QA_REFL_IMPLEMENTATION.md +++ /dev/null @@ -1,220 +0,0 @@ -# REFL Implementation Q&A - -**Date:** February 6, 2026 - ---- - -## Q1: Do we need changes to aggregator/trainer code files? - -**Answer: NO - No changes needed to main aggregator/trainer code.** - -### Why Not? - -Flame's architecture provides clean abstractions: -- **Selectors** are plug-and-play via `selector_provider` -- **Optimizers** are plug-and-play via `optimizer_provider` -- Main aggregator/trainer code is agnostic to which selector/optimizer is used - -### What Was Needed? - -Only **registration and extension**: -1. ✅ Register `refl_oort` in `flame/selectors.py` -2. ✅ Register `reflfedavg` in `flame/optimizers.py` -3. ✅ Extend `TrainResult` with optional timing fields (backward compatible) -4. ✅ Fix JSON config format (add `taskid`, proper channel structure) - -### Migration Plan Adjustments - -The MIGRATION_PLAN.md showed the new spawner/launcher system which **already works** with REFL: -- YAML experiment configs specify which selector/optimizer to use -- JSON aggregator configs get loaded normally -- Trainers spawn with correct configurations -- **No code changes needed** in aggregator/trainer main files - ---- - -## Q2: How does availability tracking work in REFL? - -**Answer: REFL uses ORACULAR traces - perfect future knowledge of availability.** - -### REFL's Approach (third_party/REFL/core/helper/client.py) - -```python -class Client: - def __init__(self, hostId, clientId, speed, traces=None): - self.traces = traces # Availability trace data - - def isActive(self, cur_time): - """Check if client is available at current time""" - if self.traces is None: - return True # Always available if no trace - - # Wrap time within trace duration - norm_time = cur_time % self.traces['finish_time'] - - # Check if within active window - if (self.traces['active'][i] <= norm_time <= - self.traces['inactive'][i]): - return True - return False -``` - -### Trace Format - -```python -# Example trace for one client -{ - 'duration': 211625, - 'finish_time': 518400, # Loop period (6 days in seconds) - 'active': [12788, 100044, 188992, ...], # Start times - 'inactive': [65881, 133574, 208292, ...], # End times - 'model': 'CPH1801' # Device model -} -``` - -**Key points:** -- `active[i]` to `inactive[i]` = one availability window -- Time wraps: `current_time % finish_time` -- From MobiPerf dataset (real mobile device traces) - -### Can REFL Function Without Traces? - -**NO - REFL fundamentally requires traces.** From the code: -- If `traces is None`, client is always available (trivial case) -- Priority computation requires trace data -- Deadline filtering needs availability predictions -- This is **intentional** - REFL assumes oracular knowledge for research purposes - -### Our Implementation (flame/availability/refl_tracker.py) - -```python -class REFLAvailabilityTracker: - def __init__(self, trace_file): - """Load traces - pickle or YAML format""" - self.traces = self.load_traces(trace_file) - - def is_available(self, trainer_id, current_time): - """Check if trainer available now""" - # Same logic as REFL's Client.isActive() - - def get_priority(self, trainer_id, round_start, deadline): - """Compute priority based on near-term availability""" - # Returns 0 (unavailable), 1 (partially), or 2 (fully available) -``` - -**Supported formats:** -- REFL pickle files (from third_party/REFL) -- Flame YAML traces (metadata/availability_traces/*.yaml) - -### Is This Realistic? - -**For research: Yes** - Allows controlled comparison -**For production: No** - Real systems don't have perfect future knowledge - -**Alternative approaches:** -- **Felix**: No traces needed, handles true unpredictability -- **Predictive REFL**: Could add ML-based availability prediction (future work) -- **Trace replay**: Use historical data to simulate realistic scenarios - ---- - -## Q3: Simplify the integration plan? - -**Done! ✅** - -### What Changed - -**Before:** 1,196 lines with detailed implementation code in sections 1-12 - -**After:** ~270 lines focused on: -- ✅ **Design decisions** - Why we made specific choices -- ✅ **Key concepts** - Availability tracking, priority selection, staleness -- ✅ **Configuration** - How to use the system -- ✅ **Status tracking** - What's done, what's remaining -- ✅ **Quick reference** - Files created, commands to run - -**Removed:** -- Detailed code examples for each component -- Step-by-step implementation instructions -- Redundant explanations of already-implemented features -- Speculative future work not relevant to current testing - -### What Remains - -1. **Executive Summary** - High-level overview -2. **Key Design Decisions** - Why we chose specific approaches -3. **Implementation Status** - Phase 1-3 complete, Phase 4 in progress -4. **Core Design Choices** - Availability tracking, selection, staleness -5. **Configuration Structure** - How configs work -6. **Remaining Work** - What's needed for Phase 4 -7. **Key Takeaways** - REFL's strengths and limitations -8. **Appendices** - File lists, REFL vs Felix comparison - ---- - -## Summary of Fixes (Feb 6, 2026) - -### Issue: KeyError 'taskid' - -**Root Cause:** JSON configs used simplified format, but Flame requires full schema - -**Fix Applied:** -```json -{ - "taskid": "experiment_id", // ✅ Added - "backend": "mqtt", // ✅ Added - "brokers": [...], // ✅ Added - "groupAssociation": {...}, // ✅ Added - "channels": [ // ✅ Fixed structure - { - "name": "param-channel", - "pair": ["trainer", "aggregator"], - "funcTags": {...} - } - ], - "hyperparameters": { - "aggGoal": 3, // ✅ Changed from rounds/epochs only - "trackTrainerAvail": {...} // ✅ Added availability tracking config - }, - "selector": {...}, - "optimizer": {...} -} -``` - -**Files Fixed:** -- ✅ `refl_config_test.json` -- ✅ `refl_config_n300.json` -- ✅ `refl_config_ablation_baseline.json` -- ✅ `refl_config_ablation_avail.json` -- ✅ `refl_config_ablation_staleness.json` - ---- - -## Next Steps - -1. **Test again** with fixed configs: - ```bash - cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml - ``` - -2. **Check logs** for successful initialization: - ```bash - tail -f expt_scripts_2026/agg_logs/*_aggregator.log - ``` - -3. **Verify** REFL components loaded: - - "Using selector: refl_oort" - - "Using optimizer: reflfedavg" - -4. **If successful**, proceed to ablation study - ---- - -## Key Points to Remember - -1. **No code changes** needed - selectors/optimizers are modular -2. **Traces required** - REFL cannot function without oracular availability data -3. **Synchronous FL** - REFL is not asynchronous like Felix -4. **Config format matters** - Must match Flame's schema exactly -5. **Two-level configs** - YAML experiment → JSON aggregator diff --git a/lib/python/examples/async_cifar10/experiments/README.md b/lib/python/examples/async_cifar10/experiments/README.md index 766738d2e..f1da7afee 100644 --- a/lib/python/examples/async_cifar10/experiments/README.md +++ b/lib/python/examples/async_cifar10/experiments/README.md @@ -1,205 +1,194 @@ -# Experiment Launcher - Quick Start Guide +# Experiment Quick-Start -This directory contains the programmatic experiment generation system for async_cifar10. Experiments are defined in compact YAML configs that reference metadata, rather than embedding full trainer configurations. +Experiments are defined as compact YAML files that reference shared metadata +(trainer registry, dataset splits, availability traces) rather than embedding +full trainer configs. The launcher assembles per-trainer configs at runtime. -## Quick Start +--- -### Running Experiments +## Running experiments -```bash -# From async_cifar10 directory -cd /path/to/flame/lib/python/examples/async_cifar10 +Run from the repository root: -# Run a small test (5 trainers) -python3 launch/run_experiment.py experiments/configs/test_phase3_mini.yaml +```bash +# Smoke test — 10 trainers, fast +python -m flame.launch.run_experiment \ + lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn0_smoke.yaml -# Run full-scale experiments (300 trainers, 4 experiments) -python3 launch/run_experiment.py experiments/configs/oort_n300_all4unavail.yaml +# With telemetry + data streaming enabled +python -m flame.launch.run_experiment \ + lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_telemetry_smoke.yaml ``` -### Reproducing Experiments - -Each experiment run creates an `execution_config.yaml` in its output directory with all parameters and metadata references needed for exact reproduction: +Batch files (multiple experiments in one YAML) run sequentially, each producing +its own timestamped output directory under `experiments/`. -```bash -# Reproduce from execution config -python3 launch/reproduce.py experiments/run_20260203_223834_oort_n300_alpha0.1_syn0/execution_config.yaml -``` +--- -## Config File Structure - -Experiment configs use **metadata references** (not embedded data) to keep files readable (~1-2KB): +## Experiment YAML shape ```yaml experiments: - - name: oort_n300_alpha0.1_syn0 - description: "Oort with 300 trainers, alpha=0.1, synthetic unavailability" - + - name: felix_n10_alpha100_syn0_smoke + description: "Optional human-readable description." + + baseline: felix # key into _metadata/baselines.yaml; owns the stack + trainer: - num_trainers: 300 # Number of trainers to spawn + num_trainers: 10 + start_id: 1 # optional; first trainer ID (default: 1) dataset: - dirichlet_alpha: 0.1 # Data heterogeneity (0.1 = high, 100.0 = low) + name: cifar10 + dirichlet_alpha: 100.0 # 0.1 (high het) → 100.0 (near-IID) availability: - mode: syn_0 # Availability trace key from metadata - + mode: syn_0 # syn_0 | syn_20 | syn_50 | mobiperf_2st | mobiperf_3st_50 | mobiperf_3st_75 + enable_training_delays: true + time_mode: simulated # real (wall-clock sleeps) | simulated (virtual-clock, no sleeps) + hyperparameters: # merged into every trainer's config + batchSize: 32 + learningRate: 0.001 + config_overrides: # deep-merged last; wins over baseline defaults + hyperparameters: + client_notify: + trace: syn_0 + aggregator: - config_template: expt_scripts_2026/configs/oort_n300_oracular_9may25_syn0.json - selector: oort # Selector algorithm - tracking_mode: oracular # Oracular or oblivious - agg_goal: 10 # Target trainers per aggregation round - + config_template: ../_metadata/aggregator_base.json # base template; baseline merges on top + selector: async_oort # must match the stack implied by the baseline + tracking_mode: client_notify # oracular | client_notify + agg_goal: 5 + log_to_wandb: false + config_overrides: # deep-merged into aggregator JSON last + job: + id: felix_n10_alpha100_syn0_smoke + hyperparameters: + rounds: 100 + aggGoal: 5 + selector: + kwargs: + c: 8 + aggGoal: 5 + execution: - num_gpus: 8 # GPUs for trainer distribution - sleep_between_spawns: 5.0 # Seconds between trainer spawns - aggregator_warmup_time: 600 # Seconds before first aggregation + num_gpus: 2 + sleep_between_spawns: 1.0 + aggregator_warmup_time: 15 # seconds; increase proportionally for large runs + monitoring: + enabled: false ``` -### Key Configuration Parameters - -**Trainer Section:** -- `num_trainers`: Total trainers to spawn (typically 5 for tests, 300 for production) -- `start_id`: Optional, starting trainer ID (default: 1) -- `dataset.dirichlet_alpha`: Data heterogeneity - - `0.1` = highly heterogeneous (realistic) - - `1.0` = moderate heterogeneity - - `10.0` = low heterogeneity - - `100.0` = nearly IID -- `availability.mode`: References key in `metadata/availability_traces/*.yaml` - - `syn_0` = always available (0% dropout) - - `syn_20` = 20% synthetic dropout - - `syn_50` = 50% synthetic dropout - - `mobiperf_2st` = real-world MobiPerf traces - -**Execution Timing (Critical for Success):** -- `sleep_between_spawns`: Delay between spawning each trainer (seconds) -- `aggregator_warmup_time`: How long aggregator waits before first round (seconds) - -**Timing Formula:** To ensure all trainers join before first aggregation: +### Key parameters + +**`baseline`** — resolves `aggregator_main` (determines asyncfl vs syncfl stack), +selector, optimizer, and trainer-side defaults from `_metadata/baselines.yaml`. +Available baselines: `felix`, `fedbuff`, `fedavg`, `oort`, `refl`, `feddance`. + +**`time_mode`** +- `simulated` — trainers skip `training_delay_s` sleeps; the aggregator orders + updates by a virtual clock. Availability and streaming use sim-seconds. + Fast iteration. +- `real` — trainers sleep true wall-clock delays. Use for wall-clock benchmarks. + +**`availability.mode`** — selects which pre-injected trace the trainer activates: +| Mode | Description | +|------|-------------| +| `syn_0` | Always available | +| `syn_20` | 20% synthetic unavailability | +| `syn_50` | 50% synthetic unavailability | +| `mobiperf_2st` | Real MobiPerf 2-state traces | +| `mobiperf_3st_50` | 3-state, 50% battery threshold | +| `mobiperf_3st_75` | 3-state, 75% battery threshold | + +**`aggregator_warmup_time`** — seconds the runner waits after spawning the +aggregator before spawning trainers. Rule of thumb: ``` -aggregator_warmup_time = 0.4 × sleep_between_spawns × num_trainers +warmup ≥ sleep_between_spawns × num_trainers × 0.4 ``` +Defaults of 15s (smoke) to 60s (300-trainer) are set in the smoke YAMLs. -Examples: -- 5 trainers: spawn=0.5s → warmup=1s -- 300 trainers: spawn=5.0s → warmup=600s - -## Metadata System - -Configs reference metadata files instead of embedding full data: +### Optional trainer features (via `config_overrides.hyperparameters`) -``` -metadata/ -├── trainer_registry.yaml # 300 trainers (intrinsic properties) -├── dataset_splits/ -│ ├── cifar10_alpha0.1_n300.yaml # Dirichlet alpha=0.1 -│ ├── cifar10_alpha1.0_n300.yaml # Dirichlet alpha=1.0 -│ ├── cifar10_alpha10.0_n300.yaml # Dirichlet alpha=10.0 -│ └── cifar10_alpha100.0_n300.yaml # Dirichlet alpha=100.0 -└── availability_traces/ - ├── synthetic_traces.yaml # syn_0, syn_20, syn_50 - └── mobiperf_traces.yaml # Real-world per-trainer traces +```yaml +# Gradual data reveal (disabled by default) +data_streaming: + enabled: "True" + full_data_available_after_s: 600 # sim-seconds until 100% of partition visible + +# Streamed-vs-full utility disparity telemetry (disabled by default) +util_counterfactual: + enabled: "True" + every_n_rounds: 1 + sample_size: 256 ``` -The launcher automatically loads metadata and generates full trainer configs at runtime. +--- -## Output Structure - -Each experiment run creates a timestamped directory: +## Output structure ``` -experiments/run_20260203_223834_oort_n300_alpha0.1_syn0/ -├── execution_config.yaml # Full reproducibility record -├── snapshot.yaml # Runtime state snapshot -├── aggregator_config.json # Generated aggregator config -├── 03_02_26_22_38_..._aggregator.log # Aggregator logs -└── 03_02_26_22_38_..._trainers.log # All trainer logs +experiments/run_YYYYMMDD_HHMMSS_/ + aggregator_config.json # merged aggregator config used for this run + execution_config.yaml # compact metadata refs + exact spawn commands + snapshot.yaml # git commit, branch, metadata SHA256 checksums + YYYYMMDD_*_aggregator.log # aggregator stdout/stderr + YYYYMMDD_*_trainers.log # all trainers combined (line-buffered) + YYYYMMDD_*_resources.log # RAM/GPU monitoring (if enabled) + telemetry/ + aggregator.jsonl + trainer_1.jsonl … trainer_N.jsonl + plots/ # auto-generated by post-run analysis (if run) ``` -**execution_config.yaml** contains: -- Git commit hash and branch -- Metadata file references and keys -- All experiment parameters -- Exact spawn commands used +`execution_config.yaml` records the metadata keys, git state, and exact commands +used — sufficient to reproduce the run at the same git commit. -Use this file to reproduce experiments exactly: -```bash -python3 launch/reproduce.py experiments/run_*/execution_config.yaml -``` +--- -## Running Multiple Experiments +## Post-run analysis -Configs can define multiple experiments in batch: +```bash +# Parse send/receive lag from aggregator log +python scripts/analyze_send_recv_lag.py \ + experiments/run_*/YYYYMMDD_*_aggregator.log [--warn-threshold-s 5.0] -```yaml -experiments: - - name: exp1_syn0 - # ... config ... - - - name: exp2_syn20 - # ... config ... - - - name: exp3_syn50 - # ... config ... +# Compare two runs (e.g., real vs simulated time_mode) +python scripts/compare_parity.py \ + experiments/run_A/ experiments/run_B/ ``` -The launcher runs them sequentially, each producing its own output directory. - -## Optional Features +Telemetry JSONL files are line-buffered — `tail -f` works during a live run. -### Weights & Biases Logging +--- -Add to your experiment config: +## Metadata layout -```yaml -aggregator: - log_to_wandb: true - wandb_project: "my-project" - wandb_run_name: "oort_n300_alpha0.1_syn0" ``` - -### Custom Trainer ID Ranges - -By default, trainers spawn with IDs [1, num_trainers]. To use a different range: - -```yaml -trainer: - num_trainers: 5 - start_id: 101 # Will spawn trainers 101-105 +_metadata/ (symlinked as async_cifar10/metadata/) + trainer_registry.yaml # 300 trainers: task_id, delay_s, speed_class + baselines.yaml # selector/optimizer/stack catalog + aggregator_base.json # aggregator config template + dataset_splits/ + cifar10_alpha0.1_n300.yaml # Dirichlet α=0.1 + cifar10_alpha1.0_n300.yaml + cifar10_alpha10.0_n300.yaml + cifar10_alpha100.0_n300.yaml + availability_traces/ + synthetic_traces.yaml # syn_0 / syn_20 / syn_50 + mobiperf_traces.yaml # per-device real-world traces ``` -## Troubleshooting +The launcher injects per-trainer data (indices, delays, all five trace types) +from this bundle at spawn time — no per-trainer JSON files needed. -**Issue: Aggregation completes with fewer trainers than expected** -- **Cause:** Aggregator started before trainers joined -- **Solution:** Increase `aggregator_warmup_time` using the formula: - ``` - warmup = 0.4 × spawn_delay × num_trainers - ``` +--- -**Issue: Losing trainers at startup** -- **Cause:** `sleep_between_spawns` too small for scale -- **Solution:** Use 5.0s for 300 trainers, 0.5s for small tests - -**Issue: OOM errors with many trainers** -- **Cause:** Too many trainers per GPU -- **Solution:** Increase `num_gpus` in execution section - -## Examples - -See configs in this directory: -- `test_phase3_mini.yaml` - Small test (5 trainers, 2 GPUs) -- `oort_n300_all4unavail.yaml` - Full batch (4 × 300 trainers, 8 GPUs) - -## Additional Tools - -```bash -# Validate metadata integrity -python3 scripts/validate_metadata.py +## Troubleshooting -# Extract/update metadata from existing configs (if needed) -python3 scripts/extract_metadata.py -``` +**Aggregation completes with fewer trainers than expected** — aggregator started +before trainers joined. Increase `aggregator_warmup_time`. -## Migration Note +**OOM with many trainers** — too many trainers sharing a GPU. Increase `num_gpus`. -This system replaces 5,924 static JSON trainer configs with 7 YAML metadata files (99.9% reduction). The old methodology in `expt_scripts_2026/` remains available for backwards compatibility. +**Stack validation error** — async selectors (`async_oort`, `async_random`, +`fedbuff`) require an asyncfl-stack baseline (e.g., `felix`, `fedbuff`); sync +selectors require a syncfl baseline (`oort`, `refl`, `fedavg`). diff --git a/lib/python/examples/async_cifar10/experiments/REFL_INTEGRATION_PLAN.md b/lib/python/examples/async_cifar10/experiments/REFL_INTEGRATION_PLAN.md deleted file mode 100644 index 71ed790f3..000000000 --- a/lib/python/examples/async_cifar10/experiments/REFL_INTEGRATION_PLAN.md +++ /dev/null @@ -1,1249 +0,0 @@ -# REFL Integration Plan for Flame/Felix System - -**Date:** February 5, 2026 -**Last Updated:** February 6, 2026 -**Status:** Phases 1-3 COMPLETED ✅ | Phase 4 TESTING 🔄 -**Purpose:** Integrate REFL (Resource-Efficient Federated Learning) into Flame to enable head-to-head comparison with Felix on identical workloads. - ---- - -## Executive Summary - -REFL has been successfully integrated into Flame's modular FL framework. The implementation provides three independent components that can be toggled for ablation studies: availability-aware selection, deadline-based aggregation, and staleness handling. - -### Key Design Decisions - -**1. Modular Architecture** -- **REFLAvailabilityTracker**: Standalone availability prediction (flame/availability/) -- **REFLOortSelector**: Extends existing OortSelector (flame/selector/) -- **REFLFedAvg**: New optimizer with deadline filtering (flame/optimizer/) - -**2. Oracular Availability Tracking** -- REFL assumes **perfect knowledge** of client availability patterns via traces -- Traces define availability windows: `[active_start, inactive_start]` periods -- **Cannot function without traces** - this is a fundamental REFL assumption -- Our implementation supports both REFL's pickle format and Flame's YAML format - -**3. Synchronous FL Paradigm** -- REFL was designed for **synchronous FL** (not async like Felix) -- Uses deadline filtering to handle stragglers -- Caches stale updates for reuse in future rounds - -**4. No Changes to Aggregator/Trainer Code** -- Implementation uses Flame's **selector/optimizer abstractions** -- No modifications needed to main aggregator/trainer code -- Registration in `flame/selectors.py` and `flame/optimizers.py` is sufficient - ---- - -## Implementation Status - -### ✅ Phase 1: Foundation (COMPLETED) -- REFLAvailabilityTracker with trace loading and priority computation -- Extended TrainResult with REFL timing fields -- Configuration schema support - -### ✅ Phase 2: Selector (COMPLETED) -- REFLOortSelector with 3 priority modes (0=none, 1=fill, 2=strict) -- Adaptive pacer mechanism -- Frequency-based blacklisting -- Registered as "refl_oort" - -### ✅ Phase 3: Aggregator (COMPLETED) -- REFLFedAvg optimizer with deadline filtering -- Stale update caching and lifecycle -- 4 staleness weighting strategies (Equal, Average, AdaSGD, DynSGD, REFL) -- Registered as "reflfedavg" - -### 🔄 Phase 4: Testing & Validation (IN PROGRESS) -- Config format fixed (added taskid, proper channel structure) -- Small-scale tests ready to run -- Ablation study configs prepared -- Remaining: runtime validation and comparison - ---- - -## Core Design Choices - -### 1. Availability Tracking - -**Design Question:** How to handle client availability? - -**REFL's Approach:** -- Oracular traces provide perfect future knowledge -- Traces from MobiPerf dataset (real-world mobile device availability) -- Format: `{active: [t1, t3, ...], inactive: [t2, t4, ...], finish_time: T}` - -**Our Implementation:** -```python -# REFLAvailabilityTracker -def is_available(self, trainer_id, current_time): - """Check if trainer available at given time using trace""" - # Wraps time: current_time % trace_duration - # Checks if within active window -``` - -**Key Point:** REFL **requires traces** - it's not a predictor, it's an oracle with perfect knowledge. This is intentional for their research setup but differs from Felix's approach. - -### 2. Priority-Based Selection - -**Design Question:** How to select clients in heterogeneous availability? - -**REFL's Three-Tier Strategy:** -1. **Compute UCB scores** (like Oort) based on training loss -2. **Estimate availability probability** from trace history -3. **Combine**: `priority = ucb_score * availability_probability` - -**Three Priority Modes:** -- `avail_priority=0`: Standard Oort (ignore availability) -- `avail_priority=1`: Fill mode (high priority first, then regular) -- `avail_priority=2`: Strict mode (only select high priority clients) - -**Adaptive Pacer:** -- Adjusts `round_threshold` based on utility trends -- If utility improving: lower threshold (be more conservative) -- If utility declining: raise threshold (be more aggressive) - -**Blacklisting:** -- Tracks selection frequency per client -- Blacklists over-selected clients to improve fairness -- Configurable threshold and max blacklist length - -### 3. Deadline Filtering & Staleness - -**Design Question:** How to handle stragglers in synchronous FL? - -**REFL's Approach:** -- Set deadline (fixed or moving average) -- Filter updates that arrive after deadline -- Cache "stale" updates (late arrivals) for future rounds -- Apply staleness weighting when using cached updates - -**Four Weighting Strategies:** -| Strategy | stale_factor | Formula | Use Case | -|----------|--------------|---------|----------| -| Equal | 1 | weight = 1.0 | Baseline (ignore staleness) | -| Average | -1 | weight = 1/avg_staleness | Equal contribution across rounds | -| AdaSGD | -2 | weight = 1/(staleness+1) | Polynomial decay | -| DynSGD | -3 | weight = exp(-staleness) | Exponential decay | -| REFL | -4 | Hybrid utility formula | Balances all factors | - -**Moving Average Deadline:** -- Tracks round durations from previous rounds -- Computes percentile (e.g., 80th percentile) -- Adapts to workload variations - -### 4. Integration with Flame - -**Why No Aggregator/Trainer Code Changes?** - -Flame's architecture provides clean abstractions: -- **Selectors** are plug-and-play via `selector_provider` -- **Optimizers** are plug-and-play via `optimizer_provider` -- **TrainResult** can be extended with optional fields - -**What was needed:** -1. Register new selector in `flame/selectors.py` -2. Register new optimizer in `flame/optimizers.py` -3. Extend TrainResult with timing fields (backward compatible) -4. Create proper JSON configs with all required Flame fields - -**Migration Plan Adjustments:** -- The existing spawner/launcher system works as-is -- Aggregator/trainer main files unchanged -- Only config format needed updates (taskid, channel structure) - ---- - -## Configuration Structure - -### Two-Level Config System - -**Level 1: YAML Experiment Config** (experiments/configs/*.yaml) -```yaml -experiments: - - name: refl_test - trainer: - num_trainers: 5 - availability: syn_0 - aggregator: - config_template: expt_scripts_2026/configs/refl_config_test.json - selector: refl_oort -``` - -**Level 2: JSON Aggregator Config** (expt_scripts_2026/configs/*.json) -```json -{ - "taskid": "experiment_id", - "selector": {"sort": "refl_oort", "kwargs": {...}}, - "optimizer": {"sort": "reflfedavg", "kwargs": {...}}, - "hyperparameters": {...} -} -``` - -### Feature Toggles - -**Ablation Studies via Config:** -- Baseline: `selector=oort, optimizer=fedavg` -- Availability only: `selector=refl_oort, optimizer=fedavg` -- Staleness only: `selector=oort, optimizer=reflfedavg` -- Full REFL: `selector=refl_oort, optimizer=reflfedavg` - ---- - -## Remaining Work (Phase 4) - -### Current Status: Testing Phase - -**Issue Fixed (Feb 6):** -- ✅ Added `taskid` field to all JSON configs -- ✅ Fixed channel structure to match Flame schema -- ✅ Updated hyperparameters format - -**Next Steps:** -1. Test small-scale (5 trainers) - verify components work -2. Debug any runtime errors -3. Run ablation study - verify independent toggles -4. Add round timing to top aggregator (for deadline calculation) -5. Extend traces to 300 trainers if needed -6. Head-to-head comparison with Felix - -### Testing Commands - -```bash -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 - -# Small test -python3 launch/run_experiment.py experiments/configs/refl_test_5trainers.yaml - -# Ablation study -python3 launch/run_experiment.py experiments/configs/refl_ablation_study.yaml - -# Full scale -python3 launch/run_experiment.py experiments/configs/refl_full_n300.yaml -``` - ---- - -## Key Takeaways - -### What REFL Provides -- ✅ Availability-aware selection using oracular traces -- ✅ Priority-based selection with adaptive pacer -- ✅ Deadline filtering for stragglers -- ✅ Sophisticated staleness handling -- ✅ Improved fairness through blacklisting - -### What REFL Assumes -- ⚠️ **Oracular traces required** - perfect future knowledge of availability -- ⚠️ **Synchronous FL paradigm** - not asynchronous like Felix -- ⚠️ **Trace-based** - cannot function without availability data - -### Comparison with Felix -- **Felix**: Handles true asynchrony, no traces needed, staleness via FedBuff -- **REFL**: Synchronous with deadline, requires traces, sophisticated staleness weighting -- **Oort**: Baseline utility-based selection, no availability awareness - -### Implementation Achievements -- ✅ Modular design - toggle features independently -- ✅ Clean integration - no changes to core aggregator/trainer code -- ✅ Well-documented - comprehensive usage guides -- ✅ Production-ready - registered components, error handling, configs - ---- - -## Documentation Files - -- **[PHASE4_COMMANDS.md](PHASE4_COMMANDS.md)** - Quick command reference -- **[PHASE4_TESTING_GUIDE.md](PHASE4_TESTING_GUIDE.md)** - Comprehensive testing guide -- **[../aggregator/REFL_README.md](../aggregator/REFL_README.md)** - Configuration reference - ---- - -## Success Criteria - -- [x] All REFL components implemented ✅ -- [x] Modular architecture with independent toggles ✅ -- [x] Backward compatible with existing Flame code ✅ -- [x] Configuration system supporting ablation studies ✅ -- [ ] Small-scale test passing ⏳ -- [ ] Ablation study validating independent features ⏳ -- [ ] Full-scale (300 trainers) running successfully ⏳ -- [ ] REFL results reproduce published performance ⏳ -- [ ] Head-to-head comparison with Felix completed ⏳ - ---- - -## Appendix A: Implementation Files - -### Created Files (7) -- `flame/availability/__init__.py` -- `flame/availability/refl_tracker.py` (485 lines) -- `flame/selector/refl_oort.py` (464 lines) -- `flame/optimizer/reflfedavg.py` (558 lines) -- `experiments/configs/refl_test_5trainers.yaml` -- `experiments/configs/refl_full_n300.yaml` -- `experiments/configs/refl_ablation_study.yaml` - -### Modified Files (4) -- `flame/optimizer/train_result.py` - Added REFL timing fields -- `flame/selectors.py` - Registered "refl_oort" -- `flame/optimizers.py` - Registered "reflfedavg" -- `expt_scripts_2026/configs/*.json` - Fixed config format (taskid, channels) - -### Configuration Files (5) -- `expt_scripts_2026/configs/refl_config_test.json` - Test (5 trainers) -- `expt_scripts_2026/configs/refl_config_n300.json` - Full scale -- `expt_scripts_2026/configs/refl_config_ablation_baseline.json` -- `expt_scripts_2026/configs/refl_config_ablation_avail.json` -- `expt_scripts_2026/configs/refl_config_ablation_staleness.json` - -### Documentation Files (4) -- `experiments/REFL_INTEGRATION_PLAN.md` - This file -- `experiments/PHASE4_COMMANDS.md` - Quick reference -- `experiments/PHASE4_TESTING_GUIDE.md` - Testing guide -- `aggregator/REFL_README.md` - Configuration reference - ---- - -## Appendix B: REFL vs Felix Comparison - -| Aspect | REFL | Felix | -|--------|------|-------| -| **FL Paradigm** | Synchronous with deadline | Asynchronous | -| **Availability** | Oracular traces required | No traces needed | -| **Selection** | Priority-based with UCB | Oort UCB-based | -| **Staleness** | Deadline filtering + caching | FedBuff buffering | -| **Weighting** | 4 strategies (AdaSGD, DynSGD, etc) | Polynomial/exponential | -| **Fairness** | Blacklisting mechanism | Implicit via Oort | -| **Research Focus** | Resource efficiency with intermittent clients | True asynchrony handling | - -**When to use REFL:** -- Known availability patterns (traces available) -- Intermittent client participation -- Need fairness guarantees -- Synchronous FL acceptable - -**When to use Felix:** -- Unknown availability patterns -- True asynchronous environment -- No oracular knowledge needed -- Continuous client participation - -### 1.1 Availability Tracking -**REFL Implementation (`client_manager.py`):** -- Loads device availability traces from pickle files (`device_avail_file`) -- Each client has availability periods: `[(start_time, end_time), ...]` -- Key methods: - - `isClientActive(clientId, cur_time, time_window)`: checks if client is available at a future time - - `isAvailable(clientId, cur_time, time_window, time_slots)`: checks availability across time slots - - `getPriority(clientId, cur_time, time_window)`: returns priority based on near-term availability (0-2) - - `getPeriodCount(clientId, cur_time, deadline)`: counts availability periods within deadline - -**Flame Current State:** -- Uses `TrainerAvailState` enum: `AVL_TRAIN`, `AVL_EVAL`, `UN_AVL` -- Trainers self-manage state transitions based on traces -- Already supports trace-based unavailability in `async_cifar10` - -### 1.2 Client Selection (Oort Enhancement) -**REFL Implementation (`oort.py` + `aggregator.py`):** -- UCB-based selection with exploration/exploitation -- Statistical utility: normalized loss + temporal uncertainty -- System utility: penalizes slow clients based on `round_prefer_duration` -- Pacer mechanism: adaptively adjusts `round_threshold` to control client speed filtering -- Blacklisting: excludes clients selected too frequently -- Priority-based selection (`args.avail_priority`): - - 0: No priority - - 1: Fill remaining slots from non-priority clients - - 2: Only select high-priority clients first - -**Key Selection Parameters:** -- `exploration_factor`: Initially 0.9, decays by `exploration_decay` (0.98) -- `exploration_min`: Floor at 0.2 -- `round_threshold`: Controls speed filtering (default 30%, adaptive) -- `alpha`: Weight for staleness in utility (default 2) -- `clip_bound`: Caps utility at 95th percentile -- `cut_off_util`: Prunes low-utility clients (95% of cutoff) - -**Flame Current State:** -- Has basic Oort selector (`flame/selector/oort.py`) -- Synchronized FL support in `syncfl/top_aggregator.py` -- Needs enhancements for: - - Priority-based selection using availability - - Pacer mechanism - - Blacklisting per REFL - -### 1.3 Aggregation with Staleness Handling -**REFL Implementation (`aggregator.py`):** -- Tracks stale updates in `self.staleWeights[clientId]` -- Applies **deadline filtering** (`exp_type=0` or `exp_type=2`): - - Fixed deadline: `args.deadline` - - Moving average deadline: `mov_avg_deadline` - - Clients exceeding deadline become stragglers, updates cached as "stale" -- Stale update lifecycle: - 1. Client times out → update stored in `staleWeights[clientId]` - 2. Each round: `staleRemainDuration[clientId]` decrements by `round_duration` - 3. When `staleRemainDuration <= 0` and `stale_rounds <= args.stale_update`: apply update - 4. If `stale_rounds > args.stale_update`: discard (too stale) - -**Stale Weighting Strategies (`args.stale_factor`):** -- `> 1`: Divide by constant factor -- `1`: Equal weight (baseline FedAvg) -- `-1`: Divide by average staleness across all stale updates -- `-2`: AdaSGD - divide by `(staleness + 1)` -- `-3`: DynSGD - multiply by `exp(-(staleness + 1))` -- `-4`: **REFL method** - hybrid formula: - ```python - weight *= (1 - beta) / (staleness + 1) + beta * (1 - exp(-client_ratio / max_ratio) / scale_coff) - ``` - - `beta` (`args.stale_beta`): balance between staleness and utility - - `client_ratio`: importance based on dataset size or loss - - `scale_coff`: scaling coefficient (default 10.0) - -**Aggregation Formula:** -```python -global_model += client_weight * client_importance * update -``` -Where: -- `client_weight`: normalized by dataset size -- `client_importance`: adjusted by staleness strategy -- Normalize after aggregation to maintain model scale - -**Flame Current State:** -- `FedAvg` (`flame/optimizer/fedavg.py`): Simple weighted averaging -- `FedBuff` (`flame/optimizer/fedbuff.py`): Asynchronous aggregation with staleness -- FedBuff already has staleness handling with `alpha_polynomial`, `alpha_exponential`, etc. -- Needs: REFL-specific staleness strategy and deadline-based filtering - -### 1.4 Experimental Configurations (`exp_type`) -REFL uses `exp_type` to control aggregation behavior: -- **0**: Deadline + target ratio (SAFA baseline) -- **1**: No deadline, wait for all selected clients -- **2**: Overcommitment with deadline -- **3**: Overcommitment without deadline - -For Flame integration, we'll focus on **exp_type=0** and **exp_type=2** (deadline-based) since these align with REFL's core contribution. - ---- - -## 2. Integration Strategy - -### 2.1 Modular Design Principles -1. **Component Independence**: Each REFL feature toggleable via config -2. **Backward Compatibility**: Existing Flame experiments unaffected -3. **Minimal Breaking Changes**: Leverage Flame's abstractions -4. **Approximation Where Necessary**: Document deviations from REFL's exact behavior - -### 2.2 Three-Tier Implementation - -#### Tier 1: Availability Tracking (REFL-Compatible) -**Goal:** Enable REFL-style availability priority and deadline filtering. - -**Implementation:** -1. **Extend Availability Traces:** - - Current: Trainers use trace files with state transitions - - Add: Compute availability periods from traces at aggregator side - - Location: `lib/python/flame/mode/horizontal/syncfl/top_aggregator.py` - -2. **Aggregator-Side Availability Manager:** - ```python - class REFLAvailabilityTracker: - def __init__(self, trace_file_path): - self.traces = self.load_traces(trace_file_path) - self.availability_periods = self.compute_availability_periods() - - def is_available(self, trainer_id, cur_time, time_window): - # Check if trainer available in [cur_time, cur_time + time_window] - - def get_priority(self, trainer_id, cur_time, time_window): - # Return 0-2 based on near-term availability - - def compute_availability_periods(self): - # Convert traces to [(start, end), ...] per trainer - ``` - -3. **Config Extension:** - ```yaml - refl: - enabled: true - availability_trace_file: "path/to/traces.pkl" - use_priority_selection: true # Enable priority-based selection - avail_priority: 2 # 0=none, 1=fill, 2=strict - avail_probability: 1.0 # Accuracy of predictions (0-1) - ``` - -**Flame Approximation:** -- REFL's traces are centralized at aggregator, but Flame trainers self-manage state -- **Solution:** Aggregator loads same trace file and mirrors trainer state predictions -- **Trade-off:** Slight desync possible, but acceptable for fair comparison - -#### Tier 2: REFL-Enhanced Oort Selector -**Goal:** Implement REFL's priority-based selection, pacer, and blacklisting. - -**Implementation:** -1. **Create `REFLOortSelector` (extends `OortSelector`):** - - Location: `lib/python/flame/selector/refl_oort.py` - - Inherits from `flame/selector/oort.py` - - Adds: - - **Priority selection logic** (uses `REFLAvailabilityTracker`) - - **Pacer mechanism** for adaptive `round_threshold` - - **Blacklisting** based on selection frequency - -2. **Key Methods:** - ```python - class REFLOortSelector(OortSelector): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.avail_tracker = REFLAvailabilityTracker(kwargs['trace_file']) - self.blacklist_rounds = kwargs.get('blacklist_rounds', -1) - self.blacklist_max_len = kwargs.get('blacklist_max_len', 0.3) - self.pacer_step = kwargs.get('pacer_step', 20) - self.pacer_delta = kwargs.get('pacer_delta', 5) - - def select(self, ends, channel_props, trainer_unavail_list, task, **kwargs): - # 1. Build priority lists using avail_tracker - priority_ends, remaining_ends = self.build_priority_lists(ends) - - # 2. Apply blacklist filter - available_ends = self.filter_blacklist(priority_ends + remaining_ends) - - # 3. Run Oort UCB selection on available_ends - selected = super()._select_with_ucb(available_ends, ...) - - # 4. Fill from priority first if avail_priority >= 1 - if self.avail_priority >= 1: - selected = self.fill_priority_first(priority_ends, selected) - - # 5. Run pacer to adjust round_threshold - self.pacer() - - return selected - - def pacer(self): - # Adaptive adjustment of round_threshold based on utility trends - if self.round % self.pacer_step == 0: - util_change = self.compute_utility_change() - if abs(util_change) < 0.1: - self.round_threshold = min(100, self.round_threshold + self.pacer_delta) - elif abs(util_change) > 5.0: - self.round_threshold = max(self.pacer_delta, self.round_threshold - self.pacer_delta) - - def get_blacklist(self, ends): - # Return set of end_ids selected > blacklist_rounds times - blacklist = set() - for end_id, end in ends.items(): - if end.get_property(PROP_SELECTED_COUNT) > self.blacklist_rounds: - blacklist.add(end_id) - # Cap at blacklist_max_len * total_clients - return blacklist[:int(self.blacklist_max_len * len(ends))] - ``` - -3. **Config Extension:** - ```yaml - selector: - sort: refl_oort - kwargs: - aggr_num: 10 - blacklist_rounds: -1 # -1 disables, else max selections before blacklist - blacklist_max_len: 0.3 # Max 30% of clients blacklisted - avail_priority: 2 # 0=none, 1=fill, 2=strict - avail_probability: 1.0 # Availability prediction accuracy - pacer_step: 20 # Evaluate pacer every N rounds - pacer_delta: 5 # % adjustment to round_threshold - ``` - -**Flame Approximation:** -- REFL's `ucbSampler` is tightly coupled with `clientManager` -- **Solution:** Extend Flame's `OortSelector` to call `REFLAvailabilityTracker` for priorities -- **Trade-off:** Cleaner separation, minor implementation differences - -#### Tier 3: REFL-Aware Aggregator (Staleness Handling) -**Goal:** Implement deadline filtering and REFL's stale update weighting. - -**Implementation:** -1. **Create `REFLFedAvg` Optimizer:** - - Location: `lib/python/flame/optimizer/reflfedavg.py` - - Implements: - - **Deadline-based filtering** of trainer updates - - **Stale update caching** for stragglers - - **REFL staleness weighting** (`stale_factor=-4`) - - Support for other weighting strategies (AdaSGD, DynSGD, etc.) - -2. **Key Data Structures:** - ```python - class REFLFedAvg(AbstractOptimizer): - def __init__(self, **kwargs): - super().__init__(**kwargs) - self.stale_weights = {} # {trainer_id: [stale_params]} - self.stale_remain_duration = {} # {trainer_id: remaining_time} - self.stale_rounds = {} # {trainer_id: num_rounds_stale} - - # Config - self.deadline = kwargs.get('deadline', 0) # 0 = moving avg - self.mov_avg_deadline = 0 - self.stale_update_max = kwargs.get('stale_update', -1) # -1 = no limit - self.stale_factor = kwargs.get('stale_factor', 1) - self.stale_beta = kwargs.get('stale_beta', 0.9) - self.scale_coff = kwargs.get('scale_coff', 10.0) - ``` - -3. **Aggregation Flow:** - ```python - def do(self, base_weights, cache, total, version, **kwargs): - round_duration = kwargs.get('round_duration', 0) - - # 1. Separate fast and slow trainers based on deadline - fast_trainers, slow_trainers = self.filter_by_deadline(cache, round_duration) - - # 2. Cache stale updates from slow trainers - for tres in slow_trainers: - self.stale_weights[tres.end_id] = tres.weights - self.stale_remain_duration[tres.end_id] = tres.duration - self.deadline - self.stale_rounds[tres.end_id] = 0 - - # 3. Retrieve applicable stale updates from previous rounds - applicable_stale = self.get_applicable_stale_updates(round_duration) - - # 4. Compute importance weights for fast + applicable stale - all_trainers = fast_trainers + applicable_stale - importance_weights = self.compute_importance_weights(all_trainers) - - # 5. Aggregate with weighted averaging - for tres in all_trainers: - rate = (tres.count / total) * importance_weights[tres.end_id] - self.aggregate_fn(tres, rate) - - # 6. Update moving average deadline - if self.deadline == 0: - self.mov_avg_deadline = self.compute_moving_avg_deadline(fast_trainers) - - return self.agg_weights - - def compute_importance_weights(self, trainers): - # Implements REFL's stale weighting strategies - weights = {} - for tres in trainers: - staleness = self.stale_rounds.get(tres.end_id, 0) - - if self.stale_factor == 1: # Equal weight - weights[tres.end_id] = 1.0 - elif self.stale_factor == -2: # AdaSGD - weights[tres.end_id] = 1.0 / (staleness + 1) - elif self.stale_factor == -3: # DynSGD - weights[tres.end_id] = math.exp(-(staleness + 1)) - elif self.stale_factor == -4: # REFL - client_ratio = tres.stat_utility / max_stat_utility - weights[tres.end_id] = ( - (1 - self.stale_beta) / (staleness + 1) + - self.stale_beta * (1 - math.exp(-client_ratio / max_ratio) / self.scale_coff) - ) - - # Normalize weights - total_weight = sum(weights.values()) - return {k: v / total_weight for k, v in weights.items()} - - def get_applicable_stale_updates(self, round_duration): - # Returns stale updates that are now ready to apply - applicable = [] - for trainer_id in list(self.stale_weights.keys()): - self.stale_remain_duration[trainer_id] -= round_duration - self.stale_rounds[trainer_id] += 1 - - if (self.stale_remain_duration[trainer_id] <= 0 and - (self.stale_update_max < 0 or self.stale_rounds[trainer_id] <= self.stale_update_max)): - # Apply this stale update - applicable.append(self.create_tres_from_stale(trainer_id)) - del self.stale_weights[trainer_id] - del self.stale_remain_duration[trainer_id] - del self.stale_rounds[trainer_id] - - return applicable - ``` - -4. **Top Aggregator Integration:** - - Modify `syncfl/top_aggregator.py` to: - - Track round durations (start/end timestamps) - - Pass `round_duration` to optimizer - - Apply deadline filtering before aggregation - -5. **Config Extension:** - ```yaml - optimizer: - sort: reflfedavg - kwargs: - deadline: 100 # seconds, 0 = use moving average - stale_update: -1 # max staleness rounds, -1 = no limit - stale_factor: -4 # -4=REFL, -3=DynSGD, -2=AdaSGD, 1=equal - stale_beta: 0.9 # REFL beta parameter - scale_coff: 10.0 # REFL scaling coefficient - target_ratio: 0.8 # Target fraction of selected clients to wait for - ``` - -**Flame Approximation:** -- REFL's aggregator handles raw client updates directly -- **Solution:** Adapt to Flame's `TrainResult` objects and diskcache -- **Trade-off:** Slight memory overhead, but maintains Flame's architecture - ---- - -## 3. Config-Driven Feature Toggles - -To enable modular experimentation, all REFL components are independently toggleable: - -```yaml -# Baseline FedAvg (no REFL) -selector: - sort: random -optimizer: - sort: fedavg - -# REFL Availability Only -selector: - sort: refl_oort - kwargs: - avail_priority: 2 - blacklist_rounds: -1 # Disable blacklisting - pacer_step: -1 # Disable pacer -optimizer: - sort: fedavg # No staleness handling - -# REFL Selection + Staleness -selector: - sort: refl_oort - kwargs: - avail_priority: 2 - blacklist_rounds: 50 - pacer_step: 20 -optimizer: - sort: reflfedavg - kwargs: - stale_factor: -4 # REFL weighting - -# REFL Full (All Features) -selector: - sort: refl_oort - kwargs: - avail_priority: 2 - blacklist_rounds: 50 - pacer_step: 20 -optimizer: - sort: reflfedavg - kwargs: - deadline: 100 - stale_factor: -4 -``` - ---- - -## 4. Integration with async_cifar10 - -### 4.1 Trainer Config Updates -Trainers in `async_cifar10` already support: -- Trace-based availability (`TrainerAvailState`) -- Self-managed state transitions -- Statistical utility reporting (loss, dataset size) - -**Additions Needed:** -- Report additional metrics for REFL: - - Training duration (already captured) - - Completion timestamp - - Staleness (if applicable) - -**No breaking changes required** - trainers continue operating as before. - -### 4.2 Aggregator Config Updates -Create REFL-specific configs in `/aggregator/`: - -```json -{ - "realm": "...", - "selector": { - "sort": "refl_oort", - "kwargs": { - "aggr_num": 10, - "avail_priority": 2, - "blacklist_rounds": 50, - "pacer_step": 20, - "pacer_delta": 5 - } - }, - "optimizer": { - "sort": "reflfedavg", - "kwargs": { - "deadline": 100, - "stale_update": 5, - "stale_factor": -4, - "stale_beta": 0.9, - "scale_coff": 10.0 - } - }, - "refl": { - "availability_trace_file": "metadata/availability_traces/mobiperf_traces.yaml", - "use_priority_selection": true - } -} -``` - -### 4.3 Metadata Integration -REFL experiments require availability traces. Extend `experiments/metadata/availability_traces/`: - -```yaml -# synthetic_traces.yaml (matching REFL's exp_type configs) -syn_0: - description: "Always available (0% dropout)" - traces: - 1: {periods: [[0, 1e12]], duration: 1e12} - 2: {periods: [[0, 1e12]], duration: 1e12} - # ... for all trainers - -syn_20: - description: "20% synthetic dropout" - # Generated from REFL's trace generation logic - -mobiperf_2st: - description: "Real-world MobiPerf traces (2-state)" - # Converted from REFL's pickle format -``` - ---- - -## 5. Implementation Phases - -### ✅ Phase 1: Foundation (COMPLETED) - -**Status:** All components implemented and registered ✅ - -1. **✅ Created `REFLAvailabilityTracker`** (`flame/availability/refl_tracker.py`) - - Loads availability traces from pickle or YAML formats - - Implements `is_available()`, `is_client_active()`, `get_priority()` - - Implements `split_by_priority()` for priority-based selection - - Supports both REFL pickle format and Flame YAML format - - Handles trace duration wrapping and period calculations - -2. **✅ Extended TrainResult** (`flame/optimizer/train_result.py`) - - Added REFL-specific fields: `completion_time`, `round_duration`, `staleness`, `end_id` - - Backward compatible with existing code - - Fields are optional (default to None/0) - -3. **✅ Config Schema Updates** - - Created example configs with REFL parameters - - Documented all parameters in REFL_README.md - - Modular design allows independent feature toggles - -**Deliverables:** -- ✅ `flame/availability/refl_tracker.py` (485 lines) -- ✅ `flame/availability/__init__.py` -- ✅ Modified `flame/optimizer/train_result.py` - -### ✅ Phase 2: Selector (COMPLETED) - -**Status:** REFLOortSelector fully implemented and registered ✅ - -1. **✅ Created `REFLOortSelector`** (`flame/selector/refl_oort.py`) - - Extends base `OortSelector` class - - Implements three priority modes (0=none, 1=fill, 2=strict) - - Adaptive pacer mechanism adjusts `round_threshold` based on utility trends - - Blacklisting prevents over-selection (configurable threshold and max length) - - Integrates with `REFLAvailabilityTracker` for availability predictions - -2. **✅ Registration & Integration** - - Registered in `flame/selectors.py` as `"refl_oort"` - - Compatible with existing selector interface - - Can be used with any optimizer - -3. **✅ Config Templates** - - Test config: `aggregator/refl_config_test.json` - - Full scale: `aggregator/refl_config_n300.json` - - Ablation configs for each feature combination - -**Deliverables:** -- ✅ `flame/selector/refl_oort.py` (464 lines) -- ✅ Modified `flame/selectors.py` (registered selector) -- ✅ Configuration files with REFL selector examples - -### ✅ Phase 3: Aggregator (COMPLETED) - -**Status:** REFLFedAvg optimizer fully implemented and registered ✅ - -1. **✅ Created `REFLFedAvg` Optimizer** (`flame/optimizer/reflfedavg.py`) - - Deadline-based filtering (fixed or moving average) - - Stale update caching with lifecycle management - - Four staleness weighting strategies: - - Equal (stale_factor=1) - Standard FedAvg - - Average (stale_factor=-1) - Divide by average staleness - - AdaSGD (stale_factor=-2) - Divide by (staleness + 1) - - DynSGD (stale_factor=-3) - Exponential decay - - REFL (stale_factor=-4) - Hybrid utility-based formula - - Moving average deadline with configurable target percentile - - Statistics tracking (cached/applied/discarded stale updates) - -2. **✅ Registration & Integration** - - Registered in `flame/optimizers.py` as `"reflfedavg"` - - Compatible with existing optimizer interface - - Works with both PyTorch and TensorFlow backends - -3. **✅ Configuration Examples** - - All ablation configs created - - Examples for each staleness strategy - - Comparison configs for REFL vs Felix - -**Deliverables:** -- ✅ `flame/optimizer/reflfedavg.py` (558 lines) -- ✅ Modified `flame/optimizers.py` (registered optimizer) -- ✅ Ablation configuration files (3 configs) -- ✅ Comprehensive usage guide: `aggregator/REFL_README.md` - -### 📋 Phase 4: Validation & Tuning (PENDING) - -**Status:** Ready to begin ⏳ - -1. **Unit Tests** (Not Started) - - Test `REFLAvailabilityTracker` methods - - Test staleness weighting formulas - - Test deadline filtering logic - - Test pacer mechanism - -2. **Integration Tests** (Not Started) - - Small-scale experiments (5 trainers) - - Enable REFL features one at a time - - Compare against non-REFL baseline - -3. **REFL Validation** (Not Started) - - Reproduce key REFL results on CIFAR-10 - - Validate convergence speed and resource efficiency - - Compare fairness metrics - -4. **Syncfl Top Aggregator Integration** (Not Started) - - Add round timing metadata to `syncfl/top_aggregator.py` - - Pass `round_duration` and `cur_time` to optimizer - - Ensure TrainResult objects populated correctly - -5. **Availability Trace Generation** (Not Started) - - Extend synthetic_traces.yaml to 300 trainers - - Convert REFL's MobiPerf traces to YAML format - - Create trainer-to-trace mappings - -6. **Head-to-Head Comparison** (Not Started) - - Run REFL vs Felix on identical workloads - - Document performance differences - - Create comparison plots and analysis - -**Next Steps:** -1. Create unit tests for core REFL components -2. Modify `syncfl/top_aggregator.py` to provide round timing -3. Generate full availability traces for 300 trainers -4. Run small-scale integration tests -5. Validate against REFL's published results - ---- - -## 6. Testing Strategy - -### 6.1 Unit Tests -- **Availability Tracker:** - - Test `is_available` logic with synthetic traces - - Validate priority computation -- **REFL Selector:** - - Test UCB selection matches expected distributions - - Verify blacklist enforcement - - Confirm pacer adjustments -- **REFL Optimizer:** - - Test staleness weighting formulas - - Validate deadline filtering - - Ensure stale cache lifecycle correct - -### 6.2 Integration Tests -- **Small-Scale Experiments (5 trainers):** - - Use `test_phase3_mini.yaml` as baseline - - Enable REFL features one at a time - - Compare against non-REFL baseline - -### 6.3 Validation Against REFL -- **Reproduce Key Results:** - - Use REFL's published hyperparameters - - Run on CIFAR-10 with `n=300`, `alpha=0.1` - - Compare: - - Convergence speed (rounds to target accuracy) - - Resource efficiency (compute + communication) - - Fairness metrics (Gini coefficient, KL divergence) - ---- - -## 7. Known Deviations & Approximations - -### 7.1 Centralized vs. Distributed Availability -**REFL:** Aggregator has centralized view of all client availability traces. -**Flame:** Trainers self-manage state, aggregator observes. -**Approximation:** Aggregator loads same trace file and predicts trainer states. -**Impact:** Minimal - acceptable for controlled experiments. - -### 7.2 Event-Driven vs. Round-Based Timing -**REFL:** Uses event queue with virtual clock for simulation. -**Flame:** Real-time system with actual network delays. -**Approximation:** Track actual round durations, apply deadline filtering post-facto. -**Impact:** Moderate - may affect deadline tuning. - -### 7.3 Model Update Format -**REFL:** Direct NumPy arrays in memory. -**Flame:** PyTorch tensors via `TrainResult` + diskcache. -**Approximation:** Convert between formats as needed. -**Impact:** Negligible - performance overhead only. - -### 7.4 Executor Model -**REFL:** Explicit executor processes managed by aggregator. -**Flame:** Trainers spawn independently, communicate via channels. -**Approximation:** Maintain REFL's aggregator-centric selection logic. -**Impact:** None - selection algorithm unchanged. - ---- - -## 8. Success Criteria - -### 8.1 Functional Completeness -- [x] All three REFL components implemented (availability, selection, aggregation) ✅ -- [x] All staleness weighting strategies available (`-4`, `-3`, `-2`, `-1`, `1`) ✅ -- [x] Priority-based selection working with configurable `avail_priority` ✅ -- [x] Deadline filtering operational with both fixed and moving average ✅ -- [x] Stale update caching and lifecycle management correct ✅ -- [x] Blacklisting and pacer mechanisms implemented ✅ - -### 8.2 Experimental Validation (Pending Phase 4) -- [ ] Reproduce REFL's reported accuracy on CIFAR-10 (within ±2%) -- [ ] Confirm resource efficiency gains (compute + communication) -- [ ] Validate fairness metrics match REFL's trends -- [ ] Unit tests for all core components -- [ ] Integration tests with async_cifar10 - -### 8.3 Fair Comparison Setup (Pending Phase 4) -- [x] Configuration system supports identical setups ✅ -- [ ] Run REFL and Felix on identical: - - Datasets (same splits, same Dirichlet alpha) - - Availability traces (same trainer-to-trace mappings) - - Hyperparameters (batch size, learning rate, etc.) - - Evaluation protocol (same test sets, same metrics) - -### 8.4 Documentation & Usability -- [x] README with REFL integration guide (`aggregator/REFL_README.md`) ✅ -- [x] Example configs for all feature toggle combinations ✅ -- [x] Troubleshooting guide for common issues ✅ -- [x] Performance tuning recommendations ✅ -- [x] Comprehensive integration plan with checkpoint ✅ - ---- - -## 9. Config Examples - -### Example 1: REFL Full Configuration (300 trainers, alpha=0.1, syn_0) -```yaml -# experiments/configs/refl_n300_alpha0.1_syn0.yaml -experiments: - - name: refl_n300_alpha0.1_syn0 - description: "REFL with 300 trainers, alpha=0.1, always available" - - trainer: - num_trainers: 300 - dataset: - dirichlet_alpha: 0.1 - availability: - mode: syn_0 - - aggregator: - ✅ Completed (Feb 5, 2026) - -**Phases 1-3 COMPLETED** - All core REFL components implemented - -**Implementation Summary:** -- **Lines of Code:** ~1,500+ lines of new implementation -- **Files Created:** 10 new files (7 implementation + 3 configs) -- **Files Modified:** 4 files (registration and extensions) -- **Documentation:** Comprehensive usage guide and integration plan - -**Key Achievements:** -1. ✅ Modular architecture - each REFL component can be toggled independently -2. ✅ Multiple staleness strategies - Equal, AdaSGD, DynSGD, REFL hybrid -3. ✅ Priority-based selection - 3 modes with availability awareness -4. ✅ Adaptive deadline - fixed or moving average -5. ✅ Backward compatible - existing Flame experiments unaffected -6. ✅ Well-documented - usage guide with examples and troubleshooting - -### 📋 Remaining Work (Phase 4+) - -**Testing & Integration:** -- Syncfl top aggregator modifications for round timing -- Unit tests for all REFL components -- Integration tests with async_cifar10 -- Availability trace generation for 300 trainers - -**Validation:** -- Reproduce REFL's CIFAR-10 results -- Head-to-head comparison: REFL vs Felix -- Performance analysis and documentation - -**Estimated Timeline:** 2-3 weeks -- Week 1: Testing and top aggregator integration -- Week 2: Trace generation and reproduction experiments -- Week 3: Comparison experiments and documentation - availability_trace_file: "metadata/availability_traces/synthetic_traces.yaml" - use_priority_selection: true - - execution: - num_gpus: 8 - sleep_between_spawns: 5.0 - aggregator_warmup_time: 600 -``` - -### Example 2: Comparison - Felix vs. REFL -```yaml -# experiments/configs/comparison_felix_vs_refl.yaml -experiments: - - name: felix_n300_alpha0.1_syn0 - # ... Felix config ... - aggregator: - selector: - sort: async_oort # Felix selector - optimizer: - sort: fedbuff # Felix optimizer - - - name: refl_n300_alpha0.1_syn0 - # ... REFL config (as above) ... -``` - -### Example 3: Ablation - Availability Only -```yaml -# experiments/configs/ablation_avail_only.yaml -experiments: - - name: refl_avail_only - aggregator: - selector: - sort: refl_oort - kwargs: - avail_priority: 2 - blacklist_rounds: -1 # Disable - pacer_step: -1 # Disable - optimizer: - sort: fedavg # Standard FedAvg, no staleness -``` - ---- - -## 10. Implementation Roadmap - -### Week 1-2: Foundation -- **Deliverables:** - - `lib/python/flame/availability/refl_tracker.py` - - Updated `TrainResult` with timing fields - - Config schema extensions - - Unit tests for availability tracker - -### Week 3-4: Selector -- **Deliverables:** - - `lib/python/flame/selector/refl_oort.py` - - Integration with `syncfl/top_aggregator.py` - - Selector unit tests - - Test configs for priority selection - -### Week 5-6: Aggregator -- **Deliverables:** - - `lib/python/flame/optimizer/reflfedavg.py` - - Deadline filtering in top aggregator - - Stale update lifecycle tests - - End-to-end integration tests - -### Week 7-8: Validation -- **Deliverables:** - - Reproduced REFL experiments on CIFAR-10 - - Comparison report: REFL vs. Felix - - Documentation updates - - Example configs and launch scripts - ---- - -## 11. Open Questions & Future Work - -### 11.1 Open Questions -1. **Moving Average Deadline:** REFL uses `target_ratio` to compute deadline - should we replicate exact formula or use Flame's existing approach? -2. **Overcommitment vs. Aggr Goal:** REFL selects `overcommitment × aggr_goal` clients - should this be selector or aggregator responsibility? -3. **Stale Update Expiry:** REFL uses `args.stale_skip_round` flag - do we need this for Flame? - -### 11.2 Future Enhancements -1. **SAFA Integration:** REFL compares against SAFA (exp_type=0) - consider implementing SAFA as another baseline. -2. **Additional Datasets:** Extend beyond CIFAR-10 to Google Speech, OpenImage. -3. **Adaptive Hyperparameters:** Auto-tune `deadline`, `stale_beta` based on workload characteristics. -4. **Visualization Tools:** Dashboards showing availability patterns, staleness distributions, etc. - ---- - -## 12. References - -### REFL Papers -- **REFL arXiv:** https://arxiv.org/abs/2111.01108 -- **REFL EuroSys'23:** ACM EuroSys 2023 proceedings - -### REFL Codebase -- **Location:** `/home/dgarg39/flame/third_party/REFL` -- **Key Files:** - - `core/aggregator.py`: Main aggregation logic - - `core/client_managdocuments the **successful implementation** of REFL within Flame's abstraction framework. By leveraging Flame's modular design, we have implemented REFL's availability tracking, client selection, and aggregation algorithms as **composable components** that can be toggled independently. - -### Implementation Achievements - -**Phases 1-3 COMPLETED (Feb 5, 2026):** -1. ✅ **Correctness:** Faithfully implemented REFL's core algorithms -2. ✅ **Modularity:** Each component independently toggleable for ablation studies -3. ✅ **Minimal Disruption:** Maintained backward compatibility with existing Flame experiments -4. ✅ **Well-Documented:** Comprehensive usage guide with examples and troubleshooting -5. ✅ **Production-Ready:** Registered components, configuration examples, error handling - -### Code Statistics - -- **New Files:** 10 (7 implementation, 3 documentation) -- **Modified Files:** 4 (registrations and extensions) -- **Lines of Code:** ~1,500+ new implementation -- **Configuration Files:** 5 test/ablation configs -- **Documentation:** 2 comprehensive guides - -### Remaining Work (Phase 4) - -**Critical for Experiments:** -1. Top aggregator integration for round timing metadata -2. Availability trace generation for 300 trainers -3. Unit and integration tests - -**For Validation:** -1. Reproduce REFL's published results -2. Head-to-head comparison with Felix -3. Performance analysis and documentation - -### Quick Start for Next Session - -```bash -# Test the implementation -cd /home/dgarg39/flame/lib/python/examples/async_cifar10 -python3 launch/run_experiment.py aggregator/refl_config_test.json - -# Run ablation studies -python3 launch/run_experiment.py aggregator/refl_config_ablation_baseline.json -python3 launch/run_experiment.py aggregator/refl_config_ablation_avail.json -python3 launch/run_experiment.py aggregator/refl_config_ablation_staleness.json -``` - -**Next Steps for Phase 4:** -1. Add round timing to syncfl/top_aggregator.py -2. Create unit tests for REFL components -3. Generate full availability traces (300 trainers) -4. Run integration tests -5. Validate against REFL's published results -6. Execute REFL vs Felix comparison experiments ---- - -## Conclusion - -This integration plan provides a **comprehensive roadmap** to implement REFL within Flame's abstraction framework. By leveraging Flame's modular design, we can implement REFL's availability tracking, client selection, and aggregation algorithms as **composable components** that can be toggled independently. - -The plan prioritizes: -1. **Correctness:** Faithfully implementing REFL's algorithms -2. **Modularity:** Each component toggleable for ablation studies -3. **Fair Comparison:** Ensuring identical experimental conditions for REFL vs. Felix -4. **Minimal Disruption:** Maintaining backward compatibility with existing Flame experiments - -**Next Steps:** Review this plan, identify any concerns or needed clarifications, then proceed with Phase 1 implementation. diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_ablation_study.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_ablation_study.yaml index 64127e9d7..ec29afbf5 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_ablation_study.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_ablation_study.yaml @@ -14,7 +14,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 # 20% unavailability - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_config_ablation_baseline.json @@ -41,7 +40,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_config_ablation_avail.json @@ -68,7 +66,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_config_ablation_staleness.json @@ -95,7 +92,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_config_n300.json diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_full_n300.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_full_n300.yaml index a6a408c6a..3d9436b17 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_full_n300.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_full_n300.yaml @@ -13,7 +13,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 # Synthetic trace - 0% unavailability (always available) - speedup_factor: 1.0 enable_training_delays: true # Enable per-trainer training delays from metadata aggregator: diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_monitoring_examples.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_monitoring_examples.yaml index e2e8cb7b0..abc7c07e6 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_monitoring_examples.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_monitoring_examples.yaml @@ -15,7 +15,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 # 0% unavailability (always available) - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_n300_syn0.json @@ -55,7 +54,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_n300_syn0.json @@ -90,7 +88,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_10trainer_demo.json diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn20_syn50_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn20_syn50_comparison.yaml index 7e96d5a7f..c3e3117a8 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn20_syn50_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn20_syn50_comparison.yaml @@ -43,7 +43,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 # 20% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 # Trainer batch size (matching REFL default for CIFAR-10) @@ -84,7 +83,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_50 # 50% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn50_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn50_comparison.yaml index bd3197941..498f84d6b 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn50_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.7_syn50_comparison.yaml @@ -43,7 +43,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_50 # 50% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn20_syn50_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn20_syn50_comparison.yaml index dfc145be2..8af66f709 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn20_syn50_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn20_syn50_comparison.yaml @@ -43,7 +43,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 # 20% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 # Trainer batch size (matching REFL default for CIFAR-10) @@ -84,7 +83,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_50 # 50% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn50_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn50_comparison.yaml index 3b8082aaa..8df54a60b 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn50_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_prob0.9_syn50_comparison.yaml @@ -43,7 +43,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_50 # 50% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn20_alpha0.1_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn20_alpha0.1_comparison.yaml index 991a256d3..83fc0f2f3 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn20_alpha0.1_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn20_alpha0.1_comparison.yaml @@ -42,7 +42,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 # 20% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 # Trainer batch size (matching REFL default for CIFAR-10) @@ -83,7 +82,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_20 # 20% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn50_alpha0.1_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn50_alpha0.1_comparison.yaml index 8467e3ad0..ea2dfe9eb 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn50_alpha0.1_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_n300_syn50_alpha0.1_comparison.yaml @@ -42,7 +42,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_50 # 50% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 # Trainer batch size (matching REFL default for CIFAR-10) @@ -83,7 +82,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_50 # 50% unavailability - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_test_5trainers.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_test_5trainers.yaml index c28cea269..08112ea1e 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_test_5trainers.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_test_5trainers.yaml @@ -13,7 +13,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 # Always available for quick testing - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_config_test.json diff --git a/lib/python/examples/async_cifar10/experiments/configs/refl_vs_oort_n300_comparison.yaml b/lib/python/examples/async_cifar10/experiments/configs/refl_vs_oort_n300_comparison.yaml index 57af71f73..33916daa6 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/refl_vs_oort_n300_comparison.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/refl_vs_oort_n300_comparison.yaml @@ -45,7 +45,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 # 0% unavailability (always available) - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 # Trainer batch size (matching REFL default for CIFAR-10) @@ -86,7 +85,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 @@ -127,7 +125,6 @@ experiments: dirichlet_alpha: 0.1 availability: mode: syn_0 - speedup_factor: 1.0 enable_training_delays: true hyperparameters: batchSize: 10 diff --git a/lib/python/examples/async_cifar10/experiments/configs/test_oort_syn0.yaml b/lib/python/examples/async_cifar10/experiments/configs/test_oort_syn0.yaml index b9b8a8d12..97d493cd2 100644 --- a/lib/python/examples/async_cifar10/experiments/configs/test_oort_syn0.yaml +++ b/lib/python/examples/async_cifar10/experiments/configs/test_oort_syn0.yaml @@ -14,7 +14,6 @@ experiments: availability: mode: syn_0 # Always available for quick testing battery_threshold: 50 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/oort_n300_oracular_9may25_syn0.json diff --git a/lib/python/examples/async_cifar10/experiments/refl_parity.md b/lib/python/examples/async_cifar10/experiments/refl_parity.md deleted file mode 100644 index d87b5af7a..000000000 --- a/lib/python/examples/async_cifar10/experiments/refl_parity.md +++ /dev/null @@ -1,1550 +0,0 @@ -# REFL Implementation Parity Analysis - -**Date:** February 9, 2026 -**Purpose:** Verify that the Flame REFL integration correctly implements the original REFL design and behavior - ---- - -## Part 1: REFL Design & Flow (Conceptual Overview) - -### 1.1 What is REFL? - -REFL (Resource-Efficient Federated Learning) is a synchronous FL system designed to handle **intermittent client availability** and **heterogeneous device capabilities**. Unlike asynchronous systems (like Felix), REFL operates in synchronous rounds but uses intelligent mechanisms to handle stragglers and late arrivals. - -**Core Philosophy:** -- **Predictive Availability**: Use oracular traces to know *in advance* which clients will be available -- **Deadline-Based Tolerance**: Wait for a subset of clients, not all (synchronous FL with deadline) -- **Stale Update Recycling**: Cache late arrivals (stragglers) and reuse them later with staleness weighting - ---- - -### 1.2 Three Pillars of REFL - -#### **Pillar 1: Availability-Aware Client Selection** - -**Problem:** In mobile/edge FL, clients have unpredictable availability patterns (devices go offline/online) - -**REFL's Solution:** -1. **Oracular Traces**: Load real-world availability traces from MobiPerf dataset - - Format: `{active: [t1, t3, ...], inactive: [t2, t4, ...], finish_time: T}` - - Traces wrap cyclically: `current_time % finish_time` - -2. **Priority-Based Selection**: - - Compute **UCB scores** (like Oort) for exploration/exploitation trade-off - - Compute **availability probability** from trace data - - **Priority = UCB × Availability** - -3. **Three Priority Modes**: - - `avail_priority=0`: Ignore availability (standard Oort) - - `avail_priority=1`: **Fill mode** - select high-priority first, then fill remaining slots - - `avail_priority=2`: **Strict mode** - only select high-priority clients - -4. **Adaptive Pacer**: - - Tracks utility trends over time - - If utility declining → raise threshold (be more aggressive in selection) - - If utility improving → lower threshold (be more conservative) - -5. **Blacklisting for Fairness**: - - Track how often each client is selected - - Temporarily blacklist over-selected clients - - Ensures diverse participation - -**Key Insight:** By knowing which clients will be available *in the near future*, REFL can preferentially select clients that are both high-utility AND reliably available. - ---- - -#### **Pillar 2: Deadline-Based Aggregation** - -**Problem:** In synchronous FL, waiting for all clients creates long rounds due to stragglers - -**REFL's Solution:** -1. **Set Deadline**: Either fixed or moving average (e.g., 80th percentile of past round times) - -2. **Fast vs Slow Classification**: - - Collect trainer updates as they arrive - - At aggregation time, split into: - - **Fast trainers**: arrived before deadline - - **Slow trainers (stragglers)**: arrived after deadline - -3. **Immediate Aggregation**: - - Aggregate fast trainers immediately - - Don't wait for slow trainers - - Move to next round quickly - -**Key Insight:** Deadline filtering provides **bounded staleness** - we know the maximum age of any stale update (1 round). - ---- - -#### **Pillar 3: Stale Update Management** - -**Problem:** Discarding straggler updates wastes computation and data - -**REFL's Solution:** -1. **Stale Update Caching**: - - When a trainer arrives after deadline → cache their update - - Store: `{trainer_id: {weights, staleness, utility, timestamp}}` - -2. **Lifecycle Management**: - - Track staleness age: how many rounds old is this update? - - Apply `stale_update_max` limit (e.g., discard if > 5 rounds old) - - Clean up expired stale updates - -3. **Reuse in Future Rounds**: - - In next round, if we don't have enough fast updates: - - Check stale cache for applicable updates - - Apply staleness weighting - - Include in aggregation - -4. **Four Staleness Weighting Strategies**: - - | Strategy | Formula | Description | - |----------|---------|-------------| - | **Equal** (`stale_factor=1`) | `weight = 1.0` | No penalty - treat like fresh update | - | **Average** (`stale_factor=-1`) | `weight = 1/avg_staleness` | Inverse of average staleness | - | **AdaSGD** (`stale_factor=-2`) | `weight = 1/(staleness+1)` | Polynomial decay (from AdaSGD paper) | - | **DynSGD** (`stale_factor=-3`) | `weight = exp(-staleness)` | Exponential decay (from DynSGD paper) | - | **REFL** (`stale_factor=-4`) | Hybrid formula | Balances loss, staleness, and utility | - - **REFL Weighting Formula**: - ``` - statistical_utility = loss_i / sum(losses) # Higher loss = more to learn - system_utility = 1 / (staleness_i + 1) # Fresher = better - - weight_i = β * statistical_utility + (1-β) * system_utility - - where β (beta) = 0.9 (emphasize data utility over freshness) - ``` - - **Critical Design Choice: β=0.9** - - REFL heavily weights **statistical utility (90%)** over **staleness penalty (10%)**. This means: - - A stale update from a high-loss client gets nearly full weight - - A fresh update from a low-loss client gets discounted - - Staleness has minimal impact on aggregation weights - - **Rationale:** REFL assumes that data distribution information (captured by loss) is more valuable than freshness, especially for rare/infrequently-available clients whose data distributions might not be well-represented otherwise. - -**Key Insight:** Stale updates still contain **valuable gradient information** from data distributions, especially from infrequently-available clients. REFL prioritizes data diversity over update freshness. - ---- - -### 1.3 REFL Training Flow - -Here's the complete flow of a REFL training round: - -``` -┌─────────────────────────────────────────────────────────┐ -│ ROUND N STARTS │ -│ Aggregator has: global_model_v(N) │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 1. CLIENT SELECTION (REFLOortSelector) │ -│ │ -│ a. Get connected trainers from channel │ -│ b. Compute UCB scores for each trainer │ -│ c. Check availability from traces (oracular) │ -│ d. Compute priority = UCB × availability │ -│ e. Apply adaptive pacer (adjust threshold) │ -│ f. Filter blacklisted trainers │ -│ g. Select K trainers (priority first, then remaining) │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 2. WEIGHT DISTRIBUTION │ -│ │ -│ - Send global_model_v(N) to selected trainers │ -│ - Record round_start_time for each trainer │ -│ - Trainers begin local training │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 3. TRAINER EXECUTION (Parallel) │ -│ │ -│ Each Trainer: │ -│ - Load local data partition │ -│ - Train for E local epochs │ -│ - Compute local update = model - global_model │ -│ - Compute utility = loss-based metric │ -│ - Send {update, utility, metadata} back │ -│ │ -│ Note: Trainers complete at different times! │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 4. UPDATE COLLECTION (As they arrive) │ -│ │ -│ Aggregator receives updates one-by-one: │ -│ - Update from Trainer A at t1 │ -│ - Update from Trainer C at t2 │ -│ - Update from Trainer B at t3 │ -│ - ... │ -│ │ -│ Tracks: arrival_time - round_start_time = duration │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 5. DEADLINE FILTERING (REFLFedAvg) │ -│ │ -│ When ready to aggregate (K updates or timeout): │ -│ │ -│ fast_trainers = [] │ -│ slow_trainers = [] │ -│ │ -│ for each trainer_update: │ -│ if duration < deadline: │ -│ fast_trainers.append(trainer_update) │ -│ else: │ -│ slow_trainers.append(trainer_update) │ -│ # Cache for later - stale update │ -│ stale_cache[trainer_id] = { │ -│ weights: update, │ -│ staleness: 1, # Fresh straggler │ -│ utility: utility, │ -│ round: N │ -│ } │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 6. STALE UPDATE AUGMENTATION │ -│ │ -│ if len(fast_trainers) < target: │ -│ # Not enough fast updates - use stale cache │ -│ for trainer_id in stale_cache: │ -│ if stale_cache[trainer_id].staleness <= max_stale: │ -│ # Compute staleness weight │ -│ if stale_factor == -4: # REFL weighting │ -│ stat_util = loss / sum_losses │ -│ sys_util = 1 / (staleness + 1) │ -│ weight = beta*stat_util + (1-beta)*sys_util │ -│ elif stale_factor == -2: # AdaSGD │ -│ weight = 1 / (staleness + 1) │ -│ # ... other strategies ... │ -│ │ -│ fast_trainers.append({ │ -│ update: stale_cache[trainer_id].weights, │ -│ weight: weight, │ -│ is_stale: True │ -│ }) │ -│ │ -│ # Age existing stale updates │ -│ for trainer_id in stale_cache: │ -│ stale_cache[trainer_id].staleness += 1 │ -│ │ -│ # Cleanup expired stale updates │ -│ remove entries where staleness > stale_update_max │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 7. FEDERATED AVERAGING │ -│ │ -│ global_model_v(N+1) = Σ (weight_i × update_i) │ -│ │ -│ where: │ -│ - weight_i = staleness_weight × data_weight │ -│ - data_weight = num_samples_i / total_samples │ -│ - staleness_weight = from step 6 │ -│ │ -│ Note: Fast updates get full weight (1.0) │ -│ Stale updates get discounted weight │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────┐ -│ 8. POST-AGGREGATION │ -│ │ -│ - Update model version: N → N+1 │ -│ - Update deadline (if using moving average) │ -│ - Log metrics (accuracy, loss, round time) │ -│ - Update pacer state based on utility trends │ -│ - Update blacklist based on selection counts │ -└─────────────────────────────────────────────────────────┘ - │ - ▼ - ROUND N+1 STARTS -``` - ---- - -### 1.4 Key Differences from Other FL Systems - -| Aspect | Traditional Sync FL | Asynchronous FL (Felix) | REFL | -|--------|---------------------|-------------------------|------| -| **Paradigm** | Synchronous | Asynchronous | Synchronous with deadline | -| **Stragglers** | Wait for all | No concept (async) | Deadline filtering | -| **Staleness** | No staleness | Unbounded, dynamically weighted | Bounded (1 round initially) | -| **Selection** | Random/utility | **Availability-aware** (reactive notify) | **Availability-aware** (oracular traces) | -| **Availability Type** | N/A | Reactive (trainers notify) | Predictive (oracular traces) | -| **Round Duration** | Slowest client | N/A (continuous) | Deadline-bounded | -| **Update Caching** | No | Buffering (FedBuff) | Stale update cache | -| **Staleness Weighting** | N/A | Polynomial/exponential decay | 90% data utility, 10% staleness | -| **Fairness** | Random ensures | Via async nature | Blacklisting mechanism | - ---- - -### 1.5 Critical Assumptions in REFL - -1. **Oracular Availability Traces Required**: - - REFL cannot function without availability data - - Assumes perfect knowledge of future availability - - In real systems, would need predictive models - -2. **Synchronous FL Paradigm**: - - Despite handling stragglers, still fundamentally synchronous - - All fast trainers train on same global model version - - Different from truly asynchronous systems - -3. **Bounded Staleness**: - - Stale updates are 1+ rounds old (not arbitrary age) - - Staleness increases by 1 each round - - Old stale updates eventually discarded - -4. **Deadline Setting**: - - Fixed deadline: must be tuned for workload - - Moving average: reactive, not predictive - - Trade-off: short deadline = more stragglers, long deadline = slow rounds - -5. **Data Utility > Freshness (β=0.9)**: - - Assumes data distribution information is 9x more valuable than freshness - - Works well when: diverse data distributions, rare clients, non-IID data - - May struggle when: model evolves rapidly, fresh gradients critical - ---- - -### 1.6 REFL vs Felix: Staleness Weighting Deep Dive - -#### **Fundamental Difference in Philosophy** - -| Aspect | REFL Approach | Felix Approach | -|--------|---------------|----------------| -| **Weighting Basis** | Data utility (loss) > Staleness | Staleness-aware with configurable strategies | -| **Weight Formula** | `0.9 * (loss/Σloss) + 0.1 * (1/staleness)` | Polynomial or exponential decay | -| **Staleness Impact** | **Minimal (10%)** | **Significant (primary factor)** | -| **Design Goal** | Maximize data diversity | Minimize gradient staleness | -| **Best For** | Non-IID with rare clients | Fast-evolving models | - -#### **When REFL's β=0.9 Works Well** - -1. **Highly Non-IID Data**: - - Each client has unique data distribution - - Rare participation patterns - - Example: Medical FL with specialized hospitals - -2. **Slow Model Evolution**: - - Gradients remain relevant across rounds - - Model converges slowly - - Example: Large language models with stable objectives - -3. **Intermittent High-Value Clients**: - - Some clients have high-quality data but low availability - - Worth waiting/caching their updates - - Example: Edge devices with unique sensor data - -#### **When Felix's Staleness-First Works Better** - -1. **Fast Model Evolution**: - - Model changes significantly each round - - Old gradients mislead optimization - - Example: RL agents in dynamic environments - -2. **IID or Mildly Non-IID Data**: - - Data distributions similar across clients - - No unique "rare" information - - Example: ImageNet classification - -3. **High Client Availability**: - - Many clients available at any time - - Fresh updates readily available - - Example: Server-grade FL (not edge/mobile) - -#### **Asynchronous Setting: Which Performs Better?** - -**In truly asynchronous settings, Felix likely has advantages:** - -1. **Natural Async Fit**: - - Felix designed for async from ground up - - No artificial round boundaries - - Continuous aggregation as updates arrive - -2. **FedBuff Mechanism**: - - Dynamically manages buffer of updates - - Intelligently selects which updates to aggregate - - Adapts to arrival patterns automatically - -3. **Reactive Availability**: - - Trainers notify when available/unavailable - - No need for oracular predictions - - Works in unpredictable environments - -4. **Staleness Decay**: - - Exponential/polynomial decay naturally handles varying staleness - - Fresh updates preferred when available - - Old updates gracefully degraded - -**REFL's advantages in sync-with-deadline setting:** - -1. **Predictive Selection**: - - Can avoid selecting clients that will become unavailable - - Better resource utilization - - Fewer wasted selections - -2. **Data Diversity Priority**: - - Ensures rare clients contribute - - Better handling of non-IID data - - Fairness in representation - -3. **Round Structure Benefits**: - - Clear checkpoints for evaluation - - Easier to reason about convergence - - Model versioning simpler - -**Hybrid Approach Potential:** - -Combining REFL's availability prediction with Felix's async aggregation could be optimal: -- Use REFL's oracular selection for predictive client selection -- Use Felix's FedBuff for async aggregation -- Adjust β based on model evolution speed (high β for slow evolution, low β for fast) - ---- - -### 1.7 REFL/Oort vs Felix: Client Selection Philosophy (Exploration vs Exploitation) - -#### **Fundamental Difference: Concentration vs Parallelism** - -| Aspect | REFL/Oort Approach | Felix Approach | -|--------|-------------------|----------------| -| **Selection Size** | Small subset per round (e.g., 30 of 300) | Large parallel pool (all available) | -| **Exploration/Exploitation** | **Exploitation-heavy** (90% exploit, 10% explore) | **Lightweight parallel exploration** | -| **Reuse Pattern** | Concentrate on high-utility trainers | Broadly sample available trainers | -| **Selection Mechanism** | UCB-based utility optimization | Availability-driven opportunistic | -| **Trainer Utilization** | ~25% of trainers used frequently | ~70-100% of trainers used over time | - -#### **REFL/Oort's Exploitation-Heavy Strategy** - -**How It Works:** -1. **Small Selection Window**: In each round, select only `aggr_num` clients (e.g., 30 out of 300) - - With overcommitment: `1.3 × 30 = 39` selected, wait for 30 to complete - -2. **Exploitation Dominance**: - ```python - exploration_len = int(30 * 0.9) = 27 # 90% exploitation - exploitation_len = 3 # 10% exploration - ``` - - **27 clients** selected by UCB utility scores (exploit best performers) - - **3 clients** randomly selected from unexplored trainers (explore new options) - -3. **UCB Selection Formula**: - ```python - score = (utility - min_utility) / range_utility + # Statistical utility - sqrt(0.1 * log(round) / last_selected_round) + # Temporal uncertainty - (round_duration / preferred_duration)^(-alpha) # System utility - ``` - - High-utility trainers: Get selected nearly every round - - Medium-utility: Selected occasionally - - Low-utility: Rarely/never selected - -**Result: Concentration on Top Performers** - -Example from actual run (300 trainers, 104 rounds): -- **Only 74 trainers used** out of 300 deployed (24.7%) -- **19 trainers** selected in 51 rounds (nearly every round - top performers) -- **10 trainers** selected 40-50 times (frequent participants) -- **45 trainers** selected 1-10 times (rare participants) -- **226 trainers** never selected (not yet explored) - -**Distribution:** -``` -Selection Frequency: -51 times: ████████████████████ (19 trainers) <- Top performers -40-50: ████████████ (10 trainers) -30-40: ███████ (8 trainers) -20-30: ████ (7 trainers) -10-20: ███ (12 trainers) -1-10: █ (18 trainers) -0: (226 trainers) <- Never explored -``` - -**Rationale:** -- **Convergence Speed**: Focus resources on trainers that provide highest utility -- **Resource Efficiency**: Don't waste time on low-utility trainers -- **Proven by Oort**: Achieves faster convergence in original Oort paper - -#### **Felix's Lightweight Parallel Exploration Strategy** - -**How It Works:** -1. **Opportunistic Parallelism**: All available trainers can be active simultaneously - - No artificial limit on concurrent trainers - - Each trainer works at their own pace - -2. **Availability-Driven Selection**: - - When a trainer becomes available → assign work immediately - - No waiting for "optimal" trainer - - No UCB scoring or utility-based filtering - -3. **Asynchronous Aggregation (FedBuff)**: - - Maintain buffer of updates from different trainers - - Aggregate when buffer reaches threshold (e.g., 30 updates) - - Updates come from diverse set of trainers over time - -**Result: Broad Trainer Coverage** - -Example Felix behavior (300 trainers, continuous operation): -- **200-250 trainers used** over equivalent time period (67-83%) -- More uniform distribution of participation -- Lower variance in selection frequency -- Better representation of all data distributions - -**Rationale:** -- **Robustness**: Don't over-rely on subset of trainers -- **Fairness**: All trainers get opportunity to contribute -- **Adaptability**: Handles dynamic availability naturally -- **Heterogeneity**: Captures diverse data distributions - -#### **When Each Approach Excels** - -**REFL/Oort's Exploitation Works Better When:** - -1. **Fast Convergence is Critical**: - - Training budget limited - - Need results quickly - - High-utility trainers are reliably available - -2. **Utility Gap is Large**: - - Some trainers have much better data/updates than others - - Clear winners in UCB scoring - - Example: Medical FL with specialized hospitals vs general clinics - -3. **Homogeneous Availability**: - - All trainers equally available - - No intermittency concerns - - Example: Datacenter FL with stable nodes - -**Felix's Exploration Works Better When:** - -1. **High Device Heterogeneity**: - - Compute speeds vary widely - - Network conditions unpredictable - - **Different trainers available at different times** - -2. **Non-IID Data with Long Tail**: - - Many trainers have unique data distributions - - Rare trainers have valuable minority class data - - Need diversity more than top utility - - Example: Mobile keyboard prediction (rare words from diverse users) - -3. **Fairness Requirements**: - - All participants should benefit from model - - Regulatory requirements for inclusive participation - - Example: Cross-silo FL with contractual participation guarantees - -4. **Dynamic Availability Patterns**: - - Trainers come online/offline unpredictably - - No oracular availability traces available - - Need to adapt to whoever is available - - Example: Mobile edge devices with user-driven availability - -#### **Impact on Heterogeneous Scenarios** - -**Why Felix Likely Performs Better in Practice:** - -1. **Captures Long-Tail Distributions**: - ``` - REFL/Oort selects: [Top 30 utility trainers] - Felix explores: [All available trainers over time] - - In non-IID: - - REFL misses: 70% of data distributions (from 226 unused trainers) - - Felix covers: 80%+ of data distributions (200+ trainers used) - ``` - -2. **Adapts to Dynamic Behavior**: - - High-utility trainer goes offline → REFL struggles, Felix adapts - - New trainer with rare data comes online → Felix uses immediately, REFL needs rounds to discover - -3. **Avoids Overfitting to Subset**: - - REFL: Model optimized for top 74 trainers' data - - Felix: Model trained on broader distribution - - Result: Felix generalizes better to held-out trainers - -4. **Better Statistical Representation**: - ``` - Central Limit Theorem: Variance ∝ 1/n - - REFL: n_effective = 74 trainers → higher variance - Felix: n_effective = 200+ trainers → lower variance - ``` - -**Trade-off Summary:** - -| Metric | REFL/Oort | Felix | -|--------|-----------|-------| -| **Convergence Speed** | ✅ Faster | ⚠️ Slower | -| **Final Accuracy (IID)** | ✅ Higher | ⚠️ Lower | -| **Final Accuracy (Non-IID)** | ⚠️ Lower | ✅ Higher | -| **Robustness to Stragglers** | ⚠️ Brittle | ✅ Resilient | -| **Fairness (Participation)** | ❌ Poor | ✅ Good | -| **Data Distribution Coverage** | ❌ 25% trainers | ✅ 70%+ trainers | -| **Rare Class Performance** | ❌ Poor | ✅ Good | -| **Scalability** | ✅ Efficient | ⚠️ Overhead | - -#### **Recommendation for Heterogeneous Scenarios** - -**For mobile/edge FL with heterogeneous devices and non-IID data, Felix's approach is preferable:** - -1. Realistic device heterogeneity means utility scores are noisy and change over time -2. Non-IID data means rare trainers have irreplaceable information -3. Dynamic availability makes concentration risky (what if your top 30 all go offline?) -4. Fairness and robustness matter in production systems - -**REFL/Oort best for:** -- Controlled environments (cross-silo FL) -- Homogeneous, reliable clients -- IID data where any K clients are representative -- Research settings where convergence speed is primary metric - -**Potential Hybrid:** -- Use Felix's parallel exploration for trainer discovery and data coverage -- Incorporate REFL's β=0.9 staleness weighting for rare clients -- Use adaptive exploration factor that increases under heterogeneity - ---- - -## Part 2: Parity Verification (Evidence-Based Analysis) - -### 2.1 Availability Tracking Parity - -#### **Original REFL (third_party/REFL)** - -**File:** `core/helper/client.py` - -**Key Method:** -```python -def isActive(self, cur_time): - if self.traces is None: - return True - - # Wrap time cyclically - norm_time = cur_time % self.traces['finish_time'] - - # Update behavior index - if norm_time > self.traces['inactive'][self.behavior_index]: - self.behavior_index += 1 - self.behavior_index %= len(self.traces['active']) - - # Check if in active window - if (self.traces['active'][self.behavior_index] <= norm_time <= - self.traces['inactive'][self.behavior_index]): - return True - return False -``` - -**Trace Format:** -```python -{ - 'duration': 211625, - 'finish_time': 518400, # 6 days in seconds - 'active': [12788, 100044, 188992, ...], # Start times - 'inactive': [65881, 133574, 208292, ...], # End times - 'model': 'CPH1801' # Device model -} -``` - -**Priority Computation (client_manager.py):** -```python -def getPriority(self, clientId, cur_time, time_window, lookup_timeslots=2): - priority = 0 - for i in range(lookup_timeslots, 0, -1): - if self.isAvailable(clientId, cur_time, time_window, i): - priority = lookup_timeslots - i - break - return priority -``` - ---- - -#### **Flame Implementation** - -**File:** `flame/availability/refl_tracker.py` - -**Evidence from Logs:** -``` -2026-02-06 16:48:56,806 | refl_tracker.py:72 | INFO | MainThread | load_trainer_registry | - Loaded trainer ID mapping for 300 trainers - -2026-02-06 16:48:56,819 | refl_tracker.py:105 | INFO | MainThread | load_traces | - Detected pattern-based trace file with 3 patterns - -2026-02-06 16:48:56,819 | refl_tracker.py:132 | INFO | MainThread | compute_availability_periods | - Pattern mode enabled - availability will be computed on-demand -``` - -**Parity Assessment:** ✅ **CONFIRMED** - -**Evidence:** -1. **Trace Loading**: Successfully loaded 300 trainer mappings -2. **Pattern Detection**: Handles both individual traces and pattern-based traces -3. **Architecture**: Flame's `REFLAvailabilityTracker` encapsulates availability logic -4. **Format Support**: Supports both REFL's pickle format and Flame's YAML format - -**Note:** While original REFL computes availability in `Client.isActive()`, Flame centralizes this in `REFLAvailabilityTracker`. This is an architectural difference, not a functional difference. - ---- - -### 2.2 Client Selection Parity - -#### **Original REFL** - -**File:** `core/aggregator.py` + Oort (`thirdparty/oort`) - -**Selection Flow:** -1. Call `client_manager.resampleClients()` with deadline info -2. Oort computes UCB scores -3. REFL adds availability-based priority filtering -4. Returns (selected_clients, priority_clients) - -**Key Parameters:** -- `avail_priority`: 0/1/2 mode -- `blacklist_rounds`: Blacklist threshold -- `pacer_step`: Adaptive pacer interval - ---- - -#### **Flame Implementation** - -**File:** `flame/selector/refl_oort.py` - -**Evidence from Logs:** - -**Initialization:** -``` -2026-02-06 16:48:56,819 | refl_oort.py:84 | INFO | MainThread | __init__ | - REFLOortSelector initialized: avail_priority=1, blacklist_rounds=50, pacer_step=20 -``` - -**Selection in Action:** -``` -# Round 0 - Single trainer available -2026-02-06 16:49:01,518 | refl_oort.py:123 | INFO | Thread-3 |select | - REFL Oort selecting 1 ends for round 0, task: train, avail_priority=1 - -2026-02-06 16:49:01,519 | refl_oort.py:143 | INFO | Thread-3 | select | - Priority split: 0 priority, 1 remaining, 0 blacklisted - -2026-02-06 16:49:01,519 | refl_oort.py:184 | INFO | Thread-3 | select | - Selected 1 ends: {'505f9fc483cf4df68a2409257b5fad7d3c580370'} - -# Round 2 - Multiple trainers available -2026-02-06 16:49:10,985 | refl_oort.py:123 | INFO | Thread-3 | select | - REFL Oort selecting 6 ends for round 2, task: train, avail_priority=1 - -2026-02-06 16:49:10,985 | refl_oort.py:143 | INFO | Thread-3 | select | - Priority split: 0 priority, 10 remaining, 0 blacklisted - -2026-02-06 16:49:10,985 | refl_oort.py:184 | INFO | Thread-3 | select | - Selected 6 ends: {'505f9fc483cf4df68a2409257b5fad7d3c580374', - '505f9fc483cf4df68a2409257b5fad7d3c580378', - '505f9fc483cf4df68a2409257b5fad7d3c580376', - '505f9fc483cf4df68a2409257b5fad7d3c580372', - '505f9fc483cf4df68a2409257b5fad7d3c580379', - '505f9fc483cf4df68a2409257b5fad7d3c580370'} -``` - -**Parity Assessment:** ✅ **CONFIRMED** - -**Evidence:** -1. **Initialization**: All REFL parameters correctly set (avail_priority=1, blacklist_rounds=50, pacer_step=20) -2. **Selection Logic**: Correctly splits clients into priority/remaining/blacklisted buckets -3. **Dynamic Selection**: Adapts to available trainers (1 in round 0, 6 in later rounds) -4. **Logging**: Provides same visibility as original REFL - -**Observations:** -- `avail_priority=1` (fill mode) is active -- No priority clients in these rounds (0 priority) - all clients have similar availability in test -- No blacklisting yet (0 blacklisted) - early rounds, no over-selection -- Selection count scales appropriately with available trainers - ---- - -### 2.3 Deadline Filtering & Aggregation Parity - -#### **Original REFL** - -**File:** `core/aggregator.py` - -**Deadline Logic (Simplified):** -```python -# Split updates by deadline -fast_clients = [] -slow_clients = [] - -for client_id, update_data in client_updates.items(): - completion_time = virtual_client_clock[client_id] - - if completion_time <= deadline: - fast_clients.append((client_id, update_data)) - else: - slow_clients.append((client_id, update_data)) - # Cache stale update - staleWeights[client_id] = copy.deepcopy(update_data['update_weight']) -``` - ---- - -#### **Flame Implementation** - -**File:** `flame/optimizer/reflfedavg.py` - -**Evidence from Logs:** - -**Optimizer Initialization:** -``` -2026-02-06 16:48:57,550 | reflfedavg.py:87 | INFO | MainThread | __init__ | - REFLFedAvg initialized: deadline=100, stale_update_max=5, stale_factor=-4 -``` - -**Deadline Filtering in Action:** -``` -# Round 1 -2026-02-06 16:49:08,099 | reflfedavg.py:144 | INFO | MainThread | do | - Deadline filtering: 1 fast, 0 slow (deadline=100.00s) - -2026-02-06 16:49:08,100 | reflfedavg.py:155 | INFO | MainThread | do | - Stale updates: 0 applicable, 0 still cached - -2026-02-06 16:49:08,100 | reflfedavg.py:177 | INFO | MainThread | do | - Aggregated 1 trainers (1 fast + 0 stale) - -# Round 2 -2026-02-06 16:49:25,756 | reflfedavg.py:144 | INFO | MainThread | do | - Deadline filtering: 5 fast, 0 slow (deadline=100.00s) - -2026-02-06 16:49:25,756 | reflfedavg.py:155 | INFO | MainThread | do | - Stale updates: 0 applicable, 0 still cached - -2026-02-06 16:49:25,758 | reflfedavg.py:177 | INFO | MainThread | do | - Aggregated 5 trainers (5 fast + 0 stale) - -# Round 6 -2026-02-06 16:50:27,275 | reflfedavg.py:144 | INFO | MainThread | do | - Deadline filtering: 5 fast, 0 slow (deadline=100.00s) - -2026-02-06 16:50:27,275 | reflfedavg.py:155 | INFO | MainThread | do | - Stale updates: 0 applicable, 0 still cached - -2026-02-06 16:50:27,277 | reflfedavg.py:177 | INFO | MainThread | do | - Aggregated 5 trainers (5 fast + 0 stale) -``` - -**Parity Assessment:** ✅ **CONFIRMED** - -**Evidence:** -1. **Deadline Set**: Fixed deadline of 100 seconds matches configuration -2. **Fast/Slow Split**: Correctly categorizes all updates (0 slow in these rounds) -3. **Stale Cache**: Tracking stale updates (0 applicable/cached in clean run) -4. **Aggregation Count**: Matches selected trainers (1, then 5, consistently) - -**Observations:** -- In this test run, all trainers completed before deadline (no stragglers) -- This is expected for small-scale test (10 trainers, CIFAR-10) -- Stale update infrastructure is active but unused (0 applicable, 0 cached) -- Would see stale updates in larger scale or with deadline pressure - ---- - -### 2.4 Round Duration Tracking - -#### **Evidence from Logs** - -**Round Start Times (Tracked per Trainer):** -``` -# Round 2, Trainer ...370 -2026-02-06 16:49:01,519 | channel.py:139 | INFO | MainThread | set_end_property | - SET property round_start_time with val (1, datetime.datetime(2026, 2, 6, 16, 49, 1, 519577)) - for end_id 505f9fc483cf4df68a2409257b5fad7d3c580370 -``` - -**Round Completion and Duration:** -``` -# Round 2, Trainer ...370 completes -2026-02-06 16:49:15,389 | channel.py:150 | INFO | MainThread | get_end_property | - GOT property round_start_time with val (2, datetime.datetime(2026, 2, 6, 16, 49, 11, 47835)) - for end_id 505f9fc483cf4df68a2409257b5fad7d3c580370 - -2026-02-06 16:49:15,389 | channel.py:139 | INFO | MainThread | set_end_property | - SET property round_duration with val 0:00:04.336114 for end_id 505f9fc483cf4df68a2409257b5fad7d3c580370 -``` - -**Duration Examples:** -- Trainer ...370: 4.34 seconds (fast) -- Trainer ...376: 11.59 seconds (slower) -- Trainer ...379: 12.43 seconds (slower) -- Trainer ...372: 13.24 seconds (slowest) - -**Parity Assessment:** ✅ **CONFIRMED** - -**Evidence:** -1. **Round Timing**: Tracks start time per trainer (not global) -2. **Duration Computation**: Correctly computes elapsed time -3. **Heterogeneity**: Captures device heterogeneity (4s to 13s range) -4. **Deadline Check**: All under 100s deadline → all classified as fast - -**Note:** This per-trainer timing is compatible with REFL's deadline filtering approach. - ---- - -### 2.5 Utility Computation - -#### **Evidence from Logs** - -**Utility Values:** -``` -# Round 2 -2026-02-06 16:49:15,390 | channel.py:139 | INFO | MainThread | set_end_property | - SET property stat_utility with val 122.65815663356189 for end_id ...370 - -2026-02-06 16:49:22,607 | channel.py:139 | INFO | MainThread | set_end_property | - SET property stat_utility with val 205.9889943692114 for end_id ...376 - -2026-02-06 16:49:23,475 | channel.py:139 | INFO | MainThread | set_end_property | - SET property stat_utility with val 106.59921677616596 for end_id ...379 - -2026-02-06 16:49:24,277 | channel.py:139 | INFO | MainThread | set_end_property | - SET property stat_utility with val 0.0 for end_id ...372 -``` - -**Parity Assessment:** ✅ **CONFIRMED** - -**Evidence:** -1. **Utility Tracking**: Each trainer sends utility metric -2. **Value Range**: Varies across trainers (0 to 205) - reflects loss differences -3. **Zero Utility**: Trainer ...372 has 0.0 utility (possible convergence or error) -4. **Storage**: Stored as end property for future selection decisions - -**Note:** These utilities feed into Oort's UCB selection mechanism for next round. - ---- - -### 2.6 Aggregation Completion - -#### **Evidence from Logs** - -``` -2026-02-06 16:49:25,760 | top_aggregator.py:140 | INFO | MainThread | _aggregate_weights | - ====== aggregation finished for round 2, - self._updates_recevied: {'505f9fc483cf4df68a2409257b5fad7d3c580370': 2, - '505f9fc483cf4df68a2409257b5fad7d3c580376': 1, - '505f9fc483cf4df68a2409257b5fad7d3c580379': 1, - '505f9fc483cf4df68a2409257b5fad7d3c580372': 1, - '505f9fc483cf4df68a2409257b5fad7d3c580378': 1} - -2026-02-06 16:49:25,760 | runtime.py:76 | INFO | MainThread | wrapper | - Runtime of aggregate is 14.703593254089355 -``` - -**Parity Assessment:** ✅ **CONFIRMED** - -**Evidence:** -1. **Update Counting**: Tracks how many updates received from each trainer -2. **Aggregation Time**: Records time taken for aggregation step -3. **Completion**: Clean completion without errors - ---- - -### 2.7 Missing Evidence & Logging Gaps - -While the core REFL functionality is confirmed, there are some aspects we cannot fully verify from current logs: - -#### **1. Stale Update Weighting (Cannot Verify)** - -**Reason:** No stragglers in this test run (all trainers completed before deadline) - -**What we'd need to see:** -``` -# Expected log (not present): -reflfedavg.py | Deadline filtering: 3 fast, 2 slow (deadline=100.00s) -reflfedavg.py | Stale updates: 2 applicable, 2 still cached -reflfedavg.py | Applying staleness weights: trainer_X (staleness=2, weight=0.333) -reflfedavg.py | Aggregated 5 trainers (3 fast + 2 stale) -``` - -**Recommendation:** No code changes needed. To verify: -1. Run with tighter deadline (e.g., deadline=5 seconds) -2. Or run with more heterogeneous devices -3. Or artificially delay some trainers in test - -#### **2. Priority Client Selection (Cannot Verify)** - -**Reason:** All trainers have similar availability in test (0 priority clients) - -**What we'd need to see:** -``` -# Expected log (not present): -refl_oort.py | Priority split: 3 priority, 7 remaining, 0 blacklisted -refl_oort.py | Selected 6 ends: 3 priority + 3 regular -``` - -**Recommendation:** No code changes needed. To verify: -1. Use traces with varied availability patterns -2. Or set `avail_priority=2` (strict mode) to force priority requirement - -#### **3. Blacklisting (Cannot Verify)** - -**Reason:** Too few rounds to trigger over-selection threshold - -**What we'd need to see:** -``` -# Expected log (not present): -refl_oort.py | Priority split: 2 priority, 6 remaining, 2 blacklisted -refl_oort.py | Blacklisted trainers: [...372, ...375] (selected 51 times) -``` - -**Recommendation:** No code changes needed. To verify: -1. Run longer experiment (50+ rounds) -2. Or lower `blacklist_rounds` threshold to 5 - -#### **4. Adaptive Pacer (Cannot Verify)** - -**Reason:** No explicit pacer adjustment logs - -**What we'd need to see:** -``` -# Expected log (not present): -refl_oort.py | Pacer adjustment at round 20: threshold 0.8 → 0.75 (utility improving) -``` - -**Recommendation:** **ADD LOGGING** in `flame/selector/refl_oort.py`: -```python -# In _update_pacer() method -logger.info( - f"Pacer adjustment at round {round_num}: " - f"threshold {old_threshold:.3f} → {new_threshold:.3f} " - f"(utility {'improving' if delta > 0 else 'declining'})" -) -``` - -#### **5. Availability Prediction Usage (Cannot Verify)** - -**Reason**: No runtime availability check logs - -**What we'd need to see:** -``` -# Expected log (not present): -refl_tracker.py | Checking availability for trainer_372 at time=12345.6 -refl_tracker.py | Trainer_372 priority: 2 (fully available in next 2 timeslots) -``` - -**Recommendation:** **ADD LOGGING** in `flame/availability/refl_tracker.py`: -```python -# In get_priority() method -logger.debug( - f"Availability check: trainer={trainer_id}, priority={priority}, " - f"near_term_windows={available_windows}" -) -``` - ---- - -## Part 3: Code-Level Parity Verification - -### 3.1 Availability Tracking - -| Aspect | Original REFL | Flame REFL | Status | -|--------|---------------|------------|--------| -| **Trace Loading** | `pickle.load()` | Supports pickle + YAML | ✅ Enhanced | -| **Time Wrapping** | `cur_time % finish_time` | Same logic | ✅ Parity | -| **Active Check** | `Client.isActive()` | `REFLAvailabilityTracker.is_available()` | ✅ Parity | -| **Priority Computation** | `ClientManager.getPriority()` | `REFLAvailabilityTracker.get_priority()` | ✅ Parity | -| **Trainer Registry** | N/A | Flame-specific mapping | ✅ Enhanced | - -### 3.2 Client Selection - -| Aspect | Original REFL | Flame REFL | Status | -|--------|---------------|------------|--------| -| **Base Selector** | Oort (external) | `OortSelector` (Flame) | ✅ Parity | -| **Priority Modes** | 0/1/2 | 0/1/2 | ✅ Parity | -| **Blacklisting** | `blacklist_rounds` | `blacklist_rounds` | ✅ Parity | -| **Adaptive Pacer** | Manual threshold | `pacer_step`, `pacer_delta` | ✅ Parity | -| **UCB Scoring** | Oort implementation | Inherited from `OortSelector` | ✅ Parity | - -### 3.3 Aggregation - -| Aspect | Original REFL | Flame REFL | Status | -|--------|---------------|------------|--------| -| **Deadline Type** | Fixed or moving avg | Fixed or moving avg | ✅ Parity | -| **Fast/Slow Split** | Manual loop | `_filter_by_deadline()` | ✅ Parity | -| **Stale Caching** | `staleWeights dict` | `stale_weights dict` | ✅ Parity | -| **Staleness Aging** | Manual increment | Automatic aging | ✅ Parity | -| **Weighting Strategies** | 4 strategies | 4 strategies | ✅ Parity | -| **REFL Formula** | Beta=0.9, scale=10 | Beta=0.9, scale=10 | ✅ Parity | - -### 3.4 Overall Architecture - -| Aspect | Original REFL | Flame REFL | Status | -|--------|---------------|------------|--------| -| **Synchronous FL** | Yes | Yes | ✅ Parity | -| **Oracular Traces** | Required | Required | ✅ Parity | -| **Modular Design** | Monolithic aggregator | Selector + Optimizer separation | ✅ Enhanced | -| **Configuration** | Command-line args | JSON config | ✅ Enhanced | -| **Logging** | Basic Python logging | Structured logging | ✅ Enhanced | - ---- - -## Part 4: Conclusions & Recommendations - -### 4.1 Comparison with Felix: Key Distinctions - -#### **1. Client Selection Philosophy: Exploitation vs Exploration (CRITICAL DIFFERENCE)** - -| Aspect | REFL/Oort | Felix | -|--------|-----------|-------| -| **Selection Strategy** | **Exploitation-heavy** (90% top utility) | **Parallel exploration** (all available) | -| **Clients per Round** | Small subset (e.g., 30 of 300) | Opportunistic (all available) | -| **Coverage** | ~25% trainers used over time | ~70-100% trainers used over time | -| **Optimization Goal** | Maximize convergence speed | Maximize robustness & fairness | -| **Best for** | IID data, stable clients, speed priority | Non-IID, heterogeneous devices, fairness | - -**Key Finding from Logs:** -- REFL run with 300 trainers over 104 rounds: **Only 74 trainers (25%) ever used** - - 19 trainers selected 51 times (top performers - nearly every round) - - 226 trainers never selected (not yet explored) -- This is **by design** in Oort: concentrate on high-utility trainers for faster convergence - -**Why This Matters for Heterogeneous Scenarios:** - -Felix's lightweight parallel exploration fundamentally differs from REFL/Oort's approach: - -1. **Data Distribution Coverage:** - - REFL: Trains on data from ~25% of trainers (misses 75% of distributions) - - Felix: Trains on data from 70-100% of trainers (captures long-tail distributions) - -2. **Handling Heterogeneity:** - - REFL: Brittle - relies on top 30 trainers being reliably available - - Felix: Resilient - adapts to whoever is available at any time - -3. **Non-IID Performance:** - - REFL: Optimized for top performers' data, may not generalize - - Felix: Better statistical representation from broader sampling - -4. **Fairness & Robustness:** - - REFL: 226 trainers contribute nothing (wasted deployment) - - Felix: All trainers have opportunity to contribute - -**Recommendation:** For realistic heterogeneous FL scenarios (mobile devices, non-IID data, dynamic availability), **Felix's exploration approach is superior** to REFL/Oort's exploitation-heavy strategy. The concentration on top performers sacrifices robustness and fairness for convergence speed. - -#### **2. Availability Awareness (Both are availability-aware!)** - -| System | Availability Type | Mechanism | Knowledge | -|--------|------------------|-----------|------------| -| **Felix** | Reactive | Trainer notify callbacks | Current state | -| **REFL** | Predictive | Oracular traces | Future availability | - -**Key Point:** Both systems are availability-aware, but differ in how: -- Felix: Trainers actively notify aggregator when they transition between available/unavailable/training/eval states -- REFL: Aggregator predicts availability using traces, selects clients likely to complete - -#### **3. Staleness Weighting Philosophy** - -**REFL's β=0.9 Choice:** -- **Pros:** - - Prioritizes data diversity over freshness - - Rare clients with unique data get full weight even if stale - - Works well for highly non-IID scenarios - - Fairness in data representation - -- **Cons:** - - Stale gradients may mislead optimization if model evolves fast - - Less effective when data distributions are similar - - Assumes loss correlates with data value (may not hold) - -**Felix's Staleness-First Approach:** -- **Pros:** - - Fresh updates prioritized = better gradient quality - - Naturally handles fast-evolving models - - Simpler: staleness is objective measure - - Proven effective in async settings - -- **Cons:** - - May underweight valuable rare clients - - Requires many available clients for freshness - - Less explicit data diversity guarantees - -#### **Expected Performance in Async Settings** - -**Scenario 1: High Availability, IID Data** -- **Winner: Felix** -- Reason: Many fresh updates available, staleness matters more than data diversity - -**Scenario 2: Low Availability, Highly Non-IID** -- **Winner: REFL** (if adapted to async) -- Reason: Data diversity critical, rare clients valuable - -**Scenario 3: Fast Model Evolution** -- **Winner: Felix** -- Reason: Stale gradients harmful, need freshness - -**Scenario 4: Slow Convergence, Intermittent Clients** -- **Winner: REFL** (if adapted to async) -- Reason: Worth waiting for diverse data, gradients remain relevant - -**Overall for Async:** Felix likely performs better in most practical async scenarios because: -1. Designed for async from ground up -2. Natural staleness handling without rounds -3. Works without oracular knowledge -4. FedBuff's dynamic buffering is well-suited to async arrival patterns - -**But:** REFL's β=0.9 weighting could be **incorporated into Felix** for scenarios where data diversity > freshness. - -### 4.2 Parity Status: **CONFIRMED** ✅ - -The Flame implementation of REFL demonstrates **functional parity** with the original REFL system across all three core pillars: - -1. **Availability-Aware Selection**: ✅ Confirmed via logs - - Trace loading working - - Priority computation active - - Selection logic matches original - -2. **Deadline-Based Aggregation**: ✅ Confirmed via logs - - Deadline filtering operational - - Fast/slow split working - - Clean aggregation completion - -3. **Stale Update Management**: ⚠️ Implemented but not exercised - - Code present and correct - - Not triggered in test (no stragglers) - - Need stress test to fully verify - -### 4.3 Architectural Improvements in Flame - -The Flame implementation makes several **design improvements** over original REFL: - -1. **Modular Components**: - - Original: Monolithic `Aggregator` class (1400+ lines) - - Flame: Clean separation into `REFLOortSelector` + `REFLFedAvg` + `REFLAvailabilityTracker` - -2. **Pluggable Architecture**: - - Can toggle REFL features independently: - - Selector only: availability-aware selection - - Optimizer only: deadline + staleness - - Both: full REFL - -3. **Configuration Management**: - - Original: Command-line args only - - Flame: Structured JSON configs with validation - -4. **Format Flexibility**: - - Original: Pickle traces only - - Flame: Pickle + YAML traces - -### 4.4 Remaining Verification Tasks - -To achieve **100% confidence**, we need to verify these scenarios: - -#### **Priority 1: Stale Update Handling (High Importance)** - -**Why:** Core REFL contribution, not yet exercised in logs - -**How to Verify:** -1. Reduce deadline to 5 seconds - ```json - "deadline": 5 - ``` -2. Run 10+ rounds -3. Check logs for: - ``` - Deadline filtering: X fast, Y slow (deadline=5.00s) - Stale updates: Y applicable, Z still cached - Aggregated N trainers (X fast + Y stale) - ``` - -#### **Priority 2: Priority Selection (Medium Importance)** - -**Why:** Should see different behavior with varied availability - -**How to Verify:** -1. Use traces with high/low availability clients -2. Set `avail_priority=2` (strict mode) -3. Check logs for: - ``` - Priority split: X priority, Y remaining, Z blacklisted - ``` - -#### **Priority 3: Blacklisting (Low Importance)** - -**Why:** Long-term fairness mechanism - -**How to Verify:** -1. Run 50-100 rounds -2. Or lower threshold: `blacklist_rounds: 5` -3. Check logs for blacklist entries - -### 4.5 Recommended Logging Enhancements - -To improve observability, add these logs: - -```python -# In flame/selector/refl_oort.py - -# 1. Availability usage -def _compute_priorities(self, ...): - for trainer_id in trainers: - priority = self.avail_tracker.get_priority(trainer_id, ...) - logger.debug(f"Trainer {trainer_id}: priority={priority}, available={...}") - -# 2. Pacer adjustments -def _update_pacer(self, ...): - logger.info( - f"Pacer: round={round_num}, threshold={old} → {new}, " - f"utility_trend={'up' if improving else 'down'}" - ) - -# 3. UCB score details -def select(self, ...): - logger.debug(f"UCB scores: {ucb_scores}") - logger.debug(f"Final priorities: {final_priorities}") -``` - -```python -# In flame/optimizer/reflfedavg.py - -# 1. Stale weight computation details -def _compute_stale_weight(self, ...): - logger.debug( - f"Stale weight for {trainer_id}: staleness={s}, " - f"stat_util={stat}, sys_util={sys}, final_weight={weight}" - ) - -# 2. Cache management -def _update_stale_cache(self, ...): - logger.info(f"Stale cache: added={len(added)}, aged={len(aged)}, removed={len(removed)}") -``` - -### 4.6 Integration Quality: **Production-Ready** ✅ - -**Status:** The Flame REFL implementation is **production-ready** with excellent engineering quality: - -✅ **Functional Parity**: All three REFL pillars correctly implemented -✅ **Clean Architecture**: Modular, testable, maintainable -✅ **Configuration**: Flexible JSON-based config -✅ **Logging**: Comprehensive observability -✅ **Verification**: Confirmed through multiple log analyses - -**Deployment Confidence:** High - ready for experiments and production use - ---- - -## Part 5: Key Takeaways & Strategic Recommendations - -### 5.1 Critical Finding: REFL/Oort's Exploitation-Heavy Design - -**The Most Important Discovery from This Analysis:** - -REFL and Oort fundamentally differ from Felix in client selection philosophy. This is **not** a bug or implementation issue - it's a **core design choice** with significant implications: - -**REFL/Oort Philosophy:** -- Select small subset per round (e.g., 30 of 300 trainers) -- 90% exploitation: repeatedly select top-utility trainers -- 10% exploration: occasionally try new trainers -- Result: Only 25% of trainers used over 100+ rounds - -**Felix Philosophy:** -- Opportunistic selection: use all available trainers -- Parallel lightweight exploration -- No artificial selection limit -- Result: 70-100% of trainers used over equivalent time - -### 5.2 Why This Matters for Heterogeneous FL - -**In heterogeneous scenarios with non-IID data, Felix's approach has fundamental advantages:** - -1. **Data Distribution Coverage:** - - Missing 75% of trainers = missing 75% of data distributions - - Long-tail classes and rare patterns not captured - - Model doesn't generalize to full trainer population - -2. **Robustness to Dynamics:** - - If top 30 trainers become unavailable → REFL performance degrades - - Felix adapts seamlessly to whoever is available - - No dependence on specific "elite" trainers - -3. **Fairness & Utilization:** - - Deployed 300 trainers but only use 74 = wasted resources - - 226 trainers never contribute = no benefit from model - - Ethical concerns in participatory FL settings - -4. **Statistical Representation:** - - Central Limit Theorem: variance ∝ 1/n - - Larger effective sample (200 vs 74) = better convergence guarantees - - More reliable estimates of population gradient - -### 5.3 When to Use Which System - -#### **Use REFL/Oort When:** - -✅ **Cross-Silo FL** (organizations, not devices) -✅ **Homogeneous clients** (datacenter nodes, cloud servers) -✅ **IID or mildly non-IID data** -✅ **Fast convergence is primary goal** -✅ **High client reliability** (low churn, stable availability) -✅ **Research benchmarks** (controlled experiments) - -**Example:** Hospital collaboration with 50 sites, stable infrastructure, speed critical - -#### **Use Felix When:** - -✅ **Cross-Device FL** (mobile phones, IoT devices) -✅ **Heterogeneous devices** (different compute/network capabilities) -✅ **Highly non-IID data** (personalized user data) -✅ **Fairness requirements** (all participants should benefit) -✅ **Dynamic availability** (unpredictable user behavior) -✅ **Production systems** (robustness over speed) -✅ **Long-tail data distributions** (rare classes matter) - -**Example:** Mobile keyboard prediction with millions of users, diverse languages/typing patterns - -### 5.4 Potential Hybrid Approaches - -**Best of Both Worlds:** - -1. **Adaptive Exploration Factor:** - ```python - # Start exploitative for fast initial progress - exploration_factor = 0.1 # 90% exploitation - - # Increase exploration as heterogeneity detected - if detect_high_heterogeneity(): - exploration_factor = 0.5 # 50/50 split - - # Pure exploration for rare data capture - if rare_classes_missing(): - exploration_factor = 0.9 # 90% exploration - ``` - -2. **Felix + REFL Staleness Weighting:** - - Use Felix's parallel opportunistic selection - - Apply REFL's β=0.9 weighting for rare clients - - Combine broad coverage with data diversity priority - -3. **Hierarchical Selection:** - - Tier 1: Always-available high-utility trainers (REFL/Oort) - - Tier 2: Opportunistic exploration of others (Felix) - - Ensures base performance while capturing diversity - -### 5.5 Final Recommendations - -#### **For Flame Development:** - -1. **Add Exploration Factor Config:** - ```json - "selector": { - "sort": "refl_oort", - "kwargs": { - "exploration_factor": 0.5, // Make tunable - "adaptive_exploration": true // Auto-adjust based on heterogeneity - } - } - ``` - -2. **Implement Diversity Metrics:** - - Track which trainers are used/unused - - Log distribution coverage statistics - - Alert when large trainer subsets never selected - -3. **Hybrid Selector Mode:** - ```python - class HybridSelector: - def select(self, ...): - # Core set from Oort - core = oort_select(k=20) - - # Exploration set from availability - explore = opportunistic_select(k=10) - - return core + explore - ``` - -#### **For Felix vs REFL Experiments:** - -1. **Report Trainer Utilization:** - - Don't just compare final accuracy - - Report: % trainers used, selection frequency distribution - - Show generalization to held-out trainers - -2. **Test Heterogeneity Scenarios:** - - Vary device speeds widely (10x - 100x range) - - Use highly non-IID data (alpha=0.01) - - Include rare classes (1% of trainers have critical data) - -3. **Measure Fairness:** - - Per-trainer accuracy improvement - - Participation equity (Gini coefficient) - - Representation of minority groups - -#### **For Research Community:** - -**Key Message:** The exploitation vs exploration trade-off in federated client selection deserves more attention. REFL/Oort's 90/10 split optimizes for convergence speed but sacrifices robustness, fairness, and data coverage. For real-world heterogeneous FL, more balanced or exploration-heavy strategies (like Felix) may be preferable despite slower initial convergence. - ---- - -### 5.6 Conclusion - -The Flame REFL implementation achieves **full functional parity** with the original REFL system. However, this analysis reveals that **REFL's core design philosophy differs fundamentally from Felix** in ways that significantly impact performance in heterogeneous scenarios. - -**REFL is correctly implemented** - the concentration on 25% of trainers is **intentional**, not a bug. But for realistic heterogeneous federated learning with non-IID data and dynamic device availability, **Felix's lightweight parallel exploration approach is more robust and fair**. - -The future of production FL systems likely lies in **hybrid approaches** that combine: -- Felix's broad exploration and opportunistic selection -- REFL's predictive availability awareness -- Adaptive strategies that balance speed, fairness, and robustness - -**Bottom Line:** Know your scenario, choose your strategy accordingly, and consider hybrid approaches for production deployments. - ---- - -*Document Version: 2.0* -*Last Updated: February 11, 2026* -*Analysis Based On: Flame REFL logs, original REFL source code, Felix architecture* - -**Strengths:** -- ✅ Clean modular design -- ✅ Proper error handling -- ✅ Comprehensive configuration -- ✅ Backward compatible (optional REFL features) -- ✅ No changes to core aggregator/trainer code - -**Minor Gaps:** -- ⚠️ Limited logging for debugging stale updates -- ⚠️ Limited logging for priority/availability decisions -- ⚠️ No unit tests added yet (integration tests work) - -**Overall:** The implementation is **correct and production-ready**, with minor logging enhancements recommended for easier debugging. - ---- - -## Summary - -### Part 1 Takeaways:** -- REFL is a synchronous FL system with three key innovations -- Uses oracular traces for availability-aware selection -- Applies deadline filtering to handle stragglers without blocking -- Caches and reweights stale updates for resource efficiency - -### Part 2 Takeaways:** -- ✅ Flame implementation shows functional parity with original REFL -- ✅ All core mechanisms present and operational -- ⚠️ Some features not exercised in test logs (stale updates, priority selection) -- ✅ Architectural improvements make Flame version more maintainable - -### Next Steps:** -1. Run stress test with tight deadline to trigger stale updates -2. Add recommended logging for better observability -3. Run 50+ round experiment to verify blacklisting -4. Optional: Add unit tests for stale weighting strategies - -**Conclusion:** The Flame REFL integration is **functionally correct** and demonstrates **parity with the original implementation**. The modular architecture enables clean ablation studies and easier maintenance than the original monolithic design. diff --git a/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/aggregator_config.json b/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/aggregator_config.json deleted file mode 100644 index 45b1b1b36..000000000 --- a/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/aggregator_config.json +++ /dev/null @@ -1,95 +0,0 @@ -{ - "taskid": "49d06b7526964db86cf37c70e8e0cdb6bd7aa742", - "backend": "mqtt", - "brokers": [ - { - "host": "localhost", - "sort": "mqtt" - }, - { - "host": "localhost:10104", - "sort": "p2p" - } - ], - "groupAssociation": { - "param-channel": "default" - }, - "channels": [ - { - "description": "Model update is sent from trainer to aggregator and vice-versa", - "groupBy": { - "type": "tag", - "value": [ - "default" - ] - }, - "name": "param-channel", - "pair": [ - "trainer", - "aggregator" - ], - "funcTags": { - "aggregator": [ - "distribute", - "aggregate" - ], - "trainer": [ - "fetch", - "upload" - ] - } - } - ], - "dataset": "https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz", - "dependencies": [ - "numpy >= 1.2.0" - ], - "hyperparameters": { - "batchSize": 32, - "learningRate": 0.001, - "rounds": 200, - "aggGoal": 5, - "trackTrainerAvail": { - "enabled": "False", - "type": "NA" - } - }, - "baseModel": { - "name": "", - "version": 2 - }, - "job": { - "id": "felix_n10_alpha100_syn20_smoke", - "name": "cifar-10" - }, - "registry": { - "sort": "dummy", - "uri": "" - }, - "selector": { - "sort": "async_oort", - "kwargs": { - "evalGoalFactor": 1.0, - "selectType": "default", - "roundNudgeType": "last_train", - "c": 8, - "aggGoal": 5 - } - }, - "optimizer": { - "sort": "fedbuff", - "kwargs": { - "use_oort_lr": "True", - "dataset_name": "cifar-10", - "agg_rate_conf": { - "type": "new", - "scale": 0.4, - "a_exp": 0.25, - "b_exp": 0.1 - } - } - }, - "maxRunTime": 600, - "realm": "default", - "role": "aggregator" -} \ No newline at end of file diff --git a/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/execution_config.yaml b/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/execution_config.yaml deleted file mode 100644 index e293d40ee..000000000 --- a/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/execution_config.yaml +++ /dev/null @@ -1,57 +0,0 @@ -execution_timestamp: '2026-05-22T09:18:03.670370' -git_info: - commit: b703d021eb3bbdbe9c762bd8536e0b1e9b0f05c5 - branch: dg/feddance_impl - dirty: true -metadata_refs: - trainer_registry: metadata/trainer_registry.yaml - dataset_split_key: cifar10_alpha100.0_n10 - availability_trace_key: syn_20 -aggregator: - config_file: aggregator_config.json - selector: async_oort - tracking_mode: client_notify - agg_goal: 5 -experiment: - name: felix_n10_alpha100_syn20_smoke - description: Felix smoke (10 trainers, alpha=100, syn_20). Validates baseline merge - + launch pipeline. - trainer: - num_trainers: 10 - start_id: 1 - trainer_id_range: - - 1 - - 10 - dataset: - name: cifar10 - alpha: 100.0 - split_key: cifar10_alpha100.0_n10 - availability: - mode: syn_20 - trace_key: syn_20 - battery_threshold: 50 - speedup_factor: 1.0 - execution: - num_gpus: 2 - sleep_between_spawns: 1.0 - aggregator_warmup_time: 15 -spawn_commands: - aggregator: - - /coc/scratch/dgarg/miniconda3/envs/dg_flame/bin/python - - /home/dgarg39/flame/lib/python/examples/async_cifar10/aggregator/pytorch/main_asyncfl_agg.py - - --config-json - - - trainers: - - python3 - - -m - - flame.launch.spawn_trainer_cli - - --alpha - - '100.0' - - --availability - - syn_20 - - --num-trainers - - '10' - - --start-id - - '1' - - --num-gpus - - '2' diff --git a/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/snapshot.yaml b/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/snapshot.yaml deleted file mode 100644 index cb9c3702a..000000000 --- a/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/snapshot.yaml +++ /dev/null @@ -1,63 +0,0 @@ -snapshot_version: '1.0' -timestamp: '2026-05-22T09:18:03.711924' -hostname: jayne.cc.gatech.edu -experiment: - name: felix_n10_alpha100_syn20_smoke - description: Felix smoke (10 trainers, alpha=100, syn_20). Validates baseline merge - + launch pipeline. - trainer: - num_trainers: 10 - start_id: 1 - dataset: - name: cifar10 - dirichlet_alpha: 100.0 - availability: - mode: syn_20 - battery_threshold: 50 - speedup_factor: 1.0 - aggregator: - config_template: ../_metadata/aggregator_base.json - selector: async_oort - tracking_mode: client_notify - agg_goal: 5 - execution: - num_gpus: 2 - sleep_between_spawns: 1.0 - aggregator_warmup_time: 15 -spawn_commands: - aggregator: - - /coc/scratch/dgarg/miniconda3/envs/dg_flame/bin/python - - /home/dgarg39/flame/lib/python/examples/async_cifar10/aggregator/pytorch/main_asyncfl_agg.py - - --config-json - - - trainers: - - python3 - - -m - - flame.launch.spawn_trainer_cli - - --alpha - - '100.0' - - --availability - - syn_20 - - --num-trainers - - '10' - - --start-id - - '1' - - --num-gpus - - '2' -git_info: - commit: b703d021eb3bbdbe9c762bd8536e0b1e9b0f05c5 - branch: dg/feddance_impl - clean: false - uncommitted_changes: true -metadata_location: /home/dgarg39/flame/lib/python/examples/async_cifar10/metadata -metadata_checksums: - trainer_registry: cb3a3cfb73edeb6dd423f7583f17ac125e26b14acd09886a40dea40dfa7f9922 - dataset_splits: - cifar10_alpha0.1_n300.yaml: dcbc7ad04b53add63287e2855101806a62018eeb7452556d604d59e91adef8b9 - cifar10_alpha1.0_n300.yaml: 8ba4d3bde49ab0c7cee8a6c3ffaf79f3ab9e2f357ffbcd937a40be9bffeb8be8 - cifar10_alpha10.0_n300.yaml: eb019d6692ad566e8a93c6e709660520e0a48a9d52359ba788fa1dd54b0dcc77 - cifar10_alpha100.0_n300.yaml: 13cbbbad0e9600bfc4847b216e70030bd42e46ec38c58b1a1515ec454cf9bb5c - availability_traces: - mobiperf_traces.yaml: f1e2da681116f0cdd6b08aa327f2b485990b9aba93f4a25949c5db7bf0355fa2 - synthetic_traces.yaml: b33d0220be78cf1a389ccc912981c15a92e52cb56860eafdfc5bd57d6a5727ca -aggregator_config: /home/dgarg39/flame/lib/python/examples/async_cifar10/experiments/run_20260522_091748_felix_n10_alpha100_syn20_smoke/aggregator_config.json diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/fedavg_n10_alpha100_smoke.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/fedavg_n10_alpha100_smoke.yaml index 982e6dc3e..ca2549f01 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/fedavg_n10_alpha100_smoke.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/fedavg_n10_alpha100_smoke.yaml @@ -19,7 +19,6 @@ experiments: availability: mode: syn_0 enable_training_delays: true - speedup_factor: 1.0 hyperparameters: batchSize: 32 learningRate: 0.001 diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/fedbuff_n10_alpha100_smoke.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/fedbuff_n10_alpha100_smoke.yaml index ab469e216..aec86d327 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/fedbuff_n10_alpha100_smoke.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/fedbuff_n10_alpha100_smoke.yaml @@ -19,7 +19,6 @@ experiments: availability: mode: syn_0 enable_training_delays: true - speedup_factor: 1.0 hyperparameters: batchSize: 32 learningRate: 0.001 diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_smoke.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_smoke.yaml index 1744d01a1..cfd26dac2 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_smoke.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_smoke.yaml @@ -20,7 +20,6 @@ experiments: availability: mode: syn_20 enable_training_delays: true - speedup_factor: 1.0 hyperparameters: batchSize: 32 learningRate: 0.001 diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_telemetry_smoke.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_telemetry_smoke.yaml new file mode 100644 index 000000000..758f9749f --- /dev/null +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_telemetry_smoke.yaml @@ -0,0 +1,78 @@ +# Felix TELEMETRY smoke test: 10 trainers, homogeneous (alpha=100), syn_20. +# Same as felix_n10_alpha100_syn20_smoke.yaml but additionally enables: +# - data_streaming : reveal samples over sim time (so visible < total early) +# - util_counterfactual : emit streamed-prefix vs full-pool statistical utility +# so the post-run analyzer produces the streamed-vs-full disparity plots in +# addition to the standard selector/aggregator/trainer telemetry plots. +# +# Telemetry itself is auto-enabled by the launcher for ANY run; these two +# trainer settings only drive the extra disparity plot. +# +# Run: +# python -m flame.launch.run_experiment \ +# lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn20_telemetry_smoke.yaml + +experiments: + - name: felix_n10_alpha100_syn20_telemetry_smoke + description: "Felix telemetry smoke (10 trainers, alpha=100, syn_20) with data streaming + utility counterfactual." + + baseline: felix + + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: cifar10 + dirichlet_alpha: 100.0 + availability: + mode: syn_20 + enable_training_delays: true + # simulated mode: no sleeps; the aggregator orders updates by a virtual + # clock fed by trainer-reported completion times. Availability + streaming + # are evaluated in sim-seconds against that clock (no speedup_factor). + time_mode: simulated + hyperparameters: + batchSize: 32 + learningRate: 0.001 + config_overrides: + hyperparameters: + client_notify: + trace: syn_20 + # Reveal data gradually over the first 600 SIM-seconds so the + # streamed-vs-full disparity is observable across the run. (Most runs + # use the 24h default = 86400 sim-s; this is the smoke-scaled value.) + data_streaming: + enabled: "True" + full_data_available_after_s: 600 + # Emit streamed-vs-full statistical-utility disparity telemetry. + # Default-style: subsample (256) every round; gated on here. + util_counterfactual: + enabled: "True" + every_n_rounds: 1 + sample_size: 256 + + aggregator: + config_template: ../_metadata/aggregator_base.json + selector: async_oort + tracking_mode: client_notify + agg_goal: 5 + log_to_wandb: false + config_overrides: + job: + id: felix_n10_alpha100_syn20_telemetry_smoke + hyperparameters: + batchSize: 32 + learningRate: 0.001 + rounds: 200 + aggGoal: 5 + selector: + kwargs: + c: 8 + aggGoal: 5 + + execution: + num_gpus: 2 + sleep_between_spawns: 1.0 + aggregator_warmup_time: 15 + monitoring: + enabled: false diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_SIMULATED.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_SIMULATED.yaml new file mode 100644 index 000000000..aea61fc7d --- /dev/null +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_SIMULATED.yaml @@ -0,0 +1,47 @@ +experiments: + - name: felix_n10_parity_SIMULATED + description: "Parity test: time_mode=simulated (no sleeps; virtual-clock ordering)." + baseline: felix + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: cifar10 + dirichlet_alpha: 100.0 + availability: + mode: syn_0 + enable_training_delays: true + time_mode: simulated + hyperparameters: + batchSize: 32 + learningRate: 0.001 + config_overrides: + hyperparameters: + client_notify: + trace: syn_0 + data_streaming: + enabled: "False" + aggregator: + config_template: ../_metadata/aggregator_base.json + selector: async_oort + tracking_mode: client_notify + agg_goal: 5 + log_to_wandb: false + config_overrides: + job: + id: felix_n10_parity_simulated + hyperparameters: + batchSize: 32 + learningRate: 0.001 + rounds: 100 + aggGoal: 5 + selector: + kwargs: + c: 8 + aggGoal: 5 + execution: + num_gpus: 2 + sleep_between_spawns: 1.0 + aggregator_warmup_time: 15 + monitoring: + enabled: false diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_real_vs_sim.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_real_vs_sim.yaml new file mode 100644 index 000000000..16a74132f --- /dev/null +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_real_vs_sim.yaml @@ -0,0 +1,113 @@ +# Task 2 verification: real vs simulated time_mode parity (felix, async_oort). +# Two experiments run SEQUENTIALLY (real first, then simulated). They are +# identical except time_mode, so you can compare per-round selection / staleness +# / learning and confirm: simulated reproduces real's behavior much faster. +# +# Scope notes: +# - availability = syn_0 (always available) on purpose: isolates the time-mode +# change and avoids the known sim-mode stall on all-UN_AVL windows (the +# virtual clock only advances on commits). Availability dynamics are a +# separate follow-up. +# - data streaming OFF: both train on full data so learning curves are directly +# comparable. +# - exact per-round parity is NOT expected (Oort selection uses unseeded RNG + +# real GPU jitter). Look for STATISTICAL parity + the speedup (see below). +# +# Run (sequential, ~20-30 min real + a few min simulated): +# python -m flame.launch.run_experiment \ +# lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_parity_real_vs_sim.yaml + +experiments: + - name: felix_n10_parity_REAL + description: "Parity baseline: time_mode=real (sleeps modeled delays at true pace)." + baseline: felix + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: cifar10 + dirichlet_alpha: 100.0 + availability: + mode: syn_0 + enable_training_delays: true + time_mode: real + hyperparameters: + batchSize: 32 + learningRate: 0.001 + config_overrides: + hyperparameters: + client_notify: + trace: syn_0 + data_streaming: + enabled: "False" + aggregator: + config_template: ../_metadata/aggregator_base.json + selector: async_oort + tracking_mode: client_notify + agg_goal: 5 + log_to_wandb: false + config_overrides: + job: + id: felix_n10_parity_real + hyperparameters: + batchSize: 32 + learningRate: 0.001 + rounds: 100 + aggGoal: 5 + selector: + kwargs: + c: 8 + aggGoal: 5 + execution: + num_gpus: 2 + sleep_between_spawns: 1.0 + aggregator_warmup_time: 15 + monitoring: + enabled: false + + - name: felix_n10_parity_SIMULATED + description: "Parity test: time_mode=simulated (no sleeps; virtual-clock ordering)." + baseline: felix + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: cifar10 + dirichlet_alpha: 100.0 + availability: + mode: syn_0 + enable_training_delays: true + time_mode: simulated + hyperparameters: + batchSize: 32 + learningRate: 0.001 + config_overrides: + hyperparameters: + client_notify: + trace: syn_0 + data_streaming: + enabled: "False" + aggregator: + config_template: ../_metadata/aggregator_base.json + selector: async_oort + tracking_mode: client_notify + agg_goal: 5 + log_to_wandb: false + config_overrides: + job: + id: felix_n10_parity_simulated + hyperparameters: + batchSize: 32 + learningRate: 0.001 + rounds: 100 + aggGoal: 5 + selector: + kwargs: + c: 8 + aggGoal: 5 + execution: + num_gpus: 2 + sleep_between_spawns: 1.0 + aggregator_warmup_time: 15 + monitoring: + enabled: false diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n300_alpha100_syn20.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n300_alpha100_syn20.yaml index a52156e21..615cd97e7 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n300_alpha100_syn20.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/felix_n300_alpha100_syn20.yaml @@ -22,7 +22,6 @@ experiments: availability: mode: syn_20 enable_training_delays: true - speedup_factor: 1.0 hyperparameters: batchSize: 32 learningRate: 0.001 diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/oort_n10_alpha100_syn0_smoke.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/oort_n10_alpha100_syn0_smoke.yaml index 47bde705d..31b78a669 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/oort_n10_alpha100_syn0_smoke.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/oort_n10_alpha100_syn0_smoke.yaml @@ -19,7 +19,6 @@ experiments: availability: mode: syn_0 enable_training_delays: true - speedup_factor: 1.0 hyperparameters: batchSize: 32 learningRate: 0.001 diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/oort_oracular_10trainer_demo.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/oort_oracular_10trainer_demo.yaml index e8ebafe2f..59ffd2b58 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/oort_oracular_10trainer_demo.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/oort_oracular_10trainer_demo.yaml @@ -27,7 +27,6 @@ experiments: availability: mode: syn_0 # Always available (no straggler failures) battery_threshold: 50 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/oort_oracular_10trainer_demo.json @@ -53,7 +52,6 @@ experiments: availability: mode: syn_20 # 20% unavailability (simulates stragglers) battery_threshold: 50 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/oort_oracular_10trainer_demo_syn20.json diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/refl_10trainer_demo.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/refl_10trainer_demo.yaml index 6a4fe02d1..0ce513bd0 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/refl_10trainer_demo.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/refl_10trainer_demo.yaml @@ -27,7 +27,6 @@ experiments: availability: mode: syn_0 # Always available (no straggler failures) battery_threshold: 50 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_10trainer_demo.json @@ -53,7 +52,6 @@ experiments: availability: mode: syn_20 # 20% unavailability (simulates stragglers) battery_threshold: 50 - speedup_factor: 1.0 aggregator: config_template: expt_scripts_2026/configs/refl_10trainer_demo_syn20.json diff --git a/lib/python/examples/async_cifar10/expt_scripts_2026/refl_n10_alpha100_syn0_smoke.yaml b/lib/python/examples/async_cifar10/expt_scripts_2026/refl_n10_alpha100_syn0_smoke.yaml index 738269594..8e5fdea08 100644 --- a/lib/python/examples/async_cifar10/expt_scripts_2026/refl_n10_alpha100_syn0_smoke.yaml +++ b/lib/python/examples/async_cifar10/expt_scripts_2026/refl_n10_alpha100_syn0_smoke.yaml @@ -19,7 +19,6 @@ experiments: availability: mode: syn_0 enable_training_delays: true - speedup_factor: 1.0 hyperparameters: batchSize: 32 learningRate: 0.001 diff --git a/lib/python/examples/async_cifar10/scripts/analyze_send_recv_lag.py b/lib/python/examples/async_cifar10/scripts/analyze_send_recv_lag.py new file mode 100644 index 000000000..d6180d450 --- /dev/null +++ b/lib/python/examples/async_cifar10/scripts/analyze_send_recv_lag.py @@ -0,0 +1,84 @@ +""" +Post-processing script: parse SEND_RECV_LAG entries from an aggregator log and +report per-trainer lag statistics. Raises a warning if any trainer's median lag +exceeds a configurable threshold. + +Usage: + python analyze_send_recv_lag.py [--warn-threshold-s 5.0] +""" + +import argparse +import re +import sys +from collections import defaultdict +import statistics + + +LAG_RE = re.compile( + r"\[SEND_RECV_LAG\] end=(\S+) version=(\d+) wall_lag_s=([0-9.]+)" +) + + +def parse_lags(log_path: str) -> dict[str, list[float]]: + lags: dict[str, list[float]] = defaultdict(list) + with open(log_path) as f: + for line in f: + m = LAG_RE.search(line) + if m: + end_id, _version, lag_s = m.group(1), m.group(2), float(m.group(3)) + lags[end_id].append(lag_s) + return dict(lags) + + +def report(lags: dict[str, list[float]], warn_threshold_s: float) -> None: + if not lags: + print("No SEND_RECV_LAG entries found — was the run built with instrumentation?") + return + + all_lags = [v for vals in lags.values() for v in vals] + print(f"\n{'='*60}") + print(f"Global n={len(all_lags):4d} " + f"min={min(all_lags):.2f}s " + f"median={statistics.median(all_lags):.2f}s " + f"p95={sorted(all_lags)[int(len(all_lags)*0.95)]:.2f}s " + f"max={max(all_lags):.2f}s") + print(f"{'='*60}") + + warnings = [] + for end_id, vals in sorted(lags.items()): + med = statistics.median(vals) + p95 = sorted(vals)[int(len(vals) * 0.95)] if len(vals) >= 20 else max(vals) + short_id = end_id[-8:] + print(f" trainer ...{short_id} n={len(vals):4d} " + f"min={min(vals):.2f}s median={med:.2f}s " + f"p95={p95:.2f}s max={max(vals):.2f}s") + if med > warn_threshold_s: + warnings.append((short_id, med)) + + if warnings: + print(f"\n{'!'*60}") + print("WARNING: the following trainers have median lag above " + f"{warn_threshold_s}s — possible MQTT broker backlog or slow " + "model delivery:") + for short_id, med in warnings: + print(f" trainer ...{short_id} median={med:.2f}s") + print("Consider increasing the inter-send sleep (0.2s → 0.5s) or " + "checking broker load.") + print(f"{'!'*60}") + sys.exit(1) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("log", help="Path to aggregator log file") + parser.add_argument( + "--warn-threshold-s", type=float, default=5.0, + help="Warn if any trainer's median lag exceeds this (seconds)" + ) + args = parser.parse_args() + lags = parse_lags(args.log) + report(lags, args.warn_threshold_s) + + +if __name__ == "__main__": + main() diff --git a/lib/python/examples/async_cifar10/scripts/compare_parity.py b/lib/python/examples/async_cifar10/scripts/compare_parity.py new file mode 100644 index 000000000..b3f829426 --- /dev/null +++ b/lib/python/examples/async_cifar10/scripts/compare_parity.py @@ -0,0 +1,517 @@ +""" +Parity comparator: real run vs simulated run. + +Checks: + 1. Selection parity — per round: same trainers selected? same order? + 2. Statistical utility — per-trainer utility distributions match? + 3. Aggregation sequence — same contributing-trainer order per round? + 4. Staleness — distributions match? + 5. Participation counts — each trainer used similar # of times? + 6. Accuracy / loss — convergence curves from agg_eval events + +Usage: + python compare_parity.py \\ + --real experiments/run_20260529_101200.../telemetry/aggregator_*.jsonl \\ + --sim experiments/run_20260529_152224.../telemetry/aggregator_*.jsonl \\ + [--rounds 100] [--strict] + +Exit 0 = all checks passed (or within tolerance) +Exit 1 = one or more checks failed +""" + +import argparse +import collections +import json +import math +import sys +from pathlib import Path +from typing import Optional + + +# ── helpers ────────────────────────────────────────────────────────────────── + +def short(end_id: str) -> str: + return end_id[-4:] if end_id else "None" + + +def load_agg_jsonl(path: str) -> dict: + """Parse an aggregator telemetry JSONL into typed lists.""" + selection_train: list[dict] = [] # task=train selection events + agg_rounds: list[dict] = [] # per-update aggregation events + agg_evals: list[dict] = [] # per-round evaluation events + + with open(path) as f: + for line in f: + e = json.loads(line.strip()) + ev = e.get("event") + if ev == "selection" and e.get("task") == "train": + selection_train.append(e) + elif ev == "agg_round": + agg_rounds.append(e) + elif ev == "agg_eval": + agg_evals.append(e) + + # Sort by round then ts + selection_train.sort(key=lambda x: (x["round"], x["ts"])) + agg_rounds.sort(key=lambda x: (x["round"], x.get("agg_goal_count", 0), x["ts"])) + agg_evals.sort(key=lambda x: x["round"]) + + return { + "selection_train": selection_train, + "agg_rounds": agg_rounds, + "agg_evals": agg_evals, + } + + +def ks_stat(a: list[float], b: list[float]) -> float: + """Two-sample Kolmogorov–Smirnov statistic (no scipy needed).""" + if not a or not b: + return float("nan") + combined = sorted(set(a + b)) + na, nb = len(a), len(b) + sa, sb = sorted(a), sorted(b) + ia = ib = 0 + d = 0.0 + for v in combined: + while ia < na and sa[ia] <= v: + ia += 1 + while ib < nb and sb[ib] <= v: + ib += 1 + d = max(d, abs(ia / na - ib / nb)) + return d + + +def mean_std(vals: list[float]) -> tuple[float, float]: + if not vals: + return float("nan"), float("nan") + m = sum(vals) / len(vals) + v = sum((x - m) ** 2 for x in vals) / len(vals) + return m, math.sqrt(v) + + +def jaccard(a: set, b: set) -> float: + if not a and not b: + return 1.0 + return len(a & b) / len(a | b) + + +# ── check functions ─────────────────────────────────────────────────────────── + +PASS = "PASS" +FAIL = "FAIL" +WARN = "WARN" + + +def check_selection_parity(real: dict, sim: dict, max_rounds: Optional[int]) -> dict: + """Check 1: per-round selection — same trainers chosen?""" + # group by outer round + def by_round(sel_events): + d = collections.defaultdict(set) + for e in sel_events: + r = e["round"] + d[r].update(e.get("chosen", [])) + return d + + r_sel = by_round(real["selection_train"]) + s_sel = by_round(sim["selection_train"]) + + rounds = sorted(set(r_sel) & set(s_sel)) + if max_rounds: + rounds = [r for r in rounds if r <= max_rounds] + + jaccards = [] + exact_match = 0 + per_round_detail = [] + for r in rounds: + rs, ss = r_sel[r], s_sel[r] + j = jaccard(rs, ss) + jaccards.append(j) + exact = rs == ss + if exact: + exact_match += 1 + per_round_detail.append((r, rs, ss, j, exact)) + + mean_j = sum(jaccards) / len(jaccards) if jaccards else float("nan") + exact_pct = exact_match / len(rounds) * 100 if rounds else 0.0 + + status = PASS if mean_j >= 0.7 else (WARN if mean_j >= 0.4 else FAIL) + + # Find rounds with lowest Jaccard + worst = sorted(per_round_detail, key=lambda x: x[3])[:5] + + return { + "status": status, + "rounds_compared": len(rounds), + "exact_match_pct": round(exact_pct, 1), + "mean_jaccard": round(mean_j, 3), + "worst_rounds": [ + {"round": r, "real": sorted(short(t) for t in rs), + "sim": sorted(short(t) for t in ss), "jaccard": round(j, 3)} + for r, rs, ss, j, _ in worst + ], + } + + +def check_utility_parity(real: dict, sim: dict) -> dict: + """Check 2: per-trainer stat_utility distributions match?""" + def per_trainer_utils(agg_rounds): + d = collections.defaultdict(list) + for e in agg_rounds: + for t, u in zip(e.get("contributing_trainers", []), e.get("stat_utility", [])): + if u is not None: + d[t].append(u) + return d + + r_utils = per_trainer_utils(real["agg_rounds"]) + s_utils = per_trainer_utils(sim["agg_rounds"]) + + all_trainers = sorted(set(r_utils) | set(s_utils)) + rows = [] + ks_stats = [] + mean_diffs = [] + for t in all_trainers: + ru = r_utils.get(t, []) + su = s_utils.get(t, []) + ks = ks_stat(ru, su) + rm, rs_ = mean_std(ru) + sm, ss_ = mean_std(su) + diff = abs(rm - sm) if not math.isnan(rm) and not math.isnan(sm) else float("nan") + rows.append({ + "trainer": short(t), + "real_n": len(ru), "sim_n": len(su), + "real_mean": round(rm, 2), "sim_mean": round(sm, 2), + "mean_diff": round(diff, 2), + "ks_stat": round(ks, 3) if not math.isnan(ks) else None, + }) + if not math.isnan(ks): + ks_stats.append(ks) + if not math.isnan(diff): + mean_diffs.append(diff) + + max_ks = max(ks_stats) if ks_stats else float("nan") + avg_mean_diff = sum(mean_diffs) / len(mean_diffs) if mean_diffs else float("nan") + + # KS < 0.2 = distributions match well; < 0.3 = acceptable + status = PASS if max_ks < 0.2 else (WARN if max_ks < 0.35 else FAIL) + + return { + "status": status, + "max_ks_stat": round(max_ks, 3) if not math.isnan(max_ks) else None, + "avg_mean_utility_diff": round(avg_mean_diff, 2) if not math.isnan(avg_mean_diff) else None, + "per_trainer": rows, + } + + +def check_aggregation_sequence(real: dict, sim: dict, max_rounds: Optional[int]) -> dict: + """Check 3: per-round aggregation sequence — same contributing trainers?""" + def by_round(agg_rounds): + d = collections.defaultdict(list) + for e in agg_rounds: + r = e["round"] + if max_rounds and r > max_rounds: + continue + t = e["contributing_trainers"][0] if e["contributing_trainers"] else None + d[r].append(t) + return d + + r_seq = by_round(real["agg_rounds"]) + s_seq = by_round(sim["agg_rounds"]) + + rounds = sorted(set(r_seq) & set(s_seq)) + exact_rounds = 0 + set_match_rounds = 0 + per_round_detail = [] + for r in rounds: + rs, ss = r_seq[r], s_seq[r] + exact = rs == ss + set_match = set(t for t in rs if t) == set(t for t in ss if t) + if exact: + exact_rounds += 1 + if set_match: + set_match_rounds += 1 + if not exact: + per_round_detail.append({ + "round": r, + "real_seq": [short(t) for t in rs], + "sim_seq": [short(t) for t in ss], + "set_match": set_match, + }) + + exact_pct = exact_rounds / len(rounds) * 100 if rounds else 0.0 + set_pct = set_match_rounds / len(rounds) * 100 if rounds else 0.0 + + status = PASS if set_pct >= 70 else (WARN if set_pct >= 40 else FAIL) + + return { + "status": status, + "rounds_compared": len(rounds), + "exact_sequence_match_pct": round(exact_pct, 1), + "set_match_pct": round(set_pct, 1), + "mismatched_rounds_sample": per_round_detail[:10], + } + + +def check_staleness(real: dict, sim: dict) -> dict: + """Check 4: staleness distributions match?""" + def get_staleness(agg_rounds): + vals = [] + for e in agg_rounds: + vals.extend(e.get("staleness", [])) + return vals + + rs = get_staleness(real["agg_rounds"]) + ss = get_staleness(sim["agg_rounds"]) + + r_counts = collections.Counter(rs) + s_counts = collections.Counter(ss) + all_k = sorted(set(r_counts) | set(s_counts)) + + ks = ks_stat([float(x) for x in rs], [float(x) for x in ss]) + r_mean, _ = mean_std([float(x) for x in rs]) + s_mean, _ = mean_std([float(x) for x in ss]) + + status = PASS if abs(r_mean - s_mean) < 1.0 else (WARN if abs(r_mean - s_mean) < 2.0 else FAIL) + + return { + "status": status, + "real_mean_staleness": round(r_mean, 3), + "sim_mean_staleness": round(s_mean, 3), + "mean_staleness_diff": round(abs(r_mean - s_mean), 3), + "ks_stat": round(ks, 3) if not math.isnan(ks) else None, + "real_distribution": {k: r_counts.get(k, 0) for k in all_k}, + "sim_distribution": {k: s_counts.get(k, 0) for k in all_k}, + } + + +def check_participation(real: dict, sim: dict) -> dict: + """Check 5: per-trainer participation counts match?""" + def counts(agg_rounds): + c = collections.Counter() + for e in agg_rounds: + for t in e.get("contributing_trainers", []): + c[t] += 1 + return c + + rc = counts(real["agg_rounds"]) + sc = counts(sim["agg_rounds"]) + trainers = sorted(set(rc) | set(sc)) + + rows = [] + diffs = [] + for t in trainers: + r, s = rc.get(t, 0), sc.get(t, 0) + d = abs(r - s) + diffs.append(d) + rows.append({"trainer": short(t), "real": r, "sim": s, "abs_diff": d}) + + rows.sort(key=lambda x: -x["abs_diff"]) + avg_diff = sum(diffs) / len(diffs) if diffs else 0.0 + max_diff = max(diffs) if diffs else 0 + + total_r = sum(rc.values()) + total_s = sum(sc.values()) + + status = PASS if avg_diff < 10 else (WARN if avg_diff < 25 else FAIL) + + return { + "status": status, + "total_updates_real": total_r, + "total_updates_sim": total_s, + "avg_participation_diff": round(avg_diff, 1), + "max_participation_diff": max_diff, + "per_trainer": rows, + } + + +def check_convergence(real: dict, sim: dict) -> dict: + """Check 6: accuracy/loss curves.""" + def curve(agg_evals): + return {e["round"]: {"acc": e.get("test-accuracy"), "loss": e.get("test-loss")} + for e in agg_evals} + + rc = curve(real["agg_evals"]) + sc = curve(real["agg_evals"]) # same source intentional for self-check + # use actual sim curve + sc = curve(sim["agg_evals"]) + + rounds = sorted(set(rc) & set(sc)) + if not rounds: + return {"status": WARN, "note": "no overlapping eval rounds"} + + acc_diffs = [] + loss_diffs = [] + rows = [] + for r in rounds: + ra, rl = rc[r].get("acc"), rc[r].get("loss") + sa, sl = sc[r].get("acc"), sc[r].get("loss") + if ra is not None and sa is not None: + acc_diffs.append(abs(ra - sa)) + if rl is not None and sl is not None: + loss_diffs.append(abs(rl - sl)) + rows.append({ + "round": r, + "real_acc": round(ra, 4) if ra is not None else None, + "sim_acc": round(sa, 4) if sa is not None else None, + "real_loss": round(rl, 4) if rl is not None else None, + "sim_loss": round(sl, 4) if sl is not None else None, + }) + + avg_acc_diff = sum(acc_diffs) / len(acc_diffs) if acc_diffs else float("nan") + avg_loss_diff = sum(loss_diffs) / len(loss_diffs) if loss_diffs else float("nan") + + # Within 5pp accuracy = PASS; within 10pp = WARN + status = PASS + if not math.isnan(avg_acc_diff): + status = PASS if avg_acc_diff < 0.05 else (WARN if avg_acc_diff < 0.10 else FAIL) + + return { + "status": status, + "eval_rounds_compared": len(rounds), + "avg_accuracy_diff": round(avg_acc_diff, 4) if not math.isnan(avg_acc_diff) else None, + "avg_loss_diff": round(avg_loss_diff, 4) if not math.isnan(avg_loss_diff) else None, + "curve": rows, + } + + +# ── reporting ───────────────────────────────────────────────────────────────── + +def _status_icon(s: str) -> str: + return {"PASS": "[OK]", "WARN": "[!!]", "FAIL": "[XX]"}.get(s, "[??]") + + +def print_report(results: dict, strict: bool) -> bool: + overall_ok = True + print(f"\n{'='*70}") + print(" PARITY REPORT: real vs simulated") + print(f"{'='*70}\n") + + checks = [ + ("1. Selection parity (trainer sets per round)", "selection"), + ("2. Statistical utility distributions", "utility"), + ("3. Aggregation sequence (per-round order)", "aggregation"), + ("4. Staleness distribution", "staleness"), + ("5. Trainer participation counts", "participation"), + ("6. Convergence (accuracy / loss)", "convergence"), + ] + + for title, key in checks: + r = results.get(key, {}) + status = r.get("status", "UNKNOWN") + icon = _status_icon(status) + print(f" {icon} [{status:4s}] {title}") + + if key == "selection": + print(f" rounds compared: {r['rounds_compared']}") + print(f" exact match: {r['exact_match_pct']}%") + print(f" mean Jaccard: {r['mean_jaccard']} (1.0 = identical sets)") + if r.get("worst_rounds"): + print(" worst rounds:") + for w in r["worst_rounds"][:3]: + print(f" round {w['round']:3d}: J={w['jaccard']:.3f} " + f"real={w['real']} sim={w['sim']}") + + elif key == "utility": + print(f" max KS stat: {r.get('max_ks_stat')} (0=identical, <0.2=good)") + print(f" avg mean-util diff: {r.get('avg_mean_utility_diff')}") + print(" per trainer:") + for row in r.get("per_trainer", []): + print(f" ...{row['trainer']}: real_mean={row['real_mean']} " + f"sim_mean={row['sim_mean']} diff={row['mean_diff']} " + f"KS={row['ks_stat']} (n_real={row['real_n']} n_sim={row['sim_n']})") + + elif key == "aggregation": + print(f" rounds compared: {r['rounds_compared']}") + print(f" exact-sequence match: {r['exact_sequence_match_pct']}%") + print(f" set match: {r['set_match_pct']}% (same trainers, any order)") + if r.get("mismatched_rounds_sample"): + print(" sample mismatches:") + for m in r["mismatched_rounds_sample"][:3]: + print(f" round {m['round']:3d}: real={m['real_seq']} " + f"sim={m['sim_seq']} set_match={m['set_match']}") + + elif key == "staleness": + print(f" real mean: {r.get('real_mean_staleness')} " + f"sim mean: {r.get('sim_mean_staleness')} " + f"diff: {r.get('mean_staleness_diff')}") + print(f" real dist: {r.get('real_distribution')}") + print(f" sim dist: {r.get('sim_distribution')}") + + elif key == "participation": + print(f" total updates real: {r.get('total_updates_real')} " + f"sim: {r.get('total_updates_sim')}") + print(f" avg participation diff per trainer: {r.get('avg_participation_diff')}") + print(f" max participation diff: {r.get('max_participation_diff')}") + print(" per trainer (sorted by diff):") + for row in r.get("per_trainer", []): + print(f" ...{row['trainer']}: real={row['real']:3d} sim={row['sim']:3d} " + f"diff={row['abs_diff']}") + + elif key == "convergence": + print(f" eval rounds compared: {r.get('eval_rounds_compared')}") + print(f" avg accuracy diff: {r.get('avg_accuracy_diff')}") + print(f" avg loss diff: {r.get('avg_loss_diff')}") + if r.get("curve"): + print(" curve (sample):") + for row in r["curve"][:5]: + print(f" round {row['round']:3d}: " + f"acc real={row['real_acc']} sim={row['sim_acc']} " + f"loss real={row['real_loss']} sim={row['sim_loss']}") + + if status == "FAIL" or (strict and status == "WARN"): + overall_ok = False + print() + + verdict = "ALL CHECKS PASSED" if overall_ok else "ONE OR MORE CHECKS FAILED" + print(f"{'='*70}") + print(f" {_status_icon('PASS' if overall_ok else 'FAIL')} {verdict}") + print(f"{'='*70}\n") + return overall_ok + + +# ── main ────────────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser(description="Compare real vs simulated run parity") + parser.add_argument("--real", required=True, help="Path to real aggregator JSONL") + parser.add_argument("--sim", required=True, help="Path to simulated aggregator JSONL") + parser.add_argument("--rounds", type=int, default=None, + help="Only compare up to this round number") + parser.add_argument("--strict", action="store_true", + help="Treat WARN as FAIL") + parser.add_argument("--json-out", default=None, + help="Write full results as JSON to this path") + args = parser.parse_args() + + print(f"Loading real: {args.real}") + real = load_agg_jsonl(args.real) + print(f"Loading sim: {args.sim}") + sim = load_agg_jsonl(args.sim) + + print(f" real: {len(real['agg_rounds'])} agg_round events, " + f"{len(real['selection_train'])} train-selection events, " + f"{len(real['agg_evals'])} eval events") + print(f" sim: {len(sim['agg_rounds'])} agg_round events, " + f"{len(sim['selection_train'])} train-selection events, " + f"{len(sim['agg_evals'])} eval events") + + results = { + "selection": check_selection_parity(real, sim, args.rounds), + "utility": check_utility_parity(real, sim), + "aggregation": check_aggregation_sequence(real, sim, args.rounds), + "staleness": check_staleness(real, sim), + "participation": check_participation(real, sim), + "convergence": check_convergence(real, sim), + } + + ok = print_report(results, args.strict) + + if args.json_out: + with open(args.json_out, "w") as f: + json.dump(results, f, indent=2) + print(f"Full results written to {args.json_out}") + + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/lib/python/examples/async_cifar10/trainer/pytorch/main.py b/lib/python/examples/async_cifar10/trainer/pytorch/main.py index d2cbbab3f..0b6e5d66e 100644 --- a/lib/python/examples/async_cifar10/trainer/pytorch/main.py +++ b/lib/python/examples/async_cifar10/trainer/pytorch/main.py @@ -24,6 +24,7 @@ import ast import calendar import gc +import hashlib import logging import os import sys @@ -38,6 +39,12 @@ import torchvision.transforms as transforms from flame.config import Config, TrainerAvailState from flame.mode.horizontal.trainer import Trainer +from flame import telemetry +from flame.telemetry.events import ( + build_avail_change, + build_trainer_round, + build_util_disparity, +) from torchvision.datasets import CIFAR10 from memory_profiler import MemoryProfiler @@ -73,7 +80,7 @@ def forward(self, x): class PyTorchCifar10Trainer(Trainer): """PyTorch CIFAR-10 Trainer.""" - def __init__(self, config: Config, battery_threshold, speedup_factor) -> None: + def __init__(self, config: Config, battery_threshold, time_mode="simulated") -> None: """Initialize a class instance.""" self.config = config self.dataset_size = 0 @@ -148,13 +155,24 @@ def __init__(self, config: Config, battery_threshold, speedup_factor) -> None: # Check if client will notify aggregator of its availability self.client_notify = self.config.hyperparameters.client_notify - # Check if client will emulate delays in training time - self.training_delay_enabled = self.config.hyperparameters.training_delay_enabled + # Check if client will emulate delays in training time. + # Normalize to bool: config.py types this as Optional[bool] but it can + # arrive as a Python bool (from JSON true) or the string "True"/"False" + # depending on the config path. Unify here so comparisons are consistent. + _tde = self.config.hyperparameters.training_delay_enabled + self.training_delay_enabled = ( + _tde if isinstance(_tde, bool) else str(_tde).strip().lower() == "true" + ) self.training_delay_s = float(self.config.hyperparameters.training_delay_s) - # Set speedup factor to accelerate all events and training/ - # eval durations - self.speedup_factor = speedup_factor + # Simulation time mode. "real": sleep modeled delays at true pace. + # "simulated": skip sleeps; report a modeled completion time so the + # aggregator orders updates by a virtual clock. No speedup_factor. + self.time_mode = str(time_mode) + self.simulated = self.time_mode == "simulated" + # sim-time stamped by the aggregator on the current task (simulated + # mode); set when weights are received. None until first task. + self._sim_send_ts = None # Use the battery_threshold to determine the # avl_events_3_state config. Default to 50 if not provided @@ -231,7 +249,33 @@ def parse_trace(value): # flag to decide whether the trainer upon unavailability will wait or exit self.wait_until_next_avl = self.config.hyperparameters.wait_until_next_avl - + + # Data streaming: reveal samples over time (see MIGRATING_TO_LAUNCHER.md) + ds_cfg = getattr(self.config.hyperparameters, "data_streaming", None) or {} + self.data_streaming_enabled = str(ds_cfg.get("enabled", "False")) == "True" + self.data_streaming_full_after_s = float( + ds_cfg.get("full_data_available_after_s", 0) + ) + logger.info( + f"Trainer {self.trainer_id}: data streaming " + f"{'ENABLED' if self.data_streaming_enabled else 'DISABLED'} " + f"(full_data_available_after_s={self.data_streaming_full_after_s})" + ) + + # Streamed-vs-full statistical-utility counterfactual telemetry (opt-in). + # Config: util_counterfactual: {enabled, every_n_rounds, sample_size}. + uc_cfg = getattr(self.config.hyperparameters, "util_counterfactual", None) or {} + self.util_cf_enabled = str(uc_cfg.get("enabled", "False")) == "True" + self.util_cf_every_n = int(uc_cfg.get("every_n_rounds", 1) or 1) + _ss = uc_cfg.get("sample_size", 256) + self.util_cf_sample_size = int(_ss) if _ss not in (None, "None", "") else None + self._pool_tensor_cache = None # lazily materialized full-pool tensors + logger.info( + f"Trainer {self.trainer_id}: util counterfactual " + f"{'ENABLED' if self.util_cf_enabled else 'DISABLED'} " + f"(every_n_rounds={self.util_cf_every_n}, sample_size={self.util_cf_sample_size})" + ) + # Initialize memory profiler self.memory_profiler = MemoryProfiler( trainer_id=str(self.trainer_id), @@ -243,13 +287,42 @@ def check_and_sleep(self): """Induce transient unavailability""" pass + def _sim_now(self) -> float: + """Current simulated time (sim-seconds since trainer start). + + real mode: wall-clock elapsed (1:1, no speedup). simulated mode: the + sim-time the aggregator stamped on the most recent task (advances only + as tasks arrive), so availability/streaming track the virtual clock. + """ + if self.simulated: + return float(self._sim_send_ts) if self._sim_send_ts is not None else 0.0 + return time.time() - self.trainer_start_ts + + def _refresh_avl_for_sim(self) -> None: + """Catch availability state up to the current sim-time (simulated mode). + + In real mode the notify thread advances state in wall-clock; in + simulated mode sim-time jumps with each task, so pop all events due by + _sim_now() before a task decides on availability. + """ + if not self.simulated: + return + guard = 0 + while ( + self.state_avl_event_ts + and self._sim_now() >= self.state_avl_event_ts[0][0] + and guard < 100000 + ): + self.check_and_update_state_avl() + guard += 1 + def check_and_update_state_avl(self): if hasattr(self, "cm") and self.cm is not None: if len(self.state_avl_event_ts) > 0: - next_event_ts = self.trainer_start_ts + ( - self.state_avl_event_ts[0][0] / self.speedup_factor - ) - if time.time() >= next_event_ts: + # event timestamps are in sim-seconds since start; compare to + # the current sim-time (no speedup_factor in either mode). + sim_elapsed = self._sim_now() + if sim_elapsed >= self.state_avl_event_ts[0][0]: state_to_set = self.state_avl_event_ts.pop(0)[1] old_status = self.avl_state.value try: @@ -263,6 +336,13 @@ def check_and_update_state_avl(self): logger.info( f"Changed the availability status of trainer {self.trainer_id} from {old_status} to {new_status}" ) + if telemetry.is_enabled(): + ev, fields = build_avail_change( + round_num=int(getattr(self, "_round", 0)), + old_state=str(old_status), + new_state=str(new_status), + ) + telemetry.emit(ev, **fields) if self.client_notify["enabled"] == "True": self._perform_channel_state_update( tag="upload", @@ -330,70 +410,72 @@ def load_data(self) -> None: indices = torch.tensor(self.trainer_indices_list) dataset = data_utils.Subset(dataset, indices) - - # GPU pre-loading optimization for small datasets - # This significantly reduces CPU RAM usage by keeping data on GPU + + # GPU pre-load for small datasets (cuts CPU RAM). Full pool is retained; + # the loader is (re)built from a prefix in _rebuild_stream_loader. dataset_size = len(indices) gpu_preload_threshold = 2000 # Adjust based on GPU memory availability - + if dataset_size <= gpu_preload_threshold and self.device is not None: logger.info( f"Trainer {self.trainer_id}: Pre-loading {dataset_size} samples to GPU " f"to reduce CPU RAM usage" ) - + # Load all data to GPU at once temp_loader = torch.utils.data.DataLoader( dataset, batch_size=dataset_size, shuffle=False ) - + all_data = [] all_targets = [] for data, target in temp_loader: all_data.append(data.to(self.device)) all_targets.append(target.to(self.device)) - - # Create TensorDataset on GPU - gpu_dataset = data_utils.TensorDataset( - torch.cat(all_data), torch.cat(all_targets) - ) - - train_kwargs = { + + # Retain the full GPU pool; loader built below from a prefix + self._stream_gpu = True + self._stream_all_data = torch.cat(all_data) + self._stream_all_targets = torch.cat(all_targets) + self._stream_train_kwargs = { "batch_size": self.batch_size, - "drop_last": False, # Keep incomplete batches for small datasets in FL + "drop_last": False, # keep incomplete batches in FL "shuffle": True, - "num_workers": 0, # No workers needed - data already on GPU + "num_workers": 0, # data already on GPU } - - self.train_loader = torch.utils.data.DataLoader(gpu_dataset, **train_kwargs) - - # Release temporary loader and CPU dataset - del temp_loader, all_data, all_targets - + + del temp_loader, all_data, all_targets, dataset logger.info( f"Trainer {self.trainer_id}: Successfully pre-loaded data to GPU" ) else: - # Standard loading for larger datasets - train_kwargs = { + # Standard loading for larger datasets; retain full Subset + self._stream_gpu = False + self._stream_full_dataset = dataset + self._stream_train_kwargs = { "batch_size": self.batch_size, - "drop_last": False, # Keep incomplete batches for small datasets in FL + "drop_last": False, # keep incomplete batches in FL "shuffle": True, - "num_workers": 0, # Changed from 2 to 0 - reduces CPU RAM usage from worker processes - "pin_memory": True, # Use pinned memory for faster CPU->GPU transfers + "num_workers": 0, # reduces CPU RAM from worker processes + "pin_memory": True, # faster CPU->GPU transfers } - - self.train_loader = torch.utils.data.DataLoader(dataset, **train_kwargs) - logger.info( f"Trainer {self.trainer_id}: Using standard loading " f"({dataset_size} samples exceeds GPU pre-load threshold)" ) - # Release the memory of the full dataset - del dataset + # Fixed shuffle of the pool (seeded by trainer_id) = data arrival + # order; streaming reveals a growing prefix of it. + self._stream_total = dataset_size + seed = int(hashlib.sha256(str(self.trainer_id).encode()).hexdigest(), 16) % (2**31) + self._stream_order = torch.randperm( + self._stream_total, generator=torch.Generator().manual_seed(seed) + ) + + # Build initial loader (full pool unless streaming is enabled) + self._rebuild_stream_loader() gc.collect() - + # Log DataLoader memory info dataloader_info = self.memory_profiler.get_dataloader_memory(self.train_loader) logger.info( @@ -410,6 +492,119 @@ def load_data(self) -> None: f"{time.time()}" ) + def _visible_sample_count(self) -> int: + """Samples unlocked so far: linear in sim-time, full after X sim-sec.""" + if not self.data_streaming_enabled or self.data_streaming_full_after_s <= 0: + return self._stream_total + # full_data_available_after_s is in sim-seconds; _sim_now() is sim-time + # (wall-clock in real mode, stamped task time in simulated mode). + frac = min(1.0, self._sim_now() / self.data_streaming_full_after_s) + n = math.floor(frac * self._stream_total) + return min(self._stream_total, max(1, n)) # >=1 so loader is non-empty + + def _rebuild_stream_loader(self) -> None: + """Rebuild train_loader over the currently-visible prefix of the pool.""" + n = self._visible_sample_count() + full = n >= self._stream_total + if self._stream_gpu: + if full: + subset = data_utils.TensorDataset( + self._stream_all_data, self._stream_all_targets + ) + else: + # align index with data device (order is built on CPU) + pos = self._stream_order[:n].to(self._stream_all_data.device) + subset = data_utils.TensorDataset( + self._stream_all_data[pos], self._stream_all_targets[pos] + ) + else: + pos = self._stream_order if full else self._stream_order[:n] + subset = data_utils.Subset(self._stream_full_dataset, pos.tolist()) + self.train_loader = torch.utils.data.DataLoader( + subset, **self._stream_train_kwargs + ) + + def _pool_tensors(self): + """Return (data, targets) tensors for the full sample pool on device. + + Reuses the GPU-preloaded pool when available; otherwise materializes + the CPU Subset once and caches it. Only used for opt-in counterfactual + telemetry, so the one-time cost is acceptable. + """ + if getattr(self, "_stream_gpu", False): + return self._stream_all_data, self._stream_all_targets + if self._pool_tensor_cache is not None: + return self._pool_tensor_cache + loader = torch.utils.data.DataLoader( + self._stream_full_dataset, batch_size=512, shuffle=False + ) + datas, targets = [], [] + for d, t in loader: + datas.append(d) + targets.append(t) + data = torch.cat(datas) + target = torch.cat(targets) + self._pool_tensor_cache = (data, target) + return self._pool_tensor_cache + + def _oort_utility(self, data, targets, norm_n, sample_size=None): + """Oort statistical utility over a (sampled) set: N * sqrt(mean(loss^2)). + + Mirrors the trainer's Oort utility but computed with no_grad over an + arbitrary index set, so the streamed-prefix and full-pool values are + directly comparable. Returns (utility, n_used). + """ + n = data.shape[0] + if n == 0: + return 0.0, 0 + if sample_size is not None and n > sample_size: + sel = torch.randperm(n)[:sample_size] + data = data[sel.to(data.device)] + targets = targets[sel.to(targets.device)] + criterion = self.loss_fn(reduction="none") + self.model.eval() + with torch.no_grad(): + data = data.to(self.device) + targets = targets.to(self.device) + output = self.model(data) + per_sample = criterion(output, targets) + sumsq = torch.square(per_sample).sum().item() + n_used = data.shape[0] + utility = norm_n * math.sqrt(sumsq / n_used) if n_used > 0 else 0.0 + return utility, n_used + + def _emit_util_disparity(self, round_num, elapsed_s): + """Compute and emit streamed-prefix vs full-pool utility (opt-in).""" + if not (self.util_cf_enabled and telemetry.is_enabled()): + return + if self.util_cf_every_n > 1 and (round_num % self.util_cf_every_n != 0): + return + try: + data, targets = self._pool_tensors() + total = self._stream_total + visible_n = self._visible_sample_count() + order = self._stream_order.to(data.device) + ss = self.util_cf_sample_size + util_streamed, _ = self._oort_utility( + data[order[:visible_n]], targets[order[:visible_n]], + norm_n=visible_n, sample_size=ss, + ) + util_full, n_used = self._oort_utility( + data[order], targets[order], norm_n=total, sample_size=ss, + ) + ev, fields = build_util_disparity( + round_num=int(round_num), + elapsed_s=elapsed_s, + visible_samples=int(visible_n), + total_samples=int(total), + utility_streamed=util_streamed, + utility_full=util_full, + sample_size_used=n_used, + ) + telemetry.emit(ev, **fields) + except Exception as e: # telemetry must never break training + logger.debug(f"util disparity emit failed: {e}") + def train(self) -> None: logger.info(f"Entered train method for {self.trainer_id}") @@ -427,15 +622,31 @@ def train(self) -> None: if self.task_to_perform != "train": logger.info(f"Trainer {self.trainer_id} is not required to train") return + # telemetry: measure time spent waiting on availability (vs. computing) + _wait_time_s = 0.0 + # simulated mode: refresh availability from the trace at this task's + # sim-time before deciding (sim-time advances only with new tasks). + self._refresh_avl_for_sim() # don't enter the if condition if the three_state_avl switch is off # if we are checking for three_state_avl - check if the mechanism is to wait or exit if self.avl_state != TrainerAvailState.AVL_TRAIN: + if self.simulated: + # sim-time can't advance while we block, so a real-time wait + # would hang. The aggregator selects available trainers; being + # unavailable here means skip this task (it will re-select). + logger.info( + f"Trainer id {self.trainer_id} not available to train " + f"(simulated, sim_t={self._sim_now()}); skipping task." + ) + return if self.wait_until_next_avl == "True": logger.info( f"Trainer id {self.trainer_id} is not available to train. Waiting for it to be available" ) + _wait_start = time.time() while self.avl_state != TrainerAvailState.AVL_TRAIN: time.sleep(1) + _wait_time_s = time.time() - _wait_start else: logger.info( f"Trainer id {self.trainer_id} is not available to train. Exiting training." @@ -444,6 +655,14 @@ def train(self) -> None: logger.info(f"Trainer {self.trainer_id} available to train") + # Refresh visible data for this selection (no-op if streaming off) + if self.data_streaming_enabled: + self._rebuild_stream_loader() + logger.info( + f"Trainer {self.trainer_id} streaming: " + f"{self._visible_sample_count()}/{self._stream_total} samples visible" + ) + """Train a model.""" self.criterion = torch.nn.CrossEntropyLoss() @@ -467,22 +686,35 @@ def train(self) -> None: # reset stat utility for OORT self.reset_stat_utility() - # Log training start with comprehensive info + # Log training start with comprehensive info and the expected cycle time. + # real mode: wall-clock will be ~GPU + D (modeled delay). + # simulated mode: wall-clock will be ~GPU only; D is reported to the + # aggregator via sim_completion_ts so it can order updates correctly. num_batches = len(self.train_loader) dataset_size = len(self.train_loader.dataset) + _D = self.training_delay_s if self.training_delay_enabled else 0.0 + if self.simulated: + _expected_wallclock_hint = f"~GPU only wall-clock; reports GPU+D to aggregator (D={_D:.1f}s)" + else: + _expected_wallclock_hint = f"~GPU + {_D:.1f}s sleep = ~{_D:.1f}s+ wall-clock" logger.info( f"[TRAIN_START] Trainer {self.trainer_id} starting training with " f"model_version={self._round}, dataset_size={dataset_size}, " - f"num_batches={num_batches}, batch_size={self.batch_size}, epochs={self.epochs}" + f"num_batches={num_batches}, batch_size={self.batch_size}, epochs={self.epochs}, " + f"time_mode={self.time_mode}, expected_cycle_time={_expected_wallclock_hint}" ) + _cycle_start = time.time() total_batches_processed = 0 final_loss = None + _gpu_start = time.time() for epoch in range(1, self.epochs + 1): epoch_batches, epoch_loss = self._train_epoch(epoch) total_batches_processed += epoch_batches if epoch_loss is not None: final_loss = epoch_loss + # real GPU/compute time for this round, excluding any simulated delay + _real_gpu_time_s = time.time() - _gpu_start # Log training completion summary loss_str = f"{final_loss:.6f}" if final_loss is not None else "N/A" @@ -509,13 +741,77 @@ def train(self) -> None: # Log memory after training round self.memory_profiler.log_memory_after_round() - # emulate delays in training (due to compute resource and/or - # dataset size and/or network latency) if enabled - if self.training_delay_enabled == "True": - time.sleep(self.training_delay_s / self.speedup_factor) + # Round duration = actual GPU compute + modeled device delay D. + # In real mode, the aggregator measures recv_ts - sent_ts ≈ GPU + D + MQTT. + # In simulated mode we skip the sleep but must still report GPU + D so that: + # 1. SIM_ROUND_DURATION fed to OORT matches what real mode would measure + # (giving OORT real per-trainer speed heterogeneity, not a uniform D). + # 2. sim_completion_ts = sim_send_ts + GPU + D correctly orders trainers by + # actual completion time — trainers with lighter data / faster GPU finish + # first, replicating real arrival order. + # Without _real_gpu_time_s, all trainers report identical D, making + # sim_completion_ts equal for everyone selected in the same loop iteration, + # which causes deterministic OORT over-selection of the top stat-utility set. + _modeled_delay_s = self.training_delay_s if self.training_delay_enabled else 0.0 + sim_round_duration = _real_gpu_time_s + _modeled_delay_s + self._sim_round_duration = sim_round_duration + self._sim_completion_ts = ( + (self._sim_send_ts if self._sim_send_ts is not None else self._sim_now()) + + sim_round_duration + ) + + # Telemetry: per-round timing/availability + streamed-vs-full disparity. + if telemetry.is_enabled(): + visible = ( + self._visible_sample_count() + if self.data_streaming_enabled + else self._stream_total + ) + ev, fields = build_trainer_round( + round_num=int(getattr(self, "_round", 0)), + real_gpu_time_s=_real_gpu_time_s, + sim_round_duration_s=sim_round_duration, + wait_time_s=_wait_time_s, + avail_state=self.avl_state.value, + visible_samples=int(visible), + total_samples=int(self._stream_total), + dataset_size=int(dataset_size), + stat_utility=float(self._stat_utility) + if isinstance(self._stat_utility, (int, float)) + else float(getattr(self._stat_utility, "item", lambda: 0.0)()), + final_loss=final_loss, + extra={"sim_completion_ts": self._sim_completion_ts, + "time_mode": self.time_mode}, + ) + telemetry.emit(ev, **fields) + self._emit_util_disparity( + int(getattr(self, "_round", 0)), self._sim_now() + ) + + # real mode: sleep D so a slow device actually takes that long in + # wall-clock. simulated mode: skip the sleep (the aggregator uses the + # reported sim_completion_ts instead). + if self.training_delay_enabled and not self.simulated: + time.sleep(self.training_delay_s) + + # Validate actual cycle time vs expected, and log the summary. + _cycle_elapsed = time.time() - _cycle_start + if self.simulated: + # Expected: ~GPU only. Log actual so we can confirm no sleep crept in. logger.info( - f"Delayed training time for trainer " - f"{self.trainer_id} by {self.training_delay_s}s" + f"[TRAIN_CYCLE] Trainer {self.trainer_id} round={self._round} " + f"time_mode=simulated: actual_wall={_cycle_elapsed:.2f}s " + f"(GPU={_real_gpu_time_s:.2f}s + D={_modeled_delay_s:.1f}s = reported {sim_round_duration:.2f}s), " + f"sim_completion_ts={self._sim_completion_ts:.2f}" + ) + else: + _expected = _real_gpu_time_s + sim_round_duration + _ok = _cycle_elapsed >= sim_round_duration # cycle must be at least D + logger.info( + f"[TRAIN_CYCLE] Trainer {self.trainer_id} round={self._round} " + f"time_mode=real: actual_wall={_cycle_elapsed:.2f}s " + f"(GPU={_real_gpu_time_s:.2f}s + sleep={sim_round_duration:.1f}s = expected~{_expected:.1f}s) " + f"{'OK' if _ok else 'WARN: actual < D, sleep may not have fired'}" ) def _train_epoch(self, epoch): @@ -605,13 +901,25 @@ def evaluate(self) -> None: ) return + # simulated mode: refresh availability at this task's sim-time. + self._refresh_avl_for_sim() if self.avl_state == TrainerAvailState.UN_AVL: + if self.simulated: + logger.info( + f"Trainer id {self.trainer_id} unavailable for eval " + f"(simulated, sim_t={self._sim_now()}); skipping." + ) + return logger.warning( f"Trainer id {self.trainer_id} is not available to perform forward pass evaluate. Waiting for it to be available" ) while self.avl_state == TrainerAvailState.UN_AVL: time.sleep(1) + # Use the same currently-visible data as training (no-op if off) + if self.data_streaming_enabled: + self._rebuild_stream_loader() + logger.info(f"Starting eval (forward pass) for trainer id {self.trainer_id}") for epoch in range(1, self.epochs + 1): for batch_idx, (data, target) in enumerate(self.train_loader): @@ -637,15 +945,12 @@ def evaluate(self) -> None: # normalize statistical utility of a trainer based on the size # of the dataset self.normalize_stat_utility(epoch) - if self.training_delay_enabled == "True": - # Updated eval duration to be one-third of training - # duration since it is evidenced on text and through - # profiling - # Eval is 3X faster than training on CPU - # Eval is 10-50X faster than training on CPUs due to NPUs - # not supporting training. We take 20X + # Eval is ~20x faster than training (NPUs don't support training), so + # the modeled eval delay is training_delay_s/20. real mode sleeps it; + # simulated mode skips it (folded into the reported sim duration). + if self.training_delay_enabled and not self.simulated: eval_delay = math.floor(self.training_delay_s / 20.0) - time.sleep(eval_delay / self.speedup_factor) + time.sleep(eval_delay) logger.debug( f"Delayed eval time for trainer " f"{self.trainer_id} by {eval_delay}s" ) @@ -707,12 +1012,14 @@ def main(): required=False, ) - # Add argument to speed up client's timescale by a factor + # Simulation time mode (replaces the removed speedup_factor). parser.add_argument( - "--speedup_factor", - type=float, - default=1.0, - help="Speedup factor to accelarate all events and training/ eval durations from the trainer. Default- no acceleration", + "--time_mode", + type=str, + choices=["real", "simulated"], + default="simulated", + help="'real': sleep modeled delays at true pace. 'simulated': skip " + "sleeps; aggregator orders updates by a virtual clock.", required=False, ) @@ -747,17 +1054,21 @@ def main(): raise ValueError("Must provide either --config or --config-json") print(f"[TRAINER STARTUP] Creating trainer object...") - t = PyTorchCifar10Trainer(config, args.battery_threshold, args.speedup_factor) + t = PyTorchCifar10Trainer(config, args.battery_threshold, args.time_mode) + + # Structured telemetry (no-op unless $FLAME_TELEMETRY_DIR is set by the + # launcher). One JSONL file per trainer process. + telemetry.configure(role="trainer", end_id=str(t.trainer_id)) print(f"[TRAINER STARTUP] Trainer created - ID: {t.trainer_id}, Job: {t.config.job.job_id}") logger.info(f"========== TRAINER STARTED: ID={t.trainer_id}, PID={os.getpid()} ==========") print( - f"# Trainer id: {t.trainer_id}, has heartbeats_enabled: " - f"{t.heartbeats_enabled}, has client_notify: " - f"{t.client_notify['enabled']}, has " + f"# Trainer id: {t.trainer_id}, time_mode: {t.time_mode}, " + f"has heartbeats_enabled: {t.heartbeats_enabled}, " + f"has client_notify: {t.client_notify['enabled']}, " f"training_delay_enabled: {t.training_delay_enabled}, " - f"with training_delay_s: {t.training_delay_s}" + f"training_delay_s: {t.training_delay_s}" ) # Register exit handler to generate memory report diff --git a/lib/python/examples/async_cifar10/trainer/pytorch/main_oort_trainer.py b/lib/python/examples/async_cifar10/trainer/pytorch/main_oort_trainer.py deleted file mode 100644 index 23c227248..000000000 --- a/lib/python/examples/async_cifar10/trainer/pytorch/main_oort_trainer.py +++ /dev/null @@ -1,366 +0,0 @@ -# Copyright 2022 Cisco Systems, Inc. and its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); you -# may not use this file except in compliance with the License. You may -# obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or -# implied. See the License for the specific language governing -# permissions and limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 -"""CIFAR-10 horizontal FL, OORT trainer for PyTorch. - -The example below is implemented based on the following example from -pytorch: -https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html. -""" - -import ast -import calendar -import logging -import time - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -import torch.utils.data as data_utils -import torchvision -import torchvision.transforms as transforms -from flame.config import Config -from flame.mode.horizontal.oort.trainer import Trainer -from torch import Tensor -from torchvision.datasets import CIFAR10 - -logger = logging.getLogger(__name__) - - -class Net(nn.Module): - """Net class.""" - - def __init__(self): - """Initialize.""" - super(Net, self).__init__() - self.conv1 = nn.Conv2d(3, 64, 3) - self.conv2 = nn.Conv2d(64, 128, 3) - self.conv3 = nn.Conv2d(128, 256, 3) - self.pool = nn.MaxPool2d(2, 2) - self.fc1 = nn.Linear(64 * 4 * 4, 128) - self.fc2 = nn.Linear(128, 256) - self.fc3 = nn.Linear(256, 10) - - def forward(self, x): - """Forward.""" - x = self.pool(F.relu(self.conv1(x))) - x = self.pool(F.relu(self.conv2(x))) - x = self.pool(F.relu(self.conv3(x))) - x = x.view(-1, 64 * 4 * 4) - x = F.relu(self.fc1(x)) - x = F.relu(self.fc2(x)) - x = self.fc3(x) - return F.log_softmax(x, dim=1) - - -class PyTorchCifar10Trainer(Trainer): - """PyTorch CIFAR-10 Trainer.""" - - def __init__(self, config: Config) -> None: - """Initialize a class instance.""" - self.config = config - self.dataset_size = 0 - self.model = None - # Oort requires its loss function to have 'reduction' - # parameter - self.loss_fn = torch.nn.CrossEntropyLoss - - self.device = None - self.train_loader = None - - self.learning_rate = self.config.hyperparameters.learning_rate - self.epochs = self.config.hyperparameters.epochs - self.batch_size = self.config.hyperparameters.batch_size or 16 - self.trainer_id = self.config.task_id - - self.criterion = None - - # TODO: Remove the hard requirement for config to include - # trainer_indices_list and two_state_unavl_durations_s Setting the - # indices used by the trainer - self.trainer_indices_list = self.config.hyperparameters.trainer_indices_list - # Loading the failure durations for trainers - self.trainer_start_ts = time.time() - self.two_state_unavl_durations_s = ast.literal_eval( - self.config.hyperparameters.two_state_unavl_durations_s - ) - self.timestamp_next_sleep_s = calendar.timegm( - time.strptime("Dec 31, 2030 @ 23:59:59 UTC", "%b %d, %Y @ %H:%M:%S UTC") - ) - if len(self.two_state_unavl_durations_s) > 0: - self.timestamp_next_sleep_s = ( - self.trainer_start_ts + self.two_state_unavl_durations_s[0][0] - ) - - def check_and_sleep(self): - curr_time = time.time() - - if (curr_time >= self.timestamp_next_sleep_s) and ( - len(self.two_state_unavl_durations_s) > 0 - ): - # pop leftmost element - sleep_config_tuple = self.two_state_unavl_durations_s.pop(0) - - # get the duration of sleep and set the params for next - # sleep - sleep_start_ts_from_trainer_init = ( - self.trainer_start_ts + sleep_config_tuple[0] - ) - sleep_duration_s = sleep_config_tuple[1] - - # remaining sleep = trainer_start_ts + actual_sleep_start - # + actual_sleep_duration - current_ts - remaining_sleep_duration_s = ( - sleep_start_ts_from_trainer_init + sleep_duration_s - curr_time - ) - logger.debug( - f"Task_id: {self.trainer_id} given_sleep_duration_s: {sleep_duration_s} with remaining_sleep_duration_s: {remaining_sleep_duration_s} at timestamp: {curr_time}" - ) - - if remaining_sleep_duration_s <= 0: - logger.info( - f"Task_id: {self.trainer_id} got -ve remaining sleep at timestamp: {curr_time}" - ) - # Need to pop out failure intervals that occur in the - # past - time_elapsed_from_start = curr_time - self.trainer_start_ts - while time_elapsed_from_start > ( - self.two_state_unavl_durations_s[0][0] - + self.two_state_unavl_durations_s[0][1] - ): - self.two_state_unavl_durations_s.pop(0) - if len(self.two_state_unavl_durations_s) == 0: - break - else: - logger.info( - f"Task_id: {self.trainer_id} going to sleep up at timestamp: {time.time()}" - ) - time.sleep(remaining_sleep_duration_s) - logger.info( - f"Task_id: {self.trainer_id} woke up at timestamp: {time.time()}" - ) - - # check if failure_list is now empty, if yes, reset - # ts_next_sleep_s if not empty, set it to the next value - if len(self.two_state_unavl_durations_s) > 0: - self.timestamp_next_sleep_s = ( - self.trainer_start_ts + self.two_state_unavl_durations_s[0][0] - ) - if self.timestamp_next_sleep_s < time.time(): - logger.info( - f"Task_id: {self.trainer_id} ERROR - JUST SET NEXT self.timestamp_next_sleep_s < curr_time" - ) - else: - self.timestamp_next_sleep_s = calendar.timegm( - time.strptime( - "Dec 31, 2030 @ 23:59:59 UTC", "%b %d, %Y @ %H:%M:%S UTC" - ) - ) - logger.info(f"Task_id: {self.trainer_id} no more sleep for trainer") - - logger.debug( - f"Task_id: {self.trainer_id} check_and_sleep completed at timestamp: {time.time()}" - ) - - def initialize(self) -> None: - """Initialize role.""" - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - self.model = Net().to(self.device) - logger.debug( - f"Task_id: {self.trainer_id} initialize completed at timestamp: {time.time()}" - ) - - def load_data(self) -> None: - """Load data.""" - transform_train = transforms.Compose( - [ - transforms.RandomCrop(32, padding=4), - transforms.RandomHorizontalFlip(), - transforms.ToTensor(), - transforms.Normalize( - (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010) - ), - ] - ) - - dataset = CIFAR10( - "/home/dgarg39/flame/lib/python/examples/async_cifar10/data", - train=True, - download=True, - transform=transform_train, - ) - - # create indices into a list and convert to tensor - indices = torch.tensor(self.trainer_indices_list) - - dataset = data_utils.Subset(dataset, indices) - - # GPU pre-loading optimization for small datasets - # This significantly reduces CPU RAM usage by keeping data on GPU - dataset_size = len(indices) - gpu_preload_threshold = 2000 # Adjust based on GPU memory availability - - if dataset_size <= gpu_preload_threshold and self.device is not None: - logger.info( - f"Trainer {self.trainer_id}: Pre-loading {dataset_size} samples to GPU " - f"to reduce CPU RAM usage" - ) - - # Load all data to GPU at once - temp_loader = torch.utils.data.DataLoader( - dataset, batch_size=dataset_size, shuffle=False - ) - - all_data = [] - all_targets = [] - for data, target in temp_loader: - all_data.append(data.to(self.device)) - all_targets.append(target.to(self.device)) - - # Create TensorDataset on GPU - gpu_dataset = data_utils.Subset(dataset, indices) - gpu_dataset = torch.utils.data.TensorDataset( - torch.cat(all_data), torch.cat(all_targets) - ) - - train_kwargs = { - "batch_size": self.batch_size, - "drop_last": False, # Keep incomplete batches for small datasets in FL - "shuffle": True, - "num_workers": 0, # No workers needed - data already on GPU - } - - self.train_loader = torch.utils.data.DataLoader(gpu_dataset, **train_kwargs) - - # Release temporary loader - del temp_loader, all_data, all_targets - - logger.info( - f"Trainer {self.trainer_id}: Successfully pre-loaded data to GPU" - ) - else: - # Standard loading for larger datasets - train_kwargs = { - "batch_size": self.batch_size, - "drop_last": False, # Keep incomplete batches for small datasets in FL - "shuffle": True, - "num_workers": 0, # Changed from 2 to 0 - reduces CPU RAM usage from worker processes - "pin_memory": True, # Use pinned memory for faster CPU->GPU transfers - } - - self.train_loader = torch.utils.data.DataLoader(dataset, **train_kwargs) - - logger.info( - f"Trainer {self.trainer_id}: Using standard loading " - f"({dataset_size} samples exceeds GPU pre-load threshold)" - ) - - logger.debug( - f"Task_id: {self.trainer_id} load_data completed at timestamp: {time.time()}" - ) - - def train(self) -> None: - """Train a model.""" - self.criterion = torch.nn.CrossEntropyLoss() - self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate) - - # Log training start with comprehensive info - num_batches = len(self.train_loader) - dataset_size = len(self.train_loader.dataset) - logger.info( - f"[TRAIN_START] Trainer {self.trainer_id} starting training with " - f"model_version={self._round}, dataset_size={dataset_size}, " - f"num_batches={num_batches}, batch_size={self.batch_size}, epochs={self.epochs}" - ) - - self.reset_stat_utility() - total_batches_processed = 0 - final_loss = None - for epoch in range(1, self.epochs + 1): - epoch_batches, epoch_loss = self._train_epoch(epoch) - total_batches_processed += epoch_batches - if epoch_loss is not None: - final_loss = epoch_loss - - # Log training completion summary - loss_str = f"{final_loss:.6f}" if final_loss is not None else "N/A" - logger.info( - f"[TRAIN_COMPLETE] Trainer {self.trainer_id} completed training with " - f"model_version={self._round}, dataset_size={dataset_size}, " - f"total_batches_processed={total_batches_processed}, final_loss={loss_str}" - ) - - # save dataset size so that the info can be shared with - # aggregator - self.dataset_size = len(self.train_loader.dataset) - - def _train_epoch(self, epoch): - self.model.train() - - batches_processed = 0 - last_loss = None - - for batch_idx, (data, target) in enumerate(self.train_loader): - data, target = data.to(self.device), target.to(self.device) - self.optimizer.zero_grad() - output = self.model(data) - # OLD CODE loss BEFORE OORT loss = F.nll_loss(output, - # target) calculate statistical utility of a trainer while - # calculating loss - loss = self.oort_loss(output, target, epoch, batch_idx) - loss.backward() - self.optimizer.step() - batches_processed += 1 - - # Log every batch for small trainers, every 100 for large trainers - num_batches = len(self.train_loader) - should_log = (num_batches <= 10) or (batch_idx % 100 == 0) - if should_log: - done = batch_idx * len(data) - total = len(self.train_loader.dataset) - percent = 100.0 * batch_idx / len(self.train_loader) - loss_val = loss.item() - last_loss = loss_val - logger.info( - f"epoch: {epoch} [{done}/{total} ({percent:.0f}%)]" - f"\tloss: {loss_val:.6f}" - ) - - # normalize statistical utility of a trainer based on the size - # of the dataset - self.normalize_stat_utility(epoch) - - return batches_processed, last_loss - - def evaluate(self) -> None: - """Evaluate a model.""" - # Implement this if testing is needed in trainer - pass - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="") - parser.add_argument("config", nargs="?", default="./config.json") - - args = parser.parse_args() - config = Config(args.config) - - t = PyTorchCifar10Trainer(config) - t.compose() - t.run() diff --git a/lib/python/flame/channel.py b/lib/python/flame/channel.py index 361ef3d47..95c731931 100644 --- a/lib/python/flame/channel.py +++ b/lib/python/flame/channel.py @@ -428,7 +428,7 @@ async def _get(): return msg, timestamp def recv_fifo( - self, end_ids: list[str], first_k: int = 0 + self, end_ids: list[str], first_k: int = 0, timeout: float = None ) -> tuple[Any, tuple[str, datetime]]: """Receive a message per end from a list of ends. @@ -445,6 +445,10 @@ def recv_fifo( means that we'd like to receive messages from all ends in the list. If first_k > len(end_ids), first_k is set to len(end_ids). + timeout: optional per-message wait budget in seconds. If no message + arrives within it, yield (None, ("", now)) and stop, so a + caller never blocks forever on in-flight ends that have gone + quiet (unavailable / departed). Default None = block (legacy). Returns ------- @@ -483,8 +487,19 @@ async def _get_message_inner(): # the _get_message_inner() coroutine fetches a message from # the temp queue; we call this coroutine first_k times for _ in range(first_k): - result, status = run_async(_get_message_inner(), self._backend.loop()) + result, status = run_async( + _get_message_inner(), self._backend.loop(), timeout=timeout + ) logger.info(f"After getting message, status: {status}") + # timeout (or any non-delivery): don't index into a None result; + # signal "no message" to the caller and stop yielding. + if not status or result is None: + logger.info( + f"recv_fifo: no message within timeout={timeout}s; " + f"yielding None and stopping" + ) + yield None, ("", datetime.now()) + return (end_id, payload) = result logger.info(f"get payload for {end_id}") diff --git a/lib/python/flame/common/util.py b/lib/python/flame/common/util.py index 42719208a..759c942a0 100644 --- a/lib/python/flame/common/util.py +++ b/lib/python/flame/common/util.py @@ -135,6 +135,7 @@ def run_async(coro, loop, timeout=None): try: return fut.result(timeout), True except concurrent.futures.TimeoutError: + fut.cancel() # remove stale waiter from _rx_queue so it can't steal future messages return None, False diff --git a/lib/python/flame/config.py b/lib/python/flame/config.py index e08555e9f..3587ab742 100644 --- a/lib/python/flame/config.py +++ b/lib/python/flame/config.py @@ -148,6 +148,7 @@ class Hyperparameters(FlameSchema, extra=Extra.allow): rounds: int epochs: int aggregation_goal: t.Optional[int] = Field(alias="aggGoal", default=None) + eval_every_n_rounds: t.Optional[int] = Field(alias="evalEveryNRounds", default=10) eval_goal_factor: t.Optional[float] = Field(alias="evalGoalFactor", default=None) round_nudge_type: t.Optional[str] = Field( alias="roundNudgeType", default="last_train" diff --git a/lib/python/flame/launch/execution_config_generator.py b/lib/python/flame/launch/execution_config_generator.py index aa2cf37ff..40dc19bfb 100644 --- a/lib/python/flame/launch/execution_config_generator.py +++ b/lib/python/flame/launch/execution_config_generator.py @@ -115,7 +115,7 @@ def create_execution_config( "trace_key": availability_trace_key, }, "battery_threshold": exp_config.trainer.battery_threshold, - "speedup_factor": exp_config.trainer.speedup_factor, + "time_mode": exp_config.trainer.time_mode, }, "execution": { "num_gpus": exp_config.execution.num_gpus, diff --git a/lib/python/flame/launch/experiment_config.py b/lib/python/flame/launch/experiment_config.py index 111efdfc3..19b65311d 100644 --- a/lib/python/flame/launch/experiment_config.py +++ b/lib/python/flame/launch/experiment_config.py @@ -36,7 +36,10 @@ class TrainerConfig: dataset: DatasetConfig = field(default_factory=DatasetConfig) availability: AvailabilityConfig = field(default_factory=AvailabilityConfig) battery_threshold: int = 50 # For 3-state modes - speedup_factor: float = 1.0 + # Simulation time mode: "real" (trainer sleeps modeled delays at true pace) + # or "simulated" (no sleep; aggregator orders by a virtual clock). Replaces + # the removed speedup_factor. + time_mode: str = "simulated" enable_training_delays: bool = True # Enable per-trainer training delays hyperparameters: Optional[dict] = None # Trainer-specific hyperparameters (e.g., batchSize, learningRate) config_overrides: Optional[dict] = None # Deep-merged into per-trainer config last (wins over baseline) @@ -187,7 +190,7 @@ def from_yaml(cls, yaml_path: Path) -> "ExperimentBatch": else AvailabilityConfig() ), battery_threshold=trainer_data.get("battery_threshold", 50), - speedup_factor=trainer_data.get("speedup_factor", 1.0), + time_mode=trainer_data.get("time_mode", "simulated"), enable_training_delays=trainer_data.get("enable_training_delays", True), hyperparameters=trainer_data.get("hyperparameters"), config_overrides=trainer_data.get("config_overrides"), diff --git a/lib/python/flame/launch/run_experiment.py b/lib/python/flame/launch/run_experiment.py index a6c212912..ce0959d7e 100755 --- a/lib/python/flame/launch/run_experiment.py +++ b/lib/python/flame/launch/run_experiment.py @@ -18,7 +18,23 @@ def _autodetect_example_dir(yaml_path: Path) -> Path: return p.parent +def _force_utf8_stdio() -> None: + """Make console output locale-independent. + + The launcher prints status glyphs (checkmarks, warnings). On hosts whose + locale is latin-1 (LANG=en_US without .UTF-8), printing those crashes with + UnicodeEncodeError. Reconfigure our streams to UTF-8 (replace on failure) + rather than scrubbing every glyph. + """ + for stream in (sys.stdout, sys.stderr): + try: + stream.reconfigure(encoding="utf-8", errors="replace") + except Exception: + pass + + def main(argv: list[str] | None = None) -> int: + _force_utf8_stdio() parser = argparse.ArgumentParser(description="Run federated experiments from YAML.") parser.add_argument("config", type=Path, help="experiment YAML file") parser.add_argument( diff --git a/lib/python/flame/launch/runner.py b/lib/python/flame/launch/runner.py index 979a2aad5..ee7c9c917 100644 --- a/lib/python/flame/launch/runner.py +++ b/lib/python/flame/launch/runner.py @@ -10,8 +10,10 @@ """ import json +import os import re import signal +import subprocess import sys from datetime import datetime from pathlib import Path @@ -118,6 +120,9 @@ def run_experiment(self, exp: ExperimentConfig) -> None: agg_cfg, agg_provenance = self._build_aggregator_config( exp, agg_config_path, baseline_entry ) + # Propagate the simulation time mode to the aggregator (it needs to + # know whether to order updates by a virtual clock or by arrival). + agg_cfg.setdefault("hyperparameters", {})["time_mode"] = exp.trainer.time_mode agg_job_id = agg_cfg.get("job", {}).get("id") agg_job_name = agg_cfg.get("job", {}).get("name") if not agg_job_id: @@ -147,12 +152,22 @@ def run_experiment(self, exp: ExperimentConfig) -> None: trainers_log = self.current_exp_dir / f"{log_prefix}_trainers.log" monitor_log = self.current_exp_dir / f"{log_prefix}_resources.log" + # Telemetry: every spawned process inherits this dir and writes its + # own JSONL stream. Set before spawning so children pick it up. + self.telemetry_dir = self.current_exp_dir / "telemetry" + os.environ["FLAME_TELEMETRY_DIR"] = str(self.telemetry_dir) + # UTF-8 child stdio so status glyphs don't crash on latin-1 locales. + os.environ.setdefault("PYTHONIOENCODING", "utf-8") + self.aggregator_spawner = AggregatorSpawner(log_file=agg_log) self.trainer_spawner = TrainerSpawner( config_gen, num_gpus=exp.execution.num_gpus, sleep_between_spawns=exp.execution.sleep_between_spawns, log_file=trainers_log, + # CLI-only knobs passed on the trainer command line. + time_mode=exp.trainer.time_mode, + battery_threshold=exp.trainer.battery_threshold, ) if exp.execution.monitoring.enabled and create_monitor_from_config is not None: @@ -239,9 +254,21 @@ def run_experiment(self, exp: ExperimentConfig) -> None: ) print(f"\nexperiment running. logs: {agg_log}, {trainers_log}") - self.trainer_spawner.wait_all() + # Wait for the aggregator to finish all rounds first, then give + # trainers a short grace window to process the EOT broadcast and + # exit cleanly. Without this, wait_all()'s per-trainer timeout fires + # immediately after spawn and kills trainers every 30s regardless of + # whether training is still in progress. + print(" waiting for aggregator to finish...") + self.aggregator_spawner.process.wait() + print(" aggregator done, waiting for trainers to exit...") + self.trainer_spawner.wait_all(timeout_per_trainer=30.0) print("\nexperiment completed.") + # Auto post-run analysis: parse the telemetry JSONL and emit plots. + # Best-effort; a plotting failure must not fail the experiment. + self._run_post_analysis() + except Exception as e: import traceback @@ -371,6 +398,33 @@ def _signal_handler(self, signum, frame): self._cleanup() sys.exit(130) + def _run_post_analysis(self): + """Run telemetry analysis to generate plots into /plots/. + + Best-effort: never raise. Skips silently if no telemetry was produced. + """ + telemetry_dir = getattr(self, "telemetry_dir", None) + if not telemetry_dir or not telemetry_dir.exists(): + return + # repo root: lib/python/flame/launch/runner.py -> parents[4] + analyzer = ( + Path(__file__).resolve().parents[4] + / "scripts" + / "analysis" + / "analyze_run.py" + ) + if not analyzer.exists(): + print(f" (skipping analysis: {analyzer} not found)") + return + try: + print(f"\nrunning telemetry analysis on {telemetry_dir} ...") + subprocess.run( + [sys.executable, str(analyzer), str(telemetry_dir)], + check=False, + ) + except Exception as e: + print(f" (telemetry analysis failed: {e})") + def _cleanup(self): if self.resource_monitor: self.resource_monitor.stop() diff --git a/lib/python/flame/launch/snapshot.py b/lib/python/flame/launch/snapshot.py index 00248151f..68d240058 100644 --- a/lib/python/flame/launch/snapshot.py +++ b/lib/python/flame/launch/snapshot.py @@ -89,7 +89,7 @@ def _serialize_experiment_config(self, exp_config: ExperimentConfig) -> Dict: "mode": exp_config.trainer.availability.mode, }, "battery_threshold": exp_config.trainer.battery_threshold, - "speedup_factor": exp_config.trainer.speedup_factor, + "time_mode": exp_config.trainer.time_mode, }, "aggregator": ( { diff --git a/lib/python/flame/launch/spawner.py b/lib/python/flame/launch/spawner.py index 756d80a6c..c2ce29e54 100644 --- a/lib/python/flame/launch/spawner.py +++ b/lib/python/flame/launch/spawner.py @@ -188,11 +188,16 @@ def __init__( num_gpus: int = 8, sleep_between_spawns: float = 1.0, log_file: Optional[Path] = None, + time_mode: str = "simulated", + battery_threshold: int = 50, ): self.config_gen = config_generator self.num_gpus = num_gpus self.sleep_between_spawns = sleep_between_spawns self.log_file = log_file + # CLI-only trainer knobs: read from argv, not config JSON. + self.time_mode = time_mode + self.battery_threshold = battery_threshold self.processes = [] self._log_handle = None @@ -242,6 +247,10 @@ def spawn_trainer( str(trainer_main_path), "--config-json", config_json, + "--time_mode", + str(self.time_mode), + "--battery_threshold", + str(self.battery_threshold), ] # Determine stdout/stderr handling @@ -303,10 +312,37 @@ def spawn_all( print(f"\n✓ Spawned {len(self.processes)} trainers") - def wait_all(self): - """Wait for all trainer processes to complete.""" + def wait_all(self, timeout_per_trainer: float = 30.0): + """Wait for all trainer processes to complete. + + Each trainer gets ``timeout_per_trainer`` seconds after the aggregator + exits to process the EOT broadcast and self-terminate. Trainers that + are blocked in ``await_join`` (waiting for the next task from an + aggregator that has already left) will never self-exit, so we force- + terminate them after the window. Without this the runner hangs + indefinitely and the next sequential experiment never starts. + """ + import signal as _signal + + deadline_per = timeout_per_trainer for proc_info in self.processes: - proc_info["process"].wait() + proc = proc_info["process"] + try: + proc.wait(timeout=deadline_per) + except Exception: + # Timeout or other error — gracefully terminate then kill. + print( + f" (trainer PID {proc.pid} did not exit within " + f"{deadline_per}s; terminating)" + ) + try: + proc.terminate() + proc.wait(timeout=5) + except Exception: + try: + proc.kill() + except Exception: + pass def terminate_all(self): """Terminate all trainer processes.""" diff --git a/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py b/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py index 7f3495f8f..fdde156de 100644 --- a/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py +++ b/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py @@ -17,10 +17,11 @@ import logging import time -from datetime import datetime +from datetime import datetime, timedelta import numpy as np from flame.channel import VAL_CH_STATE_HTBT_RECV, VAL_CH_STATE_RECV, VAL_CH_STATE_SEND +from flame.end import KEY_END_STATE, VAL_END_STATE_NONE from flame.common.constants import DeviceType from flame.common.util import weights_to_device, weights_to_model_device from flame.mode.composer import CloneComposer @@ -33,6 +34,10 @@ from flame.mode.message import MessageType from flame.mode.tasklet import Loop, Tasklet from flame.optimizer.train_result import TrainResult +from flame import telemetry +from flame.telemetry.events import build_agg_round +from flame.sim import SimReorderBuffer +from flame.selector.properties import PROP_SIM_SEND_TS from flame.selector.oort import ( PROP_DATASET_SIZE, PROP_LAST_SELECTED_ROUND, @@ -47,6 +52,10 @@ SEND_TIMEOUT_WAIT_S = 90 # 90 seconds timeout +# Max wall-clock to block on one async receive before skipping the cycle and +# re-selecting; guards against hanging when all in-flight trainers go quiet. +RECV_TIMEOUT_WAIT_S = 30 + class TopAggregator(SyncTopAgg): """Asynchronous top level Aggregator implements an ML aggregation @@ -71,6 +80,14 @@ def internal_init(self) -> None: self._per_trainer_staleness_track = {} self._track_trainer_version_duration_s = {} + # Simulated-time receive: reorder buffer (keyed by end, ordered by + # sim_completion_ts), committed-end set for the current agg-goal window, + # and pending-commit set for cross-round stragglers (blocked from + # re-selection until their buffer entry is committed). + self._sim_buffer = SimReorderBuffer() + self._sim_committed: set = set() + self._sim_pending_commit: set = set() + # check if distribute_weights was successful self._prev_distribute_weights_success = False @@ -111,6 +128,9 @@ def _reset_agg_goal_variables(self): f"{self._agg_goal_weights}" ) + if self.simulated: + self._sim_committed.clear() + # TODO: (DG) Need to update or delete, not used right now def _read_heartbeat(self, tag: str) -> None: """Receive trainer heartbeat messaages asynchronously. @@ -173,6 +193,54 @@ def _process_trainer_heartbeat(self, msg, end) -> None: else: logger.warning(f"Got invalid {msg} while processing heartbeat") + # Short per-end probe timeout: how long we wait for one trainer's MQTT + # message before moving on. Not the training wait — just queue delivery jitter. + SIM_RECV_FILL_TIMEOUT_S = 0.5 + + def _sim_recv_min(self, channel, recv_ends): + """Fill the reorder buffer from un-buffered, un-committed ends; pop and + commit the entry with the smallest sim_completion_ts.""" + to_probe = [ + e for e in recv_ends + if not self._sim_buffer.has(e) and e not in self._sim_committed + ] + for e in to_probe: + msg, metadata = next( + channel.recv_fifo([e], 1, timeout=self.SIM_RECV_FILL_TIMEOUT_S) + ) + if msg is not None: + # metadata[0] is the actual sender; may differ from probed end + # if a stale recv task delivered a different end's message first. + actual_end = metadata[0] + sct = msg.get(MessageType.SIM_COMPLETION_TS) + if sct is None: + sct = self._vclock.now + self._sim_buffer.add(actual_end, float(sct), (msg, metadata)) + + # Pop the minimum regardless of recv_ends membership so buffered updates + # are not lost when an end is cleaned up before its commit. + popped = self._sim_buffer.pop_min() + if popped is None: + time.sleep(0.5) + return None, ("", datetime.now()) + _end, sct, (m, md) = popped + self._vclock.advance(sct) + self._sim_committed.add(_end) + logger.debug( + f"[SIM_RECV] committed end={_end[-4:]} sct={sct:.1f} " + f"T_v={self._vclock.now:.1f} buf={len(self._sim_buffer)}" + ) + # Release trainer that was blocked waiting for this cross-round commit. + if _end in self._sim_pending_commit: + self._sim_pending_commit.discard(_end) + sel = channel._selector + if _end in sel.all_selected: + del sel.all_selected[_end] + if channel.has(_end): + channel._ends[_end].set_property(KEY_END_STATE, VAL_END_STATE_NONE) + logger.info(f"[SIM_PENDING_COMMIT] released {_end[-4:]} sct={sct:.1f}") + return m, md + def _aggregate_weights(self, tag: str) -> None: """Aggregate local model weights asynchronously. @@ -185,12 +253,24 @@ def _aggregate_weights(self, tag: str) -> None: logger.debug("No channel found") return - logger.debug(f"Channel {channel} found for tag {tag}") - # receive local model parameters from a trainer who arrives - # first NOTE: (DG) Right now, the leave notifications also - # cause a message to be processed and yield (None,None) from - # recv_fifo(). - msg, metadata = next(channel.recv_fifo(channel.ends(VAL_CH_STATE_RECV), 1)) + # Filter to live ends; drop ghosts that left after selection to avoid + # blocking recv_fifo on an empty queue. + recv_ends = channel.ends(VAL_CH_STATE_RECV) + if recv_ends: + recv_ends = [e for e in recv_ends if channel.has(e)] + if not recv_ends: + if self.simulated and len(self._sim_buffer) > 0: + recv_ends = [] # buffer still has entries to drain — don't block + else: + logger.debug(f"[AGG_RECV] no live recv ends (round={self._round}); skipping") + time.sleep(0.5) + return + if self.simulated: + msg, metadata = self._sim_recv_min(channel, recv_ends) + else: + msg, metadata = next( + channel.recv_fifo(recv_ends, 1, timeout=RECV_TIMEOUT_WAIT_S) + ) end, _ = metadata if not msg: logger.debug(f"[AGG_RECV] No data from {end}; skipping it, agg_model_version={self._round}") @@ -349,6 +429,18 @@ def _aggregate_weights(self, tag: str) -> None: recv_wts_version ] = recv_wts_ts + wall_lag_s = (recv_wts_ts - sent_wts_ts).total_seconds() + logger.info( + f"[SEND_RECV_LAG] end={end} version={recv_wts_version} " + f"wall_lag_s={wall_lag_s:.3f}" + ) + _lag_warn_threshold_s = 30.0 if not self.simulated else 10.0 + if wall_lag_s > _lag_warn_threshold_s: + logger.warning( + f"[SEND_RECV_LAG_HIGH] end={end} version={recv_wts_version} " + f"wall_lag_s={wall_lag_s:.1f}s — possible MQTT backlog" + ) + # TODO: (DG) Can pass a flag for this later. allow_updates_more_than_timeout_old = True @@ -407,7 +499,17 @@ def _aggregate_weights(self, tag: str) -> None: curr_cumulative_training_s = self._track_trainer_version_duration_s[ end ]["total_training_time_s"] - curr_round_time_s = (recv_wts_ts - sent_wts_ts).total_seconds() + # round duration drives Oort's speed/system-utility. simulated + # mode: use the trainer-reported modeled duration D (the real + # recv-send delta is ~GPU time and misrepresents speed). real + # mode: the wall-clock delta (the sleep makes it ~= D). + if self.simulated: + round_duration_td = timedelta( + seconds=float(msg.get(MessageType.SIM_ROUND_DURATION, 0.0)) + ) + else: + round_duration_td = recv_wts_ts - sent_wts_ts + curr_round_time_s = round_duration_td.total_seconds() new_cumulative_training_s = ( curr_cumulative_training_s + curr_round_time_s ) @@ -427,11 +529,10 @@ def _aggregate_weights(self, tag: str) -> None: # version to that trainer. logger.debug( f"Setting channel property {PROP_ROUND_DURATION} for " - f"end {end} with duration " - f"{recv_wts_ts - sent_wts_ts}" + f"end {end} with duration {round_duration_td}" ) channel.set_end_property( - end, PROP_ROUND_DURATION, recv_wts_ts - sent_wts_ts + end, PROP_ROUND_DURATION, round_duration_td ) # capture telemetry on trainer participation in rounds @@ -492,11 +593,23 @@ def _aggregate_weights(self, tag: str) -> None: # Populate round statistics vars self._round_update_values["staleness"].append(update_staleness_val) self._round_update_values["stat_utility"].append(stat_utility) - self._round_update_values["trainer_speed"].append( - channel.get_end_property( - end_id=end, key=PROP_ROUND_DURATION - ).total_seconds() - ) + _round_dur = channel.get_end_property(end_id=end, key=PROP_ROUND_DURATION) + _trainer_speed_s = _round_dur.total_seconds() if _round_dur is not None else 0.0 + self._round_update_values["trainer_speed"].append(_trainer_speed_s) + + # Telemetry: one record per processed train update. + if telemetry.is_enabled(): + ev, fields = build_agg_round( + round_num=self._round, + agg_goal=self._agg_goal, + agg_goal_count=self._agg_goal_cnt, + updates_in_queue=self._updates_in_queue, + staleness=[update_staleness_val], + stat_utility=[stat_utility], + trainer_speed_s=[_trainer_speed_s], + contributing_trainers=[end], + ) + telemetry.emit(ev, **fields) # capture per trainer staleness if end in self._per_trainer_staleness_track.keys(): @@ -663,6 +776,35 @@ def _aggregate_weights(self, tag: str) -> None: logger.debug("Agg goal reached, so resetting trainer end states in the channel") channel.cleanup_recvd_ends() + if self.simulated: + sel = channel._selector + requester = sel.requester + pending_in_buffer = set(self._sim_buffer.pending_ends()) + + # Release all_selected trainers with no buffer entry yet (GPU still + # running — rare). They'll be probed next round's fill pass. + for end_id in [e for e in list(sel.all_selected.keys()) if e not in pending_in_buffer]: + del sel.all_selected[end_id] + if channel.has(end_id): + channel._ends[end_id].set_property(KEY_END_STATE, VAL_END_STATE_NONE) + if requester in sel.selected_ends and end_id in sel.selected_ends[requester]: + sel.selected_ends[requester].discard(end_id) + + # Block every trainer with a pending buffer entry from re-selection. + # cleanup_recvd_ends may have freed them early; re-block here so the + # real-mode invariant holds: one in-flight update per trainer at a time. + for end_id in pending_in_buffer: + self._sim_pending_commit.add(end_id) + if end_id not in sel.all_selected: + sel.all_selected[end_id] = time.time() + if requester in sel.selected_ends and end_id in sel.selected_ends[requester]: + sel.selected_ends[requester].discard(end_id) + if pending_in_buffer: + logger.info( + f"[SIM_PENDING] round={self._round} blocked {len(pending_in_buffer)} " + f"trainer(s) pending buffer commit: {[e[-4:] for e in pending_in_buffer]}" + ) + def oracular_trainer_avail_check(self, end: str) -> bool: logger.debug("In oracular_trainer_avail_check") @@ -922,8 +1064,11 @@ def _distribute_weights(self, tag: str, task_to_perform: str = "train") -> None: # Add small delay between sends to distribute MQTT broker load # This prevents overwhelming the broker with many concurrent large messages # and allows the event loop to process keepalive packets - if idx < len(ends_list) - 1: # Don't sleep after last send - time.sleep(0.5) + # Stagger sends to avoid overwhelming the MQTT broker with concurrent + # large weight payloads. Use a shorter interval in simulated mode since + # trainers cycle faster and broker load is lower. + if idx < len(ends_list) - 1: + time.sleep(0.2 if self.simulated else 0.5) def compose(self) -> None: """Compose role with tasklets.""" diff --git a/lib/python/flame/mode/horizontal/syncfl/top_aggregator.py b/lib/python/flame/mode/horizontal/syncfl/top_aggregator.py index d054ec6fc..bd98a526a 100644 --- a/lib/python/flame/mode/horizontal/syncfl/top_aggregator.py +++ b/lib/python/flame/mode/horizontal/syncfl/top_aggregator.py @@ -50,6 +50,10 @@ PROP_ROUND_START_TIME, PROP_STAT_UTILITY, ) +from flame import telemetry +from flame.telemetry.events import build_agg_eval, build_agg_round +from flame.sim import VirtualClock +from flame.selector.properties import PROP_SIM_SEND_TS, PROP_SIM_COMPLETION_TS logger = logging.getLogger(__name__) @@ -119,6 +123,15 @@ def internal_init(self) -> None: self._rounds = self.config.hyperparameters.rounds self._work_done = False + # Simulation time mode (replaces speedup_factor). "simulated": order + # updates by a virtual clock fed by trainer-reported completion times; + # "real": order by physical arrival (legacy/authentic baseline). + self.time_mode = getattr( + self.config.hyperparameters, "time_mode", "simulated" + ) + self.simulated = self.time_mode == "simulated" + self._vclock = VirtualClock() + self.framework = get_ml_framework_in_use() if self.framework == MLFramework.UNKNOWN: raise NotImplementedError( @@ -303,6 +316,19 @@ def _aggregate_weights(self, tag: str) -> None: logger.debug(f"received {len(self.cache)} trainer updates in cache") + if telemetry.is_enabled(): + ev, fields = build_agg_round( + round_num=self._round, + in_flight=len(channel.ends()), + staleness=list(self._round_update_values.get("staleness", [])), + stat_utility=list(self._round_update_values.get("stat_utility", [])), + trainer_speed_s=list( + self._round_update_values.get("trainer_speed", []) + ), + contributing_trainers=list(self.cache.keys()), + ) + telemetry.emit(ev, **fields) + self._compute_aggregator_stats() if self._round % 5 == 0: logger.info(f"_agg_training_stats: {self._agg_training_stats}") @@ -362,18 +388,20 @@ def _distribute_weights(self, tag: str, task_to_perform: str = "train") -> None: logger.info( f"sending weights to {end} with model_version: {self._round} for task: {task_to_perform}" ) - channel.send( - end, - { - MessageType.WEIGHTS: weights_to_device( - self.weights, DeviceType.CPU - ), - MessageType.ROUND: self._round, - MessageType.DATASAMPLER_METADATA: datasampler_metadata, - MessageType.MODEL_VERSION: self._round, - MessageType.TASK_TO_PERFORM: task_to_perform, - }, - ) + msg = { + MessageType.WEIGHTS: weights_to_device(self.weights, DeviceType.CPU), + MessageType.ROUND: self._round, + MessageType.DATASAMPLER_METADATA: datasampler_metadata, + MessageType.MODEL_VERSION: self._round, + MessageType.TASK_TO_PERFORM: task_to_perform, + } + # simulated mode: stamp the virtual send time so the trainer can + # report sim_completion_ts = sim_send_ts + D back to us. + if self.simulated: + sim_send_ts = self._vclock.now + msg[MessageType.SIM_SEND_TS] = sim_send_ts + channel.set_end_property(end, PROP_SIM_SEND_TS, sim_send_ts) + channel.send(end, msg) # register round start time on each end for round duration # measurement. channel.set_end_property( @@ -456,6 +484,10 @@ def save_model(self): def update_metrics(self, metrics: dict[str, float]): """Update metrics.""" self.metrics = self.metrics | metrics + # Telemetry: aggregator eval metrics (generic hook for all examples). + if telemetry.is_enabled(): + ev, fields = build_agg_eval(round_num=self._round, metrics=metrics) + telemetry.emit(ev, **fields) def _update_model(self): if self.framework == MLFramework.PYTORCH: diff --git a/lib/python/flame/mode/horizontal/syncfl/trainer.py b/lib/python/flame/mode/horizontal/syncfl/trainer.py index 24e4573a6..0b63739b6 100644 --- a/lib/python/flame/mode/horizontal/syncfl/trainer.py +++ b/lib/python/flame/mode/horizontal/syncfl/trainer.py @@ -223,6 +223,12 @@ def _fetch_weights(self, tag: str) -> None: self.weights = weights_to_model_device(msg[MessageType.WEIGHTS], self.model) self._update_model() + # simulated-time mode: capture the virtual send time stamped by the + # aggregator on this task (used to compute sim_completion_ts). No-op + # for trainers/aggregators that don't use it. + if MessageType.SIM_SEND_TS in msg: + self._sim_send_ts = msg[MessageType.SIM_SEND_TS] + if MessageType.EOT in msg: self._work_done = msg[MessageType.EOT] @@ -361,6 +367,16 @@ def _send_weights(self, tag: str) -> None: MessageType.LOCAL_ACCURACY: self._local_accuracy, } + # simulated-time mode only: report the modeled completion time/duration + # so the aggregator orders this update by a virtual clock. No-op in real + # mode and for trainers that don't set these (other examples). + _sim_completion = getattr(self, "_sim_completion_ts", None) + if getattr(self, "simulated", False) and _sim_completion is not None: + msg[MessageType.SIM_COMPLETION_TS] = _sim_completion + msg[MessageType.SIM_ROUND_DURATION] = getattr( + self, "_sim_round_duration", 0.0 + ) + channel.send(end, msg) if self.task_to_perform == "train": diff --git a/lib/python/flame/mode/message.py b/lib/python/flame/mode/message.py index ed66788ab..cb2f57c2b 100644 --- a/lib/python/flame/mode/message.py +++ b/lib/python/flame/mode/message.py @@ -74,3 +74,11 @@ class MessageType(Enum): JVP_FOR_SNR_CHECK = 32 LOCAL_ACCURACY = 33 # trainer's local training accuracy (used by FedDance A_m) + + # Simulated-time fields (time_mode="simulated"): aggregator stamps the + # virtual send time on the distributed task; trainer echoes its modeled + # completion time and duration on the update. Used to order updates by a + # virtual clock instead of physical arrival. + SIM_SEND_TS = 34 # virtual clock T_v at weight distribution + SIM_COMPLETION_TS = 35 # sim_send_ts + real_gpu_time + modeled delay D + SIM_ROUND_DURATION = 36 # real_gpu_time + modeled delay D (matches real-mode recv_ts - sent_ts) diff --git a/lib/python/flame/selector/__init__.py b/lib/python/flame/selector/__init__.py index 2556f1a1c..f08286dfc 100644 --- a/lib/python/flame/selector/__init__.py +++ b/lib/python/flame/selector/__init__.py @@ -16,12 +16,19 @@ """selector abstract class.""" from abc import ABC, abstractmethod -from typing import Tuple, Union +from typing import Optional, Tuple, Union import logging import time +from .. import telemetry from ..common.typing import Scalar from ..end import End +from ..telemetry.events import build_selection +from .properties import ( + PROP_AVL_STATE, + PROP_ROUND_DURATION, + PROP_STAT_UTILITY, +) SelectorReturnType = dict[str, Union[None, Tuple[str, Scalar]]] @@ -70,6 +77,81 @@ def select( used/created during selection process; value can be none """ + def emit_selection( + self, + round_num: int, + task: str, + ends: dict[str, End], + eligible_ids, + chosen_ids, + per_trainer_extra: Optional[dict] = None, + extra: Optional[dict] = None, + ) -> None: + """Emit a structured selector-decision event (no-op if telemetry off). + + Centralized here so every selector produces an identical schema, which + is what makes cross-selector comparison possible. ``ends`` is the full + candidate pool; availability composition and per-trainer utility/speed + are derived from end properties. + """ + if not telemetry.is_enabled(): + return + try: + chosen_set = set(chosen_ids) + avail_composition: dict[str, int] = {} + per_trainer: dict[str, dict] = {} + for end_id, end in ends.items(): + state = end.get_property(PROP_AVL_STATE) + state_name = getattr(state, "value", None) or ( + str(state) if state is not None else "UNKNOWN" + ) + avail_composition[state_name] = ( + avail_composition.get(state_name, 0) + 1 + ) + util = end.get_property(PROP_STAT_UTILITY) + speed = end.get_property(PROP_ROUND_DURATION) + entry = { + "utility": util, + "speed_s": speed.total_seconds() + if hasattr(speed, "total_seconds") + else speed, + "selected": end_id in chosen_set, + } + if per_trainer_extra and end_id in per_trainer_extra: + entry.update(per_trainer_extra[end_id]) + per_trainer[end_id] = entry + + # in-flight count: selected_ends is a set/list for most selectors, + # but a {requester: set(ends)} dict for fedbuff-style selectors. + sel = self.selected_ends + if isinstance(sel, dict): + vals = list(sel.values()) + in_flight = ( + sum(len(v) for v in vals) + if vals and all(isinstance(v, (set, list)) for v in vals) + else len(sel) + ) + elif isinstance(sel, (set, list)): + in_flight = len(sel) + else: + in_flight = 0 + + ev, fields = build_selection( + round_num=int(round_num), + task=task, + selector=type(self).__name__, + num_candidates=len(ends), + num_eligible=len(set(eligible_ids)), + avail_composition=avail_composition, + chosen=list(chosen_set), + in_flight=in_flight, + per_trainer=per_trainer, + extra=extra, + ) + telemetry.emit(ev, **fields) + except Exception as e: # telemetry must never break selection + logger.debug(f"emit_selection failed: {e}") + def on_update_received( self, end_id: str, msg: dict, round_num: int ) -> None: diff --git a/lib/python/flame/selector/async_oort.py b/lib/python/flame/selector/async_oort.py index a5569a5db..a0cb15005 100644 --- a/lib/python/flame/selector/async_oort.py +++ b/lib/python/flame/selector/async_oort.py @@ -389,6 +389,19 @@ def select( ) self._select_run_counter = 0 + self.emit_selection( + channel_props.get("round", 0), + task_to_perform, + ends, + eligible_ends.keys(), + list(results.keys()), + extra={ + "concurrency": concurrency, + "effective_c": effective_c, + "requester": self.requester, + }, + ) + elif channel_props[KEY_CH_STATE] == VAL_CH_STATE_RECV: # TODO: (DG) See if eligible_ends should be passed here # too in place of ends @@ -1042,6 +1055,16 @@ def _cleanup_removed_ends(self, end_id): f"Need to check" ) + # A departed end must not linger in selected_ends (the in-flight set + # returned for recv), else the aggregator waits on a gone trainer and + # wastes a concurrency slot. selected_ends is {requester: set(ends)}. + for _req, _ends in self.selected_ends.items(): + if end_id in _ends: + _ends.discard(end_id) + logger.debug( + f"Removed ghost end_id {end_id} from selected_ends[{_req}]" + ) + # Invoked when selection mode is default i.e. of oort which trades # off exploitation/exploration and speed/stat_utility def _select_candidates_using_default( diff --git a/lib/python/flame/selector/fedbuff.py b/lib/python/flame/selector/fedbuff.py index 072e376f1..4c8e62b65 100644 --- a/lib/python/flame/selector/fedbuff.py +++ b/lib/python/flame/selector/fedbuff.py @@ -196,6 +196,14 @@ def select( results = {} if channel_props[KEY_CH_STATE] == VAL_CH_STATE_SEND: results = self._handle_send_state(eligible_ends, concurrency) + self.emit_selection( + channel_props.get("round", 0), + "train", + ends, + eligible_ends.keys(), + list(results.keys()), + extra={"concurrency": concurrency, "requester": self.requester}, + ) elif channel_props[KEY_CH_STATE] == VAL_CH_STATE_RECV: # TODO: (DG) See if eligible_ends should be passed here diff --git a/lib/python/flame/selector/feddance.py b/lib/python/flame/selector/feddance.py index 551338a17..93eb3784d 100644 --- a/lib/python/flame/selector/feddance.py +++ b/lib/python/flame/selector/feddance.py @@ -146,6 +146,10 @@ def select( ) self.round = round_num + self.emit_selection( + round_num, task_to_perform, ends, eligible.keys(), selected, + extra={"num_unavail": len(unavail)}, + ) return {key: None for key in selected} def _compute_utilities( diff --git a/lib/python/flame/selector/oort.py b/lib/python/flame/selector/oort.py index 99d0c2459..d5b9d709c 100644 --- a/lib/python/flame/selector/oort.py +++ b/lib/python/flame/selector/oort.py @@ -167,6 +167,9 @@ def select( f"let's select {num_of_ends} ends for new round {round}, task: {task_to_perform}" ) + # full candidate pool, captured before any filtering for telemetry + all_ends = dict(ends) + if round <= self.round and len(self.selected_ends) != 0: return {key: None for key in self.selected_ends} @@ -219,7 +222,12 @@ def select( # been measured; Then, perform random selection if len(utility_list) == 0 and len(self.selected_ends) == 0: self.round = round - return self.select_random(ends, num_of_ends) + result = self.select_random(ends, num_of_ends) + self.emit_selection( + round, task_to_perform, all_ends, ends.keys(), + self.selected_ends, extra={"mode": "random_first_round"}, + ) + return result # Not the first round, performing Oort-based selection # Calculate number of ends to select for exploration and @@ -233,7 +241,12 @@ def select( if len(utility_list) == 0: self.round = round - return self.select_random(ends, num_of_ends) + result = self.select_random(ends, num_of_ends) + self.emit_selection( + round, task_to_perform, all_ends, ends.keys(), + self.selected_ends, extra={"mode": "random_no_utility"}, + ) + return result utility_list = self.calculate_total_utility(utility_list, ends, round) cutoff_utility = self.cutoff_util(utility_list, num_of_ends) @@ -285,6 +298,15 @@ def select( ) self._select_run_counter = 0 + self.emit_selection( + round, task_to_perform, all_ends, eligible_ends.keys(), + self.selected_ends, + extra={ + "exploration_factor": self.exploration_factor, + "explore_ids": list(explore_end_ids), + "exploit_ids": list(exploit_end_ids), + }, + ) return {key: None for key in self.selected_ends} def cutoff_util( diff --git a/lib/python/flame/selector/properties.py b/lib/python/flame/selector/properties.py index 349f5f09b..1b5cb2f7a 100644 --- a/lib/python/flame/selector/properties.py +++ b/lib/python/flame/selector/properties.py @@ -15,6 +15,11 @@ PROP_ROUND_END_TIME = "round_end_time" PROP_ROUND_DURATION = "round_duration" +# Simulated time (time_mode="simulated"): trainer-reported virtual completion +# time, used to order updates by a virtual clock and to source round duration. +PROP_SIM_SEND_TS = "sim_send_ts" +PROP_SIM_COMPLETION_TS = "sim_completion_ts" + # Training metadata PROP_DATASET_SIZE = "dataset_size" PROP_STAT_UTILITY = "stat_utility" # I_m / Oort utility diff --git a/lib/python/flame/sim/__init__.py b/lib/python/flame/sim/__init__.py new file mode 100644 index 000000000..85f568cf6 --- /dev/null +++ b/lib/python/flame/sim/__init__.py @@ -0,0 +1,7 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Simulated-time primitives for ``time_mode="simulated"``.""" + +from .virtual_clock import SimReorderBuffer, VirtualClock, sim_ordered_ends + +__all__ = ["VirtualClock", "SimReorderBuffer", "sim_ordered_ends"] diff --git a/lib/python/flame/sim/virtual_clock.py b/lib/python/flame/sim/virtual_clock.py new file mode 100644 index 000000000..96d1113ad --- /dev/null +++ b/lib/python/flame/sim/virtual_clock.py @@ -0,0 +1,97 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Virtual (simulated) clock + reorder buffer for simulated time mode. + +These are pure, dependency-free helpers so they can be unit-tested in +isolation. The async aggregator composes them to order trainer updates by a +*simulated* completion time instead of physical arrival, which is what makes +``simulated`` mode reproduce ``real`` mode's per-round behavior without sleeping. +""" + +from __future__ import annotations + +from typing import Any, Optional + + +class VirtualClock: + """Monotone simulated clock measured in simulated seconds. + + The aggregator stamps `now` on each distributed task (`sim_send_ts`) and + advances to a committed update's `sim_completion_ts`. It never moves + backward, so committing updates in ascending completion order yields a + monotone timeline. + """ + + def __init__(self) -> None: + self._t: float = 0.0 + + @property + def now(self) -> float: + return self._t + + def advance(self, ts: float) -> float: + """Advance to ``ts`` if it is in the future; return the new time.""" + ts = float(ts) + if ts > self._t: + self._t = ts + return self._t + + def reset(self) -> None: + self._t = 0.0 + + +def sim_ordered_ends(end_to_completion: dict[str, float]) -> list[str]: + """Return ends sorted by ascending simulated completion time. + + Ties broken by end id for determinism (so ordering is reproducible across + runs and independent of dict/arrival order). + """ + return sorted(end_to_completion, key=lambda e: (end_to_completion[e], str(e))) + + +class SimReorderBuffer: + """Collects arrived-but-uncommitted updates and pops them in simulated + completion order. + + In simulated mode trainers do not sleep, so all in-flight updates land in a + tiny physical window; buffering them and popping the minimum + `sim_completion_ts` reconstructs the order they *would* have arrived in real + mode — independent of physical arrival jitter. + """ + + def __init__(self) -> None: + # end_id -> (sim_completion_ts, payload) + self._items: dict[str, tuple[float, Any]] = {} + + def add(self, end_id: str, sim_completion_ts: float, payload: Any = None) -> None: + self._items[end_id] = (float(sim_completion_ts), payload) + + def has(self, end_id: str) -> bool: + return end_id in self._items + + def pending_ends(self) -> set[str]: + return set(self._items) + + def __len__(self) -> int: + return len(self._items) + + def peek_min_ts(self) -> Optional[float]: + if not self._items: + return None + return min(ts for ts, _ in self._items.values()) + + def pop_min(self) -> Optional[tuple[str, float, Any]]: + """Remove and return ``(end_id, sim_completion_ts, payload)`` with the + smallest completion time (ties by end id). ``None`` if empty.""" + if not self._items: + return None + end_id = min(self._items, key=lambda e: (self._items[e][0], str(e))) + ts, payload = self._items.pop(end_id) + return end_id, ts, payload + + def discard(self, end_id: str) -> None: + self._items.pop(end_id, None) + + def clear(self) -> None: + """Discard all buffered entries (e.g. at round end to prevent cross-round stale commits).""" + self._items.clear() diff --git a/lib/python/flame/telemetry/__init__.py b/lib/python/flame/telemetry/__init__.py new file mode 100644 index 000000000..51a1bf9cc --- /dev/null +++ b/lib/python/flame/telemetry/__init__.py @@ -0,0 +1,171 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Structured, process-local telemetry for FLAME experiments. + +Each process (aggregator, each trainer) writes typed JSONL event records to a +per-run directory so post-run analysis can compare selectors / aggregators / +trainers on an identical schema -- no log-regex scraping. + +The writer is a process-global singleton configured once at startup. When no +run directory is configured (neither ``configure(run_dir=...)`` nor the +``FLAME_TELEMETRY_DIR`` env var), every :func:`emit` is a cheap no-op, so +existing runs and unit tests are unaffected unless telemetry is explicitly +enabled. + +Generic by design: nothing here is example-specific. The async_cifar10 example +wires it first, but any example/aggregator/selector/trainer inherits the same +hooks. +""" + +from __future__ import annotations + +import datetime as _dt +import json +import logging +import os +import threading +import time +from typing import Any, Optional + +from .events import ( # noqa: F401 (re-exported for convenience) + EVENT_AGG_EVAL, + EVENT_AGG_ROUND, + EVENT_AVAIL_CHANGE, + EVENT_RUN_META, + EVENT_SELECTION, + EVENT_TRAINER_ROUND, + EVENT_UTIL_DISPARITY, + KNOWN_EVENTS, +) + +logger = logging.getLogger(__name__) + +ENV_DIR = "FLAME_TELEMETRY_DIR" + +_writer: Optional["TelemetryWriter"] = None +_config_lock = threading.Lock() + + +def _json_default(obj: Any) -> Any: + """Best-effort JSON encoder for the value types telemetry carries.""" + if isinstance(obj, (set, frozenset)): + return sorted(obj, key=str) + if isinstance(obj, _dt.timedelta): + return obj.total_seconds() + if isinstance(obj, _dt.datetime): + return obj.isoformat() + # numpy scalars / arrays without importing numpy + if hasattr(obj, "item") and not hasattr(obj, "__len__"): + try: + return obj.item() + except Exception: + pass + if hasattr(obj, "tolist"): + try: + return obj.tolist() + except Exception: + pass + return str(obj) + + +class TelemetryWriter: + """Append-only JSONL writer. Thread-safe (trainers emit from threads).""" + + def __init__(self, path: str, role: str, end_id: Optional[str] = None) -> None: + self.path = path + self.role = role + self.end_id = end_id + self._lock = threading.Lock() + # line-buffered so a killed process still leaves complete records + self._fh = open(path, "a", buffering=1) + + def emit(self, event: str, **fields: Any) -> None: + rec: dict[str, Any] = { + "ts": time.time(), + "role": self.role, + "event": event, + } + if self.end_id is not None: + rec["end_id"] = self.end_id + rec.update(fields) + line = json.dumps(rec, default=_json_default) + with self._lock: + self._fh.write(line + "\n") + + def close(self) -> None: + with self._lock: + try: + self._fh.flush() + self._fh.close() + except Exception: + pass + + +def configure( + role: str, + end_id: Optional[str] = None, + run_dir: Optional[str] = None, + filename: Optional[str] = None, +) -> Optional[TelemetryWriter]: + """Initialize the process-global telemetry writer. + + Parameters + ---------- + role: short role label, e.g. "aggregator" or "trainer". + end_id: per-process identity (e.g. trainer/task id); used in the filename. + run_dir: output directory. Falls back to ``$FLAME_TELEMETRY_DIR``. If still + unset, telemetry stays disabled and this returns ``None``. + filename: override the default ``[_].jsonl`` filename. + + Returns the writer, or ``None`` if telemetry is disabled. + """ + global _writer + run_dir = run_dir or os.environ.get(ENV_DIR) + if not run_dir: + return None + + if filename is None: + safe_id = str(end_id).replace("/", "_") if end_id is not None else None + filename = f"{role}_{safe_id}.jsonl" if safe_id is not None else f"{role}.jsonl" + + try: + os.makedirs(run_dir, exist_ok=True) + path = os.path.join(run_dir, filename) + with _config_lock: + if _writer is not None: + _writer.close() + _writer = TelemetryWriter(path, role, end_id) + logger.info(f"telemetry enabled: writing to {path}") + return _writer + except Exception as e: # never let telemetry setup break a run + logger.warning(f"telemetry disabled (configure failed): {e}") + _writer = None + return None + + +def emit(event: str, **fields: Any) -> None: + """Emit one event. No-op when telemetry is disabled. Never raises.""" + w = _writer + if w is None: + return + try: + w.emit(event, **fields) + except Exception as e: # telemetry must never crash the experiment + logger.debug(f"telemetry emit failed for {event}: {e}") + + +def is_enabled() -> bool: + return _writer is not None + + +def get_run_dir() -> Optional[str]: + w = _writer + return os.path.dirname(w.path) if w is not None else None + + +def shutdown() -> None: + global _writer + with _config_lock: + if _writer is not None: + _writer.close() + _writer = None diff --git a/lib/python/flame/telemetry/events.py b/lib/python/flame/telemetry/events.py new file mode 100644 index 000000000..cea3c3a8d --- /dev/null +++ b/lib/python/flame/telemetry/events.py @@ -0,0 +1,202 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Telemetry event types and typed builders. + +Centralizing the event names + field schemas here is what makes runs +comparable across selectors / aggregators / examples. Emitters should call the +``build_*`` helpers (or :func:`flame.telemetry.emit` with these constants) so +every run produces the same columns. +""" + +from __future__ import annotations + +from typing import Any, Optional + +# ---- Event type constants ------------------------------------------------- + +EVENT_RUN_META = "run_meta" # one-time run identity / config snapshot +EVENT_SELECTION = "selection" # selector decision for a round +EVENT_AGG_EVAL = "agg_eval" # aggregator test loss/accuracy +EVENT_AGG_ROUND = "agg_round" # aggregation step: staleness/agg-goal/participation +EVENT_TRAINER_ROUND = "trainer_round" # per-round trainer timing/availability +EVENT_UTIL_DISPARITY = "util_disparity" # streamed-prefix vs full-pool utility +EVENT_AVAIL_CHANGE = "avail_change" # trainer availability state transition + +KNOWN_EVENTS = frozenset( + { + EVENT_RUN_META, + EVENT_SELECTION, + EVENT_AGG_EVAL, + EVENT_AGG_ROUND, + EVENT_TRAINER_ROUND, + EVENT_UTIL_DISPARITY, + EVENT_AVAIL_CHANGE, + } +) + + +# ---- Typed builders ------------------------------------------------------- +# Each returns (event_type, fields_dict). Callers do: +# telemetry.emit(*build_selection(...)) -> emit(event, **fields) +# but emit() takes (event, **fields), so callers use: +# ev, f = build_selection(...); telemetry.emit(ev, **f) + + +def build_selection( + *, + round_num: int, + task: str, + selector: str, + num_candidates: int, + num_eligible: int, + avail_composition: dict[str, int], + chosen: list[str], + in_flight: int, + per_trainer: Optional[dict[str, dict[str, Any]]] = None, + extra: Optional[dict[str, Any]] = None, +) -> tuple[str, dict[str, Any]]: + """Selector decision record. + + avail_composition: counts keyed by availability-state name + (e.g. {"AVL_TRAIN": 30, "AVL_EVAL": 5, "UN_AVL": 65}). + per_trainer: optional {end_id: {"utility": .., "speed_s": .., "selected": bool}}. + extra: selector-specific fields (e.g. explore/exploit split, cutoff). + """ + fields: dict[str, Any] = { + "round": round_num, + "task": task, + "selector": selector, + "num_candidates": num_candidates, + "num_eligible": num_eligible, + "avail_composition": avail_composition, + "chosen": list(chosen), + "num_chosen": len(chosen), + "in_flight": in_flight, + } + if per_trainer is not None: + fields["per_trainer"] = per_trainer + if extra: + fields.update(extra) + return EVENT_SELECTION, fields + + +def build_agg_eval( + *, round_num: int, metrics: dict[str, float] +) -> tuple[str, dict[str, Any]]: + """Aggregator evaluation metrics (loss/accuracy/...).""" + fields = {"round": round_num} + fields.update(metrics) + return EVENT_AGG_EVAL, fields + + +def build_agg_round( + *, + round_num: int, + agg_goal: Optional[int] = None, + agg_goal_count: Optional[int] = None, + in_flight: Optional[int] = None, + updates_in_queue: Optional[int] = None, + staleness: Optional[list[float]] = None, + stat_utility: Optional[list[float]] = None, + trainer_speed_s: Optional[list[float]] = None, + contributing_trainers: Optional[list[str]] = None, + extra: Optional[dict[str, Any]] = None, +) -> tuple[str, dict[str, Any]]: + """Aggregation-step record (one per completed aggregation).""" + fields: dict[str, Any] = {"round": round_num} + for k, v in ( + ("agg_goal", agg_goal), + ("agg_goal_count", agg_goal_count), + ("in_flight", in_flight), + ("updates_in_queue", updates_in_queue), + ("staleness", staleness), + ("stat_utility", stat_utility), + ("trainer_speed_s", trainer_speed_s), + ("contributing_trainers", contributing_trainers), + ): + if v is not None: + fields[k] = v + if extra: + fields.update(extra) + return EVENT_AGG_ROUND, fields + + +def build_trainer_round( + *, + round_num: int, + real_gpu_time_s: float, + sim_round_duration_s: Optional[float] = None, + wait_time_s: Optional[float] = None, + avail_state: Optional[str] = None, + visible_samples: Optional[int] = None, + total_samples: Optional[int] = None, + dataset_size: Optional[int] = None, + stat_utility: Optional[float] = None, + final_loss: Optional[float] = None, + extra: Optional[dict[str, Any]] = None, +) -> tuple[str, dict[str, Any]]: + """Per-round trainer timing/availability record.""" + fields: dict[str, Any] = { + "round": round_num, + "real_gpu_time_s": real_gpu_time_s, + } + for k, v in ( + ("sim_round_duration_s", sim_round_duration_s), + ("wait_time_s", wait_time_s), + ("avail_state", avail_state), + ("visible_samples", visible_samples), + ("total_samples", total_samples), + ("dataset_size", dataset_size), + ("stat_utility", stat_utility), + ("final_loss", final_loss), + ): + if v is not None: + fields[k] = v + if extra: + fields.update(extra) + return EVENT_TRAINER_ROUND, fields + + +def build_util_disparity( + *, + round_num: int, + elapsed_s: float, + visible_samples: int, + total_samples: int, + utility_streamed: float, + utility_full: float, + sample_size_used: Optional[int] = None, +) -> tuple[str, dict[str, Any]]: + """Streamed-prefix vs full-dataset statistical-utility comparison.""" + visible_fraction = ( + visible_samples / total_samples if total_samples else None + ) + ratio = ( + utility_streamed / utility_full + if utility_full not in (0, None) + else None + ) + fields: dict[str, Any] = { + "round": round_num, + "elapsed_s": elapsed_s, + "visible_samples": visible_samples, + "total_samples": total_samples, + "visible_fraction": visible_fraction, + "utility_streamed": utility_streamed, + "utility_full": utility_full, + "utility_ratio": ratio, + } + if sample_size_used is not None: + fields["sample_size_used"] = sample_size_used + return EVENT_UTIL_DISPARITY, fields + + +def build_avail_change( + *, round_num: Optional[int], old_state: str, new_state: str +) -> tuple[str, dict[str, Any]]: + """Trainer availability state transition.""" + return EVENT_AVAIL_CHANGE, { + "round": round_num, + "old_state": old_state, + "new_state": new_state, + } diff --git a/nonEvalSelectorTrainerLogs.log b/lib/python/tests/mode/__init__.py similarity index 100% rename from nonEvalSelectorTrainerLogs.log rename to lib/python/tests/mode/__init__.py diff --git a/lib/python/tests/mode/test_async_sim_ordering.py b/lib/python/tests/mode/test_async_sim_ordering.py new file mode 100644 index 000000000..18fa7462b --- /dev/null +++ b/lib/python/tests/mode/test_async_sim_ordering.py @@ -0,0 +1,134 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Aggregator-level test for simulated-mode receive ordering. + +Exercises the real ``TopAggregator._sim_recv_min`` against a fake channel to +prove that, regardless of the physical arrival order, the async aggregator +commits in-flight updates in ascending virtual-completion order and advances +its virtual clock accordingly. +""" + +import itertools + +import pytest + +from flame.mode.horizontal.asyncfl.top_aggregator import TopAggregator +from flame.mode.message import MessageType +from flame.sim import SimReorderBuffer, VirtualClock + + +class FakeChannel: + """Minimal channel: queued (end -> msg) delivered in a fixed arrival order. + + Uses a list-based queue that is consumed eagerly (not via a generator) + so that calling ``next(recv_fifo(...))`` removes the message immediately. + Mirrors the one-message-per-call pattern _sim_recv_min uses. + """ + + def __init__(self, inflight, arrival_order): + self._inflight = set(inflight) + self._queue = list(arrival_order) # list of (end_id, sct) + + def has(self, end_id): + return end_id in self._inflight + + def ends(self, state=None): + return list(self._inflight) + + def recv_fifo(self, end_ids, first_k=0, timeout=None): + # Find and remove the first queued message whose end is in end_ids, + # then yield it. The eager pop means next() leaves the queue consistent. + end_ids = set(end_ids) + for i, (end_id, sct) in enumerate(self._queue): + if end_id in end_ids: + self._queue.pop(i) + yield ( + {MessageType.WEIGHTS: f"w_{end_id}", + MessageType.SIM_COMPLETION_TS: sct}, + (end_id, None), + ) + return + # Nothing found: yield nothing (caller gets StopIteration on next()) + return + yield # make this a generator + + +class _ConcreteAgg(TopAggregator): + """Concrete stub so we can instantiate without the abstract methods.""" + + def check_and_sleep(self): + pass + + def evaluate(self): + pass + + def initialize(self): + pass + + def load_data(self): + pass + + def train(self): + pass + + +def _make_agg(): + """A TopAggregator with only the state _sim_recv_min needs.""" + agg = _ConcreteAgg.__new__(_ConcreteAgg) + agg._vclock = VirtualClock() + agg._sim_buffer = SimReorderBuffer() + agg._sim_committed = set() + return agg + + +def _drain(agg, channel): + """Repeatedly call _sim_recv_min until all messages committed. + + Each call: fill buffer from all currently-receivable ends (one probe each), + then pop and commit the minimum. The test provides instantaneous delivery + (no real timeout needed), so we keep calling until the queue and buffer + are both empty. + """ + committed = [] + max_iters = (len(channel.ends()) + 1) * 3 + for _ in range(max_iters): + recv_ends = [e for e in channel.ends() if channel.has(e)] + msg, (end, _) = agg._sim_recv_min(channel, recv_ends) + if msg is not None: + committed.append((end, msg[MessageType.SIM_COMPLETION_TS])) + if not channel._queue and not agg._sim_buffer.pending_ends(): + break + return committed, agg._vclock.now + + +class TestSimRecvMin: + def test_commits_in_completion_order(self): + durations = {"t1": 10.0, "t2": 5.0, "t3": 25.0, "t4": 15.0} + expected = [("t2", 5.0), ("t1", 10.0), ("t4", 15.0), ("t3", 25.0)] + agg = _make_agg() + # arrival order deliberately scrambled vs completion order + channel = FakeChannel( + inflight=set(durations), + arrival_order=[("t3", 25.0), ("t1", 10.0), ("t4", 15.0), ("t2", 5.0)], + ) + committed, t_v = _drain(agg, channel) + assert committed == expected + assert t_v == 25.0 # advanced to the last committed completion + + def test_independent_of_arrival_order(self): + durations = {"a": 3.0, "b": 1.0, "c": 2.0, "d": 4.0} + expected = [("b", 1.0), ("c", 2.0), ("a", 3.0), ("d", 4.0)] + for arrival in itertools.permutations(durations.items()): + agg = _make_agg() + channel = FakeChannel(set(durations), list(arrival)) + committed, _ = _drain(agg, channel) + assert committed == expected + + def test_virtual_clock_monotone(self): + durations = {"x": 8.0, "y": 2.0, "z": 5.0} + agg = _make_agg() + channel = FakeChannel(set(durations), list(durations.items())) + committed, t_v = _drain(agg, channel) + times = [c for _, c in committed] + assert times == sorted(times) # committed in nondecreasing sim time + assert t_v == max(times) diff --git a/lib/python/tests/sim/__init__.py b/lib/python/tests/sim/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/lib/python/tests/sim/test_virtual_clock.py b/lib/python/tests/sim/test_virtual_clock.py new file mode 100644 index 000000000..6387ed3b5 --- /dev/null +++ b/lib/python/tests/sim/test_virtual_clock.py @@ -0,0 +1,143 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the simulated-time primitives. + +These validate the ordering logic that makes simulated mode reproduce real +mode's per-round behavior: a monotone virtual clock, deterministic +completion-time ordering, and an arrival-order-independent reorder buffer. +""" + +import random + +import pytest + +from flame.sim import SimReorderBuffer, VirtualClock, sim_ordered_ends + + +class TestVirtualClock: + def test_starts_at_zero(self): + assert VirtualClock().now == 0.0 + + def test_advances_forward_only(self): + c = VirtualClock() + assert c.advance(5.0) == 5.0 + assert c.now == 5.0 + # a past timestamp must not move the clock backward + assert c.advance(3.0) == 5.0 + assert c.now == 5.0 + assert c.advance(7.5) == 7.5 + + def test_reset(self): + c = VirtualClock() + c.advance(10) + c.reset() + assert c.now == 0.0 + + +class TestSimOrderedEnds: + def test_ascending_by_completion(self): + order = sim_ordered_ends({"a": 30.0, "b": 10.0, "c": 20.0}) + assert order == ["b", "c", "a"] + + def test_ties_broken_by_end_id_deterministically(self): + m = {"t3": 5.0, "t1": 5.0, "t2": 5.0} + assert sim_ordered_ends(m) == ["t1", "t2", "t3"] + + def test_independent_of_insertion_order(self): + base = {"a": 3.0, "b": 1.0, "c": 2.0} + items = list(base.items()) + for _ in range(20): + random.shuffle(items) + assert sim_ordered_ends(dict(items)) == ["b", "c", "a"] + + +class TestSimReorderBuffer: + def test_pop_min_orders_by_completion(self): + buf = SimReorderBuffer() + buf.add("a", 30.0, "pa") + buf.add("b", 10.0, "pb") + buf.add("c", 20.0, "pc") + assert buf.pop_min() == ("b", 10.0, "pb") + assert buf.pop_min() == ("c", 20.0, "pc") + assert buf.pop_min() == ("a", 30.0, "pa") + assert buf.pop_min() is None + + def test_pop_order_independent_of_add_order(self): + durations = {"t0": 4.0, "t1": 1.0, "t2": 3.0, "t3": 2.0} + expected = ["t1", "t3", "t2", "t0"] # ascending by completion + for _ in range(20): + buf = SimReorderBuffer() + ends = list(durations) + random.shuffle(ends) # simulate arbitrary physical arrival order + for e in ends: + buf.add(e, durations[e], None) + popped = [] + while len(buf): + popped.append(buf.pop_min()[0]) + assert popped == expected + + def test_has_pending_discard(self): + buf = SimReorderBuffer() + buf.add("a", 1.0) + assert buf.has("a") and not buf.has("b") + assert buf.pending_ends() == {"a"} + buf.discard("a") + assert not buf.has("a") + assert len(buf) == 0 + + def test_peek_min_ts(self): + buf = SimReorderBuffer() + assert buf.peek_min_ts() is None + buf.add("a", 5.0) + buf.add("b", 2.0) + assert buf.peek_min_ts() == 2.0 + # peek does not remove + assert len(buf) == 2 + + def test_ties_broken_by_end_id(self): + buf = SimReorderBuffer() + buf.add("t2", 1.0) + buf.add("t1", 1.0) + assert buf.pop_min()[0] == "t1" + + +class TestStalenessReconstruction: + """A scripted scenario: committing buffered updates in completion order and + advancing a virtual clock reproduces a hand-computed staleness sequence, + independent of the order updates were buffered (arrival order).""" + + def test_staleness_matches_reference_independent_of_arrival(self): + # trainer -> (model_version_when_sent, modeled completion time) + # all sent at version 1; they complete at different sim times. + scenario = { + "t1": (1, 10.0), + "t2": (1, 5.0), + "t3": (1, 25.0), + "t4": (1, 15.0), + } + agg_goal = 2 # model version increments every 2 commits + + def run(arrival_order): + clock = VirtualClock() + buf = SimReorderBuffer() + for e in arrival_order: + buf.add(e, scenario[e][1], scenario[e][0]) + model_version = 1 + committed = 0 + staleness_seq = [] + while len(buf): + end, ts, sent_version = buf.pop_min() + clock.advance(ts) + staleness_seq.append((end, model_version - sent_version)) + committed += 1 + if committed % agg_goal == 0: + model_version += 1 + return staleness_seq + + # commit order is always by completion time: t2(5),t1(10),t4(15),t3(25) + # versions: t2,t1 at v1 (staleness 0); then v->2; t4,t3 at v2 (staleness 1) + reference = [("t2", 0), ("t1", 0), ("t4", 1), ("t3", 1)] + import itertools + + for perm in itertools.permutations(scenario.keys()): + assert run(list(perm)) == reference diff --git a/lib/python/tests/telemetry/__init__.py b/lib/python/tests/telemetry/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/lib/python/tests/telemetry/test_event_schema.py b/lib/python/tests/telemetry/test_event_schema.py new file mode 100644 index 000000000..d6423cd45 --- /dev/null +++ b/lib/python/tests/telemetry/test_event_schema.py @@ -0,0 +1,165 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Schema tests for telemetry event builders and selector emission. + +The cross-selector comparison feature relies on every selector emitting the +same selection-event schema, so these tests assert required fields are present +and consistent regardless of which selector produced them. +""" + +import json +from datetime import timedelta + +import pytest + +from flame import telemetry +from flame.selector.properties import ( + PROP_AVL_STATE, + PROP_ROUND_DURATION, + PROP_STAT_UTILITY, +) +from flame.telemetry.events import ( + EVENT_SELECTION, + EVENT_UTIL_DISPARITY, + build_agg_round, + build_selection, + build_trainer_round, + build_util_disparity, +) + + +@pytest.fixture(autouse=True) +def _reset_telemetry(): + telemetry.shutdown() + yield + telemetry.shutdown() + + +def _read(path): + with open(path) as fh: + return [json.loads(line) for line in fh if line.strip()] + + +# --- builder schema -------------------------------------------------------- + + +class TestBuilders: + def test_selection_required_fields(self): + ev, f = build_selection( + round_num=2, + task="train", + selector="OortSelector", + num_candidates=10, + num_eligible=7, + avail_composition={"AVL_TRAIN": 7, "UN_AVL": 3}, + chosen=["a", "b"], + in_flight=2, + ) + assert ev == EVENT_SELECTION + for key in ( + "round", "task", "selector", "num_candidates", "num_eligible", + "avail_composition", "chosen", "num_chosen", "in_flight", + ): + assert key in f + assert f["num_chosen"] == 2 + + def test_trainer_round_drops_none(self): + ev, f = build_trainer_round(round_num=1, real_gpu_time_s=3.0) + assert "real_gpu_time_s" in f + assert "wait_time_s" not in f # None omitted + + def test_agg_round_lists(self): + ev, f = build_agg_round(round_num=5, staleness=[0, 1, 2], agg_goal=3) + assert f["staleness"] == [0, 1, 2] + assert f["agg_goal"] == 3 + + def test_util_disparity_ratio_and_fraction(self): + ev, f = build_util_disparity( + round_num=1, + elapsed_s=10.0, + visible_samples=50, + total_samples=200, + utility_streamed=4.0, + utility_full=8.0, + ) + assert ev == EVENT_UTIL_DISPARITY + assert f["visible_fraction"] == pytest.approx(0.25) + assert f["utility_ratio"] == pytest.approx(0.5) + + def test_util_disparity_handles_zero_full(self): + _, f = build_util_disparity( + round_num=1, elapsed_s=1.0, visible_samples=1, total_samples=0, + utility_streamed=1.0, utility_full=0.0, + ) + assert f["utility_ratio"] is None + assert f["visible_fraction"] is None + + +# --- streaming-utility disparity direction --------------------------------- + + +class TestUtilDisparityDirection: + def test_full_pool_higher_utility_gives_ratio_below_one(self): + """Oort utility = N * sqrt(mean(loss^2)). With comparable mean squared + loss, the full pool (larger N) yields higher raw utility than the + streamed prefix, so streamed/full < 1 -- the hypothesized disparity.""" + import math + + mean_sq = 0.5 + n_streamed, n_full = 50, 200 + util_streamed = n_streamed * math.sqrt(mean_sq) + util_full = n_full * math.sqrt(mean_sq) + _, f = build_util_disparity( + round_num=1, elapsed_s=5.0, + visible_samples=n_streamed, total_samples=n_full, + utility_streamed=util_streamed, utility_full=util_full, + ) + assert f["utility_ratio"] < 1.0 + + +# --- selector emission produces a consistent schema ------------------------ + + +def _make_oort(): + from flame.selector.oort import OortSelector + + return OortSelector(aggr_num=3) + + +def _make_feddance(): + from flame.selector.feddance import FedDanceSelector + + return FedDanceSelector(aggr_num=3) + + +@pytest.mark.parametrize("factory", [_make_oort, _make_feddance]) +def test_selector_emits_selection_event(factory, make_ends, channel_props, tmp_path): + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + selector = factory() + + ends = make_ends(count=10, prefix="t") + # tag availability so composition is meaningful. Leave stat-utility unset + # so both selectors take their first-round cold-start path (the realistic + # round-1 behavior); the selection event is emitted either way. + for i, end in enumerate(ends.values()): + end.set_property(PROP_AVL_STATE, "AVL_TRAIN" if i < 7 else "UN_AVL") + end.set_property(PROP_ROUND_DURATION, timedelta(seconds=i + 1)) + + selector.select( + ends, channel_props, trainer_unavail_list=[], task_to_perform="train" + ) + telemetry.shutdown() + + recs = [r for r in _read(tmp_path / "aggregator.jsonl") if r["event"] == EVENT_SELECTION] + assert recs, "expected at least one selection event" + r = recs[-1] + # required, selector-agnostic schema + for key in ( + "round", "task", "selector", "num_candidates", "num_eligible", + "avail_composition", "chosen", "num_chosen", "in_flight", + ): + assert key in r, f"{key} missing for {r['selector']}" + assert r["num_candidates"] == 10 + # availability composition was derived from end properties + assert r["avail_composition"].get("AVL_TRAIN", 0) == 7 + assert r["avail_composition"].get("UN_AVL", 0) == 3 diff --git a/lib/python/tests/telemetry/test_telemetry_writer.py b/lib/python/tests/telemetry/test_telemetry_writer.py new file mode 100644 index 000000000..a45eb055e --- /dev/null +++ b/lib/python/tests/telemetry/test_telemetry_writer.py @@ -0,0 +1,85 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the process-global telemetry writer.""" + +import json +from datetime import timedelta + +import pytest + +from flame import telemetry + + +@pytest.fixture(autouse=True) +def _reset_telemetry(): + """Ensure each test starts/ends with telemetry disabled.""" + telemetry.shutdown() + yield + telemetry.shutdown() + + +def _read_jsonl(path): + with open(path) as fh: + return [json.loads(line) for line in fh if line.strip()] + + +class TestDisabledByDefault: + def test_emit_is_noop_when_unconfigured(self, monkeypatch): + monkeypatch.delenv(telemetry.ENV_DIR, raising=False) + assert telemetry.is_enabled() is False + # must not raise + telemetry.emit("selection", round=1) + + def test_configure_returns_none_without_dir(self, monkeypatch): + monkeypatch.delenv(telemetry.ENV_DIR, raising=False) + assert telemetry.configure(role="trainer", end_id="1") is None + assert telemetry.is_enabled() is False + + +class TestWriting: + def test_configure_and_emit_roundtrip(self, tmp_path): + w = telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + assert w is not None + assert telemetry.is_enabled() + telemetry.emit("agg_eval", round=3, **{"test-accuracy": 0.5}) + telemetry.shutdown() + + recs = _read_jsonl(tmp_path / "aggregator.jsonl") + assert len(recs) == 1 + r = recs[0] + assert r["event"] == "agg_eval" + assert r["role"] == "aggregator" + assert r["round"] == 3 + assert r["test-accuracy"] == 0.5 + assert "ts" in r + + def test_end_id_in_filename_and_record(self, tmp_path): + telemetry.configure(role="trainer", end_id="t7", run_dir=str(tmp_path)) + telemetry.emit("trainer_round", round=1, real_gpu_time_s=2.0) + telemetry.shutdown() + recs = _read_jsonl(tmp_path / "trainer_t7.jsonl") + assert recs[0]["end_id"] == "t7" + + def test_env_var_configures_dir(self, tmp_path, monkeypatch): + monkeypatch.setenv(telemetry.ENV_DIR, str(tmp_path)) + w = telemetry.configure(role="aggregator") + assert w is not None + assert telemetry.get_run_dir() == str(tmp_path) + + def test_serializes_timedelta_set_and_numpy(self, tmp_path): + import numpy as np + + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + telemetry.emit( + "agg_round", + round=1, + duration=timedelta(seconds=2.5), + chosen={"b", "a"}, + util=np.float64(1.25), + ) + telemetry.shutdown() + recs = _read_jsonl(tmp_path / "aggregator.jsonl") + r = recs[0] + assert r["duration"] == 2.5 + assert sorted(r["chosen"]) == ["a", "b"] + assert r["util"] == 1.25 diff --git a/scripts/analysis/analyze_run.py b/scripts/analysis/analyze_run.py new file mode 100644 index 000000000..77696e989 --- /dev/null +++ b/scripts/analysis/analyze_run.py @@ -0,0 +1,386 @@ +#!/usr/bin/env python +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Post-run telemetry analyzer for FLAME experiments. + +Reads the JSONL event streams written by :mod:`flame.telemetry` (one file per +process in a run's telemetry directory) and emits a bundle of PNG plots plus a +short text summary. Schema-driven, so it works identically across selectors / +aggregators -- enabling apples-to-apples comparison. + +Usage +----- + python analyze_run.py [--out ] + python analyze_run.py --compare ... [--labels a b ...] + +```` contains ``aggregator_*.jsonl`` and ``trainer_*.jsonl``. +Default output is ``/../plots``. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +from collections import Counter, defaultdict +from typing import Optional + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import plot_helpers as ph # noqa: E402 + +# Event-type names (kept in sync with flame/telemetry/events.py). Imported from +# the package when available, else hardcoded so the analyzer also runs stand-alone. +try: + sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "..", "lib", "python" + ), + ) + from flame.telemetry.events import ( # noqa: E402 + EVENT_AGG_EVAL, + EVENT_AGG_ROUND, + EVENT_AVAIL_CHANGE, + EVENT_SELECTION, + EVENT_TRAINER_ROUND, + EVENT_UTIL_DISPARITY, + ) +except Exception: # pragma: no cover - fallback for standalone use + EVENT_SELECTION = "selection" + EVENT_AGG_EVAL = "agg_eval" + EVENT_AGG_ROUND = "agg_round" + EVENT_TRAINER_ROUND = "trainer_round" + EVENT_UTIL_DISPARITY = "util_disparity" + EVENT_AVAIL_CHANGE = "avail_change" + + +def load_events(telemetry_dir: str) -> list[dict]: + """Load all JSONL records from a telemetry directory.""" + records: list[dict] = [] + for path in sorted(glob.glob(os.path.join(telemetry_dir, "*.jsonl"))): + with open(path) as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + records.append(json.loads(line)) + except json.JSONDecodeError: + continue # skip a torn final line from a killed process + return records + + +def by_event(records: list[dict], event: str) -> list[dict]: + return [r for r in records if r.get("event") == event] + + +# --- individual plots ------------------------------------------------------- + + +def plot_accuracy_loss(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_AGG_EVAL) + if not rows: + return [] + rows.sort(key=lambda r: r.get("round", 0)) + rounds = [r.get("round") for r in rows] + saved = [] + # accept either "test-accuracy"/"test-loss" or generic keys + acc = [r.get("test-accuracy") for r in rows] + loss = [r.get("test-loss") for r in rows] + if any(a is not None for a in acc): + rs = [r for r, a in zip(rounds, acc) if a is not None] + av = [a for a in acc if a is not None] + p = ph.line_plot( + {"test-accuracy": (rs, av)}, "round", "accuracy", + "Aggregator test accuracy over rounds", out_dir, "accuracy_over_rounds.png", + ) + if p: + saved.append(p) + if any(x is not None for x in loss): + rs = [r for r, x in zip(rounds, loss) if x is not None] + lv = [x for x in loss if x is not None] + p = ph.line_plot( + {"test-loss": (rs, lv)}, "round", "loss", + "Aggregator test loss over rounds", out_dir, "loss_over_rounds.png", + ) + if p: + saved.append(p) + return saved + + +def plot_staleness(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_AGG_ROUND) + if not rows: + return [] + saved = [] + stale_vals = [] + for r in rows: + for s in (r.get("staleness") or []): + stale_vals.append(s) + if stale_vals: + p = ph.cdf_plot( + stale_vals, "update staleness (rounds)", "Update staleness CDF", + out_dir, "staleness_cdf.png", + ) + if p: + saved.append(p) + # in-flight / queue timeline (by event arrival order) + inflight = [(i, r.get("updates_in_queue")) for i, r in enumerate(rows) + if r.get("updates_in_queue") is not None] + if inflight: + xs = [i for i, _ in inflight] + ys = [v for _, v in inflight] + p = ph.line_plot( + {"updates_in_queue": (xs, ys)}, "aggregation step", "updates in queue", + "Async queue depth over aggregation steps", out_dir, "queue_depth.png", + ) + if p: + saved.append(p) + return saved + + +def plot_avail_composition(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_SELECTION) + if not rows: + return [] + rows.sort(key=lambda r: r.get("round", 0)) + # one composition per round (last selection in that round wins) + per_round: dict[int, dict] = {} + for r in rows: + comp = r.get("avail_composition") + if comp: + per_round[r.get("round", 0)] = comp + if not per_round: + return [] + rounds = sorted(per_round.keys()) + state_names = sorted({s for c in per_round.values() for s in c.keys()}) + series = {s: [per_round[rd].get(s, 0) for rd in rounds] for s in state_names} + p = ph.stacked_area( + rounds, series, "round", "trainer count", + "Availability composition over rounds (selector view)", + out_dir, "availability_composition.png", + ) + return [p] if p else [] + + +def plot_utility_speed_scatter(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_SELECTION) + xs, ys, sel = [], [], [] + for r in rows: + per = r.get("per_trainer") or {} + for _, info in per.items(): + u = info.get("utility") + s = info.get("speed_s") + if u is None or s is None: + continue + xs.append(s) + ys.append(u) + sel.append(bool(info.get("selected"))) + if not xs: + return [] + p = ph.scatter_plot( + xs, ys, sel, "round duration / speed (s)", "statistical utility", + "Utility vs speed: selected vs eligible", out_dir, "utility_vs_speed.png", + ) + return [p] if p else [] + + +def plot_selection_frequency(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_SELECTION) + counter: Counter = Counter() + for r in rows: + for eid in (r.get("chosen") or []): + counter[str(eid)] += 1 + if not counter: + return [] + items = counter.most_common() + cats = [k for k, _ in items] + vals = [v for _, v in items] + p = ph.bar_plot( + cats, vals, "times selected", "Selection frequency per trainer (fairness)", + out_dir, "selection_frequency.png", + ) + return [p] if p else [] + + +def plot_trainer_time_breakdown(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_TRAINER_ROUND) + if not rows: + return [] + agg: dict[str, dict[str, list]] = defaultdict( + lambda: {"gpu": [], "sim": [], "wait": []} + ) + for r in rows: + tid = str(r.get("end_id", "?")) + agg[tid]["gpu"].append(r.get("real_gpu_time_s") or 0.0) + agg[tid]["sim"].append(r.get("sim_round_duration_s") or 0.0) + agg[tid]["wait"].append(r.get("wait_time_s") or 0.0) + cats = sorted(agg.keys()) + + def mean(xs): + return sum(xs) / len(xs) if xs else 0.0 + + segments = { + "real_gpu_time_s": [mean(agg[c]["gpu"]) for c in cats], + "sim_round_duration_s": [mean(agg[c]["sim"]) for c in cats], + "wait_time_s": [mean(agg[c]["wait"]) for c in cats], + } + p = ph.stacked_bar( + cats, segments, "mean seconds per round", + "Trainer time breakdown: real GPU vs simulated delay vs wait", + out_dir, "trainer_time_breakdown.png", + ) + return [p] if p else [] + + +def plot_util_disparity(records, out_dir) -> list[str]: + rows = by_event(records, EVENT_UTIL_DISPARITY) + if not rows: + return [] + saved = [] + # ratio over elapsed time, one series per trainer + by_trainer: dict[str, list[dict]] = defaultdict(list) + for r in rows: + by_trainer[str(r.get("end_id", "?"))].append(r) + ratio_series = {} + streamed_series = {} + full_series = {} + for tid, rs in by_trainer.items(): + rs.sort(key=lambda r: r.get("elapsed_s", 0)) + xs = [r.get("elapsed_s") for r in rs] + ratio_series[tid] = (xs, [r.get("utility_ratio") for r in rs]) + streamed_series["%s streamed" % tid] = (xs, [r.get("utility_streamed") for r in rs]) + full_series["%s full" % tid] = (xs, [r.get("utility_full") for r in rs]) + p = ph.line_plot( + ratio_series, "elapsed sim time (s)", "streamed / full utility ratio", + "Streamed-vs-full utility ratio over time (1.0 = no disparity)", + out_dir, "util_disparity_ratio.png", + ) + if p: + saved.append(p) + # absolute streamed vs full (combine; can be busy with many trainers) + combined = {} + combined.update(streamed_series) + combined.update(full_series) + p = ph.line_plot( + combined, "elapsed sim time (s)", "statistical utility", + "Streamed-prefix vs full-dataset utility over time", + out_dir, "util_disparity_absolute.png", + ) + if p: + saved.append(p) + return saved + + +def write_summary(records, out_dir, telemetry_dir) -> str: + counts = Counter(r.get("event") for r in records) + n_trainers = len({r.get("end_id") for r in records if r.get("role") == "trainer"}) + lines = [ + "Telemetry summary for: %s" % telemetry_dir, + "total events: %d" % len(records), + "trainers seen: %d" % n_trainers, + "event counts:", + ] + for ev, c in counts.most_common(): + lines.append(" %-16s %d" % (ev, c)) + text = "\n".join(lines) + "\n" + os.makedirs(out_dir, exist_ok=True) + path = os.path.join(out_dir, "summary.txt") + with open(path, "w") as fh: + fh.write(text) + return path + + +def analyze(telemetry_dir: str, out_dir: Optional[str] = None) -> list[str]: + records = load_events(telemetry_dir) + if out_dir is None: + out_dir = os.path.join(os.path.dirname(os.path.abspath(telemetry_dir)), "plots") + if not records: + print("no telemetry events found in %s" % telemetry_dir) + return [] + saved: list[str] = [] + plotters = [ + plot_accuracy_loss, + plot_staleness, + plot_avail_composition, + plot_utility_speed_scatter, + plot_selection_frequency, + plot_trainer_time_breakdown, + plot_util_disparity, + ] + for fn in plotters: + try: + saved.extend(fn(records, out_dir)) + except Exception as e: # one bad plot must not stop the rest + print(" (plot %s failed: %s)" % (fn.__name__, e)) + saved.append(write_summary(records, out_dir, telemetry_dir)) + print("wrote %d artifact(s) to %s" % (len(saved), out_dir)) + for p in saved: + print(" %s" % p) + return saved + + +def compare(dirs: list[str], labels: Optional[list[str]], out_dir: str) -> list[str]: + """Overlay accuracy / staleness across runs (one per selector).""" + if labels is None or len(labels) != len(dirs): + labels = [os.path.basename(os.path.dirname(os.path.abspath(d))) or d for d in dirs] + acc_series = {} + stale_series_vals = {} + for label, d in zip(labels, dirs): + recs = load_events(d) + ev = by_event(recs, EVENT_AGG_EVAL) + ev.sort(key=lambda r: r.get("round", 0)) + rs = [r.get("round") for r in ev if r.get("test-accuracy") is not None] + av = [r.get("test-accuracy") for r in ev if r.get("test-accuracy") is not None] + if rs: + acc_series[label] = (rs, av) + stale = [s for r in by_event(recs, EVENT_AGG_ROUND) for s in (r.get("staleness") or [])] + if stale: + stale_series_vals[label] = stale + saved = [] + p = ph.line_plot( + acc_series, "round", "accuracy", + "Accuracy comparison across runs", out_dir, "compare_accuracy.png", + ) + if p: + saved.append(p) + # overlay staleness CDFs by plotting each as its own line + if stale_series_vals: + import numpy as np + + series = {} + for label, vals in stale_series_vals.items(): + arr = np.sort(np.asarray(vals, dtype=float)) + y = np.arange(1, len(arr) + 1) / len(arr) + series[label] = (arr, y) + p = ph.line_plot( + series, "staleness (rounds)", "CDF", + "Staleness CDF comparison across runs", out_dir, "compare_staleness_cdf.png", + ) + if p: + saved.append(p) + print("wrote %d comparison artifact(s) to %s" % (len(saved), out_dir)) + return saved + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("telemetry_dir", nargs="?", help="run telemetry directory") + parser.add_argument("--out", help="output plots directory") + parser.add_argument("--compare", nargs="+", help="telemetry dirs to compare") + parser.add_argument("--labels", nargs="+", help="labels for --compare dirs") + args = parser.parse_args() + + if args.compare: + out = args.out or "compare_plots" + compare(args.compare, args.labels, out) + return + if not args.telemetry_dir: + parser.error("provide a telemetry_dir or --compare dirs...") + analyze(args.telemetry_dir, args.out) + + +if __name__ == "__main__": + main() diff --git a/scripts/analysis/plot_helpers.py b/scripts/analysis/plot_helpers.py new file mode 100644 index 000000000..225776a40 --- /dev/null +++ b/scripts/analysis/plot_helpers.py @@ -0,0 +1,204 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Reusable, dependency-light matplotlib helpers for telemetry analysis. + +These are the plotting primitives the run analyzer composes. They are kept +generic (lists / arrays in, PNG out) so future comparative plots reuse them +rather than duplicating matplotlib boilerplate. The CDF renderer mirrors the +percentile-annotated style of the legacy ``scripts/plotters/cdf_plot.py``. +""" + +from __future__ import annotations + +import os +from typing import Optional, Sequence + +import matplotlib + +matplotlib.use("Agg") # headless; never needs a display +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + + +def _save(fig, out_dir: str, file_name: str) -> str: + os.makedirs(out_dir, exist_ok=True) + path = os.path.join(out_dir, file_name) + fig.tight_layout() + fig.savefig(path, dpi=130, bbox_inches="tight") + plt.close(fig) + return path + + +def cdf_plot( + values: Sequence[float], + x_label: str, + title: str, + out_dir: str, + file_name: str, +) -> Optional[str]: + """Percentile-annotated CDF (P50/P90/P99). Returns saved path or None.""" + vals = [v for v in values if v is not None] + if not vals: + return None + arr = np.asarray(vals, dtype=float) + count, bins = np.histogram(arr, bins=min(5000, max(10, len(arr)))) + cdf = np.cumsum(count / count.sum()) + + fig, ax = plt.subplots(figsize=(8, 5)) + ax.plot(bins[1:], cdf, color="tab:red", linewidth=2) + for q, marker, label in [(0.5, "^", "P50"), (0.9, "o", "P90"), (0.99, "x", "P99")]: + xv = float(np.quantile(arr, q)) + ax.plot(xv, q, marker=marker, color="black", markersize=8) + ax.annotate( + f"{label}={xv:.3g}", + (xv, q), + textcoords="offset points", + xytext=(6, -10), + fontsize=8, + ) + ax.set_xlabel(x_label, fontweight="bold") + ax.set_ylabel("CDF", fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.grid(True, alpha=0.3) + ax.text( + 0.98, + 0.02, + f"n={len(arr)}", + transform=ax.transAxes, + ha="right", + va="bottom", + fontsize=8, + color="gray", + ) + return _save(fig, out_dir, file_name) + + +def line_plot( + series: dict[str, tuple[Sequence[float], Sequence[float]]], + x_label: str, + y_label: str, + title: str, + out_dir: str, + file_name: str, +) -> Optional[str]: + """Overlay multiple (x, y) series. ``series`` maps label -> (xs, ys).""" + if not series: + return None + fig, ax = plt.subplots(figsize=(9, 5)) + plotted = False + for label, (xs, ys) in series.items(): + if xs is None or ys is None or len(xs) == 0: + continue + ax.plot(xs, ys, marker=".", markersize=3, linewidth=1.2, label=label) + plotted = True + if not plotted: + plt.close(fig) + return None + ax.set_xlabel(x_label, fontweight="bold") + ax.set_ylabel(y_label, fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.grid(True, alpha=0.3) + if len(series) > 1: + ax.legend(fontsize=8) + return _save(fig, out_dir, file_name) + + +def stacked_area( + x: Sequence[float], + series: dict[str, Sequence[float]], + x_label: str, + y_label: str, + title: str, + out_dir: str, + file_name: str, +) -> Optional[str]: + """Stacked area chart, e.g. availability composition over rounds.""" + if not series or len(x) == 0: + return None + labels = list(series.keys()) + ys = [np.asarray(series[k], dtype=float) for k in labels] + fig, ax = plt.subplots(figsize=(9, 5)) + ax.stackplot(x, *ys, labels=labels, alpha=0.85) + ax.set_xlabel(x_label, fontweight="bold") + ax.set_ylabel(y_label, fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.legend(fontsize=8, loc="upper right") + ax.grid(True, alpha=0.3) + return _save(fig, out_dir, file_name) + + +def scatter_plot( + x: Sequence[float], + y: Sequence[float], + groups: Optional[Sequence[bool]], + x_label: str, + y_label: str, + title: str, + out_dir: str, + file_name: str, +) -> Optional[str]: + """Scatter; if ``groups`` (bool per point) given, color selected vs not.""" + if len(x) == 0: + return None + fig, ax = plt.subplots(figsize=(7, 6)) + x = np.asarray(x, dtype=float) + y = np.asarray(y, dtype=float) + if groups is not None: + g = np.asarray(groups, dtype=bool) + ax.scatter(x[~g], y[~g], s=18, alpha=0.5, color="tab:gray", label="eligible") + ax.scatter(x[g], y[g], s=28, alpha=0.8, color="tab:red", label="selected") + ax.legend(fontsize=8) + else: + ax.scatter(x, y, s=18, alpha=0.6, color="tab:blue") + ax.set_xlabel(x_label, fontweight="bold") + ax.set_ylabel(y_label, fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.grid(True, alpha=0.3) + return _save(fig, out_dir, file_name) + + +def stacked_bar( + categories: Sequence[str], + segments: dict[str, Sequence[float]], + y_label: str, + title: str, + out_dir: str, + file_name: str, +) -> Optional[str]: + """Stacked bar per category, e.g. trainer real-GPU vs sim vs wait time.""" + if not segments or len(categories) == 0: + return None + fig, ax = plt.subplots(figsize=(max(7, len(categories) * 0.5), 5)) + bottom = np.zeros(len(categories), dtype=float) + for label, vals in segments.items(): + v = np.asarray(vals, dtype=float) + ax.bar(range(len(categories)), v, bottom=bottom, label=label) + bottom += v + ax.set_xticks(range(len(categories))) + ax.set_xticklabels([str(c) for c in categories], rotation=60, ha="right", fontsize=7) + ax.set_ylabel(y_label, fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.legend(fontsize=8) + ax.grid(True, axis="y", alpha=0.3) + return _save(fig, out_dir, file_name) + + +def bar_plot( + categories: Sequence[str], + values: Sequence[float], + y_label: str, + title: str, + out_dir: str, + file_name: str, +) -> Optional[str]: + """Simple bar chart, e.g. selection frequency per trainer.""" + if len(categories) == 0: + return None + fig, ax = plt.subplots(figsize=(max(7, len(categories) * 0.4), 5)) + ax.bar(range(len(categories)), values, color="tab:blue", alpha=0.8) + ax.set_xticks(range(len(categories))) + ax.set_xticklabels([str(c) for c in categories], rotation=60, ha="right", fontsize=7) + ax.set_ylabel(y_label, fontweight="bold") + ax.set_title(title, fontweight="bold") + ax.grid(True, axis="y", alpha=0.3) + return _save(fig, out_dir, file_name)