Persist test metrics, log CPU usage, and make run identity window-aware

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)
 

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

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):