smartsim-driver.py

### Note: This is a modified version of provided driver code. This code has not been tested, and it is highly probable that won't work. However it can be used as a guide on how to incorporate MOM6 with SmartSim to run on a single machine (single server or a laptop)
### Note: Do not forget to change MOM6_EXE value

"""
Driver script for the execution of SmartSim enabled
MOM6 simulations with the OM4_05 configuration.

Note: certain parameters and calls will need to
be updated depending on the resources of the
system.

This script assumes launching on a slurm cluster
with at least
   - 228 CPU nodes with 96 cpus (including hyperthreads)
   - 16 nodes with P100 GPUs and 36 cpu cores
This can be changed to suit your system with the parameters
listed below

To run the exact same experiment as our paper, increase
the time in both batch jobs and the number of days
to 10 years.
"""

from glob import glob
from smartsim import Experiment
from smartsim.database import Orchestrator
from smartsim.settings import RunSettings


# experiment parameters
MOM6_EXE = "/bin/echo"  # TODO: This should be replaced with the actual path to the binary of MOM6

# name of experiment where output will be placed
experiment = Experiment("AI-EKE-MOM6", launcher="local")


mom_settings = RunSettings(MOM6_EXE)

# define batch parameters for entire ensemble
# batch_opts = {
#     "mincpus": 96,
#     "ntasks-per-node": 48,
#     "exclusive": None
# }
# ensemble_batch = SbatchSettings(nodes=ENSEMBLE_NODES, time="10:00:00", batch_args=batch_opts)

mom_ensemble = experiment.create_ensemble("MOM", run_settings=mom_settings, replicas=1)

# Attach input files and configuration files to each
# MOM6 simulation
mom_ensemble.attach_generator_files(
    to_configure=glob("../MOM6_config/configurable_files/*"),
    to_copy="../MOM6_config/OM4_025",
    to_symlink="../MOM6_config/INPUT",
)

# configs to write into 'to_configure' files listed
# above. If you change the number of processors for
# each MOM6 simulation, you will need to change this.
MOM6_config = {
    "SIM_DAYS": 1, # length of simlations 
    "DOMAIN_LAYOUT": "32,36",
    "MASKTABLE": "mask_table.242.32x36"
    }
for model in mom_ensemble:
    model.params = MOM6_config

# register models so keys don't overwrite each other
# in the database
for model in mom_ensemble:
    model.register_incoming_entity(model)

# creation of ML database specific to Slurm.
# there are also PBS, Cobalt, and local variants
db = Orchestrator()

# generate run directories and write configurations
experiment.generate(mom_ensemble, db, overwrite=True)

# start the database and ensemble batch jobs.
experiment.start(mom_ensemble, db, summary=True)

# print a summary of the run.
print(experiment.summary())