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/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 = []

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

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)