Source code for idmtools_platform_comps.utils.download_experiment

"""idmtools download experiment tools.

This allow downloading experiments for local testing.

    - We need some details around this somewhere. Maybe some documentation?

Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
import os
import stat
import sys
from concurrent.futures._base import as_completed
from concurrent.futures.thread import ThreadPoolExecutor
from contextlib import suppress
from idmtools import IdmConfigParser
from idmtools.core.context import get_current_platform
from idmtools.entities.experiment import Experiment
from idmtools.entities.simulation import Simulation

[docs]def get_script_extension(): """Determine extension to write out file as.""" if sys.platform in ["linux", "darwin"]: return "sh" else: return "bat"
[docs]def download_asset(asset, path): """Download a single asset.""" os.makedirs(path, exist_ok=True) asset.download_to_path(path)
[docs]def write_script(simulation: Simulation, path): """ Writes a shell script to execute simulation. Args: simulation: path: Returns: None """ command = str(simulation.task.command) sp = os.path.join(path, f"idmtools_run.{get_script_extension()}") with open(sp, "w") as sfile: sfile.write(command) if sys.platform in ["linux", "darwin"]: st = os.stat(sp) os.chmod(sp, st.st_mode | stat.S_IEXEC)
[docs]def write_experiment_script(experiment: Experiment, path: str): """ Write an experiment script. Args: experiment: path: Returns: None """ sp = os.path.join(path, f"idmtools_run.{get_script_extension()}") with open(sp, 'w') as sout: for sim in experiment.simulations: sout.write(f".{os.path.sep}{}{os.path.sep}idmtools_run.{get_script_extension()}") if sys.platform == "linux": st = os.stat(sp) os.chmod(sp, st.st_mode | stat.S_IEXEC)
[docs]def download_experiment(experiment: Experiment, destination: str): """ Downloads experiment to local directory. Useful for troubleshooting experiments Args: experiment: Experiment to download destination: Destionation Directory Returns: None """ asset_dir = os.path.join(destination, "Assets") os.makedirs(asset_dir, exist_ok=True) pool = ThreadPoolExecutor() futures = [] for asset in experiment.assets: futures.append(pool.submit(download_asset, asset, asset_dir)) for sim in experiment.simulations: sim_path = os.path.join(destination, with suppress(FileExistsError): os.makedirs(sim_path, exist_ok=True) for output in get_current_platform()._simulations.all_files(sim): futures.append(pool.submit(download_asset, output, sim_path)) write_script(sim, sim_path) if sys.platform == "linux" and not os.path.exists(os.path.join(sim_path, "Assets")): os.symlink("../Assets", os.path.join(sim_path, "Assets")) write_experiment_script(experiment, destination) if IdmConfigParser.is_progress_bar_disabled(): items = as_completed(futures) else: from tqdm import tqdm items = tqdm(as_completed(futures), total=len(futures), unit="files") for _future in items: pass