idmtools-test 0.0.0.dev0__py3-none-any.whl → 0.0.2__py3-none-any.whl

idmtools_test/utils/comps.py

@@ -0,0 +1,152 @@
+import functools
+import os
+import subprocess
+import sys
+import tempfile
+from pathlib import PurePath
+from COMPS import Data
+from COMPS.Data import AssetCollection as CompsAssetCollection
+from COMPS.Data import QueryCriteria, Simulation as COMPSSimulation, Experiment as COMPSExperiment
+from idmtools import __version__ as core_version
+from idmtools.builders import SimulationBuilder
+from idmtools.core.enums import EntityStatus
+from idmtools.entities.experiment import Experiment
+from idmtools.entities.iplatform import IPlatform
+from idmtools_models.templated_script_task import get_script_wrapper_unix_task
+from idmtools_platform_comps import __version__ as platform_comps_version
+
+CURRENT_DIR = PurePath(__file__).parent
+COMPS_VERSION = platform_comps_version.replace('nightly.0', 'nightly')
+if COMPS_VERSION.endswith(".0") and len(COMPS_VERSION) == 7:
+    COMPS_VERSION = ".".join(COMPS_VERSION.split(".")[:3])
+CORE_VERSION = core_version.replace('nightly.0', 'nightly')
+if CORE_VERSION.endswith(".0") and len(CORE_VERSION) == 7:
+    CORE_VERSION = ".".join(CORE_VERSION.split(".")[:3])
+COMPS_PACKAGE_FILENAME = f"idmtools_platform_comps-{COMPS_VERSION}.tar.gz"
+CORE_PACKAGE_FILENAME = f"idmtools-{CORE_VERSION}.tar.gz"
+COMPS_LOCAL_PACKAGE = CURRENT_DIR.parent.parent.parent.joinpath("idmtools_platform_comps", "dist", COMPS_PACKAGE_FILENAME)
+CORE_LOCAL_PACKAGE = CURRENT_DIR.parent.parent.parent.joinpath("idmtools_core", "dist", CORE_PACKAGE_FILENAME)
+COMPS_LOAD_SSMT_PACKAGES_WRAPPER = f"""
+set -o noglob
+echo Running $@ 
+
+echo after install of newer idmtools
+
+export PYTHONPATH=$(pwd)/Assets/site-packages:$(pwd)/Assets/:$PYTHONPATH
+
+echo "Installing updated versions of idmtools packages"
+pip install Assets/{COMPS_PACKAGE_FILENAME} --index-url=https://packages.idmod.org/api/pypi/pypi-production/simple
+pip install Assets/{CORE_PACKAGE_FILENAME} --force --index-url=https://packages.idmod.org/api/pypi/pypi-production/simple
+
+$@
+"""
+
+
+@functools.lru_cache(1)
+def write_wrapper_script():
+    f = tempfile.NamedTemporaryFile(suffix='.sh', mode='wb', delete=False)
+    f.write(COMPS_LOAD_SSMT_PACKAGES_WRAPPER.replace("\r", "").encode('utf-8'))
+    f.flush()
+    return f.name
+
+
+def load_library_dynamically(item, platform: IPlatform):
+    fn = write_wrapper_script()
+    for file in [COMPS_LOCAL_PACKAGE, CORE_LOCAL_PACKAGE]:
+        item.assets.add_asset(file)
+    item.task = get_script_wrapper_unix_task(task=item.task, template_content=COMPS_LOAD_SSMT_PACKAGES_WRAPPER)
+    item.task.gather_common_assets()
+    item.task.pre_creation(item, platform)
+    # item.assets.add_assets(item.task.gather_common_assets(), fail_on_duplicate=False)
+
+
+def run_package_dists():
+    mk = "pymake" if sys.platform == "win32" else "make"
+    print("Running Dist for core")
+    subprocess.call(f"{mk} dist", cwd=CORE_LOCAL_PACKAGE.parent.parent, shell=True)
+    print("Running Dist for comps")
+    subprocess.call(f"{mk} dist", cwd=COMPS_LOCAL_PACKAGE.parent.parent, shell=True)
+
+
+def get_asset_collection_id_for_simulation_id(sim_id):
+    """
+    Obtains COMPS AssetCollection id from a given simulation id.
+    :param sim_id: A simulation id to retrieve assetcollection id from
+    :return: COMPS AssetCollection id
+    """
+    simulation = COMPSSimulation.get(sim_id, query_criteria=QueryCriteria().select(
+        ['id', 'experiment_id']).select_children(
+        ["files", "configuration"]))
+
+    if simulation.configuration is None:
+        # check experiment
+        experiment = COMPSExperiment.get(simulation.experiment_id, query_criteria=QueryCriteria().select(
+            ['id']).select_children("configuration")
+        )
+        collection_id = experiment.configuration.asset_collection_id
+    else:
+        collection_id = simulation.configuration.asset_collection_id
+    return collection_id
+
+
+def get_asset_collection_by_id(collection_id, query_criteria=None) -> CompsAssetCollection:
+    """
+    Obtains COMPS AssetCollection from a given collection id.
+    :param collection_id: An asset collection id to retrieve assetcollection from
+    :param query_criteria: query_criteria
+    :return: COMPS AssetCollection
+    """
+    query_criteria = query_criteria or QueryCriteria().select_children('assets')
+    try:
+        return Data.AssetCollection.get(collection_id, query_criteria)
+    except (RuntimeError, ValueError):
+        return None
+
+
+def sims_from_experiment(e):
+    o = e
+    if isinstance(e, Experiment):
+        o = e.get_platform_object()
+    return o.get_simulations(QueryCriteria().select(['id', 'state']).select_children('hpc_jobs'))
+
+
+def workdirs_from_simulations(sims):
+    return {str(sim.id): sim.hpc_jobs[-1].working_directory for sim in sims if sim.hpc_jobs}
+
+
+def get_simulation_path(simulation):
+    path = workdirs_from_simulations([simulation])[str(simulation.id)]
+    return path
+
+
+def get_simulation_by_id(sim_id, query_criteria=None):
+    return COMPSSimulation.get(id=sim_id, query_criteria=query_criteria)
+
+
+def assure_running_then_wait_til_done(tst, experiment):
+    tst.platform.run_items(items=[experiment])
+    tst.platform.refresh_status(item=experiment)
+    tst.assertFalse(experiment.done)
+    tst.assertTrue(all([s.status == EntityStatus.RUNNING for s in experiment.simulations]))
+    # Wait till done
+    import time
+    start_time = time.time()
+    while time.time() - start_time < 180:
+        tst.platform.refresh_status(item=experiment)
+        if experiment.done:
+            break
+        time.sleep(3)
+    tst.assertTrue(experiment.done)
+
+
+def setup_test_with_platform_and_simple_sweep(tst):
+    from idmtools.core.platform_factory import Platform
+    tst.platform = Platform('SlurmStage')
+    print(tst.case_name)
+
+    def setP(simulation, p):
+        return simulation.task.set_parameter("P", p)
+
+    tst.builder = SimulationBuilder()
+    tst.builder.add_sweep_definition(setP, [1, 2, 3])
+

