profile_auxiliary_logarithmic_stress.py

Go to the documentation of this file.

#!/usr/bin/env python3
"""Profile auxiliary explicit/condensed and direct condensed Cook solves.
 
Collect wall timings, standard PETSc events and convergence while other builds
and tests are idle. The supplied environment must already support ep, including
its Python modules. Fixture and examples paths default to the ep directory.
"""
 
import argparse
import copy
import hashlib
import json
import math
import os
from pathlib import Path
import re
import shutil
import statistics
import subprocess
import time
 
 
def parse_arguments():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--ep", type=Path, required=True)
    parser.add_argument("--mpi", default="mpiexec")
    parser.add_argument("--mpi-arg", action="append", default=[],
                        help="launcher argument; repeat, e.g. --mpi-arg=--oversubscribe")
    parser.add_argument("--ranks", type=int, choices=(1, 2), default=1)
    parser.add_argument("--repetitions", type=int, default=3,
                        help="measured repetitions after one warm-up (default: 3)")
    parser.add_argument("--output", type=Path, required=True,
                        help="new or empty directory for all run artifacts")
    parser.add_argument("--fixture", type=Path)
    parser.add_argument("--examples", type=Path)
    parser.add_argument("--param-file", type=Path)
    args = parser.parse_args()
    if args.repetitions < 1:
        parser.error("--repetitions must be positive")
    args.ep = args.ep.expanduser().resolve()
    args.fixture = (args.fixture or args.ep.parent /
                    "ep_cook_auxiliary_logarithmic_stress.json").expanduser().resolve()
    args.examples = (args.examples or args.ep.parent / "examples").expanduser().resolve()
    args.param_file = (args.param_file or args.ep.parent /
                       "param_file.petsc").expanduser().resolve()
    for path in (args.ep, args.fixture, args.param_file):
        if not path.is_file():
            parser.error(f"required file does not exist: {path}")
    if not os.access(args.ep, os.X_OK):
        parser.error(f"ep is not executable: {args.ep}")
    if not args.examples.is_dir():
        parser.error(f"examples directory does not exist: {args.examples}")
    launcher = shutil.which(args.mpi)
    if launcher is None:
        parser.error(f"MPI launcher is not executable: {args.mpi}")
    args.mpi = str(Path(launcher).resolve())
    args.output = args.output.expanduser().resolve()
    if args.output.exists() and (not args.output.is_dir() or any(args.output.iterdir())):
        parser.error("--output must be a new or empty directory")
    return args
 
 
def file_hash(path):
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()
 
 
def write_json(path, value):
    path.write_text(json.dumps(value, indent=2) + "\n")
 
 
def configure_case(configuration, case, ranks):
    configuration = copy.deepcopy(configuration)
    mesh = "cook.cub" if ranks == 1 else "cook_2p.h5m"
    configuration["meshes"][0]["file_name"] = f"examples/cook_cantilever/{mesh}"
    configuration["mofem"]["auxiliary_logarithmic_stress"] = case != "direct_condensed"
    elastic = configuration["petsc"]["elastic"]
    elastic["snes"]["converged_reason"] = True
    elastic["mat_mumps_icntl_20"] = 0
    elastic["mat_mumps_icntl_14"] = 800
    elastic["ksp"] = {"type": "preonly", "converged_reason": True}
    elastic["pc"] = {"type": "lu", "factor_mat_solver_type": "mumps"}
    if case.endswith("_condensed"):
        elastic["ksp"].update(type="fgmres", rtol=1e-12, atol=1e-12)
        elastic["pc"] = {"type": "fieldsplit", "fieldsplit_type": "schur"}
        elastic["fieldsplit_0"] = {"ksp": {"type": "preonly"}}
        elastic["fieldsplit_1"] = {
            "pc": {"type": "ksp"},
            "ksp": {
                "type": "preonly",
                "pc": {"type": "lu", "factor_mat_solver_type": "mumps"},
                "mat_mumps_icntl_20": 0,
                "mat_mumps_icntl_14": 800,
            },
        }
    return configuration
 
 
def log_diagnostics(log):
    steps = re.findall(r"Step (\d+) time\s+(\S+) strain energy (\S+)", log)
    return {
        "accepted_steps": [
            {"step": int(step), "time": float(t), "physical_energy": float(energy)}
            for step, t, energy in steps
        ],
        "nonlinear_solves": [
            {"reason": reason, "iterations": int(iterations)}
            for reason, iterations in re.findall(
                r"Nonlinear (?:elastic_ )?solve converged due to (\S+) iterations (\d+)", log)
        ],
        "linear_iterations": [int(value) for value in re.findall(
            r"Linear (?:elastic_ )?solve converged due to \S+ iterations (\d+)", log)],
    }
 
 
