Fix/atera

scripts/create_resources/spatial/process_10x_atera_nebius.sh

@@ -8,7 +8,9 @@ cd "$REPO_ROOT"
 
 set -e
 
-publish_dir="s3://openproblems-data/resources/datasets"
+# store the loader output locally, mirroring the process_datasets layout ($id/)
+# but under a sibling raw/ folder
+publish_dir="/scratch/task_ist_preprocessing/raw"
 
 cat > /tmp/params_atera.yaml << HERE
 param_list:

src/base/labels_nebius.config

@@ -47,20 +47,26 @@ process {
   withLabel: midcpu { cpus = 15 }
   withLabel: highcpu { cpus = 30 }
   withLabel: lowmem {
-    memory = { get_memory( 20.GB * task.attempt ) }
+    // Nebius 48vcpu-192gb nodes: 188 GB allocatable
+    memory = { get_memory( 25.GB * task.attempt ) }
     disk = { 50.GB * task.attempt }
+    pod = [[nodeSelector: 'nebius.com/node-group-id=mk8snodegroup-e00p5wjvb7bc55b2js']]
   }
   withLabel: midmem {
+    // Nebius 64vcpu-256gb nodes: 251 GB allocatable
     memory = { get_memory( 50.GB * task.attempt ) }
     disk = { 100.GB * task.attempt }
+    pod = [[nodeSelector: 'nebius.com/node-group-id=mk8snodegroup-e00wvqfyde1nfwezs0']]
   }
   withLabel: highmem {
-    memory = { task.attempt == 1 ? 200.GB : 500.GB }
+    // Nebius 128vcpu-512gb nodes: 503 GB allocatable — cap well below node max
+    memory = { task.attempt == 1 ? 200.GB : 450.GB }
     disk = 200.GB
     pod = [[nodeSelector: 'nebius.com/node-group-id=mk8snodegroup-e00hnqhyfdcsy9m09n']]
   }
   withLabel: veryhighmem {
-    memory = { get_memory( 200.GB * task.attempt ) }
+    // Nebius 128vcpu-512gb nodes: 503 GB allocatable — cap at 450 GB to leave headroom
+    memory = { [ 200.GB * task.attempt, 450.GB ].min() }
     disk = { 400.GB * task.attempt }
     pod = [[nodeSelector: 'nebius.com/node-group-id=mk8snodegroup-e00hnqhyfdcsy9m09n']]
   }
@@ -115,6 +121,17 @@ process {
     containerOptions = { workflow.containerEngine == "singularity" ? '--nv':
        ( workflow.containerEngine == "docker" ? '--gpus all': null ) }
   }
+  withLabel: gpuhighmem {
+    // Nebius high-mem gpu-node-group: NVIDIA L40S, 1 GPU, 40 CPUs, 160 GiB RAM
+    // ~156 GiB allocatable — cap memory at 150 GB to leave headroom
+    cpus = 30
+    accelerator = 1
+    memory = { [ 100.GB * task.attempt, 150.GB ].min() }
+    disk = 200.GB
+    pod = [[nodeSelector: 'nebius.com/node-group-id=mk8snodegroup-e00dhcgx1xqjskycvc']]
+    containerOptions = { workflow.containerEngine == "singularity" ? '--nv':
+       ( workflow.containerEngine == "docker" ? '--gpus all': null ) }
+  }
 
 withLabel: hightime    { time = 12.h }
 withLabel: veryhightime { time = 24.h }

src/datasets/loaders/allen_brain_cell_atlas_merfish/script.py

@@ -1,6 +1,7 @@
 import math
 import os
 import re
+import time
 import urllib.request
 from datetime import datetime
 from pathlib import Path
@@ -137,8 +138,38 @@ def close(self):
         self.fileptr.close()
 
 
+def robust_urlretrieve(url, path, retries=5, backoff=5):
+    """urlretrieve that retries on transient network errors.
+
+    Returns True on success, False if the file is genuinely absent (404) or
+    every retry was exhausted. Partial downloads are cleaned up so a later
+    ``.exists()`` check does not treat a truncated file as complete.
+    """
+    for attempt in range(1, retries + 1):
+        try:
+            urllib.request.urlretrieve(url, path)
+            return True
+        except urllib.error.HTTPError as e:
+            if e.code == 404:
+                print(datetime.now() - t0, f"Not found (404): {url}", flush=True)
+                return False
+            print(datetime.now() - t0,
+                  f"HTTP {e.code} (attempt {attempt}/{retries}) for {url}", flush=True)
+        except (urllib.error.URLError, ConnectionError, TimeoutError) as e:
+            print(datetime.now() - t0,
+                  f"Download failed (attempt {attempt}/{retries}) for {url}: {e}", flush=True)
+        Path(path).unlink(missing_ok=True)
+        if attempt < retries:
+            time.sleep(backoff * attempt)
+    return False
+
+
 def download_fov_image(fov_idx, n_format, prefix, suffix, base_url, frames, stain):
-    """Download one FOV dax file and return a (n_z, px, py) uint16 array."""
+    """Download one FOV dax file and return a (n_z, px, py) uint16 array.
+
+    Returns None if the FOV image cannot be fetched (missing file or persistent
+    network failure) so the caller can skip it instead of crashing the run.
+    """
     fmt = f"{fov_idx:03d}" if n_format == 3 else f"{fov_idx:04d}"
     stem = prefix + fmt + suffix
     dax_path = TMP_DIR / (stem + ".dax")
@@ -149,8 +180,11 @@ def download_fov_image(fov_idx, n_format, prefix, suffix, base_url, frames, stai
         return tifffile.imread(tif_path)
 
     if not dax_path.exists():
-        urllib.request.urlretrieve(base_url + stem + ".dax", dax_path)
-        urllib.request.urlretrieve(base_url + stem + ".inf", inf_path)
+        if not robust_urlretrieve(base_url + stem + ".dax", dax_path):
+            return None
+        if not robust_urlretrieve(base_url + stem + ".inf", inf_path):
+            dax_path.unlink(missing_ok=True)
+            return None
 
     reader = DaxReader(str(dax_path))
     img = np.array([reader.load_frame(f) for f in frames])
@@ -288,6 +322,9 @@ def um_to_px_y(y):
     yp, yp_max = int(row["y_min"]), int(row["y_max"])
 
     fov_img = download_fov_image(fov, n_format, dapi_file_prefix, download_suffix, image_base_url, dapi_frames, "DAPI")
+    if fov_img is None:
+        print(datetime.now() - t0, f"Skipping FOV {fov}: image unavailable", flush=True)
+        continue
 
     # Find neighbors for linear blending at overlaps
     fov_right = fov_df[
@@ -304,22 +341,29 @@ def um_to_px_y(y):
     right = len(fov_right) > 0 and fov_right[0] not in missing_fovs
     bottom = len(fov_bottom) > 0 and fov_bottom[0] not in missing_fovs
 
+    bot_img = right_img = None
     if bottom:
         bot_img = download_fov_image(int(fov_bottom[0]), n_format, dapi_file_prefix, download_suffix, image_base_url, dapi_frames, "DAPI")
+        bottom = bot_img is not None
+    if bottom:
         fov_img[:, FRAME_SIZE_PIXELS - OVERLAP_SIZE:, :] = (
             np.multiply(fov_img[:, FRAME_SIZE_PIXELS - OVERLAP_SIZE:, :], w1) +
             np.multiply(bot_img[:, :OVERLAP_SIZE, :], w2)
         ) / (w1 + w2)
 
     if right:
         right_img = download_fov_image(int(fov_right[0]), n_format, dapi_file_prefix, download_suffix, image_base_url, dapi_frames, "DAPI")
+        right = right_img is not None
+    if right:
         fov_img[:, :, FRAME_SIZE_PIXELS - OVERLAP_SIZE:] = (
             np.multiply(fov_img[:, :, FRAME_SIZE_PIXELS - OVERLAP_SIZE:], w1.T) +
             np.multiply(right_img[:, :, :OVERLAP_SIZE], w2.T)
         ) / (w1.T + w2.T)
 
+    br_img = None
     if len(fov_br) > 0 and fov_br[0] not in missing_fovs:
         br_img = download_fov_image(int(fov_br[0]), n_format, dapi_file_prefix, download_suffix, image_base_url, dapi_frames, "DAPI")
+    if br_img is not None:
         sl_x = slice(FRAME_SIZE_PIXELS - OVERLAP_SIZE, None)
         sl_y = slice(FRAME_SIZE_PIXELS - OVERLAP_SIZE, None)
         if bottom and right:

src/datasets/loaders/tenx_atera/config.vsh.yaml

@@ -68,4 +68,4 @@ runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [midmem, midcpu, midtime]
+      label: [highmem, midcpu, midtime]

src/methods_segmentation/cellpose/config.vsh.yaml

@@ -93,4 +93,4 @@ runners:
   - type: executable
   - type: nextflow
     directives:
-      label: [ midtime, midcpu, highmem, gpu ]
+      label: [ midtime, midcpu, highmem, gpuhighmem ]