Codecov Report

Attention: Patch coverage is 79.22948% with 124 lines in your changes missing coverage. Please review.

Project coverage is 80.10%. Comparing base (c97655d) to head (d80ff0f).

@@            Coverage Diff             @@
##             main    #2932      +/-   ##
==========================================
- Coverage   80.14%   80.10%   -0.05%     
==========================================
  Files         193      194       +1     
  Lines       17121    17604     +483     
==========================================
+ Hits        13722    14101     +379     
- Misses       3399     3503     +104     

Hi Ori,
I gave it a try to this branch. First, I needed to install psutil. Second, when I ran the tutorial, I found an error:

Code:

model = scvi.model.SCVI(adata=None, registry=datamodule.registry, datamodule=datamodule)
Traceback (most recent call last):
File "", line 1, in
File "/allen/programs/celltypes/workgroups/rnaseqanalysis/Mariano/Anaconda3/envs/scvi-env-largedata/lib/python3.12/site-packages/scvi/model/_scvi.py", line 184, in init
self.module = self._module_cls(
^^^^^^^^^^^^^^^^^
TypeError: VAE.init() got an unexpected keyword argument 'datamodule'

I fixed it.

In the tutorial file: scvi-tools/docs/user_guide/use_case /custom_dataloaders.md

Replace the line:
model = scvi.model.SCVI(adata=None, registry=datamodule.registry, datamodule=datamodule)

with:
model = scvi.model.SCVI(adata=None, registry=datamodule.registry)

@marianogabitto thanks. I fixed those things.
Also see the test_custom_dataloder.py file and the corresponding tutorial on this to know the full capabilities we have currently on c.dataloders (scverse/scvi-tutorials#425).

As you also saw, my concern right now is the slowness of the process, in the data loading between batches. Let's focus on this.
Did you find/ have any suggestions on how we could make it faster?

feel free to fork the branch and add your commits.

Hi Ori,
I am adding the point that I brought in discourse: I compared head to head the TileDB and the regular anndata loaders and Tile is 50% slower. I tested the TileDB dataloader ir regular and DDP mode and what is causing the delay is the slow access to data. GPU peaks and process super fast but in between batches there is a long waiting time

Let me tell you how I did this comparison. On the one hand, I save the TileDB experiment as an anndata and I run scvi regularly. On the other hand, I grab the data loader code from scvi (splitting, Anndataloader, AnnDataset) and create an scvi external AnnData loader (just to be sure there was no difference between running the anndata in regular scvi versus passing it as a new anndataloader).

2 questions:
1) If you can wait until Wednesday, I am going to talk to people in my High Performance Computer Cluster and representatives from TileDB and try to debug if there is any way to accelerate the data loader with multiple workers or options. It is expected that the data loader that access disk is slower, I want to see how slow.
2) what is the data loader that is used in regular scvi ? Should I also use the following? wouldn't the train dataloder only have an amount of data equivalent to train_size ?
inference_dataloader = (
inference_datamodule.on_before_batch_transfer(batch, None)
for batch in inference_datamodule.train_dataloader()
)

M

Just checking is it still faster if you load the AmnData in disk backed mode? This would indeed be surprising while the other overhead could come from loading from disk? Any chance to use a fast SSD storage of the data as usually recommended for from-disk loading. Data could also be scattered across SSDs and there tileDB might have suggestions how to optimize this - not perfect randomness is not really an issue if they first randomize order on disks (not a single experiment on one SSD).

Speeds are:
anndata in memory > anndata backed >~ tileDB .

yes, I am testing two things: I am talking with my support team about a partition with fast disk access and touching base with tiled about how to optimize this.

@marianogabitto I will try to do some testing on my end.
Actually, the whole part of "on_before_batch_transfer" was something that I inherited from the census implementation long ago, but I never had the time to check it fully (see https://github.com/chanzuckerberg/cellxgene-census/blob/756708e9aa18791b7bae3712e9dd66d2b6ce9d75/api/python/notebooks/experimental/pytorch_scvi.ipynb)
and it makes sense to me, it might be the bottleneck.
We are not in a rush to release it. we can wait for the inputs from support team.

But you need the "on_before_batch" to move the batch to PyTorch, no ? isn't the reason of that callback to get the X, batch, labels data and to move it to tensors? what happens if they are already tensors?

One more... how about making the anndata_manager and the registry tutorial ?

We will deal with everything, but let's take it step by step. on_before_batch is needed of course, but lets see how we can improve it.

Your suggestion worked and I do see much better performance in train time when using

datamodule.setup()
model = scvi.model.SCVI(
    adata=None,
    registry=datamodule.registry,
    n_layers=n_layers,
    n_latent=n_latent,
    gene_likelihood="nb",
    encode_covariates=False,
)
# creating the dataloader for trainset
training_dataloader = (
    datamodule.on_before_batch_transfer(batch, None)
    for batch in datamodule.train_dataloader()
)
import time
start = time.time()
model.train(
    datamodule=training_dataloader,
    #datamodule=datamodule,
    max_epochs=100,
    batch_size=1024,
    # accelerator="gpu",
    # devices=-1,
    # strategy="ddp_find_unused_parameters_true",
)
end = time.time()
print(f"Elapsed time: {end - start:.2f} seconds")

See my updated tutorial on the other branch

Ori, I got lost ... what branch should I look for the tutorial? What branch should I check for testing your commits?

This are the tutorials:
https://app.reviewnb.com/scverse/scvi-tutorials/pull/425/

tests are in the current branch

Ori, this is not working for me. When I invoke in the notebook:
training_dataloader = (
datamodule.on_before_batch_transfer(batch, None) for batch in datamodule.train_dataloader()
)
I get:
switching torch multiprocessing start method from "fork" to "spawn"
and then errors out....

Ori, all the examples that I am listing below are run by removing the code ".on_before_batch_transfer()". The way I posted before.

1. When num_workers=0, I can train with low speeds, defined as below. When I fix number of workers, like num_workers=4,12 or 24. The trainer takes forever to initialize and then is even slower than below.

2. Can you monitor your GPU usage with nvitop or nvtop? Let me tell you my head-to-head comparisons.

• TIleDB from cell census. I believe that this is reading from S3, so it is never actually copy data to disk.
  It takes 120 sec/it to train. I see GPU activity almost zero all the time except at moments when it picks to 100%.

• TIleDB from anndata created from the query . This is reading from a local disk directory.
  It takes 11 sec/it to train. I see GPU activity almost zero all the time except at moments when it picks to 100%.

• Regular way of loading anndata into memory. It takes 1.2 sec/it to train. I see GPU activity at 40% all the time.

These led me to believe that we are not loading data into GPU memory fast enough.

3. I forgot to tell you but the TileDB representative send me this as reference. It is different from the way we run because they launch the processes.
   https://github.com/single-cell-data/TileDB-SOMA-ML/tree/rw/cli/src/tiledbsoma_ml/cli#example-invocation