PR #45682: pjrt: propagate Prepare failure on non-two-phase multi-device Execute - #123626

PR #45682: pjrt: propagate Prepare failure on non-two-phase multi-device Execute

Imported from GitHub PR openxla/xla#45682

The bug

CommonPjRtLoadedExecutable::Execute's multi-device path calls ExecutePrepareWithOomRetries per device, then — only when supports_two_phase_launch() is true — checks the returned status inside the two-phase barrier block. CpuClient overrides supports_two_phase_launch() to false, so on a multi-device CpuClient a Prepare failure is never checked and execution falls through to ExecuteLaunch(*launch_args, ...) with launch_args.executable still null → null-deref crash instead of a clean error.

The fix

Add an else if (!launch_status.ok()) branch after the two-phase block: record the failure in results[i], decrement launching under mu, return. Same shape as the two-phase failure path minus the barrier it doesn't participate in.

Scope

CpuClient is the only in-tree client with supports_two_phase_launch()==false. Two-phase clients (GPU/TPU stream-executor) are unchanged — the diff is purely additive.

Test

MultiDevicePrepareFailurePropagatesError in [cpu_client_test.cc](http://cpu_client_test.cc/): 2-device CpuClient, 2-replica identity program, wrong-size buffer on device 1 so CheckBufferCompatibilities rejects it during Prepare. Asserts Execute() returns an error (not a crash) and that the error comes from the per-device Prepare path.
Copybara import of the project:

--
1e96c1e75c504728fe20c950b4d46ea7ad464b59 by Nikhil Bhargava [email protected]:

pjrt: propagate Prepare failure on non-two-phase multi-device Execute

When supports_two_phase_launch()==false (CpuClient), the multi-device
Execute path never checked the status returned by
ExecutePrepareWithOomRetries — the only check was inside the two-phase
branch. A Prepare failure fell through to ExecuteLaunch with
launch_args.executable still null and crashed.

Add the missing status check on the non-two-phase path: record the
failure in results[i], decrement the launching counter under mu, and
return — mirroring the two-phase failure handling without the barrier.

Test: MultiDevicePrepareFailurePropagatesError in cpu_client_test.cc
creates a 2-device CpuClient, passes a wrong-size buffer on one device
so CheckBufferCompatibilities rejects it during Prepare, and asserts
Execute() returns an error rather than crashing.

Merging this change closes #45682

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45682 from nbhargava:pjrt-non-two-phase-prepare-error 1e96c1e75c504728fe20c950b4d46ea7ad464b59