1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/core/kernels/data/experimental/auto_shard_dataset_op.h"
16
17 #include "tensorflow/core/data/rewrite_utils.h"
18 #include "tensorflow/core/protobuf/rewriter_config.pb.h"
19
20 namespace tensorflow {
21 namespace data {
22 namespace experimental {
23
24 /* static */ constexpr const char* const AutoShardDatasetOp::kAutoShardPolicy;
25 /* static */ constexpr const char* const AutoShardDatasetOp::kDatasetType;
26 /* static */ constexpr const char* const AutoShardDatasetOp::kInputDataset;
27 /* static */ constexpr const char* const AutoShardDatasetOp::kNumWorkers;
28 /* static */ constexpr const char* const AutoShardDatasetOp::kNumReplicas;
29 /* static */ constexpr const char* const AutoShardDatasetOp::kIndex;
30 /* static */ constexpr const char* const AutoShardDatasetOp::kOutputTypes;
31 /* static */ constexpr const char* const AutoShardDatasetOp::kOutputShapes;
32
33 constexpr char kOptimizerName[] = "tf_auto_shard";
34
AutoShardDatasetOp(OpKernelConstruction * ctx)35 AutoShardDatasetOp::AutoShardDatasetOp(OpKernelConstruction* ctx)
36 : UnaryDatasetOpKernel(ctx), auto_shard_policy_(0) {
37 if (ctx->HasAttr(kAutoShardPolicy)) {
38 OP_REQUIRES_OK(ctx, ctx->GetAttr(kAutoShardPolicy, &auto_shard_policy_));
39 }
40 if (ctx->HasAttr(kNumReplicas)) {
41 OP_REQUIRES_OK(ctx, ctx->GetAttr(kNumReplicas, &num_replicas_));
42 }
43 }
44
MakeDataset(OpKernelContext * ctx,DatasetBase * input,DatasetBase ** output)45 void AutoShardDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
46 DatasetBase** output) {
47 int64_t index, num_workers, auto_shard_policy, num_replicas;
48 OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, kNumWorkers, &num_workers));
49 OP_REQUIRES(
50 ctx, num_workers > 0,
51 errors::InvalidArgument("num_workers must be greater than zero."));
52
53 OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, kIndex, &index));
54 OP_REQUIRES(
55 ctx, index >= 0 && index < num_workers,
56 errors::InvalidArgument("index must be between 0 and ", num_workers - 1));
57 auto_shard_policy = auto_shard_policy_;
58 if (input->options().distribute_options().auto_shard_policy() !=
59 AutoShardPolicy::AUTO) {
60 auto_shard_policy =
61 input->options().distribute_options().auto_shard_policy();
62 }
63 num_replicas = num_replicas_;
64
65 auto config_factory = [num_workers, index, auto_shard_policy,
66 num_replicas]() {
67 return CreateConfig(num_workers, index, auto_shard_policy, num_replicas);
68 };
69
70 // We only want to optimize functions for some particular datasets like
71 // FlatMapDataset, InterleaveDataset etc. So we disable generalized
72 // function optimization and explicitly handle function modifications
73 // for those datasets in the rewrite.
74 core::RefCountPtr<DatasetBase> rewritten;
75 OP_REQUIRES_OK(ctx, RewriteDataset(ctx, input, std::move(config_factory),
76 /*record_fingerprint=*/false, &rewritten));
77 *output = rewritten.release();
78 }
79
CreateConfig(int64_t num_workers,int64_t index,int64_t auto_shard_policy,int64_t num_replicas)80 RewriterConfig AutoShardDatasetOp::CreateConfig(int64_t num_workers,
81 int64_t index,
82 int64_t auto_shard_policy,
83 int64_t num_replicas) {
84 RewriterConfig rewriter_config;
85 rewriter_config.set_fail_on_optimizer_errors(true);
86 rewriter_config.set_meta_optimizer_iterations(RewriterConfig::ONE);
87
88 rewriter_config.add_optimizers(kOptimizerName);
89 auto custom_optimizer = rewriter_config.add_custom_optimizers();
90 custom_optimizer->set_name(kOptimizerName);
91
92 const std::array<std::pair<const char* const, int64_t>, 4> attr_pairs = {
93 {{kNumWorkers, num_workers},
94 {kIndex, index},
95 {kAutoShardPolicy, auto_shard_policy},
96 {kNumReplicas, num_replicas}}};
97
98 for (const auto& pair : attr_pairs) {
99 AttrValue attr;
100 attr.set_i(pair.second);
101 (*custom_optimizer->mutable_parameter_map())[pair.first] = attr;
102 }
103
104 return rewriter_config;
105 }
106
107 namespace {
108 REGISTER_KERNEL_BUILDER(Name("AutoShardDataset").Device(DEVICE_CPU),
109 AutoShardDatasetOp);
110 REGISTER_KERNEL_BUILDER(Name("ExperimentalAutoShardDataset").Device(DEVICE_CPU),
111 AutoShardDatasetOp);
112 } // anonymous namespace
113 } // namespace experimental
114 } // namespace data
115 } // namespace tensorflow
116