Dump the compile traceback before the runner kills the process

autofit/non_linear/jax_compile.py

@@ -9,10 +9,15 @@
 # How often a compile that is still running reports that it is still alive.
 DEFAULT_HEARTBEAT_SECONDS = 30.0
 
-# How long a first compile may run under CI before it dumps its own traceback.
-# Off by default off-CI: an interactive user watching a slow compile does not
-# want a traceback on stderr, they want the heartbeat above.
-DEFAULT_CI_DUMP_SECONDS = 300.0
+# How long a first compile may run under CI before it dumps its own traceback,
+# when the runner has not told us its own cap. Off by default off-CI: an
+# interactive user watching a slow compile does not want a traceback on stderr,
+# they want the heartbeat above.
+DEFAULT_CI_DUMP_SECONDS = 240.0
+
+# Fraction of the runner's per-script cap at which to dump. The dump is only
+# ever useful STRICTLY BEFORE the kill -- see `dump_traceback_seconds`.
+DUMP_FRACTION_OF_CAP = 0.8
 
 
 def _env_seconds(name, default):
@@ -60,8 +65,22 @@ def dump_traceback_seconds():
     runner -- the alternative, threading an environment variable through each
     workspace's `config/build/env_vars_*.yaml`, is more repos touched for the
     same effect, and the workspace scripts are user-facing documentation.
+
+    Under CI the default is derived from `BUILD_SCRIPT_TIMEOUT`, the per-script
+    cap the workspace runners and PyAutoHands both enforce. **A dump scheduled
+    at or after that cap never happens**: the runner SIGKILLs the process group
+    on expiry, and a killed process writes no traceback. The first CI use of
+    this watchdog hit exactly that -- a flat 300s default against a 300s smoke
+    cap produced heartbeats from 20 stalled runs and not one stack
+    (autolens_workspace_test#271). Dumping at a fraction of the cap leaves the
+    traceback time to reach stderr before the kill lands.
     """
-    default = DEFAULT_CI_DUMP_SECONDS if os.environ.get("CI") else 0.0
+    if not os.environ.get("CI"):
+        default = 0.0
+    else:
+        cap = _env_seconds("BUILD_SCRIPT_TIMEOUT", 0.0)
+        default = cap * DUMP_FRACTION_OF_CAP if cap > 0 else DEFAULT_CI_DUMP_SECONDS
+
     return _env_seconds("PYAUTOFIT_JAX_COMPILE_DUMP_SECS", default)
 
 

test_autofit/non_linear/test_jax_compile.py

@@ -256,6 +256,7 @@ def test_the_heartbeat_interval_comes_from_the_environment(monkeypatch):
 
 def test_the_traceback_dump_defaults_on_under_ci_and_off_elsewhere(monkeypatch):
     monkeypatch.delenv("PYAUTOFIT_JAX_COMPILE_DUMP_SECS", raising=False)
+    monkeypatch.delenv("BUILD_SCRIPT_TIMEOUT", raising=False)
 
     monkeypatch.delenv("CI", raising=False)
     assert jax_compile.dump_traceback_seconds() == 0.0
@@ -272,6 +273,45 @@ def test_the_traceback_dump_defaults_on_under_ci_and_off_elsewhere(monkeypatch):
     assert jax_compile.dump_traceback_seconds() == 0.0
 
 
+def test_the_dump_lands_strictly_before_the_runner_kills_the_process(monkeypatch):
+    monkeypatch.delenv("PYAUTOFIT_JAX_COMPILE_DUMP_SECS", raising=False)
+    monkeypatch.setenv("CI", "true")
+
+    # The defect this pins (autolens_workspace_test#271): a flat 300s default
+    # against a 300s smoke cap meant the runner's SIGKILL always beat the dump,
+    # so 20 stalled CI runs produced heartbeats and not one traceback. A killed
+    # process writes no stack, so the threshold MUST be under the cap.
+    for cap in ("300", "1800", "60"):
+        monkeypatch.setenv("BUILD_SCRIPT_TIMEOUT", cap)
+        assert jax_compile.dump_traceback_seconds() < float(cap)
+
+    monkeypatch.setenv("BUILD_SCRIPT_TIMEOUT", "300")
+    assert jax_compile.dump_traceback_seconds() == 240.0
+
+    monkeypatch.setenv("BUILD_SCRIPT_TIMEOUT", "1800")
+    assert jax_compile.dump_traceback_seconds() == 1440.0
+
+
+def test_an_unusable_runner_cap_falls_back_to_the_flat_ci_default(monkeypatch):
+    monkeypatch.delenv("PYAUTOFIT_JAX_COMPILE_DUMP_SECS", raising=False)
+    monkeypatch.setenv("CI", "true")
+
+    # No cap advertised, or a meaningless one: there is nothing to derive from,
+    # so use the flat default rather than computing a fraction of zero (which
+    # would silently disable the dump).
+    for cap in ("", "0", "not-a-number"):
+        monkeypatch.setenv("BUILD_SCRIPT_TIMEOUT", cap)
+        assert jax_compile.dump_traceback_seconds() == jax_compile.DEFAULT_CI_DUMP_SECONDS
+
+
+def test_an_explicit_dump_threshold_still_wins_over_the_derived_one(monkeypatch):
+    monkeypatch.setenv("CI", "true")
+    monkeypatch.setenv("BUILD_SCRIPT_TIMEOUT", "1800")
+    monkeypatch.setenv("PYAUTOFIT_JAX_COMPILE_DUMP_SECS", "90")
+
+    assert jax_compile.dump_traceback_seconds() == 90.0
+
+
 def test_a_malformed_interval_falls_back_rather_than_raising(monkeypatch):
     monkeypatch.setenv("PYAUTOFIT_JAX_COMPILE_HEARTBEAT_SECS", "soon")
     assert jax_compile.heartbeat_seconds() == jax_compile.DEFAULT_HEARTBEAT_SECONDS