Add support for motioncor2/motioncor3's Serial Flag for large runs / batch multiple images into 1 command invocation

[jpellman]

motioncor2 is primarily I/O bound. In cursory/non-rigorous tests, close to around 1/3 of the total runtime involves loading the gain, dark, defect and input images. The total walltime for a motioncor2 run will decrease by around 1/3 by warming the page cache / pre-loading these images.

In order to make more efficient usage of our GPUs, we'd like to exploit this fact. Here's what we've tried so far:

• Slurm GRES Sharding: We can have two motioncor2 Dask workers warm the cache and each invocation of motioncor2 take a lock on a GPU before running to prevent more than one motioncor2 command from running on the GPU at a time (motioncor2 does not support running multiple motioncor2 processes on the same GPU simultaneously as discovered via testing). With the specific version of Slurm we use (20.11.7+really20.11.4-2+deb11u1 Debian 11), this does not work for the most part. We generally get this error on the nodes in the slurmd logs if we update gres.conf explicity:

fatal: Invalid GRES record for shard, count does not match File value

Someone else had the same issue, but there are no responses there. Without updating gres.conf, sharding appears to work, but the GPU device is not added to the batch job's cgroup (as confirmed via an interactive session via salloc). The Slurm docs mention the CUDA_VISIBLE_DEVICES environment variable being set, which makes me think that maybe cgroups aren't supported as an isolation mechanism when using shards.

• dask-jobqueue allows you to run 2 workers in parallel within the same Slurm job by setting the processes keyword arg in the SLURMCluster constructor (see here). I updated the app logic itself to warm the cache within the pipeline and run motioncor2 with a GPU lock. This actually worsened performance somehow (~30s to ~60s runtime) and one or two images did not process properly (possibly due to a bug I introduced).

Here's what we haven't tried, but should try:

motioncor2/motioncor3 provide built-in logic for overlapping I/O with computation so that processing is performed while the inputs for the next image to be processed are being loaded in the background. Specifically, there is a Serial flag that implements this. For this to work, we'd need to group together inputs with common gains/dark references/defect maps into subdirectories, ideally also imposing a batch size limit (maybe 10 images per subdirectory?). We would then ideally want to move outputs back to the main processing directory and then have the post-task link them to the Leginon directory, so we'd need some sort of map.

[jpellman]

