diff --git a/v4/classes/profiles/fundus_images.py b/v4/classes/profiles/fundus_images.py index e475a0c..109af42 100644 --- a/v4/classes/profiles/fundus_images.py +++ b/v4/classes/profiles/fundus_images.py @@ -124,71 +124,47 @@ def compute_geometry_features(disc_mask: np.ndarray, cup_mask: np.ndarray) -> np # --------------------------------------------------------------------------- class GTGeometryLoader: - """Pre-computes per-eye geometry vectors from PAPILA GT contour annotations. + """Per-eye 5-feature CDR vectors from PAPILA GT contour annotations. - File naming: RET{pid:03d}{eye}_{disc|cup}_exp{n}.txt - Averages exp1 and exp2 when both are present; zero vector for missing entries. - - Usage: - loader = GTGeometryLoader(contour_dir) - loader.precompute(df, patient_col="Patient ID") - vecs = loader.all_vectors() # {(pid, eye): ndarray} + Loads contour coordinates from per-expert text files, rasterises them at + the original image's pixel space (so polygons aren't clipped), then + computes the CDR feature vector. Expert masks are merged before feature + computation to match the seg-map path. """ - _EXPERTS = (1, 2) - feature_dim = _FEATURE_DIM feature_names = _FEATURE_NAMES - def __init__(self, contour_dir: str | Path) -> None: - self._dir = Path(contour_dir) + def __init__(self, contour_dir: str | Path, *, mask_size: int = _MASK_SIZE[0]) -> None: + self._dir = Path(contour_dir) + self._mask_size = mask_size self._cache: dict[tuple, np.ndarray] = {} - def precompute(self, df, patient_col: str = "Patient ID") -> None: + def reset_cache(self) -> None: + self._cache.clear() + + def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None: n_ok = 0 - for _, row in df.iterrows(): - pid = int(row[patient_col]) - eye = str(row.get("eyeID", "OD")) + for pid, eye, image_path in samples: key = (pid, eye) if key in self._cache: continue - vec = self._compute(pid, eye) - self._cache[key] = vec if vec is not None else np.zeros(self.feature_dim, dtype=np.float32) - if vec is not None: + with Image.open(image_path) as img: + image_size = img.size # (W, H) + res = _papila_disc_cup_masks( + pid, eye, self._dir, image_size, self._mask_size, + ) + if res is None: + self._cache[key] = np.zeros(self.feature_dim, dtype=np.float32) + else: + disc, cup = res + self._cache[key] = compute_geometry_features(disc, cup) n_ok += 1 print(f"[GTGeometryLoader] {n_ok}/{len(self._cache)} geometry vectors computed", flush=True) def all_vectors(self) -> dict: return dict(self._cache) - def _compute(self, pid: int, eye: str) -> "np.ndarray | None": - stem = f"RET{pid:03d}{eye}" - vecs: list[np.ndarray] = [] - for exp in self._EXPERTS: - disc_path = self._dir / f"{stem}_disc_exp{exp}.txt" - cup_path = self._dir / f"{stem}_cup_exp{exp}.txt" - if not disc_path.exists(): - continue - try: - disc_c = np.loadtxt(disc_path) - if disc_c.ndim == 1: - disc_c = disc_c.reshape(-1, 2) - disc_mask = _contour_to_mask(disc_c, _MASK_SIZE) - if cup_path.exists(): - cup_c = np.loadtxt(cup_path) - if cup_c.ndim == 1: - cup_c = cup_c.reshape(-1, 2) - cup_mask = _contour_to_mask(cup_c, _MASK_SIZE) - else: - cup_mask = np.zeros((_MASK_SIZE[1], _MASK_SIZE[0]), dtype=np.uint8) - cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8) - vecs.append(compute_geometry_features(disc_mask, cup_mask)) - except Exception: - continue - if not vecs: - return None - return np.stack(vecs).mean(axis=0).astype(np.float32) - # --------------------------------------------------------------------------- # Seg-map utilities (shared by GT and UNet seg-map loaders) @@ -411,12 +387,103 @@ class _UNetFTDataset(Dataset): return self._imgs[idx], self._masks[idx] +class _PapilaUNetMaskPipeline: + """Per-fold UNet mask producer. + + Owns a UNetSegmenter, handles the per-fold lifecycle: + - reset_weights() restores REFUGE base state (call at start of each fold) + - finetune(train_samples) fine-tunes on the train split's GT contours + - predict(samples) returns {(pid, eye): (disc_mask, cup_mask)} raw masks + + Downstream loaders interpret the raw masks differently — seg-map loader + crops/resizes/encodes for a CNN, geometry-vector loader computes 5 CDR + features. Sharing this pipeline avoids duplicating UNet load + finetune + when both seg-map and feature-vector outputs are needed in one config. + """ + + def __init__( + self, + weights_path: str | Path, + *, + contour_dir: str | Path, + unet_size: int = 512, + normalize: str = "per_image", + threshold: float = 0.5, + finetune_epochs: int = 0, + finetune_lr: float = 1e-5, + finetune_batch_size: int = 4, + device: str | None = None, + ) -> None: + from v4.classes.accessory.unet import UNetSegmenter + + self._contour_dir = Path(contour_dir) + self._threshold = threshold + self._ft_epochs = finetune_epochs + self._ft_lr = finetune_lr + self._ft_batch_size = finetune_batch_size + + self._segmenter = UNetSegmenter( + target_size=unet_size, normalize=normalize, device=device, + ).load_weights(Path(weights_path)) + self._base_state = copy.deepcopy(self._segmenter.model.state_dict()) + + def reset_weights(self) -> None: + """Restore base REFUGE weights (undo any prior fine-tuning).""" + self._segmenter.model.load_state_dict(copy.deepcopy(self._base_state)) + + def finetune(self, train_samples: list) -> None: + """Fine-tune the UNet on the training fold's GT contours.""" + if self._ft_epochs <= 0: + return + ds = _UNetFTDataset(train_samples, self._contour_dir, self._segmenter) + loader = DataLoader( + ds, batch_size=self._ft_batch_size, shuffle=True, num_workers=0, + ) + print( + f"[UNetMaskPipeline] fine-tuning UNet for {self._ft_epochs} epochs " + f"on {len(train_samples)} samples (lr={self._ft_lr}, " + f"batch_size={self._ft_batch_size})", + flush=True, + ) + self._segmenter.finetune( + loader, epochs=self._ft_epochs, lr=self._ft_lr, + log_prefix="[UNet ft]", + ) + + def predict( + self, samples: Iterable[Tuple[int, str, Path]], + ) -> dict[tuple, Tuple[np.ndarray, np.ndarray]]: + """Run inference. Returns {(pid, eye): (disc_mask, cup_mask)} raw uint8.""" + import time + samples = list(samples) + if not samples: + return {} + print( + f"[UNetMaskPipeline] running UNet inference on {len(samples)} images...", + flush=True, + ) + result: dict[tuple, Tuple[np.ndarray, np.ndarray]] = {} + t0 = time.time() + report = max(1, len(samples) // 4) + for i, (pid, eye, image_path) in enumerate(samples, 1): + with Image.open(image_path) as raw: + disc, cup = self._segmenter.predict(raw, threshold=self._threshold) + result[(pid, eye)] = (disc, cup) + if i % report == 0 or i == len(samples): + print( + f" [UNet inf] {i}/{len(samples)} ({time.time() - t0:.1f}s)", + flush=True, + ) + return result + + class UNetSegMapLoader: - """Pre-computes per-eye seg maps via a REFUGE-pretrained UNet. + """Per-eye CNN-ready seg maps via a REFUGE-pretrained UNet. - Optionally fine-tunes the UNet per fold on the training split's GT contours. + Wraps a `_PapilaUNetMaskPipeline` and post-processes raw masks into + (C, H, W) float32 arrays sized for a downstream CNN. - Output: dict {(pid, eye): np.ndarray (C, H, W) float32} cached for the fold. + Output: dict {(pid, eye): np.ndarray (C, H, W) float32} """ def __init__( @@ -435,22 +502,20 @@ class UNetSegMapLoader: finetune_batch_size: int = 4, device: str | None = None, ) -> None: - from v4.classes.accessory.unet import UNetSegmenter - - self._weights_path = Path(weights_path) - self._contour_dir = Path(contour_dir) - self._channels = channels - self._target_size = target_size - self._threshold = threshold - self._crop = crop_to_disc - self._ft_epochs = finetune_epochs - self._ft_lr = finetune_lr - self._ft_batch_size = finetune_batch_size - - self._segmenter = UNetSegmenter( - target_size=unet_size, normalize=normalize, device=device, - ).load_weights(self._weights_path) - self._base_state = copy.deepcopy(self._segmenter.model.state_dict()) + self._pipeline = _PapilaUNetMaskPipeline( + weights_path, + contour_dir=contour_dir, + unet_size=unet_size, + normalize=normalize, + threshold=threshold, + finetune_epochs=finetune_epochs, + finetune_lr=finetune_lr, + finetune_batch_size=finetune_batch_size, + device=device, + ) + self._channels = channels + self._target_size = target_size + self._crop = crop_to_disc self._cache: dict[tuple, np.ndarray] = {} @property @@ -458,89 +523,130 @@ class UNetSegMapLoader: return (self._channels, self._target_size, self._target_size) def reset_cache(self) -> None: - """Clear cached seg maps (call between folds).""" self._cache.clear() def reset_weights(self) -> None: - """Restore base REFUGE weights (undo any prior fine-tuning).""" - self._segmenter.model.load_state_dict(copy.deepcopy(self._base_state)) + self._pipeline.reset_weights() def finetune(self, train_samples: list) -> None: - """Fine-tune the UNet on the training fold's GT contours. - - train_samples: list of (pid, eye, image_path) tuples — train split only. - """ - if self._ft_epochs <= 0: - return - ds = _UNetFTDataset(train_samples, self._contour_dir, self._segmenter) - loader = DataLoader( - ds, batch_size=self._ft_batch_size, shuffle=True, num_workers=0, - ) - print( - f"[UNetSegMapLoader] fine-tuning UNet for {self._ft_epochs} epochs " - f"on {len(train_samples)} samples (lr={self._ft_lr}, " - f"batch_size={self._ft_batch_size})", - flush=True, - ) - self._segmenter.finetune( - loader, epochs=self._ft_epochs, lr=self._ft_lr, - log_prefix="[UNet ft]", - ) + self._pipeline.finetune(train_samples) def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None: - """Run UNet inference on every sample and cache the resulting seg map.""" - import time - samples = list(samples) - todo = [s for s in samples if (s[0], s[1]) not in self._cache] - if not todo: - return - print( - f"[UNetSegMapLoader] running UNet inference on {len(todo)} images...", - flush=True, - ) - t0 = time.time() - report = max(1, len(todo) // 4) - for i, (pid, eye, image_path) in enumerate(todo, 1): - with Image.open(image_path) as raw: - disc, cup = self._segmenter.predict(raw, threshold=self._threshold) + todo = [s for s in samples if (s[0], s[1]) not in self._cache] + masks = self._pipeline.predict(todo) + for (pid, eye), (disc, cup) in masks.items(): seg = _combine_disc_cup(disc, cup) if self._crop: seg = _crop_to_disc_bbox(seg) self._cache[(pid, eye)] = _seg_map_to_array( seg, self._channels, self._target_size, ) - if i % report == 0 or i == len(todo): - print( - f" [UNet inf] {i}/{len(todo)} ({time.time() - t0:.1f}s)", - flush=True, - ) - print( - f"[UNetSegMapLoader] {len(todo)} seg maps cached via UNet " - f"(channels={self._channels}, target={self._target_size})", - flush=True, - ) + if masks: + print( + f"[UNetSegMapLoader] {len(masks)} seg maps cached " + f"(channels={self._channels}, target={self._target_size})", + flush=True, + ) def all_seg_maps(self) -> dict: return self._cache +class UNetGeometryLoader: + """Per-eye 5-feature CDR vectors derived from UNet-predicted masks. + + Same UNet lifecycle as `UNetSegMapLoader` but the output is a 5-vector + (compute_geometry_features over the predicted disc/cup masks) rather + than a CNN-ready seg map. Designed to slot into ImageEncoder's + geometry_source mechanism for vector-style geometry injection. + """ + + feature_dim = _FEATURE_DIM + feature_names = _FEATURE_NAMES + + def __init__( + self, + weights_path: str | Path, + *, + contour_dir: str | Path, + unet_size: int = 512, + normalize: str = "per_image", + threshold: float = 0.5, + finetune_epochs: int = 0, + finetune_lr: float = 1e-5, + finetune_batch_size: int = 4, + device: str | None = None, + ) -> None: + self._pipeline = _PapilaUNetMaskPipeline( + weights_path, + contour_dir=contour_dir, + unet_size=unet_size, + normalize=normalize, + threshold=threshold, + finetune_epochs=finetune_epochs, + finetune_lr=finetune_lr, + finetune_batch_size=finetune_batch_size, + device=device, + ) + self._cache: dict[tuple, np.ndarray] = {} + + def reset_cache(self) -> None: + self._cache.clear() + + def reset_weights(self) -> None: + self._pipeline.reset_weights() + + def finetune(self, train_samples: list) -> None: + self._pipeline.finetune(train_samples) + + def precompute(self, samples: Iterable[Tuple[int, str, Path]]) -> None: + todo = [s for s in samples if (s[0], s[1]) not in self._cache] + masks = self._pipeline.predict(todo) + for (pid, eye), (disc, cup) in masks.items(): + self._cache[(pid, eye)] = compute_geometry_features(disc, cup) + if masks: + print( + f"[UNetGeometryLoader] {len(masks)} geometry vectors cached " + f"(dim={self.feature_dim})", + flush=True, + ) + + def all_vectors(self) -> dict: + return dict(self._cache) + + # --------------------------------------------------------------------------- # Factories # --------------------------------------------------------------------------- def build_geometry_loader(source: str, **kwargs): - """Return the appropriate geometry-vector loader for the given source string. + """Return the appropriate geometry-vector loader for the given source. Parameters ---------- - source : "gt" | "unet" - contour_dir : (gt) path to contour annotation directory + source : "gt" | "unet" + + GT kwargs: + contour_dir + UNet kwargs: + weights_path, contour_dir, unet_size=512, normalize="per_image", + threshold=0.5, finetune_epochs=0, finetune_lr=1e-5, + finetune_batch_size=4, device=None """ if source == "gt": contour_dir = kwargs.get("contour_dir") if contour_dir is None: raise ValueError("build_geometry_loader source='gt' requires contour_dir") return GTGeometryLoader(contour_dir) + if source == "unet": + if "weights_path" not in kwargs: + raise ValueError("build_geometry_loader source='unet' requires weights_path") + if "contour_dir" not in kwargs: + raise ValueError( + "build_geometry_loader source='unet' requires contour_dir " + "(needed for per-fold fine-tuning, even if finetune_epochs=0)" + ) + return UNetGeometryLoader(**kwargs) raise NotImplementedError(f"build_geometry_loader: source={source!r} not implemented") diff --git a/v4/classes/profiles/v4papila.py b/v4/classes/profiles/v4papila.py index 759e4d7..8e19273 100644 --- a/v4/classes/profiles/v4papila.py +++ b/v4/classes/profiles/v4papila.py @@ -507,6 +507,25 @@ class PapilaBundle: """ return [self._bundle.patient_col, "eyeID"] + # ── Sample collection (for tower early_pass) ───────────────────────────── + + def collect_samples(self, df: pd.DataFrame | None) -> list[tuple]: + """Build (pid, eye, image_path) tuples from a split DataFrame. + + Used by tower early_pass implementations that need per-eye image paths + (UNet inference, contour rasterisation, etc.). Returns [] for an + empty/None df. + """ + if df is None or len(df) == 0: + return [] + pc = self._bundle.patient_col + out: list[tuple] = [] + for _, row in df.iterrows(): + pid = int(row[pc]) + eye = str(row.get("eyeID", "OD")) + out.append((pid, eye, self.image.get_image_path(pid, eye))) + return out + # ── Backward-compat delegates ──────────────────────────────────────────── @property diff --git a/v4/classes/towers/geometry_tower.py b/v4/classes/towers/geometry_tower.py index 421a92a..caa35e1 100644 --- a/v4/classes/towers/geometry_tower.py +++ b/v4/classes/towers/geometry_tower.py @@ -126,16 +126,14 @@ class GeometrySegEncoder(TowerBase): # ── EPC early_pass ─────────────────────────────────────────────────────── def early_pass(self, context) -> None: - data = context.require("data") - split = context.require("split") + data = context.require("data") + split = context.require("split") - train_samples = self._collect_samples(split.train, data) - all_samples = self._collect_samples(split.train, data) - all_samples += self._collect_samples(split.val, data) + train_samples = data.collect_samples(split.train) + all_samples = train_samples + data.collect_samples(split.val) if split.test is not None: - all_samples += self._collect_samples(split.test, data) + all_samples += data.collect_samples(split.test) - # Reset per-fold state if loader supports it (UNet only). if hasattr(self._loader, "reset_cache"): self._loader.reset_cache() if hasattr(self._loader, "reset_weights"): @@ -172,18 +170,6 @@ class GeometrySegEncoder(TowerBase): # ── Internals ──────────────────────────────────────────────────────────── - def _collect_samples(self, df, data) -> list: - """Build (pid, eye, image_path) tuples from a split DataFrame.""" - if df is None or len(df) == 0: - return [] - pc = data.patient_col - out = [] - for _, row in df.iterrows(): - pid = int(row[pc]) - eye = str(row.get("eyeID", "OD")) - out.append((pid, eye, data.image.get_image_path(pid, eye))) - return out - @staticmethod def _augment_array(arr: np.ndarray) -> np.ndarray: """Random flip + 90° rotation on a (C, H, W) seg-map array.""" diff --git a/v4/classes/towers/image_tower.py b/v4/classes/towers/image_tower.py index 8ee87ad..cadf79a 100644 --- a/v4/classes/towers/image_tower.py +++ b/v4/classes/towers/image_tower.py @@ -151,11 +151,12 @@ class ImageEncoder(TowerBase): def early_pass(self, context) -> None: """Per-fold setup: warm tensor cache (if enabled), publish geometry vectors.""" - data = context.require("data") + data = context.require("data") + split = context.require("split") if self._cache_transformed: self._tensor_cache.clear() - n = self._warm_tensor_cache(data, context.require("split")) + n = self._warm_tensor_cache(data, split) print( f"[ImageEncoder] warmed transformed-tensor cache for {n} entries " f"({self._name})", @@ -164,7 +165,20 @@ class ImageEncoder(TowerBase): if self._geom_loader is None: return - self._geom_loader.precompute(data.df, patient_col=data.patient_col) + + train_samples = data.collect_samples(split.train) + all_samples = train_samples + data.collect_samples(split.val) + if split.test is not None: + all_samples += data.collect_samples(split.test) + + if hasattr(self._geom_loader, "reset_cache"): + self._geom_loader.reset_cache() + if hasattr(self._geom_loader, "reset_weights"): + self._geom_loader.reset_weights() + if hasattr(self._geom_loader, "finetune"): + self._geom_loader.finetune(train_samples) + + self._geom_loader.precompute(all_samples) vecs = self._geom_loader.all_vectors() context.put(self.EPC_GEOMETRY_KEY, vecs) print( diff --git a/v4/configs/ensemble_fused_geom_unet.json b/v4/configs/ensemble_fused_geom_unet.json new file mode 100644 index 0000000..35a5d79 --- /dev/null +++ b/v4/configs/ensemble_fused_geom_unet.json @@ -0,0 +1,142 @@ +{ + "_notes": [ + "Geometry vector injection via UNet (no GT, no seg-CNN tower).", + "img tower runs UNet per-fold to predict disc/cup masks, computes 5 CDR", + "features per eye, publishes to EPC; cd tower consumes them.", + "Same shape as ensemble_fused.json — just adds the UNet-derived geometry hook." + ], + "run_name": "v4/ensemble_fused_geom_unet", + "num_classes": 2, + "label_filter": [0, 1], + "split_identity_level": 1, + "eval_stage": "hb", + "save_predictions": true, + "seed": 1234, + "folds": 5, + "fold_seed": 100, + "output_root": "v4/results", + "out_dir_tags": ["binary"], + + "data": { + "module": "v4.classes.profiles.v4papila", + "args": { + "image_dir": "Papila/FundusImages", + "clinical_dir": "Papila/ClinicalData", + "label_col": "Diagnosis", + "iop_corr_method": "ratio", + "iop_drop_raw": true, + "exclude_cols": ["Axial_Length"], + "in_memory_cache": true + } + }, + + "towers": [ + { + "name": "img", + "module": "v4.classes.towers.image_tower", + "class": "ImageEncoder", + "data_source": "image", + "epc_supplies": ["geometry_vectors"], + "args": { + "backbone": "refugelike", + "freeze_ratio": 0.0, + "augment": true, + "geometry_source": "unet", + "weights_path": "models/v2/refuge/segmentation/per_image/best.pt", + "contour_dir": "Papila/ExpertsSegmentations/Contours", + "unet_size": 512, + "normalize": "per_image", + "threshold": 0.5, + "finetune_epochs": 10, + "finetune_lr": 1e-5, + "finetune_batch_size": 4 + } + }, + { + "name": "cd", + "module": "v4.classes.towers.clinical_tower", + "class": "ClinicalEncoder", + "data_source": "matrix", + "epc_requests": ["geometry_vectors"], + "args": { + "hidden_dim": 128, + "geom_dim": 5 + } + } + ], + + "stages": [ + { + "name": "cd_warm", + "type": "warm", + "tower": "cd", + "head_name": "cd_aux", + "level": "eye", + "epochs": 40 + }, + { + "name": "img_aux", + "type": "head", + "input": "img", + "train_with": "nt", + "bcd": true + }, + { + "name": "cd_aux", + "type": "head", + "input": "cd", + "train_with": "nt", + "bcd": true + }, + { + "name": "nt", + "type": "fusion", + "module": "v4.classes.bridges.fusion_bridge", + "class": "FusionBridge", + "inputs": ["img", "cd"], + "level": "eye", + "epochs": 36, + "train_towers": true, + "warmup": { + "tower_epochs": 3, + "fused_epochs": 3 + }, + "args": { + "fusion_dim": 256 + } + }, + { + "name": "nt_head", + "type": "head", + "input": "nt", + "train_with": "nt" + }, + { + "name": "hb", + "type": "fusion", + "module": "v4.classes.bridges.hyperbridge", + "class": "HyperBridge", + "inputs": { "a": "nt", "b": "nt" }, + "level": "patient", + "epochs": 10, + "args": { + "hidden_dim": 256, + "mode": "embedding_mlp" + } + }, + { + "name": "hb_head", + "type": "head", + "input": "hb", + "train_with": "hb", + "args": { "dropout": 0.3 } + } + ], + + "training": { + "lr": 1e-4, + "batch_size": 8, + "bcd_prob": 0.5, + "tune_binary_threshold": true + } +} diff --git a/v4/distributed/cli.py b/v4/distributed/cli.py index 81c584b..74f4e3d 100644 --- a/v4/distributed/cli.py +++ b/v4/distributed/cli.py @@ -126,11 +126,12 @@ def _clients_table(api: _API) -> str: c["hostname"], c["gpu_info"][:30], s["state"], + (s.get("job_id") or "-")[:12], s.get("run_name") or "-", prog, _ago(c["last_seen"]), ]) - headers = ["ID", "HOST", "GPU", "STATE", "RUN", "PROGRESS", "SEEN"] + headers = ["ID", "HOST", "GPU", "STATE", "JOB_ID", "RUN", "PROGRESS", "SEEN"] widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))] sep = " " lines = [] diff --git a/v4/distributed/client.py b/v4/distributed/client.py index 1445c07..b3347f7 100644 --- a/v4/distributed/client.py +++ b/v4/distributed/client.py @@ -232,6 +232,20 @@ def _run_job(job: JobSpec, server: _Server, log_dir.mkdir(parents=True, exist_ok=True) log_file = log_dir / f"job_{job.job_id}.log" ctx: dict = {} + last_push = [time.time()] # mutable holder so _tail and _heartbeat share it + + def _push(): + server.push_status(StatusPush( + state="running", + job_id=job.job_id, + run_name=job.run_name, + fold=ctx.get("fold"), + stage=ctx.get("stage"), + epoch=ctx.get("epoch"), + total_epochs=ctx.get("total_epochs"), + last_val_auc=ctx.get("last_val_auc"), + )) + last_push[0] = time.time() def _tail(path: Path): with open(path, "r") as f: @@ -242,16 +256,7 @@ def _run_job(job: JobSpec, server: _Server, info = _parse_line(raw) ctx.update(info) if "epoch" in info: - server.push_status(StatusPush( - state="running", - job_id=job.job_id, - run_name=job.run_name, - fold=ctx.get("fold"), - stage=ctx.get("stage"), - epoch=ctx.get("epoch"), - total_epochs=ctx.get("total_epochs"), - last_val_auc=ctx.get("last_val_auc"), - )) + _push() elif proc.poll() is not None: for raw in f: print(raw, end="", flush=True) @@ -259,6 +264,18 @@ def _run_job(job: JobSpec, server: _Server, else: time.sleep(0.05) + def _heartbeat(): + # Push a status update every ~30s even when no log line is parsed. + # Prevents the server's reaper from declaring this client stale during + # long deterministic blocks (UNet fine-tune, data load, etc.). + while proc.poll() is None: + time.sleep(5) + if time.time() - last_push[0] >= 30: + try: + _push() + except Exception: + pass + with open(log_file, "w") as logf: proc = subprocess.Popen( cmd, @@ -268,10 +285,13 @@ def _run_job(job: JobSpec, server: _Server, start_new_session=True, ) - tailer = threading.Thread(target=_tail, args=(log_file,), daemon=True) + tailer = threading.Thread(target=_tail, args=(log_file,), daemon=True) + heartbeat = threading.Thread(target=_heartbeat, daemon=True) tailer.start() + heartbeat.start() proc.wait() tailer.join(timeout=5) + heartbeat.join(timeout=5) log_file.unlink(missing_ok=True) success = proc.returncode == 0 diff --git a/v4/distributed/jobs.db b/v4/distributed/jobs.db index 624a4d7..107ec39 100644 Binary files a/v4/distributed/jobs.db and b/v4/distributed/jobs.db differ diff --git a/v4/distributed/server.py b/v4/distributed/server.py index c483c5b..7859f8b 100644 --- a/v4/distributed/server.py +++ b/v4/distributed/server.py @@ -48,13 +48,20 @@ _DB_PATH: Path = Path("v4/distributed/jobs.db") _REPO_ROOT: Path = Path.cwd() _CLIENT_TTL: int = 120 # seconds before a client is considered gone _MAX_ATTEMPTS: int = 3 # max times a job is retried before being left as failed +_SERVER_START_TS: float = 0.0 # set in main(); used as a reaper grace window _clients: dict[str, ClientInfo] = {} _clients_lock = threading.Lock() def _reap_stale_clients(): - """Background thread: remove silent clients and re-queue their running jobs.""" + """Background thread: remove silent clients and re-queue their running jobs. + + On startup, the in-memory `_clients` dict is empty until clients re-register + via the `please_reregister` mechanism. We skip the running-job re-queue pass + for the first `_CLIENT_TTL` seconds after startup so still-alive clients have + time to come back; otherwise the reaper would orphan their jobs. + """ while True: time.sleep(30) cutoff = datetime.now(timezone.utc).timestamp() - _CLIENT_TTL @@ -73,6 +80,10 @@ def _reap_stale_clients(): del _clients[cid] known_ids = set(_clients.keys()) + in_grace = (time.time() - _SERVER_START_TS) < _CLIENT_TTL + if in_grace: + continue + with _db() as conn: rows = conn.execute( "SELECT job_id, assigned_to FROM jobs WHERE state='running'" @@ -454,12 +465,13 @@ def main(): if not args.token: ap.error("--token is required (or set HT_TOKEN)") - global _TOKEN, _DB_PATH, _REPO_ROOT, _CLIENT_TTL, _MAX_ATTEMPTS + global _TOKEN, _DB_PATH, _REPO_ROOT, _CLIENT_TTL, _MAX_ATTEMPTS, _SERVER_START_TS _TOKEN = args.token _DB_PATH = Path(args.db) _REPO_ROOT = Path(args.root).resolve() if args.root else Path.cwd() _CLIENT_TTL = args.client_ttl _MAX_ATTEMPTS = args.max_attempts + _SERVER_START_TS = time.time() _init_db() reaper = threading.Thread(target=_reap_stale_clients, daemon=True) diff --git a/v4/scripts/experiments/tri_v1/vec_geom_gt.json b/v4/scripts/experiments/tri_v1/vec_geom_gt.json new file mode 100644 index 0000000..e49e5db --- /dev/null +++ b/v4/scripts/experiments/tri_v1/vec_geom_gt.json @@ -0,0 +1,7 @@ +[ + { + "_note": "GT-derived 5-feature CDR vector injected into cd via EPC. Re-run with bug fix (rasterisation now in original image space, not 512x512 canvas). 3 reps.", + "run_name": "experiments/tri_v1/geom_vec_gt", + "reps": 3 + } +] diff --git a/v4/scripts/experiments/tri_v1/vec_geom_unet.json b/v4/scripts/experiments/tri_v1/vec_geom_unet.json new file mode 100644 index 0000000..957332e --- /dev/null +++ b/v4/scripts/experiments/tri_v1/vec_geom_unet.json @@ -0,0 +1,7 @@ +[ + { + "_note": "UNet-derived 5-feature CDR vector injected into cd via EPC. 3 reps.", + "run_name": "experiments/tri_v1/geom_vec_unet", + "reps": 3 + } +]