Source code for

"""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

    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')

[docs]def get_argument_parser(): p = get_common_parser("Download") p.add_argument("--zip-name", default="") return p
[docs]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):"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)