Another possibility would be to implement a fused_motioncor function like in e83e288 (pretty sure there's an issue there with passing outputs to the post-task, so definitely not production-ready). Ultimately I think that something like that will add way too much complexity that we don't want for marginal benefit, so it's probably better to stick with the Serial flag.

[jpellman]

More thoughts:

• We could revise the pre-task to not save the dark, defect, and fmintfiles using a naming convention that's based on the name of the input file name. Instead these files could be named according to a record ID and/or a hash of parameters used to generate their contents. For instance, the defect map file could use the primary key of the corrector plan that it is derived from.
• This would require us to restructure the imgmetadata object to include the corrector plan ID (or the corrector plan data could be a dictionary that's nested within imgmetadata).
• If multiple tables are involved in generating an input (e.g., CameraEMData and CorrectorPlanData are both used to make defect maps) both primary keys could be incorporated into the file name (e.g., defectmap_correctorplan-12345_camera-12345.mrc). Command line arguments could also be added to the filenames.
• If an input file exists, the pre-task won't generate it.

[jpellman]

I don't really see a good reason to support a workflow where we use motioncor2 without the -Serial flag. If we only want to process 1 image at time, we can always set up the run directory so that the input folders only have one image in them, and that will effectively be the same.

Earlier today, I performed a few benchmarks on an 8 image input folder. I passed the input folder into motioncor2 with -Serial 1 set and then also processed each image individually in a for loop without -Serial 1. I also attempted to gauge the effects of warming the page cache by running cat ~/inputs/* > /dev/null both before and in parallel with the motioncor2 command (results in table below only reflect the motioncor2 command runtime; the rows with Before have longer total runtime since they don't include the time spent warming the cache). motioncor2 commands were run using a single NVIDIA Quadro RTX 5000.

Page Cache Warming	Serial Flag	For Loop	Time (s)
None	Yes	No	48
Before	Yes	No	38
Parallel	Yes	No	40
None	No	Yes	76
Before	No	Yes	67
Parallel	No	Yes	68

Overall, these results suggest that we can cut processing time in half by using the Serial flag with a function that warms up the page cache in parallel in a separate process.

Note: The page cache was cleared between runs with echo 1 > /proc/sys/vm/drop_caches.

[jpellman]

Another observation:

The -Serial flag kind of obviates the need for a task queue that runs on top of Slurm. When I first decided to incorporate a task queue into our motion correction jobs, my assumption was that a task queue worker would hold onto a GPU for an extended period of time, execute motioncor2 commands rapidly one after another and help us avoid the issue of lots of small Slurm jobs coming and going (and preempting) on the cluster. The -Serial flag does exactly what I had hoped a task queue worker would do, so there isn't really a good reason to not execute functions within individual Slurm jobs instead of as Dask worker tasks. In fact, switching away from dask-jobqueue would probably eliminate a whole lot of unnecessary complexity (e.g., spillover directory issues, workers not disappearing when they should) and make our lives easier. Maybe time to take another look at Parsl?

[jpellman]

Took a look at Parsl. It seems to be more oriented toward multicomputing / message passing jobs, which isn't really what we want.

I struggled to get it to not ask for whole node allocations (even when I got it to remove the --exclusive flag by updating the exclusive param in the Slurm provider, it still wouldn't work). I also couldn't find a way to get it to stop passing in the --nodes and --tasks-per-node flags to sbatch, since it rather annoyingly hard-codes in a template (I'd prefer it if I could specify the batch template myself). Eventually I was able to get jobs to work as expected by passing in the following (with core count explicitly specified to prevent impossible configurations):

provider = SlurmProvider(partition='preempt', nodes_per_block=1, init_blocks=1, min_blocks=0, exclusive=False, max_blocks=32, walltime="00:15:00", scheduler_options="#SBATCH --gres=gpu:1\n#SBATCH -c 4", mem_per_node=64, status_batch_size=32, worker_init="source /h2/jpellman/.pyenv/versions/parsl/bin/activate", launcher=SrunLauncher())
executor = HighThroughputExecutor(label="semc-cluster", address=address_by_interface('eth0'), provider=provider)
config = Config(executors=[executor])

With this particular setup, Slurm jobs don't seem to exit until they hit their walltime, which is a bit annoying but might be a feature for the high throughput MPI jobs that Parsl seems better suited for.

I tested out submitit and this seems to be way more in line with what we might want, and is also way simpler / less complex. Although its documentation is a bit lacking I was able to get it working pretty quickly with minimal pain. The main caveat (which also applies to dask-jobqueue and parsl using the config from above) is that it doesn't use array jobs (which increases load on the scheduler), but I'm not as concerned about this for our cluster since we're relatively small scale.

[jpellman]

The Serial flag seems to cause motioncor3 to segfault according to czimaginginstitute/MotionCor3#8. Might be worth keeping in mind.

It's also worth remembering that Serial only supports MRC files. For EERs and Tiffs we won't see any benefit (and indeed, Serial will just not work outright). It might be worth considering how many of our sessions use Tiffs/EERs vs MRCs, and whether or not the speed-up from using Serial is worth the added complexity of sorting images with the same input parameters into batches and then pushing those all through motioncor2 in a single run (vs invoking motioncor2 once for each image and just adding more GPUs to our cluster).

[jpellman]

Just looking at the frames currently on storage suggest that the breakdown by image format is roughly as follows:

Format	Percentage
EER	36.82%
MRC	1.45%
TIFF	61.74%

So it looks like we would only benefit from using Serial in a minority of cases.

Of course this is only one sample (current data is roughly around a month's worth of acquisitions based on our data lifecycle policy)- for a better idea of the breakdown we'd need to have a longer sample (all of the data we acquired throughout the year, for instance).