Skip to content

Instantly share code, notes, and snippets.

@alitoufighi

alitoufighi/smartsim-driver.py

Created April 28, 2024 23:58
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save alitoufighi/48025d6e10d2215b8dad2ed5f784b77f to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save alitoufighi/48025d6e10d2215b8dad2ed5f784b77f to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. alitoufighi created this gist Apr 28, 2024.
    84 changes: 84 additions & 0 deletions smartsim-driver.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,84 @@
    ### 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())