scdataloader 2.0.2__tar.gz → 2.0.3__tar.gz

{scdataloader-2.0.2 → scdataloader-2.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scdataloader
-Version: 2.0.2
+Version: 2.0.3
 Summary: a dataloader for single cell data in lamindb
 Project-URL: repository, https://github.com/jkobject/scDataLoader
 Author-email: jkobject <jkobject@gmail.com>

{scdataloader-2.0.2 → scdataloader-2.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "scdataloader"
-version = "2.0.2"
+version = "2.0.3"
 description = "a dataloader for single cell data in lamindb"
 authors = [
     {name = "jkobject", email = "jkobject@gmail.com"}
@@ -15,6 +15,7 @@ dependencies = [
     "cellxgene-census>=0.1.0",
     "torch>=2.2.0",
     "pytorch-lightning>=2.3.0",
+    "lightning>=2.3.0",
     "anndata>=0.9.0",
     "zarr>=2.10.0",
     "matplotlib>=3.5.0",
@@ -27,8 +28,6 @@ dependencies = [
     "django>=4.0.0",
     "scikit-misc>=0.5.0",
     "jupytext>=1.16.0",
-    "lightning>=2.3.0",
-    "pytorch-lightning>=2.3.0",
 ]
 
 [project.optional-dependencies]

{scdataloader-2.0.2 → scdataloader-2.0.3}/scdataloader/datamodule.py

@@ -65,6 +65,7 @@ class DataModule(L.LightningDataModule):
         genedf: Optional[pd.DataFrame] = None,
         n_bins: int = 0,
         curiculum: int = 0,
+        start_at: int = 0,
         **kwargs,
     ):
         """
@@ -162,6 +163,7 @@ class DataModule(L.LightningDataModule):
         self.sampler_chunk_size = sampler_chunk_size
         self.store_location = store_location
         self.nnz = None
+        self.start_at = start_at
         self.idx_full = None
         self.max_len = max_len
         self.test_datasets = []
@@ -324,9 +326,9 @@ class DataModule(L.LightningDataModule):
                 len_test = self.test_split
             else:
                 len_test = int(self.n_samples * self.test_split)
-            assert len_test + len_valid < self.n_samples, (
-                "test set + valid set size is configured to be larger than entire dataset."
-            )
+            assert (
+                len_test + len_valid < self.n_samples
+            ), "test set + valid set size is configured to be larger than entire dataset."
 
             idx_full = []
             if len(self.assays_to_drop) > 0:
@@ -461,7 +463,7 @@ class DataModule(L.LightningDataModule):
             dataset = None
         else:
             dataset = Subset(self.dataset, self.idx_full)
-            train_sampler = RankShardSampler(len(dataset))
+            train_sampler = RankShardSampler(len(dataset), start_at=self.start_at)
         current_loader_kwargs = kwargs.copy()
         current_loader_kwargs.update(self.kwargs)
         return DataLoader(
@@ -492,8 +494,8 @@ class DataModule(L.LightningDataModule):
     def predict_dataloader(self):
         subset = Subset(self.dataset, self.idx_full)
         return DataLoader(
-            subset,
-            sampler=RankShardSampler(len(subset)),
+            self.dataset,
+            sampler=RankShardSampler(len(subset), start_at=self.start_at),
             **self.kwargs,
         )
 
@@ -667,7 +669,9 @@ class LabelWeightedSampler(Sampler[int]):
         unique_samples, sample_counts = torch.unique(sample_labels, return_counts=True)
 
         # Initialize result tensor
-        result_indices_list = []  # Changed name to avoid conflict if you had result_indices elsewhere
+        result_indices_list = (
+            []
+        )  # Changed name to avoid conflict if you had result_indices elsewhere
 
         # Process only the classes that were actually sampled
         for i, (label, count) in tqdm(
@@ -850,8 +854,9 @@ class RankShardSampler(Sampler[int]):
     """Shards a dataset contiguously across ranks without padding or duplicates.
     Preserves the existing order (e.g., your pre-shuffled idx_full)."""
 
-    def __init__(self, data_len: int):
+    def __init__(self, data_len: int, start_at: int = 0) -> None:
         self.data_len = data_len
+        self.start_at = start_at
         if torch.distributed.is_available() and torch.distributed.is_initialized():
             self.rank = torch.distributed.get_rank()
             self.world_size = torch.distributed.get_world_size()
@@ -859,9 +864,16 @@ class RankShardSampler(Sampler[int]):
             self.rank, self.world_size = 0, 1
 
         # contiguous chunk per rank (last rank may be shorter)
-        per_rank = math.ceil(self.data_len / self.world_size)
-        self.start = self.rank * per_rank
+        if self.start_at > 0:
+            print(
+                "!!!!ATTTENTION: make sure that you are running on the exact same \
+                    number of GPU as your previous run!!!!!"
+            )
+        print(f"Sharding data of size {data_len} over {self.world_size} ranks")
+        per_rank = math.ceil((self.data_len - self.start_at) / self.world_size)
+        self.start = int((self.start_at / self.world_size) + (self.rank * per_rank))
         self.end = min(self.start + per_rank, self.data_len)
+        print(f"Rank {self.rank} processing indices from {self.start} to {self.end}")
 
     def __iter__(self):
         return iter(range(self.start, self.end))