Sweep jobs/multi run jobs

[bruAristimunha]

I noticed that Exca has a very nice Asynchronous job configuration that lets you run and submit multiple jobs in parallel within submitit.

I think the API is really cool, and many thanks for pushing this.

However, I'd like to know if there's a better way to parallelize the parameters in each job.

Something very similar to the multirun/joblib plugin within Hydra, which would allow me to extract more from my GPU. I designed something that allows me to use Exca and Hydra, but I feel like I'm going against Exca's philosophy.

My jobs can be very small sometimes, consuming only 2GB, and I have a hard limit on how many jobs I can do, especially considering my GPU usage per hour, so a new job isn't very appealing to me. Do you have any suggestions in this context, for those of us with limited GPUs?

[bruAristimunha]

any input here @jrapin? 🙏🏽

[jrapin]

Hello
Not sure to understand exactly what is your use case here can you clarify? (eg, with TaskInfra or MapInfra?)

[kingjr]

@bruAristimunha glad to see you're giving exca a try.

I think you're looking for MapInfra's max_jobs and min_samples_per_job

[PierreGtch]

Thanks @kingjr for the pointer this is helping us!
(I misread the doc and thought that a method decorated by MapInfra.apply() had to receive Iterable[int] necessary)

Is this how you would typically use MapInfra for this purpose? You do you have something else in mind?

from pathlib import Path
from collections.abc import Iterable

from pydantic import BaseModel
from exca import TaskInfra, MapInfra
from exca.task import LocalJob
from submitit import Job


class ModelTrainer(BaseModel):
    seed: int = 0
    infra: TaskInfra = TaskInfra(
        cluster=None,
        folder=Path("~/.cache/exca").expanduser(),
    )

    @infra.apply()
    def train(self) -> None:
        pass  # training look for one small model (multiple would fit in parallel on the same GPU)


class WithinNodeParallelTrainer(BaseModel):
    infra: MapInfra = MapInfra(
        cluster="slurm",
        folder=None,
        min_samples_per_job=3,
    )

    @infra.apply(item_uid=lambda item: item.infra.uid())
    def train(self, items: Iterable[ModelTrainer]) -> Iterable[Job | LocalJob]:
        for item in items:
            yield item.infra.job()


trainer_configs = [ModelTrainer(seed=i) for i in range(6)]
parallel_trainer = WithinNodeParallelTrainer()

# this should request 2 nodes, and launch 3 training jobs in parallel on each node:
jobs = parallel_trainer.train(trainer_configs)

# this should wait for all the training jobs to finish:
trained_models = [job.result() for job in jobs]

[PierreGtch]

A limitation I see with this implementation is that WithinNodeParallelTrainer and ModelTrainer can not share the same cache.

It could be interesting to have a helper class that allows easily parallelising multiple TaskInfras on the same node, while using the original TaskInfra caches. Maybe something like this:

class ParallelTask(BaseModel):
    infra: MapInfra

    @infra.apply(
        item_uid=lambda item: item.infra.uid(),
        cache_type="exca",
    )
    def run(self, items: Iterable[TaskInfra]) -> Iterable[Any]:
        jobs = [infra.job() for infra in items]
        for job in jobs:
            yield job.result()

[kingjr]

In this case, I would recommend doing

trainer = ModelTrainer(infra={"folder": tmp_path})

with trainer.infra.job_array() as tasks:  
    for seed in range(6):
        task = task.infra.clone_obj({"seed": seed})
        tasks.append(task)

[PierreGtch]

With your solution, I would expect all tasks to be submitted in different jobs.

The interest of the map infra was to have multiple trainers running in parallel in each job. The use case would look like this:

trainer = ModelTrainer(infra={"folder": tmp_path})
trainers = [trainer.infra.clone_obj({"seed": seed}) for seed in range(6)]

paralleliser = ParallelTask(infra={"cluster": "slurm", "min_samples_per_job": 3})

paralleliser.run(trainers)  # this should only launch two slurm jobs, not six

But the cache created/read via ParallelTask would ideally be at the same location and use the same file structure as through your solution (a regular TaskInfra)

[kingjr]

Ok, thanks for clarifying, yes you could do a two class approach in this case.

[bruAristimunha]

Hey @jrapin, @jrapin and @hubertjb!

At the end, neither solution worked as expected, and we implemented this on our fine-tune paper, @PierreGtch can provide more input.

https://github.com/braindecode/OpenEEGBench/blob/75971ba03bd9adccce5021d39dca2b2b611a3ee9/open_eeg_bench/helpers.py#L24-L53

And looking at this function here at the neural train, I feel that you are also trying to find a way to run a grid into a "single" job.

https://github.com/facebookresearch/neuroai/blob/722555a64a3f84cd2d422eba4d1271f403197036/neuraltrain-repo/neuraltrain/utils.py#L117

Can we create something more robust here at Exca?

In the context of most academics or people without multiple small GPUs, we really need to use them in parallel within a single job. Usually, we can get one big GPU at once (h100), not multiple GPUs; it is really a block. In my experience with Jeanzay (France) and Access (USA), it is really necessary.

happy to discuss alternatives too!