[tf.data] DeadlineExceededError from input pipeline graph rewrites should not prevent further computation.

PiperOrigin-RevId: 296513646
Change-Id: I9492c986567502f52bce63c04f9f7bb6d9da76ac
This commit is contained in:
Jiri Simsa 2020-02-21 14:35:31 -08:00 committed by TensorFlower Gardener
parent 36b569fcd5
commit e406991c37
1 changed files with 9 additions and 3 deletions

View File

@ -55,9 +55,15 @@ void OptimizeDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
auto config_factory = [this, &optimizations]() {
return CreateConfig(optimizations, optimization_configs_);
};
OP_REQUIRES_OK(ctx, RewriteDataset(ctx, input, std::move(config_factory),
/*optimize_function_library=*/true,
/*record_fingerprint=*/true, output));
Status s = RewriteDataset(ctx, input, std::move(config_factory),
/*optimize_function_library=*/true,
/*record_fingerprint=*/true, output);
if (errors::IsDeadlineExceeded(s)) {
// Ignore DeadlineExceeded as it implies that the attempted rewrite took too
// long which should not prevent further computation.
return;
}
OP_REQUIRES_OK(ctx, s);
}
RewriterConfig OptimizeDatasetOp::CreateConfig(