fix[next]: stop KeyError from masking compiled-program failures

[havogt]

Description

Draft — split out of #2731 (now closed) so the fix and the error-message improvement can be
reviewed separately. #2734 stacks on top with the message change.

When a compiled program cannot be obtained, the traceback leads with a KeyError on the
dispatch dict even though the actual cause is something else. Real incident (ICON/icon4py,
10 MPI ranks, shared OTF cache on Lustre scratch):

ERROR - A Python error occurred: [Errno 116] Stale file handle
Traceback (most recent call last):
  File ".../gt4py/next/otf/compiled_program.py", line 419, in __call__
    compiled_program = self.compiled_programs[key]
KeyError: ((np.int32(557), np.int32(72619), np.int32(0), np.int32(80)), -5810085652480803948, None)
During handling of the above exception, another exception occurred:
...

The genuine cause was OSError: [Errno 116] Stale file handle while loading the artifact, but
the KeyError with its opaque key tuple is what everyone reads first — it looks like a
cache-key/dispatch bug, and it is not. That cost hours of misdirected debugging.

CompiledProgramsPool.__call__ did all cache-miss handling inside the except KeyError
handler, so everything raised there — future.result() re-raising a worker exception,
artifact.load() doing I/O — was chained to the KeyError.

Changes

• The dispatch lookup is now self.compiled_programs.get(key) plus an if ... is None. Nothing
  is raised on a miss, so there is no exception context for the real failure to be chained to.
  The miss branch itself is unchanged.
• _load_artifact wraps artifact.load() failures as
  RuntimeError("Failed to load the compiled program '<name>'.") from <cause>, so the OSError
  is what the reader sees. Worker exceptions from future.result() are deliberately not
  wrapped — they propagate as themselves.

Hot path

__call__ runs per program invocation (thousands per second in ICON), so the cost of the
successful lookup matters. Isolated dict lookup, key present:

	3.10	3.12	3.13
try: d[key] except KeyError:	30.0 ns	15.7 ns	16.1 ns
d.get(key)	30.6 ns	21.7 ns	20.1 ns
key in d + d[key]	43.1 ns	—	—

So .get() costs ~5 ns on ≥3.11, where the non-raising try became free, and nothing on 3.10.
That is 0.2 % of the ~2800 ns a full dispatch takes; a whole-dispatch microbenchmark (program
call with the executable replaced by a no-op, under -O) cannot resolve it above run-to-run
noise. Taken deliberately in exchange for a much smaller diff — an earlier revision kept the
try/except and restructured __call__ so the handler was left before any miss handling
(identical hit-path bytecode, zero cost); see
the discussion with
@egparedes.

Note that just adding raise ... from e at the raise sites does not fix this: from e sets
__suppress_context__ on the outer exception, but the intermediate OSError still carries
__context__ = KeyError and is printed as the cause, so the KeyError still leads the
traceback. Hiding it that way would require setting inner.__context__ = None at every site.

Open question

_finish_compilation_job still has assert key not in self.compiled_programs, which never runs
under icon4py's PYTHONOPTIMIZE=2. Left as an assert: it is only reachable right after a
confirmed dict miss, so it cannot be violated by user input.

Out of scope: the underlying Lustre ESTALE, which belongs with the OTF cache write/replace
logic (#2691). This is only about which error the user is shown.

Requirements

• All fixes and/or new features come with corresponding tests.
  Three unit tests in tests/next_tests/unit_tests/otf_tests/test_compiled_program.py
  (load OSError, worker exception, genuine miss); all three fail on main and pass here.
• Important design decisions have been documented in the appropriate ADR.

[Copilot]

🟢 Ready to approve

The changes are localized, align with the stated failure-mode goal, and are covered by targeted regression tests that validate the traceback/cause behavior.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

This PR improves error reporting in gt4py.next’s OTF compiled-program dispatch so that real compilation/load failures are no longer masked by an initial KeyError from the dispatch dictionary lookup. This directly targets clearer tracebacks and faster diagnosis of underlying I/O / artifact-load problems in production deployments.

Changes:

• Refactors CompiledProgramsPool.__call__ to avoid KeyError-based miss handling (uses dict.get() + None check), preventing unrelated exceptions from being chained to a dict-miss KeyError.
• Adds _load_artifact() to wrap artifact.load() failures with a targeted RuntimeError(... ) from <cause> while preserving worker exceptions from future.result() as-is.
• Adds unit tests ensuring (a) load errors surface as the cause, (b) worker exceptions propagate without KeyError context, and (c) non-JIT cache misses don’t show KeyError in the traceback.

File summaries

File	Description
src/gt4py/next/otf/compiled_program.py	Removes KeyError-driven miss path and wraps artifact load failures to prevent KeyError from dominating tracebacks.
tests/next_tests/unit_tests/otf_tests/test_compiled_program.py	Adds regression tests for traceback chaining behavior across load failures, worker failures, and non-JIT misses.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Low

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.