fix modular_filesystem call itself - #45682

…ice 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

Reverts changelist 949495726

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

…ice 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

Reverts changelist 949495726

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

…ice 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

Reverts changelist 949495726

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

…ice 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

Reverts changelist 949495726

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

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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
PiperOrigin-RevId: 950753335

…ice 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

PiperOrigin-RevId: 953266494