idmtools-platform-comps 0.0.0.dev0__py3-none-any.whl → 0.0.2__py3-none-any.whl

idmtools_platform_comps/utils/download/download.py

@@ -0,0 +1,178 @@
+"""idmtools download work item output.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import os
+import sys
+import warnings
+import zipfile
+from dataclasses import dataclass, field
+from enum import Enum
+from logging import getLogger, DEBUG
+from pathlib import PurePath
+from uuid import UUID
+from COMPS.Data import WorkItem
+from tqdm import tqdm
+
+from idmtools import IdmConfigParser
+from idmtools.assets.file_list import FileList
+from idmtools.core import EntityStatus
+from idmtools.entities.iplatform import IPlatform
+from idmtools_platform_comps.utils.file_filter_workitem import FileFilterWorkItem
+from idmtools_platform_comps.utils.general import get_file_as_generator
+
+logger = getLogger(__name__)
+user_logger = getLogger('user')
+
+
+class CompressType(Enum):
+    """Defines the compression types we support.
+
+    lzma is the best balance between speed and compression ratio typically
+    """
+    lzma = 'lzma'
+    deflate = 'deflate'
+    bz = "bz"
+
+
+@dataclass(repr=False)
+class DownloadWorkItem(FileFilterWorkItem):
+    """
+    DownloadWorkItem provides a utility to download items through a workitem with compression.
+
+    The main advantage of this over Analyzers is the compression. This is most effective when the targets to download have many
+    items that are similar to download. For example, an experiment with 1000 simulations with similar output can greatly benefit
+    from downloading through this method.
+
+    Notes:
+        - TODO Link examples here.
+    """
+    output_path: str = field(default_factory=os.getcwd)
+    extract_after_download: bool = field(default=True)
+    delete_after_download: bool = field(default=True)
+    zip_name: str = field(default='output.zip')
+    compress_type: CompressType = field(default=None)
+
+    def __post_init__(self, item_name: str, asset_collection_id: UUID, asset_files: FileList, user_files: FileList,
+                      command: str):
+        """
+        Constructor for DownloadWorkItem.
+
+        Args:
+            item_name: item name
+            asset_collection_id: AssetCollection id
+            asset_files: Asset files
+            user_files: user asset files
+            command: command
+
+        Returns:
+            None
+        """
+        self._ssmt_script = str(PurePath(__file__).parent.joinpath("download_ssmt.py"))
+        if self.compress_type is None:
+            if (self.extract_after_download and self.delete_after_download) or sys.platform != 'win32':
+                self.compress_type = CompressType.lzma
+            elif sys.platform == 'win32':
+                self.compress_type = CompressType.deflate
+            else:
+                self.compress_type = CompressType.lzma
+        super().__post_init__(item_name, asset_collection_id, asset_files, user_files, command)
+        warnings.warn(
+            message="DownloadWorkItem is deprecated. Until a replacement is added, please use native COMPS scripts. See "
+                    "'https://github.com/InstituteforDiseaseModeling/pyCOMPS/blob/master/examples/retrieve_simulation_outputs_for_experiment.py'",
+            category=FutureWarning)
+
+    def _extra_command_args(self, command: str) -> str:
+        """
+        Add specific additional arguments for the download command. In this case, only the zip name and compression type can be changed.
+
+        Args:
+            command: Command to append additional items to
+
+        Returns:
+            Updated command
+        """
+        if self.zip_name != "output.zip":
+            command += f" --zip-name {self.zip_name}"
+        if self.compress_type != "lzma":
+            command += f" --compress-type {self.compress_type.value}"
+        return command
+
+    def wait(self, wait_on_done_progress: bool = True, timeout: int = None, refresh_interval=None,
+             platform: 'IPlatform' = None) -> None:
+        """
+        Waits on Download WorkItem to finish. This first waits on any dependent items to finish(Experiment/Simulation/WorkItems).
+
+        Args:
+            wait_on_done_progress: When set to true, a progress bar will be shown from the item
+            timeout: Timeout for waiting on item. If none, wait will be forever
+            refresh_interval: How often to refresh progress
+            platform: Platform
+
+        Returns:
+            AssetCollection created if item succeeds
+        """
+        # wait on related items before we wait on our item
+        p = super()._check_for_platform_from_context(platform)
+        opts = dict(wait_on_done_progress=wait_on_done_progress, timeout=timeout, refresh_interval=refresh_interval,
+                    platform=p)
+        self._wait_on_children(**opts)
+
+        super().wait(**opts)
+        if self.status == EntityStatus.SUCCEEDED and not self.dry_run:
+            # Download our zip
+            po: WorkItem = self.get_platform_object(platform=self.platform)
+            if self._uid:
+                oi = po.retrieve_output_file_info([self.zip_name])
+                zip_name = PurePath(self.output_path).joinpath(self.zip_name)
+                with tqdm(total=oi[0].length, unit='B', unit_scale=True, unit_divisor=1024,
+                          desc="Downloading Files") as pbar:
+                    self.__download_file(oi, pbar, zip_name)
+                    if self.extract_after_download:
+                        self.__extract_output(zip_name)
+
+                if self.delete_after_download:
+                    if self.extract_after_download:
+                        if IdmConfigParser.is_output_enabled():
+                            user_logger.debug(f"Removing {zip_name}")
+                        os.remove(zip_name)
+                    if IdmConfigParser.is_output_enabled():
+                        user_logger.debug(f'Deleting workitem {self.uid}')
+                    po.delete()
+                    self.uid = None
+
+    def __extract_output(self, zip_name):
+        """
+        Extra output from our zip file.
+
+        Args:
+            zip_name: Zip file to extract
+
+        Returns:
+            None
+        """
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f"Extracting {zip_name}")
+        with zipfile.ZipFile(zip_name, 'r') as zin:
+            zin.extractall(self.output_path)
+
+    def __download_file(self, oi, pbar, zip_name):
+        """
+        Download our file tracking progress as we go.
+
+        Args:
+            oi: Stream to download
+            pbar: Prograss Bar
+            zip_name: Zip file to save to
+
+        Returns:
+            None
+        """
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f"Downloading file to {zip_name}")
+        parent_dir = PurePath(zip_name).parent
+        os.makedirs(parent_dir, exist_ok=True)
+        with open(zip_name, 'wb') as zo:
+            for chunk in get_file_as_generator(oi[0]):
+                pbar.update(len(chunk))
+                zo.write(chunk)

idmtools_platform_comps/utils/download/download_ssmt.py

@@ -0,0 +1,81 @@
+"""idmtools download workitem ssmt script.
+
+This script is meant to be ran remotely on SSMT, not locally.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+# flake8: noqa F405 F403
+import sys
+import zipfile
+from logging import getLogger
+import os
+from COMPS.Data import WorkItem
+
+# we have to support two ways to load the utils. The first is load from Assets
+# the fallback is the installed package
+# this allows us to move forward with changes to utils without need for new images
+from idmtools_platform_comps.utils.ssmt_utils.file_filter import parse_filter_args_common, filter_files_and_assets
+
+try:
+    from file_filter import *
+    from common import *
+except (FileNotFoundError, ImportError):
+    from idmtools_platform_comps.utils.ssmt_utils.file_filter import *
+    from idmtools_platform_comps.utils.ssmt_utils.common import *
+
+logger = getLogger(__name__)
+user_logger = getLogger('user')
+
+
+def get_argument_parser():
+    p = get_common_parser("Download")
+    p.add_argument("--zip-name", default="output.zip")
+    return p
+
+
+def create_archive_from_files(args: Namespace, files, files_from_ac, compress_type: str = "lzma"):
+    if compress_type == "lzma":
+        compress_type = zipfile.ZIP_LZMA
+    elif compress_type == "deflate":
+        compress_type = zipfile.ZIP_DEFLATED
+    elif compress_type == "bz":
+        compress_type = zipfile.ZIP_BZIP2
+    with zipfile.ZipFile(args.zip_name, mode="w", compression=compress_type) as zo:
+        for f in tqdm(sorted(files, key=lambda x: x[1]), total=len(files), mininterval=5, maxinterval=15):
+            logger.info(f"Adding {f[0].encode('ascii', 'ignore').decode('utf-8')} to {f[1].encode('ascii', 'ignore').decode('utf-8')}")
+            zo.write(f[0], f[1].encode('ascii', 'ignore').decode('utf-8'))
+        for f in tqdm(files_from_ac, total=len(files_from_ac), mininterval=5, maxinterval=15):
+            fn = PurePath(f.relative_path).joinpath(f.file_name) if f.relative_path else f.file_name
+            if logger.isEnabledFor(DEBUG):
+                logger.debug(f"Adding {f[0]} to {fn}")
+            zo.write(f.uri, fn)
+
+
+if __name__ == "__main__":  # pragma: no cover
+    # build our argument parser and then parse the command line
+    parser = get_argument_parser()
+    parser.add_argument("--compress-type", choices=["lzma", "deflate", "bz2"], default="lzma")
+    args = parser.parse_args()
+
+    # Set our JOB config global with config provided
+    JOB_CONFIG = vars(args)
+    # register our error handler
+    sys.excepthook = get_error_handler_dump_config_and_error(JOB_CONFIG)
+
+    # Parse the common arguments common to filter scripts
+    entity_filter_func, fn_format_func = parse_filter_args_common(args)
+
+    # login
+    client = login_to_env()
+
+    # Load our work-item
+    wi = WorkItem.get(os.environ['COMPS_WORKITEM_GUID'])
+
+    # Gather all our files in Experiments, Simulations, and Asset Collections
+    files, files_from_ac = filter_files_and_assets(args, entity_filter_func, wi, fn_format_func)
+    ensure_no_duplicates(files_from_ac, files)
+
+    if args.dry_run:
+        print_results(files_from_ac, files)
+    else:
+        create_archive_from_files(args, files, files_from_ac)

idmtools_platform_comps/utils/download_experiment.py

@@ -0,0 +1,116 @@
+"""idmtools download experiment tools.
+
+This allow downloading experiments for local testing.
+
+Notes:
+    - 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
+
+
+def get_script_extension():
+    """Determine extension to write out file as."""
+    if sys.platform in ["linux", "darwin"]:
+        return "sh"
+    else:
+        return "bat"
+
+
+def download_asset(asset, path):
+    """Download a single asset."""
+    os.makedirs(path, exist_ok=True)
+    asset.download_to_path(path)
+
+
+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)
+
+
+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}{sim.id}{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)
+
+
+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, sim.id)
+        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