From cf8bcfd40a21fd6f6a975c11bd824c608ab8d367 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9Csambasy=E2=80=9D?=
 <“ibtisamssaleh@outlook.com”>
Date: Sat, 19 Jul 2025 00:17:18 +0300
Subject: [PATCH 01/42] added backend mode to file reading (1st function)

---
 cellarium/ml/data/fileio.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index c1a3229e..6af1e54e 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -1,6 +1,6 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
-
+from typing import Literal
 import re
 import shutil
 import tempfile
@@ -12,7 +12,7 @@
 url_schemes = ("http:", "https:", "ftp:")
 
 
-def read_h5ad_gcs(filename: str, storage_client: Client | None = None) -> AnnData:
+def read_h5ad_gcs(filename: str, storage_client: Client | None = None, backed: Literal['r', 'r+'] | bool | None = None) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the Google Cloud Storage.
 
@@ -22,6 +22,9 @@ def read_h5ad_gcs(filename: str, storage_client: Client | None = None) -> AnnDat
 
     Args:
         filename: Path to the data file in Cloud Storage.
+        backed: If 'r', load in backed mode instead of fully loading into memory.
+               If 'r+', load in backed mode with write access (only X can be modified).
+               If True, equivalent to 'r'. Default is None (load into memory).
     """
     if not filename.startswith("gs:"):
         raise ValueError("The filename must start with 'gs:' protocol name.")
@@ -36,7 +39,7 @@ def read_h5ad_gcs(filename: str, storage_client: Client | None = None) -> AnnDat
     blob = bucket.blob(blob_name)
 
     with blob.open("rb") as f:
-        return read_h5ad(f)
+        return read_h5ad(f,backed=backed)
 
 
 def read_h5ad_url(filename: str) -> AnnData:

From febd1a9298b07e34a1c467b5fb3088b1ccdefdfa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9Csambasy=E2=80=9D?=
 <“ibtisamssaleh@outlook.com”>
Date: Sat, 19 Jul 2025 00:21:01 +0300
Subject: [PATCH 02/42] bacekdn mode for the 2nd function

---
 cellarium/ml/data/fileio.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 6af1e54e..baf6b9a6 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -42,7 +42,7 @@ def read_h5ad_gcs(filename: str, storage_client: Client | None = None, backed: L
         return read_h5ad(f,backed=backed)
 
 
-def read_h5ad_url(filename: str) -> AnnData:
+def read_h5ad_url(filename: str,backed: Literal['r', 'r+'] | bool | None = None) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the URL.
 
@@ -51,16 +51,23 @@ def read_h5ad_url(filename: str) -> AnnData:
         >>> adata = read_h5ad_url(
         ...     "https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad"
         ... )
+        >>> adata = read_h5ad_url(
+        ...     "https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
+        ...     backed='r'
+        ... )
 
     Args:
         filename: URL of the data file.
+         backed: If 'r', load in backed mode instead of fully loading into memory.
+               If 'r+', load in backed mode with write access (only X can be modified).
+               If True, equivalent to 'r'. Default is None (load into memory).
     """
     if not any(filename.startswith(scheme) for scheme in url_schemes):
         raise ValueError("The filename must start with 'http:', 'https:', or 'ftp:' protocol name.")
     with urllib.request.urlopen(filename) as response:
         with tempfile.TemporaryFile() as tmp_file:
             shutil.copyfileobj(response, tmp_file)
-            return read_h5ad(tmp_file)
+            return read_h5ad(tmp_file,backed=backed)
 
 
 def read_h5ad_local(filename: str) -> AnnData:

From c28abc5bafe0c8bf7695d762a25b66b92d543aa1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9Csambasy=E2=80=9D?=
 <“ibtisamssaleh@outlook.com”>
Date: Sat, 19 Jul 2025 00:23:22 +0300
Subject: [PATCH 03/42] backend mode for the third function

---
 cellarium/ml/data/fileio.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index baf6b9a6..222e981d 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -70,17 +70,21 @@ def read_h5ad_url(filename: str,backed: Literal['r', 'r+'] | bool | None = None)
             return read_h5ad(tmp_file,backed=backed)
 
 
-def read_h5ad_local(filename: str) -> AnnData:
+def read_h5ad_local(filename: str,str,backed: Literal['r', 'r+'] | bool | None = None) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the local disk.
 
     Args:
         filename: Path to the local data file.
+        backed: If 'r', load in backed mode instead of fully loading into memory.
+               If 'r+', load in backed mode with write access (only X can be modified).
+               If True, equivalent to 'r'. Default is None (load into memory).
+
     """
     if not filename.startswith("file:"):
         raise ValueError("The filename must start with 'file:' protocol name.")
     filename = re.sub(r"^file://?", "", filename)
-    return read_h5ad(filename)
+    return read_h5ad(filename,backed=backed)
 
 
 def read_h5ad_file(filename: str, **kwargs) -> AnnData:

From 0957c4c08ff5c94e07008f45a61f9211d11307cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9Csambasy=E2=80=9D?=
 <“ibtisamssaleh@outlook.com”>
Date: Sat, 19 Jul 2025 00:24:28 +0300
Subject: [PATCH 04/42] backend mode for 4th function

---
 cellarium/ml/data/fileio.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 222e981d..fc09f8d2 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -87,12 +87,16 @@ def read_h5ad_local(filename: str,str,backed: Literal['r', 'r+'] | bool | None =
     return read_h5ad(filename,backed=backed)
 
 
-def read_h5ad_file(filename: str, **kwargs) -> AnnData:
+def read_h5ad_file(filename: str, backed: Literal['r', 'r+'] | bool | None = None, **kwargs) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from a filename.
 
     Args:
         filename: Path to the data file.
+        backed: If 'r', load in backed mode instead of fully loading into memory.
+               If 'r+', load in backed mode with write access (only X can be modified).
+               If True, equivalent to 'r'. Default is None (load into memory).
+
     """
     if filename.startswith("gs:"):
         return read_h5ad_gcs(filename, **kwargs)
@@ -103,4 +107,4 @@ def read_h5ad_file(filename: str, **kwargs) -> AnnData:
     if any(filename.startswith(scheme) for scheme in url_schemes):
         return read_h5ad_url(filename)
 
-    return read_h5ad(filename)
+    return read_h5ad(filename,backed=backed)

From ea601cbb93bde7785fed65e012b633c2ccebe820 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Mon, 28 Jul 2025 13:19:44 -0400
Subject: [PATCH 05/42] make format, make typecheck

---
 cellarium/ml/data/fileio.py | 23 ++++++++++++++---------
 1 file changed, 14 insertions(+), 9 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index fc09f8d2..57eeeaaf 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -1,10 +1,11 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
-from typing import Literal
+
 import re
 import shutil
 import tempfile
 import urllib.request
+from typing import Literal
 
 from anndata import AnnData, read_h5ad
 from google.cloud.storage import Client
@@ -12,7 +13,11 @@
 url_schemes = ("http:", "https:", "ftp:")
 
 
-def read_h5ad_gcs(filename: str, storage_client: Client | None = None, backed: Literal['r', 'r+'] | bool | None = None) -> AnnData:
+def read_h5ad_gcs(
+    filename: str,
+    storage_client: Client | None = None,
+    backed: Literal["r", "r+"] | bool | None = None,
+) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the Google Cloud Storage.
 
@@ -39,10 +44,10 @@ def read_h5ad_gcs(filename: str, storage_client: Client | None = None, backed: L
     blob = bucket.blob(blob_name)
 
     with blob.open("rb") as f:
-        return read_h5ad(f,backed=backed)
+        return read_h5ad(f, backed=backed)
 
 
-def read_h5ad_url(filename: str,backed: Literal['r', 'r+'] | bool | None = None) -> AnnData:
+def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = None) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the URL.
 
@@ -67,10 +72,10 @@ def read_h5ad_url(filename: str,backed: Literal['r', 'r+'] | bool | None = None)
     with urllib.request.urlopen(filename) as response:
         with tempfile.TemporaryFile() as tmp_file:
             shutil.copyfileobj(response, tmp_file)
-            return read_h5ad(tmp_file,backed=backed)
+            return read_h5ad(tmp_file, backed=backed)
 
 
-def read_h5ad_local(filename: str,str,backed: Literal['r', 'r+'] | bool | None = None) -> AnnData:
+def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = None) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the local disk.
 
