From 53adbfd41f61d3dda4518a928ff3034fb3808905 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 22 Apr 2026 23:17:06 +0200
Subject: [PATCH 01/49] feat: implement Virchow2 tile embedding extraction
 pipeline

Adds preprocessing/embeddings.py with Virchow2 model wrapper and
per-slide TileDataset that reads tiles from WSIs via OpenSlide. Tiles
and slide metadata are fetched from the tiling MLflow run; embeddings
are saved as per-slide parquet files and logged to MLflow.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 .../preprocessing/embeddings_05mpp.yaml       |   7 +
 configs/preprocessing/embeddings.yaml         |  26 +++
 preprocessing/embeddings.py                   | 204 ++++++++++++++++++
 pyproject.toml                                |   4 +
 scripts/submit_embeddings.py                  |  18 ++
 5 files changed, 259 insertions(+)
 create mode 100644 configs/experiment/preprocessing/embeddings_05mpp.yaml
 create mode 100644 configs/preprocessing/embeddings.yaml
 create mode 100644 preprocessing/embeddings.py
 create mode 100644 scripts/submit_embeddings.py

diff --git a/configs/experiment/preprocessing/embeddings_05mpp.yaml b/configs/experiment/preprocessing/embeddings_05mpp.yaml
new file mode 100644
index 0000000..2466041
--- /dev/null
+++ b/configs/experiment/preprocessing/embeddings_05mpp.yaml
@@ -0,0 +1,7 @@
+# @package _global_
+
+defaults:
+  - /data: dataset
+  - _self_
+
+tiling_run_id: ???
diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
new file mode 100644
index 0000000..9fa562b
--- /dev/null
+++ b/configs/preprocessing/embeddings.yaml
@@ -0,0 +1,26 @@
+# @package _global_
+
+mlflow_artifact_path: embeddings
+tiling_run_id: ???
+tile_size: 224
+
+dataloader:
+  batch_size: 2048  # 2048 for H100, 1024 for A40, 512 for mig-2g.20gb
+  num_workers: 8
+  persistent_workers: true
+
+metadata:
+  run_name: Virchow2 Embeddings ${dataset.name}
+  description: |
+    Embeddings preprocessing experiment config.
+    The embeddings pipeline:
+    * downloads slides/tiles parquet files from the tiling MLflow run
+    * initializes the Virchow2 tile encoder
+    * processes slides in batches through the encoder
+    * saves embeddings as parquet files with tile coordinates (one file per slide)
+    * logs artifacts to MLflow
+  hyperparams:
+    tiling_run_id: ${tiling_run_id}
+    tile_size: ${tile_size}
+    batch_size: ${dataloader.batch_size}
+    num_workers: ${dataloader.num_workers}
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
new file mode 100644
index 0000000..26d80b7
--- /dev/null
+++ b/preprocessing/embeddings.py
@@ -0,0 +1,204 @@
+import tempfile
+from pathlib import Path
+from typing import cast
+
+import hydra
+import mlflow
+import mlflow.artifacts
+import numpy as np
+import openslide
+import pandas as pd
+import timm
+import torch
+import torchvision.transforms.v2 as T
+from omegaconf import DictConfig
+from rationai.mlkit import autolog, with_cli_args
+from rationai.mlkit.lightning.loggers import MLFlowLogger
+from timm.layers.mlp import SwiGLUPacked
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+
+class Virchow2(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+        # Requires HF_TOKEN env variable; model is gated on HuggingFace.
+        self.module = timm.create_model(
+            "hf-hub:paige-ai/Virchow2",
+            pretrained=True,
+            mlp_layer=SwiGLUPacked,
+            act_layer=torch.nn.SiLU,
+        ).eval()
+        self.embed_dim: int = cast("int", self.module.embed_dim) * 2
+
+    @torch.inference_mode()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = cast("torch.Tensor", self.module(x))  # B x 261 x 1280
+
+        class_token = output[:, 0]  # B x 1280
+        patch_tokens = output[:, 5:]  # B x 256 x 1280; tokens 1-4 are register tokens
+
+        return torch.cat([class_token, patch_tokens.mean(1)], dim=-1)  # B x 2560
+
+
+class TileDataset(Dataset):
+    """PyTorch Dataset that reads tiles from a single WSI on-the-fly via OpenSlide."""
+
+    def __init__(
+        self,
+        slide_path: str,
+        tiles: np.ndarray,
+        level: int,
+        tile_size: int,
+        transform: T.Compose,
+    ) -> None:
+        self.slide_path = slide_path
+        self.tiles = tiles  # shape (N, 2): columns are [x, y]
+        self.level = level
+        self.tile_size = tile_size
+        self.transform = transform
+
+        # Opened lazily so each DataLoader worker gets its own file handle.
+        self._slide: openslide.OpenSlide | None = None
+        self._downsample: float | None = None
+
+    def _open_slide(self) -> tuple[openslide.OpenSlide, float]:
+        if self._slide is None:
+            self._slide = openslide.OpenSlide(self.slide_path)
+            self._downsample = self._slide.level_downsamples[self.level]
+        return self._slide, cast("float", self._downsample)
+
+    def __len__(self) -> int:
+        return len(self.tiles)
+
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, dict[str, int]]:
+        x, y = int(self.tiles[idx, 0]), int(self.tiles[idx, 1])
+        slide, downsample = self._open_slide()
+
+        # OpenSlide read_region expects level-0 coordinates.
+        location = (round(x * downsample), round(y * downsample))
+        region = slide.read_region(location, self.level, (self.tile_size, self.tile_size))
+        image = np.array(region.convert("RGB"))
+
+        return self.transform(image), {"x": x, "y": y}
+
+
+def compute_slide_embeddings(
+    slide_path: str,
+    slide_tiles: pd.DataFrame,
+    level: int,
+    model: Virchow2,
+    transform: T.Compose,
+    tile_size: int,
+    batch_size: int,
+    num_workers: int,
+    persistent_workers: bool,
+    device: torch.device,
+) -> pd.DataFrame:
+    dataset = TileDataset(
+        slide_path=slide_path,
+        tiles=slide_tiles[["x", "y"]].values,
+        level=level,
+        tile_size=tile_size,
+        transform=transform,
+    )
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        persistent_workers=persistent_workers and num_workers > 0,
+    )
+
+    all_embeddings = torch.zeros((len(dataset), model.embed_dim), dtype=torch.float32)
+    all_x = torch.zeros(len(dataset), dtype=torch.int32)
+    all_y = torch.zeros(len(dataset), dtype=torch.int32)
+
+    for i, (images, metadata) in enumerate(dataloader):
+        images = images.to(device)
+        batch_embeddings = model(images)
+
+        start = i * batch_size
+        end = start + batch_embeddings.size(0)
+        all_embeddings[start:end] = batch_embeddings.cpu()
+        all_x[start:end] = metadata["x"].to(torch.int32)
+        all_y[start:end] = metadata["y"].to(torch.int32)
+
+    return pd.DataFrame(
+        {
+            "x": all_x.numpy(),
+            "y": all_y.numpy(),
+            "embedding": [emb.numpy() for emb in all_embeddings],
+        }
+    )
+
+
+@with_cli_args(["+preprocessing=embeddings"])
+@hydra.main(config_path="../configs", config_name="preprocessing", version_base=None)
+@autolog
+def main(config: DictConfig, logger: MLFlowLogger) -> None:
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = Virchow2().to(device)
+
+    transform = T.Compose(
+        [
+            T.ToImage(),
+            T.ToDtype(torch.float32, scale=True),
+            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    for split_name in ["train", "test"]:
+        slides_df = pd.read_parquet(
+            mlflow.artifacts.download_artifacts(
+                run_id=config.tiling_run_id,
+                artifact_path=f"{split_name}_split/slides.parquet",
+            )
+        )
+        tiles_df = pd.read_parquet(
+            mlflow.artifacts.download_artifacts(
+                run_id=config.tiling_run_id,
+                artifact_path=f"{split_name}_split/tiles.parquet",
+            )
+        )
+
+        tiles_with_path = tiles_df.merge(
+            slides_df[["id", "path", "level"]],
+            left_on="slide_id",
+            right_on="id",
+            how="left",
+        )
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir)
+
+            for slide_id, slide_tiles in tqdm(
+                tiles_with_path.groupby("slide_id"),
+                desc=f"Slides ({split_name})",
+            ):
+                slide_row = slide_tiles.iloc[0]
+
+                try:
+                    df = compute_slide_embeddings(
+                        slide_path=str(slide_row["path"]),
+                        slide_tiles=slide_tiles,
+                        level=int(slide_row["level"]),
+                        model=model,
+                        transform=transform,
+                        tile_size=config.tile_size,
+                        batch_size=config.dataloader.batch_size,
+                        num_workers=config.dataloader.num_workers,
+                        persistent_workers=config.dataloader.persistent_workers,
+                        device=device,
+                    )
+                    df.to_parquet(tmp_path / f"{slide_id}.parquet", index=False)
+                except Exception as e:
+                    print(f"Error processing slide {slide_id}: {e}")
+
+            mlflow.log_artifacts(str(tmp_path), f"{split_name}_split/embeddings")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index bcf708f..f197343 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,6 +23,10 @@ dependencies = [
     "numpy>=2.3.5",
     "rationai-tiling>=1.1.1",
     "tifffile>=2025.12.20",
+    "torch>=2.0.0",
+    "torchvision>=0.15.0",
+    "timm>=1.0.0",
+    "einops>=0.8.0",
 ]
 
 [dependency-groups]
diff --git a/scripts/submit_embeddings.py b/scripts/submit_embeddings.py
new file mode 100644
index 0000000..8e1b088
--- /dev/null
+++ b/scripts/submit_embeddings.py
@@ -0,0 +1,18 @@
+from kube_jobs import storage, submit_job
+
+
+submit_job(
+    job_name="tissue-classification-embeddings",
+    username=...,
+    cpu=8,
+    memory="32Gi",
+    gpu=1,
+    public=False,
+    script=[
+        "git clone https://gitlab.ics.muni.cz/rationai/digital-pathology/pathology/tissue-classification.git workdir",
+        "cd workdir",
+        "uv sync",
+        "uv run -m preprocessing.embeddings +experiment=preprocessing/embeddings_05mpp tiling_run_id=...",
+    ],
+    storage=[storage.secure.DATA],
+)

From 53526ea57b3de1981eea1d89fc442f2e4a792141 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 22 Apr 2026 23:59:55 +0200
Subject: [PATCH 02/49] fix(submit): align embeddings job script with project
 conventions

Add base image, HF_TOKEN export, --frozen sync, and PROJECTS storage mount.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 scripts/submit_embeddings.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/scripts/submit_embeddings.py b/scripts/submit_embeddings.py
index 8e1b088..3c00308 100644
--- a/scripts/submit_embeddings.py
+++ b/scripts/submit_embeddings.py
@@ -4,15 +4,17 @@
 submit_job(
     job_name="tissue-classification-embeddings",
     username=...,
-    cpu=8,
-    memory="32Gi",
-    gpu=1,
+    image="cerit.io/rationai/base:2.0.6",
+    cpu=24,
+    memory="64Gi",
+    gpu="A40",
     public=False,
     script=[
         "git clone https://gitlab.ics.muni.cz/rationai/digital-pathology/pathology/tissue-classification.git workdir",
         "cd workdir",
-        "uv sync",
+        "export HF_TOKEN=<YOUR_HF_TOKEN>",
+        "uv sync --frozen",
         "uv run -m preprocessing.embeddings +experiment=preprocessing/embeddings_05mpp tiling_run_id=...",
     ],
-    storage=[storage.secure.DATA],
+    storage=[storage.secure.DATA, storage.secure.PROJECTS],
 )

From db5c017d6287f054f0ac02d6d2c0024ba715e429 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 23 Apr 2026 00:03:32 +0200
Subject: [PATCH 03/49] fix: move tiling_run_id into experiment config

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 configs/experiment/preprocessing/embeddings_05mpp.yaml | 2 +-
 scripts/submit_embeddings.py                           | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/configs/experiment/preprocessing/embeddings_05mpp.yaml b/configs/experiment/preprocessing/embeddings_05mpp.yaml
index 2466041..da8b92d 100644
--- a/configs/experiment/preprocessing/embeddings_05mpp.yaml
+++ b/configs/experiment/preprocessing/embeddings_05mpp.yaml
@@ -4,4 +4,4 @@ defaults:
   - /data: dataset
   - _self_
 
-tiling_run_id: ???
+tiling_run_id: "???"  # MLflow run ID of the tiling_05mpp run
diff --git a/scripts/submit_embeddings.py b/scripts/submit_embeddings.py
index 3c00308..9309256 100644
--- a/scripts/submit_embeddings.py
+++ b/scripts/submit_embeddings.py
@@ -4,7 +4,6 @@
 submit_job(
     job_name="tissue-classification-embeddings",
     username=...,
-    image="cerit.io/rationai/base:2.0.6",
     cpu=24,
     memory="64Gi",
     gpu="A40",
@@ -13,8 +12,8 @@
         "git clone https://gitlab.ics.muni.cz/rationai/digital-pathology/pathology/tissue-classification.git workdir",
         "cd workdir",
         "export HF_TOKEN=<YOUR_HF_TOKEN>",
-        "uv sync --frozen",
-        "uv run -m preprocessing.embeddings +experiment=preprocessing/embeddings_05mpp tiling_run_id=...",
+        "uv sync",
+        "uv run -m preprocessing.embeddings +experiment=preprocessing/embeddings_05mpp",
     ],
     storage=[storage.secure.DATA, storage.secure.PROJECTS],
 )

From b5ff14150f7dd989c59eeececf47384130f59176 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 23 Apr 2026 00:05:04 +0200
Subject: [PATCH 04/49] fix: move tiling_run_id into dataset.mlflow_artifacts

Consistent with how all other preprocessing run IDs are stored.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 configs/data/dataset.yaml                              | 1 +
 configs/experiment/preprocessing/embeddings_05mpp.yaml | 2 --
 configs/preprocessing/embeddings.yaml                  | 3 +--
 preprocessing/embeddings.py                            | 4 ++--
 4 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/configs/data/dataset.yaml b/configs/data/dataset.yaml
index ef6cbe7..625d07d 100644
--- a/configs/data/dataset.yaml
+++ b/configs/data/dataset.yaml
@@ -12,6 +12,7 @@ dataset:
     train_test_split_run_id: "90d7c88418ae4b879c3febbb70e9e741"
     train_split_filename: "split_mapping/train_split.csv"
     test_split_filename: "split_mapping/test_split.csv"
+    tiling_run_id: "fdf7550a2004474f8c7a05dc0cf1fd86"
 
   exclusions:
     bad_slides:
diff --git a/configs/experiment/preprocessing/embeddings_05mpp.yaml b/configs/experiment/preprocessing/embeddings_05mpp.yaml
index da8b92d..3891b97 100644
--- a/configs/experiment/preprocessing/embeddings_05mpp.yaml
+++ b/configs/experiment/preprocessing/embeddings_05mpp.yaml
@@ -3,5 +3,3 @@
 defaults:
   - /data: dataset
   - _self_
-
-tiling_run_id: "???"  # MLflow run ID of the tiling_05mpp run
diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 9fa562b..88f08b9 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -1,7 +1,6 @@
 # @package _global_
 
 mlflow_artifact_path: embeddings
-tiling_run_id: ???
 tile_size: 224
 
 dataloader:
@@ -20,7 +19,7 @@ metadata:
     * saves embeddings as parquet files with tile coordinates (one file per slide)
     * logs artifacts to MLflow
   hyperparams:
