Skip to content

fix(executorch): drain the stream when a device-to-host copy fails - #4471

Open
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:trt-drain-on-copy-failure
Open

fix(executorch): drain the stream when a device-to-host copy fails#4471
shoumikhin wants to merge 1 commit into
pytorch:mainfrom
shoumikhin:trt-drain-on-copy-failure

Conversation

@shoumikhin

Copy link
Copy Markdown
Contributor

The problem

When a device-to-host copy of a delegate output fails, execute() returns straight away.
By that point enqueueV3() has already succeeded, so the engine is still running on the
stream, and the early return skips two things:

      if (cuda_err != cudaSuccess) {
        ET_LOG(Error, "D2H copy failed for output %zu: %s", ...);
        return Error::InvalidProgram;      // <- leaves here
      }
    }
    cuda_err = cudaStreamSynchronize(stream);   // never runs
    engine->inflight_pending = false;           // never runs

TensorRT forbids mutating an execution context while one of its enqueues is in flight. The
next call runs setInputShape and setTensorAddress, both host-side, on a context that may
still be executing.

The in-flight guard at the top of execute() cannot save it, because the flag it reads was
left stale:

  if (engine->inflight_pending) {          // false, so no wait happens
    cuda_err = cudaEventSynchronize(engine->inflight_event);

So the failure is not just an un-awaited copy. It also disarms the mechanism that exists to
prevent exactly this.

The fix

Record the failure, finish the cleanup, then return:

    Error copy_err = Error::Ok;
    for (auto& output : outputs_needing_copy) {
      ...
      if (cuda_err != cudaSuccess) {
        ET_LOG(Error, "D2H copy failed for output %zu: %s", ...);
        copy_err = Error::InvalidProgram;
        break;
      }
    }
    cuda_err = cudaStreamSynchronize(stream);
    engine->inflight_pending = false;
    if (cuda_err != cudaSuccess) { ... return Error::InvalidProgram; }
    if (copy_err != Error::Ok) {
      return copy_err;
    }

A copy failure now leaves the engine in the same state a successful call does, so the next
call behaves predictably instead of depending on whether the previous one happened to fail.
The reported error is unchanged.

Note the ordering: a synchronize failure is reported ahead of the copy failure, because a
stream that cannot be drained is the more serious of the two and the copy error is already
logged by then.

Testing

Built the delegate and ran a program with host-backed outputs, so the copy path is
exercised on every call, and confirmed repeated calls still produce output matching the
eager model. Compiled against TensorRT 11.1 with no new warnings.

I have not built a fault injector for the copy itself, so the recovery path is reasoned from
the code rather than observed. Happy to add a test if there is a good way to force a
cudaMemcpyAsync failure in this suite.

When a device-to-host copy of a delegate output fails, execute() returns
immediately. By then enqueueV3() has already succeeded, so the engine is still
running on the stream, and two pieces of cleanup are skipped:

  * cudaStreamSynchronize, so the work is never awaited
  * inflight_pending = false

TensorRT forbids mutating an execution context while one of its enqueues is in
flight. The next call reaches setInputShape and setTensorAddress on a context
that may still be executing, and its own in-flight guard cannot help because
inflight_pending was left stale.

Record the failure, still drain the stream and clear the flag, then return. A
copy failure now leaves the engine in the same state a successful call does.
@meta-cla meta-cla Bot added the cla signed label Aug 8, 2026
@github-actions github-actions Bot added the component: api [C++] Issues re: C++ API label Aug 8, 2026
@github-actions
github-actions Bot requested a review from narendasan August 8, 2026 23:57
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant