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

idmtools_platform_comps/utils/python_requirements_ac/install_requirements.py

@@ -0,0 +1,109 @@
+"""idmtools script to run on Slurm to install python files.
+
+This is part of the RequirementsToAssetCollection tool. This will run on the HPC in an Experiment to install the python requirements
+as output that will be converted to an AssetCollection later.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import compileall
+import glob
+import os
+import subprocess
+import sys
+import time
+import traceback
+from concurrent.futures.thread import ThreadPoolExecutor
+from datetime import datetime
+
+CURRENT_DIRECTORY = os.getcwd()
+LIBRARY_ROOT = 'L'
+LIBRARY_PATH = os.path.join(CURRENT_DIRECTORY, LIBRARY_ROOT)
+REQUIREMENT_FILE = 'requirements_updated.txt'
+INDEX_URL = 'https://packages.idmod.org/artifactory/api/pypi/pypi-production/simple'
+
+
+def install_packages_from_requirements(python_paths=None):
+    """
+    Install our packages to a local directory.
+
+    Args:
+        python_paths: system Python path
+    Returns: None
+    """
+    if python_paths is None:
+        env = dict()
+    else:
+        if type(python_paths) is not list:
+            python_paths = [python_paths]
+
+        env = dict(os.environ)
+        env['PYTHONPATH'] = os.pathsep.join(python_paths)
+
+    print("Running pip install -r {} to tmp directory".format(REQUIREMENT_FILE))
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "-t", LIBRARY_PATH, "-r", f"Assets/{REQUIREMENT_FILE}", "-i",
+         f"{INDEX_URL}"], env=env)
+
+
+def set_python_dates():
+    """
+    Set python to the same dates so we don't create pyc files with differing dates.
+
+    Pyc embed the date, so this is a workaround for that behaviour.
+    """
+    print("Updating file dates")
+    pool = ThreadPoolExecutor()
+    date = datetime(year=2020, month=1, day=1, hour=0, minute=0, second=0, microsecond=0)
+    mod_time = time.mktime(date.timetuple())
+    for filename in glob.glob(f"{LIBRARY_PATH}{os.path.sep}**/*.py", recursive=True):
+        print(f"Updating date on {filename}")
+        pool.submit(os.utime, filename, (mod_time, mod_time))
+    pool.shutdown(True)
+
+
+def compile_all(python_paths=None):
+    """
+    Compile all the python files to pyc.
+
+    This is useful to reduce how often this happens since python will be an asset
+    """
+    print("Compiling pyc files")
+    if python_paths is None:
+        env = dict()
+    else:
+        if type(python_paths) is not list:
+            python_paths = [python_paths]
+
+        env = dict(os.environ)
+        env['PYTHONPATH'] = os.pathsep.join(python_paths)
+    print(f'Compiling {LIBRARY_PATH}')
+    compileall.compile_dir(os.path.relpath(LIBRARY_PATH).strip(os.path.sep), force=True)
+    print(f'Pyc Files Generated: {len(glob.glob(f"{LIBRARY_PATH}{os.path.sep}**/*.pyc", recursive=True))}')
+
+
+if __name__ == "__main__":
+    print('CURRENT_DIRECTORY: \n', CURRENT_DIRECTORY)
+    print('LIBRARY_PATH: \n', LIBRARY_PATH)
+
+    if sys.platform == "win32":
+        full_path = os.path.join(LIBRARY_PATH, 'lib', 'site-packages')
+    else:
+        full_path = os.path.join(LIBRARY_PATH, 'lib', 'python{}'.format(sys.version[:3]), 'site-packages')
+
+    if not os.path.exists(full_path):
+        os.makedirs(full_path)
+
+    print("Adding {} to the system path".format(full_path))
+    sys.path.insert(1, full_path)
+
+    tb = None
+    try:
+        install_packages_from_requirements(sys.path)
+        set_python_dates()
+        compile_all(sys.path)
+    except Exception:
+        tb = traceback.format_exc()
+        print(tb)
+    finally:
+        if tb:
+            sys.exit(-1)

idmtools_platform_comps/utils/python_requirements_ac/requirements_to_asset_collection.py

@@ -0,0 +1,374 @@
+"""idmtools requirements to asset collection.
+
+This is the entry point for users to use RequirementsToAssetCollection tool.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import os
+import hashlib
+from dataclasses import dataclass, field
+from logging import getLogger, DEBUG
+from typing import List
+from packaging.requirements import Requirement
+from COMPS.Data import QueryCriteria
+from COMPS.Data.AssetCollection import AssetCollection as COMPSAssetCollection
+from idmtools.assets import Asset, AssetCollection
+from idmtools.core import ItemType
+from idmtools.entities.experiment import Experiment
+from idmtools_models.python.json_python_task import JSONConfiguredPythonTask
+from idmtools_platform_comps.comps_platform import COMPSPlatform, SLURM_ENVS
+from idmtools_platform_comps.utils.package_version import get_highest_version
+
+CURRENT_DIRECTORY = os.path.dirname(__file__)
+REQUIREMENT_FILE = 'requirements_updated.txt'
+MODEL_LOAD_LIB = "install_requirements.py"
+MODEL_CREATE_AC = 'create_asset_collection.py'
+MD5_KEY = 'idmtools-requirements-md5-{}'
+logger = getLogger(__name__)
+user_logger = getLogger("user")
+
+
+@dataclass(repr=False)
+class RequirementsToAssetCollection:
+    """
+    RequirementsToAssetCollection provides a utility to install python packages into an asset collection.
+
+    Notes:
+        - TODO - Incorporate example in this docs
+    """
+    #: Platform object
+    platform: COMPSPlatform = field(default=None)
+    #: Name of experiment when installing requirements
+    name: str = field(default="install custom requirements")
+    #: Path to requirements file
+    requirements_path: str = field(default=None)
+    #: list of packages
+    pkg_list: list = field(default=None)
+    #: list of wheel files locally to upload and install
+    local_wheels: list = field(default=None)
+    # User tags
+    asset_tags: dict = field(default=None)
+    #: Internal checksum to calculate unique requirements set has be ran before
+    _checksum: str = field(default=None, init=False)
+    #: Calculated requirements including versions
+    _requirements: List[str] = field(default=None, init=False)
+    #: Since requirements vary by os, target it on the platform as well
+    _os_target: str = field(default=None, init=False)
+    #: Reserved tags
+    __reserved_tag: list = field(default=None, init=False)
+
+    def __post_init__(self):
+        """
+        Constructor.
+
+        Raises:
+            ValueError - if  requirements_path, pkg_list, and local_wheels are empty.
+        """
+        if not any([self.requirements_path, self.pkg_list, self.local_wheels]):
+            raise ValueError("Impossible to proceed without either requirements path or package list or local wheels!")
+
+        self.requirements_path = os.path.abspath(self.requirements_path) if self.requirements_path else None
+        self.pkg_list = self.pkg_list or []
+        self.local_wheels = [os.path.abspath(whl) for whl in self.local_wheels] if self.local_wheels else []
+        self.asset_tags = self.asset_tags or {}
+
+    @property
+    def checksum(self):
+        """
+        Calculate checksum on the requirements file.
+
+        Returns:
+            The md5 of the requirements.
+        """
+        if not self._checksum:
+            req_content = '\n'.join(self.requirements)
+            self._checksum = hashlib.md5(req_content.encode('utf-8')).hexdigest()
+
+        return self._checksum
+
+    @property
+    def md5_tag(self):
+        """
+        Get unique key for our requirements + target.
+
+        Returns:
+            The md5 tag.
+        """
+        self.init_platform()
+        return {MD5_KEY.format(self._os_target): self.checksum}
+
+    @property
+    def requirements(self):
+        """
+        Requirements property. We calculate this using consolidate_requirements.
+
+        Returns:
+            Consolidated requirements.
+        """
+        if not self._requirements:
+            self._requirements = self.consolidate_requirements()
+
+        return self._requirements
+
+    def init_platform(self):
+        """Initialize the platform."""
+        if self.platform is None:
+            # Try to detect platform
+            from idmtools.core.context import get_current_platform
+            p = get_current_platform()
+            if p is not None:
+                self.platform = p
+
+        self._os_target = "win" if "slurm" not in self.platform.environment.lower() and self.platform.environment.lower() not in SLURM_ENVS else "linux"
+        self.__reserved_tag = ['idmtools', 'task_type', MD5_KEY.format(self._os_target)]
+
+    def run(self, rerun=False):
+        """
+        Run our utility.
+
+        The working logic of this utility:
+            1. check if asset collection exists for given requirements, return ac id if exists
+            2. create an Experiment to install the requirements on COMPS
+            3. create a WorkItem to create a Asset Collection
+
+        Returns: return ac id based on the requirements if Experiment and WorkItem Succeeded
+
+        Raises:
+            Exception - If an error happens on workitem
+
+        Notes:
+            - TODO The exceptions here should be rewritten to parse errors from remote system like AssetizeOutputs
+        """
+        # Late validation
+        self.init_platform()
+
+        # Check if ac with md5 exists
+        ac = self.retrieve_ac_by_tag()
+
+        if ac and not rerun:
+            return ac.id
+
+        # Create Experiment to install custom requirements
+        exp = self.run_experiment_to_install_lib()
+        if exp is None:
+            if logger.isEnabledFor(DEBUG):
+                logger.debug('Failed to install requirements!')
+            raise Exception('Failed to install requirements!')
+
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f'\nexp: {exp.uid}')
+
+        # Create a WorkItem to create asset collection
+        wi = self.run_wi_to_create_ac(exp.uid)
+        if wi is None:
+            if logger.isEnabledFor(DEBUG):
+                logger.debug(f'Failed to create asset collection from experiment: {exp.uid}')
+            raise Exception(f'Failed to create asset collection from experiment: {exp.uid}')
+
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f'\nwi: {wi.uid}')
+
+        # get ac or return ad_id
+        ac = self.retrieve_ac_from_wi(wi)
+
+        if ac:
+            return ac.id
+
+    def save_updated_requirements(self):
+        """
+        Save consolidated requirements to a file requirements_updated.txt.
+
+        Returns:
+            None
+        """
+        user_logger.info(
+            f"Creating an updated requirements file ensuring all versions are specified at {REQUIREMENT_FILE}")
+        req_content = '\n'.join(self.requirements)
+        with open(REQUIREMENT_FILE, 'w') as outfile:
+            outfile.write(req_content)
+
+    def retrieve_ac_by_tag(self, md5_check=None):
+        """
+        Retrieve comps asset collection given ac tag.
+
+        Args:
+            md5_check: also can use custom md5 string as search tag
+        Returns: comps asset collection
+        """
+        # Late validation
+        self.init_platform()
+
+        md5_str = md5_check or self.checksum
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f'md5_str: {md5_str}')
+
+        # check if ac with tag idmtools-requirements-md5 = my_md5 exists
+        ac_list = COMPSAssetCollection.get(
+            query_criteria=QueryCriteria().select_children('tags').where_tag(
+                [f'{MD5_KEY.format(self._os_target)}={md5_str}']))
+
+        # if exists, get ac and return it
+        if len(ac_list) > 0:
+            ac_list = sorted(ac_list, key=lambda t: t.date_created, reverse=True)
+            user_logger.info(f"Found existing requirements assets at {ac_list[0].id}")
+            return ac_list[0]
+
+    def retrieve_ac_from_wi(self, wi):
+        """
+        Retrieve ac id from file ac_info.txt saved by WI.
+
+        Args:
+            wi: SSMTWorkItem (which was used to create ac from library)
+        Returns: COMPS asset collection
+        """
+        ac_file = "ac_info.txt"
+
+        # retrieve ac file
+        ret = self.platform.get_files_by_id(wi.uid, ItemType.WORKFLOW_ITEM, [ac_file])
+
+        # get file content
+        ac_id_bytes = ret[ac_file]
+
+        # convert bytes to string
+        ac_id_str = ac_id_bytes.decode('utf-8')
+
+        # return comps ac
+        return self.platform.get_item(ac_id_str, ItemType.ASSETCOLLECTION, raw=True)
+
+    def add_wheels_to_assets(self, experiment):
+        """
+        Add wheels to assets of our experiment.
+
+        Args:
+            experiment: Experiment to add assets to
+
+        Returns:
+            None
+        """
+        for whl in self.local_wheels:
+            a = Asset(filename=os.path.basename(whl), absolute_path=whl)
+            experiment.add_asset(a)
+
+    def run_experiment_to_install_lib(self):
+        """
+        Create an Experiment which will run another py script to install requirements.
+
+        Returns: Experiment created
+        """
+        self.save_updated_requirements()
+
+        task = JSONConfiguredPythonTask(script_path=os.path.join(CURRENT_DIRECTORY, MODEL_LOAD_LIB))
+        experiment = Experiment(name=self.name, simulations=[task.to_simulation()])
+        experiment.add_asset(Asset(REQUIREMENT_FILE))
+        experiment.tags = self.md5_tag
+
+        # Avoid conflict to reserved tag
+        if len(set(self.asset_tags).intersection(self.__reserved_tag)) > 0:
+            raise Exception(f"{self.__reserved_tag} are reserved tags, please use other tags!")
+
+        # Remove conflicts in case
+        for tag in self.__reserved_tag:
+            self.asset_tags.pop(tag, None)
+
+        # Update experiment's tags
+        experiment.tags.update(self.asset_tags)
+
+        self.add_wheels_to_assets(experiment)
+        user_logger.info("Run install of python requirements on COMPS. To view the details, see the experiment below")
+        experiment.run(wait_until_done=True, platform=self.platform, use_short_path=True, num_cores=1)
+
+        if experiment.succeeded:
+            return experiment
+
+    def run_wi_to_create_ac(self, exp_id):
+        """
+        Create an WorkItem which will run another py script to create new asset collection.
+
+        Args:
+            exp_id: the Experiment id (which installed requirements)
+        Returns: work item created
+        """
+        from idmtools_platform_comps.ssmt_work_items.comps_workitems import SSMTWorkItem
+
+        md5_str = self.checksum
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f'md5_str: {md5_str}')
+
+        wi_name = "wi to create ac"
+        command = f"python3 {MODEL_CREATE_AC} {exp_id} {self.platform.endpoint} {self._os_target}"
+
+        # Update tags
+        tags = {MD5_KEY.format(self._os_target): self.checksum}
+        tags.update(self.asset_tags)
+
+        user_logger.info(
+            "Converting Python Packages to an Asset Collection. This may take some time for large dependency lists")
+        wi = SSMTWorkItem(name=wi_name, command=command,
+                          transient_assets=AssetCollection([os.path.join(CURRENT_DIRECTORY, MODEL_CREATE_AC)]),
+                          tags=tags, related_experiments=[exp_id])
+
+        wi.run(wait_until_done=True, platform=self.platform)
+
+        if wi.succeeded:
+            # make ac as related_asset_collection to wi
+            from COMPS.Data.WorkItem import RelationType
+            comps_ac = self.retrieve_ac_from_wi(wi)
+            comps_wi = self.platform.get_item(wi.uid, ItemType.WORKFLOW_ITEM, raw=True)
+            comps_wi.add_related_asset_collection(comps_ac.id, relation_type=RelationType.Created)
+            comps_wi.save()
+            return wi
+        else:
+            user_logger.warning("Work item failed. See logs")
+            try:
+                files = self.platform.get_files_by_id(wi.uid, wi.item_type, ["stderr.txt"])
+                user_logger.error(f'Server Error Log: {files["stderr.txt"].decode("utf-8")}')
+            except:  # noqa: E722
+                pass
+
+    def consolidate_requirements(self):
+        """
+        Combine requirements and dynamic requirements (a list).
+
+        We do the following:
+          - get the latest version of package if version is not provided
+          - dynamic requirements will overwrites the requirements file
+
+        Returns: the consolidated requirements (as a list)
+        """
+        req_dict = {}
+        comment_list = []
+
+        if self.requirements_path:
+            with open(self.requirements_path, 'r') as fd:
+                for _cnt, line in enumerate(fd):
+                    line = line.strip()
+                    if line == '':
+                        continue
+
+                    if line.startswith('#'):
+                        comment_list.append(line)
+                        continue
+
+                    req = Requirement(line)
+                    req_dict[req.name] = req.specifier
+
+        # pkg_list will overwrite pkg in requirement file
+        if self.pkg_list:
+            for pkg in self.pkg_list:
+                req = Requirement(pkg)
+                req_dict[req.name] = req.specifier
+
+        req_list = []
+        for k, v in req_dict.items():
+            pkg_requirement = f'{k}{str(v)}'
+            req_list.append(f'{k}=={get_highest_version(pkg_requirement)}')
+
+        wheel_list = []
+        if self.local_wheels:
+            wheel_list.extend([f"Assets/{os.path.basename(whl)}" for whl in self.local_wheels])
+
+        req_list = sorted(req_list, reverse=False)
+        wheel_list = sorted(wheel_list, reverse=False)
+        update_req_list = req_list + wheel_list
+
+        return update_req_list

