feat: tile embeddings pipeline with lazy parquet loading

configs/experiment/preprocessing/embeddings_virchow2_05mpp.yaml

@@ -0,0 +1,7 @@
+# @package _global_
+
+defaults:
+  - /data: dataset
+  - _self_
+
+model: virchow2

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

configs/preprocessing/embeddings.yaml

@@ -0,0 +1,15 @@
+# @package _global_
+
+model: ???
+output_dir: ${project_path}/embeddings
+concurrency: 512
+block_size: 2048
+rows_per_file: 5000
+
+metadata:
+  run_name: "Embeddings: ${model}"
+  description: Tile embeddings using model ${model}
+  hyperparams:
+    model: ${model}
+    concurrency: ${concurrency}
+    block_size: ${block_size}

configs/preprocessing/tile_masks.yaml

@@ -1,7 +1,5 @@
 # @package _global_
 
-slides_artifact_path: ???
-tiles_artifact_path: ???
 tile_percentage_cols: ???
 
 max_concurrent: 1

preprocessing/embeddings.py

@@ -0,0 +1,139 @@
+import asyncio
+import shutil
+import time
+from pathlib import Path
+from typing import Any
+
+import httpx
+import hydra
+import mlflow.artifacts
+import pandas as pd
+import pyarrow.dataset as pads
+import ray
+from omegaconf import DictConfig
+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
+from ray.data.expressions import col
+
+
+class EmbedTiles:
+    def __init__(self, model: str, concurrency: int) -> None:
+        self.model = model
+        self.client = AsyncClient(
+            limits=httpx.Limits(
+                max_connections=concurrency, max_keepalive_connections=concurrency
+            ),
+            timeout=200,
+        )
+
+    async def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
+        try:
+            last_exc: httpx.HTTPError | None = None
+            for attempt in range(6):
+                try:
+                    embedding = (
+                        (await self.client.models.embed_image(self.model, row["tile"]))
+                        .reshape(-1)
+                        .tolist()
+                    )
+                    break
+                except httpx.HTTPError as exc:
+                    last_exc = exc
+                    if attempt < 5:
+                        await asyncio.sleep(min(60, 2**attempt))
+            else:
+                raise RuntimeError("embed_image failed after 6 attempts") from last_exc
+        finally:
+            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:
+    for name in ["train", "test"]:
+        split_folder = Path(
+            mlflow.artifacts.download_artifacts(
+                run_id=config.dataset.mlflow_artifacts.tiling_run_id,
+                artifact_path=f"{name}_split",
+            )
+        )
+        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")
+
+        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)
+
+        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"]]},
+            fn_kwargs={"si": slide_info},
+        )
+        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,  # 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),
+            ),
+            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"
+        if tiles_parquet_dir.exists():
+            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))
+
+
+if __name__ == "__main__":
+    ctx = ray.data.DataContext.get_current()
+    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()

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()

pyproject.toml

@@ -25,6 +25,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",
     "matplotlib>=3.10.7",
     "pyarrow>=19.0.1",
     "datasets>=4.0.0",
@@ -40,7 +44,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]

scripts/submit_embeddings.py

@@ -0,0 +1,18 @@
+from kube_jobs import storage, submit_job
+
+
+submit_job(
+    job_name="tissue-classification-embeddings",
+    username=...,
+    public=False,
+    cpu=8,
+    memory="64Gi",
+    shm="24Gi",
+    script=[
+        "git clone https://github.com/RationAI/tissue-classification.git workdir",
+        "cd workdir",
+        "uv sync",
+        "uv run -m preprocessing.embeddings +experiment=...",
+    ],
+    storage=[storage.secure.DATA, storage.secure.PROJECTS],
+)