@@ -84,10 +89,10 @@ def read_h5ad_local(filename: str,str,backed: Literal['r', 'r+'] | bool | None =
     if not filename.startswith("file:"):
         raise ValueError("The filename must start with 'file:' protocol name.")
     filename = re.sub(r"^file://?", "", filename)
-    return read_h5ad(filename,backed=backed)
+    return read_h5ad(filename, backed=backed)
 
 
-def read_h5ad_file(filename: str, backed: Literal['r', 'r+'] | bool | None = None, **kwargs) -> AnnData:
+def read_h5ad_file(filename: str, backed: Literal["r", "r+"] | bool | None = None, **kwargs) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from a filename.
 
@@ -107,4 +112,4 @@ def read_h5ad_file(filename: str, backed: Literal['r', 'r+'] | bool | None = Non
     if any(filename.startswith(scheme) for scheme in url_schemes):
         return read_h5ad_url(filename)
 
-    return read_h5ad(filename,backed=backed)
+    return read_h5ad(filename, backed=backed)

From a91efd940ea08a8a2bc038789bdea89868ca2e31 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Mon, 28 Jul 2025 16:16:54 -0400
Subject: [PATCH 06/42] change default to backed='r'

---
 cellarium/ml/data/fileio.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 57eeeaaf..5a4c9c4d 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -11,12 +11,13 @@
 from google.cloud.storage import Client
 
 url_schemes = ("http:", "https:", "ftp:")
+backed_mode_default: Literal["r", "r+"] | bool | None = "r"
 
 
 def read_h5ad_gcs(
     filename: str,
     storage_client: Client | None = None,
-    backed: Literal["r", "r+"] | bool | None = None,
+    backed: Literal["r", "r+"] | bool | None = backed_mode_default,
 ) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the Google Cloud Storage.
@@ -47,7 +48,7 @@ def read_h5ad_gcs(
         return read_h5ad(f, backed=backed)
 
 
-def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = None) -> AnnData:
+def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the URL.
 
@@ -75,7 +76,7 @@ def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = None
             return read_h5ad(tmp_file, backed=backed)
 
 
-def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = None) -> AnnData:
+def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the local disk.
 
@@ -92,7 +93,7 @@ def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = No
     return read_h5ad(filename, backed=backed)
 
 
-def read_h5ad_file(filename: str, backed: Literal["r", "r+"] | bool | None = None, **kwargs) -> AnnData:
+def read_h5ad_file(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default, **kwargs) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from a filename.
 

From c20823112b98730fa0c2c89239a377ce3b101abf Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Mon, 28 Jul 2025 16:41:25 -0400
Subject: [PATCH 07/42] fix for URL and GCS readers in backed mode

---
 cellarium/ml/data/fileio.py | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 5a4c9c4d..a1111849 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -1,6 +1,7 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import os
 import re
 import shutil
 import tempfile
@@ -44,8 +45,17 @@ def read_h5ad_gcs(
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(blob_name)
 
-    with blob.open("rb") as f:
-        return read_h5ad(f, backed=backed)
+    # write to a named temporary file
+    with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
+        temp_path = tmp_file.name
+        blob.download_to_file(tmp_file)
+        try:
+            return read_h5ad(temp_path, backed=backed)
+        finally:
+            try:
+                os.unlink(temp_path)  # clean up the temp file
+            except OSError:
+                pass  # if there's an error during cleanup, continue
 
 
 def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default) -> AnnData:
@@ -70,10 +80,19 @@ def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = back
     """
     if not any(filename.startswith(scheme) for scheme in url_schemes):
         raise ValueError("The filename must start with 'http:', 'https:', or 'ftp:' protocol name.")
-    with urllib.request.urlopen(filename) as response:
-        with tempfile.TemporaryFile() as tmp_file:
+
+    # write to a named temporary file
+    with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
+        temp_path = tmp_file.name
+        with urllib.request.urlopen(filename) as response:
             shutil.copyfileobj(response, tmp_file)
-            return read_h5ad(tmp_file, backed=backed)
+        try:
+            return read_h5ad(temp_path, backed=backed)
+        finally:
+            try:
+                os.unlink(temp_path)  # clean up the temp file
+            except OSError:
+                pass  # if there's an error during cleanup, continue
 
 
 def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default) -> AnnData:

From 905ea23f7c1c23704af5823e46d3d3188bed90ef Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 14:00:16 -0400
Subject: [PATCH 08/42] testing anndata is in memory for comparison

---
 tests/dataloader/test_distributed_anndata.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/dataloader/test_distributed_anndata.py b/tests/dataloader/test_distributed_anndata.py
index dad39235..31c4769f 100644
--- a/tests/dataloader/test_distributed_anndata.py
+++ b/tests/dataloader/test_distributed_anndata.py
@@ -70,7 +70,7 @@ def adatas_path(tmp_path: Path):
 @pytest.fixture
 def adt(adatas_path: Path):
     # single anndata
-    adt = read_h5ad_file(str(os.path.join(adatas_path, "adata.h5ad")))
+    adt = read_h5ad_file(str(os.path.join(adatas_path, "adata.h5ad")), backed=None)
     return adt
 
 
@@ -132,6 +132,7 @@ def test_init_shard_size(adatas_path: Path, num_shards: int, last_shard_size: in
     ids=["one adata", "two adatas", "sorted two adatas", "unsorted three adatas"],
 )
 @pytest.mark.parametrize("vidx", [slice(0, 2), [3, 1, 0]])
+# @pytest.mark.parametrize("adt", [None, "r"], ids=["memory", "backed"], indirect=True)
 def test_indexing(
     adt: AnnData,
     dat: DistributedAnnDataCollection,

From ca09d140387217b23e0a81a7d4a32297368379ea Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 14:04:15 -0400
Subject: [PATCH 09/42] allow backed mode CellariumAnnDataDataModule load to
 pass

---
 tests/dataloader/test_datamodule.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index 8923872f..2c19a9f8 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -178,11 +178,12 @@ def _check_transform_lists_match(iterable_modules1, iterable_modules2, message):
 )
 @pytest.mark.parametrize("batch_size", [50, None])
 def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) -> None:
+    dadc = DistributedAnnDataCollection(
+        filenames="https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
+        shard_size=100,
+    )
     datamodule = CellariumAnnDataDataModule(
-        DistributedAnnDataCollection(
-            filenames="https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
-            shard_size=100,
-        ),
+        dadc=dadc,
         batch_size=100,
         batch_keys={
             "x_ng": AnnDataField(attr="X", convert_fn=densify),
@@ -193,15 +194,14 @@ def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) ->
     trainer.fit(module, datamodule)
 
     ckpt_path = str(tmp_path / "lightning_logs/version_0/checkpoints/epoch=0-step=1.ckpt")
-    adata = datamodule.dadc.adatas[0].adata
-    kwargs = {"dadc": adata}
+    kwargs = {"dadc": dadc}
     if batch_size is not None:
         kwargs["batch_size"] = batch_size
     loaded_datamodule = CellariumAnnDataDataModule.load_from_checkpoint(ckpt_path, **kwargs)
 
     assert loaded_datamodule.batch_keys == datamodule.batch_keys
     assert loaded_datamodule.batch_size == batch_size or datamodule.batch_size
-    assert loaded_datamodule.dadc is adata
+    assert loaded_datamodule.dadc is dadc
 
 
 @pytest.mark.parametrize(

From de242043bf04e6cee7d5b96cc68fea8724c1d0cb Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 16:40:38 -0400
Subject: [PATCH 10/42] mypy

---
 tests/dataloader/test_datamodule.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index 2c19a9f8..fbbd4643 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from pathlib import Path
+from typing import Any
 
 import lightning.pytorch as pl
 import pytest
@@ -194,7 +195,7 @@ def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) ->
     trainer.fit(module, datamodule)
 
     ckpt_path = str(tmp_path / "lightning_logs/version_0/checkpoints/epoch=0-step=1.ckpt")
-    kwargs = {"dadc": dadc}
+    kwargs: dict[str, Any] = {"dadc": dadc}
     if batch_size is not None:
         kwargs["batch_size"] = batch_size
     loaded_datamodule = CellariumAnnDataDataModule.load_from_checkpoint(ckpt_path, **kwargs)

From d2ee9f8580d0e98b97a2760ac3ce4c173a99d822 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 18:01:24 -0400
Subject: [PATCH 11/42] typing for backed mode

---
 cellarium/ml/data/fileio.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index a1111849..27de26d2 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -12,7 +12,8 @@
 from google.cloud.storage import Client
 
 url_schemes = ("http:", "https:", "ftp:")
-backed_mode_default: Literal["r", "r+"] | bool | None = "r"
+backed_mode_type = Literal["r"] | bool | None
+backed_mode_default: backed_mode_type = "r"
 
 
 def read_h5ad_gcs(

From a634b92f34581533fe297c5c483b171eef81ab61 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 18:01:53 -0400
Subject: [PATCH 12/42] parameterize backed mode in LazyAnnData and
 DistributedAnnDataCollection

---
 cellarium/ml/data/distributed_anndata.py     | 42 +++++++++++++++-----
 tests/dataloader/test_distributed_anndata.py | 16 ++++++--
 2 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/cellarium/ml/data/distributed_anndata.py b/cellarium/ml/data/distributed_anndata.py
index 60eae53b..4cf92b4e 100644
--- a/cellarium/ml/data/distributed_anndata.py
+++ b/cellarium/ml/data/distributed_anndata.py
@@ -18,7 +18,7 @@
 from boltons.cacheutils import LRU
 from braceexpand import braceexpand
 
-from cellarium.ml.data.fileio import read_h5ad_file
+from cellarium.ml.data.fileio import backed_mode_default, backed_mode_type, read_h5ad_file
 from cellarium.ml.data.schema import AnnDataSchema
 
 
@@ -28,6 +28,8 @@ class getattr_mode:
 
 _GETATTR_MODE = getattr_mode()
 
+allowed_backed_modes = [None, True, False, "r"]
+
 
 @contextmanager
 def lazy_getattr():
@@ -155,6 +157,12 @@ class DistributedAnnDataCollection(AnnCollection):
         obs_columns_to_validate:
             Subset of columns to validate in the :attr:`obs` attribute.
             If ``None``, all columns are validated.
+        backed:
+            Optional backing mode for the h5ad files. ``'r'`` will leave count matrices
+            on disk until specific cell indices are queried, enabling the use of very large
+            h5ad files, while ``None`` will load entire count matrices from individual h5ad files
+            into cached memory as needed: a strategy that necessitates smaller chunked h5ad files.
+            See :func:`anndata.read_h5ad` for details on backing modes.
     """
 
     def __init__(
@@ -171,6 +179,7 @@ def __init__(
         convert: ConvertType | None = None,
         indices_strict: bool = True,
         obs_columns_to_validate: Sequence[str] | None = None,
+        backed: backed_mode_type = backed_mode_default,
     ):
         self.filenames = list(braceexpand(filenames) if isinstance(filenames, str) else filenames)
         if (shard_size is None) and (last_shard_size is not None):
@@ -192,8 +201,11 @@ def __init__(
         self.cache = LRU(max_cache_size)
         self.max_cache_size = max_cache_size
         self.cache_size_strictly_enforced = cache_size_strictly_enforced
+        if backed not in allowed_backed_modes:
+            raise ValueError(f"Invalid backed mode: {backed}. Choose from {allowed_backed_modes}")
+        self.backed = backed
         # schema
-        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0])
+        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0], backed=backed)
         if len(adata0) != limits[0]:
             raise ValueError(
                 f"The number of cells in the first anndata file ({len(adata0)}) "
@@ -203,7 +215,7 @@ def __init__(
         self.schema = AnnDataSchema(adata0, obs_columns_to_validate)
         # lazy anndatas
         lazy_adatas = [
-            LazyAnnData(filename, (start, end), self.schema, self.cache)
+            LazyAnnData(filename, (start, end), self.schema, self.cache, backed=backed)
             for start, end, filename in zip([0] + limits, limits, self.filenames)
         ]
         # use filenames as default keys
@@ -298,10 +310,10 @@ def __getstate__(self):
     def __setstate__(self, state):
         self.__dict__.update(state)
         self.cache = LRU(self.max_cache_size)
-        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0])
+        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0], backed=self.backed)
         self.schema = AnnDataSchema(adata0, self.obs_columns_to_validate)
         self.adatas = [
-            LazyAnnData(filename, (start, end), self.schema, self.cache)
+            LazyAnnData(filename, (start, end), self.schema, self.cache, backed=self.backed)
             for start, end, filename in zip([0] + self.limits, self.limits, self.filenames)
         ]
         self.obs_names = pd.Index([f"cell_{i}" for i in range(self.limits[-1])])
@@ -323,6 +335,11 @@ class LazyAnnData:
             Schema used as a reference for lazy attributes.
         cache:
             Shared LRU cache storing buffered anndatas.
