LLNL / merlin

Machine Learning for HPC Workflows
MIT License
118 stars 26 forks source link

example/stop_ensemble #458

Open dylancliche opened 9 months ago

dylancliche commented 9 months ago

A new example for having the workflow itself stop the workers assigned to it.

lucpeterson commented 4 months ago

@bgunnar5 what do you think about creating a custom stop worker command that also shuts down a worker’s batch job (and any dependencies)? We could register a custom command that workers can listen for

https://docs.celeryq.dev/en/latest/userguide/workers.html#id18

This would mean that $(MERLIN_STOP_WORKERS) could shut down batch jobs too and make it easy for any workflow to shut down its workers and allocation at the end

bgunnar5 commented 4 months ago

@lucpeterson yeah I think ultimately baking this into Merlin rather than having it as a standalone example is probably the way to go. This also means that I could expand it to work for more than just slurm too.

I'm not sure if this would need to be an entirely new custom command or if we could just add it on to the existing merlin stop-workers command.

lucpeterson commented 4 months ago

Yeah the stop workers command issues a broadcast call to shutdown workers remotely. If we made a new command that we could broadcast, this would slip in nicely. You could kill workers and batch jobs from either the command line or a workflow.

dylancliche commented 3 months ago

@bgunnar5 @lucpeterson, I personally think that having a new command in the form of something along the lines $(MERLIN_STOP_WORKFLOW) which accomplishes this example would be the best way to move forward. From a users perspective, adding a single line command to execute the scripts within this example would be a lot better than copying and pasting a whole step into a workflow.

dylancliche commented 3 months ago

@bgunnar5 Let me know if you would like some help in accomplishing this. I can work on this a little at a time if needed.

bgunnar5 commented 3 months ago

Ideally we could be able to do something like merlin stop-workers --all-dependencies if we create the custom celery command that Luc was linking. This would then lead us to be able to create a new step return variable like you mentioned @dylancliche (see the big if/elif chain of how results are processed in merlin/common/tasks.py::merlin_step for how this could be implemented; specifically see how workers are currently stopped when STOP_WORKERS is provided). With this implementation users could either invoke the stop command from the command line or from within their workflow.

We'll need to figure out how to register this custom command based on the scheduler so there will likely be some work to do within the script adapters themselves. Maybe each script adapter can have its own kill_dependent_jobs method (naming conventions up for debate on this) and that can be linked to the custom command? I'm not sure the best way to go about this as of now.

If you have time to work on it then I definitely won't stop you!

dylancliche commented 3 months ago

@bgunnar5 Thanks for pointing me to the right locations of where STOP_WORKERS takes place as well as the adapters. I will first take a look at using the adapters with this example to help make this scheduler independent. We'll go from there.