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):