From 103002aaa6847588817290e9a9be7e9fe9dbe1a8 Mon Sep 17 00:00:00 2001 From: zhimding Date: Wed, 19 Aug 2026 05:28:19 +0000 Subject: [PATCH 01/11] [AOT] Move parallel scheduling into FlyDSL Centralize fork-based AOT job orchestration and environment controls so downstream integrations no longer maintain their own process pools. Co-authored-by: Cursor --- python/flydsl/__init__.py | 2 +- python/flydsl/utils/env.py | 58 ++++++- python/flydsl/utils/parallel.py | 274 ++++++++++++++++++++++++++++++++ tests/README.md | 8 + tests/unit/test_parallel.py | 228 ++++++++++++++++++++++++++ 5 files changed, 568 insertions(+), 2 deletions(-) create mode 100644 python/flydsl/utils/parallel.py create mode 100644 tests/unit/test_parallel.py diff --git a/python/flydsl/__init__.py b/python/flydsl/__init__.py index b0c97b06d..056285c32 100644 --- a/python/flydsl/__init__.py +++ b/python/flydsl/__init__.py @@ -2,7 +2,7 @@ # Copyright (c) 2025 FlyDSL Project Contributors # ruff: noqa: I001 -__version__ = "0.3.1" +__version__ = "0.3.2" from .autotune import Config as Config, autotune as autotune diff --git a/python/flydsl/utils/env.py b/python/flydsl/utils/env.py index 1839dd7df..732d241a6 100644 --- a/python/flydsl/utils/env.py +++ b/python/flydsl/utils/env.py @@ -106,6 +106,35 @@ def parse_value(self, raw: str) -> int: return int(raw) +class OptFloat(EnvOption[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, + ): + validator = None + if min_value is not None or max_value is not None: + + def validator(v: float) -> bool: + if min_value is not None and v < min_value: + return False + if max_value is not None and v > max_value: + return False + return True + + super().__init__(default, env_var, description, validator) + self.min_value = min_value + self.max_value = max_value + + def parse_value(self, raw: str) -> float: + return float(raw) + + class OptStr(EnvOption[str]): """String environment option with optional ``choices`` validation.""" @@ -224,6 +253,31 @@ 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; when unset, use the CPU and available-memory based automatic limit" + ), + ) + 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 +350,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..d1da9326e --- /dev/null +++ b/python/flydsl/utils/parallel.py @@ -0,0 +1,274 @@ +# 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 tempfile +import time +from collections.abc import Callable +from multiprocessing.connection import wait as wait_for_sentinels +from pathlib import Path +from typing import Any + +from .env import aot +from .file import atomic_write + +_DEFAULT_MAX_WORKERS = 64 + + +def _run_one_to_file( + worker: Callable[..., dict[str, Any]], + kwargs: dict[str, Any], + out_path: str, +) -> None: + result = worker(**kwargs) + with atomic_write(Path(out_path), mode="w", encoding="utf-8") as output: + json.dump(result, output) + + +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_cap(default_workers: int) -> int: + per_worker_gb = aot.mem_per_worker_gb + if per_worker_gb <= 0: + return default_workers + + try: + import psutil + + available_gb = psutil.virtual_memory().available / (1024**3) + except Exception: # noqa: BLE001 + return default_workers + return min(default_workers, max(1, int(available_gb / per_worker_gb))) + + +def _get_max_workers(num_jobs: int) -> int: + if "FLYDSL_AOT_WORKERS" in os.environ: + max_workers = max(aot.workers, 1) + 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 _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] | None]: + """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 + progress_stride = max(1, num_jobs // 20) + + queue = list(range(num_jobs)) + queue.reverse() + running: dict[Any, tuple[int, float | None]] = {} + + def launch() -> None: + while queue and len(running) < max_workers: + index = queue.pop() + 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() -> None: + nonlocal completed + completed += 1 + if completed % progress_stride == 0 or completed == num_jobs: + print(f" ... {completed}/{num_jobs} jobs done", flush=True) + + def retry_or_drop(index: int, reason: str) -> None: + nonlocal retries_used + if attempts[index] < max_retries: + attempts[index] += 1 + retries_used += 1 + queue.append(index) + print( + f"[flydsl] AOT job {_job_label(jobs[index])} {reason}; retry {attempts[index]}/{max_retries}", + flush=True, + ) + else: + note_done() + + def reap(process: Any) -> None: + index, _ = running.pop(process) + out_path = os.path.join(result_dir, f"k{index}.json") + try: + if process.exitcode != 0: + retry_or_drop( + index, + f"worker crashed (exitcode={process.exitcode})", + ) + return + + result: dict[str, Any] | None = None + if os.path.isfile(out_path): + try: + with open(out_path, encoding="utf-8") as result_file: + loaded = json.load(result_file) + if isinstance(loaded, dict): + result = loaded + except Exception: # noqa: BLE001 + result = None + if result is None: + result = { + "kernel_name": _job_label(jobs[index]), + "compile_time": None, + } + results[index] = result + note_done() + 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): + index, deadline = running[process] + if deadline is not None and now > deadline and process.is_alive(): + process.kill() + process.join() + running.pop(process) + process.close() + retry_or_drop( + index, + f"exceeded per-job timeout ({kernel_timeout:.0f}s); killed", + ) + + 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() + + if retries_used: + print( + f"[flydsl] AOT: {retries_used} retr{'y' if retries_used == 1 else 'ies'} after abnormal worker exits", + flush=True, + ) + return results + + +def run_jobs_parallel( + 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. Abnormal exits and timeouts are retried before + being normalized to the same failure shape. + + This executor uses the Linux ``fork`` multiprocessing context because it is + intended for compile-only workers that inherit the initialized compiler. + """ + if not jobs: + return [] + + max_workers = _get_max_workers(len(jobs)) + print( + f"[flydsl] AOT: {len(jobs)} jobs, {max_workers} worker processes", + flush=True, + ) + + result_dir = tempfile.mkdtemp(prefix="flydsl_aot_results_") + try: + raw_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) + + results: list[dict[str, Any]] = [] + for job, result in zip(jobs, raw_results): + if result is None: + results.append( + { + "kernel_name": _job_label(job), + "compile_time": None, + } + ) + else: + results.append(result) + return results + + +__all__ = ["run_jobs_parallel"] diff --git a/tests/README.md b/tests/README.md index efc4ed7a9..47401e37d 100644 --- a/tests/README.md +++ b/tests/README.md @@ -52,6 +52,14 @@ 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` (unset: 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_jobs_parallel` uses the Linux `fork` multiprocessing +context and is intended for compile-only AOT jobs. It is not a device-runtime +executor. Session-level pytest options are supported in `tests/conftest.py`: diff --git a/tests/unit/test_parallel.py b/tests/unit/test_parallel.py new file mode 100644 index 000000000..854ae6dd4 --- /dev/null +++ b/tests/unit/test_parallel.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +import fcntl +import json +import os +import time +from pathlib import Path + +import pytest + +import flydsl.utils.parallel as parallel +from flydsl.utils.env import aot +from flydsl.utils.parallel import run_jobs_parallel + +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): + 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} + + +@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") + + +def test_empty_jobs_do_not_invoke_worker(monkeypatch): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "invalid") + assert run_jobs_parallel(_success_worker, []) == [] + + +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_jobs_parallel(_success_worker, jobs) + + assert [result["index"] for result in results] == [0, 1, 2] + + +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_jobs_parallel(_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_jobs_parallel( + _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_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_jobs_parallel( + _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 == [{"kernel_name": "dead-job", "compile_time": None}] + + +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_jobs_parallel( + _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 + + +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_jobs_parallel( + _sleep_worker, + [{"kernel_name": "hung-job", "delay": 60}], + ) + + assert time.monotonic() - started < 5 + assert results == [{"kernel_name": "hung-job", "compile_time": None}] + + +def test_temporary_result_directory_is_removed(monkeypatch, tmp_path): + result_dir = tmp_path / "results" + + def make_result_dir(prefix): + assert prefix == "flydsl_aot_results_" + result_dir.mkdir() + return str(result_dir) + + monkeypatch.setattr(parallel.tempfile, "mkdtemp", make_result_dir) + + run_jobs_parallel( + _success_worker, + [{"kernel_name": "one", "index": 0}], + ) + + assert not result_dir.exists() + + +@pytest.mark.parametrize( + ("variable", "accessor"), + [ + ("FLYDSL_AOT_WORKERS", lambda: aot.workers), + ("FLYDSL_AOT_MEM_PER_WORKER_GB", lambda: aot.mem_per_worker_gb), + ("FLYDSL_AOT_TIMEOUT", lambda: aot.timeout), + ("FLYDSL_AOT_MAX_RETRIES", lambda: aot.max_retries), + ], +) +def test_invalid_environment_value_raises(monkeypatch, variable, accessor): + monkeypatch.setenv(variable, "invalid") + with pytest.raises(ValueError, match=variable): + accessor() From 5ab7607ec37fd4a9773d9d4c41df17fda6a637f5 Mon Sep 17 00:00:00 2001 From: zhimding Date: Wed, 19 Aug 2026 08:58:19 +0000 Subject: [PATCH 02/11] [AOT] Clarify parallel job scheduler API Rename the helper for clearer call sites, keep the package version unchanged, and strengthen failure and environment-control coverage. Co-authored-by: Cursor --- python/flydsl/__init__.py | 2 +- python/flydsl/utils/parallel.py | 4 +- tests/README.md | 2 +- ...test_parallel.py => test_parallel_jobs.py} | 64 ++++++++++++++++--- 4 files changed, 59 insertions(+), 13 deletions(-) rename tests/unit/{test_parallel.py => test_parallel_jobs.py} (78%) diff --git a/python/flydsl/__init__.py b/python/flydsl/__init__.py index 056285c32..b0c97b06d 100644 --- a/python/flydsl/__init__.py +++ b/python/flydsl/__init__.py @@ -2,7 +2,7 @@ # Copyright (c) 2025 FlyDSL Project Contributors # ruff: noqa: I001 -__version__ = "0.3.2" +__version__ = "0.3.1" from .autotune import Config as Config, autotune as autotune diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index d1da9326e..2a2d4b3d6 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -220,7 +220,7 @@ def reap(process: Any) -> None: return results -def run_jobs_parallel( +def run_parallel_jobs( worker: Callable[..., dict[str, Any]], jobs: list[dict[str, Any]], ) -> list[dict[str, Any]]: @@ -271,4 +271,4 @@ def run_jobs_parallel( return results -__all__ = ["run_jobs_parallel"] +__all__ = ["run_parallel_jobs"] diff --git a/tests/README.md b/tests/README.md index 47401e37d..f3d841f72 100644 --- a/tests/README.md +++ b/tests/README.md @@ -57,7 +57,7 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. | 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_jobs_parallel` uses the Linux `fork` multiprocessing +`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. diff --git a/tests/unit/test_parallel.py b/tests/unit/test_parallel_jobs.py similarity index 78% rename from tests/unit/test_parallel.py rename to tests/unit/test_parallel_jobs.py index 854ae6dd4..0320336c2 100644 --- a/tests/unit/test_parallel.py +++ b/tests/unit/test_parallel_jobs.py @@ -4,14 +4,16 @@ import fcntl import json import os +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_jobs_parallel +from flydsl.utils.parallel import run_parallel_jobs pytestmark = [pytest.mark.l0_backend_agnostic] @@ -72,6 +74,15 @@ def _sleep_worker(kernel_name, delay): return {"kernel_name": kernel_name, "compile_time": 0.01} +def _mixed_outcome_worker(kernel_name, outcome): + if outcome == "crash": + os._exit(17) + return { + "kernel_name": kernel_name, + "compile_time": None if outcome == "compile-error" else 0.01, + } + + @pytest.fixture(autouse=True) def _parallel_env(monkeypatch): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") @@ -82,7 +93,7 @@ def _parallel_env(monkeypatch): def test_empty_jobs_do_not_invoke_worker(monkeypatch): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "invalid") - assert run_jobs_parallel(_success_worker, []) == [] + assert run_parallel_jobs(_success_worker, []) == [] def test_results_follow_input_order(): @@ -92,7 +103,7 @@ def test_results_follow_input_order(): {"kernel_name": "last", "index": 2, "delay": 0.01}, ] - results = run_jobs_parallel(_success_worker, jobs) + results = run_parallel_jobs(_success_worker, jobs) assert [result["index"] for result in results] == [0, 1, 2] @@ -113,7 +124,7 @@ def test_worker_count_is_bounded(monkeypatch, tmp_path): for index in range(6) ] - results = run_jobs_parallel(_tracked_worker, jobs) + results = run_parallel_jobs(_tracked_worker, jobs) state = json.loads(state_path.read_text(encoding="utf-8")) assert state == {"active": 0, "peak": 2} @@ -126,7 +137,7 @@ def test_crashed_worker_is_retried(monkeypatch, tmp_path): attempt_path = tmp_path / "attempt.txt" attempt_path.write_text("0", encoding="utf-8") - results = run_jobs_parallel( + results = run_parallel_jobs( _crash_then_succeed, [ { @@ -147,7 +158,7 @@ def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): attempt_path = tmp_path / "attempt.txt" attempt_path.write_text("0", encoding="utf-8") - results = run_jobs_parallel( + results = run_parallel_jobs( _crash_then_succeed, [ { @@ -162,12 +173,36 @@ def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): assert results == [{"kernel_name": "dead-job", "compile_time": None}] +def test_failed_jobs_do_not_stop_remaining_jobs(): + jobs = [ + {"kernel_name": "first", "outcome": "success"}, + {"kernel_name": "crashed", "outcome": "crash"}, + {"kernel_name": "compile-error", "outcome": "compile-error"}, + {"kernel_name": "last", "outcome": "success"}, + ] + + results = run_parallel_jobs(_mixed_outcome_worker, jobs) + + assert [result["kernel_name"] for result in results] == [ + "first", + "crashed", + "compile-error", + "last", + ] + assert [result["compile_time"] for result in results] == [ + 0.01, + None, + None, + 0.01, + ] + + 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_jobs_parallel( + results = run_parallel_jobs( _deterministic_failure, [ { @@ -186,7 +221,7 @@ def test_timed_out_worker_is_killed(monkeypatch): monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.05") started = time.monotonic() - results = run_jobs_parallel( + results = run_parallel_jobs( _sleep_worker, [{"kernel_name": "hung-job", "delay": 60}], ) @@ -205,7 +240,7 @@ def make_result_dir(prefix): monkeypatch.setattr(parallel.tempfile, "mkdtemp", make_result_dir) - run_jobs_parallel( + run_parallel_jobs( _success_worker, [{"kernel_name": "one", "index": 0}], ) @@ -213,6 +248,17 @@ def make_result_dir(prefix): assert not result_dir.exists() +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 + + @pytest.mark.parametrize( ("variable", "accessor"), [ From ef4e14beb8cad7c0b45deb7008882e87dcb4727f Mon Sep 17 00:00:00 2001 From: zhimding Date: Wed, 19 Aug 2026 15:40:32 +0000 Subject: [PATCH 03/11] [AOT] Harden parallel job failure handling Make memory limiting fail safe, back off after possible OOM kills, and preserve structured failure causes so large AOT builds remain diagnosable. Co-authored-by: Cursor --- MANIFEST.in | 1 + python/flydsl/utils/parallel.py | 250 ++++++++++++++++++++++++++----- requirements.txt | 1 + tests/README.md | 15 +- tests/unit/test_parallel_jobs.py | 168 ++++++++++++++++++++- 5 files changed, 390 insertions(+), 45 deletions(-) create mode 100644 requirements.txt diff --git a/MANIFEST.in b/MANIFEST.in index a7c4ebf0d..11807f054 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,7 @@ include LICENSE include README.md include pyproject.toml +include requirements.txt include setup.py recursive-include scripts *.sh diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index 2a2d4b3d6..408de6dc1 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -8,8 +8,10 @@ import multiprocessing import os import shutil +import signal import tempfile import time +import traceback from collections.abc import Callable from multiprocessing.connection import wait as wait_for_sentinels from pathlib import Path @@ -19,6 +21,13 @@ from .file import atomic_write _DEFAULT_MAX_WORKERS = 64 +_POSSIBLE_OOM_EXITCODES = {-signal.SIGKILL, 128 + 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 _run_one_to_file( @@ -26,9 +35,21 @@ def _run_one_to_file( kwargs: dict[str, Any], out_path: str, ) -> None: - result = worker(**kwargs) - with atomic_write(Path(out_path), mode="w", encoding="utf-8") as output: - json.dump(result, output) + try: + result = worker(**kwargs) + _write_json_file(out_path, result) + except Exception as error: + _write_json_file( + out_path, + { + _WORKER_EXCEPTION_KEY: { + "kind": "worker_exception", + "reason": f"{type(error).__name__}: {error}", + "traceback": traceback.format_exc(), + } + }, + ) + raise def _affinity_aware_cpu_count() -> int: @@ -55,10 +76,19 @@ def _memory_worker_cap(default_workers: int) -> int: try: import psutil + except ImportError as error: + raise RuntimeError( + "psutil is required for automatic AOT worker memory limiting; " + "install FlyDSL runtime dependencies or set FLYDSL_AOT_WORKERS explicitly" + ) from error + try: available_gb = psutil.virtual_memory().available / (1024**3) - except Exception: # noqa: BLE001 - return default_workers + except Exception as error: + raise RuntimeError( + "failed to query available memory for the AOT worker limit; " + "set FLYDSL_AOT_WORKERS explicitly to bypass automatic detection" + ) from error return min(default_workers, max(1, int(available_gb / per_worker_gb))) @@ -75,6 +105,32 @@ 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 _run_file_pool( worker: Callable[..., dict[str, Any]], jobs: list[dict[str, Any]], @@ -91,6 +147,7 @@ def _run_file_pool( attempts = [0] * num_jobs retries_used = 0 completed = 0 + failed_jobs = 0 progress_stride = max(1, num_jobs // 20) queue = list(range(num_jobs)) @@ -114,13 +171,46 @@ def launch() -> None: deadline = time.monotonic() + kernel_timeout if kernel_timeout > 0 else None running[process] = (index, deadline) - def note_done() -> None: - nonlocal completed + def note_done(*, is_failure: bool = False) -> None: + nonlocal completed, failed_jobs completed += 1 + failed_jobs += int(is_failure) if completed % progress_stride == 0 or completed == num_jobs: - print(f" ... {completed}/{num_jobs} jobs done", flush=True) + print( + f" ... {completed}/{num_jobs} jobs finished ({failed_jobs} failed)", + flush=True, + ) - def retry_or_drop(index: int, reason: str) -> None: + 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, + ) + print( + f"[flydsl] AOT job {_job_label(jobs[index])} {reason}; not retrying", + flush=True, + ) + 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 @@ -131,35 +221,108 @@ def retry_or_drop(index: int, reason: str) -> None: flush=True, ) else: - note_done() + finish_failure( + index, + kind=kind, + reason=reason, + exitcode=exitcode, + traceback_text=traceback_text, + ) def reap(process: Any) -> None: + nonlocal max_workers 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" + + if process.exitcode in _POSSIBLE_OOM_EXITCODES: + previous_max_workers = max_workers + max_workers = max(1, max_workers // 2) + reason = ( + "worker killed by SIGKILL (possible OOM)" + if process.exitcode == -signal.SIGKILL + else f"worker exited with code {process.exitcode} (possible OOM)" + ) + if max_workers == previous_max_workers: + finish_failure( + index, + kind="possible_oom", + reason=f"{reason} at the minimum worker limit", + exitcode=process.exitcode, + ) + else: + retry_or_drop( + index, + kind="possible_oom", + reason=(f"{reason}; reduced worker limit {previous_max_workers}->{max_workers}"), + exitcode=process.exitcode, + ) + return + if process.exitcode != 0: + worker_exception = loaded.get(_WORKER_EXCEPTION_KEY) if isinstance(loaded, dict) else None + if isinstance(worker_exception, dict): + 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 + kind = "worker_exception" + else: + reason = f"worker crashed (exitcode={process.exitcode})" + traceback_text = None + kind = "worker_crash" retry_or_drop( index, - f"worker crashed (exitcode={process.exitcode})", + kind=kind, + reason=reason, + exitcode=process.exitcode, + traceback_text=traceback_text, ) return - result: dict[str, Any] | None = None - if os.path.isfile(out_path): - try: - with open(out_path, encoding="utf-8") as result_file: - loaded = json.load(result_file) - if isinstance(loaded, dict): - result = loaded - except Exception: # noqa: BLE001 - result = None - if result is None: - result = { - "kernel_name": _job_label(jobs[index]), - "compile_time": None, - } + if load_error is not None: + finish_failure( + index, + kind="invalid_result", + reason=load_error, + exitcode=process.exitcode, + ) + return + if not isinstance(loaded, dict) or _WORKER_EXCEPTION_KEY in loaded: + finish_failure( + index, + kind="invalid_result", + reason=(f"worker returned {type(loaded).__name__}, expected a dictionary"), + exitcode=process.exitcode, + ) + return + + 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 - note_done() + note_done(is_failure=result.get("compile_time") is None) finally: process.close() @@ -190,10 +353,13 @@ def reap(process: Any) -> None: process.kill() process.join() running.pop(process) + exitcode = process.exitcode process.close() retry_or_drop( index, - f"exceeded per-job timeout ({kernel_timeout:.0f}s); killed", + kind="timeout", + reason=(f"exceeded per-job timeout ({kernel_timeout:.0f}s); killed"), + exitcode=exitcode, ) launch() @@ -212,11 +378,13 @@ def reap(process: Any) -> None: pass running.clear() - if retries_used: - print( - f"[flydsl] AOT: {retries_used} retr{'y' if retries_used == 1 else 'ies'} after abnormal worker exits", - flush=True, - ) + retry_label = "retry" if retries_used == 1 else "retries" + print( + f"[flydsl] AOT: {num_jobs - failed_jobs} succeeded, " + f"{failed_jobs} failed; {retries_used} {retry_label} " + "after abnormal worker exits", + flush=True, + ) return results @@ -229,11 +397,16 @@ def run_parallel_jobs( 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. Abnormal exits and timeouts are retried before - being normalized to the same failure shape. + 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 ``traceback``. This executor uses the Linux ``fork`` multiprocessing context because it is intended for compile-only workers that inherit the initialized compiler. + A possible OOM exit (``-SIGKILL`` or shell-style ``137``) halves the + concurrency limit before retrying and is not retried once that limit + reaches one. """ if not jobs: return [] @@ -261,10 +434,11 @@ def run_parallel_jobs( for job, result in zip(jobs, raw_results): if result is None: results.append( - { - "kernel_name": _job_label(job), - "compile_time": None, - } + _failure_result( + job, + kind="scheduler_error", + reason="scheduler returned no result for the job", + ) ) else: results.append(result) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 000000000..a4d92cc08 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +psutil diff --git a/tests/README.md b/tests/README.md index f3d841f72..4d4a02cec 100644 --- a/tests/README.md +++ b/tests/README.md @@ -59,7 +59,20 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. `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. +executor. Its automatic memory cap requires the runtime `psutil` dependency and +fails with an actionable error instead of silently disabling the cap when +memory cannot be queried. A worker killed by `SIGKILL` (`exitcode=-9`, or shell +form `137`; possible OOM) halves the worker limit before retrying and is not +retried once the limit reaches one. +Progress and final summary lines report terminal failure counts separately from +the number of jobs that have finished. + +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`, +`invalid_result`, and `scheduler_error`; `reason` and `attempts` are always +included when known, while process failures also carry `exitcode` and uncaught +Python exceptions carry `traceback`. 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 index 0320336c2..652ca3917 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -4,6 +4,7 @@ import fcntl import json import os +import signal import sys import time from pathlib import Path @@ -62,6 +63,17 @@ def _crash_then_succeed(kernel_name, attempt_path, crashes): return {"kernel_name": kernel_name, "compile_time": 0.01} +def _sigkill_once_worker(kernel_name, attempt_path=None, delay=0.0): + if attempt_path is not None: + path = Path(attempt_path) + attempt = int(path.read_text(encoding="utf-8")) + 1 + path.write_text(str(attempt), encoding="utf-8") + if attempt == 1: + os.kill(os.getpid(), signal.SIGKILL) + time.sleep(delay) + 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 @@ -77,10 +89,22 @@ def _sleep_worker(kernel_name, delay): def _mixed_outcome_worker(kernel_name, outcome): if outcome == "crash": os._exit(17) - return { - "kernel_name": kernel_name, - "compile_time": None if outcome == "compile-error" else 0.01, - } + 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 == "compile-error": + return { + "kernel_name": kernel_name, + "compile_time": None, + "failure": { + "kind": "compile_error", + "reason": "synthetic codegen failure", + }, + } + return {"kernel_name": kernel_name, "compile_time": 0.01} @pytest.fixture(autouse=True) @@ -152,6 +176,55 @@ def test_crashed_worker_is_retried(monkeypatch, tmp_path): assert results[0]["compile_time"] is not None +def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, capsys): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") + 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( + _sigkill_once_worker, + [ + { + "kernel_name": "oom-job", + "attempt_path": str(attempt_path), + }, + { + "kernel_name": "companion", + "delay": 0.1, + }, + ], + ) + + assert attempt_path.read_text(encoding="utf-8") == "2" + assert all(result["compile_time"] is not None for result in results) + assert "possible OOM); reduced worker limit 2->1; retry 1/1" in capsys.readouterr().out + + +def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, capsys): + 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( + _sigkill_once_worker, + [ + { + "kernel_name": "oom-job", + "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"] == "possible_oom" + assert results[0]["failure"]["exitcode"] == -signal.SIGKILL + assert results[0]["failure"]["attempts"] == 1 + assert "at the minimum worker limit; not retrying" in capsys.readouterr().out + + def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") @@ -170,7 +243,73 @@ def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): ) assert attempt_path.read_text(encoding="utf-8") == "2" - assert results == [{"kernel_name": "dead-job", "compile_time": None}] + 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_final_logs_report_permanent_failures(monkeypatch, capsys): + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") + jobs = [ + {"kernel_name": "success-0", "outcome": "success"}, + {"kernel_name": "failed-0", "outcome": "crash"}, + {"kernel_name": "failed-1", "outcome": "crash"}, + {"kernel_name": "success-1", "outcome": "success"}, + {"kernel_name": "failed-2", "outcome": "crash"}, + ] + + results = run_parallel_jobs(_mixed_outcome_worker, jobs) + output = capsys.readouterr().out + + assert [result["compile_time"] for result in results] == [ + 0.01, + None, + None, + 0.01, + None, + ] + for kernel_name in ("failed-0", "failed-1", "failed-2"): + assert f"AOT job {kernel_name} worker crashed (exitcode=17); not retrying" in output + assert "... 5/5 jobs finished (3 failed)" in output + assert ("[flydsl] AOT: 2 succeeded, 3 failed; 6 retries after abnormal worker exits") in output + + +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"] == "possible_oom" + assert failures[1]["reason"].startswith("worker exited with code 137 (possible OOM)") + assert failures[1]["attempts"] == 1 + assert failures[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 "TypeError: synthetic worker type error" in failures[2]["traceback"] + 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_failed_jobs_do_not_stop_remaining_jobs(): @@ -214,6 +353,11 @@ def test_deterministic_failure_is_not_retried(monkeypatch, tmp_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): @@ -227,7 +371,10 @@ def test_timed_out_worker_is_killed(monkeypatch): ) assert time.monotonic() - started < 5 - assert results == [{"kernel_name": "hung-job", "compile_time": None}] + 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 def test_temporary_result_directory_is_removed(monkeypatch, tmp_path): @@ -259,6 +406,15 @@ def test_mem_per_worker_env_caps_automatic_workers(monkeypatch): assert parallel._get_max_workers(num_jobs=100) == 4 +def test_missing_psutil_does_not_silently_disable_memory_cap(monkeypatch): + monkeypatch.delenv("FLYDSL_AOT_WORKERS") + monkeypatch.setenv("FLYDSL_AOT_MEM_PER_WORKER_GB", "2") + monkeypatch.setitem(sys.modules, "psutil", None) + + with pytest.raises(RuntimeError, match="psutil is required"): + parallel._get_max_workers(num_jobs=100) + + @pytest.mark.parametrize( ("variable", "accessor"), [ From 66e1468a592f49bd54e8a95e0c1140a89d5f2780 Mon Sep 17 00:00:00 2001 From: zhimding Date: Thu, 20 Aug 2026 00:50:54 +0000 Subject: [PATCH 04/11] [AOT] Refine scheduler configuration and retries Align automatic worker semantics, queue retries fairly behind pending jobs, and route precise scheduler diagnostics through the FlyDSL logger. Co-authored-by: Cursor --- python/flydsl/utils/env.py | 69 ++++++++++++++++---------- python/flydsl/utils/parallel.py | 57 +++++++++++++--------- tests/README.md | 9 ++-- tests/unit/test_parallel_jobs.py | 84 +++++++++++++++++++++++++++++--- 4 files changed, 159 insertions(+), 60 deletions(-) diff --git a/python/flydsl/utils/env.py b/python/flydsl/utils/env.py index 732d241a6..17d14c62f 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,22 @@ 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, ): 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,14 +101,36 @@ 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 - def parse_value(self, raw: str) -> int: - return int(raw) + def parse_value(self, raw: str) -> NumberT: + return self.parser(raw) -class OptFloat(EnvOption[float]): +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, + ): + super().__init__( + default, + int, + env_var, + description, + min_value, + max_value, + ) + + +class OptFloat(_OptNumber[float]): """Floating-point environment option with optional min/max validation.""" def __init__( @@ -117,22 +141,14 @@ def __init__( min_value: Optional[float] = None, max_value: Optional[float] = None, ): - validator = None - if min_value is not None or max_value is not None: - - def validator(v: float) -> bool: - if min_value is not None and v < min_value: - return False - if max_value is not None and v > max_value: - return False - return True - - super().__init__(default, env_var, description, validator) - self.min_value = min_value - self.max_value = max_value - - def parse_value(self, raw: str) -> float: - return float(raw) + super().__init__( + default, + float, + env_var, + description, + min_value, + max_value, + ) class OptStr(EnvOption[str]): @@ -261,7 +277,8 @@ class AotEnvManager(EnvManager): workers = OptInt( 0, description=( - "Maximum concurrent worker processes; when unset, use the CPU and available-memory based automatic limit" + "Maximum concurrent worker processes; unset, empty, or non-positive values use the CPU and " + "available-memory based automatic limit" ), ) mem_per_worker_gb = OptFloat( diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index 408de6dc1..5294d3c1b 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -12,6 +12,7 @@ import tempfile import time import traceback +from collections import deque from collections.abc import Callable from multiprocessing.connection import wait as wait_for_sentinels from pathlib import Path @@ -19,6 +20,7 @@ from .env import aot from .file import atomic_write +from .logger import log _DEFAULT_MAX_WORKERS = 64 _POSSIBLE_OOM_EXITCODES = {-signal.SIGKILL, 128 + signal.SIGKILL} @@ -93,8 +95,10 @@ def _memory_worker_cap(default_workers: int) -> int: def _get_max_workers(num_jobs: int) -> int: - if "FLYDSL_AOT_WORKERS" in os.environ: - max_workers = max(aot.workers, 1) + workers_raw = os.environ.get("FLYDSL_AOT_WORKERS", "").strip() + configured_workers = aot.workers if workers_raw else 0 + 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) @@ -150,13 +154,12 @@ def _run_file_pool( failed_jobs = 0 progress_stride = max(1, num_jobs // 20) - queue = list(range(num_jobs)) - queue.reverse() + queue = deque(range(num_jobs)) running: dict[Any, tuple[int, float | None]] = {} def launch() -> None: while queue and len(running) < max_workers: - index = queue.pop() + index = queue.popleft() out_path = os.path.join(result_dir, f"k{index}.json") try: os.remove(out_path) @@ -176,9 +179,11 @@ def note_done(*, is_failure: bool = False) -> None: completed += 1 failed_jobs += int(is_failure) if completed % progress_stride == 0 or completed == num_jobs: - print( - f" ... {completed}/{num_jobs} jobs finished ({failed_jobs} failed)", - flush=True, + log().info( + "... %d/%d jobs finished (%d failed)", + completed, + num_jobs, + failed_jobs, ) def finish_failure( @@ -197,9 +202,10 @@ def finish_failure( exitcode=exitcode, traceback_text=traceback_text, ) - print( - f"[flydsl] AOT job {_job_label(jobs[index])} {reason}; not retrying", - flush=True, + log().warning( + "AOT job %s %s; not retrying", + _job_label(jobs[index]), + reason, ) note_done(is_failure=True) @@ -216,9 +222,12 @@ def retry_or_drop( attempts[index] += 1 retries_used += 1 queue.append(index) - print( - f"[flydsl] AOT job {_job_label(jobs[index])} {reason}; retry {attempts[index]}/{max_retries}", - flush=True, + log().warning( + "AOT job %s %s; retry %d/%d", + _job_label(jobs[index]), + reason, + attempts[index], + max_retries, ) else: finish_failure( @@ -358,7 +367,7 @@ def reap(process: Any) -> None: retry_or_drop( index, kind="timeout", - reason=(f"exceeded per-job timeout ({kernel_timeout:.0f}s); killed"), + reason=f"exceeded per-job timeout ({kernel_timeout:g}s); killed", exitcode=exitcode, ) @@ -379,11 +388,12 @@ def reap(process: Any) -> None: running.clear() retry_label = "retry" if retries_used == 1 else "retries" - print( - f"[flydsl] AOT: {num_jobs - failed_jobs} succeeded, " - f"{failed_jobs} failed; {retries_used} {retry_label} " - "after abnormal worker exits", - flush=True, + log().info( + "AOT: %d succeeded, %d failed; %d %s after abnormal worker exits", + num_jobs - failed_jobs, + failed_jobs, + retries_used, + retry_label, ) return results @@ -412,9 +422,10 @@ def run_parallel_jobs( return [] max_workers = _get_max_workers(len(jobs)) - print( - f"[flydsl] AOT: {len(jobs)} jobs, {max_workers} worker processes", - flush=True, + log().info( + "AOT: %d jobs, %d worker processes", + len(jobs), + max_workers, ) result_dir = tempfile.mkdtemp(prefix="flydsl_aot_results_") diff --git a/tests/README.md b/tests/README.md index 4d4a02cec..8e917015a 100644 --- a/tests/README.md +++ b/tests/README.md @@ -52,7 +52,7 @@ 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` (unset: automatic CPU/memory limit) | +| 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`) | @@ -63,9 +63,12 @@ executor. Its automatic memory cap requires the runtime `psutil` dependency and fails with an actionable error instead of silently disabling the cap when memory cannot be queried. A worker killed by `SIGKILL` (`exitcode=-9`, or shell form `137`; possible OOM) halves the worker limit before retrying and is not -retried once the limit reaches one. +retried once the limit reaches one. Retries are appended behind jobs that have +not started yet instead of jumping to the front of the queue. Progress and final summary lines report terminal failure counts separately from -the number of jobs that have finished. +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`, diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py index 652ca3917..f441e6aa5 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -54,7 +54,10 @@ def _tracked_worker(kernel_name, index, state_path, lock_path, delay): fcntl.flock(lock_file, fcntl.LOCK_UN) -def _crash_then_succeed(kernel_name, attempt_path, crashes): +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") @@ -115,11 +118,35 @@ def _parallel_env(monkeypatch): 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_empty_jobs_do_not_invoke_worker(monkeypatch): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "invalid") assert run_parallel_jobs(_success_worker, []) == [] +@pytest.mark.parametrize("workers", [None, "", "0", "-2"]) +def test_non_positive_or_empty_workers_use_automatic_limit(monkeypatch, workers): + if workers is None: + monkeypatch.delenv("FLYDSL_AOT_WORKERS", raising=False) + else: + monkeypatch.setenv("FLYDSL_AOT_WORKERS", workers) + monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 8) + monkeypatch.setattr(parallel, "_memory_worker_cap", lambda workers: workers) + + assert parallel._get_max_workers(num_jobs=100) == 8 + + def test_results_follow_input_order(): jobs = [ {"kernel_name": "slow", "index": 0, "delay": 0.1}, @@ -176,7 +203,47 @@ def test_crashed_worker_is_retried(monkeypatch, tmp_path): assert results[0]["compile_time"] is not None -def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, capsys): +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_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") attempt_path = tmp_path / "attempt.txt" @@ -198,10 +265,10 @@ def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, capsys assert attempt_path.read_text(encoding="utf-8") == "2" assert all(result["compile_time"] is not None for result in results) - assert "possible OOM); reduced worker limit 2->1; retry 1/1" in capsys.readouterr().out + assert "possible OOM); reduced worker limit 2->1; retry 1/1" in "\n".join(log_messages) -def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, capsys): +def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") attempt_path = tmp_path / "attempt.txt" @@ -222,7 +289,7 @@ def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, c assert results[0]["failure"]["kind"] == "possible_oom" assert results[0]["failure"]["exitcode"] == -signal.SIGKILL assert results[0]["failure"]["attempts"] == 1 - assert "at the minimum worker limit; not retrying" in capsys.readouterr().out + assert "at the minimum worker limit; not retrying" in "\n".join(log_messages) def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): @@ -252,7 +319,7 @@ def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): } -def test_final_logs_report_permanent_failures(monkeypatch, capsys): +def test_final_logs_report_permanent_failures(monkeypatch, log_messages): monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") jobs = [ {"kernel_name": "success-0", "outcome": "success"}, @@ -263,7 +330,7 @@ def test_final_logs_report_permanent_failures(monkeypatch, capsys): ] results = run_parallel_jobs(_mixed_outcome_worker, jobs) - output = capsys.readouterr().out + output = "\n".join(log_messages) assert [result["compile_time"] for result in results] == [ 0.01, @@ -275,7 +342,7 @@ def test_final_logs_report_permanent_failures(monkeypatch, capsys): for kernel_name in ("failed-0", "failed-1", "failed-2"): assert f"AOT job {kernel_name} worker crashed (exitcode=17); not retrying" in output assert "... 5/5 jobs finished (3 failed)" in output - assert ("[flydsl] AOT: 2 succeeded, 3 failed; 6 retries after abnormal worker exits") in output + assert "AOT: 2 succeeded, 3 failed; 6 retries after abnormal worker exits" in output def test_failure_results_preserve_distinct_causes(): @@ -375,6 +442,7 @@ def test_timed_out_worker_is_killed(monkeypatch): 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_temporary_result_directory_is_removed(monkeypatch, tmp_path): From 91e311d324ed78f9ec7e9b9877729d4137196994 Mon Sep 17 00:00:00 2001 From: zhimding Date: Thu, 20 Aug 2026 01:02:47 +0000 Subject: [PATCH 05/11] [AOT] Keep psutil optional Avoid adding a package-wide runtime dependency while retaining explicit guidance for automatic AOT memory limiting. Co-authored-by: Cursor --- MANIFEST.in | 1 - python/flydsl/utils/parallel.py | 2 +- requirements.txt | 1 - tests/README.md | 13 +++++++------ 4 files changed, 8 insertions(+), 9 deletions(-) delete mode 100644 requirements.txt diff --git a/MANIFEST.in b/MANIFEST.in index 11807f054..a7c4ebf0d 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,7 +1,6 @@ include LICENSE include README.md include pyproject.toml -include requirements.txt include setup.py recursive-include scripts *.sh diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index 5294d3c1b..e75b01c19 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -81,7 +81,7 @@ def _memory_worker_cap(default_workers: int) -> int: except ImportError as error: raise RuntimeError( "psutil is required for automatic AOT worker memory limiting; " - "install FlyDSL runtime dependencies or set FLYDSL_AOT_WORKERS explicitly" + "install psutil or set FLYDSL_AOT_WORKERS explicitly" ) from error try: diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index a4d92cc08..000000000 --- a/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -psutil diff --git a/tests/README.md b/tests/README.md index 8e917015a..ddd2fc82f 100644 --- a/tests/README.md +++ b/tests/README.md @@ -59,12 +59,13 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. `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 runtime `psutil` dependency and -fails with an actionable error instead of silently disabling the cap when -memory cannot be queried. A worker killed by `SIGKILL` (`exitcode=-9`, or shell -form `137`; possible OOM) halves the worker limit before retrying and is not -retried once the limit reaches one. Retries are appended behind jobs that have -not started yet instead of jumping to the front of the queue. +executor. Its automatic memory cap requires the optional `psutil` package and +fails with an actionable error instead of silently disabling the cap when the +optional package is not installed or memory cannot be queried. A worker killed +by `SIGKILL` (`exitcode=-9`, or shell form `137`; possible OOM) halves the worker +limit before retrying and is not retried once the limit reaches one. Retries are +appended behind jobs that have not started yet instead of jumping to the front +of the queue. 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 From 156d394d123f21a153d8699f951b2f25f7885925 Mon Sep 17 00:00:00 2001 From: zhimding Date: Fri, 21 Aug 2026 02:38:12 +0000 Subject: [PATCH 06/11] [AOT] Fix scheduler result and backoff handling Trust atomically committed results, bound exception diagnostics, and make OOM/timeout backoff generation-aware with additive recovery. Co-authored-by: Cursor --- python/flydsl/utils/parallel.py | 262 +++++++++++++++--------- tests/README.md | 15 +- tests/unit/test_parallel_jobs.py | 333 ++++++++++++++++++++++++++++++- 3 files changed, 507 insertions(+), 103 deletions(-) diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index e75b01c19..0a054634a 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -16,14 +16,16 @@ from collections.abc import Callable from multiprocessing.connection import wait as wait_for_sentinels from pathlib import Path -from typing import Any +from typing import Any, cast from .env import aot from .file import atomic_write from .logger import log _DEFAULT_MAX_WORKERS = 64 -_POSSIBLE_OOM_EXITCODES = {-signal.SIGKILL, 128 + signal.SIGKILL} +_MAX_FAILURE_REASON_CHARS = 2 * 1024 +_MAX_TRACEBACK_CHARS = 16 * 1024 +_OOM_EXITCODE = -signal.SIGKILL _WORKER_EXCEPTION_KEY = "__flydsl_worker_exception__" @@ -32,6 +34,40 @@ def _write_json_file(out_path: str, payload: Any) -> None: 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], @@ -40,18 +76,20 @@ def _run_one_to_file( try: result = worker(**kwargs) _write_json_file(out_path, result) - except Exception as error: - _write_json_file( - out_path, - { - _WORKER_EXCEPTION_KEY: { - "kind": "worker_exception", - "reason": f"{type(error).__name__}: {error}", - "traceback": traceback.format_exc(), - } - }, - ) - raise + 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: @@ -143,7 +181,7 @@ def _run_file_pool( kernel_timeout: float, max_retries: int, result_dir: str, -) -> list[dict[str, Any] | None]: +) -> 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) @@ -153,9 +191,12 @@ def _run_file_pool( completed = 0 failed_jobs = 0 progress_stride = max(1, num_jobs // 20) + initial_max_workers = max_workers + launch_epoch = 0 + successful_since_backoff = 0 queue = deque(range(num_jobs)) - running: dict[Any, tuple[int, float | None]] = {} + running: dict[Any, tuple[int, float | None, int]] = {} def launch() -> None: while queue and len(running) < max_workers: @@ -172,7 +213,7 @@ def launch() -> None: ) process.start() deadline = time.monotonic() + kernel_timeout if kernel_timeout > 0 else None - running[process] = (index, deadline) + running[process] = (index, deadline, launch_epoch) def note_done(*, is_failure: bool = False) -> None: nonlocal completed, failed_jobs @@ -186,6 +227,32 @@ def note_done(*, is_failure: bool = False) -> None: failed_jobs, ) + def note_success(worker_epoch: int) -> None: + nonlocal max_workers, successful_since_backoff + if max_workers >= initial_max_workers or worker_epoch != launch_epoch: + return + successful_since_backoff += 1 + if successful_since_backoff < max_workers: + return + previous_max_workers = max_workers + max_workers += 1 + successful_since_backoff = 0 + log().info( + "AOT worker limit recovered %d->%d after healthy completions", + previous_max_workers, + max_workers, + ) + + def backoff_worker_limit(worker_epoch: int, reason: str) -> str: + nonlocal launch_epoch, max_workers, successful_since_backoff + if worker_epoch != launch_epoch: + return f"{reason}; worker limit already reduced to {max_workers} for this failure wave" + previous_max_workers = max_workers + max_workers = max(1, max_workers // 2) + launch_epoch += 1 + successful_since_backoff = 0 + return f"{reason}; reduced worker limit {previous_max_workers}->{max_workers}" + def finish_failure( index: int, *, @@ -238,9 +305,9 @@ def retry_or_drop( traceback_text=traceback_text, ) - def reap(process: Any) -> None: - nonlocal max_workers - index, _ = running.pop(process) + def reap(process: Any, *, timeout_reason: str | None = None) -> None: + nonlocal launch_epoch, max_workers, successful_since_backoff + index, _, worker_epoch = running.pop(process) out_path = os.path.join(result_dir, f"k{index}.json") try: loaded: Any = None @@ -254,15 +321,73 @@ def reap(process: Any) -> None: else: load_error = "worker produced no result file" - if process.exitcode in _POSSIBLE_OOM_EXITCODES: - previous_max_workers = max_workers - max_workers = max(1, max_workers // 2) - reason = ( - "worker killed by SIGKILL (possible OOM)" - if process.exitcode == -signal.SIGKILL - else f"worker exited with code {process.exitcode} (possible OOM)" + worker_exception = loaded.get(_WORKER_EXCEPTION_KEY) if isinstance(loaded, dict) 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 worker_exception is None: + 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) + if not is_failure: + note_success(worker_epoch) + return + + # A structured Python exception is deterministic for the same job. + # Preserve it even if teardown later changes the process exit code. + if isinstance(worker_exception, dict): + 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: + if attempts[index] < max_retries and max_workers > 1: + timeout_reason = backoff_worker_limit( + worker_epoch, + timeout_reason, + ) + retry_or_drop( + index, + kind="timeout", + reason=timeout_reason, + exitcode=process.exitcode, ) - if max_workers == previous_max_workers: + return + + if process.exitcode == _OOM_EXITCODE: + reason = "worker killed by SIGKILL (possible OOM)" + can_retry = attempts[index] < max_retries + if not can_retry: + finish_failure( + index, + kind="possible_oom", + reason=reason, + exitcode=process.exitcode, + ) + elif max_workers <= 1: finish_failure( index, kind="possible_oom", @@ -270,37 +395,21 @@ def reap(process: Any) -> None: exitcode=process.exitcode, ) else: + reason = backoff_worker_limit(worker_epoch, reason) retry_or_drop( index, kind="possible_oom", - reason=(f"{reason}; reduced worker limit {previous_max_workers}->{max_workers}"), + reason=reason, exitcode=process.exitcode, ) return if process.exitcode != 0: - worker_exception = loaded.get(_WORKER_EXCEPTION_KEY) if isinstance(loaded, dict) else None - if isinstance(worker_exception, dict): - 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 - kind = "worker_exception" - else: - reason = f"worker crashed (exitcode={process.exitcode})" - traceback_text = None - kind = "worker_crash" retry_or_drop( index, - kind=kind, - reason=reason, + kind="worker_crash", + reason=f"worker crashed (exitcode={process.exitcode})", exitcode=process.exitcode, - traceback_text=traceback_text, ) return @@ -312,7 +421,7 @@ def reap(process: Any) -> None: exitcode=process.exitcode, ) return - if not isinstance(loaded, dict) or _WORKER_EXCEPTION_KEY in loaded: + if not isinstance(loaded, dict): finish_failure( index, kind="invalid_result", @@ -320,18 +429,6 @@ def reap(process: Any) -> None: exitcode=process.exitcode, ) return - - 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 - note_done(is_failure=result.get("compile_time") is None) finally: process.close() @@ -339,7 +436,7 @@ def reap(process: Any) -> None: launch() while running: if kernel_timeout > 0: - nearest_deadline = min(deadline for _, deadline in running.values() if deadline is not None) + 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 @@ -357,19 +454,12 @@ def reap(process: Any) -> None: if kernel_timeout > 0: now = time.monotonic() for process in list(running): - index, deadline = running[process] + _, 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() - running.pop(process) - exitcode = process.exitcode - process.close() - retry_or_drop( - index, - kind="timeout", - reason=f"exceeded per-job timeout ({kernel_timeout:g}s); killed", - exitcode=exitcode, - ) + reap(process, timeout_reason=timeout_reason) launch() finally: @@ -395,7 +485,9 @@ def reap(process: Any) -> None: retries_used, retry_label, ) - return results + if any(result is None for result in results): + raise RuntimeError("internal AOT scheduler error: unfinished jobs remain") + return cast(list[dict[str, Any]], results) def run_parallel_jobs( @@ -410,13 +502,14 @@ def run_parallel_jobs( 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 ``traceback``. + ``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. - A possible OOM exit (``-SIGKILL`` or shell-style ``137``) halves the - concurrency limit before retrying and is not retried once that limit - reaches one. + An OOM-like ``-SIGKILL`` halves the concurrency limit once per launch wave; + healthy completions then restore it additively. OOM exits are not retried + once the limit reaches one. """ if not jobs: return [] @@ -430,7 +523,7 @@ def run_parallel_jobs( result_dir = tempfile.mkdtemp(prefix="flydsl_aot_results_") try: - raw_results = _run_file_pool( + results = _run_file_pool( worker, jobs, max_workers=max_workers, @@ -440,19 +533,6 @@ def run_parallel_jobs( ) finally: shutil.rmtree(result_dir, ignore_errors=True) - - results: list[dict[str, Any]] = [] - for job, result in zip(jobs, raw_results): - if result is None: - results.append( - _failure_result( - job, - kind="scheduler_error", - reason="scheduler returned no result for the job", - ) - ) - else: - results.append(result) return results diff --git a/tests/README.md b/tests/README.md index ddd2fc82f..a9a6af8e2 100644 --- a/tests/README.md +++ b/tests/README.md @@ -62,10 +62,10 @@ 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 fails with an actionable error instead of silently disabling the cap when the optional package is not installed or memory cannot be queried. A worker killed -by `SIGKILL` (`exitcode=-9`, or shell form `137`; possible OOM) halves the worker -limit before retrying and is not retried once the limit reaches one. Retries are -appended behind jobs that have not started yet instead of jumping to the front -of the queue. +by `SIGKILL` (`exitcode=-9`; possible OOM) halves the worker limit at most once +per launch wave before retrying; timeout retries use the same backoff. Healthy +completions restore the limit additively. No backoff is applied when retries are +disabled, and retries are appended behind jobs that have not started yet. 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 @@ -74,9 +74,10 @@ 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`, -`invalid_result`, and `scheduler_error`; `reason` and `attempts` are always -included when known, while process failures also carry `exitcode` and uncaught -Python exceptions carry `traceback`. +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 index f441e6aa5..9cb7fa0b3 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -54,6 +54,13 @@ def _tracked_worker(kernel_name, index, state_path, lock_path, delay): fcntl.flock(lock_file, fcntl.LOCK_UN) +def _init_tracking_state(tmp_path): + 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") + return state_path, lock_path + + 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: @@ -77,6 +84,50 @@ def _sigkill_once_worker(kernel_name, attempt_path=None, delay=0.0): return {"kernel_name": kernel_name, "compile_time": 0.01} +def _oom_then_track_worker( + kernel_name, + attempt_path, + crashes, + state_path, + lock_path, + delay, +): + 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.kill(os.getpid(), signal.SIGKILL) + return _tracked_worker( + kernel_name, + attempt, + state_path, + lock_path, + delay, + ) + + +def _timeout_then_track_worker( + kernel_name, + attempt_path, + state_path, + lock_path, + delay, + timeout_delay, +): + path = Path(attempt_path) + attempt = int(path.read_text(encoding="utf-8")) + 1 + path.write_text(str(attempt), encoding="utf-8") + if attempt == 1: + time.sleep(timeout_delay) + return _tracked_worker( + kernel_name, + attempt, + state_path, + lock_path, + delay, + ) + + def _deterministic_failure(kernel_name, attempt_path): path = Path(attempt_path) attempt = int(path.read_text(encoding="utf-8")) + 1 @@ -98,6 +149,10 @@ def _mixed_outcome_worker(kernel_name, outcome): os._exit(137) if outcome == "type-error": raise TypeError("synthetic worker type error") + if outcome == "system-exit": + raise SystemExit(3) + if outcome == "large-error": + raise RuntimeError("X" * (parallel._MAX_TRACEBACK_CHARS * 4)) if outcome == "compile-error": return { "kernel_name": kernel_name, @@ -110,6 +165,27 @@ def _mixed_outcome_worker(kernel_name, outcome): return {"kernel_name": kernel_name, "compile_time": 0.01} +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 _exit_without_result(worker, kwargs, out_path): + return None + + +def _non_dict_worker(kernel_name): + return ["not", "a", "dictionary"] + + @pytest.fixture(autouse=True) def _parallel_env(monkeypatch): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") @@ -159,6 +235,54 @@ def test_results_follow_input_order(): 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" @@ -268,6 +392,87 @@ def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, log_me assert "possible OOM); reduced worker limit 2->1; retry 1/1" in "\n".join(log_messages) +def test_simultaneous_ooms_back_off_once_per_launch_wave(monkeypatch, tmp_path, log_messages): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") + state_path, lock_path = _init_tracking_state(tmp_path) + attempt_paths = [tmp_path / f"oom-attempt-{index}.txt" for index in range(4)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + jobs = [ + { + "kernel_name": f"oom-{index}", + "attempt_path": str(attempt_path), + "crashes": 1, + "state_path": str(state_path), + "lock_path": str(lock_path), + "delay": 0.1, + } + for index, attempt_path in enumerate(attempt_paths) + ] + + results = run_parallel_jobs(_oom_then_track_worker, jobs) + output = "\n".join(log_messages) + + assert all(result["compile_time"] is not None for result in results) + assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 2 + assert output.count("reduced worker limit 4->2") == 1 + assert "reduced worker limit 2->1" not in output + + +def test_oom_without_retries_does_not_reduce_pending_concurrency(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "0") + state_path, lock_path = _init_tracking_state(tmp_path) + attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(8)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + jobs = [ + { + "kernel_name": f"job-{index}", + "attempt_path": str(attempt_path), + "crashes": int(index == 0), + "state_path": str(state_path), + "lock_path": str(lock_path), + "delay": 0.2, + } + for index, attempt_path in enumerate(attempt_paths) + ] + + results = run_parallel_jobs(_oom_then_track_worker, jobs) + + assert results[0]["failure"]["kind"] == "possible_oom" + assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 + + +def test_worker_limit_recovers_after_healthy_completions(monkeypatch, tmp_path, log_messages): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") + state_path, lock_path = _init_tracking_state(tmp_path) + attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(13)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + jobs = [ + { + "kernel_name": f"job-{index}", + "attempt_path": str(attempt_path), + "crashes": int(index == 0), + "state_path": str(state_path), + "lock_path": str(lock_path), + "delay": 0.08, + } + for index, attempt_path in enumerate(attempt_paths) + ] + + results = run_parallel_jobs(_oom_then_track_worker, jobs) + output = "\n".join(log_messages) + + assert all(result["compile_time"] is not None for result in results) + assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 + assert "worker limit recovered 2->3" in output + assert "worker limit recovered 3->4" in output + + def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") @@ -362,15 +567,19 @@ def test_failure_results_preserve_distinct_causes(): "attempts": 1, "exitcode": -signal.SIGTERM, } - assert failures[1]["kind"] == "possible_oom" - assert failures[1]["reason"].startswith("worker exited with code 137 (possible OOM)") - assert failures[1]["attempts"] == 1 - assert failures[1]["exitcode"] == 137 + 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 "TypeError: synthetic worker type error" in failures[2]["traceback"] + 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", @@ -379,6 +588,91 @@ def test_failure_results_preserve_distinct_causes(): 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_system_exit_is_preserved_and_not_retried(monkeypatch): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") + + result = run_parallel_jobs( + _mixed_outcome_worker, + [{"kernel_name": "system-exit", "outcome": "system-exit"}], + )[0] + + assert result["failure"]["kind"] == "worker_exception" + assert result["failure"]["reason"] == "SystemExit: 3" + assert result["failure"]["attempts"] == 1 + assert result["failure"]["exitcode"] == 3 + assert result["failure"]["traceback"].endswith("SystemExit: 3\n") + + +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_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_failed_jobs_do_not_stop_remaining_jobs(): jobs = [ {"kernel_name": "first", "outcome": "success"}, @@ -445,6 +739,35 @@ def test_timed_out_worker_is_killed(monkeypatch): assert results[0]["failure"]["reason"] == "exceeded per-job timeout (0.05s); killed" +def test_timeout_wave_backs_off_once_before_retry(monkeypatch, tmp_path, log_messages): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") + monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.2") + state_path, lock_path = _init_tracking_state(tmp_path) + attempt_paths = [tmp_path / f"timeout-attempt-{index}.txt" for index in range(4)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + jobs = [ + { + "kernel_name": f"timeout-{index}", + "attempt_path": str(attempt_path), + "state_path": str(state_path), + "lock_path": str(lock_path), + "delay": 0.1, + "timeout_delay": 1.0, + } + for index, attempt_path in enumerate(attempt_paths) + ] + + results = run_parallel_jobs(_timeout_then_track_worker, jobs) + output = "\n".join(log_messages) + + assert all(result["compile_time"] is not None for result in results) + assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 2 + assert output.count("reduced worker limit 4->2") == 1 + assert "reduced worker limit 2->1" not in output + + def test_temporary_result_directory_is_removed(monkeypatch, tmp_path): result_dir = tmp_path / "results" From 85865086fd0493063206107f2e83029ae4ce39b2 Mon Sep 17 00:00:00 2001 From: zhimding Date: Fri, 21 Aug 2026 03:04:29 +0000 Subject: [PATCH 07/11] [AOT] Fix epoch backoff routing and recovery Keep same-wave OOM siblings retryable, recover from all healthy completions, leave timeouts concurrency-neutral, and allow AOT to continue without psutil. Co-authored-by: Cursor --- python/flydsl/utils/parallel.py | 37 +++++++------- tests/README.md | 14 +++--- tests/unit/test_parallel_jobs.py | 82 +++++++++++++++++++++++++++++--- 3 files changed, 103 insertions(+), 30 deletions(-) diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index 0a054634a..d4e22787f 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -116,19 +116,20 @@ def _memory_worker_cap(default_workers: int) -> int: try: import psutil - except ImportError as error: - raise RuntimeError( - "psutil is required for automatic AOT worker memory limiting; " - "install psutil or set FLYDSL_AOT_WORKERS explicitly" - ) from error + except ImportError: + log().warning( + "psutil is not installed; AOT memory limiting is disabled and the CPU-based worker limit will be used" + ) + return default_workers try: available_gb = psutil.virtual_memory().available / (1024**3) except Exception as error: - raise RuntimeError( - "failed to query available memory for the AOT worker limit; " - "set FLYDSL_AOT_WORKERS explicitly to bypass automatic detection" - ) from error + log().warning( + "failed to query available memory for AOT worker limiting (%s); the CPU-based worker limit will be used", + error, + ) + return default_workers return min(default_workers, max(1, int(available_gb / per_worker_gb))) @@ -227,9 +228,9 @@ def note_done(*, is_failure: bool = False) -> None: failed_jobs, ) - def note_success(worker_epoch: int) -> None: + def note_success() -> None: nonlocal max_workers, successful_since_backoff - if max_workers >= initial_max_workers or worker_epoch != launch_epoch: + if max_workers >= initial_max_workers: return successful_since_backoff += 1 if successful_since_backoff < max_workers: @@ -339,7 +340,7 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: is_failure = result.get("compile_time") is None note_done(is_failure=is_failure) if not is_failure: - note_success(worker_epoch) + note_success() return # A structured Python exception is deterministic for the same job. @@ -364,11 +365,6 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: return if timeout_reason is not None: - if attempts[index] < max_retries and max_workers > 1: - timeout_reason = backoff_worker_limit( - worker_epoch, - timeout_reason, - ) retry_or_drop( index, kind="timeout", @@ -387,6 +383,13 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: reason=reason, exitcode=process.exitcode, ) + elif worker_epoch != launch_epoch: + retry_or_drop( + index, + kind="possible_oom", + reason=(f"{reason}; worker limit already reduced to {max_workers} for this failure wave"), + exitcode=process.exitcode, + ) elif max_workers <= 1: finish_failure( index, diff --git a/tests/README.md b/tests/README.md index a9a6af8e2..eb20b22d3 100644 --- a/tests/README.md +++ b/tests/README.md @@ -60,12 +60,14 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. `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 -fails with an actionable error instead of silently disabling the cap when the -optional package is not installed or memory cannot be queried. A worker killed -by `SIGKILL` (`exitcode=-9`; possible OOM) halves the worker limit at most once -per launch wave before retrying; timeout retries use the same backoff. Healthy -completions restore the limit additively. No backoff is applied when retries are -disabled, and retries are appended behind jobs that have not started yet. +logs a warning before falling back to the CPU-based worker limit when the +optional package is unavailable or memory cannot be queried. A worker killed by +`SIGKILL` (`exitcode=-9`; possible OOM) halves the worker limit at most once per +launch wave before retrying. Every healthy completion after a backoff helps +restore the limit additively, including work launched before that backoff. +Timeouts do not change the global worker limit. No backoff is applied when +retries are disabled, and retries are appended behind jobs that have not +started yet. 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 diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py index 9cb7fa0b3..ed6cc731d 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -392,6 +392,31 @@ def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, log_me assert "possible OOM); reduced worker limit 2->1; retry 1/1" in "\n".join(log_messages) +def test_same_wave_oom_siblings_retry_at_reduced_limit(monkeypatch, tmp_path, log_messages): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") + attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(2)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + + results = run_parallel_jobs( + _sigkill_once_worker, + [ + { + "kernel_name": f"oom-{index}", + "attempt_path": str(attempt_path), + } + for index, attempt_path in enumerate(attempt_paths) + ], + ) + output = "\n".join(log_messages) + + assert all(result["compile_time"] is not None for result in results) + assert [path.read_text(encoding="utf-8") for path in attempt_paths] == ["2", "2"] + assert output.count("reduced worker limit 2->1") == 1 + assert "worker limit already reduced to 1 for this failure wave; retry 1/3" in output + + def test_simultaneous_ooms_back_off_once_per_launch_wave(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") @@ -473,6 +498,33 @@ def test_worker_limit_recovers_after_healthy_completions(monkeypatch, tmp_path, assert "worker limit recovered 3->4" in output +def test_pre_backoff_inflight_successes_contribute_to_recovery(monkeypatch, tmp_path, log_messages): + monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") + monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") + state_path, lock_path = _init_tracking_state(tmp_path) + attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(4)] + for attempt_path in attempt_paths: + attempt_path.write_text("0", encoding="utf-8") + jobs = [ + { + "kernel_name": f"job-{index}", + "attempt_path": str(attempt_path), + "crashes": int(index == 0), + "state_path": str(state_path), + "lock_path": str(lock_path), + "delay": 0.15, + } + for index, attempt_path in enumerate(attempt_paths) + ] + + results = run_parallel_jobs(_oom_then_track_worker, jobs) + output = "\n".join(log_messages) + + assert all(result["compile_time"] is not None for result in results) + assert attempt_paths[0].read_text(encoding="utf-8") == "2" + assert "worker limit recovered 2->3" in output + + def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") @@ -739,7 +791,7 @@ def test_timed_out_worker_is_killed(monkeypatch): assert results[0]["failure"]["reason"] == "exceeded per-job timeout (0.05s); killed" -def test_timeout_wave_backs_off_once_before_retry(monkeypatch, tmp_path, log_messages): +def test_timeout_wave_retries_without_reducing_concurrency(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.2") @@ -763,9 +815,8 @@ def test_timeout_wave_backs_off_once_before_retry(monkeypatch, tmp_path, log_mes output = "\n".join(log_messages) assert all(result["compile_time"] is not None for result in results) - assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 2 - assert output.count("reduced worker limit 4->2") == 1 - assert "reduced worker limit 2->1" not in output + assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 + assert "reduced worker limit" not in output def test_temporary_result_directory_is_removed(monkeypatch, tmp_path): @@ -797,13 +848,30 @@ def test_mem_per_worker_env_caps_automatic_workers(monkeypatch): assert parallel._get_max_workers(num_jobs=100) == 4 -def test_missing_psutil_does_not_silently_disable_memory_cap(monkeypatch): +def test_missing_psutil_warns_and_uses_cpu_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) + + assert parallel._get_max_workers(num_jobs=100) == 16 + assert "psutil is not installed; AOT memory limiting is disabled" in "\n".join(log_messages) + + +def test_memory_query_failure_warns_and_uses_cpu_limit(monkeypatch, log_messages): + def fail_memory_query(): + raise OSError("memory query unavailable") - with pytest.raises(RuntimeError, match="psutil is required"): - parallel._get_max_workers(num_jobs=100) + fake_psutil = SimpleNamespace(virtual_memory=fail_memory_query) + monkeypatch.delenv("FLYDSL_AOT_WORKERS") + monkeypatch.setenv("FLYDSL_AOT_MEM_PER_WORKER_GB", "2") + monkeypatch.setitem(sys.modules, "psutil", fake_psutil) + monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 16) + + assert parallel._get_max_workers(num_jobs=100) == 16 + assert ("failed to query available memory for AOT worker limiting (memory query unavailable)") in "\n".join( + log_messages + ) @pytest.mark.parametrize( From e9e0ca054660a73c5364f461d821c1ff7a46e2cc Mon Sep 17 00:00:00 2001 From: zhimding Date: Fri, 21 Aug 2026 03:29:55 +0000 Subject: [PATCH 08/11] [AOT] Remove adaptive concurrency backoff Keep retries concurrency-neutral, use explicit success accounting, and fall back to a conservative four-worker limit when memory information is unavailable. Co-authored-by: Cursor --- python/flydsl/utils/parallel.py | 129 +++++++++++-------------------- tests/README.md | 13 ++-- tests/unit/test_parallel_jobs.py | 99 ++++++------------------ 3 files changed, 74 insertions(+), 167 deletions(-) diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index d4e22787f..1475fad70 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -12,6 +12,7 @@ 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 @@ -23,6 +24,7 @@ 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 @@ -109,6 +111,17 @@ 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) + warnings.warn(message, RuntimeWarning, stacklevel=2) + return fallback_workers + + def _memory_worker_cap(default_workers: int) -> int: per_worker_gb = aot.mem_per_worker_gb if per_worker_gb <= 0: @@ -117,19 +130,18 @@ def _memory_worker_cap(default_workers: int) -> int: try: import psutil except ImportError: - log().warning( - "psutil is not installed; AOT memory limiting is disabled and the CPU-based worker limit will be used" + return _memory_worker_fallback( + default_workers, + "psutil is not installed; automatic AOT memory limiting is unavailable", ) - return default_workers try: available_gb = psutil.virtual_memory().available / (1024**3) except Exception as error: - log().warning( - "failed to query available memory for AOT worker limiting (%s); the CPU-based worker limit will be used", - error, + return _memory_worker_fallback( + default_workers, + f"failed to query available memory for AOT worker limiting ({error})", ) - return default_workers return min(default_workers, max(1, int(available_gb / per_worker_gb))) @@ -190,14 +202,12 @@ def _run_file_pool( attempts = [0] * num_jobs retries_used = 0 completed = 0 + succeeded_jobs = 0 failed_jobs = 0 progress_stride = max(1, num_jobs // 20) - initial_max_workers = max_workers - launch_epoch = 0 - successful_since_backoff = 0 queue = deque(range(num_jobs)) - running: dict[Any, tuple[int, float | None, int]] = {} + running: dict[Any, tuple[int, float | None]] = {} def launch() -> None: while queue and len(running) < max_workers: @@ -214,12 +224,15 @@ def launch() -> None: ) process.start() deadline = time.monotonic() + kernel_timeout if kernel_timeout > 0 else None - running[process] = (index, deadline, launch_epoch) + running[process] = (index, deadline) def note_done(*, is_failure: bool = False) -> None: - nonlocal completed, failed_jobs + nonlocal completed, failed_jobs, succeeded_jobs completed += 1 - failed_jobs += int(is_failure) + 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)", @@ -228,32 +241,6 @@ def note_done(*, is_failure: bool = False) -> None: failed_jobs, ) - def note_success() -> None: - nonlocal max_workers, successful_since_backoff - if max_workers >= initial_max_workers: - return - successful_since_backoff += 1 - if successful_since_backoff < max_workers: - return - previous_max_workers = max_workers - max_workers += 1 - successful_since_backoff = 0 - log().info( - "AOT worker limit recovered %d->%d after healthy completions", - previous_max_workers, - max_workers, - ) - - def backoff_worker_limit(worker_epoch: int, reason: str) -> str: - nonlocal launch_epoch, max_workers, successful_since_backoff - if worker_epoch != launch_epoch: - return f"{reason}; worker limit already reduced to {max_workers} for this failure wave" - previous_max_workers = max_workers - max_workers = max(1, max_workers // 2) - launch_epoch += 1 - successful_since_backoff = 0 - return f"{reason}; reduced worker limit {previous_max_workers}->{max_workers}" - def finish_failure( index: int, *, @@ -307,8 +294,7 @@ def retry_or_drop( ) def reap(process: Any, *, timeout_reason: str | None = None) -> None: - nonlocal launch_epoch, max_workers, successful_since_backoff - index, _, worker_epoch = running.pop(process) + index, _ = running.pop(process) out_path = os.path.join(result_dir, f"k{index}.json") try: loaded: Any = None @@ -339,8 +325,6 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: results[index] = result is_failure = result.get("compile_time") is None note_done(is_failure=is_failure) - if not is_failure: - note_success() return # A structured Python exception is deterministic for the same job. @@ -374,37 +358,12 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: return if process.exitcode == _OOM_EXITCODE: - reason = "worker killed by SIGKILL (possible OOM)" - can_retry = attempts[index] < max_retries - if not can_retry: - finish_failure( - index, - kind="possible_oom", - reason=reason, - exitcode=process.exitcode, - ) - elif worker_epoch != launch_epoch: - retry_or_drop( - index, - kind="possible_oom", - reason=(f"{reason}; worker limit already reduced to {max_workers} for this failure wave"), - exitcode=process.exitcode, - ) - elif max_workers <= 1: - finish_failure( - index, - kind="possible_oom", - reason=f"{reason} at the minimum worker limit", - exitcode=process.exitcode, - ) - else: - reason = backoff_worker_limit(worker_epoch, reason) - retry_or_drop( - index, - kind="possible_oom", - reason=reason, - exitcode=process.exitcode, - ) + retry_or_drop( + index, + kind="possible_oom", + reason="worker killed by SIGKILL (possible OOM)", + exitcode=process.exitcode, + ) return if process.exitcode != 0: @@ -439,7 +398,7 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: launch() while running: if kernel_timeout > 0: - nearest_deadline = min(deadline for _, deadline, _ in running.values() if deadline is not None) + 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 @@ -457,7 +416,7 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: if kernel_timeout > 0: now = time.monotonic() for process in list(running): - _, deadline, _ = running[process] + _, 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() @@ -480,16 +439,21 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: pass running.clear() + 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", - num_jobs - failed_jobs, + succeeded_jobs, failed_jobs, retries_used, retry_label, ) - if any(result is None for result in results): - raise RuntimeError("internal AOT scheduler error: unfinished jobs remain") return cast(list[dict[str, Any]], results) @@ -510,9 +474,8 @@ def run_parallel_jobs( This executor uses the Linux ``fork`` multiprocessing context because it is intended for compile-only workers that inherit the initialized compiler. - An OOM-like ``-SIGKILL`` halves the concurrency limit once per launch wave; - healthy completions then restore it additively. OOM exits are not retried - once the limit reaches one. + 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 [] diff --git a/tests/README.md b/tests/README.md index eb20b22d3..cf23eae42 100644 --- a/tests/README.md +++ b/tests/README.md @@ -60,14 +60,11 @@ Use the same names as [`python/flydsl/utils/env.py`](../python/flydsl/utils/env. `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 -logs a warning before falling back to the CPU-based worker limit when the -optional package is unavailable or memory cannot be queried. A worker killed by -`SIGKILL` (`exitcode=-9`; possible OOM) halves the worker limit at most once per -launch wave before retrying. Every healthy completion after a backoff helps -restore the limit additively, including work launched before that backoff. -Timeouts do not change the global worker limit. No backoff is applied when -retries are disabled, and retries are appended behind jobs that have not -started yet. +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. +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. 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 diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py index ed6cc731d..9b957bab0 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -367,7 +367,7 @@ def test_retries_wait_behind_pending_jobs(monkeypatch, tmp_path): ] -def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, log_messages): +def test_sigkill_retries_without_changing_worker_limit(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") attempt_path = tmp_path / "attempt.txt" @@ -389,10 +389,12 @@ def test_sigkill_reduces_worker_limit_before_retry(monkeypatch, tmp_path, log_me assert attempt_path.read_text(encoding="utf-8") == "2" assert all(result["compile_time"] is not None for result in results) - assert "possible OOM); reduced worker limit 2->1; retry 1/1" in "\n".join(log_messages) + output = "\n".join(log_messages) + assert "possible OOM); retry 1/1" in output + assert "worker limit" not in output -def test_same_wave_oom_siblings_retry_at_reduced_limit(monkeypatch, tmp_path, log_messages): +def test_same_wave_oom_siblings_all_retry(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(2)] @@ -413,11 +415,11 @@ def test_same_wave_oom_siblings_retry_at_reduced_limit(monkeypatch, tmp_path, lo assert all(result["compile_time"] is not None for result in results) assert [path.read_text(encoding="utf-8") for path in attempt_paths] == ["2", "2"] - assert output.count("reduced worker limit 2->1") == 1 - assert "worker limit already reduced to 1 for this failure wave; retry 1/3" in output + assert output.count("possible OOM); retry 1/3") == 2 + assert "worker limit" not in output -def test_simultaneous_ooms_back_off_once_per_launch_wave(monkeypatch, tmp_path, log_messages): +def test_simultaneous_oom_retries_keep_configured_concurrency(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") state_path, lock_path = _init_tracking_state(tmp_path) @@ -440,9 +442,8 @@ def test_simultaneous_ooms_back_off_once_per_launch_wave(monkeypatch, tmp_path, output = "\n".join(log_messages) assert all(result["compile_time"] is not None for result in results) - assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 2 - assert output.count("reduced worker limit 4->2") == 1 - assert "reduced worker limit 2->1" not in output + assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 + assert "worker limit" not in output def test_oom_without_retries_does_not_reduce_pending_concurrency(monkeypatch, tmp_path): @@ -470,62 +471,7 @@ def test_oom_without_retries_does_not_reduce_pending_concurrency(monkeypatch, tm assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 -def test_worker_limit_recovers_after_healthy_completions(monkeypatch, tmp_path, log_messages): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") - state_path, lock_path = _init_tracking_state(tmp_path) - attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(13)] - for attempt_path in attempt_paths: - attempt_path.write_text("0", encoding="utf-8") - jobs = [ - { - "kernel_name": f"job-{index}", - "attempt_path": str(attempt_path), - "crashes": int(index == 0), - "state_path": str(state_path), - "lock_path": str(lock_path), - "delay": 0.08, - } - for index, attempt_path in enumerate(attempt_paths) - ] - - results = run_parallel_jobs(_oom_then_track_worker, jobs) - output = "\n".join(log_messages) - - assert all(result["compile_time"] is not None for result in results) - assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 - assert "worker limit recovered 2->3" in output - assert "worker limit recovered 3->4" in output - - -def test_pre_backoff_inflight_successes_contribute_to_recovery(monkeypatch, tmp_path, log_messages): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") - state_path, lock_path = _init_tracking_state(tmp_path) - attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(4)] - for attempt_path in attempt_paths: - attempt_path.write_text("0", encoding="utf-8") - jobs = [ - { - "kernel_name": f"job-{index}", - "attempt_path": str(attempt_path), - "crashes": int(index == 0), - "state_path": str(state_path), - "lock_path": str(lock_path), - "delay": 0.15, - } - for index, attempt_path in enumerate(attempt_paths) - ] - - results = run_parallel_jobs(_oom_then_track_worker, jobs) - output = "\n".join(log_messages) - - assert all(result["compile_time"] is not None for result in results) - assert attempt_paths[0].read_text(encoding="utf-8") == "2" - assert "worker limit recovered 2->3" in output - - -def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, log_messages): +def test_sigkill_at_one_worker_still_uses_retry_budget(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") attempt_path = tmp_path / "attempt.txt" @@ -541,12 +487,11 @@ def test_sigkill_at_minimum_worker_limit_is_not_retried(monkeypatch, tmp_path, l ], ) - assert attempt_path.read_text(encoding="utf-8") == "1" - assert results[0]["compile_time"] is None - assert results[0]["failure"]["kind"] == "possible_oom" - assert results[0]["failure"]["exitcode"] == -signal.SIGKILL - assert results[0]["failure"]["attempts"] == 1 - assert "at the minimum worker limit; not retrying" in "\n".join(log_messages) + assert attempt_path.read_text(encoding="utf-8") == "2" + assert results[0]["compile_time"] is not None + output = "\n".join(log_messages) + assert "possible OOM); retry 1/2" in output + assert "worker limit" not in output def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): @@ -848,17 +793,18 @@ def test_mem_per_worker_env_caps_automatic_workers(monkeypatch): assert parallel._get_max_workers(num_jobs=100) == 4 -def test_missing_psutil_warns_and_uses_cpu_limit(monkeypatch, log_messages): +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) - assert parallel._get_max_workers(num_jobs=100) == 16 - assert "psutil is not installed; AOT memory limiting is disabled" in "\n".join(log_messages) + 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) -def test_memory_query_failure_warns_and_uses_cpu_limit(monkeypatch, log_messages): +def test_memory_query_failure_warns_and_uses_conservative_limit(monkeypatch, log_messages): def fail_memory_query(): raise OSError("memory query unavailable") @@ -868,7 +814,8 @@ def fail_memory_query(): monkeypatch.setitem(sys.modules, "psutil", fake_psutil) monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 16) - assert parallel._get_max_workers(num_jobs=100) == 16 + with pytest.warns(RuntimeWarning, match="limiting AOT concurrency to 4 workers"): + assert parallel._get_max_workers(num_jobs=100) == 4 assert ("failed to query available memory for AOT worker limiting (memory query unavailable)") in "\n".join( log_messages ) From dd9358fe429c2d9322af84bc4293a37027dcc211 Mon Sep 17 00:00:00 2001 From: zhimding Date: Fri, 21 Aug 2026 06:22:38 +0000 Subject: [PATCH 09/11] [AOT] Close scheduler edge-case coverage gaps Localize malformed worker payloads, validate summaries before logging, clarify fork constraints, and add exact boundary and retry-exhaustion tests. Co-authored-by: Cursor --- python/flydsl/utils/parallel.py | 93 ++++++++++++++------- tests/README.md | 5 ++ tests/unit/test_parallel_jobs.py | 136 ++++++++++++++++++++++++++++--- 3 files changed, 192 insertions(+), 42 deletions(-) diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index 1475fad70..4ff604724 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -118,7 +118,9 @@ def _memory_worker_fallback(default_workers: int, reason: str) -> int: ) message = f"{reason}; limiting AOT concurrency to {fallback_workers} worker{'s' if fallback_workers != 1 else ''}" log().warning(message) - warnings.warn(message, RuntimeWarning, stacklevel=2) + # Public path: user -> run_parallel_jobs -> _get_max_workers -> + # _memory_worker_cap -> this helper. + warnings.warn(message, RuntimeWarning, stacklevel=5) return fallback_workers @@ -186,6 +188,33 @@ def _failure_result( } +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]], @@ -308,11 +337,12 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: else: load_error = "worker produced no result file" - worker_exception = loaded.get(_WORKER_EXCEPTION_KEY) if isinstance(loaded, dict) else None + 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 worker_exception is None: + if isinstance(loaded, dict) and not has_worker_exception_marker: result = loaded if result.get("compile_time") is None: failure = result.get("failure") @@ -329,7 +359,17 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: # A structured Python exception is deterministic for the same job. # Preserve it even if teardown later changes the process exit code. - if isinstance(worker_exception, dict): + 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", @@ -357,20 +397,16 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: ) return - if process.exitcode == _OOM_EXITCODE: - retry_or_drop( - index, - kind="possible_oom", - reason="worker killed by SIGKILL (possible OOM)", - exitcode=process.exitcode, - ) - return - if process.exitcode != 0: + possible_oom = process.exitcode == _OOM_EXITCODE retry_or_drop( index, - kind="worker_crash", - reason=f"worker crashed (exitcode={process.exitcode})", + 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 @@ -439,22 +475,14 @@ def reap(process: Any, *, timeout_reason: str | None = None) -> None: pass running.clear() - 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 _finalize_pool_results( + results, + num_jobs=num_jobs, + completed=completed, + succeeded_jobs=succeeded_jobs, + failed_jobs=failed_jobs, + retries_used=retries_used, ) - return cast(list[dict[str, Any]], results) def run_parallel_jobs( @@ -474,6 +502,11 @@ def run_parallel_jobs( 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. """ diff --git a/tests/README.md b/tests/README.md index cf23eae42..d88dba2e4 100644 --- a/tests/README.md +++ b/tests/README.md @@ -62,9 +62,14 @@ 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 diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py index 9b957bab0..89ff21a99 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -7,6 +7,7 @@ import signal import sys import time +import warnings from pathlib import Path from types import SimpleNamespace @@ -178,6 +179,13 @@ def _exception_marker_then_exit_zero(worker, kwargs, out_path): ) +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 @@ -186,6 +194,13 @@ 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") @@ -211,7 +226,7 @@ def test_empty_jobs_do_not_invoke_worker(monkeypatch): assert run_parallel_jobs(_success_worker, []) == [] -@pytest.mark.parametrize("workers", [None, "", "0", "-2"]) +@pytest.mark.parametrize("workers", [None, "", " ", "0", "-2"]) def test_non_positive_or_empty_workers_use_automatic_limit(monkeypatch, workers): if workers is None: monkeypatch.delenv("FLYDSL_AOT_WORKERS", raising=False) @@ -391,7 +406,6 @@ def test_sigkill_retries_without_changing_worker_limit(monkeypatch, tmp_path, lo assert all(result["compile_time"] is not None for result in results) output = "\n".join(log_messages) assert "possible OOM); retry 1/1" in output - assert "worker limit" not in output def test_same_wave_oom_siblings_all_retry(monkeypatch, tmp_path, log_messages): @@ -416,10 +430,9 @@ def test_same_wave_oom_siblings_all_retry(monkeypatch, tmp_path, log_messages): assert all(result["compile_time"] is not None for result in results) assert [path.read_text(encoding="utf-8") for path in attempt_paths] == ["2", "2"] assert output.count("possible OOM); retry 1/3") == 2 - assert "worker limit" not in output -def test_simultaneous_oom_retries_keep_configured_concurrency(monkeypatch, tmp_path, log_messages): +def test_simultaneous_oom_retries_keep_configured_concurrency(monkeypatch, tmp_path): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") state_path, lock_path = _init_tracking_state(tmp_path) @@ -439,11 +452,8 @@ def test_simultaneous_oom_retries_keep_configured_concurrency(monkeypatch, tmp_p ] results = run_parallel_jobs(_oom_then_track_worker, jobs) - output = "\n".join(log_messages) - assert all(result["compile_time"] is not None for result in results) assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 - assert "worker limit" not in output def test_oom_without_retries_does_not_reduce_pending_concurrency(monkeypatch, tmp_path): @@ -471,6 +481,30 @@ def test_oom_without_retries_does_not_reduce_pending_concurrency(monkeypatch, tm assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 +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_sigkill_at_one_worker_still_uses_retry_budget(monkeypatch, tmp_path, log_messages): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") @@ -491,7 +525,6 @@ def test_sigkill_at_one_worker_still_uses_retry_budget(monkeypatch, tmp_path, lo assert results[0]["compile_time"] is not None output = "\n".join(log_messages) assert "possible OOM); retry 1/2" in output - assert "worker limit" not in output def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): @@ -547,6 +580,20 @@ def test_final_logs_report_permanent_failures(monkeypatch, log_messages): assert "AOT: 2 succeeded, 3 failed; 6 retries after abnormal worker exits" in output +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"}, @@ -616,6 +663,30 @@ def test_system_exit_is_preserved_and_not_retried(monkeypatch): assert result["failure"]["traceback"].endswith("SystemExit: 3\n") +@pytest.mark.parametrize( + ("code", "expected"), + [ + (256, 0), + (512, 0), + (-1, 255), + ], +) +def test_system_exit_code_is_normalized_to_process_status(code, expected): + assert parallel._exception_exitcode(SystemExit(code)) == expected + + +def test_traceback_truncation_exact_boundary(): + at_limit = "A" * parallel._MAX_TRACEBACK_CHARS + over_limit = "B" * (parallel._MAX_TRACEBACK_CHARS + 1) + + assert parallel._truncate_traceback(at_limit) == at_limit + truncated = parallel._truncate_traceback(over_limit) + assert len(truncated) == parallel._MAX_TRACEBACK_CHARS + assert "... traceback truncated ..." in truncated + assert truncated.startswith("B" * 32) + assert truncated.endswith("B" * 32) + + def test_large_exception_diagnostics_are_bounded(): result = run_parallel_jobs( _mixed_outcome_worker, @@ -652,6 +723,26 @@ def test_exception_marker_is_preserved_when_process_exits_zero(monkeypatch): } +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, @@ -736,7 +827,7 @@ def test_timed_out_worker_is_killed(monkeypatch): assert results[0]["failure"]["reason"] == "exceeded per-job timeout (0.05s); killed" -def test_timeout_wave_retries_without_reducing_concurrency(monkeypatch, tmp_path, log_messages): +def test_timeout_wave_retries_without_reducing_concurrency(monkeypatch, tmp_path): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.2") @@ -757,11 +848,8 @@ def test_timeout_wave_retries_without_reducing_concurrency(monkeypatch, tmp_path ] results = run_parallel_jobs(_timeout_then_track_worker, jobs) - output = "\n".join(log_messages) - assert all(result["compile_time"] is not None for result in results) assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 - assert "reduced worker limit" not in output def test_temporary_result_directory_is_removed(monkeypatch, tmp_path): @@ -804,6 +892,30 @@ def test_missing_psutil_warns_and_uses_conservative_limit(monkeypatch, log_messa assert "psutil is not installed; automatic AOT memory limiting is unavailable" in "\n".join(log_messages) +def test_missing_psutil_warning_points_to_public_caller_and_deduplicates( + monkeypatch, +): + 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: 1) + + def invoke_scheduler(): + return run_parallel_jobs( + _success_worker, + [{"kernel_name": "one", "index": 0}], + ) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("default") + for _ in range(5): + invoke_scheduler() + + runtime_warnings = [warning for warning in caught if issubclass(warning.category, RuntimeWarning)] + assert len(runtime_warnings) == 1 + assert runtime_warnings[0].filename == __file__ + + def test_memory_query_failure_warns_and_uses_conservative_limit(monkeypatch, log_messages): def fail_memory_query(): raise OSError("memory query unavailable") From de01efb4d0af6c77f80b9cd0d844b71902bc3cb3 Mon Sep 17 00:00:00 2001 From: zhimding Date: Fri, 21 Aug 2026 06:48:48 +0000 Subject: [PATCH 10/11] [Test] Focus parallel scheduler coverage on critical paths Remove lower-priority and overlapping cases while retaining the 19 tests that protect correctness, failure isolation, and memory safety. Co-authored-by: Cursor --- tests/unit/test_parallel_jobs.py | 397 +------------------------------ 1 file changed, 1 insertion(+), 396 deletions(-) diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py index 89ff21a99..1a1ed75de 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -7,14 +7,12 @@ import signal import sys import time -import warnings 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] @@ -55,13 +53,6 @@ def _tracked_worker(kernel_name, index, state_path, lock_path, delay): fcntl.flock(lock_file, fcntl.LOCK_UN) -def _init_tracking_state(tmp_path): - 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") - return state_path, lock_path - - 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: @@ -74,61 +65,6 @@ def _crash_then_succeed(kernel_name, attempt_path, crashes, order_path=None): return {"kernel_name": kernel_name, "compile_time": 0.01} -def _sigkill_once_worker(kernel_name, attempt_path=None, delay=0.0): - if attempt_path is not None: - path = Path(attempt_path) - attempt = int(path.read_text(encoding="utf-8")) + 1 - path.write_text(str(attempt), encoding="utf-8") - if attempt == 1: - os.kill(os.getpid(), signal.SIGKILL) - time.sleep(delay) - return {"kernel_name": kernel_name, "compile_time": 0.01} - - -def _oom_then_track_worker( - kernel_name, - attempt_path, - crashes, - state_path, - lock_path, - delay, -): - 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.kill(os.getpid(), signal.SIGKILL) - return _tracked_worker( - kernel_name, - attempt, - state_path, - lock_path, - delay, - ) - - -def _timeout_then_track_worker( - kernel_name, - attempt_path, - state_path, - lock_path, - delay, - timeout_delay, -): - path = Path(attempt_path) - attempt = int(path.read_text(encoding="utf-8")) + 1 - path.write_text(str(attempt), encoding="utf-8") - if attempt == 1: - time.sleep(timeout_delay) - return _tracked_worker( - kernel_name, - attempt, - state_path, - lock_path, - delay, - ) - - def _deterministic_failure(kernel_name, attempt_path): path = Path(attempt_path) attempt = int(path.read_text(encoding="utf-8")) + 1 @@ -142,16 +78,12 @@ def _sleep_worker(kernel_name, delay): def _mixed_outcome_worker(kernel_name, outcome): - if outcome == "crash": - os._exit(17) 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 == "system-exit": - raise SystemExit(3) if outcome == "large-error": raise RuntimeError("X" * (parallel._MAX_TRACEBACK_CHARS * 4)) if outcome == "compile-error": @@ -163,7 +95,7 @@ def _mixed_outcome_worker(kernel_name, outcome): "reason": "synthetic codegen failure", }, } - return {"kernel_name": kernel_name, "compile_time": 0.01} + raise ValueError(f"unknown synthetic outcome: {outcome}") def _exception_marker_then_exit_zero(worker, kwargs, out_path): @@ -221,23 +153,6 @@ def record(message, *args): return messages -def test_empty_jobs_do_not_invoke_worker(monkeypatch): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "invalid") - assert run_parallel_jobs(_success_worker, []) == [] - - -@pytest.mark.parametrize("workers", [None, "", " ", "0", "-2"]) -def test_non_positive_or_empty_workers_use_automatic_limit(monkeypatch, workers): - if workers is None: - monkeypatch.delenv("FLYDSL_AOT_WORKERS", raising=False) - else: - monkeypatch.setenv("FLYDSL_AOT_WORKERS", workers) - monkeypatch.setattr(parallel, "_affinity_aware_cpu_count", lambda: 8) - monkeypatch.setattr(parallel, "_memory_worker_cap", lambda workers: workers) - - assert parallel._get_max_workers(num_jobs=100) == 8 - - def test_results_follow_input_order(): jobs = [ {"kernel_name": "slow", "index": 0, "delay": 0.1}, @@ -382,105 +297,6 @@ def test_retries_wait_behind_pending_jobs(monkeypatch, tmp_path): ] -def test_sigkill_retries_without_changing_worker_limit(monkeypatch, tmp_path, log_messages): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") - 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( - _sigkill_once_worker, - [ - { - "kernel_name": "oom-job", - "attempt_path": str(attempt_path), - }, - { - "kernel_name": "companion", - "delay": 0.1, - }, - ], - ) - - assert attempt_path.read_text(encoding="utf-8") == "2" - assert all(result["compile_time"] is not None for result in results) - output = "\n".join(log_messages) - assert "possible OOM); retry 1/1" in output - - -def test_same_wave_oom_siblings_all_retry(monkeypatch, tmp_path, log_messages): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "2") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") - attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(2)] - for attempt_path in attempt_paths: - attempt_path.write_text("0", encoding="utf-8") - - results = run_parallel_jobs( - _sigkill_once_worker, - [ - { - "kernel_name": f"oom-{index}", - "attempt_path": str(attempt_path), - } - for index, attempt_path in enumerate(attempt_paths) - ], - ) - output = "\n".join(log_messages) - - assert all(result["compile_time"] is not None for result in results) - assert [path.read_text(encoding="utf-8") for path in attempt_paths] == ["2", "2"] - assert output.count("possible OOM); retry 1/3") == 2 - - -def test_simultaneous_oom_retries_keep_configured_concurrency(monkeypatch, tmp_path): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") - state_path, lock_path = _init_tracking_state(tmp_path) - attempt_paths = [tmp_path / f"oom-attempt-{index}.txt" for index in range(4)] - for attempt_path in attempt_paths: - attempt_path.write_text("0", encoding="utf-8") - jobs = [ - { - "kernel_name": f"oom-{index}", - "attempt_path": str(attempt_path), - "crashes": 1, - "state_path": str(state_path), - "lock_path": str(lock_path), - "delay": 0.1, - } - for index, attempt_path in enumerate(attempt_paths) - ] - - results = run_parallel_jobs(_oom_then_track_worker, jobs) - assert all(result["compile_time"] is not None for result in results) - assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 - - -def test_oom_without_retries_does_not_reduce_pending_concurrency(monkeypatch, tmp_path): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "0") - state_path, lock_path = _init_tracking_state(tmp_path) - attempt_paths = [tmp_path / f"attempt-{index}.txt" for index in range(8)] - for attempt_path in attempt_paths: - attempt_path.write_text("0", encoding="utf-8") - jobs = [ - { - "kernel_name": f"job-{index}", - "attempt_path": str(attempt_path), - "crashes": int(index == 0), - "state_path": str(state_path), - "lock_path": str(lock_path), - "delay": 0.2, - } - for index, attempt_path in enumerate(attempt_paths) - ] - - results = run_parallel_jobs(_oom_then_track_worker, jobs) - - assert results[0]["failure"]["kind"] == "possible_oom" - assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 - - 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") @@ -505,28 +321,6 @@ def test_oom_retry_exhaustion_returns_possible_oom(monkeypatch, tmp_path, log_me assert "not retrying" in output -def test_sigkill_at_one_worker_still_uses_retry_budget(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") - - results = run_parallel_jobs( - _sigkill_once_worker, - [ - { - "kernel_name": "oom-job", - "attempt_path": str(attempt_path), - } - ], - ) - - assert attempt_path.read_text(encoding="utf-8") == "2" - assert results[0]["compile_time"] is not None - output = "\n".join(log_messages) - assert "possible OOM); retry 1/2" in output - - def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") @@ -554,32 +348,6 @@ def test_retry_exhaustion_returns_failure(monkeypatch, tmp_path): } -def test_final_logs_report_permanent_failures(monkeypatch, log_messages): - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "2") - jobs = [ - {"kernel_name": "success-0", "outcome": "success"}, - {"kernel_name": "failed-0", "outcome": "crash"}, - {"kernel_name": "failed-1", "outcome": "crash"}, - {"kernel_name": "success-1", "outcome": "success"}, - {"kernel_name": "failed-2", "outcome": "crash"}, - ] - - results = run_parallel_jobs(_mixed_outcome_worker, jobs) - output = "\n".join(log_messages) - - assert [result["compile_time"] for result in results] == [ - 0.01, - None, - None, - 0.01, - None, - ] - for kernel_name in ("failed-0", "failed-1", "failed-2"): - assert f"AOT job {kernel_name} worker crashed (exitcode=17); not retrying" in output - assert "... 5/5 jobs finished (3 failed)" in output - assert "AOT: 2 succeeded, 3 failed; 6 retries after abnormal worker exits" in output - - def test_incomplete_pool_results_raise_before_summary(log_messages): with pytest.raises(RuntimeError, match="0/1 jobs completed"): parallel._finalize_pool_results( @@ -647,46 +415,6 @@ def test_python_exception_is_not_retried_or_duplicated_to_stderr(monkeypatch, ca assert capfd.readouterr().err == "" -def test_system_exit_is_preserved_and_not_retried(monkeypatch): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "1") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") - - result = run_parallel_jobs( - _mixed_outcome_worker, - [{"kernel_name": "system-exit", "outcome": "system-exit"}], - )[0] - - assert result["failure"]["kind"] == "worker_exception" - assert result["failure"]["reason"] == "SystemExit: 3" - assert result["failure"]["attempts"] == 1 - assert result["failure"]["exitcode"] == 3 - assert result["failure"]["traceback"].endswith("SystemExit: 3\n") - - -@pytest.mark.parametrize( - ("code", "expected"), - [ - (256, 0), - (512, 0), - (-1, 255), - ], -) -def test_system_exit_code_is_normalized_to_process_status(code, expected): - assert parallel._exception_exitcode(SystemExit(code)) == expected - - -def test_traceback_truncation_exact_boundary(): - at_limit = "A" * parallel._MAX_TRACEBACK_CHARS - over_limit = "B" * (parallel._MAX_TRACEBACK_CHARS + 1) - - assert parallel._truncate_traceback(at_limit) == at_limit - truncated = parallel._truncate_traceback(over_limit) - assert len(truncated) == parallel._MAX_TRACEBACK_CHARS - assert "... traceback truncated ..." in truncated - assert truncated.startswith("B" * 32) - assert truncated.endswith("B" * 32) - - def test_large_exception_diagnostics_are_bounded(): result = run_parallel_jobs( _mixed_outcome_worker, @@ -761,30 +489,6 @@ def test_invalid_result_reasons_distinguish_missing_and_non_dict(monkeypatch): assert missing["failure"]["reason"] == "worker produced no result file" -def test_failed_jobs_do_not_stop_remaining_jobs(): - jobs = [ - {"kernel_name": "first", "outcome": "success"}, - {"kernel_name": "crashed", "outcome": "crash"}, - {"kernel_name": "compile-error", "outcome": "compile-error"}, - {"kernel_name": "last", "outcome": "success"}, - ] - - results = run_parallel_jobs(_mixed_outcome_worker, jobs) - - assert [result["kernel_name"] for result in results] == [ - "first", - "crashed", - "compile-error", - "last", - ] - assert [result["compile_time"] for result in results] == [ - 0.01, - None, - None, - 0.01, - ] - - def test_deterministic_failure_is_not_retried(monkeypatch, tmp_path): monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "3") attempt_path = tmp_path / "attempt.txt" @@ -827,49 +531,6 @@ def test_timed_out_worker_is_killed(monkeypatch): assert results[0]["failure"]["reason"] == "exceeded per-job timeout (0.05s); killed" -def test_timeout_wave_retries_without_reducing_concurrency(monkeypatch, tmp_path): - monkeypatch.setenv("FLYDSL_AOT_WORKERS", "4") - monkeypatch.setenv("FLYDSL_AOT_MAX_RETRIES", "1") - monkeypatch.setenv("FLYDSL_AOT_TIMEOUT", "0.2") - state_path, lock_path = _init_tracking_state(tmp_path) - attempt_paths = [tmp_path / f"timeout-attempt-{index}.txt" for index in range(4)] - for attempt_path in attempt_paths: - attempt_path.write_text("0", encoding="utf-8") - jobs = [ - { - "kernel_name": f"timeout-{index}", - "attempt_path": str(attempt_path), - "state_path": str(state_path), - "lock_path": str(lock_path), - "delay": 0.1, - "timeout_delay": 1.0, - } - for index, attempt_path in enumerate(attempt_paths) - ] - - results = run_parallel_jobs(_timeout_then_track_worker, jobs) - assert all(result["compile_time"] is not None for result in results) - assert json.loads(state_path.read_text(encoding="utf-8"))["peak"] == 4 - - -def test_temporary_result_directory_is_removed(monkeypatch, tmp_path): - result_dir = tmp_path / "results" - - def make_result_dir(prefix): - assert prefix == "flydsl_aot_results_" - result_dir.mkdir() - return str(result_dir) - - monkeypatch.setattr(parallel.tempfile, "mkdtemp", make_result_dir) - - run_parallel_jobs( - _success_worker, - [{"kernel_name": "one", "index": 0}], - ) - - assert not result_dir.exists() - - def test_mem_per_worker_env_caps_automatic_workers(monkeypatch): gib = 1024**3 fake_psutil = SimpleNamespace(virtual_memory=lambda: SimpleNamespace(available=8 * gib)) @@ -890,59 +551,3 @@ def test_missing_psutil_warns_and_uses_conservative_limit(monkeypatch, log_messa 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) - - -def test_missing_psutil_warning_points_to_public_caller_and_deduplicates( - monkeypatch, -): - 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: 1) - - def invoke_scheduler(): - return run_parallel_jobs( - _success_worker, - [{"kernel_name": "one", "index": 0}], - ) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("default") - for _ in range(5): - invoke_scheduler() - - runtime_warnings = [warning for warning in caught if issubclass(warning.category, RuntimeWarning)] - assert len(runtime_warnings) == 1 - assert runtime_warnings[0].filename == __file__ - - -def test_memory_query_failure_warns_and_uses_conservative_limit(monkeypatch, log_messages): - def fail_memory_query(): - raise OSError("memory query unavailable") - - fake_psutil = SimpleNamespace(virtual_memory=fail_memory_query) - monkeypatch.delenv("FLYDSL_AOT_WORKERS") - monkeypatch.setenv("FLYDSL_AOT_MEM_PER_WORKER_GB", "2") - monkeypatch.setitem(sys.modules, "psutil", fake_psutil) - 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 ("failed to query available memory for AOT worker limiting (memory query unavailable)") in "\n".join( - log_messages - ) - - -@pytest.mark.parametrize( - ("variable", "accessor"), - [ - ("FLYDSL_AOT_WORKERS", lambda: aot.workers), - ("FLYDSL_AOT_MEM_PER_WORKER_GB", lambda: aot.mem_per_worker_gb), - ("FLYDSL_AOT_TIMEOUT", lambda: aot.timeout), - ("FLYDSL_AOT_MAX_RETRIES", lambda: aot.max_retries), - ], -) -def test_invalid_environment_value_raises(monkeypatch, variable, accessor): - monkeypatch.setenv(variable, "invalid") - with pytest.raises(ValueError, match=variable): - accessor() From 09f20909339539b4051bdabe44607440d62ad93a Mon Sep 17 00:00:00 2001 From: zhimding Date: Fri, 21 Aug 2026 07:05:19 +0000 Subject: [PATCH 11/11] [AOT] Treat empty worker configuration as automatic Make the typed environment option return its default for empty numeric values so scheduler semantics match the documented contract directly. Co-authored-by: Cursor --- python/flydsl/utils/env.py | 9 +++++++++ python/flydsl/utils/parallel.py | 3 +-- tests/unit/test_parallel_jobs.py | 10 ++++++++++ 3 files changed, 20 insertions(+), 2 deletions(-) diff --git a/python/flydsl/utils/env.py b/python/flydsl/utils/env.py index 17d14c62f..13e47ac0a 100644 --- a/python/flydsl/utils/env.py +++ b/python/flydsl/utils/env.py @@ -89,6 +89,7 @@ def __init__( description: str = "", 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: @@ -104,8 +105,11 @@ def validator(v: NumberT) -> bool: 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) @@ -119,6 +123,7 @@ def __init__( description: str = "", min_value: Optional[int] = None, max_value: Optional[int] = None, + empty_is_default: bool = False, ): super().__init__( default, @@ -127,6 +132,7 @@ def __init__( description, min_value, max_value, + empty_is_default, ) @@ -140,6 +146,7 @@ def __init__( description: str = "", min_value: Optional[float] = None, max_value: Optional[float] = None, + empty_is_default: bool = False, ): super().__init__( default, @@ -148,6 +155,7 @@ def __init__( description, min_value, max_value, + empty_is_default, ) @@ -280,6 +288,7 @@ class AotEnvManager(EnvManager): "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, diff --git a/python/flydsl/utils/parallel.py b/python/flydsl/utils/parallel.py index 4ff604724..5d623cc6c 100644 --- a/python/flydsl/utils/parallel.py +++ b/python/flydsl/utils/parallel.py @@ -148,8 +148,7 @@ def _memory_worker_cap(default_workers: int) -> int: def _get_max_workers(num_jobs: int) -> int: - workers_raw = os.environ.get("FLYDSL_AOT_WORKERS", "").strip() - configured_workers = aot.workers if workers_raw else 0 + configured_workers = aot.workers if configured_workers > 0: max_workers = configured_workers else: diff --git a/tests/unit/test_parallel_jobs.py b/tests/unit/test_parallel_jobs.py index 1a1ed75de..badf83cb4 100644 --- a/tests/unit/test_parallel_jobs.py +++ b/tests/unit/test_parallel_jobs.py @@ -13,6 +13,7 @@ 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] @@ -153,6 +154,15 @@ def record(message, *args): 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},