| 533 | } |
| 534 | |
| 535 | StatusOr<GlobalDataHandle> Service::ExecuteAndRegisterResult( |
| 536 | Executable* executable, |
| 537 | absl::Span<const std::vector<const ShapedBuffer*>> arguments, |
| 538 | Backend* backend, const DeviceHandle& device_handle, |
| 539 | const string& result_tag, ExecutionProfile* profile) { |
| 540 | // Set up streams. |
| 541 | std::vector<StreamPool::Ptr> streams; |
| 542 | |
| 543 | TF_ASSIGN_OR_RETURN(auto replicas, Replicas(*backend, device_handle)); |
| 544 | TF_RET_CHECK(!replicas.empty()); |
| 545 | for (se::StreamExecutor* executor : replicas) { |
| 546 | TF_ASSIGN_OR_RETURN(StreamPool::Ptr stream, |
| 547 | backend->BorrowStream(executor)); |
| 548 | streams.push_back(std::move(stream)); |
| 549 | } |
| 550 | |
| 551 | DeviceAssignment device_assignment(options_.number_of_replicas(), |
| 552 | /*computation_count=*/1); |
| 553 | for (int64 replica = 0; replica < replicas.size(); ++replica) { |
| 554 | device_assignment(replica, 0) = replicas[replica]->device_ordinal(); |
| 555 | } |
| 556 | |
| 557 | // Set up run options. |
| 558 | std::vector<ServiceExecutableRunOptions> run_options; |
| 559 | for (const StreamPool::Ptr& stream : streams) { |
| 560 | ExecutableRunOptions options; |
| 561 | options.set_stream(stream.get()); |
| 562 | options.set_device_ordinal(stream->parent()->device_ordinal()); |
| 563 | options.set_allocator(backend->memory_allocator()); |
| 564 | options.set_intra_op_thread_pool( |
| 565 | backend->eigen_intra_op_thread_pool_device()); |
| 566 | options.set_device_assignment(&device_assignment); |
| 567 | options.set_execution_profile(profile); |
| 568 | run_options.emplace_back(options, backend->StreamBorrower()); |
| 569 | } |
| 570 | |
| 571 | if (options_.number_of_replicas() == 1) { |
| 572 | TF_ASSIGN_OR_RETURN(auto result, executable->ExecuteOnStreamWrapper( |
| 573 | &run_options[0], arguments[0])); |
| 574 | return allocation_tracker_.Register(std::move(result), result_tag); |
| 575 | } |
| 576 | |
| 577 | // TODO(b/69985541): Support profiling also on this path. |
| 578 | |
| 579 | std::vector<absl::Span<const ShapedBuffer* const>> replicated_arguments; |
| 580 | for (const auto& arg : arguments) { |
| 581 | replicated_arguments.push_back(arg); |
| 582 | } |
| 583 | |
| 584 | TF_ASSIGN_OR_RETURN(auto results, executable->ExecuteOnStreams( |
| 585 | run_options, replicated_arguments)); |
| 586 | TF_RET_CHECK(!results.empty()); |
| 587 | return allocation_tracker_.RegisterReplicatedBuffers(std::move(results), |
| 588 | result_tag); |
| 589 | } |
| 590 | |
| 591 | StatusOr<std::vector<se::StreamExecutor*>> Service::GetExecutors( |
| 592 | const ExecutionOptions& execution_options, int64 requests_size, |
nothing calls this directly
no test coverage detected