experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #41676 from samikama:Transactions_part8

Browse Source 
PiperOrigin-RevId: 323944140
Change-Id: I33073187fbd6cfb01e551381cf076455a1a3f841
This commit is contained in:
TensorFlower Gardener 2020-07-29 23:21:51 -07:00
commit d365a03c9a
1 changed files with 211 additions and 140 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

149
tensorflow/python/lib/io/file_io_wrapper.cc
View File

@ -32,25 +32,53 @@ limitations under the License.
#include "tensorflow/python/lib/core/pybind11_absl.h" #include "tensorflow/python/lib/core/pybind11_absl.h"
#include "tensorflow/python/lib/core/pybind11_status.h" #include "tensorflow/python/lib/core/pybind11_status.h"
namespace tensorflow {
struct PyTransactionToken {
TransactionToken* token_;
};
inline TransactionToken* TokenFromPyToken(PyTransactionToken* t) {
return (t ? t->token_ : nullptr);
}
} // namespace tensorflow
namespace { namespace {
namespace py = pybind11; namespace py = pybind11;
PYBIND11_MODULE(_pywrap_file_io, m) { PYBIND11_MODULE(_pywrap_file_io, m) {
m.def("FileExists", [](const std::string& filename) { using tensorflow::PyTransactionToken;
using tensorflow::TransactionToken;
py::class_<PyTransactionToken>(m, "TransactionToken")
.def("__repr__", [](const PyTransactionToken* t) {
if (t->token_) {
return std::string(t->token_->owner->DecodeTransaction(t->token_));
}
return std::string("Invalid token!");
});
m.def(
"FileExists",
[](const std::string& filename, PyTransactionToken* token) {
tensorflow::Status status; tensorflow::Status status;
{ {
py::gil_scoped_release release; py::gil_scoped_release release;
status = tensorflow::Env::Default()->FileExists(filename); status = tensorflow::Env::Default()->FileExists(filename);
} }
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("DeleteFile", [](const std::string& filename) { py::arg("filename"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"DeleteFile",
[](const std::string& filename, PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
tensorflow::Status status = tensorflow::Status status =
tensorflow::Env::Default()->DeleteFile(filename); tensorflow::Env::Default()->DeleteFile(filename);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("ReadFileToString", [](const std::string& filename) { py::arg("filename"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"ReadFileToString",
[](const std::string& filename, PyTransactionToken* token) {
std::string data; std::string data;
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = const auto status =
@ -58,15 +86,22 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
pybind11::gil_scoped_acquire acquire; pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status); tensorflow::MaybeRaiseRegisteredFromStatus(status);
return py::bytes(data); return py::bytes(data);
}); },
m.def("WriteStringToFile", py::arg("filename"), py::arg("token") = (PyTransactionToken*)nullptr);
[](const std::string& filename, tensorflow::StringPiece data) { m.def(
"WriteStringToFile",
[](const std::string& filename, tensorflow::StringPiece data,
PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = const auto status =
WriteStringToFile(tensorflow::Env::Default(), filename, data); WriteStringToFile(tensorflow::Env::Default(), filename, data);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("GetChildren", [](const std::string& dirname) { py::arg("filename"), py::arg("data"),
py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"GetChildren",
[](const std::string& dirname, PyTransactionToken* token) {
std::vector<std::string> results; std::vector<std::string> results;
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = const auto status =
@ -74,8 +109,11 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
pybind11::gil_scoped_acquire acquire; pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status); tensorflow::MaybeRaiseRegisteredFromStatus(status);
return results; return results;
}); },
m.def("GetMatchingFiles", [](const std::string& pattern) { py::arg("dirname"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"GetMatchingFiles",
[](const std::string& pattern, PyTransactionToken* token) {
std::vector<std::string> results; std::vector<std::string> results;
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = const auto status =
@ -83,23 +121,32 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
pybind11::gil_scoped_acquire acquire; pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status); tensorflow::MaybeRaiseRegisteredFromStatus(status);
return results; return results;
}); },
m.def("CreateDir", [](const std::string& dirname) { py::arg("pattern"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"CreateDir",
[](const std::string& dirname, PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = tensorflow::Env::Default()->CreateDir(dirname); const auto status = tensorflow::Env::Default()->CreateDir(dirname);
if (tensorflow::errors::IsAlreadyExists(status)) { if (tensorflow::errors::IsAlreadyExists(status)) {
return; return;
} }
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("RecursivelyCreateDir", [](const std::string& dirname) { py::arg("dirname"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"RecursivelyCreateDir",
[](const std::string& dirname, PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = const auto status =
tensorflow::Env::Default()->RecursivelyCreateDir(dirname); tensorflow::Env::Default()->RecursivelyCreateDir(dirname);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("CopyFile", py::arg("dirname"), py::arg("token") = (PyTransactionToken*)nullptr);
[](const std::string& src, const std::string& target, bool overwrite) { m.def(
"CopyFile",
[](const std::string& src, const std::string& target, bool overwrite,
PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
auto* env = tensorflow::Env::Default(); auto* env = tensorflow::Env::Default();
tensorflow::Status status; tensorflow::Status status;
@ -109,9 +156,13 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
status = env->CopyFile(src, target); status = env->CopyFile(src, target);
} }
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("RenameFile", py::arg("src"), py::arg("target"), py::arg("overwrite"),
[](const std::string& src, const std::string& target, bool overwrite) { py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"RenameFile",
[](const std::string& src, const std::string& target, bool overwrite,
PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
auto* env = tensorflow::Env::Default(); auto* env = tensorflow::Env::Default();
tensorflow::Status status; tensorflow::Status status;
@ -121,20 +172,27 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
status = env->RenameFile(src, target); status = env->RenameFile(src, target);
} }
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("DeleteRecursively", [](const std::string& dirname) { py::arg("src"), py::arg("target"), py::arg("overwrite"),
py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"DeleteRecursively",
[](const std::string& dirname, PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
tensorflow::int64 undeleted_files; tensorflow::int64 undeleted_files;
tensorflow::int64 undeleted_dirs; tensorflow::int64 undeleted_dirs;
auto status = tensorflow::Env::Default()->DeleteRecursively( auto status = tensorflow::Env::Default()->DeleteRecursively(
dirname, &undeleted_files, &undeleted_dirs); dirname, &undeleted_files, &undeleted_dirs);
if (status.ok() && (undeleted_files > 0 || undeleted_dirs > 0)) { if (status.ok() && (undeleted_files > 0 || undeleted_dirs > 0)) {
status = status = tensorflow::errors::PermissionDenied(
tensorflow::errors::PermissionDenied("could not fully delete dir"); "could not fully delete dir");
} }
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}); },
m.def("IsDirectory", [](const std::string& dirname) { py::arg("dirname"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def(
"IsDirectory",
[](const std::string& dirname, PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
const auto status = tensorflow::Env::Default()->IsDirectory(dirname); const auto status = tensorflow::Env::Default()->IsDirectory(dirname);
// FAILED_PRECONDITION response means path exists but isn't a dir. // FAILED_PRECONDITION response means path exists but isn't a dir.
@ -144,7 +202,8 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
return true; return true;
}); },
py::arg("dirname"), py::arg("token") = (PyTransactionToken*)nullptr);
m.def("HasAtomicMove", [](const std::string& path) { m.def("HasAtomicMove", [](const std::string& path) {
py::gil_scoped_release release; py::gil_scoped_release release;
bool has_atomic_move; bool has_atomic_move;
@ -159,19 +218,24 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
.def_readonly("mtime_nsec", &tensorflow::FileStatistics::mtime_nsec) .def_readonly("mtime_nsec", &tensorflow::FileStatistics::mtime_nsec)
.def_readonly("is_directory", &tensorflow::FileStatistics::is_directory); .def_readonly("is_directory", &tensorflow::FileStatistics::is_directory);
m.def("Stat", [](const std::string& filename) { m.def(
"Stat",
[](const std::string& filename, PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
std::unique_ptr<tensorflow::FileStatistics> self( std::unique_ptr<tensorflow::FileStatistics> self(
new tensorflow::FileStatistics); new tensorflow::FileStatistics);
const auto status = tensorflow::Env::Default()->Stat(filename, self.get()); const auto status =
tensorflow::Env::Default()->Stat(filename, self.get());
py::gil_scoped_acquire acquire; py::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status); tensorflow::MaybeRaiseRegisteredFromStatus(status);
return self.release(); return self.release();
}); },
py::arg("filename"), py::arg("token") = (PyTransactionToken*)nullptr);
using tensorflow::WritableFile; using tensorflow::WritableFile;
py::class_<WritableFile>(m, "WritableFile") py::class_<WritableFile>(m, "WritableFile")
.def(py::init([](const std::string& filename, const std::string& mode) { .def(py::init([](const std::string& filename, const std::string& mode,
PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
auto* env = tensorflow::Env::Default(); auto* env = tensorflow::Env::Default();
std::unique_ptr<WritableFile> self; std::unique_ptr<WritableFile> self;
@ -181,7 +245,9 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
py::gil_scoped_acquire acquire; py::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status); tensorflow::MaybeRaiseRegisteredFromStatus(status);
return self.release(); return self.release();
})) }),
py::arg("filename"), py::arg("mode"),
py::arg("token") = (PyTransactionToken*)nullptr)
.def("append", .def("append",
[](WritableFile* self, tensorflow::StringPiece data) { [](WritableFile* self, tensorflow::StringPiece data) {
const auto status = self->Append(data); const auto status = self->Append(data);
@ -209,19 +275,24 @@ PYBIND11_MODULE(_pywrap_file_io, m) {
using tensorflow::io::BufferedInputStream; using tensorflow::io::BufferedInputStream;
py::class_<BufferedInputStream>(m, "BufferedInputStream") py::class_<BufferedInputStream>(m, "BufferedInputStream")
.def(py::init([](const std::string& filename, size_t buffer_size) { .def(py::init([](const std::string& filename, size_t buffer_size,
PyTransactionToken* token) {
py::gil_scoped_release release; py::gil_scoped_release release;
std::unique_ptr<tensorflow::RandomAccessFile> file; std::unique_ptr<tensorflow::RandomAccessFile> file;
const auto status = const auto status =
tensorflow::Env::Default()->NewRandomAccessFile(filename, &file); tensorflow::Env::Default()->NewRandomAccessFile(filename,
&file);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status); tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
std::unique_ptr<tensorflow::io::RandomAccessInputStream> input_stream( std::unique_ptr<tensorflow::io::RandomAccessInputStream>
new tensorflow::io::RandomAccessInputStream(file.release(), input_stream(new tensorflow::io::RandomAccessInputStream(
file.release(),
/*owns_file=*/true)); /*owns_file=*/true));
py::gil_scoped_acquire acquire; py::gil_scoped_acquire acquire;
return new BufferedInputStream(input_stream.release(), buffer_size, return new BufferedInputStream(input_stream.release(), buffer_size,
/*owns_input_stream=*/true); /*owns_input_stream=*/true);
})) }),
py::arg("filename"), py::arg("buffer_size"),
py::arg("token") = (PyTransactionToken*)nullptr)
.def("read", .def("read",
[](BufferedInputStream* self, tensorflow::int64 bytes_to_read) { [](BufferedInputStream* self, tensorflow::int64 bytes_to_read) {
py::gil_scoped_release release; py::gil_scoped_release release;