| 524 | } |
| 525 | |
| 526 | Status PyLocalBuffer::BlockHostUntilReady() { |
| 527 | tensorflow::profiler::TraceMe traceme("PyLocalBuffer::BlockHostUntilReady"); |
| 528 | std::shared_ptr<SharedDeviceBuffer> device_buffer = DeviceBuffer(); |
| 529 | if (!device_buffer) { |
| 530 | return InvalidArgument("BlockHostUntilReady() called on invalid buffer."); |
| 531 | } |
| 532 | |
| 533 | // This code waits at least until the buffer is ready, but it may wait longer |
| 534 | // if there are other device to host transfers scheduled. If this proves to |
| 535 | // be an issue, we could either use a separate stream for this purpose, or |
| 536 | // poll for the buffer definition events. |
| 537 | se::Stream* stream = |
| 538 | client_->device_state(device_->local_device_state()->device_ordinal()) |
| 539 | .GetDeviceToHostStream(); |
| 540 | WaitForBufferDefinitionEventsOnStream(*device_buffer, stream); |
| 541 | return stream->BlockHostUntilDone(); |
| 542 | } |
| 543 | |
| 544 | static std::shared_ptr<Device> LookupDevice(const PyLocalClient& client, |
| 545 | int device_id) { |
no test coverage detected