Source code for idmtools_platform_slurm.platform_operations.experiment_operations

Here we implement the SlurmPlatform experiment operations.

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
import os
import shutil
from pathlib import Path
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, List, Type, Dict, Optional, Any
from idmtools.assets import Asset, AssetCollection
from idmtools.core import EntityStatus
from idmtools.core import ItemType
from idmtools.entities import Suite
from idmtools.entities.experiment import Experiment
from idmtools.entities.iplatform_ops.iplatform_experiment_operations import IPlatformExperimentOperations
from idmtools_platform_slurm.platform_operations.utils import SlurmExperiment, SlurmSimulation, SlurmSuite, \
from logging import getLogger

logger = getLogger(__name__)
user_logger = getLogger('user')

    from idmtools_platform_slurm.slurm_platform import SlurmPlatform

[docs]@dataclass class SlurmPlatformExperimentOperations(IPlatformExperimentOperations): platform: 'SlurmPlatform' # noqa: F821 platform_type: Type = field(default=SlurmExperiment)
[docs] def get(self, experiment_id: str, **kwargs) -> Dict: """ Gets an experiment from the Slurm platform. Args: experiment_id: experiment id kwargs: keyword arguments used to expand functionality Returns: Slurm Experiment object """ metas = self.platform._metas.filter(item_type=ItemType.EXPERIMENT, property_filter={'id': str(experiment_id)}) if len(metas) > 0: return SlurmExperiment(metas[0]) else: raise RuntimeError(f"Not found Experiment with id '{experiment_id}'")
[docs] def platform_create(self, experiment: Experiment, **kwargs) -> SlurmExperiment: """ Creates an experiment on Slurm Platform. Args: experiment: idmtools experiment kwargs: keyword arguments used to expand functionality Returns: Slurm Experiment object created """ # ensure experiment's parent experiment.parent_id = experiment.parent_id or experiment.suite_id if experiment.parent_id is None: suite = add_dummy_suite(experiment) self.platform._suites.platform_create(suite) # update parent experiment.parent = suite # Generate Suite/Experiment/Simulation folder structure self.platform._op_client.mk_directory(experiment, exist_ok=False) meta = self.platform._metas.dump(experiment) self.platform._assets.dump_assets(experiment) self.platform._op_client.create_batch_file(experiment, **kwargs) # Copy file run_simulation_script = Path(__file__).parent.parent.joinpath('assets/') dest_script = Path(self.platform._op_client.get_directory(experiment)).joinpath('') shutil.copy(str(run_simulation_script), str(dest_script)) # Make executable self.platform._op_client.update_script_mode(dest_script) # Return Slurm Experiment return SlurmExperiment(meta)
[docs] def get_children(self, experiment: SlurmExperiment, parent: Experiment = None, raw=True, **kwargs) -> List[Any]: """ Fetch slurm experiment's children. Args: experiment: Slurm experiment raw: True/False parent: the parent of the simulations kwargs: keyword arguments used to expand functionality Returns: List of slurm simulations """ sim_list = [] sim_meta_list = self.platform._metas.get_children(parent) for meta in sim_meta_list: slurm_sim = SlurmSimulation(meta) slurm_sim.status = self.platform._op_client.get_simulation_status( if raw: sim_list.append(slurm_sim) else: sim = self.platform._simulations.to_entity(slurm_sim, parent=parent) sim_list.append(sim) return sim_list
[docs] def get_parent(self, experiment: SlurmExperiment, **kwargs) -> SlurmSuite: """ Fetches the parent of an experiment. Args: experiment: Slurm experiment kwargs: keyword arguments used to expand functionality Returns: The Suite being the parent of this experiment. """ if experiment.parent_id is None: return None else: return self.platform._suites.get(experiment.parent_id, raw=True, **kwargs)
[docs] def platform_run_item(self, experiment: Experiment, dry_run: bool = False, **kwargs): """ Run experiment. Args: experiment: idmtools Experiment dry_run: True/False kwargs: keyword arguments used to expand functionality Returns: None """ # Ensure parent experiment.parent.add_experiment(experiment) self.platform._metas.dump(experiment.parent) # Generate/update metadata self.platform._metas.dump(experiment) # Commission if not dry_run: self.platform._op_client.submit_job(experiment, **kwargs) suite_id = experiment.parent_id or experiment.suite_id #'job_id: {slurm_job_id}')'job_directory: {Path(self.platform.job_directory).resolve()}')'suite: {str(suite_id)}')'experiment: {}')
[docs] def send_assets(self, experiment: Experiment, **kwargs): """ Copy our experiment assets. Replaced by self.platform._assets.dump_assets(experiment) Args: experiment: idmtools Experiment kwargs: keyword arguments used to expand functionality Returns: None """ pass
[docs] def list_assets(self, experiment: Experiment, **kwargs) -> List[Asset]: """ List assets for an experiment. Args: experiment: Experiment to get assets for kwargs: Returns: List[Asset] """ assets = self.platform._assets.list_assets(experiment, **kwargs) return assets
[docs] def get_assets_from_slurm_experiment(self, experiment: SlurmExperiment) -> AssetCollection: """ Get assets for a comps experiment. Args: experiment: Experiment to get asset collection for. Returns: AssetCollection if configuration is set and configuration.asset_collection_id is set. """ assets = AssetCollection() assets_dir = Path(self.platform._op_client.get_directory_by_id(, ItemType.EXPERIMENT), 'Assets') assets_list = AssetCollection.assets_from_directory(assets_dir, recursive=True) for a in assets_list: assets.add_asset(a) return assets
[docs] def to_entity(self, slurm_exp: SlurmExperiment, parent: Optional[Suite] = None, children: bool = True, **kwargs) -> Experiment: """ Convert a SlurmExperiment to idmtools Experiment. Args: slurm_exp: simulation to convert parent: optional experiment object children: bool kwargs: Returns: Experiment object """ if parent is None: parent = self.platform.get_item(slurm_exp.parent_id, ItemType.SUITE, force=True) exp = Experiment() exp.platform = self.platform exp.uid = slurm_exp.uid = exp.parent_id = exp.parent = parent exp.tags = slurm_exp.tags exp._platform_object = slurm_exp exp.simulations = [] exp.assets = self.get_assets_from_slurm_experiment(slurm_exp) if exp.assets is None: exp.assets = AssetCollection() if children: exp.simulations = self.get_children(slurm_exp, parent=exp, raw=False) return exp
[docs] def refresh_status(self, experiment: Experiment, **kwargs): """ Refresh status of experiment. Args: experiment: idmtools Experiment kwargs: keyword arguments used to expand functionality Returns: Dict of simulation id as key and working dir as value """ # Check if file job_id.txt exists job_id_path = self.platform._op_client.get_directory(experiment).joinpath('job_id.txt') if not job_id_path.exists(): logger.debug(f'job_id is not available for experiment: {}') return # Refresh status for each simulation for sim in experiment.simulations: sim.status = self.platform._op_client.get_simulation_status(, **kwargs)
[docs] def create_sim_directory_map(self, experiment_id: str) -> Dict: """ Build simulation working directory mapping. Args: experiment_id: experiment id Returns: Dict of simulation id as key and working dir as value """ exp = self.platform.get_item(experiment_id, ItemType.EXPERIMENT, raw=False) sims = exp.simulations return { str(self.platform._op_client.get_directory(sim)) for sim in sims}
[docs] def platform_delete(self, experiment_id: str) -> None: """ Delete platform experiment. Args: experiment_id: platform experiment id Returns: None """ exp = self.platform.get_item(experiment_id, ItemType.EXPERIMENT, raw=False) try: shutil.rmtree(self.platform._op_client.get_directory(exp)) except RuntimeError:"Could not delete the associated experiment...") return
[docs] def platform_cancel(self, experiment_id: str, force: bool = True) -> None: """ Cancel platform experiment's slurm job. Args: experiment_id: experiment id force: bool, True/False Returns: Any """ experiment = self.platform.get_item(experiment_id, ItemType.EXPERIMENT, raw=False) if force or experiment.status == EntityStatus.RUNNING: logger.debug(f"cancel slurm job for experiment: {experiment_id}...") job_id = self.platform._op_client.get_job_id(experiment_id, ItemType.EXPERIMENT) if job_id is None: logger.debug(f"Slurm job for experiment: {experiment_id} is not available!") else: result = self.platform._op_client.cancel_job(job_id)"Cancel Experiment {experiment_id}: {result}") else:"Experiment {experiment_id} is not running, no cancel needed...")
[docs] def post_run_item(self, experiment: Experiment, **kwargs): """ Trigger right after commissioning experiment on platform. Args: experiment: Experiment just commissioned kwargs: keyword arguments used to expand functionality Returns: None """ super().post_run_item(experiment, **kwargs) job_ids = self.platform._op_client.get_job_id(, ItemType.EXPERIMENT) if job_ids is None: logger.debug(f"Slurm job for experiment: {} is not available!")"Slurm Job Ids: None") else: job_ids = [f'{" ".ljust(3)}{id}' for id in job_ids]"Slurm Job Ids ({len(job_ids)}):")'\n'.join(job_ids)) f'\nYou may try the following command to check simulations running status: \n idmtools slurm {os.path.abspath(self.platform.job_directory)} status --exp-id {}')