Include batching of multiple prediction jobs in a single SLURM job

[epkumpu]

Hi!

When running a large set of short predictions (only inference), a significant amount of time is used in SLURM queue on a cluster system. For example, I could be running many 15 min predictions, request 1h for each in SLURM, and wait in queue for each of them. OR I could request 2h of compute time and batch together 4-6 15 min predictions and waste no time in queue in between. Usually there's only a minor difference in queue time between 1h and 2h requests, so at least in our setup this would speed things up considerably.

Would it be possible to include this batching of multiple inference jobs into one SLURM job in snakemake? Maybe give this an explicit parameter to manually set batch size?

Cheers,
EP

[DimaMolod]

Thanks @epkumpu, that's a sensible request, and most of what's needed already exists.

Under the hood, this workflow calls AlphaPulldown's command-line tool, run_structure_prediction.py, once per fold. That tool already accepts multiple folds in a single run. You give it a list of inputs and matching output directories, and it loads the model once and predicts them one after another in the same process, writing each fold's results as it finishes. So batching would save both the SLURM queue wait and the repeated model-loading/compilation cost, and if a job dies partway through, the folds it has already finished are kept.

What's missing is only the Snakemake side, which currently submits one job per fold. There's even a half-finished start on this already: workflow/scripts/cluster_sequence_length.py groups folds by sequence length, and the prediction rule already knows how to handle several folds at once.

There are two ways to wire it up:

1. Snakemake's group directive is the simplest. Folds in a group share one SLURM allocation, so you wait in the queue once. But each fold still runs as its own process, so you don't save the model-loading cost.

2. Group folds by size into batches (finishing the cluster_sequence_length.py idea) - one SLURM job per size-bin, handing the whole list to the AlphaPulldown CLI in a single run. This gets both savings. The trade-off: the job must be sized (walltime/memory) for the largest fold in the bin, and if one fold fails, the whole bin reruns — grouping by similar size keeps that contained.

I'd lean toward option 2 with an explicit opt-in batch_size setting that defaults to 1 (today's behaviour).
I'll mark this as an enhancement for now. In the meantime, you can batch folds yourself by calling the AlphaPulldown CLI (run_structure_prediction.py) directly with a list of inputs and matching output directories.

[DimaMolod]

Hi @epkumpu
I implemented the prototype of the pipeline that submits predictions in batches for both AF2&AF3 backends:
https://github.com/KosinskiLab/AlphaPulldownSnakemake/tree/feature/batch-small-jobs
I would really appreciate it if you could clone/install and test it on your short inference use cases and let me know if it works as expected.

Many thanks,
Dima