+        backed:
+            Optional backing mode for the anndata. ``'r'`` will leave count matrix
+            on disk, while ``None`` will load count matrix in memory (when the anndata is
+            cached by calling the `.adata` property).
+            See :func:`anndata.read_h5ad` for details on backing modes.
     """
 
     _lazy_attrs = ["obs", "obsm", "layers", "var", "varm", "varp", "var_names"]
@@ -343,10 +360,14 @@ def __init__(
         limits: tuple[int, int],
         schema: AnnDataSchema,
         cache: LRU | None = None,
+        backed: backed_mode_type = backed_mode_default,
     ):
         self.filename = filename
         self.limits = limits
         self.schema = schema
+        if backed not in allowed_backed_modes:
+            raise ValueError(f"Invalid backed mode: {backed}. Choose from {allowed_backed_modes}")
+        self.backed = backed
         if cache is None:
             cache = LRU()
         self.cache = cache
@@ -382,16 +403,16 @@ def cached(self) -> bool:
 
     @property
     def adata(self) -> AnnData:
-        """Return backed anndata from the filename"""
+        """Return anndata from the filename"""
         try:
             adata = self.cache[self.filename]
         except KeyError:
             # fetch anndata
-            adata = read_h5ad_file(self.filename)
+            adata = read_h5ad_file(self.filename, backed=self.backed)
             # validate anndata
             if self.n_obs != adata.n_obs:
                 raise ValueError(
-                    "Expected `n_obs` for LazyAnnData object and backed anndata to match "
+                    "Expected `n_obs` for LazyAnnData object and loaded anndata to match "
                     f"but found {self.n_obs} and {adata.n_obs}, respectively."
                 )
             self.schema.validate_anndata(adata)
@@ -426,8 +447,9 @@ def __repr__(self) -> str:
             buffered = "Cached "
         else:
             buffered = ""
-        backed_at = f" backed at {str(self.filename)!r}"
-        descr = f"{buffered}LazyAnnData object with n_obs × n_vars = {self.n_obs} × {self.n_vars}{backed_at}"
+        located_at = f" referencing {str(self.filename)!r}"
+        backed = " in backed mode" if (self.backed in [True, "r"]) else " in memory mode"
+        descr = f"{buffered}LazyAnnData object with n_obs × n_vars = {self.n_obs} × {self.n_vars}{located_at}{backed}"
         if self.cached:
             for attr in self._all_attrs:
                 keys = getattr(self, attr).keys()
diff --git a/tests/dataloader/test_distributed_anndata.py b/tests/dataloader/test_distributed_anndata.py
index 31c4769f..4244a4c7 100644
--- a/tests/dataloader/test_distributed_anndata.py
+++ b/tests/dataloader/test_distributed_anndata.py
@@ -69,22 +69,31 @@ def adatas_path(tmp_path: Path):
 
 @pytest.fixture
 def adt(adatas_path: Path):
-    # single anndata
+    # single anndata in memory
     adt = read_h5ad_file(str(os.path.join(adatas_path, "adata.h5ad")), backed=None)
     return adt
 
 
-@pytest.fixture(params=[(i, j) for i in (1, 2, 3) for j in (True, False)])
+@pytest.fixture(
+    params=[(i, j, k) for i in (1, 2, 3) for j in (True, False) for k in (None, "r")],
+    ids=[
+        f"cache{i}-enforce{j}-{'memory' if k is None else 'backed'}"
+        for i in (1, 2, 3)
+        for j in (True, False)
+        for k in (None, "r")
+    ],
+)
 def dat(adatas_path: Path, request: pytest.FixtureRequest):
     # distributed anndata
     filenames = str(os.path.join(adatas_path, "adata.{000..002}.h5ad"))
     limits = [2, 5, 10]
-    max_cache_size, cache_size_strictly_enforced = request.param
+    max_cache_size, cache_size_strictly_enforced, backed_mode = request.param
     dat = DistributedAnnDataCollection(
         filenames,
         limits,
         max_cache_size=max_cache_size,
         cache_size_strictly_enforced=cache_size_strictly_enforced,
+        backed=backed_mode,
     )
     return dat
 
@@ -132,7 +141,6 @@ def test_init_shard_size(adatas_path: Path, num_shards: int, last_shard_size: in
     ids=["one adata", "two adatas", "sorted two adatas", "unsorted three adatas"],
 )
 @pytest.mark.parametrize("vidx", [slice(0, 2), [3, 1, 0]])
-# @pytest.mark.parametrize("adt", [None, "r"], ids=["memory", "backed"], indirect=True)
 def test_indexing(
     adt: AnnData,
     dat: DistributedAnnDataCollection,

From 83c514f2c6742f61441ac98a4e2fe11a4668e7fc Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 18:06:17 -0400
Subject: [PATCH 13/42] use backed_mode_type

---
 cellarium/ml/data/fileio.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 27de26d2..86b481aa 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -19,7 +19,7 @@
 def read_h5ad_gcs(
     filename: str,
     storage_client: Client | None = None,
-    backed: Literal["r", "r+"] | bool | None = backed_mode_default,
+    backed: backed_mode_type = backed_mode_default,
 ) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the Google Cloud Storage.
@@ -59,7 +59,7 @@ def read_h5ad_gcs(
                 pass  # if there's an error during cleanup, continue
 
 
-def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default) -> AnnData:
+def read_h5ad_url(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the URL.
 
@@ -96,7 +96,7 @@ def read_h5ad_url(filename: str, backed: Literal["r", "r+"] | bool | None = back
                 pass  # if there's an error during cleanup, continue
 
 
-def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default) -> AnnData:
+def read_h5ad_local(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the local disk.
 
@@ -113,7 +113,7 @@ def read_h5ad_local(filename: str, backed: Literal["r", "r+"] | bool | None = ba
     return read_h5ad(filename, backed=backed)
 
 
-def read_h5ad_file(filename: str, backed: Literal["r", "r+"] | bool | None = backed_mode_default, **kwargs) -> AnnData:
+def read_h5ad_file(filename: str, backed: backed_mode_type = backed_mode_default, **kwargs) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from a filename.
 

From 1b8142868048915eafc06df6405ba832a14d4a3f Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 30 Jul 2025 18:11:25 -0400
Subject: [PATCH 14/42] docstring update to reflect lack of "r+" mode

---
 cellarium/ml/data/fileio.py | 26 ++++++++++++--------------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 86b481aa..01a3c1c6 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -30,9 +30,9 @@ def read_h5ad_gcs(
 
     Args:
         filename: Path to the data file in Cloud Storage.
-        backed: If 'r', load in backed mode instead of fully loading into memory.
-               If 'r+', load in backed mode with write access (only X can be modified).
-               If True, equivalent to 'r'. Default is None (load into memory).
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
     """
     if not filename.startswith("gs:"):
         raise ValueError("The filename must start with 'gs:' protocol name.")
@@ -75,9 +75,9 @@ def read_h5ad_url(filename: str, backed: backed_mode_type = backed_mode_default)
 
     Args:
         filename: URL of the data file.
-         backed: If 'r', load in backed mode instead of fully loading into memory.
-               If 'r+', load in backed mode with write access (only X can be modified).
-               If True, equivalent to 'r'. Default is None (load into memory).
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
     """
     if not any(filename.startswith(scheme) for scheme in url_schemes):
         raise ValueError("The filename must start with 'http:', 'https:', or 'ftp:' protocol name.")
@@ -102,10 +102,9 @@ def read_h5ad_local(filename: str, backed: backed_mode_type = backed_mode_defaul
 
     Args:
         filename: Path to the local data file.
-        backed: If 'r', load in backed mode instead of fully loading into memory.
-               If 'r+', load in backed mode with write access (only X can be modified).
-               If True, equivalent to 'r'. Default is None (load into memory).
-
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
     """
     if not filename.startswith("file:"):
         raise ValueError("The filename must start with 'file:' protocol name.")
