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 #ifndef TENSORFLOW_CORE_DATA_SERVICE_DATASET_STORE_H_ 16 #define TENSORFLOW_CORE_DATA_SERVICE_DATASET_STORE_H_ 17 18 #include "absl/container/flat_hash_map.h" 19 #include "tensorflow/core/data/service/dispatcher_state.h" 20 #include "tensorflow/core/lib/core/status.h" 21 #include "tensorflow/core/lib/io/record_reader.h" 22 #include "tensorflow/core/lib/io/record_writer.h" 23 #include "tensorflow/core/platform/env.h" 24 25 namespace tensorflow { 26 namespace data { 27 28 // An interface for storing and getting dataset definitions. 29 class DatasetStore { 30 public: 31 virtual ~DatasetStore() = default; 32 33 // Stores the given dataset under the given key. Returns ALREADY_EXISTS if the 34 // key already exists. 35 virtual Status Put(const std::string& key, const DatasetDef& dataset) = 0; 36 // Gets the dataset for the given key, storing the dataset in `dataset_def`. 37 virtual Status Get(const std::string& key, 38 std::shared_ptr<const DatasetDef>& dataset_def) = 0; 39 }; 40 41 // Dataset store which reads and writes datasets within a directory. 42 // The dataset with key `key` is stored at the path "datasets_dir/key". 43 class FileSystemDatasetStore : public DatasetStore { 44 public: 45 explicit FileSystemDatasetStore(const std::string& datasets_dir); 46 FileSystemDatasetStore(const FileSystemDatasetStore&) = delete; 47 FileSystemDatasetStore& operator=(const FileSystemDatasetStore&) = delete; 48 49 Status Put(const std::string& key, const DatasetDef& dataset) override; 50 Status Get(const std::string& key, 51 std::shared_ptr<const DatasetDef>& dataset_def) override; 52 53 private: 54 const std::string datasets_dir_; 55 }; 56 57 // DatasetStore which stores all datasets in memory. This is useful when the 58 // dispatcher doesn't have a work directory configured. 59 class MemoryDatasetStore : public DatasetStore { 60 public: 61 MemoryDatasetStore(); 62 MemoryDatasetStore(const MemoryDatasetStore&) = delete; 63 MemoryDatasetStore& operator=(const MemoryDatasetStore&) = delete; 64 65 Status Put(const std::string& key, const DatasetDef& dataset) override; 66 Status Get(const std::string& key, 67 std::shared_ptr<const DatasetDef>& dataset_def) override; 68 69 private: 70 // Mapping from key to dataset definition. 71 absl::flat_hash_map<std::string, std::shared_ptr<const DatasetDef>> datasets_; 72 }; 73 74 } // namespace data 75 } // namespace tensorflow 76 77 #endif // TENSORFLOW_CORE_DATA_SERVICE_DATASET_STORE_H_ 78