-    tiling_run_id: ${tiling_run_id}
+    tiling_run_id: ${dataset.mlflow_artifacts.tiling_run_id}
     tile_size: ${tile_size}
     batch_size: ${dataloader.batch_size}
     num_workers: ${dataloader.num_workers}
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 26d80b7..952beb7 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -153,13 +153,13 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
     for split_name in ["train", "test"]:
         slides_df = pd.read_parquet(
             mlflow.artifacts.download_artifacts(
-                run_id=config.tiling_run_id,
+                run_id=config.dataset.mlflow_artifacts.tiling_run_id,
                 artifact_path=f"{split_name}_split/slides.parquet",
             )
         )
         tiles_df = pd.read_parquet(
             mlflow.artifacts.download_artifacts(
-                run_id=config.tiling_run_id,
+                run_id=config.dataset.mlflow_artifacts.tiling_run_id,
                 artifact_path=f"{split_name}_split/tiles.parquet",
             )
         )

From e50f6fb1eb2ae98b6e2ae48812060a2fa6188479 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 23 Apr 2026 11:07:05 +0200
Subject: [PATCH 05/49] fix(submit): update clone URL to GitHub

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 scripts/submit_embeddings.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/submit_embeddings.py b/scripts/submit_embeddings.py
index 9309256..cbe0206 100644
--- a/scripts/submit_embeddings.py
+++ b/scripts/submit_embeddings.py
@@ -9,7 +9,7 @@
     gpu="A40",
     public=False,
     script=[
-        "git clone https://gitlab.ics.muni.cz/rationai/digital-pathology/pathology/tissue-classification.git workdir",
+        "git clone https://github.com/RationAI/tissue-classification.git workdir",
         "cd workdir",
         "export HF_TOKEN=<YOUR_HF_TOKEN>",
         "uv sync",

From 7698c7611fe632d2c16142b401521da05cd37b92 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Mon, 27 Apr 2026 16:47:12 +0200
Subject: [PATCH 06/49] refactor: use configs that support different models

---
 .../embeddings_virchow2_05mpp.yaml            |  7 +++++
 configs/preprocessing/embeddings.yaml         | 29 ++++++-------------
 2 files changed, 16 insertions(+), 20 deletions(-)
 create mode 100644 configs/experiment/preprocessing/embeddings_virchow2_05mpp.yaml

diff --git a/configs/experiment/preprocessing/embeddings_virchow2_05mpp.yaml b/configs/experiment/preprocessing/embeddings_virchow2_05mpp.yaml
new file mode 100644
index 0000000..8a3e120
--- /dev/null
+++ b/configs/experiment/preprocessing/embeddings_virchow2_05mpp.yaml
@@ -0,0 +1,7 @@
+# @package _global_
+
+defaults:
+  - /data: dataset
+  - _self_
+
+model: virchow2
diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 88f08b9..0d706af 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -1,25 +1,14 @@
 # @package _global_
 
-mlflow_artifact_path: embeddings
-tile_size: 224
-
-dataloader:
-  batch_size: 2048  # 2048 for H100, 1024 for A40, 512 for mig-2g.20gb
-  num_workers: 8
-  persistent_workers: true
+model: ???
+output_dir: ${project_dir}/embeddings
+concurrency: 1024
+batch_size: 2048
 
 metadata:
-  run_name: Virchow2 Embeddings ${dataset.name}
-  description: |
-    Embeddings preprocessing experiment config.
-    The embeddings pipeline:
-    * downloads slides/tiles parquet files from the tiling MLflow run
-    * initializes the Virchow2 tile encoder
-    * processes slides in batches through the encoder
-    * saves embeddings as parquet files with tile coordinates (one file per slide)
-    * logs artifacts to MLflow
+  run_name: "Embeddings: ${model}"
+  description: Tile embeddings using model ${model}
   hyperparams:
-    tiling_run_id: ${dataset.mlflow_artifacts.tiling_run_id}
-    tile_size: ${tile_size}
-    batch_size: ${dataloader.batch_size}
-    num_workers: ${dataloader.num_workers}
+    model: ${model}
+    concurrency: ${concurrency}
+    batch_size: ${batch_size}

From f08ef005e4474576cd1900b69eb8eeb8a0d47195 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Mon, 27 Apr 2026 16:47:59 +0200
Subject: [PATCH 07/49] refactor: use new approach for generating embeddings

---
 preprocessing/embeddings.py | 257 +++++++++++-------------------------
 1 file changed, 79 insertions(+), 178 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 952beb7..bcb0c39 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,204 +1,105 @@
-import tempfile
+import shutil
 from pathlib import Path
-from typing import cast
 
 import hydra
-import mlflow
 import mlflow.artifacts
-import numpy as np
-import openslide
 import pandas as pd
-import timm
-import torch
-import torchvision.transforms.v2 as T
+import pyarrow as pa
+import ray
 from omegaconf import DictConfig
+from PIL import Image
+from rationai import AsyncClient
 from rationai.mlkit import autolog, with_cli_args
 from rationai.mlkit.lightning.loggers import MLFlowLogger
-from timm.layers.mlp import SwiGLUPacked
-from torch.utils.data import DataLoader, Dataset
-from tqdm import tqdm
-
-
-class Virchow2(torch.nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-        # Requires HF_TOKEN env variable; model is gated on HuggingFace.
-        self.module = timm.create_model(
-            "hf-hub:paige-ai/Virchow2",
-            pretrained=True,
-            mlp_layer=SwiGLUPacked,
-            act_layer=torch.nn.SiLU,
-        ).eval()
-        self.embed_dim: int = cast("int", self.module.embed_dim) * 2
-
-    @torch.inference_mode()
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        output = cast("torch.Tensor", self.module(x))  # B x 261 x 1280
-
-        class_token = output[:, 0]  # B x 1280
-        patch_tokens = output[:, 5:]  # B x 256 x 1280; tokens 1-4 are register tokens
-
-        return torch.cat([class_token, patch_tokens.mean(1)], dim=-1)  # B x 2560
-
-
-class TileDataset(Dataset):
-    """PyTorch Dataset that reads tiles from a single WSI on-the-fly via OpenSlide."""
-
-    def __init__(
-        self,
-        slide_path: str,
-        tiles: np.ndarray,
-        level: int,
-        tile_size: int,
-        transform: T.Compose,
-    ) -> None:
-        self.slide_path = slide_path
-        self.tiles = tiles  # shape (N, 2): columns are [x, y]
-        self.level = level
-        self.tile_size = tile_size
-        self.transform = transform
-
-        # Opened lazily so each DataLoader worker gets its own file handle.
-        self._slide: openslide.OpenSlide | None = None
-        self._downsample: float | None = None
-
-    def _open_slide(self) -> tuple[openslide.OpenSlide, float]:
-        if self._slide is None:
-            self._slide = openslide.OpenSlide(self.slide_path)
-            self._downsample = self._slide.level_downsamples[self.level]
-        return self._slide, cast("float", self._downsample)
-
-    def __len__(self) -> int:
-        return len(self.tiles)
-
-    def __getitem__(self, idx: int) -> tuple[torch.Tensor, dict[str, int]]:
-        x, y = int(self.tiles[idx, 0]), int(self.tiles[idx, 1])
-        slide, downsample = self._open_slide()
-
-        # OpenSlide read_region expects level-0 coordinates.
-        location = (round(x * downsample), round(y * downsample))
-        region = slide.read_region(location, self.level, (self.tile_size, self.tile_size))
-        image = np.array(region.convert("RGB"))
-
-        return self.transform(image), {"x": x, "y": y}
-
-
-def compute_slide_embeddings(
-    slide_path: str,
-    slide_tiles: pd.DataFrame,
-    level: int,
-    model: Virchow2,
-    transform: T.Compose,
-    tile_size: int,
-    batch_size: int,
-    num_workers: int,
-    persistent_workers: bool,
-    device: torch.device,
-) -> pd.DataFrame:
-    dataset = TileDataset(
-        slide_path=slide_path,
-        tiles=slide_tiles[["x", "y"]].values,
-        level=level,
-        tile_size=tile_size,
-        transform=transform,
-    )
+from ratiopath.tiling.read_slide_tiles import read_openslide_tiles
+from ratiopath.tiling.types import ReadTilesArguments
+
+
+def read_tiles_batch(batch: pd.DataFrame) -> pd.DataFrame:
+    tile_images: dict[int, Image.Image] = {}
+    batch = batch.reset_index(drop=True)
+
+    for path, group in batch.groupby("path"):
+        assert isinstance(path, str)
+        kwargs: ReadTilesArguments = {
+            "tile_x": pa.array(group["x"].tolist()),
+            "tile_y": pa.array(group["y"].tolist()),
+            "tile_extent_x": pa.array(group["tile_extent_x"].tolist()),
+            "tile_extent_y": pa.array(group["tile_extent_y"].tolist()),
+            "level": pa.array(group["level"].tolist()),
+        }
+        tiles = read_openslide_tiles(path, **kwargs)
+        for i, idx in enumerate(group.index):
+            tile_images[idx] = Image.fromarray(tiles[i])
 
-    dataloader = DataLoader(
-        dataset,
-        batch_size=batch_size,
-        num_workers=num_workers,
-        persistent_workers=persistent_workers and num_workers > 0,
+    return batch.drop(columns=["path", "level", "tile_extent_x", "tile_extent_y"]).assign(
+        tile=pd.Series(tile_images)
     )
 
-    all_embeddings = torch.zeros((len(dataset), model.embed_dim), dtype=torch.float32)
-    all_x = torch.zeros(len(dataset), dtype=torch.int32)
-    all_y = torch.zeros(len(dataset), dtype=torch.int32)
-
-    for i, (images, metadata) in enumerate(dataloader):
-        images = images.to(device)
-        batch_embeddings = model(images)
-
-        start = i * batch_size
-        end = start + batch_embeddings.size(0)
-        all_embeddings[start:end] = batch_embeddings.cpu()
-        all_x[start:end] = metadata["x"].to(torch.int32)
-        all_y[start:end] = metadata["y"].to(torch.int32)
-
-    return pd.DataFrame(
-        {
-            "x": all_x.numpy(),
-            "y": all_y.numpy(),
-            "embedding": [emb.numpy() for emb in all_embeddings],
-        }
-    )
+
+class EmbedTiles:
+    def __init__(self, model: str) -> None:
+        self.model = model
+        self.client = AsyncClient()
+
+    async def __call__(self, row: dict) -> dict:
+        embedding = (
+            (await self.client.models.embed_image(self.model, row["tile"]))
+            .reshape(-1)
+            .tolist()
+        )
+        del row["tile"]
+        row["embedding"] = embedding
+        return row
 
 
 @with_cli_args(["+preprocessing=embeddings"])
 @hydra.main(config_path="../configs", config_name="preprocessing", version_base=None)
 @autolog
 def main(config: DictConfig, logger: MLFlowLogger) -> None:
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    for name, split_uri in config.dataset.mlflow_uris.tiling.items():
+        folder = Path(mlflow.artifacts.download_artifacts(split_uri))
+        slides = pd.read_parquet(folder / "slides.parquet")
+        tiles = pd.read_parquet(folder / "tiles.parquet")
 
-    model = Virchow2().to(device)
+        slide_info = slides.set_index("id")[["path", "level", "tile_extent_x", "tile_extent_y"]]
+        tiles_enriched = tiles.join(slide_info, on="slide_id")
 
-    transform = T.Compose(
-        [
-            T.ToImage(),
-            T.ToDtype(torch.float32, scale=True),
-            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-        ]
-    )
-
-    for split_name in ["train", "test"]:
-        slides_df = pd.read_parquet(
-            mlflow.artifacts.download_artifacts(
-                run_id=config.dataset.mlflow_artifacts.tiling_run_id,
-                artifact_path=f"{split_name}_split/slides.parquet",
-            )
+        ds = ray.data.from_pandas(tiles_enriched).repartition(
+            target_num_rows_per_block=config.batch_size
         )
-        tiles_df = pd.read_parquet(
-            mlflow.artifacts.download_artifacts(
-                run_id=config.dataset.mlflow_artifacts.tiling_run_id,
-                artifact_path=f"{split_name}_split/tiles.parquet",
-            )
+        ds = ds.map_batches(
+            read_tiles_batch,  # type: ignore[arg-type]
+            batch_format="pandas",
+            batch_size=config.batch_size,
+            memory=config.batch_size * 224 * 224 * 3 * 10,
         )
-
-        tiles_with_path = tiles_df.merge(
-            slides_df[["id", "path", "level"]],
-            left_on="slide_id",
-            right_on="id",
-            how="left",
+        ds = ds.map(
+            EmbedTiles,  # type: ignore[arg-type]
+            fn_constructor_args=(config.model,),
+            compute=ray.data.ActorPoolStrategy(
+                max_tasks_in_flight_per_actor=config.concurrency
+            ),
+            memory=224 * 224 * 224 * 3 * 10,
+            max_concurrency=config.concurrency,
         )
 
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            tmp_path = Path(tmp_dir)
-
-            for slide_id, slide_tiles in tqdm(
-                tiles_with_path.groupby("slide_id"),
-                desc=f"Slides ({split_name})",
-            ):
-                slide_row = slide_tiles.iloc[0]
-
-                try:
-                    df = compute_slide_embeddings(
-                        slide_path=str(slide_row["path"]),
-                        slide_tiles=slide_tiles,
-                        level=int(slide_row["level"]),
-                        model=model,
-                        transform=transform,
-                        tile_size=config.tile_size,
-                        batch_size=config.dataloader.batch_size,
-                        num_workers=config.dataloader.num_workers,
-                        persistent_workers=config.dataloader.persistent_workers,
-                        device=device,
-                    )
-                    df.to_parquet(tmp_path / f"{slide_id}.parquet", index=False)
-                except Exception as e:
-                    print(f"Error processing slide {slide_id}: {e}")
-
-            mlflow.log_artifacts(str(tmp_path), f"{split_name}_split/embeddings")
+        split_dir = Path(config.output_dir) / str(name)
+        split_dir.mkdir(parents=True, exist_ok=True)
+        tiles_parquet_dir = split_dir / "tiles.parquet"
+        if tiles_parquet_dir.exists():
+            shutil.rmtree(tiles_parquet_dir)
+
+        slides.to_parquet(split_dir / "slides.parquet", index=False)
+        ds.write_parquet(str(tiles_parquet_dir))
+
+        logger.log_artifacts(str(split_dir), str(name))
 
 
 if __name__ == "__main__":
-    main()
+    ctx = ray.data.DataContext.get_current()
+    ctx.enable_rich_progress_bars = True
+    ctx.use_ray_tqdm = False
+
+    with ray.init(runtime_env={"excludes": [".git", ".venv"]}):
+        main()

From ba96068fdac2eb4cc580e5be0c29d78d8595d0cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 29 Apr 2026 22:45:17 +0200
Subject: [PATCH 08/49] refactor: make some update and update dependencies

---
 configs/preprocessing/embeddings.yaml |   5 +-
 preprocessing/embeddings.py           |  79 ++++++------
 uv.lock                               | 174 ++++++++++++++++++++++++++
 3 files changed, 213 insertions(+), 45 deletions(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 0d706af..efcc704 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -3,7 +3,8 @@
 model: ???
 output_dir: ${project_dir}/embeddings
 concurrency: 1024
-batch_size: 2048
+block_size: 2048
+rows_per_file: 65536
 
 metadata:
   run_name: "Embeddings: ${model}"
@@ -11,4 +12,4 @@ metadata:
   hyperparams:
     model: ${model}
     concurrency: ${concurrency}
-    batch_size: ${batch_size}
+    block_size: ${block_size}
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index bcb0c39..e8fda1c 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,48 +1,31 @@
 import shutil
 from pathlib import Path
+from typing import Any
 
+import httpx
 import hydra
 import mlflow.artifacts
 import pandas as pd
 import pyarrow as pa
 import ray
 from omegaconf import DictConfig
-from PIL import Image
 from rationai import AsyncClient
 from rationai.mlkit import autolog, with_cli_args
 from rationai.mlkit.lightning.loggers import MLFlowLogger
-from ratiopath.tiling.read_slide_tiles import read_openslide_tiles
-from ratiopath.tiling.types import ReadTilesArguments
-
-
-def read_tiles_batch(batch: pd.DataFrame) -> pd.DataFrame:
-    tile_images: dict[int, Image.Image] = {}
-    batch = batch.reset_index(drop=True)
-
-    for path, group in batch.groupby("path"):
-        assert isinstance(path, str)
-        kwargs: ReadTilesArguments = {
-            "tile_x": pa.array(group["x"].tolist()),
-            "tile_y": pa.array(group["y"].tolist()),
-            "tile_extent_x": pa.array(group["tile_extent_x"].tolist()),
-            "tile_extent_y": pa.array(group["tile_extent_y"].tolist()),
-            "level": pa.array(group["level"].tolist()),
-        }
-        tiles = read_openslide_tiles(path, **kwargs)
-        for i, idx in enumerate(group.index):
-            tile_images[idx] = Image.fromarray(tiles[i])
-
-    return batch.drop(columns=["path", "level", "tile_extent_x", "tile_extent_y"]).assign(
-        tile=pd.Series(tile_images)
-    )
+from ratiopath.tiling.read_slide_tiles import read_slide_tiles
+from ray.data.expressions import col
 
 
 class EmbedTiles:
-    def __init__(self, model: str) -> None:
+    def __init__(self, model: str, concurrency: int) -> None:
         self.model = model
-        self.client = AsyncClient()
+        self.client = AsyncClient(
+            limits=httpx.Limits(
+                max_connections=concurrency, max_keepalive_connections=concurrency
+            )
+        )
 
-    async def __call__(self, row: dict) -> dict:
+    async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         embedding = (
             (await self.client.models.embed_image(self.model, row["tile"]))
             .reshape(-1)
@@ -62,36 +45,46 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         slides = pd.read_parquet(folder / "slides.parquet")
         tiles = pd.read_parquet(folder / "tiles.parquet")
 
-        slide_info = slides.set_index("id")[["path", "level", "tile_extent_x", "tile_extent_y"]]
+        slide_info = slides.set_index("id")[
+            ["path", "level", "tile_extent_x", "tile_extent_y"]
+        ]
         tiles_enriched = tiles.join(slide_info, on="slide_id")
 
-        ds = ray.data.from_pandas(tiles_enriched).repartition(
-            target_num_rows_per_block=config.batch_size
-        )
-        ds = ds.map_batches(
-            read_tiles_batch,  # type: ignore[arg-type]
-            batch_format="pandas",
-            batch_size=config.batch_size,
-            memory=config.batch_size * 224 * 224 * 3 * 10,
+        ds = ray.data.from_arrow(
+            pa.Table.from_pandas(tiles_enriched, preserve_index=False)
+        ).repartition(target_num_rows_per_block=config.block_size)
+        ds = ds.with_column(
+            "tile",
+            read_slide_tiles(  # pyright: ignore[reportCallIssue]
+                col("path"),
+                col("x"),
+                col("y"),
+                col("tile_extent_x"),
+                col("tile_extent_y"),
+                col("level"),
+            ),
+            num_cpus=1,
+            memory=4 * 1024**3,
         )
+        ds = ds.drop_columns(["path", "level", "tile_extent_x", "tile_extent_y"])
         ds = ds.map(
-            EmbedTiles,  # type: ignore[arg-type]
-            fn_constructor_args=(config.model,),
+            EmbedTiles,  # pyright: ignore[reportArgumentType]
+            fn_constructor_args=(config.model, config.concurrency),
             compute=ray.data.ActorPoolStrategy(
-                max_tasks_in_flight_per_actor=config.concurrency
+                max_size=4,
+                max_tasks_in_flight_per_actor=config.concurrency // 4,
             ),
-            memory=224 * 224 * 224 * 3 * 10,
             max_concurrency=config.concurrency,
         )
 
         split_dir = Path(config.output_dir) / str(name)
         split_dir.mkdir(parents=True, exist_ok=True)
-        tiles_parquet_dir = split_dir / "tiles.parquet"
+        tiles_parquet_dir = split_dir / "tiles"
         if tiles_parquet_dir.exists():
             shutil.rmtree(tiles_parquet_dir)
 
         slides.to_parquet(split_dir / "slides.parquet", index=False)
-        ds.write_parquet(str(tiles_parquet_dir))
+        ds.write_parquet(str(tiles_parquet_dir), min_rows_per_file=config.rows_per_file)
 
         logger.log_artifacts(str(split_dir), str(name))
 
diff --git a/uv.lock b/uv.lock
index 9a2416d..4c405d0 100644
--- a/uv.lock
+++ b/uv.lock
@@ -355,6 +355,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/0c/d5/c5db1ea3394c6e1732fb3286b3bd878b59507a8f77d32a2cebda7d7b7cd4/donfig-0.8.1.post1-py3-none-any.whl", hash = "sha256:2a3175ce74a06109ff9307d90a230f81215cbac9a751f4d1c6194644b8204f9d", size = 21592, upload-time = "2024-05-23T14:13:55.283Z" },
 ]
 
+[[package]]
+name = "einops"
+version = "0.8.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" },
+]
+
 [[package]]
 name = "fastapi"
 version = "0.121.2"
@@ -595,6 +604,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" },
 ]
 
+[[package]]
+name = "hf-xet"
+version = "1.4.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/53/92/ec9ad04d0b5728dca387a45af7bc98fbb0d73b2118759f5f6038b61a57e8/hf_xet-1.4.3.tar.gz", hash = "sha256:8ddedb73c8c08928c793df2f3401ec26f95be7f7e516a7bee2fbb546f6676113", size = 670477, upload-time = "2026-03-31T22:40:07.874Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ac/9f/9c23e4a447b8f83120798f9279d0297a4d1360bdbf59ef49ebec78fe2545/hf_xet-1.4.3-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:d0da85329eaf196e03e90b84c2d0aca53bd4573d097a75f99609e80775f98025", size = 3805048, upload-time = "2026-03-31T22:39:53.105Z" },
+    { url = "https://files.pythonhosted.org/packages/0b/f8/7aacb8e5f4a7899d39c787b5984e912e6c18b11be136ef13947d7a66d265/hf_xet-1.4.3-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:e23717ce4186b265f69afa66e6f0069fe7efbf331546f5c313d00e123dc84583", size = 3562178, upload-time = "2026-03-31T22:39:51.295Z" },
+    { url = "https://files.pythonhosted.org/packages/df/9a/a24b26dc8a65f0ecc0fe5be981a19e61e7ca963b85e062c083f3a9100529/hf_xet-1.4.3-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc360b70c815bf340ed56c7b8c63aacf11762a4b099b2fe2c9bd6d6068668c08", size = 4212320, upload-time = "2026-03-31T22:39:42.922Z" },
+    { url = "https://files.pythonhosted.org/packages/53/60/46d493db155d2ee2801b71fb1b0fd67696359047fdd8caee2c914cc50c79/hf_xet-1.4.3-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:39f2d2e9654cd9b4319885733993807aab6de9dfbd34c42f0b78338d6617421f", size = 3991546, upload-time = "2026-03-31T22:39:41.335Z" },
+    { url = "https://files.pythonhosted.org/packages/bc/f5/067363e1c96c6b17256910830d1b54099d06287e10f4ec6ec4e7e08371fc/hf_xet-1.4.3-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:49ad8a8cead2b56051aa84d7fce3e1335efe68df3cf6c058f22a65513885baac", size = 4193200, upload-time = "2026-03-31T22:40:01.936Z" },
+    { url = "https://files.pythonhosted.org/packages/42/4b/53951592882d9c23080c7644542fda34a3813104e9e11fa1a7d82d419cb8/hf_xet-1.4.3-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7716d62015477a70ea272d2d68cd7cad140f61c52ee452e133e139abfe2c17ba", size = 4429392, upload-time = "2026-03-31T22:40:03.492Z" },
+    { url = "https://files.pythonhosted.org/packages/8a/21/75a6c175b4e79662ad8e62f46a40ce341d8d6b206b06b4320d07d55b188c/hf_xet-1.4.3-cp37-abi3-win_amd64.whl", hash = "sha256:6b591fcad34e272a5b02607485e4f2a1334aebf1bc6d16ce8eb1eb8978ac2021", size = 3677359, upload-time = "2026-03-31T22:40:13.619Z" },
+    { url = "https://files.pythonhosted.org/packages/8a/7c/44314ecd0e89f8b2b51c9d9e5e7a60a9c1c82024ac471d415860557d3cd8/hf_xet-1.4.3-cp37-abi3-win_arm64.whl", hash = "sha256:7c2c7e20bcfcc946dc67187c203463f5e932e395845d098cc2a93f5b67ca0b47", size = 3533664, upload-time = "2026-03-31T22:40:12.152Z" },
+]
+
 [[package]]
 name = "httpcore"
 version = "1.0.9"
@@ -623,6 +648,26 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517, upload-time = "2024-12-06T15:37:21.509Z" },
 ]
 
+[[package]]
+name = "huggingface-hub"
+version = "1.12.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "filelock" },
+    { name = "fsspec" },
+    { name = "hf-xet", marker = "platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" },
+    { name = "httpx" },
+    { name = "packaging" },
+    { name = "pyyaml" },
+    { name = "tqdm" },
+    { name = "typer" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/3e/9f/3fda8b014db3ae239addc9b48b35c2cf7d318950b430712f34a2473ef81d/huggingface_hub-1.12.2.tar.gz", hash = "sha256:282c4999e641c89affdc4c02c265eddea944c1390dc19e89dac8ad3ae76dbdaf", size = 763393, upload-time = "2026-04-29T09:45:09.202Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/71/c1/1fa4162f6dd53259daf2ad31385273341821fa0acce164cd03971937a60e/huggingface_hub-1.12.2-py3-none-any.whl", hash = "sha256:7968e897fdbc6343c871c240d87d4434efe0ad9f80d57daa1cc5678c6d148529", size = 647757, upload-time = "2026-04-29T09:45:07.63Z" },
+]
+
 [[package]]
 name = "hydra-core"
 version = "1.3.2"
@@ -850,6 +895,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/de/1f/77fa3081e4f66ca3576c896ae5d31c3002ac6607f9747d2e3aa49227e464/markdown-3.10.2-py3-none-any.whl", hash = "sha256:e91464b71ae3ee7afd3017d9f358ef0baf158fd9a298db92f1d4761133824c36", size = 108180, upload-time = "2026-02-09T14:57:25.787Z" },
 ]
 