@@ -119,10 +118,9 @@ def read_h5ad_file(filename: str, backed: backed_mode_type = backed_mode_default
 
     Args:
         filename: Path to the data file.
-        backed: If 'r', load in backed mode instead of fully loading into memory.
-               If 'r+', load in backed mode with write access (only X can be modified).
-               If True, equivalent to 'r'. Default is None (load into memory).
-
+        backed: See :func:`anndata.read_h5ad` for details on backed mode.
+            ['r', True] will load in backed mode instead of fully loading into memory.
+            [False, None] will use in-memory mode.
     """
     if filename.startswith("gs:"):
         return read_h5ad_gcs(filename, **kwargs)

From 983178aeb7c0c4bf6f16278fb8d7cc12751f57aa Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 31 Jul 2025 10:33:38 -0400
Subject: [PATCH 15/42] test memory pressure in backed mode with massive faked
 h5ad

---
 tests/dataloader/test_datamodule.py | 68 +++++++++++++++++++++++++++--
 1 file changed, 65 insertions(+), 3 deletions(-)

diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index fbbd4643..d386b76a 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -177,15 +177,15 @@ def _check_transform_lists_match(iterable_modules1, iterable_modules2, message):
     "accelerator",
     ["cpu", pytest.param("gpu", marks=pytest.mark.skipif(not USE_CUDA, reason="requires_cuda"))],
 )
-@pytest.mark.parametrize("batch_size", [50, None])
-def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) -> None:
+@pytest.mark.parametrize("batch_size", [50, 100])
+def test_datamodule(tmp_path: Path, batch_size: int, accelerator: str) -> None:
     dadc = DistributedAnnDataCollection(
         filenames="https://storage.googleapis.com/dsp-cellarium-cas-public/test-data/test_0.h5ad",
         shard_size=100,
     )
     datamodule = CellariumAnnDataDataModule(
         dadc=dadc,
-        batch_size=100,
+        batch_size=batch_size,
         batch_keys={
             "x_ng": AnnDataField(attr="X", convert_fn=densify),
         },
@@ -205,6 +205,68 @@ def test_datamodule(tmp_path: Path, batch_size: int | None, accelerator: str) ->
     assert loaded_datamodule.dadc is dadc
 
 
+@pytest.fixture
+def fake_massive_dense_h5ad(tmp_path: Path) -> Path:
+    import h5py
+    import numpy as np
+    
+    # Create a dataset that CLAIMS to be ~40GB but uses almost no disk space
+    n_obs = 2_000_000  # 2 million cells  
+    n_vars = 5_000     # 5k genes
+    
+    h5ad_path = tmp_path / "massive_fake.h5ad"
+    
+    with h5py.File(h5ad_path, "w") as f:
+        # Create X dataset with claimed huge size but minimal actual storage
+        # Using fillvalue=0.0 with chunking - chunks are only allocated when written to
+        f.create_dataset(
+            "X", 
+            shape=(n_obs, n_vars),
+            dtype=np.float32,
+            fillvalue=0.0,
+            chunks=True,  # Enable chunking so not all data needs to be stored
+            compression=None  # No compression to keep it simple
+        )
+        
+        # Create minimal obs metadata - just the index is required
+        obs_group = f.create_group("obs")
+        # Create a small obs index but tell HDF5 it could expand to n_obs
+        obs_index_data = np.array([f"CELL_{i:07d}".encode('utf-8') for i in range(n_obs)])
+        obs_group.create_dataset("_index", data=obs_index_data, maxshape=(n_obs,), dtype="S12")
+        
+        # Create minimal var metadata - just the index is required  
+        var_group = f.create_group("var")
+        var_index_data = np.array([f"GENE_{i:05d}".encode('utf-8') for i in range(n_vars)])
+        var_group.create_dataset("_index", data=var_index_data, dtype="S10")
+        
+        # Set minimal h5ad format attributes that anndata expects
+        f.attrs["encoding-type"] = "anndata"
+        f.attrs["encoding-version"] = "0.1.0"
+        
+    return h5ad_path
+
+
+def test_datamodule_massive_h5ad_backed(tmp_path: Path, fake_massive_dense_h5ad: Path) -> None:
+    # try training using a massive (faked) h5ad file which should only succeed if backed mode works
+    dadc = DistributedAnnDataCollection(
+        filenames=str(fake_massive_dense_h5ad),  # Use full path instead of just name
+        shard_size=2_000_000,
+        backed=True,
+    )
+    datamodule = CellariumAnnDataDataModule(
+        dadc=dadc,
+        batch_size=100,
+        batch_keys={
+            "x_ng": AnnDataField(attr="X", convert_fn=None),  # already dense
+        },
+    )
+    module = CellariumModule(model=BoringModel())
+    trainer = pl.Trainer(accelerator='cpu', devices=1, max_steps=1, default_root_dir=tmp_path)
+    trainer.fit(module, datamodule)
+    # the idea is if this can run without a memory overflow, backed mode is implemented correctly
+    # we have separately verified that backed=False will crash due to 40GB memory use
+
+
 @pytest.mark.parametrize(
     "n_samples, train_size, val_size, n_train_expected, n_val_expected",
     [

From 39bc58252a577b7bcd75718f1a36d2fc00c58110 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 31 Jul 2025 10:34:06 -0400
Subject: [PATCH 16/42] format

---
 tests/dataloader/test_datamodule.py | 30 ++++++++++++++---------------
 1 file changed, 15 insertions(+), 15 deletions(-)

diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index d386b76a..e7acd85c 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -209,40 +209,40 @@ def test_datamodule(tmp_path: Path, batch_size: int, accelerator: str) -> None:
 def fake_massive_dense_h5ad(tmp_path: Path) -> Path:
     import h5py
     import numpy as np
-    
+
     # Create a dataset that CLAIMS to be ~40GB but uses almost no disk space
-    n_obs = 2_000_000  # 2 million cells  
-    n_vars = 5_000     # 5k genes
-    
+    n_obs = 2_000_000  # 2 million cells
+    n_vars = 5_000  # 5k genes
+
     h5ad_path = tmp_path / "massive_fake.h5ad"
-    
+
     with h5py.File(h5ad_path, "w") as f:
         # Create X dataset with claimed huge size but minimal actual storage
         # Using fillvalue=0.0 with chunking - chunks are only allocated when written to
         f.create_dataset(
-            "X", 
+            "X",
             shape=(n_obs, n_vars),
             dtype=np.float32,
             fillvalue=0.0,
             chunks=True,  # Enable chunking so not all data needs to be stored
-            compression=None  # No compression to keep it simple
+            compression=None,  # No compression to keep it simple
         )
-        
+
         # Create minimal obs metadata - just the index is required
         obs_group = f.create_group("obs")
         # Create a small obs index but tell HDF5 it could expand to n_obs
-        obs_index_data = np.array([f"CELL_{i:07d}".encode('utf-8') for i in range(n_obs)])
+        obs_index_data = np.array([f"CELL_{i:07d}".encode("utf-8") for i in range(n_obs)])
         obs_group.create_dataset("_index", data=obs_index_data, maxshape=(n_obs,), dtype="S12")
-        
-        # Create minimal var metadata - just the index is required  
+
+        # Create minimal var metadata - just the index is required
         var_group = f.create_group("var")
-        var_index_data = np.array([f"GENE_{i:05d}".encode('utf-8') for i in range(n_vars)])
+        var_index_data = np.array([f"GENE_{i:05d}".encode("utf-8") for i in range(n_vars)])
         var_group.create_dataset("_index", data=var_index_data, dtype="S10")
-        
+
         # Set minimal h5ad format attributes that anndata expects
         f.attrs["encoding-type"] = "anndata"
         f.attrs["encoding-version"] = "0.1.0"
-        
+
     return h5ad_path
 
 
@@ -261,7 +261,7 @@ def test_datamodule_massive_h5ad_backed(tmp_path: Path, fake_massive_dense_h5ad:
         },
     )
     module = CellariumModule(model=BoringModel())
-    trainer = pl.Trainer(accelerator='cpu', devices=1, max_steps=1, default_root_dir=tmp_path)
+    trainer = pl.Trainer(accelerator="cpu", devices=1, max_steps=1, default_root_dir=tmp_path)
     trainer.fit(module, datamodule)
     # the idea is if this can run without a memory overflow, backed mode is implemented correctly
     # we have separately verified that backed=False will crash due to 40GB memory use

From a370420656f5f324d81d289af942fac28e3b37d6 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Tue, 7 Apr 2026 18:16:10 -0400
Subject: [PATCH 17/42] remove tests with repeat indices

---
 tests/dataloader/test_distributed_anndata.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/dataloader/test_distributed_anndata.py b/tests/dataloader/test_distributed_anndata.py
index 4244a4c7..98710b3e 100644
--- a/tests/dataloader/test_distributed_anndata.py
+++ b/tests/dataloader/test_distributed_anndata.py
@@ -137,7 +137,7 @@ def test_init_shard_size(adatas_path: Path, num_shards: int, last_shard_size: in
 
 @pytest.mark.parametrize(
     "row_select",
-    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 4, 4], 2), ([6, 1, 3], 3)],
+    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 3, 4], 2), ([6, 1, 3], 3)],
     ids=["one adata", "two adatas", "sorted two adatas", "unsorted three adatas"],
 )
 @pytest.mark.parametrize("vidx", [slice(0, 2), [3, 1, 0]])
@@ -186,7 +186,7 @@ def test_pickle(dat: DistributedAnnDataCollection):
 
 @pytest.mark.parametrize(
     "row_select",
-    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 4, 4], 2), ([6, 1, 3], 3)],
+    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 3, 4], 2), ([6, 1, 3], 3)],
     ids=["one adata", "two adatas", "sorted two adatas", "unsorted three adatas"],
 )
 def test_indexing_dataset(

From 4b6a9e582400d2925f4bde78e82a63f5e447110d Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Tue, 7 Apr 2026 18:22:16 -0400
Subject: [PATCH 18/42] revert that change

---
 tests/dataloader/test_distributed_anndata.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/dataloader/test_distributed_anndata.py b/tests/dataloader/test_distributed_anndata.py
index 98710b3e..4244a4c7 100644
--- a/tests/dataloader/test_distributed_anndata.py
+++ b/tests/dataloader/test_distributed_anndata.py
@@ -137,7 +137,7 @@ def test_init_shard_size(adatas_path: Path, num_shards: int, last_shard_size: in
 
 @pytest.mark.parametrize(
     "row_select",
-    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 3, 4], 2), ([6, 1, 3], 3)],
+    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 4, 4], 2), ([6, 1, 3], 3)],
     ids=["one adata", "two adatas", "sorted two adatas", "unsorted three adatas"],
 )
 @pytest.mark.parametrize("vidx", [slice(0, 2), [3, 1, 0]])
@@ -186,7 +186,7 @@ def test_pickle(dat: DistributedAnnDataCollection):
 
 @pytest.mark.parametrize(
     "row_select",
-    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 3, 4], 2), ([6, 1, 3], 3)],
+    [(slice(0, 2), 1), (slice(1, 4), 2), ([1, 2, 4, 4], 2), ([6, 1, 3], 3)],
     ids=["one adata", "two adatas", "sorted two adatas", "unsorted three adatas"],
 )
 def test_indexing_dataset(

From 3e5c22fa31232106ba26c9b00453c1a9a2e987ce Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Tue, 7 Apr 2026 18:31:27 -0400
Subject: [PATCH 19/42] upgrade anndata which requires python 3.12

---
 .github/workflows/ci.yml       | 12 ++++++------
 .github/workflows/pypi.yml     |  2 +-
 .github/workflows/testpypi.yml |  2 +-
 pyproject.toml                 |  4 ++--
 4 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 03bf8f4e..82dcee8f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -15,7 +15,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -35,7 +35,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -56,7 +56,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
@@ -77,7 +77,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
         devices: ["1", "2", "3"]
 
     env:
@@ -101,7 +101,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
         devices: ["1", "2", "3"]
 
     env:
@@ -125,7 +125,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     steps:
     - uses: actions/checkout@v2
diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
index 4e3249b3..747ab8af 100644
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -22,7 +22,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "3.12"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/.github/workflows/testpypi.yml b/.github/workflows/testpypi.yml
index cfc99a27..920936d5 100644
--- a/.github/workflows/testpypi.yml
+++ b/.github/workflows/testpypi.yml
@@ -22,7 +22,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v3
       with:
-        python-version: "3.10"
+        python-version: "3.12"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
diff --git a/pyproject.toml b/pyproject.toml
index 48b3e6fc..1a7d20f6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [
 ]
 description = "Machine learning library for single-cell data analysis"
 readme = "README.rst"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 license = {file = "LICENSE.md"}
 classifiers = [
   "Intended Audience :: Developers",
@@ -21,7 +21,7 @@ classifiers = [
 ]
 
 dependencies = [
-  "anndata",
+  "anndata>=0.12.4",
   "boltons",
   "braceexpand",
   "cosine_annealing_warmup @ https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup/tarball/master#egg=cosine_annealing_warmupmatplotlib",

From 60584d82a222397be5193b499ee24051907257b6 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Tue, 7 Apr 2026 18:50:54 -0400
Subject: [PATCH 20/42] update jsonargparse and fix its usage

---
 cellarium/ml/cli.py | 4 ++--
 pyproject.toml      | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 9570df75..27d6d566 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -19,7 +19,7 @@
 import torch
 import yaml
 from jsonargparse import Namespace, class_from_function
-from jsonargparse._loaders_dumpers import DefaultLoader
+from jsonargparse._loaders_dumpers import get_yaml_default_loader
 from jsonargparse._util import import_object
 from lightning.pytorch.cli import ArgsType, LightningArgumentParser, LightningCLI
 from torch._subclasses.fake_tensor import FakeCopyMode, FakeTensorMode
@@ -137,7 +137,7 @@ def checkpoint_loader_constructor(loader: yaml.SafeLoader, node: yaml.nodes.Mapp
     return CheckpointLoader(**loader.construct_mapping(node))  # type: ignore[arg-type]
 
 
-loader = DefaultLoader
+loader = get_yaml_default_loader()
 loader.add_constructor("!FileLoader", file_loader_constructor)
 loader.add_constructor("!CheckpointLoader", checkpoint_loader_constructor)
 
diff --git a/pyproject.toml b/pyproject.toml
index 1a7d20f6..f7e35136 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,7 @@ dependencies = [
   "cosine_annealing_warmup @ https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup/tarball/master#egg=cosine_annealing_warmupmatplotlib",
   "crick>=0.0.4",
   "google-cloud-storage",
-  "jsonargparse[signatures]==4.27.7",
+  "jsonargparse[signatures]>=4.28.0",
   "lightning>=2.2.0,<=2.5.2",
   "numpy<=1.26.4",
   "pyro-ppl>=1.9.1",

From 4dc3bf5c0b6b2e133cde27bb8b222084e9738d0a Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Tue, 7 Apr 2026 19:10:26 -0400
Subject: [PATCH 21/42] allow newer lightning and fix weights_only=True default
 loader issue

---
 cellarium/ml/cli.py | 5 ++++-
 pyproject.toml      | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 27d6d566..89917ab3 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -124,8 +124,11 @@ class CheckpointLoader(FileLoader):
     convert_fn: Callable[[Any], Any] | str | None = None
 
     def __new__(cls, file_path, attr=None, key=None, convert_fn=None):
-        return super().__new__(cls, file_path, CellariumModule.load_from_checkpoint, attr, key, convert_fn)
+        return super().__new__(cls, file_path, _load_cellarium_module, attr, key, convert_fn)
 
+def _load_cellarium_module(path: str) -> CellariumModule:
+    """Wraps load_from_checkpoint, injecting weights_only=False, due to pytorch change in 2.4.0."""
+    return CellariumModule.load_from_checkpoint(path, weights_only=False)
 
 def file_loader_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> FileLoader:
     """Construct an object from a file."""
diff --git a/pyproject.toml b/pyproject.toml
index f7e35136..f0076c2c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
   "crick>=0.0.4",
   "google-cloud-storage",
   "jsonargparse[signatures]>=4.28.0",
-  "lightning>=2.2.0,<=2.5.2",
+  "lightning>=2.2.0",
   "numpy<=1.26.4",
   "pyro-ppl>=1.9.1",
   "pytest",

From 9d188f5d38e60bfd549c211a7fdb51d19f5c17bc Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Tue, 7 Apr 2026 19:10:42 -0400
Subject: [PATCH 22/42] linting

---
 cellarium/ml/cli.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 89917ab3..77ea76c9 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -126,10 +126,12 @@ class CheckpointLoader(FileLoader):
     def __new__(cls, file_path, attr=None, key=None, convert_fn=None):
         return super().__new__(cls, file_path, _load_cellarium_module, attr, key, convert_fn)
 
+
 def _load_cellarium_module(path: str) -> CellariumModule:
     """Wraps load_from_checkpoint, injecting weights_only=False, due to pytorch change in 2.4.0."""
     return CellariumModule.load_from_checkpoint(path, weights_only=False)
 
+
 def file_loader_constructor(loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> FileLoader:
     """Construct an object from a file."""
     return FileLoader(**loader.construct_mapping(node))  # type: ignore[arg-type]

From 21f06ba4200eb58aff22ac0af712e801189a9c7d Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 8 Apr 2026 00:50:24 -0400
Subject: [PATCH 23/42] override torch 2.4 default weights_only=True with False

---
 cellarium/ml/cli.py                 |  6 ++++++
 cellarium/ml/core/module.py         | 20 ++++++++++++++++++++
 tests/dataloader/test_datamodule.py |  2 +-
 tests/dataloader/test_dataset.py    |  4 ++--
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 77ea76c9..3b8979d5 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -340,6 +340,12 @@ def _add_instantiators(self) -> None:
             # https://github.com/Lightning-AI/pytorch-lightning/pull/18105
             pass
 
+        def _prepare_subcommand_parser(self, klass, subcommand, **kwargs):
+            """Override the default checkpoint loading with weights_only=True in torch 2.4.0+"""
+            parser = super()._prepare_subcommand_parser(klass, subcommand, **kwargs)
+            parser.set_defaults({"weights_only": False})
+            return parser
+
         def before_instantiate_classes(self):
             # issue a UserWarning if the subcommand is predict and return_predictions is not set to False
             if self.subcommand == "predict":
diff --git a/cellarium/ml/core/module.py b/cellarium/ml/core/module.py
index 260c6ec6..f0429653 100644
--- a/cellarium/ml/core/module.py
+++ b/cellarium/ml/core/module.py
@@ -79,6 +79,26 @@ def __init__(
             # Thus, we need to use manual optimization for the No Optimizer case.
             self.automatic_optimization = False
 
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint_path,
+        map_location=None,
+        hparams_file=None,
+        strict=None,
+        weights_only=False,
+        **kwargs,
+    ):
+        """Replace the torch 2.4.0+ default weights_only=True with False"""
+        return super().load_from_checkpoint(
+            checkpoint_path,
+            map_location=map_location,
+            hparams_file=hparams_file,
+            strict=strict,
+            weights_only=weights_only,
+            **kwargs,
+        )
+
     def configure_model(self) -> None:
         """
         .. note::
diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index e7acd85c..083ebc43 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -140,7 +140,7 @@ def _check_transform_lists_match(iterable_modules1, iterable_modules2, message):
     module = _new_module()
     trainer = pl.Trainer(accelerator=accelerator, devices=1, max_steps=2, default_root_dir=tmp_path)
     print("Training one more step starting from a checkpoint...")
-    trainer.fit(module, datamodule, ckpt_path=ckpt_path)
+    trainer.fit(module, datamodule, ckpt_path=ckpt_path, weights_only=False)
     print("    ... ✓")
 
     print("Ensuring the pipeline is correctly configured in the restarted Trainer... ")
diff --git a/tests/dataloader/test_dataset.py b/tests/dataloader/test_dataset.py
index 8d3a4899..ed606c9f 100644
--- a/tests/dataloader/test_dataset.py
+++ b/tests/dataloader/test_dataset.py
@@ -370,10 +370,10 @@ def test_load_from_checkpoint(
     )
     try:
         ckpt_path = tmp_path / f"checkpoints/epoch=0-step={resume_step}.ckpt"
-        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path)
+        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path, weights_only=False)
     except FileNotFoundError:
         ckpt_path = tmp_path / f"checkpoints/epoch=1-step={resume_step}.ckpt"
