| 289 | } |
| 290 | |
| 291 | void IteratorHandleOp::Compute(OpKernelContext* context) LOCKS_EXCLUDED(mu_) { |
| 292 | { |
| 293 | mutex_lock l(mu_); |
| 294 | if (resource_ == nullptr) { |
| 295 | FunctionLibraryRuntime* flr; |
| 296 | std::unique_ptr<DeviceMgr> device_mgr(nullptr); |
| 297 | std::unique_ptr<FunctionLibraryDefinition> flib_def(nullptr); |
| 298 | std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(nullptr); |
| 299 | // If the iterator is shared then we construct a new FLR, and pass that |
| 300 | // in. NOTE(mrry,rohanj): In this case it is not possible to call remote |
| 301 | // functions from the iterator. We may add this functionality if there |
| 302 | // is sufficient demand, but it will require a significant refactoring. |
| 303 | if (!name_.empty()) { |
| 304 | flr = CreatePrivateFLR(context, &device_mgr, &flib_def, &pflr); |
| 305 | } else { |
| 306 | OP_REQUIRES_OK(context, context->function_library()->Clone( |
| 307 | &flib_def, &pflr, &flr, true)); |
| 308 | } |
| 309 | |
| 310 | ResourceMgr* mgr = context->resource_manager(); |
| 311 | #ifdef TENSORFLOW_USE_ELASTIC_SERVER |
| 312 | OP_REQUIRES_OK(context, cinfo_.Init(mgr, def(), true)); |
| 313 | #else |
| 314 | OP_REQUIRES_OK(context, cinfo_.Init(mgr, def(), false)); |
| 315 | #endif |
| 316 | |
| 317 | IteratorResource* resource; |
| 318 | OP_REQUIRES_OK( |
| 319 | context, |
| 320 | mgr->LookupOrCreate<IteratorResource>( |
| 321 | cinfo_.container(), cinfo_.name(), &resource, |
| 322 | [context, flr, &device_mgr, &flib_def, &pflr, |
| 323 | this](IteratorResource** ret) EXCLUSIVE_LOCKS_REQUIRED(mu_) { |
| 324 | *ret = new IteratorResource( |
| 325 | context->env(), output_dtypes_, output_shapes_, |
| 326 | graph_def_version_, std::move(device_mgr), |
| 327 | std::move(flib_def), std::move(pflr), flr); |
| 328 | return Status::OK(); |
| 329 | })); |
| 330 | |
| 331 | Status s = VerifyResource(resource); |
| 332 | if (TF_PREDICT_FALSE(!s.ok())) { |
| 333 | resource->Unref(); |
| 334 | context->SetStatus(s); |
| 335 | return; |
| 336 | } |
| 337 | |
| 338 | resource_ = resource; |
| 339 | } |
| 340 | } |
| 341 | OP_REQUIRES_OK(context, MakeResourceHandleToOutput( |
| 342 | context, 0, cinfo_.container(), cinfo_.name(), |
| 343 | MakeTypeIndex<IteratorResource>())); |
| 344 | } |
| 345 | |
| 346 | Status IteratorHandleOp::VerifyResource(IteratorResource* resource) { |
| 347 | TF_RETURN_IF_ERROR( |
nothing calls this directly
no test coverage detected