| 786 | } |
| 787 | |
| 788 | Status TryInit(OpKernelContext* ctx, IteratorResource** iterator, |
| 789 | ContainerInfo* cinfo) { |
| 790 | #ifdef TENSORFLOW_USE_ELASTIC_SERVER |
| 791 | TF_RETURN_IF_ERROR(cinfo->Init(ctx->resource_manager(), def(), true)); |
| 792 | #else |
| 793 | TF_RETURN_IF_ERROR(cinfo->Init(ctx->resource_manager(), def(), false)); |
| 794 | #endif |
| 795 | |
| 796 | FunctionLibraryRuntime* flr; |
| 797 | std::unique_ptr<FunctionLibraryDefinition> flib_def(nullptr); |
| 798 | std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(nullptr); |
| 799 | TF_RETURN_IF_ERROR( |
| 800 | ctx->function_library()->Clone(&flib_def, &pflr, &flr, true)); |
| 801 | |
| 802 | // Create an IteratorResource that will hold the iterator for this op. |
| 803 | TF_RETURN_IF_ERROR( |
| 804 | ctx->resource_manager()->LookupOrCreate<IteratorResource>( |
| 805 | cinfo->container(), cinfo->name(), iterator, |
| 806 | [ctx, flr, this, &flib_def, &pflr](IteratorResource** ret) |
| 807 | EXCLUSIVE_LOCKS_REQUIRED(mu_) { |
| 808 | *ret = new IteratorResource( |
| 809 | ctx->env(), output_dtypes_, output_shapes_, |
| 810 | graph_def_version_, nullptr, std::move(flib_def), |
| 811 | std::move(pflr), flr); |
| 812 | return Status::OK(); |
| 813 | })); |
| 814 | |
| 815 | core::ScopedUnref unref_iterator(*iterator); |
| 816 | |
| 817 | TF_RETURN_IF_ERROR( |
| 818 | VerifyTypesMatch(output_dtypes_, (*iterator)->output_dtypes())); |
| 819 | TF_RETURN_IF_ERROR( |
| 820 | VerifyShapesCompatible(output_shapes_, (*iterator)->output_shapes())); |
| 821 | |
| 822 | // Call the dataset_factory_func_ to create a new dataset, |
| 823 | // over which this op will iterate. |
| 824 | FunctionLibraryRuntime::Handle f_handle; |
| 825 | TF_RETURN_IF_ERROR(ctx->function_library()->Instantiate( |
| 826 | dataset_factory_func_.name(), AttrSlice(&dataset_factory_func_.attr()), |
| 827 | &f_handle)); |
| 828 | FunctionLibraryRuntime::Options opts; |
| 829 | opts.cancellation_manager = ctx->cancellation_manager(); |
| 830 | ScopedStepContainer step_container(opts.step_id, [ctx](const string& name) { |
| 831 | ctx->resource_manager()->Cleanup(name).IgnoreError(); |
| 832 | }); |
| 833 | opts.step_container = &step_container; |
| 834 | opts.runner = ctx->runner(); |
| 835 | Notification n; |
| 836 | Status factory_status; |
| 837 | std::vector<Tensor> return_values; |
| 838 | ctx->function_library()->Run(opts, f_handle, {}, &return_values, |
| 839 | [&n, &factory_status](Status s) { |
| 840 | factory_status.Update(s); |
| 841 | n.Notify(); |
| 842 | }); |
| 843 | n.WaitForNotification(); |
| 844 | TF_RETURN_IF_ERROR(factory_status); |
| 845 | if (return_values.size() != 1 || return_values[0].dtype() != DT_VARIANT || |
nothing calls this directly
no test coverage detected