diff --git a/toolchain/mfc/args.py b/toolchain/mfc/args.py index 65f3adba04..a7193a6f0f 100644 --- a/toolchain/mfc/args.py +++ b/toolchain/mfc/args.py @@ -65,11 +65,14 @@ def add_common_arguments(p, mask = None): for target in DEPENDENCY_TARGETS: p.add_argument(f"--no-{target.name}", action="store_true", help=f"Do not build the {target.name} dependency. Use the system's instead.") + if "g" not in mask: + p.add_argument("-g", "--gpus", nargs="+", type=int, default=[0], help="(GPU) List of GPU #s to use.") + # === BUILD === - add_common_arguments(build) + add_common_arguments(build, "g") # === CLEAN === - add_common_arguments(clean, "j") + add_common_arguments(clean, "jg") binaries = [ b.bin for b in BINARIES ] @@ -82,7 +85,6 @@ def add_common_arguments(p, mask = None): test.add_argument("-b", "--binary", choices=binaries, type=str, default=None, help="(Serial) Override MPI execution binary") test.add_argument("-r", "--relentless", action="store_true", default=False, help="Run all tests, even if multiple fail.") test.add_argument("-a", "--test-all", action="store_true", default=False, help="Run the Post Process Tests too.") - test.add_argument("-g", "--gpus", type=str, default="0", help="(GPU) Comma separated list of GPU #s to use.") test.add_argument("-%", "--percent", type=int, default=100, help="Percentage of tests to run.") test.add_argument("-m", "--max-attempts", type=int, default=3, help="Maximum number of attempts to run a test.") @@ -120,7 +122,7 @@ def add_common_arguments(p, mask = None): add_common_arguments(bench, "t") # === COUNT === - add_common_arguments(count) + add_common_arguments(count, "g") args: dict = vars(parser.parse_args()) @@ -150,8 +152,4 @@ def append_defaults_to_data(name: str, parser): if args[e] is not None: args[e] = os.path.abspath(args[e]) - # Turn GPU ID list into a comma separated string - if "gpus" in args: - args["gpus"] = [int(g) for g in args["gpus"].split(",")] - return args diff --git a/toolchain/mfc/common.py b/toolchain/mfc/common.py index 888733f7a9..d0e718dbd8 100644 --- a/toolchain/mfc/common.py +++ b/toolchain/mfc/common.py @@ -29,9 +29,12 @@ class MFCException(Exception): pass -def system(command: typing.List[str], no_exception: bool = False, exception_text=None, on_error=lambda: None, cwd=None, stdout=None, stderr=None) -> int: +def system(command: typing.List[str], no_exception: bool = False, exception_text=None, on_error=lambda: None, cwd=None, stdout=None, stderr=None, env: dict = None) -> int: cmd = [ str(x) for x in command if not isspace(str(x)) ] + if env is None: + env = os.environ.copy() + if stdout != subprocess.DEVNULL: cons.print(no_indent=True) @@ -40,7 +43,7 @@ def system(command: typing.List[str], no_exception: bool = False, exception_text if stdout != subprocess.DEVNULL: cons.print(no_indent=True) - r = subprocess.run(cmd, cwd=cwd, stdout=stdout, stderr=stderr) + r = subprocess.run(cmd, cwd=cwd, stdout=stdout, stderr=stderr, env=env) if r.returncode != 0: on_error() diff --git a/toolchain/mfc/count.py b/toolchain/mfc/count.py index 865a911e44..4ffabdf289 100644 --- a/toolchain/mfc/count.py +++ b/toolchain/mfc/count.py @@ -1,6 +1,7 @@ import os, glob, typing, typing -from .common import MFC_ROOTDIR +from .state import ARG +from .common import MFC_ROOTDIR, format_list_to_string from .printer import cons import rich.table @@ -20,12 +21,13 @@ def handle_dir(dirpath: str) -> typing.Tuple[typing.List[typing.Tuple[str, int]] return (files, total) def count(): - cons.print("[bold]Counting lines of code in [magenta]MFC[/magenta][/bold] (excluding whitespace lines)") - cons.print() + target_str_list = format_list_to_string(ARG('targets'), 'magenta') + + cons.print(f"[bold]Counting lines of code in {target_str_list}[/bold] (excluding whitespace lines)") cons.indent() total = 0 - for codedir in ['common', 'pre_process', 'simulation', 'post_process']: + for codedir in ['common'] + ARG("targets"): dirfiles, dircount = handle_dir(os.path.join(MFC_ROOTDIR, 'src', codedir)) table = rich.table.Table(show_header=True, box=rich.table.box.SIMPLE) table.add_column(f"File (in [magenta]{codedir}[/magenta])", justify="left") @@ -38,7 +40,7 @@ def count(): cons.raw.print(table) - cons.print(f"[bold]Total MFC lines: [bold cyan]{total}[/bold cyan].[/bold]") + cons.print(f"[bold]Total {target_str_list} lines: [bold cyan]{total}[/bold cyan].[/bold]") cons.print() cons.unindent() diff --git a/toolchain/mfc/run/engines.py b/toolchain/mfc/run/engines.py index 066c044469..b8c45f772f 100644 --- a/toolchain/mfc/run/engines.py +++ b/toolchain/mfc/run/engines.py @@ -74,7 +74,6 @@ def get_args(self) -> str: def get_exec_cmd(self, target: MFCTarget) -> typing.List[str]: cmd = [] - if ARG("mpi"): cmd += [self.mpibin.bin] + self.mpibin.gen_params() + ARG("flags")[:] @@ -148,7 +147,13 @@ def run(self, targets: typing.List[MFCTarget]) -> None: if not ARG("dry_run"): start_time = time.monotonic() - system(self.get_exec_cmd(target), cwd=self.input.case_dirpath) + system( + self.get_exec_cmd(target), cwd=self.input.case_dirpath, + env={ + **os.environ.copy(), + 'CUDA_VISIBLE_DEVICES': ','.join([str(_) for _ in ARG('gpus')]) + } + ) end_time = time.monotonic() cons.print(no_indent=True) diff --git a/toolchain/mfc/run/run.py b/toolchain/mfc/run/run.py index a1ea94e1f0..47fcd2fe22 100644 --- a/toolchain/mfc/run/run.py +++ b/toolchain/mfc/run/run.py @@ -67,3 +67,12 @@ def run_target(target: MFCTarget): def run() -> None: run_targets([ get_target(_) for _ in ARG("targets")]) + + +def run_targets_with(targets: typing.List[MFCTarget]): + pass + + +def run_target_with(target: MFCTarget): + run_targets_with([target]) + diff --git a/toolchain/mfc/test/case.py b/toolchain/mfc/test/case.py index adc868878f..d839274928 100644 --- a/toolchain/mfc/test/case.py +++ b/toolchain/mfc/test/case.py @@ -78,7 +78,6 @@ 'sigV' : 0.1, 'rhoRV' : 0.0, - 'Monopole' : 'F', 'num_mono' : 1, 'Mono(1)%loc(1)' : 0.5,