From af813bbb62d3b97a7056129ee68700fa22eade75 Mon Sep 17 00:00:00 2001 From: rpotter6298 Date: Wed, 29 Apr 2026 11:05:19 +0200 Subject: [PATCH] Refactor geometry feature loaders and enhance distributed client status reporting - Updated GTGeometryLoader to streamline geometry vector computation and caching. - Introduced UNetGeometryLoader for UNet-derived geometry vectors. - Added sample collection method in PapilaBundle for better data handling. - Refined GeometrySegEncoder and ImageEncoder to utilize new sample collection. - Enhanced distributed client with heartbeat mechanism for improved job tracking. - Added new configuration files for UNet-derived geometry integration. --- v4/classes/profiles/fundus_images.py | 340 ++++++++++++------ v4/classes/profiles/v4papila.py | 19 + v4/classes/towers/geometry_tower.py | 24 +- v4/classes/towers/image_tower.py | 20 +- v4/configs/ensemble_fused_geom_unet.json | 142 ++++++++ v4/distributed/cli.py | 3 +- v4/distributed/client.py | 42 ++- v4/distributed/jobs.db | Bin 69632 -> 69632 bytes v4/distributed/server.py | 16 +- .../experiments/tri_v1/vec_geom_gt.json | 7 + .../experiments/tri_v1/vec_geom_unet.json | 7 + 11 files changed, 467 insertions(+), 153 deletions(-) create mode 100644 v4/configs/ensemble_fused_geom_unet.json create mode 100644 v4/scripts/experiments/tri_v1/vec_geom_gt.json create mode 100644 v4/scripts/experiments/tri_v1/vec_geom_unet.json 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 624a4d722de58f4ca8c88c522408052316999327..107ec39350a02bb56af8c423e42522af29ef5e0b 100644 GIT binary patch literal 69632 zcmeI*Pi)&%90zbaPQ0dV^0q_O!`LffrfiZ=cza$`5kX#Nnp zrIQ#^4{(435;rusAt50Sa6*&dg2WXD5;qVMlMoz02x)>tC-I(RXGxuPbp5jxee;u3 z$NBk}_dd_D`kbmcJ2P3a3~8a>tZJ5&OK^t8k>kFs4NGXZfp zKEzkMt2LKu#hO+%Sf1!qTvy6XX_s?1L}vWGN;aC6db46(VkHB;OS&oTlJ17cQ-U}g z8RE6>R+^S(*-iBKu4g~rE@nT9ye{$khKGjuHxFCd{Gz=#Ymv{XP#Aw~GH>%pNBgC9 zH=0n>Fsg--|?T)9t|HO15|73nvay!~CAbVd{GtEl5X6Qw$-fFiurPMU& z(V!RU40v7Fyt%X&&8}Q;N~vCLEE-M{lk%n*&1St>teR!ko!%*$Wf|3m)f%qbcerbu zB1e4U@FPP!XCLh5%ZqewU(}YYy8XEIZdcebE6@`So8NPx@EJ}Fu5rON!i81_KNz^y zzs5b>Gj{mggz^x800bZa0SG_<0y{0RdYLDqQzuvB4<770G&MzkBUDN`Rn4eIR4=8C zm4?x*RE?Tt#;sM!i~GHcG{%nqe)cVx>j;%e7&~mQ%$GmuR+Cf5m7%KN5|W z>a~SRdE}ThLT3c62WDKam<`P;T`=_coodrrL_?Ni7tMNYu)dOV#vo>XEb3m7^s!c~AS~Jw`Ver*t}%Qc_Nz=7jcp&pGyT zx!i(gDA8nA-l$Kz^4{Cj@s~vcP^}l<&MC zDY6ZnSMvgFL+7P1zr9PxbtHjpP&zJX_=qqX?7Yg`n*U$r$dBY(@)fyCZjj#wudxV; zcAN;YSg_;05R32ZSPRhzAF|_xFpD?r_&m(w7Z&6JJ7x)sSM7M8&=`34?t@(R#x9Qw zUJe2ffB*y_009U<00Izz00bbgdjbQ(LEiaW6R~4G(6R3CSQps3^S1-Oe&JxyS#Qn% zuXALL{7wEMzms3d_3+rg{k01O=gw$+&G}Y5P$##AOHafKmY;| zfB*y_009VWTEPDNKXj7|e;A$(-6S88v*9)V`>n>Nn}Ho5009U<00Izz00bZa0SG`~ zCj|=r)iHjn-`BD38!N=^Eb9LI`@nVb1cJ0R`{{CkE+M)M&}EP=VY(2y z?4t|I!3_csfB*y_009U<00Izz00bZa0SG_<0uX=z1Rwwb2tWV=5P$##AOHafKmY;| zfB*y_009U<00Izz00bZa0SG_<0uX=z1Rwwb2tWV=5P$##AOHafKmY;|fB*y_009U< z00Izz00bZa0SG_<0uX=z1Rwwb2tWV=5P$##AOHafKmY;|fB*y_009U<00Izz00bZa d0SG_<0uX=z1Rwwb2tWV=5P$##An?Bw_!q{{1m*w$ literal 69632 zcmeHQdu$xXdB5ZFd3U>O{~^nPUF7tjBuZLxc4uaHhk8FKiQ+@zOA_^B_o>j9NSdS+ zyNQD^sEahQV;F7PIzZy4j#ImZ>o`DB6iL&#c3TupQpd4_)=29njbWrs4Wy~vq^aA^ z+;LCCB~9(xw-jv;5B&7$<9FZg?>FDf&g|^W4h{5<&-$fca{7`zTY{MzGTCfqZ>f~Y z)ctZMlev-p=IL*a{x;CxdiopwWJBf6|D`e7uY6KcewJx0yqb}=D}Pt~LGk|LPWm75 zBNz}22nGZLf&syRU_dY+7!VAk%s_8rbN^-koE&O*%7U_Vk4rrZgWI z?AkZnRT}B<9vJB=b@w0aI#pWWRXWyRVxC*rSK0;|8rOEL$zEC*ZTj-W*o1w_XMTF4d5sq3uYnac$e0tiA9oXJ+l$@+sCY z9k2ZNYOwPE+U_Z3o7S#bll{boS=+f-{x&sh`}X&CmFX#M$(JHO)NyRwn>P-3 zof$UP*a{Drc6|FPjx|447|V&6-X-sOuvANtLGS-DluY++n# zR@KtcAzH*2$7g02ozmpx*{RF3W8U~Q^<41P(!s7n`$l?)OPjBtj%k19^2OPi&5O^o zLg`ZHc{RY4KjDo}Tv+1UVR@Fn-`zjlb+~J=6auXV0E-W6&&-TpnDD)^*~!XzFLrUK zeVPs47##sepe^>LD5e+qE_QJzFHK$a=bac+i(UNb>B;G_OEVW3oTXjt*;)V6)NG}= z-W>3Pa%$UL*SL1anrx<=>@#;g#vZ_Cq}qE-BQzDK$@e@Xg~^hd?dNq?UIPO-aqd*Mq`y)>G?zvcOsms?&h zeptG#@TtOc(x!Z8@!N$h#my~GOW!X%RQ#!A78?p{3nc%;{Mr0#g)@c!EHoA_=9{Gd zkk&}g7N5xfZt9~5Cb>lTCI3s6 zGXPKbE#GuhfYAwNunZ%xq3x)SW{)s~rt8>+9#pa64KsreZQnF>R|T3o#0=2a2>rEH z%hd;&!3`V_n3|~~7YvlY;?R(~3JkM{8BEXA5i$c68Rk)Dz?hx@jjZ~n)m@IE zBa?myt4`n@VFu3%uyd<2QA z1Fax5| zb;=R_3dYRf0|OzKvhFzsVg?;sHnl+&=>#%^YhxR+puU53W&kFp`QZ`OF^R?uSf`l{ z14{+2511iE!*HO!kfUP;)7u@(kqV%bC4} z2Xo3Z$`_T-$m8-kc|h)xiLA=&1F93q;E=3Nq;OoCVf_VK)Oe|OS&Z4 z(umY8?UFR9ReGN!N!j9S#UB;_vH0!cUl#wg_;~U2#ZMQ1y?A%=j-p>2E%p`n6>+hx z_)Enb>79a~7yiBQ{lW`{zbX7#;mN`yg$Lt)l)D#ycsxQE38Ck;l+#af`bkbd!RaHMKFsMuoIc3u1Dt-G(~og_ zKd1L`dM~H@Iryu0>1Dw8})Aw=uUQXY`>AN|77pL#!^c|eu z$?1gCElxK%-QaZ0>4?)Ir|X=qaXR4i4o+|9bd}S$bNV(;-^%H2oW6zAH*@+XPH*M( zjhw!L)7L9cWLw{nm2xFdnj>sob8b!6o1E}V$~QA>7M(=+S9&KW0~Y;tc|%^gJ+oKd zmQe!5R?aFXl|iLXIieg;b}K|dN{6ylX;n(fEy_*GjY?5zQtIS4zx; zNq$j&L4H>LTluf#r{u57eQ} zCXO(1n2AG7bTM&|i33dRXJQ`{dzsk7#BL^bF|m`09ZYmGL71?ZFqtrzz)T<}AQL(h z8WVtt4kp@}P?^}y#5N|jGSSAw7A7_`v5ARRCN?s$fr<4@tYf0Y1RI@yIw9rhyXOm& znZ23UGKHQ3y`__XDgTxHZ|7(7$MW0qO)cMRd9vl#TF$h9mWJjRo1bX@)#lOW_U2sE z-#0zp^s%OsP1~EYjo)qj!^V#`jx=s<_`ind8y;);NW)OWmijm9pR51<`VZF+)NjiD zBKK_Wk=(E3j^#Gi{k-m5br09wS=U#$KKozUZ)888y^`(8mgot*oljDlw&q5&{+(0) z^!O!zVs@ruc6xm5iq>&qdfe-9T+hT~?gzE8iCJizkMDXezU$ezt~x#w-}Q8S*U`AH z8a@@@^<;e46LDNYxd=w$yAH>99g6F!3_;fg`E7WBq< z?TPPtG_I?TyW_haiSK$iuB(O*#dqzB?|LwfYo!((i0`^TzU#iYt~%Zu-*r!X*WGbl zHM}dn>(2PDJFejhuU-o}o8jrguOuB(Pod{-FXRgdFZsRdemR}kN| zBd)8C+vB^c@m;sab=B~;_^w;yyS8n}^(@sRa_kxZ;`oF=MtewGV)$&1;j<}*PiqXH zjWK*S#PC@k!)IL#pHd7TItHkJAZc$^_zMOE1A>867?AR}W(v9Unf(2-DcLQDnx~pR z)A&ros|`2T@6VmDd{Q~E2r3wO*E29@w<)coyXJPRyQMkPN8vW}Ys8kh1yvq)us*Mb zub_7KBHhH)pJ~5>#y)Tty=^r+`9XjB%z9OICntjO3+uO+*3+F=EB~44@WyAR>{<5% zzSj}&2z1ZZDFE7Uo0*(gf9~oQuJcbH>OeQ>3Q$!yhk?b%n zwD#z&HWp7KI0pF>tF*rVS85K!{3TV1!gs3)0P$6#Z-gQQ0B|)W|bWN`wqS_Qe z#VxPk&7>p?3vu9u8#n+odYT?a(~xwCkCC^Y)869!q!XKntf@7Sr&&|IR@s2)x@cUgz8P}DW>~^%tPYqDedQKo!DrHm|ExC zc$zZbUAeSbsUy|V17($m*x`!GLrT|59akJjFd!HZ3z4=F5 zo@xH)rvIs&OZN9N-{8+m=z@XwCIfS}8ky7Q_l{lS)@Qh^q&lasr#`2z<8m!ud?R$3 zIen~YKo4*4+S?--`R#C56~@JrGRt5@7QLf8KNH#XTbbwIZ5cdgw@1p2-{^Q%m0O~z z2Mp;HayZv&(=BNhJX2$G8!^YhxMI)LB&vT%H!O4h=@XilLHfBlXB!bFA-6sEwU}fJ zTxipTO(LGYLzQe zqFQePf?)WRnGGWt`F-J66~;uh-mr9JRLfD1VC1*DUsV_r<*I?PZp=RoF=uNLjQsBX z*TPs?1ytjQ1{d|upiypf66Gp=mxgYH&p6p2f{}kA!K%WTC|5DDh*|9+c0@4p?^sw> z7!&0xHUWJlL)flvi(us60TBmdy6q}Ao==pk2xF+tUmxdt8SF8D8|d$i>GwLs5d;H* z0l|P^KrkQ}5DW+g1OtKr!GK^uFd!IMX$JBQzmd7M?v6~uuKHK%AE@7vds+DdWmx`S zx`+Ru^qTZL(vIRk73Yc@3x8d>tsv#Ulpkt&z2(7{ow<+aw$Z)(?pX43U!9qKF?)Bm zlzDn3<}p^9GjY7P!9XpZX5lvFT5zYfGmYTkEmD$3@bKQGRny?})BoX3L35;jPJOFS z5G`6#A1;#JP}8X4BGvUZjT$nruBK7Lj40JKYM2UikEq;ZZ=-^VuL=eP11rHm8f|vi zwisyy51SZV)8LV*eUwJA1S{lJaT0k1XLs~j&8a&cBX-^|~ zNQ;_A@Q{}6H4PrAMBCB`9@4Tkjo=|IZ8hx_h5I*cscF=3|D??|jT-J>w5g_1!~KI= zYZ^7&zh`5ua)`d;?_Rk>dQQ45-CX=i z@q95`_(I`u{$KO==2 zXC_nk$+Rn~h%Xq3!9XqYoO^rAf)4okG=hhf=v*4X!?xvYO@l|8+%ste4{14_M(~i9 z(V7O2G`XkJ2p-aMGL7INEhlOge2IR`NE*RIT87gI9?~*Y)8LUxG?+&4kd}cof`_ym zuW9f|B|4Ty@Q{}NG=hh;^wlDGbz~7)P|%x3@Q{|CG=hh;9Ia{aNG0k{BX~&5ku-vb zv>dK!@JJ;(lt%E7mM)3*%5KbjFQeQVUH9K~b0cavUVkz7@wzW$zfN7=aX)kR;mG<8 z&kC^X2Shb=ZJG5O)ryuj2((#v>dvC#kYW}J2(!; z)Yhm}F-8(P9!8BWr>d?s>WW|_f#(6N*{bfsRfRFESFtc^mH~_^#)A=zo^9#`*^X)x zvZ^p9%2m3*Ji5Aa`#=OEB9`x&8dP=Py%xqwZ#1Q~f8|_VqQ4uj9O0H zc1-|0Pt`4dRbfn&tAJpnSJ!avi(rHXA`k&rwV=JKFeb`XU=Uzd+ts}hjJ`_(WCloe zO)n0{bgvyOS0nwWSh)(2PD?Ue-?%4&5!oQ{5hSV>II9X{vRu`0b=T%Odv^q*Zy{{D zcA%1L_W>#_r`_hvmt%rlh3&*Lky%}*x+{W_Xg=|Lhp3is#3{#g`v8o^5KL{pGG5+%umk)kp3K3vL|8Xjfr7t&X%|4Kb1kMqq%z2hC4zTV}iDZVg#B2t7GF65sZLd7%ks)XnVT4 z$}vG(gW3%f!s^?;?2!mY%P;~P((%8eEq7fkvs@Xurdf`aBP7pN0~tE0?%XyU!Kmpv zcA*DVY%DjXT6v6#a@Ei*6NK-gvxg!WeL5sI4c%3Nw%k3AmB*MUS25A5Z>*ZL2O}7v zuMzrdtCp)@t6fcXOjTv^t&EK^P1fq0677KqMmKOgU}~m{T(GJzCdyT;A%eqqT-wJY z7=4GfIJ)DjKD~wf=6yJh<+R(J zKouCP3uB^Og^1o(toB!rMlhP5$zBK^U#`vJ&LHjjIDdUil&iX7XvTcI%EteV z`9~Mb|G(oV`o#$g1_T3w0l|P^KrkQ}5DdKQ7?`uqNBZKJKspm=LNy@EjT_7NIM^Eq zQX3m{M@T;2un2}$b=-I^f)Qz&Pwx<56)pRG9CrpOjW@U>lrbh6Z_t?xF6U7MW`vGhnsUBSJSITH4pTcbw|GW1<)n<*I2xQ?HKyM