def validate_convergence(record, configuration):
    steps = [step for step in record["accepted_steps"] if step["step"] > 0]
    solves = record["nonlinear_solves"]
    ts = configuration["petsc"]["elastic"]["ts"]
    expected_time = min(ts["max_time"], ts["dt"] * ts["max_steps"])
    if (not steps or not solves or len(steps) != len(solves)
            or not math.isclose(steps[-1]["time"], expected_time,
                                rel_tol=0., abs_tol=1e-12)
            or any(not math.isfinite(step["physical_energy"]) for step in steps)
            or any(not solve["reason"].startswith("CONVERGED_") for solve in solves)):
        raise RuntimeError(f"Incomplete or invalid solve; see {record['directory']}/run.log")
 
 
def run_case(args, case, configuration, environment, summary):
    case_directory = args.output / case
    case_directory.mkdir()
    write_json(case_directory / "input.json", configuration)
    mesh = args.examples / Path(configuration["meshes"][0]["file_name"]).relative_to("examples")
    case_record = {
        "case": case,
        "configuration_sha256": file_hash(case_directory / "input.json"),
        "mesh": str(mesh),
        "mesh_sha256": file_hash(mesh),
        "runs": [],
    }
    summary["cases"].append(case_record)
    for repetition in range(args.repetitions + 1):
        run_name = "warmup" if repetition == 0 else f"run_{repetition}"
        directory = case_directory / run_name
        directory.mkdir()
        (directory / "examples").symlink_to(args.examples, target_is_directory=True)
        shutil.copy2(case_directory / "input.json", directory / "input.json")
        shutil.copy2(args.output / "param_file.petsc", directory / "param_file.petsc")
        command = [
            args.mpi, *args.mpi_arg, "-np", str(args.ranks), str(args.output / "ep"),
            "-json_config", "input.json", "-log_view", ":petsc_profile.txt",
            "-elastic_snes_view_solution", "binary:accepted.vec",
        ]
        record = {"directory": str(directory), "warmup": repetition == 0,
                  "command": command}
        case_record["runs"].append(record)
        start = time.monotonic()
        with (directory / "run.log").open("w") as output:
            result = subprocess.run(command, cwd=directory, env=environment,
                                    stdout=output, stderr=subprocess.STDOUT)
        record.update(exit_code=result.returncode, elapsed_seconds=time.monotonic() - start)
        record.update(log_diagnostics((directory / "run.log").read_text(errors="replace")))
        for artifact in ("run.log", "petsc_profile.txt", "accepted.vec"):
            path = directory / artifact
            if path.is_file():
                record.setdefault("artifact_sha256", {})[artifact] = file_hash(path)
        write_json(directory / "result.json", record)
        write_json(args.output / "profile.json", summary)
        if result.returncode:
            raise RuntimeError(f"{case}/{run_name} failed ({result.returncode}); see {directory / 'run.log'}")
        if not all((directory / name).is_file() for name in ("petsc_profile.txt", "accepted.vec")):
            raise RuntimeError(f"{case}/{run_name} did not write its PETSc profile and final vector")
        validate_convergence(record, configuration)
        print(f"{case} {run_name}: {record['elapsed_seconds']:.3f} s", flush=True)
    case_record["median_elapsed_seconds"] = statistics.median(
        run["elapsed_seconds"] for run in case_record["runs"] if not run["warmup"])
    write_json(args.output / "profile.json", summary)
 
 
def main():
    args = parse_arguments()
    configuration = json.loads(args.fixture.read_text())
    args.output.mkdir(parents=True, exist_ok=True)
    shutil.copy2(args.ep, args.output / "ep")
    shutil.copy2(args.fixture, args.output / "fixture.json")
    shutil.copy2(args.param_file, args.output / "param_file.petsc")
    environment = dict(os.environ)
    # Keep PYTHONPATH and runtime library paths supplied by the build environment.
    thread_options = {name: "1" for name in (
        "OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS",
        "BLIS_NUM_THREADS", "VECLIB_MAXIMUM_THREADS")}
    environment.update(thread_options)
    summary = {
        "ranks": args.ranks, "measured_repetitions": args.repetitions,
        "source_executable": str(args.ep), "thread_environment": thread_options,
        "input_sha256": {name: file_hash(args.output / name)
                         for name in ("ep", "fixture.json", "param_file.petsc")},
        "vector_contents": "One final accepted solution per run; PETSc overwrites earlier steps.",
        "cases": [],
    }
    write_json(args.output / "profile.json", summary)
    for case in ("auxiliary_explicit", "auxiliary_condensed", "direct_condensed"):
        run_case(args, case, configure_case(configuration, case, args.ranks),
                 environment, summary)
    print(f"Profiles and median timings: {args.output / 'profile.json'}")
 
 
if __name__ == "__main__":
    try:
        main()
    except (OSError, ValueError, KeyError, RuntimeError) as error:
        raise SystemExit(str(error)) from error