1 /* Copyright 2020 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 16 #ifndef TENSORFLOW_CORE_KERNELS_DATA_EXPERIMENTAL_DATA_SERVICE_OPS_H_ 17 #define TENSORFLOW_CORE_KERNELS_DATA_EXPERIMENTAL_DATA_SERVICE_OPS_H_ 18 19 #include <string> 20 21 #include "tensorflow/core/framework/dataset.h" 22 #include "tensorflow/core/framework/op_kernel.h" 23 24 namespace tensorflow { 25 namespace data { 26 27 // Registers a dataset with the tf.data service. 28 // 29 // The address and protocol inputs are used to connect to the dispatcher. 30 // The external state policy attribute determines whether to ignore, warn, or 31 // error out when the dataset contains external state. 32 // The op produces a dataset id for identifying the registered dataset. 33 class RegisterDatasetOp : public OpKernel { 34 public: 35 static constexpr const char* const kAddress = "address"; 36 static constexpr const char* const kProtocol = "protocol"; 37 static constexpr const char* const kExternalStatePolicy = 38 "external_state_policy"; 39 static constexpr const char* const kElementSpec = "element_spec"; 40 static constexpr const char* const kMetadata = "metadata"; 41 static constexpr const char* const kRequestedDatasetId = 42 "requested_dataset_id"; 43 static constexpr const char* const kTimeoutMs = "timeout_ms"; 44 45 explicit RegisterDatasetOp(OpKernelConstruction* ctx); 46 47 void Compute(OpKernelContext* ctx) override; 48 49 private: 50 int op_version_; 51 SerializationContext::ExternalStatePolicy external_state_policy_; 52 std::string element_spec_; 53 std::string serialized_metadata_; 54 std::string requested_dataset_id_; 55 }; 56 57 } // namespace data 58 } // namespace tensorflow 59 #endif // TENSORFLOW_CORE_KERNELS_DATA_EXPERIMENTAL_DATA_SERVICE_OPS_H_ 60