wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/lib/BUILD
+21
View File
@@ -0,0 +1,21 @@
# python/lib package
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
visibility = [
"//tensorflow:__subpackages__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = visibility,
licenses = ["notice"],
)
py_library(
name = "__init__",
srcs = ["__init__.py"],
strict_deps = True,
visibility = ["//tensorflow:__subpackages__"],
deps = [],
)
tensorflow/python/lib/__init__.py
View File
tensorflow/python/lib/core/BUILD
+313
View File
@@ -0,0 +1,313 @@
# python/lib/core package
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "tf_external_workspace_visible", "tf_python_pybind_extension")
visibility = [
"//third_party/cloud_tpu/convergence_tools:__subpackages__",
"//third_party/mlperf:__subpackages__",
"//tensorflow:internal",
"//tensorflow/lite/toco/python:__pkg__",
"//tensorflow_models:__subpackages__",
"//tensorflow_model_optimization:__subpackages__",
"//third_party/py/cleverhans:__subpackages__",
"//third_party/py/launchpad:__subpackages__",
"//third_party/py/reverb:__subpackages__",
"//third_party/py/neural_structured_learning:__subpackages__",
"//third_party/py/tensorflow_examples:__subpackages__",
"//third_party/py/tf_agents:__subpackages__", # For benchmarks.
"//third_party/py/tf_slim:__subpackages__",
"//third_party/py/tensorflow_docs:__subpackages__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = visibility,
licenses = ["notice"],
)
cc_library(
name = "ndarray_tensor_bridge",
srcs = ["ndarray_tensor_bridge.cc"],
hdrs = ["ndarray_tensor_bridge.h"],
visibility = tf_external_workspace_visible(visibility + [
"//tensorflow:ndarray_tensor_allow_list",
]),
deps = [
":py_util",
"//tensorflow/c:c_api_no_xla",
"//tensorflow/c:tf_datatype_hdrs",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@xla//xla/tsl/python/lib/core:ml_dtypes_lib",
"@xla//xla/tsl/python/lib/core:numpy",
],
)
cc_library(
name = "py_exception_registry",
srcs = ["py_exception_registry.cc"],
hdrs = ["py_exception_registry.h"],
deps = [
"//tensorflow/c:tf_status_headers",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@xla//third_party/python_runtime:headers",
],
alwayslink = 1,
)
cc_library(
name = "pybind11_absl",
hdrs = ["pybind11_absl.h"],
features = ["-parse_headers"],
visibility = tf_external_workspace_visible(visibility),
deps = [
"//tensorflow/core/platform:stringpiece",
"@pybind11",
],
)
cc_library(
name = "pybind11_lib",
hdrs = ["pybind11_lib.h"],
compatible_with = get_compatible_with_portable(),
features = ["-parse_headers"],
visibility = tf_external_workspace_visible(visibility),
deps = [
"@pybind11",
],
)
cc_library(
name = "pybind11_status_headers",
hdrs = [
"py_exception_registry.h",
"pybind11_status.h",
"//tensorflow/c:headers",
"//tensorflow/c:tf_status_internal_headers",
"//tensorflow/c/eager:headers",
],
features = [
"-parse_headers",
],
visibility = tf_external_workspace_visible(visibility),
deps = [
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"@com_google_absl//absl/status",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
cc_library(
name = "pybind11_status",
hdrs = [
"py_exception_registry.h",
"pybind11_status.h",
],
features = ["-parse_headers"],
visibility = tf_external_workspace_visible(visibility),
deps = [
":pybind11_status_headers",
"//tensorflow/core:lib",
"@com_google_absl//absl/status",
],
)
cc_library(
name = "pybind11_proto",
hdrs = ["pybind11_proto.h"],
features = ["-parse_headers"],
visibility = tf_external_workspace_visible(visibility),
deps = [
"@com_google_absl//absl/strings",
"@pybind11",
],
)
filegroup(
name = "py_exception_registry_hdr",
srcs = [
"py_exception_registry.h",
],
visibility = ["//visibility:public"],
)
filegroup(
name = "ndarray_tensor_hdr",
srcs = ["ndarray_tensor.h"],
)
filegroup(
name = "basic_hdrs",
srcs = [
"ndarray_tensor.h",
"ndarray_tensor_bridge.h",
"py_exception_registry.h",
"pybind11_status.h",
"safe_pyobject_ptr.h",
],
)
cc_library(
name = "py_func_lib",
srcs = ["py_func.cc"],
hdrs = ["py_func.h"],
visibility = ["//visibility:public"],
deps = [
":ndarray_tensor",
":ndarray_tensor_bridge",
":py_util",
"//tensorflow/c:safe_ptr",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:tfe_context_internal",
"//tensorflow/c/eager:tfe_tensorhandle_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:script_ops_op_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:tensor_handle",
"//tensorflow/python/eager:pywrap_tfe_lib",
"//third_party/py/numpy:headers",
"@xla//third_party/python_runtime:headers",
"@xla//xla/tsl/python/lib/core:numpy",
],
alwayslink = 1,
)
tf_python_pybind_extension(
name = "_pywrap_py_func",
srcs = ["py_func_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_py_func.pyi",
],
deps = [
"//tensorflow/python:py_func_headers_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
cc_library(
name = "safe_pyobject_ptr",
srcs = ["safe_pyobject_ptr.cc"],
hdrs = ["safe_pyobject_ptr.h"],
deps = [
"@xla//third_party/python_runtime:headers",
],
)
filegroup(
name = "safe_pyobject_ptr_required_hdrs",
srcs = ["safe_pyobject_ptr.h"],
)
cc_library(
name = "ndarray_tensor_headers",
hdrs = [
"ndarray_tensor.h",
"ndarray_tensor_bridge.h",
"safe_pyobject_ptr.h",
"//tensorflow/c:headers",
"//tensorflow/c:safe_ptr_hdr",
"//tensorflow/c/eager:headers",
"@xla//xla/tsl/python/lib/core:numpy_hdr",
],
features = [
"-parse_headers",
],
visibility = tf_external_workspace_visible(visibility + [
"//tensorflow:ndarray_tensor_allow_list",
]),
deps = [
"//tensorflow/c:pywrap_required_hdrs",
"//tensorflow/c:tf_status_headers",
"//tensorflow/core:framework_internal_headers_lib",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//third_party/py/numpy:headers",
"@xla//third_party/python_runtime:headers",
"@xla//xla/tsl/python/lib/core:numpy",
],
)
cc_library(
name = "ndarray_tensor",
srcs = ["ndarray_tensor.cc"],
hdrs = ["ndarray_tensor.h"],
visibility = tf_external_workspace_visible(visibility + [
"//tensorflow:ndarray_tensor_allow_list",
]),
deps = [
":ndarray_tensor_bridge",
"//tensorflow/c:c_api_internal",
"//tensorflow/c:safe_ptr",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/c/eager:tfe_context_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/python/lib/core:safe_pyobject_ptr",
"@xla//xla/tsl/python/lib/core:ml_dtypes_lib",
"@xla//xla/tsl/python/lib/core:numpy",
],
)
cc_library(
name = "py_seq_tensor",
srcs = ["py_seq_tensor.cc"],
hdrs = ["py_seq_tensor.h"],
features = ["-parse_headers"],
deps = [
":ndarray_tensor",
":ndarray_tensor_bridge",
":py_util",
":safe_pyobject_ptr",
"//tensorflow/c:safe_ptr",
"//tensorflow/c:tensor_interface",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/c/eager:c_api_internal",
"//tensorflow/c/eager:tfe_context_internal",
"//tensorflow/c/eager:tfe_tensorhandle_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@xla//third_party/python_runtime:headers", # build_cleaner: keep; DNR: b/35864863
"@xla//xla/tsl/python/lib/core:numpy",
],
)
cc_library(
name = "py_util",
srcs = ["py_util.cc"],
hdrs = ["py_util.h"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:script_ops_op_lib",
"//tensorflow/core/platform:logging",
"@xla//third_party/python_runtime:headers",
],
)
# copybara:uncomment_begin(google-only)
# cc_library(
# name = "pyclif_tensor",
# srcs = ["pyclif_tensor.cc"],
# hdrs = ["pyclif_tensor.h"],
# visibility = ["//visibility:public"],
# deps = [
# ":ndarray_tensor",
# ":ndarray_tensor_bridge",
# ":safe_pyobject_ptr",
# "@com_google_absl//absl/log:check",
# "@com_google_absl//absl/status",
# "//third_party/clif/python:clif",
# "@xla//third_party/python_runtime:headers",
# "@xla//xla/tsl/python/lib/core:numpy",
# "//tensorflow/core:framework",
# ],
# )
# copybara:uncomment_end
tensorflow/python/lib/core/__init__.py
View File
tensorflow/python/lib/core/_pywrap_py_func.pyi
+16
View File
@@ -0,0 +1,16 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
def initialize_py_trampoline(arg0: object) -> None: ...
tensorflow/python/lib/core/ndarray_tensor.cc
+638
View File
@@ -0,0 +1,638 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// Must be included first
// clang-format off
#include "tensorflow/c/tf_datatype.h"
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
// clang-format on
#include "tensorflow/python/lib/core/ndarray_tensor.h"
#include <cstring> // NOLINT
#include <optional> // NOLINT
#include "tensorflow/c/eager/tfe_context_internal.h"
#include "tensorflow/c/tf_tensor_internal.h"
#include "xla/tsl/python/lib/core/ml_dtypes.h"
#include "tensorflow/core/lib/core/coding.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/util/port.h"
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
namespace tensorflow {
namespace {
char const* numpy_type_name(int numpy_type) {
switch (numpy_type) {
#define TYPE_CASE(s) \
case s: \
return #s
TYPE_CASE(NPY_BOOL);
TYPE_CASE(NPY_BYTE);
TYPE_CASE(NPY_UBYTE);
TYPE_CASE(NPY_SHORT);
TYPE_CASE(NPY_USHORT);
TYPE_CASE(NPY_INT);
TYPE_CASE(NPY_UINT);
TYPE_CASE(NPY_LONG);
TYPE_CASE(NPY_ULONG);
TYPE_CASE(NPY_LONGLONG);
TYPE_CASE(NPY_ULONGLONG);
TYPE_CASE(NPY_FLOAT);
TYPE_CASE(NPY_DOUBLE);
TYPE_CASE(NPY_LONGDOUBLE);
TYPE_CASE(NPY_CFLOAT);
TYPE_CASE(NPY_CDOUBLE);
TYPE_CASE(NPY_CLONGDOUBLE);
TYPE_CASE(NPY_OBJECT);
TYPE_CASE(NPY_STRING);
TYPE_CASE(NPY_UNICODE);
TYPE_CASE(NPY_VOID);
TYPE_CASE(NPY_DATETIME);
TYPE_CASE(NPY_TIMEDELTA);
TYPE_CASE(NPY_HALF);
TYPE_CASE(NPY_NOTYPE);
TYPE_CASE(NPY_CHAR);
TYPE_CASE(NPY_USERDEF);
default:
return "not a numpy type";
}
}
#if NPY_ABI_VERSION < 0x02000000
#define PyDataType_FIELDS(descr) ((descr)->fields)
#endif // NPY_ABI_VERSION < 0x02000000
absl::Status PyArrayDescr_to_TF_DataType(PyArray_Descr* descr,
TF_DataType* out_tf_datatype) {
PyObject* key;
PyObject* value;
Py_ssize_t pos = 0;
// Return an error if the fields attribute is null.
// Occurs with an improper conversion attempt to resource.
if (PyDataType_FIELDS(descr) == nullptr) {
return absl::InternalError("Unexpected numpy data type");
}
if (PyDict_Next(PyDataType_FIELDS(descr), &pos, &key, &value)) {
// In Python 3, the keys of numpy custom struct types are unicode, unlike
// Python 2, where the keys are bytes.
const char* key_string =
PyBytes_Check(key) ? PyBytes_AsString(key)
: PyBytes_AsString(PyUnicode_AsASCIIString(key));
if (!key_string) {
return absl::InternalError("Corrupt numpy type descriptor");
}
std::string key = key_string;
// The typenames here should match the field names in the custom struct
// types constructed in test_util.py.
// TODO(mrry,keveman): Investigate Numpy type registration to replace this
// hard-coding of names.
if (key == "quint8") {
*out_tf_datatype = TF_QUINT8;
} else if (key == "qint8") {
*out_tf_datatype = TF_QINT8;
} else if (key == "qint16") {
*out_tf_datatype = TF_QINT16;
} else if (key == "quint16") {
*out_tf_datatype = TF_QUINT16;
} else if (key == "qint32") {
*out_tf_datatype = TF_QINT32;
} else if (key == "resource") {
*out_tf_datatype = TF_RESOURCE;
} else {
return absl::InternalError("Unsupported numpy data type");
}
return absl::OkStatus();
}
return absl::InternalError("Unsupported numpy data type");
}
absl::Status PyArray_TYPE_to_TF_DataType(PyArrayObject* array,
TF_DataType* out_tf_datatype) {
const tsl::ml_dtypes::NumpyDtypes& custom_dtypes =
tsl::ml_dtypes::GetNumpyDtypes();
int pyarray_type = PyArray_TYPE(array);
PyArray_Descr* descr = PyArray_DESCR(array);
switch (pyarray_type) {
case NPY_FLOAT16:
*out_tf_datatype = TF_HALF;
break;
case NPY_FLOAT32:
*out_tf_datatype = TF_FLOAT;
break;
case NPY_FLOAT64:
*out_tf_datatype = TF_DOUBLE;
break;
case NPY_INT32:
*out_tf_datatype = TF_INT32;
break;
case NPY_UINT8:
*out_tf_datatype = TF_UINT8;
break;
case NPY_UINT16:
*out_tf_datatype = TF_UINT16;
break;
case NPY_UINT32:
*out_tf_datatype = TF_UINT32;
break;
case NPY_UINT64:
*out_tf_datatype = TF_UINT64;
break;
case NPY_INT8:
*out_tf_datatype = TF_INT8;
break;
case NPY_INT16:
*out_tf_datatype = TF_INT16;
break;
case NPY_INT64:
*out_tf_datatype = TF_INT64;
break;
case NPY_BOOL:
*out_tf_datatype = TF_BOOL;
break;
case NPY_COMPLEX64:
*out_tf_datatype = TF_COMPLEX64;
break;
case NPY_COMPLEX128:
*out_tf_datatype = TF_COMPLEX128;
break;
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE:
*out_tf_datatype = TF_STRING;
break;
case NPY_VOID:
// Quantized types are currently represented as custom struct types.
// PyArray_TYPE returns NPY_VOID for structs, and we should look into
// descr to derive the actual type.
// Direct feeds of certain types of ResourceHandles are represented as a
// custom struct type.
return PyArrayDescr_to_TF_DataType(descr, out_tf_datatype);
default:
if (pyarray_type == custom_dtypes.bfloat16) {
*out_tf_datatype = TF_BFLOAT16;
break;
} else if (pyarray_type == NPY_ULONGLONG) {
// NPY_ULONGLONG is equivalent to NPY_UINT64, while their enum values
// might be different on certain platforms.
*out_tf_datatype = TF_UINT64;
break;
} else if (pyarray_type == NPY_LONGLONG) {
// NPY_LONGLONG is equivalent to NPY_INT64, while their enum values
// might be different on certain platforms.
*out_tf_datatype = TF_INT64;
break;
} else if (pyarray_type == NPY_INT) {
// NPY_INT is equivalent to NPY_INT32, while their enum values might be
// different on certain platforms.
*out_tf_datatype = TF_INT32;
break;
} else if (pyarray_type == NPY_UINT) {
// NPY_UINT is equivalent to NPY_UINT32, while their enum values might
// be different on certain platforms.
*out_tf_datatype = TF_UINT32;
break;
} else if (pyarray_type == custom_dtypes.float8_e5m2) {
*out_tf_datatype = TF_FLOAT8_E5M2;
break;
} else if (pyarray_type == custom_dtypes.float8_e4m3fn) {
*out_tf_datatype = TF_FLOAT8_E4M3FN;
break;
} else if (pyarray_type == custom_dtypes.float8_e4m3fnuz) {
*out_tf_datatype = TF_FLOAT8_E4M3FNUZ;
break;
} else if (pyarray_type == custom_dtypes.float8_e4m3b11fnuz) {
*out_tf_datatype = TF_FLOAT8_E4M3B11FNUZ;
break;
} else if (pyarray_type == custom_dtypes.float8_e5m2fnuz) {
*out_tf_datatype = TF_FLOAT8_E5M2FNUZ;
break;
} else if (pyarray_type == custom_dtypes.float4_e2m1fn) {
*out_tf_datatype = TF_FLOAT4_E2M1FN;
break;
} else if (pyarray_type == custom_dtypes.int4) {
*out_tf_datatype = TF_INT4;
break;
} else if (pyarray_type == custom_dtypes.uint4) {
*out_tf_datatype = TF_UINT4;
break;
} else if (pyarray_type == custom_dtypes.int2) {
*out_tf_datatype = TF_INT2;
break;
} else if (pyarray_type == custom_dtypes.uint2) {
*out_tf_datatype = TF_UINT2;
break;
}
return absl::InternalError(absl::StrCat("Unsupported numpy type: ",
numpy_type_name(pyarray_type)));
}
return absl::OkStatus();
}
absl::Status PyObjectToString(PyObject* obj, const char** ptr, Py_ssize_t* len,
PyObject** ptr_owner) {
*ptr_owner = nullptr;
if (PyBytes_Check(obj)) {
char* buf;
if (PyBytes_AsStringAndSize(obj, &buf, len) != 0) {
return absl::InternalError("Unable to get element as bytes.");
}
*ptr = buf;
return absl::OkStatus();
} else if (PyUnicode_Check(obj)) {
#if (PY_MAJOR_VERSION > 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION >= 3))
*ptr = PyUnicode_AsUTF8AndSize(obj, len);
if (*ptr != nullptr) return absl::OkStatus();
#else
PyObject* utemp = PyUnicode_AsUTF8String(obj);
char* buf;
if (utemp != nullptr && PyBytes_AsStringAndSize(utemp, &buf, len) != -1) {
*ptr = buf;
*ptr_owner = utemp;
return OkStatus();
}
Py_XDECREF(utemp);
#endif
return absl::InternalError("Unable to convert element to UTF-8");
} else {
return absl::InternalError(
absl::StrCat("Unsupported object type ", obj->ob_type->tp_name));
}
}
// Iterate over the string array 'array', extract the ptr and len of each string
// element and call f(ptr, len).
template <typename F>
absl::Status PyBytesArrayMap(PyArrayObject* array, F f) {
Safe_PyObjectPtr iter = tensorflow::make_safe(
PyArray_IterNew(reinterpret_cast<PyObject*>(array)));
while (PyArray_ITER_NOTDONE(iter.get())) {
auto item = tensorflow::make_safe(PyArray_GETITEM(
array, static_cast<char*>(PyArray_ITER_DATA(iter.get()))));
if (!item) {
return absl::InternalError(
"Unable to get element from the feed - no item.");
}
Py_ssize_t len;
const char* ptr;
PyObject* ptr_owner = nullptr;
TF_RETURN_IF_ERROR(PyObjectToString(item.get(), &ptr, &len, &ptr_owner));
f(ptr, len);
Py_XDECREF(ptr_owner);
PyArray_ITER_NEXT(iter.get());
}
return absl::OkStatus();
}
// Encode the strings in 'array' into a contiguous buffer and return the base of
// the buffer. The caller takes ownership of the buffer.
absl::Status EncodePyBytesArray(PyArrayObject* array, int64_t nelems,
size_t* size, void** buffer) {
// Encode all strings.
*size = nelems * sizeof(tensorflow::tstring);
std::unique_ptr<tensorflow::tstring[]> base_ptr(
new tensorflow::tstring[nelems]);
tensorflow::tstring* dst = base_ptr.get();
TF_RETURN_IF_ERROR(
PyBytesArrayMap(array, [&dst](const char* ptr, Py_ssize_t len) {
dst->assign(ptr, len);
dst++;
}));
*buffer = base_ptr.release();
return absl::OkStatus();
}
absl::Status CopyTF_TensorStringsToPyArray(const TF_Tensor* src,
uint64_t nelems,
PyArrayObject* dst) {
const void* tensor_data = TF_TensorData(src);
DCHECK(tensor_data != nullptr);
DCHECK_EQ(TF_STRING, TF_TensorType(src));
const tstring* tstr = static_cast<const tstring*>(tensor_data);
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
TF_NewStatus(), TF_DeleteStatus);
auto iter = make_safe(PyArray_IterNew(reinterpret_cast<PyObject*>(dst)));
for (int64_t i = 0; i < static_cast<int64_t>(nelems); ++i) {
const tstring& tstr_i = tstr[i];
auto py_string =
make_safe(PyBytes_FromStringAndSize(tstr_i.data(), tstr_i.size()));
if (py_string == nullptr) {
return absl::InternalError(absl::StrCat(
"failed to create a python byte array when converting element #", i,
" of a TF_STRING tensor to a numpy ndarray"));
}
if (PyArray_SETITEM(dst, static_cast<char*>(PyArray_ITER_DATA(iter.get())),
py_string.get()) != 0) {
return absl::InternalError(
absl::StrCat("Error settings element #", i, " in the numpy ndarray"));
}
PyArray_ITER_NEXT(iter.get());
}
return absl::OkStatus();
}
// Determine the dimensions of a numpy ndarray to be created to represent an
// output Tensor.
absl::Status GetPyArrayDimensionsForTensor(
const TF_Tensor* tensor, absl::InlinedVector<npy_intp, 4UL>* dims,
int64_t* nelems) {
dims->clear();
const int ndims = TF_NumDims(tensor);
if (TF_TensorType(tensor) == TF_RESOURCE) {
if (ndims != 0) {
return absl::InvalidArgumentError(
"Fetching of non-scalar resource tensors is not supported.");
}
dims->push_back(TF_TensorByteSize(tensor));
*nelems = dims->back();
} else {
*nelems = 1;
for (int i = 0; i < ndims; ++i) {
dims->push_back(TF_Dim(tensor, i));
*nelems *= dims->back();
}
}
return absl::OkStatus();
}
// Determine the type description (PyArray_Descr) of a numpy ndarray to be
// created to represent an output Tensor.
absl::Status GetPyArrayDescrForTensor(const TF_Tensor* tensor,
PyArray_Descr** descr) {
if (TF_TensorType(tensor) == TF_RESOURCE) {
PyObject* field = PyTuple_New(3);
#if PY_MAJOR_VERSION < 3
PyTuple_SetItem(field, 0, PyBytes_FromString("resource"));
#else
PyTuple_SetItem(field, 0, PyUnicode_FromString("resource"));
#endif
PyTuple_SetItem(field, 1, PyArray_TypeObjectFromType(NPY_UBYTE));
PyTuple_SetItem(field, 2, PyLong_FromLong(1));
PyObject* fields = PyList_New(1);
PyList_SetItem(fields, 0, field);
int convert_result = PyArray_DescrConverter(fields, descr);
Py_CLEAR(fields);
if (convert_result != 1) {
return absl::InternalError(
absl::StrCat("Failed to create numpy array description for ",
"TF_RESOURCE-type tensor"));
}
} else {
int type_num = -1;
TF_RETURN_IF_ERROR(
TF_DataType_to_PyArray_TYPE(TF_TensorType(tensor), &type_num));
*descr = PyArray_DescrFromType(type_num);
}
return absl::OkStatus();
}
inline void FastMemcpy(void* dst, const void* src, size_t size) {
// clang-format off
switch (size) {
// Most compilers will generate inline code for fixed sizes,
// which is significantly faster for small copies.
case 1: memcpy(dst, src, 1); break;
case 2: memcpy(dst, src, 2); break;
case 3: memcpy(dst, src, 3); break;
case 4: memcpy(dst, src, 4); break;
case 5: memcpy(dst, src, 5); break;
case 6: memcpy(dst, src, 6); break;
case 7: memcpy(dst, src, 7); break;
case 8: memcpy(dst, src, 8); break;
case 9: memcpy(dst, src, 9); break;
case 10: memcpy(dst, src, 10); break;
case 11: memcpy(dst, src, 11); break;
case 12: memcpy(dst, src, 12); break;
case 13: memcpy(dst, src, 13); break;
case 14: memcpy(dst, src, 14); break;
case 15: memcpy(dst, src, 15); break;
case 16: memcpy(dst, src, 16); break;
#if defined(PLATFORM_GOOGLE) || defined(PLATFORM_POSIX) && \
!defined(IS_MOBILE_PLATFORM)
// On Linux, memmove appears to be faster than memcpy for
// large sizes, strangely enough.
default: memmove(dst, src, size); break;
#else
default: memcpy(dst, src, size); break;
#endif
}
// clang-format on
}
} // namespace
// TODO(slebedev): revise TF_TensorToPyArray usages and switch to the
// aliased version where appropriate.
absl::Status TF_TensorToMaybeAliasedPyArray(Safe_TF_TensorPtr tensor,
PyObject** out_ndarray) {
auto dtype = TF_TensorType(tensor.get());
if (dtype == TF_STRING || dtype == TF_RESOURCE) {
return TF_TensorToPyArray(std::move(tensor), out_ndarray);
}
TF_Tensor* moved = tensor.release();
int64_t nelems = -1;
absl::InlinedVector<npy_intp, 4UL> dims;
TF_RETURN_IF_ERROR(GetPyArrayDimensionsForTensor(moved, &dims, &nelems));
return ArrayFromMemory(
dims.size(), dims.data(), TF_TensorData(moved),
static_cast<DataType>(dtype), [moved] { TF_DeleteTensor(moved); },
out_ndarray);
}
// Converts the given TF_Tensor to a numpy ndarray.
// If the returned status is OK, the caller becomes the owner of *out_array.
absl::Status TF_TensorToPyArray(Safe_TF_TensorPtr tensor,
PyObject** out_ndarray) {
// A fetched operation will correspond to a null tensor, and a None
// in Python.
if (tensor == nullptr) {
Py_INCREF(Py_None);
*out_ndarray = Py_None;
return absl::OkStatus();
}
int64_t nelems = -1;
absl::InlinedVector<npy_intp, 4UL> dims;
TF_RETURN_IF_ERROR(
GetPyArrayDimensionsForTensor(tensor.get(), &dims, &nelems));
// If the type is neither string nor resource we can reuse the Tensor memory.
TF_Tensor* original = tensor.get();
TF_Tensor* moved = TF_TensorMaybeMove(tensor.release());
if (moved != nullptr) {
if (ArrayFromMemory(
dims.size(), dims.data(), TF_TensorData(moved),
static_cast<DataType>(TF_TensorType(moved)),
[moved] { TF_DeleteTensor(moved); }, out_ndarray)
.ok()) {
return absl::OkStatus();
}
}
tensor.reset(original);
// Copy the TF_TensorData into a newly-created ndarray and return it.
PyArray_Descr* descr = nullptr;
TF_RETURN_IF_ERROR(GetPyArrayDescrForTensor(tensor.get(), &descr));
Safe_PyObjectPtr safe_out_array =
tensorflow::make_safe(PyArray_Empty(dims.size(), dims.data(), descr, 0));
if (!safe_out_array) {
return absl::InternalError("Could not allocate ndarray");
}
PyArrayObject* py_array =
reinterpret_cast<PyArrayObject*>(safe_out_array.get());
if (TF_TensorType(tensor.get()) == TF_STRING) {
absl::Status s =
CopyTF_TensorStringsToPyArray(tensor.get(), nelems, py_array);
if (!s.ok()) {
return s;
}
} else if ((!IsZenDnnEnabled()) &&
(static_cast<size_t>(PyArray_NBYTES(py_array)) !=
TF_TensorByteSize(tensor.get()))) {
return errors::Internal("ndarray was ", PyArray_NBYTES(py_array),
" bytes but TF_Tensor was ",
TF_TensorByteSize(tensor.get()), " bytes");
}
// Default TF implementation compares ndarray bytes with TF_Tensor bytes
// which is allocated using allocate_output API for fixed size.
// ZenDNN mempool optimization utilizes pre-allocated buffers with TF_Tensor.
// Hence, the amount of allocated TF_Tensor bytes is always greater than or
// equal to the number of bytes requested by the op execution.
else if (IsZenDnnEnabled() && (static_cast<size_t>(PyArray_NBYTES(py_array)) >
TF_TensorByteSize(tensor.get()))) {
return errors::Internal("ndarray was ", PyArray_NBYTES(py_array),
" bytes but TF_Tensor was ",
TF_TensorByteSize(tensor.get()), " bytes");
} else {
FastMemcpy(PyArray_DATA(py_array), TF_TensorData(tensor.get()),
PyArray_NBYTES(py_array));
}
*out_ndarray = safe_out_array.release();
return absl::OkStatus();
}
absl::Status NdarrayToTensor(TFE_Context* ctx, PyObject* ndarray,
Safe_TF_TensorPtr* ret) {
DCHECK(ret != nullptr);
// Make sure we dereference this array object in case of error, etc.
Safe_PyObjectPtr array_safe(make_safe(
PyArray_FromAny(ndarray, nullptr, 0, 0, NPY_ARRAY_CARRAY_RO, nullptr)));
if (!array_safe) return absl::InvalidArgumentError("Not a ndarray.");
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());
// Convert numpy dtype to TensorFlow dtype.
TF_DataType dtype = TF_FLOAT;
TF_RETURN_IF_ERROR(PyArray_TYPE_to_TF_DataType(array, &dtype));
int64_t nelems = 1;
absl::InlinedVector<int64_t, 4UL> dims;
for (int i = 0; i < PyArray_NDIM(array); ++i) {
dims.push_back(PyArray_SHAPE(array)[i]);
nelems *= dims[i];
}
// Create a TF_Tensor based on the fed data. In the case of non-string data
// type, this steals a reference to array, which will be relinquished when
// the underlying buffer is deallocated. For string, a new temporary buffer
// is allocated into which the strings are encoded.
if (dtype == TF_RESOURCE) {
size_t size = PyArray_NBYTES(array);
array_safe.release();
if (ctx) {
*ret = make_safe(new TF_Tensor{tensorflow::unwrap(ctx)->CreateTensor(
static_cast<tensorflow::DataType>(dtype), {}, 0, PyArray_DATA(array),
size, &DelayedNumpyDecref, array)});
} else {
*ret = make_safe(TF_NewTensor(dtype, {}, 0, PyArray_DATA(array), size,
&DelayedNumpyDecref, array));
}
} else if (dtype != TF_STRING) {
size_t size = PyArray_NBYTES(array);
array_safe.release();
if (ctx) {
*ret = make_safe(new TF_Tensor{tensorflow::unwrap(ctx)->CreateTensor(
static_cast<tensorflow::DataType>(dtype), dims.data(), dims.size(),
PyArray_DATA(array), size, &DelayedNumpyDecref, array)});
} else {
*ret = make_safe(TF_NewTensor(dtype, dims.data(), dims.size(),
PyArray_DATA(array), size,
&DelayedNumpyDecref, array));
}
} else {
size_t size = 0;
void* encoded = nullptr;
TF_RETURN_IF_ERROR(EncodePyBytesArray(array, nelems, &size, &encoded));
if (ctx) {
*ret = make_safe(new TF_Tensor{tensorflow::unwrap(ctx)->CreateTensor(
static_cast<tensorflow::DataType>(dtype), dims.data(), dims.size(),
encoded, size,
[](void* data, size_t len, void* arg) {
delete[] reinterpret_cast<tensorflow::tstring*>(data);
},
nullptr)});
} else {
*ret = make_safe(TF_NewTensor(
dtype, dims.data(), dims.size(), encoded, size,
[](void* data, size_t len, void* arg) {
delete[] reinterpret_cast<tensorflow::tstring*>(data);
},
nullptr));
}
}
return absl::OkStatus();
}
absl::Status TF_TensorToTensor(const TF_Tensor* src, Tensor* dst);
TF_Tensor* TF_TensorFromTensor(const tensorflow::Tensor& src,
absl::Status* status);
absl::Status NdarrayToTensor(PyObject* obj, Tensor* ret) {
Safe_TF_TensorPtr tf_tensor = make_safe(static_cast<TF_Tensor*>(nullptr));
absl::Status s = NdarrayToTensor(nullptr /*ctx*/, obj, &tf_tensor);
if (!s.ok()) {
return s;
}
return TF_TensorToTensor(tf_tensor.get(), ret);
}
absl::Status TensorToNdarray(const Tensor& t, PyObject** ret) {
absl::Status status;
Safe_TF_TensorPtr tf_tensor = make_safe(TF_TensorFromTensor(t, &status));
if (!status.ok()) {
return status;
}
return TF_TensorToMaybeAliasedPyArray(std::move(tf_tensor), ret);
}
} // namespace tensorflow
tensorflow/python/lib/core/ndarray_tensor.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_NDARRAY_TENSOR_H_
#define TENSORFLOW_PYTHON_LIB_CORE_NDARRAY_TENSOR_H_
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/safe_ptr.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
namespace tensorflow {
absl::Status TF_TensorToMaybeAliasedPyArray(Safe_TF_TensorPtr tensor,
PyObject** out_ndarray);
absl::Status TF_TensorToPyArray(Safe_TF_TensorPtr tensor,
PyObject** out_ndarray);
// Creates a tensor in 'ret' from the input `ndarray`. The returned TF_Tensor
// in `ret` may have its own Python reference to `ndarray`s data. After `ret`
// is destroyed, this reference must (eventually) be decremented via
// ClearDecrefCache().
ABSL_MUST_USE_RESULT
absl::Status NdarrayToTensor(TFE_Context* ctx, PyObject* ndarray,
Safe_TF_TensorPtr* ret);
// Creates a tensor in 'ret' from the input Ndarray.
// TODO(kkb): This is an old conversion function that does not support TFRT.
// Currently it's used for session, py_func, and an internal project. Migrate
// them.
ABSL_MUST_USE_RESULT
absl::Status NdarrayToTensor(PyObject* obj, Tensor* ret);
// Creates a numpy array in 'ret' which either aliases the content of 't' or has
// a copy.
absl::Status TensorToNdarray(const Tensor& t, PyObject** ret);
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_NDARRAY_TENSOR_H_
tensorflow/python/lib/core/ndarray_tensor_bridge.cc
+285
View File
@@ -0,0 +1,285 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// clang-format off
// Must be included first.
#include "tensorflow/c/tf_datatype.h"
#include "xla/tsl/python/lib/core/numpy.h"
// clang-format on
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
#include <vector>
#include "tensorflow/c/c_api.h"
#include "xla/tsl/python/lib/core/ml_dtypes.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/python/lib/core/py_util.h"
namespace tensorflow {
// Mutex used to serialize accesses to cached vector of pointers to python
// arrays to be dereferenced.
static mutex* DelayedDecrefLock() {
static mutex* decref_lock = new mutex;
return decref_lock;
}
// Caches pointers to numpy arrays which need to be dereferenced.
static std::vector<void*>* DecrefCache() {
static std::vector<void*>* decref_cache = new std::vector<void*>;
return decref_cache;
}
// Destructor passed to TF_NewTensor when it reuses a numpy buffer. Stores a
// pointer to the pyobj in a buffer to be dereferenced later when we're actually
// holding the GIL.
void DelayedNumpyDecref(void* data, size_t len, void* obj) {
mutex_lock ml(*DelayedDecrefLock());
DecrefCache()->push_back(obj);
}
// Actually dereferences cached numpy arrays. REQUIRES being called while
// holding the GIL.
void ClearDecrefCache() {
std::vector<void*> cache_copy;
{
mutex_lock ml(*DelayedDecrefLock());
cache_copy.swap(*DecrefCache());
}
for (void* obj : cache_copy) {
Py_DECREF(reinterpret_cast<PyObject*>(obj));
}
}
// Structure which keeps a reference to a Tensor alive while numpy has a pointer
// to it.
struct TensorReleaser {
// Python macro to include standard members.
PyObject_HEAD
// Destructor responsible for releasing the memory.
std::function<void()>* destructor;
};
extern PyTypeObject TensorReleaserType;
static void TensorReleaser_dealloc(PyObject* pself) {
TensorReleaser* self = reinterpret_cast<TensorReleaser*>(pself);
(*self->destructor)();
delete self->destructor;
TensorReleaserType.tp_free(pself);
}
// clang-format off
PyTypeObject TensorReleaserType = {
PyVarObject_HEAD_INIT(nullptr, 0) /* head init */
"tensorflow_wrapper", /* tp_name */
sizeof(TensorReleaser), /* tp_basicsize */
0, /* tp_itemsize */
/* methods */
TensorReleaser_dealloc, /* tp_dealloc */
#if PY_VERSION_HEX < 0x03080000
nullptr, /* tp_print */
#else
0, /* tp_vectorcall_offset */
#endif
nullptr, /* tp_getattr */
nullptr, /* tp_setattr */
nullptr, /* tp_compare */
nullptr, /* tp_repr */
nullptr, /* tp_as_number */
nullptr, /* tp_as_sequence */
nullptr, /* tp_as_mapping */
nullptr, /* tp_hash */
nullptr, /* tp_call */
nullptr, /* tp_str */
nullptr, /* tp_getattro */
nullptr, /* tp_setattro */
nullptr, /* tp_as_buffer */
Py_TPFLAGS_DEFAULT, /* tp_flags */
"Wrapped TensorFlow Tensor", /* tp_doc */
nullptr, /* tp_traverse */
nullptr, /* tp_clear */
nullptr, /* tp_richcompare */
};
// clang-format on
absl::Status TF_DataType_to_PyArray_TYPE(TF_DataType tf_datatype,
int* out_pyarray_type) {
const tsl::ml_dtypes::NumpyDtypes& custom_dtypes =
tsl::ml_dtypes::GetNumpyDtypes();
switch (tf_datatype) {
case TF_HALF:
*out_pyarray_type = NPY_FLOAT16;
break;
case TF_FLOAT:
*out_pyarray_type = NPY_FLOAT32;
break;
case TF_DOUBLE:
*out_pyarray_type = NPY_FLOAT64;
break;
case TF_INT32:
*out_pyarray_type = NPY_INT32;
break;
case TF_UINT32:
*out_pyarray_type = NPY_UINT32;
break;
case TF_UINT8:
*out_pyarray_type = NPY_UINT8;
break;
case TF_UINT16:
*out_pyarray_type = NPY_UINT16;
break;
case TF_INT8:
*out_pyarray_type = NPY_INT8;
break;
case TF_INT16:
*out_pyarray_type = NPY_INT16;
break;
case TF_INT64:
*out_pyarray_type = NPY_INT64;
break;
case TF_UINT64:
*out_pyarray_type = NPY_UINT64;
break;
case TF_BOOL:
*out_pyarray_type = NPY_BOOL;
break;
case TF_COMPLEX64:
*out_pyarray_type = NPY_COMPLEX64;
break;
case TF_COMPLEX128:
*out_pyarray_type = NPY_COMPLEX128;
break;
case TF_STRING:
*out_pyarray_type = NPY_OBJECT;
break;
case TF_RESOURCE:
*out_pyarray_type = NPY_VOID;
break;
// TODO(keveman): These should be changed to NPY_VOID, and the type used for
// the resulting numpy array should be the custom struct types that we
// expect for quantized types.
case TF_QINT8:
*out_pyarray_type = NPY_INT8;
break;
case TF_QUINT8:
*out_pyarray_type = NPY_UINT8;
break;
case TF_QINT16:
*out_pyarray_type = NPY_INT16;
break;
case TF_QUINT16:
*out_pyarray_type = NPY_UINT16;
break;
case TF_QINT32:
*out_pyarray_type = NPY_INT32;
break;
case TF_BFLOAT16:
*out_pyarray_type = custom_dtypes.bfloat16;
break;
case TF_FLOAT8_E5M2:
*out_pyarray_type = custom_dtypes.float8_e5m2;
break;
case TF_FLOAT8_E4M3FN:
*out_pyarray_type = custom_dtypes.float8_e4m3fn;
break;
case TF_FLOAT8_E4M3FNUZ:
*out_pyarray_type = custom_dtypes.float8_e4m3fnuz;
break;
case TF_FLOAT8_E4M3B11FNUZ:
*out_pyarray_type = custom_dtypes.float8_e4m3b11fnuz;
break;
case TF_FLOAT8_E5M2FNUZ:
*out_pyarray_type = custom_dtypes.float8_e5m2fnuz;
break;
case TF_FLOAT4_E2M1FN:
*out_pyarray_type = custom_dtypes.float4_e2m1fn;
break;
case TF_INT4:
*out_pyarray_type = custom_dtypes.int4;
break;
case TF_UINT4:
*out_pyarray_type = custom_dtypes.uint4;
break;
case TF_INT2:
*out_pyarray_type = custom_dtypes.int2;
break;
case TF_UINT2:
*out_pyarray_type = custom_dtypes.uint2;
break;
default:
return absl::InternalError(absl::StrCat(
"Tensorflow type ", tf_datatype, " not convertible to numpy dtype."));
}
return absl::OkStatus();
}
absl::Status ArrayFromMemory(int dim_size, npy_intp* dims, void* data,
DataType dtype, std::function<void()> destructor,
PyObject** result) {
if (dtype == DT_STRING || dtype == DT_RESOURCE) {
return absl::FailedPreconditionError(
"Cannot convert string or resource Tensors.");
}
int type_num = -1;
absl::Status s =
TF_DataType_to_PyArray_TYPE(static_cast<TF_DataType>(dtype), &type_num);
if (!s.ok()) {
return s;
}
if (dim_size > NPY_MAXDIMS) {
return absl::InvalidArgumentError(absl::StrCat(
"Cannot convert tensor with ", dim_size,
" dimensions to NumPy array. NumPy arrays can have at most ",
NPY_MAXDIMS, " dimensions"));
}
auto* np_array = reinterpret_cast<PyArrayObject*>(
PyArray_SimpleNewFromData(dim_size, dims, type_num, data));
if (np_array == nullptr) {
std::string shape_str = absl::StrJoin(
absl::Span<npy_intp>{dims, static_cast<size_t>(dim_size)}, ", ");
if (PyErr_Occurred()) {
std::string exception_str = PyExceptionFetch();
PyErr_Clear();
return absl::InvalidArgumentError(
absl::StrCat("Failed to create numpy array from tensor of shape [",
shape_str, "]. Numpy error: ", exception_str));
}
return absl::InternalError(absl::StrCat(
"Failed to create numpy array from tensor of shape [", shape_str, "]"));
}
PyArray_CLEARFLAGS(np_array, NPY_ARRAY_OWNDATA);
if (PyType_Ready(&TensorReleaserType) == -1) {
return absl::UnknownError("Python type initialization failed.");
}
auto* releaser = reinterpret_cast<TensorReleaser*>(
TensorReleaserType.tp_alloc(&TensorReleaserType, 0));
releaser->destructor = new std::function<void()>(std::move(destructor));
if (PyArray_SetBaseObject(np_array, reinterpret_cast<PyObject*>(releaser)) ==
-1) {
Py_DECREF(releaser);
return absl::UnknownError("Python array refused to use memory.");
}
*result = reinterpret_cast<PyObject*>(np_array);
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/python/lib/core/ndarray_tensor_bridge.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_NDARRAY_TENSOR_BRIDGE_H_
#define TENSORFLOW_PYTHON_LIB_CORE_NDARRAY_TENSOR_BRIDGE_H_
// Must be included first
// clang-format off
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
// clang-format on
#include <functional>
#include "tensorflow/c/c_api.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// Destructor passed to TF_NewTensor when it reuses a numpy buffer. Stores a
// pointer to the pyobj in a buffer to be dereferenced later when we're actually
// holding the GIL. Data and len are ignored.
void DelayedNumpyDecref(void* data, size_t len, void* obj);
// Actually dereferences cached numpy arrays. REQUIRES being called while
// holding the GIL.
void ClearDecrefCache();
// Creates a numpy array with shapes specified by dim_size and dims and content
// in data. The array does not own the memory, and destructor will be called to
// release it. If the status is not ok the caller is responsible for releasing
// the memory.
absl::Status ArrayFromMemory(int dim_size, npy_intp* dims, void* data,
DataType dtype, std::function<void()> destructor,
PyObject** result);
// Converts TF_DataType to the corresponding numpy type.
absl::Status TF_DataType_to_PyArray_TYPE(TF_DataType tf_datatype,
int* out_pyarray_type);
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_NDARRAY_TENSOR_BRIDGE_H_
tensorflow/python/lib/core/py_exception_registry.cc
+76
View File
@@ -0,0 +1,76 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/python/lib/core/py_exception_registry.h"
#include <Python.h>
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
PyExceptionRegistry* PyExceptionRegistry::singleton_ = nullptr;
void PyExceptionRegistry::Init(PyObject* code_to_exc_type_map) {
CHECK(singleton_ == nullptr) << "PyExceptionRegistry::Init() already called";
singleton_ = new PyExceptionRegistry;
CHECK(PyDict_Check(code_to_exc_type_map));
PyObject* key;
PyObject* value;
Py_ssize_t pos = 0;
while (PyDict_Next(code_to_exc_type_map, &pos, &key, &value)) {
singleton_->exc_types_.emplace(static_cast<TF_Code>(PyLong_AsLong(key)),
value);
// The exception classes should also have the lifetime of the process, but
// incref just in case.
Py_INCREF(value);
}
static const TF_Code kAllCodes[] = {TF_CANCELLED,
TF_UNKNOWN,
TF_INVALID_ARGUMENT,
TF_DEADLINE_EXCEEDED,
TF_NOT_FOUND,
TF_ALREADY_EXISTS,
TF_PERMISSION_DENIED,
TF_UNAUTHENTICATED,
TF_RESOURCE_EXHAUSTED,
TF_FAILED_PRECONDITION,
TF_ABORTED,
TF_OUT_OF_RANGE,
TF_UNIMPLEMENTED,
TF_INTERNAL,
TF_UNAVAILABLE,
TF_DATA_LOSS};
for (TF_Code code : kAllCodes) {
CHECK(singleton_->exc_types_.find(code) != singleton_->exc_types_.end())
<< error::Code_Name(static_cast<error::Code>(code))
<< " is not registered";
}
}
PyObject* PyExceptionRegistry::Lookup(TF_Code code) {
CHECK(singleton_ != nullptr) << "Must call PyExceptionRegistry::Init() "
"before PyExceptionRegistry::Lookup()";
CHECK_NE(code, TF_OK);
auto it = singleton_->exc_types_.find(code);
CHECK(it != singleton_->exc_types_.end())
<< "Unknown error code passed to PyExceptionRegistry::Lookup: " << code;
return it->second;
}
} // namespace tensorflow
tensorflow/python/lib/core/py_exception_registry.h
+74
View File
@@ -0,0 +1,74 @@
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PY_EXCEPTION_REGISTRY_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PY_EXCEPTION_REGISTRY_H_
#include <Python.h>
#include <map>
#include "tensorflow/c/tf_status.h"
#include "tensorflow/core/lib/core/error_codes.pb.h"
namespace tensorflow {
// Global registry mapping C API error codes to the corresponding custom Python
// exception type. This is used to expose the exception types to C extension
// code (i.e. so we can raise custom exceptions via SWIG).
//
// Init() must be called exactly once at the beginning of the process before
// Lookup() can be used.
//
// Example usage:
// TF_Status* status = TF_NewStatus();
// TF_Foo(..., status);
//
// if (TF_GetCode(status) != TF_OK) {
// PyObject* exc_type = PyExceptionRegistry::Lookup(TF_GetCode(status));
// // Arguments to OpError base class. Set `node_def` and `op` to None.
// PyObject* args =
// Py_BuildValue("sss", nullptr, nullptr, TF_Message(status));
// PyErr_SetObject(exc_type, args);
// Py_DECREF(args);
// TF_DeleteStatus(status);
// return NULL;
// }
class PyExceptionRegistry {
public:
// Initializes the process-wide registry. Should be called exactly once near
// the beginning of the process. The arguments are the various Python
// exception types (e.g. `cancelled_exc` corresponds to
// errors.CancelledError).
static void Init(PyObject* code_to_exc_type_map);
// Returns the Python exception type corresponding to `code`. Init() must be
// called before using this function. `code` should not be TF_OK.
static PyObject* Lookup(TF_Code code);
static inline PyObject* Lookup(int code) {
return Lookup(static_cast<TF_Code>(code));
}
private:
static PyExceptionRegistry* singleton_;
PyExceptionRegistry() = default;
// Maps error codes to the corresponding Python exception type.
std::map<TF_Code, PyObject*> exc_types_;
};
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_PY_EXCEPTION_REGISTRY_H_
tensorflow/python/lib/core/py_func.cc
+436
View File
@@ -0,0 +1,436 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// clang-format: off
// Must be included first.
#include "tensorflow/python/lib/core/py_func.h"
#include "xla/tsl/python/lib/core/numpy.h"
// clang-format: on
#include <Python.h>
#include <array>
#include "numpy/arrayobject.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/tfe_context_internal.h"
#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
#include "tensorflow/c/safe_ptr.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
#include "tensorflow/core/framework/allocation_description.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/python/eager/pywrap_tfe.h"
#include "tensorflow/python/lib/core/ndarray_tensor.h"
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
#include "tensorflow/python/lib/core/py_util.h"
namespace tensorflow {
namespace {
static mutex mu(LINKER_INITIALIZED);
static PyObject* py_trampoline TF_GUARDED_BY(mu) = nullptr;
// Returns the py_trampoline that is used to pass the control to the
// python runtime.
PyObject* GetPyTrampoline() {
mutex_lock l(mu);
return py_trampoline;
}
// A call to the registered python function.
struct PyCall {
// Passed to python runtime to call the python function registered
// with this "token".
std::string token;
// The device on which Tensors are stored; only used for EagerPyFunc.
Device* device = nullptr;
// True if the call is associated with an EagerPyFunc.
bool eager = false;
// True if the call is running under eager async mode.
bool eager_async = false;
// Inputs and outputs of this function invocation.
std::vector<Tensor> ins;
std::vector<Tensor> out;
};
bool IsCPUDevice(const Device* d) {
return d == nullptr || d->tensorflow_accelerator_device_info() == nullptr;
}
// Given the 'call', prepares the token and inputs as a python tuple that is
// appropriate for calling the trampoline.
absl::Status MakeArgTuple(const PyCall* call, TFE_Context* ctx,
PyObject** tuple) {
int64_t n = call->ins.size();
PyObject* lst = PyList_New(n);
CHECK(lst);
// TFE_TensorHandle assumes that CPU is identified by nullptr.
//
// Set device name to be empty if the device is CPU.
const char* device_name = nullptr;
if (call->device != nullptr && !IsCPUDevice(call->device))
device_name = call->device->name().c_str();
for (int64_t i = 0; i < n; ++i) {
PyObject* arg = nullptr;
if (call->eager) {
Tensor t = call->ins[i];
arg = EagerTensorFromHandle(tensorflow::wrap(
tensorflow::unwrap(ctx)->CreateLocalHandleFromTFTensor(t,
device_name)));
if (arg == nullptr) {
Py_DECREF(lst);
return absl::InternalError(
"Unable to procure EagerTensor from Tensor.");
}
} else {
absl::Status s = TensorToNdarray(call->ins[i], &arg);
if (!s.ok()) {
Py_DECREF(lst);
return s;
}
arg = PyArray_Return(reinterpret_cast<PyArrayObject*>(arg));
}
PyList_SetItem(lst, i, arg);
}
*tuple = Py_BuildValue("(ssN)", call->token.c_str(), device_name, lst);
if (*tuple == nullptr) {
return absl::InternalError(
"Failed to create python tuple. Please make sure `token` is a "
"well-formed UTF-8 string.");
}
return absl::OkStatus();
}
bool IsSingleNone(PyObject* obj) {
if (!PyArray_Check(obj)) {
return false;
}
PyArrayObject* array_obj = reinterpret_cast<PyArrayObject*>(obj);
if (PyArray_NDIM(array_obj) != 0 || PyArray_SIZE(array_obj) != 1) {
return false;
}
std::array<npy_intp, 0> indices;
char* item_ptr =
static_cast<char*>(PyArray_GetPtr(array_obj, indices.data()));
PyObject* item = PyArray_GETITEM(array_obj, item_ptr);
CHECK(item);
return item == Py_None;
}
// Retrieves a Tensor from `eager_tensor` and stores it in `output_tensor`.
// Validates that `output_tensor` is backed by memory in `expected_device`
// (which is assumed to be a local device, one on which the kernel was
// executed.)
//
// It may be nice to copy the tensor to the right device instead of failing if
// it isn't already there. This is left as a future exercise. The required
// device-copying logic is implemented in Python at the moment.
absl::Status ExtractTensorFromEagerTensor(const PyObject* eager_tensor,
TFE_Context* ctx,
const Device* expected_device,
const Tensor** output_tensor) {
tensorflow::TensorHandle* handle = absl::down_cast<TensorHandle*>(
tensorflow::unwrap(ctx)->TFTensorHandleFromInterface(
tensorflow::unwrap(EagerTensor_Handle(eager_tensor))));
Device* actual_device = handle->device();
TF_RETURN_IF_ERROR(handle->Tensor(output_tensor));
// actual_device may be nullptr, which implies local CPU.
if (expected_device == actual_device) return absl::OkStatus();
const std::string& expected_device_name =
expected_device->attributes().name();
if (actual_device == nullptr) {
if (!IsCPUDevice(expected_device)) {
return absl::InternalError(absl::StrCat(
"Expected the py_func to return a Tensor backed by memory in ",
expected_device_name,
", but is actually backed by local host memory. This is a bug."));
}
return absl::OkStatus();
}
// NOTE(ebrevdo): Here we could try comparing "actual_device_name"
// (actual_device->attributes()->name()) to expected_device_name and ensure
// they're the same. However, this comparison fails if we create a ClusterDef
// on localhost, mainly because the Device created by Eager code doesn't match
// the device created by a session. In this case, expected_device_name may
// contain "worker" but the Eager device name contains "localhost". Since we
// can't easily access the true underlying device of "worker" here, we are not
// able to perform a proper comparison. Furthermore, we can't check
// IsCPUDevice(actual_device) because the kernel's device may indeed be a
// GPU device (the python interpreter doesn't use it, however).
return absl::OkStatus();
}
// Calls the registered py function through the trampoline.
absl::Status DoCallPyFunc(PyCall* call, bool* out_log_on_error) {
*out_log_on_error = true;
PyObject* trampoline = GetPyTrampoline();
if (trampoline == nullptr) {
return absl::InvalidArgumentError(
"Missing py trampoline. Most likely, it is a link error.");
}
// Prepare the argument.
PyObject* args = nullptr;
std::unique_ptr<EagerExecutor> new_executor = nullptr;
EagerExecutor* old_executor = nullptr;
if (call->eager) {
// See FuncRegistry._ctx.
PyObject* ctx_handle = PyObject_GetAttrString(trampoline, "_ctx");
TFE_Context* ctx = reinterpret_cast<TFE_Context*>(
PyCapsule_GetPointer(ctx_handle, "TFE_Context"));
Py_XDECREF(ctx_handle);
CHECK_NE(ctx, nullptr);
TF_RETURN_IF_ERROR(MakeArgTuple(call, ctx, &args));
new_executor.reset(new EagerExecutor(call->eager_async));
old_executor = &(tensorflow::unwrap(ctx)->Executor());
tensorflow::unwrap(ctx)->SetExecutorForThread(new_executor.get());
} else {
TF_RETURN_IF_ERROR(MakeArgTuple(call, nullptr, &args));
}
CHECK(args);
// Invokes the trampoline.
PyObject* result = PyObject_Call(trampoline, args, nullptr);
Py_DECREF(args);
absl::Status s = absl::OkStatus();
if (result == nullptr) {
if (PyErr_Occurred()) {
if (PyErr_ExceptionMatches(PyExc_ValueError) ||
PyErr_ExceptionMatches(PyExc_TypeError)) {
s = absl::InvalidArgumentError(PyExceptionFetch());
} else if (PyErr_ExceptionMatches(PyExc_StopIteration)) {
*out_log_on_error = false;
s = absl::OutOfRangeError(PyExceptionFetch());
} else if (PyErr_ExceptionMatches(PyExc_MemoryError)) {
s = absl::ResourceExhaustedError(PyExceptionFetch());
} else if (PyErr_ExceptionMatches(PyExc_NotImplementedError)) {
s = absl::UnimplementedError(PyExceptionFetch());
} else {
// TODO(ebrevdo): Check if exception is an OpError and use the
// OpError.error_code property to map it back in the Status.
s = absl::UnknownError(PyExceptionFetch());
}
} else {
s = absl::InternalError(absl::StrCat("Failed to run py callback ",
call->token, ": see error log."));
}
}
PyObject* ctx_handle = PyObject_GetAttrString(trampoline, "_ctx");
TFE_Context* ctx = reinterpret_cast<TFE_Context*>(
PyCapsule_GetPointer(ctx_handle, "TFE_Context"));
Py_XDECREF(ctx_handle);
if (new_executor != nullptr) {
s.Update(new_executor->WaitForAllPendingNodes());
tensorflow::unwrap(ctx)->SetExecutorForThread(old_executor);
}
TF_RETURN_IF_ERROR(s);
// Process the return values and convert them to TF Tensors.
if (PyList_Check(result)) {
// `result` is a Python list; if this operation is an `EagerPyFunc`, then
// every item in the list must be an `EagerTensor`; otherwise, every element
// must be a NumPy array.
call->out.clear();
for (int i = 0; i < PyList_Size(result); ++i) {
Tensor t;
if (call->eager) {
const PyObject* item = PyList_GetItem(result, i);
if (EagerTensor_CheckExact(item)) {
const Tensor* tensor = nullptr;
s = ExtractTensorFromEagerTensor(item, ctx, call->device, &tensor);
if (s.ok()) t = *tensor;
} else {
s = absl::FailedPreconditionError(
absl::StrCat("Expected EagerTensor, found PyObject of type: ",
Py_TYPE(item)->tp_name));
}
} else {
s = NdarrayToTensor(PyList_GetItem(result, i), &t);
}
if (!s.ok()) {
break;
}
call->out.push_back(t);
}
} else if (EagerTensor_CheckExact(result) || result == Py_None) {
// result is an `EagerTensor` or `None`.
DCHECK(call->eager);
if (result != Py_None) {
const Tensor* t = nullptr;
s = ExtractTensorFromEagerTensor(result, ctx, call->device, &t);
if (s.ok()) call->out.push_back(*t);
}
} else if (PyArray_Check(result)) {
// `result` is a NumPy array.
DCHECK(!call->eager);
if (!IsSingleNone(result)) {
Tensor t;
s = NdarrayToTensor(result, &t);
if (s.ok()) {
call->out.push_back(t);
}
}
} else {
s = absl::InternalError(absl::StrCat("Unexpected PyObject was returned: ",
Py_TYPE(result)->tp_name));
}
Py_DECREF(result);
return s;
}
} // end namespace
void InitializePyTrampoline(PyObject* trampoline) {
mutex_lock l(mu);
if (py_trampoline == nullptr) {
py_trampoline = trampoline;
Py_INCREF(py_trampoline);
} else {
LOG(WARNING) << "InitializeCallback should only be called once";
}
}
class PyFuncOp : public OpKernel {
public:
explicit PyFuncOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("token", &token_));
eager_ = type_string() == "EagerPyFunc";
if (eager_) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("is_async", &eager_async_));
}
}
bool IsExpensive() override { return true; }
void Compute(OpKernelContext* ctx) override {
PyCall call;
call.token = token_;
call.eager = eager_;
if (call.eager) {
// Eager's C API uses `Device`, whereas `OpKernelContext` stores a
// `DeviceBase`; attempt to downcast.
call.device = dynamic_cast<Device*>(ctx->device());
if (call.device == nullptr) {
ctx->CtxFailureWithWarning(absl::InternalError(absl::StrCat(
"Unrecognized device class: ", ctx->device()->name())));
return;
}
call.eager_async = eager_async_;
}
VLOG(1) << "PyFuncOp of token " << call.token << "is called.";
for (int i = 0; i < ctx->num_inputs(); ++i) {
call.ins.push_back(ctx->input(i));
}
// NOTE(mrry): There is a potential time-of-check-to-time-of-use race here.
// because it is possible that `Py_Finalize()` could be called in another
// thread between this check and the call to `PyGILState_Ensure()`, which
// will abort the process if `Py_Finalize()` has been called. A more robust
// solution would be welcome, but it is not obvious how to make this work
// using the current Python C API.
OP_REQUIRES(ctx, Py_IsInitialized(),
absl::FailedPreconditionError(
"Python interpreter state is not initialized. "
"The process may be terminated."));
PyGILState_STATE py_threadstate;
py_threadstate = PyGILState_Ensure();
bool log_on_error;
absl::Status s = DoCallPyFunc(&call, &log_on_error);
// Sometimes py_funcs can be called without a session and leak memory. This
// ensures we clear the decref cache so this doesn't happen.
ClearDecrefCache();
PyGILState_Release(py_threadstate);
// Ensures that GIL is released even when !s.ok().
if (!s.ok()) {
if (log_on_error) {
ctx->CtxFailureWithWarning(s);
} else {
ctx->CtxFailure(s);
}
return;
}
OP_REQUIRES(
ctx, static_cast<int32_t>(call.out.size()) == ctx->num_outputs(),
absl::InvalidArgumentError(absl::StrCat(
token_, " returns ", call.out.size(),
" values, but expects to see ", ctx->num_outputs(), " values.")));
for (size_t i = 0; i < call.out.size(); ++i) {
const auto& t = call.out[i];
OP_REQUIRES(ctx, t.dtype() == output_type(i),
absl::InvalidArgumentError(
absl::StrCat(i, "-th value returned by ", token_, " is ",
DataTypeString(t.dtype()), ", but expects ",
DataTypeString(output_type(i)))));
ctx->set_output(i, t);
}
}
private:
std::string token_;
// True if and only if this op should execute the python function eagerly,
// i.e., if and only if the eager attribute is set.
bool eager_;
bool eager_async_;
PyFuncOp(const PyFuncOp&) = delete;
void operator=(const PyFuncOp&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("PyFunc").Device(DEVICE_CPU), PyFuncOp);
REGISTER_KERNEL_BUILDER(Name("PyFuncStateless").Device(DEVICE_CPU), PyFuncOp);
REGISTER_KERNEL_BUILDER(Name("EagerPyFunc").Device(DEVICE_CPU), PyFuncOp);
DataType gpu_types[] = {
// No strings and int32s, no ref types and no resource/variant types.
DT_FLOAT, DT_DOUBLE, DT_UINT8, DT_INT16, DT_INT8,
DT_COMPLEX64, DT_INT64, DT_BOOL, DT_QINT8, DT_QUINT8,
DT_QINT32, DT_BFLOAT16, DT_QINT16, DT_QUINT16, DT_UINT16,
DT_COMPLEX128, DT_HALF, DT_UINT32, DT_UINT64,
};
REGISTER_KERNEL_BUILDER(Name("EagerPyFunc")
.Device(DEVICE_DEFAULT)
.TypeConstraint("Tin", gpu_types)
.TypeConstraint("Tout", gpu_types),
PyFuncOp);
} // end namespace tensorflow
tensorflow/python/lib/core/py_func.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PY_FUNC_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PY_FUNC_H_
#include <Python.h>
namespace tensorflow {
// Called by python code on initialization.
//
// "trampoline" must represent a python function which has the
// following signature:
// (string, list(ndarray)) | (string, list(EagerTensor)) ->
// ndarray | list(ndarray) | python scalar |
// EagerTensor | list(EagerTensor) | None
//
// The trampoline takes two arguments, the first is a string token
// used by the python frontend's dispatching logic; the second is a
// list of numpy ndarrays or EagerTensor objects. It can return a
// single numpy ndarray, a list of numpy ndarrays, a python scalar, an
// EagerTensor, a list of EagerTensors, or None.
//
// PyFunc requires inputs and outputs to be ndarrays. EagerPyFunc requires
// inputs to be a list of EagerTensors and outputs to be an EagerTensor, a list
// of EagerTensors, or None.
//
// The C++ runtime converts outputs back to Tensor objects.
//
// This function is called by script_ops.py during its module initialization.
//
// TODO(zhifengc): Support distributed runtime.
void InitializePyTrampoline(PyObject* trampoline);
} // end namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_PY_FUNC_H_
tensorflow/python/lib/core/py_func_wrapper.cc
+25
View File
@@ -0,0 +1,25 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/python/lib/core/py_func.h"
namespace py = pybind11;
PYBIND11_MODULE(_pywrap_py_func, m) {
m.def("initialize_py_trampoline", [](py::object trampoline) {
return tensorflow::InitializePyTrampoline(trampoline.ptr());
});
}
tensorflow/python/lib/core/py_seq_tensor.cc
+927
View File
@@ -0,0 +1,927 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// Must be included first
// clang-format off
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
// clang-format on
#include "tensorflow/python/lib/core/py_seq_tensor.h"
#include "tensorflow/c/eager/tfe_context_internal.h"
#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
#include "tensorflow/c/safe_ptr.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/c/tf_tensor_internal.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/python/lib/core/ndarray_tensor.h"
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
#include "tensorflow/python/lib/core/py_util.h"
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
namespace tensorflow {
namespace {
inline bool PyIsInstance(PyObject* obj, PyTypeObject* t) {
return PyObject_IsInstance(obj, reinterpret_cast<PyObject*>(t));
}
inline PyObject* PyType(PyObject* obj) {
return reinterpret_cast<PyObject*>(obj->ob_type);
}
bool IsPyString(PyObject* obj) {
return PyBytes_Check(obj) || PyUnicode_Check(obj);
}
bool IsPyInt(PyObject* obj) {
#if PY_MAJOR_VERSION >= 3
return PyLong_Check(obj) ||
PyIsInstance(obj, &PyIntegerArrType_Type); // NumPy integers
#else
return PyInt_Check(obj) || PyLong_Check(obj) ||
PyIsInstance(obj, &PyIntegerArrType_Type); // NumPy integers
#endif
}
bool IsPyDouble(PyObject* obj) {
return PyIsInstance(obj, &PyDoubleArrType_Type); // NumPy double type.
}
bool IsNumpyHalf(PyObject* obj) {
return PyIsInstance(obj, &PyHalfArrType_Type);
}
bool IsPyFloat(PyObject* obj) {
return PyFloat_Check(obj) ||
PyIsInstance(obj, &PyFloatingArrType_Type); // NumPy float types
}
struct ConverterState {
// The inferred tensor shape.
absl::InlinedVector<int64_t, 4UL> inferred_shape;
// The inferred tensor data type.
DataType inferred_dtype;
// The following fields are used by ZeroDimArrayToScalar.
// We cache the last result of the check for a zero dimensional array: the
// function is called many times in a conversion, and most of the time is
// to check for the same type. This cache can reduce the conversion time by
// about 25%.
PyTypeObject* last_zerodim_type;
bool last_zerodim_check;
ConverterState() : inferred_dtype(DT_INVALID), last_zerodim_type(nullptr) {}
};
// If the input is a zero dimensional PyArray return it converted to a scalar.
// Otherwise return the input and increment its reference count.
// Users must Py_DECREF the output of this method.
PyObject* ZeroDimArrayToScalar(PyObject* obj, ConverterState* state) {
auto type = Py_TYPE(obj);
auto pyarray_obj = reinterpret_cast<PyArrayObject*>(obj);
if (type != state->last_zerodim_type) {
state->last_zerodim_type = type;
state->last_zerodim_check =
PyObject_TypeCheck(obj, &PyArray_Type) &&
!PyObject_TypeCheck(obj, &PyGenericArrType_Type);
}
if (state->last_zerodim_check && PyArray_NDIM(pyarray_obj) == 0) {
obj = PyArray_ToScalar(PyArray_DATA(pyarray_obj), pyarray_obj);
} else {
Py_INCREF(obj);
}
return obj;
}
// Sets *elem to a NEW reference to an element in seq on success.
// REQUIRES: PySequence_Check(seq) && PySequence_Length(seq) > 0.
absl::Status SampleElementFromSequence(PyObject* seq, PyObject** elem) {
*elem = PySequence_GetItem(seq, 0);
if (*elem != nullptr) return absl::OkStatus();
// seq may implement the sequence protocol (i.e., implement __getitem__)
// but may legitimately not have a 0-th element (__getitem__(self, 0)
// raises a KeyError). For example:
// seq = pandas.Series([0, 1, 2], index=[2, 4, 6])
//
// We don't actually care for the element at key 0, any element will do
// for inferring the element types. All elements are expected to
// have the same type, and this will be validated when converting
// to an EagerTensor.
PyErr_Clear();
Safe_PyObjectPtr iter(PyObject_GetIter(seq));
if (PyErr_Occurred()) {
return absl::InvalidArgumentError(
absl::StrCat("Cannot infer dtype of a ", Py_TYPE(seq)->tp_name,
" object: ", PyExceptionFetch()));
}
*elem = PyIter_Next(iter.get());
if (PyErr_Occurred()) {
return absl::InvalidArgumentError(absl::StrCat(
"Cannot infer dtype of a ", Py_TYPE(seq)->tp_name,
" object, as iter(<object>).next() failed: ", PyExceptionFetch()));
}
if (*elem == nullptr) {
return absl::InvalidArgumentError(
absl::StrCat("Cannot infer dtype of a ", Py_TYPE(seq)->tp_name,
" object since it is an empty sequence"));
}
return absl::OkStatus();
}
tstring PyRepr(PyObject* obj);
bool IsPyDimension(PyObject* obj);
absl::Status InferShapeAndType(PyObject* obj, ConverterState* state) {
std::vector<Safe_PyObjectPtr> refs_to_clean;
while (true) {
// Convert any zero dimensional numpy arrays to scalars first of all.
// We also have to make sure a reference to the safe_obj is kept.
obj = ZeroDimArrayToScalar(obj, state);
refs_to_clean.push_back(make_safe(obj));
// We test strings first, in case a string is considered a sequence.
if (IsPyString(obj)) {
state->inferred_dtype = DT_STRING;
} else if (PySequence_Check(obj)) {
auto length = PySequence_Length(obj);
if (length > 0) {
state->inferred_shape.push_back(length);
PyObject* elem = nullptr;
TF_RETURN_IF_ERROR(SampleElementFromSequence(obj, &elem));
obj = elem;
refs_to_clean.push_back(make_safe(obj));
continue;
} else if (length == 0) {
state->inferred_shape.push_back(length);
state->inferred_dtype = DT_INVALID; // Invalid dtype for empty tensors.
} else {
// The sequence does not have a valid length (PySequence_Length < 0).
if (PyErr_Occurred()) {
// PySequence_Length failed and set an exception. Fetch the message
// and convert it to a failed status.
return absl::InvalidArgumentError(PyExceptionFetch());
} else {
// This is almost certainly dead code: PySequence_Length failed but
// did not set an exception.
return absl::InvalidArgumentError(
"Attempted to convert an invalid sequence to a Tensor.");
}
}
} else if (IsPyDouble(obj)) {
state->inferred_dtype = DT_DOUBLE;
} else if (IsNumpyHalf(obj)) {
state->inferred_dtype = DT_HALF;
} else if (IsPyFloat(obj)) {
state->inferred_dtype = DT_FLOAT;
} else if (PyBool_Check(obj) || PyIsInstance(obj, &PyBoolArrType_Type)) {
// Have to test for bool before int, since IsInt(True/False) == true.
state->inferred_dtype = DT_BOOL;
} else if (IsPyInt(obj)) {
state->inferred_dtype = DT_INT64;
} else if (IsPyDimension(obj)) {
state->inferred_dtype = DT_INT64;
} else if (PyComplex_Check(obj) ||
PyIsInstance(obj, &PyComplexFloatingArrType_Type)) { // NumPy
state->inferred_dtype = DT_COMPLEX128;
} else {
return absl::InvalidArgumentError(
absl::StrCat("Attempt to convert a value (", PyRepr(obj),
") with an unsupported type (", PyRepr(PyType(obj)),
") to a Tensor."));
}
return absl::OkStatus();
}
}
// Error messages
const char ErrorConverting[] =
"Error while converting Python sequence to Tensor.";
const char ErrorRectangular[] =
"Can't convert non-rectangular Python sequence to Tensor.";
const char ErrorMixedTypes[] =
"Can't convert Python sequence with mixed types to Tensor.";
const char ErrorOutOfRange[] =
"Can't convert Python sequence with out-of-range integer to Tensor.";
const char ErrorOutOfRangeDouble[] =
"Can't convert Python sequence with a value out of range for a "
"double-precision float.";
const char ErrorConvertingUnicodeString[] =
"Error converting unicode string while converting Python sequence to "
"Tensor.";
const char ErrorFoundInt64[] =
"Can't convert Python sequence with out-of-range integer to int32 Tensor.";
const char ErrorFoundFloat[] =
"Can't convert Python sequence with floating point values to integer "
"Tensor.";
// Defines a converter that recursively converts an object into
// an array of type T using the conversion function defined by the
// traits class in a ConvertScalar function.
//
// Note that these helper functions require shape.dims() >= 1.
template <class T>
struct ConverterTraits {
static const tensorflow::DataType kTypeEnum;
static const char* ConvertScalar(PyObject* v, T* out);
};
template <class T>
struct Converter {
static const char* Helper(PyObject* obj, int depth, ConverterState* state,
T** buf) {
if (TF_PREDICT_FALSE(obj == nullptr)) {
return ErrorConverting;
}
Safe_PyObjectPtr seq = make_safe(PySequence_Fast(obj, ""));
if (TF_PREDICT_FALSE(seq == nullptr)) return ErrorRectangular;
const int64_t s = state->inferred_shape[depth];
if (TF_PREDICT_FALSE(s != PySequence_Fast_GET_SIZE(seq.get()))) {
return ErrorRectangular;
}
if (state->inferred_shape.size() - depth > 1) {
/* Iterate over outer dim, and recursively convert each element. */
for (int64_t i = 0; i < s; ++i) {
const char* error = Helper(PySequence_Fast_GET_ITEM(seq.get(), i),
depth + 1, state, buf);
if (TF_PREDICT_FALSE(error != nullptr)) return error;
}
} else {
PyObject** l = PySequence_Fast_ITEMS(seq.get());
for (int64_t i = 0; i < s; ++i) {
auto scalar = ZeroDimArrayToScalar(l[i], state);
const char* error = ConverterTraits<T>::ConvertScalar(scalar, *buf);
Py_DECREF(scalar);
if (TF_PREDICT_FALSE(error != nullptr)) return error;
++*buf;
}
}
return nullptr;
}
static absl::Status Convert(TFE_Context* ctx, PyObject* obj,
ConverterState* state, TFE_TensorHandle** h,
const char** error) {
// TODO(josh11b): Allocator & attributes
AbstractTensorInterface* t;
if (state->inferred_shape.empty()) { /* Scalar case */
T value;
auto scalar = ZeroDimArrayToScalar(obj, state);
*error = ConverterTraits<T>::ConvertScalar(scalar, &value);
Py_DECREF(scalar);
if (*error != nullptr) return absl::InvalidArgumentError(*error);
t = ConverterTraits<T>::CreateScalar(ctx, value);
if (t == nullptr) {
return absl::InternalError("Cannot create tensor.");
}
} else {
t = ConverterTraits<T>::CreateTensor(ctx, state->inferred_shape);
if (t == nullptr) {
return absl::InternalError("Cannot create tensor.");
}
if (t->NumElements() > 0) {
T* buf = static_cast<T*>(t->Data());
*error = Helper(obj, 0, state, &buf);
if (*error != nullptr) {
t->Release();
return absl::InvalidArgumentError(*error);
}
}
}
*h = tensorflow::wrap(tensorflow::unwrap(ctx)->CreateLocalHandle(t));
t->Release();
return absl::OkStatus();
}
};
// Int support
template <>
struct ConverterTraits<int64_t> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx,
int64_t value) {
return tensorflow::unwrap(ctx)->CreateInt64Scalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_INT64, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, int64_t* out) {
#if PY_MAJOR_VERSION < 3
if (TF_PREDICT_TRUE(PyInt_Check(v))) {
*out = PyInt_AS_LONG(v);
return nullptr;
}
#endif
if (TF_PREDICT_TRUE(PyLong_Check(v) || IsPyDimension(v))) {
int overflow = 0;
// Have to use LongLong for 64 bits, since long is 32 bits on Windows.
*out = PyLong_AsLongLongAndOverflow(v, &overflow);
if (TF_PREDICT_FALSE(overflow)) return ErrorOutOfRange;
return nullptr;
}
if (PyIsInstance(v, &PyIntegerArrType_Type)) { // NumPy integers
#if PY_MAJOR_VERSION < 3
Safe_PyObjectPtr as_int = make_safe(PyNumber_Int(v));
#else
Safe_PyObjectPtr as_int = make_safe(PyNumber_Long(v));
#endif
return ConvertScalar(as_int.get(), out);
}
if (IsPyFloat(v)) return ErrorFoundFloat;
return ErrorMixedTypes;
}
};
typedef Converter<int64_t> Int64Converter;
template <>
struct ConverterTraits<uint64_t> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx,
uint64_t value) {
return tensorflow::unwrap(ctx)->CreateUint64Scalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_UINT64, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, uint64_t* out) {
#if PY_MAJOR_VERSION < 3
if (TF_PREDICT_TRUE(PyInt_Check(v))) {
*out = PyInt_AsUnsignedLongLongMask(v);
return nullptr;
}
#endif
if (TF_PREDICT_TRUE(PyLong_Check(v) || IsPyDimension(v))) {
*out = PyLong_AsUnsignedLongLong(v);
return nullptr;
}
if (PyIsInstance(v, &PyIntegerArrType_Type)) { // NumPy integers
#if PY_MAJOR_VERSION < 3
Safe_PyObjectPtr as_int = make_safe(PyNumber_Int(v));
#else
Safe_PyObjectPtr as_int = make_safe(PyNumber_Long(v));
#endif
return ConvertScalar(as_int.get(), out);
}
if (IsPyFloat(v)) return ErrorFoundFloat;
return ErrorMixedTypes;
}
};
typedef Converter<uint64_t> UInt64Converter;
template <>
struct ConverterTraits<int32_t> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx,
int32_t value) {
return tensorflow::unwrap(ctx)->CreateInt32Scalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_INT32, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, int32_t* out) {
int64_t i;
#if PY_MAJOR_VERSION < 3
if (TF_PREDICT_TRUE(PyInt_Check(v))) {
i = PyInt_AS_LONG(v);
} else
#endif
if (PyLong_Check(v) || IsPyDimension(v)) {
int overflow = 0;
// Have to use LongLong for 64 bits, since long is 32 bits on Windows.
i = PyLong_AsLongLongAndOverflow(v, &overflow);
if (TF_PREDICT_FALSE(overflow)) return ErrorOutOfRange;
} else if (PyIsInstance(v, &PyIntegerArrType_Type)) { // NumPy integers
#if PY_MAJOR_VERSION < 3
Safe_PyObjectPtr as_int = make_safe(PyNumber_Int(v));
#else
Safe_PyObjectPtr as_int = make_safe(PyNumber_Long(v));
#endif
return ConvertScalar(as_int.get(), out);
} else if (IsPyFloat(v)) {
return ErrorFoundFloat;
} else {
return ErrorMixedTypes;
}
*out = static_cast<uint32_t>(static_cast<uint64_t>(i));
// Check for 32-bit overflow.
if (TF_PREDICT_FALSE(i != *out)) return ErrorFoundInt64;
return nullptr;
}
};
typedef Converter<int32_t> Int32Converter;
// Floating-point support
// Returns `true` if `out` overflows when converted from `as_double`.
template <class T>
static inline bool CheckForOverflow(double as_double, T* out) {
return (sizeof(T) < sizeof(double) && std::isinf(*out) &&
std::isfinite(as_double));
}
// There is no `std::isinf` that takes `Eigen::half` as argument but Eigen
// provides `Eigen::numext::isinf` instead.
template <>
inline bool CheckForOverflow<Eigen::half>(double as_double, Eigen::half* out) {
return (Eigen::numext::isinf(*out) && std::isfinite(as_double));
}
template <class T>
static const char* ConvertOneFloat(PyObject* v, T* out) {
if (PyErr_Occurred()) {
return nullptr;
}
if (TF_PREDICT_TRUE(PyFloat_Check(v))) {
const double as_double = PyFloat_AS_DOUBLE(v);
*out = static_cast<T>(as_double);
// Check for overflow
if (TF_PREDICT_FALSE(CheckForOverflow<T>(as_double, out))) {
return ErrorOutOfRangeDouble;
}
return nullptr;
}
#if PY_MAJOR_VERSION < 3
if (PyInt_Check(v)) {
*out = static_cast<T>(PyInt_AS_LONG(v));
return nullptr;
}
#endif
if (PyLong_Check(v)) {
*out = static_cast<T>(PyLong_AsDouble(v));
if (PyErr_Occurred()) return ErrorOutOfRangeDouble;
return nullptr;
}
if (PyIsInstance(v, &PyFloatingArrType_Type)) { // NumPy float types
Safe_PyObjectPtr as_float = make_safe(PyNumber_Float(v));
if (PyErr_Occurred()) {
return nullptr;
}
return ConvertOneFloat<T>(as_float.get(), out);
}
if (PyIsInstance(v, &PyIntegerArrType_Type)) { // NumPy integers
#if PY_MAJOR_VERSION < 3
Safe_PyObjectPtr as_int = make_safe(PyNumber_Int(v));
#else
Safe_PyObjectPtr as_int = make_safe(PyNumber_Long(v));
#endif
if (PyErr_Occurred()) {
return nullptr;
}
return ConvertOneFloat<T>(as_int.get(), out);
}
return ErrorMixedTypes;
}
template <>
struct ConverterTraits<float> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx, float value) {
return tensorflow::unwrap(ctx)->CreateFloatScalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_FLOAT, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, float* out) {
return ConvertOneFloat<float>(v, out);
}
};
template <>
struct ConverterTraits<double> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx, double value) {
return tensorflow::unwrap(ctx)->CreateDoubleScalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_DOUBLE, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, double* out) {
return ConvertOneFloat<double>(v, out);
}
};
typedef Converter<double> DoubleConverter;
typedef Converter<float> FloatConverter;
template <>
struct ConverterTraits<Eigen::half> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx,
Eigen::half value) {
return tensorflow::unwrap(ctx)->CreateHalfScalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_HALF, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, Eigen::half* out) {
return ConvertOneFloat<Eigen::half>(v, out);
}
};
typedef Converter<Eigen::half> NumpyHalfConverter;
// String support
template <>
struct ConverterTraits<tstring> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx,
tstring value) {
return tensorflow::unwrap(ctx)->CreateStringScalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_STRING, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, tstring* out) {
if (PyBytes_Check(v)) {
out->assign(PyBytes_AS_STRING(v), PyBytes_GET_SIZE(v));
return nullptr;
}
if (PyUnicode_Check(v)) {
#if PY_MAJOR_VERSION >= 3
Py_ssize_t size;
const char* str = PyUnicode_AsUTF8AndSize(v, &size);
if (str == nullptr) return ErrorConvertingUnicodeString;
out->assign(str, size);
return nullptr;
#else
PyObject* py_str = PyUnicode_AsUTF8String(v);
if (py_str == nullptr) return ErrorConvertingUnicodeString;
out->assign(PyBytes_AS_STRING(py_str), PyBytes_GET_SIZE(py_str));
Py_DECREF(py_str);
return nullptr;
#endif
}
return ErrorMixedTypes;
}
};
typedef Converter<tstring> StringConverter;
// Converts Python object `c` that should hold a Python string into a
// C++ string in *out. Returns nullptr on success, or a message on error.
// Defined below, but forward declared here for use in PyRepr.
tstring PyRepr(PyObject* obj) {
if (obj == nullptr) {
return "<null>";
}
Safe_PyObjectPtr repr_obj = make_safe(PyObject_Repr(obj));
if (repr_obj) {
tstring repr_str;
if (ConverterTraits<tstring>::ConvertScalar(repr_obj.get(), &repr_str) ==
nullptr) {
return repr_str;
}
}
return "<error computing repr()>";
}
bool IsPyDimension(PyObject* obj) {
const char* tp_name = obj->ob_type->tp_name;
if (strcmp(tp_name, "Dimension") != 0) return false;
bool ret =
absl::EndsWith(PyRepr(PyType(obj)),
"tensorflow.python.framework.tensor_shape.Dimension'>");
return ret;
}
// Complex support
template <>
struct ConverterTraits<complex128> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx,
complex128 value) {
return tensorflow::unwrap(ctx)->CreateComplex128Scalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_COMPLEX128, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, complex128* out) {
if (PyComplex_Check(v)) {
*out = complex128(PyComplex_RealAsDouble(v), PyComplex_ImagAsDouble(v));
return nullptr;
} else if (PyIsInstance(v, &PyComplexFloatingArrType_Type)) { // NumPy
auto as_complex = PyComplex_AsCComplex(v);
*out = complex128(as_complex.real, as_complex.imag);
return nullptr;
}
double as_double;
auto error = ConvertOneFloat<double>(v, &as_double);
if (error != nullptr) return error;
*out = complex128(as_double, 0.0);
return nullptr;
}
};
typedef Converter<complex128> Complex128Converter;
// Bool support
template <>
struct ConverterTraits<bool> {
static AbstractTensorInterface* CreateScalar(TFE_Context* ctx, bool value) {
return tensorflow::unwrap(ctx)->CreateBoolScalar(value);
}
static AbstractTensorInterface* CreateTensor(
TFE_Context* ctx, absl::Span<const int64_t> dim_sizes) {
return tensorflow::unwrap(ctx)->CreateTensor(DT_BOOL, dim_sizes);
}
static const char* ConvertScalar(PyObject* v, bool* out) {
if (v == Py_True) {
*out = true;
} else if (v == Py_False) {
*out = false;
} else if (PyIsInstance(v, &PyBoolArrType_Type)) { // NumPy
*out = PyObject_IsTrue(v);
} else {
return ErrorMixedTypes;
}
return nullptr;
}
};
typedef Converter<bool> BoolConverter;
// Convert a Python numpy.ndarray object to a TFE_TensorHandle.
// The two may share underlying storage so changes to one may reflect in the
// other.
TFE_TensorHandle* NumpyToTFE_TensorHandle(TFE_Context* ctx, PyObject* obj) {
Safe_TF_TensorPtr tf_tensor = make_safe(static_cast<TF_Tensor*>(nullptr));
absl::Status status = tensorflow::NdarrayToTensor(ctx, obj, &tf_tensor);
if (TF_PREDICT_FALSE(!status.ok())) {
PyErr_SetString(
PyExc_ValueError,
absl::StrCat("Failed to convert a NumPy array to a Tensor (",
status.message(), ").")
.c_str());
return nullptr;
}
return tensorflow::wrap(
tensorflow::unwrap(ctx)->CreateLocalHandle(tf_tensor->tensor));
}
} // namespace
// TODO(b/147743551): This function handles enough conversions to justify
// promoting to something like PyObjectToTensorHandle.
// TODO(b/147828820): Handle Tensors properly.
TFE_TensorHandle* PySeqToTFE_TensorHandle(TFE_Context* ctx, PyObject* obj,
DataType dtype) {
// Shortcut: __array__ objects (such as Pandas data frames).
// These objects are efficiently handled by Numpy. We transform them into
// Numpy arrays and handle them in the Numpy case below. Note that Tensors
// implement the __array__ function, and will be handled in this shortcut.
// We used to call PyArray_FromArrayAttr here, but NumPy 2.0 changed its
// semantics such that it errors if a copy of the array is required.
// (Ideally no copy would be needed here, but that would be a larger change.)
Safe_PyObjectPtr array;
if (PyObject_HasAttrString(obj, "__array__")) {
array = make_safe(PyObject_CallMethod(obj, "__array__", nullptr));
if (array == nullptr) {
return nullptr;
}
if (!PyArray_Check(array.get())) {
PyErr_SetString(PyExc_ValueError,
"Value returned by __array__ is not a NumPy array");
return nullptr;
}
}
if (!array) {
// Try __array_interface__ objects (such as PIL Image).
array = make_safe(PyArray_FromInterface(obj));
if (array == nullptr) {
return nullptr;
}
if (array.get() == Py_NotImplemented) {
array.release();
} else {
obj = array.get();
}
} else {
// PyArray_FromArrayAttr ensures that `array` is a PyArrayObject, so all
// we have to do is replace `obj` with it and continue.
obj = array.get();
}
// Shortcut: Numpy arrays.
if (PyArray_Check(obj)) {
int desired_np_dtype = -1;
if (dtype != tensorflow::DT_INVALID) {
if (!tensorflow::TF_DataType_to_PyArray_TYPE(
static_cast<TF_DataType>(dtype), &desired_np_dtype)
.ok()) {
PyErr_SetString(
PyExc_TypeError,
absl::StrCat("Invalid dtype argument value ", dtype).c_str());
return nullptr;
}
}
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(obj);
int array_dtype = PyArray_TYPE(array);
Safe_PyObjectPtr safe_value(nullptr);
// Use Numpy to convert between types if needed.
if ((desired_np_dtype >= 0 && desired_np_dtype != array_dtype) ||
!PyArray_ISCARRAY(array)) {
int new_dtype = desired_np_dtype >= 0 ? desired_np_dtype : array_dtype;
safe_value = tensorflow::make_safe(
PyArray_FromAny(obj, PyArray_DescrFromType(new_dtype), 0, 0,
NPY_ARRAY_CARRAY_RO | NPY_ARRAY_FORCECAST, nullptr));
if (PyErr_Occurred()) return nullptr;
if (safe_value == nullptr) {
PyErr_SetString(PyExc_ValueError, "Error while casting a numpy value");
}
obj = safe_value.get();
}
return NumpyToTFE_TensorHandle(ctx, obj);
}
ConverterState state;
absl::Status status = InferShapeAndType(obj, &state);
if (!status.ok()) {
PyErr_SetString(PyExc_ValueError, absl::StatusMessageAsCStr(status));
return nullptr;
}
DataType requested_dtype = DT_INVALID;
if (dtype != DT_INVALID) {
requested_dtype = dtype;
}
// NOTE(josh11b): If don't successfully convert to the requested type,
// we just try instead to create a tensor of the inferred type and
// let the caller convert it to the requested type using a cast
// operation.
const char* error = nullptr;
TFE_TensorHandle* handle = nullptr;
status = absl::UnimplementedError(
absl::StrCat("Missing Python -> Tensor conversion for ",
DataTypeString(state.inferred_dtype)));
switch (requested_dtype) {
case DT_FLOAT:
status = FloatConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_DOUBLE:
status = DoubleConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_HALF:
status = NumpyHalfConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_INT64:
status = Int64Converter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_INT32:
status = Int32Converter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_UINT64:
status = UInt64Converter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_COMPLEX128:
status = Complex128Converter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_STRING:
status = StringConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_BOOL:
status = BoolConverter::Convert(ctx, obj, &state, &handle, &error);
break;
default:
break;
}
if (status.ok()) return handle;
switch (state.inferred_dtype) {
case DT_FLOAT:
// TODO(josh11b): Handle mixed floats and complex numbers?
if (requested_dtype == DT_INVALID) {
// TensorFlow uses float32s to represent floating point numbers
// by default (for space and speed over using doubles).
status = FloatConverter::Convert(ctx, obj, &state, &handle, &error);
} else {
// We are going to do a cast to the user's requested dtype
// after this. We use doubles for this intermediate result so
// we don't lose precision that might be representable in the
// final type.
status = DoubleConverter::Convert(ctx, obj, &state, &handle, &error);
}
break;
case DT_DOUBLE:
status = DoubleConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_HALF:
status = NumpyHalfConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_INT64:
if (requested_dtype == DT_INVALID) {
status = Int32Converter::Convert(ctx, obj, &state, &handle, &error);
if (error == ErrorFoundInt64) {
status = Int64Converter::Convert(ctx, obj, &state, &handle, &error);
}
if (error == ErrorFoundFloat) {
status = FloatConverter::Convert(ctx, obj, &state, &handle, &error);
}
// TODO(josh11b): May also want to fall back to using doubles if
// error == ErrorOutOfRange?
} else {
status = Int64Converter::Convert(ctx, obj, &state, &handle, &error);
if (error == ErrorFoundFloat) {
status = DoubleConverter::Convert(ctx, obj, &state, &handle, &error);
}
}
break;
case DT_STRING:
status = StringConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_COMPLEX128:
status = Complex128Converter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_BOOL:
status = BoolConverter::Convert(ctx, obj, &state, &handle, &error);
break;
case DT_INVALID: // Only occurs for empty tensors.
{
AbstractTensorInterface* t = tensorflow::unwrap(ctx)->CreateTensor(
requested_dtype == DT_INVALID ? DT_FLOAT : requested_dtype,
state.inferred_shape);
auto* result =
tensorflow::wrap(tensorflow::unwrap(ctx)->CreateLocalHandle(t));
t->Release();
return result;
}
default:
break;
}
if (!status.ok()) {
PyErr_SetString(PyExc_ValueError, absl::StatusMessageAsCStr(status));
return nullptr;
}
return handle;
}
} // namespace tensorflow
tensorflow/python/lib/core/py_seq_tensor.h
+42
View File
@@ -0,0 +1,42 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PY_SEQ_TENSOR_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PY_SEQ_TENSOR_H_
#include <Python.h>
#include "tensorflow/c/eager/c_api_internal.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
// Converts Python object `obj` representing a rectangular array of
// Python values (a scalar, a sequence of scalars, a sequence of
// sequences, etc.) into a TFE_TensorHandle.
// If dtype is not None it should by a Python integer
// representing the desired dtype of the resulting Tensor.
// This is used only as a hint, *ret may not have that dtype on
// success and may require a cast.
//
// If an error occurs, this return nullptr and sets the python error indicator
// with PyErr_SetString.
TFE_TensorHandle* PySeqToTFE_TensorHandle(TFE_Context* ctx, PyObject* obj,
DataType dtype);
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_PY_SEQ_TENSOR_H_
tensorflow/python/lib/core/py_util.cc
+127
View File
@@ -0,0 +1,127 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/python/lib/core/py_util.h"
// Place `<locale>` before <Python.h> to avoid build failure in macOS.
#include <locale>
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/strings/strcat.h"
namespace tensorflow {
namespace {
// py.__class__.__name__
const char* ClassName(PyObject* py) {
/* PyPy doesn't have a separate C API for old-style classes. */
#if PY_MAJOR_VERSION < 3 && !defined(PYPY_VERSION)
if (PyClass_Check(py))
return PyString_AS_STRING(
CHECK_NOTNULL(reinterpret_cast<PyClassObject*>(py)->cl_name));
if (PyInstance_Check(py))
return PyString_AS_STRING(CHECK_NOTNULL(
reinterpret_cast<PyInstanceObject*>(py)->in_class->cl_name));
#endif
if (Py_TYPE(py) == &PyType_Type) {
return reinterpret_cast<PyTypeObject*>(py)->tp_name;
}
return Py_TYPE(py)->tp_name;
}
} // end namespace
// Returns a PyObject containing a string, or null
void TryAppendTraceback(PyObject* ptype, PyObject* pvalue, PyObject* ptraceback,
std::string* out) {
// The "traceback" module is assumed to be imported already by script_ops.py.
PyObject* tb_module = PyImport_AddModule("traceback");
if (!tb_module) {
return;
}
PyObject* format_exception =
PyObject_GetAttrString(tb_module, "format_exception");
if (!format_exception) {
return;
}
if (!PyCallable_Check(format_exception)) {
Py_DECREF(format_exception);
return;
}
PyObject* ret_val = PyObject_CallFunctionObjArgs(format_exception, ptype,
pvalue, ptraceback, nullptr);
Py_DECREF(format_exception);
if (!ret_val) {
return;
}
if (!PyList_Check(ret_val)) {
Py_DECREF(ret_val);
return;
}
Py_ssize_t n = PyList_GET_SIZE(ret_val);
for (Py_ssize_t i = 0; i < n; ++i) {
PyObject* v = PyList_GET_ITEM(ret_val, i);
#if PY_MAJOR_VERSION < 3
strings::StrAppend(out, PyString_AS_STRING(v), "\n");
#else
absl::StrAppend(out, PyUnicode_AsUTF8(v), "\n");
#endif
}
// Iterate through ret_val.
Py_DECREF(ret_val);
}
std::string PyExceptionFetch() {
CHECK(PyErr_Occurred())
<< "Must only call PyExceptionFetch after an exception.";
PyObject* ptype;
PyObject* pvalue;
PyObject* ptraceback;
PyErr_Fetch(&ptype, &pvalue, &ptraceback);
PyErr_NormalizeException(&ptype, &pvalue, &ptraceback);
std::string err = ClassName(ptype);
if (pvalue) {
PyObject* str = PyObject_Str(pvalue);
if (str) {
#if PY_MAJOR_VERSION < 3
strings::StrAppend(&err, ": ", PyString_AS_STRING(str), "\n");
#else
absl::StrAppend(&err, ": ", PyUnicode_AsUTF8(str), "\n");
#endif
Py_DECREF(str);
} else {
absl::StrAppend(&err, "(unknown error message)\n");
}
TryAppendTraceback(ptype, pvalue, ptraceback, &err);
Py_DECREF(pvalue);
}
Py_DECREF(ptype);
Py_XDECREF(ptraceback);
return err;
}
} // end namespace tensorflow
tensorflow/python/lib/core/py_util.h
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PY_UTIL_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PY_UTIL_H_
#include <Python.h>
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
// Fetch the exception message as a string. An exception must be set
// (PyErr_Occurred() must be true).
std::string PyExceptionFetch();
// Assert that Python GIL is held.
inline void DCheckPyGilState() {
#if PY_MAJOR_VERSION >= 3 && PY_MINOR_VERSION >= 4
DCHECK(PyGILState_Check());
#endif
}
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_PY_UTIL_H_
tensorflow/python/lib/core/pybind11_absl.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_ABSL_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_ABSL_H_
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/platform/stringpiece.h"
#ifndef ABSL_USES_STD_STRING_VIEW
namespace pybind11 {
namespace detail {
// Convert between tensorflow::StringPiece (aka absl::string_view) and Python.
//
// pybind11 supports std::string_view, and absl::string_view is meant to be a
// drop-in replacement for std::string_view, so we can just use the built in
// implementation.
template <>
struct type_caster<tensorflow::StringPiece>
: string_caster<tensorflow::StringPiece, true> {};
} // namespace detail
} // namespace pybind11
#endif // ABSL_USES_STD_STRING_VIEW
#endif // TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_ABSL_H_
tensorflow/python/lib/core/pybind11_lib.h
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_LIB_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_LIB_H_
namespace py = pybind11;
// SWIG struct so pybind11 can handle SWIG objects returned by tf_session
// until that is converted over to pybind11.
// This type is intended to be layout-compatible with an initial sequence of
// certain objects pointed to by a PyObject pointer. The intended use is to
// first check dynamically that a given PyObject* py has the correct type,
// and then use `reinterpret_cast<SwigPyObject*>(py)` to retrieve the member
// `ptr` for further, custom use. SWIG wrapped objects' layout is documented
// here: http://www.swig.org/Doc4.0/Python.html#Python_nn28
typedef struct {
PyObject_HEAD void* ptr; // This is the pointer to the actual C++ obj.
void* ty;
int own;
PyObject* next;
PyObject* dict;
} SwigPyObject;
namespace tensorflow {
// Convert PyObject* to py::object with no error handling.
inline py::object Pyo(PyObject* ptr) {
return py::reinterpret_steal<py::object>(ptr);
}
// Raise an exception if the PyErrOccurred flag is set or else return the Python
// object.
inline py::object PyoOrThrow(PyObject* ptr) {
if (PyErr_Occurred() || ptr == nullptr) {
throw py::error_already_set();
}
return Pyo(ptr);
}
[[noreturn]] inline void ThrowTypeError(const char* error_message) {
PyErr_SetString(PyExc_TypeError, error_message);
throw pybind11::error_already_set();
}
[[noreturn]] inline void ThrowValueError(const char* error_message) {
PyErr_SetString(PyExc_ValueError, error_message);
throw pybind11::error_already_set();
}
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_LIB_H_
tensorflow/python/lib/core/pybind11_proto.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_PROTO_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_PROTO_H_
#include "absl/strings/str_cat.h"
#include "pybind11/pybind11.h" // from @pybind11
namespace tensorflow {
inline void CheckProtoType(const pybind11::handle& py_object,
const std::string expected_proto_type) {
// Check the name of the proto object.
if (pybind11::hasattr(py_object, "DESCRIPTOR")) {
pybind11::handle descriptor = pybind11::getattr(py_object, "DESCRIPTOR");
std::string py_object_type =
pybind11::cast<std::string>(pybind11::getattr(descriptor, "full_name"));
if (py_object_type == expected_proto_type) {
return;
}
throw pybind11::type_error(absl::StrCat("Expected an ", expected_proto_type,
" proto, but got ",
py_object_type));
}
throw pybind11::type_error(absl::StrCat(
std::string(py_object.get_type().str()), " is not a valid proto."));
}
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_PROTO_H_
tensorflow/python/lib/core/pybind11_status.h
+221
View File
@@ -0,0 +1,221 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_STATUS_H_
#define TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_STATUS_H_
#include <Python.h>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "pybind11/cast.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/python/lib/core/py_exception_registry.h"
namespace tsl {
namespace internal {
inline PyObject* CodeToPyExc(const int code) {
switch (code) {
case error::Code::INVALID_ARGUMENT:
return PyExc_ValueError;
case error::Code::OUT_OF_RANGE:
return PyExc_IndexError;
case error::Code::UNIMPLEMENTED:
return PyExc_NotImplementedError;
default:
return PyExc_RuntimeError;
}
}
inline PyObject* StatusToPyExc(const absl::Status& status) {
return CodeToPyExc(status.raw_code());
}
inline PyObject* TFStatusToPyExc(const TF_Status* status) {
return CodeToPyExc(TF_GetCode(status));
}
inline pybind11::dict StatusPayloadToDict(const absl::Status& status) {
pybind11::dict dict;
const auto& payloads = errors::GetPayloads(status);
for (auto& pair : payloads) {
dict[PyBytes_FromString(pair.first.c_str())] =
PyBytes_FromString(pair.second.c_str());
}
return dict;
}
inline pybind11::dict TFStatusPayloadToDict(TF_Status* status) {
return StatusPayloadToDict(status->status);
}
} // namespace internal
inline void MaybeRaiseFromStatus(const absl::Status& status) {
if (!status.ok()) {
absl::string_view status_message = status.message();
pybind11::set_error(
internal::StatusToPyExc(status),
pybind11::str(status_message.data(), status_message.size()));
throw pybind11::error_already_set();
}
}
inline void SetRegisteredErrFromStatus(const absl::Status& status) {
PyErr_SetObject(
tensorflow::PyExceptionRegistry::Lookup(status.raw_code()),
pybind11::make_tuple(pybind11::none(), pybind11::none(), status.message(),
internal::StatusPayloadToDict(status))
.ptr());
}
inline void SetRegisteredErrFromTFStatus(TF_Status* status) {
PyErr_SetObject(tensorflow::PyExceptionRegistry::Lookup(TF_GetCode(status)),
pybind11::make_tuple(pybind11::none(), pybind11::none(),
TF_Message(status),
internal::TFStatusPayloadToDict(status))
.ptr());
}
inline void MaybeRaiseRegisteredFromStatus(const absl::Status& status) {
if (!status.ok()) {
SetRegisteredErrFromStatus(status);
throw pybind11::error_already_set();
}
}
inline void MaybeRaiseRegisteredFromStatusWithGIL(const absl::Status& status) {
if (!status.ok()) {
// Acquire GIL for throwing exception.
pybind11::gil_scoped_acquire acquire;
SetRegisteredErrFromStatus(status);
throw pybind11::error_already_set();
}
}
inline void MaybeRaiseFromTFStatus(TF_Status* status) {
TF_Code code = TF_GetCode(status);
if (code != TF_OK) {
PyErr_SetString(internal::TFStatusToPyExc(status), TF_Message(status));
throw pybind11::error_already_set();
}
}
inline void MaybeRaiseRegisteredFromTFStatus(TF_Status* status) {
TF_Code code = TF_GetCode(status);
if (code != TF_OK) {
SetRegisteredErrFromTFStatus(status);
throw pybind11::error_already_set();
}
}
inline void MaybeRaiseRegisteredFromTFStatusWithGIL(TF_Status* status) {
TF_Code code = TF_GetCode(status);
if (code != TF_OK) {
// Acquire GIL for throwing exception.
pybind11::gil_scoped_acquire acquire;
SetRegisteredErrFromTFStatus(status);
throw pybind11::error_already_set();
}
}
} // namespace tsl
namespace tensorflow {
using tsl::MaybeRaiseFromStatus;
using tsl::MaybeRaiseFromTFStatus;
using tsl::MaybeRaiseRegisteredFromStatus;
using tsl::MaybeRaiseRegisteredFromStatusWithGIL;
using tsl::MaybeRaiseRegisteredFromTFStatus;
using tsl::MaybeRaiseRegisteredFromTFStatusWithGIL;
using tsl::SetRegisteredErrFromStatus;
using tsl::SetRegisteredErrFromTFStatus;
} // namespace tensorflow
namespace pybind11 {
namespace detail {
// Convert tensorflow::Status
//
// Raise an exception if a given status is not OK, otherwise return None.
//
// The correspondence between status codes and exception classes is given
// by PyExceptionRegistry. Note that the registry should be initialized
// in order to be used, see PyExceptionRegistry::Init.
template <>
struct type_caster<absl::Status> {
public:
PYBIND11_TYPE_CASTER(absl::Status, _("Status"));
static handle cast(absl::Status status, return_value_policy, handle) {
tensorflow::MaybeRaiseFromStatus(status);
return none().inc_ref();
}
};
// Convert tensorflow::StatusOr
//
// Uses the same logic as the Abseil implementation: raise an exception if the
// status is not OK, otherwise return its payload.
template <typename PayloadType>
struct type_caster<tensorflow::StatusOr<PayloadType>> {
public:
using PayloadCaster = make_caster<PayloadType>;
using StatusCaster = make_caster<absl::Status>;
static constexpr auto name = PayloadCaster::name;
static handle cast(const tensorflow::StatusOr<PayloadType>* src,
return_value_policy policy, handle parent) {
if (!src) return none().release();
return cast_impl(*src, policy, parent);
}
static handle cast(const tensorflow::StatusOr<PayloadType>& src,
return_value_policy policy, handle parent) {
return cast_impl(src, policy, parent);
}
static handle cast(tensorflow::StatusOr<PayloadType>&& src,
return_value_policy policy, handle parent) {
return cast_impl(std::move(src), policy, parent);
}
private:
template <typename CType>
static handle cast_impl(CType&& src, return_value_policy policy,
handle parent) {
if (src.ok()) {
// Convert and return the payload.
return PayloadCaster::cast(std::forward<CType>(src).value(), policy,
parent);
} else {
// Convert and return the error.
return StatusCaster::cast(std::forward<CType>(src).status(),
return_value_policy::move, parent);
}
}
};
} // namespace detail
} // namespace pybind11
#endif // TENSORFLOW_PYTHON_LIB_CORE_PYBIND11_STATUS_H_
tensorflow/python/lib/core/safe_pyobject_ptr.cc
+24
View File
@@ -0,0 +1,24 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
namespace tensorflow {
Safe_PyObjectPtr make_safe(PyObject* object) {
return Safe_PyObjectPtr(object);
}
} // namespace tensorflow
tensorflow/python/lib/core/safe_pyobject_ptr.h
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_LIB_CORE_SAFE_PYOBJECT_PTR_H_
#define TENSORFLOW_PYTHON_LIB_CORE_SAFE_PYOBJECT_PTR_H_
#include <Python.h>
#include <memory>
namespace tensorflow {
namespace detail {
struct PyDecrefDeleter {
void operator()(PyObject* p) const { Py_DECREF(p); }
};
} // namespace detail
// Safe container for an owned PyObject. On destruction, the reference count of
// the contained object will be decremented.
using Safe_PyObjectPtr = std::unique_ptr<PyObject, detail::PyDecrefDeleter>;
Safe_PyObjectPtr make_safe(PyObject* o);
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_LIB_CORE_SAFE_PYOBJECT_PTR_H_
tensorflow/python/lib/io/BUILD
+166
View File
@@ -0,0 +1,166 @@
# python/lib/io package
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "if_oss")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test", "tf_python_pybind_extension")
# copybara:uncomment_begin(google-only)
# load("//third_party/zlib:BUILD_defs.bzl", "brittle_test_relying_on_stable_zlib_output")
# copybara:uncomment_end
visibility = [
"//tensorflow:__subpackages__",
]
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = visibility,
licenses = ["notice"],
)
tf_python_pybind_extension(
name = "_pywrap_file_io",
srcs = ["file_io_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_file_io.pyi",
],
deps = [
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:portable_jpeg_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:file_statistics",
"//tensorflow/python/lib/core:pybind11_absl",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@pybind11",
],
)
tf_python_pybind_extension(
name = "_pywrap_record_io",
srcs = ["record_io_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_record_io.pyi",
],
deps = [
"//tensorflow/core:core_stringpiece",
"//tensorflow/core:lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:types",
"//tensorflow/python/lib/core:pybind11_absl",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@pybind11",
],
)
py_library(
name = "file_io",
srcs = ["file_io.py"],
strict_deps = True,
# copybara:uncomment_begin(google-only)
# visibility = [
# "//learning/serving/servables/wiz/lora:__subpackages__",
# "//third_party/cloud_tpu/convergence_tools:__subpackages__",
# "//third_party/py/tf_slim:__subpackages__",
# "//tensorflow:__subpackages__",
# "//tensorflow:internal",
# ],
# copybara:uncomment_end_and_comment_begin
visibility = [
"//visibility:public",
],
# copybara:comment_end
deps = [
":_pywrap_file_io",
"//tensorflow/python/framework:errors",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
"@six_archive//:six",
],
)
py_library(
name = "tf_record",
srcs = ["tf_record.py"],
strict_deps = True,
visibility = [
"//tensorflow:__subpackages__",
"//tensorflow:internal",
"//third_party/py/tf_agents:__subpackages__",
],
deps = [
":_pywrap_record_io",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "python_io",
srcs = ["python_io.py"],
strict_deps = True,
visibility = [
"//tensorflow:__subpackages__",
"//tensorflow:internal",
],
deps = [":tf_record"],
)
tf_py_strict_test(
name = "file_io_test",
size = "small",
srcs = ["file_io_test.py"],
tags = [
"no_rocm",
"no_windows",
],
deps = [
":file_io",
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"@absl_py//absl/testing:parameterized",
],
)
tf_py_strict_test(
name = if_oss("tf_record_test", "_tf_record_test"),
size = "small",
srcs = ["tf_record_test.py"],
main = "tf_record_test.py",
tags = [
# copybara:uncomment_begin(google-only)
# "manual",
# "notap",
# copybara:uncomment_end
],
deps = [
":tf_record",
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/util:compat",
],
)
# _tf_record_test relies on stable zlib output only during its execution.
# It should not be brittle with respect to upgrades of zlib software or
# different computers using different zlib software.
# copybara:uncomment_begin(google-only)
# brittle_test_relying_on_stable_zlib_output(
# name = "tf_record_test",
# test_target = ":_tf_record_test",
# )
# copybara:uncomment_end
tensorflow/python/lib/io/__init__.py
View File
tensorflow/python/lib/io/_pywrap_file_io.pyi
+53
View File
@@ -0,0 +1,53 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
class BufferedInputStream:
def __init__(self, filename: str, buffer_size: int) -> None: ...
def read(self, arg0: int) -> bytes: ...
def readline(self) -> bytes: ...
def seek(self, arg0: int) -> None: ...
def tell(self) -> int: ...
class FileStatistics:
def __init__(self, *args, **kwargs) -> None: ...
@property
def is_directory(self) -> bool: ...
@property
def length(self) -> int: ...
@property
def mtime_nsec(self) -> int: ...
class WritableFile:
def __init__(self, filename: str, mode: str) -> None: ...
def append(self, arg0: str) -> None: ...
def close(self) -> None: ...
def flush(self) -> None: ...
def tell(self) -> int: ...
def CopyFile(src: str, target: str, overwrite: bool) -> None: ...
def CreateDir(dirname: str) -> None: ...
def DeleteFile(filename: str) -> None: ...
def DeleteRecursively(dirname: str) -> None: ...
def FileExists(filename: str) -> None: ...
def GetChildren(dirname: str) -> list[str]: ...
def GetMatchingFiles(pattern: str) -> list[str]: ...
def GetRegisteredSchemes() -> list[str]: ...
def HasAtomicMove(arg0: str) -> bool: ...
def IsDirectory(dirname: str) -> bool: ...
def ReadFileToString(filename: str) -> bytes: ...
def RecursivelyCreateDir(dirname: str) -> None: ...
def RenameFile(src: str, target: str, overwrite: bool) -> None: ...
def Stat(filename: str) -> FileStatistics: ...
def WriteStringToFile(filename: str, data: str) -> None: ...
tensorflow/python/lib/io/_pywrap_record_io.pyi
+52
View File
@@ -0,0 +1,52 @@
# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
class RandomRecordReader:
def __init__(self, arg0: str) -> None: ...
def close(self) -> None: ...
def read(self, arg0: int) -> tuple: ...
class RecordIterator:
def __init__(self, arg0: str, arg1: str) -> None: ...
def close(self) -> None: ...
def reopen(self) -> None: ...
def __iter__(self) -> object: ...
def __next__(self) -> bytes: ...
class RecordWriter:
def __init__(self, arg0: str, arg1: RecordWriterOptions) -> None: ...
def close(self) -> None: ...
def flush(self) -> None: ...
def write(self, record: str) -> None: ...
def __enter__(self) -> object: ...
def __exit__(self, *args) -> None: ...
class RecordWriterOptions:
def __init__(self, arg0: str) -> None: ...
@property
def compression_type(self): ...
@property
def zlib_options(self) -> ZlibCompressionOptions: ...
class ZlibCompressionOptions:
compression_level: int
compression_method: int
compression_strategy: int
flush_mode: int
input_buffer_size: int
mem_level: int
output_buffer_size: int
window_bits: int
def __init__(self, *args, **kwargs) -> None: ...
tensorflow/python/lib/io/file_io.py
+1004
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/lib/io/file_io_test.py
+698
View File
@@ -0,0 +1,698 @@
# This Python file uses the following encoding: utf-8
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================
"""Testing File IO operations in file_io.py."""
import os.path
from absl.testing import parameterized
import numpy as np
from tensorflow.python.framework import errors
from tensorflow.python.lib.io import file_io
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
class PathLike(object):
"""Backport of pathlib.Path for Python < 3.6"""
def __init__(self, name):
self.name = name
def __fspath__(self):
return self.name
def __str__(self):
return self.name
run_all_path_types = parameterized.named_parameters(
("str", file_io.join),
("pathlike", lambda *paths: PathLike(file_io.join(*paths))))
class FileIoTest(test.TestCase, parameterized.TestCase):
def setUp(self):
self._base_dir = file_io.join(self.get_temp_dir(), "base_dir")
file_io.create_dir(self._base_dir)
def tearDown(self):
file_io.delete_recursively(self._base_dir)
def testEmptyFilename(self):
f = file_io.FileIO("", mode="r")
with self.assertRaises(errors.NotFoundError):
_ = f.read()
def testJoinUrlLike(self):
"""file_io.join joins url-like filesystems with '/' on all platform."""
for fs in ("ram://", "gcs://", "file://"):
expected = fs + "exists/a/b/c.txt"
self.assertEqual(file_io.join(fs, "exists", "a", "b", "c.txt"), expected)
self.assertEqual(file_io.join(fs + "exists", "a", "b", "c.txt"), expected)
self.assertEqual(file_io.join(fs, "exists/a", "b", "c.txt"), expected)
self.assertEqual(file_io.join(fs, "exists", "a", "b/c.txt"), expected)
def testJoinFilesystem(self):
"""file_io.join respects the os.path.join behavior for native filesystems."""
for sep in ("/", "\\", os.sep):
self.assertEqual(os.path.join("a", "b", "c"), file_io.join("a", "b", "c"))
self.assertEqual(
os.path.join(sep + "a", "b", "c"), file_io.join(sep + "a", "b", "c"))
self.assertEqual(
os.path.join("a", sep + "b", "c"), file_io.join("a", sep + "b", "c"))
self.assertEqual(
os.path.join("a", "b", sep + "c"), file_io.join("a", "b", sep + "c"))
self.assertEqual(
os.path.join("a", "b", "c" + sep), file_io.join("a", "b", "c" + sep))
@run_all_path_types
def testFileDoesntExist(self, join):
file_path = join(self._base_dir, "temp_file")
self.assertFalse(file_io.file_exists(file_path))
with self.assertRaises(errors.NotFoundError):
_ = file_io.read_file_to_string(file_path)
@run_all_path_types
def testWriteToString(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.write_string_to_file(file_path, "testing")
self.assertTrue(file_io.file_exists(file_path))
file_contents = file_io.read_file_to_string(file_path)
self.assertEqual("testing", file_contents)
def testAtomicWriteStringToFile(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.atomic_write_string_to_file(file_path, "testing")
self.assertTrue(file_io.file_exists(file_path))
file_contents = file_io.read_file_to_string(file_path)
self.assertEqual("testing", file_contents)
def testAtomicWriteStringToFileOverwriteFalse(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.atomic_write_string_to_file(file_path, "old", overwrite=False)
with self.assertRaises(errors.AlreadyExistsError):
file_io.atomic_write_string_to_file(file_path, "new", overwrite=False)
file_contents = file_io.read_file_to_string(file_path)
self.assertEqual("old", file_contents)
file_io.delete_file(file_path)
file_io.atomic_write_string_to_file(file_path, "new", overwrite=False)
file_contents = file_io.read_file_to_string(file_path)
self.assertEqual("new", file_contents)
@run_all_path_types
def testReadBinaryMode(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.write_string_to_file(file_path, "testing")
with file_io.FileIO(file_path, mode="rb") as f:
self.assertEqual(b"testing", f.read())
@run_all_path_types
def testWriteBinaryMode(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.FileIO(file_path, "wb").write("testing")
with file_io.FileIO(file_path, mode="r") as f:
self.assertEqual("testing", f.read())
def testAppend(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="w") as f:
f.write("begin\n")
with file_io.FileIO(file_path, mode="a") as f:
f.write("a1\n")
with file_io.FileIO(file_path, mode="a") as f:
f.write("a2\n")
with file_io.FileIO(file_path, mode="r") as f:
file_contents = f.read()
self.assertEqual("begin\na1\na2\n", file_contents)
def testMultipleFiles(self):
file_prefix = file_io.join(self._base_dir, "temp_file")
for i in range(5000):
f = file_io.FileIO(file_prefix + str(i), mode="w+")
f.write("testing")
f.flush()
self.assertEqual("testing", f.read())
f.close()
def testMultipleWrites(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="w") as f:
f.write("line1\n")
f.write("line2")
file_contents = file_io.read_file_to_string(file_path)
self.assertEqual("line1\nline2", file_contents)
def testFileWriteBadMode(self):
file_path = file_io.join(self._base_dir, "temp_file")
with self.assertRaises(errors.PermissionDeniedError):
file_io.FileIO(file_path, mode="r").write("testing")
def testFileReadBadMode(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
self.assertTrue(file_io.file_exists(file_path))
with self.assertRaises(errors.PermissionDeniedError):
file_io.FileIO(file_path, mode="w").read()
@run_all_path_types
def testFileDelete(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
file_io.delete_file(file_path)
self.assertFalse(file_io.file_exists(file_path))
def testFileDeleteFail(self):
file_path = file_io.join(self._base_dir, "temp_file")
with self.assertRaises(errors.NotFoundError):
file_io.delete_file(file_path)
def testGetMatchingFiles(self):
dir_path = file_io.join(self._base_dir, "temp_dir")
file_io.create_dir(dir_path)
files = ["file1.txt", "file2.txt", "file3.txt", "file*.txt"]
for name in files:
file_path = file_io.join(dir_path, name)
file_io.FileIO(file_path, mode="w").write("testing")
expected_match = [file_io.join(dir_path, name) for name in files]
self.assertItemsEqual(
file_io.get_matching_files(file_io.join(dir_path, "file*.txt")),
expected_match)
self.assertItemsEqual(file_io.get_matching_files(tuple()), [])
files_subset = [
file_io.join(dir_path, files[0]),
file_io.join(dir_path, files[2])
]
self.assertItemsEqual(
file_io.get_matching_files(files_subset), files_subset)
file_io.delete_recursively(dir_path)
self.assertFalse(file_io.file_exists(file_io.join(dir_path, "file3.txt")))
def testGetMatchingFilesWhenParentDirContainsParantheses(self):
dir_path = file_io.join(self._base_dir, "dir_(special)")
file_io.create_dir(dir_path)
files = ["file1.txt", "file(2).txt"]
for name in files:
file_path = file_io.join(dir_path, name)
file_io.FileIO(file_path, mode="w").write("testing")
expected_match = [file_io.join(dir_path, name) for name in files]
glob_pattern = file_io.join(dir_path, "*")
self.assertItemsEqual(
file_io.get_matching_files(glob_pattern), expected_match)
@run_all_path_types
def testCreateRecursiveDir(self, join):
dir_path = join(self._base_dir, "temp_dir/temp_dir1/temp_dir2")
file_io.recursive_create_dir(dir_path)
file_io.recursive_create_dir(dir_path) # repeat creation
file_path = file_io.join(str(dir_path), "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
self.assertTrue(file_io.file_exists(file_path))
file_io.delete_recursively(file_io.join(self._base_dir, "temp_dir"))
self.assertFalse(file_io.file_exists(file_path))
@run_all_path_types
def testCopy(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
copy_path = join(self._base_dir, "copy_file")
file_io.copy(file_path, copy_path)
self.assertTrue(file_io.file_exists(copy_path))
f = file_io.FileIO(file_path, mode="r")
self.assertEqual("testing", f.read())
self.assertEqual(7, f.tell())
def testCopyOverwrite(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
copy_path = file_io.join(self._base_dir, "copy_file")
file_io.FileIO(copy_path, mode="w").write("copy")
file_io.copy(file_path, copy_path, overwrite=True)
self.assertTrue(file_io.file_exists(copy_path))
self.assertEqual("testing", file_io.FileIO(file_path, mode="r").read())
def testCopyOverwriteFalse(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
copy_path = file_io.join(self._base_dir, "copy_file")
file_io.FileIO(copy_path, mode="w").write("copy")
with self.assertRaises(errors.AlreadyExistsError):
file_io.copy(file_path, copy_path, overwrite=False)
@run_all_path_types
def testRename(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
rename_path = join(self._base_dir, "rename_file")
file_io.rename(file_path, rename_path)
self.assertTrue(file_io.file_exists(rename_path))
self.assertFalse(file_io.file_exists(file_path))
def testRenameOverwrite(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
rename_path = file_io.join(self._base_dir, "rename_file")
file_io.FileIO(rename_path, mode="w").write("rename")
file_io.rename(file_path, rename_path, overwrite=True)
self.assertTrue(file_io.file_exists(rename_path))
self.assertFalse(file_io.file_exists(file_path))
def testRenameOverwriteFalse(self):
file_path = file_io.join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
rename_path = file_io.join(self._base_dir, "rename_file")
file_io.FileIO(rename_path, mode="w").write("rename")
with self.assertRaises(errors.AlreadyExistsError):
file_io.rename(file_path, rename_path, overwrite=False)
self.assertTrue(file_io.file_exists(rename_path))
self.assertTrue(file_io.file_exists(file_path))
def testDeleteRecursivelyFail(self):
fake_dir_path = file_io.join(self._base_dir, "temp_dir")
with self.assertRaises(errors.NotFoundError):
file_io.delete_recursively(fake_dir_path)
@run_all_path_types
def testIsDirectory(self, join):
dir_path = join(self._base_dir, "test_dir")
# Failure for a non-existing dir.
self.assertFalse(file_io.is_directory(dir_path))
file_io.create_dir(dir_path)
self.assertTrue(file_io.is_directory(dir_path))
file_path = join(str(dir_path), "test_file")
file_io.FileIO(file_path, mode="w").write("test")
# False for a file.
self.assertFalse(file_io.is_directory(file_path))
# Test that the value returned from `stat()` has `is_directory` set.
file_statistics = file_io.stat(dir_path)
self.assertTrue(file_statistics.is_directory)
@run_all_path_types
def testListDirectory(self, join):
dir_path = join(self._base_dir, "test_dir")
file_io.create_dir(dir_path)
files = ["file1.txt", "file2.txt", "file3.txt"]
for name in files:
file_path = join(str(dir_path), name)
file_io.FileIO(file_path, mode="w").write("testing")
subdir_path = join(str(dir_path), "sub_dir")
file_io.create_dir(subdir_path)
subdir_file_path = join(str(subdir_path), "file4.txt")
file_io.FileIO(subdir_file_path, mode="w").write("testing")
dir_list = file_io.list_directory(dir_path)
self.assertItemsEqual(files + ["sub_dir"], dir_list)
def testListDirectoryFailure(self):
dir_path = file_io.join(self._base_dir, "test_dir")
with self.assertRaises(errors.NotFoundError):
file_io.list_directory(dir_path)
def _setupWalkDirectories(self, dir_path):
# Creating a file structure as follows
# test_dir -> file: file1.txt; dirs: subdir1_1, subdir1_2, subdir1_3
# subdir1_1 -> file: file3.txt
# subdir1_2 -> dir: subdir2
file_io.create_dir(dir_path)
file_io.FileIO(
file_io.join(dir_path, "file1.txt"), mode="w").write("testing")
sub_dirs1 = ["subdir1_1", "subdir1_2", "subdir1_3"]
for name in sub_dirs1:
file_io.create_dir(file_io.join(dir_path, name))
file_io.FileIO(
file_io.join(dir_path, "subdir1_1/file2.txt"),
mode="w").write("testing")
file_io.create_dir(file_io.join(dir_path, "subdir1_2/subdir2"))
@run_all_path_types
def testWalkInOrder(self, join):
dir_path_str = file_io.join(self._base_dir, "test_dir")
dir_path = join(self._base_dir, "test_dir")
self._setupWalkDirectories(dir_path_str)
# Now test the walk (in_order = True)
all_dirs = []
all_subdirs = []
all_files = []
for (w_dir, w_subdirs, w_files) in file_io.walk(dir_path, in_order=True):
all_dirs.append(w_dir)
all_subdirs.append(w_subdirs)
all_files.append(w_files)
self.assertItemsEqual(all_dirs, [dir_path_str] + [
file_io.join(dir_path_str, item) for item in
["subdir1_1", "subdir1_2", "subdir1_2/subdir2", "subdir1_3"]
])
self.assertEqual(dir_path_str, all_dirs[0])
self.assertLess(
all_dirs.index(file_io.join(dir_path_str, "subdir1_2")),
all_dirs.index(file_io.join(dir_path_str, "subdir1_2/subdir2")))
self.assertItemsEqual(all_subdirs[1:5], [[], ["subdir2"], [], []])
self.assertItemsEqual(all_subdirs[0],
["subdir1_1", "subdir1_2", "subdir1_3"])
self.assertItemsEqual(all_files, [["file1.txt"], ["file2.txt"], [], [], []])
self.assertLess(
all_files.index(["file1.txt"]), all_files.index(["file2.txt"]))
def testWalkPostOrder(self):
dir_path = file_io.join(self._base_dir, "test_dir")
self._setupWalkDirectories(dir_path)
# Now test the walk (in_order = False)
all_dirs = []
all_subdirs = []
all_files = []
for (w_dir, w_subdirs, w_files) in file_io.walk(dir_path, in_order=False):
all_dirs.append(w_dir)
all_subdirs.append(w_subdirs)
all_files.append(w_files)
self.assertItemsEqual(all_dirs, [
file_io.join(dir_path, item) for item in
["subdir1_1", "subdir1_2/subdir2", "subdir1_2", "subdir1_3"]
] + [dir_path])
self.assertEqual(dir_path, all_dirs[4])
self.assertLess(
all_dirs.index(file_io.join(dir_path, "subdir1_2/subdir2")),
all_dirs.index(file_io.join(dir_path, "subdir1_2")))
self.assertItemsEqual(all_subdirs[0:4], [[], [], ["subdir2"], []])
self.assertItemsEqual(all_subdirs[4],
["subdir1_1", "subdir1_2", "subdir1_3"])
self.assertItemsEqual(all_files, [["file2.txt"], [], [], [], ["file1.txt"]])
self.assertLess(
all_files.index(["file2.txt"]), all_files.index(["file1.txt"]))
def testWalkFailure(self):
dir_path = file_io.join(self._base_dir, "test_dir")
# Try walking a directory that wasn't created.
all_dirs = []
all_subdirs = []
all_files = []
for (w_dir, w_subdirs, w_files) in file_io.walk(dir_path, in_order=False):
all_dirs.append(w_dir)
all_subdirs.append(w_subdirs)
all_files.append(w_files)
self.assertItemsEqual(all_dirs, [])
self.assertItemsEqual(all_subdirs, [])
self.assertItemsEqual(all_files, [])
@run_all_path_types
def testStat(self, join):
file_path = join(self._base_dir, "temp_file")
file_io.FileIO(file_path, mode="w").write("testing")
file_statistics = file_io.stat(file_path)
os_statistics = os.stat(str(file_path))
self.assertEqual(7, file_statistics.length)
self.assertEqual(
int(os_statistics.st_mtime), int(file_statistics.mtime_nsec / 1e9))
self.assertFalse(file_statistics.is_directory)
def testReadLine(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="r+") as f:
f.write("testing1\ntesting2\ntesting3\n\ntesting5")
self.assertEqual(36, f.size())
self.assertEqual("testing1\n", f.readline())
self.assertEqual("testing2\n", f.readline())
self.assertEqual("testing3\n", f.readline())
self.assertEqual("\n", f.readline())
self.assertEqual("testing5", f.readline())
self.assertEqual("", f.readline())
def testRead(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="r+") as f:
f.write("testing1\ntesting2\ntesting3\n\ntesting5")
self.assertEqual(36, f.size())
self.assertEqual("testing1\n", f.read(9))
self.assertEqual("testing2\n", f.read(9))
self.assertEqual("t", f.read(1))
self.assertEqual("esting3\n\ntesting5", f.read())
def testReadErrorReacquiresGil(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="r+") as f:
f.write("testing1\ntesting2\ntesting3\n\ntesting5")
with self.assertRaises(errors.InvalidArgumentError):
# At present, this is sufficient to convince ourselves that the change
# fixes the problem. That is, this test will seg fault without the change,
# and pass with it. Unfortunately, this is brittle, as it relies on the
# Python layer to pass the argument along to the wrapped C++ without
# checking the argument itself.
f.read(-2)
def testTell(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="r+") as f:
f.write("testing1\ntesting2\ntesting3\n\ntesting5")
self.assertEqual(0, f.tell())
self.assertEqual("testing1\n", f.readline())
self.assertEqual(9, f.tell())
self.assertEqual("testing2\n", f.readline())
self.assertEqual(18, f.tell())
self.assertEqual("testing3\n", f.readline())
self.assertEqual(27, f.tell())
self.assertEqual("\n", f.readline())
self.assertEqual(28, f.tell())
self.assertEqual("testing5", f.readline())
self.assertEqual(36, f.tell())
self.assertEqual("", f.readline())
self.assertEqual(36, f.tell())
def testSeek(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="r+") as f:
f.write("testing1\ntesting2\ntesting3\n\ntesting5")
self.assertEqual("testing1\n", f.readline())
self.assertEqual(9, f.tell())
# Seek to 18
f.seek(18)
self.assertEqual(18, f.tell())
self.assertEqual("testing3\n", f.readline())
# Seek back to 9
f.seek(9)
self.assertEqual(9, f.tell())
self.assertEqual("testing2\n", f.readline())
f.seek(0)
self.assertEqual(0, f.tell())
self.assertEqual("testing1\n", f.readline())
with self.assertRaises(errors.InvalidArgumentError):
f.seek(-1)
with self.assertRaises(TypeError):
f.seek()
# TODO(jhseu): Delete after position deprecation.
with self.assertRaises(TypeError):
f.seek(offset=0, position=0)
f.seek(position=9)
self.assertEqual(9, f.tell())
self.assertEqual("testing2\n", f.readline())
def testSeekFromWhat(self):
file_path = file_io.join(self._base_dir, "temp_file")
with file_io.FileIO(file_path, mode="r+") as f:
f.write("testing1\ntesting2\ntesting3\n\ntesting5")
self.assertEqual("testing1\n", f.readline())
self.assertEqual(9, f.tell())
# Seek to 18
f.seek(9, 1)
self.assertEqual(18, f.tell())
self.assertEqual("testing3\n", f.readline())
# Seek back to 9
f.seek(9, 0)
self.assertEqual(9, f.tell())
self.assertEqual("testing2\n", f.readline())
f.seek(-f.size(), 2)
self.assertEqual(0, f.tell())
self.assertEqual("testing1\n", f.readline())
with self.assertRaises(errors.InvalidArgumentError):
f.seek(0, 3)
def testReadingIterator(self):
file_path = file_io.join(self._base_dir, "temp_file")
data = ["testing1\n", "testing2\n", "testing3\n", "\n", "testing5"]
with file_io.FileIO(file_path, mode="r+") as f:
f.write("".join(data))
actual_data = []
for line in f:
actual_data.append(line)
self.assertSequenceEqual(actual_data, data)
def testReadlines(self):
file_path = file_io.join(self._base_dir, "temp_file")
data = ["testing1\n", "testing2\n", "testing3\n", "\n", "testing5"]
f = file_io.FileIO(file_path, mode="r+")
f.write("".join(data))
f.flush()
lines = f.readlines()
self.assertSequenceEqual(lines, data)
def testUTF8StringPath(self):
file_path = file_io.join(self._base_dir, "UTF8测试_file")
file_io.write_string_to_file(file_path, "testing")
with file_io.FileIO(file_path, mode="rb") as f:
self.assertEqual(b"testing", f.read())
def testEof(self):
"""Test that reading past EOF does not raise an exception."""
file_path = file_io.join(self._base_dir, "temp_file")
f = file_io.FileIO(file_path, mode="r+")
content = "testing"
f.write(content)
f.flush()
self.assertEqual(content, f.read(len(content) + 1))
@run_all_path_types
def testUTF8StringPathExists(self, join):
file_path = join(self._base_dir, "UTF8测试_file_exist")
file_io.write_string_to_file(file_path, "testing")
v = file_io.file_exists(file_path)
self.assertEqual(v, True)
def testFilecmp(self):
file1 = file_io.join(self._base_dir, "file1")
file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
file2 = file_io.join(self._base_dir, "file2")
file_io.write_string_to_file(file2, "This is another sentence\n" * 100)
file3 = file_io.join(self._base_dir, "file3")
file_io.write_string_to_file(file3, u"This is another sentence\n" * 100)
self.assertFalse(file_io.filecmp(file1, file2))
self.assertTrue(file_io.filecmp(file2, file3))
def testFilecmpSameSize(self):
file1 = file_io.join(self._base_dir, "file1")
file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
file2 = file_io.join(self._base_dir, "file2")
file_io.write_string_to_file(file2, "This is b sentence\n" * 100)
file3 = file_io.join(self._base_dir, "file3")
file_io.write_string_to_file(file3, u"This is b sentence\n" * 100)
self.assertFalse(file_io.filecmp(file1, file2))
self.assertTrue(file_io.filecmp(file2, file3))
def testFilecmpBinary(self):
file1 = file_io.join(self._base_dir, "file1")
file_io.FileIO(file1, "wb").write("testing\n\na")
file2 = file_io.join(self._base_dir, "file2")
file_io.FileIO(file2, "wb").write("testing\n\nb")
file3 = file_io.join(self._base_dir, "file3")
file_io.FileIO(file3, "wb").write("testing\n\nb")
file4 = file_io.join(self._base_dir, "file4")
file_io.FileIO(file4, "wb").write("testing\n\ntesting")
self.assertFalse(file_io.filecmp(file1, file2))
self.assertFalse(file_io.filecmp(file1, file4))
self.assertTrue(file_io.filecmp(file2, file3))
def testFileCrc32(self):
file1 = file_io.join(self._base_dir, "file1")
file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
crc1 = file_io.file_crc32(file1)
file2 = file_io.join(self._base_dir, "file2")
file_io.write_string_to_file(file2, "This is another sentence\n" * 100)
crc2 = file_io.file_crc32(file2)
file3 = file_io.join(self._base_dir, "file3")
file_io.write_string_to_file(file3, "This is another sentence\n" * 100)
crc3 = file_io.file_crc32(file3)
self.assertTrue(crc1 != crc2)
self.assertEqual(crc2, crc3)
def testFileCrc32WithBytes(self):
file1 = file_io.join(self._base_dir, "file1")
file_io.write_string_to_file(file1, "This is a sentence\n" * 100)
crc1 = file_io.file_crc32(file1, block_size=24)
file2 = file_io.join(self._base_dir, "file2")
file_io.write_string_to_file(file2, "This is another sentence\n" * 100)
crc2 = file_io.file_crc32(file2, block_size=24)
file3 = file_io.join(self._base_dir, "file3")
file_io.write_string_to_file(file3, "This is another sentence\n" * 100)
crc3 = file_io.file_crc32(file3, block_size=-1)
self.assertTrue(crc1 != crc2)
self.assertEqual(crc2, crc3)
def testFileCrc32Binary(self):
file1 = file_io.join(self._base_dir, "file1")
file_io.FileIO(file1, "wb").write("testing\n\n")
crc1 = file_io.file_crc32(file1)
file2 = file_io.join(self._base_dir, "file2")
file_io.FileIO(file2, "wb").write("testing\n\n\n")
crc2 = file_io.file_crc32(file2)
file3 = file_io.join(self._base_dir, "file3")
file_io.FileIO(file3, "wb").write("testing\n\n\n")
crc3 = file_io.file_crc32(file3)
self.assertTrue(crc1 != crc2)
self.assertEqual(crc2, crc3)
def testFileSeekableWithZip(self):
# Note: Test case for GitHub issue 27276, issue only exposed in python 3.7+.
filename = file_io.join(self._base_dir, "a.npz")
np.savez_compressed(filename, {"a": 1, "b": 2})
with gfile.GFile(filename, "rb") as f:
info = np.load(f, allow_pickle=True) # pylint: disable=unexpected-keyword-arg
_ = [i for i in info.items()]
def testHasAtomicMove(self):
self.assertTrue(file_io.has_atomic_move("/a/b/c"))
def testGetRegisteredSchemes(self):
expected = ["", "file", "ram"]
actual = file_io.get_registered_schemes()
# Be flexible about additional schemes that may sometimes be registered when
# this test is run, while still verifying each scheme appears just once.
maybe_expected = ["gs", "hypercomputer"]
for scheme in maybe_expected:
if scheme in actual:
expected.append(scheme)
self.assertCountEqual(expected, actual)
def testReadWriteWithEncoding(self):
file_path = file_io.join(self._base_dir, "temp_file")
with open(file_path, mode="w", encoding="cp932") as f:
f.write("今日はいい天気")
with file_io.FileIO(file_path, mode="r", encoding="cp932") as f:
self.assertEqual(f.read(), "今日はいい天気")
with file_io.FileIO(file_path, mode="w", encoding="cp932") as f:
f.write("今日はいい天気")
with open(file_path, mode="r", encoding="cp932") as f:
self.assertEqual(f.read(), "今日はいい天気")
if __name__ == "__main__":
test.main()
tensorflow/python/lib/io/file_io_wrapper.cc
+310
View File
@@ -0,0 +1,310 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/core/lib/core/error_codes.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/io/buffered_inputstream.h"
#include "tensorflow/core/lib/io/random_inputstream.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_statistics.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/python/lib/core/pybind11_absl.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace tensorflow {
} // namespace tensorflow
namespace {
namespace py = pybind11;
PYBIND11_MODULE(_pywrap_file_io, m) {
m.def(
"FileExists",
[](const std::string& filename) {
absl::Status status;
{
py::gil_scoped_release release;
status = tensorflow::Env::Default()->FileExists(filename);
}
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("filename"));
m.def(
"DeleteFile",
[](const std::string& filename) {
py::gil_scoped_release release;
absl::Status status = tensorflow::Env::Default()->DeleteFile(filename);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("filename"));
m.def(
"ReadFileToString",
[](const std::string& filename) {
std::string data;
py::gil_scoped_release release;
const auto status =
ReadFileToString(tensorflow::Env::Default(), filename, &data);
pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status);
return py::bytes(data);
},
py::arg("filename"));
m.def(
"WriteStringToFile",
[](const std::string& filename, absl::string_view data) {
py::gil_scoped_release release;
const auto status =
WriteStringToFile(tensorflow::Env::Default(), filename, data);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("filename"), py::arg("data"));
m.def(
"GetChildren",
[](const std::string& dirname) {
std::vector<std::string> results;
py::gil_scoped_release release;
const auto status =
tensorflow::Env::Default()->GetChildren(dirname, &results);
pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status);
return results;
},
py::arg("dirname"));
m.def(
"GetMatchingFiles",
[](const std::string& pattern) {
std::vector<std::string> results;
py::gil_scoped_release release;
const auto status =
tensorflow::Env::Default()->GetMatchingPaths(pattern, &results);
pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status);
return results;
},
py::arg("pattern"));
m.def(
"CreateDir",
[](const std::string& dirname) {
py::gil_scoped_release release;
const auto status = tensorflow::Env::Default()->CreateDir(dirname);
if (absl::IsAlreadyExists(status)) {
return;
}
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("dirname"));
m.def(
"RecursivelyCreateDir",
[](const std::string& dirname) {
py::gil_scoped_release release;
const auto status =
tensorflow::Env::Default()->RecursivelyCreateDir(dirname);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("dirname"));
m.def(
"CopyFile",
[](const std::string& src, const std::string& target, bool overwrite) {
py::gil_scoped_release release;
auto* env = tensorflow::Env::Default();
absl::Status status;
if (!overwrite && env->FileExists(target).ok()) {
status = absl::AlreadyExistsError("file already exists");
} else {
status = env->CopyFile(src, target);
}
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("src"), py::arg("target"), py::arg("overwrite"));
m.def(
"RenameFile",
[](const std::string& src, const std::string& target, bool overwrite) {
py::gil_scoped_release release;
auto* env = tensorflow::Env::Default();
absl::Status status;
if (!overwrite && env->FileExists(target).ok()) {
status = absl::AlreadyExistsError("file already exists");
} else {
status = env->RenameFile(src, target);
}
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("src"), py::arg("target"), py::arg("overwrite"));
m.def(
"DeleteRecursively",
[](const std::string& dirname) {
py::gil_scoped_release release;
int64_t undeleted_files;
int64_t undeleted_dirs;
auto status = tensorflow::Env::Default()->DeleteRecursively(
dirname, &undeleted_files, &undeleted_dirs);
if (status.ok() && (undeleted_files > 0 || undeleted_dirs > 0)) {
status = absl::PermissionDeniedError("could not fully delete dir");
}
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
},
py::arg("dirname"));
m.def(
"IsDirectory",
[](const std::string& dirname) {
py::gil_scoped_release release;
const auto status = tensorflow::Env::Default()->IsDirectory(dirname);
// FAILED_PRECONDITION response means path exists but isn't a dir.
if (absl::IsFailedPrecondition(status)) {
return false;
}
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
return true;
},
py::arg("dirname"));
m.def("HasAtomicMove", [](const std::string& path) {
py::gil_scoped_release release;
bool has_atomic_move;
const auto status =
tensorflow::Env::Default()->HasAtomicMove(path, &has_atomic_move);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
return has_atomic_move;
});
py::class_<tensorflow::FileStatistics>(m, "FileStatistics")
.def_readonly("length", &tensorflow::FileStatistics::length)
.def_readonly("mtime_nsec", &tensorflow::FileStatistics::mtime_nsec)
.def_readonly("is_directory", &tensorflow::FileStatistics::is_directory);
m.def(
"Stat",
[](const std::string& filename) {
py::gil_scoped_release release;
std::unique_ptr<tensorflow::FileStatistics> self(
new tensorflow::FileStatistics);
const auto status =
tensorflow::Env::Default()->Stat(filename, self.get());
py::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status);
return self.release();
},
py::arg("filename"));
m.def("GetRegisteredSchemes", []() {
std::vector<std::string> results;
py::gil_scoped_release release;
const auto status =
tensorflow::Env::Default()->GetRegisteredFileSystemSchemes(&results);
pybind11::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status);
return results;
});
using tensorflow::WritableFile;
py::class_<WritableFile>(m, "WritableFile")
.def(py::init([](const std::string& filename, const std::string& mode) {
py::gil_scoped_release release;
auto* env = tensorflow::Env::Default();
std::unique_ptr<WritableFile> self;
const auto status = mode.find('a') == std::string::npos
? env->NewWritableFile(filename, &self)
: env->NewAppendableFile(filename, &self);
py::gil_scoped_acquire acquire;
tensorflow::MaybeRaiseRegisteredFromStatus(status);
return self.release();
}),
py::arg("filename"), py::arg("mode"))
.def("append",
[](WritableFile* self, absl::string_view data) {
const auto status = self->Append(data);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
})
// TODO(slebedev): Make WritableFile::Tell const and change self
// to be a reference.
.def("tell",
[](WritableFile* self) {
int64_t pos = -1;
py::gil_scoped_release release;
const auto status = self->Tell(&pos);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
return pos;
})
.def("flush",
[](WritableFile* self) {
py::gil_scoped_release release;
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(self->Flush());
})
.def("close", [](WritableFile* self) {
py::gil_scoped_release release;
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(self->Close());
});
using tensorflow::io::BufferedInputStream;
py::class_<BufferedInputStream>(m, "BufferedInputStream")
.def(py::init([](const std::string& filename, size_t buffer_size) {
py::gil_scoped_release release;
std::unique_ptr<tensorflow::RandomAccessFile> file;
const auto status =
tensorflow::Env::Default()->NewRandomAccessFile(filename,
&file);
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
std::unique_ptr<tensorflow::io::RandomAccessInputStream>
input_stream(new tensorflow::io::RandomAccessInputStream(
file.release(),
/*owns_file=*/true));
py::gil_scoped_acquire acquire;
return new BufferedInputStream(input_stream.release(), buffer_size,
/*owns_input_stream=*/true);
}),
py::arg("filename"), py::arg("buffer_size"))
.def("read",
[](BufferedInputStream* self, int64_t bytes_to_read) {
py::gil_scoped_release release;
tensorflow::tstring result;
const auto status = self->ReadNBytes(bytes_to_read, &result);
if (!status.ok() && !absl::IsOutOfRange(status)) {
result.clear();
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(status);
}
py::gil_scoped_acquire acquire;
return py::bytes(result);
})
.def("readline",
[](BufferedInputStream* self) {
py::gil_scoped_release release;
auto output = self->ReadLineAsString();
py::gil_scoped_acquire acquire;
return py::bytes(output);
})
.def("seek",
[](BufferedInputStream* self, int64_t pos) {
py::gil_scoped_release release;
tensorflow::MaybeRaiseRegisteredFromStatusWithGIL(self->Seek(pos));
})
.def("tell", [](BufferedInputStream* self) {
py::gil_scoped_release release;
return self->Tell();
});
}
} // namespace
tensorflow/python/lib/io/python_io.py
+24
View File
@@ -0,0 +1,24 @@
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Python functions for directly manipulating TFRecord-formatted files.
API docstring: tensorflow.python_io
"""
# go/tf-wildcard-import
# pylint: disable=wildcard-import
from tensorflow.python.lib.io.tf_record import *
# pylint: enable=wildcard-import
tensorflow/python/lib/io/record_io_wrapper.cc
+426
View File
@@ -0,0 +1,426 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/lib/io/zlib_compression_options.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/python/lib/core/pybind11_absl.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace {
namespace py = ::pybind11;
class PyRecordReader {
public:
// NOTE(sethtroisi): At this time PyRecordReader doesn't benefit from taking
// RecordReaderOptions, if this changes the API can be updated at that time.
static absl::Status New(const std::string& filename,
const std::string& compression_type,
PyRecordReader** out) {
auto tmp = new PyRecordReader(filename, compression_type);
TF_RETURN_IF_ERROR(tmp->Reopen());
*out = tmp;
return absl::OkStatus();
}
PyRecordReader() = delete;
~PyRecordReader() { Close(); }
absl::Status ReadNextRecord(tensorflow::tstring* out) {
tensorflow::mutex_lock lock(mutex_);
if (IsClosedInternal()) {
return absl::FailedPreconditionError("Reader is closed.");
}
return reader_->ReadRecord(&offset_, out);
}
bool IsClosed() const {
tensorflow::mutex_lock lock(mutex_);
return IsClosedInternal();
}
void Close() {
tensorflow::mutex_lock lock(mutex_);
reader_ = nullptr;
file_ = nullptr;
}
// Reopen a closed writer by re-opening the file and re-creating the reader,
// but preserving the prior read offset. If not closed, returns an error.
//
// This is useful to allow "refreshing" the underlying file handle, in cases
// where the file was replaced with a newer version containing additional data
// that otherwise wouldn't be available via the existing file handle. This
// allows the file to be polled continuously using the same iterator, even as
// it grows, which supports use cases such as TensorBoard.
absl::Status Reopen() {
tensorflow::mutex_lock lock(mutex_);
if (!IsClosedInternal()) {
return absl::FailedPreconditionError("Reader is not closed.");
}
TF_RETURN_IF_ERROR(
tensorflow::Env::Default()->NewRandomAccessFile(filename_, &file_));
reader_ =
std::make_unique<tensorflow::io::RecordReader>(file_.get(), options_);
return absl::OkStatus();
}
private:
bool IsClosedInternal() const {
return file_ == nullptr && reader_ == nullptr;
}
static constexpr uint64_t kReaderBufferSize = 16 * 1024 * 1024;
PyRecordReader(const std::string& filename,
const std::string& compression_type)
: filename_(filename),
options_(CreateOptions(compression_type)),
offset_(0),
file_(nullptr),
reader_(nullptr) {}
static tensorflow::io::RecordReaderOptions CreateOptions(
const std::string& compression_type) {
auto options =
tensorflow::io::RecordReaderOptions::CreateRecordReaderOptions(
compression_type);
options.buffer_size = kReaderBufferSize;
return options;
}
const std::string filename_;
const tensorflow::io::RecordReaderOptions options_;
uint64_t offset_;
std::unique_ptr<tensorflow::RandomAccessFile> file_;
std::unique_ptr<tensorflow::io::RecordReader> reader_;
mutable tensorflow::mutex mutex_;
PyRecordReader(const PyRecordReader&) = delete;
void operator=(const PyRecordReader&) = delete;
};
class PyRecordRandomReader {
public:
static absl::Status New(const std::string& filename,
PyRecordRandomReader** out) {
std::unique_ptr<tensorflow::RandomAccessFile> file;
TF_RETURN_IF_ERROR(
tensorflow::Env::Default()->NewRandomAccessFile(filename, &file));
auto options =
tensorflow::io::RecordReaderOptions::CreateRecordReaderOptions("");
options.buffer_size = kReaderBufferSize;
auto reader =
std::make_unique<tensorflow::io::RecordReader>(file.get(), options);
*out = new PyRecordRandomReader(std::move(file), std::move(reader));
return absl::OkStatus();
}
PyRecordRandomReader() = delete;
~PyRecordRandomReader() { Close(); }
absl::Status ReadRecord(uint64_t* offset, tensorflow::tstring* out) {
tensorflow::mutex_lock lock(mutex_);
if (IsClosedInternal()) {
return absl::FailedPreconditionError("Random TFRecord Reader is closed.");
}
return reader_->ReadRecord(offset, out);
}
bool IsClosed() const {
tensorflow::mutex_lock lock(mutex_);
return IsClosedInternal();
}
void Close() {
tensorflow::mutex_lock lock(mutex_);
reader_ = nullptr;
file_ = nullptr;
}
private:
bool IsClosedInternal() const {
return file_ == nullptr && reader_ == nullptr;
}
static constexpr uint64_t kReaderBufferSize = 16 * 1024 * 1024;
PyRecordRandomReader(std::unique_ptr<tensorflow::RandomAccessFile> file,
std::unique_ptr<tensorflow::io::RecordReader> reader)
: file_(std::move(file)), reader_(std::move(reader)) {}
std::unique_ptr<tensorflow::RandomAccessFile> file_;
std::unique_ptr<tensorflow::io::RecordReader> reader_;
mutable tensorflow::mutex mutex_;
PyRecordRandomReader(const PyRecordRandomReader&) = delete;
void operator=(const PyRecordRandomReader&) = delete;
};
class PyRecordWriter {
public:
static absl::Status New(const std::string& filename,
const tensorflow::io::RecordWriterOptions& options,
PyRecordWriter** out) {
std::unique_ptr<tensorflow::WritableFile> file;
TF_RETURN_IF_ERROR(
tensorflow::Env::Default()->NewWritableFile(filename, &file));
auto writer =
std::make_unique<tensorflow::io::RecordWriter>(file.get(), options);
*out = new PyRecordWriter(std::move(file), std::move(writer));
return absl::OkStatus();
}
PyRecordWriter() = delete;
~PyRecordWriter() { (void)Close(); }
absl::Status WriteRecord(absl::string_view record) {
tensorflow::mutex_lock lock(mutex_);
if (IsClosedInternal()) {
return absl::FailedPreconditionError("Writer is closed.");
}
return writer_->WriteRecord(record);
}
absl::Status Flush() {
tensorflow::mutex_lock lock(mutex_);
if (IsClosedInternal()) {
return absl::FailedPreconditionError("Writer is closed.");
}
auto status = writer_->Flush();
if (status.ok()) {
// Per the RecordWriter contract, flushing the RecordWriter does not
// flush the underlying file. Here we need to do both.
return file_->Flush();
}
return status;
}
bool IsClosed() const {
tensorflow::mutex_lock lock(mutex_);
return IsClosedInternal();
}
absl::Status Close() {
tensorflow::mutex_lock lock(mutex_);
if (writer_ != nullptr) {
auto status = writer_->Close();
writer_ = nullptr;
if (!status.ok()) return status;
}
if (file_ != nullptr) {
auto status = file_->Close();
file_ = nullptr;
if (!status.ok()) return status;
}
return absl::OkStatus();
}
private:
bool IsClosedInternal() const { return writer_ == nullptr; }
PyRecordWriter(std::unique_ptr<tensorflow::WritableFile> file,
std::unique_ptr<tensorflow::io::RecordWriter> writer)
: file_(std::move(file)), writer_(std::move(writer)) {}
std::unique_ptr<tensorflow::WritableFile> file_;
std::unique_ptr<tensorflow::io::RecordWriter> writer_;
mutable tensorflow::mutex mutex_;
PyRecordWriter(const PyRecordWriter&) = delete;
void operator=(const PyRecordWriter&) = delete;
};
PYBIND11_MODULE(_pywrap_record_io, m, pybind11::mod_gil_not_used()) {
py::class_<PyRecordReader>(m, "RecordIterator")
.def(py::init(
[](const std::string& filename, const std::string& compression_type) {
absl::Status status;
PyRecordReader* self = nullptr;
{
py::gil_scoped_release release;
status = PyRecordReader::New(filename, compression_type, &self);
}
tsl::MaybeRaiseRegisteredFromStatus(status);
return self;
}))
.def("__iter__", [](const py::object& self) { return self; })
.def("__next__",
[](PyRecordReader* self) {
bool is_closed;
{
py::gil_scoped_release release;
is_closed = self->IsClosed();
}
if (is_closed) {
throw py::stop_iteration();
}
tensorflow::tstring record;
absl::Status status;
{
py::gil_scoped_release release;
status = self->ReadNextRecord(&record);
}
if (absl::IsOutOfRange(status)) {
// Don't close because the file being read could be updated
// in-between
// __next__ calls.
throw py::stop_iteration();
}
tsl::MaybeRaiseRegisteredFromStatus(status);
return py::bytes(record);
})
.def("close",
[](PyRecordReader* self) {
py::gil_scoped_release release;
self->Close();
})
.def("reopen", [](PyRecordReader* self) {
absl::Status status;
{
py::gil_scoped_release release;
status = self->Reopen();
}
tsl::MaybeRaiseRegisteredFromStatus(status);
});
py::class_<PyRecordRandomReader>(m, "RandomRecordReader")
.def(py::init([](const std::string& filename) {
absl::Status status;
PyRecordRandomReader* self = nullptr;
{
py::gil_scoped_release release;
status = PyRecordRandomReader::New(filename, &self);
}
tsl::MaybeRaiseRegisteredFromStatus(status);
return self;
}))
.def("read",
[](PyRecordRandomReader* self, uint64_t offset) {
uint64_t temp_offset = offset;
tensorflow::tstring record;
absl::Status status;
{
py::gil_scoped_release release;
status = self->ReadRecord(&temp_offset, &record);
}
if (absl::IsOutOfRange(status)) {
throw py::index_error(
absl::StrCat("Out of range at reading offset ", offset));
}
tsl::MaybeRaiseRegisteredFromStatus(status);
return py::make_tuple(py::bytes(record), temp_offset);
})
.def("close", [](PyRecordRandomReader* self) {
py::gil_scoped_release release;
self->Close();
});
using tensorflow::io::ZlibCompressionOptions;
py::class_<ZlibCompressionOptions>(m, "ZlibCompressionOptions")
.def_readwrite("flush_mode", &ZlibCompressionOptions::flush_mode)
.def_readwrite("input_buffer_size",
&ZlibCompressionOptions::input_buffer_size)
.def_readwrite("output_buffer_size",
&ZlibCompressionOptions::output_buffer_size)
.def_readwrite("window_bits", &ZlibCompressionOptions::window_bits)
.def_readwrite("compression_level",
&ZlibCompressionOptions::compression_level)
.def_readwrite("compression_method",
&ZlibCompressionOptions::compression_method)
.def_readwrite("mem_level", &ZlibCompressionOptions::mem_level)
.def_readwrite("compression_strategy",
&ZlibCompressionOptions::compression_strategy);
using tensorflow::io::RecordWriterOptions;
py::class_<RecordWriterOptions>(m, "RecordWriterOptions")
.def(py::init(&RecordWriterOptions::CreateRecordWriterOptions))
.def_readonly("compression_type", &RecordWriterOptions::compression_type)
.def_readonly("zlib_options", &RecordWriterOptions::zlib_options);
using tensorflow::MaybeRaiseRegisteredFromStatus;
py::class_<PyRecordWriter>(m, "RecordWriter")
.def(py::init(
[](const std::string& filename, const RecordWriterOptions& options) {
PyRecordWriter* self = nullptr;
absl::Status status;
{
py::gil_scoped_release release;
status = PyRecordWriter::New(filename, options, &self);
}
tsl::MaybeRaiseRegisteredFromStatus(status);
return self;
}))
.def("__enter__", [](const py::object& self) { return self; })
.def("__exit__",
[](PyRecordWriter* self, py::args) {
absl::Status status;
{
py::gil_scoped_release release;
status = self->Close();
}
tsl::MaybeRaiseRegisteredFromStatus(status);
})
.def(
"write",
[](PyRecordWriter* self, absl::string_view record) {
absl::Status status;
{
py::gil_scoped_release release;
status = self->WriteRecord(record);
}
tsl::MaybeRaiseRegisteredFromStatus(status);
},
py::arg("record"))
.def("flush",
[](PyRecordWriter* self) {
absl::Status status;
{
py::gil_scoped_release release;
status = self->Flush();
}
tsl::MaybeRaiseRegisteredFromStatus(status);
})
.def("close", [](PyRecordWriter* self) {
absl::Status status;
{
py::gil_scoped_release release;
status = self->Close();
}
tsl::MaybeRaiseRegisteredFromStatus(status);
});
}
} // namespace
tensorflow/python/lib/io/tf_record.py
+318
View File
@@ -0,0 +1,318 @@
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""For reading and writing TFRecords files."""
from tensorflow.python.lib.io import _pywrap_record_io
from tensorflow.python.util import compat
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
@tf_export(
v1=["io.TFRecordCompressionType", "python_io.TFRecordCompressionType"])
@deprecation.deprecated_endpoints("io.TFRecordCompressionType",
"python_io.TFRecordCompressionType")
class TFRecordCompressionType(object):
"""The type of compression for the record."""
NONE = 0
ZLIB = 1
GZIP = 2
@tf_export(
"io.TFRecordOptions",
v1=["io.TFRecordOptions", "python_io.TFRecordOptions"])
@deprecation.deprecated_endpoints("python_io.TFRecordOptions")
class TFRecordOptions(object):
"""Options used for manipulating TFRecord files."""
compression_type_map = {
TFRecordCompressionType.ZLIB: "ZLIB",
TFRecordCompressionType.GZIP: "GZIP",
TFRecordCompressionType.NONE: ""
}
def __init__(self,
compression_type=None,
flush_mode=None,
input_buffer_size=None,
output_buffer_size=None,
window_bits=None,
compression_level=None,
compression_method=None,
mem_level=None,
compression_strategy=None):
# pylint: disable=line-too-long
"""Creates a `TFRecordOptions` instance.
Options only effect TFRecordWriter when compression_type is not `None`.
Documentation, details, and defaults can be found in
[`zlib_compression_options.h`](https://www.tensorflow.org/code/tensorflow/core/lib/io/zlib_compression_options.h)
and in the [zlib manual](http://www.zlib.net/manual.html).
Leaving an option as `None` allows C++ to set a reasonable default.
Args:
compression_type: `"GZIP"`, `"ZLIB"`, or `""` (no compression).
flush_mode: flush mode or `None`, Default: Z_NO_FLUSH.
input_buffer_size: int or `None`.
output_buffer_size: int or `None`.
window_bits: int or `None`.
compression_level: 0 to 9, or `None`.
compression_method: compression method or `None`.
mem_level: 1 to 9, or `None`.
compression_strategy: strategy or `None`. Default: Z_DEFAULT_STRATEGY.
Returns:
A `TFRecordOptions` object.
Raises:
ValueError: If compression_type is invalid.
"""
# pylint: enable=line-too-long
# Check compression_type is valid, but for backwards compatibility don't
# immediately convert to a string.
self.get_compression_type_string(compression_type)
self.compression_type = compression_type
self.flush_mode = flush_mode
self.input_buffer_size = input_buffer_size
self.output_buffer_size = output_buffer_size
self.window_bits = window_bits
self.compression_level = compression_level
self.compression_method = compression_method
self.mem_level = mem_level
self.compression_strategy = compression_strategy
@classmethod
def get_compression_type_string(cls, options):
"""Convert various option types to a unified string.
Args:
options: `TFRecordOption`, `TFRecordCompressionType`, or string.
Returns:
Compression type as string (e.g. `'ZLIB'`, `'GZIP'`, or `''`).
Raises:
ValueError: If compression_type is invalid.
"""
if not options:
return ""
elif isinstance(options, TFRecordOptions):
return cls.get_compression_type_string(options.compression_type)
elif isinstance(options, TFRecordCompressionType):
return cls.compression_type_map[options]
elif options in TFRecordOptions.compression_type_map:
return cls.compression_type_map[options]
elif options in TFRecordOptions.compression_type_map.values():
return options
else:
raise ValueError('Not a valid compression_type: "{}"'.format(options))
def _as_record_writer_options(self):
"""Convert to RecordWriterOptions for use with PyRecordWriter."""
options = _pywrap_record_io.RecordWriterOptions(
compat.as_bytes(
self.get_compression_type_string(self.compression_type)))
if self.flush_mode is not None:
options.zlib_options.flush_mode = self.flush_mode
if self.input_buffer_size is not None:
options.zlib_options.input_buffer_size = self.input_buffer_size
if self.output_buffer_size is not None:
options.zlib_options.output_buffer_size = self.output_buffer_size
if self.window_bits is not None:
options.zlib_options.window_bits = self.window_bits
if self.compression_level is not None:
options.zlib_options.compression_level = self.compression_level
if self.compression_method is not None:
options.zlib_options.compression_method = self.compression_method
if self.mem_level is not None:
options.zlib_options.mem_level = self.mem_level
if self.compression_strategy is not None:
options.zlib_options.compression_strategy = self.compression_strategy
return options
@tf_export(v1=["io.tf_record_iterator", "python_io.tf_record_iterator"])
@deprecation.deprecated(
date=None,
instructions=("Use eager execution and: \n"
"`tf.data.TFRecordDataset(path)`"))
def tf_record_iterator(path, options=None):
"""An iterator that read the records from a TFRecords file.
Args:
path: The path to the TFRecords file.
options: (optional) A TFRecordOptions object.
Returns:
An iterator of serialized TFRecords.
Raises:
IOError: If `path` cannot be opened for reading.
"""
compression_type = TFRecordOptions.get_compression_type_string(options)
return _pywrap_record_io.RecordIterator(path, compression_type)
def tf_record_random_reader(path):
"""Creates a reader that allows random-access reads from a TFRecords file.
The created reader object has the following method:
- `read(offset)`, which returns a tuple of `(record, ending_offset)`, where
`record` is the TFRecord read at the offset, and
`ending_offset` is the ending offset of the read record.
The method throws a `tf.errors.DataLossError` if data is corrupted at
the given offset. The method throws `IndexError` if the offset is out of
range for the TFRecords file.
Usage example:
```py
reader = tf_record_random_reader(file_path)
record_1, offset_1 = reader.read(0) # 0 is the initial offset.
# offset_1 is the ending offset of the 1st record and the starting offset of
# the next.
record_2, offset_2 = reader.read(offset_1)
# offset_2 is the ending offset of the 2nd record and the starting offset of
# the next.
# We can jump back and read the first record again if so desired.
reader.read(0)
```
Args:
path: The path to the TFRecords file.
Returns:
An object that supports random-access reading of the serialized TFRecords.
Raises:
IOError: If `path` cannot be opened for reading.
"""
return _pywrap_record_io.RandomRecordReader(path)
@tf_export(
"io.TFRecordWriter", v1=["io.TFRecordWriter", "python_io.TFRecordWriter"])
@deprecation.deprecated_endpoints("python_io.TFRecordWriter")
class TFRecordWriter(_pywrap_record_io.RecordWriter):
"""A class to write records to a TFRecords file.
[TFRecords tutorial](https://www.tensorflow.org/tutorials/load_data/tfrecord)
TFRecords is a binary format which is optimized for high throughput data
retrieval, generally in conjunction with `tf.data`. `TFRecordWriter` is used
to write serialized examples to a file for later consumption. The key steps
are:
Ahead of time:
- [Convert data into a serialized format](
https://www.tensorflow.org/tutorials/load_data/tfrecord#tfexample)
- [Write the serialized data to one or more files](
https://www.tensorflow.org/tutorials/load_data/tfrecord#tfrecord_files_in_python)
During training or evaluation:
- [Read serialized examples into memory](
https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file)
- [Parse (deserialize) examples](
https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file)
A minimal example is given below:
>>> import tempfile
>>> example_path = os.path.join(tempfile.gettempdir(), "example.tfrecords")
>>> np.random.seed(0)
>>> # Write the records to a file.
... with tf.io.TFRecordWriter(example_path) as file_writer:
... for _ in range(4):
... x, y = np.random.random(), np.random.random()
...
... record_bytes = tf.train.Example(features=tf.train.Features(feature={
... "x": tf.train.Feature(float_list=tf.train.FloatList(value=[x])),
... "y": tf.train.Feature(float_list=tf.train.FloatList(value=[y])),
... })).SerializeToString()
... file_writer.write(record_bytes)
>>> # Read the data back out.
>>> def decode_fn(record_bytes):
... return tf.io.parse_single_example(
... # Data
... record_bytes,
...
... # Schema
... {"x": tf.io.FixedLenFeature([], dtype=tf.float32),
... "y": tf.io.FixedLenFeature([], dtype=tf.float32)}
... )
>>> for batch in tf.data.TFRecordDataset([example_path]).map(decode_fn):
... print("x = {x:.4f}, y = {y:.4f}".format(**batch))
x = 0.5488, y = 0.7152
x = 0.6028, y = 0.5449
x = 0.4237, y = 0.6459
x = 0.4376, y = 0.8918
This class implements `__enter__` and `__exit__`, and can be used
in `with` blocks like a normal file. (See the usage example above.)
"""
# TODO(josh11b): Support appending?
def __init__(self, path, options=None):
"""Opens file `path` and creates a `TFRecordWriter` writing to it.
Args:
path: The path to the TFRecords file.
options: (optional) String specifying compression type,
`TFRecordCompressionType`, or `TFRecordOptions` object.
Raises:
IOError: If `path` cannot be opened for writing.
ValueError: If valid compression_type can't be determined from `options`.
"""
if not isinstance(options, TFRecordOptions):
options = TFRecordOptions(compression_type=options)
# pylint: disable=protected-access
super(TFRecordWriter, self).__init__(
compat.as_bytes(path), options._as_record_writer_options())
# pylint: enable=protected-access
# TODO(slebedev): The following wrapper methods are there to compensate
# for lack of signatures in pybind11-generated classes. Switch to
# __text_signature__ when TensorFlow drops Python 2.X support.
# See https://github.com/pybind/pybind11/issues/945
# pylint: disable=useless-super-delegation
def write(self, record):
"""Write a string record to the file.
Args:
record: str
"""
super(TFRecordWriter, self).write(record)
def flush(self):
"""Flush the file."""
super(TFRecordWriter, self).flush()
def close(self):
"""Close the file."""
super(TFRecordWriter, self).close()
# pylint: enable=useless-super-delegation
tensorflow/python/lib/io/tf_record_test.py
+646
View File
@@ -0,0 +1,646 @@
# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for tf_record.TFRecordWriter and tf_record.tf_record_iterator."""
import gzip
import os
import random
import string
import zlib
import six
from tensorflow.python.framework import errors_impl
from tensorflow.python.lib.io import tf_record
from tensorflow.python.platform import test
from tensorflow.python.util import compat
TFRecordCompressionType = tf_record.TFRecordCompressionType
# Edgar Allan Poe's 'Eldorado'
_TEXT = b"""Gaily bedight,
A gallant knight,
In sunshine and in shadow,
Had journeyed long,
Singing a song,
In search of Eldorado.
But he grew old
This knight so bold
And o'er his heart a shadow
Fell as he found
No spot of ground
That looked like Eldorado.
And, as his strength
Failed him at length,
He met a pilgrim shadow
'Shadow,' said he,
'Where can it be
This land of Eldorado?'
'Over the Mountains
Of the Moon'
Down the Valley of the Shadow,
Ride, boldly ride,'
The shade replied,
'If you seek for Eldorado!'
"""
class TFCompressionTestCase(test.TestCase):
"""TFCompression Test"""
def setUp(self):
super(TFCompressionTestCase, self).setUp()
self._num_files = 2
self._num_records = 7
def _Record(self, f, r):
return compat.as_bytes("Record %d of file %d" % (r, f))
def _CreateFiles(self, options=None, prefix=""):
filenames = []
for i in range(self._num_files):
name = prefix + "tfrecord.%d.txt" % i
records = [self._Record(i, j) for j in range(self._num_records)]
fn = self._WriteRecordsToFile(records, name, options)
filenames.append(fn)
return filenames
def _WriteRecordsToFile(self, records, name="tfrecord", options=None):
fn = os.path.join(self.get_temp_dir(), name)
with tf_record.TFRecordWriter(fn, options=options) as writer:
for r in records:
writer.write(r)
return fn
def _ZlibCompressFile(self, infile, name="tfrecord.z"):
# zlib compress the file and write compressed contents to file.
with open(infile, "rb") as f:
cdata = zlib.compress(f.read())
zfn = os.path.join(self.get_temp_dir(), name)
with open(zfn, "wb") as f:
f.write(cdata)
return zfn
def _GzipCompressFile(self, infile, name="tfrecord.gz"):
# gzip compress the file and write compressed contents to file.
with open(infile, "rb") as f:
cdata = f.read()
gzfn = os.path.join(self.get_temp_dir(), name)
with gzip.GzipFile(gzfn, "wb") as f:
f.write(cdata)
return gzfn
def _ZlibDecompressFile(self, infile, name="tfrecord"):
with open(infile, "rb") as f:
cdata = zlib.decompress(f.read())
fn = os.path.join(self.get_temp_dir(), name)
with open(fn, "wb") as f:
f.write(cdata)
return fn
def _GzipDecompressFile(self, infile, name="tfrecord"):
with gzip.GzipFile(infile, "rb") as f:
cdata = f.read()
fn = os.path.join(self.get_temp_dir(), name)
with open(fn, "wb") as f:
f.write(cdata)
return fn
class TFRecordWriterTest(TFCompressionTestCase):
"""TFRecordWriter Test"""
def _AssertFilesEqual(self, a, b, equal):
for an, bn in zip(a, b):
with open(an, "rb") as af, open(bn, "rb") as bf:
if equal:
self.assertEqual(af.read(), bf.read())
else:
self.assertNotEqual(af.read(), bf.read())
def _CompressionSizeDelta(self, records, options_a, options_b):
"""Validate compression with options_a and options_b and return size delta.
Compress records with options_a and options_b. Uncompress both compressed
files and assert that the contents match the original records. Finally
calculate how much smaller the file compressed with options_a was than the
file compressed with options_b.
Args:
records: The records to compress
options_a: First set of options to compress with, the baseline for size.
options_b: Second set of options to compress with.
Returns:
The difference in file size when using options_a vs options_b. A positive
value means options_a was a better compression than options_b. A negative
value means options_b had better compression than options_a.
"""
fn_a = self._WriteRecordsToFile(records, "tfrecord_a", options=options_a)
test_a = list(tf_record.tf_record_iterator(fn_a, options=options_a))
self.assertEqual(records, test_a, options_a)
fn_b = self._WriteRecordsToFile(records, "tfrecord_b", options=options_b)
test_b = list(tf_record.tf_record_iterator(fn_b, options=options_b))
self.assertEqual(records, test_b, options_b)
# Negative number => better compression.
return os.path.getsize(fn_a) - os.path.getsize(fn_b)
def testWriteReadZLibFiles(self):
"""test Write Read ZLib Files"""
# Write uncompressed then compress manually.
options = tf_record.TFRecordOptions(TFRecordCompressionType.NONE)
files = self._CreateFiles(options, prefix="uncompressed")
zlib_files = [
self._ZlibCompressFile(fn, "tfrecord_%s.z" % i)
for i, fn in enumerate(files)
]
self._AssertFilesEqual(files, zlib_files, False)
# Now write compressd and verify same.
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
compressed_files = self._CreateFiles(options, prefix="compressed")
self._AssertFilesEqual(compressed_files, zlib_files, True)
# Decompress compress and verify same.
uncompressed_files = [
self._ZlibDecompressFile(fn, "tfrecord_%s.z" % i)
for i, fn in enumerate(compressed_files)
]
self._AssertFilesEqual(uncompressed_files, files, True)
def testWriteReadGzipFiles(self):
"""test Write Read Gzip Files"""
# Write uncompressed then compress manually.
options = tf_record.TFRecordOptions(TFRecordCompressionType.NONE)
files = self._CreateFiles(options, prefix="uncompressed")
gzip_files = [
self._GzipCompressFile(fn, "tfrecord_%s.gz" % i)
for i, fn in enumerate(files)
]
self._AssertFilesEqual(files, gzip_files, False)
# Now write compressd and verify same.
options = tf_record.TFRecordOptions(TFRecordCompressionType.GZIP)
compressed_files = self._CreateFiles(options, prefix="compressed")
# Note: Gzips written by TFRecordWriter add 'tfrecord_0' so
# compressed_files can't be compared with gzip_files
# Decompress compress and verify same.
uncompressed_files = [
self._GzipDecompressFile(fn, "tfrecord_%s.gz" % i)
for i, fn in enumerate(compressed_files)
]
self._AssertFilesEqual(uncompressed_files, files, True)
def testNoCompressionType(self):
"""test No Compression Type"""
self.assertEqual(
"",
tf_record.TFRecordOptions.get_compression_type_string(
tf_record.TFRecordOptions()))
self.assertEqual(
"",
tf_record.TFRecordOptions.get_compression_type_string(
tf_record.TFRecordOptions("")))
with self.assertRaises(ValueError):
tf_record.TFRecordOptions(5)
with self.assertRaises(ValueError):
tf_record.TFRecordOptions("BZ2")
def testZlibCompressionType(self):
"""test Zlib Compression Type"""
zlib_t = tf_record.TFRecordCompressionType.ZLIB
self.assertEqual(
"ZLIB",
tf_record.TFRecordOptions.get_compression_type_string(
tf_record.TFRecordOptions("ZLIB")))
self.assertEqual(
"ZLIB",
tf_record.TFRecordOptions.get_compression_type_string(
tf_record.TFRecordOptions(zlib_t)))
self.assertEqual(
"ZLIB",
tf_record.TFRecordOptions.get_compression_type_string(
tf_record.TFRecordOptions(tf_record.TFRecordOptions(zlib_t))))
def testCompressionOptions(self):
"""Create record with mix of random and repeated data to test compression on."""
rnd = random.Random(123)
random_record = compat.as_bytes(
"".join(rnd.choice(string.digits) for _ in range(10000)))
repeated_record = compat.as_bytes(_TEXT)
for _ in range(10000):
start_i = rnd.randint(0, len(_TEXT))
length = rnd.randint(10, 200)
repeated_record += _TEXT[start_i:start_i + length]
records = [random_record, repeated_record, random_record]
tests = [
("compression_level", 2, "LE"), # Lower compression is worse or equal.
("compression_level", 6, 0), # Default compression_level is equal.
("flush_mode", zlib.Z_FULL_FLUSH, 1), # A few less bytes.
("flush_mode", zlib.Z_NO_FLUSH, 0), # NO_FLUSH is the default.
("input_buffer_size", 4096, 0), # Increases time not size.
("output_buffer_size", 4096, 0), # Increases time not size.
("window_bits", 8, -1), # Smaller than default window increases size.
("compression_strategy", zlib.Z_HUFFMAN_ONLY, -1), # Worse.
("compression_strategy", zlib.Z_FILTERED, "LE"), # Worse or equal.
]
compression_type = tf_record.TFRecordCompressionType.ZLIB
options_a = tf_record.TFRecordOptions(compression_type)
for prop, value, delta_sign in tests:
options_b = tf_record.TFRecordOptions(
compression_type=compression_type, **{prop: value})
delta = self._CompressionSizeDelta(records, options_a, options_b)
if delta_sign == "LE":
self.assertLessEqual(
delta, 0,
"Setting {} = {}, file was {} smaller didn't match sign of {}"
.format(prop, value, delta, delta_sign))
else:
self.assertTrue(
delta == 0 if delta_sign == 0 else delta // delta_sign > 0,
"Setting {} = {}, file was {} smaller didn't match sign of {}"
.format(prop, value, delta, delta_sign))
class TFRecordWriterZlibTest(TFCompressionTestCase):
"""TFRecordWriter Zlib test"""
def testZLibFlushRecord(self):
"""test ZLib Flush Record"""
original = [b"small record"]
fn = self._WriteRecordsToFile(original, "small_record")
with open(fn, "rb") as h:
buff = h.read()
# creating more blocks and trailing blocks shouldn't break reads
compressor = zlib.compressobj(9, zlib.DEFLATED, zlib.MAX_WBITS)
output = b""
for c in buff:
if isinstance(c, int):
c = six.int2byte(c)
output += compressor.compress(c)
output += compressor.flush(zlib.Z_FULL_FLUSH)
output += compressor.flush(zlib.Z_FULL_FLUSH)
output += compressor.flush(zlib.Z_FULL_FLUSH)
output += compressor.flush(zlib.Z_FINISH)
# overwrite the original file with the compressed data
with open(fn, "wb") as h:
h.write(output)
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
actual = list(tf_record.tf_record_iterator(fn, options=options))
self.assertEqual(actual, original)
def testZlibReadWrite(self):
"""Verify that files produced are zlib compatible."""
original = [b"foo", b"bar"]
fn = self._WriteRecordsToFile(original, "zlib_read_write.tfrecord")
zfn = self._ZlibCompressFile(fn, "zlib_read_write.tfrecord.z")
# read the compressed contents and verify.
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
actual = list(tf_record.tf_record_iterator(zfn, options=options))
self.assertEqual(actual, original)
def testZlibReadWriteLarge(self):
"""Verify that writing large contents also works."""
# Make it large (about 5MB)
original = [_TEXT * 10240]
fn = self._WriteRecordsToFile(original, "zlib_read_write_large.tfrecord")
zfn = self._ZlibCompressFile(fn, "zlib_read_write_large.tfrecord.z")
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
actual = list(tf_record.tf_record_iterator(zfn, options=options))
self.assertEqual(actual, original)
def testGzipReadWrite(self):
"""Verify that files produced are gzip compatible."""
original = [b"foo", b"bar"]
fn = self._WriteRecordsToFile(original, "gzip_read_write.tfrecord")
gzfn = self._GzipCompressFile(fn, "tfrecord.gz")
options = tf_record.TFRecordOptions(TFRecordCompressionType.GZIP)
actual = list(tf_record.tf_record_iterator(gzfn, options=options))
self.assertEqual(actual, original)
class TFRecordIteratorTest(TFCompressionTestCase):
"""TFRecordIterator test"""
def setUp(self):
super(TFRecordIteratorTest, self).setUp()
self._num_records = 7
def testIterator(self):
"""test Iterator"""
records = [self._Record(0, i) for i in range(self._num_records)]
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
fn = self._WriteRecordsToFile(records, "compressed_records", options)
reader = tf_record.tf_record_iterator(fn, options)
for expected in records:
record = next(reader)
self.assertEqual(expected, record)
with self.assertRaises(StopIteration):
record = next(reader)
def testWriteZlibRead(self):
"""Verify compression with TFRecordWriter is zlib library compatible."""
original = [b"foo", b"bar"]
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
fn = self._WriteRecordsToFile(original, "write_zlib_read.tfrecord.z",
options)
zfn = self._ZlibDecompressFile(fn, "write_zlib_read.tfrecord")
actual = list(tf_record.tf_record_iterator(zfn))
self.assertEqual(actual, original)
def testWriteZlibReadLarge(self):
"""Verify compression for large records is zlib library compatible."""
# Make it large (about 5MB)
original = [_TEXT * 10240]
options = tf_record.TFRecordOptions(TFRecordCompressionType.ZLIB)
fn = self._WriteRecordsToFile(original, "write_zlib_read_large.tfrecord.z",
options)
zfn = self._ZlibDecompressFile(fn, "write_zlib_read_large.tfrecord")
actual = list(tf_record.tf_record_iterator(zfn))
self.assertEqual(actual, original)
def testWriteGzipRead(self):
original = [b"foo", b"bar"]
options = tf_record.TFRecordOptions(TFRecordCompressionType.GZIP)
fn = self._WriteRecordsToFile(original, "write_gzip_read.tfrecord.gz",
options)
gzfn = self._GzipDecompressFile(fn, "write_gzip_read.tfrecord")
actual = list(tf_record.tf_record_iterator(gzfn))
self.assertEqual(actual, original)
def testReadGrowingFile_preservesReadOffset(self):
"""Verify that tf_record_iterator preserves read offset even after EOF.
When a file is iterated to EOF, the iterator should raise StopIteration but
not actually close the reader. Then if later new data is appended, the
iterator should start returning that new data on the next call to next(),
preserving the read offset. This behavior is required by TensorBoard.
"""
# Start the file with a good record.
fn = os.path.join(self.get_temp_dir(), "file.tfrecord")
with tf_record.TFRecordWriter(fn) as writer:
writer.write(b"one")
writer.write(b"two")
writer.flush()
iterator = tf_record.tf_record_iterator(fn)
self.assertEqual(b"one", next(iterator))
self.assertEqual(b"two", next(iterator))
# Iterating at EOF results in StopIteration repeatedly.
with self.assertRaises(StopIteration):
next(iterator)
with self.assertRaises(StopIteration):
next(iterator)
# Retrying after adding a new record successfully returns the new record,
# preserving the prior read offset.
writer.write(b"three")
writer.flush()
self.assertEqual(b"three", next(iterator))
with self.assertRaises(StopIteration):
next(iterator)
def testReadTruncatedFile_preservesReadOffset(self):
"""Verify that tf_record_iterator throws an exception on bad TFRecords.
When a truncated record is completed, the iterator should return that new
record on the next attempt at iteration, preserving the read offset. This
behavior is required by TensorBoard.
"""
# Write out a record and read it back it to get the raw bytes.
fn = os.path.join(self.get_temp_dir(), "temp_file")
with tf_record.TFRecordWriter(fn) as writer:
writer.write(b"truncated")
with open(fn, "rb") as f:
record_bytes = f.read()
# Start the file with a good record.
fn_truncated = os.path.join(self.get_temp_dir(), "truncated_file")
with tf_record.TFRecordWriter(fn_truncated) as writer:
writer.write(b"good")
with open(fn_truncated, "ab", buffering=0) as f:
# Cause truncation by omitting the last byte from the record.
f.write(record_bytes[:-1])
iterator = tf_record.tf_record_iterator(fn_truncated)
# Good record appears first.
self.assertEqual(b"good", next(iterator))
# Truncated record repeatedly causes DataLossError upon iteration.
with self.assertRaises(errors_impl.DataLossError):
next(iterator)
with self.assertRaises(errors_impl.DataLossError):
next(iterator)
# Retrying after completing the record successfully returns the rest of
# the file contents, preserving the prior read offset.
f.write(record_bytes[-1:])
self.assertEqual(b"truncated", next(iterator))
with self.assertRaises(StopIteration):
next(iterator)
def testReadReplacedFile_preservesReadOffset_afterReopen(self):
"""Verify that tf_record_iterator allows reopening at the same read offset.
In some cases, data will be logically "appended" to a file by replacing the
entire file with a new version that includes the additional data. For
example, this can happen with certain GCS implementations (since GCS has no
true append operation), or when using rsync without the `--inplace` option
to transfer snapshots of a growing file. Since the iterator retains a handle
to a stale version of the file, it won't return any of the new data.
To force this to happen, callers can check for a replaced file (e.g. via a
stat call that reflects an increased file size) and opt to close and reopen
the iterator. When iteration is next attempted, this should result in
reading from the newly opened file, while preserving the read offset. This
behavior is required by TensorBoard.
"""
def write_records_to_file(filename, records):
writer = tf_record.TFRecordWriter(filename)
for record in records:
writer.write(record)
writer.close()
fn = os.path.join(self.get_temp_dir(), "orig_file")
write_records_to_file(fn, [b"one", b"two"])
iterator = tf_record.tf_record_iterator(fn)
self.assertEqual(b"one", next(iterator))
self.assertEqual(b"two", next(iterator))
# Iterating at EOF results in StopIteration repeatedly.
with self.assertRaises(StopIteration):
next(iterator)
with self.assertRaises(StopIteration):
next(iterator)
# Add a new record to the end of the file by overwriting it.
fn2 = os.path.join(self.get_temp_dir(), "new_file")
write_records_to_file(fn2, [b"one", b"two", b"three"])
# Windows disallows replacing files while in use, so close iterator early.
if os.name == "nt":
iterator.close()
os.replace(fn2, fn)
# Iterating at EOF still results in StopIteration; new data is not shown.
with self.assertRaises(StopIteration):
next(iterator)
with self.assertRaises(StopIteration):
next(iterator)
# Retrying after close and reopen successfully returns the new record,
# preserving the prior read offset.
iterator.close()
iterator.reopen()
self.assertEqual(b"three", next(iterator))
with self.assertRaises(StopIteration):
next(iterator)
class TFRecordRandomReaderTest(TFCompressionTestCase):
def testRandomReaderReadingWorks(self):
"""Test read access to random offsets in the TFRecord file."""
records = [self._Record(0, i) for i in range(self._num_records)]
fn = self._WriteRecordsToFile(records, "uncompressed_records")
reader = tf_record.tf_record_random_reader(fn)
offset = 0
offsets = [offset]
# Do a pass of forward reading.
for i in range(self._num_records):
record, offset = reader.read(offset)
self.assertEqual(record, records[i])
offsets.append(offset)
# Reading off the bound should lead to error.
with self.assertRaisesRegex(IndexError, r"Out of range.*offset"):
reader.read(offset)
# Do a pass of backward reading.
for i in range(self._num_records - 1, 0, -1):
record, offset = reader.read(offsets[i])
self.assertEqual(offset, offsets[i + 1])
self.assertEqual(record, records[i])
def testRandomReaderThrowsErrorForInvalidOffset(self):
records = [self._Record(0, i) for i in range(self._num_records)]
fn = self._WriteRecordsToFile(records, "uncompressed_records")
reader = tf_record.tf_record_random_reader(fn)
with self.assertRaisesRegex(errors_impl.DataLossError, r"corrupted record"):
reader.read(1) # 1 is guaranteed to be an invalid offset.
def testClosingRandomReaderCausesErrorsForFurtherReading(self):
records = [self._Record(0, i) for i in range(self._num_records)]
fn = self._WriteRecordsToFile(records, "uncompressed_records")
reader = tf_record.tf_record_random_reader(fn)
reader.close()
with self.assertRaisesRegex(errors_impl.FailedPreconditionError, r"closed"):
reader.read(0)
class TFRecordWriterCloseAndFlushTests(test.TestCase):
"""TFRecordWriter close and flush tests"""
# pylint: disable=arguments-differ
def setUp(self, compression_type=TFRecordCompressionType.NONE):
super(TFRecordWriterCloseAndFlushTests, self).setUp()
self._fn = os.path.join(self.get_temp_dir(), "tf_record_writer_test.txt")
self._options = tf_record.TFRecordOptions(compression_type)
self._writer = tf_record.TFRecordWriter(self._fn, self._options)
self._num_records = 20
def _Record(self, r):
return compat.as_bytes("Record %d" % r)
def testWriteAndLeaveOpen(self):
records = list(map(self._Record, range(self._num_records)))
for record in records:
self._writer.write(record)
# Verify no segfault if writer isn't explicitly closed.
def testWriteAndRead(self):
records = list(map(self._Record, range(self._num_records)))
for record in records:
self._writer.write(record)
self._writer.close()
actual = list(tf_record.tf_record_iterator(self._fn, self._options))
self.assertListEqual(actual, records)
def testFlushAndRead(self):
records = list(map(self._Record, range(self._num_records)))
for record in records:
self._writer.write(record)
self._writer.flush()
actual = list(tf_record.tf_record_iterator(self._fn, self._options))
self.assertListEqual(actual, records)
def testDoubleClose(self):
self._writer.write(self._Record(0))
self._writer.close()
self._writer.close()
def testFlushAfterCloseIsError(self):
self._writer.write(self._Record(0))
self._writer.close()
with self.assertRaises(errors_impl.FailedPreconditionError):
self._writer.flush()
def testWriteAfterCloseIsError(self):
self._writer.write(self._Record(0))
self._writer.close()
with self.assertRaises(errors_impl.FailedPreconditionError):
self._writer.write(self._Record(1))
class TFRecordWriterCloseAndFlushGzipTests(TFRecordWriterCloseAndFlushTests):
# pylint: disable=arguments-differ
def setUp(self):
super(TFRecordWriterCloseAndFlushGzipTests,
self).setUp(TFRecordCompressionType.GZIP)
class TFRecordWriterCloseAndFlushZlibTests(TFRecordWriterCloseAndFlushTests):
# pylint: disable=arguments-differ
def setUp(self):
super(TFRecordWriterCloseAndFlushZlibTests,
self).setUp(TFRecordCompressionType.ZLIB)
if __name__ == "__main__":
test.main()