Added transactional checkpointing example and possible implementation changes

Author: samikama

rfcs/20200505-transactional-fs.md

@@ -697,7 +697,7 @@ Status MergeFiles(const string& fname, const string& dirname,
 
 ## Final design proposal
 
-During the review meeting it has been decided to merge wrapped filesystem approach and stateful token approach. Then the final proposal is as shown below. The `WrappedFileSystem` and `Filesystem` classes described above are to be extended with three new methods,
+During the review meeting it has been decided to merge wrapped filesystem approach and stateful token approach. Then the final proposal is as shown below. The `WrappedFileSystem` and `Filesystem` classes described above are to be extended with four new methods,
 
 ```cpp
 class WrappedFileSystem : public Filesystem {
@@ -707,6 +707,11 @@ class WrappedFileSystem : public Filesystem {
     return fs_->AddToTransaction(uri, (token ? token : token_));
   }
 
+  virtual Status GetTransactionForPath(const string& uri,
+                                            TransactionToken*& token) {
+    return fs_->GetTransactionForPath(uri, token);
+  }
+
   virtual Status GetTokenOrStartTransaction(const string& uri,
                                             TransactionToken*& token) {
     return fs_->GetTokenOrStartTransaction(uri, token);
@@ -726,18 +731,120 @@ Then the current C API's `TF_FilesystemOps` table is to be extended by 6 new fun
 struct TF_FilesystemOps{
   // Existing members are not modified
   // Transaction management
-  void (*const StartTransaction)(TF_Filesystem*, TransactionToken**);
-  void (*const EndTransaction)(TF_Filesystem*, TransactionToken*);
-  void (*const AddToTransaction)(TF_Filesystem* fs, const char* file_name, TransactionToken** token);
-  void (*const GetTransactionTokenForFile)(TF_Filesystem* fs, const char* file_name, TransactionToken** token);
-  void (*const GetOrStartTransactionTokenForFile)(TF_Filesystem* fs, const char* file_name, TransactionToken** token);
+  void (*start_transaction)(TF_Filesystem*, TF_TransactionToken**, TF_Status );
+  void (*end_transaction)(TF_Filesystem*, TF_TransactionToken*);
+  void (*add_to_transaction)(TF_Filesystem* fs, const char* path, TF_TransactionToken* token);
+  void (*get_transaction_for_path)(TF_Filesystem* fs, const char* path, TF_TransactionToken** token);
+  void (*get_or_start_transaction_for_path)(TF_Filesystem* fs, const char* path, TF_TransactionToken** token);
   // Optional Transaction Debugging
-  void (*const DecodeTransactionToken)(const TF_Filesystem*, const TransactionToken*, char**);
+  char* (*decode_transaction_token)(const TF_Filesystem*, const TF_TransactionToken*);
 };
 ```
 
-The new functions will be null pointers until respective plugins implement them. `ModularFilesystem` implementation will check whether a plugin implements the transactions and will ignore the transaction if it is not implemented, possibly after producing a log message, thus falling back to current transactionless state. Since these function pointers will be added after existing pointers, already compiled plugins will keep functioning and they can be gradually start supporting transactions. Any filesystem plugin that start supporting transaction will be used by the framework.
+The new functions will be null pointers until respective plugins implement them. `ModularFilesystem` implementation will check whether a plugin implements the transactions and will ignore the transaction if it is not implemented, possibly after producing a log message, thus falling back to
+current transactionless state. Since these function pointers will be added after existing pointers, already compiled plugins will keep functioning and they can be gradually start supporting transactions. Any filesystem plugin that start supporting transaction will be used by the framework.
+
+## Example use
+
+With these final modifications, there is no need to carry transaction tokens through different compilation units or ops in the graph. For example current checkpointing logic involve adding one or more `SaveV2` ops and a `MergeV2Checkpoints` op to the graph. In order to prevent
+corrupt checkpoints in case of errors and optimize i/o, SaveV2 ops write their outputs to the temporary directories, given as constant input arguments, and MergeV2Checkpoints op is given the names of the temporary save files generated by all SaveV2 ops which reads and merges
+them into a final file and then deletes temporary files. Just by adding a few lines to SaveV2 to and MergeV2Checkpoints ops, these operations can be completed transactionally. In this case overview of the operations would be,
+
+- SaveV2 op uses `Env::GetTokenOrStartTransaction(base_dir)` call to start a new transaction or get an existing transaction on the base output directory.
+- SaveV2 op adds files it generates to the transaction.
+- MergeV2Checkpoint op uses `Env::GetTransactionTokenForPath(base_dir)` to get the transaction started by SaveV2 ops and adds its output to the same transaction.
+- MergeV2Checkpoint op removes intermediate files and calls `EndTransaction()` to finalize the transaction.
+
+Then a transactional file system may operate in a cache and move files to the final destination only at the end of the transaction which would ensure that the files in the checkpoint directory are consistent. The
+implementation of how this is achieved may be different for each filesystem but the end result should be consistent.
+
+An example implementation of this could be as shown in the diff set below.
+```diff
+--- a/tensorflow/core/kernels/save_restore_v2_ops.cc
++++ b/tensorflow/core/kernels/save_restore_v2_ops.cc
+@@ -237,6 +237,17 @@ class MergeV2Checkpoints : public OpKernel {
+         gtl::ArraySlice<tstring>(checkpoint_prefixes.flat<tstring>());
+     Env* env = Env::Default();
+     const string& merged_prefix = destination_prefix.scalar<tstring>()();
++    auto token_deleter = [env](TransactionToken* token) {
++      if (token) {
++        env->EndTransaction(token);
++      }
++    };
++    TransactionToken* token = nullptr;
++    env->GetTokenOrStartTransaction(string(io::Dirname(input_prefixes[0])), &token);
++    auto token_scope =
++        std::unique_ptr<TransactionToken, decltype(token_deleter)>(
++            token, token_deleter);
++    env->AddToTransaction(string(io::Dirname(merged_prefix)), token);
+     OP_REQUIRES_OK(
+         context, tensorflow::MergeBundles(env, input_prefixes, merged_prefix));
+ 
+--- a/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
++++ b/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
+@@ -421,9 +421,10 @@ BundleWriter::BundleWriter(Env* env, StringPiece prefix, const Options& options)
+   if (!status_.ok() && !errors::IsAlreadyExists(status_)) {
+     return;
+   }
+-
++  TransactionToken* token=nullptr;
++  env->GetTokenOrStartTransaction(string(io::Dirname(prefix_)),&token).IgnoreError();
+   std::unique_ptr<WritableFile> wrapper;
+-  status_ = env_->NewWritableFile(data_path_, &wrapper);
++  status_ = env_->NewWritableFile(data_path_, &wrapper,token);
+   if (!status_.ok()) return;
+   out_ = std::unique_ptr<FileOutputBuffer>(
+       new FileOutputBuffer(wrapper.release(), 8 << 20 /* 8MB write buffer */));
+@@ -527,7 +528,9 @@ Status BundleWriter::Finish() {
+   if (!status_.ok()) return status_;
+   // Build key -> BundleEntryProto table.
+   std::unique_ptr<WritableFile> file;
+-  status_ = env_->NewWritableFile(metadata_path_, &file);
++  TransactionToken* token;
++  env_->GetTokenOrStartTransaction(string(io::Dirname(prefix_)),&token).IgnoreError();
++  status_ = env_->NewWritableFile(metadata_path_, &file, token);
+   if (!status_.ok()) return status_;
+   {
+     // N.B.: the default use of Snappy compression may not be supported on all
+@@ -554,7 +557,7 @@ Status BundleWriter::Finish() {
+   }
+   status_.Update(file->Close());
+   if (!status_.ok()) {
+-    Env::Default()->DeleteFile(metadata_path_).IgnoreError();
++    Env::Default()->DeleteFile(metadata_path_, token).IgnoreError();
+     return status_;
+   } else if (use_temp_file_) {
+     status_ = Env::Default()->RenameFile(metadata_path_, MetaFilename(prefix_));
+@@ -590,6 +593,8 @@ static Status MergeOneBundle(Env* env, StringPiece prefix,
+                              MergeState* merge_state) {
+   VLOG(1) << "Merging bundle:" << prefix;
+   const string filename = MetaFilename(prefix);
++  TransactionToken* token=nullptr;
++  env->GetTokenOrStartTransaction(string(filename),&token).IgnoreError();
+   uint64 file_size;
+   TF_RETURN_IF_ERROR(env->GetFileSize(filename, &file_size));
+   std::unique_ptr<RandomAccessFile> file;
+@@ -690,7 +695,9 @@ Status MergeBundles(Env* env, gtl::ArraySlice<tstring> prefixes,
+   // Merges all metadata tables.
+   // TODO(zhifengc): KeyValue sorter if it becomes too big.
+   MergeState merge;
+-  Status status = env->CreateDir(string(io::Dirname(merged_prefix)));
++  TransactionToken *token=nullptr;
++  env->GetTokenOrStartTransaction(string(io::Dirname(prefixes[0])),&token);
++  Status status = env->CreateDir(string(io::Dirname(merged_prefix)),token);
+   if (!status.ok() && !errors::IsAlreadyExists(status)) return status;
+   for (int i = 0; i < prefixes.size(); ++i) {
+     TF_RETURN_IF_ERROR(MergeOneBundle(env, prefixes[i], &merge));
+@@ -708,7 +715,7 @@ Status MergeBundles(Env* env, gtl::ArraySlice<tstring> prefixes,
+   // Writes the final metadata table under the merged prefix.
+   std::unique_ptr<WritableFile> merged_metadata;
+   TF_RETURN_IF_ERROR(
+-      env->NewWritableFile(MetaFilename(merged_prefix), &merged_metadata));
++      env->NewWritableFile(MetaFilename(merged_prefix), &merged_metadata,token));
+   {
+     table::TableBuilder builder(TableBuilderOptions(), merged_metadata.get());
+     // Header entry.
 
-With these final modifications, there is no need to carry transaction tokens through different compilation units or ops in the graph. For example checkpointing logic can be implemented in an atomic-like way by simply modifying the save and merge ops to use same transaction using accessed file names and directories.
+```
 
 [filesystem_plugin]: https://github.com/tensorflow/community/blob/master/rfcs/20190506-filesystem-plugin-modular-tensorflow.md