From 3874d9f0cecc839c45bf54576db64d0fabb207b9 Mon Sep 17 00:00:00 2001 From: Dhruv Garg Date: Mon, 1 Jun 2026 17:08:31 -0400 Subject: [PATCH] Fix sim/real parity: unified max(gpu,D) timing model + GPU contention detection (#62) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Unified timing model (trainer, aggregator, OORT): - sim_round_duration = max(gpu_time, D) in both modes, mirroring real recv_ts-sent_ts. No contention: max(~0.5s,D)=D; overrun: gpu>D so OORT correctly deprioritizes contended trainers. - Real mode sleep = max(0, D-gpu); sim mode virtual advance = max(gpu,D). - Adds MessageType.TRAINING_BUDGET_S (37) to carry D to the aggregator. GPU contention detection: - Trainer logs [TIMING_OVERRUN] when gpu > D with advice to reduce trainers-per-GPU or add GPUs. - Aggregator logs [TIMING_OVERRUN_AGG] from budget vs elapsed in both modes. OORT crash fix (sample_by_util ValueError): - Root cause: timedelta(0) round duration -> system_utility=0 for all trainers -> np.random.choice(replace=False, p=[...]) fails. Fixed at source (max(gpu,D) is always >0) plus a defensive zero-prob filter in sample_by_util. compare_parity.py enhancements: - Check 9: GPU contention analysis (overrun fraction; FAIL>25%, WARN>10%). - --plot-out PATH: 2-panel timing sanity plot (mean GPU vs budget bars, round-by-round deviation). Graceful on old runs without training_budget_s. Docs + first-task guides: - TELEMETRY_AND_SPEEDUP_PLAN.md: new 2026-06-01 update section; Q2 corrected to reflect max(gpu,D) not just D. - david_first_task.md: add per-trainer lat/long mobility trace, emit as telemetry each round (real mode). - seshu_first_task.md: sinusoidal training_delay_s variation ±20% of mean, 0.5s floor, config-gated (real mode). - README.md: next-steps table with all three first-task files. Co-Authored-By: Claude Sonnet 4.6 --- lib/python/examples/async_cifar10/README.md | 13 +- .../TELEMETRY_AND_SPEEDUP_PLAN.md | 22 +- .../async_cifar10/david_first_task.md | 149 +++++++ .../async_cifar10/scripts/compare_parity.py | 390 +++++++++++++++++- .../async_cifar10/seshu_first_task.md | 157 +++++++ .../async_cifar10/trainer/pytorch/main.py | 128 ++---- .../mode/horizontal/asyncfl/top_aggregator.py | 249 ++++------- .../flame/mode/horizontal/syncfl/trainer.py | 23 +- lib/python/flame/mode/message.py | 7 +- lib/python/flame/selector/async_oort.py | 29 +- lib/python/flame/telemetry/events.py | 26 ++ 11 files changed, 914 insertions(+), 279 deletions(-) create mode 100644 lib/python/examples/async_cifar10/david_first_task.md create mode 100644 lib/python/examples/async_cifar10/seshu_first_task.md diff --git a/lib/python/examples/async_cifar10/README.md b/lib/python/examples/async_cifar10/README.md index b09526e9b..7ac3f3c18 100644 --- a/lib/python/examples/async_cifar10/README.md +++ b/lib/python/examples/async_cifar10/README.md @@ -148,8 +148,17 @@ tail -f eurosys26_expts/trainer_logs/log_trainer_*.log ## Next steps -- First task for a new contributor: [`ava_first_task.md`](ava_first_task.md) - (split the trainer delay into compute + RTT, then make RTT time-varying). +First-task guides for new contributors (each is self-contained; start from a +passing smoke test run): + +| File | Task summary | +|------|-------------| +| [`ava_first_task.md`](ava_first_task.md) | Split trainer delay into compute + RTT; make RTT time-varying | +| [`david_first_task.md`](david_first_task.md) | Add per-trainer lat/long from a mobility trace; emit location as telemetry each round | +| [`seshu_first_task.md`](seshu_first_task.md) | Make `training_delay_s` sinusoidally time-varying (±20% of mean, 0.5 s floor) | + +All three tasks target **`time_mode: real`** (wall-clock mode). Simulation mode +is still being validated and is not the target for these tasks. ## References diff --git a/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md b/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md index 4cfebf95a..2fe364d81 100644 --- a/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md +++ b/lib/python/examples/async_cifar10/TELEMETRY_AND_SPEEDUP_PLAN.md @@ -40,9 +40,27 @@ Implemented and unit-verified (156 tests green incl. `tests/sim` + `tests/mode`) - 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. +## Update (2026-06-01): unified timing model + GPU contention tracking + +**Unified `max(gpu, D)` round duration** ([trainer/pytorch/main.py](trainer/pytorch/main.py), [asyncfl/top_aggregator.py](../../flame/mode/horizontal/asyncfl/top_aggregator.py)): +- Previously both modes used a fixed `D` for round duration. Now `sim_round_duration = max(gpu_time, D)` in all modes: no contention → `D` (correct device speed); GPU overrun → `gpu > D` (OORT correctly deprioritizes the contended trainer). `sim_completion_ts = sim_send_ts + max(gpu, D)` mirrors real `recv_ts − sent_ts = max(gpu, D)` exactly. +- Real mode sleep: `time.sleep(max(0, D − gpu))` — absorbs remaining budget. No sleep on overrun. +- Aggregator: `PROP_ROUND_DURATION = timedelta(SIM_ROUND_DURATION)` in sim mode; `recv_ts − sent_ts` in real mode. Both yield `max(gpu, D)`. + +**Budget tracking + overrun detection** (`MessageType.TRAINING_BUDGET_S = 37`): +- Trainer sends `training_delay_s` (D) as `TRAINING_BUDGET_S` on every update (both modes). +- Aggregator: `[TIMING_OVERRUN_AGG]` warning when virtual elapsed ≤ 0 (sim) or wall lag > budget (real). Trainer: `[TIMING_OVERRUN]` warning when `gpu > D` with advice to reduce trainers-per-GPU or add GPUs. + +**OORT crash fix** ([selector/async_oort.py](../../flame/selector/async_oort.py)): +- Root cause: `timedelta(0)` round duration → `system_utility = 0` for all trainers → `np.random.choice(replace=False, p=[...])` fails when all probabilities are zero. Fixed at the source (overrun now yields `gpu > D > 0`, never zero duration) plus a defensive zero-probability filter in `sample_by_util`. + +**Compare-parity enhancements** ([scripts/compare_parity.py](scripts/compare_parity.py)): +- Check 9: GPU contention analysis — per-trainer overrun fraction vs budget; FAIL >25%, WARN >10%; graceful on old runs without `training_budget_s`. +- `--plot-out PATH`: 2-panel timing sanity plot — per-trainer mean GPU bar + budget marker (top), round-by-round deviation mean/max (bottom). Overrun zone highlighted in red. + 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. +- **End-to-end two-run verification** (Q1: `real` vs `simulated` per-round parity + speedup) needs a GPU/MQTT run — in progress (initial runs done; parity investigation ongoing). - **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). @@ -72,7 +90,7 @@ This is more efficient than scaling sleeps: wall-clock is bounded by *real GPU t **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.) +**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_ROUND_DURATION = max(gpu, D)`. In `real` mode the real delta `recv_ts − sent_ts = max(gpu, D)` (sleep absorbs remaining budget). Both modes deliver the same quantity to OORT: correct device speed when no contention, true overrun cost when GPU is shared. (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`. diff --git a/lib/python/examples/async_cifar10/david_first_task.md b/lib/python/examples/async_cifar10/david_first_task.md new file mode 100644 index 000000000..d9e301bde --- /dev/null +++ b/lib/python/examples/async_cifar10/david_first_task.md @@ -0,0 +1,149 @@ +# First task: add time-varying trainer location (lat/long) from a mobility trace + +This builds on the async_cifar10 example. Follow [README](README.md) first and +confirm you can run the felix smoke test end-to-end before starting here. + +The goal is to give each trainer a location identity that moves over time — like a +real mobile device. The trainer should read a per-device lat/long trace, interpolate +its current coordinates based on wall-clock time elapsed since it started, and emit +location as structured telemetry each round. + +The smoke test to target is the **real-mode** felix smoke: +```bash +python -m flame.launch.run_experiment \ + lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn0_smoke.yaml +``` +(Set `time_mode: real` in the YAML. Simulation mode is still being verified; do +not target it for this task.) + +--- + +## Background: how trainer metadata flows today + +1. **Registry** — `examples/_metadata/trainer_registry.yaml`. Each entry already + has a `mobiperf_device_id` (e.g. `device_001`). This is the logical handle + that links a trainer to a device profile. + +2. **Spawner injection** — `flame/launch/spawner.py` (~line 118-160) reads + `trainer_meta` from the registry and injects into each trainer's config: + `training_delay_s`, `trainer_indices_list`, all availability traces. You will + follow the same pattern to inject a location trace. + +3. **Schema** — `flame/config.py`, `Hyperparameters` class defines what fields the + trainer config carries. You will add a new field here. + +4. **Trainer** — `trainer/pytorch/main.py`. Reads from `config.hyperparameters` at + `__init__`. Emits structured telemetry via `telemetry.emit(ev, **fields)` in the + `train()` method (~line 721). You will add location reading and emission here. + +5. **Telemetry events** — `flame/telemetry/events.py` defines typed event builders + (e.g. `build_trainer_round`). You can add a field to `build_trainer_round` or + add it to the `extra={}` dict — pick whichever keeps the schema clean. + +--- + +## Task A — Create a per-device location trace and wire it to the registry + +Each trainer needs a sequence of `(elapsed_s, lat, lon)` waypoints that represent +where the device is at a given number of seconds after the experiment starts. + +Suggested subtasks: + +- [ ] **Design the trace format.** A simple YAML or JSON list per device works: + ```yaml + device_001: + - {elapsed_s: 0, lat: 37.7749, lon: -122.4194} # San Francisco + - {elapsed_s: 300, lat: 37.3861, lon: -122.0839} # Palo Alto + - {elapsed_s: 600, lat: 37.3382, lon: -121.8863} # San Jose + ``` + Keep the data small (3-5 waypoints per device) for now. You can use any + geographically plausible coordinates — real-world city pairs work well. + For the 10 smoke-test trainers (device_001 … device_010), 10 entries suffice. + +- [ ] **Write a generator script.** Add a script under + `examples/_metadata/` (e.g. `generate_location_traces.py`) that produces a + `location_traces.yaml` file in the same directory. Use random or hand-chosen + routes; document what you chose and why. Check it in alongside the generated + file so it can be re-run if routes need changing. + +- [ ] **Add to registry.** No need to embed full trace data in + `trainer_registry.yaml`; the spawner already loads traces separately. Store the + mapping `trainer_id → device_id` (already there as `mobiperf_device_id`) and + let the spawner resolve it. + +--- + +## Task B — Inject the location trace via the spawner + +Follow the same pattern as availability traces (~line 136-158 in `spawner.py`). + +- [ ] **`ExperimentMetadata`.** In `flame/launch/metadata.py` (or wherever + `get_mobiperf_trace` is implemented), add a `get_location_trace(trainer_id)` + method that loads the device's waypoint list from `location_traces.yaml`. + +- [ ] **`spawner.py`.** After the mobility trace injection block, add: + ```python + config["hyperparameters"]["location_trace"] = self.metadata.get_location_trace(trainer_id) + ``` + +- [ ] **`flame/config.py`.** Add `location_trace: Optional[List[Dict]] = None` + to `Hyperparameters` (or `Optional[Any]` if the schema is flexible). Keep the + field optional so existing configs without it still load. + +--- + +## Task C — Implement time-varying location in the trainer + +- [ ] **Read trace at `__init__`.** Parse `config.hyperparameters.location_trace` + into a list of `(elapsed_s, lat, lon)` tuples. Store as `self.location_trace`. + If absent or empty, set `self.location_trace = None`. + +- [ ] **Interpolate current location.** Write a small helper + `_current_location(elapsed_s)` that linearly interpolates `(lat, lon)` from the + waypoint list. Use the last waypoint once past the final entry (no extrapolation). + ```python + def _current_location(self, elapsed_s: float): + if not self.location_trace: + return None, None + # find bounding waypoints and interpolate + ... + ``` + +- [ ] **Emit each round.** In `train()`, after computing `_cycle_start`, compute: + ```python + elapsed = time.time() - self._experiment_start_time + lat, lon = self._current_location(elapsed) + ``` + Add `"lat": lat, "lon": lon` to the `extra={}` dict in the `build_trainer_round` + telemetry call (~line 741). Log it at DEBUG level too. + +- [ ] **`_experiment_start_time`.** Set `self._experiment_start_time = time.time()` + once in `__init__` (or on the first train call). Don't reset it per round. + +--- + +## Task D — Verify + +- [ ] Run the real-mode smoke test (`time_mode: real`, 10 trainers). Confirm the + trainer JSONL telemetry files contain `lat`/`lon` fields in `trainer_round` + events. +- [ ] Write a short script (or add to `compare_parity.py` as a helper) that reads + the trainer telemetry and plots each trainer's path (lat vs lon) over time. Two + trainers on the same device_id should trace identical paths. +- [ ] Confirm `lat`/`lon` are `None` for trainers whose config has no + `location_trace` (graceful no-op). + +--- + +## Notes & gotchas + +- The trainer already has `_sim_now()` which returns wall-clock-since-start in real + mode. You can reuse that instead of a separate `_experiment_start_time` if it + already tracks what you need. +- The location feature should be **read-only** — it must not affect selection, + training, or aggregation. It is purely telemetry/metadata. +- Keep trace files small (< 50 KB total). If you want more realistic movement, + sample from publicly available GPS datasets; just keep the format consistent. +- The spawner runs inside the launcher process; the trainer runs as a subprocess. + Data passes via the generated config JSON — it must be JSON-serializable (list of + dicts, no tuples at the top level). diff --git a/lib/python/examples/async_cifar10/scripts/compare_parity.py b/lib/python/examples/async_cifar10/scripts/compare_parity.py index b3f829426..825d755b5 100644 --- a/lib/python/examples/async_cifar10/scripts/compare_parity.py +++ b/lib/python/examples/async_cifar10/scripts/compare_parity.py @@ -2,18 +2,23 @@ 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 + 1. Selection parity — per round: same trainers selected? + 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. Convergence — accuracy / loss from agg_eval events + 7. Round duration parity — per-trainer sim_round_duration_s distributions match? + 8. sim_send_ts presence — sim mode: vclock stamps non-zero; real mode: null + 9. GPU contention — actual GPU time vs budget per trainer; detect overruns Usage: python compare_parity.py \\ - --real experiments/run_20260529_101200.../telemetry/aggregator_*.jsonl \\ - --sim experiments/run_20260529_152224.../telemetry/aggregator_*.jsonl \\ - [--rounds 100] [--strict] + --real experiments/run_.../telemetry/aggregator_*.jsonl \\ + --sim experiments/run_.../telemetry/aggregator_*.jsonl \\ + [--real-trainer-dir experiments/run_.../telemetry/] \\ + [--sim-trainer-dir experiments/run_.../telemetry/] \\ + [--rounds 100] [--strict] [--plot-out timing.png] [--json-out results.json] Exit 0 = all checks passed (or within tolerance) Exit 1 = one or more checks failed @@ -34,6 +39,33 @@ def short(end_id: str) -> str: return end_id[-4:] if end_id else "None" +def load_trainer_jsonl_dir(telemetry_dir: Optional[str]) -> dict: + """Load all trainer_*.jsonl files from a telemetry dir. + + Returns {short_id: {"task_recv": [...], "trainer_round": [...]}} + """ + if not telemetry_dir: + return {} + d = Path(telemetry_dir) + result = {} + for f in sorted(d.glob("trainer_*.jsonl")): + short_id = f.stem[-4:] # last 4 hex chars + task_recv_evs, trainer_round_evs = [], [] + with open(f) as fp: + for line in fp: + try: + e = json.loads(line.strip()) + except json.JSONDecodeError: + continue + ev = e.get("event") + if ev == "task_recv": + task_recv_evs.append(e) + elif ev == "trainer_round": + trainer_round_evs.append(e) + result[short_id] = {"task_recv": task_recv_evs, "trainer_round": trainer_round_evs} + return result + + def load_agg_jsonl(path: str) -> dict: """Parse an aggregator telemetry JSONL into typed lists.""" selection_train: list[dict] = [] # task=train selection events @@ -373,6 +405,281 @@ def curve(agg_evals): } +def check_round_duration_parity(real_trainers: dict, sim_trainers: dict) -> dict: + """Check 7: per-trainer sim_round_duration_s distributions match? + + Sim mode: sim_round_duration_s = remaining_time = max(0, D - gpu_time) ≈ D - gpu. + Real mode: sim_round_duration_s = gpu + remaining_time = max(gpu, D) = D (no overrun). + Expected mean diff ≈ gpu_time (small, ~0.5s for fast trainers). + Large diffs indicate overruns or dataset-size divergence. + """ + all_ids = sorted(set(real_trainers) | set(sim_trainers)) + if not all_ids: + return {"status": WARN, "note": "no trainer telemetry dirs provided"} + + rows, ks_stats, mean_diffs = [], [], [] + for tid in all_ids: + r_evs = real_trainers.get(tid, {}).get("trainer_round", []) + s_evs = sim_trainers.get(tid, {}).get("trainer_round", []) + r_durs = [e["sim_round_duration_s"] for e in r_evs if "sim_round_duration_s" in e] + s_durs = [e["sim_round_duration_s"] for e in s_evs if "sim_round_duration_s" in e] + ks = ks_stat(r_durs, s_durs) + rm, _ = mean_std(r_durs) + sm, _ = mean_std(s_durs) + diff = abs(rm - sm) if not (math.isnan(rm) or math.isnan(sm)) else float("nan") + rows.append({"trainer": tid, "real_mean_s": round(rm, 3), "sim_mean_s": round(sm, 3), + "mean_diff_s": round(diff, 3), "ks_stat": round(ks, 3) if not math.isnan(ks) else None, + "n_real": len(r_durs), "n_sim": len(s_durs)}) + 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_diff = sum(mean_diffs) / len(mean_diffs) if mean_diffs else float("nan") + status = PASS if (math.isnan(max_ks) or max_ks < 0.2) else (WARN if max_ks < 0.4 else FAIL) + return {"status": status, "max_ks_stat": round(max_ks, 3) if not math.isnan(max_ks) else None, + "avg_mean_diff_s": round(avg_diff, 3) if not math.isnan(avg_diff) else None, + "per_trainer": rows} + + +def check_sim_send_ts(real_trainers: dict, sim_trainers: dict) -> dict: + """Check 8: sim_send_ts correctness. + + Real mode: task_recv.sim_send_ts should be None (aggregator doesn't stamp vclock there). + Sim mode: task_recv.sim_send_ts must be non-None and must increase over rounds, + confirming the SIM_SEND_TS fix is in effect. A flat 0.0 across all rounds + indicates the bug is still present (trainer using _sim_now()=0 as base). + """ + issues = [] + + # real: all sim_send_ts should be null + for tid, data in real_trainers.items(): + non_null = [e["sim_send_ts"] for e in data.get("task_recv", []) if e.get("sim_send_ts") is not None] + if non_null: + issues.append(f"real/{tid}: unexpected non-null sim_send_ts values: {non_null[:3]}") + + # sim: sim_send_ts should be non-null and non-trivially non-zero after round 1 + sim_ok, sim_all_zero, sim_missing = 0, 0, 0 + for tid, data in sim_trainers.items(): + evs = [e for e in data.get("task_recv", []) if e.get("round", 0) > 1] + if not evs: + sim_missing += 1 + continue + vals = [e.get("sim_send_ts") for e in evs] + nulls = [v for v in vals if v is None] + zeros = [v for v in vals if v is not None and v == 0.0] + non_zeros = [v for v in vals if v is not None and v > 0.0] + if nulls: + issues.append(f"sim/{tid}: {len(nulls)} null sim_send_ts values — SIM_SEND_TS fix may not be active") + sim_all_zero += 1 + elif not non_zeros: + issues.append(f"sim/{tid}: all sim_send_ts==0.0 (rounds>1) — vclock not advancing") + sim_all_zero += 1 + else: + sim_ok += 1 + + status = PASS if not issues else FAIL + return {"status": status, "issues": issues, "sim_trainers_ok": sim_ok, + "sim_trainers_flat_zero": sim_all_zero, "sim_trainers_no_data": sim_missing} + + +def check_gpu_contention(real_trainers: dict, sim_trainers: dict) -> dict: + """Check 9: GPU contention — does actual GPU time respect the per-trainer budget? + + Reads training_budget_s and real_gpu_time_s from trainer_round events. + Overrun fraction = rounds where gpu_time > budget. + FAIL if mean overrun fraction > 25%; WARN if > 10%. + Old runs without training_budget_s in telemetry are skipped with WARN. + """ + def _extract(trainers): + stats = {} + for tid, d in trainers.items(): + evs = [e for e in d.get("trainer_round", []) + if "real_gpu_time_s" in e and e.get("training_budget_s", 0) > 0] + if not evs: + continue + gpu = [e["real_gpu_time_s"] for e in evs] + bgt = [e["training_budget_s"] for e in evs] + overruns = [g > b for g, b in zip(gpu, bgt)] + stats[tid] = { + "n": len(evs), + "budget_s": round(bgt[0], 2), + "mean_gpu_s": round(sum(gpu) / len(gpu), 3), + "max_gpu_s": round(max(gpu), 3), + "overrun_frac": round(sum(overruns) / len(overruns), 3), + "overran_rounds": int(sum(overruns)), + } + return stats + + real_st = _extract(real_trainers) + sim_st = _extract(sim_trainers) + + if not real_st and not sim_st: + return {"status": WARN, + "note": "no training_budget_s in telemetry (old run without timing model changes)"} + + def _agg(st): + if not st: + return None + fracs = [v["overrun_frac"] for v in st.values()] + return { + "n_trainers": len(st), + "mean_overrun_frac": round(sum(fracs) / len(fracs), 3), + "max_overrun_frac": round(max(fracs), 3), + "trainers_with_any_overrun": int(sum(1 for f in fracs if f > 0)), + } + + real_agg = _agg(real_st) + sim_agg = _agg(sim_st) + + worst = max( + (real_agg or {}).get("mean_overrun_frac", 0), + (sim_agg or {}).get("mean_overrun_frac", 0), + ) + status = FAIL if worst > 0.25 else (WARN if worst > 0.10 else PASS) + + all_tids = sorted(set(real_st) | set(sim_st)) + rows = [] + for tid in all_tids: + r, s = real_st.get(tid, {}), sim_st.get(tid, {}) + rows.append({ + "trainer": tid, + "budget_s": (r or s).get("budget_s"), + "real_mean_gpu_s": r.get("mean_gpu_s"), + "real_max_gpu_s": r.get("max_gpu_s"), + "real_overrun_frac": r.get("overrun_frac"), + "sim_mean_gpu_s": s.get("mean_gpu_s"), + "sim_max_gpu_s": s.get("max_gpu_s"), + "sim_overrun_frac": s.get("overrun_frac"), + }) + rows.sort(key=lambda x: max( + x.get("real_overrun_frac") or 0, + x.get("sim_overrun_frac") or 0, + ), reverse=True) + + return { + "status": status, + "real_summary": real_agg, + "sim_summary": sim_agg, + "per_trainer": rows, + } + + +def plot_timing_sanity(real_trainers: dict, sim_trainers: dict, out_path: str) -> None: + """Plot actual GPU time vs expected budget per trainer per round. + + Layout (2 rows × N_modes cols): + Top: per-trainer bar (mean GPU time) + red budget marker per trainer + Bottom: round-by-round deviation (GPU − budget); positive = contention + + If contention is present the bars will exceed the red markers and the + bottom panel will show positive spikes. If everything is healthy the + bars stay well below the red markers and the bottom panel stays negative. + """ + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError: + print("[plot] matplotlib not available; skipping timing plot") + return + + def _extract(trainers): + data = {} + for tid, d in trainers.items(): + evs = sorted( + [e for e in d.get("trainer_round", []) + if "real_gpu_time_s" in e and e.get("training_budget_s", 0) > 0], + key=lambda x: x.get("round", 0), + ) + if not evs: + continue + data[tid] = { + "rounds": [e["round"] for e in evs], + "gpu": [e["real_gpu_time_s"] for e in evs], + "budget": evs[0]["training_budget_s"], + } + return data + + real_data = _extract(real_trainers) if real_trainers else {} + sim_data = _extract(sim_trainers) if sim_trainers else {} + + modes = [(lbl, d) for lbl, d in [("Real", real_data), ("Sim", sim_data)] if d] + if not modes: + print("[plot] No training_budget_s data found; skipping timing plot") + return + + ncols = len(modes) + fig, axes = plt.subplots(2, ncols, figsize=(7 * ncols, 9), squeeze=False) + fig.suptitle("GPU Training Time vs Budget (Contention Sanity Check)", fontsize=13) + cmap = plt.cm.tab10 + + for col, (label, data) in enumerate(modes): + tids = sorted(data.keys(), key=lambda t: data[t]["budget"]) + colors = {t: cmap(i % 10) for i, t in enumerate(tids)} + xs = list(range(len(tids))) + + # ── top: per-trainer mean GPU bar + budget marker ────────────────── + ax_t = axes[0][col] + mean_gpus = [sum(data[t]["gpu"]) / len(data[t]["gpu"]) for t in tids] + budgets = [data[t]["budget"] for t in tids] + + ax_t.bar(xs, mean_gpus, color=[colors[t] for t in tids], alpha=0.75, zorder=2, + label="mean GPU time") + for x, b in zip(xs, budgets): + ax_t.plot([x - 0.38, x + 0.38], [b, b], color="red", linewidth=2.5, zorder=3) + # dummy line for legend + ax_t.plot([], [], color="red", linewidth=2.5, label="budget (D)") + + for x, g, b in zip(xs, mean_gpus, budgets): + if g > b: + ax_t.annotate("OVER", xy=(x, g), ha="center", va="bottom", + color="red", fontsize=7, fontweight="bold") + + ax_t.set_xticks(xs) + ax_t.set_xticklabels([f"...{t}" for t in tids], rotation=45, ha="right", fontsize=8) + ax_t.set_ylabel("seconds") + ax_t.set_title(f"{label}: mean GPU time vs budget", fontsize=10) + ax_t.set_ylim(bottom=0) + ax_t.legend(fontsize=8) + ax_t.grid(axis="y", alpha=0.3) + + # ── bottom: round-by-round deviation (mean + max across trainers) ── + ax_b = axes[1][col] + devs_per_round: dict = collections.defaultdict(list) + for t in tids: + bgt = data[t]["budget"] + for r, g in zip(data[t]["rounds"], data[t]["gpu"]): + devs_per_round[r].append(g - bgt) + + if devs_per_round: + rds = sorted(devs_per_round.keys()) + means = [sum(devs_per_round[r]) / len(devs_per_round[r]) for r in rds] + maxes = [max(devs_per_round[r]) for r in rds] + + ax_b.fill_between(rds, means, maxes, alpha=0.20, color="orange") + ax_b.plot(rds, means, color="steelblue", linewidth=1.5, label="mean deviation") + ax_b.plot(rds, maxes, color="darkorange", linewidth=1.0, alpha=0.8, + label="max deviation") + ax_b.fill_between(rds, 0, maxes, + where=[m > 0 for m in maxes], + alpha=0.10, color="red") + ax_b.axhline(0, color="red", linestyle="--", linewidth=1.5, + label="budget boundary (0 = on budget)") + + ax_b.set_xlabel("round") + ax_b.set_ylabel("GPU time − budget (s)") + ax_b.set_title(f"{label}: GPU contention deviation per round", fontsize=10) + ax_b.legend(fontsize=8, loc="upper left") + ax_b.grid(alpha=0.3) + + plt.tight_layout() + plt.savefig(out_path, dpi=120, bbox_inches="tight") + plt.close() + print(f"[plot] Timing sanity plot → {out_path}") + + # ── reporting ───────────────────────────────────────────────────────────────── def _status_icon(s: str) -> str: @@ -392,6 +699,9 @@ def print_report(results: dict, strict: bool) -> bool: ("4. Staleness distribution", "staleness"), ("5. Trainer participation counts", "participation"), ("6. Convergence (accuracy / loss)", "convergence"), + ("7. Round duration parity (sim_round_duration_s)", "round_duration"), + ("8. sim_send_ts presence / correctness", "sim_send_ts"), + ("9. GPU contention (actual GPU time vs budget)", "gpu_contention"), ] for title, key in checks: @@ -457,6 +767,51 @@ def print_report(results: dict, strict: bool) -> bool: f"acc real={row['real_acc']} sim={row['sim_acc']} " f"loss real={row['real_loss']} sim={row['sim_loss']}") + elif key == "round_duration": + note = r.get("note") + if note: + print(f" note: {note}") + else: + print(f" max KS stat: {r.get('max_ks_stat')} (<0.2=good)") + print(f" avg mean diff (s): {r.get('avg_mean_diff_s')}") + for row in r.get("per_trainer", []): + print(f" ...{row['trainer']}: real={row['real_mean_s']}s " + f"sim={row['sim_mean_s']}s diff={row['mean_diff_s']}s " + f"KS={row['ks_stat']} (n_real={row['n_real']} n_sim={row['n_sim']})") + + elif key == "sim_send_ts": + print(f" sim trainers OK (non-zero vclock): {r.get('sim_trainers_ok')}") + print(f" sim trainers flat-zero (bug present): {r.get('sim_trainers_flat_zero')}") + print(f" sim trainers no task_recv data: {r.get('sim_trainers_no_data')}") + for issue in r.get("issues", [])[:5]: + print(f" [!] {issue}") + + elif key == "gpu_contention": + note = r.get("note") + if note: + print(f" note: {note}") + else: + ra, sa = r.get("real_summary"), r.get("sim_summary") + if ra: + print(f" real: {ra['n_trainers']} trainers, " + f"mean_overrun_frac={ra['mean_overrun_frac']:.1%} " + f"max_overrun_frac={ra['max_overrun_frac']:.1%} " + f"trainers_with_overruns={ra['trainers_with_any_overrun']}") + if sa: + print(f" sim: {sa['n_trainers']} trainers, " + f"mean_overrun_frac={sa['mean_overrun_frac']:.1%} " + f"max_overrun_frac={sa['max_overrun_frac']:.1%} " + f"trainers_with_overruns={sa['trainers_with_any_overrun']}") + print(" per trainer (sorted by worst overrun fraction):") + for row in r.get("per_trainer", []): + r_str = (f"real: gpu={row['real_mean_gpu_s']}s max={row['real_max_gpu_s']}s " + f"overrun={row['real_overrun_frac']:.1%}" + if row.get("real_mean_gpu_s") is not None else "real: n/a") + s_str = (f"sim: gpu={row['sim_mean_gpu_s']}s max={row['sim_max_gpu_s']}s " + f"overrun={row['sim_overrun_frac']:.1%}" + if row.get("sim_mean_gpu_s") is not None else "sim: n/a") + print(f" ...{row['trainer']} D={row['budget_s']}s | {r_str} | {s_str}") + if status == "FAIL" or (strict and status == "WARN"): overall_ok = False print() @@ -474,12 +829,18 @@ 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("--real-trainer-dir", default=None, + help="Dir containing real trainer_*.jsonl files (for checks 7-8)") + parser.add_argument("--sim-trainer-dir", default=None, + help="Dir containing sim trainer_*.jsonl files (for checks 7-8)") 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") + parser.add_argument("--plot-out", default=None, + help="Write timing sanity plot (GPU vs budget) to this PNG path") args = parser.parse_args() print(f"Loading real: {args.real}") @@ -494,6 +855,11 @@ def main(): f"{len(sim['selection_train'])} train-selection events, " f"{len(sim['agg_evals'])} eval events") + real_trainers = load_trainer_jsonl_dir(args.real_trainer_dir) + sim_trainers = load_trainer_jsonl_dir(args.sim_trainer_dir) + if real_trainers or sim_trainers: + print(f" real trainer files: {len(real_trainers)}, sim trainer files: {len(sim_trainers)}") + results = { "selection": check_selection_parity(real, sim, args.rounds), "utility": check_utility_parity(real, sim), @@ -501,10 +867,16 @@ def main(): "staleness": check_staleness(real, sim), "participation": check_participation(real, sim), "convergence": check_convergence(real, sim), + "round_duration": check_round_duration_parity(real_trainers, sim_trainers), + "sim_send_ts": check_sim_send_ts(real_trainers, sim_trainers), + "gpu_contention": check_gpu_contention(real_trainers, sim_trainers), } ok = print_report(results, args.strict) + if args.plot_out: + plot_timing_sanity(real_trainers, sim_trainers, args.plot_out) + if args.json_out: with open(args.json_out, "w") as f: json.dump(results, f, indent=2) diff --git a/lib/python/examples/async_cifar10/seshu_first_task.md b/lib/python/examples/async_cifar10/seshu_first_task.md new file mode 100644 index 000000000..f5b44f313 --- /dev/null +++ b/lib/python/examples/async_cifar10/seshu_first_task.md @@ -0,0 +1,157 @@ +# First task: make training delay time-varying (sinusoidal model) + +This builds on the async_cifar10 example. Follow [README](README.md) first and +confirm you can run the felix smoke test end-to-end before starting here. + +Today each trainer has a fixed `training_delay_s` (D) that never changes. Real +devices vary: a phone on charge trains faster, a busy device trains slower. The +goal is to make D vary over wall-clock time using a sinusoidal model, so each +trainer's effective budget oscillates around its configured mean value. + +The smoke test to target is the **real-mode** felix smoke: +```bash +python -m flame.launch.run_experiment \ + lib/python/examples/async_cifar10/expt_scripts_2026/felix_n10_alpha100_syn0_smoke.yaml +``` +(Set `time_mode: real` in the YAML. Simulation mode is still being verified; do +not target it for this task.) + +--- + +## The model + +``` +D(t) = clamp(mean_D + amplitude × sin(2π × t / period_s), min=0.5) +``` + +where: +- `mean_D` = `training_delay_s` from the registry (the existing per-trainer value) +- `amplitude` = `mean_D × amplitude_fraction` (default `0.2` → ±20% of mean) +- `period_s` = oscillation period in seconds (default `120s`) +- `min` = hard floor of `0.5s` (never less than half a second) +- `t` = wall-clock seconds elapsed since the trainer started + +With the default values a trainer with `mean_D = 4s` oscillates between `3.2s` +and `4.8s`; a trainer with `mean_D = 16s` oscillates between `12.8s` and `19.2s`. + +Both min and max stay within 20% of the mean, satisfying the original spec. + +--- + +## Background: how training_delay_s works today + +1. **Registry** — `examples/_metadata/trainer_registry.yaml`. Each entry has a + `training_delay_s` value (e.g. `"4.0"` for fast, `"16.0"` for very_slow). + +2. **Injection** — `flame/launch/spawner.py` (~line 130): + ```python + config["hyperparameters"]["training_delay_s"] = trainer_meta["training_delay_s"] + ``` + +3. **Schema** — `flame/config.py`, `Hyperparameters`: + `training_delay_s` (alias `trainingDelaySeconds`), `training_delay_enabled`. + +4. **Application** — `trainer/pytorch/main.py`: reads `self.training_delay_s` at + `__init__`; uses it as the budget D each round in the timing block (~line 696): + ```python + _modeled_delay_s = self.training_delay_s if self.training_delay_enabled else 0.0 + _remaining_time = max(0.0, _modeled_delay_s - _real_gpu_time_s) + ``` + +Your task is to make `_modeled_delay_s` time-varying instead of fixed. + +--- + +## Task A — Add variation config fields to the schema + +- [ ] **`flame/config.py`.** Add a nested config block to `Hyperparameters`: + ```python + training_delay_variation: Optional[Dict] = None + # Expected keys: enabled (bool), period_s (float), amplitude_fraction (float) + # Defaults: enabled=False, period_s=120, amplitude_fraction=0.2 + ``` + Keep it optional so all existing configs without it continue to work unchanged. + +- [ ] **Smoke-test YAML.** In + `expt_scripts_2026/felix_n10_alpha100_syn0_smoke.yaml` (or a copy), add under + `trainer.hyperparameters`: + ```yaml + training_delay_variation: + enabled: true + period_s: 120 + amplitude_fraction: 0.2 + ``` + +--- + +## Task B — Implement time-varying delay in the trainer + +- [ ] **Read config at `__init__`** (`trainer/pytorch/main.py`). Parse + `config.hyperparameters.training_delay_variation` and store: + ```python + _var = config.hyperparameters.training_delay_variation or {} + self.delay_variation_enabled = str(_var.get("enabled", "False")).lower() == "true" + self.delay_variation_period_s = float(_var.get("period_s", 120)) + self.delay_variation_amplitude = float(_var.get("amplitude_fraction", 0.2)) * self.training_delay_s + ``` + +- [ ] **Record start time.** Add `self._experiment_start_time = time.time()` in + `__init__` (set once; never reset per round). + +- [ ] **Compute D(t) each round.** Add a helper method: + ```python + def _effective_delay_s(self) -> float: + if not self.delay_variation_enabled: + return self.training_delay_s + t = time.time() - self._experiment_start_time + raw = self.training_delay_s + self.delay_variation_amplitude * math.sin( + 2 * math.pi * t / self.delay_variation_period_s + ) + return max(0.5, raw) + ``` + +- [ ] **Use it in the timing block.** Replace the fixed read with the helper in + the `train()` method (~line 696): + ```python + _modeled_delay_s = self._effective_delay_s() if self.training_delay_enabled else 0.0 + ``` + No other lines need to change — `_remaining_time`, overrun logic, and + `sim_round_duration` all derive from `_modeled_delay_s`. + +- [ ] **Log it.** Add `effective_delay_s` and `delay_phase_t` to the telemetry + `extra={}` dict so you can plot D(t) vs round from the trainer JSONL. + +--- + +## Task C — Verify + +- [ ] Run the real-mode smoke test (10 trainers, `time_mode: real`, variation + enabled). Confirm from trainer logs that `[TRAIN_CYCLE]` lines show varying + `budget=` values over rounds for each trainer. + +- [ ] Write a short script that reads `trainer_N.jsonl` files and plots + `effective_delay_s` vs round number (or elapsed time) for all trainers. You + should see sinusoidal oscillation around each trainer's `mean_D`. + +- [ ] Confirm that with `enabled: false` (or config absent) the behavior is + identical to the original fixed-delay path. + +- [ ] Confirm the 0.5s floor: set `amplitude_fraction: 1.5` for a trainer with + `mean_D = 0.3s`; verify logged `effective_delay_s` never drops below 0.5s. + +--- + +## Notes & gotchas + +- `math.sin` is already imported in `trainer/pytorch/main.py`. +- The variation parameters are **per-trainer-globally-fixed** (read once from + config), but `_effective_delay_s()` returns a **per-call varying value** because + `t = time.time() - start` changes each call. This is intentional. +- In sim mode `self._experiment_start_time` (wall-clock) is not the same as + sim-time. For now scope this task to `time_mode: real` only. Add a + `TODO: sim-time support` comment if you want to mark the gap. +- The `training_delay_s` in the registry stays as the **mean** value; you are not + changing the registry. Variation is additive on top of that mean. +- Eval path also uses `training_delay_s` (search for `eval_delay` or a second + sleep in `evaluate()`). Check if it should also vary — for this task, leave it + fixed and note the inconsistency. diff --git a/lib/python/examples/async_cifar10/trainer/pytorch/main.py b/lib/python/examples/async_cifar10/trainer/pytorch/main.py index 0b6e5d66e..1ea4247bd 100644 --- a/lib/python/examples/async_cifar10/trainer/pytorch/main.py +++ b/lib/python/examples/async_cifar10/trainer/pytorch/main.py @@ -97,52 +97,31 @@ def __init__(self, config: Config, battery_threshold, time_mode="simulated") -> self.batch_size = self.config.hyperparameters.batch_size or 16 self.trainer_id = self.config.task_id - # Learning rate decay configuration (optional, for REFL) - # REFL uses decay_factor=0.98 every 10 rounds, Oort doesn't use decay self.lr_decay_enabled = getattr(self.config.hyperparameters, 'lr_decay_enabled', False) self.lr_decay_factor = getattr(self.config.hyperparameters, 'lr_decay_factor', 0.98) self.lr_decay_epoch = getattr(self.config.hyperparameters, 'lr_decay_epoch', 10) self.min_learning_rate = getattr(self.config.hyperparameters, 'min_learning_rate', 1e-4) - - logger.info( - f"Trainer {self.trainer_id}: LR decay {'ENABLED' if self.lr_decay_enabled else 'DISABLED'} " - f"(factor={self.lr_decay_factor}, epoch={self.lr_decay_epoch}, min_lr={self.min_learning_rate})" - ) self.criterion = None self.task_to_perform = "train" - # Enable/disable use of oort_loss fromt he config. Needed for - # oort and asyncOORT. self.use_oort_loss_fn = self.config.hyperparameters.use_oort_loss_fn - logger.info( - f"Trainer: {self.trainer_id} has " - f"use_oort_loss_fn: {self.use_oort_loss_fn}" - ) - - # TODO: (DG) Remove the hard requirement for config to include - # trainer_indices_list - # Setting the indices used by the trainer self.trainer_indices_list = self.config.hyperparameters.trainer_indices_list self.trainer_start_ts = time.time() - # sending heartbeats to aggregator - if "enabled" in self.config.hyperparameters.heartbeats.keys(): + if "enabled" in self.config.hyperparameters.heartbeats: self.heartbeats_enabled = self.config.hyperparameters.heartbeats["enabled"] else: self.heartbeats_enabled = False - if "frequency_s" in self.config.hyperparameters.heartbeats.keys(): + if "frequency_s" in self.config.hyperparameters.heartbeats: self.heartbeats_second_freq = self.config.hyperparameters.heartbeats[ "frequency_s" ] else: self.heartbeats_second_freq = 99999 - # TODO: (DG) self.timestamp_next_heartbeat_s might not be - # getting used. Remove? if heartbeats are enabled, compute - # first heartbeat time if self.heartbeats_enabled is True: self.timestamp_next_heartbeat_s = ( self.trainer_start_ts + self.heartbeats_second_freq @@ -152,27 +131,18 @@ def __init__(self, config: Config, battery_threshold, time_mode="simulated") -> time.strptime("Dec 31, 2030 @ 23:59:59 UTC", "%b %d, %Y @ %H:%M:%S UTC") ) - # 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. - # 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. + # Normalize to bool: may arrive as Python bool or string "True"/"False". _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) - # 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 + self._sim_send_ts = None # set by aggregator stamp on each task (sim mode) # Use the battery_threshold to determine the # avl_events_3_state config. Default to 50 if not provided @@ -181,7 +151,6 @@ def __init__(self, config: Config, battery_threshold, time_mode="simulated") -> f"Trainer id {self.trainer_id} has battery threshold set to {self.event_battery_threshold}" ) - # Helper function to handle both string and list formats def parse_trace(value): if isinstance(value, list): return value # Already parsed (from JSON config) @@ -250,7 +219,6 @@ 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( @@ -262,8 +230,6 @@ def parse_trace(value): 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) @@ -288,23 +254,13 @@ def check_and_sleep(self): 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. - """ + """Wall-elapsed (real) or last-task sim_send_ts (simulated).""" 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. - """ + """Advance availability state to current sim-time (sim mode only).""" if not self.simulated: return guard = 0 @@ -686,17 +642,13 @@ def train(self) -> None: # reset stat utility for OORT self.reset_stat_utility() - # 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)" + _expected_wallclock_hint = f"~GPU wall-clock only; virtual_advance=max(gpu,D={_D:.1f}s)" else: - _expected_wallclock_hint = f"~GPU + {_D:.1f}s sleep = ~{_D:.1f}s+ wall-clock" + _expected_wallclock_hint = f"~max(gpu,D={_D:.1f}s) wall-clock; sleep=max(0,D-gpu)" logger.info( f"[TRAIN_START] Trainer {self.trainer_id} starting training with " f"model_version={self._round}, dataset_size={dataset_size}, " @@ -741,26 +693,28 @@ def train(self) -> None: # Log memory after training round self.memory_profiler.log_memory_after_round() - # 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 + _remaining_time = max(0.0, _modeled_delay_s - _real_gpu_time_s) + _overran = self.training_delay_enabled and _real_gpu_time_s > _modeled_delay_s + self._training_budget_s = _modeled_delay_s + + if _overran: + logger.warning( + f"[TIMING_OVERRUN] Trainer {self.trainer_id} round={self._round} " + f"exceeded budget: gpu={_real_gpu_time_s:.2f}s > budget={_modeled_delay_s:.2f}s " + f"(excess={_real_gpu_time_s - _modeled_delay_s:.2f}s). " + f"Reduce trainers-per-GPU or add GPUs." + ) + + # max(gpu, D): no contention → D; overrun → gpu > D (OORT sees trainer as slow). + sim_round_duration = _real_gpu_time_s + _remaining_time # = max(gpu, D) + 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() @@ -780,38 +734,40 @@ def train(self) -> None: 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}, + extra={ + "sim_completion_ts": self._sim_completion_ts, + "sim_send_ts": float(self._sim_send_ts) if self._sim_send_ts is not None else None, + "time_mode": self.time_mode, + "training_budget_s": _modeled_delay_s, + "remaining_time_s": _remaining_time, + "overran": _overran, + }, ) 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) + if not self.simulated and _remaining_time > 0: + time.sleep(_remaining_time) - # 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"[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}" + f"time_mode=simulated: wall={_cycle_elapsed:.2f}s " + f"GPU={_real_gpu_time_s:.2f}s budget={_modeled_delay_s:.1f}s " + f"virtual_advance={sim_round_duration:.2f}s " + f"{'OVERRUN' if _overran else 'OK'} " + f"sct={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'}" + f"time_mode=real: wall={_cycle_elapsed:.2f}s " + f"GPU={_real_gpu_time_s:.2f}s budget={_modeled_delay_s:.1f}s " + f"sleep={_remaining_time:.2f}s total={sim_round_duration:.1f}s " + f"{'OVERRUN' if _overran else 'OK'}" ) def _train_epoch(self, epoch): diff --git a/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py b/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py index fdde156de..cd6298c20 100644 --- a/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py +++ b/lib/python/flame/mode/horizontal/asyncfl/top_aggregator.py @@ -37,7 +37,7 @@ 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.properties import PROP_SIM_SEND_TS, PROP_SIM_COMPLETION_TS from flame.selector.oort import ( PROP_DATASET_SIZE, PROP_LAST_SELECTED_ROUND, @@ -80,40 +80,27 @@ 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 - # variables related to checking trainer availability self._per_trainer_last_heartbeat_ts = {} - if "heartbeat_freq_s" in self.config.hyperparameters.track_trainer_avail.keys(): + if "heartbeat_freq_s" in self.config.hyperparameters.track_trainer_avail: self._trainer_heartbeat_freq_s = ( self.config.hyperparameters.track_trainer_avail["heartbeat_freq_s"] ) else: self._trainer_heartbeat_freq_s = 99999 - if ( - "max_allowed_miss_heartbeats" - in self.config.hyperparameters.track_trainer_avail.keys() - ): + if "max_allowed_miss_heartbeats" in self.config.hyperparameters.track_trainer_avail: self._trainer_max_miss_heartbeats = ( - self.config.hyperparameters.track_trainer_avail[ - "max_allowed_miss_heartbeats" - ] + self.config.hyperparameters.track_trainer_avail["max_allowed_miss_heartbeats"] ) else: self._trainer_max_miss_heartbeats = 99999 - # maintain a set of all trainers that have sent heartbeats - # previously self.all_trainers = set() def _reset_agg_goal_variables(self): @@ -290,31 +277,18 @@ def _aggregate_weights(self, tag: str) -> None: f"agg_current_version={self._round}" ) - # For OORT selector NOTE: (DG) Last selected round should - # have ideally been set in distribute weights. But it was - # here in the old oort code and ive kept it. Instead of - # PROP_LAST_SELECTED_ROUND, it should have been - # PROP_LAST_UPDATE_RECVD_ROUND. channel.set_end_property( end, PROP_LAST_SELECTED_ROUND, msg[MessageType.MODEL_VERSION] ) - # Set last eval round for the trainer since training also - # means that eval was done for the same round. channel.set_end_property( end, PROP_LAST_EVAL_ROUND, msg[MessageType.MODEL_VERSION] ) - # calculate round duration for this end, if the round - # number information is identical with round_start_time - logger.debug( - f"Getting channel property {PROP_ROUND_START_TIME} for " f"end {end}" - ) round_start_time_tup = channel.get_end_property(end, PROP_ROUND_START_TIME) end = metadata[0] timestamp = metadata[1] logger.debug( - f"Returned round_start_time_tup: {round_start_time_tup} for " - f"end {end} and timestamp {timestamp}" + f"round_start_time_tup={round_start_time_tup} end={end} ts={timestamp}" ) # TODO: (DG) Also set the end property for task=eval done @@ -441,6 +415,28 @@ def _aggregate_weights(self, tag: str) -> None: f"wall_lag_s={wall_lag_s:.1f}s — possible MQTT backlog" ) + _budget_s = float(msg.get(MessageType.TRAINING_BUDGET_S, 0.0)) + if _budget_s > 0: + if self.simulated: + _sst = channel.get_end_property(end, PROP_SIM_SEND_TS) + _sct = msg.get(MessageType.SIM_COMPLETION_TS) + if _sst is not None and _sct is not None: + _virt_elapsed = float(_sct) - float(_sst) + if _virt_elapsed <= 0.0: + logger.warning( + f"[TIMING_OVERRUN_AGG] {end[-4:]} ver={recv_wts_version} " + f"budget={_budget_s:.1f}s overrun: virtual_elapsed={_virt_elapsed:.2f}s. " + f"Reduce trainers-per-GPU or add GPUs." + ) + else: + if wall_lag_s > _budget_s: + logger.warning( + f"[TIMING_OVERRUN_AGG] {end[-4:]} ver={recv_wts_version} " + f"budget={_budget_s:.1f}s overrun: wall_lag={wall_lag_s:.2f}s " + f"(excess={wall_lag_s - _budget_s:.2f}s). " + f"Reduce trainers-per-GPU or add GPUs." + ) + # TODO: (DG) Can pass a flag for this later. allow_updates_more_than_timeout_old = True @@ -499,10 +495,7 @@ def _aggregate_weights(self, tag: str) -> None: curr_cumulative_training_s = self._track_trainer_version_duration_s[ end ]["total_training_time_s"] - # 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). + # Both modes: round_duration = max(gpu, D); mirrors recv_ts-sent_ts for OORT utility. if self.simulated: round_duration_td = timedelta( seconds=float(msg.get(MessageType.SIM_ROUND_DURATION, 0.0)) @@ -535,21 +528,10 @@ def _aggregate_weights(self, tag: str) -> None: end, PROP_ROUND_DURATION, round_duration_td ) - # capture telemetry on trainer participation in rounds channel._selector.ordered_updates_recv_ends.append(end) - logger.debug( - f"After appending {end} to ordered_updates_recv_ends: " - f"{channel._selector.ordered_updates_recv_ends}" - ) - self._updates_in_queue += 1 - - self._per_round_update_list.append( - end - ) # SC_TS: PER Round! In Async FwdLLM -> we need to decide whether per data bin or per iteration! - - # SC_TS: what is even the diff between this and self._updates_in_queue += 1??!! - if end not in self._updates_recevied.keys(): + self._per_round_update_list.append(end) + if end not in self._updates_recevied: self._updates_recevied[end] = 1 else: self._updates_recevied[end] += 1 @@ -597,8 +579,8 @@ def _aggregate_weights(self, tag: str) -> None: _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(): + _sct_recv = msg.get(MessageType.SIM_COMPLETION_TS) ev, fields = build_agg_round( round_num=self._round, agg_goal=self._agg_goal, @@ -608,6 +590,10 @@ def _aggregate_weights(self, tag: str) -> None: stat_utility=[stat_utility], trainer_speed_s=[_trainer_speed_s], contributing_trainers=[end], + extra={ + "sim_completion_ts_recv": float(_sct_recv) if _sct_recv is not None else None, + "vclock_now": self._vclock.now if self.simulated else None, + }, ) telemetry.emit(ev, **fields) @@ -649,7 +635,6 @@ def _aggregate_weights(self, tag: str) -> None: # discarding") return logger.info("proceeding to agg weights") - # SC_TS: append weights to this list, till agg goal reached! self._agg_goal_weights = self.optimizer.do( self._agg_goal_weights, self.cache, @@ -661,16 +646,10 @@ def _aggregate_weights(self, tag: str) -> None: self._agg_goal_cnt += 1 if self._agg_goal_cnt < self._agg_goal: - # didn't reach the aggregation goal; return logger.debug( - f"didn't reach agg goal. _agg_goal_cnt: {self._agg_goal_cnt} while _agg_goal is {self._agg_goal}" - ) - - # Set trainer participation count property here to be used - # later in selection. - channel.set_end_property( - end, PROP_UPDATE_COUNT, self._updates_recevied[end] + f"agg_goal_cnt={self._agg_goal_cnt} < agg_goal={self._agg_goal}, waiting for more" ) + channel.set_end_property(end, PROP_UPDATE_COUNT, self._updates_recevied[end]) return if self._agg_goal_weights is None: @@ -682,15 +661,8 @@ def _aggregate_weights(self, tag: str) -> None: # aggregation goal if self._agg_goal_cnt == self._agg_goal: logger.info( - f"reached agg goal since _agg_goal_cnt: {self._agg_goal_cnt} and _agg_goal is: {self._agg_goal}" - ) - logger.debug( - f"Reached agg_goal {self._agg_goal}, " - f"current _updates_in_queue: {self._updates_in_queue}, " - f"current round before agg: {self._round}" + f"agg_goal={self._agg_goal} reached, round={self._round}" ) - - # update per-trainer participation in round agg for trainer_update in self._per_round_update_list: if ( trainer_update @@ -709,10 +681,6 @@ def _aggregate_weights(self, tag: str) -> None: self._round - 1 ] = 1 - # Computing rate: Not used anywhere right now rate = 1 / - # math.sqrt(1 + self._round - tres.version) logger.debug(f" - # rate at top_agg: {rate}") - self.weights = self.optimizer.scale_add_agg_weights( self.weights, self._agg_goal_weights, self._agg_goal ) @@ -920,153 +888,92 @@ def check_trainer_availability(self, end: str) -> bool: return picked_trainer_is_available def _distribute_weights(self, tag: str, task_to_perform: str = "train") -> None: - """Distribute a global model in asynchronous FL fashion. - - This method is overridden from one in synchronous top - aggregator (..top_aggregator). - """ + """Distribute a global model in asynchronous FL fashion.""" channel = self.cm.get_by_tag(tag) if not channel: logger.debug(f"channel not found for tag {tag}") return - # this call waits for at least one peer to join this channel channel.await_join() - - # before distributing weights, update it from global model self._update_weights() - # busy wait for 0.1 seconds before proceeding. This is to wait - # on distribute_weights to let the system state get updated - # before selector is invoked again - logger.debug(f"Starting busy wait at time {time.time()}") + # Brief pause so channel state settles before selector runs. time.sleep(0.1) - logger.debug(f"Ended busy wait at time {time.time()}") - # before invoking channel.ends() to select, set the - # trainer_unavail if it isn't None if self.trainer_event_dict is not None: curr_unavail_trainer_list = self.get_curr_unavail_trainers() channel.set_curr_unavailable_trainers( trainer_unavail_list=curr_unavail_trainer_list ) - logger.debug( - f"Passed curr_unavail_trainer_list: " - f"{curr_unavail_trainer_list} to channel" - ) else: - # Handling the case for oort's selector since it expects 3 - # arguments channel.set_curr_unavailable_trainers(trainer_unavail_list=[]) - # check if there are any ends to send weights to + # Expose current vclock to selector so it can attach it to selection events. + if self.simulated: + channel.properties["vclock_now"] = self._vclock.now - logger.debug( - f"Sending weights to trainers with task_to_perform = {task_to_perform}" - ) ends = channel.ends(VAL_CH_STATE_SEND, task_to_perform) - # NRL TODO: else will take care of randomly selecting x - # trainers for "eval only" operation if not ends: - logger.debug( - f"No trainers found for tag {tag}, will " - f"move to get() for fetch weights from trainers" - ) + logger.debug(f"No trainers found for tag {tag}") return - # send out global model parameters to trainers ends_list = list(ends) for idx, end in enumerate(ends_list): - # AsyncFL checks: Similar to FedBuff, prevent sending weights to trainers that: - # 1. Already have been sent weights for current version and haven't responded - # 2. Are already selected in the current round - - # Check 1: Has trainer been sent this version but not responded yet? - if end in self._track_trainer_version_duration_s.keys(): + if end in self._track_trainer_version_duration_s: sent_versions = self._track_trainer_version_duration_s[end]["sent_wts_version_ts"] recv_versions = self._track_trainer_version_duration_s[end]["recv_wts_version_ts"] - - # If current version was sent but not received, skip if self._round in sent_versions and self._round not in recv_versions: logger.warning( f"[SELECTION_CHECK] Skipping {end}: already sent model_version={self._round} " - f"but no response received yet. Sent at {sent_versions[self._round]}, " - f"sent_versions={list(sent_versions.keys())}, recv_versions={list(recv_versions.keys())}" + f"but no response received yet." ) continue - - # Check 2: Are there any unreturned versions (sent but not received)? - unreturned_versions = [v for v in sent_versions.keys() if v not in recv_versions.keys()] - if unreturned_versions: + unreturned = [v for v in sent_versions if v not in recv_versions] + if unreturned: logger.warning( - f"[SELECTION_CHECK] Trainer {end} has {len(unreturned_versions)} unreturned versions: " - f"{unreturned_versions}. Current version to send: {self._round}. " - f"This may indicate concurrent selection - proceeding but this could cause issues." + f"[SELECTION_CHECK] {end} has {len(unreturned)} unreturned versions: {unreturned}" ) - - # Send shouldn't be allowed if already sent to a trainer - # in that same round - logger.info( - f"sending weights to {end} with model_version: {self._round} for task: {task_to_perform}" - ) - # setting start time for OORT TODO: (DG) round_start_time - # for all trainers in the same round may not be the same - logger.debug( - f"Setting channel property {PROP_ROUND_START_TIME} for " - f"end {end}. For round {self._round} at time: {datetime.now()}" + logger.info( + f"sending weights to {end} model_version={self._round} task={task_to_perform}" ) channel.set_end_property( end, PROP_ROUND_START_TIME, (self._round, datetime.now()) ) - # we use _round to indicate a model version - channel.send( - end, - { - MessageType.WEIGHTS: weights_to_device( - self.weights, DeviceType.CPU - ), - MessageType.ROUND: self._round, - MessageType.MODEL_VERSION: self._round, - MessageType.TASK_TO_PERFORM: task_to_perform, - }, - ) - - # Update send_time in training_duration_s - if end not in self._track_trainer_version_duration_s.keys(): + msg = { + MessageType.WEIGHTS: weights_to_device(self.weights, DeviceType.CPU), + MessageType.ROUND: self._round, + MessageType.MODEL_VERSION: self._round, + MessageType.TASK_TO_PERFORM: task_to_perform, + } + # Stamp virtual send-time so trainer can compute sim_completion_ts correctly. + # Without this, _sim_send_ts stays None on the trainer → _sim_now()=0 → + # sim_completion_ts = D for all trainers regardless of when task was dispatched, + # collapsing the reorder buffer to sort by tiny GPU-time differences only. + 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) logger.debug( - f"{end} not in _track_trainer_version_duration_s, " f"will add" + f"[SIM_SEND] end={end[-4:]} round={self._round} sim_send_ts={sim_send_ts:.2f}" ) - self._track_trainer_version_duration_s[end] = dict() - self._track_trainer_version_duration_s[end]["last_send_wts_ts"] = -1 - - # sent_wts_version_ts, recv_wts_version_ts is a dict - # of version sent/recv and its timestamp. This will be - # primarily used by AsyncOORT selector since it needs - # round_duration times. TODO: (DG) Right now the dict - # maintains ALL sent/recv versions and timestamps for - # all trainers. For thousands of trainers it might - # incur memory-bloat. Can optimize to retain just the - # versions and timestamps of those that were sent but - # not received back for the trainer. - self._track_trainer_version_duration_s[end]["sent_wts_version_ts"] = {} - self._track_trainer_version_duration_s[end]["recv_wts_version_ts"] = {} - self._track_trainer_version_duration_s[end][ - "total_training_time_s" - ] = -1 - # Update sent_wts_version_ts with version and timestamp + channel.send(end, msg) + + if end not in self._track_trainer_version_duration_s: + self._track_trainer_version_duration_s[end] = { + "last_send_wts_ts": -1, + "sent_wts_version_ts": {}, + "recv_wts_version_ts": {}, + "total_training_time_s": -1, + } + self._track_trainer_version_duration_s[end]["sent_wts_version_ts"][ self._round ] = datetime.now() - - # 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 - # 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. + + # Stagger sends: shorter interval in sim (no real sleeps) vs real mode. if idx < len(ends_list) - 1: time.sleep(0.2 if self.simulated else 0.5) diff --git a/lib/python/flame/mode/horizontal/syncfl/trainer.py b/lib/python/flame/mode/horizontal/syncfl/trainer.py index 0b63739b6..14be7c3df 100644 --- a/lib/python/flame/mode/horizontal/syncfl/trainer.py +++ b/lib/python/flame/mode/horizontal/syncfl/trainer.py @@ -43,6 +43,8 @@ from flame.optimizers import optimizer_provider from flame.privacies import privacy_provider from flame.registries import registry_provider +from flame import telemetry +from flame.telemetry.events import build_task_recv # TODO: (DG) torch is needed for asyncoort in oort_loss() function, # but need to comment / uncomment based on the backend used. If it is @@ -223,12 +225,23 @@ 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. + # Capture virtual send-time stamped by aggregator (sim mode); used for sim_completion_ts. if MessageType.SIM_SEND_TS in msg: self._sim_send_ts = msg[MessageType.SIM_SEND_TS] + if telemetry.is_enabled(): + _sim_send_ts_val = getattr(self, "_sim_send_ts", None) + _time_mode = getattr(self, "time_mode", "real") + _avl = getattr(getattr(self, "avl_state", None), "value", None) + ev, fields = build_task_recv( + round_num=int(self._round), + trainer_id=str(getattr(self, "trainer_id", "")), + time_mode=_time_mode, + sim_send_ts=float(_sim_send_ts_val) if _sim_send_ts_val is not None else None, + avl_state=_avl, + ) + telemetry.emit(ev, **fields) + if MessageType.EOT in msg: self._work_done = msg[MessageType.EOT] @@ -377,6 +390,10 @@ def _send_weights(self, tag: str) -> None: self, "_sim_round_duration", 0.0 ) + _budget = getattr(self, "_training_budget_s", None) + if _budget is not None: + msg[MessageType.TRAINING_BUDGET_S] = float(_budget) + 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 cb2f57c2b..f086746c0 100644 --- a/lib/python/flame/mode/message.py +++ b/lib/python/flame/mode/message.py @@ -79,6 +79,7 @@ class MessageType(Enum): # 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) + SIM_SEND_TS = 34 # virtual clock T_v at weight distribution + SIM_COMPLETION_TS = 35 # sim_send_ts + max(gpu, D); reorder buffer commits by this + SIM_ROUND_DURATION = 36 # max(gpu_time, D) = true round cost; mirrors real recv_ts-sent_ts for OORT speed utility + TRAINING_BUDGET_S = 37 # trainer's configured total round budget (training_delay_s) diff --git a/lib/python/flame/selector/async_oort.py b/lib/python/flame/selector/async_oort.py index a0cb15005..c9b5ae2fa 100644 --- a/lib/python/flame/selector/async_oort.py +++ b/lib/python/flame/selector/async_oort.py @@ -389,16 +389,27 @@ def select( ) self._select_run_counter = 0 + per_trainer_extra = { + eid: { + "in_all_selected": eid in self.all_selected, + "in_pending_commit": eid in getattr(self, "_agg_pending_commit_ref", set()), + "last_eval_round": ends[eid].get_property(PROP_LAST_EVAL_ROUND), + } + for eid in ends + } self.emit_selection( channel_props.get("round", 0), task_to_perform, ends, eligible_ends.keys(), list(results.keys()), + per_trainer_extra=per_trainer_extra, extra={ "concurrency": concurrency, "effective_c": effective_c, "requester": self.requester, + "vclock_now": channel_props.get("vclock_now"), + "exploration_factor": self.exploration_factor, }, ) @@ -465,11 +476,23 @@ def sample_by_util( for prob_idx in range(len(over_cutoff_utility_probs)): over_cutoff_utility_probs[prob_idx] /= over_cutoff_utility_sum + # Exclude zero-probability entries; np.random.choice(replace=False) requires ≥size non-zero. + nz_pairs = [ + (e, p) + for e, p in zip(over_cutoff_utility_end_ids, over_cutoff_utility_probs) + if p > 0 + ] + if not nz_pairs: + return [] + nz_ends, nz_probs = zip(*nz_pairs) + nz_total = sum(nz_probs) + nz_probs = [p / nz_total for p in nz_probs] + selected_ends = np.random.choice( - over_cutoff_utility_end_ids, - size=min(len(over_cutoff_utility_end_ids), num_of_ends), + list(nz_ends), + size=min(len(nz_ends), num_of_ends), replace=False, - p=over_cutoff_utility_probs, + p=nz_probs, ) return selected_ends diff --git a/lib/python/flame/telemetry/events.py b/lib/python/flame/telemetry/events.py index cea3c3a8d..841cc3912 100644 --- a/lib/python/flame/telemetry/events.py +++ b/lib/python/flame/telemetry/events.py @@ -21,6 +21,7 @@ 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 +EVENT_TASK_RECV = "task_recv" # trainer received a task from aggregator KNOWN_EVENTS = frozenset( { @@ -31,6 +32,7 @@ EVENT_TRAINER_ROUND, EVENT_UTIL_DISPARITY, EVENT_AVAIL_CHANGE, + EVENT_TASK_RECV, } ) @@ -200,3 +202,27 @@ def build_avail_change( "old_state": old_state, "new_state": new_state, } + + +def build_task_recv( + *, + round_num: int, + trainer_id: str, + time_mode: str, + sim_send_ts: Optional[float] = None, + avl_state: Optional[str] = None, +) -> tuple[str, dict[str, Any]]: + """Trainer received a task (weights) from the aggregator. + + sim_send_ts: the virtual clock value stamped by the aggregator (sim mode only). + Emitting None in real mode for both sim_send_ts and vclock makes the per-round + trainer state directly comparable between real and sim telemetry. + """ + fields: dict[str, Any] = { + "round": round_num, + "trainer_id": trainer_id, + "time_mode": time_mode, + "sim_send_ts": sim_send_ts, + "avl_state": avl_state, + } + return EVENT_TASK_RECV, fields