wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 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/client/BUILD
+549
View File
@@ -0,0 +1,549 @@
# Contains targets that expose client session APIs
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "tf_cuda_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test", "tf_python_pybind_extension")
load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap", "if_static")
load(
"//tensorflow/tools/test:performance.bzl",
"cuda_py_benchmark_test",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
py_library(
name = "pywrap_tf_session",
srcs = ["pywrap_tf_session.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":_pywrap_tf_session",
"//tensorflow/python:pywrap_tensorflow",
"//tensorflow/python/util:tf_stack",
],
)
tf_python_pybind_extension(
name = "_pywrap_tf_session",
srcs = ["tf_session_wrapper.cc"],
hdrs = if_pywrap(
if_false = [
"tf_session_helper.h",
"//tensorflow/c:headers",
"//tensorflow/c:safe_ptr_hdr",
"//tensorflow/c/eager:headers",
"//tensorflow/c/eager:pywrap_required_hdrs",
"//tensorflow/c/experimental/ops:pywrap_required_hdrs",
"@xla//xla/tsl/distributed_runtime:pywrap_required_hdrs",
"@xla//xla/tsl/distributed_runtime/coordination:pywrap_required_hdrs",
"@xla//xla/tsl/python/lib/core:numpy_hdr",
"//tensorflow/core/common_runtime/eager:pywrap_required_hdrs",
"//tensorflow/core/distributed_runtime:pywrap_required_hdrs",
"//tensorflow/core/distributed_runtime/coordination:pywrap_required_hdrs",
"//tensorflow/core/distributed_runtime/eager:pywrap_required_hdrs",
"//tensorflow/python/lib/core:safe_pyobject_ptr_required_hdrs",
],
if_true = [],
),
additional_exported_symbols = [
"_TF_SetTarget",
"_TF_SetConfig",
"_TF_NewSessionOptions",
],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tf_session.pyi",
],
deps = [
"//tensorflow/c:pywrap_required_hdrs",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/config:flags_headers",
"//tensorflow/core/framework:pywrap_required_hdrs",
"//tensorflow/core/lib/llvm_rtti",
"//tensorflow/core/public:release_version",
"//tensorflow/core/util:version_info",
"//tensorflow/python/lib/core:pybind11_lib",
"//tensorflow/python/lib/core:pybind11_status",
"//tensorflow/python/lib/core:safe_pyobject_ptr",
"//third_party/py/numpy:headers",
"@com_google_absl//absl/types:optional",
"@eigen_archive//:eigen3",
"@pybind11",
"@pybind11_abseil//pybind11_abseil:absl_casters",
"@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
"@xla//third_party/python_runtime:headers",
"@xla//xla/tsl/python/lib/core:numpy",
] + if_pywrap([
"//tensorflow/c:safe_ptr",
"//tensorflow/c:c_api_experimental",
"//tensorflow/python/client:tf_session_helper",
"//tensorflow/c:python_api",
"//tensorflow/core/common_runtime:core_cpu_lib",
]) + if_static(
extra_deps = [
"//tensorflow/core/protobuf:eager_service_proto_cc",
"//tensorflow/core/protobuf:master_proto_cc",
"//tensorflow/core/protobuf:worker_proto_cc",
"//tensorflow/core:version_lib",
"@xla//xla/tsl/protobuf:coordination_service_proto_cc",
],
otherwise = [
"//tensorflow/core/protobuf:eager_service_proto_cc_headers_only",
"//tensorflow/core/protobuf:master_proto_cc_headers_only",
"//tensorflow/core/protobuf:worker_proto_cc_headers_only",
"@xla//xla/tsl/protobuf:coordination_service_proto_cc_headers_only",
],
),
)
tf_python_pybind_extension(
name = "_pywrap_debug_events_writer",
srcs = ["debug_events_writer_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_debug_events_writer.pyi",
],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/python/lib/core:pybind11_absl",
"//tensorflow/python/lib/core:pybind11_proto",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
tf_python_pybind_extension(
name = "_pywrap_events_writer",
srcs = ["events_writer_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_events_writer.pyi",
],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/python/lib/core:pybind11_absl",
"//tensorflow/python/lib/core:pybind11_proto",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
py_library(
name = "client",
srcs = ["client_lib.py"],
strict_deps = True,
deps = [
":_pywrap_device_lib",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:pywrap_tensorflow",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:build_info",
"//tensorflow/python/platform:tf_logging",
],
)
tf_py_strict_test(
name = "events_writer_test",
size = "small",
srcs = ["events_writer_test.py"],
deps = [
":_pywrap_events_writer",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:pywrap_tensorflow",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/lib/io:tf_record",
"//tensorflow/python/platform:test",
"//tensorflow/python/util:compat",
],
)
py_library(
name = "device_lib",
srcs = ["device_lib.py"],
strict_deps = True,
visibility = [
"//tensorflow:internal",
"//third_party/py/cleverhans:__subpackages__",
],
deps = [
":_pywrap_device_lib",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:pywrap_tensorflow",
],
)
tf_python_pybind_extension(
name = "_pywrap_device_lib",
srcs = ["device_lib_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_device_lib.pyi",
],
deps = [
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework_internal_headers_lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/common_runtime:device",
"//tensorflow/core/common_runtime:device_factory",
"//tensorflow/python/lib/core:pybind11_proto",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
cuda_py_strict_test(
name = "device_lib_test",
size = "small",
srcs = [
"device_lib_test.py",
],
deps = [
":device_lib",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
],
)
cc_library(
name = "session_ref",
srcs = ["session_ref.cc"],
hdrs = ["session_ref.h"],
deps = [
"//tensorflow/core:core_cpu",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/protobuf:replay_log_proto_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
] + if_static(
extra_deps = [
"//tensorflow/core/protobuf:master_proto_cc",
],
otherwise = [
"//tensorflow/core/protobuf:master_proto_cc_headers_only",
],
),
)
tf_cuda_library(
name = "tf_session_helper",
srcs = ["tf_session_helper.cc"],
hdrs = ["tf_session_helper.h"],
compatible_with = [],
deps = [
":construction_fails_op",
":session_ref",
"//tensorflow/c:c_api",
"//tensorflow/c:c_api_internal",
"//tensorflow/c:safe_ptr",
"//tensorflow/c:tf_buffer_internal",
"//tensorflow/c:tf_status_helper",
"//tensorflow/core",
"//tensorflow/core:all_kernels",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:direct_session",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/python/framework:test_ops_kernels",
"//tensorflow/python/lib/core:ndarray_tensor",
"//tensorflow/python/lib/core:ndarray_tensor_bridge",
"//tensorflow/python/lib/core:safe_pyobject_ptr",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/python/lib/core:numpy",
],
alwayslink = 1,
)
py_library(
name = "session",
srcs = ["session.py"],
strict_deps = True,
# copybara:uncomment_begin(google-only)
# visibility = [
# "//third_party/cloud_tpu/convergence_tools:__subpackages__",
# "//third_party/py/tf_slim:__subpackages__",
# "//tensorflow:internal",
# ],
# copybara:uncomment_end_and_comment_begin
visibility = [
"//visibility:public",
],
# copybara:comment_end
deps = [
":pywrap_tf_session",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:error_interpolation",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:stack",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:session_ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/training/experimental:mixed_precision_global_state",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:numpy_compat",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
"@pypi//wrapt",
],
)
py_library(
name = "timeline",
srcs = ["timeline.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/platform:build_info",
"//tensorflow/python/platform:tf_logging",
],
)
# Just used by tests.
tf_cuda_library(
name = "construction_fails_op",
srcs = ["test_construction_fails_op.cc"],
deps = [
"//tensorflow/core",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
],
alwayslink = 1,
)
tf_py_strict_test(
name = "session_test",
size = "medium",
srcs = ["session_test.py"],
grpc_enabled = True,
tags = [
"no_gpu", # b/127001953
"no_oss", # b/198485096
"no_pip_gpu", # testInteractivePlacePrunedGraph fails on invalid assumption about GPU ops.
"no_windows",
],
deps = [
":session",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:indexed_slices",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:sparse_tensor",
"//tensorflow/python/framework:stack",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:versions",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:control_flow_ops_gen",
"//tensorflow/python/ops:data_flow_ops",
"//tensorflow/python/ops:gradients",
"//tensorflow/python/ops:gradients_impl",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:test",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/util:compat",
"//third_party/py/numpy",
"@six_archive//:six",
],
)
tf_py_strict_test(
name = "session_clusterspec_prop_test",
size = "small",
srcs = ["session_clusterspec_prop_test.py"],
grpc_enabled = True,
tags = [
"no_gpu",
"no_oss",
"no_pip_gpu",
"notap",
],
deps = [
":session",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//tensorflow/python/training:server_lib",
"//third_party/py/numpy",
],
)
tf_py_strict_test(
name = "session_list_devices_test",
size = "small",
srcs = ["session_list_devices_test.py"],
grpc_enabled = True,
tags = [
"no_gpu",
"no_pip_gpu",
],
deps = [
":pywrap_tf_session",
":session",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:test",
"//tensorflow/python/training:server_lib",
],
)
tf_py_strict_test(
name = "session_partial_run_test",
size = "small",
srcs = ["session_partial_run_test.py"],
grpc_enabled = True,
tags = [
"no_gpu",
"no_rocm",
"no_windows",
],
deps = [
":session",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:test",
"//tensorflow/python/training:server_lib",
],
)
cuda_py_strict_test(
name = "timeline_test",
size = "small",
srcs = ["timeline_test.py"],
tags = [
"gpu_cupti",
"no_gpu", # b/154742661
],
xla_enable_strict_auto_jit = False, # Graph structure is different with autojit
deps = [
":session",
":timeline",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "virtual_gpu_test",
size = "small",
srcs = ["virtual_gpu_test.py"],
tags = [
"no_gpu", # b/127386241
"no_windows_gpu",
],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:tf_logging",
],
)
cuda_py_benchmark_test(
name = "session_benchmark",
srcs = ["session_benchmark.py"],
grpc_enabled = True,
main = "session_benchmark.py",
deps = [
":session",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:server_lib",
"//third_party/py/numpy",
],
)
tensorflow/python/client/__init__.py
View File
tensorflow/python/client/_pywrap_debug_events_writer.pyi
+26
View File
@@ -0,0 +1,26 @@
# 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 Close(arg0: str) -> None: ...
def FlushExecutionFiles(arg0: str) -> None: ...
def FlushNonExecutionFiles(arg0: str) -> None: ...
def Init(arg0: str, arg1: str, arg2: int) -> None: ...
def RegisterDeviceAndGetId(arg0: str, arg1: str) -> int: ...
def WriteDebuggedGraph(arg0: str, arg1: object) -> None: ...
def WriteExecution(arg0: str, arg1: object) -> None: ...
def WriteGraphExecutionTrace(arg0: str, arg1: object) -> None: ...
def WriteGraphOpCreation(arg0: str, arg1: object) -> None: ...
def WriteSourceFile(arg0: str, arg1: object) -> None: ...
def WriteStackFrameWithId(arg0: str, arg1: object) -> None: ...
tensorflow/python/client/_pywrap_device_lib.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 list_devices(arg0: object) -> list: ...
tensorflow/python/client/_pywrap_events_writer.pyi
+25
View File
@@ -0,0 +1,25 @@
# 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 EventsWriter:
def __init__(self, arg0: str) -> None: ...
def Close(self) -> Status: ...
def FileName(self) -> str: ...
def Flush(self) -> Status: ...
def InitWithSuffix(self, arg0: str) -> Status: ...
def WriteEvent(self, arg0: object) -> None: ...
class Status:
def __init__(self, *args, **kwargs) -> None: ...
tensorflow/python/client/_pywrap_tf_session.pyi
+394
View File
@@ -0,0 +1,394 @@
# 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.
# ==============================================================================
from typing import ClassVar, Iterator, overload
TF_ABORTED: TF_Code
TF_BFLOAT16: TF_DataType
TF_BOOL: TF_DataType
TF_CANCELLED: TF_Code
TF_COMPLEX: TF_DataType
TF_COMPLEX128: TF_DataType
TF_COMPLEX64: TF_DataType
TF_DATA_LOSS: TF_Code
TF_DEADLINE_EXCEEDED: TF_Code
TF_DOUBLE: TF_DataType
TF_FAILED_PRECONDITION: TF_Code
TF_FLOAT: TF_DataType
TF_HALF: TF_DataType
TF_INT16: TF_DataType
TF_INT32: TF_DataType
TF_INT64: TF_DataType
TF_INT8: TF_DataType
TF_INTERNAL: TF_Code
TF_INVALID_ARGUMENT: TF_Code
TF_OK: TF_Code
TF_OUT_OF_RANGE: TF_Code
TF_PERMISSION_DENIED: TF_Code
TF_QINT16: TF_DataType
TF_QINT32: TF_DataType
TF_QINT8: TF_DataType
TF_QUINT16: TF_DataType
TF_QUINT8: TF_DataType
TF_RESOURCE: TF_DataType
TF_RESOURCE_EXHAUSTED: TF_Code
TF_STRING: TF_DataType
TF_UINT16: TF_DataType
TF_UINT32: TF_DataType
TF_UINT64: TF_DataType
TF_UINT8: TF_DataType
TF_UNAUTHENTICATED: TF_Code
TF_UNIMPLEMENTED: TF_Code
TF_UNKNOWN: TF_Code
TF_VARIANT: TF_DataType
class ItemsView:
def __init__(self, *args, **kwargs) -> None: ...
def __iter__(self) -> Iterator: ...
def __len__(self) -> int: ...
class KeysView:
def __init__(self, *args, **kwargs) -> None: ...
def __contains__(self, arg0: object) -> bool: ...
def __iter__(self) -> Iterator: ...
def __len__(self) -> int: ...
class OpsById:
def __init__(self) -> None: ...
def items(self) -> ItemsView: ...
def keys(self) -> KeysView: ...
def values(self) -> ValuesView: ...
def __bool__(self) -> bool: ...
@overload
def __contains__(self, arg0: int) -> bool: ...
@overload
def __contains__(self, arg0: object) -> bool: ...
def __delitem__(self, arg0: int) -> None: ...
def __getitem__(self, arg0: int) -> object: ...
def __iter__(self) -> Iterator[int]: ...
def __len__(self) -> int: ...
def __setitem__(self, arg0: int, arg1: object) -> None: ...
class OpsByName:
def __init__(self) -> None: ...
def items(self) -> ItemsView: ...
def keys(self) -> KeysView: ...
def values(self) -> ValuesView: ...
def __bool__(self) -> bool: ...
@overload
def __contains__(self, arg0: str) -> bool: ...
@overload
def __contains__(self, arg0: object) -> bool: ...
def __delitem__(self, arg0: str) -> None: ...
def __getitem__(self, arg0: str) -> object: ...
def __iter__(self) -> Iterator[str]: ...
def __len__(self) -> int: ...
def __setitem__(self, arg0: str, arg1: object) -> None: ...
class PyGraph:
@classmethod
def __init__(cls, *args, **kwargs) -> None: ...
@classmethod
def Dismantle(cls, *args, **kwargs): ...
@classmethod
def get_operations(cls, *args, **kwargs): ...
@classmethod
def new_operations(cls, *args, **kwargs): ...
@classmethod
def num_operations(cls, *args, **kwargs): ...
@property
def operations(self) -> list: ...
@property
def version(self) -> int: ...
class PyOperation:
graph: object
@classmethod
def __init__(cls, *args, **kwargs) -> None: ...
@property
def control_inputs(self) -> list: ...
@property
def device(self) -> str: ...
@property
def name(self) -> str: ...
@property
def outputs(self) -> list: ...
@property
def type(self) -> str: ...
class PyTensor:
@classmethod
def __init__(cls, *args, **kwargs) -> None: ...
@classmethod
def consumers(cls, *args, **kwargs): ...
@property
def device(self) -> str: ...
@property
def graph(self) -> object: ...
@property
def ndim(self) -> int: ...
@property
def op(self) -> object: ...
@property
def value_index(self) -> int: ...
class TF_ApiDefMap:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Buffer:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Code:
__members__: ClassVar[dict] = ... # read-only
TF_ABORTED: ClassVar[TF_Code] = ...
TF_CANCELLED: ClassVar[TF_Code] = ...
TF_DATA_LOSS: ClassVar[TF_Code] = ...
TF_DEADLINE_EXCEEDED: ClassVar[TF_Code] = ...
TF_FAILED_PRECONDITION: ClassVar[TF_Code] = ...
TF_INTERNAL: ClassVar[TF_Code] = ...
TF_INVALID_ARGUMENT: ClassVar[TF_Code] = ...
TF_OK: ClassVar[TF_Code] = ...
TF_OUT_OF_RANGE: ClassVar[TF_Code] = ...
TF_PERMISSION_DENIED: ClassVar[TF_Code] = ...
TF_RESOURCE_EXHAUSTED: ClassVar[TF_Code] = ...
TF_UNAUTHENTICATED: ClassVar[TF_Code] = ...
TF_UNIMPLEMENTED: ClassVar[TF_Code] = ...
TF_UNKNOWN: ClassVar[TF_Code] = ...
__entries: ClassVar[dict] = ...
def __init__(self, value: int) -> None: ...
def __eq__(self, other: object) -> bool: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __int__(self) -> int: ...
def __ne__(self, other: object) -> bool: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class TF_DataType:
__members__: ClassVar[dict] = ... # read-only
TF_BFLOAT16: ClassVar[TF_DataType] = ...
TF_BOOL: ClassVar[TF_DataType] = ...
TF_COMPLEX: ClassVar[TF_DataType] = ...
TF_COMPLEX128: ClassVar[TF_DataType] = ...
TF_COMPLEX64: ClassVar[TF_DataType] = ...
TF_DOUBLE: ClassVar[TF_DataType] = ...
TF_FLOAT: ClassVar[TF_DataType] = ...
TF_HALF: ClassVar[TF_DataType] = ...
TF_INT16: ClassVar[TF_DataType] = ...
TF_INT32: ClassVar[TF_DataType] = ...
TF_INT64: ClassVar[TF_DataType] = ...
TF_INT8: ClassVar[TF_DataType] = ...
TF_QINT16: ClassVar[TF_DataType] = ...
TF_QINT32: ClassVar[TF_DataType] = ...
TF_QINT8: ClassVar[TF_DataType] = ...
TF_QUINT16: ClassVar[TF_DataType] = ...
TF_QUINT8: ClassVar[TF_DataType] = ...
TF_RESOURCE: ClassVar[TF_DataType] = ...
TF_STRING: ClassVar[TF_DataType] = ...
TF_UINT16: ClassVar[TF_DataType] = ...
TF_UINT32: ClassVar[TF_DataType] = ...
TF_UINT64: ClassVar[TF_DataType] = ...
TF_UINT8: ClassVar[TF_DataType] = ...
TF_VARIANT: ClassVar[TF_DataType] = ...
__entries: ClassVar[dict] = ...
def __init__(self, value: int) -> None: ...
def __eq__(self, other: object) -> bool: ...
def __hash__(self) -> int: ...
def __index__(self) -> int: ...
def __int__(self) -> int: ...
def __ne__(self, other: object) -> bool: ...
@property
def name(self) -> str: ...
@property
def value(self) -> int: ...
class TF_DeprecatedSession:
def __init__(self, *args, **kwargs) -> None: ...
class TF_ImportGraphDefOptions:
def __init__(self, *args, **kwargs) -> None: ...
class TF_ImportGraphDefResults:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Input:
index: int
oper: TF_Operation
def __init__(self) -> None: ...
class TF_Library:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Operation:
def __init__(self, *args, **kwargs) -> None: ...
class TF_OperationDescription:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Output:
index: int
oper: TF_Operation
def __init__(self) -> None: ...
class TF_Server:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Session:
def __init__(self, *args, **kwargs) -> None: ...
class TF_SessionOptions:
def __init__(self, *args, **kwargs) -> None: ...
class TF_Status:
def __init__(self, *args, **kwargs) -> None: ...
class ValuesView:
def __init__(self, *args, **kwargs) -> None: ...
def __iter__(self) -> Iterator: ...
def __len__(self) -> int: ...
def AddWhileInputHack(arg0: PyGraph, arg1: TF_Output, arg2: TF_Operation) -> None: ...
def ClearAttr(arg0: PyGraph, arg1: TF_Operation, arg2: str) -> None: ...
@overload
def EqualAttrValueWrapper(arg0: str, arg1: str) -> str: ...
@overload
def EqualAttrValueWrapper(arg0: str, arg1: str) -> str: ...
def EqualGraphDefWrapper(arg0: str, arg1: str) -> str: ...
def ExtendSession(arg0: TF_Session) -> None: ...
def GetHandleShapeAndType(arg0: PyGraph, arg1: TF_Output) -> bytes: ...
def GetOperationInputs(arg0: TF_Operation) -> list[TF_Output]: ...
def SetAttr(arg0: PyGraph, arg1: TF_Operation, arg2: str, arg3: TF_Buffer) -> None: ...
def SetFullType(arg0: PyGraph, arg1: TF_Operation, arg2: TF_Buffer) -> None: ...
def SetHandleShapeAndType(arg0: PyGraph, arg1: TF_Output, arg2: bytes) -> None: ...
def TF_AddControlInput(arg0: TF_OperationDescription, arg1: TF_Operation) -> None: ...
def TF_AddInput(arg0: TF_OperationDescription, arg1: TF_Output) -> None: ...
def TF_AddInputList(arg0: TF_OperationDescription, arg1: object) -> None: ...
def TF_ApiDefMapGet(arg0: TF_ApiDefMap, arg1: str, arg2: int) -> TF_Buffer: ...
def TF_ApiDefMapPut(arg0: TF_ApiDefMap, arg1: str, arg2: int) -> None: ...
def TF_CloseSession(arg0: TF_Session) -> None: ...
def TF_CreatePlaceholders(arg0: PyGraph, arg1: object, arg2: str) -> list[TF_Output]: ...
def TF_DeleteApiDefMap(arg0: TF_ApiDefMap) -> None: ...
def TF_DeleteBuffer(arg0: TF_Buffer) -> None: ...
@overload
def TF_DeleteDeviceList(arg0: TF_DeviceList) -> None: ...
@overload
def TF_DeleteDeviceList(arg0: TF_DeviceList) -> None: ...
def TF_DeleteFunction(arg0: TF_Function) -> None: ...
def TF_DeleteImportGraphDefOptions(arg0: TF_ImportGraphDefOptions) -> None: ...
def TF_DeleteImportGraphDefResults(arg0: TF_ImportGraphDefResults) -> None: ...
def TF_DeleteLibraryHandle(arg0: TF_Library) -> None: ...
def TF_DeleteSession(arg0: TF_Session) -> None: ...
def TF_DeleteSessionOptions(arg0: TF_SessionOptions) -> None: ...
def TF_DeleteStatus(arg0: TF_Status) -> None: ...
def TF_DeviceListCount(arg0: TF_DeviceList) -> int: ...
def TF_DeviceListIncarnation(arg0: TF_DeviceList, arg1: int) -> int: ...
def TF_DeviceListMemoryBytes(arg0: TF_DeviceList, arg1: int) -> int: ...
def TF_DeviceListName(arg0: TF_DeviceList, arg1: int) -> str: ...
def TF_DeviceListType(arg0: TF_DeviceList, arg1: int) -> str: ...
def TF_FinishOperation(arg0: TF_OperationDescription) -> TF_Operation: ...
def TF_FunctionImportFunctionDef(arg0: bytes) -> TF_Function: ...
def TF_FunctionImportFunctionDefNoSerialization(arg0) -> TF_Function: ...
def TF_FunctionSetAttrValueProto(arg0: TF_Function, arg1: str, arg2: bytes) -> None: ...
def TF_FunctionToFunctionDef(arg0: TF_Function, arg1: TF_Buffer) -> None: ...
def TF_GetAllOpList() -> TF_Buffer: ...
def TF_GetAllRegisteredKernels() -> TF_Buffer: ...
def TF_GetBuffer(arg0: TF_Buffer) -> object: ...
def TF_GetCode(arg0: TF_Status) -> TSL_Code: ...
def TF_GetOpList(arg0: TF_Library) -> object: ...
def TF_GetRegisteredKernelsForOp(arg0: str) -> TF_Buffer: ...
def TF_GetXlaAutoJitEnabled() -> int: ...
def TF_GetXlaConstantFoldingDisabled() -> int: ...
def TF_GraphCopyFunction(arg0: PyGraph, arg1: TF_Function, arg2: TF_Function) -> None: ...
def TF_GraphImportGraphDefWithResults(arg0: PyGraph, arg1: TF_Buffer, arg2: TF_ImportGraphDefOptions) -> TF_ImportGraphDefResults: ...
def TF_GraphImportGraphDefWithResultsNoSerialization(arg0: PyGraph, arg1, arg2: TF_ImportGraphDefOptions) -> TF_ImportGraphDefResults: ...
def TF_GraphNextOperation(arg0: PyGraph, arg1: int) -> tuple: ...
def TF_GraphRemoveFunction(arg0: PyGraph, arg1: str) -> None: ...
def TF_GraphSetOutputHandleShapesAndTypes_wrapper(arg0: PyGraph, arg1: TF_Output, arg2: list[list[int] | None], arg3: list[int], arg4: object) -> None: ...
def TF_GraphToFunction_wrapper(arg0: PyGraph, arg1: str, arg2: bool, arg3: list[TF_Operation] | None, arg4: list[TF_Output], arg5: list[TF_Output], arg6: list[bytes], arg7: list[TF_Operation], arg8: list[bytes], arg9: None, arg10: str) -> TF_Function: ...
def TF_GraphToGraphDef(arg0: PyGraph, arg1: TF_Buffer) -> None: ...
def TF_GraphToGraphDefPybind(*args, **kwargs): ...
def TF_ImportGraphDefOptionsAddInputMapping(arg0: TF_ImportGraphDefOptions, arg1: str, arg2: int, arg3: TF_Output) -> None: ...
def TF_ImportGraphDefOptionsAddReturnOperation(arg0: TF_ImportGraphDefOptions, arg1: str) -> None: ...
def TF_ImportGraphDefOptionsAddReturnOutput(arg0: TF_ImportGraphDefOptions, arg1: str, arg2: int) -> None: ...
def TF_ImportGraphDefOptionsRemapControlDependency(arg0: TF_ImportGraphDefOptions, arg1: str, arg2: TF_Operation) -> None: ...
def TF_ImportGraphDefOptionsSetPrefix(arg0: TF_ImportGraphDefOptions, arg1: str) -> None: ...
def TF_ImportGraphDefOptionsSetPropagateDeviceSpec(arg0: TF_ImportGraphDefOptions, arg1: int) -> None: ...
def TF_ImportGraphDefOptionsSetUniquifyNames(arg0: TF_ImportGraphDefOptions, arg1: int) -> None: ...
def TF_ImportGraphDefOptionsSetValidateColocationConstraints(arg0: TF_ImportGraphDefOptions, arg1: int) -> None: ...
def TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper(arg0: TF_ImportGraphDefResults) -> list[str]: ...
def TF_ImportGraphDefResultsReturnOperations(arg0: TF_ImportGraphDefResults) -> list: ...
def TF_ImportGraphDefResultsReturnOutputs(arg0: TF_ImportGraphDefResults) -> list: ...
def TF_LoadLibrary(arg0: str) -> TF_Library: ...
def TF_LoadPluggableDeviceLibrary(arg0: str) -> TF_Library: ...
def TF_NewApiDefMap(arg0: TF_Buffer) -> TF_ApiDefMap: ...
def TF_NewBuffer() -> TF_Buffer: ...
def TF_NewBufferFromString(arg0: bytes) -> TF_Buffer: ...
def TF_NewImportGraphDefOptions() -> TF_ImportGraphDefOptions: ...
def TF_NewOperation(arg0: PyGraph, arg1: str, arg2: str) -> TF_OperationDescription: ...
def TF_NewServer(arg0: bytes) -> TF_Server: ...
def TF_NewSession(arg0: PyGraph, arg1: TF_SessionOptions) -> TF_Session: ...
def TF_NewSessionRef(arg0: PyGraph, arg1: TF_SessionOptions) -> TF_Session: ...
def TF_NewStatus() -> TF_Status: ...
def TF_OperationDevice(arg0: TF_Operation) -> str: ...
def TF_OperationGetAttrBool(arg0: TF_Operation, arg1: str) -> object: ...
def TF_OperationGetAttrInt(arg0: TF_Operation, arg1: str) -> object: ...
def TF_OperationGetAttrType(arg0: TF_Operation, arg1: str) -> TF_DataType: ...
def TF_OperationGetAttrValueProto(arg0: TF_Operation, arg1: str, arg2: TF_Buffer) -> None: ...
def TF_OperationGetControlOutputs_wrapper(arg0: TF_Operation) -> list[TF_Operation]: ...
def TF_OperationGetStackTrace(arg0: TF_Operation) -> object: ...
def TF_OperationInputType(arg0: TF_Input) -> TF_DataType: ...
def TF_OperationName(arg0: TF_Operation) -> str: ...
def TF_OperationNumInputs(arg0: TF_Operation) -> int: ...
def TF_OperationNumOutputs(arg0: TF_Operation) -> int: ...
def TF_OperationOpType(arg0: TF_Operation) -> str: ...
def TF_OperationOutputType(arg0: TF_Output) -> TF_DataType: ...
def TF_OperationToNodeDef(arg0: TF_Operation, arg1: TF_Buffer) -> None: ...
def TF_PluggableDeviceLibraryHandle(arg0: TF_Library) -> None: ...
def TF_RegisterFilesystemPlugin(arg0: str) -> None: ...
def TF_Reset_wrapper(arg0: TF_SessionOptions, arg1: list[bytes]) -> None: ...
def TF_ServerJoin(arg0: TF_Server) -> None: ...
def TF_ServerStart(arg0: TF_Server) -> None: ...
def TF_ServerStop(arg0: TF_Server) -> None: ...
def TF_ServerTarget(arg0: TF_Server) -> str: ...
def TF_SessionListDevices(arg0: TF_Session) -> TF_DeviceList: ...
def TF_SessionMakeCallable(arg0: TF_Session, arg1: TF_Buffer) -> int: ...
def TF_SessionPRunSetup_wrapper(arg0: TF_Session, arg1: list[TF_Output], arg2: list[TF_Output], arg3: list[TF_Operation]) -> str: ...
def TF_SessionPRun_wrapper(arg0: TF_Session, arg1: str, arg2: object, arg3: list[TF_Output]) -> object: ...
def TF_SessionReleaseCallable(arg0: TF_Session, arg1: int) -> None: ...
def TF_SessionRunCallable(arg0: TF_Session, arg1: int, arg2: object, arg3: TF_Buffer) -> list: ...
def TF_SessionRun_wrapper(arg0: TF_Session, arg1: TF_Buffer, arg2: object, arg3: list[TF_Output], arg4: list[TF_Operation], arg5: TF_Buffer) -> object: ...
def TF_SetAttrValueProto(arg0: TF_OperationDescription, arg1: str, arg2: bytes) -> None: ...
def TF_SetDevice(arg0: TF_OperationDescription, arg1: str) -> None: ...
def TF_SetOpStackTrace(arg0: TF_Operation, arg1) -> None: ...
def TF_SetTfXlaCpuGlobalJit(arg0: int) -> int: ...
def TF_SetXlaAutoJitMode(arg0: str) -> None: ...
def TF_SetXlaConstantFoldingDisabled(arg0: int) -> None: ...
def TF_SetXlaEnableLazyCompilation(arg0: int) -> int: ...
def TF_SetXlaMinClusterSize(arg0: int) -> None: ...
def TF_TryEvaluateConstant_wrapper(arg0: PyGraph, arg1: TF_Output) -> object: ...
def UpdateEdge(arg0: PyGraph, arg1: TF_Output, arg2: TF_Input) -> None: ...
def get_compiler_version() -> str: ...
def get_cxx11_abi_flag() -> int: ...
def get_cxx_version() -> int: ...
def get_eigen_max_align_bytes() -> int: ...
def get_git_version() -> str: ...
def get_graph_def_version() -> int: ...
def get_graph_def_version_min_consumer() -> int: ...
def get_graph_def_version_min_producer() -> int: ...
def get_monolithic_build() -> int: ...
def get_tensor_handle_key() -> str: ...
def get_version() -> str: ...
tensorflow/python/client/client_lib.py
+28
View File
@@ -0,0 +1,28 @@
# 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.
# ==============================================================================
"""Support for launching graphs and executing operations.
See the [Client](https://www.tensorflow.org/guide/graphs) guide.
"""
# pylint: disable=unused-import
from tensorflow.python.client.session import InteractiveSession
from tensorflow.python.client.session import Session
from tensorflow.python.framework import errors
from tensorflow.python.framework.errors import OpError
from tensorflow.python.framework.ops import get_default_session
tensorflow/python/client/debug_events_writer_wrapper.cc
+126
View File
@@ -0,0 +1,126 @@
/* 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 <string>
#include "absl/strings/str_format.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/util/debug_events_writer.h"
#include "tensorflow/python/lib/core/pybind11_absl.h"
#include "tensorflow/python/lib/core/pybind11_proto.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
PYBIND11_MODULE(_pywrap_debug_events_writer, m) {
namespace py = pybind11;
using namespace tensorflow; // NOLINT(build/namespaces)
using namespace tensorflow::tfdbg; // NOLINT(build/namespaces)
m.def("Init",
[](const std::string& dump_root, const std::string& tfdbg_run_id,
const int64_t circular_buffer_size) {
DebugEventsWriter* writer = DebugEventsWriter::GetDebugEventsWriter(
dump_root, tfdbg_run_id, circular_buffer_size);
if (!writer->Init().ok()) {
throw py::value_error(absl::StrFormat(
"Failed to initialize debug events writer at: %s",
dump_root.c_str()));
}
});
m.def("WriteSourceFile",
[](const std::string& dump_root, const py::object obj) {
CheckProtoType(obj, "tensorflow.DebugEvent");
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(
DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
writer->WriteSerializedNonExecutionDebugEvent(
obj.attr("SerializeToString")().cast<std::string>(),
tfdbg::DebugEventFileType::SOURCE_FILES);
});
m.def("WriteStackFrameWithId",
[](const std::string& dump_root, const py::object& obj) {
CheckProtoType(obj, "tensorflow.DebugEvent");
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(
DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
writer->WriteSerializedNonExecutionDebugEvent(
obj.attr("SerializeToString")().cast<std::string>(),
tfdbg::DebugEventFileType::STACK_FRAMES);
});
m.def("WriteGraphOpCreation",
[](const std::string& dump_root, const py::object& obj) {
CheckProtoType(obj, "tensorflow.DebugEvent");
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(
DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
writer->WriteSerializedNonExecutionDebugEvent(
obj.attr("SerializeToString")().cast<std::string>(),
tfdbg::DebugEventFileType::GRAPHS);
});
m.def("WriteDebuggedGraph",
[](const std::string& dump_root, const py::object& obj) {
CheckProtoType(obj, "tensorflow.DebugEvent");
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(
DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
writer->WriteSerializedNonExecutionDebugEvent(
obj.attr("SerializeToString")().cast<std::string>(),
tfdbg::DebugEventFileType::GRAPHS);
});
m.def("WriteExecution",
[](const std::string& dump_root, const py::object& obj) {
CheckProtoType(obj, "tensorflow.DebugEvent");
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(
DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
writer->WriteSerializedExecutionDebugEvent(
obj.attr("SerializeToString")().cast<std::string>(),
tfdbg::DebugEventFileType::EXECUTION);
});
m.def("WriteGraphExecutionTrace",
[](const std::string& dump_root, const py::object& obj) {
CheckProtoType(obj, "tensorflow.DebugEvent");
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(
DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
writer->WriteSerializedExecutionDebugEvent(
obj.attr("SerializeToString")().cast<std::string>(),
tfdbg::DebugEventFileType::GRAPH_EXECUTION_TRACES);
});
m.def("RegisterDeviceAndGetId", [](const std::string& dump_root,
const std::string& device_name) {
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
return writer->RegisterDeviceAndGetId(device_name);
});
m.def("FlushNonExecutionFiles", [](const std::string& dump_root) {
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
(void)writer->FlushNonExecutionFiles();
});
m.def("FlushExecutionFiles", [](const std::string& dump_root) {
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
(void)writer->FlushExecutionFiles();
});
m.def("Close", [](const std::string& dump_root) {
DebugEventsWriter* writer = nullptr;
TF_CHECK_OK(DebugEventsWriter::LookUpDebugEventsWriter(dump_root, &writer));
(void)writer->Close();
});
};
tensorflow/python/client/device_lib.py
+42
View File
@@ -0,0 +1,42 @@
# 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.
# ==============================================================================
"""A Python interface for creating TensorFlow servers."""
from tensorflow.core.framework import device_attributes_pb2
# pylint: disable=invalid-import-order, g-bad-import-order, wildcard-import, unused-import, undefined-variable
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.client import _pywrap_device_lib
def list_local_devices(session_config=None):
"""List the available devices available in the local process.
Args:
session_config: a session config proto or None to use the default config.
Returns:
A list of `DeviceAttribute` protocol buffers.
"""
def _convert(pb_str):
m = device_attributes_pb2.DeviceAttributes()
m.ParseFromString(pb_str)
return m
serialized_config = None
if session_config is not None:
serialized_config = session_config.SerializeToString()
return [
_convert(s) for s in _pywrap_device_lib.list_devices(serialized_config)
]
tensorflow/python/client/device_lib_test.py
+42
View File
@@ -0,0 +1,42 @@
# 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.
# ==============================================================================
"""Tests for the SWIG-wrapped device lib."""
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import device_lib
from tensorflow.python.framework import test_util
from tensorflow.python.platform import googletest
from tensorflow.python.platform import test
class DeviceLibTest(test_util.TensorFlowTestCase):
def testListLocalDevices(self):
devices = device_lib.list_local_devices()
self.assertGreater(len(devices), 0)
self.assertEqual(devices[0].device_type, "CPU")
devices = device_lib.list_local_devices(config_pb2.ConfigProto())
self.assertGreater(len(devices), 0)
self.assertEqual(devices[0].device_type, "CPU")
# GPU test
if test.is_gpu_available():
self.assertGreater(len(devices), 1)
self.assertIn("GPU", [d.device_type for d in devices])
if __name__ == "__main__":
googletest.main()
tensorflow/python/client/device_lib_wrapper.cc
+61
View File
@@ -0,0 +1,61 @@
/* 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 <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/python/lib/core/pybind11_proto.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = ::pybind11;
PYBIND11_MODULE(_pywrap_device_lib, m, pybind11::mod_gil_not_used()) {
m.def("list_devices", [](py::object serialized_config) {
tensorflow::ConfigProto config;
if (!serialized_config.is_none()) {
config.ParseFromString(
static_cast<std::string>(serialized_config.cast<py::bytes>()));
}
tensorflow::SessionOptions options;
options.config = config;
std::vector<std::unique_ptr<tensorflow::Device>> devices;
tensorflow::MaybeRaiseFromStatus(tensorflow::DeviceFactory::AddDevices(
options, /*name_prefix=*/"", &devices));
py::list results;
std::string serialized_attr;
for (const auto& device : devices) {
if (!device->attributes().SerializeToString(&serialized_attr)) {
tensorflow::MaybeRaiseFromStatus(absl::InternalError(
"Could not serialize DeviceAttributes to bytes"));
}
// The default type caster for std::string assumes its contents
// is UTF8-encoded.
results.append(py::bytes(serialized_attr));
}
return results;
});
}
tensorflow/python/client/events_writer_test.py
+79
View File
@@ -0,0 +1,79 @@
# 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 the SWIG-wrapped events writer."""
import os.path
from tensorflow.core.framework import summary_pb2
from tensorflow.core.util import event_pb2
# pylint: disable=invalid-import-order, g-bad-import-order, wildcard-import, unused-import, undefined-variable
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.client import _pywrap_events_writer
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.lib.io import tf_record
from tensorflow.python.platform import googletest
from tensorflow.python.util import compat
class PywrapeventsWriterTest(test_util.TensorFlowTestCase):
def testWriteEvents(self):
file_prefix = os.path.join(self.get_temp_dir(), "events")
writer = _pywrap_events_writer.EventsWriter(compat.as_bytes(file_prefix))
filename = compat.as_text(writer.FileName())
event_written = event_pb2.Event(
wall_time=123.45,
step=67,
summary=summary_pb2.Summary(
value=[summary_pb2.Summary.Value(
tag="foo", simple_value=89.0)]))
writer.WriteEvent(event_written)
writer.Flush()
writer.Close()
with self.assertRaises(errors.NotFoundError):
for r in tf_record.tf_record_iterator(filename + "DOES_NOT_EXIST"):
self.assertTrue(False)
reader = tf_record.tf_record_iterator(filename)
event_read = event_pb2.Event()
event_read.ParseFromString(next(reader))
self.assertTrue(event_read.HasField("file_version"))
event_read.ParseFromString(next(reader))
# Second event
self.assertProtoEquals("""
wall_time: 123.45 step: 67
summary { value { tag: 'foo' simple_value: 89.0 } }
""", event_read)
with self.assertRaises(StopIteration):
next(reader)
def testWriteEventInvalidType(self):
class _Invalid(object):
def __str__(self):
return "Invalid"
with self.assertRaisesRegex(TypeError, "Invalid"):
_pywrap_events_writer.EventsWriter(b"foo").WriteEvent(_Invalid())
if __name__ == "__main__":
googletest.main()
tensorflow/python/client/events_writer_wrapper.cc
+53
View File
@@ -0,0 +1,53 @@
/* 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 <string>
#include "absl/status/status.h"
#include "pybind11/attr.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/util/events_writer.h"
#include "tensorflow/python/lib/core/pybind11_absl.h"
#include "tensorflow/python/lib/core/pybind11_proto.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
PYBIND11_MODULE(_pywrap_events_writer, m, pybind11::mod_gil_not_used()) {
py::class_<absl::Status> Status(m, "Status", py::module_local());
py::class_<tensorflow::EventsWriter> events_writer_class(m, "EventsWriter");
events_writer_class.def(py::init<const std::string&>())
.def("InitWithSuffix",
[](tensorflow::EventsWriter& self, const std::string& suffix) {
return self.InitWithSuffix(suffix);
})
.def("FileName",
[](tensorflow::EventsWriter& self) { return self.FileName(); })
.def("_WriteSerializedEvent",
[](tensorflow::EventsWriter& self, const std::string& event_str) {
self.WriteSerializedEvent(event_str);
})
.def("Flush", [](tensorflow::EventsWriter& self) { return self.Flush(); })
.def("Close", [](tensorflow::EventsWriter& self) { return self.Close(); })
.def("WriteEvent",
[](tensorflow::EventsWriter& self, const py::object obj) {
// Verify the proto type is an event prior to writing.
tensorflow::CheckProtoType(obj, "tensorflow.Event");
self.WriteSerializedEvent(
obj.attr("SerializeToString")().cast<std::string>());
});
};
tensorflow/python/client/notebook.py
+125
View File
@@ -0,0 +1,125 @@
# 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.
# ==============================================================================
"""Notebook front-end to TensorFlow.
When you run this binary, you'll see something like below, which indicates
the serving URL of the notebook:
The IPython Notebook is running at: http://127.0.0.1:8888/
Press "Shift+Enter" to execute a cell
Press "Enter" on a cell to go into edit mode.
Press "Escape" to go back into command mode and use arrow keys to navigate.
Press "a" in command mode to insert cell above or "b" to insert cell below.
Your root notebooks directory is FLAGS.notebook_dir
"""
import argparse
import os
import socket
import sys
from absl import app
# pylint: disable=g-import-not-at-top
# Official recommended way of turning on fast protocol buffers as of 10/21/14
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "cpp"
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION_VERSION"] = "2"
FLAGS = None
ORIG_ARGV = sys.argv
# Main notebook process calls itself with argv[1]="kernel" to start kernel
# subprocesses.
IS_KERNEL = len(sys.argv) > 1 and sys.argv[1] == "kernel"
def main(unused_argv):
sys.argv = ORIG_ARGV
if not IS_KERNEL:
# Drop all flags.
sys.argv = [sys.argv[0]]
# NOTE(sadovsky): For some reason, putting this import at the top level
# breaks inline plotting. It's probably a bug in the stone-age version of
# matplotlib.
from IPython.html.notebookapp import NotebookApp # pylint: disable=g-import-not-at-top
notebookapp = NotebookApp.instance()
notebookapp.open_browser = True
# password functionality adopted from quality/ranklab/main/tools/notebook.py
# add options to run with "password"
if FLAGS.password:
from IPython.lib import passwd # pylint: disable=g-import-not-at-top
notebookapp.ip = "0.0.0.0"
notebookapp.password = passwd(FLAGS.password)
else:
print("\nNo password specified; Notebook server will only be available"
" on the local machine.\n")
notebookapp.initialize(argv=["--notebook-dir", FLAGS.notebook_dir])
if notebookapp.ip == "0.0.0.0":
proto = "https" if notebookapp.certfile else "http"
url = "%s://%s:%d%s" % (proto, socket.gethostname(), notebookapp.port,
notebookapp.base_project_url)
print("\nNotebook server will be publicly available at: %s\n" % url)
notebookapp.start()
return
# Drop the --flagfile flag so that notebook doesn't complain about an
# "unrecognized alias" when parsing sys.argv.
sys.argv = ([sys.argv[0]] +
[z for z in sys.argv[1:] if not z.startswith("--flagfile")])
from IPython.kernel.zmq.kernelapp import IPKernelApp # pylint: disable=g-import-not-at-top
kernelapp = IPKernelApp.instance()
kernelapp.initialize()
# Enable inline plotting. Equivalent to running "%matplotlib inline".
ipshell = kernelapp.shell
ipshell.enable_matplotlib("inline")
kernelapp.start()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--password",
type=str,
default=None,
help="""\
Password to require. If set, the server will allow public access. Only
used if notebook config file does not exist.\
""")
parser.add_argument(
"--notebook_dir",
type=str,
default="experimental/brain/notebooks",
help="root location where to store notebooks")
# When the user starts the main notebook process, we don't touch sys.argv.
# When the main process launches kernel subprocesses, it writes all flags
# to a tmpfile and sets --flagfile to that tmpfile, so for kernel
# subprocesses here we drop all flags *except* --flagfile, then call
# app.run(), and then (in main) restore all flags before starting the
# kernel app.
if IS_KERNEL:
# Drop everything except --flagfile.
sys.argv = (
[sys.argv[0]] + [x for x in sys.argv[1:] if x.startswith("--flagfile")])
FLAGS, unparsed = parser.parse_known_args()
app.run(main=main, argv=[sys.argv[0]] + unparsed)
tensorflow/python/client/pywrap_tf_session.py
+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.
# ==============================================================================
"""Python module for Session ops, vars, and functions exported by pybind11."""
# pylint: disable=invalid-import-order,g-bad-import-order, wildcard-import, unused-import
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.client._pywrap_tf_session import *
from tensorflow.python.client._pywrap_tf_session import _TF_SetTarget
from tensorflow.python.client._pywrap_tf_session import _TF_SetConfig
from tensorflow.python.client._pywrap_tf_session import _TF_NewSessionOptions
# Register pybind11 type caster for StackTraceWrapper/AbstractStackTrace
from tensorflow.python.util import tf_stack
# Convert versions to strings for Python2 and keep api_compatibility_test green.
# We can remove this hack once we remove Python2 presubmits. pybind11 can only
# return unicode for Python2 even with py::str.
# https://pybind11.readthedocs.io/en/stable/advanced/cast/strings.html#returning-c-strings-to-python
# pylint: disable=undefined-variable
__version__ = str(get_version())
__git_version__ = str(get_git_version())
__compiler_version__ = str(get_compiler_version())
__cxx11_abi_flag__ = get_cxx11_abi_flag()
__cxx_version__ = get_cxx_version()
__monolithic_build__ = get_monolithic_build()
# User getters to hold attributes rather than pybind11's m.attr due to
# b/145559202.
GRAPH_DEF_VERSION = get_graph_def_version()
GRAPH_DEF_VERSION_MIN_CONSUMER = get_graph_def_version_min_consumer()
GRAPH_DEF_VERSION_MIN_PRODUCER = get_graph_def_version_min_producer()
TENSOR_HANDLE_KEY = get_tensor_handle_key()
# pylint: enable=undefined-variable
# Disable pylint invalid name warnings for legacy functions.
# pylint: disable=invalid-name
def TF_NewSessionOptions(target=None, config=None):
# NOTE: target and config are validated in the session constructor.
opts = _TF_NewSessionOptions()
if target is not None:
_TF_SetTarget(opts, target)
if config is not None:
config_str = config.SerializeToString()
_TF_SetConfig(opts, config_str)
return opts
# Disable pylind undefined-variable as the variable is exported in the shared
# object via pybind11.
# pylint: disable=undefined-variable
def TF_Reset(target, containers=None, config=None):
opts = TF_NewSessionOptions(target=target, config=config)
try:
TF_Reset_wrapper(opts, containers)
finally:
TF_DeleteSessionOptions(opts)
tensorflow/python/client/session.py
+1835
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/client/session_benchmark.py
+215
View File
@@ -0,0 +1,215 @@
# 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.
# ==============================================================================
"""Tests and benchmarks for interacting with the `tf.compat.v1.Session`."""
import time
import numpy as np
from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
class SessionBenchmark(test.Benchmark):
"""Tests and benchmarks for interacting with the `tf.compat.v1.Session`."""
def _benchmarkFeed(self, name, target, size, iters):
"""Runs a microbenchmark to measure the cost of feeding a tensor.
Reports the median cost of feeding a tensor of `size` * `sizeof(float)`
bytes.
Args:
name: A human-readable name for logging the output.
target: The session target to use for the benchmark.
size: The number of floating-point numbers to be feed.
iters: The number of iterations to perform.
"""
feed_val = np.random.rand(size).astype(np.float32)
times = []
with ops.Graph().as_default():
p = array_ops.placeholder(dtypes.float32, shape=[size])
# Fetch the operation rather than the tensor, to avoid measuring the time
# to fetch back the value.
no_op = array_ops.identity(p).op
with session.Session(target) as sess:
sess.run(no_op, feed_dict={p: feed_val}) # Warm-up run.
for _ in range(iters):
start_time = time.time()
sess.run(no_op, feed_dict={p: feed_val})
end_time = time.time()
times.append(end_time - start_time)
print("%s %d %f" % (name, size, np.median(times)))
self.report_benchmark(iters=1, wall_time=np.median(times), name=name)
def _benchmarkFetch(self, name, target, size, iters):
"""Runs a microbenchmark to measure the cost of fetching a tensor.
Reports the median cost of fetching a tensor of `size` * `sizeof(float)`
bytes.
Args:
name: A human-readable name for logging the output.
target: The session target to use for the benchmark.
size: The number of floating-point numbers to be fetched.
iters: The number of iterations to perform.
"""
times = []
with ops.Graph().as_default():
# Define the tensor to be fetched as a variable, to avoid
# constant-folding.
v = variables.Variable(random_ops.random_normal([size]))
with session.Session(target) as sess:
sess.run(v.initializer)
sess.run(v) # Warm-up run.
for _ in range(iters):
start_time = time.time()
sess.run(v)
end_time = time.time()
times.append(end_time - start_time)
print("%s %d %f" % (name, size, np.median(times)))
self.report_benchmark(iters=1, wall_time=np.median(times), name=name)
def _benchmarkFetchPrebuilt(self, name, target, size, iters):
"""Runs a microbenchmark to measure the cost of fetching a tensor.
Reports the median cost of fetching a tensor of `size` * `sizeof(float)`
bytes.
Args:
name: A human-readable name for logging the output.
target: The session target to use for the benchmark.
size: The number of floating-point numbers to be fetched.
iters: The number of iterations to perform.
"""
times = []
with ops.Graph().as_default():
# Define the tensor to be fetched as a variable, to avoid
# constant-folding.
v = variables.Variable(random_ops.random_normal([size]))
with session.Session(target) as sess:
sess.run(v.initializer)
runner = sess.make_callable(v)
runner() # Warm-up run.
for _ in range(iters):
start_time = time.time()
runner()
end_time = time.time()
times.append(end_time - start_time)
print("%s %d %f" % (name, size, np.median(times)))
self.report_benchmark(iters=1, wall_time=np.median(times), name=name)
def _benchmarkRunOp(self, name, target, iters):
"""Runs a microbenchmark to measure the cost of running an op.
Reports the median cost of running a trivial (Variable) op.
Args:
name: A human-readable name for logging the output.
target: The session target to use for the benchmark.
iters: The number of iterations to perform.
"""
times = []
with ops.Graph().as_default():
# Define the op to be run as a variable, to avoid
# constant-folding.
v = variables.Variable(random_ops.random_normal([]))
with session.Session(target) as sess:
sess.run(v.initializer)
sess.run(v.op) # Warm-up run.
for _ in range(iters):
start_time = time.time()
sess.run(v.op)
end_time = time.time()
times.append(end_time - start_time)
print("%s %f" % (name, np.median(times)))
self.report_benchmark(iters=1, wall_time=np.median(times), name=name)
def _benchmarkRunOpPrebuilt(self, name, target, iters):
"""Runs a microbenchmark to measure the cost of running an op.
Reports the median cost of running a trivial (Variable) op.
Args:
name: A human-readable name for logging the output.
target: The session target to use for the benchmark.
iters: The number of iterations to perform.
"""
times = []
with ops.Graph().as_default():
# Define the op to be run as a variable, to avoid
# constant-folding.
v = variables.Variable(random_ops.random_normal([]))
with session.Session(target) as sess:
sess.run(v.initializer)
runner = sess.make_callable(v.op)
runner() # Warm-up run.
for _ in range(iters):
start_time = time.time()
runner()
end_time = time.time()
times.append(end_time - start_time)
print("%s %f" % (name, np.median(times)))
self.report_benchmark(iters=1, wall_time=np.median(times), name=name)
def benchmarkGrpcSession(self):
server = server_lib.Server.create_local_server()
self._benchmarkFeed("benchmark_session_feed_grpc_4B", server.target, 1,
30000)
session.Session.reset(server.target)
self._benchmarkFeed("benchmark_session_feed_grpc_4MB", server.target,
1 << 20, 25000)
session.Session.reset(server.target)
self._benchmarkFetch("benchmark_session_fetch_grpc_4B", server.target, 1,
40000)
session.Session.reset(server.target)
self._benchmarkFetch("benchmark_session_fetch_grpc_4MB", server.target,
1 << 20, 20000)
session.Session.reset(server.target)
self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_grpc_4B",
server.target, 1, 50000)
session.Session.reset(server.target)
self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_grpc_4MB",
server.target, 1 << 20, 50000)
session.Session.reset(server.target)
self._benchmarkRunOp("benchmark_session_runop_grpc", server.target, 50000)
session.Session.reset(server.target)
self._benchmarkRunOpPrebuilt("benchmark_session_runopprebuilt_grpc",
server.target, 100000)
session.Session.reset(server.target)
def benchmarkDirectSession(self):
self._benchmarkFeed("benchmark_session_feed_direct_4B", "", 1, 80000)
self._benchmarkFeed("benchmark_session_feed_direct_4MB", "", 1 << 20, 20000)
self._benchmarkFetch("benchmark_session_fetch_direct_4B", "", 1, 100000)
self._benchmarkFetch("benchmark_session_fetch_direct_4MB", "", 1 << 20,
20000)
self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_direct_4B",
"", 1, 200000)
self._benchmarkFetchPrebuilt("benchmark_session_fetchprebuilt_direct_4MB",
"", 1 << 20, 200000)
self._benchmarkRunOp("benchmark_session_runop_direct", "", 200000)
self._benchmarkRunOpPrebuilt("benchmark_session_runopprebuilt_direct", "",
200000)
if __name__ == "__main__":
test.main()
tensorflow/python/client/session_clusterspec_prop_test.py
+562
View File
@@ -0,0 +1,562 @@
# 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 tensorflow.python.client.session.Session's ClusterSpec Propagation.
These tests exercise the ClusterSpec Propagation capabilities of distributed
Sessions.
"""
import numpy as np
from tensorflow.core.protobuf import cluster_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
# Import resource_variable_ops for the variables-to-tensor implicit conversion.
from tensorflow.python.ops import resource_variable_ops # pylint: disable=unused-import
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import googletest
from tensorflow.python.platform import test
from tensorflow.python.training import server_lib
class SessionClusterSpecPropagationTest(test_util.TensorFlowTestCase):
def testClusterSpecPropagationSimple(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
const = constant_op.constant(17)
sess = session.Session(server1.target, config=config)
output = self.evaluate(const)
self.assertEqual(17, output)
def testClusterSpecPropagationWorker2Placement(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.Graph().as_default() as g, ops.device('/job:worker/task:1'):
with ops.device('/cpu:0'):
const = constant_op.constant(17)
sess = session.Session(server1.target, config=config, graph=g)
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
output = sess.run(const, options=run_options, run_metadata=run_metadata)
self.assertEqual(17, output)
self.assertEqual(1,
len([
node_stats
for dev_stats in run_metadata.step_stats.dev_stats
for node_stats in dev_stats.node_stats
if '/job:worker/replica:0/task:1/device:CPU:0' ==
dev_stats.device and 'Const' == node_stats.node_name
]))
def testClusterSpecPropagationWorker1Placement(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.Graph().as_default() as g, ops.device('/job:worker/task:0'):
const = constant_op.constant(17)
with session.Session(server1.target, config=config, graph=g):
output = self.evaluate(const)
self.assertEqual(17, output)
def testCanonicalDeviceNames(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.Graph().as_default() as g, ops.device(
'/job:worker/task:1/device:CPU:0'):
const = constant_op.constant(17)
sess = session.Session(server1.target, config=config, graph=g)
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
output = sess.run(const, options=run_options, run_metadata=run_metadata)
self.assertEqual(17, output)
self.assertEqual(1,
len([
node_stats
for dev_stats in run_metadata.step_stats.dev_stats
for node_stats in dev_stats.node_stats
if '/job:worker/replica:0/task:1/device:CPU:0' ==
dev_stats.device and 'Const' == node_stats.node_name
]))
def testFullDeviceNames(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'renamed_worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.Graph().as_default() as g, ops.device(
'/job:renamed_worker/replica:0/task:1/device:CPU:0'):
const = constant_op.constant(17)
sess = session.Session(server1.target, config=config, graph=g)
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
output = sess.run(const, options=run_options, run_metadata=run_metadata)
self.assertEqual(17, output)
self.assertEqual(1,
len([
node_stats
for dev_stats in run_metadata.step_stats.dev_stats
for node_stats in dev_stats.node_stats
if '/job:renamed_worker/replica:0/task:1/device:CPU:0'
== dev_stats.device and 'Const' == node_stats.node_name
]))
def testMultipleLocalDevices(self):
# Note: CPU->CPU transfers have a fast-path in
# BaseRemoteRendezvous::SameWorkerRecvDone that means the test doesn't
# actually capture the motivating bug unless run on a GPU machine.
#
# Example error message (before bugfix -- line breaks added because lint):
#
# W0718 17:14:41.521534 190121 device_mgr.cc:107] Unknown device:
# /job:worker/replica:0/task:0/device:CPU:0 all devices:
# /job:local/replica:0/task:0/device:GPU:0,
# /job:local/replica:0/task:0/device:GPU:0,
# /job:local/replica:0/task:0/cpu:1, CPU:0, GPU:0,
# /job:local/replica:0/task:0/device:CPU:1,
# /job:local/replica:0/task:0/device:CPU:0, CPU:1,
# /job:local/replica:0/task:0/cpu:0
server_config = config_pb2.ConfigProto(device_count={'CPU': 2})
server1 = server_lib.Server.create_local_server(config=server_config)
server2 = server_lib.Server.create_local_server(config=server_config)
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.Graph().as_default() as g:
with ops.device('/job:worker/task:1/cpu:1'):
input1 = constant_op.constant(17, dtypes.float32)
with ops.device('/job:worker/task:0/cpu:1'):
input2 = constant_op.constant(3, dtypes.float32)
with ops.device('/job:worker/task:1/cpu:0'):
sum1 = input1 + input2
if test.is_gpu_available():
device_str = '/job:worker/task:0/device:GPU:0'
else:
device_str = '/job:worker/task:0/cpu:1'
with ops.device(device_str):
sum2 = input2 + input1
with ops.device('/job:worker/task:0/cpu:0'):
sum3 = sum1 + sum2
with session.Session(server1.target, config=config, graph=g):
output = self.evaluate(sum3)
self.assertEqual(40, output)
def testLegacyDeviceNames(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.Graph().as_default() as g, ops.device('/job:worker/task:1/cpu:0'):
const = constant_op.constant(17)
sess = session.Session(server1.target, config=config, graph=g)
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
output = sess.run(const, options=run_options, run_metadata=run_metadata)
self.assertEqual(17, output)
self.assertEqual(1,
len([
node_stats
for dev_stats in run_metadata.step_stats.dev_stats
for node_stats in dev_stats.node_stats
if '/job:worker/replica:0/task:1/device:CPU:0' ==
dev_stats.device and 'Const' == node_stats.node_name
]))
def testClusterSpecPropagationThreeServers2Graphs(self):
"""Boots 3 servers, creates 2 sessions, ensures appropriate operations.
We create 2 clusterspecs:
1. server2 as the master, server1 as a worker
2. server2 as the master, server3 as a worker
We ensure that variables on the workers are independent.
"""
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
server3 = server_lib.Server.create_local_server()
cluster_def1 = cluster_pb2.ClusterDef()
job1 = cluster_def1.job.add()
job1.name = 'worker1'
job1.tasks[0] = server2.target[len('grpc://'):]
job1.tasks[1] = server1.target[len('grpc://'):]
cluster_def2 = cluster_pb2.ClusterDef()
job2 = cluster_def2.job.add()
job2.name = 'worker2'
job2.tasks[0] = server2.target[len('grpc://'):]
job2.tasks[1] = server3.target[len('grpc://'):]
config1 = config_pb2.ConfigProto(cluster_def=cluster_def1)
config2 = config_pb2.ConfigProto(cluster_def=cluster_def2)
with ops.Graph().as_default() as g1:
with ops.device('/job:worker1/task:1'):
var1 = variables.Variable(array_ops.zeros([2]), name='var1')
update_op1 = state_ops.assign_add(
var1, array_ops.ones([2]), name='var1_assign_add')
init1 = variables.global_variables_initializer()
with ops.Graph().as_default() as g2:
with ops.device('/job:worker2/task:1'):
var2 = variables.Variable(array_ops.zeros([2]), name='var2')
update_op2 = state_ops.assign_add(
var2, array_ops.ones([2]), name='var2_assign_add')
init2 = variables.global_variables_initializer()
sess1 = session.Session(server2.target, graph=g1, config=config1)
sess2 = session.Session(server2.target, graph=g2, config=config2)
init1.run(session=sess1)
init2.run(session=sess2)
expected_zeros = np.zeros([2])
expected_ones = np.ones([2])
self.assertAllEqual(expected_zeros, sess1.run(var1))
self.assertAllEqual(expected_zeros, sess2.run(var2))
self.assertAllEqual(expected_ones, sess1.run(update_op1))
self.assertAllEqual(expected_ones, sess1.run(var1))
self.assertAllEqual(expected_zeros, sess2.run(var2))
self.assertAllEqual(expected_ones, sess2.run(update_op2))
self.assertAllEqual(expected_ones + expected_ones, sess1.run(update_op1))
self.assertAllEqual(expected_ones, sess2.run(var2))
self.assertAllEqual(expected_ones + expected_ones, sess1.run(var1))
def testClusterSpecPropagationThreeServers(self):
"""Boots 3 servers, creates 2 sessions, ensures appropriate operations.
We create 2 clusterspecs:
1. server2 as the master, server1 as a worker
2. server2 as the master, server3 as a worker
We ensure that variables on the workers are independent.
"""
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
server3 = server_lib.Server.create_local_server()
cluster_def1 = cluster_pb2.ClusterDef()
job1 = cluster_def1.job.add()
job1.name = 'worker'
job1.tasks[0] = server2.target[len('grpc://'):]
job1.tasks[1] = server1.target[len('grpc://'):]
cluster_def2 = cluster_pb2.ClusterDef()
job2 = cluster_def2.job.add()
job2.name = 'worker'
job2.tasks[0] = server2.target[len('grpc://'):]
job2.tasks[1] = server3.target[len('grpc://'):]
config1 = config_pb2.ConfigProto(cluster_def=cluster_def1)
config2 = config_pb2.ConfigProto(cluster_def=cluster_def2)
with ops.device('/job:worker/task:1'):
var = variables.Variable(array_ops.zeros([2]), name='var')
feed = array_ops.placeholder(dtypes.float32, shape=(2))
update_op = var.assign_add(feed)
sess1 = session.Session(server2.target, config=config1)
sess2 = session.Session(server2.target, config=config2)
variables.global_variables_initializer().run(session=sess1)
variables.global_variables_initializer().run(session=sess2)
expected_zeros = np.zeros([2])
expected_ones = np.ones([2])
self.assertAllEqual(expected_zeros, sess1.run(var))
self.assertAllEqual(expected_zeros, sess2.run(var))
self.assertAllEqual(expected_ones,
sess1.run(update_op, feed_dict={feed: expected_ones}))
self.assertAllEqual(expected_ones, sess1.run(var))
self.assertAllEqual(expected_zeros, sess2.run(var))
self.assertAllEqual(expected_ones,
sess2.run(update_op, feed_dict={feed: expected_ones}))
self.assertAllEqual(expected_ones + expected_ones,
sess1.run(update_op, feed_dict={feed: expected_ones}))
self.assertAllEqual(expected_ones, sess2.run(var))
self.assertAllEqual(expected_ones + expected_ones, sess1.run(var))
def testClusterSpecPropagationThreeServersOneCluster(self):
"""Boots 3 servers, ensures appropriate communication across workers.
Additionally, in this cluster, we ensure the master is not the 0-th worker.
Note: this test only uses one session.
"""
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
server3 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server3.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
job.tasks[2] = server1.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
# Add ops to the devices in non-linear order.
with ops.device('/job:worker/task:1'):
feed1 = array_ops.placeholder(dtypes.float32, shape=(2))
const1 = constant_op.constant(2.0)
mul1 = const1 * feed1
with ops.device('/job:worker/task:2'):
feed2 = array_ops.placeholder(dtypes.float32, shape=(2))
const2 = constant_op.constant(2.0)
mul2 = const2 * feed2
with ops.device('/job:worker/task:0'):
feed0 = array_ops.placeholder(dtypes.float32, shape=(2))
const0 = constant_op.constant(2.0)
mul0 = const0 * feed0
sum_op = mul0 + mul1 + mul2
ones = np.ones([2])
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
# Run!
with session.Session(server1.target, config=config) as sess:
output = sess.run(
sum_op,
options=run_options,
run_metadata=run_metadata,
feed_dict={feed1: ones,
feed2: ones,
feed0: ones})
self.assertAllEqual(6 * ones, output)
self.assertEqual(
3,
len([
dev_stats.device
for dev_stats in run_metadata.step_stats.dev_stats
for node_stats in dev_stats.node_stats
if '/job:worker/replica:0/task:' in dev_stats.device and
node_stats.node_name.startswith('Const')
]), run_metadata)
def testClusterSpecPropagationIsolation(self):
"""Test that two sessions using ClusterSpec propagation are isolated."""
server = server_lib.Server.create_local_server()
init_value = array_ops.placeholder(dtypes.int32, shape=[])
v = variables.Variable(init_value)
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
sess1 = session.Session(server.target, config=config)
sess2 = session.Session(server.target, config=config)
# Initially, the variable is uninitialized in both sessions.
with self.assertRaises(errors.FailedPreconditionError):
sess1.run(v)
with self.assertRaises(errors.FailedPreconditionError):
sess2.run(v)
# An update in sess1 should be visible in sess1 only.
sess1.run(v.initializer, feed_dict={init_value: 37})
self.assertEqual(37, sess1.run(v))
with self.assertRaises(errors.FailedPreconditionError):
sess2.run(v)
# An update in sess2 should be visible in sess2 only.
sess2.run(v.initializer, feed_dict={init_value: 86})
self.assertEqual(37, sess1.run(v))
self.assertEqual(86, sess2.run(v))
# Closing sess2 has no effect on the state of sess1.
sess2.close()
self.assertEqual(37, sess1.run(v))
# Subsequent sessions will not see the state of existing sessions.
sess3 = session.Session(server.target, config=config)
self.assertEqual(37, sess1.run(v))
with self.assertRaises(errors.FailedPreconditionError):
sess3.run(v)
def testClusterSpecPropagationNonIsolation(self):
"""Test that two sessions using ClusterSpec propagation shares state.
For example, the updated Variable value are visible among all worker
sessions registered in the same server.
"""
server = server_lib.Server.create_local_server()
init_value = array_ops.placeholder(dtypes.int32, shape=[])
v = variables.Variable(init_value)
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
config.experimental.share_session_state_in_clusterspec_propagation = True
sess1 = session.Session(server.target, config=config)
sess2 = session.Session(server.target, config=config)
# Initially, the variable is uninitialized in both sessions.
with self.assertRaises(errors.FailedPreconditionError):
sess1.run(v)
with self.assertRaises(errors.FailedPreconditionError):
sess2.run(v)
# An update in sess1 should be visible in sess2.
sess1.run(v.initializer, feed_dict={init_value: 37})
self.assertEqual(37, sess1.run(v))
self.assertEqual(37, sess2.run(v))
# Closing sess2 has no effect on the state of sess1.
sess2.close()
self.assertEqual(37, sess1.run(v))
# Subsequent sessions should see the state of existing sessions.
sess3 = session.Session(server.target, config=config)
self.assertEqual(37, sess1.run(v))
self.assertEqual(37, sess3.run(v))
def testClusterSpecPropagationNonIsolation2Graphs(self):
"""Creates 2 sessions with each own graph, ensures appropriate operations.
We ensure that variables on the workers shares state.
"""
server = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
config.experimental.share_session_state_in_clusterspec_propagation = True
with ops.Graph().as_default() as g1:
var1 = variables.Variable(array_ops.zeros([2]), name='var')
update_op1 = state_ops.assign_add(
var1, array_ops.ones([2]), name='var1_assign_add')
init1 = variables.global_variables_initializer()
with ops.Graph().as_default() as g2:
var2 = variables.Variable(array_ops.zeros([2]), name='var')
update_op2 = state_ops.assign_add(
var2, array_ops.ones([2]), name='var2_assign_add')
sess1 = session.Session(server.target, graph=g1, config=config)
sess2 = session.Session(server.target, graph=g2, config=config)
expected_zeros = np.zeros([2])
expected_ones = np.ones([2])
init1.run(session=sess1)
self.assertAllEqual(expected_zeros, sess1.run(var1))
self.assertAllEqual(expected_zeros, sess2.run(var2))
self.assertAllEqual(expected_ones, sess1.run(update_op1))
self.assertAllEqual(expected_ones, sess1.run(var1))
self.assertAllEqual(expected_ones, sess2.run(var2))
self.assertAllEqual(expected_ones + expected_ones, sess2.run(update_op2))
self.assertAllEqual(expected_ones + expected_ones, sess2.run(var2))
self.assertAllEqual(expected_ones + expected_ones, sess1.run(var1))
def testClusterSpecPropagationPartialRun(self):
"""Test successful partial run with ClusterSpec propagation."""
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with ops.device('/job:worker/task:0'):
a = array_ops.placeholder(dtypes.float32, shape=[])
with ops.device('/job:worker/task:1'):
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
with ops.device('/job:worker/task:0'):
r2 = math_ops.multiply(r1, c)
with session.Session(server1.target, config=config) as sess:
h = sess.partial_run_setup([r1, r2], [a, b, c])
res = sess.partial_run(h, r1, feed_dict={a: 1, b: 2})
self.assertEqual(3, res)
res = sess.partial_run(h, r2, feed_dict={c: 3})
self.assertEqual(9, res)
if __name__ == '__main__':
googletest.main()
tensorflow/python/client/session_list_devices_test.py
+89
View File
@@ -0,0 +1,89 @@
# 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 tensorflow.python.client.session.Session's list_devices API."""
from tensorflow.core.protobuf import cluster_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import pywrap_tf_session as tf_session
from tensorflow.python.client import session
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.platform import googletest
from tensorflow.python.training import server_lib
class SessionListDevicesTest(test_util.TensorFlowTestCase):
def testListDevices(self):
with session.Session() as sess:
devices = sess.list_devices()
self.assertIn(
'/job:localhost/replica:0/task:0/device:CPU:0',
{d.name for d in devices},
devices,
)
# All valid device incarnations must be non-zero.
self.assertTrue(all(d.incarnation != 0 for d in devices))
def testInvalidDeviceNumber(self):
opts = tf_session.TF_NewSessionOptions()
with ops.get_default_graph()._c_graph.get() as c_graph:
c_session = tf_session.TF_NewSession(c_graph, opts)
raw_device_list = tf_session.TF_SessionListDevices(c_session)
size = tf_session.TF_DeviceListCount(raw_device_list)
with self.assertRaises(errors.InvalidArgumentError):
tf_session.TF_DeviceListMemoryBytes(raw_device_list, size)
tf_session.TF_DeleteDeviceList(raw_device_list)
tf_session.TF_CloseSession(c_session)
def testTFGetBufferNull(self):
with self.assertRaises((ValueError, TypeError)):
tf_session.TF_GetBuffer(None)
def testListDevicesGrpcSession(self):
server = server_lib.Server.create_local_server()
with session.Session(server.target) as sess:
devices = sess.list_devices()
self.assertIn(
'/job:localhost/replica:0/task:0/device:CPU:0',
{d.name for d in devices},
devices,
)
# All valid device incarnations must be non-zero.
self.assertTrue(all(d.incarnation != 0 for d in devices))
def testListDevicesClusterSpecPropagation(self):
server1 = server_lib.Server.create_local_server()
server2 = server_lib.Server.create_local_server()
cluster_def = cluster_pb2.ClusterDef()
job = cluster_def.job.add()
job.name = 'worker'
job.tasks[0] = server1.target[len('grpc://'):]
job.tasks[1] = server2.target[len('grpc://'):]
config = config_pb2.ConfigProto(cluster_def=cluster_def)
with session.Session(server1.target, config=config) as sess:
devices = sess.list_devices()
device_names = {d.name for d in devices}
self.assertIn('/job:worker/replica:0/task:0/device:CPU:0', device_names)
self.assertIn('/job:worker/replica:0/task:1/device:CPU:0', device_names)
# All valid device incarnations must be non-zero.
self.assertTrue(all(d.incarnation != 0 for d in devices))
if __name__ == '__main__':
googletest.main()
tensorflow/python/client/session_partial_run_test.py
+299
View File
@@ -0,0 +1,299 @@
# 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 tensorflow.python.client.session.Session's partial run APIs."""
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import googletest
from tensorflow.python.training import server_lib
class PartialRunTest(test_util.TensorFlowTestCase):
def RunTestPartialRun(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
r2 = math_ops.multiply(r1, c)
h = sess.partial_run_setup([r1, r2], [a, b, c])
res = sess.partial_run(h, r1, feed_dict={a: 1, b: 2})
self.assertEqual(3, res)
temp = res * 17
res = sess.partial_run(h, r2, feed_dict={c: temp})
self.assertEqual(153, res)
# Call again on the same graph.
h2 = sess.partial_run_setup([r1, r2], [a, b, c])
res = sess.partial_run(h2, r1, feed_dict={a: 1, b: 2})
self.assertEqual(3, res)
temp = res * 18
res = sess.partial_run(h2, r2, feed_dict={c: temp})
self.assertEqual(162, res)
def RunTestPartialRunIncomplete(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
r2 = math_ops.multiply(r1, c)
h = sess.partial_run_setup([r1, r2], [a, b, c])
res = sess.partial_run(h, r1, feed_dict={a: 1, b: 2})
self.assertEqual(3, res)
def RunTestConcurrentPartialRun(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
r2 = math_ops.multiply(r1, c)
h1 = sess.partial_run_setup([r1], [a, b, c])
h2 = sess.partial_run_setup([r1, r2], [a, b, c])
res = sess.partial_run(h1, r1, feed_dict={a: 1, b: 2})
self.assertEqual(3, res)
temp = res * 19
res = sess.partial_run(h2, r1, feed_dict={a: temp, b: 9})
self.assertEqual(66, res)
res = sess.partial_run(h2, r2, feed_dict={c: 7})
self.assertEqual(462, res)
def RunTestManyPartialRun(self, sess):
steps = 200
inputs = []
outputs = []
a = constant_op.constant(2.0, dtypes.float32)
for i in range(steps):
inputs.append(array_ops.placeholder(dtypes.float32, shape=[]))
a = math_ops.multiply(a, inputs[i])
outputs.append(a)
h = sess.partial_run_setup(outputs, inputs)
for i in range(steps):
res = sess.partial_run(h, outputs[i], feed_dict={inputs[i]: 1.0})
self.assertEqual(2.0, res)
feed_dict = {}
for i in range(steps):
feed_dict[inputs[i]] = 1.0
res = sess.run(outputs, feed_dict)
self.assertEqual(steps, len(res))
self.assertEqual(2.0, res[-1])
def RunTestRunAndPartialRun(self, sess):
a = constant_op.constant(2.0, dtypes.float32)
b = a * 2
c = b * 3
r1 = self.evaluate([b, c])
h = sess.partial_run_setup([b, c], [])
r2 = sess.partial_run(h, [b, c])
self.assertEqual(r1, r2)
def RunTestPartialRunMissingPlaceholderFeedException(self, sess):
x = array_ops.placeholder(dtypes.float32, shape=())
fetches = [x * 2, x * 3]
handle = sess.partial_run_setup(fetches=fetches, feeds=[])
with self.assertRaisesRegex(errors.InvalidArgumentError,
'You must feed a value for placeholder'):
sess.partial_run(handle, fetches[0])
def RunTestPartialRunUnspecifiedFeed(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
h = sess.partial_run_setup([r1], [a, b])
with self.assertRaisesRegex(errors.InvalidArgumentError,
'was not specified in partial_run_setup.$'):
sess.partial_run(h, r1, feed_dict={a: 1, b: 2, c: 3})
def RunTestPartialRunUnspecifiedFetch(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
r2 = math_ops.multiply(a, c)
h = sess.partial_run_setup([r1], [a, b, c])
with self.assertRaisesRegex(errors.InvalidArgumentError,
'was not specified in partial_run_setup.$'):
sess.partial_run(h, r2, feed_dict={a: 1, c: 3})
def RunTestPartialRunAlreadyFed(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
r2 = math_ops.multiply(a, c)
h = sess.partial_run_setup([r1, r2], [a, b, c])
sess.partial_run(h, r1, feed_dict={a: 1, b: 2})
with self.assertRaisesRegex(errors.InvalidArgumentError,
'has already been fed.$'):
sess.partial_run(h, r2, feed_dict={a: 1, c: 3})
def RunTestPartialRunAlreadyFetched(self, sess):
a = array_ops.placeholder(dtypes.float32, shape=[])
b = array_ops.placeholder(dtypes.float32, shape=[])
c = array_ops.placeholder(dtypes.float32, shape=[])
r1 = math_ops.add(a, b)
r2 = math_ops.multiply(a, c)
h = sess.partial_run_setup([r1, r2], [a, b, c])
sess.partial_run(h, r1, feed_dict={a: 1, b: 2})
with self.assertRaisesRegex(errors.InvalidArgumentError,
'has already been fetched.$'):
sess.partial_run(h, r1, feed_dict={c: 3})
def RunTestPartialRunEmptyFetches(self, sess):
a = array_ops.placeholder(dtypes.float32)
b = a * 2.0
h = sess.partial_run_setup(fetches=[b], feeds=[a])
sess.partial_run(h, [], {a: 3.0})
r = sess.partial_run(h, [b], {})
self.assertEqual([6.0], r)
@test_util.run_deprecated_v1
def testInvalidPartialRunSetup(self):
sess = session.Session()
x = array_ops.placeholder(dtypes.float32, shape=[])
with self.assertRaisesRegex(
errors.InvalidArgumentError,
'specify at least one target to fetch or execute.'):
sess.partial_run_setup(fetches=[], feeds=[x])
@test_util.run_deprecated_v1
def testPartialRunSetupNoFeedsPassed(self):
sess = session.Session()
r1 = constant_op.constant([6.0])
h = sess.partial_run_setup([r1])
result1 = sess.partial_run(h, r1)
self.assertEqual([6.0], result1)
@test_util.run_deprecated_v1
def testPartialRunDirect(self):
self.RunTestPartialRun(session.Session())
@test_util.run_deprecated_v1
def testPartialRunIncompleteDirect(self):
self.RunTestPartialRunIncomplete(session.Session())
@test_util.run_deprecated_v1
def testConcurrentPartialRunDirect(self):
self.RunTestConcurrentPartialRun(session.Session())
@test_util.run_deprecated_v1
def testManyPartialRunDirect(self):
self.RunTestManyPartialRun(session.Session())
@test_util.run_deprecated_v1
def testRunAndPartialRunDirect(self):
self.RunTestRunAndPartialRun(session.Session())
@test_util.run_deprecated_v1
def testPartialRunMissingPlaceholderFeedExceptionDirect(self):
self.RunTestPartialRunMissingPlaceholderFeedException(session.Session())
@test_util.run_deprecated_v1
def testPartialRunUnspecifiedFeedDirect(self):
self.RunTestPartialRunUnspecifiedFeed(session.Session())
@test_util.run_deprecated_v1
def testPartialRunUnspecifiedFetchDirect(self):
self.RunTestPartialRunUnspecifiedFetch(session.Session())
@test_util.run_deprecated_v1
def testPartialRunAlreadyFedDirect(self):
self.RunTestPartialRunAlreadyFed(session.Session())
@test_util.run_deprecated_v1
def testPartialRunAlreadyFetchedDirect(self):
self.RunTestPartialRunAlreadyFetched(session.Session())
@test_util.run_deprecated_v1
def testPartialRunEmptyFetchesDirect(self):
self.RunTestPartialRunEmptyFetches(session.Session())
@test_util.run_deprecated_v1
def testPartialRunDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRun(session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunIncompleteDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRunIncomplete(session.Session(server.target))
@test_util.run_deprecated_v1
def testConcurrentPartialRunDist(self):
server = server_lib.Server.create_local_server()
self.RunTestConcurrentPartialRun(session.Session(server.target))
@test_util.run_deprecated_v1
def testManyPartialRunDist(self):
server = server_lib.Server.create_local_server()
self.RunTestManyPartialRun(session.Session(server.target))
@test_util.run_deprecated_v1
def testRunAndPartialRunDist(self):
server = server_lib.Server.create_local_server()
self.RunTestRunAndPartialRun(session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunMissingPlaceholderFeedExceptionDist(self):
self.skipTest('Flaky test. Short term b/278768411, long term b/280102873')
server = server_lib.Server.create_local_server()
self.RunTestPartialRunMissingPlaceholderFeedException(
session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunUnspecifiedFeedDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRunUnspecifiedFeed(session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunUnspecifiedFetchDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRunUnspecifiedFetch(session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunAlreadyFedDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRunAlreadyFed(session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunAlreadyFetchedDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRunAlreadyFetched(session.Session(server.target))
@test_util.run_deprecated_v1
def testPartialRunEmptyFetchesDist(self):
server = server_lib.Server.create_local_server()
self.RunTestPartialRunEmptyFetches(session.Session(server.target))
if __name__ == '__main__':
googletest.main()
tensorflow/python/client/session_ref.cc
+541
View File
@@ -0,0 +1,541 @@
/* 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/client/session_ref.h"
#include <stdlib.h>
#include <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/log/log.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/strings/str_format.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/lib/strings/stringprintf.h"
#include "tensorflow/core/protobuf/master.pb.h"
#include "tensorflow/core/protobuf/named_tensor.pb.h"
#include "tensorflow/core/protobuf/replay_log.pb.h"
namespace tensorflow {
namespace {
// Scope helper to track active calls and manage session lifetime.
// SessionRef blocks closing until all active calls complete or are cancelled.
struct RunCounter {
std::shared_ptr<Session> session;
uint64_t* value;
mutex* m;
condition_variable* cv;
explicit RunCounter(const std::shared_ptr<Session>& s, uint64_t* v, mutex* m,
condition_variable* cv)
: value(v), m(m), cv(cv) {
mutex_lock l(*m);
if (s != nullptr) {
session = s;
++*value;
}
}
~RunCounter() {
if (session != nullptr) {
mutex_lock l(*m);
if (--*value == 0) {
cv->notify_all();
}
}
}
};
std::string SessionToHandle(Session* session) {
return absl::StrFormat("%llu", static_cast<unsigned long long>(
reinterpret_cast<uintptr_t>(session)));
}
// The Session interface has many methods of the form:
//
// X(a, b);
// X(RunOptions, a, b);
//
// Not all sessions support the second case (with an empty RunOptions()).
// We use this variable as a sentinel to dispatch to the correct call.
RunOptions* kEmptyRunOptions() {
static RunOptions* options = new RunOptions();
return options;
}
} // namespace
// Run the given session operation, recording start and end timestamps.
// If the operation returns a bad status, return after flushing the current
// log request. This should be run _after_ all request information has been
// added to the current op.
#define RUN_WITH_TIMESTAMP(OpName, ...) \
op.set_start_time_us(Env::Default()->NowMicros()); \
Status status = session->OpName(__VA_ARGS__); \
op.set_end_time_us(Env::Default()->NowMicros()); \
if (!status.ok()) { \
Flush(op).IgnoreError(); \
return status; \
}
// Records requests (and optionally responses) performed against a session.
// The resulting replay log can be used with the `tf_replay` tool to replicate
// the operations against a simulated environment, without requiring the
// original code or cluster setup.
//
// Session logging by setting the TF_REPLAY_LOG_FILE environment variable.
class SessionLogger {
public:
SessionLogger() {
const char* log_file_env = getenv("TF_REPLAY_LOG_FILE");
std::string log_name = log_file_env ? std::string(log_file_env) : ".";
LOG(INFO) << "Constructing new session logger for " << log_name;
TF_CHECK_OK(Env::Default()->RecursivelyCreateDir(
std::string(io::Dirname(log_name))));
Env::Default()->DeleteFile(log_name).IgnoreError();
TF_CHECK_OK(Env::Default()->NewWritableFile(log_name, &log_file_));
log_writer_ = absl::make_unique<io::RecordWriter>(log_file_.get());
}
~SessionLogger() {
log_writer_->Close().IgnoreError();
log_writer_.release();
log_file_->Close().IgnoreError();
}
absl::Status RecordNewSession(Session* session) {
ReplayOp op;
NewReplaySession* req = op.mutable_new_replay_session();
req->set_session_handle(SessionToHandle(session));
return Flush(op);
}
absl::Status RecordRun(
Session* session,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_node_names,
std::vector<Tensor>* outputs) {
return RecordRun(session, *kEmptyRunOptions(), inputs, output_tensor_names,
target_node_names, outputs, nullptr);
}
absl::Status RecordRun(
Session* session, const RunOptions& run_options,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_node_names,
std::vector<Tensor>* outputs, RunMetadata* run_metadata) {
ReplayOp op;
RunStepRequest* req = op.mutable_run_step();
RunStepResponse* resp = op.mutable_run_step_response();
req->set_session_handle(SessionToHandle(session));
*req->mutable_options() = run_options;
for (const auto& it : inputs) {
NamedTensorProto* feed = req->add_feed();
feed->set_name(it.first);
it.second.AsProtoField(feed->mutable_tensor());
}
// Build an index from fetch tensor name to first index in
// output_tensor_names.
std::unordered_map<std::string, int> output_name_to_offset;
for (int i = 0, end = output_tensor_names.size(); i < end; ++i) {
const std::string& name = output_tensor_names[i];
if (output_name_to_offset.insert(std::make_pair(name, i)).second) {
req->add_fetch(name);
}
}
for (const std::string& target : target_node_names) {
req->add_target(target);
}
if (&run_options == kEmptyRunOptions()) {
RUN_WITH_TIMESTAMP(Run, inputs, output_tensor_names, target_node_names,
outputs);
} else {
RUN_WITH_TIMESTAMP(Run, run_options, inputs, output_tensor_names,
target_node_names, outputs, run_metadata);
}
for (size_t i = 0; i < outputs->size(); ++i) {
const Tensor& tensor = (*outputs)[i];
NamedTensorProto* tproto = resp->add_tensor();
tensor.AsProtoField(tproto->mutable_tensor());
tproto->set_name(output_tensor_names[i]);
}
if (run_metadata) {
*resp->mutable_metadata() = *run_metadata;
}
return Flush(op);
}
absl::Status RecordCreate(Session* session, const GraphDef& graph) {
return RecordCreate(session, *kEmptyRunOptions(), graph);
}
// N.B. RunOptions is not stored (it has no entry in CreateRequest)
absl::Status RecordCreate(Session* session, const RunOptions& run_options,
const GraphDef& graph) {
ReplayOp op;
CreateSessionRequest* req = op.mutable_create_session();
*req->mutable_graph_def() = graph;
CreateSessionResponse* resp = op.mutable_create_session_response();
if (&run_options == kEmptyRunOptions()) {
RUN_WITH_TIMESTAMP(Create, graph);
} else {
RUN_WITH_TIMESTAMP(Create, run_options, graph);
}
resp->set_session_handle(SessionToHandle(session));
return Flush(op);
}
absl::Status RecordExtend(Session* session, const GraphDef& graph) {
return RecordExtend(session, *kEmptyRunOptions(), graph);
}
// N.B. RunOptions is not stored (it has no entry in ExtendRequest)
absl::Status RecordExtend(Session* session, const RunOptions& run_options,
const GraphDef& graph) {
ReplayOp op;
ExtendSessionRequest* req = op.mutable_extend_session();
op.mutable_extend_session_response();
req->set_session_handle(SessionToHandle(session));
*req->mutable_graph_def() = graph;
if (&run_options == kEmptyRunOptions()) {
RUN_WITH_TIMESTAMP(Extend, graph);
} else {
RUN_WITH_TIMESTAMP(Extend, run_options, graph);
}
return Flush(op);
}
absl::Status RecordClose(Session* session) {
return RecordClose(session, *kEmptyRunOptions());
}
// N.B. RunOptions is not stored (it has no entry in CloseRequest)
absl::Status RecordClose(Session* session, const RunOptions& run_options) {
ReplayOp op;
CloseSessionRequest* req = op.mutable_close_session();
req->set_session_handle(SessionToHandle(session));
op.mutable_close_session_response();
if (&run_options == kEmptyRunOptions()) {
RUN_WITH_TIMESTAMP(Close);
} else {
RUN_WITH_TIMESTAMP(Close, run_options);
}
return Flush(op);
}
absl::Status RecordListDevices(Session* session,
std::vector<DeviceAttributes>* response) {
ReplayOp op;
ListDevicesRequest* req = op.mutable_list_devices();
ListDevicesResponse* resp = op.mutable_list_devices_response();
req->set_session_handle(SessionToHandle(session));
RUN_WITH_TIMESTAMP(ListDevices, response);
// TODO(power) -- local vs remote device distinction is lost here!
*resp->mutable_local_device() = {response->begin(), response->end()};
return Flush(op);
}
absl::Status RecordPRunSetup(Session* session,
const std::vector<std::string>& input_names,
const std::vector<std::string>& output_names,
const std::vector<std::string>& target_nodes,
std::string* handle) {
ReplayOp op;
PartialRunSetupRequest* req = op.mutable_partial_run_setup();
req->set_session_handle(SessionToHandle(session));
for (auto& input : input_names) {
req->add_feed(input);
}
for (auto& output : output_names) {
req->add_fetch(output);
}
for (auto& target : target_nodes) {
req->add_target(target);
}
RUN_WITH_TIMESTAMP(PRunSetup, input_names, output_names, target_nodes,
handle);
op.mutable_partial_run_setup_response()->set_partial_run_handle(*handle);
return Flush(op);
}
absl::Status RecordPRun(
Session* session, const std::string& handle,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_names,
std::vector<Tensor>* outputs) {
ReplayOp op;
RunStepRequest* req = op.mutable_run_step();
RunStepResponse* resp = op.mutable_run_step_response();
req->set_session_handle(SessionToHandle(session));
// Mark this step as a partial run for replay.
req->set_partial_run_handle(handle);
for (auto& input : inputs) {
auto* feed = req->add_feed();
feed->set_name(input.first);
input.second.AsProtoField(feed->mutable_tensor());
}
for (auto& output : output_names) {
req->add_fetch(output);
}
RUN_WITH_TIMESTAMP(PRun, handle, inputs, output_names, outputs);
for (size_t i = 0; i < outputs->size(); ++i) {
const Tensor& tensor = (*outputs)[i];
NamedTensorProto* tproto = resp->add_tensor();
tensor.AsProtoField(tproto->mutable_tensor());
tproto->set_name(output_names[i]);
}
return Flush(op);
}
absl::Status RecordMakeCallable(Session* session,
const CallableOptions& callable_options,
Session::CallableHandle* handle) {
ReplayOp op;
MakeCallableRequest* req = op.mutable_make_callable();
req->set_session_handle(SessionToHandle(session));
*req->mutable_options() = callable_options;
RUN_WITH_TIMESTAMP(MakeCallable, callable_options, handle);
MakeCallableResponse* resp = op.mutable_make_callable_response();
resp->set_handle(*handle);
return Flush(op);
}
absl::Status RecordRunCallable(Session* session,
Session::CallableHandle handle,
const std::vector<Tensor>& feed_tensors,
std::vector<Tensor>* fetch_tensors,
RunMetadata* run_metadata) {
ReplayOp op;
RunCallableRequest* req = op.mutable_run_callable();
req->set_session_handle(SessionToHandle(session));
req->set_handle(handle);
for (auto& tensor : feed_tensors) {
tensor.AsProtoField(req->add_feed());
}
RUN_WITH_TIMESTAMP(RunCallable, handle, feed_tensors, fetch_tensors,
run_metadata);
RunCallableResponse* resp = op.mutable_run_callable_response();
if (run_metadata) {
*resp->mutable_metadata() = *run_metadata;
}
for (const Tensor& tensor : *fetch_tensors) {
tensor.AsProtoTensorContent(resp->add_fetch());
}
return Flush(op);
}
absl::Status RecordReleaseCallable(Session* session,
Session::CallableHandle handle) {
ReplayOp op;
ReleaseCallableRequest* req = op.mutable_release_callable();
req->set_session_handle(SessionToHandle(session));
req->set_handle(handle);
RUN_WITH_TIMESTAMP(ReleaseCallable, handle);
return Flush(op);
}
private:
absl::Status Flush(const ReplayOp& op) {
mutex_lock l(log_mutex_);
std::string buf;
op.SerializeToString(&buf);
TF_RETURN_IF_ERROR(log_writer_->WriteRecord(buf));
// TODO(b/116624106): Not all file-systems respect calls to `Sync()`
return log_file_->Sync();
}
std::unique_ptr<WritableFile> log_file_;
std::unique_ptr<io::RecordWriter> log_writer_;
mutex log_mutex_;
};
static SessionLogger* global_session_logger() {
static SessionLogger* logger = new SessionLogger();
return logger;
}
SessionRef::SessionRef(Session* session) : session_(session) {
if (getenv("TF_REPLAY_LOG_FILE") != nullptr) {
logger_ = global_session_logger();
logger_->RecordNewSession(this->session_.get()).IgnoreError();
} else {
logger_ = nullptr;
}
}
SessionRef::~SessionRef() = default;
// If logging is active, log the start and end time of the operation along with
// the request and response.
#define LOG_AND_RUN_OPERATION(OpName, ...) \
RunCounter rc(session_, &run_count_, &run_lock_, &run_finished_); \
if (rc.session == nullptr) { \
return absl::CancelledError("Session has been closed."); \
} \
if (!logger_) { \
return rc.session->OpName(__VA_ARGS__); \
} \
return logger_->Record##OpName(rc.session.get(), __VA_ARGS__);
absl::Status SessionRef::Run(
const RunOptions& run_options,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_node_names,
std::vector<Tensor>* outputs, RunMetadata* run_metadata) {
LOG_AND_RUN_OPERATION(Run, run_options, inputs, output_tensor_names,
target_node_names, outputs, run_metadata);
}
absl::Status SessionRef::Run(
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_node_names,
std::vector<Tensor>* outputs) {
LOG_AND_RUN_OPERATION(Run, inputs, output_tensor_names, target_node_names,
outputs);
}
absl::Status SessionRef::Create(const GraphDef& graph) {
LOG_AND_RUN_OPERATION(Create, graph);
}
absl::Status SessionRef::Create(const RunOptions& run_options,
const GraphDef& graph) {
LOG_AND_RUN_OPERATION(Create, run_options, graph);
}
absl::Status SessionRef::Extend(const RunOptions& run_options,
const GraphDef& graph) {
LOG_AND_RUN_OPERATION(Extend, run_options, graph);
}
absl::Status SessionRef::Extend(const GraphDef& graph) {
LOG_AND_RUN_OPERATION(Extend, graph);
}
absl::Status SessionRef::ListDevices(std::vector<DeviceAttributes>* response) {
LOG_AND_RUN_OPERATION(ListDevices, response);
}
absl::Status SessionRef::PRunSetup(const std::vector<std::string>& input_names,
const std::vector<std::string>& output_names,
const std::vector<std::string>& target_nodes,
std::string* handle) {
LOG_AND_RUN_OPERATION(PRunSetup, input_names, output_names, target_nodes,
handle);
}
absl::Status SessionRef::PRun(
const std::string& handle,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_names,
std::vector<Tensor>* outputs) {
LOG_AND_RUN_OPERATION(PRun, handle, inputs, output_names, outputs);
}
absl::Status SessionRef::MakeCallable(const CallableOptions& callable_options,
CallableHandle* out_handle) {
LOG_AND_RUN_OPERATION(MakeCallable, callable_options, out_handle);
}
absl::Status SessionRef::RunCallable(CallableHandle handle,
const std::vector<Tensor>& feed_tensors,
std::vector<Tensor>* fetch_tensors,
RunMetadata* run_metadata) {
LOG_AND_RUN_OPERATION(RunCallable, handle, feed_tensors, fetch_tensors,
run_metadata);
}
absl::Status SessionRef::ReleaseCallable(CallableHandle handle) {
{
mutex_lock l(run_lock_);
if (session_ == nullptr) {
// Session already closed. Do nothing.
return absl::OkStatus();
}
}
LOG_AND_RUN_OPERATION(ReleaseCallable, handle);
}
absl::Status SessionRef::Close(const RunOptions& run_options) {
mutex_lock l(run_lock_);
if (session_ == nullptr) {
return absl::CancelledError("Session has been closed.");
}
absl::Status status;
if (logger_) {
status = logger_->RecordClose(session_.get(), run_options);
} else {
status = session_->Close(run_options);
}
session_.reset();
while (run_count_ > 0) {
run_finished_.wait(l);
}
return status;
}
absl::Status SessionRef::Close() {
mutex_lock l(run_lock_);
if (session_ == nullptr) {
return absl::CancelledError("Session has been closed.");
}
absl::Status status;
if (logger_) {
status = logger_->RecordClose(session_.get());
} else {
status = session_->Close();
}
session_.reset();
while (run_count_ > 0) {
run_finished_.wait(l);
}
return status;
}
} // namespace tensorflow
tensorflow/python/client/session_ref.h
+96
View File
@@ -0,0 +1,96 @@
/* 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_CLIENT_SESSION_REF_H_
#define TENSORFLOW_PYTHON_CLIENT_SESSION_REF_H_
#include <memory>
#include "absl/status/status.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/session.h"
namespace tensorflow {
class SessionLogger;
// A `SessionRef` manages the lifetime of a wrapped `Session` pointer.
//
// SessionRef blocks the return of Close() until all pending operations have
// been completed or cancelled and underlying session has been freed. Any
// subsequent operations on the SessionRef object will return errors::Cancelled.
class SessionRef : public Session {
public:
explicit SessionRef(Session* session);
~SessionRef() override;
absl::Status Create(const GraphDef& graph) override;
absl::Status Extend(const GraphDef& graph) override;
absl::Status Create(const RunOptions& run_options,
const GraphDef& graph) override;
absl::Status Extend(const RunOptions& run_options,
const GraphDef& graph) override;
absl::Status Run(const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_node_names,
std::vector<Tensor>* outputs) override;
absl::Status ListDevices(std::vector<DeviceAttributes>* response) override;
absl::Status Close() override;
absl::Status Close(const RunOptions& run_options) override;
absl::Status Run(const RunOptions& run_options,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_node_names,
std::vector<Tensor>* outputs,
RunMetadata* run_metadata) override;
absl::Status PRunSetup(const std::vector<std::string>& input_names,
const std::vector<std::string>& output_names,
const std::vector<std::string>& target_nodes,
std::string* handle) override;
absl::Status PRun(const std::string& handle,
const std::vector<std::pair<std::string, Tensor> >& inputs,
const std::vector<std::string>& output_names,
std::vector<Tensor>* outputs) override;
absl::Status MakeCallable(const CallableOptions& callable_options,
CallableHandle* out_handle) override;
absl::Status RunCallable(CallableHandle handle,
const std::vector<Tensor>& feed_tensors,
std::vector<Tensor>* fetch_tensors,
RunMetadata* run_metadata) override;
absl::Status ReleaseCallable(CallableHandle handle) override;
private:
mutex run_lock_;
condition_variable run_finished_;
uint64_t run_count_ TF_GUARDED_BY(run_lock_) = {0};
std::shared_ptr<Session> session_;
// Borrowed reference to global session logger.
SessionLogger* logger_;
};
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_CLIENT_SESSION_REF_H_
tensorflow/python/client/session_test.py
+2117
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/client/test_construction_fails_op.cc
+38
View File
@@ -0,0 +1,38 @@
/* 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 "absl/status/status.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
REGISTER_OP("ConstructionFails");
class ConstructionFailsOp : public OpKernel {
public:
explicit ConstructionFailsOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES(ctx, false,
absl::InvalidArgumentError("Failure during construction."));
}
void Compute(OpKernelContext* ctx) override {}
};
REGISTER_KERNEL_BUILDER(Name("ConstructionFails").Device(DEVICE_CPU),
ConstructionFailsOp);
} // end namespace tensorflow
tensorflow/python/client/tf_session_helper.cc
+734
View File
@@ -0,0 +1,734 @@
/* 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.
==============================================================================*/
#include "tensorflow/python/client/tf_session_helper.h"
#include <cstdint>
#include <cstring>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/c_api_internal.h"
#include "tensorflow/c/safe_ptr.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_buffer_internal.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/attr_value_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/tensor_id.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringprintf.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/equal_graph_def.h"
#include "tensorflow/python/client/session_ref.h"
#include "tensorflow/python/lib/core/ndarray_tensor.h"
#include "tensorflow/python/lib/core/ndarray_tensor_bridge.h"
#include "tensorflow/python/lib/core/safe_pyobject_ptr.h"
namespace tensorflow {
namespace {
static const char* kFeedDictErrorMsg =
"feed_dict must be a dictionary mapping strings to NumPy arrays.";
} // end namespace
TF_Session* TF_NewSessionRef(TF_Graph* graph, const TF_SessionOptions* opts,
TF_Status* status) {
TF_Session* tf_session = TF_NewSession(graph, opts, status);
if (tf_session == nullptr) {
return nullptr;
}
Session* session = reinterpret_cast<Session*>(tf_session->session);
SessionRef* session_ref = new SessionRef(session);
tf_session->session = session_ref;
return tf_session;
}
void TF_Run_wrapper_helper(TF_DeprecatedSession* session, const char* handle,
const TF_Buffer* run_options, PyObject* feed_dict,
const NameVector& output_names,
const NameVector& target_nodes,
TF_Status* out_status, PyObjectVector* out_values,
TF_Buffer* run_outputs) {
// 1. Convert the feed inputs to the appropriate form for TF_Run.
if (!PyDict_Check(feed_dict)) {
tsl::Set_TF_Status_from_Status(
out_status, absl::InvalidArgumentError(kFeedDictErrorMsg));
return;
}
NameVector input_names;
std::vector<Safe_TF_TensorPtr> inputs_safe; // Used to delete tensors.
TF_TensorVector inputs_unsafe; // Used to contain the arg to TF_Run.
PyObject* key;
PyObject* value;
Py_ssize_t pos = 0;
int index = 0;
absl::Status s;
while (PyDict_Next(feed_dict, &pos, &key, &value)) {
char* key_string = PyBytes_AsString(key);
if (!key_string) {
tsl::Set_TF_Status_from_Status(
out_status, absl::InvalidArgumentError(kFeedDictErrorMsg));
return;
}
input_names.push_back(key_string);
inputs_safe.emplace_back(make_safe(static_cast<TF_Tensor*>(nullptr)));
s = NdarrayToTensor(nullptr /*ctx*/, value, &inputs_safe.back());
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
inputs_unsafe.push_back(inputs_safe.back().get());
++index;
}
// 2. Allocate a container for the output data.
TF_TensorVector outputs(output_names.size());
// In case any tensors were leftover from previous runs we might as well clear
// them here.
ClearDecrefCache();
// 3. Actually call TF_Run().
Py_BEGIN_ALLOW_THREADS;
if (handle == nullptr) {
TF_Run(session, run_options, input_names.data(), inputs_unsafe.data(),
input_names.size(), const_cast<const char**>(output_names.data()),
outputs.data(), output_names.size(),
const_cast<const char**>(target_nodes.data()), target_nodes.size(),
run_outputs, out_status);
} else {
TF_PRun(session, handle, input_names.data(), inputs_unsafe.data(),
input_names.size(), const_cast<const char**>(output_names.data()),
outputs.data(), output_names.size(),
const_cast<const char**>(target_nodes.data()), target_nodes.size(),
out_status);
}
Py_END_ALLOW_THREADS;
// Decref any numpy arrays we are not using anymore.
ClearDecrefCache();
if (TF_GetCode(out_status) != TF_OK) {
return;
}
// 4. We now own the fetched tensors, so set up a safe container to
// delete them when we exit this scope.
std::vector<Safe_TF_TensorPtr> tf_outputs_safe;
for (const auto& output : outputs) {
tf_outputs_safe.emplace_back(make_safe(output));
}
// 5. Convert the fetched tensors into numpy ndarrays. Store them in a safe
// container so that we do not leak
std::vector<Safe_PyObjectPtr> py_outputs_safe;
for (size_t i = 0; i < output_names.size(); ++i) {
PyObject* py_array;
s = TF_TensorToPyArray(std::move(tf_outputs_safe[i]), &py_array);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
py_outputs_safe.emplace_back(
make_safe(PyArray_Return(reinterpret_cast<PyArrayObject*>(py_array))));
}
// 6. If we reach this point, we have successfully built a list of objects
// so we can release them from the safe container.
for (auto& output : py_outputs_safe) {
out_values->push_back(output.release());
}
}
// Wrapper for TF_Run that converts the arguments to appropriate types.
// If *out_status is OK, the caller becomes the owner of the PyObjects
// in *out_values.
void TF_Run_wrapper(TF_DeprecatedSession* session, const TF_Buffer* run_options,
PyObject* feed_dict, const NameVector& output_names,
const NameVector& target_nodes, TF_Status* out_status,
PyObjectVector* out_values, TF_Buffer* run_outputs) {
TF_Run_wrapper_helper(session, nullptr, run_options, feed_dict, output_names,
target_nodes, out_status, out_values, run_outputs);
ClearDecrefCache();
}
namespace {
void MakeCallableHelper(tensorflow::Session* session,
const TF_Buffer* callable_options, int64_t* out_handle,
TF_Status* out_status) {
tensorflow::CallableOptions callable_options_proto;
if (callable_options != nullptr &&
!callable_options_proto.ParseFromArray(callable_options->data,
callable_options->length)) {
tsl::Set_TF_Status_from_Status(
out_status,
absl::InvalidArgumentError("Unparsable CallableOptions proto"));
return;
}
tensorflow::Session::CallableHandle handle;
absl::Status s = session->MakeCallable(callable_options_proto, &handle);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
*out_handle = handle;
}
} // namespace
void TF_DeprecatedSessionMakeCallable(TF_DeprecatedSession* session,
const TF_Buffer* callable_options,
int64_t* out_handle, TF_Status* status) {
MakeCallableHelper(session->session, callable_options, out_handle, status);
}
void TF_SessionMakeCallable(TF_Session* session,
const TF_Buffer* callable_options,
int64_t* out_handle, TF_Status* status) {
MakeCallableHelper(session->session, callable_options, out_handle, status);
}
namespace {
void RunCallableHelper(tensorflow::Session* session, int64_t handle,
PyObject* feed_values, TF_Status* out_status,
PyObjectVector* out_values, TF_Buffer* run_metadata) {
// Convert feed values to a vector of tensorflow::Tensor objects.
std::vector<Tensor> input_tensors;
absl::Status s;
{
feed_values =
PySequence_Fast(feed_values, "feed_values must be a sequence");
if (feed_values == nullptr) return;
Safe_PyObjectPtr feed_values_holder(make_safe(feed_values));
Py_ssize_t len = PySequence_Fast_GET_SIZE(feed_values);
input_tensors.reserve(len);
for (Py_ssize_t i = 0; i < len; ++i) {
PyObject* elem = PySequence_Fast_GET_ITEM(feed_values, i);
if (!elem) {
tsl::Set_TF_Status_from_Status(
out_status,
absl::InternalError(absl::StrCat("Could not get feed value ", i)));
return;
}
Tensor t;
s = NdarrayToTensor(elem, &t);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
input_tensors.push_back(std::move(t));
}
}
RunMetadata run_metadata_proto;
// Run the callable.
std::vector<Tensor> output_tensors;
Py_BEGIN_ALLOW_THREADS;
s = session->RunCallable(handle, input_tensors, &output_tensors,
&run_metadata_proto);
Py_END_ALLOW_THREADS;
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
// If requested, serialize the RunMetadata to pass it back to the caller.
if (run_metadata != nullptr) {
s = MessageToBuffer(run_metadata_proto, run_metadata);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
}
// Convert results to NumPy arrays. Since this can fail, stage the
// results via a safe container that takes care of decreasing the
// reference count on failure.
std::vector<Safe_PyObjectPtr> py_outputs_safe;
py_outputs_safe.reserve(output_tensors.size());
for (const Tensor& output : output_tensors) {
PyObject* py_array;
s = TensorToNdarray(output, &py_array);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
py_outputs_safe.push_back(
make_safe(PyArray_Return(reinterpret_cast<PyArrayObject*>(py_array))));
}
// If we reach this point, we have successfully built a list of objects
// so we can release them from the safe container.
out_values->reserve(py_outputs_safe.size());
for (auto& output : py_outputs_safe) {
out_values->push_back(output.release());
}
}
} // namespace
void TF_DeprecatedSessionRunCallable(TF_DeprecatedSession* session,
int64_t handle, PyObject* feed_values,
PyObjectVector* out_values,
TF_Buffer* run_metadata,
TF_Status* status) {
RunCallableHelper(session->session, handle, feed_values, status, out_values,
run_metadata);
ClearDecrefCache();
}
void TF_SessionRunCallable(TF_Session* session, int64_t handle,
PyObject* feed_values, PyObjectVector* out_values,
TF_Buffer* run_metadata, TF_Status* status) {
RunCallableHelper(session->session, handle, feed_values, status, out_values,
run_metadata);
ClearDecrefCache();
}
void TF_DeprecatedSessionReleaseCallable(TF_DeprecatedSession* session,
int64_t handle, TF_Status* status) {
tsl::Set_TF_Status_from_Status(status,
session->session->ReleaseCallable(handle));
}
void TF_SessionReleaseCallable(TF_Session* session, int64_t handle,
TF_Status* status) {
tsl::Set_TF_Status_from_Status(status,
session->session->ReleaseCallable(handle));
}
// Wrapper for TF_PRunSetup that converts the arguments to appropriate types.
// If *out_status is OK, the caller becomes the owner of *out_handle.
void TF_PRunSetup_wrapper(TF_DeprecatedSession* session,
const NameVector& input_names,
const NameVector& output_names,
const NameVector& target_nodes, TF_Status* out_status,
const char** out_handle) {
Py_BEGIN_ALLOW_THREADS;
TF_PRunSetup(
session, const_cast<const char**>(input_names.data()), input_names.size(),
const_cast<const char**>(output_names.data()), output_names.size(),
const_cast<const char**>(target_nodes.data()), target_nodes.size(),
out_handle, out_status);
Py_END_ALLOW_THREADS;
}
// Wrapper for TF_PRun that converts the arguments to appropriate types.
// If *out_status is OK, the caller becomes the owner of the PyObjects
// in *out_values.
void TF_PRun_wrapper(TF_DeprecatedSession* session, const char* handle,
PyObject* feed_dict, const NameVector& output_names,
TF_Status* out_status, PyObjectVector* out_values) {
TF_Run_wrapper_helper(session, handle, nullptr, feed_dict, output_names,
NameVector(), out_status, out_values, nullptr);
ClearDecrefCache();
}
// Wrapper for TF_Reset that converts the string vectors to character arrays.
void TF_Reset_wrapper(const TF_SessionOptions* opt,
const NameVector& containers, TF_Status* status) {
TF_Reset(opt, const_cast<const char**>(containers.data()), containers.size(),
status);
}
void TF_SessionRun_wrapper_helper(TF_Session* session, const char* handle,
const TF_Buffer* run_options,
const std::vector<TF_Output>& inputs,
const std::vector<PyObject*>& input_ndarrays,
const std::vector<TF_Output>& outputs,
const std::vector<TF_Operation*>& targets,
TF_Buffer* run_metadata,
TF_Status* out_status,
std::vector<PyObject*>* py_outputs) {
DCHECK_EQ(inputs.size(), input_ndarrays.size());
DCHECK(py_outputs != nullptr);
DCHECK(py_outputs->empty());
absl::Status s;
// Convert input ndarray PyObjects to TF_Tensors. We maintain a continuous
// array of TF_Tensor*s as well as scoped containers to make sure they're
// cleaned up properly.
//
// Memory management:
// NdarrayToTensor() creates a new ndarray PyObject from the input
// ndarray. We manage the new ndarray's lifetime in order to keep the
// underlying data buffer alive (the new ndarray also guarantees a contiguous
// data buffer). The new ndarray's data buffer is used to create the
// corresponding TF_Tensor. The TF_Tensor's deallocator will queue the new
// ndarray to be decref'd by the next ClearDecrefCache() call (we can't call
// Py_DECREF in the deallocator directly because the GIL must be held).
//
// Note that TF_Tensor may directly delegate its data and deallocator to a
// TensorBuffer, which may outlive the TF_Tensor (e.g. if the tensor gets
// queued or assigned to a variable).
TF_TensorVector input_vals;
std::vector<Safe_TF_TensorPtr> input_vals_safe;
for (PyObject* ndarray : input_ndarrays) {
input_vals_safe.emplace_back(make_safe(static_cast<TF_Tensor*>(nullptr)));
s = NdarrayToTensor(nullptr, ndarray, &input_vals_safe.back());
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
input_vals.push_back(input_vals_safe.back().get());
}
// Allocate space for output TF_Tensor*s
TF_TensorVector output_vals(outputs.size());
// Clear up any unused memory leftover from previous runs
ClearDecrefCache();
// Call TF_SessionRun() (and release GIL during execution)
Py_BEGIN_ALLOW_THREADS;
if (handle == nullptr) {
TF_SessionRun(session, run_options, inputs.data(), input_vals.data(),
inputs.size(), outputs.data(), output_vals.data(),
outputs.size(), targets.data(), targets.size(), run_metadata,
out_status);
} else {
TF_SessionPRun(session, handle, inputs.data(), input_vals.data(),
inputs.size(), outputs.data(), output_vals.data(),
outputs.size(), targets.data(), targets.size(), out_status);
}
Py_END_ALLOW_THREADS;
// Create scoped containers for output tensors
std::vector<Safe_TF_TensorPtr> output_vals_safe;
for (TF_Tensor* output : output_vals) {
output_vals_safe.emplace_back(make_safe(output));
}
// Convert outputs to ndarrays (in scoped containers)
std::vector<Safe_PyObjectPtr> py_outputs_safe;
for (size_t i = 0; i < outputs.size(); ++i) {
PyObject* py_array;
s = TF_TensorToPyArray(std::move(output_vals_safe[i]), &py_array);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(out_status, s);
return;
}
py_outputs_safe.emplace_back(
make_safe(PyArray_Return(reinterpret_cast<PyArrayObject*>(py_array))));
}
// If we reach this point, we have successfully built a list of objects so we
// can release them from the safe container into the return vector.
for (size_t i = 0; i < outputs.size(); ++i) {
py_outputs->push_back(py_outputs_safe[i].release());
}
}
void TF_SessionRun_wrapper(TF_Session* session, const TF_Buffer* run_options,
const std::vector<TF_Output>& inputs,
const std::vector<PyObject*>& input_ndarrays,
const std::vector<TF_Output>& outputs,
const std::vector<TF_Operation*>& targets,
TF_Buffer* run_metadata, TF_Status* out_status,
std::vector<PyObject*>* py_outputs) {
TF_SessionRun_wrapper_helper(session, nullptr, run_options, inputs,
input_ndarrays, outputs, targets, run_metadata,
out_status, py_outputs);
// Release any unused ndarray references (see memory management comment in
// TF_SessionRun_wrapper_helper)
ClearDecrefCache();
}
std::string EqualGraphDefWrapper(const std::string& actual,
const std::string& expected) {
GraphDef actual_def;
if (!actual_def.ParseFromString(actual)) {
return "actual is not a valid serialized GraphDef";
}
GraphDef expected_def;
if (!expected_def.ParseFromString(expected)) {
return "expected is not a valid serialized GraphDef";
}
std::string diff;
return EqualGraphDef(actual_def, expected_def, &diff) ? "" : diff;
}
std::string EqualAttrValueWrapper(const std::string& actual,
const std::string& expected) {
AttrValue actual_attr_value;
if (!actual_attr_value.ParseFromString(actual)) {
return "actual is not a valid serialized AttrValue";
}
AttrValue expected_attr_value;
if (!expected_attr_value.ParseFromString(expected)) {
return "expected is not a valid serialized AttrValue";
}
std::string diff;
if (!AreAttrValuesEqual(actual_attr_value, expected_attr_value)) {
diff = absl::StrFormat(
"Actual AttrValue %s does not match Expected AttrValue %s.",
SummarizeAttrValue(actual_attr_value).c_str(),
SummarizeAttrValue(expected_attr_value).c_str());
}
return diff;
}
// Return value set to 6 inlined elements so it fits in a 64-byte cache line.
absl::InlinedVector<int64_t, 6UL> TF_GraphGetTensorShapeHelper(
TF_Graph* graph, TF_Output output, TF_Status* out_status,
bool* unknown_shape) {
// Allocate a single variable for holding the result for RVO.
absl::InlinedVector<int64_t, 6UL> result;
*unknown_shape = false;
int num_dims = TF_GraphGetTensorNumDims(graph, output, out_status);
if (TF_GetCode(out_status) != TF_OK) {
return result;
}
// If shape is unknown, set boolean and return.
if (num_dims == -1) {
*unknown_shape = true;
return result;
}
// If shape is a scalar, avoid another C call and just return {}.
if (num_dims == 0) {
return result;
}
result.resize(num_dims);
TF_GraphGetTensorShape(graph, output, result.data(), num_dims, out_status);
return result;
}
void TF_SessionPRunSetup_wrapper(TF_Session* session,
const std::vector<TF_Output>& inputs,
const std::vector<TF_Output>& outputs,
const std::vector<TF_Operation*>& targets,
const char** out_handle,
TF_Status* out_status) {
// Call TF_SessionPRunSetup() (and release GIL during execution)
Py_BEGIN_ALLOW_THREADS;
TF_SessionPRunSetup(session, inputs.data(), inputs.size(), outputs.data(),
outputs.size(), targets.data(), targets.size(),
out_handle, out_status);
Py_END_ALLOW_THREADS;
}
void TF_SessionPRun_wrapper(TF_Session* session, const char* handle,
const std::vector<TF_Output>& inputs,
const std::vector<PyObject*>& input_ndarrays,
const std::vector<TF_Output>& outputs,
TF_Status* out_status,
std::vector<PyObject*>* py_outputs) {
const std::vector<TF_Operation*> targets;
TF_SessionRun_wrapper_helper(session, handle,
nullptr, // run_options
inputs, input_ndarrays, outputs, targets,
nullptr, // run_metadata
out_status, py_outputs);
// Release any unused ndarray references (see memory management comment in
// TF_SessionRun_wrapper_helper)
ClearDecrefCache();
}
std::vector<TF_Output> GetOperationInputs(TF_Operation* oper) {
int num_inputs = TF_OperationNumInputs(oper);
std::vector<TF_Output> inputs(num_inputs);
TF_OperationAllInputs(oper, inputs.data(), inputs.size());
return inputs;
}
std::vector<TF_Operation*> TF_OperationGetControlInputs_wrapper(
TF_Operation* oper) {
std::vector<TF_Operation*> control_inputs(TF_OperationNumControlInputs(oper));
TF_OperationGetControlInputs(oper, control_inputs.data(),
control_inputs.size());
return control_inputs;
}
std::vector<TF_Operation*> TF_OperationGetControlOutputs_wrapper(
TF_Operation* oper) {
std::vector<TF_Operation*> control_outputs(
TF_OperationNumControlOutputs(oper));
TF_OperationGetControlOutputs(oper, control_outputs.data(),
control_outputs.size());
return control_outputs;
}
std::vector<const char*> TF_OperationOutputConsumers_wrapper(
TF_Output oper_out) {
int num_consumers = TF_OperationOutputNumConsumers(oper_out);
std::vector<TF_Input> consumers(num_consumers);
TF_OperationOutputConsumers(oper_out, consumers.data(), num_consumers);
std::vector<const char*> consumer_names(num_consumers);
for (int i = 0; i < num_consumers; ++i) {
consumer_names[i] = TF_OperationName(consumers[i].oper);
}
return consumer_names;
}
TF_Function* TF_GraphToFunction_wrapper(
const TF_Graph* fn_body, const char* fn_name, bool append_hash_to_fn_name,
const std::vector<TF_Operation*>* opers,
const std::vector<TF_Output>& inputs, const std::vector<TF_Output>& outputs,
const NameVector& output_names,
const std::vector<TF_Operation*>* control_outputs,
const NameVector& control_output_names, const TF_FunctionOptions* opts,
const char* description, TF_Status* out_status) {
if (!output_names.empty() && output_names.size() != outputs.size()) {
tsl::Set_TF_Status_from_Status(
out_status,
absl::InvalidArgumentError(absl::StrCat(
"output names must be either empty or equal in size to outputs. ",
"output names size = ", output_names.size(),
" outputs size = ", outputs.size())));
return nullptr;
}
int nopers = -1;
const TF_Operation* const* opers_array = nullptr;
if (opers != nullptr) {
nopers = opers->size();
opers_array = opers->data();
}
const char** output_names_ptr =
output_names.empty() ? nullptr
: const_cast<const char**>(output_names.data());
const char** control_output_names_ptr =
control_output_names.empty()
? nullptr
: const_cast<const char**>(control_output_names.data());
return TF_GraphToFunctionWithControlOutputs(
fn_body, fn_name, append_hash_to_fn_name, nopers, opers_array,
inputs.size(), inputs.data(), outputs.size(), outputs.data(),
output_names_ptr,
control_outputs == nullptr ? 0 : control_outputs->size(),
control_outputs == nullptr ? nullptr : control_outputs->data(),
control_output_names_ptr, opts, description, out_status);
}
void TF_GraphSetOutputHandleShapesAndTypes_wrapper(
TF_Graph* graph, TF_Output output,
const std::vector<std::vector<int64_t>>& shapes,
const std::vector<int>& ranks, const std::vector<TF_DataType>& types,
TF_Status* status) {
std::vector<const int64_t*> shapes_pointers(shapes.size());
for (int i = 0; i < shapes.size(); ++i) {
shapes_pointers[i] = ranks[i] <= 0 ? nullptr : &shapes[i][0];
}
TF_GraphSetOutputHandleShapesAndTypes(graph, output, shapes.size(),
shapes_pointers.data(), ranks.data(),
types.data(), status);
}
void CreatePlaceholder(TF_Graph* graph, TF_Status* s, std::string&& name,
TF_DataType dtype, TF_Output* output) {
TF_OperationDescription* desc =
TF_NewOperation(graph, "Placeholder", name.data());
TF_SetAttrType(desc, "dtype", dtype);
TF_Operation* op = TF_FinishOperation(desc, s);
output->oper = op;
output->index = 0;
}
std::vector<TF_Output> TF_CreatePlaceholders(TF_Graph* graph, PyObject* dtypes,
const char* prefix,
TF_Status* status) {
std::vector<TF_Output> outputs;
dtypes = PySequence_Fast(dtypes, "dtypes must be a sequence");
if (dtypes == nullptr) {
tsl::Set_TF_Status_from_Status(status,
absl::InternalError("dtypes is nullptr"));
return outputs;
}
Safe_PyObjectPtr dtypes_holder(make_safe(dtypes));
Py_ssize_t len = PySequence_Fast_GET_SIZE(dtypes);
outputs.reserve(len);
for (size_t i = 0; i < len; i++) {
PyObject* dtype = PySequence_Fast_GET_ITEM(dtypes, i);
if (!dtype) {
tsl::Set_TF_Status_from_Status(
status, absl::InternalError(absl::StrCat("Could not get dtype ", i)));
return outputs;
}
#if PY_MAJOR_VERSION >= 3
TF_DataType tf_datatype = static_cast<TF_DataType>(PyLong_AsLong(dtype));
#else
TF_DataType tf_datatype = static_cast<TF_DataType>(PyInt_AsLong(dtype));
#endif
outputs.push_back(TF_Output());
CreatePlaceholder(graph, status, absl::StrCat(prefix, i), tf_datatype,
&outputs.back());
if (TF_GetCode(status) != TF_OK) break;
}
return outputs;
}
void TF_GraphSetTensorShape_wrapper(TF_Graph* graph, TF_Output output,
const std::vector<int64_t>& dims,
bool unknown_shape, TF_Status* status) {
if (unknown_shape) {
TF_GraphSetTensorShape(graph, output, nullptr, -1, status);
return;
}
TF_GraphSetTensorShape(graph, output, dims.data(), dims.size(), status);
}
std::vector<std::string>
TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper(
TF_ImportGraphDefResults* results) {
int num_missing_unused_input_mappings;
const char** src_names;
int* src_indexes;
TF_ImportGraphDefResultsMissingUnusedInputMappings(
results, &num_missing_unused_input_mappings, &src_names, &src_indexes);
std::vector<std::string> input_strs(num_missing_unused_input_mappings);
for (int i = 0; i < num_missing_unused_input_mappings; ++i) {
input_strs[i] = TensorId(src_names[i], src_indexes[i]).ToString();
}
return input_strs;
}
PyObject* TF_TryEvaluateConstant_wrapper(TF_Graph* graph, TF_Output output,
TF_Status* status) {
TF_Tensor* result_tensor;
bool evaluated =
TF_TryEvaluateConstant(graph, output, &result_tensor, status);
if (!evaluated || TF_GetCode(status) != TF_OK) Py_RETURN_NONE;
Safe_TF_TensorPtr safe_result_tensor(result_tensor);
PyObject* out;
absl::Status s = TF_TensorToPyArray(std::move(safe_result_tensor), &out);
tsl::Set_TF_Status_from_Status(status, s);
if (!s.ok()) Py_RETURN_NONE;
return PyArray_Return(reinterpret_cast<PyArrayObject*>(out));
}
} // namespace tensorflow
tensorflow/python/client/tf_session_helper.h
+260
View File
@@ -0,0 +1,260 @@
/* 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_CLIENT_TF_SESSION_HELPER_H_
#define TENSORFLOW_PYTHON_CLIENT_TF_SESSION_HELPER_H_
// Must be included first
// clang-format off
#include "xla/tsl/python/lib/core/numpy.h" //NOLINT
// clang-format on
#include "tensorflow/c/c_api.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
namespace tensorflow {
// Container types for the various arguments and temporary values used
// in the wrapper.
// A NameVector is a vector of tensor or operation names, as borrowed
// C strings.
typedef absl::InlinedVector<const char*, 8UL> NameVector;
// A PyObjectVector is a vector of borrowed pointers to PyObjects.
typedef absl::InlinedVector<PyObject*, 8UL> PyObjectVector;
// A TF_TensorVector is a vector of borrowed pointers to TF_Tensors.
typedef absl::InlinedVector<TF_Tensor*, 8UL> TF_TensorVector;
TF_Session* TF_NewSessionRef(TF_Graph* graph, const TF_SessionOptions* opts,
TF_Status* status);
// Run the graph associated with the session starting with the
// supplied inputs[]. Regardless of success or failure, inputs[] are
// stolen by the implementation (i.e. the implementation will
// eventually call Py_DECREF on each array input).
//
// The PyObject* feed_dict must be a dictionary mapping strings to
// NumPy arrays. This function does not modify its reference count.
//
// On success, the tensors corresponding to output_names[0,noutputs-1]
// are placed in out_values[], and these outputs[] become the property
// of the caller (the caller must eventually call Py_DECREF on them).
//
// On failure, out_status contains a tensorflow::Status with an error
// message.
void TF_Run_wrapper(TF_DeprecatedSession* session, const TF_Buffer* run_options,
PyObject* feed_dict, const NameVector& output_names,
const NameVector& target_nodes, TF_Status* out_status,
PyObjectVector* out_values, TF_Buffer* run_outputs);
// Python wrappers for the `Session::MakeCallable()` API.
void TF_DeprecatedSessionMakeCallable(TF_DeprecatedSession* session,
const TF_Buffer* callable_options,
int64_t* out_handle, TF_Status* status);
void TF_SessionMakeCallable(TF_Session* session,
const TF_Buffer* callable_options,
int64_t* out_handle, TF_Status* status);
// Python wrappers for the `Session::RunCallable()` API.
void TF_DeprecatedSessionRunCallable(TF_DeprecatedSession* session,
int64_t handle, PyObject* feed_values,
PyObjectVector* out_values,
TF_Buffer* run_metadata,
TF_Status* status);
void TF_SessionRunCallable(TF_Session* session, int64_t handle,
PyObject* feed_values, PyObjectVector* out_values,
TF_Buffer* run_metadata, TF_Status* status);
// Python wrappers for the `Session::ReleaseCallable()` API.
void TF_DeprecatedSessionReleaseCallable(TF_DeprecatedSession* session,
int64_t handle, TF_Status* status);
void TF_SessionReleaseCallable(TF_Session* session, int64_t handle,
TF_Status* status);
// Set up the graph with the intended feeds and fetches for partial run.
// *out_handle is owned by the caller.
//
// On success, returns a handle that is used for subsequent PRun calls.
//
// On failure, out_status contains a tensorflow::Status with an error
// message.
void TF_PRunSetup_wrapper(TF_DeprecatedSession* session,
const NameVector& input_names,
const NameVector& output_names,
const NameVector& target_nodes, TF_Status* out_status,
const char** out_handle);
// Continue to run the graph with additional feeds and fetches. The
// execution state is uniquely identified by the handle.
//
// The PyObject* feed_dict must be a dictionary mapping strings to
// NumPy arrays. This function does not modify its reference count.
//
// On success, the tensors corresponding to output_names[0,noutputs-1]
// are placed in out_values[], and these outputs[] become the property
// of the caller (the caller must eventually call Py_DECREF on them).
//
// On failure, out_status contains a tensorflow::Status with an error
// message.
void TF_PRun_wrapper(TF_DeprecatedSession* session, const char* handle,
PyObject* feed_dict, const NameVector& output_names,
TF_Status* out_status, PyObjectVector* out_values);
// Wrapper for TF_Reset that converts the string vectors to character arrays.
void TF_Reset_wrapper(const TF_SessionOptions* opt,
const NameVector& containers, TF_Status* status);
// Convenience wrapper around EqualGraphDef to make it easier to wrap.
// Returns an explanation if a difference is found, or the empty string
// for no difference.
std::string EqualGraphDefWrapper(const std::string& actual,
const std::string& expected);
// Convenience wrapper around AreAttrValuesEqual to make it easier to wrap.
// The actual and expected strings must correspond to a serialized binary
// representation of two AttrValue proto instances.
// Returns an explanation if a difference is found, or the empty string
// for no difference.
std::string EqualAttrValueWrapper(const std::string& actual,
const std::string& expected);
// Gets shape from C API Graph object.
//
// If shape is known, returns shape vector where -1 means "unknown
// dimension". Sets unknown_shape to false.
//
// If shape is unknown, sets unknown_shape to true.
absl::InlinedVector<int64_t, 6UL> TF_GraphGetTensorShapeHelper(
TF_Graph* graph, TF_Output output, TF_Status* status, bool* unknown_shape);
// Runs the graph associated with the session starting with the supplied inputs.
// On success, `py_outputs` is populated with a numpy ndarray for each output
// (the caller must decref these ndarrays, although this will likely be handled
// by the Python gc). `session`, `out_status`, and `py_outputs` must be
// non-null. `py_outputs` should be empty.
void TF_SessionRun_wrapper(TF_Session* session, const TF_Buffer* run_options,
const std::vector<TF_Output>& inputs,
const std::vector<PyObject*>& input_ndarrays,
const std::vector<TF_Output>& outputs,
const std::vector<TF_Operation*>& targets,
TF_Buffer* run_metadata, TF_Status* status,
std::vector<PyObject*>* py_outputs);
// Set up the graph with the intended feeds (inputs) and fetches (output) for
// a sequence of partial run calls.
//
// On success, returns a handle that can be used for subsequent PRun calls. The
// handle is owned by the caller and should be deleted with TF_DeletePRunHandle
// when it is no longer needed.
//
// On failure, out_status contains a tensorflow::Status with an error
// message.
void TF_SessionPRunSetup_wrapper(TF_Session* session,
const std::vector<TF_Output>& inputs,
const std::vector<TF_Output>& outputs,
const std::vector<TF_Operation*>& targets,
const char** out_handle, TF_Status* status);
// Continue to run the graph with additional feeds and fetches. The
// execution state is uniquely identified by the handle.
//
// On success, `py_outputs` is populated with a numpy ndarray for each output
// (the caller must decref these ndarrays, although this will likely be handled
// by the Python gc). `session`, `handle`, `out_status`, and `py_outputs` must
// be non-null. `py_outputs` should be empty.
//
// On failure, out_status contains a tensorflow::Status with an error
// message.
void TF_SessionPRun_wrapper(TF_Session* session, const char* handle,
const std::vector<TF_Output>& inputs,
const std::vector<PyObject*>& input_ndarrays,
const std::vector<TF_Output>& outputs,
TF_Status* status,
std::vector<PyObject*>* py_outputs);
// Retrieves the inputs of this operation.
std::vector<TF_Output> GetOperationInputs(TF_Operation* oper);
// Retrieves the control inputs of this operation.
std::vector<TF_Operation*> TF_OperationGetControlInputs_wrapper(
TF_Operation* oper);
// Retrieves the control outputs of this operation.
std::vector<TF_Operation*> TF_OperationGetControlOutputs_wrapper(
TF_Operation* oper);
// Retrieves the op names of the consumers of `oper_out`. The returned strings
// have the lifetime of the underlying TF_Graph.
std::vector<const char*> TF_OperationOutputConsumers_wrapper(
TF_Output oper_out);
// `opers` equaling NULL are converted to `nopers = -1`.
// `output_names` must be empty or have the same length as `outputs`.
TF_Function* TF_GraphToFunction_wrapper(
const TF_Graph* fn_body, const char* fn_name, bool append_hash_to_fn_name,
const std::vector<TF_Operation*>* opers,
const std::vector<TF_Output>& inputs, const std::vector<TF_Output>& outputs,
const NameVector& output_names,
const std::vector<TF_Operation*>* control_outputs,
const NameVector& control_output_names, const TF_FunctionOptions* opts,
const char* description, TF_Status* status);
// Set the shapes and types for the output's handle.
//
// The sizes of 'shapes', 'ranks', and 'types' must be equal; `shapes[i]`
// contains the shape of the handle's i-th value, `ranks[i]` contains the i-th
// shape's rank, and `types[i]` contains the i-th value's dtype. If the i-th
// shape is unknown, then `ranks[i]` must be equal to -1.
//
// The space between the double angle brackets below looks extraneous, but
// our version of SWIG cannot parse ">>".
void TF_GraphSetOutputHandleShapesAndTypes_wrapper(
TF_Graph* graph, TF_Output output,
const std::vector<std::vector<int64_t> >& shapes,
const std::vector<int>& ranks, const std::vector<TF_DataType>& types,
TF_Status* status);
// Creates Placeholders with specified types in the Graph.
//
// This is an internal API used to speed up creation of unused placeholders
// in while_v2 cond graph and is subject to change/removal.
std::vector<TF_Output> TF_CreatePlaceholders(TF_Graph* graph, PyObject* dtypes,
const char* prefix,
TF_Status* status);
// Set the shape of output. If unknown is true, `num_dims` must be set to
// -1 and `dims` is set to nullptr.
void TF_GraphSetTensorShape_wrapper(TF_Graph* graph, TF_Output output,
const std::vector<int64_t>& dims,
bool unknown_shape, TF_Status* status);
// Returns the string representations of the missing unused input mappings.
std::vector<std::string>
TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper(
TF_ImportGraphDefResults* results);
// If evaluation was possible, returns the numpy ndarray of the evaluated
// result. Otherwise returns None.
PyObject* TF_TryEvaluateConstant_wrapper(TF_Graph* graph, TF_Output output,
TF_Status* status);
} // namespace tensorflow
#endif // TENSORFLOW_PYTHON_CLIENT_TF_SESSION_HELPER_H_
tensorflow/python/client/tf_session_wrapper.cc
+2238
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/client/timeline.py
+856
View File
@@ -0,0 +1,856 @@
# 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.
# ==============================================================================
"""Timeline visualization for TensorFlow using Chrome Trace Format."""
import collections
import copy
import json
import re
from typing import Any, Dict, List, Optional, Tuple, Union
from tensorflow.core.framework import step_stats_pb2
# The timeline target is usually imported as part of BUILD target
# "platform_test", which includes also includes the "platform"
# dependency. This is why the logging import here is okay.
from tensorflow.python.platform import build_info
from tensorflow.python.platform import tf_logging as logging
class AllocationMaximum(
collections.namedtuple(
'AllocationMaximum', ('timestamp', 'num_bytes', 'tensors')
)
):
"""Stores the maximum allocation for a given allocator within the timelne.
Parameters:
timestamp: `tensorflow::Env::NowMicros()` when this maximum was reached.
num_bytes: the total memory used at this time.
tensors: the set of tensors allocated at this time.
"""
class StepStatsAnalysis(
collections.namedtuple(
'StepStatsAnalysis', ('chrome_trace', 'allocator_maximums')
)
):
"""Stores the step stats analysis output.
Parameters:
chrome_trace: A dict containing the chrome trace analysis.
allocator_maximums: A dict mapping allocator names to AllocationMaximum.
"""
class _ChromeTraceFormatter(object):
"""A helper class for generating traces in Chrome Trace Format."""
def __init__(self, show_memory: bool = False) -> None:
"""Constructs a new Chrome Trace formatter."""
self._show_memory = show_memory
self._events = []
self._metadata = []
def _create_event(
self,
ph: str,
category: str,
name: str,
pid: int,
tid: int,
timestamp: int,
) -> Dict[str, Union[str, int]]:
"""Creates a new Chrome Trace event.
For details of the file format, see:
https://chromium.googlesource.com/catapult/+/refs/heads/main/tracing/README.md
Args:
ph: The type of event - usually a single character.
category: The event category as a string.
name: The event name as a string.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
timestamp: The timestamp of this event as a long integer.
Returns:
A JSON compatible event object.
"""
event = {}
event['ph'] = ph
event['cat'] = category
event['name'] = name
event['pid'] = pid
event['tid'] = tid
event['ts'] = timestamp
return event
def emit_pid(self, name: str, pid: int) -> None:
"""Adds a process metadata event to the trace.
Args:
name: The process name as a string.
pid: Identifier of the process as an integer.
"""
event = {}
event['name'] = 'process_name'
event['ph'] = 'M'
event['pid'] = pid
event['args'] = {'name': name}
self._metadata.append(event)
def emit_tid(self, name, pid, tid):
"""Adds a thread metadata event to the trace.
Args:
name: The thread name as a string.
pid: Identifier of the process as an integer.
tid: Identifier of the thread as an integer.
"""
event = {}
event['name'] = 'thread_name'
event['ph'] = 'M'
event['pid'] = pid
event['tid'] = tid
event['args'] = {'name': name}
self._metadata.append(event)
def emit_region(
self,
timestamp: int,
duration: int,
pid: int,
tid: int,
category: str,
name: str,
args: Dict[str, Any],
) -> None:
"""Adds a region event to the trace.
Args:
timestamp: The start timestamp of this region as a long integer.
duration: The duration of this region as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
category: The event category as a string.
name: The event name as a string.
args: A JSON-compatible dictionary of event arguments.
"""
event = self._create_event('X', category, name, pid, tid, timestamp)
event['dur'] = duration
event['args'] = args
self._events.append(event)
def emit_obj_create(
self,
category: str,
name: str,
timestamp: int,
pid: int,
tid: int,
object_id: int,
) -> None:
"""Adds an object creation event to the trace.
Args:
category: The event category as a string.
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
object_id: Identifier of the object as an integer.
"""
event = self._create_event('N', category, name, pid, tid, timestamp)
event['id'] = object_id
self._events.append(event)
def emit_obj_delete(
self,
category: str,
name: str,
timestamp: int,
pid: int,
tid: int,
object_id: int,
) -> None:
"""Adds an object deletion event to the trace.
Args:
category: The event category as a string.
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
object_id: Identifier of the object as an integer.
"""
event = self._create_event('D', category, name, pid, tid, timestamp)
event['id'] = object_id
self._events.append(event)
def emit_obj_snapshot(
self,
category: str,
name: str,
timestamp: int,
pid: int,
tid: int,
object_id: int,
snapshot: Dict[str, Any],
) -> None:
"""Adds an object snapshot event to the trace.
Args:
category: The event category as a string.
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
object_id: Identifier of the object as an integer.
snapshot: A JSON-compatible representation of the object.
"""
event = self._create_event('O', category, name, pid, tid, timestamp)
event['id'] = object_id
event['args'] = {'snapshot': snapshot}
self._events.append(event)
def emit_flow_start(
self, name: str, timestamp: int, pid: int, tid: int, flow_id: int
) -> None:
"""Adds a flow start event to the trace.
When matched with a flow end event (with the same 'flow_id') this will
cause the trace viewer to draw an arrow between the start and end events.
Args:
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
flow_id: Identifier of the flow as an integer.
"""
event = self._create_event('s', 'DataFlow', name, pid, tid, timestamp)
event['id'] = flow_id
self._events.append(event)
def emit_flow_end(
self, name: str, timestamp: int, pid: int, tid: int, flow_id: int
) -> None:
"""Adds a flow end event to the trace.
When matched with a flow start event (with the same 'flow_id') this will
cause the trace viewer to draw an arrow between the start and end events.
Args:
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
flow_id: Identifier of the flow as an integer.
"""
event = self._create_event('t', 'DataFlow', name, pid, tid, timestamp)
event['id'] = flow_id
self._events.append(event)
def emit_counter(
self,
category: str,
name: str,
pid: int,
timestamp: int,
counter: str,
value: int,
) -> None:
"""Emits a record for a single counter.
Args:
category: The event category as a string.
name: The event name as a string.
pid: Identifier of the process generating this event as an integer.
timestamp: The timestamp of this event as a long integer.
counter: Name of the counter as a string.
value: Value of the counter as an integer.
"""
event = self._create_event('C', category, name, pid, 0, timestamp)
event['args'] = {counter: value}
self._events.append(event)
def emit_counters(self, category, name, pid, timestamp, counters):
"""Emits a counter record for the dictionary 'counters'.
Args:
category: The event category as a string.
name: The event name as a string.
pid: Identifier of the process generating this event as an integer.
timestamp: The timestamp of this event as a long integer.
counters: Dictionary of counter values.
"""
event = self._create_event('C', category, name, pid, 0, timestamp)
event['args'] = counters.copy()
self._events.append(event)
def format_to_string(self, pretty: bool = False) -> str:
"""Formats the chrome trace to a string.
Args:
pretty: (Optional.) If True, produce human-readable JSON output.
Returns:
A JSON-formatted string in Chrome Trace format.
"""
trace = {}
trace['traceEvents'] = self._metadata + self._events
if pretty:
return json.dumps(trace, indent=4, separators=(',', ': '))
else:
return json.dumps(trace, separators=(',', ':'))
class _TensorTracker(object):
"""An internal class to track the lifetime of a Tensor."""
def __init__(
self,
name: str,
object_id: int,
timestamp: int,
pid: int,
allocator: str,
num_bytes: int,
) -> None:
"""Creates an object to track tensor references.
This class is not thread safe and is intended only for internal use by
the 'Timeline' class in this file.
Args:
name: The name of the Tensor as a string.
object_id: Chrome Trace object identifier assigned for this Tensor.
timestamp: The creation timestamp of this event as a long integer.
pid: Process identifier of the associated device, as an integer.
allocator: Name of the allocator used to create the Tensor.
num_bytes: Number of bytes allocated (long integer).
Returns:
A 'TensorTracker' object.
"""
self._name = name
self._pid = pid
self._object_id = object_id
self._create_time = timestamp
self._allocator = allocator
self._num_bytes = num_bytes
self._ref_times = []
self._unref_times = []
@property
def name(self) -> str:
"""Name of this tensor."""
return self._name
@property
def pid(self) -> int:
"""ID of the process which created this tensor (an integer)."""
return self._pid
@property
def create_time(self) -> int:
"""Timestamp when this tensor was created (long integer)."""
return self._create_time
@property
def object_id(self) -> int:
"""Returns the object identifier of this tensor (integer)."""
return self._object_id
@property
def num_bytes(self) -> int:
"""Size of this tensor in bytes (long integer)."""
return self._num_bytes
@property
def allocator(self) -> str:
"""Name of the allocator used to create this tensor (string)."""
return self._allocator
@property
def last_unref(self) -> int:
"""Last unreference timestamp of this tensor (long integer)."""
return max(self._unref_times)
def add_ref(self, timestamp: int) -> None:
"""Adds a reference to this tensor with the specified timestamp.
Args:
timestamp: Timestamp of object reference as an integer.
"""
self._ref_times.append(timestamp)
def add_unref(self, timestamp: int) -> None:
"""Adds an unref to this tensor with the specified timestamp.
Args:
timestamp: Timestamp of object unreference as an integer.
"""
self._unref_times.append(timestamp)
class Timeline(object):
"""A class for visualizing execution timelines of TensorFlow steps."""
def __init__(
self, step_stats: step_stats_pb2.StepStats, graph: Optional[Any] = None
) -> None:
"""Constructs a new Timeline.
A 'Timeline' is used for visualizing the execution of a TensorFlow
computation. It shows the timings and concurrency of execution at
the granularity of TensorFlow Ops.
This class is not thread safe.
Args:
step_stats: The 'step_stats_pb2.StepStats' proto recording execution
times.
graph: (Optional) The 'Graph' that was executed.
"""
self._origin_step_stats = step_stats
self._step_stats = None
self._graph = graph
self._chrome_trace = _ChromeTraceFormatter()
self._next_pid = 0
self._device_pids = {} # device name -> pid for compute activity.
self._tensor_pids = {} # device name -> pid for tensors.
self._tensors = {} # tensor_name -> TensorTracker
self._next_flow_id = 0
self._flow_starts = {} # tensor_name -> (timestamp, pid, tid)
self._alloc_times = {} # tensor_name -> ( time, allocator, size )
self._allocator_maximums = {} # allocator name => maximum bytes long
def _alloc_pid(self) -> int:
"""Allocate a process Id."""
pid = self._next_pid
self._next_pid += 1
return pid
def _alloc_flow_id(self) -> int:
"""Allocate a flow Id."""
flow_id = self._next_flow_id
self._next_flow_id += 1
return flow_id
def _parse_op_label(
self, label: str
) -> Tuple[str, str, List[str]]:
"""Parses the fields in a node timeline label."""
# Expects labels of the form: name = op(arg, arg, ...).
match = re.match(r'(.*) = (.*)\((.*)\)', label)
if match is None:
return 'unknown', 'unknown', []
nn, op, inputs = match.groups()
if not inputs:
inputs = []
else:
inputs = inputs.split(', ')
return nn, op, inputs
def _parse_kernel_label(self, label, node_name):
"""Parses the fields in a node timeline label."""
# Expects labels of the form: retval (arg) detail @@annotation
start = label.find('@@')
end = label.find('#')
if start >= 0 and end >= 0 and start + 2 < end:
node_name = label[start + 2 : end]
# Node names should always have the form 'name:op'.
fields = node_name.split(':') + ['unknown']
name, op = fields[:2]
return name, op
def _assign_lanes(self) -> None:
"""Assigns non-overlapping lanes for the activities on each device."""
for device_stats in self._step_stats.dev_stats:
# TODO(pbar): Genuine thread IDs in step_stats_pb2.NodeExecStats
# might be helpful.
lanes = [0]
for ns in device_stats.node_stats:
l = -1
for i, lts in enumerate(lanes):
if ns.all_start_micros > lts:
l = i
lanes[l] = ns.all_start_micros + ns.all_end_rel_micros
break
if l < 0:
l = len(lanes)
lanes.append(ns.all_start_micros + ns.all_end_rel_micros)
ns.thread_id = l
def _emit_op(
self, nodestats: step_stats_pb2.NodeExecStats, pid: int, is_gputrace: bool
) -> None:
"""Generates a Chrome Trace event to show Op execution.
Args:
nodestats: The 'step_stats_pb2.NodeExecStats' proto recording op
execution.
pid: The pid assigned for the device where this op ran.
is_gputrace: If True then this op came from the GPUTracer.
"""
node_name = nodestats.node_name
start = nodestats.all_start_micros
duration = nodestats.all_end_rel_micros
tid = nodestats.thread_id
inputs = []
if is_gputrace:
node_name, op = self._parse_kernel_label(
nodestats.timeline_label, node_name
)
elif node_name == 'RecvTensor':
# RPC tracing does not use the standard timeline_label format.
op = 'RecvTensor'
else:
_, op, inputs = self._parse_op_label(nodestats.timeline_label)
args = {'name': node_name, 'op': op}
if build_info.build_info['is_rocm_build']:
args['kernel'] = nodestats.timeline_label.split('@@')[0]
for i, iname in enumerate(inputs):
args['input%d' % i] = iname
self._chrome_trace.emit_region(start, duration, pid, tid, 'Op', op, args)
def _emit_tensor_snapshot(
self,
tensor: _TensorTracker,
timestamp: int,
pid: int,
tid: int,
value: step_stats_pb2.NodeOutput,
) -> None:
"""Generate Chrome Trace snapshot event for a computed Tensor.
Args:
tensor: A 'TensorTracker' object.
timestamp: The timestamp of this snapshot as a long integer.
pid: The pid assigned for showing the device where this op ran.
tid: The tid of the thread computing the tensor snapshot.
value: A JSON-compliant snapshot of the object.
"""
desc = str(value.tensor_description).replace('"', '')
snapshot = {'tensor_description': desc}
self._chrome_trace.emit_obj_snapshot(
'Tensor', tensor.name, timestamp, pid, tid, tensor.object_id, snapshot
)
def _produce_tensor(
self,
name: str,
timestamp: int,
tensors_pid: int,
allocator: str,
num_bytes: int,
) -> _TensorTracker:
"""Creates a new tensor tracker."""
object_id = len(self._tensors)
tensor = _TensorTracker(
name, object_id, timestamp, tensors_pid, allocator, num_bytes
)
self._tensors[name] = tensor
return tensor
def _is_gputrace_device(self, device_name: str) -> bool:
"""Returns true if this device is part of the GPUTracer logging."""
return '/stream:' in device_name or '/memcpy' in device_name
def _allocate_pids(self) -> None:
"""Allocate fake process ids for each device in the step_stats_pb2.StepStats."""
self._allocators_pid = self._alloc_pid()
self._chrome_trace.emit_pid('Allocators', self._allocators_pid)
# Add processes in the Chrome trace to show compute and data activity.
for dev_stats in self._step_stats.dev_stats:
device_pid = self._alloc_pid()
self._device_pids[dev_stats.device] = device_pid
tensors_pid = self._alloc_pid()
self._tensor_pids[dev_stats.device] = tensors_pid
self._chrome_trace.emit_pid(dev_stats.device + ' Compute', device_pid)
self._chrome_trace.emit_pid(dev_stats.device + ' Tensors', tensors_pid)
def _analyze_tensors(self, show_memory: bool) -> None:
"""Analyze tensor references to track dataflow."""
for dev_stats in self._step_stats.dev_stats:
device_pid = self._device_pids[dev_stats.device]
tensors_pid = self._tensor_pids[dev_stats.device]
for node_stats in dev_stats.node_stats:
tid = node_stats.thread_id
node_name = node_stats.node_name
start_time = node_stats.all_start_micros
end_time = node_stats.all_start_micros + node_stats.all_end_rel_micros
for index, output in enumerate(node_stats.output):
if index:
output_name = '%s:%d' % (node_name, index)
else:
output_name = node_name
allocation = output.tensor_description.allocation_description
num_bytes = allocation.requested_bytes
allocator_name = allocation.allocator_name
tensor = self._produce_tensor(
output_name, start_time, tensors_pid, allocator_name, num_bytes
)
tensor.add_ref(start_time)
tensor.add_unref(end_time)
self._flow_starts[output_name] = (end_time, device_pid, tid)
if show_memory:
self._chrome_trace.emit_obj_create(
'Tensor',
output_name,
start_time,
tensors_pid,
tid,
tensor.object_id,
)
self._emit_tensor_snapshot(
tensor, end_time - 1, tensors_pid, tid, output
)
def _show_compute(self, show_dataflow: bool) -> None:
"""Visualize the computation activity."""
for dev_stats in self._step_stats.dev_stats:
device_name = dev_stats.device
device_pid = self._device_pids[device_name]
is_gputrace = self._is_gputrace_device(device_name)
for node_stats in dev_stats.node_stats:
tid = node_stats.thread_id
start_time = node_stats.all_start_micros
end_time = node_stats.all_start_micros + node_stats.all_end_rel_micros
self._emit_op(node_stats, device_pid, is_gputrace)
if is_gputrace or node_stats.node_name == 'RecvTensor':
continue
_, _, inputs = self._parse_op_label(node_stats.timeline_label)
for input_name in inputs:
if input_name not in self._tensors:
# This can happen when partitioning has inserted a Send/Recv.
# We remove the numeric suffix so that the dataflow appears to
# come from the original node. Ideally, the StepStats would
# contain logging for the Send and Recv nodes.
index = input_name.rfind('/_')
if index > 0:
input_name = input_name[:index]
if input_name in self._tensors:
tensor = self._tensors[input_name]
tensor.add_ref(start_time)
tensor.add_unref(end_time - 1)
if show_dataflow:
# We use a different flow ID for every graph edge.
create_time, create_pid, create_tid = self._flow_starts[
input_name
]
# Don't add flows when producer and consumer ops are on the same
# pid/tid since the horizontal arrows clutter the visualization.
if create_pid != device_pid or create_tid != tid:
flow_id = self._alloc_flow_id()
self._chrome_trace.emit_flow_start(
input_name, create_time, create_pid, create_tid, flow_id
)
self._chrome_trace.emit_flow_end(
input_name, start_time, device_pid, tid, flow_id
)
else:
logging.vlog(
1, "Can't find tensor %s - removed by CSE?", input_name
)
def _show_memory_counters(self) -> None:
"""Produce a counter series for each memory allocator."""
# Iterate over all tensor trackers to build a list of allocations and
# frees for each allocator. Then sort the lists and emit a cumulative
# counter series for each allocator.
allocations = {}
for name in self._tensors:
tensor = self._tensors[name]
self._chrome_trace.emit_obj_delete(
'Tensor', name, tensor.last_unref, tensor.pid, 0, tensor.object_id
)
allocator = tensor.allocator
if allocator not in allocations:
allocations[allocator] = []
num_bytes = tensor.num_bytes
allocations[allocator].append((tensor.create_time, num_bytes, name))
allocations[allocator].append((tensor.last_unref, -num_bytes, name))
alloc_maxes = {}
# Generate a counter series showing total allocations for each allocator.
for allocator in allocations:
alloc_list = allocations[allocator]
alloc_list.sort()
total_bytes = 0
alloc_tensor_set = set()
alloc_maxes[allocator] = AllocationMaximum(
timestamp=0, num_bytes=0, tensors=set()
)
for time, num_bytes, name in sorted(
alloc_list, key=lambda allocation: allocation[0]
):
total_bytes += num_bytes
if num_bytes < 0:
alloc_tensor_set.discard(name)
else:
alloc_tensor_set.add(name)
if total_bytes > alloc_maxes[allocator].num_bytes:
alloc_maxes[allocator] = AllocationMaximum(
timestamp=time,
num_bytes=total_bytes,
tensors=copy.deepcopy(alloc_tensor_set),
)
self._chrome_trace.emit_counter(
'Memory',
allocator,
self._allocators_pid,
time,
allocator,
total_bytes,
)
self._allocator_maximums = alloc_maxes
def _preprocess_op_time(self, op_time: str) -> None:
"""Update the start and end time of ops in step stats.
Args:
op_time: How the execution time of op is shown in timeline. Possible
values are "schedule", "gpu" and "all". "schedule" will show op from
the time it is scheduled to the end of the scheduling. Notice by the end
of its scheduling its async kernels may not start yet. It is shown using
the default value from step_stats. "gpu" will show op with the execution
time of its kernels on GPU. "all" will show op from the start of its
scheduling to the end of its last kernel.
"""
if op_time == 'schedule':
self._step_stats = self._origin_step_stats
return
self._step_stats = copy.deepcopy(self._origin_step_stats)
# Separate job task and gpu tracer stream
stream_all_stats = []
job_stats = []
for stats in self._step_stats.dev_stats:
if '/stream:all' in stats.device:
stream_all_stats.append(stats)
elif '/job' in stats.device:
job_stats.append(stats)
# Record the start time of the first kernel and the end time of
# the last gpu kernel for all ops.
op_gpu_start = {}
op_gpu_end = {}
for stats in stream_all_stats:
for kernel in stats.node_stats:
name, _ = self._parse_kernel_label(
kernel.timeline_label, kernel.node_name
)
start = kernel.all_start_micros
end = kernel.all_start_micros + kernel.all_end_rel_micros
if name in op_gpu_start:
op_gpu_start[name] = min(op_gpu_start[name], start)
op_gpu_end[name] = max(op_gpu_end[name], end)
else:
op_gpu_start[name] = start
op_gpu_end[name] = end
# Update the start and end time of each op according to the op_time
for stats in job_stats:
for op in stats.node_stats:
if op.node_name in op_gpu_start:
end = max(
op_gpu_end[op.node_name],
op.all_start_micros + op.all_end_rel_micros,
)
if op_time == 'gpu':
op.all_start_micros = op_gpu_start[op.node_name]
op.all_end_rel_micros = end - op.all_start_micros
def analyze_step_stats(
self,
show_dataflow: bool = True,
show_memory: bool = True,
op_time: str = 'schedule',
) -> StepStatsAnalysis:
"""Analyze the step stats and format it into Chrome Trace Format.
Args:
show_dataflow: (Optional.) If True, add flow events to the trace
connecting producers and consumers of tensors.
show_memory: (Optional.) If True, add object snapshot events to the trace
showing the sizes and lifetimes of tensors.
op_time: (Optional.) How the execution time of op is shown in timeline.
Possible values are "schedule", "gpu" and "all". "schedule" will show op
from the time it is scheduled to the end of the scheduling. Notice by
the end of its scheduling its async kernels may not start yet. It is
shown using the default value from step_stats. "gpu" will show op with
the execution time of its kernels on GPU. "all" will show op from the
start of its scheduling to the end of its last kernel.
Returns:
A 'StepStatsAnalysis' object.
"""
self._preprocess_op_time(op_time)
self._allocate_pids()
self._assign_lanes()
self._analyze_tensors(show_memory)
self._show_compute(show_dataflow)
if show_memory:
self._show_memory_counters()
return StepStatsAnalysis(
chrome_trace=self._chrome_trace,
allocator_maximums=self._allocator_maximums,
)
def generate_chrome_trace_format(
self,
show_dataflow: bool = True,
show_memory: bool = False,
op_time: str = 'schedule',
) -> str:
# pyformat: disable
"""Produces a trace in Chrome Trace Format.
Args:
show_dataflow: (Optional.) If True, add flow events to the trace
connecting producers and consumers of tensors.
show_memory: (Optional.) If True, add object snapshot events to the trace
showing the sizes and lifetimes of tensors.
op_time: (Optional.) How the execution time of op is shown in timeline.
Possible values are "schedule", "gpu" and "all".
"schedule" will show op from the time it is scheduled to the end of
the scheduling.
Notice by the end of its scheduling its async kernels may not start
yet. It is shown using the default value from step_stats.
"gpu" will show op with the execution time of its kernels on GPU.
"all" will show op from the start of its scheduling to the end of
its last kernel.
Returns:
A JSON formatted string in Chrome Trace format.
"""
# pyformat: enable
step_stats_analysis = self.analyze_step_stats(
show_dataflow=show_dataflow, show_memory=show_memory, op_time=op_time
)
return step_stats_analysis.chrome_trace.format_to_string(pretty=True)
tensorflow/python/client/timeline_test.py
+201
View File
@@ -0,0 +1,201 @@
# 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.
# ==============================================================================
"""Tests for tensorflow.python.client.Timeline."""
import json
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.client import session
from tensorflow.python.client import timeline
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import test_util
from tensorflow.python.framework import ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
class TimelineTest(test.TestCase):
def _validateTrace(self, chrome_trace_format):
# Check that the supplied string is valid JSON.
trace = json.loads(chrome_trace_format)
# It should have a top-level key containing events.
self.assertTrue('traceEvents' in trace)
# Every event in the list should have a 'ph' field.
for event in trace['traceEvents']:
self.assertTrue('ph' in event)
def testSimpleTimeline(self):
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
with ops.device('/cpu:0'):
with session.Session() as sess:
sess.run(constant_op.constant(1.0),
options=run_options,
run_metadata=run_metadata)
self.assertTrue(run_metadata.HasField('step_stats'))
tl = timeline.Timeline(run_metadata.step_stats)
ctf = tl.generate_chrome_trace_format()
self._validateTrace(ctf)
@test_util.deprecated_graph_mode_only
def testTimelineCpu(self):
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
with self.session(use_gpu=False) as sess:
const1 = constant_op.constant(1.0, name='const1')
const2 = constant_op.constant(2.0, name='const2')
result = math_ops.add(const1, const2) + const1 * const2
sess.run(result, options=run_options, run_metadata=run_metadata)
self.assertTrue(run_metadata.HasField('step_stats'))
step_stats = run_metadata.step_stats
devices = [d.device for d in step_stats.dev_stats]
self.assertTrue('/job:localhost/replica:0/task:0/device:CPU:0' in devices)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format()
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(show_dataflow=False)
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(show_memory=False)
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(
show_memory=False, show_dataflow=False)
self._validateTrace(ctf)
@test_util.deprecated_graph_mode_only
def testTimelineGpu(self):
if not test.is_gpu_available(cuda_only=True):
return
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
with self.session(force_gpu=True) as sess:
const1 = constant_op.constant(1.0, name='const1')
const2 = constant_op.constant(2.0, name='const2')
result = math_ops.add(const1, const2) + const1 * const2
sess.run(result, options=run_options, run_metadata=run_metadata)
self.assertTrue(run_metadata.HasField('step_stats'))
step_stats = run_metadata.step_stats
devices = [d.device for d in step_stats.dev_stats]
self.assertTrue('/job:localhost/replica:0/task:0/device:GPU:0' in devices)
self.assertIn('/device:GPU:0/stream:all', devices)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format()
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(show_dataflow=False)
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(show_memory=False)
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(
show_memory=False, show_dataflow=False)
self._validateTrace(ctf)
def testTimelineWithRPCs(self):
"""Tests that Timeline can handle RPC tracing."""
metadata = config_pb2.RunMetadata()
step_stats = metadata.step_stats
dev_stats = step_stats.dev_stats.add()
dev_stats.device = '/job:worker/replica:0/task:0/cpu:0'
node_stats = dev_stats.node_stats.add()
node_stats.node_name = 'RecvTensor'
node_stats.all_start_micros = 12345
node_stats.op_end_rel_micros = 42
node_stats.timeline_label = ('[1024B] edge_160_conv2/biases/read from '
'/job:ps/replica:0/task:3/cpu:0 to '
'/job:worker/replica:0/task:0/cpu:0')
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format()
self._validateTrace(ctf)
def testAnalysisAndAllocations(self):
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
config = config_pb2.ConfigProto(device_count={'CPU': 3})
with session.Session(config=config) as sess:
with ops.device('/cpu:0'):
num1 = variables.Variable(1.0, name='num1')
with ops.device('/cpu:1'):
num2 = variables.Variable(2.0, name='num2')
with ops.device('/cpu:2'):
result = num1 + num2 + num1 * num2
self.evaluate(variables.global_variables_initializer())
sess.run(result, options=run_options, run_metadata=run_metadata)
self.assertTrue(run_metadata.HasField('step_stats'))
tl = timeline.Timeline(run_metadata.step_stats)
step_analysis = tl.analyze_step_stats()
ctf = step_analysis.chrome_trace.format_to_string()
self._validateTrace(ctf)
maximums = step_analysis.allocator_maximums
cpuname = 'mklcpu' if test_util.IsMklEnabled() else 'cpu'
self.assertTrue(cpuname in maximums)
cpu_max = maximums[
'cuda_host_bfc'] if 'cuda_host_bfc' in maximums else maximums[cpuname]
# At least num1 + num2, both float32s (4 bytes each)
self.assertGreaterEqual(cpu_max.num_bytes, 8)
self.assertGreater(cpu_max.timestamp, 0)
def testManyCPUs(self):
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
config = config_pb2.ConfigProto(device_count={'CPU': 3})
with session.Session(config=config) as sess:
with ops.device('/cpu:0'):
num1 = variables.Variable(1.0, name='num1')
with ops.device('/cpu:1'):
num2 = variables.Variable(2.0, name='num2')
with ops.device('/cpu:2'):
result = num1 + num2 + num1 * num2
self.evaluate(variables.global_variables_initializer())
sess.run(result, options=run_options, run_metadata=run_metadata)
self.assertTrue(run_metadata.HasField('step_stats'))
step_stats = run_metadata.step_stats
devices = [d.device for d in step_stats.dev_stats]
self.assertTrue('/job:localhost/replica:0/task:0/device:CPU:0' in devices)
self.assertTrue('/job:localhost/replica:0/task:0/device:CPU:1' in devices)
self.assertTrue('/job:localhost/replica:0/task:0/device:CPU:2' in devices)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format()
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(show_dataflow=False)
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(show_memory=False)
self._validateTrace(ctf)
tl = timeline.Timeline(step_stats)
ctf = tl.generate_chrome_trace_format(
show_memory=False, show_dataflow=False)
self._validateTrace(ctf)
if __name__ == '__main__':
test.main()
tensorflow/python/client/virtual_gpu_test.py
+242
View File
@@ -0,0 +1,242 @@
# 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.
# ==============================================================================
"""Tests for multiple virtual GPU support."""
import random
import numpy as np
from google.protobuf import text_format
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
class VirtualGpuTestUtil(object):
def __init__(self,
dim=1000,
num_ops=100,
virtual_devices_per_gpu=None,
device_probabilities=None):
self._dim = dim
self._num_ops = num_ops
if virtual_devices_per_gpu is None:
self._virtual_devices_per_gpu = [3]
else:
self._virtual_devices_per_gpu = virtual_devices_per_gpu
self._visible_device_list = [
i for i in range(len(self._virtual_devices_per_gpu))
]
gpu_devices = [
('/gpu:' + str(i)) for i in range(sum(self._virtual_devices_per_gpu))
]
self.devices = ['/cpu:0'] + gpu_devices
self._num_devices = len(self.devices)
# Each virtual device gets 2GB memory.
self._mem_limits_mb = [
([1 << 11] * i) for i in self._virtual_devices_per_gpu
]
self.config = self._GetSessionConfig()
if device_probabilities is not None:
self._device_probabilities = list(device_probabilities) # Deep copy
for i in range(1, self._num_devices):
self._device_probabilities[i] += self._device_probabilities[i - 1]
else:
# Each device gets same probability to be assigned an operation.
step = 1.0 / self._num_devices
self._device_probabilities = [
(x + 1) * step for x in range(self._num_devices)
]
# To prevent rounding error causing problems.
self._device_probabilities[self._num_devices - 1] = 1.1
logging.info('dim: %d', self._dim)
logging.info('num_ops: %d', self._num_ops)
logging.info('visible_device_list: %s', str(self._visible_device_list))
logging.info('virtual_devices_per_gpu: %s',
str(self._virtual_devices_per_gpu))
logging.info('mem_limits: %s', str(self._mem_limits_mb))
logging.info('devices: %s', str(self.devices))
logging.info('config: %s', text_format.MessageToString(self.config))
logging.info('device_probabilities: %s', str(self._device_probabilities))
# Creates virtual GPU devices
def _GetSessionConfig(self):
virtual_device_gpu_options = config_pb2.GPUOptions(
visible_device_list=','.join(str(d) for d in self._visible_device_list),
experimental=config_pb2.GPUOptions.Experimental(virtual_devices=[
config_pb2.GPUOptions.Experimental.VirtualDevices(
memory_limit_mb=i) for i in self._mem_limits_mb
]))
return config_pb2.ConfigProto(gpu_options=virtual_device_gpu_options)
# Generates a list of 3-tuples, each tuple contains the source and destination
# device index for a binary operation like 'add', like:
# (src_device_1, src_device_2, dst_device)
def _GenerateOperationPlacement(self):
result = []
for unused_i in range(self._num_ops):
op_device = ()
for unused_j in range(3):
random_num = random.random()
for device_index in range(self._num_devices):
if self._device_probabilities[device_index] > random_num:
op_device += (device_index,)
break
result.append(op_device)
return result
# Logs part of the matrix for debugging purposes.
def _LogMatrix(self, mat, dim):
logging.info('---- printing the first 10*10 submatrix ----')
for i in range(min(10, dim)):
row = ''
for j in range(min(10, dim)):
row += ' ' + str(mat[i][j])
logging.info(row)
# Runs a list of 'add' operations where each operation satisfies the device
# placement constraints in `op_placement`, and returns the result.
def _TestRandomGraphWithDevices(self,
sess,
seed,
op_placement,
devices,
debug_mode=False):
data = []
shape = (self._dim, self._dim)
feed_dict = {}
# Initialize the matrices
for i in range(len(devices)):
with ops.device(devices[i]):
var = array_ops.placeholder(dtypes.float32, shape=shape)
np.random.seed(seed + i)
feed_dict[var] = np.random.uniform(
low=0, high=0.1, size=shape).astype(np.float32)
data.append(var)
# Run the 'add' operations on those matrices
for op in op_placement:
with ops.device(devices[op[2]]):
data[op[2]] = math_ops.add(data[op[0]], data[op[1]])
with ops.device('/cpu:0'):
s = data[0]
for i in range(1, len(data)):
s = math_ops.add(s, data[i])
if debug_mode:
logging.info(ops.get_default_graph().as_graph_def())
result = sess.run(s, feed_dict=feed_dict)
self._LogMatrix(result, self._dim)
return result
# Generates a random graph with `self._num_ops` 'add' operations with each
# operation placed on different virtual device, test that the result is
# identical to the result obtained by running the same graph on cpu only.
def TestRandomGraph(self, sess, op_placement=None, random_seed=None):
debug_mode = False
if op_placement is None:
op_placement = self._GenerateOperationPlacement()
else:
debug_mode = True
if random_seed is None:
random_seed = random.randint(0, 1 << 31)
else:
debug_mode = True
logging.info('Virtual gpu functional test for random graph...')
logging.info('operation placement: %s', str(op_placement))
logging.info('random seed: %d', random_seed)
# Run with multiple virtual gpus.
result_vgd = self._TestRandomGraphWithDevices(
sess, random_seed, op_placement, self.devices, debug_mode=debug_mode)
# Run with single cpu.
result_cpu = self._TestRandomGraphWithDevices(
sess,
random_seed,
op_placement, ['/cpu:0'] * self._num_devices,
debug_mode=debug_mode)
# Test the result
for i in range(self._dim):
for j in range(self._dim):
if result_vgd[i][j] != result_cpu[i][j]:
logging.error(
'Result mismatch at row %d column %d: expected %f, actual %f', i,
j, result_cpu[i][j], result_vgd[i][j])
logging.error('Devices: %s', self.devices)
logging.error('Memory limits (in MB): %s', self._mem_limits_mb)
return False
return True
class VirtualGpuTest(test_util.TensorFlowTestCase):
def __init__(self, method_name):
super(VirtualGpuTest, self).__init__(method_name)
self._util = VirtualGpuTestUtil()
@test_util.deprecated_graph_mode_only
def testStatsContainAllDeviceNames(self):
with self.session(config=self._util.config) as sess:
# TODO(laigd): b/70811538. The is_gpu_available() call will invoke
# DeviceFactory::AddDevices() with a default SessionOption, which prevents
# adding virtual devices in the future, thus must be called within a
# context of a session within which virtual devices are created. Same in
# the following test case.
if not test.is_gpu_available(cuda_only=True):
self.skipTest('No GPU available')
run_options = config_pb2.RunOptions(
trace_level=config_pb2.RunOptions.FULL_TRACE)
run_metadata = config_pb2.RunMetadata()
mat_shape = [10, 10]
data = []
for d in self._util.devices:
with ops.device(d):
var = variables.Variable(random_ops.random_uniform(mat_shape))
self.evaluate(var.initializer)
data.append(var)
s = data[0]
for i in range(1, len(data)):
s = math_ops.add(s, data[i])
sess.run(s, options=run_options, run_metadata=run_metadata)
self.assertTrue(run_metadata.HasField('step_stats'))
step_stats = run_metadata.step_stats
devices = [d.device for d in step_stats.dev_stats]
self.assertTrue('/job:localhost/replica:0/task:0/device:CPU:0' in devices)
self.assertTrue('/job:localhost/replica:0/task:0/device:GPU:0' in devices)
self.assertTrue('/job:localhost/replica:0/task:0/device:GPU:1' in devices)
self.assertTrue('/job:localhost/replica:0/task:0/device:GPU:2' in devices)
@test_util.deprecated_graph_mode_only
def testLargeRandomGraph(self):
with self.session(config=self._util.config) as sess:
if not test.is_gpu_available(cuda_only=True):
self.skipTest('No GPU available')
for _ in range(5):
if not self._util.TestRandomGraph(sess):
return
if __name__ == '__main__':
test.main()