diff --git a/python/flydsl/utils/env.py b/python/flydsl/utils/env.py index 1839dd7df..13e47ac0a 100644 --- a/python/flydsl/utils/env.py +++ b/python/flydsl/utils/env.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Dict, Generic, Optional, TypeVar T = TypeVar("T") +NumberT = TypeVar("NumberT", int, float) class EnvOption(Generic[T]): @@ -77,21 +78,23 @@ def parse_value(self, raw: str) -> bool: return raw.lower() in ("1", "true", "yes", "on") -class OptInt(EnvOption[int]): - """Integer environment option with optional min/max validation.""" +class _OptNumber(EnvOption[NumberT]): + """Numeric environment option with optional min/max validation.""" def __init__( self, - default: int = 0, + default: NumberT, + parser: Callable[[str], NumberT], env_var: Optional[str] = None, description: str = "", - min_value: Optional[int] = None, - max_value: Optional[int] = None, + min_value: Optional[NumberT] = None, + max_value: Optional[NumberT] = None, + empty_is_default: bool = False, ): validator = None if min_value is not None or max_value is not None: - def validator(v: int) -> bool: + def validator(v: NumberT) -> bool: if min_value is not None and v < min_value: return False if max_value is not None and v > max_value: @@ -99,11 +102,61 @@ def validator(v: int) -> bool: return True super().__init__(default, env_var, description, validator) + self.parser = parser self.min_value = min_value self.max_value = max_value + self.empty_is_default = empty_is_default + + def parse_value(self, raw: str) -> NumberT: + if self.empty_is_default and not raw.strip(): + return self.default + return self.parser(raw) + + +class OptInt(_OptNumber[int]): + """Integer environment option with optional min/max validation.""" + + def __init__( + self, + default: int = 0, + env_var: Optional[str] = None, + description: str = "", + min_value: Optional[int] = None, + max_value: Optional[int] = None, + empty_is_default: bool = False, + ): + super().__init__( + default, + int, + env_var, + description, + min_value, + max_value, + empty_is_default, + ) - def parse_value(self, raw: str) -> int: - return int(raw) + +class OptFloat(_OptNumber[float]): + """Floating-point environment option with optional min/max validation.""" + + def __init__( + self, + default: float = 0.0, + env_var: Optional[str] = None, + description: str = "", + min_value: Optional[float] = None, + max_value: Optional[float] = None, + empty_is_default: bool = False, + ): + super().__init__( + default, + float, + env_var, + description, + min_value, + max_value, + empty_is_default, + ) class OptStr(EnvOption[str]): @@ -224,6 +277,33 @@ class AutotuneEnvManager(EnvManager): config_dir = OptStr("", description="Directory for offline config artifacts; empty disables artifacts") +class AotEnvManager(EnvManager): + """AOT job options (``FLYDSL_AOT_*`` environment variables).""" + + env_prefix = "AOT" + + workers = OptInt( + 0, + description=( + "Maximum concurrent worker processes; unset, empty, or non-positive values use the CPU and " + "available-memory based automatic limit" + ), + empty_is_default=True, + ) + mem_per_worker_gb = OptFloat( + 2.0, + description="Assumed GiB per worker for the automatic memory cap; non-positive disables the cap", + ) + timeout = OptFloat( + 1200.0, + description="Per-job wall-clock timeout in seconds; non-positive disables the timeout", + ) + max_retries = OptInt( + 2, + description="Retries after an abnormal worker exit or timeout; negative values clamp to zero", + ) + + class CompileEnvManager(EnvManager): """Compile-time options (``FLYDSL_COMPILE_*`` environment variables).""" @@ -296,16 +376,18 @@ class RuntimeEnvManager(EnvManager): enable_cache = OptBool(True, description="Enable kernel caching") run_only = OptBool( False, - description=("Skip JIT compilation; only load AOT cache. " "Raise RuntimeError on cache miss."), + description=("Skip JIT compilation; only load AOT cache. Raise RuntimeError on cache miss."), ) +aot = AotEnvManager() autotune = AutotuneEnvManager() compile = CompileEnvManager() debug = DebugEnvManager() runtime = RuntimeEnvManager() __all__ = [ + "aot", "autotune", "compile", "debug", diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py new file mode 100644 index 000000000..5d623cc6c --- /dev/null +++ b/python/flydsl/utils/parallel.py @@ -0,0 +1,537 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +# Copyright (c) 2025 FlyDSL Project Contributors + +from __future__ import annotations + +import json +import multiprocessing +import os +import shutil +import signal +import tempfile +import time +import traceback +import warnings +from collections import deque +from collections.abc import Callable +from multiprocessing.connection import wait as wait_for_sentinels +from pathlib import Path +from typing import Any, cast + +from .env import aot +from .file import atomic_write +from .logger import log + +_DEFAULT_MAX_WORKERS = 64 +_FALLBACK_MAX_WORKERS_WITHOUT_MEMORY_INFO = 4 +_MAX_FAILURE_REASON_CHARS = 2 * 1024 +_MAX_TRACEBACK_CHARS = 16 * 1024 +_OOM_EXITCODE = -signal.SIGKILL +_WORKER_EXCEPTION_KEY = "__flydsl_worker_exception__" + + +def _write_json_file(out_path: str, payload: Any) -> None: + with atomic_write(Path(out_path), mode="w", encoding="utf-8") as output: + json.dump(payload, output) + + +def _truncate_text(text: str, max_chars: int, marker: str) -> str: + if len(text) <= max_chars: + return text + remaining = max_chars - len(marker) + head = remaining // 2 + tail = remaining - head + return text[:head] + marker + text[-tail:] + + +def _truncate_failure_reason(reason: str) -> str: + return _truncate_text( + reason, + _MAX_FAILURE_REASON_CHARS, + "\n... failure reason truncated ...\n", + ) + + +def _truncate_traceback(traceback_text: str) -> str: + return _truncate_text( + traceback_text, + _MAX_TRACEBACK_CHARS, + "\n... traceback truncated ...\n", + ) + + +def _exception_exitcode(error: BaseException) -> int: + if isinstance(error, SystemExit): + if error.code is None: + return 0 + if isinstance(error.code, int): + return error.code & 0xFF + return 1 + + +def _run_one_to_file( + worker: Callable[..., dict[str, Any]], + kwargs: dict[str, Any], + out_path: str, +) -> None: + try: + result = worker(**kwargs) + _write_json_file(out_path, result) + except BaseException as error: + try: + _write_json_file( + out_path, + { + _WORKER_EXCEPTION_KEY: { + "kind": "worker_exception", + "reason": _truncate_failure_reason(f"{type(error).__name__}: {error}"), + "traceback": _truncate_traceback(traceback.format_exc()), + } + }, + ) + finally: + os._exit(_exception_exitcode(error)) + + +def _affinity_aware_cpu_count() -> int: + """Return the CPU count available to this process.""" + try: + count = len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + count = os.cpu_count() or 0 + return max(count, 1) + + +def _get_kernel_timeout() -> float: + return max(aot.timeout, 0.0) + + +def _get_max_retries() -> int: + return max(aot.max_retries, 0) + + +def _memory_worker_fallback(default_workers: int, reason: str) -> int: + fallback_workers = min( + default_workers, + _FALLBACK_MAX_WORKERS_WITHOUT_MEMORY_INFO, + ) + message = f"{reason}; limiting AOT concurrency to {fallback_workers} worker{'s' if fallback_workers != 1 else ''}" + log().warning(message) + # Public path: user -> run_parallel_jobs -> _get_max_workers -> + # _memory_worker_cap -> this helper. + warnings.warn(message, RuntimeWarning, stacklevel=5) + return fallback_workers + + +def _memory_worker_cap(default_workers: int) -> int: + per_worker_gb = aot.mem_per_worker_gb + if per_worker_gb <= 0: + return default_workers + + try: + import psutil + except ImportError: + return _memory_worker_fallback( + default_workers, + "psutil is not installed; automatic AOT memory limiting is unavailable", + ) + + try: + available_gb = psutil.virtual_memory().available / (1024**3) + except Exception as error: + return _memory_worker_fallback( + default_workers, + f"failed to query available memory for AOT worker limiting ({error})", + ) + return min(default_workers, max(1, int(available_gb / per_worker_gb))) + + +def _get_max_workers(num_jobs: int) -> int: + configured_workers = aot.workers + if configured_workers > 0: + max_workers = configured_workers + else: + max_workers = min(_affinity_aware_cpu_count(), _DEFAULT_MAX_WORKERS) + max_workers = _memory_worker_cap(max_workers) + return min(max_workers, num_jobs) + + +def _job_label(job: dict[str, Any]) -> str: + return str(job.get("kernel_name", "?")) + + +def _failure_result( + job: dict[str, Any], + *, + kind: str, + reason: str, + attempts: int | None = None, + exitcode: int | None = None, + traceback_text: str | None = None, +) -> dict[str, Any]: + failure: dict[str, Any] = { + "kind": kind, + "reason": reason, + } + if attempts is not None: + failure["attempts"] = attempts + if exitcode is not None: + failure["exitcode"] = exitcode + if traceback_text is not None: + failure["traceback"] = traceback_text + return { + "kernel_name": _job_label(job), + "compile_time": None, + "failure": failure, + } + + +def _finalize_pool_results( + results: list[dict[str, Any] | None], + *, + num_jobs: int, + completed: int, + succeeded_jobs: int, + failed_jobs: int, + retries_used: int, +) -> list[dict[str, Any]]: + if completed != num_jobs or succeeded_jobs + failed_jobs != completed or any(result is None for result in results): + raise RuntimeError( + "internal AOT scheduler error: " + f"{completed}/{num_jobs} jobs completed, " + f"{succeeded_jobs} succeeded, {failed_jobs} failed" + ) + + retry_label = "retry" if retries_used == 1 else "retries" + log().info( + "AOT: %d succeeded, %d failed; %d %s after abnormal worker exits", + succeeded_jobs, + failed_jobs, + retries_used, + retry_label, + ) + return cast(list[dict[str, Any]], results) + + +def _run_file_pool( + worker: Callable[..., dict[str, Any]], + jobs: list[dict[str, Any]], + *, + max_workers: int, + kernel_timeout: float, + max_retries: int, + result_dir: str, +) -> list[dict[str, Any]]: + """Run jobs through a Linux fork pool using files for result transport.""" + ctx = multiprocessing.get_context("fork") + num_jobs = len(jobs) + results: list[dict[str, Any] | None] = [None] * num_jobs + attempts = [0] * num_jobs + retries_used = 0 + completed = 0 + succeeded_jobs = 0 + failed_jobs = 0 + progress_stride = max(1, num_jobs // 20) + + queue = deque(range(num_jobs)) + running: dict[Any, tuple[int, float | None]] = {} + + def launch() -> None: + while queue and len(running) < max_workers: + index = queue.popleft() + out_path = os.path.join(result_dir, f"k{index}.json") + try: + os.remove(out_path) + except OSError: + pass + + process = ctx.Process( + target=_run_one_to_file, + args=(worker, jobs[index], out_path), + ) + process.start() + deadline = time.monotonic() + kernel_timeout if kernel_timeout > 0 else None + running[process] = (index, deadline) + + def note_done(*, is_failure: bool = False) -> None: + nonlocal completed, failed_jobs, succeeded_jobs + completed += 1 + if is_failure: + failed_jobs += 1 + else: + succeeded_jobs += 1 + if completed % progress_stride == 0 or completed == num_jobs: + log().info( + "... %d/%d jobs finished (%d failed)", + completed, + num_jobs, + failed_jobs, + ) + + def finish_failure( + index: int, + *, + kind: str, + reason: str, + exitcode: int | None = None, + traceback_text: str | None = None, + ) -> None: + results[index] = _failure_result( + jobs[index], + kind=kind, + reason=reason, + attempts=attempts[index] + 1, + exitcode=exitcode, + traceback_text=traceback_text, + ) + log().warning( + "AOT job %s %s; not retrying", + _job_label(jobs[index]), + reason, + ) + note_done(is_failure=True) + + def retry_or_drop( + index: int, + *, + kind: str, + reason: str, + exitcode: int | None = None, + traceback_text: str | None = None, + ) -> None: + nonlocal retries_used + if attempts[index] < max_retries: + attempts[index] += 1 + retries_used += 1 + queue.append(index) + log().warning( + "AOT job %s %s; retry %d/%d", + _job_label(jobs[index]), + reason, + attempts[index], + max_retries, + ) + else: + finish_failure( + index, + kind=kind, + reason=reason, + exitcode=exitcode, + traceback_text=traceback_text, + ) + + def reap(process: Any, *, timeout_reason: str | None = None) -> None: + index, _ = running.pop(process) + out_path = os.path.join(result_dir, f"k{index}.json") + try: + loaded: Any = None + load_error: str | None = None + if os.path.isfile(out_path): + try: + with open(out_path, encoding="utf-8") as result_file: + loaded = json.load(result_file) + except Exception as error: # noqa: BLE001 + load_error = f"failed to read worker result: {type(error).__name__}: {error}" + else: + load_error = "worker produced no result file" + + has_worker_exception_marker = isinstance(loaded, dict) and _WORKER_EXCEPTION_KEY in loaded + worker_exception = loaded.get(_WORKER_EXCEPTION_KEY) if has_worker_exception_marker else None + + # The atomic result file is the worker's commit point. Once a valid + # result is present, teardown-time signals must not discard it. + if isinstance(loaded, dict) and not has_worker_exception_marker: + result = loaded + if result.get("compile_time") is None: + failure = result.get("failure") + if not isinstance(failure, dict): + failure = {} + result["failure"] = failure + failure.setdefault("kind", "compile_error") + failure.setdefault("reason", "worker returned compile_time=None") + failure.setdefault("attempts", attempts[index] + 1) + results[index] = result + is_failure = result.get("compile_time") is None + note_done(is_failure=is_failure) + return + + # A structured Python exception is deterministic for the same job. + # Preserve it even if teardown later changes the process exit code. + if has_worker_exception_marker: + if not isinstance(worker_exception, dict): + finish_failure( + index, + kind="invalid_result", + reason=( + f"worker exception marker must contain a dictionary, got {type(worker_exception).__name__}" + ), + exitcode=process.exitcode, + ) + return + reason = str( + worker_exception.get( + "reason", + f"worker raised an exception (exitcode={process.exitcode})", + ) + ) + traceback_text = worker_exception.get("traceback") + if not isinstance(traceback_text, str): + traceback_text = None + finish_failure( + index, + kind="worker_exception", + reason=reason, + exitcode=process.exitcode, + traceback_text=traceback_text, + ) + return + + if timeout_reason is not None: + retry_or_drop( + index, + kind="timeout", + reason=timeout_reason, + exitcode=process.exitcode, + ) + return + + if process.exitcode != 0: + possible_oom = process.exitcode == _OOM_EXITCODE + retry_or_drop( + index, + kind="possible_oom" if possible_oom else "worker_crash", + reason=( + "worker killed by SIGKILL (possible OOM)" + if possible_oom + else f"worker crashed (exitcode={process.exitcode})" + ), + exitcode=process.exitcode, + ) + return + + if load_error is not None: + finish_failure( + index, + kind="invalid_result", + reason=load_error, + exitcode=process.exitcode, + ) + return + if not isinstance(loaded, dict): + finish_failure( + index, + kind="invalid_result", + reason=(f"worker returned {type(loaded).__name__}, expected a dictionary"), + exitcode=process.exitcode, + ) + return + finally: + process.close() + + try: + launch() + while running: + if kernel_timeout > 0: + nearest_deadline = min(deadline for _, deadline in running.values() if deadline is not None) + wait_timeout: float | None = max(0.0, nearest_deadline - time.monotonic()) + else: + wait_timeout = None + + wait_for_sentinels( + [process.sentinel for process in running], + timeout=wait_timeout, + ) + + for process in list(running): + if not process.is_alive(): + process.join() + reap(process) + + if kernel_timeout > 0: + now = time.monotonic() + for process in list(running): + _, deadline = running[process] + if deadline is not None and now > deadline and process.is_alive(): + timeout_reason = f"exceeded per-job timeout ({kernel_timeout:g}s); killed" + process.kill() + process.join() + reap(process, timeout_reason=timeout_reason) + + launch() + finally: + for process in list(running): + try: + if process.is_alive(): + process.kill() + process.join() + except Exception: # noqa: BLE001 + pass + finally: + try: + process.close() + except Exception: # noqa: BLE001 + pass + running.clear() + + return _finalize_pool_results( + results, + num_jobs=num_jobs, + completed=completed, + succeeded_jobs=succeeded_jobs, + failed_jobs=failed_jobs, + retries_used=retries_used, + ) + + +def run_parallel_jobs( + worker: Callable[..., dict[str, Any]], + jobs: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Run independent AOT jobs in parallel and return results in input order. + + Each job is expanded into ``worker(**job)``. Workers must return a + JSON-serializable dictionary containing ``compile_time``; deterministic + compile errors should be represented by ``compile_time=None`` rather than + escaping as an exception. Every failed result keeps ``compile_time=None`` + for compatibility and adds a ``failure`` dictionary with machine-readable + ``kind``, ``reason``, and ``attempts`` fields. Process failures also include + ``exitcode``; uncaught Python exceptions include a bounded ``traceback`` and + are not retried. + + This executor uses the Linux ``fork`` multiprocessing context because it is + intended for compile-only workers that inherit the initialized compiler. + Callers must ensure no compiler or MLIR work is active on other native + threads when forking; ``threading.active_count()`` does not account for all + native threads. Workers must not create long-lived child processes because + timeout and cleanup handling supervise only the direct worker process. + + An OOM-like ``-SIGKILL`` is reported as ``possible_oom`` and retried according + to ``FLYDSL_AOT_MAX_RETRIES`` without changing global concurrency. + """ + if not jobs: + return [] + + max_workers = _get_max_workers(len(jobs)) + log().info( + "AOT: %d jobs, %d worker processes", + len(jobs), + max_workers, + ) + + result_dir = tempfile.mkdtemp(prefix="flydsl_aot_results_") + try: + results = _run_file_pool( + worker, + jobs, + max_workers=max_workers, + kernel_timeout=_get_kernel_timeout(), + max_retries=_get_max_retries(), + result_dir=result_dir, + ) + finally: + shutil.rmtree(result_dir, ignore_errors=True) + return results + + +__all__ = ["run_parallel_jobs"] diff --git a/tests/README.md b/tests/README.md index efc4ed7a9..d88dba2e4 100644 --- a/tests/README.md +++ b/tests/README.md @@ -52,6 +52,36 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. | IR dump | `FLYDSL_DUMP_IR`, `FLYDSL_DUMP_DIR` | | Device runtime kind | `FLYDSL_RUNTIME_KIND` | | ROCm arch hints (detection helpers) | `FLYDSL_GPU_ARCH`, `HSA_OVERRIDE_GFX_VERSION` | +| AOT worker-process limit | `FLYDSL_AOT_WORKERS` (positive: explicit limit; unset, empty, zero, or negative: automatic CPU/memory limit) | +| AOT automatic memory cap | `FLYDSL_AOT_MEM_PER_WORKER_GB` (default `2.0`; non-positive disables) | +| AOT per-job timeout | `FLYDSL_AOT_TIMEOUT` (default `1200` seconds; non-positive disables) | +| AOT abnormal-exit retries | `FLYDSL_AOT_MAX_RETRIES` (default `2`) | + +`flydsl.utils.parallel.run_parallel_jobs` uses the Linux `fork` multiprocessing +context and is intended for compile-only AOT jobs. It is not a device-runtime +executor. Its automatic memory cap requires the optional `psutil` package and +emits both a warning and `RuntimeWarning` before falling back to at most four +workers when the optional package is unavailable or memory cannot be queried. +The `RuntimeWarning` points to the public caller and follows Python's standard +once-per-call-site filtering; the logger records each fallback. +OOM-like `SIGKILL` (`exitcode=-9`) and timeout failures are classified +separately and retried according to `FLYDSL_AOT_MAX_RETRIES`; neither changes +global concurrency. Retries are appended behind jobs that have not started yet. +Call the scheduler only when no compiler/MLIR work is active on other native +threads, and do not let workers create long-lived child processes: the +fork-based pool supervises only its direct workers. +Progress and final summary lines report terminal failure counts separately from +the number of jobs that have finished. Scheduler messages use the FlyDSL logger; +set `FLYDSL_DEBUG_LOG_TO_CONSOLE=1` and `FLYDSL_DEBUG_LOG_LEVEL=INFO` to emit +progress and summaries to the console. + +Failed results retain `compile_time=None` and include a machine-readable +`failure` mapping. Its `kind` distinguishes `compile_error`, +`worker_exception`, `worker_crash`, `possible_oom`, `timeout`, +and `invalid_result`; `reason` and `attempts` are always included when known, +while process failures also carry `exitcode`. Uncaught Python exceptions are +not retried and carry bounded `reason` and `traceback` strings. A valid atomic +result remains authoritative if the process is killed during teardown. Session-level pytest options are supported in `tests/conftest.py`: diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py new file mode 100644 index 000000000..badf83cb4 --- /dev/null +++ b/tests/unit/test_parallel_jobs.py @@ -0,0 +1,563 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import fcntl +import json +import os +import signal +import sys +import time +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import flydsl.utils.parallel as parallel +from flydsl.utils.env import aot +from flydsl.utils.parallel import run_parallel_jobs + +pytestmark = [pytest.mark.l0_backend_agnostic] + + +def _success_worker(kernel_name, index, delay=0.0): + time.sleep(delay) + return { + "kernel_name": kernel_name, + "index": index, + "compile_time": 0.01, + } + + +def _tracked_worker(kernel_name, index, state_path, lock_path, delay): + state_file = Path(state_path) + with open(lock_path, "a+", encoding="utf-8") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + state = json.loads(state_file.read_text(encoding="utf-8")) + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + state_file.write_text(json.dumps(state), encoding="utf-8") + fcntl.flock(lock_file, fcntl.LOCK_UN) + + try: + time.sleep(delay) + return { + "kernel_name": kernel_name, + "index": index, + "compile_time": 0.01, + } + finally: + with open(lock_path, "a+", encoding="utf-8") as lock_file: + fcntl.flock(lock_file, fcntl.LOCK_EX) + state = json.loads(state_file.read_text(encoding="utf-8")) + state["active"] -= 1 + state_file.write_text(json.dumps(state), encoding="utf-8") + fcntl.flock(lock_file, fcntl.LOCK_UN) + + +def _crash_then_succeed(kernel_name, attempt_path, crashes, order_path=None): + if order_path is not None: + with open(order_path, "a", encoding="utf-8") as order_file: + order_file.write(f"{kernel_name}\n") + path = Path(attempt_path) + attempt = int(path.read_text(encoding="utf-8")) + 1 + path.write_text(str(attempt), encoding="utf-8") + if attempt <= crashes: + os._exit(17) + return {"kernel_name": kernel_name, "compile_time": 0.01} + + +def _deterministic_failure(kernel_name, attempt_path): + path = Path(attempt_path) + attempt = int(path.read_text(encoding="utf-8")) + 1 + path.write_text(str(attempt), encoding="utf-8") + return {"kernel_name": kernel_name, "compile_time": None} + + +def _sleep_worker(kernel_name, delay): + time.sleep(delay) + return {"kernel_name": kernel_name, "compile_time": 0.01} + + +def _mixed_outcome_worker(kernel_name, outcome): + if outcome == "signal": + os.kill(os.getpid(), signal.SIGTERM) + if outcome == "exit-137": + os._exit(137) + if outcome == "type-error": + raise TypeError("synthetic worker type error") + if outcome == "large-error": + raise RuntimeError("X" * (parallel._MAX_TRACEBACK_CHARS * 4)) + if outcome == "compile-error": + return { + "kernel_name": kernel_name, + "compile_time": None, + "failure": { + "kind": "compile_error", + "reason": "synthetic codegen failure", + }, + } + raise ValueError(f"unknown synthetic outcome: {outcome}") + + +def _exception_marker_then_exit_zero(worker, kwargs, out_path): + parallel._write_json_file( + out_path, + { + parallel._WORKER_EXCEPTION_KEY: { + "kind": "worker_exception", + "reason": "TypeError: preserved reason", + "traceback": "Traceback (most recent call last): preserved frame", + } + }, + ) + + +def _malformed_exception_marker_then_exit_zero(worker, kwargs, out_path): + parallel._write_json_file( + out_path, + {parallel._WORKER_EXCEPTION_KEY: "not-a-dictionary"}, + ) + + +def _exit_without_result(worker, kwargs, out_path): + return None + + +def _non_dict_worker(kernel_name): + return ["not", "a", "dictionary"] + + +def _always_oom_worker(kernel_name, attempt_path): + path = Path(attempt_path) + attempt = int(path.read_text(encoding="utf-8")) + 1 + path.write_text(str(attempt), encoding="utf-8") + os.kill(os.getpid(), signal.SIGKILL) + + +@pytest.fixture(autouse=True) +def _parallel_env(monkeypatch): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") + monkeypatch.setenv("FLYDSL_AOT_MEM_PER_WORKER_GB", "0") + monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "5") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "0") + + +@pytest.fixture +def log_messages(monkeypatch): + messages = [] + + def record(message, *args): + messages.append(message % args if args else message) + + logger = SimpleNamespace(info=record, warning=record) + monkeypatch.setattr(parallel, "log", lambda: logger) + return messages + + +def test_workers_option_whitespace_uses_automatic_default(monkeypatch): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", " ") + monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 8) + monkeypatch.setattr(parallel, "_memory_worker_cap", lambda workers: workers) + + assert aot.workers == 0 + assert parallel._get_max_workers(num_jobs=100) == 8 + + +def test_results_follow_input_order(): + jobs = [ + {"kernel_name": "slow", "index": 0, "delay": 0.1}, + {"kernel_name": "fast", "index": 1, "delay": 0.0}, + {"kernel_name": "last", "index": 2, "delay": 0.01}, + ] + + results = run_parallel_jobs(_success_worker, jobs) + + assert [result["index"] for result in results] == [0, 1, 2] + + +def test_committed_result_wins_over_teardown_sigkill(monkeypatch): + original_write = parallel._write_json_file + + def write_then_sigkill(out_path, payload): + original_write(out_path, payload) + os.kill(os.getpid(), signal.SIGKILL) + + monkeypatch.setattr(parallel, "_write_json_file", write_then_sigkill) + + results = run_parallel_jobs( + _success_worker, + [{"kernel_name": "committed", "index": 0}], + ) + + assert results == [ + { + "kernel_name": "committed", + "index": 0, + "compile_time": 0.01, + } + ] + + +def test_committed_result_wins_over_teardown_timeout(monkeypatch): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.05") + original_write = parallel._write_json_file + + def write_then_wait(out_path, payload): + original_write(out_path, payload) + time.sleep(60) + + monkeypatch.setattr(parallel, "_write_json_file", write_then_wait) + + results = run_parallel_jobs( + _success_worker, + [{"kernel_name": "committed", "index": 0}], + ) + + assert results == [ + { + "kernel_name": "committed", + "index": 0, + "compile_time": 0.01, + } + ] + + +def test_worker_count_is_bounded(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") + state_path = tmp_path / "state.json" + lock_path = tmp_path / "state.lock" + state_path.write_text(json.dumps({"active": 0, "peak": 0}), encoding="utf-8") + jobs = [ + { + "kernel_name": f"job-{index}", + "index": index, + "state_path": str(state_path), + "lock_path": str(lock_path), + "delay": 0.1, + } + for index in range(6) + ] + + results = run_parallel_jobs(_tracked_worker, jobs) + + state = json.loads(state_path.read_text(encoding="utf-8")) + assert state == {"active": 0, "peak": 2} + assert all(result["compile_time"] is not None for result in results) + + +def test_crashed_worker_is_retried(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") + attempt_path = tmp_path / "attempt.txt" + attempt_path.write_text("0", encoding="utf-8") + + results = run_parallel_jobs( + _crash_then_succeed, + [ + { + "kernel_name": "retry-job", + "attempt_path": str(attempt_path), + "crashes": 2, + } + ], + ) + + assert attempt_path.read_text(encoding="utf-8") == "3" + assert results[0]["compile_time"] is not None + + +def test_retries_wait_behind_pending_jobs(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") + order_path = tmp_path / "order.txt" + order_path.write_text("", encoding="utf-8") + attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(3)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + jobs = [ + { + "kernel_name": "retry", + "attempt_path": str(attempt_paths[0]), + "crashes": 1, + "order_path": str(order_path), + }, + { + "kernel_name": "pending-0", + "attempt_path": str(attempt_paths[1]), + "crashes": 0, + "order_path": str(order_path), + }, + { + "kernel_name": "pending-1", + "attempt_path": str(attempt_paths[2]), + "crashes": 0, + "order_path": str(order_path), + }, + ] + + results = run_parallel_jobs(_crash_then_succeed, jobs) + + assert all(result["compile_time"] is not None for result in results) + assert order_path.read_text(encoding="utf-8").splitlines() == [ + "retry", + "pending-0", + "pending-1", + "retry", + ] + + +def test_oom_retry_exhaustion_returns_possible_oom(monkeypatch, tmp_path, log_messages): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") + attempt_path = tmp_path / "attempt.txt" + attempt_path.write_text("0", encoding="utf-8") + + result = run_parallel_jobs( + _always_oom_worker, + [{"kernel_name": "always-oom", "attempt_path": str(attempt_path)}], + )[0] + output = "\n".join(log_messages) + + assert attempt_path.read_text(encoding="utf-8") == "3" + assert result["failure"] == { + "kind": "possible_oom", + "reason": "worker killed by SIGKILL (possible OOM)", + "attempts": 3, + "exitcode": -signal.SIGKILL, + } + assert "retry 1/2" in output + assert "retry 2/2" in output + assert "not retrying" in output + + +def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") + attempt_path = tmp_path / "attempt.txt" + attempt_path.write_text("0", encoding="utf-8") + + results = run_parallel_jobs( + _crash_then_succeed, + [ + { + "kernel_name": "dead-job", + "attempt_path": str(attempt_path), + "crashes": 3, + } + ], + ) + + assert attempt_path.read_text(encoding="utf-8") == "2" + assert results[0]["compile_time"] is None + assert results[0]["failure"] == { + "kind": "worker_crash", + "reason": "worker crashed (exitcode=17)", + "attempts": 2, + "exitcode": 17, + } + + +def test_incomplete_pool_results_raise_before_summary(log_messages): + with pytest.raises(RuntimeError, match="0/1 jobs completed"): + parallel._finalize_pool_results( + [None], + num_jobs=1, + completed=0, + succeeded_jobs=0, + failed_jobs=0, + retries_used=0, + ) + + assert not any("succeeded" in message for message in log_messages) + + +def test_failure_results_preserve_distinct_causes(): + jobs = [ + {"kernel_name": "signal", "outcome": "signal"}, + {"kernel_name": "exit-137", "outcome": "exit-137"}, + {"kernel_name": "type-error", "outcome": "type-error"}, + {"kernel_name": "compile-error", "outcome": "compile-error"}, + ] + + results = run_parallel_jobs(_mixed_outcome_worker, jobs) + failures = [result["failure"] for result in results] + + assert failures[0] == { + "kind": "worker_crash", + "reason": f"worker crashed (exitcode={-signal.SIGTERM})", + "attempts": 1, + "exitcode": -signal.SIGTERM, + } + assert failures[1] == { + "kind": "worker_crash", + "reason": "worker crashed (exitcode=137)", + "attempts": 1, + "exitcode": 137, + } + assert failures[2]["kind"] == "worker_exception" + assert failures[2]["reason"] == "TypeError: synthetic worker type error" + assert failures[2]["attempts"] == 1 + assert failures[2]["exitcode"] == 1 + assert failures[2]["traceback"].startswith("Traceback (most recent call last):") + assert "test_parallel_jobs.py" in failures[2]["traceback"] + assert failures[2]["traceback"].endswith("TypeError: synthetic worker type error\n") + assert failures[3] == { + "kind": "compile_error", + "reason": "synthetic codegen failure", + "attempts": 1, + } + assert len({json.dumps(failure, sort_keys=True) for failure in failures}) == 4 + + +def test_python_exception_is_not_retried_or_duplicated_to_stderr(monkeypatch, capfd): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") + + result = run_parallel_jobs( + _mixed_outcome_worker, + [{"kernel_name": "type-error", "outcome": "type-error"}], + )[0] + + assert result["failure"]["kind"] == "worker_exception" + assert result["failure"]["attempts"] == 1 + assert result["failure"]["traceback"].startswith("Traceback (most recent call last):") + assert capfd.readouterr().err == "" + + +def test_large_exception_diagnostics_are_bounded(): + result = run_parallel_jobs( + _mixed_outcome_worker, + [{"kernel_name": "large-error", "outcome": "large-error"}], + )[0] + failure = result["failure"] + + assert len(failure["reason"]) <= parallel._MAX_FAILURE_REASON_CHARS + assert "... failure reason truncated ..." in failure["reason"] + assert len(failure["traceback"]) <= parallel._MAX_TRACEBACK_CHARS + assert failure["traceback"].startswith("Traceback (most recent call last):") + assert "... traceback truncated ..." in failure["traceback"] + assert failure["traceback"].endswith("X" * 32 + "\n") + + +def test_exception_marker_is_preserved_when_process_exits_zero(monkeypatch): + monkeypatch.setattr( + parallel, + "_run_one_to_file", + _exception_marker_then_exit_zero, + ) + + result = run_parallel_jobs( + _success_worker, + [{"kernel_name": "marker", "index": 0}], + )[0] + + assert result["failure"] == { + "kind": "worker_exception", + "reason": "TypeError: preserved reason", + "attempts": 1, + "exitcode": 0, + "traceback": "Traceback (most recent call last): preserved frame", + } + + +def test_malformed_exception_marker_is_localized_to_invalid_result(monkeypatch): + monkeypatch.setattr( + parallel, + "_run_one_to_file", + _malformed_exception_marker_then_exit_zero, + ) + + result = run_parallel_jobs( + _success_worker, + [{"kernel_name": "malformed-marker", "index": 0}], + )[0] + + assert result["failure"] == { + "kind": "invalid_result", + "reason": ("worker exception marker must contain a dictionary, got str"), + "attempts": 1, + "exitcode": 0, + } + + +def test_invalid_result_reasons_distinguish_missing_and_non_dict(monkeypatch): + non_dict = run_parallel_jobs( + _non_dict_worker, + [{"kernel_name": "non-dict"}], + )[0] + + monkeypatch.setattr(parallel, "_run_one_to_file", _exit_without_result) + missing = run_parallel_jobs( + _success_worker, + [{"kernel_name": "missing", "index": 0}], + )[0] + + assert non_dict["failure"]["kind"] == "invalid_result" + assert non_dict["failure"]["reason"] == "worker returned list, expected a dictionary" + assert missing["failure"]["kind"] == "invalid_result" + assert missing["failure"]["reason"] == "worker produced no result file" + + +def test_deterministic_failure_is_not_retried(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") + attempt_path = tmp_path / "attempt.txt" + attempt_path.write_text("0", encoding="utf-8") + + results = run_parallel_jobs( + _deterministic_failure, + [ + { + "kernel_name": "compile-error", + "attempt_path": str(attempt_path), + } + ], + ) + + assert attempt_path.read_text(encoding="utf-8") == "1" + assert results[0]["compile_time"] is None + assert results[0]["failure"] == { + "kind": "compile_error", + "reason": "worker returned compile_time=None", + "attempts": 1, + } + + +def test_timed_out_worker_is_killed(monkeypatch): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.05") + + started = time.monotonic() + results = run_parallel_jobs( + _sleep_worker, + [{"kernel_name": "hung-job", "delay": 60}], + ) + + assert time.monotonic() - started < 5 + assert results[0]["compile_time"] is None + assert results[0]["failure"]["kind"] == "timeout" + assert results[0]["failure"]["exitcode"] == -signal.SIGKILL + assert results[0]["failure"]["attempts"] == 1 + assert results[0]["failure"]["reason"] == "exceeded per-job timeout (0.05s); killed" + + +def test_mem_per_worker_env_caps_automatic_workers(monkeypatch): + gib = 1024**3 + fake_psutil = SimpleNamespace(virtual_memory=lambda: SimpleNamespace(available=8 * gib)) + monkeypatch.delenv("FLYDSL_AOT_WORKERS") + monkeypatch.setenv("FLYDSL_AOT_MEM_PER_WORKER_GB", "2") + monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 16) + monkeypatch.setitem(sys.modules, "psutil", fake_psutil) + + assert parallel._get_max_workers(num_jobs=100) == 4 + + +def test_missing_psutil_warns_and_uses_conservative_limit(monkeypatch, log_messages): + monkeypatch.delenv("FLYDSL_AOT_WORKERS") + monkeypatch.setenv("FLYDSL_AOT_MEM_PER_WORKER_GB", "2") + monkeypatch.setitem(sys.modules, "psutil", None) + monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 16) + + with pytest.warns(RuntimeWarning, match="limiting AOT concurrency to 4 workers"): + assert parallel._get_max_workers(num_jobs=100) == 4 + assert "psutil is not installed; automatic AOT memory limiting is unavailable" in "\n".join(log_messages)