-        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path)
+        trainer2.fit(module2, datamodule2, ckpt_path=ckpt_path, weights_only=False)
 
     iter_data1 = collate_fn(module1.model.iter_data)
     iter_data2 = collate_fn(module2.model.iter_data)

From 8476e547bd824ca1775fb0cdf3319e774384eaab Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 8 Apr 2026 10:55:39 -0400
Subject: [PATCH 24/42] fix meta device use for compute_var_names_g

---
 cellarium/ml/cli.py | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 3b8979d5..38f07f03 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -17,6 +17,7 @@
 import numpy as np
 import pandas as pd
 import torch
+from transformers import pipeline
 import yaml
 from jsonargparse import Namespace, class_from_function
 from jsonargparse._loaders_dumpers import get_yaml_default_loader
@@ -251,9 +252,11 @@ def compute_var_names_g(
     pipeline = CellariumPipeline(cpu_transforms) + CellariumPipeline(transforms)
     with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode:
         fake_batch = collate_fn([batch])
-        with FakeCopyMode(fake_mode):
-            fake_pipeline = copy.deepcopy(pipeline)
-        output = fake_pipeline(fake_batch)
+        fake_batch = {k: v.to("meta") if isinstance(v, torch.Tensor) else v for k, v in fake_batch.items()}
+        with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode:
+            with FakeCopyMode(fake_mode):
+                fake_pipeline = copy.deepcopy(pipeline)
+            output = fake_pipeline(fake_batch)
     return output["var_names_g"]
 
 
@@ -345,6 +348,30 @@ def _prepare_subcommand_parser(self, klass, subcommand, **kwargs):
             parser = super()._prepare_subcommand_parser(klass, subcommand, **kwargs)
             parser.set_defaults({"weights_only": False})
             return parser
+        
+        def _parse_ckpt_path(self) -> None:
+            from pathlib import Path
+            if not self.config.get("subcommand"):
+                return
+            ckpt_path = self.config[self.config.subcommand].get("ckpt_path")
+            if ckpt_path and Path(ckpt_path).is_file():
+                ckpt = torch.load(ckpt_path, weights_only=False, map_location="cpu")
+                hparams = ckpt.get("hyper_parameters", {})
+                hparams.pop("_instantiator", None)
+                if not hparams:
+                    return
+                if "_class_path" in hparams:
+                    hparams = {
+                        "class_path": hparams.pop("_class_path"),
+                        "dict_kwargs": hparams,
+                    }
+                hparams = {self.config.subcommand: {"model": hparams}}
+                try:
+                    self.config = self.parser.parse_object(hparams, self.config)
+                except SystemExit:
+                    import sys
+                    sys.stderr.write("Parsing of ckpt_path hyperparameters failed!\n")
+                    raise
 
         def before_instantiate_classes(self):
             # issue a UserWarning if the subcommand is predict and return_predictions is not set to False

From c8755130361ce04b5180d65d8cb61a54dfaaeea5 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 8 Apr 2026 10:59:39 -0400
Subject: [PATCH 25/42] set weights_only=False in test

---
 tests/dataloader/test_datamodule.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index 083ebc43..ab40a8f0 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -195,7 +195,7 @@ def test_datamodule(tmp_path: Path, batch_size: int, accelerator: str) -> None:
     trainer.fit(module, datamodule)
 
     ckpt_path = str(tmp_path / "lightning_logs/version_0/checkpoints/epoch=0-step=1.ckpt")
-    kwargs: dict[str, Any] = {"dadc": dadc}
+    kwargs: dict[str, Any] = {"dadc": dadc, "weights_only": False}
     if batch_size is not None:
         kwargs["batch_size"] = batch_size
     loaded_datamodule = CellariumAnnDataDataModule.load_from_checkpoint(ckpt_path, **kwargs)

From b661b5fee545963938434d0727a998ba4ab9d6e9 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 8 Apr 2026 11:05:19 -0400
Subject: [PATCH 26/42] linting

---
 cellarium/ml/cli.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 38f07f03..ced1d5ea 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -17,7 +17,6 @@
 import numpy as np
 import pandas as pd
 import torch
-from transformers import pipeline
 import yaml
 from jsonargparse import Namespace, class_from_function
 from jsonargparse._loaders_dumpers import get_yaml_default_loader
@@ -348,9 +347,10 @@ def _prepare_subcommand_parser(self, klass, subcommand, **kwargs):
             parser = super()._prepare_subcommand_parser(klass, subcommand, **kwargs)
             parser.set_defaults({"weights_only": False})
             return parser
-        
+
         def _parse_ckpt_path(self) -> None:
             from pathlib import Path
+
             if not self.config.get("subcommand"):
                 return
             ckpt_path = self.config[self.config.subcommand].get("ckpt_path")
@@ -370,6 +370,7 @@ def _parse_ckpt_path(self) -> None:
                     self.config = self.parser.parse_object(hparams, self.config)
                 except SystemExit:
                     import sys
+
                     sys.stderr.write("Parsing of ckpt_path hyperparameters failed!\n")
                     raise
 

From ab9af9d474967d9b57baf1e292f05b1e7e88730a Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 8 Apr 2026 13:35:23 -0400
Subject: [PATCH 27/42] better fix for CLI compute_var_names_g

---
 cellarium/ml/cli.py | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index ced1d5ea..845e9b48 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -249,13 +249,23 @@ def compute_var_names_g(
     adata = data.dadc[0]
     batch = tree_map(lambda field: field(adata), data.batch_keys)
     pipeline = CellariumPipeline(cpu_transforms) + CellariumPipeline(transforms)
+    # Normalize all pipeline tensors to meta device before entering FakeTensorMode.
+    # This is necessary because CheckpointLoader-loaded modules carry real cpu tensors,
+    # while freshly instantiated modules have meta tensors. FakeCopyMode would then
+    # produce a mix of cpu and meta FakeTensors, causing _find_common_device to raise.
+    # Only deepcopy when cpu tensors are present (e.g. CheckpointLoader transforms);
+    # in the common case all tensors are already on meta so the copy is unnecessary.
+    if any(p.device.type == "cpu" for p in pipeline.parameters()):
+        pipeline_meta = copy.deepcopy(pipeline)
+        pipeline_meta.to_empty(device="meta")
+    else:
+        pipeline_meta = pipeline
+    fake_batch = collate_fn([batch])
+    fake_batch = {k: v.to("meta") if isinstance(v, torch.Tensor) else v for k, v in fake_batch.items()}
     with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode:
-        fake_batch = collate_fn([batch])
-        fake_batch = {k: v.to("meta") if isinstance(v, torch.Tensor) else v for k, v in fake_batch.items()}
-        with FakeTensorMode(allow_non_fake_inputs=True) as fake_mode:
-            with FakeCopyMode(fake_mode):
-                fake_pipeline = copy.deepcopy(pipeline)
-            output = fake_pipeline(fake_batch)
+        with FakeCopyMode(fake_mode):
+            fake_pipeline = copy.deepcopy(pipeline_meta)
+        output = fake_pipeline(fake_batch)
     return output["var_names_g"]
 
 

From 1b5f363e82c346515f5110e2574c9e3e03deb12b Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 8 Apr 2026 14:07:13 -0400
Subject: [PATCH 28/42] fix compute_var_names_g again for meta cpu conflicts

---
 cellarium/ml/cli.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index 845e9b48..d9005526 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -255,7 +255,7 @@ def compute_var_names_g(
     # produce a mix of cpu and meta FakeTensors, causing _find_common_device to raise.
     # Only deepcopy when cpu tensors are present (e.g. CheckpointLoader transforms);
     # in the common case all tensors are already on meta so the copy is unnecessary.
-    if any(p.device.type == "cpu" for p in pipeline.parameters()):
+    if any(t.device.type == "cpu" for t in [*pipeline.parameters(), *pipeline.buffers()]):
         pipeline_meta = copy.deepcopy(pipeline)
         pipeline_meta.to_empty(device="meta")
     else:

From 591573743220c1b35342651bffc9b49432f25b26 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Fri, 19 Jun 2026 01:43:52 -0400
Subject: [PATCH 29/42] future-proof anndata 0.13.x

---
 cellarium/ml/data/distributed_anndata.py     |  8 +++++++-
 tests/dataloader/test_distributed_anndata.py | 10 +++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/data/distributed_anndata.py b/cellarium/ml/data/distributed_anndata.py
index 4cf92b4e..503c5e49 100644
--- a/cellarium/ml/data/distributed_anndata.py
+++ b/cellarium/ml/data/distributed_anndata.py
@@ -9,7 +9,13 @@
 import pandas as pd
 from anndata import AnnData, concat
 from anndata._core.index import _normalize_indices
-from anndata.compat import Index, Index1D
+
+try:
+    from anndata.compat import Index, Index1D
+except ImportError:
+    from anndata.typing import Index, Index1D
+
+    dtype_deprecated = True
 from anndata.experimental.multi_files._anncollection import (
     AnnCollection,
     AnnCollectionView,
diff --git a/tests/dataloader/test_distributed_anndata.py b/tests/dataloader/test_distributed_anndata.py
index 4244a4c7..d8ada2a6 100644
--- a/tests/dataloader/test_distributed_anndata.py
+++ b/tests/dataloader/test_distributed_anndata.py
@@ -1,6 +1,7 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import importlib.util
 import os
 import pickle
 from pathlib import Path
@@ -10,6 +11,12 @@
 import pytest
 from anndata import AnnData
 
+if importlib.util.find_spec("anndata.compat.Index") is None:
+    # anndata 0.13.x +
+    dtype_deprecated = True
+else:
+    dtype_deprecated = False
+
 from cellarium.ml.data import (
     DistributedAnnDataCollection,
     IterableDistributedAnnDataCollectionDataset,
@@ -49,9 +56,10 @@ def adatas_path(tmp_path: Path):
         },
         index=[f"gene{i:03d}" for i in range(n_gene)],
     )
+    dtype_kwargs = {"dtype": X.dtype} if not dtype_deprecated else {}
     adata = AnnData(
         X,
-        dtype=X.dtype,
+        **dtype_kwargs,
         obs=obs,
         var=var,
         layers={"L": L},

From 3a24697e0690a9bc4d041b18cb34b80e9e6d8e5e Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Fri, 19 Jun 2026 01:44:34 -0400
Subject: [PATCH 30/42] future-proof anndata 0.13.x

---
 cellarium/ml/data/distributed_anndata.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/cellarium/ml/data/distributed_anndata.py b/cellarium/ml/data/distributed_anndata.py
index 503c5e49..d33a50da 100644
--- a/cellarium/ml/data/distributed_anndata.py
+++ b/cellarium/ml/data/distributed_anndata.py
@@ -9,12 +9,10 @@
 import pandas as pd
 from anndata import AnnData, concat
 from anndata._core.index import _normalize_indices
-
 try:
     from anndata.compat import Index, Index1D
 except ImportError:
     from anndata.typing import Index, Index1D
-
     dtype_deprecated = True
 from anndata.experimental.multi_files._anncollection import (
     AnnCollection,

From 3ff559413e9461f01b7efafe90dbae13ad769d2f Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Fri, 19 Jun 2026 01:44:48 -0400
Subject: [PATCH 31/42] linting

---
 cellarium/ml/data/distributed_anndata.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cellarium/ml/data/distributed_anndata.py b/cellarium/ml/data/distributed_anndata.py
index d33a50da..503c5e49 100644
--- a/cellarium/ml/data/distributed_anndata.py
+++ b/cellarium/ml/data/distributed_anndata.py
@@ -9,10 +9,12 @@
 import pandas as pd
 from anndata import AnnData, concat
 from anndata._core.index import _normalize_indices
+
 try:
     from anndata.compat import Index, Index1D
 except ImportError:
     from anndata.typing import Index, Index1D
+
     dtype_deprecated = True
 from anndata.experimental.multi_files._anncollection import (
     AnnCollection,

From 5eb734c96580f5cf3fca51265e51f11e0d9ede98 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Fri, 19 Jun 2026 01:45:05 -0400
Subject: [PATCH 32/42] mypy fix

---
 cellarium/ml/cli.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cellarium/ml/cli.py b/cellarium/ml/cli.py
index d9005526..28d82a03 100644
--- a/cellarium/ml/cli.py
+++ b/cellarium/ml/cli.py
@@ -28,7 +28,7 @@
 from cellarium.ml import CellariumAnnDataDataModule, CellariumModule, CellariumPipeline
 from cellarium.ml.utilities.data import AnnDataField, collate_fn
 
-cached_loaders = {}
+cached_loaders: dict[Callable, Callable] = {}
 
 
 @dataclass
@@ -147,7 +147,7 @@ def checkpoint_loader_constructor(loader: yaml.SafeLoader, node: yaml.nodes.Mapp
 loader.add_constructor("!CheckpointLoader", checkpoint_loader_constructor)
 
 
-REGISTERED_MODELS = {}
+REGISTERED_MODELS: dict[str, Callable[[ArgsType], None]] = {}
 
 
 def register_model(model: Callable[[ArgsType], None]):

From 7d10999f900db964a8f73e057292b2033624c3b7 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 1 Jul 2026 18:59:18 -0400
Subject: [PATCH 33/42] pin scverse-misc due to anndata

---
 pyproject.toml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index d38cfe25..7469290b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,6 +32,7 @@ dependencies = [
   "pyro-ppl>=1.9.1",
   "pytest",
   "scikit-misc==0.5.1",
+  "scverse-misc-0.1.1",
   "torch>=2.2.0",
   "transformers",
   "zuko==1.6.0",
@@ -50,7 +51,7 @@ test = [
   "tensorboard",
   "torchvision",
   "scvi-tools==1.3.3",
-  "zarr==2.18.3",
+  "zarr==3.1.5",
 ]
 docs = [
   "nbsphinx",

From 901bfee4b38f55b48a3ef7a7ea64f16ffb3ac63e Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 1 Jul 2026 19:31:38 -0400
Subject: [PATCH 34/42] pin scverse-misc

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7469290b..b549f22b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,7 +32,7 @@ dependencies = [
   "pyro-ppl>=1.9.1",
   "pytest",
   "scikit-misc==0.5.1",
-  "scverse-misc-0.1.1",
+  "scverse-misc==0.1.1",
   "torch>=2.2.0",
   "transformers",
   "zuko==1.6.0",

From 3a9a3cb5e1a1d5ac06e4c14eab7f9a8fb4bc2747 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Wed, 1 Jul 2026 19:31:47 -0400
Subject: [PATCH 35/42] fix h5py pickling issue

---
 cellarium/ml/core/datamodule.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/cellarium/ml/core/datamodule.py b/cellarium/ml/core/datamodule.py
index b5fb437d..f0aee7ce 100644
--- a/cellarium/ml/core/datamodule.py
+++ b/cellarium/ml/core/datamodule.py
@@ -126,9 +126,7 @@ def __init__(
         pin_memory: bool = False,
     ) -> None:
         super().__init__()
-        self.save_hyperparameters(logger=False)
-        # Don't save dadc to the checkpoint
-        self.hparams["dadc"] = None
+        self.save_hyperparameters(logger=False, ignore=["dadc"])
 
         self.dadc = dadc
         # IterableDistributedAnnDataCollectionDataset args

From 70dc6153baa1a5f14a9328b18c5d90e51ce84de2 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 2 Jul 2026 00:12:19 -0400
Subject: [PATCH 36/42] fix scvi test backed mode issue

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

diff --git a/tests/test_scvi.py b/tests/test_scvi.py
index 68388e8b..6e6e4bb7 100644
--- a/tests/test_scvi.py
+++ b/tests/test_scvi.py
@@ -2375,8 +2375,8 @@ def test_predict_reconstructed_counts_realdata(train_cellarium_model, train_scvi
     print(f"training data measured counts normalized to {reconstructed_library_size}")
     df = pd.DataFrame(
         (
-            np.array(train_adata_cellarium.X[:, :n_genes_to_reconstruct].todense()).squeeze()
-            / np.array(train_adata_cellarium.X.sum(axis=1)).squeeze()[:, None]
+            np.array(train_adata_cellarium.X[:, :n_genes_to_reconstruct].toarray()).squeeze()
+            / np.asarray(train_adata_cellarium.X[:].sum(axis=1)).squeeze()[:, None]  # materialize for backed case
             * reconstructed_library_size
         ),
         index=train_adata_cellarium.obs_names,

From 33bd58f1502bf1f81b00a7749b4fbb3ade5660a1 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 2 Jul 2026 17:02:23 -0400
Subject: [PATCH 37/42] backed default False

---
 cellarium/ml/data/fileio.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 01a3c1c6..05a49809 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -13,7 +13,7 @@
 
 url_schemes = ("http:", "https:", "ftp:")
 backed_mode_type = Literal["r"] | bool | None
-backed_mode_default: backed_mode_type = "r"
+backed_mode_default: backed_mode_type = False
 
 
 def read_h5ad_gcs(

From 1a1acc18e2bcb31c287294a0d96684fbcf698d1a Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 2 Jul 2026 17:23:32 -0400
Subject: [PATCH 38/42] fix bug with gcs read and test it

---
 cellarium/ml/data/fileio.py         | 14 +++--
 tests/dataloader/test_datamodule.py | 86 +++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index 05a49809..a0116e20 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -50,13 +50,15 @@ def read_h5ad_gcs(
     with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
         temp_path = tmp_file.name
         blob.download_to_file(tmp_file)
+        tmp_file.flush()
+
+    try:
+        return read_h5ad(temp_path, backed=backed)
+    finally:
         try:
-            return read_h5ad(temp_path, backed=backed)
-        finally:
-            try:
-                os.unlink(temp_path)  # clean up the temp file
-            except OSError:
-                pass  # if there's an error during cleanup, continue
+            os.unlink(temp_path)
+        except OSError:
+            pass
 
 
 def read_h5ad_url(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index ab40a8f0..ef496a33 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -267,6 +267,92 @@ def test_datamodule_massive_h5ad_backed(tmp_path: Path, fake_massive_dense_h5ad:
     # we have separately verified that backed=False will crash due to 40GB memory use
 
 
+@pytest.mark.parametrize("backed", [False, True], ids=["in_memory", "backed"])
+def test_datamodule_gcs_paths(backed: bool) -> None:
+    """End-to-end regression test: dataloader works with real GCS bucket paths.
+
+    Covers backed=False (in-memory) and backed=True so that both the write-buffer
+    flush fix and any backed-mode interaction with temporary files are exercised
+    against the actual GCS storage layer.
+    """
+    gcs_filenames = [
+        "gs://dsp-cellarium-cas-public/test-data/test_0.h5ad",
+        "gs://dsp-cellarium-cas-public/test-data/test_1.h5ad",
+        "gs://dsp-cellarium-cas-public/test-data/test_2.h5ad",
+    ]
+    dadc = DistributedAnnDataCollection(
+        filenames=gcs_filenames,
+        shard_size=100,
+        backed=backed,
+    )
+    datamodule = CellariumAnnDataDataModule(
+        dadc=dadc,
+        batch_size=50,
+        batch_keys={
+            "x_ng": AnnDataField(attr="X", convert_fn=densify),
+            "var_names_g": AnnDataField(attr="var_names"),
+        },
+        num_workers=0,
+    )
+    datamodule.setup("fit")
+    batch = next(iter(datamodule.train_dataloader()))
+
+    assert "x_ng" in batch
+    assert "var_names_g" in batch
+    x_ng = batch["x_ng"]
+    var_names_g = batch["var_names_g"]
+    assert x_ng.shape == (50, len(var_names_g)), (
+        f"Unexpected batch shape {x_ng.shape}; expected (50, {len(var_names_g)})"
+    )
+
+
+def test_read_h5ad_gcs_flushes_write_buffer(tmp_path: Path) -> None:
+    """Regression test: read_h5ad_gcs must flush the write buffer before reading.
+
+    The real GCS client streams data in chunks. Python's BufferedWriter accumulates
+    those chunks; when the buffer fills (every ~8 KB) it is auto-flushed to the OS,
+    but the final partial buffer fill stays in userspace memory until flush()/close()
+    is called. If read_h5ad opens the same path before that flush, h5py sees a
+    truncated file and raises OSError ("truncated file: eof = X, stored_eof = Y").
+
+    A single large write would bypass the buffer entirely, so this test writes in
+    1 KB chunks without flushing — exactly replicating the real client's behavior.
+    """
+    from unittest.mock import MagicMock
+
+    import anndata as ad
+    import numpy as np
+
+    from cellarium.ml.data.fileio import read_h5ad_gcs
+
+    rng = np.random.default_rng(42)
+    adata = ad.AnnData(X=rng.random((4, 3)).astype(np.float32))
+    small_h5ad = tmp_path / "small.h5ad"
+    adata.write_h5ad(small_h5ad)
+    h5ad_bytes = small_h5ad.read_bytes()
+
+    # The real GCS client streams in chunks; a single large write bypasses the buffer
+    # entirely and reaches disk immediately (no flush needed). Writing in 1 KB chunks
+    # means the last partial buffer fill (<8 KB) stays in Python's write buffer, so
+    # h5py would see a truncated file without an explicit flush().
+    chunk_size = 1024
+
+    def download_to_file_chunked(f: object) -> None:
+        for i in range(0, len(h5ad_bytes), chunk_size):
+            f.write(h5ad_bytes[i : i + chunk_size])  # type: ignore[attr-defined]
+
+    mock_blob = MagicMock()
+    mock_blob.download_to_file.side_effect = download_to_file_chunked
+    mock_bucket = MagicMock()
+    mock_bucket.blob.return_value = mock_blob
+    mock_client = MagicMock()
+    mock_client.bucket.return_value = mock_bucket
+
+    result = read_h5ad_gcs("gs://fake-bucket/fake.h5ad", storage_client=mock_client)
+
+    assert result.shape == adata.shape
+
+
 @pytest.mark.parametrize(
     "n_samples, train_size, val_size, n_train_expected, n_val_expected",
     [

From 6583fefc9dbf333c894792db86569ced324efa1b Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 2 Jul 2026 17:53:13 -0400
Subject: [PATCH 39/42] try default google credentials

---
 cellarium/ml/data/fileio.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index a0116e20..a1dde43d 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -9,6 +9,7 @@
 from typing import Literal
 
 from anndata import AnnData, read_h5ad
+from google.auth.exceptions import DefaultCredentialsError
 from google.cloud.storage import Client
 
 url_schemes = ("http:", "https:", "ftp:")
@@ -41,7 +42,10 @@ def read_h5ad_gcs(
     bucket_name, blob_name = filename.split("/", 1)
 
     if storage_client is None:
-        storage_client = Client()
+        try:
+            storage_client = Client()
+        except DefaultCredentialsError:
+            storage_client = Client.create_anonymous_client()
 
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(blob_name)

From 848bf96eb51b693c5653aaeefbae859b84076487 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Thu, 2 Jul 2026 23:00:12 -0400
Subject: [PATCH 40/42] bypass downloads to preserve disk in non-backed mode

---
 cellarium/ml/data/fileio.py | 36 +++++++++++++++++++++++++++---------
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index a1dde43d..c68d5d3b 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -50,7 +50,16 @@ def read_h5ad_gcs(
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(blob_name)
 
-    # write to a named temporary file
+    if backed not in [True, "r"]:
+        # Stream directly into memory — no temp file, no disk I/O.
+        with blob.open("rb") as f:
+            return read_h5ad(f)
+
+    # Backed mode requires h5py to have a real seekable file path. The file is
+    # flushed and closed before h5py opens it to avoid Python's write buffer
+    # leaving the last chunk off disk (which causes h5py to report a truncated
+    # file). After os.unlink the directory entry is gone but the inode stays
+    # allocated until h5py closes its fd (i.e. when the AnnData is GC'd).
     with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
         temp_path = tmp_file.name
         blob.download_to_file(tmp_file)
@@ -88,18 +97,27 @@ def read_h5ad_url(filename: str, backed: backed_mode_type = backed_mode_default)
     if not any(filename.startswith(scheme) for scheme in url_schemes):
         raise ValueError("The filename must start with 'http:', 'https:', or 'ftp:' protocol name.")
 
-    # write to a named temporary file
+    if backed not in [True, "r"]:
+        # Anonymous TemporaryFile: no path, no flush needed, deleted automatically.
+        with urllib.request.urlopen(filename) as response:
+            with tempfile.TemporaryFile() as tmp_file:
+                shutil.copyfileobj(response, tmp_file)
+                return read_h5ad(tmp_file)
+
+    # Backed mode needs a real path; flush and close before h5py opens it.
     with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
         temp_path = tmp_file.name
         with urllib.request.urlopen(filename) as response:
             shutil.copyfileobj(response, tmp_file)
+        tmp_file.flush()
+
+    try:
+        return read_h5ad(temp_path, backed=backed)
+    finally:
         try:
-            return read_h5ad(temp_path, backed=backed)
-        finally:
-            try:
-                os.unlink(temp_path)  # clean up the temp file
-            except OSError:
-                pass  # if there's an error during cleanup, continue
+            os.unlink(temp_path)
+        except OSError:
+            pass
 
 
 def read_h5ad_local(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
@@ -132,7 +150,7 @@ def read_h5ad_file(filename: str, backed: backed_mode_type = backed_mode_default
         return read_h5ad_gcs(filename, **kwargs)
 
     if filename.startswith("file:"):
-        return read_h5ad_local(filename)
+        return read_h5ad_local(filename, backed=backed)
 
     if any(filename.startswith(scheme) for scheme in url_schemes):
         return read_h5ad_url(filename)

From 8b312fb133e0412aebce56c0ad88c63f446976d3 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Fri, 3 Jul 2026 00:55:32 -0400
Subject: [PATCH 41/42] fix test

---
 tests/dataloader/test_datamodule.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/dataloader/test_datamodule.py b/tests/dataloader/test_datamodule.py
index ef496a33..0647a8ca 100644
--- a/tests/dataloader/test_datamodule.py
+++ b/tests/dataloader/test_datamodule.py
@@ -315,6 +315,10 @@ def test_read_h5ad_gcs_flushes_write_buffer(tmp_path: Path) -> None:
     is called. If read_h5ad opens the same path before that flush, h5py sees a
     truncated file and raises OSError ("truncated file: eof = X, stored_eof = Y").
 
+    This flush risk only applies to the backed-mode code path, which downloads to a
+    named temp file. Non-backed mode streams directly via blob.open() with no temp
+    file involved. The test therefore uses backed='r'.
+
     A single large write would bypass the buffer entirely, so this test writes in
     1 KB chunks without flushing — exactly replicating the real client's behavior.
     """
@@ -348,7 +352,7 @@ def download_to_file_chunked(f: object) -> None:
     mock_client = MagicMock()
     mock_client.bucket.return_value = mock_bucket
 
-    result = read_h5ad_gcs("gs://fake-bucket/fake.h5ad", storage_client=mock_client)
+    result = read_h5ad_gcs("gs://fake-bucket/fake.h5ad", storage_client=mock_client, backed="r")
 
     assert result.shape == adata.shape
 

From 09ec3e4fa9b48a05be8383fa1b205117ec81a286 Mon Sep 17 00:00:00 2001
From: Stephen Fleming <sjfleming@users.noreply.github.com>
Date: Fri, 3 Jul 2026 01:52:50 -0400
Subject: [PATCH 42/42] cache streamed gcs files for future local reading

---
 cellarium/ml/data/distributed_anndata.py | 29 ++++++++---
 cellarium/ml/data/fileio.py              | 63 +++++++++++++++++++-----
 2 files changed, 74 insertions(+), 18 deletions(-)

diff --git a/cellarium/ml/data/distributed_anndata.py b/cellarium/ml/data/distributed_anndata.py
index 503c5e49..116c6810 100644
--- a/cellarium/ml/data/distributed_anndata.py
+++ b/cellarium/ml/data/distributed_anndata.py
@@ -24,7 +24,7 @@
 from boltons.cacheutils import LRU
 from braceexpand import braceexpand
 
-from cellarium.ml.data.fileio import backed_mode_default, backed_mode_type, read_h5ad_file
+from cellarium.ml.data.fileio import GCS_CACHE_DIR, backed_mode_default, backed_mode_type, read_h5ad_file
 from cellarium.ml.data.schema import AnnDataSchema
 
 
@@ -169,6 +169,13 @@ class DistributedAnnDataCollection(AnnCollection):
             h5ad files, while ``None`` will load entire count matrices from individual h5ad files
             into cached memory as needed: a strategy that necessitates smaller chunked h5ad files.
             See :func:`anndata.read_h5ad` for details on backing modes.
+        cache_dir:
+            Directory for caching GCS-downloaded h5ad files on local disk. On first access each
+            shard is downloaded and written here; subsequent accesses (including after a checkpoint
+            restart on the same machine) read from disk instead of re-downloading. If a write fails
+            (e.g. disk full) the shard is streamed from GCS instead. Set to ``None`` to always
+            stream GCS files directly into memory with no local disk usage. Only affects
+            ``gs://`` paths; local and HTTP paths are unaffected.
     """
 
     def __init__(
@@ -186,6 +193,7 @@ def __init__(
         indices_strict: bool = True,
         obs_columns_to_validate: Sequence[str] | None = None,
         backed: backed_mode_type = backed_mode_default,
+        cache_dir: str | None = str(GCS_CACHE_DIR),
     ):
         self.filenames = list(braceexpand(filenames) if isinstance(filenames, str) else filenames)
         if (shard_size is None) and (last_shard_size is not None):
@@ -210,8 +218,11 @@ def __init__(
         if backed not in allowed_backed_modes:
             raise ValueError(f"Invalid backed mode: {backed}. Choose from {allowed_backed_modes}")
         self.backed = backed
+        self.cache_dir = cache_dir
         # schema
-        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0], backed=backed)
+        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(
+            self.filenames[0], backed=backed, cache_dir=cache_dir
+        )
         if len(adata0) != limits[0]:
             raise ValueError(
                 f"The number of cells in the first anndata file ({len(adata0)}) "
@@ -221,7 +232,7 @@ def __init__(
         self.schema = AnnDataSchema(adata0, obs_columns_to_validate)
         # lazy anndatas
         lazy_adatas = [
-            LazyAnnData(filename, (start, end), self.schema, self.cache, backed=backed)
+            LazyAnnData(filename, (start, end), self.schema, self.cache, backed=backed, cache_dir=cache_dir)
             for start, end, filename in zip([0] + limits, limits, self.filenames)
         ]
         # use filenames as default keys
@@ -316,10 +327,14 @@ def __getstate__(self):
     def __setstate__(self, state):
         self.__dict__.update(state)
         self.cache = LRU(self.max_cache_size)
-        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(self.filenames[0], backed=self.backed)
+        adata0 = self.cache[self.filenames[0]] = read_h5ad_file(
+            self.filenames[0], backed=self.backed, cache_dir=self.cache_dir
+        )
         self.schema = AnnDataSchema(adata0, self.obs_columns_to_validate)
         self.adatas = [
-            LazyAnnData(filename, (start, end), self.schema, self.cache, backed=self.backed)
+            LazyAnnData(
+                filename, (start, end), self.schema, self.cache, backed=self.backed, cache_dir=self.cache_dir
+            )
             for start, end, filename in zip([0] + self.limits, self.limits, self.filenames)
         ]
         self.obs_names = pd.Index([f"cell_{i}" for i in range(self.limits[-1])])
@@ -367,6 +382,7 @@ def __init__(
         schema: AnnDataSchema,
         cache: LRU | None = None,
         backed: backed_mode_type = backed_mode_default,
+        cache_dir: str | None = str(GCS_CACHE_DIR),
     ):
         self.filename = filename
         self.limits = limits
@@ -374,6 +390,7 @@ def __init__(
         if backed not in allowed_backed_modes:
             raise ValueError(f"Invalid backed mode: {backed}. Choose from {allowed_backed_modes}")
         self.backed = backed
+        self.cache_dir = cache_dir
         if cache is None:
             cache = LRU()
         self.cache = cache
@@ -414,7 +431,7 @@ def adata(self) -> AnnData:
             adata = self.cache[self.filename]
         except KeyError:
             # fetch anndata
-            adata = read_h5ad_file(self.filename, backed=self.backed)
+            adata = read_h5ad_file(self.filename, backed=self.backed, cache_dir=self.cache_dir)
             # validate anndata
             if self.n_obs != adata.n_obs:
                 raise ValueError(
diff --git a/cellarium/ml/data/fileio.py b/cellarium/ml/data/fileio.py
index c68d5d3b..888367f7 100644
--- a/cellarium/ml/data/fileio.py
+++ b/cellarium/ml/data/fileio.py
@@ -1,11 +1,13 @@
 # Copyright Contributors to the Cellarium project.
 # SPDX-License-Identifier: BSD-3-Clause
 
+import contextlib
 import os
 import re
 import shutil
 import tempfile
 import urllib.request
+from pathlib import Path
 from typing import Literal
 
 from anndata import AnnData, read_h5ad
@@ -16,11 +18,18 @@
 backed_mode_type = Literal["r"] | bool | None
 backed_mode_default: backed_mode_type = False
 
+# Default persistent cache for GCS-downloaded h5ad files. Shards are written
+# here on first access and reused on subsequent accesses or after a checkpoint
+# restart on the same machine. Pass cache_dir=None to stream directly from GCS
+# with no local disk usage.
+GCS_CACHE_DIR: Path = Path.home() / ".cache" / "cellarium_gcs_cache"
+
 
 def read_h5ad_gcs(
     filename: str,
     storage_client: Client | None = None,
     backed: backed_mode_type = backed_mode_default,
+    cache_dir: Path | str | None = GCS_CACHE_DIR,
 ) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from the Google Cloud Storage.
@@ -34,6 +43,10 @@ def read_h5ad_gcs(
         backed: See :func:`anndata.read_h5ad` for details on backed mode.
             ['r', True] will load in backed mode instead of fully loading into memory.
             [False, None] will use in-memory mode.
+        cache_dir: Directory for caching downloaded files on local disk. On first
+            access the shard is saved here; subsequent accesses read from disk instead
+            of re-downloading. If a write fails (e.g. disk full), falls back to
+            streaming from GCS. Set to ``None`` to always stream with no disk usage.
     """
     if not filename.startswith("gs:"):
         raise ValueError("The filename must start with 'gs:' protocol name.")
@@ -50,16 +63,36 @@ def read_h5ad_gcs(
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(blob_name)
 
+    if cache_dir is not None:
+        local_path = Path(cache_dir) / bucket_name / blob_name
+        if local_path.exists():
+            return read_h5ad(str(local_path), backed=backed)
+        # Cache miss — download to a staging file then rename atomically so that
+        # a concurrent worker or an interrupted run never sees a partial file.
+        local_path.parent.mkdir(parents=True, exist_ok=True)
+        staging = local_path.with_suffix(local_path.suffix + ".download")
+        try:
+            with open(staging, "wb") as f:
+                blob.download_to_file(f)
+                f.flush()
+            staging.rename(local_path)
+            return read_h5ad(str(local_path), backed=backed)
+        except OSError:
+            # Disk full or permission error — discard the staging file and fall
+            # through to the no-cache path below.
+            with contextlib.suppress(OSError):
+                staging.unlink()
+
+    # No cache (or cache write failed) — stream without leaving anything on disk.
     if backed not in [True, "r"]:
-        # Stream directly into memory — no temp file, no disk I/O.
         with blob.open("rb") as f:
             return read_h5ad(f)
 
-    # Backed mode requires h5py to have a real seekable file path. The file is
-    # flushed and closed before h5py opens it to avoid Python's write buffer
-    # leaving the last chunk off disk (which causes h5py to report a truncated
-    # file). After os.unlink the directory entry is gone but the inode stays
-    # allocated until h5py closes its fd (i.e. when the AnnData is GC'd).
+    # Backed mode without a persistent cache: download to an anonymous temp file.
+    # Flushed and closed before h5py opens it to avoid truncation from an
+    # unflushed write buffer. The unlink removes the directory entry immediately;
+    # on Linux/macOS the inode persists until h5py closes its fd (when the
+    # AnnData is GC'd or evicted from the LRU cache).
     with tempfile.NamedTemporaryFile(suffix=".h5ad", delete=False) as tmp_file:
         temp_path = tmp_file.name
         blob.download_to_file(tmp_file)
@@ -68,10 +101,8 @@ def read_h5ad_gcs(
     try:
         return read_h5ad(temp_path, backed=backed)
     finally:
-        try:
+        with contextlib.suppress(OSError):
             os.unlink(temp_path)
-        except OSError:
-            pass
 
 
 def read_h5ad_url(filename: str, backed: backed_mode_type = backed_mode_default) -> AnnData:
@@ -136,7 +167,12 @@ def read_h5ad_local(filename: str, backed: backed_mode_type = backed_mode_defaul
     return read_h5ad(filename, backed=backed)
 
 
-def read_h5ad_file(filename: str, backed: backed_mode_type = backed_mode_default, **kwargs) -> AnnData:
+def read_h5ad_file(
+    filename: str,
+    backed: backed_mode_type = backed_mode_default,
+    cache_dir: Path | str | None = GCS_CACHE_DIR,
+    **kwargs,
+) -> AnnData:
     r"""
     Read ``.h5ad``-formatted hdf5 file from a filename.
 
@@ -145,14 +181,17 @@ def read_h5ad_file(filename: str, backed: backed_mode_type = backed_mode_default
         backed: See :func:`anndata.read_h5ad` for details on backed mode.
             ['r', True] will load in backed mode instead of fully loading into memory.
             [False, None] will use in-memory mode.
+        cache_dir: Directory for caching GCS-downloaded files on local disk.
+            Only used when ``filename`` starts with ``gs://``. Set to ``None`` to
+            stream GCS files directly into memory with no disk usage.
     """
     if filename.startswith("gs:"):
-        return read_h5ad_gcs(filename, **kwargs)
+        return read_h5ad_gcs(filename, backed=backed, cache_dir=cache_dir, **kwargs)
 
     if filename.startswith("file:"):
         return read_h5ad_local(filename, backed=backed)
 
     if any(filename.startswith(scheme) for scheme in url_schemes):
-        return read_h5ad_url(filename)
+        return read_h5ad_url(filename, backed=backed)
 
     return read_h5ad(filename, backed=backed)