cancel_experiment.py

import argparse

import utilities.generic

import models.nemo.utilities.namelist
import models.nemo.utilities.experiment
import models.nemo.utilities.experiment_plan
import models.nemo.utilities.argument_parser
import models.nemo.utilities.global_information


#############
# functions #
#############

def cancel_experiment(version, configuration, name, keep_run_directories):
  '''
  Cancel all running and submitted legs of an experiment.
  '''

  print('Canceling experiment "%s", in configuration "%s", in version "%s"...\n'%(name,
    configuration, version))

  # create the experiment object
  experiment = models.nemo.utilities.experiment.Experiment(version=version,
    configuration=configuration, name=name)
  experiment_plan = experiment.get_experiment_plan(for_modification=True)

  # cancel the experiment legs
  number_of_canceled_legs = 0

  for leg in experiment_plan.legs:
    if leg.status in (
      models.nemo.utilities.experiment_plan.EXPERIMENT_LEG_STATUS_SUBMITTED,
      models.nemo.utilities.experiment_plan.EXPERIMENT_LEG_STATUS_RUNNING,
    ):
      leg.cancel(delete_run_directory=(not keep_run_directories))
      number_of_canceled_legs += 1

  if number_of_canceled_legs > 0:
    print()

  # write the updated experiment plan
  experiment_plan.rewrite()


#################
# main function #
#################

def main():

  print()

  parser = argparse.ArgumentParser(description=cancel_experiment.__doc__,
    formatter_class=argparse.RawTextHelpFormatter)
  models.nemo.utilities.argument_parser.add_argument_version(parser=parser, required=False)
  models.nemo.utilities.argument_parser.add_argument_configuration(parser=parser, required=False)
  models.nemo.utilities.argument_parser.add_argument_name(parser=parser)
  parser.add_argument(
    '-krd', '--keep_run_directories',
    action='store_true',
    help='keep the leg run directories upon cancellation\n(default: False)',
  )
  arguments = vars(parser.parse_args())

  try:
    arguments['version'], arguments['configuration'] = (
      models.nemo.utilities.global_information.get_version_and_configuration(
      tentative_version=arguments['version'], tentative_configuration=arguments['configuration']))
  except utilities.generic.ClimateToolboxException as e:
    e.display()

  try:
    cancel_experiment(**arguments)
  except utilities.generic.ClimateToolboxException as e:
    e.display()


if __name__ == '__main__':
  main()