From 1b611d76194a62488abd70b1dbe1d145fde6a61b Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 27 Aug 2026 02:14:47 +0000 Subject: [PATCH] fix: report which half of a JAX first call is stalled The first call to a jax.jit/vmap/grad wrapper contains two waits -- trace/lower/compile, then jax.block_until_ready execution -- and the instrumentation could not tell a stalled run which one it was in. Two defects, both of which made a stall report the wrong diagnosis: 1. The heartbeat hardcoded "still compiling". A stalled CI run therefore logged `JAX jit still compiling ... 1770s elapsed` while its captured faulthandler stack sat in jax.block_until_ready. That is positive evidence for the wrong cause, and five quarantine markers across the two test workspaces were written against it, all calling this an "intermittent XLA compile stall". It is not one: compilation completes, in ~16s, and execution never returns. 2. The compile/execute split line was emitted only after BOTH halves finished. A stalled run never finishes the second, so it reported neither number and the split characterised only the healthy case -- twenty stalled runs could not say that compilation had in fact completed. Now: the compile half is logged the instant func() returns, the heartbeat moves to naming the materialize half and repeats the compile time on every beat, and the materialize half is logged on its own when it lands. A run SIGKILLed at any point after tracing has already said on stderr which half it was in and how long the other took. Also extracts _block_until_ready() so the execution half is substitutable in tests without JAX installed, and so the broad `except Exception` guards only the optional jax import rather than the wait itself. No call site changes: all four (Fitness._vmap/_jit/_grad, analysis/latent) inherit this through log_on_first_compile. For PyAutoFit#1528 (jax-compile-stall phase 3); follows #1517/#1518. Co-Authored-By: Claude --- autofit/non_linear/jax_compile.py | 110 +++++++++++++--- test_autofit/non_linear/test_jax_compile.py | 139 +++++++++++++++++++- 2 files changed, 224 insertions(+), 25 deletions(-) diff --git a/autofit/non_linear/jax_compile.py b/autofit/non_linear/jax_compile.py index 83e63c1e3..017b4bfb0 100644 --- a/autofit/non_linear/jax_compile.py +++ b/autofit/non_linear/jax_compile.py @@ -84,9 +84,23 @@ def dump_traceback_seconds(): return _env_seconds("PYAUTOFIT_JAX_COMPILE_DUMP_SECS", default) +# The two halves of a first call. They are separately nameable, and a heartbeat +# that cannot tell them apart reports the wrong one: before this existed a +# stalled run logged `still compiling ... 1770s elapsed` while its stack sat in +# `jax.block_until_ready`, and every marker written off that log calls this an +# "XLA compile stall" (autolens_workspace_test#271). +COMPILING = "compiling" +MATERIALIZING = "materializing" + + class _Heartbeat: """ - Log that a compile is still running, on an interval, until stopped. + Log that a first call is still running, on an interval, until stopped. + + The message names which half it is in. A phase-blind heartbeat is worse + than none for the case it exists to diagnose: it is *positive evidence* + for the wrong diagnosis, and three quarantines were written against + exactly that. Without this a long compile is indistinguishable from a stopped one: the process emits the "compiling..." line and then nothing at all until it @@ -101,15 +115,42 @@ def __init__(self, description, start, interval=None): self.description = description self.start_time = start self.interval = heartbeat_seconds() if interval is None else interval + self.phase = COMPILING + self.compiled_in = None self._stop = threading.Event() self._thread = None + def materializing(self, compiled_in): + """ + Move to the execution half, having compiled in `compiled_in` seconds. + + `compiled_in` is written BEFORE `phase` so the beat thread can never + observe `MATERIALIZING` with no compile time to report. Both are plain + attribute assignments, which are atomic under the GIL; the ordering is + the whole synchronisation and it is deliberate. + """ + self.compiled_in = compiled_in + self.phase = MATERIALIZING + def _beat(self): while not self._stop.wait(self.interval): - logger.info( - f"JAX jit still compiling {self.description}, " - f"{time.time() - self.start_time:.0f}s elapsed..." - ) + elapsed = time.time() - self.start_time + + if self.phase == COMPILING: + logger.info( + f"JAX jit still compiling {self.description}, " + f"{elapsed:.0f}s elapsed..." + ) + else: + # The compile time rides on every beat rather than only on the + # summary line, because a stalled run never reaches a summary + # line -- the runner SIGKILLs it. Whatever the last beat to + # reach stderr was, it carries the split. + logger.info( + f"JAX jit compiled {self.description} in " + f"{self.compiled_in:.1f}s and is still waiting for the " + f"result to materialize, {elapsed:.0f}s elapsed..." + ) def start(self): if self.interval <= 0: @@ -191,6 +232,23 @@ def cancel(self): self._armed = False +def _block_until_ready(result): + """ + Wait for `result` to materialize, if JAX is installed. + + Separated from the wrapper for two reasons: the execution half is the one + that hangs, so it has to be substitutable in a test that must not depend on + JAX being installed; and keeping the broad `except` around the import alone + stops it swallowing failures from the wait itself. + """ + try: + import jax + except Exception: + return + + jax.block_until_ready(result) + + def log_on_first_compile(func, description): """ Wrap `func` so that its first invocation reports the JAX compilation it @@ -211,11 +269,19 @@ def log_on_first_compile(func, description): 2. `jax.block_until_ready(result)` -- execution, since JAX dispatches asynchronously. - One summary line covering both cannot say which half a stalled run is stuck - in, and that ambiguity is why the same intermittent stall has been - quarantined three times across the test workspaces without a diagnosis. A - heartbeat reports liveness while either half is in flight, and a - `faulthandler` watchdog dumps a traceback if the whole thing overruns. + Both are reported AS THEY FINISH, never in one line at the end. A stalled + run is killed by the runner's cap and reaches no end, so a summary emitted + there describes only the runs that did not stall -- which is how the same + intermittent stall was quarantined five times across the test workspaces + under the name "XLA compile stall" while its stack sat in + `jax.block_until_ready`, the *execution* half, with compilation long since + finished (autolens_workspace_test#271, PyAutoFit#1528). + + So: the compile half is logged the instant `func` returns; the heartbeat + then names the materialize half it has moved into and repeats the compile + time on every beat; and a `faulthandler` watchdog dumps a traceback if the + whole thing overruns. A run killed at any point after tracing has already + said on stderr which half it was in, and how long the other one took. The one-shot flag lives in a closure rather than on an object because the callers cache these wrappers on attributes that are stripped for pickling; a @@ -254,16 +320,23 @@ def wrapper(*args, **kwargs): result = func(*args, **kwargs) compiled_at = time.time() + # Reported HERE, not once both halves are done. A stalled run never + # reaches the end of this function, so a split emitted after the + # wait characterises only the healthy case -- which is why the logs + # from 20 stalled runs could not say that compilation had in fact + # finished in ~16s. Emitted before the wait, it says so. + heartbeat.materializing(compiled_at - start) + logger.info( + f"JAX jit compilation of {description}: traced, lowered and " + f"compiled in {compiled_at - start:.1f} seconds; waiting for " + f"the result to materialize..." + ) + # JAX dispatches asynchronously, so `result` may be a future that is # not yet materialized. Block once, on this first call only, so the # duration logged below is the wait the user actually sat through # rather than the dispatch latency. - try: - import jax - - jax.block_until_ready(result) - except Exception: - pass + _block_until_ready(result) materialized_at = time.time() finally: @@ -272,9 +345,8 @@ def wrapper(*args, **kwargs): state["compiled"] = True logger.info( - f"JAX jit compilation of {description}: traced, lowered and " - f"compiled in {compiled_at - start:.1f} seconds, result " - f"materialized in {materialized_at - compiled_at:.1f} seconds." + f"JAX jit compilation of {description}: result materialized in " + f"{materialized_at - compiled_at:.1f} seconds." ) logger.info( diff --git a/test_autofit/non_linear/test_jax_compile.py b/test_autofit/non_linear/test_jax_compile.py index 0bbd95eef..fe6e0ff7c 100644 --- a/test_autofit/non_linear/test_jax_compile.py +++ b/test_autofit/non_linear/test_jax_compile.py @@ -218,12 +218,139 @@ def test_the_compile_wait_and_the_execution_wait_are_reported_separately(caplog) messages = [record.getMessage() for record in caplog.records] - # One summary line covering both waits cannot say which half a stalled run - # is stuck in -- the ambiguity this whole module exists to remove. - assert any( - "traced, lowered and compiled in" in m and "result materialized in" in m - for m in messages - ) + compiled = [i for i, m in enumerate(messages) if "traced, lowered and compiled in" in m] + materialized = [i for i, m in enumerate(messages) if "result materialized in" in m] + + assert compiled and materialized + + # Two lines, in order -- not one line carrying both. The order is the + # assertion that matters: it is what lets a run that is killed between them + # be read as "compiled, then stuck in execution". + assert compiled[0] < materialized[0] + + +def test_the_compile_half_is_reported_before_the_execution_half_is_waited_on(caplog): + """ + The defect this pins: the split used to be logged once BOTH halves finished, + so a stalled run -- which never finishes the second -- reported neither, and + twenty of them could not say that compilation had in fact completed. + """ + blocked = threading.Event() + reached = threading.Event() + + def never_returns(_result): + reached.set() + blocked.wait(timeout=10) + + wrapped = log_on_first_compile(lambda: None, "vectorized (vmap) likelihood function") + + with caplog.at_level(logging.INFO, logger=LOGGER): + with mock.patch.object(jax_compile, "_block_until_ready", never_returns): + worker = threading.Thread(target=wrapped, daemon=True) + worker.start() + assert reached.wait(timeout=5) + + # Still inside the execution half, exactly where a stalled CI run + # sits when the runner kills it. + messages = [record.getMessage() for record in caplog.records] + assert any("traced, lowered and compiled in" in m for m in messages) + assert not any("result materialized in" in m for m in messages) + + blocked.set() + worker.join(timeout=5) + + assert not worker.is_alive() + + +def test_the_heartbeat_names_the_half_it_is_actually_in(caplog): + """ + A phase-blind heartbeat logged `still compiling ... 1770s elapsed` while the + process sat in `jax.block_until_ready`. Every marker written off that log + calls this an "XLA compile stall"; it is not one. + """ + blocked = threading.Event() + reached = threading.Event() + + def never_returns(_result): + reached.set() + blocked.wait(timeout=10) + + wrapped = log_on_first_compile(lambda: None, "vectorized (vmap) likelihood function") + + with caplog.at_level(logging.INFO, logger=LOGGER): + with mock.patch.object(jax_compile, "heartbeat_seconds", return_value=0.05): + with mock.patch.object(jax_compile, "_block_until_ready", never_returns): + worker = threading.Thread(target=wrapped, daemon=True) + worker.start() + assert reached.wait(timeout=5) + time.sleep(0.25) + + beats = [ + m + for m in (r.getMessage() for r in caplog.records) + if "elapsed..." in m + ] + + blocked.set() + worker.join(timeout=5) + + assert beats + + # Every beat emitted while parked in the execution half must say so, and + # none may still claim to be compiling. + stuck = [m for m in beats if "waiting for the result to materialize" in m] + assert stuck + assert not any("still compiling" in m for m in beats[-len(stuck):]) + + +def test_a_heartbeat_in_the_execution_half_carries_the_compile_time(caplog): + """ + A stalled run is SIGKILLed and reaches no summary line, so the compile/execute + split has to ride on the beats themselves to survive to the captured tail. + """ + blocked = threading.Event() + reached = threading.Event() + + def never_returns(_result): + reached.set() + blocked.wait(timeout=10) + + def slow_compile(): + time.sleep(0.2) + + wrapped = log_on_first_compile(slow_compile, "vectorized (vmap) likelihood function") + + with caplog.at_level(logging.INFO, logger=LOGGER): + with mock.patch.object(jax_compile, "heartbeat_seconds", return_value=0.05): + with mock.patch.object(jax_compile, "_block_until_ready", never_returns): + worker = threading.Thread(target=wrapped, daemon=True) + worker.start() + assert reached.wait(timeout=5) + time.sleep(0.2) + + blocked.set() + worker.join(timeout=5) + + # "is still waiting" is the beat; the wrapper's own hand-off line says + # "waiting" without the "is still", and carries no elapsed counter. + stuck = [ + m + for m in (r.getMessage() for r in caplog.records) + if "is still waiting for the result to materialize" in m + ] + + assert stuck + assert all("compiled vectorized (vmap) likelihood function in" in m for m in stuck) + + +def test_the_execution_half_is_skipped_without_jax_and_does_not_raise(): + """ + `_block_until_ready` exists to be substitutable, but its real body must stay + a no-op when JAX is absent -- the library suite is numpy-only, and losing the + wait is a lost diagnostic, never a failed fit. + """ + with mock.patch.dict("sys.modules", {"jax": None}): + jax_compile._block_until_ready(object()) def test_the_split_is_not_reported_when_the_first_call_raises(caplog):