+[[package]]
+name = "markdown-it-py"
+version = "4.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mdurl" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" },
+]
+
 [[package]]
 name = "markupsafe"
 version = "3.0.3"
@@ -895,6 +952,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/de/77/ef1fc78bfe99999b2675435cc52120887191c566b25017d78beaabef7f2d/matplotlib-3.10.7-cp312-cp312-win_arm64.whl", hash = "sha256:5f3f6d315dcc176ba7ca6e74c7768fb7e4cf566c49cb143f6bc257b62e634ed8", size = 7992812, upload-time = "2025-10-09T00:26:54.882Z" },
 ]
 
+[[package]]
+name = "mdurl"
+version = "0.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" },
+]
+
 [[package]]
 name = "mlflow"
 version = "2.22.4"
@@ -1496,6 +1562,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" },
 ]
 
+[[package]]
+name = "pygments"
+version = "2.20.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" },
+]
+
 [[package]]
 name = "pyogrio"
 version = "0.12.1"
@@ -1786,6 +1861,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" },
 ]
 
+[[package]]
+name = "rich"
+version = "15.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markdown-it-py" },
+    { name = "pygments" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" },
+]
+
 [[package]]
 name = "rpds-py"
 version = "0.29.0"
@@ -1847,6 +1935,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/e5/80/69756670caedcf3b9be597a6e12276a6cf6197076eb62aad0c608f8efce0/ruff-0.14.5-py3-none-win_arm64.whl", hash = "sha256:4b700459d4649e2594b31f20a9de33bc7c19976d4746d8d0798ad959621d64a4", size = 13433331, upload-time = "2025-11-13T19:58:48.434Z" },
 ]
 
