Source code for spacr.sweep_child

"""Run one regression trial in a fresh worker process.

    python -m spacr.sweep_child <settings.json> <result.json>

The parent may place this process in a kernel-enforced resource scope. Thread
limits are set before NumPy imports so BLAS libraries initialize with one
thread per trial. Results and failures are returned through a JSON file,
including when no in-memory return value is available.
"""

import os

for _variable in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS",
                  "NUMEXPR_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_variable, "1")

try:
    with open(f"/proc/{os.getpid()}/oom_score_adj", "w") as _handle:
        _handle.write("800")
except OSError:
    pass

import json  # noqa: E402
import sys  # noqa: E402
import time  # noqa: E402
import traceback  # noqa: E402




def main(argv=None) -> int:
    """Run one sweep trial in this process and write its result as JSON.

    Parameters
    ----------
    argv : sequence of str, optional
        Settings and output JSON paths. Defaults to ``sys.argv[1:]``.

    Returns
    -------
    int
        ``0`` once a result document was written, including when that document
        reports a failed trial, or ``2`` when the argument count is invalid.

    Raises
    ------
    OSError
        If the settings document cannot be read or the result document cannot
        be written.
    json.JSONDecodeError
        If the settings document is not valid JSON.
    """
    argv = list(sys.argv[1:] if argv is None else argv)
    if len(argv) != 2:
        print("usage: python -m spacr.sweep_child <settings.json> <out.json>",
              file=sys.stderr)
        return 2
    settings_path, out_path = argv

    with open(settings_path) as handle:
        payload = json.load(handle)
    settings = payload.get("settings", payload)

    trial_id = payload.get("trial_id")
    from .fit_resources import _worker_stamp

    result = {
        "status": "failed",
        "trial_id": trial_id,
        "_resource_worker": _worker_stamp(
            "parameter_sweep_trial", trial_id),
    }
    began = time.time()
    try:
        import matplotlib
        matplotlib.use("Agg")

        from .parameter_sweep import _pin_threads
        _pin_threads()

        from .figure_font import _open_sans_if_a_run_started_this
        from .ml import perform_regression

        with _open_sans_if_a_run_started_this():
            output = perform_regression(dict(settings))
        result["status"] = "ok"

        from .trial_metrics import summarise_trial
        result.update(summarise_trial(output, settings))

        controls = payload.get("controls") or {}
        if controls:
            try:
                import pandas as pd

                from .parameter_sweep import _named_control_rows

                results = output.get("results") \
                    if hasattr(output, "get") else None
                if isinstance(results, pd.DataFrame):
                    result.update(_named_control_rows(results, controls))
            except Exception:  # noqa: BLE001 - an alias must not sink a trial
                pass
    except BaseException as error:  # noqa: BLE001 - a failed trial is a result
        from .runctx import _is_overload_failure

        result["status"] = "failed"
        result["_overload"] = _is_overload_failure(error)
        result["error_type"] = type(error).__name__
        result["error"] = (str(error).splitlines() or [""])[0][:400]
        folder = settings.get("src")
        if folder:
            try:
                with open(os.path.join(folder, "error.txt"), "w",
                          encoding="utf-8") as handle:
                    handle.write(traceback.format_exc())
            except OSError:
                pass
    result["seconds"] = round(time.time() - began, 2)

    with open(out_path, "w") as handle:
        json.dump(result, handle, indent=2, default=str)
    return 0



if __name__ == "__main__":
    raise SystemExit(main())