diff --git a/lib/pavilion/cmd_utils.py b/lib/pavilion/cmd_utils.py index 10acb3727..d6c343328 100644 --- a/lib/pavilion/cmd_utils.py +++ b/lib/pavilion/cmd_utils.py @@ -7,6 +7,7 @@ import logging import sys import time +import os from pathlib import Path from typing import List, TextIO, Union, Iterator, Optional from collections import defaultdict @@ -523,18 +524,31 @@ def get_last_test_id(pav_cfg: "PavConfig", errfile: TextIO) -> Optional[TestID]: if last_series is None: return None - test_ids = list(last_series.tests.keys()) + id_pairs = list(last_series.tests.keys()) - if len(test_ids) == 0: + if len(id_pairs) == 0: output.fprint( errfile, f"Most recent series contains no tests.") return None - if len(test_ids) > 1: + if len(id_pairs) > 1: output.fprint( errfile, f"Multiple tests exist in last series. Could not unambiguously identify last test.") return None - return TestID(test_ids[0]) + return TestID(str(id_pairs[0][1])) + + +def list_files(path: Path, include_root: bool = False) -> Iterator[Path]: + """Recursively list all files in a directory, optionally including the directory itself.""" + + for root, dirs, files in os.walk(path): + if include_root: + yield Path(root) + + for fname in files: + yield Path(root) / fname + for dname in dirs: + yield Path(root) / dname diff --git a/lib/pavilion/commands/__init__.py b/lib/pavilion/commands/__init__.py index 54309f773..7c76e2815 100644 --- a/lib/pavilion/commands/__init__.py +++ b/lib/pavilion/commands/__init__.py @@ -23,6 +23,7 @@ 'config': ('config', 'ConfigCommand'), 'graph': ('graph', 'GraphCommand'), 'group': ('group', 'GroupCommand'), + 'isolate': ('isolate', 'IsolateCommand'), 'list': ('list_cmd', 'ListCommand'), 'log': ('log', 'LogCommand'), 'ls': ('ls', 'LSCommand'), diff --git a/lib/pavilion/commands/isolate.py b/lib/pavilion/commands/isolate.py new file mode 100644 index 000000000..8e7c9ceb9 --- /dev/null +++ b/lib/pavilion/commands/isolate.py @@ -0,0 +1,216 @@ +from argparse import ArgumentParser, Namespace, Action +from pathlib import Path +import tarfile +import sys +import shutil +import tempfile +from typing import Iterable + +from pavilion import output +from pavilion import schedulers +from pavilion.config import PavConfig +from pavilion.test_run import TestRun +from pavilion.test_ids import TestID +from pavilion.cmd_utils import get_last_test_id, get_tests_by_id, list_files +from pavilion.utils import copytree_resolved +from pavilion.scriptcomposer import ScriptComposer +from pavilion.errors import SchedulerPluginError +from pavilion.schedulers.config import validate_config, calc_node_range +from .base_classes import Command + + +class IsolateCommand(Command): + """Isolates an existing test run in a form that can be run without Pavilion.""" + + IGNORE_FILES = ("series", "job") + KICKOFF_FN = "kickoff.isolated" + + def __init__(self): + super().__init__( + "isolate", + "Isolate an existing test run.", + short_help="Isolate a test run." + ) + + def _setup_arguments(self, parser: ArgumentParser) -> None: + """Setup the argument parser for the isolate command.""" + + parser.add_argument( + "test_id", + type=TestID, + nargs="?", + help="test ID" + ) + + parser.add_argument( + "path", + type=Path, + help="isolation path" + ) + + parser.add_argument( + "-a", + "--archive", + action="store_true", + default=False, + help="archive the test" + ) + + parser.add_argument( + "-z", + "--zip", + default=False, + help="compress the test archive", + action="store_true" + ) + + def run(self, pav_cfg: PavConfig, args: Namespace) -> int: + """Run the isolate command.""" + + if args.zip and not args.archive: + output.fprint(self.errfile, "--archive must be specified to use --zip.") + + return 1 + + test_id = args.test_id + + if args.test_id is None: + test_id = get_last_test_id(pav_cfg, self.errfile) + + if test_id is None: + output.fprint(self.errfile, "No last test found.", color=output.RED) + + return 2 + + tests = get_tests_by_id(pav_cfg, [test_id], self.errfile) + + if len(tests) == 0: + output.fprint(self.errfile, "Could not find test '{}'".format(test_id)) + + return 3 + + elif len(tests) > 1: + output.fprint( + self.errfile, "Matched multiple tests. Printing file contents for first " + "test only (test {})".format(tests[0].full_id), + color=output.YELLOW) + + return 4 + + test = next(iter(tests)) + + return self._isolate(pav_cfg, test, args.path, args.archive, args.zip) + + @classmethod + def _isolate(cls, pav_cfg: PavConfig, test: TestRun, dest: Path, archive: bool, + zip: bool) -> int: + """Given a test run and a destination path, isolate that test run, optionally + creating a tarball.""" + + if not test.path.is_dir(): + output.fprint(sys.stderr, "Directory '{}' does not exist." + .format(test.path.as_posix()), color=output.RED) + + return 5 + + if dest.exists(): + output.fprint( + sys.stderr, + f"Unable to isolate test {test.id}. Destination {dest} already exists.", + color=output.RED) + + return 6 + + if archive: + cls._write_tarball(pav_cfg, + test, + dest, + zip, + cls.IGNORE_FILES) + + else: + try: + copytree_resolved(test.path, dest, ignore_files=cls.IGNORE_FILES) + except OSError as err: + output.fprint( + sys.stderr, + f"Unable to isolate test {test.id} at {dest}: {err}", + color=output.RED) + + return 8 + + pav_lib_bash = pav_cfg.pav_root / 'bin' / TestRun.PAV_LIB_FN + shutil.copyfile(pav_lib_bash, dest / TestRun.PAV_LIB_FN) + + cls._write_kickoff_script(pav_cfg, test, dest / cls.KICKOFF_FN) + + return 0 + + @classmethod + def _write_tarball(cls, pav_cfg: PavConfig, test: TestRun, dest: Path, zip: bool, + ignore_files: Iterable[str]) -> None: + """Given a test run object, create a tarball of its run directory in the specified + location.""" + + if zip: + if len(dest.suffixes) == 0: + dest = dest.with_suffix(".tgz") + + modestr = "w:gz" + else: + if len(dest.suffixes) == 0: + dest = dest.with_suffix(".tar") + + modestr = "w:" + + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + tmp_dest = tmp / dest.stem + tmp_dest.mkdir() + copytree_resolved(test.path, tmp_dest, ignore_files=ignore_files) + + # Copy Pavilion bash library into tarball + pav_lib_bash = pav_cfg.pav_root / 'bin' / TestRun.PAV_LIB_FN + shutil.copyfile(pav_lib_bash, tmp_dest / TestRun.PAV_LIB_FN) + + cls._write_kickoff_script(pav_cfg, test, tmp_dest / cls.KICKOFF_FN) + + try: + with tarfile.open(dest, modestr) as tarf: + for fname in list_files(tmp): + tarf.add( + fname, + arcname=fname.relative_to(tmp), + recursive=False) + except (tarfile.TarError, OSError): + output.fprint( + sys.stderr, + f"Unable to isolate test {test.id} at {dest}.", + color=output.RED) + + return 7 + + @classmethod + def _write_kickoff_script(cls, pav_cfg: PavConfig, test: TestRun, script_path: Path) -> None: + """Write a special kickoff script that can be used to run the given test independently of + Pavilion.""" + + try: + sched = schedulers.get_plugin(test.scheduler) + except SchedulerPluginError: + output.fprint( + sys.stderr, + f"Unable to generate kickoff script for test {test_id}: unable to load scheduler" + f" {test.scheduler}." + ) + return 9 + + sched_config = validate_config(test.config['schedule']) + node_range = calc_node_range(sched_config, sched_config['cluster_info']['node_count']) + + script = sched.create_kickoff_script( + pav_cfg, + test, + isolate=True) + + script.write(script_path) diff --git a/lib/pavilion/schedulers/advanced.py b/lib/pavilion/schedulers/advanced.py index 35f673b47..ee1227895 100644 --- a/lib/pavilion/schedulers/advanced.py +++ b/lib/pavilion/schedulers/advanced.py @@ -4,11 +4,14 @@ import collections import pprint from abc import ABC -from typing import Tuple, List, Any, Union, Dict, FrozenSet, NewType +from pathlib import Path +from typing import Tuple, List, Any, Union, Dict, FrozenSet, NewType, Optional +from pavilion.config import PavConfig from pavilion.jobs import Job, JobError from pavilion.status_file import STATES from pavilion.test_run import TestRun +from pavilion.scriptcomposer import ScriptComposer from pavilion.types import NodeInfo, Nodes, NodeList, NodeSet, NodeRange from .config import validate_config, AVAILABLE, BACKFILL, calc_node_range from .scheduler import SchedulerPlugin @@ -542,19 +545,7 @@ def _schedule_shared(self, pav_cfg, tests: List[TestRun], node_range: NodeRange, # Clear the node range - it's only used for flexible scheduling. node_range = None - - job_name = 'pav_{}'.format(','.join(test.name for test in tests[:4])) - if len(tests) > 4: - job_name += ' ...' - script = self._create_kickoff_script_stub(pav_cfg, job_name, job.kickoff_log, - base_sched_config, nodes=picked_nodes, - node_range=node_range, - shebang=base_test.shebang) - - # Run each test via pavilion - script.command('echo "Starting {} tests - $(date)"'.format(len(tests))) - script.command('pav _run {}'.format(" ".join(test.full_id for test in tests))) - + script = self.create_kickoff_script(pav_cfg, tests, job.kickoff_log, nodes=picked_nodes) script.write(job.kickoff_path) # Create symlinks for each test to the one test with the kickoff script and @@ -567,7 +558,7 @@ def _schedule_shared(self, pav_cfg, tests: List[TestRun], node_range: NodeRange, pav_cfg=pav_cfg, job=job, sched_config=base_sched_config, - job_name=job_name, + job_name=self._job_name(tests), nodes=picked_nodes, node_range=node_range) except SchedulerPluginError as err: @@ -608,16 +599,7 @@ def _schedule_indi_flex(self, pav_cfg, tests: List[TestRun], node_range = calc_node_range(sched_config, len(chunk)) - job_name = 'pav_{}'.format(test.name) - script = self._create_kickoff_script_stub( - pav_cfg=pav_cfg, - job_name=job_name, - log_path=job.kickoff_log, - sched_config=sched_config, - node_range=node_range, - shebang=test.shebang) - - script.command('pav _run {t.full_id}'.format(t=test)) + script = self.create_kickoff_script(pav_cfg, test, job.kickoff_log) script.write(job.kickoff_path) test.job = job @@ -627,7 +609,7 @@ def _schedule_indi_flex(self, pav_cfg, tests: List[TestRun], pav_cfg=pav_cfg, job=job, sched_config=sched_config, - job_name=job_name, + job_name=self._job_name(test), node_range=node_range, ) except SchedulerPluginError as err: @@ -703,16 +685,7 @@ def _schedule_indi_chunk(self, pav_cfg, tests: List[TestRun], prior_error=err, tests=[test])) continue - job_name = 'pav_{}'.format(test.name) - script = self._create_kickoff_script_stub( - pav_cfg=pav_cfg, - job_name=job_name, - log_path=job.kickoff_log, - sched_config=sched_config, - nodes=picked_nodes, - shebang=test.shebang) - - script.command('pav _run {t.full_id}'.format(t=test)) + script = self.create_kickoff_script(pav_cfg, test, job.kickoff_log, nodes=picked_nodes) script.write(job.kickoff_path) test.job = job @@ -722,7 +695,7 @@ def _schedule_indi_chunk(self, pav_cfg, tests: List[TestRun], pav_cfg=pav_cfg, job=job, sched_config=sched_config, - job_name=job_name, + job_name=self._job_name(test), nodes=picked_nodes) except SchedulerPluginError as err: return [self._make_kickoff_error(err, [test])] @@ -737,3 +710,70 @@ def _schedule_indi_chunk(self, pav_cfg, tests: List[TestRun], .format(self.name, len(test_chunk))) return errors + + def create_kickoff_script(self, + pav_cfg: PavConfig, + tests: Union[TestRun, List[TestRun]], + log_path: Optional[Path] = None, + nodes: Optional[NodeSet] = None, + isolate: bool = False) -> ScriptComposer: + """Create the kickoff script.""" + + if not isinstance(tests, list): + tests = [tests] + + sched_config = validate_config(tests[0].config['schedule']) + time_limits = [t.config["schedule"]["time_limit"] for t in tests \ + if t.config["schedule"]["time_limit"] is not None] + + if len(time_limits) > 0: + sched_config["time_limit"] = max(time_limits) + + job_name = self._job_name(tests) + + if isolate: + job_name = job_name + self.ISOLATE_KICKOFF_SUFFIX + + if nodes is None: + node_range = calc_node_range( + sched_config, + sched_config['cluster_info']['node_count']) + else: + node_range = None + + script = self._create_kickoff_script_stub( + pav_cfg=pav_cfg, + job_name=job_name, + log_path=log_path, + sched_config=sched_config, + node_range=node_range, + nodes=nodes, + shebang=tests[0].shebang, + isolate=isolate) + + test_ids = ' '.join(test.full_id for test in tests) + + script.newline() + script.command('echo "Starting {} tests - $(date)"'.format(len(tests))) + + script.newline() + + if isolate: + script = tests[0].make_script(script, "run", isolate=True) + else: + script.command('pav _run {}'.format(test_ids)) + + return script + + def _job_name(self, tests: Union[TestRun, List[TestRun]]) -> str: + """Given a test, get the name of the job.""" + + if not isinstance(tests, list): + tests = [tests] + + job_name = 'pav_{}'.format(','.join(test.name for test in tests[:4])) + + if len(tests) > 4: + job_name += ' ...' + + return job_name diff --git a/lib/pavilion/schedulers/basic.py b/lib/pavilion/schedulers/basic.py index 381023666..de2408cfb 100644 --- a/lib/pavilion/schedulers/basic.py +++ b/lib/pavilion/schedulers/basic.py @@ -3,12 +3,15 @@ from abc import ABC from collections import defaultdict -from typing import List +from pathlib import Path +from typing import List, Dict, Tuple, Optional, Union +from pavilion.config import PavConfig from pavilion.jobs import Job, JobError from pavilion.status_file import STATES from pavilion.test_run import TestRun -from pavilion.types import NodeInfo, Nodes +from pavilion.scriptcomposer import ScriptComposer +from pavilion.types import NodeInfo, Nodes, NodeSet from .config import validate_config, calc_node_range from .scheduler import SchedulerPlugin from ..errors import SchedulerPluginError @@ -92,21 +95,10 @@ def schedule_tests(self, pav_cfg, tests: List[TestRun]) -> List[SchedulerPluginE prior_error=err, tests=[test])) continue - job_name = 'pav-{}-{}-runs'.format(self.name, test_bin[0].series) - for test in test_bin: test.job = job - script = self._create_kickoff_script_stub( - pav_cfg=pav_cfg, - job_name=job_name, - log_path=job.kickoff_log, - sched_config=sched_config, - node_range=node_range, - shebang=test.shebang) - - test_ids = ' '.join(test.full_id for test in tests) - script.command('pav _run {}'.format(test_ids)) + script = self.create_kickoff_script(pav_cfg, test_bin, job.kickoff_log) script.write(job.kickoff_path) try: @@ -114,7 +106,7 @@ def schedule_tests(self, pav_cfg, tests: List[TestRun]) -> List[SchedulerPluginE pav_cfg=pav_cfg, job=job, sched_config=sched_config, - job_name=job_name, + job_name=self._job_name(test_bin), node_range=node_range) except SchedulerPluginError as err: errors.append(self._make_kickoff_error(err, [test])) @@ -131,3 +123,58 @@ def schedule_tests(self, pav_cfg, tests: List[TestRun]) -> List[SchedulerPluginE "Test kicked off with the {} scheduler".format(self.name)) return errors + + def _job_name(self, tests: List[TestRun]) -> str: + """Given a test, get the name of the job.""" + + return 'pav-{}-{}-runs'.format(self.name, tests[0].series) + + def create_kickoff_script(self, + pav_cfg: PavConfig, + tests: Union[TestRun, List[TestRun]], + log_path: Optional[Path] = None, + nodes: Optional[NodeSet] = None, + isolate: bool = False) -> ScriptComposer: + """Create the kickoff script.""" + + if not isinstance(tests, list): + tests = [tests] + + sched_config = validate_config(tests[0].config['schedule']) + + if nodes is None: + node_range = calc_node_range( + sched_config, + sched_config['cluster_info']['node_count']) + else: + node_range = None + + job_name = self._job_name(tests) + + if isolate: + job_name = job_name + self.ISOLATE_KICKOFF_SUFFIX + + script = self._create_kickoff_script_stub( + pav_cfg=pav_cfg, + job_name=job_name, + log_path=log_path, + sched_config=sched_config, + node_range=node_range, + nodes=nodes, + shebang=tests[0].shebang, + isolate=isolate) + + test_ids = ' '.join(test.full_id for test in tests) + + # This is commented out for consistency with prior behavior and with expected output for + # logging unit tests. We may want to consider adding it. — HW + # script.command('echo "Starting {} tests - $(date)"'.format(len(tests))) + + script.newline() + + if isolate: + script = tests[0].make_script(script, "run", isolate=True) + else: + script.command('pav _run {}'.format(test_ids)) + + return script diff --git a/lib/pavilion/schedulers/scheduler.py b/lib/pavilion/schedulers/scheduler.py index ed08420b1..c938efe49 100644 --- a/lib/pavilion/schedulers/scheduler.py +++ b/lib/pavilion/schedulers/scheduler.py @@ -5,10 +5,12 @@ import inspect import os import time +from abc import abstractmethod from pathlib import Path -from typing import List, Union, Dict, NewType, Tuple, Type +from typing import List, Union, Dict, NewType, Tuple, Type, Optional import yaml_config as yc +from pavilion.config import PavConfig from pavilion.jobs import JobError, JobInfo, Job from pavilion.scriptcomposer import ScriptHeader, ScriptComposer from pavilion.status_file import STATES, TestStatusInfo @@ -49,8 +51,13 @@ def __init__(self, job_name: str, sched_config: dict, if nodes is None: self._include_nodes = self._config['include_nodes'] self._exclude_nodes = self._config['exclude_nodes'] - self._node_min = node_range[0] - self._node_max = node_range[1] + + if node_range is not None: + self._node_min = node_range[0] + self._node_max = node_range[1] + else: + self._node_min = None + self._node_max = None else: self._include_nodes = nodes # Any nodes in the exclude list will have already been filtered out. @@ -94,6 +101,10 @@ class SchedulerPlugin(IPlugin.IPlugin): KICKOFF_FN = None """If the kickoff script requires a special filename, set it here.""" + KICKOFF_LOG_DEFAULT_FN = "kickoff.log" + + ISOLATE_KICKOFF_SUFFIX = "_isolated" + VAR_CLASS = SchedulerVariables # type: Type[SchedulerVariables] """The scheduler's variable class.""" @@ -417,12 +428,26 @@ def _get_config_elems(self) -> Tuple[List[yc.ConfigElement], dict, dict]: return [], {}, {} - def _create_kickoff_script_stub(self, pav_cfg, job_name: str, log_path: Path, + @abstractmethod + def create_kickoff_script(self, + pav_cfg: PavConfig, + tests: Union[TestRun, List[TestRun]], + log_path: Optional[Path] = None, + nodes: Optional = None, + isolate: bool = False) -> ScriptComposer: + """Create the kickoff script.""" + + raise NotImplementedError + + def _create_kickoff_script_stub(self, + pav_cfg: PavConfig, + job_name: str, sched_config: dict, + log_path: Optional[Path] = None, nodes: Union[NodeList, None] = None, node_range: Union[Tuple[int, int], None] = None, - shebang: str = None)\ - -> ScriptComposer: + shebang: str = None, + isolate: bool = False) -> ScriptComposer: """Generate the kickoff script essentials preamble common to all scheduled tests. @@ -445,17 +470,24 @@ def _create_kickoff_script_stub(self, pav_cfg, job_name: str, log_path: Path, script = ScriptComposer(header=header) script.comment("Redirect all output to the kickoff log.") - script.command("exec >{} 2>&1".format(log_path.as_posix())) - # Make sure the pavilion spawned - env_changes = { - 'PATH': '{}:${{PATH}}'.format(pav_cfg.pav_root / 'bin'), - 'PAV_CONFIG_FILE': str(pav_cfg.pav_cfg_file), - } - if 'PAV_CONFIG_DIR' in os.environ: - env_changes['PAV_CONFIG_DIR'] = os.environ['PAV_CONFIG_DIR'] + if log_path is not None: + script.command(f"exec >{log_path.as_posix()} 2>&1") + else: + script.command( + f'exec > $(dirname -- ${{BASH_SOURCE[0]}})/{self.KICKOFF_LOG_DEFAULT_FN} 2>&1') + + if not isolate: + script.newline() + # Make sure the pavilion spawned + env_changes = { + 'PATH': '{}:${{PATH}}'.format(pav_cfg.pav_root / 'bin'), + 'PAV_CONFIG_FILE': str(pav_cfg.pav_cfg_file), + } + if 'PAV_CONFIG_DIR' in os.environ: + env_changes['PAV_CONFIG_DIR'] = os.environ['PAV_CONFIG_DIR'] - script.env_change(env_changes) + script.env_change(env_changes) # Run Kickoff Env setup commands for command in pav_cfg.env_setup: @@ -572,6 +604,12 @@ def convert_lists_to_tuples(obj): return tuple(key_parts) + @abstractmethod + def _job_name(self, tests: Union[TestRun, List[TestRun]]) -> str: + """Given a test, get the name of the job.""" + + raise NotImplementedError + def __reset(): """This exists for testing purposes only.""" diff --git a/lib/pavilion/test_run/test_run.py b/lib/pavilion/test_run/test_run.py index a3b5eec6c..ffb7322e2 100644 --- a/lib/pavilion/test_run/test_run.py +++ b/lib/pavilion/test_run/test_run.py @@ -100,6 +100,9 @@ class TestRun(TestAttributes): BUILD_TEMPLATE_DIR = 'templates' """Directory that holds build templates.""" + PAV_LIB_FN = "pav-lib.bash" + """Pavilion bash utilities""" + def __init__(self, pav_cfg: PavConfig, config: Dict, var_man: VariableSetManager = None, _id: int = None, rebuild: bool = False, build_only: bool = False): """Create an new TestRun object. If loading an existing test @@ -319,12 +322,12 @@ def save(self) -> None: self.status.set(STATES.CREATED, "Test directory and status file created.") + header = scriptcomposer.ScriptHeader(shebang=self.shebang) + script = scriptcomposer.ScriptComposer(header=header) + if self._build_needed(): - self._write_script( - 'build', - path=self.build_script_path, - config=self.config.get('build', {}), - module_wrappers=self.config.get('module_wrappers', {})) + script = self.make_script(script, 'build') + script.write(self.build_script_path) self.builder = self._make_builder() self.build_name = self.builder.name @@ -333,11 +336,8 @@ def save(self) -> None: # process of creating and using a builder. self._build_trivial() - self._write_script( - 'run', - path=self.run_tmpl_path, - config=self.config.get('run', {}), - module_wrappers=self.config.get('module_wrappers', {})) + script = self.make_script(script, 'run') + script.write(self.run_tmpl_path) self.save_attributes() self.status.set(STATES.CREATED, "Test directory setup complete.") @@ -521,12 +521,11 @@ def finalize(self, new_vars: VariableSetManager): self.save_attributes() - self._write_script( - 'run', - self.run_script_path, - self.config['run'], - self.config.get('module_wrappers', {}) - ) + header = scriptcomposer.ScriptHeader(shebang=self.shebang) + script = scriptcomposer.ScriptComposer(header=header) + + script = self.make_script(script, 'run') + script.write(self.run_script_path) self.status.set(STATES.FINALIZED, "Test Run Finalized.") @@ -1119,7 +1118,10 @@ def complete_time(self): .format(run_complete_path.as_posix(), err)) return None - def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: dict): + def make_script(self, + script: scriptcomposer.ScriptComposer, + stype: str, + isolate: bool = False) -> scriptcomposer.ScriptComposer: """Write a build or run script or template. The formats for each are mostly identical. :param stype: The type of script (run or build). @@ -1128,9 +1130,8 @@ def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: d :param module_wrappers: The module wrappers definition. """ - header = scriptcomposer.ScriptHeader(shebang=self.shebang) - script = scriptcomposer.ScriptComposer(header=header) - + config = self.config[stype] + module_wrappers = self.config.get('module_wrappers', {}) verbose = config.get('verbose', 'false').lower() == 'true' if verbose: @@ -1138,17 +1139,26 @@ def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: d script.command('set -v') script.newline() - pav_lib_bash = self._pav_cfg.pav_root/'bin'/'pav-lib.bash' - script.command(f'echo "(pav) Starting {stype} script"') + script.newline() + # If we include this directly, it breaks build hashing. script.comment('The first (and only) argument of the build script is ' 'the test id.') - script.env_change({ - 'TEST_ID': '${1:-0}', # Default to test id 0 if one isn't given. - 'PAV_CONFIG_FILE': self._pav_cfg['pav_cfg_file'] - }) + + env = {'TEST_ID': '${1:-0}'} # Default to test id 0 if one isn't given. + + if not isolate: + env["PAV_CONFIG_FILE"] = self._pav_cfg['pav_cfg_file'] + + script.env_change(env) + + if isolate: + pav_lib_bash = f'$( dirname -- "${{BASH_SOURCE[0]}}" )/{self.PAV_LIB_FN}' + else: + pav_lib_bash = self._pav_cfg.pav_root / 'bin' / self.PAV_LIB_FN + script.command('source {}'.format(pav_lib_bash)) if config.get('preamble', []): @@ -1160,6 +1170,7 @@ def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: d if stype == 'build' and not self.build_local: script.comment('To be built in an allocation.') + script.newline() script.command(f'echo "(pav) Setting up {stype} environment."') purge = utils.str_bool(config.get("purge_modules")) @@ -1191,8 +1202,6 @@ def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: d script.comment('List all the module modules for posterity') script.command("module -t list") script.newline() - script.comment('Output the environment for posterity') - script.command("declare -p") if self.spack_enabled(): script.command(f'echo "(pav) Setting up Spack."') @@ -1229,6 +1238,17 @@ def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: d script.command('spack load {} || exit 1' .format(package)) + if not isolate: + script.newline() + script.comment('Output the environment for posterity') + + if verbose: + script.command( + f'declare -p | tee > $( dirname -- "${{BASH_SOURCE[0]}}" )/{stype}.env.sh') + else: + script.command(f'declare -p > $( dirname -- "${{BASH_SOURCE[0]}}" )/{stype}.env.sh') + + script.newline() script.command(f'echo "(pav) Executing {stype} commands."') script.newline() cmds = config.get('cmds', []) @@ -1244,9 +1264,10 @@ def _write_script(self, stype: str, path: Path, config: dict, module_wrappers: d else: script.comment('No commands given for this script.') + script.newline() script.command(f'echo "(pav) Test {stype} commands completed without error."') - script.write(path) + return script def __repr__(self): return "TestRun({s.name}-{s.full_id})".format(s=self) diff --git a/lib/pavilion/utils.py b/lib/pavilion/utils.py index 877384803..c140126d2 100644 --- a/lib/pavilion/utils.py +++ b/lib/pavilion/utils.py @@ -13,8 +13,7 @@ import textwrap import zipfile from pathlib import Path -from typing import Iterator, Union, TextIO -from typing import List, Dict +from typing import Iterator, Union, TextIO, List, Dict, Optional, Iterable class WrappedFormatter(argparse.HelpFormatter): @@ -208,6 +207,87 @@ def copytree(src, dst, symlinks=False, ignore=None, copy_function=shutil.copy2, raise shutil.Error(errors) return dst +def copytree_resolved( + src: Path, + dest: Path, + src_root: Optional[Path] = None, + dest_root: Optional[Path] = None, + symlinks: Optional[Dict[Path, Path]] = None, + ignore_files: Optional[Iterable[str]] = None) -> None: + """Copy a directory tree to another location, such that the only symlinks that remain are + symlinks internal to the directory.""" + + src_root = src_root or src + dest_root = dest_root or dest + ignore_files = ignore_files or [] + + if symlinks is None: + symlinks = {} + + if src.name in ignore_files: + return + + if src.is_symlink(): + try: + resolved = src.resolve() + except RuntimeError: + # There is a circular symlink + return + + if resolved in symlinks: + # Create a relative symlink + dest.symlink_to( + Path( + os.path.relpath( + symlinks.get(resolved), + dest.resolve().parent))) + else: + # Only recreate symlinks if they are internal to the source directory + target_in_tree = True + + try: + rel_target = resolved.relative_to(src_root) + except ValueError: + target_in_tree = False + + if target_in_tree: + # Don't create the symlink if it points inside a directory we're ignoring + skip_link = False + + for part in rel_target.parts: + if part in ignore_files: + skip_link = True + break + + if not skip_link: + # Create a relative symlink + dest.symlink_to( + Path( + os.path.relpath( + dest_root / rel_target, + dest.resolve().parent))) + symlinks[resolved] = dest_root / rel_target + else: + symlinks[resolved] = dest + copytree_resolved(resolved, dest, src_root, dest_root, symlinks, ignore_files) + + return + + elif src.is_file(): + shutil.copy(src, dest) + + return + elif src.is_dir(): + dest.mkdir(exist_ok=True) + + # Sort for reproduceability + files = sorted(src.iterdir(), key=lambda p: p.name) + + for fname in files: + copytree_resolved(fname, dest / fname.name, src_root, dest_root, + symlinks, ignore_files) + + return def path_is_external(path: Path): """Returns True if a path contains enough back 'up-references' to escape diff --git a/test/tests/isolate_cmd_tests.py b/test/tests/isolate_cmd_tests.py new file mode 100644 index 000000000..e55944f3f --- /dev/null +++ b/test/tests/isolate_cmd_tests.py @@ -0,0 +1,98 @@ +import tempfile +import os +import tarfile +from pathlib import Path +import subprocess as sp +from typing import Iterator + +from pavilion import commands +from pavilion import arguments +from pavilion.unittest import PavTestCase +from pavilion.cmd_utils import list_files +from pavilion.test_run import TestRun + + +class IsolateCmdTests(PavTestCase): + + def test_no_archive(self): + """Test that isolating without archiving works correctly.""" + + run_cmd = commands.get_command("run") + isolate_cmd = commands.get_command("isolate") + + run_cmd.silence() + isolate_cmd.silence() + + parser = arguments.get_parser() + run_args = ["run", "-H", "this", "hello_world.hello"] + + run_cmd.run(self.pav_cfg, parser.parse_args(run_args)) + last_test = next(iter(run_cmd.last_tests)) + + last_test.wait(timeout=10) + + with tempfile.TemporaryDirectory() as dir: + isolate_args = parser.parse_args(["isolate", str(Path(dir) / "dest")]) + + self.assertEqual(isolate_cmd.run(self.pav_cfg, isolate_args), 0) + + source_files = set(map( + lambda x: x.relative_to(last_test.path), + list_files(last_test.path))) + dest_files = set(map( + lambda x: x.relative_to(Path(dir) / "dest"), + list_files(Path(dir) / "dest"))) + + self.assertIn(Path(TestRun.PAV_LIB_FN), dest_files) + self.assertIn(Path(isolate_cmd.KICKOFF_FN), dest_files) + + self.assertEqual( + {f for f in source_files if f.name not in ("series", "job")}, + {f for f in dest_files if f.name not in \ + (TestRun.PAV_LIB_FN, isolate_cmd.KICKOFF_FN) and \ + # Have to exclude build_origin because of how os.walk handles symlinks + f.parent.name != "build_origin"}) + + def test_zip_archive(self): + """Test that isolating using a compressed archive works correctly.""" + + run_cmd = commands.get_command("run") + isolate_cmd = commands.get_command("isolate") + + run_cmd.silence() + isolate_cmd.silence() + + parser = arguments.get_parser() + run_args = ["run", "-H", "this", "hello_world.hello"] + + run_cmd.run(self.pav_cfg, parser.parse_args(run_args)) + last_test = next(iter(run_cmd.last_tests)) + + with tempfile.TemporaryDirectory() as dir: + isolate_args = parser.parse_args(["isolate", + str(Path(dir) / "dest"), + "--archive", + "--zip"]) + + self.assertEqual(isolate_cmd.run(self.pav_cfg, isolate_args), 0) + + with tempfile.TemporaryDirectory() as extract_dir: + with tarfile.open(Path(dir) / "dest.tgz", "r:gz") as tf: + tf.extractall(extract_dir) + + source_files = set(map( + lambda x: Path(x).relative_to(last_test.path), + list_files(last_test.path, include_root=True))) + dest_files = set(map( + lambda x: x.relative_to(Path(extract_dir) / "dest"), + list_files(Path(extract_dir)))) + + self.assertIn(Path(TestRun.PAV_LIB_FN), dest_files) + self.assertIn(Path(isolate_cmd.KICKOFF_FN), dest_files) + + self.assertEqual( + {f for f in source_files if f.name not in ("series", "job")}, + {f for f in dest_files if f.name not in \ + ("pav-lib.bash", "kickoff.isolated") and \ + # Have to exclude build_origin because of how os.walk handles symlinks + f.parent.name != "build_origin"}) diff --git a/test/tests/utils_tests.py b/test/tests/utils_tests.py index fe0706492..877f3c453 100644 --- a/test/tests/utils_tests.py +++ b/test/tests/utils_tests.py @@ -9,6 +9,7 @@ from pavilion import unittest from pavilion import utils +from pavilion.cmd_utils import list_files class UtilsTests(unittest.PavTestCase): @@ -155,4 +156,106 @@ def test_copytree_dotfiles(self): utils.copytree(src, dest) names = (map(lambda x: x.name, dest.iterdir())) - self.assertIn(".dotfile", names) \ No newline at end of file + self.assertIn(".dotfile", names) + + def test_copytree_resolved(self): + examples = [ + { + "copy_root": "foo", + "files": [ + {"name": "foo", "dir": True, "target": None}, + {"name": "bar", "dir": False, "target": None}, + {"name": "foo/baz", "dir": False, "target": "bar"}, + ], + "expected": [ + {"name": "baz", "dir": False, "target": None}, + ] + }, + { + "copy_root": None, + "files": [ + {"name": "foo", "dir": False, "target": "bar"}, + {"name": "bar", "dir": False, "target": "foo"}, + ], + "expected": [] + }, + { + "copy_root": "foo", + "files": [ + {"name": "foo", "dir": True, "target": None}, + {"name": "foo/bar", "dir": False, "target": "foobar"}, + {"name": "foo/baz", "dir": False, "target": None}, + {"name": "foobar", "dir": False, "target": "foo/baz"} + ], + "expected": [ + {"name": "bar", "dir": False, "target": "baz"}, + {"name": "baz", "dir": False, "target": None}, + ] + }, + { + "copy_root": None, + "files": [ + {"name": "foo", "dir": True, "target": None}, + {"name": "bar", "dir": False, "target": None}, + {"name": "foo/baz", "dir": False, "target": "bar"}, + {"name": "foobar", "dir": False, "target": "bar"}, + ], + "expected": [ + {"name": "foo", "dir": True, "target": None}, + {"name": "bar", "dir": False, "target": None}, + {"name": "foo/baz", "dir": False, "target": "bar"}, + {"name": "foobar", "dir": False, "target": "bar"}, + ] + }, + { + "copy_root": "foo", + "files": [ + {"name": "foo", "dir": True, "target": None}, + {"name": "bar", "dir": False, "target": None}, + {"name": "foo/baz", "dir": False, "target": "bar"}, + {"name": "foo/foobar", "dir": False, "target": "bar"}, + ], + "expected": [ + {"name": "baz", "dir": False, "target": None}, + {"name": "foobar", "dir": False, "target": "baz"}, + ] + } + ] + + for i, ex in enumerate(examples): + with tempfile.TemporaryDirectory() as src: + src = Path(src) + + with tempfile.TemporaryDirectory() as dest: + dest = Path(dest) + + # Create the files. + for f in ex["files"]: + file_path = src / f["name"] + + if f["dir"]: + file_path.mkdir(parents=True) + else: + if f["target"] is None: + file_path.touch() + else: + target_path = src / f["target"] + file_path.symlink_to(target_path) + + if ex["copy_root"] is None: + utils.copytree_resolved(src, dest) + else: + utils.copytree_resolved(src / ex["copy_root"], dest) + + expected = set(Path(f["name"]) for f in ex["expected"]) + actual = set(p.relative_to(dest) for p in list_files(dest)) + + self.assertEqual(expected, actual) + + for f in ex["expected"]: + for g in actual: + if Path(f["name"]) == g: + if f["target"] is None: + self.assertFalse((dest / g).is_symlink()) + else: + self.assertTrue((dest / g).is_symlink())