diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 0736cb0e0..059604e65 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -43,6 +43,7 @@ import argparse import csv +import json import logging import warnings from pathlib import Path @@ -82,6 +83,8 @@ def __init__( run_name: str, tags: list[str], config: Dict[str, Any], + group: Optional[str] = None, + job_type: Optional[str] = None, ) -> None: self.enabled = enabled self._run = None @@ -94,6 +97,8 @@ def __init__( name=run_name, tags=tags, config=config, + group=group, + job_type=job_type, ) def log(self, data: Dict[str, Any], step: Optional[int] = None) -> None: @@ -391,7 +396,15 @@ def run(args: argparse.Namespace) -> Path: else None ) - exp_name = f"{args.task}_{args.model}_seed{args.seed}" + # The window belongs in the name: an observation-window arm is a different + # experiment from the full-stay run at the same task/model/seed, and without + # the suffix the two share an output directory and a W&B run name. + window_suffix = ( + f"_w{int(args.observation_window_hours)}" + if args.observation_window_hours + else "" + ) + exp_name = f"{args.task}_{args.model}_seed{args.seed}{window_suffix}" output_dir = Path(args.output_dir) wandb_logger = WandbLogger( @@ -401,6 +414,10 @@ def run(args: argparse.Namespace) -> Path: run_name=args.wandb_run_name or exp_name, tags=args.wandb_tags.split(",") if args.wandb_tags else [args.task, args.model], config=vars(args), + # Group by arm and split by backbone so a many-cell sweep is navigable + # instead of one flat list of runs. + group=f"{args.task}{window_suffix}", + job_type=args.model, ) trainer = Trainer( @@ -438,9 +455,15 @@ def run(args: argparse.Namespace) -> Path: for epoch_record in metrics_history: wandb_logger.log(epoch_record, step=epoch_record["epoch"]) - if wandb_logger.enabled and test_loader is not None: + # Test evaluation must not depend on the logger. This was gated on + # wandb_logger.enabled, so a run without --wandb never computed test + # metrics at all -- they were not merely unlogged, they were never + # calculated. + test_scores = None + if test_loader is not None: test_scores = trainer.evaluate(test_loader) - wandb_logger.log({f"test_{k}": v for k, v in test_scores.items()}) + if wandb_logger.enabled: + wandb_logger.log({f"test_{k}": v for k, v in test_scores.items()}) if test_loader is not None: inference_loader, eval_split = test_loader, "test" @@ -469,6 +492,14 @@ def run(args: argparse.Namespace) -> Path: }, ) + # metrics_history.json carries validation only, so without this the test + # numbers that go in the paper live nowhere on disk -- only in stdout and + # W&B, and are recoverable afterwards only by re-scoring predictions. + if test_scores is not None: + test_path = output_dir / exp_name / "test_metrics.json" + with open(test_path, "w") as handle: + json.dump({"eval_split": eval_split, **test_scores}, handle, indent=2) + output_csv = output_dir / exp_name / f"predictions_{args.model}.csv" _write_predictions(output_csv, patient_ids, y_true, y_prob) diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index 2085221b3..f67bcbc84 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -90,6 +90,38 @@ def _vram_stats(device: str) -> Dict[str, float]: return {"vram_allocated_mb": allocated, "vram_peak_mb": peak} +def _cpu_seconds() -> Optional[float]: + """Cumulative CPU seconds for this process and its dataloader workers. + + Self-only time badly understates a run whose cost is data loading, since + workers are separate processes. psutil is already available via wandb; the + resource fallback only counts children that have been reaped, so it reads + low while persistent workers are still alive. + """ + try: + import psutil + + proc = psutil.Process() + times = proc.cpu_times() + total = times.user + times.system + for child in proc.children(recursive=True): + try: + ctimes = child.cpu_times() + total += ctimes.user + ctimes.system + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + return total + except Exception: + try: + import resource + + me = resource.getrusage(resource.RUSAGE_SELF) + kids = resource.getrusage(resource.RUSAGE_CHILDREN) + return me.ru_utime + me.ru_stime + kids.ru_utime + kids.ru_stime + except Exception: + return None + + def get_metrics_fn(mode: str) -> Callable: if mode == "binary": return binary_metrics_fn @@ -270,6 +302,7 @@ def train( if torch.cuda.is_available() and str(self.device).startswith("cuda"): torch.cuda.reset_peak_memory_stats(self.device) epoch_start = time.perf_counter() + cpu_start = _cpu_seconds() # batch training loop logger.info("") for step_idx in trange( @@ -336,6 +369,16 @@ def train( epoch_time = time.perf_counter() - epoch_start vram = _vram_stats(self.device) + cpu = {} + cpu_end = _cpu_seconds() + if cpu_start is not None and cpu_end is not None: + cpu_s = max(cpu_end - cpu_start, 0.0) + cpu = { + "cpu_seconds": round(cpu_s, 2), + # >100% means several cores busy, which is the normal case + # with dataloader workers. + "cpu_util_pct": round(100.0 * cpu_s / max(epoch_time, 1e-9), 1), + } epochs_done = epoch + 1 epochs_left = epochs - epochs_done @@ -372,6 +415,7 @@ def train( "epoch_time_s": round(epoch_time, 3), "skipped_steps": epoch_skipped_steps, **{f"train_{k}": v for k, v in vram.items()}, + **{f"train_{k}": v for k, v in cpu.items()}, } # validation diff --git a/pyhealth/utils.py b/pyhealth/utils.py index b46c66ac4..e3653afbc 100644 --- a/pyhealth/utils.py +++ b/pyhealth/utils.py @@ -23,8 +23,11 @@ def set_seed(seed): def create_directory(directory): - if not os.path.exists(directory): - os.makedirs(directory) + # exist_ok, not a prior exists() check: two processes importing pyhealth + # for the first time both pass the check and one loses the makedirs race + # with FileExistsError. Seen on a shared cluster home, where two concurrent + # jobs both tried to create ~/.cache/pyhealth/medcode/. + os.makedirs(directory, exist_ok=True) def load_pickle(filename):