| 738 | } |
| 739 | |
| 740 | StatusOr<std::vector<std::unique_ptr<PyLocalBuffer>>> |
| 741 | PyLocalExecutable::ExecuteOnLocalDevices( |
| 742 | absl::Span<const std::vector<PyLocalBuffer*>> argument_handles) { |
| 743 | tensorflow::profiler::TraceMe traceme( |
| 744 | "LocalExecutable::ExecuteOnLocalDevices"); |
| 745 | |
| 746 | const int num_local_devices = local_devices_.size(); |
| 747 | |
| 748 | if (argument_handles.size() != num_local_devices) { |
| 749 | return InvalidArgument( |
| 750 | "Attempted to execute with %d argument lists when local device " |
| 751 | "count is %d (total replica count: %d, partition count: %d)", |
| 752 | argument_handles.size(), num_local_devices, num_replicas(), |
| 753 | num_partitions()); |
| 754 | } |
| 755 | |
| 756 | VLOG(1) << "Executing computation " << name() |
| 757 | << "; num_replicas=" << num_replicas() |
| 758 | << " num_partitions=" << num_partitions() |
| 759 | << " num_local_devices=" << num_local_devices; |
| 760 | std::vector<StatusOr<std::unique_ptr<PyLocalBuffer>>> results( |
| 761 | num_local_devices); |
| 762 | if (num_local_devices == 1) { |
| 763 | // Fast-path if there is only one device — run the computation on the |
| 764 | // current thread. |
| 765 | const int replica = local_logical_device_ids_[0].first; |
| 766 | const int partition = local_logical_device_ids_[0].second; |
| 767 | results[0] = |
| 768 | ExecuteHelper(argument_handles[0], replica, partition, RunId()); |
| 769 | } else { |
| 770 | RunId run_id; |
| 771 | absl::Mutex mu; |
| 772 | int running = num_local_devices; |
| 773 | int failed = 0; |
| 774 | Status first_failure_status; |
| 775 | |
| 776 | for (int i = 0; i < num_local_devices; ++i) { |
| 777 | const int replica = local_logical_device_ids_[i].first; |
| 778 | const int partition = local_logical_device_ids_[i].second; |
| 779 | std::shared_ptr<Device> device = local_devices_[i]; |
| 780 | const LocalDeviceState& device_state = *device->local_device_state(); |
| 781 | device_state.execute_thread()->Schedule([&, replica, partition, i] { |
| 782 | results[i] = |
| 783 | ExecuteHelper(argument_handles[i], replica, partition, run_id); |
| 784 | |
| 785 | absl::MutexLock lock(&mu); |
| 786 | --running; |
| 787 | if (!results[i].ok()) { |
| 788 | if (failed == 0) { |
| 789 | first_failure_status = results[i].status(); |
| 790 | } |
| 791 | ++failed; |
| 792 | } |
| 793 | }); |
| 794 | } |
| 795 | |
| 796 | auto done_running_or_failed = [&]() { |
| 797 | mu.AssertHeld(); |
nothing calls this directly
no test coverage detected