diff --git a/source/fab/api.py b/source/fab/api.py index e16cf681..4a18ce0c 100644 --- a/source/fab/api.py +++ b/source/fab/api.py @@ -28,6 +28,7 @@ from fab.steps.link import link_exe, link_shared_object from fab.steps.preprocess import preprocess_c, preprocess_fortran from fab.steps.psyclone import preprocess_x90, psyclone +from fab.steps.psyclone_transmute import psyclone_transmute from fab.steps.root_inc_files import root_inc_files from fab.tools.category import Category from fab.tools.compiler import Compiler, Ifort @@ -88,6 +89,7 @@ "preprocess_fortran", "preprocess_x90", "psyclone", + "psyclone_transmute", "root_inc_files", "run_mp", "Shell", diff --git a/source/fab/artefacts.py b/source/fab/artefacts.py index 5a66febc..8553e409 100644 --- a/source/fab/artefacts.py +++ b/source/fab/artefacts.py @@ -18,7 +18,7 @@ from collections import defaultdict from enum import auto, Enum from pathlib import Path -from typing import Iterable, Optional, Union +from typing import Iterable, Optional, Sequence, Union from fab.dep_tree import filter_source_tree, AnalysedDependent from fab.util import suffix_filter @@ -115,8 +115,8 @@ def copy_artefacts(self, source: Union[str, ArtefactSet], self.add(dest, self[source]) def replace(self, artefact: Union[str, ArtefactSet], - remove_files: list[Union[str, Path]], - add_files: Union[list[Union[str, Path]], dict]): + remove_files: Union[Sequence[str], Sequence[Path]], + add_files: Union[Sequence[str], Sequence[Path]]): '''Replaces artefacts in one artefact set with other artefacts. This can be used e.g to replace files that have been preprocessed and renamed. There is no requirement for these lists to have the diff --git a/source/fab/steps/psyclone.py b/source/fab/steps/psyclone.py index 8a1e2515..e87361eb 100644 --- a/source/fab/steps/psyclone.py +++ b/source/fab/steps/psyclone.py @@ -317,7 +317,7 @@ def do_one_file(arg: tuple[Path, MpCommonArgs]): psyclone = config.tool_box.get_tool(Category.PSYCLONE) if not isinstance(psyclone, Psyclone): raise RuntimeError(f"Unexpected tool '{psyclone.name}' of type " - f"'{type(psyclone)}' instead of Psyclone") + f"'{type(psyclone)}' instead of PSyclone") try: transformation_script = mp_payload.transformation_script logger.info(f"running psyclone on '{x90_file}'.") diff --git a/source/fab/steps/psyclone_transmute.py b/source/fab/steps/psyclone_transmute.py new file mode 100644 index 00000000..682acfe1 --- /dev/null +++ b/source/fab/steps/psyclone_transmute.py @@ -0,0 +1,258 @@ +# ############################################################################## +# (c) Crown copyright Met Office. All rights reserved. +# For further details please refer to the file COPYRIGHT +# which you should have received as part of this distribution +# ############################################################################## +""" +A preprocessor and code generation step using PSyclone's transmute (Fortran to +Fortran) ability. +https://github.com/stfc/PSyclone + +""" +from dataclasses import dataclass +import logging +import shutil +import warnings +from pathlib import Path +from typing import Callable, cast, Optional, Sequence, Union + + +from fab.build_config import BuildConfig + +from fab.artefacts import ArtefactSet +from fab.steps import run_mp, check_for_errors, step +from fab.tools.category import Category +from fab.tools.psyclone import Psyclone +from fab.util import (log_or_dot, input_to_output_fpath, file_checksum, + file_walk, TimerLogger, string_checksum, + log_or_dot_finish) + +logger = logging.getLogger(__name__) + + +@dataclass +class MpCommonArgs: + """ + Runtime data for child processes to read. + + Contains data used to calculate the prebuild hash. + + """ + config: BuildConfig + suffix: str + transformation_script: Optional[Callable[[Path, BuildConfig], Path]] + cli_args: list[str] + overrides_folder: Optional[Path] + # filenames (not paths) of hand crafted overrides + override_files: list[str] + + +@step +def psyclone_transmute( + config: BuildConfig, + fortran_files: Union[Sequence[Path], Sequence[Path]], + transformation_script: Optional[Callable[[Path, + BuildConfig], Path]] = None, + cli_args: Optional[list[str]] = None, + suffix: Optional[str] = None, + overrides_folder: Optional[Path] = None, + artefact_set: Optional[ArtefactSet] = None, + ): + """ + PSyclone runner step. + + .. note:: + + This step reads pre-processed Fortran files and produces replacement + Fortran files. So it must be run before the + :class:`~fab.steps.analyse.Analyse` step. + + This step stores results as prebuilds to speed up subsequent builds. + To generate the prebuild hashes, it analyses the files, storing prebuilt + results for these also. + + :param config: + The :class:`fab.build_config.BuildConfig` object where we can read + settings such as the project workspace folder or the multiprocessing + flag. + :param fortran_files: list of files to transform. + :param transformation_script: + The function to get Python transformation script. + It takes in a file path and the config object, and returns the path + of the transformation script or None. If no function is given or the + function returns None, no script will be applied and PSyclone still + runs. + :param cli_args: + Passed through to the psyclone cli tool. + :param overrides_folder: + Optional folder containing hand-crafted override files. + Must be part of the subsequently analysed source code. + Any file produced by psyclone will be deleted if there is a + corresponding file in this folder. + :param suffix: a suffix to be added to create the new filename. + :param artefact_set: an optional artefact set. If specified, the + input files names will be replaced with the newly transmuted ones. + """ + + # We need to allow "" as suffix indicating to overwrite the source file. + if suffix is None: + suffix = "_transmute" + + cli_args = cli_args or [] + + # get the data in a payload object for child processes to calculate + # prebuild hashes + mp_payload = _generate_mp_payload(config, overrides_folder, + transformation_script, cli_args, suffix) + + config.prebuild_folder.mkdir(parents=True, exist_ok=True) + + # Run PSyclone. For every file, we get back a tuple of the output file and + # the prebuild + mp_arg = [(fortran_file, mp_payload) for fortran_file in fortran_files] + with TimerLogger(f"running PSyclone transmute on {len(fortran_files)} " + f"Fortran files"): + results = run_mp(config, mp_arg, transmute_one_file) + log_or_dot_finish(logger) + outputs, prebuilds = zip(*results) if results else ((), ()) + # This call will abort in case of an error + check_for_errors(outputs, caller_label='psyclone') + + if artefact_set: + config.artefact_store.replace( + artefact_set, + remove_files=fortran_files, + add_files=outputs) + + # record the output files in the artefact store for further processing + config.artefact_store.add(ArtefactSet.FORTRAN_COMPILER_FILES, outputs) + outputs_str = "\n".join(map(str, outputs)) + logger.debug(f'psyclone outputs:\n{outputs_str}\n') + + # Mark the prebuilds as being current so the + # cleanup step doesn't delete them + config.add_current_prebuilds(prebuilds) + prebuilds_str = "\n".join(map(str, prebuilds)) + logger.debug(f'psyclone prebuilds:\n{prebuilds_str}\n') + + +def _generate_mp_payload(config, + overrides_folder, + transformation_script, + cli_args, + suffix: str) -> MpCommonArgs: + override_files: list[str] = [] + if overrides_folder: + override_files = [f.name for f in file_walk(overrides_folder)] + + return MpCommonArgs( + config=config, + transformation_script=transformation_script, + cli_args=cli_args, + overrides_folder=overrides_folder, + override_files=override_files, + suffix=suffix, + ) + + +def transmute_one_file( + arg: tuple[Path, MpCommonArgs]) -> Union[tuple[Path, Path], + tuple[Exception, None]]: + """ + Transmutes a single file. This function is called in parallel + from psyclone_transmute. + + :param arg: all required data, stored in MpCommonArgs + """ + input_file, mp_payload = arg + config = mp_payload.config + + prebuild_hash = _gen_prebuild_hash(input_file, + config, + mp_payload.cli_args, + mp_payload.transformation_script) + + # Create the output file name (with the suffix, and in the output + # folder of Fab) + output_file = input_to_output_fpath(config=config, input_path=input_file) + output_file.parent.mkdir(parents=True, exist_ok=True) + output_file = output_file.with_stem(output_file.stem + mp_payload.suffix) + + prebuild_out = (config.prebuild_folder / + f'{output_file.stem}.{prebuild_hash}{output_file.suffix}') + + # First check if we have an override file. If so, copy the override + # file as the expected output file, and delete the prebuild file. + if input_file.name in mp_payload.override_files: + # Help mypy to know that overrides_folder is not None + assert mp_payload.overrides_folder + # there is an override so delete this output file... + logger.warning(f"\nOverride found for '{output_file}'.") + shutil.copy2(mp_payload.overrides_folder / input_file.name, + output_file) + # Delete a prebuild, we do not want to store them + prebuild_out.unlink(missing_ok=True) + + elif prebuild_out.exists(): + msg = f'Found prebuild for {input_file}: {prebuild_out}' + log_or_dot(logger=logger, msg=msg) + shutil.copy2(prebuild_out, output_file) + else: + psyclone = config.tool_box.get_tool(Category.PSYCLONE) + psyclone = cast(Psyclone, psyclone) + try: + transformation_script = mp_payload.transformation_script + logger.info(f"Running PSyclone on '{input_file}'," + f" creating '{output_file}'.") + psyclone.process(config=mp_payload.config, + api=None, + x90_file=input_file, + transformed_file=output_file, + transformation_script=transformation_script, + additional_parameters=mp_payload.cli_args) + + shutil.copy2(output_file, prebuild_out) + msg = f'Created prebuilds for {input_file}: {prebuild_out}' + log_or_dot(logger=logger, msg=msg) + + except RuntimeError as err: + logger.error(err) + return err, None + + return output_file, prebuild_out + + +def _gen_prebuild_hash(input_file: Path, + config: BuildConfig, + cli_args: list[str], + script_func): + """ + Calculate the prebuild hash for this Fortran input file, based on + the source file and the transformation script. + + Changes which must trigger reprocessing of an x90 file: + - input_file source: + - transformation script + - cli args + + :param input_file: Fortran input file. + :param mp_payload: provides config file and command line args + """ + + input_hash = file_checksum(input_file).file_hash + # calculate the transformation script hash for this file + script_hash = 0 + if script_func: + script = script_func(input_file, config) + if script: + script_hash = file_checksum(script).file_hash + if script_hash == 0: + # Only a warning. Running PSyclone without script can be used to + # remove e.g. openmp directives (which PSyclone by default will do). + warnings.warn(f'No transformation script specified for {input_file}.') + + # hash everything which should trigger re-processing + # todo: hash the psyclone version? + return sum([input_hash, + string_checksum(str(cli_args)), + script_hash]) diff --git a/tests/unit_tests/steps/test_psyclone_transmute.py b/tests/unit_tests/steps/test_psyclone_transmute.py new file mode 100644 index 00000000..d2a41b0c --- /dev/null +++ b/tests/unit_tests/steps/test_psyclone_transmute.py @@ -0,0 +1,196 @@ +# ############################################################################## +# (c) Crown copyright Met Office. All rights reserved. +# For further details please refer to the file COPYRIGHT +# which you should have received as part of this distribution +# ############################################################################## + +""" +Tests the PSyclone transmutation step in Fab. It requires PSyclone to +be available (otherwise the tests will be skipped). +""" + +import logging + +from pytest import fixture, mark, raises, warns + +from fab.build_config import BuildConfig +from fab.artefacts import ArtefactSet +from fab.steps.psyclone_transmute import psyclone_transmute +from fab.tools.psyclone import Psyclone +from fab.tools.tool_box import ToolBox + + +@fixture(name="config") +def config_fixture(tmp_path): + """ + Create a fake workspace with input Fortran files. + """ + src = tmp_path / "src" + src.mkdir() + f1 = src / "a.f90" + f2 = src / "b.f90" + f1.write_text("program a\nend program") + f2.write_text("program b\nend program") + override = tmp_path / "override" + override.mkdir() + f1_override = override / "a.f90" + f1_override.write_text("program overwrite_a\nend program") + cfg = BuildConfig(project_label="test", + fab_workspace=tmp_path, + tool_box=ToolBox(), + multiprocessing=False, + ) + cfg.artefact_store.add(ArtefactSet.FORTRAN_COMPILER_FILES, [f1, f2]) + return cfg + + +@mark.skipif(not Psyclone().is_available, + reason="psyclone cli tool not available") +def test_psyclone_transmute_basic(config): + """ + Test basic behaviour, without changing any artefact set + """ + + input_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + # Make a copy to ensure changes to artefact store will be detected + input_files = input_files.copy() + + with (warns(UserWarning, + match="_metric_send_conn not set, cannot send metrics"), + warns(UserWarning, + match="No transformation script specified")): + psyclone_transmute( + config, + input_files) + + input_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + # Since we didn't specify a suffix, we should still have the same file + assert (config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] == + input_files) + + +@mark.skipif(not Psyclone().is_available, + reason="psyclone cli tool not available") +def test_psyclone_transmute_artefact_set(config): + """ + Verifies that we get the expected updated artefact set. + """ + + input_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + # Expected files will be in the build output directory and + # have the new suffix `_transmute` added. + expected = {config.build_output / '/'.join(i.parts[1:]) + for i in input_files} + expected = {i.with_stem(i.stem + "_transmute") for i in expected} + + with (warns(UserWarning, + match="_metric_send_conn not set, cannot send metrics"), + warns(UserWarning, + match="No transformation script specified")): + psyclone_transmute(config, input_files, + artefact_set=ArtefactSet.FORTRAN_COMPILER_FILES) + output_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + assert expected == output_files + + +@mark.skipif(not Psyclone().is_available, + reason="psyclone cli tool not available") +def test_psyclone_transmute_script(tmp_path, config): + """ + Check that we catch the error if the transformation script does not have + a .py extension (which is a PSyclone requirement). + """ + + input_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + script = tmp_path / "script" + script.write_text("invalid python\n") + with raises(RuntimeError) as err: + psyclone_transmute(config, input_files, + transformation_script=lambda _a, _b: script, + artefact_set=ArtefactSet.FORTRAN_COMPILER_FILES) + assert ("expected the script file \\'script\\' to have the \\'.py\\' " + "extension" in str(err)) + + +@mark.skipif(not Psyclone().is_available, + reason="psyclone cli tool not available") +def test_psyclone_transmute_prebuilt(config, caplog): + """ + Tests the handling of existing prebuild files. + """ + + input_files = \ + config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES].copy() + + # Expected files will be in the build output directory and + # have the new suffix `_transmute` added. + expected = {config.build_output / '/'.join(i.parts[1:]) + for i in input_files} + expected = {i.with_stem(i.stem + "_transmute") for i in expected} + + # Make sure debug messages are enabled so we can check for + # messages about prebuilds. logging.DEBUG is required to ensure that. + with (warns(UserWarning, + match="_metric_send_conn not set, cannot send metrics"), + warns(UserWarning, + match="No transformation script specified")): + with caplog.at_level(logging.DEBUG, + logger="fab.steps.psyclone_transmute"): + psyclone_transmute(config, input_files, + artefact_set=ArtefactSet.FORTRAN_COMPILER_FILES) + + output_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + assert expected == output_files + assert "Found prebuild for" not in caplog.text + + # Now rerun - remove the preprocessed filed from the previous step + config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] = \ + input_files.copy() + + # Now it should find prebuilds: + with (warns(UserWarning, + match="_metric_send_conn not set, cannot send metrics"), + warns(UserWarning, + match="No transformation script specified")): + with caplog.at_level(logging.DEBUG, + logger="fab.steps.psyclone_transmute"): + psyclone_transmute(config, input_files, + artefact_set=ArtefactSet.FORTRAN_COMPILER_FILES) + output_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + assert expected == output_files + assert "Found prebuild for" in caplog.text + + +@mark.skipif(not Psyclone().is_available, + reason="psyclone cli tool not available") +def test_psyclone_transmute_override(tmp_path, config): + """ + Test that the override directive works. + """ + + input_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + # Expected files will be in the build output directory and + # have the new suffix `_transmute` added. + expected = {config.build_output / '/'.join(i.parts[1:]) + for i in input_files} + expected = {i.with_stem(i.stem + "_transmute") for i in expected} + + overrides_folder = tmp_path / "override" + with (warns(UserWarning, + match="_metric_send_conn not set, cannot send metrics"), + warns(UserWarning, + match="No transformation script specified")): + psyclone_transmute( + config, + config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES], + artefact_set=ArtefactSet.FORTRAN_COMPILER_FILES, + overrides_folder=overrides_folder) + output_files = config.artefact_store[ArtefactSet.FORTRAN_COMPILER_FILES] + + assert expected == output_files