idmtools_test/utils/decorators.py

@@ -0,0 +1,208 @@
+import functools
+import shutil
+
+import tempfile
+
+import os
+import platform
+import time
+import unittest
+from functools import wraps
+from logging import getLogger
+from typing import Callable, Union, Any, Optional
+import pytest
+from idmtools import IdmConfigParser
+from idmtools_test import COMMON_INPUT_PATH
+from idmtools_test.utils.utils import is_global_configuration_enabled
+
+# The following decorators are used to control test
+# To allow for different use cases(dev, test, packaging, etc)
+# We have switches that should allow a rich set of possible
+# test combinations
+#
+# The default tests run with all the optional tests set to off(except Linux since that is auto-detected)
+# For test-external runs any tests that require external communication
+# This currently is any comps related test
+# test-docker run any tests that depend on docker locally(Mostly local runn)
+# test-all runs all tests
+
+logger = getLogger(__name__)
+
+linux_only = unittest.skipIf(
+    not platform.system() in ["Linux", "Darwin"], 'No Tests that are meant for linux'
+)
+
+windows_only = unittest.skipIf(
+    platform.system() in ["Linux", "Darwin"], 'No Tests that are meant for Windows'
+)
+
+skip_if_global_configuration_is_enabled = pytest.mark.skipif(is_global_configuration_enabled(), reason=f"Either {IdmConfigParser.get_global_configuration_name()} is set or the environment variable 'IDMTOOLS_CONFIG_FILE' is set")
+# this is mainly for docker in docker environments but also applies to environments
+# where you must use the local ip address for connectivity vs localhost
+skip_api_host = unittest.skipIf(os.getenv("API_HOST", None) is not None, "API_HOST is defined")
+
+
+def run_test_in_n_seconds(n: int, print_elapsed_time: bool = False) -> Callable:
+    """
+    Decorator that assert a test will run in N seconds. If it does not, it will fail.
+
+    THIS DOES NOT MANAGE processes and it will not stop long running processes. It simply times a test and ensure it
+    ran in less than N seconds
+
+    Args:
+        n (int): Number of seconds that is considered acceptable
+        print_elapsed_time: Will print the function name and elasped time at the end of each function it decorates
+
+    Returns:
+        (Callable) : Wrapped function
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start = time.time()
+            ret = func(*args, **kwargs)
+            end = time.time()
+            if print_elapsed_time:
+                print(f"{func.__name__} took {end - start}s to run!")
+            args[0].assertLess(end - start, n, f"{func.__name__} took {end - start}s to run!")
+            return ret
+
+        return wrapper
+
+    return decorator
+
+
+def default_fixture_pickle_save(filename, args, pickler=None, write_mode='bw'):
+    if pickler is None:
+        import pickle
+        pickler = pickle
+    with open(filename, write_mode) as o:
+        try:
+            pickler.dump(args, o)
+        except TypeError as e:
+            print("Unpicklable object!")
+            print(e)
+
+
+def dump_function_input_for_test(output_directory: Union[str, Callable[[str, Callable], str]] = None,
+                                 capture_output: bool = False, include_module_in_path: bool = True,
+                                 custom_pickler=None, save_extension: str = '.pkl', write_mode: str = 'bw',
+                                 custom_save_output_func: Optional[Callable[[str, Any], Any]] = None):
+    """
+
+    Args:
+        output_directory: Output directory of fixture
+        capture_output: Capture output of function as well
+        include_module_in_path: When using the default path scheme, should the module path be part of the
+            destination name
+        custom_pickler: When using default save function, do you want to use a custom pickler or alternative like json
+        save_extension: Save extension. Default to '.pkl'
+        custom_save_output_func: Custom save function
+
+    Examples:
+        Using with a class:
+        ```
+        class Experiment
+
+            @classmethod
+            @dump_function_input_for_test()
+            def save(cls)
+                pass
+
+            @dump_function_input_for_test()
+            def normal_funct(self, a):
+                pass
+        ```
+    Returns:
+
+    """
+    # if no output directory, set to fixture path
+    if output_directory is None:
+        output_directory = COMMON_INPUT_PATH
+
+    def decorate(func):
+        # create directory for function output
+        # if we have a function for directory naming, call it
+        if callable(output_directory):
+            out = output_directory(COMMON_INPUT_PATH, func)
+        else:
+            fname = f'{func.__module__}.{func.__name__}' if include_module_in_path else func.__name__
+            fname = fname.replace('.', os.path.sep)
+            out = os.path.abspath(os.path.join(output_directory, fname))
+
+        # create our paths to input/output files
+        out_directory = os.path.join(out, 'output')
+        input_directory = os.path.join(out, 'input')
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            call_time = int(round(time.time() * 1000))
+            # if output directory is a
+            # save keywords and args to files
+            os.makedirs(input_directory, exist_ok=True)
+            input_file_out = os.path.join(input_directory, f'{call_time}{save_extension}')
+            if custom_save_output_func:
+                custom_save_output_func(input_file_out, (args, kwargs))
+            else:
+                default_fixture_pickle_save(input_file_out, (args, kwargs), custom_pickler, write_mode)
+            result = func(*args, **kwargs)
+            if capture_output:
+                os.makedirs(out_directory, exist_ok=True)
+                output_file_out = os.path.join(out_directory, f'{call_time}{save_extension}')
+                if custom_save_output_func:
+                    custom_save_output_func(output_file_out, result)
+                else:
+                    default_fixture_pickle_save(output_file_out, result, custom_pickler, write_mode)
+            return result
+
+        return wrapper
+
+    return decorate
+
+
+def run_in_temp_dir(func):
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        current_dir = os.getcwd()
+        temp_dir = tempfile.mkdtemp()
+        try:
+            logger.debug(f"Running function in: {temp_dir}")
+            os.chdir(temp_dir)
+            func(*args, **kwargs)
+        finally:
+            os.chdir(current_dir)
+            try:
+                shutil.rmtree(temp_dir)
+            except:
+                pass
+
+    return wrapper
+
+
+def warn_amount_ssmt_image_decorator(func):
+    """
+    A decorator to warn developers about possible failures due to SSMT.
+
+    Args:
+        func: Function to wrap
+
+    Returns:
+        Wrapped function
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            logger.error(
+                "These tests can fail due to changes to idmtools-core, idmtools-models, or idmtools-platform-comps. "
+                "If you have changed the code in those libraries, you will need to build a new ssmt image, publish to staging,"
+                "then update idmtools_platform_comps/tests/idmtools.ini by uncommenting the 'docker_image' options. You should"
+                "change the value to the new version of the SSMT image. COMPS Will automatically pull the new image."
+            )
+            raise e
+
+    return wrapper

