| 745 | } |
| 746 | |
| 747 | void InstantiatedCapturedFunction::RunAsync( |
| 748 | IteratorContext* ctx, std::vector<Tensor>&& args, std::vector<Tensor>* rets, |
| 749 | FunctionLibraryRuntime::DoneCallback done, const string& prefix) const { |
| 750 | auto& info = captured_func_->short_circuit_info(); |
| 751 | if (!info.indices.empty()) { |
| 752 | // Run the `done` callback on a threadpool thread, because it will |
| 753 | // potentially do a non-trivial amount of (e.g. copying) work, and we may |
| 754 | // want to run that concurrently with the next invocation. |
| 755 | Status s = RunShortCircuit(info, std::move(args), captured_func_, rets); |
| 756 | (*ctx->runner())( |
| 757 | std::bind([s](FunctionLibraryRuntime::DoneCallback& done) { done(s); }, |
| 758 | std::move(done))); |
| 759 | return; |
| 760 | } |
| 761 | |
| 762 | // NOTE(mrry): This method does not transfer ownership of `ctx`, and it may |
| 763 | // be deleted before `done` is called. Take care not to capture `ctx` in any |
| 764 | // code that may execute asynchronously in this function. |
| 765 | OwnedArgsCallFrame* frame = new OwnedArgsCallFrame( |
| 766 | std::move(args), &captured_func_->captured_inputs(), ret_types_); |
| 767 | |
| 768 | FunctionLibraryRuntime::Options f_opts; |
| 769 | ResourceMgr* resource_mgr = lib_->device()->resource_manager(); |
| 770 | ScopedStepContainer* step_container = new ScopedStepContainer( |
| 771 | f_opts.step_id, [resource_mgr](const string& name) { |
| 772 | resource_mgr->Cleanup(name).IgnoreError(); |
| 773 | }); |
| 774 | f_opts.step_container = step_container; |
| 775 | f_opts.runner = ctx->runner(); |
| 776 | f_opts.create_rendezvous = ShouldCreateRendezvous(); |
| 777 | auto cancellation_manager = absl::make_unique<CancellationManager>(); |
| 778 | f_opts.cancellation_manager = cancellation_manager.get(); |
| 779 | std::function<void()> deregister_fn; |
| 780 | Status s = ConnectCancellationManagers( |
| 781 | ctx->cancellation_manager(), cancellation_manager.get(), &deregister_fn); |
| 782 | if (!s.ok()) { |
| 783 | done(s); |
| 784 | return; |
| 785 | } |
| 786 | |
| 787 | std::shared_ptr<SimpleStepStatsCollector> stats_collector; |
| 788 | if (ctx->model() || ctx->stats_aggregator()) { |
| 789 | stats_collector = absl::make_unique<SimpleStepStatsCollector>(); |
| 790 | } |
| 791 | f_opts.stats_collector = stats_collector.get(); |
| 792 | |
| 793 | // Transfer ownership of the cancellation manager to `callback`. |
| 794 | CancellationManager* raw_cancellation_manager = |
| 795 | cancellation_manager.release(); |
| 796 | auto callback = std::bind( |
| 797 | [this, rets, step_container, raw_cancellation_manager, frame]( |
| 798 | const FunctionLibraryRuntime::DoneCallback& done, |
| 799 | IteratorContext* ctx, const std::function<void()>& deregister_fn, |
| 800 | const string& prefix, |
| 801 | const std::shared_ptr<SimpleStepStatsCollector>& stats_collector, |
| 802 | // Begin unbound arguments. |
| 803 | Status s) { |
| 804 | delete step_container; |
no test coverage detected