Builds a string-string table of tensor names to BundleEntryProto (metadata). On construction, attempts to create a directory given by the dirname of "prefix", so "status()" must be checked before calling any member functions. All threads accessing the same BundleWriter must synchronize.
| 106 | // |
| 107 | // All threads accessing the same BundleWriter must synchronize. |
| 108 | class BundleWriter { |
| 109 | public: |
| 110 | struct Options { |
| 111 | Options() {} |
| 112 | // Alignment, in bytes, for tensor data. |
| 113 | // Must be >= 1. The default size of 1 densely packs tensors. |
| 114 | int data_alignment{1}; |
| 115 | }; |
| 116 | BundleWriter(Env* env, StringPiece prefix, |
| 117 | const Options& options = Options()); |
| 118 | |
| 119 | // Adds the tensor "val" under key "key". |
| 120 | // Across calls "key" must be unique but can be added in any order. |
| 121 | Status Add(StringPiece key, const Tensor& val); |
| 122 | |
| 123 | Status AddTensorHeader(StringPiece key, DataType dtype, TensorShape shape); |
| 124 | Status AddTensorHeader(StringPiece key, DataType dtype); |
| 125 | void FillTensorShape(TensorShape shape); |
| 126 | Status AddCompeleteData(char* content, int64 data_bytes_written); |
| 127 | Status AppendSegmentData(char* content, int64 data_bytes_written); |
| 128 | void EndSegmentData(int64 total_bytes_written, int64 end_bytes_written); |
| 129 | // Partitioned variables support. |
| 130 | // A slice of a full tensor is stored in two entries in the metadata table: |
| 131 | // |
| 132 | // full_tensor_key -> BundleEntryProto, describing all stored slices |
| 133 | // of this full tensor. Does not append to the data |
| 134 | // file. |
| 135 | // encoded slice key -> BundleEntryProto, describing one particular slice. |
| 136 | // Appends values of this slice to the data file. |
| 137 | // |
| 138 | // Slices of a full tensor can be added in any order. |
| 139 | // |
| 140 | // If a full tensor has slices placed on N devices and N BundleWriter's are |
| 141 | // concurrently used, the caller must use MergeBundles() to ensure that a |
| 142 | // consistent entry for "full_tensor_key" is produced. |
| 143 | // |
| 144 | // Returns an error if the same slice is added the second time. |
| 145 | Status AddSlice(StringPiece full_tensor_key, |
| 146 | const TensorShape& full_tensor_shape, |
| 147 | const TensorSlice& slice_spec, const Tensor& slice_tensor); |
| 148 | |
| 149 | Status AddSliceHeader( |
| 150 | string tensor_name, const TensorShape& shape, DataType type, bool is_hash, |
| 151 | TensorSliceProto** proto); |
| 152 | |
| 153 | // Finishes the writer and flushes. |
| 154 | Status Finish() TF_MUST_USE_RESULT; |
| 155 | |
| 156 | Status status() const { return status_; } |
| 157 | |
| 158 | private: |
| 159 | Env* const env_; // Not owned. |
| 160 | const Options options_; |
| 161 | const string prefix_; |
| 162 | const string tmp_metadata_path_; |
| 163 | const string tmp_data_path_; |
| 164 | std::unique_ptr<FileOutputBuffer> out_; |
| 165 | int64 size_; // Number of bytes written into out_. |