idmtools_test/utils/execute_operations/experiment_operations.py

@@ -0,0 +1,237 @@
+import hashlib
+import io
+import json
+import os
+import shutil
+from dataclasses import field, dataclass
+from functools import partial
+from logging import getLogger, DEBUG
+from threading import Lock
+from pathlib import Path
+from typing import Any, List, Type, Dict, TYPE_CHECKING, Optional
+from idmtools.assets import Asset, AssetCollection
+from idmtools.core import EntityStatus, ItemType
+from idmtools.core import IDMTOOLS_USER_HOME
+from idmtools.entities.experiment import Experiment
+from idmtools.entities.iplatform_ops.iplatform_experiment_operations import IPlatformExperimentOperations
+from idmtools.utils.file import file_content_to_generator
+from idmtools.utils.json import IDMJSONEncoder
+
+if TYPE_CHECKING:  # pragma: no cover
+    from idmtools_test.utils.test_execute_platform import TestExecutePlatform
+
+logger = getLogger(__name__)
+current_directory = os.path.dirname(os.path.realpath(__file__))
+data_path = os.path.abspath(os.path.join(current_directory, "..", "..", "data"))
+EXPERIMENTS_LOCK = Lock()
+
+
+class ExperimentDict(dict):
+    pass
+
+
+@dataclass
+class TestExecutePlatformExperimentOperation(IPlatformExperimentOperations):
+    platform: 'TestExecutePlatform'
+    platform_type: Type = field(default=ExperimentDict)
+    experiments: Dict[str, Experiment] = field(default_factory=dict, compare=False, metadata={"pickle_ignore": True})
+
+    def get(self, experiment_id: str, **kwargs) -> Any:
+        exp_path = self.get_experiment_path(experiment_id)
+        experiment_path = os.path.join(exp_path, "experiment.json")
+        if not os.path.exists(experiment_path):
+            logger.error(f"Cannot find experiment with id {experiment_id}")
+            raise FileNotFoundError(f"Cannot find experiment with id {experiment_id}")
+
+        simulation_metadata_path = os.path.join(exp_path, "simulation_index.json")
+        if not os.path.exists(simulation_metadata_path):
+            logger.error(f"Cannot find simulation index for experiment with id {experiment_id}")
+            raise FileNotFoundError(f"Cannot find simulation index for experiment with id {experiment_id}")
+
+        logger.info(f"Loading experiment metadata from {experiment_path}")
+        with open(experiment_path, 'r') as metadata_in:
+            metadata = json.load(metadata_in)
+        logger.info(f"Loading simulation metadata from {simulation_metadata_path}")
+        with open(simulation_metadata_path, 'r') as metadata_in:
+            metadata['simulations'] = json.load(metadata_in)
+        return ExperimentDict(metadata)
+
+    def platform_create(self, experiment: Experiment, **kwargs) -> Any:
+        if logger.isEnabledFor(DEBUG):
+            logger.debug('Creating Experiment')
+        EXPERIMENTS_LOCK.acquire()
+        self.experiments[experiment.uid] = experiment
+        EXPERIMENTS_LOCK.release()
+        logger.debug(f"Created Experiment {experiment.uid}")
+        self.send_assets(experiment, **kwargs)
+        return experiment
+
+    def get_children(self, experiment: ExperimentDict, **kwargs) -> List[Any]:
+        children = []
+        for sim in experiment['simulations']:
+            children.append(self.platform.get_item(sim, ItemType.SIMULATION,
+                                                   experiment_id=experiment['_uid'],
+                                                   raw=True,
+                                                   **kwargs
+                                                   )
+                            )
+        return children
+
+    def get_parent(self, experiment: Any, **kwargs) -> Experiment:
+        pass
+
+    def platform_run_item(self, experiment: Experiment, **kwargs):
+        exp_path = self.get_experiment_path(experiment.uid)
+        path = os.path.join(exp_path, "experiment.json")
+        os.makedirs(exp_path, exist_ok=True)
+        if not os.path.exists(path):
+            with open(path, 'w') as out:
+                out.write(json.dumps(experiment.to_dict(), cls=IDMJSONEncoder))
+        for sim in experiment.simulations:
+            if sim.status in [None, EntityStatus.CREATED]:
+                self.platform._simulations.run_item(sim)
+
+    def get_experiment_path(self, experiment_id: str) -> str:
+        """
+        Get path to experiment directory
+
+        Args:
+            experiment_id:
+
+        Returns:
+
+        """
+        return os.path.join(self.platform.execute_directory, str(experiment_id))
+
+    @staticmethod
+    def download_asset(path):
+        logger.info(f"Downloading asst from {path}")
+        if not os.path.exists(path):
+            logger.error(f"Cannot the asset {path}")
+            raise FileNotFoundError(f"Cannot the asset {path}")
+        with open(path, 'rb') as i:
+            while True:
+                res = i.read(128)
+                if res:
+                    yield res
+                else:
+                    break
+
+    def send_assets(self, experiment: Experiment, **kwargs):
+        # calculate total md5 of all files
+        md5 = hashlib.md5()
+        path = os.path.join(self.platform.execute_directory, str(experiment.uid), "Assets")
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f"Creating {path}")
+        os.makedirs(path, exist_ok=True)
+        for asset in experiment.assets:
+            remote_path = os.path.join(path, asset.relative_path) if asset.relative_path else path
+            remote_path = os.path.join(remote_path, asset.filename)
+            if asset.absolute_path:
+                if logger.isEnabledFor(DEBUG):
+                    logger.debug(f"Copying {asset.absolute_path} to {remote_path}")
+                with open(asset.absolute_path) as ifi:
+                    self.__calculate_partial_md5(ifi, md5)
+                shutil.copy(asset.absolute_path, remote_path)
+            else:
+                ifi = io.BytesIO(asset.content.encode('utf-8') if isinstance(asset.content, str) else asset.content)
+                self.__calculate_partial_md5(ifi, md5)
+                if logger.isEnabledFor(DEBUG):
+                    logger.debug(f"Writing {asset.absolute_path} to {remote_path}")
+                with open(remote_path, 'wb') as out:
+                    if isinstance(asset.content, str):
+                        out.write(asset.content.encode('utf-8'))
+                    else:
+                        out.write(asset.content)
+
+        experiment.assets.platform_id = md5.hexdigest()
+
+    def __calculate_partial_md5(self, ifi, md5):
+        while True:
+            chunk = ifi.read(8196)
+            if not chunk:
+                break
+            else:
+                if isinstance(chunk, bytes):
+                    md5.update(chunk)
+                else:
+                    md5.update(chunk.encode('utf-8'))
+
+    def refresh_status(self, experiment: Experiment, **kwargs):
+        if logger.isEnabledFor(DEBUG):
+            logger.debug(f'Refreshing status for Experiment: {experiment.uid}')
+        for simulation in self.platform._simulations.simulations.get(experiment.uid):
+            for esim in experiment.simulations:
+                if esim == simulation:
+                    logger.debug(f'Setting {simulation.uid} Status to {simulation.status}')
+                    esim.status = simulation.status
+                    break
+
+    def list_assets(self, experiment: Experiment, children: bool = False,
+                    **kwargs) -> List[Asset]:
+        """
+        List assets for the experiment
+
+        Args:
+            experiment:
+            children:
+            **kwargs:
+
+        Returns:
+
+        """
+        logger.info("Listing assets for experiment")
+        assets = []
+        asset_path = os.path.join(self.get_experiment_path(experiment.uid), "Assets")
+        for root, files, dirs in os.walk(asset_path):
+            for file in files:
+                fp = os.path.join(asset_path, file)
+                asset = Asset(absolute_path=fp, filename=file)
+                assets.append(asset)
+
+        if children:
+            for sim in experiment.simulations:
+                assets.extend(self.platform._simulations.list_assets(sim))
+        return assets
+
+    def to_entity(self, data: Dict[Any, Any], parent: Optional[Any] = None, children: bool = True, **kwargs) -> \
+            Experiment:
+        excluded = ['platform_id', 'item_type', 'frozen', 'simulations']
+        experiment = Experiment(**{k: v for k, v in data.items() if k not in excluded})
+        experiment.platform_metadata = data
+        experiment.task_type = data['tags']['task_type']
+        if data['assets']:
+            assets = AssetCollection()
+            exp_path = os.path.join(self.get_experiment_path(experiment.uid), "Assets")
+            for root, dirs, files in os.walk(exp_path):
+                for file in files:
+                    fp = os.path.abspath(os.path.join(root, file))
+                    asset = Asset(absolute_path=fp, filename=file)
+                    asset.download_generator_hook = partial(file_content_to_generator, fp)
+                    assets.add_asset(asset)
+            experiment.assets = assets
+        if children:
+            experiment.simulations = self.platform.get_children(
+                experiment.uid,
+                ItemType.EXPERIMENT,
+                item=experiment,
+                **kwargs
+            )
+
+        return experiment
+
+    def platform_modify_experiment(self, experiment: Experiment, regather_common_assets: bool = False, **kwargs) -> Experiment:
+        experiment.pre_creation(self.platform, gather_assets=regather_common_assets)
+        EXPERIMENTS_LOCK.acquire()
+        self.experiments[experiment.uid] = experiment
+        EXPERIMENTS_LOCK.release()
+        self.send_assets(experiment)
+        return experiment
+
+    def post_run_item(self, experiment: Experiment, **kwargs):
+        exp_path = self.get_experiment_path(experiment.uid)
+        sim_path = Path(exp_path, "simulation_index.json")
+        with open(sim_path, "w") as f:
+            json.dump([s.id for s in experiment.simulations], f)
+
+        super().post_run_item(experiment, **kwargs)