fix(workflow): size auto tiles for the segment GPU's VRAM

docs/guide/snakemake.md

@@ -56,6 +56,7 @@ shard: false                   # true → pack chunks into shards (fewer files)
 channel: 0                     # channel to segment (null = keep all)
 level: 0                       # pyramid level (0 = full resolution)
 tile_shape: "auto"             # "auto", or e.g. [16, 1024, 1024] (zyx)
+gpu_memory_gb: null            # for "auto" on SLURM: your segment GPU's VRAM
 overlap: 30                    # halo ≈ one object diameter
 skip_empty: true               # skip background tiles
 empty_threshold: null          # null → Otsu
@@ -81,6 +82,11 @@ sequential_labels: true        # renumber labels to a contiguous 1..N
     more (faster) jobs; very large 3-D tiles are slow. Keep `do_3D: false` (2-D
     per slice) if your objects segment fine per slice — it is much faster.
 
+    Tile planning runs on a **CPU** node, which cannot see the segment GPU, so
+    it logs `GPU memory query failed … using 8 GiB default` and sizes tiles for
+    8 GiB. Harmless, but to size for the real GPU set `gpu_memory_gb:` to its
+    VRAM (e.g. `24`, `40`, `80`) — or just set `tile_shape` explicitly.
+
 ## 4. Dry-run (always do this first)
 
 Check the plan without running anything:

workflow/config/config.yaml

@@ -22,6 +22,8 @@ shard: false # true → pack chunks into shards (zarr v3; fewer files)
 channel: 0 # channel to segment
 level: 0 # pyramid level to segment (0 = full res)
 tile_shape: "auto" # "auto", or e.g. [16, 1024, 1024] (zyx)
+gpu_memory_gb: null # for "auto" on a CPU prepare node: your segment GPU's VRAM
+#                     (e.g. 24, 40, 80). null → 8 GiB guess + a harmless warning
 overlap: 30 # halo (≈ one object diameter)
 skip_empty: true # skip background tiles
 empty_threshold: null # null → Otsu; or a number

workflow/scripts/prepare_tiles.py

@@ -20,12 +20,17 @@
 ts = cfg.get("tile_shape", "auto")
 if ts == "auto":
     cp = cfg["cellpose"]
+    # prepare runs on a CPU node, so the segment GPU's VRAM can't be queried
+    # here; pass gpu_memory_gb from the config to size tiles for it (avoids the
+    # "GPU memory query failed" fallback). None => the built-in 8 GiB default.
+    gpu_gb = cfg.get("gpu_memory_gb")
     tile_shape = tuple(
         partial(
             auto_tile_shape_cellpose,
             do_3D=cp.get("do_3D", False),
             use_gpu=cp.get("gpu", True),
             diameter=cp.get("diameter"),
+            gpu_memory=int(gpu_gb * 1024**3) if gpu_gb else None,
         )(image.shape, image.dtype)
     )
 else: