Cell segmentation parallelization

[quentinblampey]

For efficiency, cell segmentation must be parallelized. This is done automatically using the Snakemake pipeline, but it's possible to use your own parallelization mechanism if you prefer using the CLI or the API.

Note that Snakemake may require setting up a profile. There is already an existing one for Slurm clusters, but some other profiles can be added in the future (don't hesitate to ask for it!).

You can ask questions here, or show your parallelization technique to help other users.

[Marius1311]

An easy, and maybe a bit naive, way of doing this with dask through the python API would be

from dask.distributed import Client
from dask import delayed
import dask

# Initialize a Dask client to manage distributed workers
client = Client(n_workers=16)  # Adjust based on the number of patches you have. 

# Define a delayed function to process each patch
@delayed
def process_patch_delayed(patch_index):
    segmentation.write_patch_cells(cellpose_temp_dir, patch_index)

# Wrap patch processing in delayed objects
tasks = [process_patch_delayed(patch_index) for patch_index in range(len(sdata["sopa_patches"]))]

# Compute in parallel using Dask
dask.compute(*tasks)
client.close()

Works well in one example case, I haven't tested this extensively. Also, I could not get this to work with Baysor yet - it always runs in sequential manner. You can use the Client to monitor performance trough the dask dashboard.

[quentinblampey]

Thanks @Marius1311, it looks very interesting, I'll try it one day!

In the future, I may try to benchmark different parallelization techniques, and add the best techniques in the API

[Marius1311]

Sounds great! Have you tried parallelizing Baysor through the Python API as well? I could not get that to run yet - whatever I tried, it ended up using a single process only.

[quentinblampey]

No, I never tried, I always parallelize Baysor using multiple jobs on a Slurm cluster.
But I think this is a known issue of Bayser, and if I remember correctly, its maintainer may be working on this!

[lguerard]

I managed to get Baysor running multithreaded (as far as I understand it though...) with these lines.

def run_baysor(x, baysor_exec_path, baysor_temp_dir, thread_number):
    import os
    command = (
        f"{baysor_exec_path} run --save-polygons GeoJSON -c {baysor_temp_dir}/{x}/config.toml {baysor_temp_dir}/{x}/transcripts.csv -o {baysor_temp_dir}/{x} :cell",
    )
    joined_command = " && ".join(command)
    os.system(joined_command)

import sys
from functools import partial 
from p_tqdm import p_umap

if sys.platform == "win32" and os.cpu_count() > 61:
    num_threads = 61
else:
    num_threads = None

result_objects = p_umap(partial(run_baysor, baysor_exec_path = BAYSOR_EXECUTABLE_PATH, baysor_temp_dir = baysor_temp_dir, thread_number = num_threads), valid_indices, num_cpus=num_threads)

On the other side, your code works for me but when running it in a notebook cell, it sends an error message that it's OOM on the GPU (single GPU, 11GB VRAM), the cell is considered as finished executing but processes are still happening and new parquet files get created. Any ideas on how to improve this ? Sorry I don't know much about dask...

[quentinblampey]

I'm working on sopa==2.0.0, which will significantly improve and ease the API usage. It will notably include some parallelization backends, such as dask (adapted from @Marius1311 comment), and maybe others in the future.

Example usage:

sopa.setting.parallelization_backend = "dask" # sets globally the backend

sopa.segmentation.cellpose(sdata, "DAPI", diameter=35) # runs cellpose on all patches in parallel

What do you think?

[Marius1311]

Looks great!

[stergioc]

Hi @quentinblampey, just jumping here quickly to mention that such functionality would be also very useful on other api functions like: sopa.patches.infer.infer_wsi_patches.

[quentinblampey]

Thanks @stergioc, I can have a look! It should be very easy to change, since internally we only need to do sopa.settings._run_with_backend(functions), where functions is a list of callable.