+[[package]]
+name = "safetensors"
+version = "0.7.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" },
+    { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" },
+    { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" },
+    { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" },
+    { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" },
+    { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" },
+    { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" },
+    { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" },
+    { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" },
+    { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" },
+    { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" },
+]
+
 [[package]]
 name = "scikit-image"
 version = "0.26.0"
@@ -1942,6 +2052,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/bf/cb6c1c505cb31e818e900b9312d514f381fbfa5c4363edfce0fcc4f8c1a4/shapely-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:743044b4cfb34f9a67205cee9279feaf60ba7d02e69febc2afc609047cb49179", size = 1722861, upload-time = "2025-09-24T13:50:43.35Z" },
 ]
 
+[[package]]
+name = "shellingham"
+version = "1.5.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/58/15/8b3609fd3830ef7b27b655beb4b4e9c62313a4e8da8c676e142cc210d58e/shellingham-1.5.4.tar.gz", hash = "sha256:8dbca0739d487e5bd35ab3ca4b36e11c4078f3a234bfce294b0a0291363404de", size = 10310, upload-time = "2023-10-24T04:13:40.426Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" },
+]
+
 [[package]]
 name = "simsimd"
 version = "6.5.13"
@@ -2100,11 +2219,28 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1b/fe/e59859aa1134fac065d36864752daf13215c98b379cb5d93f954dc0ec830/tifffile-2025.12.20-py3-none-any.whl", hash = "sha256:bc0345a20675149353cfcb3f1c48d0a3654231ee26bd46beebaab4d2168feeb6", size = 232031, upload-time = "2025-12-21T06:23:53.003Z" },
 ]
 
+[[package]]
+name = "timm"
+version = "1.0.26"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "huggingface-hub" },
+    { name = "pyyaml" },
+    { name = "safetensors" },
+    { name = "torch" },
+    { name = "torchvision" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7b/1e/e924b3b2326a856aaf68586f9c52a5fc81ef45715eca408393b68c597e0e/timm-1.0.26.tar.gz", hash = "sha256:f66f082f2f381cf68431c22714c8b70f723837fa2a185b155961eab90f2d5b10", size = 2419859, upload-time = "2026-03-23T18:12:10.272Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/6f/e9/bebf3d50e3fc847378988235f87c37ad3ac26d386041ab915d15e92025cd/timm-1.0.26-py3-none-any.whl", hash = "sha256:985c330de5ccc3a2aa0224eb7272e6a336084702390bb7e3801f3c91603d3683", size = 2568766, upload-time = "2026-03-23T18:12:08.062Z" },
+]
+
 [[package]]
 name = "tissue-classification"
 version = "0.1.0"
 source = { virtual = "." }
 dependencies = [
+    { name = "einops" },
     { name = "hydra-core" },
     { name = "mlflow" },
     { name = "numpy" },
@@ -2119,6 +2255,9 @@ dependencies = [
     { name = "ray" },
     { name = "scikit-learn" },
     { name = "tifffile" },
+    { name = "timm" },
+    { name = "torch" },
+    { name = "torchvision" },
     { name = "tqdm" },
 ]
 
@@ -2132,6 +2271,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
+    { name = "einops", specifier = ">=0.8.0" },
     { name = "hydra-core", specifier = ">=1.3.2" },
     { name = "mlflow", specifier = "<3.0.0" },
     { name = "numpy", specifier = ">=2.3.5" },
@@ -2146,6 +2286,9 @@ requires-dist = [
     { name = "ray", specifier = ">=2.51.1" },
     { name = "scikit-learn", specifier = ">=1.8.0" },
     { name = "tifffile", specifier = ">=2025.12.20" },
+    { name = "timm", specifier = ">=1.0.0" },
+    { name = "torch", specifier = ">=2.0.0" },
+    { name = "torchvision", specifier = ">=0.15.0" },
     { name = "tqdm", specifier = ">=4.66.0" },
 ]
 
@@ -2208,6 +2351,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/02/21/aa0f434434c48490f91b65962b1ce863fdcce63febc166ca9fe9d706c2b6/torchmetrics-1.8.2-py3-none-any.whl", hash = "sha256:08382fd96b923e39e904c4d570f3d49e2cc71ccabd2a94e0f895d1f0dac86242", size = 983161, upload-time = "2025-09-03T14:00:51.921Z" },
 ]
 
+[[package]]
+name = "torchvision"
+version = "0.24.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+    { name = "pillow" },
+    { name = "torch" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f0/af/18e2c6b9538a045f60718a0c5a058908ccb24f88fde8e6f0fc12d5ff7bd3/torchvision-0.24.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e48bf6a8ec95872eb45763f06499f87bd2fb246b9b96cb00aae260fda2f96193", size = 1891433, upload-time = "2025-11-12T15:25:03.232Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/43/600e5cfb0643d10d633124f5982d7abc2170dfd7ce985584ff16edab3e76/torchvision-0.24.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:7fb7590c737ebe3e1c077ad60c0e5e2e56bb26e7bccc3b9d04dbfc34fd09f050", size = 2386737, upload-time = "2025-11-12T15:25:08.288Z" },
+    { url = "https://files.pythonhosted.org/packages/93/b1/db2941526ecddd84884132e2742a55c9311296a6a38627f9e2627f5ac889/torchvision-0.24.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:66a98471fc18cad9064123106d810a75f57f0838eee20edc56233fd8484b0cc7", size = 8049868, upload-time = "2025-11-12T15:25:13.058Z" },
+    { url = "https://files.pythonhosted.org/packages/69/98/16e583f59f86cd59949f59d52bfa8fc286f86341a229a9d15cbe7a694f0c/torchvision-0.24.1-cp312-cp312-win_amd64.whl", hash = "sha256:4aa6cb806eb8541e92c9b313e96192c6b826e9eb0042720e2fa250d021079952", size = 4302006, upload-time = "2025-11-12T15:25:16.184Z" },
+]
+
 [[package]]
 name = "tqdm"
 version = "4.67.1"
@@ -2228,6 +2387,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f2/50/9a8358d3ef58162c0a415d173cfb45b67de60176e1024f71fbc4d24c0b6d/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d2c6b915a03888ab931a9fd3e55ba36785e1fe70cbea0b40c6ef93b20fc85232", size = 170470207, upload-time = "2025-11-11T17:41:00.253Z" },
 ]
 
+[[package]]
+name = "typer"
+version = "0.25.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "annotated-doc" },
+    { name = "click" },
+    { name = "rich" },
+    { name = "shellingham" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7b/27/ede8cec7596e0041ba7e7b80b47d132562f56ff454313a16f6084e555c9f/typer-0.25.0.tar.gz", hash = "sha256:123eaf9f19bb40fd268310e12a542c0c6b4fab9c98d9d23342a01ff95e3ce930", size = 120150, upload-time = "2026-04-26T08:46:14.767Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9a/72/193d4e586ec5a4db834a36bbeb47641a62f951f114ffd0fe5b1b46e8d56f/typer-0.25.0-py3-none-any.whl", hash = "sha256:ac01b48823d3db9a83c9e164338057eadbb1c9957a2a6b4eeb486669c560b5dc", size = 55993, upload-time = "2026-04-26T08:46:15.889Z" },
+]
+
 [[package]]
 name = "types-pytz"
 version = "2025.2.0.20251108"

From 1689d3eec5670c007952efa9eeaf02fde009d584 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 29 Apr 2026 22:48:57 +0200
Subject: [PATCH 09/49] fix: tweak submission script and concurrency

---
 configs/preprocessing/embeddings.yaml | 6 +++---
 scripts/submit_embeddings.py          | 9 ++++-----
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index efcc704..26e3373 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -2,9 +2,9 @@
 
 model: ???
 output_dir: ${project_dir}/embeddings
-concurrency: 1024
-block_size: 2048
-rows_per_file: 65536
+concurrency: 512
+block_size: 128
+rows_per_file: 5000
 
 metadata:
   run_name: "Embeddings: ${model}"
diff --git a/scripts/submit_embeddings.py b/scripts/submit_embeddings.py
index cbe0206..8ada923 100644
--- a/scripts/submit_embeddings.py
+++ b/scripts/submit_embeddings.py
@@ -4,16 +4,15 @@
 submit_job(
     job_name="tissue-classification-embeddings",
     username=...,
-    cpu=24,
-    memory="64Gi",
-    gpu="A40",
     public=False,
+    cpu=8,
+    memory="32Gi",
+    shm="16Gi",
     script=[
         "git clone https://github.com/RationAI/tissue-classification.git workdir",
         "cd workdir",
-        "export HF_TOKEN=<YOUR_HF_TOKEN>",
         "uv sync",
-        "uv run -m preprocessing.embeddings +experiment=preprocessing/embeddings_05mpp",
+        "uv run -m preprocessing.embeddings +experiment=...",
     ],
     storage=[storage.secure.DATA, storage.secure.PROJECTS],
 )

From 195806905a104bb57bc102396083353c36e70a51 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 29 Apr 2026 22:55:24 +0200
Subject: [PATCH 10/49] fix: resolve path names

---
 configs/preprocessing/embeddings.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 26e3373..443a18b 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -1,7 +1,7 @@
 # @package _global_
 
 model: ???
-output_dir: ${project_dir}/embeddings
+output_dir: ${project_path}/embeddings
 concurrency: 512
 block_size: 128
 rows_per_file: 5000

From ea79dd44c25875242e76c4f4304a6c388da503cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 29 Apr 2026 23:02:17 +0200
Subject: [PATCH 11/49] fix: tiling splits filenames

---
 preprocessing/embeddings.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index e8fda1c..b9dc4bb 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -40,8 +40,9 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
 @hydra.main(config_path="../configs", config_name="preprocessing", version_base=None)
 @autolog
 def main(config: DictConfig, logger: MLFlowLogger) -> None:
-    for name, split_uri in config.dataset.mlflow_uris.tiling.items():
-        folder = Path(mlflow.artifacts.download_artifacts(split_uri))
+    run_id = config.dataset.mlflow_artifacts.tiling_run_id
+    for name in ["train", "test"]:
+        folder = Path(mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=f"{name}_split"))
         slides = pd.read_parquet(folder / "slides.parquet")
         tiles = pd.read_parquet(folder / "tiles.parquet")
 

From 7e32bfb86b611c375cb95bd8ade054ff2cce91a5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 29 Apr 2026 23:19:39 +0200
Subject: [PATCH 12/49] fix: point rationai sdk to the github repo

---
 pyproject.toml | 2 +-
 uv.lock        | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index f197343..402d0cc 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,7 @@ dev = [
 
 [tool.uv.sources]
 rationai-mlkit = { git = "https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/mlkit.git" }
-rationai-sdk = { git = "https://gitlab.ics.muni.cz/rationai/infrastructure/rationai-sdk-python.git" }
+rationai-sdk = { git = "https://github.com/RationAI/rationai-sdk-python.git" }
 rationai-tiling = { git = "https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/tiling.git" }
 
 [tool.uv]
diff --git a/uv.lock b/uv.lock
index 4c405d0..aa74b3d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1752,7 +1752,7 @@ dependencies = [
 [[package]]
 name = "rationai-sdk"
 version = "0.1.0"
-source = { git = "https://gitlab.ics.muni.cz/rationai/infrastructure/rationai-sdk-python.git#a4d25084850cd26678783485dda87bbeed949492" }
+source = { git = "https://github.com/RationAI/rationai-sdk-python.git#3fb74c0f867432422dacde29ec5c0a4584af4830" }
 dependencies = [
     { name = "httpx" },
     { name = "lz4" },
@@ -2280,7 +2280,7 @@ requires-dist = [
     { name = "pandas", specifier = ">=2.0.0" },
     { name = "rationai-masks" },
     { name = "rationai-mlkit", git = "https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/mlkit.git" },
-    { name = "rationai-sdk", git = "https://gitlab.ics.muni.cz/rationai/infrastructure/rationai-sdk-python.git" },
+    { name = "rationai-sdk", git = "https://github.com/RationAI/rationai-sdk-python.git" },
     { name = "rationai-tiling", git = "https://gitlab.ics.muni.cz/rationai/digital-pathology/libraries/tiling.git" },
     { name = "ratiopath", specifier = ">=1.2.0" },
     { name = "ray", specifier = ">=2.51.1" },

From 53927517f3f239c52748577febd4f1c30ca805b8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 10:55:34 +0200
Subject: [PATCH 13/49] refactor: optimize data loading

---
 preprocessing/embeddings.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index b9dc4bb..d8ed429 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -6,7 +6,6 @@
 import hydra
 import mlflow.artifacts
 import pandas as pd
-import pyarrow as pa
 import ray
 from omegaconf import DictConfig
 from rationai import AsyncClient
@@ -44,16 +43,16 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
     for name in ["train", "test"]:
         folder = Path(mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=f"{name}_split"))
         slides = pd.read_parquet(folder / "slides.parquet")
-        tiles = pd.read_parquet(folder / "tiles.parquet")
 
-        slide_info = slides.set_index("id")[
-            ["path", "level", "tile_extent_x", "tile_extent_y"]
-        ]
-        tiles_enriched = tiles.join(slide_info, on="slide_id")
+        slide_info = (
+            slides.set_index("id")[["path", "level", "tile_extent_x", "tile_extent_y"]]
+            .to_dict("index")
+        )
 
-        ds = ray.data.from_arrow(
-            pa.Table.from_pandas(tiles_enriched, preserve_index=False)
-        ).repartition(target_num_rows_per_block=config.block_size)
+        ds = ray.data.read_parquet(str(folder / "tiles.parquet")).map(
+            lambda row, si: {**row, **si[row["slide_id"]]},
+            fn_kwargs={"si": slide_info},
+        )
         ds = ds.with_column(
             "tile",
             read_slide_tiles(  # pyright: ignore[reportCallIssue]

From d5e2c02985ad3bbb18a6acfb3521c3de3154fa00 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 21:26:52 +0200
Subject: [PATCH 14/49] feat: limit RAM consumption

---
 preprocessing/embeddings.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index d8ed429..e156843 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -41,15 +41,21 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
 def main(config: DictConfig, logger: MLFlowLogger) -> None:
     run_id = config.dataset.mlflow_artifacts.tiling_run_id
     for name in ["train", "test"]:
-        folder = Path(mlflow.artifacts.download_artifacts(run_id=run_id, artifact_path=f"{name}_split"))
+        folder = Path(
+            mlflow.artifacts.download_artifacts(
+                run_id=run_id, artifact_path=f"{name}_split"
+            )
+        )
         slides = pd.read_parquet(folder / "slides.parquet")
 
-        slide_info = (
-            slides.set_index("id")[["path", "level", "tile_extent_x", "tile_extent_y"]]
-            .to_dict("index")
-        )
+        slide_info = slides.set_index("id")[
+            ["path", "level", "tile_extent_x", "tile_extent_y"]
+        ].to_dict("index")
 
-        ds = ray.data.read_parquet(str(folder / "tiles.parquet")).map(
+        ds = ray.data.read_parquet(
+            str(folder / "tiles.parquet"),
+            ray_remote_args={"memory": 8 * 1024**3},
+        ).map(
             lambda row, si: {**row, **si[row["slide_id"]]},
             fn_kwargs={"si": slide_info},
         )
@@ -94,5 +100,8 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
     ctx.enable_rich_progress_bars = True
     ctx.use_ray_tqdm = False
 
-    with ray.init(runtime_env={"excludes": [".git", ".venv"]}):
+    with ray.init(
+        runtime_env={"excludes": [".git", ".venv"]},
+        object_store_memory=16 * 1024**3,
+    ):
         main()

From b2cb642fb8376c3db2665b0bdce16d288e5b0f2c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 21:33:24 +0200
Subject: [PATCH 15/49] feat: dynamically split blocks

---
 preprocessing/embeddings.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index e156843..6acfcf2 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -99,6 +99,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
     ctx = ray.data.DataContext.get_current()
     ctx.enable_rich_progress_bars = True
     ctx.use_ray_tqdm = False
+    ctx.target_max_block_size = 64 * 1024 * 1024
 
     with ray.init(
         runtime_env={"excludes": [".git", ".venv"]},

From 4adbb56e6acf78872f432cb2a10af98618a4bb21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 21:42:55 +0200
Subject: [PATCH 16/49] feat: raise the amount of block

---
 configs/preprocessing/embeddings.yaml | 2 +-
 preprocessing/embeddings.py           | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 443a18b..19bb9ca 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -3,7 +3,7 @@
 model: ???
 output_dir: ${project_path}/embeddings
 concurrency: 512
-block_size: 128
+block_size: 2000
 rows_per_file: 5000
 
 metadata:
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 6acfcf2..dceea68 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -6,6 +6,7 @@
 import hydra
 import mlflow.artifacts
 import pandas as pd
+import pyarrow.dataset as pads
 import ray
 from omegaconf import DictConfig
 from rationai import AsyncClient
@@ -52,9 +53,14 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
             ["path", "level", "tile_extent_x", "tile_extent_y"]
         ].to_dict("index")
 
+        tiles_path = folder / "tiles.parquet"
+        num_rows = pads.dataset(str(tiles_path), format="parquet").count_rows()
+        num_blocks = max(1, num_rows // config.block_size)
+
         ds = ray.data.read_parquet(
-            str(folder / "tiles.parquet"),
+            str(tiles_path),
             ray_remote_args={"memory": 8 * 1024**3},
+            override_num_blocks=num_blocks,
         ).map(
             lambda row, si: {**row, **si[row["slide_id"]]},
             fn_kwargs={"si": slide_info},

From 85a914c18d6d7d861117deab24dfc89050ba74cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 21:49:16 +0200
Subject: [PATCH 17/49] feat: add prints

---
 preprocessing/embeddings.py | 45 ++++++++++++++++++++++++++++++++++---
 1 file changed, 42 insertions(+), 3 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index dceea68..8ebe471 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,4 +1,5 @@
 import shutil
+import time
 from pathlib import Path
 from typing import Any
 
@@ -16,6 +17,22 @@
 from ray.data.expressions import col
 
 
+class JoinSlideInfo:
+    def __init__(self, slide_info: dict) -> None:
+        self.slide_info = slide_info
+        self.count = 0
+        self.start = time.monotonic()
+
+    def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
+        if self.count == 0:
+            print(f"[JoinSlideInfo] first row at t={time.monotonic() - self.start:.2f}s")
+        self.count += 1
+        if self.count % 50000 == 0:
+            elapsed = time.monotonic() - self.start
+            print(f"[JoinSlideInfo] {self.count} rows in {elapsed:.1f}s ({self.count / elapsed:.0f} rows/s)")
+        return {**row, **self.slide_info[row["slide_id"]]}
+
+
 class EmbedTiles:
     def __init__(self, model: str, concurrency: int) -> None:
         self.model = model
@@ -24,13 +41,24 @@ def __init__(self, model: str, concurrency: int) -> None:
                 max_connections=concurrency, max_keepalive_connections=concurrency
             )
         )
+        self.count = 0
+        self.start = time.monotonic()
+        print(f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}")
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
+        if self.count == 0:
+            print(f"[EmbedTiles] first row at t={time.monotonic() - self.start:.2f}s")
+        t0 = time.monotonic()
         embedding = (
             (await self.client.models.embed_image(self.model, row["tile"]))
             .reshape(-1)
             .tolist()
         )
+        latency = time.monotonic() - t0
+        self.count += 1
+        if self.count % 100 == 0:
+            elapsed = time.monotonic() - self.start
+            print(f"[EmbedTiles] {self.count} embeddings in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, last latency={latency * 1000:.0f}ms)")
         del row["tile"]
         row["embedding"] = embedding
         return row
@@ -42,28 +70,35 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
 def main(config: DictConfig, logger: MLFlowLogger) -> None:
     run_id = config.dataset.mlflow_artifacts.tiling_run_id
     for name in ["train", "test"]:
+        t = time.monotonic()
+        print(f"[main] === split={name} ===")
         folder = Path(
             mlflow.artifacts.download_artifacts(
                 run_id=run_id, artifact_path=f"{name}_split"
             )
         )
-        slides = pd.read_parquet(folder / "slides.parquet")
+        print(f"[main] artifacts downloaded in {time.monotonic() - t:.1f}s")
 
+        t = time.monotonic()
+        slides = pd.read_parquet(folder / "slides.parquet")
         slide_info = slides.set_index("id")[
             ["path", "level", "tile_extent_x", "tile_extent_y"]
         ].to_dict("index")
+        print(f"[main] loaded {len(slide_info)} slides in {time.monotonic() - t:.1f}s")
 
+        t = time.monotonic()
         tiles_path = folder / "tiles.parquet"
         num_rows = pads.dataset(str(tiles_path), format="parquet").count_rows()
         num_blocks = max(1, num_rows // config.block_size)
+        print(f"[main] {num_rows} tile rows, {num_blocks} blocks (parquet metadata in {time.monotonic() - t:.1f}s)")
 
         ds = ray.data.read_parquet(
             str(tiles_path),
             ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
-            lambda row, si: {**row, **si[row["slide_id"]]},
-            fn_kwargs={"si": slide_info},
+            JoinSlideInfo,  # pyright: ignore[reportArgumentType]
+            fn_constructor_args=(slide_info,),
         )
         ds = ds.with_column(
             "tile",
@@ -96,7 +131,11 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
             shutil.rmtree(tiles_parquet_dir)
 
         slides.to_parquet(split_dir / "slides.parquet", index=False)
+
+        t = time.monotonic()
+        print(f"[main] starting write_parquet for split={name}")
         ds.write_parquet(str(tiles_parquet_dir), min_rows_per_file=config.rows_per_file)
+        print(f"[main] write_parquet finished in {time.monotonic() - t:.1f}s")
 
         logger.log_artifacts(str(split_dir), str(name))
 

From 06c818afa2a5f3acf59f3cb5a13354c3e7c7480a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 21:57:57 +0200
Subject: [PATCH 18/49] feat: change the cpu balance

---
 preprocessing/embeddings.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 8ebe471..e8118e2 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -110,8 +110,8 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
                 col("tile_extent_y"),
                 col("level"),
             ),
-            num_cpus=1,
-            memory=4 * 1024**3,
+            num_cpus=0.25,
+            memory=1 * 1024**3,
         )
         ds = ds.drop_columns(["path", "level", "tile_extent_x", "tile_extent_y"])
         ds = ds.map(
@@ -148,6 +148,6 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
 
     with ray.init(
         runtime_env={"excludes": [".git", ".venv"]},
-        object_store_memory=16 * 1024**3,
+        object_store_memory=32 * 1024**3,
     ):
         main()

From 014d11fff2ae7f288a3440c4b4348965dd9f4639 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 22:04:53 +0200
Subject: [PATCH 19/49] fix: change concurrency

---
 preprocessing/embeddings.py | 20 ++------------------
 1 file changed, 2 insertions(+), 18 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index e8118e2..e1bc6d8 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -17,22 +17,6 @@
 from ray.data.expressions import col
 
 
-class JoinSlideInfo:
-    def __init__(self, slide_info: dict) -> None:
-        self.slide_info = slide_info
-        self.count = 0
-        self.start = time.monotonic()
-
-    def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
-        if self.count == 0:
-            print(f"[JoinSlideInfo] first row at t={time.monotonic() - self.start:.2f}s")
-        self.count += 1
-        if self.count % 50000 == 0:
-            elapsed = time.monotonic() - self.start
-            print(f"[JoinSlideInfo] {self.count} rows in {elapsed:.1f}s ({self.count / elapsed:.0f} rows/s)")
-        return {**row, **self.slide_info[row["slide_id"]]}
-
-
 class EmbedTiles:
     def __init__(self, model: str, concurrency: int) -> None:
         self.model = model
@@ -97,8 +81,8 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
             ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
-            JoinSlideInfo,  # pyright: ignore[reportArgumentType]
-            fn_constructor_args=(slide_info,),
+            lambda row, si: {**row, **si[row["slide_id"]]},
+            fn_kwargs={"si": slide_info},
         )
         ds = ds.with_column(
             "tile",

From f3e4cf48c8a4feadbb52d1ac238d407cdf452f58 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 22:10:24 +0200
Subject: [PATCH 20/49] fix: lower concurrency

---
 configs/preprocessing/embeddings.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 19bb9ca..84b7ae0 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -2,7 +2,7 @@
 
 model: ???
 output_dir: ${project_path}/embeddings
-concurrency: 512
+concurrency: 128
 block_size: 2000
 rows_per_file: 5000
 

From 606b8e40524d8a726a3dd47114dd139fe3da1aab Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 30 Apr 2026 22:22:42 +0200
Subject: [PATCH 21/49] feat: add timout, raise concurrency

---
 configs/preprocessing/embeddings.yaml |  2 +-
 preprocessing/embeddings.py           | 20 +++++++++++++-------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 84b7ae0..19bb9ca 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -2,7 +2,7 @@
 
 model: ???
 output_dir: ${project_path}/embeddings
-concurrency: 128
+concurrency: 512
 block_size: 2000
 rows_per_file: 5000
 
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index e1bc6d8..4718c99 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -23,15 +23,20 @@ def __init__(self, model: str, concurrency: int) -> None:
         self.client = AsyncClient(
             limits=httpx.Limits(
                 max_connections=concurrency, max_keepalive_connections=concurrency
-            )
+            ),
+            timeout=200,
         )
         self.count = 0
+        self.in_flight = 0
+        self.first_row_logged = False
         self.start = time.monotonic()
         print(f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}")
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
-        if self.count == 0:
+        if not self.first_row_logged:
+            self.first_row_logged = True
             print(f"[EmbedTiles] first row at t={time.monotonic() - self.start:.2f}s")
+        self.in_flight += 1
         t0 = time.monotonic()
         embedding = (
             (await self.client.models.embed_image(self.model, row["tile"]))
@@ -40,9 +45,10 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         )
         latency = time.monotonic() - t0
         self.count += 1
-        if self.count % 100 == 0:
+        self.in_flight -= 1
+        if self.count % 50 == 0:
             elapsed = time.monotonic() - self.start
-            print(f"[EmbedTiles] {self.count} embeddings in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, last latency={latency * 1000:.0f}ms)")
+            print(f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)")
         del row["tile"]
         row["embedding"] = embedding
         return row
@@ -94,8 +100,8 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
                 col("tile_extent_y"),
                 col("level"),
             ),
-            num_cpus=0.25,
-            memory=1 * 1024**3,
+            num_cpus=1,
+            memory=4 * 1024**3,
         )
         ds = ds.drop_columns(["path", "level", "tile_extent_x", "tile_extent_y"])
         ds = ds.map(
@@ -132,6 +138,6 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
 
     with ray.init(
         runtime_env={"excludes": [".git", ".venv"]},
-        object_store_memory=32 * 1024**3,
+        object_store_memory=16 * 1024**3,
     ):
         main()

From d36dd98804d269313b1723ac8ed2a796823792e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Mon, 4 May 2026 18:08:19 +0200
Subject: [PATCH 22/49] fix: format

---
 preprocessing/embeddings.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 4718c99..9300201 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -10,7 +10,7 @@
 import pyarrow.dataset as pads
 import ray
 from omegaconf import DictConfig
-from rationai import AsyncClient
+from rationai import AsyncClient  # type: ignore[attr-defined]
 from rationai.mlkit import autolog, with_cli_args
 from rationai.mlkit.lightning.loggers import MLFlowLogger
 from ratiopath.tiling.read_slide_tiles import read_slide_tiles
@@ -30,7 +30,9 @@ def __init__(self, model: str, concurrency: int) -> None:
         self.in_flight = 0
         self.first_row_logged = False
         self.start = time.monotonic()
-        print(f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}")
+        print(
+            f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
+        )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         if not self.first_row_logged:
@@ -48,7 +50,9 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         self.in_flight -= 1
         if self.count % 50 == 0:
             elapsed = time.monotonic() - self.start
-            print(f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)")
+            print(
+                f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)"
+            )
         del row["tile"]
         row["embedding"] = embedding
         return row
@@ -80,7 +84,9 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         tiles_path = folder / "tiles.parquet"
         num_rows = pads.dataset(str(tiles_path), format="parquet").count_rows()
         num_blocks = max(1, num_rows // config.block_size)
-        print(f"[main] {num_rows} tile rows, {num_blocks} blocks (parquet metadata in {time.monotonic() - t:.1f}s)")
+        print(
+            f"[main] {num_rows} tile rows, {num_blocks} blocks (parquet metadata in {time.monotonic() - t:.1f}s)"
+        )
 
         ds = ray.data.read_parquet(
             str(tiles_path),

From 531f5f8ba2dba1bf0c8c6c13765111281b6e5f3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Mon, 4 May 2026 19:46:16 +0200
Subject: [PATCH 23/49] fix: use more memory

---
 scripts/submit_embeddings.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scripts/submit_embeddings.py b/scripts/submit_embeddings.py
index 8ada923..d2f287d 100644
--- a/scripts/submit_embeddings.py
+++ b/scripts/submit_embeddings.py
@@ -6,8 +6,8 @@
     username=...,
     public=False,
     cpu=8,
-    memory="32Gi",
-    shm="16Gi",
+    memory="64Gi",
+    shm="24Gi",
     script=[
         "git clone https://github.com/RationAI/tissue-classification.git workdir",
         "cd workdir",

From 7bd872513160ae34f8c1b0cb504f021d9fdc419a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Tue, 5 May 2026 10:50:05 +0200
Subject: [PATCH 24/49] fix: retry embed_image on failure and always free tile
 data

Wraps the embed_image call in a retry loop (up to 3 attempts with
exponential backoff) so transient network errors don't cause Ray to
retry the entire block. Re-raises after all attempts are exhausted so
failures stay visible. Moves del row["tile"] into a finally block to
free tile pixel data promptly even when an exception occurs.
---
 preprocessing/embeddings.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 9300201..619fcd3 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,3 +1,4 @@
+import asyncio
 import shutil
 import time
 from pathlib import Path
@@ -40,11 +41,24 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
             print(f"[EmbedTiles] first row at t={time.monotonic() - self.start:.2f}s")
         self.in_flight += 1
         t0 = time.monotonic()
-        embedding = (
-            (await self.client.models.embed_image(self.model, row["tile"]))
-            .reshape(-1)
-            .tolist()
-        )
+        try:
+            last_exc: Exception | None = None
+            for attempt in range(3):
+                try:
+                    embedding = (
+                        (await self.client.models.embed_image(self.model, row["tile"]))
+                        .reshape(-1)
+                        .tolist()
+                    )
+                    break
+                except Exception as exc:
+                    last_exc = exc
+                    if attempt < 2:
+                        await asyncio.sleep(2**attempt)
+            else:
+                raise RuntimeError("embed_image failed after 3 attempts") from last_exc
+        finally:
+            del row["tile"]
         latency = time.monotonic() - t0
         self.count += 1
         self.in_flight -= 1
@@ -53,7 +67,6 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
             print(
                 f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)"
             )
-        del row["tile"]
         row["embedding"] = embedding
         return row
 

From b23a0e58b72684fce082078f0c9e49f94de98b09 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Tue, 5 May 2026 15:22:17 +0200
Subject: [PATCH 25/49] feat: filter zero-coverage tiles before embedding

Skip tiles with no annotation coverage and no tissue coverage before
feeding them into the Ray pipeline, using PyArrow predicate pushdown
to avoid materialising the full 80M-row dataset in memory.

Tissue stats run ID stored in dataset.yaml; referenced from the
embeddings config via tile_filters.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 configs/data/dataset.yaml             |  1 +
 configs/preprocessing/embeddings.yaml |  4 ++
 preprocessing/embeddings.py           | 56 ++++++++++++++++++++++++++-
 3 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/configs/data/dataset.yaml b/configs/data/dataset.yaml
index 57d87c5..0303fc2 100644
--- a/configs/data/dataset.yaml
+++ b/configs/data/dataset.yaml
@@ -14,6 +14,7 @@ dataset:
     test_split_filename: "split_mapping/test_split.csv"
     tiling_run_id: "fdf7550a2004474f8c7a05dc0cf1fd86"
     tissue_masks_run_id: "52bc0924f8624b259819c480c7cf213f"
+    tissue_stats_run_id: "16ae2d003d88471b924e5f332415232a"
 
   exclusions:
     bad_slides:
diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 19bb9ca..d7b54d6 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -6,6 +6,10 @@ concurrency: 512
 block_size: 2000
 rows_per_file: 5000
 
+tile_filters:
+  tissue_stats_run_id: ${dataset.mlflow_artifacts.tissue_stats_run_id}
+  tissue_stats_artifact_path: tissue_stats
+
 metadata:
   run_name: "Embeddings: ${model}"
   description: Tile embeddings using model ${model}
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 619fcd3..9edc17d 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -8,7 +8,9 @@
 import hydra
 import mlflow.artifacts
 import pandas as pd
+import pyarrow.compute as pc
 import pyarrow.dataset as pads
+import pyarrow.parquet as pq
 import ray
 from omegaconf import DictConfig
 from rationai import AsyncClient  # type: ignore[attr-defined]
@@ -18,6 +20,56 @@
 from ray.data.expressions import col
 
 
+def filter_tiles(folder: Path, config: DictConfig, split_name: str) -> Path:
+    """Drop tiles that will never be used for training and return the parquet path to use.
+
+    Removes tiles where all annotation class coverages are zero (unlabeled) and
+    tiles with zero tissue coverage. Uses PyArrow predicate pushdown so only
+    matching rows are materialised — safe for 80M-row datasets.
+    """
+    tiles_path = folder / "tiles.parquet"
+    tiles_ds = pads.dataset(str(tiles_path), format="parquet")
+    original_count = tiles_ds.count_rows()
+
+    ann_cols = [f.name for f in tiles_ds.schema if f.name.startswith("tile_coverage_")]
+    ann_filter = None
+    if ann_cols:
+        ann_filter = pads.field(ann_cols[0]) > 0
+        for c in ann_cols[1:]:
+            ann_filter = pc.or_(ann_filter, pads.field(c) > 0)
+
+    tiles_table = tiles_ds.to_table(filter=ann_filter)
+    ann_count = len(tiles_table)
+    print(
+        f"[filter_tiles] {split_name}: annotation filter"
+        f" {original_count} → {ann_count} ({ann_count / original_count:.1%} kept)"
+    )
+
+    tissue_local = mlflow.artifacts.download_artifacts(
+        run_id=config.tile_filters.tissue_stats_run_id,
+        artifact_path=f"{config.tile_filters.tissue_stats_artifact_path}/{split_name}_tiles.parquet",
+    )
+    tissue_table = pads.dataset(tissue_local, format="parquet").to_table(
+        columns=["slide_id", "x", "y"],
+        filter=pads.field("tile_tissue_coverage") > 0,
+    )
+    tiles_table = tiles_table.join(
+        tissue_table, keys=["slide_id", "x", "y"], join_type="inner"
+    )
+    tissue_count = len(tiles_table)
+    print(
+        f"[filter_tiles] {split_name}: tissue filter"
+        f" {ann_count} → {tissue_count} ({tissue_count / ann_count:.1%} kept)"
+    )
+
+    if tissue_count == original_count:
+        return tiles_path
+
+    filtered_path = folder / "tiles_filtered.parquet"
+    pq.write_table(tiles_table, str(filtered_path))
+    return filtered_path
+
+
 class EmbedTiles:
     def __init__(self, model: str, concurrency: int) -> None:
         self.model = model
@@ -94,11 +146,11 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         print(f"[main] loaded {len(slide_info)} slides in {time.monotonic() - t:.1f}s")
 
         t = time.monotonic()
-        tiles_path = folder / "tiles.parquet"
+        tiles_path = filter_tiles(folder, config, name)
         num_rows = pads.dataset(str(tiles_path), format="parquet").count_rows()
         num_blocks = max(1, num_rows // config.block_size)
         print(
-            f"[main] {num_rows} tile rows, {num_blocks} blocks (parquet metadata in {time.monotonic() - t:.1f}s)"
+            f"[main] {num_rows} tile rows, {num_blocks} blocks (filter+metadata in {time.monotonic() - t:.1f}s)"
         )
 
         ds = ray.data.read_parquet(

From 60a997127b9d0ff2be4085ce9380f08fb77f9ca5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Tue, 5 May 2026 15:25:29 +0200
Subject: [PATCH 26/49] fix: use | operator for PyArrow expression OR, not
 pc.or_

pc.or_ is an array kernel; expression combination uses the | operator.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 9edc17d..d72f351 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -8,7 +8,6 @@
 import hydra
 import mlflow.artifacts
 import pandas as pd
-import pyarrow.compute as pc
 import pyarrow.dataset as pads
 import pyarrow.parquet as pq
 import ray
@@ -36,7 +35,7 @@ def filter_tiles(folder: Path, config: DictConfig, split_name: str) -> Path:
     if ann_cols:
         ann_filter = pads.field(ann_cols[0]) > 0
         for c in ann_cols[1:]:
-            ann_filter = pc.or_(ann_filter, pads.field(c) > 0)
+            ann_filter = ann_filter | (pads.field(c) > 0)
 
     tiles_table = tiles_ds.to_table(filter=ann_filter)
     ann_count = len(tiles_table)

From 76dafae39090bc4c087ac074066b1d7c4bf7d1cf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Tue, 5 May 2026 15:56:35 +0200
Subject: [PATCH 27/49] chore: edge cases

---
 .../preprocessing/embeddings_05mpp.yaml       |  5 ----
 preprocessing/embeddings.py                   | 27 +++++++++++++------
 2 files changed, 19 insertions(+), 13 deletions(-)
 delete mode 100644 configs/experiment/preprocessing/embeddings_05mpp.yaml

diff --git a/configs/experiment/preprocessing/embeddings_05mpp.yaml b/configs/experiment/preprocessing/embeddings_05mpp.yaml
deleted file mode 100644
index 3891b97..0000000
--- a/configs/experiment/preprocessing/embeddings_05mpp.yaml
+++ /dev/null
@@ -1,5 +0,0 @@
-# @package _global_
-
-defaults:
-  - /data: dataset
-  - _self_
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index d72f351..1c36354 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -23,19 +23,24 @@ def filter_tiles(folder: Path, config: DictConfig, split_name: str) -> Path:
     """Drop tiles that will never be used for training and return the parquet path to use.
 
     Removes tiles where all annotation class coverages are zero (unlabeled) and
-    tiles with zero tissue coverage. Uses PyArrow predicate pushdown so only
-    matching rows are materialised — safe for 80M-row datasets.
+    tiles with zero tissue coverage. Uses PyArrow predicate pushdown so the full
+    80M-row input is never loaded into memory — only the post-filter result is
+    materialised, so peak memory scales with the number of tiles that pass the
+    annotation filter, not the total tile count.
     """
     tiles_path = folder / "tiles.parquet"
     tiles_ds = pads.dataset(str(tiles_path), format="parquet")
     original_count = tiles_ds.count_rows()
 
     ann_cols = [f.name for f in tiles_ds.schema if f.name.startswith("tile_coverage_")]
-    ann_filter = None
-    if ann_cols:
-        ann_filter = pads.field(ann_cols[0]) > 0
-        for c in ann_cols[1:]:
-            ann_filter = ann_filter | (pads.field(c) > 0)
+    if not ann_cols:
+        raise RuntimeError(
+            "No tile_coverage_* columns found in tiles parquet. "
+            "Check that the tiling run used a class mapping with annotations."
+        )
+    ann_filter = pads.field(ann_cols[0]) > 0
+    for c in ann_cols[1:]:
+        ann_filter = ann_filter | (pads.field(c) > 0)
 
     tiles_table = tiles_ds.to_table(filter=ann_filter)
     ann_count = len(tiles_table)
@@ -43,6 +48,11 @@ def filter_tiles(folder: Path, config: DictConfig, split_name: str) -> Path:
         f"[filter_tiles] {split_name}: annotation filter"
         f" {original_count} → {ann_count} ({ann_count / original_count:.1%} kept)"
     )
+    if ann_count == 0:
+        raise RuntimeError(
+            f"All {original_count} tiles were filtered out by annotation coverage. "
+            "Check that the tiling run used the correct class mapping and annotation sources."
+        )
 
     tissue_local = mlflow.artifacts.download_artifacts(
         run_id=config.tile_filters.tissue_stats_run_id,
@@ -154,6 +164,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
 
         ds = ray.data.read_parquet(
             str(tiles_path),
+            columns=["slide_id", "x", "y"],
             ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
@@ -179,7 +190,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
             fn_constructor_args=(config.model, config.concurrency),
             compute=ray.data.ActorPoolStrategy(
                 max_size=4,
-                max_tasks_in_flight_per_actor=config.concurrency // 4,
+                max_tasks_in_flight_per_actor=max(1, config.concurrency // 4),
             ),
             max_concurrency=config.concurrency,
         )

From 59a6516627af7a873f8bb2d43fa513bb45ef87cc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Tue, 5 May 2026 20:26:53 +0200
Subject: [PATCH 28/49] refactor: use filter_tiles output in embeddings instead
 of inline filtering

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 configs/data/dataset.yaml             |  1 +
 configs/preprocessing/embeddings.yaml |  4 --
 preprocessing/embeddings.py           | 86 ++++-----------------------
 3 files changed, 11 insertions(+), 80 deletions(-)

diff --git a/configs/data/dataset.yaml b/configs/data/dataset.yaml
index 0303fc2..bd5afc4 100644
--- a/configs/data/dataset.yaml
+++ b/configs/data/dataset.yaml
@@ -15,6 +15,7 @@ dataset:
     tiling_run_id: "fdf7550a2004474f8c7a05dc0cf1fd86"
     tissue_masks_run_id: "52bc0924f8624b259819c480c7cf213f"
     tissue_stats_run_id: "16ae2d003d88471b924e5f332415232a"
+    filter_tiles_run_id: "???"
 
   exclusions:
     bad_slides:
diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index d7b54d6..19bb9ca 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -6,10 +6,6 @@ concurrency: 512
 block_size: 2000
 rows_per_file: 5000
 
-tile_filters:
-  tissue_stats_run_id: ${dataset.mlflow_artifacts.tissue_stats_run_id}
-  tissue_stats_artifact_path: tissue_stats
-
 metadata:
   run_name: "Embeddings: ${model}"
   description: Tile embeddings using model ${model}
diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 1c36354..3a940d8 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -9,7 +9,6 @@
 import mlflow.artifacts
 import pandas as pd
 import pyarrow.dataset as pads
-import pyarrow.parquet as pq
 import ray
 from omegaconf import DictConfig
 from rationai import AsyncClient  # type: ignore[attr-defined]
@@ -19,66 +18,6 @@
 from ray.data.expressions import col
 
 
-def filter_tiles(folder: Path, config: DictConfig, split_name: str) -> Path:
-    """Drop tiles that will never be used for training and return the parquet path to use.
-
-    Removes tiles where all annotation class coverages are zero (unlabeled) and
-    tiles with zero tissue coverage. Uses PyArrow predicate pushdown so the full
-    80M-row input is never loaded into memory — only the post-filter result is
-    materialised, so peak memory scales with the number of tiles that pass the
-    annotation filter, not the total tile count.
-    """
-    tiles_path = folder / "tiles.parquet"
-    tiles_ds = pads.dataset(str(tiles_path), format="parquet")
-    original_count = tiles_ds.count_rows()
-
-    ann_cols = [f.name for f in tiles_ds.schema if f.name.startswith("tile_coverage_")]
-    if not ann_cols:
-        raise RuntimeError(
-            "No tile_coverage_* columns found in tiles parquet. "
-            "Check that the tiling run used a class mapping with annotations."
-        )
-    ann_filter = pads.field(ann_cols[0]) > 0
-    for c in ann_cols[1:]:
-        ann_filter = ann_filter | (pads.field(c) > 0)
-
-    tiles_table = tiles_ds.to_table(filter=ann_filter)
-    ann_count = len(tiles_table)
-    print(
-        f"[filter_tiles] {split_name}: annotation filter"
-        f" {original_count} → {ann_count} ({ann_count / original_count:.1%} kept)"
-    )
-    if ann_count == 0:
-        raise RuntimeError(
-            f"All {original_count} tiles were filtered out by annotation coverage. "
-            "Check that the tiling run used the correct class mapping and annotation sources."
-        )
-
-    tissue_local = mlflow.artifacts.download_artifacts(
-        run_id=config.tile_filters.tissue_stats_run_id,
-        artifact_path=f"{config.tile_filters.tissue_stats_artifact_path}/{split_name}_tiles.parquet",
-    )
-    tissue_table = pads.dataset(tissue_local, format="parquet").to_table(
-        columns=["slide_id", "x", "y"],
-        filter=pads.field("tile_tissue_coverage") > 0,
-    )
-    tiles_table = tiles_table.join(
-        tissue_table, keys=["slide_id", "x", "y"], join_type="inner"
-    )
-    tissue_count = len(tiles_table)
-    print(
-        f"[filter_tiles] {split_name}: tissue filter"
-        f" {ann_count} → {tissue_count} ({tissue_count / ann_count:.1%} kept)"
-    )
-
-    if tissue_count == original_count:
-        return tiles_path
-
-    filtered_path = folder / "tiles_filtered.parquet"
-    pq.write_table(tiles_table, str(filtered_path))
-    return filtered_path
-
-
 class EmbedTiles:
     def __init__(self, model: str, concurrency: int) -> None:
         self.model = model
@@ -136,31 +75,26 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
 @hydra.main(config_path="../configs", config_name="preprocessing", version_base=None)
 @autolog
 def main(config: DictConfig, logger: MLFlowLogger) -> None:
-    run_id = config.dataset.mlflow_artifacts.tiling_run_id
     for name in ["train", "test"]:
-        t = time.monotonic()
-        print(f"[main] === split={name} ===")
-        folder = Path(
+        split_folder = Path(
             mlflow.artifacts.download_artifacts(
-                run_id=run_id, artifact_path=f"{name}_split"
+                run_id=config.dataset.mlflow_artifacts.tiling_run_id,
+                artifact_path=f"{name}_split",
             )
         )
-        print(f"[main] artifacts downloaded in {time.monotonic() - t:.1f}s")
-
-        t = time.monotonic()
-        slides = pd.read_parquet(folder / "slides.parquet")
+        slides = pd.read_parquet(split_folder / "slides.parquet")
         slide_info = slides.set_index("id")[
             ["path", "level", "tile_extent_x", "tile_extent_y"]
         ].to_dict("index")
-        print(f"[main] loaded {len(slide_info)} slides in {time.monotonic() - t:.1f}s")
 
-        t = time.monotonic()
-        tiles_path = filter_tiles(folder, config, name)
+        tiles_path = Path(
+            mlflow.artifacts.download_artifacts(
+                run_id=config.dataset.mlflow_artifacts.filter_tiles_run_id,
+                artifact_path=f"filter_tiles/{name}_tiles.parquet",
+            )
+        )
         num_rows = pads.dataset(str(tiles_path), format="parquet").count_rows()
         num_blocks = max(1, num_rows // config.block_size)
-        print(
-            f"[main] {num_rows} tile rows, {num_blocks} blocks (filter+metadata in {time.monotonic() - t:.1f}s)"
-        )
 
         ds = ray.data.read_parquet(
             str(tiles_path),

From d42561b13b7b9c56acc3888b46bf18a33d635913 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Tue, 5 May 2026 21:04:18 +0200
Subject: [PATCH 29/49] feat: add temporary filter run id for testing

---
 configs/data/dataset.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/configs/data/dataset.yaml b/configs/data/dataset.yaml
index bd5afc4..abb46a9 100644
--- a/configs/data/dataset.yaml
+++ b/configs/data/dataset.yaml
@@ -15,7 +15,7 @@ dataset:
     tiling_run_id: "fdf7550a2004474f8c7a05dc0cf1fd86"
     tissue_masks_run_id: "52bc0924f8624b259819c480c7cf213f"
     tissue_stats_run_id: "16ae2d003d88471b924e5f332415232a"
-    filter_tiles_run_id: "???"
+    filter_tiles_run_id: "b51fa26332d84dad94dc6e97477ed061"
 
   exclusions:
     bad_slides:

From e109dc5ad9104dc7ca20ab777bd346133419ba70 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 20:22:41 +0200
Subject: [PATCH 30/49] feat: generate tile masks from filtered tiles for both
 train and test splits

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 configs/data/dataset.yaml                     |  2 +-
 .../preprocessing/tile_masks_05mpp.yaml       |  2 -
 configs/preprocessing/tile_masks.yaml         |  2 -
 preprocessing/tile_masks.py                   | 46 ++++++++++---------
 4 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/configs/data/dataset.yaml b/configs/data/dataset.yaml
index abb46a9..d971bc1 100644
--- a/configs/data/dataset.yaml
+++ b/configs/data/dataset.yaml
@@ -15,7 +15,7 @@ dataset:
     tiling_run_id: "fdf7550a2004474f8c7a05dc0cf1fd86"
     tissue_masks_run_id: "52bc0924f8624b259819c480c7cf213f"
     tissue_stats_run_id: "16ae2d003d88471b924e5f332415232a"
-    filter_tiles_run_id: "b51fa26332d84dad94dc6e97477ed061"
+    filter_tiles_run_id: "4e8f5d3c82124ea5a8f871a42d3ed9ba"
 
   exclusions:
     bad_slides:
diff --git a/configs/experiment/preprocessing/tile_masks_05mpp.yaml b/configs/experiment/preprocessing/tile_masks_05mpp.yaml
index 8cfc5a9..4fb3ab5 100644
--- a/configs/experiment/preprocessing/tile_masks_05mpp.yaml
+++ b/configs/experiment/preprocessing/tile_masks_05mpp.yaml
@@ -4,8 +4,6 @@ defaults:
   - /data: dataset
   - _self_
 
-slides_artifact_path: train_split/slides.parquet
-tiles_artifact_path: train_split/tiles.parquet
 tile_percentage_cols:
   - tile_coverage_Nerve
   - tile_coverage_Blood
diff --git a/configs/preprocessing/tile_masks.yaml b/configs/preprocessing/tile_masks.yaml
index a287440..27fbd96 100644
--- a/configs/preprocessing/tile_masks.yaml
+++ b/configs/preprocessing/tile_masks.yaml
@@ -1,7 +1,5 @@
 # @package _global_
 
-slides_artifact_path: ???
-tiles_artifact_path: ???
 tile_percentage_cols: ???
 
 max_concurrent: 1
diff --git a/preprocessing/tile_masks.py b/preprocessing/tile_masks.py
index 9738f3d..721a1cd 100644
--- a/preprocessing/tile_masks.py
+++ b/preprocessing/tile_masks.py
@@ -74,34 +74,36 @@ def process_slide(
 @autolog
 def main(config: DictConfig, logger: MLFlowLogger) -> None:
     tiling_run_id = config.dataset.mlflow_artifacts.tiling_run_id
+    filter_tiles_run_id = config.dataset.mlflow_artifacts.filter_tiles_run_id
     tile_percentage_cols: list[str] = list(config.tile_percentage_cols)
 
-    slides_path = mlflow.artifacts.download_artifacts(
-        run_id=tiling_run_id,
-        artifact_path=config.slides_artifact_path,
-    )
-    tiles_path = mlflow.artifacts.download_artifacts(
-        run_id=tiling_run_id,
-        artifact_path=config.tiles_artifact_path,
-    )
-    slides = pd.read_parquet(slides_path)
+    for split_name in ("train", "test"):
+        slides_path = mlflow.artifacts.download_artifacts(
+            run_id=tiling_run_id,
+            artifact_path=f"{split_name}_split/slides.parquet",
+        )
+        tiles_path = mlflow.artifacts.download_artifacts(
+            run_id=filter_tiles_run_id,
+            artifact_path=f"filter_tiles/{split_name}_tiles.parquet",
+        )
+        slides = pd.read_parquet(slides_path)
+        items = cast("list[dict[str, Any]]", slides.to_dict(orient="records"))
 
-    items = cast("list[dict[str, Any]]", slides.to_dict(orient="records"))
+        with TemporaryDirectory() as output_dir:
+            Path(output_dir, "outlines").mkdir()
+            for slide in tqdm(items, desc=split_name):
+                process_slide(
+                    slide,
+                    output_dir=output_dir,
+                    tile_percentage_cols=tile_percentage_cols,
+                    tiles_path=tiles_path,
+                )
 
-    with TemporaryDirectory() as output_dir:
-        Path(output_dir, "outlines").mkdir()
-        for slide in tqdm(items):
-            process_slide(
-                slide,
-                output_dir=output_dir,
-                tile_percentage_cols=tile_percentage_cols,
-                tiles_path=tiles_path,
+            logger.log_artifacts(
+                local_dir=output_dir,
+                artifact_path=f"{config.mlflow_artifact_path}/{split_name}_split",
             )
 
-        logger.log_artifacts(
-            local_dir=output_dir, artifact_path=config.mlflow_artifact_path
-        )
-
 
 if __name__ == "__main__":
     main()

From 1d99fcb0610a869c135641ae1f9b751ecba9e90f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:27:18 +0200
Subject: [PATCH 31/49] fix: change block size to the power two

---
 configs/preprocessing/embeddings.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/configs/preprocessing/embeddings.yaml b/configs/preprocessing/embeddings.yaml
index 19bb9ca..db9ceaa 100644
--- a/configs/preprocessing/embeddings.yaml
+++ b/configs/preprocessing/embeddings.yaml
@@ -3,7 +3,7 @@
 model: ???
 output_dir: ${project_path}/embeddings
 concurrency: 512
-block_size: 2000
+block_size: 2048
 rows_per_file: 5000
 
 metadata:

From ef74b6caaa7f4fece427bbf21a99a8b7596bed7e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:30:24 +0200
Subject: [PATCH 32/49] chore: remove first-row debug log from EmbedTiles

Ray actor logs and progress bars provide cold-start visibility;
the bespoke first_row_logged flag was a tuning aid no longer needed.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 3a940d8..8717a00 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -29,16 +29,12 @@ def __init__(self, model: str, concurrency: int) -> None:
         )
         self.count = 0
         self.in_flight = 0
-        self.first_row_logged = False
         self.start = time.monotonic()
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
-        if not self.first_row_logged:
-            self.first_row_logged = True
-            print(f"[EmbedTiles] first row at t={time.monotonic() - self.start:.2f}s")
         self.in_flight += 1
         t0 = time.monotonic()
         try:

From 82163ac316b715d7968528e276e45c985869deff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:35:17 +0200
Subject: [PATCH 33/49] refactor: use tenacity for embed retries, narrow to
 httpx errors

Replaces hand-rolled retry loop in EmbedTiles with a tenacity-decorated
helper. Retries are now scoped to httpx.HTTPError (network/timeout/status)
so programming bugs surface immediately instead of being retried 3x.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 34 ++++++++++++++++++----------------
 pyproject.toml              |  1 +
 uv.lock                     |  2 ++
 3 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 8717a00..7bb73f6 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,4 +1,3 @@
-import asyncio
 import shutil
 import time
 from pathlib import Path
@@ -16,6 +15,12 @@
 from rationai.mlkit.lightning.loggers import MLFlowLogger
 from ratiopath.tiling.read_slide_tiles import read_slide_tiles
 from ray.data.expressions import col
+from tenacity import (
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 
 class EmbedTiles:
@@ -34,25 +39,22 @@ def __init__(self, model: str, concurrency: int) -> None:
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
 
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=1, max=4),
+        retry=retry_if_exception_type(httpx.HTTPError),
+        reraise=True,
+    )
+    async def _embed(self, tile: Any) -> list[float]:
+        return (
+            (await self.client.models.embed_image(self.model, tile)).reshape(-1).tolist()
+        )
+
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         self.in_flight += 1
         t0 = time.monotonic()
         try:
-            last_exc: Exception | None = None
-            for attempt in range(3):
-                try:
-                    embedding = (
-                        (await self.client.models.embed_image(self.model, row["tile"]))
-                        .reshape(-1)
-                        .tolist()
-                    )
-                    break
-                except Exception as exc:
-                    last_exc = exc
-                    if attempt < 2:
-                        await asyncio.sleep(2**attempt)
-            else:
-                raise RuntimeError("embed_image failed after 3 attempts") from last_exc
+            embedding = await self._embed(row["tile"])
         finally:
             del row["tile"]
         latency = time.monotonic() - t0
diff --git a/pyproject.toml b/pyproject.toml
index 491b9f0..a041e2c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,6 +29,7 @@ dependencies = [
     "einops>=0.8.0",
     "pyarrow>=19.0.1",
     "datasets>=4.0.0",
+    "tenacity>=9.0.0",
 ]
 
 [dependency-groups]
diff --git a/uv.lock b/uv.lock
index 7b83df0..5f30ccb 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2305,6 +2305,7 @@ dependencies = [
     { name = "ratiopath" },
     { name = "ray" },
     { name = "scikit-learn" },
+    { name = "tenacity" },
     { name = "tifffile" },
     { name = "timm" },
     { name = "torch" },
@@ -2338,6 +2339,7 @@ requires-dist = [
     { name = "ratiopath", specifier = ">=1.2.0" },
     { name = "ray", specifier = ">=2.51.1" },
     { name = "scikit-learn", specifier = ">=1.8.0" },
+    { name = "tenacity", specifier = ">=9.0.0" },
     { name = "tifffile", specifier = ">=2025.12.20" },
     { name = "timm", specifier = ">=1.0.0" },
     { name = "torch", specifier = ">=2.0.0" },

From 434aa435396ad66079614fe0607a379fa28a592c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:39:12 +0200
Subject: [PATCH 34/49] chore: drop per-actor throughput log from EmbedTiles

Ray Data progress bars surface rows/sec and in-flight counts, so the
periodic counter log and the latency/in_flight bookkeeping it required
are redundant.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 13 -------------
 1 file changed, 13 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 7bb73f6..9a9c1d6 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -32,9 +32,6 @@ def __init__(self, model: str, concurrency: int) -> None:
             ),
             timeout=200,
         )
-        self.count = 0
-        self.in_flight = 0
-        self.start = time.monotonic()
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
@@ -51,20 +48,10 @@ async def _embed(self, tile: Any) -> list[float]:
         )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
-        self.in_flight += 1
-        t0 = time.monotonic()
         try:
             embedding = await self._embed(row["tile"])
         finally:
             del row["tile"]
-        latency = time.monotonic() - t0
-        self.count += 1
-        self.in_flight -= 1
-        if self.count % 50 == 0:
-            elapsed = time.monotonic() - self.start
-            print(
-                f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)"
-            )
         row["embedding"] = embedding
         return row
 

From 287408f6499a5484f938d26a65db0efb9f7f942c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:41:26 +0200
Subject: [PATCH 35/49] chore: drop oversized memory reservation on
 read_parquet

Reading three projected columns doesn't need 8 GB; default scheduling
lets Ray pack more readers per node.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 9a9c1d6..9f93c81 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -44,7 +44,9 @@ def __init__(self, model: str, concurrency: int) -> None:
     )
     async def _embed(self, tile: Any) -> list[float]:
         return (
-            (await self.client.models.embed_image(self.model, tile)).reshape(-1).tolist()
+            (await self.client.models.embed_image(self.model, tile))
+            .reshape(-1)
+            .tolist()
         )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
@@ -84,7 +86,6 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         ds = ray.data.read_parquet(
             str(tiles_path),
             columns=["slide_id", "x", "y"],
-            ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
             lambda row, si: {**row, **si[row["slide_id"]]},

From 8130f78a3191133983c96fdddbf5cbd6a41bea83 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:45:19 +0200
Subject: [PATCH 36/49] fix: build tenacity retryer in __init__ to avoid pickle
 failure

The @retry decorator at class scope captures an AsyncRetrying instance
whose internal threading.local() makes EmbedTiles unpicklable, which
breaks Ray Data's actor serialization. Constructing the retryer in
__init__ moves that state onto the worker.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 9f93c81..ff4e5d7 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -16,7 +16,7 @@
 from ratiopath.tiling.read_slide_tiles import read_slide_tiles
 from ray.data.expressions import col
 from tenacity import (
-    retry,
+    AsyncRetrying,
     retry_if_exception_type,
     stop_after_attempt,
     wait_exponential,
@@ -32,23 +32,26 @@ def __init__(self, model: str, concurrency: int) -> None:
             ),
             timeout=200,
         )
+        self._retryer = AsyncRetrying(
+            stop=stop_after_attempt(3),
+            wait=wait_exponential(multiplier=1, min=1, max=4),
+            retry=retry_if_exception_type(httpx.HTTPError),
+            reraise=True,
+        )
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
 
-    @retry(
-        stop=stop_after_attempt(3),
-        wait=wait_exponential(multiplier=1, min=1, max=4),
-        retry=retry_if_exception_type(httpx.HTTPError),
-        reraise=True,
-    )
-    async def _embed(self, tile: Any) -> list[float]:
+    async def _embed_once(self, tile: Any) -> list[float]:
         return (
             (await self.client.models.embed_image(self.model, tile))
             .reshape(-1)
             .tolist()
         )
 
+    async def _embed(self, tile: Any) -> list[float]:
+        return await self._retryer(self._embed_once, tile)
+
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         try:
             embedding = await self._embed(row["tile"])

From aa37cf381de045d676af51386e87cd41bd096ba6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 22:55:51 +0200
Subject: [PATCH 37/49] Revert "chore: drop oversized memory reservation on
 read_parquet"

Without the 8 GB memory reservation on read_parquet, downstream stages
stall and no embeddings are produced. Restoring until we understand the
scheduling interaction.

This reverts commit 287408f.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index ff4e5d7..80ddabc 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -89,6 +89,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         ds = ray.data.read_parquet(
             str(tiles_path),
             columns=["slide_id", "x", "y"],
+            ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
             lambda row, si: {**row, **si[row["slide_id"]]},

From 27da14e2c3da3a017599ef6b1ee26ad97c7732bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 23:07:17 +0200
Subject: [PATCH 38/49] Revert "refactor: use tenacity for embed retries,
 narrow to httpx errors"

Tenacity-based retries broke embedding production under high actor
concurrency (likely shared AsyncRetrying state and narrowed exception
filter dropping retryable non-httpx errors). Restoring the manual
loop, which was last known to work.

This reverts commit 82163ac.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 39 +++++++++++++++----------------------
 pyproject.toml              |  1 -
 uv.lock                     |  2 --
 3 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 80ddabc..5b8aa78 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,3 +1,4 @@
+import asyncio
 import shutil
 import time
 from pathlib import Path
@@ -15,12 +16,6 @@
 from rationai.mlkit.lightning.loggers import MLFlowLogger
 from ratiopath.tiling.read_slide_tiles import read_slide_tiles
 from ray.data.expressions import col
-from tenacity import (
-    AsyncRetrying,
-    retry_if_exception_type,
-    stop_after_attempt,
-    wait_exponential,
-)
 
 
 class EmbedTiles:
@@ -32,29 +27,27 @@ def __init__(self, model: str, concurrency: int) -> None:
             ),
             timeout=200,
         )
-        self._retryer = AsyncRetrying(
-            stop=stop_after_attempt(3),
-            wait=wait_exponential(multiplier=1, min=1, max=4),
-            retry=retry_if_exception_type(httpx.HTTPError),
-            reraise=True,
-        )
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
 
-    async def _embed_once(self, tile: Any) -> list[float]:
-        return (
-            (await self.client.models.embed_image(self.model, tile))
-            .reshape(-1)
-            .tolist()
-        )
-
-    async def _embed(self, tile: Any) -> list[float]:
-        return await self._retryer(self._embed_once, tile)
-
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         try:
-            embedding = await self._embed(row["tile"])
+            last_exc: Exception | None = None
+            for attempt in range(3):
+                try:
+                    embedding = (
+                        (await self.client.models.embed_image(self.model, row["tile"]))
+                        .reshape(-1)
+                        .tolist()
+                    )
+                    break
+                except Exception as exc:
+                    last_exc = exc
+                    if attempt < 2:
+                        await asyncio.sleep(2**attempt)
+            else:
+                raise RuntimeError("embed_image failed after 3 attempts") from last_exc
         finally:
             del row["tile"]
         row["embedding"] = embedding
diff --git a/pyproject.toml b/pyproject.toml
index a041e2c..491b9f0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,6 @@ dependencies = [
     "einops>=0.8.0",
     "pyarrow>=19.0.1",
     "datasets>=4.0.0",
-    "tenacity>=9.0.0",
 ]
 
 [dependency-groups]
diff --git a/uv.lock b/uv.lock
index 5f30ccb..7b83df0 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2305,7 +2305,6 @@ dependencies = [
     { name = "ratiopath" },
     { name = "ray" },
     { name = "scikit-learn" },
-    { name = "tenacity" },
     { name = "tifffile" },
     { name = "timm" },
     { name = "torch" },
@@ -2339,7 +2338,6 @@ requires-dist = [
     { name = "ratiopath", specifier = ">=1.2.0" },
     { name = "ray", specifier = ">=2.51.1" },
     { name = "scikit-learn", specifier = ">=1.8.0" },
-    { name = "tenacity", specifier = ">=9.0.0" },
     { name = "tifffile", specifier = ">=2025.12.20" },
     { name = "timm", specifier = ">=1.0.0" },
     { name = "torch", specifier = ">=2.0.0" },

From ec7a63a55b11787dea14a1004093615d8ac29f7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 23:15:02 +0200
Subject: [PATCH 39/49] Revert preprocessing/embeddings.py to commit 1d99fcb

Restore the last-known-working version of EmbedTiles after several
review-driven refactors broke embedding production.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 5b8aa78..3a940d8 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -27,11 +27,20 @@ def __init__(self, model: str, concurrency: int) -> None:
             ),
             timeout=200,
         )
+        self.count = 0
+        self.in_flight = 0
+        self.first_row_logged = False
+        self.start = time.monotonic()
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
+        if not self.first_row_logged:
+            self.first_row_logged = True
+            print(f"[EmbedTiles] first row at t={time.monotonic() - self.start:.2f}s")
+        self.in_flight += 1
+        t0 = time.monotonic()
         try:
             last_exc: Exception | None = None
             for attempt in range(3):
@@ -50,6 +59,14 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
                 raise RuntimeError("embed_image failed after 3 attempts") from last_exc
         finally:
             del row["tile"]
+        latency = time.monotonic() - t0
+        self.count += 1
+        self.in_flight -= 1
+        if self.count % 50 == 0:
+            elapsed = time.monotonic() - self.start
+            print(
+                f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)"
+            )
         row["embedding"] = embedding
         return row
 

From 73061036ae63dd196531cbfb3c7ebc228b3ddf36 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Wed, 6 May 2026 23:21:18 +0200
Subject: [PATCH 40/49] chore: drop oversized memory reservation on
 read_parquet

Three-column projection doesn't justify 8 GB per task; let Ray's
default scheduling pack readers per node.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 3a940d8..51b9516 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -99,7 +99,6 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         ds = ray.data.read_parquet(
             str(tiles_path),
             columns=["slide_id", "x", "y"],
-            ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
             lambda row, si: {**row, **si[row["slide_id"]]},

From d1a22892b428983240588ac3641f8333848f664a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 7 May 2026 10:35:55 +0200
Subject: [PATCH 41/49] chore: drop first_row_logged debug print

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 51b9516..96c1402 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -29,16 +29,12 @@ def __init__(self, model: str, concurrency: int) -> None:
         )
         self.count = 0
         self.in_flight = 0
-        self.first_row_logged = False
         self.start = time.monotonic()
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
-        if not self.first_row_logged:
-            self.first_row_logged = True
-            print(f"[EmbedTiles] first row at t={time.monotonic() - self.start:.2f}s")
         self.in_flight += 1
         t0 = time.monotonic()
         try:

From ea437e51f4be0efdad0c0a460241f58d55c9c533 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 7 May 2026 10:44:18 +0200
Subject: [PATCH 42/49] refactor: use tenacity for embed retries, narrow to
 httpx errors

Build AsyncRetrying in __init__ so the actor stays picklable
(threading.local inside tenacity can't cross the wire if captured at
class scope).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 20 +++++++++-----------
 pyproject.toml              |  1 +
 uv.lock                     |  2 ++
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 96c1402..1115804 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,4 +1,3 @@
-import asyncio
 import shutil
 import time
 from pathlib import Path
@@ -16,6 +15,7 @@
 from rationai.mlkit.lightning.loggers import MLFlowLogger
 from ratiopath.tiling.read_slide_tiles import read_slide_tiles
 from ray.data.expressions import col
+from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 
 class EmbedTiles:
@@ -30,6 +30,12 @@ def __init__(self, model: str, concurrency: int) -> None:
         self.count = 0
         self.in_flight = 0
         self.start = time.monotonic()
+        self.retryer = AsyncRetrying(
+            stop=stop_after_attempt(3),
+            wait=wait_exponential(multiplier=1, min=1, max=4),
+            retry=retry_if_exception_type(httpx.HTTPError),
+            reraise=True,
+        )
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
@@ -38,21 +44,13 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         self.in_flight += 1
         t0 = time.monotonic()
         try:
-            last_exc: Exception | None = None
-            for attempt in range(3):
-                try:
+            async for attempt in self.retryer:
+                with attempt:
                     embedding = (
                         (await self.client.models.embed_image(self.model, row["tile"]))
                         .reshape(-1)
                         .tolist()
                     )
-                    break
-                except Exception as exc:
-                    last_exc = exc
-                    if attempt < 2:
-                        await asyncio.sleep(2**attempt)
-            else:
-                raise RuntimeError("embed_image failed after 3 attempts") from last_exc
         finally:
             del row["tile"]
         latency = time.monotonic() - t0
diff --git a/pyproject.toml b/pyproject.toml
index 491b9f0..a041e2c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,6 +29,7 @@ dependencies = [
     "einops>=0.8.0",
     "pyarrow>=19.0.1",
     "datasets>=4.0.0",
+    "tenacity>=9.0.0",
 ]
 
 [dependency-groups]
diff --git a/uv.lock b/uv.lock
index 7b83df0..5f30ccb 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2305,6 +2305,7 @@ dependencies = [
     { name = "ratiopath" },
     { name = "ray" },
     { name = "scikit-learn" },
+    { name = "tenacity" },
     { name = "tifffile" },
     { name = "timm" },
     { name = "torch" },
@@ -2338,6 +2339,7 @@ requires-dist = [
     { name = "ratiopath", specifier = ">=1.2.0" },
     { name = "ray", specifier = ">=2.51.1" },
     { name = "scikit-learn", specifier = ">=1.8.0" },
+    { name = "tenacity", specifier = ">=9.0.0" },
     { name = "tifffile", specifier = ">=2025.12.20" },
     { name = "timm", specifier = ">=1.0.0" },
     { name = "torch", specifier = ">=2.0.0" },

From 3a8bacd1bd2b478033bdd87e0215c5642a5b73f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 7 May 2026 11:08:27 +0200
Subject: [PATCH 43/49] Revert "refactor: use tenacity for embed retries,
 narrow to httpx errors"

This reverts commit ea437e51f4be0efdad0c0a460241f58d55c9c533.
---
 preprocessing/embeddings.py | 20 +++++++++++---------
 pyproject.toml              |  1 -
 uv.lock                     |  2 --
 3 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 1115804..96c1402 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -1,3 +1,4 @@
+import asyncio
 import shutil
 import time
 from pathlib import Path
@@ -15,7 +16,6 @@
 from rationai.mlkit.lightning.loggers import MLFlowLogger
 from ratiopath.tiling.read_slide_tiles import read_slide_tiles
 from ray.data.expressions import col
-from tenacity import AsyncRetrying, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 
 class EmbedTiles:
@@ -30,12 +30,6 @@ def __init__(self, model: str, concurrency: int) -> None:
         self.count = 0
         self.in_flight = 0
         self.start = time.monotonic()
-        self.retryer = AsyncRetrying(
-            stop=stop_after_attempt(3),
-            wait=wait_exponential(multiplier=1, min=1, max=4),
-            retry=retry_if_exception_type(httpx.HTTPError),
-            reraise=True,
-        )
         print(
             f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
         )
@@ -44,13 +38,21 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         self.in_flight += 1
         t0 = time.monotonic()
         try:
-            async for attempt in self.retryer:
-                with attempt:
+            last_exc: Exception | None = None
+            for attempt in range(3):
+                try:
                     embedding = (
                         (await self.client.models.embed_image(self.model, row["tile"]))
                         .reshape(-1)
                         .tolist()
                     )
+                    break
+                except Exception as exc:
+                    last_exc = exc
+                    if attempt < 2:
+                        await asyncio.sleep(2**attempt)
+            else:
+                raise RuntimeError("embed_image failed after 3 attempts") from last_exc
         finally:
             del row["tile"]
         latency = time.monotonic() - t0
diff --git a/pyproject.toml b/pyproject.toml
index a041e2c..491b9f0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,6 @@ dependencies = [
     "einops>=0.8.0",
     "pyarrow>=19.0.1",
     "datasets>=4.0.0",
-    "tenacity>=9.0.0",
 ]
 
 [dependency-groups]
diff --git a/uv.lock b/uv.lock
index 5f30ccb..7b83df0 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2305,7 +2305,6 @@ dependencies = [
     { name = "ratiopath" },
     { name = "ray" },
     { name = "scikit-learn" },
-    { name = "tenacity" },
     { name = "tifffile" },
     { name = "timm" },
     { name = "torch" },
@@ -2339,7 +2338,6 @@ requires-dist = [
     { name = "ratiopath", specifier = ">=1.2.0" },
     { name = "ray", specifier = ">=2.51.1" },
     { name = "scikit-learn", specifier = ">=1.8.0" },
-    { name = "tenacity", specifier = ">=9.0.0" },
     { name = "tifffile", specifier = ">=2025.12.20" },
     { name = "timm", specifier = ">=1.0.0" },
     { name = "torch", specifier = ">=2.0.0" },

From 1c2615f2965926e51cd345b2cc9d446803f2d713 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 7 May 2026 11:36:30 +0200
Subject: [PATCH 44/49] chore: drop per-actor throughput log, ray reports it

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 96c1402..058605e 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -27,16 +27,8 @@ def __init__(self, model: str, concurrency: int) -> None:
             ),
             timeout=200,
         )
-        self.count = 0
-        self.in_flight = 0
-        self.start = time.monotonic()
-        print(
-            f"[EmbedTiles] actor initialized, model={model}, concurrency={concurrency}"
-        )
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
-        self.in_flight += 1
-        t0 = time.monotonic()
         try:
             last_exc: Exception | None = None
             for attempt in range(3):
@@ -55,14 +47,6 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
                 raise RuntimeError("embed_image failed after 3 attempts") from last_exc
         finally:
             del row["tile"]
-        latency = time.monotonic() - t0
-        self.count += 1
-        self.in_flight -= 1
-        if self.count % 50 == 0:
-            elapsed = time.monotonic() - self.start
-            print(
-                f"[EmbedTiles] {self.count} done in {elapsed:.1f}s ({self.count / elapsed:.1f}/s, in_flight={self.in_flight}, latency={latency * 1000:.0f}ms)"
-            )
         row["embedding"] = embedding
         return row
 

From ad3a7c56f6a1ec8f313478e1d9b14e6894ce6337 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 7 May 2026 12:00:42 +0200
Subject: [PATCH 45/49] chore: switch ray data progress to tqdm for non-tty
 envs

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index 058605e..dd4aa1a 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -126,12 +126,12 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
 
 if __name__ == "__main__":
     ctx = ray.data.DataContext.get_current()
-    ctx.enable_rich_progress_bars = True
-    ctx.use_ray_tqdm = False
+    ctx.enable_rich_progress_bars = False
+    ctx.use_ray_tqdm = True
     ctx.target_max_block_size = 64 * 1024 * 1024
 
     with ray.init(
         runtime_env={"excludes": [".git", ".venv"]},
         object_store_memory=16 * 1024**3,
     ):
-        main()
+        main()
\ No newline at end of file

From 7a07f41abbe644bd3f87ac92071037fb4eb807c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Thu, 7 May 2026 12:10:22 +0200
Subject: [PATCH 46/49] fix: format

---
 preprocessing/embeddings.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index dd4aa1a..ca24437 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -134,4 +134,4 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         runtime_env={"excludes": [".git", ".venv"]},
         object_store_memory=16 * 1024**3,
     ):
-        main()
\ No newline at end of file
+        main()

From ac2b4b3d583c61dd357381958d0f7936cab64c4a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Fri, 8 May 2026 15:21:02 +0200
Subject: [PATCH 47/49] chore: remove duplicate mlflow run id

---
 configs/data/dataset.yaml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/configs/data/dataset.yaml b/configs/data/dataset.yaml
index 0fb10f9..e13fec8 100644
--- a/configs/data/dataset.yaml
+++ b/configs/data/dataset.yaml
@@ -16,7 +16,6 @@ dataset:
     filter_tiles_run_id: "4e8f5d3c82124ea5a8f871a42d3ed9ba"
     tissue_masks_run_id: "52bc0924f8624b259819c480c7cf213f"
     tissue_stats_run_id: "16ae2d003d88471b924e5f332415232a"
-    filter_tiles_run_id: "4e8f5d3c82124ea5a8f871a42d3ed9ba"
 
   exclusions:
     bad_slides:

From 5eae41c9571dd1dca60d4e3f5ab7905d80b0368c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Fri, 8 May 2026 16:48:14 +0200
Subject: [PATCH 48/49] fix: harden embed retries against pool timeouts

- Narrow retry to httpx.HTTPError (6 attempts, exp backoff capped 60s)
- Pin actor pool to 4 (min_size=min_size=4) so backpressure can't shrink it to 1

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index ca24437..b7b2486 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -30,8 +30,8 @@ def __init__(self, model: str, concurrency: int) -> None:
 
     async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
         try:
-            last_exc: Exception | None = None
-            for attempt in range(3):
+            last_exc: httpx.HTTPError | None = None
+            for attempt in range(6):
                 try:
                     embedding = (
                         (await self.client.models.embed_image(self.model, row["tile"]))
@@ -39,12 +39,12 @@ async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
                         .tolist()
                     )
                     break
-                except Exception as exc:
+                except httpx.HTTPError as exc:
                     last_exc = exc
-                    if attempt < 2:
-                        await asyncio.sleep(2**attempt)
+                    if attempt < 5:
+                        await asyncio.sleep(min(60, 2**attempt))
             else:
-                raise RuntimeError("embed_image failed after 3 attempts") from last_exc
+                raise RuntimeError("embed_image failed after 6 attempts") from last_exc
         finally:
             del row["tile"]
         row["embedding"] = embedding
@@ -102,6 +102,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
             EmbedTiles,  # pyright: ignore[reportArgumentType]
             fn_constructor_args=(config.model, config.concurrency),
             compute=ray.data.ActorPoolStrategy(
+                min_size=4,
                 max_size=4,
                 max_tasks_in_flight_per_actor=max(1, config.concurrency // 4),
             ),

From e4e83a2dd812147a07dcc7453946a37a995aaa16 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vojt=C4=9Bch=20C=C3=ADfka?= <550433@mail.muni.cz>
Date: Fri, 8 May 2026 16:59:25 +0200
Subject: [PATCH 49/49] fix: re-add 8GB memory reservation on read_parquet

Throttles upstream read parallelism so the object store doesn't fill
up and trigger downstream backpressure. Without it, Ray scales the
EmbedTiles actor pool from 4 down to 1 under load, collapsing
throughput (12h to 51% vs prior ~12h end-to-end).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 preprocessing/embeddings.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/preprocessing/embeddings.py b/preprocessing/embeddings.py
index b7b2486..de98143 100644
--- a/preprocessing/embeddings.py
+++ b/preprocessing/embeddings.py
@@ -79,6 +79,7 @@ def main(config: DictConfig, logger: MLFlowLogger) -> None:
         ds = ray.data.read_parquet(
             str(tiles_path),
             columns=["slide_id", "x", "y"],
+            ray_remote_args={"memory": 8 * 1024**3},
             override_num_blocks=num_blocks,
         ).map(
             lambda row, si: {**row, **si[row["slide_id"]]},