idmtools_platform_comps/utils/python_version.py

@@ -0,0 +1,40 @@
+"""idmtools special comps hooks.
+
+Notes:
+    - TODO update this to use new function plugins
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+SLURM_ENVIRONMENTS = ['slurmdev', 'slurm2', 'slurmstage', 'calculon']
+PYTHON_EXECUTABLE = 'python3'
+
+
+def platform_task_hooks(task, platform):
+    """
+    Update python task with proper python executable.
+
+    Args:
+        task: PythonTask or CommandTask
+        platform: the platform user uses
+
+    Returns: re-build task
+
+    Notes:
+        - TODO revisit with SingularityTasks later
+    """
+    try:
+        from idmtools_models.python.python_task import PythonTask
+        from idmtools.entities.command_task import CommandTask
+        if isinstance(task, PythonTask):
+            if platform.environment.lower() in SLURM_ENVIRONMENTS:
+                task.python_path = PYTHON_EXECUTABLE
+                task.command.executable = PYTHON_EXECUTABLE
+        elif isinstance(task, CommandTask) and platform.environment.lower() in SLURM_ENVIRONMENTS:
+            cmd_list = task.command.executable.split(' ')
+            if cmd_list[0].lower() == 'python':
+                cmd_list[0] = PYTHON_EXECUTABLE
+                task.command.executable = ' '.join(cmd_list)
+    except ImportError:
+        pass
+
+    return task

idmtools_platform_comps/utils/scheduling.py

@@ -0,0 +1,154 @@
+"""idmtools scheduling utils for comps.
+
+Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
+"""
+import json
+from os import PathLike
+from typing import List, Union, Dict
+from idmtools.assets import Asset
+from idmtools.entities.experiment import Experiment
+from idmtools.entities.simulation import Simulation
+from idmtools.entities.templated_simulation import TemplatedSimulations
+from logging import DEBUG
+import logging
+
+SCHEDULING_ERROR_UNSUPPORTED_TYPE = "The method only support object type: Experiment, Simulation, TemplatedSimulations!"
+SCHEDULING_ERROR_EMPTY_EXPERIMENT = "You cannot add scheduling config to an empty experiment."
+
+logger = logging.getLogger(__name__)
+
+
+def default_add_workorder_sweep_callback(simulation, file_name, file_path):
+    """
+    Utility function to add updated WorkOrder.json to each simulation as linked file via simulation task.
+
+    first loads original workorder file from local, then update Command field in it from each simulation object's
+    simulation.task.command.cmd, then write updated command to WorkOrder.json, and load this file to simulation
+
+    Args:
+        simulation: Simulation we are configuring
+        file_name: Filename to use
+        file_path: Path to file
+
+    Returns:
+        None
+    """
+    add_work_order(simulation, file_name=file_name, file_path=file_path)
+
+
+def default_add_schedule_config_sweep_callback(simulation, command: str = None, **config_opts):
+    """Default callback to be used for sweeps that affect a scheduling config."""
+    add_schedule_config(simulation, command=command, **config_opts["config_opts"])
+
+
+def scheduled(simulation: Simulation):
+    """
+    Determine if scheduling is defined on the simulation.
+
+    Args:
+        simulation: Simulation to check
+
+    Returns:
+        True if simulation.scheduling is defined and true.
+    """
+    scheduling = getattr(simulation, 'scheduling', False)
+    return scheduling
+
+
+def _add_work_order_asset(item: Union[Experiment, Simulation, TemplatedSimulations], config: Dict,
+                          file_name: str = "WorkOrder.json"):
+    """
+    Helper function to add an WorkOrder.json asset to an item.
+
+    Args:
+        item: The item to add the asset to
+        config: The configuration dictionary
+        file_name: The name of the file to create
+
+    Returns:
+        None
+    """
+
+    def _process_simulation(simulation: Simulation):
+        setattr(simulation, 'scheduling', True)
+        if hasattr(simulation.task.command, 'cmd') and len(simulation.task.command.cmd) > 0:
+            config["Command"] = simulation.task.command.cmd
+        ctn = json.dumps(config, indent=3)
+        simulation.add_asset(Asset(filename=file_name, content=ctn))
+
+    if isinstance(item, Simulation):
+        _process_simulation(item)
+    elif isinstance(item, TemplatedSimulations):
+        _process_simulation(item.base_simulation)
+    elif isinstance(item, Experiment):
+        if isinstance(item.simulations.items, TemplatedSimulations):
+            if len(item.simulations.items) == 0:
+                raise ValueError(SCHEDULING_ERROR_EMPTY_EXPERIMENT)
+            if logger.isEnabledFor(DEBUG):
+                logger.debug("Using Base task from template for WorkOrder.json assets")
+            _process_simulation(item.simulations.items.base_simulation)
+            for sim in item.simulations.items.extra_simulations():
+                _process_simulation(sim)
+        elif isinstance(item.simulations.items, List):
+            if len(item.simulations.items) == 0:
+                raise ValueError(SCHEDULING_ERROR_EMPTY_EXPERIMENT)
+            if logger.isEnabledFor(DEBUG):
+                logger.debug("Using all tasks to gather assets")
+            for sim in item.simulations.items:
+                _process_simulation(sim)
+        else:
+            raise ValueError("You cannot run an empty experiment")
+    else:
+        raise ValueError(SCHEDULING_ERROR_UNSUPPORTED_TYPE)
+
+
+def add_work_order(item: Union[Experiment, Simulation, TemplatedSimulations], file_name: str = "WorkOrder.json",
+                   file_path: Union[str, PathLike] = "./WorkOrder.json"):
+    """
+    Adds workorder.json.
+
+    Args:
+        item: Item to add work order to
+        file_name: Workorder file name
+        file_path: Path to file(locally)
+
+    Returns:
+        None
+
+    Raises:
+        ValueError - If experiment is empty
+                     If item is not an experiment, simulation, or TemplatedSimulations
+    """
+    with open(str(file_path), "r") as jsonFile:
+        config = json.loads(jsonFile.read())
+    _add_work_order_asset(item, config, file_name=file_name)
+
+
+def add_schedule_config(item: Union[Experiment, Simulation, TemplatedSimulations], command: str = None, **config_opts):
+    """
+    Add scheduling config to an Item.
+
+    Scheduling config supports adding to Experiments, Simulations, and TemplatedSimulations
+
+    Args:
+        item: Item to add scheduling config to
+        command: Command to run
+        **config_opts: Additional config options
+
+    config_opts details:
+            - Environment: Environment variables to set in the job environment; these can be dynamically expanded
+            - SingleNode (HPC only): A flag to limit all reserved cores to being on the same compute node
+            - Exclusive (HPC only): A flag that controls whether nodes should be exclusively allocated to this job
+            - EnableMpi (HPC or Slurm): A flag that controls whether to run the job with mpiexec
+            - NodeGroupName (HPC or Slurm): The cluster node-group to commission to
+            - NumCores (HPC or Slurm): The number of cores to reserve
+            - NumNodes (Slurm Only): The number of nodes to schedule
+            - NumProcesses (Slurm Only): The number of processes to execute
+            - additionalProperties (HPC or Slurm): True or False
+
+    Returns:
+        None
+    """
+    config = dict(Command=command)
+    config.update(config_opts)
+    _add_work_order_asset(item, config, file_name="WorkOrder.json")