# the yaml config file defines which the different combinations of parameters
# that will be used to fill a template file
# at its core, it is simply do the cartesian product of all the parameters and list of possible values defined for them.
# each instance of the product will define a single sbatch script, i.e.
# a single job. all the variables defined will be replaced
# with their value in the template (here, `template.sbatch`)

# there are some special variables that are used by the job manager:

# Path to the sbatch template file, this is the basic squeleton of all sbatch files
# where variables to be replaced are written as {NAME} (see Step 1)
template: template.sbatch 

# Path of the standard output file, it is important as it is used for checking
# if the job is frozen (if no change in during `check_interval_secs` secs)
# and to find if the termination string (`termination_str`) appeared in the output file, 
# this is used to stop from restarting the job forever (default behavior).
# Remember that we have a max time limit in SLURM (e.g., usually 24h), 
# so we restart the job as much as needed until we find the `termination_str`.
output_file: "{logs}/{name}/slurm.out"

# It is IMPORTANT to define the `termination_str`, it is a regexp used to detect
# if a job is finished, otherwise, it will be restarted FOREVER.
# Here, for instance, we detect a finishing job if it finishes the zero-shot 
# evaluation the latest epoch.
# ({epochs} will take the value of epochs, see section experiments below).
termination_str: "Eval Epoch: {epochs}"

# an alternative is to use `termination_cmd`, where instead a shell command
# is executed, if it returns the value 1, the job is considered as finished.
termination_cmd: ""

# one can also have start condition, where the job is launched only
# under some constraint. This can be the case for evaluations, for instance,
# as they require that checkpoints of the models do exist beforehand.
# Here, we execute the shell command 'start_condition_cmd', if it returns
# the value 1, the job is launched.
start_condition_cmd: ""

# Path of sbatch scripts that are generated from the `template`
# each experiment will have a dedicated sbatch script.
sbatch_script: "sbatch/{name}.sbatch"

# Command to run for each job.
cmd: "sbatch {sbatch_script}"

# Check the status jobs each number of secs, to restart them if needed
check_interval_secs: 600

# Each experiment will have a UNIQUE name, which we can define in any way
# we want. 
# it will be used in the template (`template.sbatch` here) but also to make 
# the sbatch script name.
# `name` is a crucial parameter. It is used to uniquely identifiy each job, and
# to handle resuming autoexperiment in a new session. When autoexperiment process
# fails for some reason, but the SLURM jobs are still running, if we relaunch
# autoexperiment it will automatically recover the running jobs in the SLURM queue
# by assigning the the job to the SLURM job with the same `name` value.
# IMPORTANT: thus, `name` has to enforce two constraints:
# - 1) it has to be unique
# - 2) SLURM job name should be exactly as `name` of the job. So, it is NECESSARY in in the sbatch template to do `#SBATCH --job-name {name}` to have the correct resuming behavior
name: "{dataset}_{model}_{epochs}"

# Above were special variables.
# Next, we define variables that can be used in the sbatch template.
# These can be named anything, and can be nested.

dataset:
  - datacomp:
      train_data: "/path/{0000000..0139827}.tar"
  - laion2b:
      train_data: "/path/{00000..23295}.tar"
model_scale:
  - s32:
      model: ViT-S-32
      batch_size: 1024
  - m32:
      model: ViT-M-32
      batch_size: 1024
epochs: 1 
logs: "logs"
nodes: 1
train_num_samples: [12_800_000]