| 373 | } |
| 374 | |
| 375 | Status OptimizeGraph(const GrapplerItem& item, int64 num_workers, int64 index, |
| 376 | GraphDef* output) { |
| 377 | if (num_workers == 1 && index == 0) { |
| 378 | return Status::OK(); |
| 379 | } |
| 380 | |
| 381 | *output = item.graph; |
| 382 | MutableGraphView graph(output); |
| 383 | FunctionLibraryDefinition flib(OpRegistry::Global(), item.graph.library()); |
| 384 | |
| 385 | NodeDef target_node; |
| 386 | absl::flat_hash_set<string> nodes_to_delete; |
| 387 | |
| 388 | // The basic approach here is to walk the graph from sink to source, and find |
| 389 | // the latest occurrence of a ReaderDataset (e.g. CSVDataset, TFRecordDataset, |
| 390 | // etc...). We then add a shard after that dataset to shard the outputs of |
| 391 | // that dataset, in effect giving a piece to each worker. Finally, we remove |
| 392 | // occurences from randomness from before that point in the graph (e.g. things |
| 393 | // like ShuffleDataset) to ensure that `shard` returns a sensible result. |
| 394 | NodeDef* sink_node; |
| 395 | TF_RETURN_IF_ERROR(graph_utils::GetFetchNode(graph, item, &sink_node)); |
| 396 | Status s = RecursivelyHandleOp(*sink_node, num_workers, index, &flib, &graph, |
| 397 | &nodes_to_delete); |
| 398 | |
| 399 | if (!s.ok() && errors::IsNotFound(s)) { |
| 400 | LOG(WARNING) << "Cannot find shardable dataset, adding a shard node at " |
| 401 | << "the end of the dataset instead. This may have performance " |
| 402 | << "implications."; |
| 403 | TF_RETURN_IF_ERROR(AddShardNode(&graph, *sink_node, num_workers, index)); |
| 404 | } else if (!s.ok()) { |
| 405 | return s; |
| 406 | } |
| 407 | |
| 408 | TF_RETURN_IF_ERROR(graph.DeleteNodes(nodes_to_delete)); |
| 409 | |
| 410 | return Status::OK(); |
| 411 | } |
| 412 | |
| 413 | } // anonymous namespace |
| 414 |
no test coverage detected