Spaces:
Paused
Paused
| """Distributed fixed-bank diagnostics for v6 dose-intervention tasks.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from pathlib import Path | |
| import lightning.pytorch as pl | |
| import torch | |
| import torch.distributed as dist | |
| from pff_pk.evaluation.bank import load_bank, record_to_batch | |
| from pff_pk.inference.protocol_validation import ( | |
| PROTOCOL_CATEGORIES, | |
| generate_protocol_draws, | |
| plot_protocol_category, | |
| protocol_seed, | |
| select_representatives, | |
| ) | |
| class ProtocolValidationDiagnostics(pl.Callback): | |
| """Persist and plot the same 42 protocol strata at fixed step intervals. | |
| Every DDP rank evaluates a disjoint slice of the fixed validation bank. The | |
| resulting CPU records are gathered on rank zero, saved with full provenance, | |
| rendered into the four protocol panels, and logged to the existing Comet run. | |
| The bank and per-stratum random seeds remain fixed across checkpoints. | |
| """ | |
| def __init__( | |
| self, | |
| bank_path: str | Path, | |
| output_directory: str | Path, | |
| *, | |
| interval: int, | |
| samples: int, | |
| ) -> None: | |
| super().__init__() | |
| self.bank_path = Path(bank_path) | |
| self.output_directory = Path(output_directory) | |
| self.interval = interval | |
| self.samples = samples | |
| self._logged: set[int] = set() | |
| self._bank: dict | None = None | |
| def _completed_milestones(self) -> set[int]: | |
| """Recover completed milestones so resumed jobs do not duplicate diagnostics.""" | |
| completed: set[int] = set() | |
| if not self.output_directory.is_dir(): | |
| return completed | |
| for manifest in self.output_directory.glob("step*/manifest.json"): | |
| try: | |
| payload = json.loads(manifest.read_text(encoding="utf-8")) | |
| step = int(payload["checkpoint_step"]) | |
| milestone = int( | |
| payload.get("trigger_milestone", step // self.interval * self.interval) | |
| ) | |
| except (KeyError, TypeError, ValueError, json.JSONDecodeError, OSError): | |
| continue | |
| if milestone > 0: | |
| completed.add(milestone) | |
| return completed | |
| def _load_bank(self) -> dict: | |
| if self._bank is None: | |
| self._bank = load_bank(self.bank_path) | |
| return self._bank | |
| def on_validation_epoch_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None: | |
| step = int(trainer.global_step) | |
| milestone = step // self.interval * self.interval | |
| if milestone == 0: | |
| return | |
| if not self._logged: | |
| self._logged.update(self._completed_milestones()) | |
| if milestone in self._logged: | |
| return | |
| bank = self._load_bank() | |
| rank = int(trainer.global_rank) | |
| world_size = int(trainer.world_size) | |
| representatives = select_representatives(bank) | |
| selected = representatives[rank::world_size] | |
| device = pl_module.device | |
| cuda_devices = [device.index] if device.type == "cuda" else [] | |
| was_training = pl_module.training | |
| local_records = [] | |
| pl_module.eval() | |
| try: | |
| with torch.random.fork_rng(devices=cuda_devices), torch.inference_mode(): | |
| for stratum, record_index, row in selected: | |
| batch = record_to_batch(bank["batches"][record_index], torch.device("cpu")) | |
| result = generate_protocol_draws( | |
| pl_module, batch, row, self.samples, protocol_seed(stratum) | |
| ) | |
| local_records.append({"stratum": stratum, **result}) | |
| finally: | |
| pl_module.train(was_training) | |
| gathered: list[list[dict] | None] | None | |
| if dist.is_available() and dist.is_initialized(): | |
| gathered = [None] * world_size if trainer.is_global_zero else None | |
| dist.gather_object(local_records, gathered, dst=0) | |
| else: | |
| gathered = [local_records] | |
| if not trainer.is_global_zero: | |
| self._logged.add(milestone) | |
| return | |
| assert gathered is not None | |
| records = { | |
| record["stratum"]: record | |
| for rank_records in gathered | |
| for record in (rank_records or []) | |
| } | |
| if len(records) != 42: | |
| raise RuntimeError(f"expected 42 protocol strata, found {len(records)}") | |
| output = self.output_directory / f"step{step:06d}" | |
| output.mkdir(parents=True, exist_ok=True) | |
| draws_path = output / "protocol_draws.pt" | |
| torch.save( | |
| { | |
| "format": "v6-protocol-validation-draws-v1", | |
| "checkpoint_step": step, | |
| "trigger_milestone": milestone, | |
| "validation_bank": str(self.bank_path.resolve()), | |
| "validation_bank_sha256": bank["sha256"], | |
| "samples_per_task": self.samples, | |
| "records": [records[key] for key in sorted(records)], | |
| }, | |
| draws_path, | |
| ) | |
| figures = [ | |
| plot_protocol_category(records, category, cases, output, step) | |
| for category, cases in PROTOCOL_CATEGORIES.items() | |
| ] | |
| manifest = { | |
| "format": "v6-protocol-validation-draws-v1", | |
| "checkpoint_step": step, | |
| "trigger_milestone": milestone, | |
| "validation_bank": str(self.bank_path.resolve()), | |
| "validation_bank_sha256": bank["sha256"], | |
| "samples_per_task": self.samples, | |
| "strata": sorted(records), | |
| "draws": draws_path.name, | |
| "figures": [path.name for pair in figures for path in pair], | |
| } | |
| manifest_path = output / "manifest.json" | |
| manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") | |
| comet_step = step + int(os.environ.get("PFFF_COMET_STEP_OFFSET", "0")) | |
| experiment = trainer.logger.experiment | |
| for png, pdf in figures: | |
| category = png.stem.removeprefix("protocol_validation_").rsplit("_step", 1)[0] | |
| experiment.log_image( | |
| image_data=str(png), | |
| name=f"validation/protocol_stratified/{category}.png", | |
| step=comet_step, | |
| ) | |
| experiment.log_asset( | |
| str(pdf), | |
| file_name=f"validation/protocol_stratified/step_{step}/{pdf.name}", | |
| step=comet_step, | |
| ) | |
| experiment.log_asset( | |
| str(draws_path), | |
| file_name=f"validation/protocol_stratified/step_{step}/protocol_draws.pt", | |
| step=comet_step, | |
| ) | |
| experiment.log_asset( | |
| str(manifest_path), | |
| file_name=f"validation/protocol_stratified/step_{step}/manifest.json", | |
| step=comet_step, | |
| ) | |
| self._logged.add(milestone) | |