Package: slurm-sweeps
slurm sweepsA simple tool to perform parameter sweeps on SLURM clusters.The main motivation was to provide a lightweightASHA implementationforSLURM clustersthat is fully compatible withpytorch-lightning's ddp.It is heavily inspired by tools likeRay TuneandOptuna.
However, on a SLURM cluster, these tools can be complicated to set up and introduce considerable overhead.Slurm sweepsis simple, lightweight, and has few dependencies.
It uses SLURM Job Steps to run the individual trials.Installationpip install slurm-sweepsDependenciescloudpicklenumpypandaspyyamlUsageYou can just run this example on your laptop.
By default, the maximum number of parallel trials equals the number of CPUs on your machine.""" Content of test_ss.py """fromtimeimportsleepimportslurm_sweepsasss# Define your train functiondeftrain(cfg:dict):forepochinrange(cfg["epochs"]):sleep(0.5)loss=(cfg["parameter"]-1)**2*epoch# log your metricsss.log({"loss":loss},epoch)# Define your experimentexperiment=ss.Experiment(train=train,cfg={"epochs":10,"parameter":ss.Uniform(0,2),},asha=ss.ASHA(metric="loss",mode="min"),)# Run your experimentdataframe=experiment.run(n_trials=1000)# Your results are stored in a pandas DataFrameprint(f"\nBest trial:\n{dataframe.sort_values('loss').iloc[0]}")Or submit it to a SLURM cluster.
Write a small SLURM scripttest_ss.slurmthat runs the code above:#!/bin/bash -l#SBATCH --nodes=2#SBATCH --ntasks-per-node=18#SBATCH --cpus-per-task=4#SBATCH --mem-per-cpu=1GBpythontest_ss.pyBy default, this will run$SLURM_NTASKStrials in parallel.
In the case above:2 nodes * 18 tasks = 36 trialsThen submit it to the queue:sbatch test_ss.slurmSee thetestsfolder for an advanced example of training a PyTorch model with Lightning's DDP.API DocumentationCLASS slurm_sweeps.Experimentdef__init__(self,train:Callable,cfg:Dict,name:str="MySweep",local_dir:Union[str,Path]="./slurm_sweeps",backend:Optional[Backend]=None,asha:Optional[ASHA]=None,restore:bool=False,overwrite:bool=False,)Set up an HPO experiment.Arguments:train- A train function that takes as input thecfgdict.cfg- A dict passed on to thetrainfunction.
It must contain the search spaces viaslurm_sweeps.Uniform,slurm_sweeps.Choice, etc.name- The name of the experiment.local_dir- Where to store and run the experiments. In this directory
we will create the databaseslurm_sweeps.dband a folder with the experiment name.backend- A backend to execute the trials. By default, we choose theSlurmBackendif Slurm is available,
otherwise we choose the standardBackendthat simply executes the trial in another process.asha- An optional ASHA instance to cancel less promising trials.restore- Restore an experiment with the same name?overwrite- Overwrite an existing experiment with the same name?Experiment.rundefrun(self,n_trials:int=1,max_concurrent_trials:Optional[int]=None,summary_interval_in_sec:float=5.0,nr_of_rows_in_summary:int=10,summarize_cfg_and_metrics:Union[bool,List[str]]=True,)->pd.DataFrameRun the experiment.Arguments:n_trials- Number of trials to run. For grid searches this parameter is ignored.max_concurrent_trials- The maximum number of trials running concurrently. By default, we will set this to
the number of cpus available, or the number of total Slurm tasks divided by the number of trial Slurm
tasks requested.summary_interval_in_sec- Print a summary of the experiment every x seconds.nr_of_rows_in_summary- How many rows of the summary table should we print?summarize_cfg_and_metrics- Should we include the cfg and the metrics in the summary table?
You can also pass in a list of strings to only select a few cfg and metric keys.Returns:A DataFrame of the database.CLASS slurm_sweeps.SlurmBackenddef__init__(self,exclusive:bool=True,nodes:int=1,ntasks:int=1,args:str="")Execute the training runs on a Slurm cluster viasrun.Pass an instance of this class to your experiment.Arguments:exclusive- Add the--exclusiveswitch.nodes- How many nodes do you request for your srun?ntasks- How many tasks do you request for your srun?args- Additional command line arguments for srun, formatted as a string.ContactDavid Carreto Fidalgo (david.carreto.fidalgo@mpcdf.mpg.de)