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

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

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/core/data/BUILD
+772
View File
@@ -0,0 +1,772 @@
load(
"//tensorflow:tensorflow.bzl",
"if_not_mobile",
"tf_cc_test",
)
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_protos_all",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
# Export files for use on Android.
exports_files([
"captured_function.cc",
"captured_function.h",
"compression_utils.cc",
"compression_utils.h",
"dataset_utils.cc",
"dataset_utils.h",
"finalization_utils.cc",
"finalization_utils.h",
"flat_map_utils.cc",
"flat_map_utils.h",
"global_shuffle_utils.cc",
"global_shuffle_utils.h",
"metric_utils.cc",
"metric_utils.h",
"name_utils.cc",
"name_utils.h",
"rewrite_utils.cc",
"rewrite_utils.h",
"root_dataset.cc",
"root_dataset.h",
"serialization_utils.cc",
"serialization_utils.h",
"split_utils.cc",
"split_utils.h",
"stats_utils.cc",
"stats_utils.h",
"tf_data_memory_logger.cc",
"tf_data_memory_logger.h",
"tfdataz_metrics.h",
"tfdataz_metrics.cc",
"unbounded_thread_pool.cc",
"unbounded_thread_pool.h",
"utils.cc",
"utils.h",
])
cc_library(
name = "captured_function",
srcs = ["captured_function.cc"],
hdrs = ["captured_function.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//visibility:public"],
deps = [
":dataset_utils",
":stats_utils",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/profiler/lib:traceme",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/time",
] + if_not_mobile([
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/optimizers:meta_optimizer",
]),
)
cc_library(
name = "compression_utils",
srcs = ["compression_utils.cc"],
hdrs = ["compression_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "compression_utils_test",
srcs = ["compression_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
tags = [
"requires-mem:24g",
],
deps = [
":compression_utils",
":dataset_test_base",
"//tensorflow/core:framework",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"@com_google_absl//absl/status:status_matchers",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
cc_library(
name = "dataset_test_base",
testonly = 1,
srcs = ["dataset_test_base.cc"],
hdrs = ["dataset_test_base.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//visibility:public"],
deps = [
":dataset_utils",
":name_utils",
":serialization_utils",
":split_utils",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:testlib",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/kernels:function_ops",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@eigen_archive//:eigen3",
"@xla//xla/tsl/framework/fixedpoint",
],
)
cc_library(
name = "dataset_utils",
srcs = ["dataset_utils.cc"],
hdrs = ["dataset_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:regexp",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "dataset_utils_test",
size = "small",
srcs = ["dataset_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":compression_utils",
":dataset_test_base",
":dataset_utils",
":serialization_utils",
":test_utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:str_util",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
"@xla//xla/tsl/util:determinism_test_util",
],
)
cc_library(
name = "flat_map_utils",
srcs = ["flat_map_utils.cc"],
hdrs = ["flat_map_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":captured_function",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core/framework:types_proto_cc",
"//tensorflow/core/kernels/data:iterator_ops",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@tsl//tsl/platform:refcount",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "global_shuffle_utils",
srcs = ["global_shuffle_utils.cc"],
hdrs = ["global_shuffle_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:framework",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "hash_utils",
srcs = ["hash_utils.cc"],
hdrs = ["hash_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":dataset_utils",
":serialization_utils",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:regexp",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "hash_utils_test",
size = "small",
srcs = ["hash_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":hash_utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/platform:test_main",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
cc_library(
name = "tf_data_memory_logger",
srcs = ["tf_data_memory_logger.cc"],
hdrs = ["tf_data_memory_logger.h"],
visibility = ["//visibility:public"],
deps = [
":tfdataz_metrics",
"//tensorflow/core/framework:model_proto_cc",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:numbers",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/base",
"@com_google_absl//absl/container:flat_hash_set",
"@xla//xla/tsl/platform:logging",
],
)
cc_library(
name = "metric_utils",
srcs = ["metric_utils.cc"],
hdrs = [
"metric_utils.h",
],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//visibility:public"],
deps = [
":tfdataz_metrics",
"//tensorflow/core:framework",
"//tensorflow/core/data:utils",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:mutex",
"//tensorflow/core/platform:thread_annotations",
"@com_google_absl//absl/time",
],
)
tf_cc_test(
name = "metric_utils_test",
size = "small",
srcs = ["metric_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":metric_utils",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/lib/monitoring:cell_reader",
"//tensorflow/core/lib/monitoring:test_utils",
"//tensorflow/core/platform:env",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/time",
],
)
cc_library(
name = "tfdataz_metrics",
srcs = ["tfdataz_metrics.cc"],
hdrs = ["tfdataz_metrics.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:mutex",
"//tensorflow/core/platform:thread_annotations",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/time",
],
)
tf_cc_test(
name = "tfdataz_metrics_test",
size = "small",
srcs = ["tfdataz_metrics_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":tfdataz_metrics",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:env",
"//tensorflow/core/util:fake_clock_env",
"@com_google_absl//absl/time",
],
)
cc_library(
name = "name_utils",
srcs = ["name_utils.cc"],
hdrs = ["name_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "name_utils_test",
size = "small",
srcs = ["name_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":name_utils",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
cc_library(
name = "rewrite_utils",
srcs = ["rewrite_utils.cc"],
hdrs = ["rewrite_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":dataset_utils",
":hash_utils",
":serialization_utils",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:grappler_item_builder",
"//tensorflow/core/grappler/clusters:virtual_cluster",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
"//tensorflow/core/grappler/optimizers:meta_optimizer",
"//tensorflow/core/grappler/optimizers/data",
"//tensorflow/core/grappler/optimizers/data:function_utils",
"//tensorflow/core/grappler/optimizers/data:graph_utils",
"//tensorflow/core/platform",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "rewrite_utils_test",
size = "small",
srcs = ["rewrite_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":rewrite_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/framework:function_proto_cc",
"//tensorflow/core/framework:function_testlib",
"//tensorflow/core/framework:graph_proto_cc",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/framework:types_proto_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/ops",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "root_dataset",
srcs = ["root_dataset.cc"],
hdrs = ["root_dataset.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":dataset_utils",
":name_utils",
":rewrite_utils",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:platform_port",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:stringprintf",
"@com_google_absl//absl/status",
],
)
cc_library(
name = "serialization_utils",
srcs = ["serialization_utils.cc"],
hdrs = ["serialization_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":compression_utils",
":dataset_utils",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/lib/core:status",
"//tensorflow/core/platform:stringpiece",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "serialization_utils_test",
size = "small",
srcs = ["serialization_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":dataset_test_base",
":dataset_utils",
":serialization_utils",
":test_utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:session_options",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:device_factory",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:str_util",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
cc_library(
name = "split_utils",
srcs = ["split_utils.cc"],
hdrs = ["split_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:mutex",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:logging",
"@xla//xla/tsl/platform:types",
],
)
tf_cc_test(
name = "split_utils_test",
size = "small",
srcs = ["split_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":dataset_test_base",
":serialization_utils",
":split_utils",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:dataset_utils",
"//tensorflow/core/framework:tensor_testutil",
],
)
cc_library(
name = "snapshot_utils",
srcs = ["snapshot_utils.cc"],
hdrs = ["snapshot_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":name_utils",
"//tensorflow/core:core_cpu_lib",
"//tensorflow/core:dataset_ops_op_lib",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:coding",
"//tensorflow/core/platform:random",
"//tensorflow/core/profiler/lib:traceme",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "snapshot_utils_test",
size = "small",
srcs = ["snapshot_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":snapshot_utils",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:test_util",
"@com_google_absl//absl/status",
],
)
cc_library(
name = "standalone",
srcs = ["standalone.cc"],
hdrs = ["standalone.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":dataset_utils",
":root_dataset",
":serialization_utils",
":tf_data_memory_logger",
":tfdataz_metrics",
":unbounded_thread_pool",
"//tensorflow/core:all_kernels",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:session_options",
"//tensorflow/core/framework:graph_proto_cc",
"@com_google_absl//absl/memory",
"@tsl//tsl/platform:refcount",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "standalone_save_restore_test",
srcs = ["standalone_save_restore_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":standalone",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:test_util",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:protos_all_cc",
] + tf_protos_all(),
)
tf_cc_test(
name = "standalone_test",
srcs = ["standalone_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":standalone",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
] + tf_protos_all(),
)
cc_library(
name = "stats_utils",
srcs = ["stats_utils.cc"],
hdrs = ["stats_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:lib",
"@com_google_absl//absl/base:core_headers",
],
)
cc_library(
name = "test_utils",
srcs = ["test_utils.cc"],
hdrs = ["test_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/framework:function_proto_cc",
"//tensorflow/core/public:version",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "unbounded_thread_pool",
srcs = ["unbounded_thread_pool.cc"],
hdrs = ["unbounded_thread_pool.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "finalization_utils",
srcs = ["finalization_utils.cc"],
hdrs = ["finalization_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":root_dataset",
"//tensorflow/core:framework",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:refcount",
"@com_google_absl//absl/status:statusor",
],
)
tf_cc_test(
name = "unbounded_thread_pool_test",
size = "small",
srcs = ["unbounded_thread_pool_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":unbounded_thread_pool",
"//tensorflow/core:lib_internal",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/status:statusor",
],
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":utils",
"//tensorflow/core:framework",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/core/data/captured_function.cc
+1055
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/captured_function.h
+341
View File
@@ -0,0 +1,341 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_CAPTURED_FUNCTION_H_
#define TENSORFLOW_CORE_DATA_CAPTURED_FUNCTION_H_
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/cancellation.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/model.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/lib/random/random.h"
#include "tensorflow/core/platform/macros.h"
namespace tensorflow {
class Device;
class OpKernelContext;
class ResourceMgr;
namespace data {
class CapturedFunction;
class InstantiatedCapturedFunction;
// Creates an iterator for a dataset which is created by applying the given
// function to the given input element.
absl::Status MakeIteratorFromInputElement(
IteratorContext* ctx, const DatasetBaseIterator* parent,
const std::vector<Tensor>& input_element, int64_t thread_index,
const InstantiatedCapturedFunction& inst_captured_func,
absl::string_view prefix, std::unique_ptr<IteratorBase>* out_iterator);
// Creates an iterator for a dataset which is created by applying the given
// function to the given input element. Pass non-null `node` to record
// processing time for modeling Iterator's GetNext() resource usage.
absl::Status MakeIteratorFromInputElement(
IteratorContext* ctx, const DatasetBaseIterator* parent,
const std::vector<Tensor>& input_element, int64_t thread_index,
const InstantiatedCapturedFunction& inst_captured_func,
absl::string_view prefix, std::unique_ptr<IteratorBase>* out_iterator,
const std::shared_ptr<model::Node>& node);
struct ShortCircuitInfo {
std::vector<int> indices;
std::vector<bool> can_move;
};
// Metadata shared across all captures of the same function.
class FunctionMetadata {
public:
struct Params {
bool use_inter_op_parallelism = true;
bool use_default_device = true;
};
// Creates a new instance of the `FunctionMetadata` class, fetching function
// from a context argument.
static absl::Status Create(tensorflow::OpKernelConstruction* ctx,
const std::string& func_name, Params params,
std::shared_ptr<FunctionMetadata>* out_metadata);
// Creates a new instance of the `FunctionMetadata` class, using the provided
// function.
static absl::Status Create(tensorflow::OpKernelConstruction* ctx,
NameAttrList&& func, Params params,
std::shared_ptr<FunctionMetadata>* out_metadata);
// Returns the named list of function arguments.
const NameAttrList& func() const { return func_; }
// Returns a borrowed pointer to the function library that contains the
// transitive closure of definitions used by the function.
const FunctionLibraryDefinition* lib_def() const { return lib_def_.get(); }
// Returns short-circuit information.
const ShortCircuitInfo& short_circuit_info() const {
return short_circuit_info_;
}
// Indicates whether a default device should be used for executing function
// ops.
bool use_default_device() const { return use_default_device_; }
// Indicates whether to use inter-op parallelism for execution of the
// function.
bool use_inter_op_parallelism() const { return use_inter_op_parallelism_; }
// Indicates whether the function should a multi-device function backend.
bool use_multi_device_function() const { return use_multi_device_function_; }
private:
FunctionMetadata(NameAttrList&& func, Params params)
: func_(std::move(func)),
use_default_device_(params.use_default_device),
use_inter_op_parallelism_(params.use_inter_op_parallelism) {}
NameAttrList func_;
std::unique_ptr<FunctionLibraryDefinition> lib_def_ = nullptr;
ShortCircuitInfo short_circuit_info_;
bool use_default_device_ = true;
bool use_inter_op_parallelism_ = true;
bool use_multi_device_function_ = true;
};
// Constructs and stores the parameters for the CapturedFunction Instantiate
// function.
struct InstantiateCapturedFunctionParams {
explicit InstantiateCapturedFunctionParams(IteratorContext* ctx) {
flr = ctx->flr();
function_handle_cache = ctx->function_handle_cache();
runner = ctx->runner();
}
explicit InstantiateCapturedFunctionParams(OpKernelContext* ctx) {
flr = ctx->function_library();
function_handle_cache = nullptr;
runner = ctx->runner();
}
FunctionLibraryRuntime* flr;
FunctionHandleCache* function_handle_cache;
std::function<void(std::function<void()>)>* runner;
};
// A `CapturedFunction` encapsulates a TensorFlow function, plus any "captured"
// arguments that it closed over in the user program.
class CapturedFunction {
public:
// Creates a new instance using a list of named attributes, fetching captured
// inputs from a context argument.
static absl::Status Create(OpKernelContext* ctx,
std::shared_ptr<const FunctionMetadata> metadata,
const std::string& argument_name,
std::unique_ptr<CapturedFunction>* out_function);
// Creates a new instance using a list of named attributes, using provided
// captured inputs.
static absl::Status Create(OpKernelContext* ctx,
std::shared_ptr<const FunctionMetadata> metadata,
std::vector<Tensor>&& captured_inputs,
std::unique_ptr<CapturedFunction>* out_function);
// Adds the definition of this captured function into the given graph,
// returning its captured inputs and types through the respective output
// arguments.
absl::Status AddToGraph(SerializationContext* ctx,
DatasetBase::DatasetGraphDefBuilder* b,
std::vector<Node*>* other_arguments,
DataTypeVector* other_arguments_types) const;
// Instantiates this function for use in the given context, providing an
// InstantiatedCapturedFunction that can be used to execute functions.
absl::Status Instantiate(IteratorContext* ctx,
std::unique_ptr<InstantiatedCapturedFunction>*
instantiated_captured_function);
absl::Status Instantiate(InstantiateCapturedFunctionParams params,
std::unique_ptr<InstantiatedCapturedFunction>*
instantiated_captured_function);
// Determines whether the captured function is stateful.
absl::Status CheckExternalState() const;
// Returns the additional captured inputs that will be passed to the function.
const std::vector<Tensor>& captured_inputs() const {
return captured_inputs_;
}
// Returns the named list of function arguments.
const NameAttrList& func() const { return metadata_->func(); }
// Returns the transitive set of function definition required to instantiate
// this function.
const FunctionLibraryDefinition* lib_def() const {
return metadata_->lib_def();
}
// If every function output corresponds to one of its inputs, the method
// returns the mapping from output indices to input indices. Otherwise, it
// returns an empty list.
const ShortCircuitInfo& short_circuit_info() const {
return metadata_->short_circuit_info();
}
// Indicates whether the function should use inter op parallelism.
bool use_inter_op_parallelism() const {
return metadata_->use_inter_op_parallelism();
}
private:
CapturedFunction(std::shared_ptr<const FunctionMetadata> metadata,
std::vector<Tensor> captured_inputs);
absl::Status IsMultiDevice(FunctionLibraryRuntime* flr,
bool* is_multi_device) const;
const std::shared_ptr<const FunctionMetadata> metadata_;
const std::vector<Tensor> captured_inputs_;
CapturedFunction(const CapturedFunction&) = delete;
void operator=(const CapturedFunction&) = delete;
};
// `InstantiatedCapturedFunction` encapsulates all the runtime support needed
// to execute a tensorflow function.
//
// While `CapturedFunction` encapsulates constant attributes of the function,
// such as its name and captured arguments, `InstantiatedCapturedFunction`
// encapsulates runtime aspects, such as `FunctionLibraryRuntime` and function
// handle.
//
// The `Iterator` related classes use `InstantiatedCapturedFunction` to execute
// functions outside of the normal `OpKernel::Compute()` context.
class InstantiatedCapturedFunction {
public:
// Runs the instantiated captured function. This method takes ownership of
// the tensors in `args`, in order to be able to deallocate them as early as
// possible. Use `RunWithBorrowedArgs()` if the caller needs to retain
// ownership of the `args`.
absl::Status Run(IteratorContext* ctx, std::vector<Tensor>&& args,
std::vector<Tensor>* rets) const;
// Runs the instantiated captured function. This method takes ownership of
// the tensors in `args`, in order to be able to deallocate them as early as
// possible. Use `RunWithBorrowedArgs()` if the caller needs to retain
// ownership of the `args`. Pass non-null `node` to record processing time
// for modeling Iterator's GetNext() resource usage. When non-null node is
// provided, the pre-requisite is that the calling thread has previously
// called `DatasetBaseIterator::RecordStart().
absl::Status Run(IteratorContext* ctx, std::vector<Tensor>&& args,
std::vector<Tensor>* rets,
const std::shared_ptr<model::Node>& node) const;
// Synchronously runs the captured function on the given `args`, and stores
// the results in `*rets`. Prefer to use `Run()` or `RunAsync()` when
// possible.
absl::Status RunWithBorrowedArgs(IteratorContext* ctx,
const std::vector<Tensor>& args,
std::vector<Tensor>* rets) const;
// Synchronously runs the captured function on the given `args`, and stores
// the results in `*rets`. Prefer to use `Run()` or `RunAsync()` when
// possible. Pass non-null `node` to record processing time for modeling
// Iterator's GetNext() resource usage. When non-null node is provided, the
// pre-requisite is that the calling thread has previously called
// `DatasetBaseIterator::RecordStart().
absl::Status RunWithBorrowedArgs(
IteratorContext* ctx, const std::vector<Tensor>& args,
std::vector<Tensor>* rets,
const std::shared_ptr<model::Node>& node) const;
// Synchronously runs the captured function on the given `args`, and stores
// the results in `*rets`. Prefer to use `Run()` or `RunAsync()` when
// possible. This can be useful for calling a captured function in cases where
// an `IteratorContext*` is not available (such as a destructor).
//
// TODO(b/144278100): Avoid running functions without IteratorContext.
absl::Status RunInstantiated(const std::vector<Tensor>& args,
std::vector<Tensor>* rets);
// Asynchronously runs the captured function on the given `args`, stores the
// results in `*rets`, and calls the given `done` callback when the function
// returns. This method takes ownership of the tensors in `args`, in order to
// be able to deallocate them as early as possible. Pass non-null `node` to
// record processing time for modeling Iterator's GetNext() resource usage.
// When non-null node is provided, the pre-requisite is that the calling
// thread has previously called `DatasetBaseIterator::RecordStart().
void RunAsync(IteratorContext* ctx, std::vector<Tensor>&& args,
std::vector<Tensor>* rets,
FunctionLibraryRuntime::DoneCallback done,
const std::shared_ptr<model::Node>& node) const {
RunAsync(*(ctx->runner()), ctx->cancellation_manager(),
ctx->collective_executor(), std::move(args), rets, done, node);
}
// A version of `RunAsync` that does not take an `IteratorContext` but a
// runner, a cancellation manager, and a collective executor.
void RunAsync(std::function<void(std::function<void()>)> runner,
CancellationManager* parent_cancellation_manager,
CollectiveExecutor* collective_executor,
std::vector<Tensor>&& args, std::vector<Tensor>* rets,
FunctionLibraryRuntime::DoneCallback done,
const std::shared_ptr<model::Node>& node) const;
std::string func_name() const { return captured_func_->func().name(); }
private:
friend class CapturedFunction;
InstantiatedCapturedFunction(
FunctionLibraryRuntime* lib, FunctionLibraryRuntime::Handle f_handle,
DataTypeVector ret_types,
std::function<void(std::function<void()>)> runner,
CapturedFunction* captured_func, bool is_multi_device);
// Determines whether a rendezvous object should be created when running the
// instantiated function.
bool ShouldCreateRendezvous() const;
FunctionLibraryRuntime* const lib_; // Not owned.
const FunctionLibraryRuntime::Handle f_handle_;
const DataTypeVector ret_types_;
// Note: We capture the runner at function instantiation time to be able to
// run the function without `IteratorContext` via `RunInstantiated`.
std::function<void(std::function<void()>)> captured_runner_;
CapturedFunction* const captured_func_; // Not owned.
const bool is_multi_device_;
InstantiatedCapturedFunction(const InstantiatedCapturedFunction&) = delete;
void operator=(const InstantiatedCapturedFunction&) = delete;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_CAPTURED_FUNCTION_H_
tensorflow/core/data/compression_utils.cc
+272
View File
@@ -0,0 +1,272 @@
/* Copyright 2020 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/core/data/compression_utils.h"
#include <cstddef>
#include <cstdint>
#include <limits>
#include <string>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/common_runtime/dma_helper.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variant_op_registry.h"
#include "tensorflow/core/platform/snappy.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace data {
namespace {
// Increment this when making changes to the `CompressedElement` proto. The
// `UncompressElement` function will determine what to read according to the
// version.
constexpr int kCompressedElementVersion = 0;
} // namespace
class Iov {
public:
explicit Iov(size_t size) : iov_(size), idx_(0), num_bytes_(0) {}
void Add(void* base, size_t len) {
iov_[idx_].iov_base = base;
iov_[idx_].iov_len = len;
num_bytes_ += len;
++idx_;
}
iovec* Data() { return iov_.data(); }
size_t NumBytes() const { return num_bytes_; }
size_t NumPieces() const { return iov_.size(); }
private:
std::vector<struct iovec> iov_;
size_t idx_;
size_t num_bytes_;
};
absl::Status CompressElement(const std::vector<Tensor>& element,
CompressedElement* out) {
// First pass: preprocess the non`memcpy`able tensors.
size_t num_string_tensors = 0;
size_t num_string_tensor_strings = 0;
std::vector<TensorProto> nonmemcpyable_components;
size_t total_nonmemcpyable_size = 0;
for (const auto& component : element) {
if (component.dtype() == DT_STRING) {
++num_string_tensors;
num_string_tensor_strings += component.NumElements();
} else if (!DataTypeCanUseMemcpy(component.dtype())) {
nonmemcpyable_components.emplace_back();
component.AsProtoTensorContent(&nonmemcpyable_components.back());
total_nonmemcpyable_size +=
nonmemcpyable_components.back().ByteSizeLong();
}
}
// Second pass: build an iov array of the tensor data.
// - `memcpy`able tensors are pointed to directly from a single iovec.
// - String tensors are pointed to directly from multiple iovecs (one for each
// string).
// - All other tensors are serialized and copied into a string (a `tstring`
// for access to `resize_unitialized`).
Iov iov{element.size() + num_string_tensor_strings - num_string_tensors};
tstring nonmemcpyable;
nonmemcpyable.resize_uninitialized(total_nonmemcpyable_size);
char* nonmemcpyable_pos = nonmemcpyable.mdata();
int nonmemcpyable_component_index = 0;
for (int i = 0; i < element.size(); ++i) {
const auto& component = element[i];
CompressedComponentMetadata* metadata =
out->mutable_component_metadata()->Add();
metadata->set_dtype(component.dtype());
component.shape().AsProto(metadata->mutable_tensor_shape());
if (DataTypeCanUseMemcpy(component.dtype())) {
const TensorBuffer* buffer = DMAHelper::buffer(&component);
if (buffer) {
iov.Add(buffer->data(), buffer->size());
metadata->add_uncompressed_bytes(buffer->size());
}
} else if (component.dtype() == DT_STRING) {
const auto& flats = component.unaligned_flat<tstring>();
for (int i = 0; i < flats.size(); ++i) {
iov.Add(const_cast<char*>(flats.data()[i].data()),
flats.data()[i].size());
metadata->add_uncompressed_bytes(flats.data()[i].size());
}
} else {
TensorProto& proto =
nonmemcpyable_components[nonmemcpyable_component_index++];
proto.SerializeToArray(nonmemcpyable_pos, proto.ByteSizeLong());
iov.Add(nonmemcpyable_pos, proto.ByteSizeLong());
nonmemcpyable_pos += proto.ByteSizeLong();
metadata->add_uncompressed_bytes(proto.ByteSizeLong());
}
}
if (iov.NumBytes() > std::numeric_limits<uint32_t>::max()) {
return absl::OutOfRangeError(
absl::StrCat("Encountered dataset element of size ", iov.NumBytes(),
", exceeding the 4GB Snappy limit."));
}
if (!port::Snappy_CompressFromIOVec(iov.Data(), iov.NumBytes(),
out->mutable_data())) {
return absl::InternalError("Failed to compress using snappy.");
}
out->set_version(kCompressedElementVersion);
VLOG(3) << "Compressed element from " << iov.NumBytes() << " bytes to "
<< out->data().size() << " bytes";
return absl::OkStatus();
}
absl::Status UncompressElement(const CompressedElement& compressed,
std::vector<Tensor>* out) {
if (compressed.version() != kCompressedElementVersion) {
return absl::InternalError(absl::StrCat(
"Unsupported compressed element version: ", compressed.version()));
}
int num_components = compressed.component_metadata_size();
out->clear();
out->reserve(num_components);
// First pass: preprocess the non`memcpy`able tensors.
size_t num_string_tensors = 0;
size_t num_string_tensor_strings = 0;
size_t total_nonmemcpyable_size = 0;
for (const auto& metadata : compressed.component_metadata()) {
if (metadata.dtype() == DT_STRING) {
++num_string_tensors;
num_string_tensor_strings += metadata.uncompressed_bytes_size();
} else {
TensorShape shape(metadata.tensor_shape());
int64_t num_elements = shape.num_elements();
if (num_elements > 0 && metadata.uncompressed_bytes_size() == 0) {
return absl::InvalidArgumentError(
"Missing uncompressed_bytes metadata for non-empty tensor");
}
if (!DataTypeCanUseMemcpy(metadata.dtype()) && num_elements > 0) {
total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
}
}
}
// Second pass: prepare the memory to be uncompressed into.
// - `memcpy`able tensors are directly uncompressed into via a single iovec.
// - String tensors are directly uncompressed into via multiple iovecs (one
// for each string).
// - All other tensors are uncompressed into a string (a `tstring` for access
// to `resize_unitialized`).
Iov iov{num_components + num_string_tensor_strings - num_string_tensors};
tstring nonmemcpyable;
nonmemcpyable.resize_uninitialized(total_nonmemcpyable_size);
char* nonmemcpyable_pos = nonmemcpyable.mdata();
for (const auto& metadata : compressed.component_metadata()) {
if (DataTypeCanUseMemcpy(metadata.dtype())) {
TensorShape shape(metadata.tensor_shape());
int64_t num_elements = shape.num_elements();
out->emplace_back(metadata.dtype(), metadata.tensor_shape());
TensorBuffer* buffer = DMAHelper::buffer(&out->back());
if (buffer && num_elements > 0) {
if (metadata.uncompressed_bytes(0) > buffer->size()) {
return absl::InvalidArgumentError(absl::StrCat(
"uncompressed_bytes (", metadata.uncompressed_bytes(0),
") exceeds allocated buffer size (", buffer->size(), ")"));
}
iov.Add(buffer->data(), metadata.uncompressed_bytes(0));
}
} else if (metadata.dtype() == DT_STRING) {
out->emplace_back(metadata.dtype(), metadata.tensor_shape());
const auto& flats = out->back().unaligned_flat<tstring>();
if (metadata.uncompressed_bytes_size() > flats.size()) {
return absl::InvalidArgumentError(absl::StrCat(
"uncompressed_bytes count (", metadata.uncompressed_bytes_size(),
") exceeds allocated string tensor elements count (", flats.size(),
")"));
}
for (int i = 0; i < metadata.uncompressed_bytes_size(); ++i) {
flats.data()[i].resize(metadata.uncompressed_bytes(i));
iov.Add(flats.data()[i].mdata(), metadata.uncompressed_bytes(i));
}
} else {
TensorShape shape(metadata.tensor_shape());
int64_t num_elements = shape.num_elements();
out->emplace_back();
if (num_elements > 0) {
iov.Add(nonmemcpyable_pos, metadata.uncompressed_bytes(0));
nonmemcpyable_pos += metadata.uncompressed_bytes(0);
}
}
}
// Step 2: Uncompress into the iovec.
const std::string& compressed_data = compressed.data();
size_t uncompressed_size;
if (!port::Snappy_GetUncompressedLength(
compressed_data.data(), compressed_data.size(), &uncompressed_size)) {
return absl::InternalError(absl::StrCat(
"Could not get snappy uncompressed length. Compressed data size: ",
compressed_data.size()));
}
if (uncompressed_size != static_cast<size_t>(iov.NumBytes())) {
return absl::InternalError(absl::StrCat(
"Uncompressed size mismatch. Snappy expects ", uncompressed_size,
" whereas the tensor metadata suggests ", iov.NumBytes()));
}
if (!port::Snappy_UncompressToIOVec(compressed_data.data(),
compressed_data.size(), iov.Data(),
iov.NumPieces())) {
return absl::InternalError("Failed to perform snappy decompression.");
}
// Third pass: deserialize nonstring, non`memcpy`able tensors.
nonmemcpyable_pos = nonmemcpyable.mdata();
for (int i = 0; i < num_components; ++i) {
const CompressedComponentMetadata& metadata =
compressed.component_metadata(i);
if (!DataTypeCanUseMemcpy(metadata.dtype()) &&
metadata.dtype() != DT_STRING) {
TensorProto tp;
if (!tp.ParseFromString(
{nonmemcpyable_pos,
static_cast<size_t>(metadata.uncompressed_bytes(0))})) {
return absl::InternalError("Could not parse TensorProto");
}
if (!out->at(i).FromProto(tp)) {
return absl::InternalError("Could not parse Tensor");
}
nonmemcpyable_pos += metadata.uncompressed_bytes(0);
}
}
return absl::OkStatus();
}
REGISTER_UNARY_VARIANT_DECODE_FUNCTION(CompressedElement,
"tensorflow.data.CompressedElement");
} // namespace data
} // namespace tensorflow
tensorflow/core/data/compression_utils.h
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2020 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_CORE_DATA_COMPRESSION_UTILS_H_
#define TENSORFLOW_CORE_DATA_COMPRESSION_UTILS_H_
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace data {
// Compresses the components of `element` into the `CompressedElement` proto.
//
// In addition to writing the actual compressed bytes, `Compress` fills
// out the per-component metadata for the `CompressedElement`.
//
// Returns an error if the uncompressed size of the element exceeds 4GB.
absl::Status CompressElement(const std::vector<Tensor>& element,
CompressedElement* out);
// Uncompresses a `CompressedElement` into a vector of tensor components.
absl::Status UncompressElement(const CompressedElement& compressed,
std::vector<Tensor>* out);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_COMPRESSION_UTILS_H_
tensorflow/core/data/compression_utils_test.cc
+117
View File
@@ -0,0 +1,117 @@
/* Copyright 2020 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/core/data/compression_utils.h"
#include <cstdint>
#include <vector>
#include <gmock/gmock.h>
#include "absl/status/status_matchers.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/dataset_test_base.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::testing::HasSubstr;
using ::tsl::testing::StatusIs;
TEST(CompressionUtilsTest, Exceeds4GB) {
std::vector<Tensor> element = {
CreateTensor<int64_t>(TensorShape{1024, 1024, 513})}; // Just over 4GB.
CompressedElement compressed;
EXPECT_THAT(
CompressElement(element, &compressed),
absl_testing::StatusIs(error::OUT_OF_RANGE,
HasSubstr("exceeding the 4GB Snappy limit")));
}
std::vector<std::vector<Tensor>> TestCases() {
return {
// Single int64.
CreateTensors<int64_t>(TensorShape{1}, {{1}}),
// Multiple int64s.
CreateTensors<int64_t>(TensorShape{1}, {{1}, {2}}),
// Single tstring.
CreateTensors<tstring>(TensorShape{1}, {{"a"}, {"b"}}),
// Multiple tstrings.
{CreateTensor<tstring>(TensorShape{1, 2}, {"abc", "xyz"}),
CreateTensor<tstring>(TensorShape{2, 1}, {"ijk", "mnk"})},
// Mix of tstring and int64.
{CreateTensor<tstring>(TensorShape{1}, {"a"}),
CreateTensor<int64_t>(TensorShape{1}, {1})},
// Empty element.
{},
// Empty tensor.
{CreateTensor<int64_t>(TensorShape{1, 0})},
// Larger int64.
{CreateTensor<int64_t>(TensorShape{128, 128}),
CreateTensor<int64_t>(TensorShape{64, 2})},
// Variants.
{
DatasetOpsTestBase::CreateTestVariantTensor(
{CreateTensor<int64_t>(TensorShape{3, 1}, {1, 2, 3}),
CreateTensor<tstring>(TensorShape{}, {"abc"})}),
DatasetOpsTestBase::CreateTestVariantTensor(
{CreateTensor<int64_t>(TensorShape{3, 1}, {10, 11, 12}),
CreateTensor<tstring>(TensorShape{}, {"xyz"})}),
},
};
}
class ParameterizedCompressionUtilsTest
: public DatasetOpsTestBase,
public ::testing::WithParamInterface<std::vector<Tensor>> {};
TEST_P(ParameterizedCompressionUtilsTest, RoundTrip) {
std::vector<Tensor> element = GetParam();
CompressedElement compressed;
TF_ASSERT_OK(CompressElement(element, &compressed));
std::vector<Tensor> round_trip_element;
TF_ASSERT_OK(UncompressElement(compressed, &round_trip_element));
TF_EXPECT_OK(
ExpectEqual(element, round_trip_element, /*compare_order=*/true));
}
TEST_P(ParameterizedCompressionUtilsTest, CompressedElementVersion) {
std::vector<Tensor> element = GetParam();
CompressedElement compressed;
TF_ASSERT_OK(CompressElement(element, &compressed));
EXPECT_EQ(0, compressed.version());
}
TEST_P(ParameterizedCompressionUtilsTest, VersionMismatch) {
std::vector<Tensor> element = GetParam();
CompressedElement compressed;
TF_ASSERT_OK(CompressElement(element, &compressed));
compressed.set_version(1);
std::vector<Tensor> round_trip_element;
EXPECT_THAT(UncompressElement(compressed, &round_trip_element),
absl_testing::StatusIs(error::INTERNAL));
}
INSTANTIATE_TEST_SUITE_P(Instantiation, ParameterizedCompressionUtilsTest,
::testing::ValuesIn(TestCases()));
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/dataset_test_base.cc
+1263
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/dataset_test_base.h
+1137
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/dataset_utils.cc
+1050
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/dataset_utils.h
+438
View File
@@ -0,0 +1,438 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_DATASET_UTILS_H_
#define TENSORFLOW_CORE_DATA_DATASET_UTILS_H_
#include <atomic>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/resource_handle.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
namespace data {
// Constant used for indicating that the argument of tf.data.Dataset.shard
// should be supplied by the auto-sharding rewrite.
constexpr int kShardHint = -1;
// Creates a resource handle with a unique name for the given resource where
// the resource is managed by the Resource Manager.
template <typename T>
absl::Status CreateWeakHandle(OpKernelContext* ctx, T* resource,
const std::string& container_name,
ResourceHandle* handle) {
static std::atomic<int64_t> resource_id_counter(0);
std::string unique_name =
absl::StrCat(container_name, resource_id_counter.fetch_add(1));
ResourceMgr* mgr = ctx->resource_manager();
TF_RETURN_IF_ERROR(mgr->Create<T>(container_name, unique_name, resource));
*handle = MakeResourceHandle(container_name, unique_name, *ctx->device(),
TypeIndex::Make<T>());
return absl::OkStatus();
}
// Creates a ref-counting resource handle for the given resource, where the
// resource is owned by the handle.
template <typename T>
absl::Status CreateHandle(OpKernelContext* ctx, T* resource,
ResourceHandle* handle) {
ResourceMgr* mgr = ctx->resource_manager();
*handle =
ResourceHandle::MakeRefCountingHandle(resource, ctx->device()->name());
TF_RETURN_IF_ERROR(
mgr->CreateUnowned<T>(handle->container(), handle->name(), resource));
return absl::OkStatus();
}
// TODO(b/198162355): Merge this class with ResourceOpKernel.
template <typename T>
class AnonymousResourceOp : public OpKernel {
public:
// Creates an AnonymousResourceOp.
// ref_counting: Determines if the Op returns a ref-counting ResourceHandle.
// ResourceHandle. See go/tf-resource-handle-ref-count.
// return_deleter: Determines if the Op outputs a deleter tensor in addition
// to the resource handle tensor.
// If the resource handle is ref-counting, a no-op deleter is returned.
explicit AnonymousResourceOp(OpKernelConstruction* context, bool ref_counting,
bool return_deleter)
: OpKernel(context),
ref_counting_(ref_counting),
return_deleter_(return_deleter) {}
void Compute(OpKernelContext* ctx) override {
FunctionLibraryRuntime* lib;
std::unique_ptr<FunctionLibraryDefinition> flib_def(nullptr);
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr(nullptr);
OP_REQUIRES_OK(
ctx, ctx->function_library()->Clone(&flib_def, &pflr, &lib, true));
T* resource;
OP_REQUIRES_OK(ctx, CreateResource(ctx, std::move(flib_def),
std::move(pflr), lib, &resource));
ResourceHandle handle;
if (ref_counting_) {
OP_REQUIRES_OK(ctx, CreateHandle(ctx, resource, &handle));
} else {
OP_REQUIRES_OK(ctx, CreateWeakHandle(ctx, resource, name(), &handle));
}
Tensor* handle_t;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &handle_t));
handle_t->scalar<ResourceHandle>()() = handle;
if (return_deleter_) {
Tensor* deleter_t;
AllocatorAttributes attr;
attr.set_on_host(true);
OP_REQUIRES_OK(
ctx, ctx->allocate_output(1, TensorShape({}), &deleter_t, attr));
// TODO(feyu): Consider returning an OptionalVariant.
if (!ref_counting_) {
// A deleter output that deletes the resource when destroyed.
deleter_t->scalar<Variant>()() =
ResourceDeleter(handle, ctx->resource_manager());
}
}
}
protected:
virtual std::string name() = 0;
virtual absl::Status CreateResource(
OpKernelContext* ctx, std::unique_ptr<FunctionLibraryDefinition> flib_def,
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr,
FunctionLibraryRuntime* lib, T** resource) = 0;
private:
const bool ref_counting_;
const bool return_deleter_;
};
// Returns OkStatus() if `expected` and `received` types match,
// errors::InvalidArgument otherwise.
absl::Status VerifyTypesMatch(const DataTypeVector& expected,
const DataTypeVector& received);
absl::Status VerifyTypesMatch(const DataTypeVector& expected,
const std::vector<Tensor>& received);
// Returns OkStatus() if `expected` and `received` shapes are compatible,
// errors::InvalidArgument otherwise.
absl::Status VerifyShapesCompatible(
const std::vector<PartialTensorShape>& expected,
const std::vector<PartialTensorShape>& received);
absl::Status VerifyShapesCompatible(
const std::vector<PartialTensorShape>& expected,
const std::vector<Tensor>& received);
// Dataset op level determinism policy.
class DeterminismPolicy {
public:
enum class Type : int {
// The op must produce elements deterministically.
kDeterministic,
// The op may relax determinism to improve performance.
kNondeterministic,
// The determinism policy is not specified at the op level. In this case we
// use the experimental_deterministic dataset option to determine the
// determinism policy.
kDefault,
};
static constexpr const char* const kDeterministic = "true";
static constexpr const char* const kNondeterministic = "false";
static constexpr const char* const kDefault = "default";
DeterminismPolicy() : determinism_(Type::kDefault) {}
explicit DeterminismPolicy(Type determinism) : determinism_(determinism) {}
// Creates a DeterminismPolicy with Type kDeterministic or
// kNondeterministic, depending on the values of `is_deterministic`.
explicit DeterminismPolicy(bool is_deterministic);
static absl::Status FromString(const std::string& s, DeterminismPolicy* out);
// Returns the string representing the determinism policy. This will be one of
// the string constants defined above.
std::string String() const;
/// Convenience methods for checking the DeterminismPolicy::Type.
bool IsDeterministic() const { return determinism_ == Type::kDeterministic; }
bool IsNondeterministic() const {
return determinism_ == Type::kNondeterministic;
}
bool IsDefault() const { return determinism_ == Type::kDefault; }
private:
Type determinism_;
};
// Resolves non-deterministic seeds if necessary, returning either the original
// seeds or the resolved seeds.
//
// By TensorFlow convention, if both seeds are 0, they should be replaced with
// non-deterministically chosen seeds.
std::pair<int64_t, int64_t> MaybeOverrideSeeds(
std::pair<int64_t, int64_t> seeds);
// Adds the functions in `to_add` to `base`. If a function with a matching
// signature already exists in `base`, replaces it with the function from
// `to_add`.
absl::Status AddToFunctionLibrary(FunctionLibraryDefinition* base,
const FunctionLibraryDefinition& to_add);
absl::Status AddToFunctionLibrary(FunctionLibraryDefinition* base,
const FunctionDefLibrary& to_add);
// Determines whether the given function is stateful.
absl::Status IsFunctionStateful(const FunctionLibraryDefinition& library,
const FunctionDef& function_def);
// Determines whether the given node is stateful.
absl::Status IsNodeStateful(const FunctionLibraryDefinition& library,
const NodeDef& node);
// Creates a runner that runs functions with limited parallelism.
std::function<void(std::function<void()>)> RunnerWithMaxParallelism(
std::function<void(std::function<void()>)> runner, int max_parallelism);
// Op for creating a typed dummy resource.
//
// This op is used to provide a resource "placeholder" for ops such as
// `CacheDatasetV2` or `ShuffleDatasetV2` that expects a resource input.
// Originally, the lifetime of the resources passed into these ops was managed
// externally. After the implementation changed to manage the lifetime of the
// resources (including creation) by the ops themselves, the resource input is
// only needed to pass a resource handle through graph rewrites. When they are
// invoked from user code, the implementation passes in a dummy resource.
template <typename ResourceType>
class DummyResourceOp : public OpKernel {
public:
explicit DummyResourceOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
void Compute(OpKernelContext* ctx) override {
Tensor* tensor;
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &tensor));
tensor->scalar<ResourceHandle>()() = MakeResourceHandle<ResourceType>(
ctx, /*container=*/"", /*name=*/"dummy_resource");
}
};
// Given an op prefix and an op to match, returns whether the op to match
// is a match for any version of the op prefix. For example,
// MatchesAnyVersion("BatchDataset", "BatchDataset") == true
// MatchesAnyVersion("BatchDataset", "BatchDatasetV2") == true
// MatchesAnyVersion("BatchDataset", "BatchDatasetV3") == true
// MatchesAnyVersion("PaddedBatchDataset", "BatchDataset") == false
bool MatchesAnyVersion(absl::string_view op_prefix,
absl::string_view op_to_match);
// Returns the index-th slice of a given tensor. If the index-th slice of
// the tensor is not aligned, returns a deep copy of the tensor.
Tensor MaybeCopySubSlice(const Tensor& tensor, int64_t index);
// Removes device placements from the ops of all functions in `library`.
void StripDevicePlacement(FunctionDefLibrary* library);
// Copies partial of the batch output.
absl::Status CopyPartialBatch(int64_t num_elements, const Tensor& value,
Tensor* output);
// Reads a batch when restoring the iterator.
absl::Status ReadBatch(IteratorContext* ctx, IteratorStateReader* reader,
int64_t batch_size, const std::string& iterator_prefix,
const std::string& batch_prefix,
std::vector<Tensor>* batch);
// Writes a batch when saving the iterator.
absl::Status WriteBatch(int64_t batch_size, int64_t num_elements,
const std::string& iterator_prefix,
const std::string& batch_prefix,
IteratorStateWriter* writer,
std::vector<Tensor>* batch);
// Reads a status when restoring the iterator.
absl::Status ReadStatus(const std::string& iterator_prefix,
const std::string& prefix, IteratorStateReader* reader,
absl::Status* status);
// Writes a status when saving the iterator.
absl::Status WriteStatus(const std::string& iterator_prefix,
const std::string& prefix, const absl::Status& status,
IteratorStateWriter* writer);
// Processes a batch to output. In the case a partial batch is encountered, copy
// only partial of the batch.
absl::Status ProcessBatch(int64_t batch_size, int64_t num_elements,
bool drop_remainder, const absl::Status& status,
IteratorContext* ctx, std::vector<Tensor>* output,
bool* end_of_sequence, std::vector<Tensor>* batch);
// Copies the input elements to a batch.
//
// The `batch_elements` argument contains the individual elements to copy into a
// batch. The `parallel_copy` argument indicates whether to parallelize the
// copy.
// The `out_tensors` argument will be used to store the resulting batch (one for
// each component of the input).
absl::Status CopyBatch(AnyContext ctx,
std::vector<std::vector<Tensor>>&& batch_elements,
bool parallel_copy, std::vector<Tensor>* out_tensors);
// Computes the set of experiments to apply based on the job name, task id,
// rollout percentage of registered experiments, and the
// TF_DATA_EXPERIMENT_OPT_IN and TF_DATA_EXPERIMENT_OPT_OUT environment
// variables.
absl::flat_hash_set<std::string> GetExperiments();
absl::flat_hash_set<std::string> GetExperiments(
const std::string& job_name, int64_t task_id,
std::function<uint64_t(const std::string&)> hash_func);
// Logs and records the experiments that will be applied.
void LogAndRecordExperiments(
const absl::flat_hash_set<std::string>& experiments);
// Computes the set of enabled, disabled, and default optimizations based on the
// given options. An optimization must be a graph optimizer name that has been
// registered with Grappler.
void GetOptimizations(const Options& options,
absl::flat_hash_set<tstring>* optimizations_enabled,
absl::flat_hash_set<tstring>* optimizations_disabled,
absl::flat_hash_set<tstring>* optimizations_default);
// Creates graph rewrite configs based on the given options. The configs will
// only be used if their corresponding optimizers registered with Grappler are
// enabled.
// A config is a string with the following format:
// <optimizer name>:<attribute name>:<attribute value>
absl::flat_hash_set<tstring> CreateGraphRewriteConfigs(const Options& options);
// Determines whether max intra-op parallelism should be configured.
bool ShouldConfigureMaxIntraOpParallelism(const Options& options);
// Determines whether private threadpool should be used.
bool ShouldUsePrivateThreadPool(const Options& options);
// Determines whether autotuning should be used.
bool ShouldUseAutotuning(const Options& options);
// Determines whether optimizations should be applied.
bool ShouldApplyOptimizations(
const Options& options,
const absl::flat_hash_set<tstring>& optimizations_enabled,
const absl::flat_hash_set<tstring>& optimizations_default);
// Returns the default CPU budget.
inline int GetCpuBudget() {
static bool in_experiment = GetExperiments().contains("tune_cpu_budget");
return (in_experiment ? 1.2 : 1.0) * port::NumSchedulableCPUs();
}
// Returns the initial value for parallelism parameter before the first Autotune
// optimization.
int64_t GetAutotuneDefaultParallelism(IteratorContext* ctx);
// Returns the minimum parallelism value for Autotune optimization.
int64_t GetAutotuneMinParallelism(IteratorContext* ctx);
// Creates an iterator context appropriate for a nested dataset's iterator. A
// nested dataset is a dataset created within another dataset, e.g. by the
// function passed to `interleave` or `flat_map`.
IteratorContext MakeNestedIteratorContext(IteratorContext* ctx);
// A `DatasetExperimentRegistry::JobSelector` that randomly selects
// `rollout_pct` percent of all jobs. `name_hash` is a hash of the experiment
// and job names.
template <int64_t rollout_pct>
bool RandomJobSamplePercentage(uint64_t name_hash) {
return name_hash % 100 < rollout_pct;
}
// A `DatasetExperimentRegistry::TaskSelector` that selects all tasks.
bool AllTasks(int64_t unused_task_id, bool unused_evens);
// A `DatasetExperimentRegistry::TaskSelector` that selects the tasks for half
// of all hosts. Typically, one or two consecutive tasks run on a single host.
// If `evens` is `true`, selects tasks 0,1,4,5,8,9,..., otherwise selects tasks
// 2,3,6,7,10,11,...
bool IndependentHostTasks(int64_t task_id, bool evens);
// Registry of tf.data experiments.
class DatasetExperimentRegistry {
public:
using JobSelector = std::function<bool(uint64_t name_hash)>;
using TaskSelector = std::function<bool(int64_t task_id, bool evens)>;
struct ExperimentSelector {
JobSelector job_selector;
TaskSelector task_selector;
};
// Registers the experiment.
static void Register(const std::string& experiment, JobSelector job_selector,
TaskSelector task_selector);
// Returns all registered experiments.
static absl::flat_hash_map<std::string, ExperimentSelector> Experiments();
};
// Helper class to register a dataset experiment.
class DatasetExperimentRegistrar {
public:
explicit DatasetExperimentRegistrar(
const std::string& experiment,
DatasetExperimentRegistry::JobSelector job_selector,
DatasetExperimentRegistry::TaskSelector task_selector) {
DatasetExperimentRegistry::Register(experiment, job_selector,
task_selector);
}
};
// Macro that can be used to register a dataset experiment.
#define REGISTER_DATASET_EXPERIMENT(experiment, job_selector, task_selector) \
REGISTER_DATASET_OP_NAME_UNIQ_HELPER(__COUNTER__, experiment, job_selector, \
task_selector)
#define REGISTER_DATASET_OP_NAME_UNIQ_HELPER(ctr, experiment, job_selector, \
task_selector) \
REGISTER_DATASET_OP_NAME_UNIQ(ctr, experiment, job_selector, task_selector)
#define REGISTER_DATASET_OP_NAME_UNIQ(ctr, experiment, job_selector, \
task_selector) \
static ::tensorflow::data::DatasetExperimentRegistrar \
registrar__body__##ctr##__object(experiment, job_selector, \
task_selector)
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_DATASET_UTILS_H_
tensorflow/core/data/dataset_utils_test.cc
+788
View File
@@ -0,0 +1,788 @@
/* 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/core/data/dataset_utils.h"
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_join.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "xla/tsl/util/determinism_test_util.h"
#include "tensorflow/core/data/compression_utils.h"
#include "tensorflow/core/data/dataset_test_base.h"
#include "tensorflow/core/data/serialization_utils.h"
#include "tensorflow/core/data/test_utils.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/util/work_sharder.h"
namespace tensorflow {
namespace data {
namespace {
using ::testing::HasSubstr;
using ::tsl::testing::StatusIs;
TEST(DatasetUtilsTest, MatchesAnyVersion) {
EXPECT_TRUE(MatchesAnyVersion("BatchDataset", "BatchDataset"));
EXPECT_TRUE(MatchesAnyVersion("BatchDataset", "BatchDatasetV2"));
EXPECT_TRUE(MatchesAnyVersion("BatchDataset", "BatchDatasetV3"));
EXPECT_FALSE(MatchesAnyVersion("BatchDataset", "BatchDatasetXV3"));
EXPECT_FALSE(MatchesAnyVersion("BatchDataset", "BatchV2Dataset"));
EXPECT_FALSE(MatchesAnyVersion("BatchDataset", "PaddedBatchDataset"));
}
TEST(DatasetUtilsTest, AddToFunctionLibrary) {
auto make_fn_a = [](const std::string& fn_name) {
return FunctionDefHelper::Create(
/*function_name=*/fn_name,
/*in_def=*/{"arg: int64"},
/*out_def=*/{"ret: int64"},
/*attr_def=*/{},
/*node_def=*/{{{"node"}, "Identity", {"arg"}, {{"T", DT_INT64}}}},
/*ret_def=*/{{"ret", "node:output:0"}});
};
auto make_fn_b = [](const std::string& fn_name) {
return FunctionDefHelper::Create(
/*function_name=*/fn_name,
/*in_def=*/{"arg: int64"},
/*out_def=*/{"ret: int64"},
/*attr_def=*/{},
/*node_def=*/
{{{"node"}, "Identity", {"arg"}, {{"T", DT_INT64}}},
{{"node2"}, "Identity", {"node:output:0"}, {{"T", DT_INT64}}}},
/*ret_def=*/{{"ret", "node2:output:0"}});
};
FunctionDefLibrary fdef_base;
*fdef_base.add_function() = make_fn_a("0");
*fdef_base.add_function() = make_fn_a("1");
*fdef_base.add_function() = make_fn_a("2");
FunctionDefLibrary fdef_to_add;
*fdef_to_add.add_function() = make_fn_b("0"); // Override
*fdef_to_add.add_function() = make_fn_a("1"); // Do nothing
*fdef_to_add.add_function() = make_fn_b("3"); // Add new function
FunctionLibraryDefinition flib_0(OpRegistry::Global(), fdef_base);
TF_ASSERT_OK(AddToFunctionLibrary(&flib_0, fdef_to_add));
FunctionLibraryDefinition flib_1(OpRegistry::Global(), fdef_base);
FunctionLibraryDefinition flib_to_add(OpRegistry::Global(), fdef_to_add);
TF_ASSERT_OK(AddToFunctionLibrary(&flib_1, flib_to_add));
for (const auto& flib : {flib_0, flib_1}) {
EXPECT_TRUE(FunctionDefsEqual(*flib.Find("0"), make_fn_b("0")));
EXPECT_TRUE(FunctionDefsEqual(*flib.Find("1"), make_fn_a("1")));
EXPECT_TRUE(FunctionDefsEqual(*flib.Find("2"), make_fn_a("2")));
EXPECT_TRUE(FunctionDefsEqual(*flib.Find("3"), make_fn_b("3")));
}
}
TEST(DatasetUtilsTest, AddToFunctionLibraryWithConflictingSignatures) {
FunctionDefLibrary fdef_base;
*fdef_base.add_function() = FunctionDefHelper::Create(
/*function_name=*/"0",
/*in_def=*/{"arg: int64"},
/*out_def=*/{"ret: int64"},
/*attr_def=*/{},
/*node_def=*/{},
/*ret_def=*/{{"ret", "arg"}});
FunctionDefLibrary fdef_to_add;
*fdef_to_add.add_function() = FunctionDefHelper::Create(
/*function_name=*/"0",
/*in_def=*/{"arg: int64"},
/*out_def=*/{"ret: int64", "ret2: int64"},
/*attr_def=*/{},
/*node_def=*/{},
/*ret_def=*/{{"ret", "arg"}, {"ret2", "arg"}});
FunctionLibraryDefinition flib_0(OpRegistry::Global(), fdef_base);
absl::Status s = AddToFunctionLibrary(&flib_0, fdef_to_add);
EXPECT_EQ(error::Code::INVALID_ARGUMENT, s.code());
EXPECT_EQ(
"Cannot add function '0' because a different function with the same "
"signature already exists.",
s.message());
FunctionLibraryDefinition flib_1(OpRegistry::Global(), fdef_base);
FunctionLibraryDefinition flib_to_add(OpRegistry::Global(), fdef_to_add);
s = AddToFunctionLibrary(&flib_1, flib_to_add);
EXPECT_EQ(error::Code::INVALID_ARGUMENT, s.code());
EXPECT_EQ(
"Cannot add function '0' because a different function with the same "
"signature already exists.",
s.message());
}
TEST(DatasetUtilsTest, StripDevicePlacement) {
FunctionDefLibrary flib;
*flib.add_function() = FunctionDefHelper::Create(
/*function_name=*/"0",
/*in_def=*/{"arg: int64"},
/*out_def=*/{"ret: int64"},
/*attr_def=*/{},
/*node_def=*/
{{{"node"},
"Identity",
{"arg"},
{{"T", DT_INT64}},
/*dep=*/{},
/*device=*/"device:CPU:0"}},
/*ret_def=*/{{"ret", "arg"}});
EXPECT_EQ(flib.function(0).node_def(0).device(), "device:CPU:0");
StripDevicePlacement(&flib);
EXPECT_EQ(flib.function(0).node_def(0).device(), "");
}
TEST(DatasetUtilsTest, RunnerWithMaxParallelism) {
auto runner =
RunnerWithMaxParallelism([](const std::function<void()> fn) { fn(); }, 2);
auto fn = []() { ASSERT_EQ(GetPerThreadMaxParallelism(), 2); };
runner(fn);
}
TEST(DatasetUtilsTest, ParseDeterminismPolicy) {
DeterminismPolicy determinism;
TF_ASSERT_OK(DeterminismPolicy::FromString("true", &determinism));
EXPECT_TRUE(determinism.IsDeterministic());
TF_ASSERT_OK(DeterminismPolicy::FromString("false", &determinism));
EXPECT_TRUE(determinism.IsNondeterministic());
TF_ASSERT_OK(DeterminismPolicy::FromString("default", &determinism));
EXPECT_TRUE(determinism.IsDefault());
}
TEST(DatasetUtilsTest, DeterminismString) {
for (auto s : {"true", "false", "default"}) {
DeterminismPolicy determinism;
TF_ASSERT_OK(DeterminismPolicy::FromString(s, &determinism));
EXPECT_TRUE(s == determinism.String());
}
}
TEST(DatasetUtilsTest, BoolConstructor) {
EXPECT_TRUE(DeterminismPolicy(true).IsDeterministic());
EXPECT_FALSE(DeterminismPolicy(true).IsNondeterministic());
EXPECT_FALSE(DeterminismPolicy(true).IsDefault());
EXPECT_TRUE(DeterminismPolicy(false).IsNondeterministic());
EXPECT_FALSE(DeterminismPolicy(false).IsDeterministic());
EXPECT_FALSE(DeterminismPolicy(false).IsDefault());
}
class TestSplitProvider : public SplitProvider {
public:
absl::Status GetNext(Tensor* split, bool* end_of_splits) override {
return absl::OkStatus();
}
absl::Status Reset() override { return absl::OkStatus(); }
absl::Status Save(std::function<std::string(std::string)> key_name_fn,
IteratorStateWriter* writer) override {
return absl::OkStatus();
}
absl::Status Restore(std::function<std::string(std::string)> key_name_fn,
IteratorStateReader* reader) override {
return absl::OkStatus();
}
};
TEST(DatasetUtilsTest, MakeNestedIteratorContext) {
TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestContext> test_ctx,
TestContext::Create());
IteratorContext::Params params(test_ctx->op_ctx());
params.split_providers.push_back(std::make_unique<TestSplitProvider>());
IteratorContext iter_ctx(params);
IteratorContext nested_ctx = MakeNestedIteratorContext(&iter_ctx);
EXPECT_FALSE(iter_ctx.split_providers().empty());
EXPECT_TRUE(nested_ctx.split_providers().empty());
}
REGISTER_DATASET_EXPERIMENT("test_only_experiment_0",
RandomJobSamplePercentage<0>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_experiment_1",
RandomJobSamplePercentage<1>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_experiment_5",
RandomJobSamplePercentage<5>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_experiment_10",
RandomJobSamplePercentage<10>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_experiment_50",
RandomJobSamplePercentage<50>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_experiment_99",
RandomJobSamplePercentage<99>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_experiment_100",
RandomJobSamplePercentage<100>, AllTasks);
REGISTER_DATASET_EXPERIMENT("test_only_task_experiment_100",
RandomJobSamplePercentage<100>,
IndependentHostTasks);
struct GetExperimentsHashTestCase {
uint64_t hash;
std::vector<std::string> expected_in;
std::vector<std::string> expected_out;
};
class GetExperimentsHashTest
: public ::testing::TestWithParam<GetExperimentsHashTestCase> {};
TEST_P(GetExperimentsHashTest, DatasetUtils) {
const GetExperimentsHashTestCase test_case = GetParam();
uint64_t hash_result = test_case.hash;
const std::string job_name = "job";
const int64_t task_id = 0;
auto hash_func = [hash_result](const std::string& str) {
return hash_result;
};
auto experiments = GetExperiments(job_name, task_id, hash_func);
absl::flat_hash_set<std::string> experiment_set(experiments.begin(),
experiments.end());
for (const auto& experiment : test_case.expected_in) {
EXPECT_TRUE(experiment_set.find(experiment) != experiment_set.end())
<< "experiment=" << experiment << " hash=" << hash_result;
}
for (const auto& experiment : test_case.expected_out) {
EXPECT_TRUE(experiment_set.find(experiment) == experiment_set.end())
<< "experiment=" << experiment << " hash=" << hash_result;
}
}
INSTANTIATE_TEST_SUITE_P(
Test, GetExperimentsHashTest,
::testing::Values<GetExperimentsHashTestCase>(
GetExperimentsHashTestCase{
/*hash=*/0,
/*expected_in=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/{"test_only_experiment_0"},
},
GetExperimentsHashTestCase{
/*hash=*/5,
/*expected_in=*/
{"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/
{
"test_only_experiment_0",
"test_only_experiment_1",
"test_only_experiment_5",
},
},
GetExperimentsHashTestCase{
/*hash=*/95,
/*expected_in=*/
{"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50"},
},
GetExperimentsHashTestCase{
/*hash=*/99,
/*expected_in=*/{"test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99"},
},
GetExperimentsHashTestCase{
/*hash=*/100,
/*expected_in=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/{"test_only_experiment_0"},
},
GetExperimentsHashTestCase{
/*hash=*/105,
/*expected_in=*/
{"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/
{
"test_only_experiment_0",
"test_only_experiment_1",
"test_only_experiment_5",
},
},
GetExperimentsHashTestCase{
/*hash=*/195,
/*expected_in=*/
{"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50"},
}));
struct GetExperimentsOptTestCase {
std::vector<std::string> opt_ins;
std::vector<std::string> opt_outs;
std::vector<std::string> expected_in;
std::vector<std::string> expected_out;
};
class GetExperimentsOptTest
: public ::testing::TestWithParam<GetExperimentsOptTestCase> {};
TEST_P(GetExperimentsOptTest, DatasetUtils) {
const GetExperimentsOptTestCase test_case = GetParam();
auto opt_ins = test_case.opt_ins;
auto opt_outs = test_case.opt_outs;
if (!opt_ins.empty()) {
setenv("TF_DATA_EXPERIMENT_OPT_IN", absl::StrJoin(opt_ins, ",").c_str(), 1);
}
if (!opt_outs.empty()) {
setenv("TF_DATA_EXPERIMENT_OPT_OUT", absl::StrJoin(opt_outs, ",").c_str(),
1);
}
const std::string job_name = "job";
const int64_t task_id = 0;
auto hash_func = [](const std::string& str) { return 0; };
auto experiments = GetExperiments(job_name, task_id, hash_func);
absl::flat_hash_set<std::string> experiment_set(experiments.begin(),
experiments.end());
for (const auto& experiment : test_case.expected_in) {
EXPECT_TRUE(experiment_set.find(experiment) != experiment_set.end())
<< "experiment=" << experiment << " opt_ins={"
<< absl::StrJoin(opt_ins, ",") << "} opt_outs={"
<< absl::StrJoin(opt_outs, ",") << "}";
}
for (const auto& experiment : test_case.expected_out) {
EXPECT_TRUE(experiment_set.find(experiment) == experiment_set.end())
<< "experiment=" << experiment << " opt_ins={"
<< absl::StrJoin(opt_ins, ",") << "} opt_outs={"
<< absl::StrJoin(opt_outs, ",") << "}";
}
if (!opt_ins.empty()) {
unsetenv("TF_DATA_EXPERIMENT_OPT_IN");
}
if (!opt_outs.empty()) {
unsetenv("TF_DATA_EXPERIMENT_OPT_OUT");
}
}
INSTANTIATE_TEST_SUITE_P(
Test, GetExperimentsOptTest,
::testing::Values<GetExperimentsOptTestCase>(
GetExperimentsOptTestCase{
/*opt_ins=*/{"all"},
/*opt_outs=*/{"all"},
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{"all"},
/*opt_outs=*/{},
/*expected_in=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100"},
/*expected_out=*/{}},
GetExperimentsOptTestCase{
/*opt_ins=*/{"all"},
/*opt_outs=*/{"test_only_experiment_1", "test_only_experiment_99"},
/*expected_in=*/
{"test_only_experiment_0", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_1", "test_only_experiment_99"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{},
/*opt_outs=*/{"all"},
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{},
/*opt_outs=*/{},
/*expected_in=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/{"test_only_experiment_0"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{},
/*opt_outs=*/{"test_only_experiment_1", "test_only_experiment_99"},
/*expected_in=*/
{"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_99"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{"test_only_experiment_0", "test_only_experiment_100"},
/*opt_outs=*/{"all"},
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{"test_only_experiment_0", "test_only_experiment_100"},
/*opt_outs=*/{"all_except_opt_in"},
/*expected_in=*/
{"test_only_experiment_0", "test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99"}},
GetExperimentsOptTestCase{
/*opt_ins=*/{"test_only_experiment_0", "test_only_experiment_100"},
/*opt_outs=*/{},
/*expected_in=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100"},
/*expected_out=*/{}},
GetExperimentsOptTestCase{
/*opt_ins=*/{"test_only_experiment_0", "test_only_experiment_100"},
/*opt_outs=*/{"test_only_experiment_1", "test_only_experiment_99"},
/*expected_in=*/
{"test_only_experiment_0", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_1", "test_only_experiment_99"}}));
struct GetExperimentsJobNameTestCase {
uint64_t hash;
std::string job_name;
int64_t task_id;
std::vector<std::string> expected_in;
std::vector<std::string> expected_out;
};
class GetExperimentsJobNameTest
: public ::testing::TestWithParam<GetExperimentsJobNameTestCase> {};
TEST_P(GetExperimentsJobNameTest, DatasetUtils) {
const GetExperimentsJobNameTestCase test_case = GetParam();
auto job_name = test_case.job_name;
auto task_id = test_case.task_id;
uint64_t hash_result = test_case.hash;
auto hash_func = [hash_result](const std::string& str) {
return hash_result;
};
auto experiments = GetExperiments(job_name, task_id, hash_func);
absl::flat_hash_set<std::string> experiment_set(experiments.begin(),
experiments.end());
for (const auto& experiment : test_case.expected_in) {
EXPECT_TRUE(experiment_set.find(experiment) != experiment_set.end())
<< "experiment=" << experiment << " job_name=" << job_name;
}
for (const auto& experiment : test_case.expected_out) {
EXPECT_TRUE(experiment_set.find(experiment) == experiment_set.end())
<< "experiment=" << experiment << " job_name=" << job_name;
}
}
// Note: These tests use (deterministic) randomness. The behavior is correct but
// this approach is generally frowned upon (see go/python-tips/048).
INSTANTIATE_TEST_SUITE_P(
Test, GetExperimentsJobNameTest,
::testing::Values(
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"",
/*task_id=*/0,
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100", "test_only_task_experiment_100"}},
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"",
/*task_id=*/-1,
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100", "test_only_task_experiment_100"}},
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"",
/*task_id=*/2,
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100", "test_only_task_experiment_100"}},
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"job_name",
/*task_id=*/-1,
/*expected_in=*/{},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_experiment_99",
"test_only_experiment_100", "test_only_task_experiment_100"}},
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"job_name",
/*task_id=*/0,
/*expected_in=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100",
"test_only_task_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0"}},
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"job_name",
/*task_id=*/1,
/*expected_in=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100",
"test_only_task_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0"}},
GetExperimentsJobNameTestCase{
/*hash=*/0,
/*job_name=*/"job_name",
/*task_id=*/2,
/*expected_in=*/
{"test_only_experiment_1", "test_only_experiment_5",
"test_only_experiment_10", "test_only_experiment_50",
"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_task_experiment_100"}},
GetExperimentsJobNameTestCase{
/*hash=*/95,
/*job_name=*/"job_name",
/*task_id=*/1,
/*expected_in=*/
{"test_only_experiment_99", "test_only_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50", "test_only_task_experiment_100"}},
GetExperimentsJobNameTestCase{
/*hash=*/95,
/*job_name=*/"job_name",
/*task_id=*/2,
/*expected_in=*/
{"test_only_experiment_99", "test_only_experiment_100",
"test_only_task_experiment_100"},
/*expected_out=*/
{"test_only_experiment_0", "test_only_experiment_1",
"test_only_experiment_5", "test_only_experiment_10",
"test_only_experiment_50"}}));
struct GetOptimizationsTestCase {
Options options;
std::vector<std::string> expected_enabled;
std::vector<std::string> expected_disabled;
std::vector<std::string> expected_default;
};
// Tests the default.
GetOptimizationsTestCase GetOptimizationTestCase1() {
return {
/*options=*/Options(),
/*expected_enabled=*/{},
/*expected_disabled=*/{},
/*expected_default=*/
{"noop_elimination", "map_and_batch_fusion", "shuffle_and_repeat_fusion",
"map_parallelization", "parallel_batch", "inject_prefetch"}};
}
// Tests disabling application of default optimizations.
GetOptimizationsTestCase GetOptimizationTestCase2() {
Options options;
options.mutable_optimization_options()->set_apply_default_optimizations(
false);
return {options, /*expected_enabled=*/{}, /*expected_disabled=*/{},
/*expected_default=*/{}};
}
// Tests explicitly enabling / disabling some default and non-default
// optimizations.
GetOptimizationsTestCase GetOptimizationTestCase3() {
Options options;
options.set_deterministic(false);
options.mutable_optimization_options()->set_map_and_batch_fusion(true);
options.mutable_optimization_options()->set_map_parallelization(false);
options.mutable_optimization_options()->set_parallel_batch(false);
return {options,
/*expected_enabled=*/{"make_sloppy", "map_and_batch_fusion"},
/*expected_disabled=*/{"parallel_batch", "map_parallelization"},
/*expected_default=*/
{"noop_elimination", "shuffle_and_repeat_fusion", "inject_prefetch"}};
}
// Test enabling all / most available optimizations.
GetOptimizationsTestCase GetOptimizationTestCase4() {
Options options;
options.set_deterministic(false);
options.mutable_optimization_options()->set_filter_fusion(true);
options.mutable_optimization_options()->set_filter_parallelization(true);
options.mutable_optimization_options()->set_map_and_batch_fusion(true);
options.mutable_optimization_options()->set_map_and_filter_fusion(true);
options.mutable_optimization_options()->set_map_fusion(true);
options.mutable_optimization_options()->set_map_parallelization(true);
options.mutable_optimization_options()->set_noop_elimination(true);
options.mutable_optimization_options()->set_parallel_batch(true);
options.mutable_optimization_options()->set_shuffle_and_repeat_fusion(true);
options.mutable_optimization_options()->set_inject_prefetch(true);
options.mutable_optimization_options()->set_seq_interleave_prefetch(true);
options.set_slack(true);
return {options,
/*expected_enabled=*/
{"filter_fusion", "filter_parallelization", "make_sloppy",
"map_and_batch_fusion", "map_and_filter_fusion", "map_fusion",
"map_parallelization", "noop_elimination", "parallel_batch",
"shuffle_and_repeat_fusion", "slack", "inject_prefetch",
"seq_interleave_prefetch"},
/*expected_disabled=*/{},
/*expected_default=*/{}};
}
class GetOptimizationsTest
: public ::testing::TestWithParam<GetOptimizationsTestCase> {};
TEST_P(GetOptimizationsTest, DatasetUtils) {
const GetOptimizationsTestCase test_case = GetParam();
auto options = test_case.options;
absl::flat_hash_set<tstring> actual_enabled, actual_disabled, actual_default;
GetOptimizations(options, &actual_enabled, &actual_disabled, &actual_default);
EXPECT_THAT(
std::vector<std::string>(actual_enabled.begin(), actual_enabled.end()),
::testing::UnorderedElementsAreArray(test_case.expected_enabled));
EXPECT_THAT(
std::vector<std::string>(actual_disabled.begin(), actual_disabled.end()),
::testing::UnorderedElementsAreArray(test_case.expected_disabled));
EXPECT_THAT(
std::vector<std::string>(actual_default.begin(), actual_default.end()),
::testing::UnorderedElementsAreArray(test_case.expected_default));
}
INSTANTIATE_TEST_SUITE_P(Test, GetOptimizationsTest,
::testing::Values(GetOptimizationTestCase1(),
GetOptimizationTestCase2(),
GetOptimizationTestCase3(),
GetOptimizationTestCase4()));
TEST(DeterministicOpsTest, GetOptimizations) {
tsl::test::DeterministicOpsScope det_scope;
Options options;
// options.deterministic should be ignored when deterministic ops are enabled.
options.set_deterministic(false);
absl::flat_hash_set<tstring> actual_enabled, actual_disabled, actual_default;
GetOptimizations(options, &actual_enabled, &actual_disabled, &actual_default);
EXPECT_THAT(
std::vector<std::string>(actual_enabled.begin(), actual_enabled.end()),
::testing::UnorderedElementsAreArray({"make_deterministic"}));
EXPECT_EQ(actual_disabled.size(), 0);
}
REGISTER_DATASET_EXPERIMENT("test_only_experiment",
RandomJobSamplePercentage<42>, AllTasks);
TEST(DatasetUtilsTest, DatasetExperimentRegistry) {
auto experiments = DatasetExperimentRegistry::Experiments();
EXPECT_TRUE(experiments.find("test_only_experiment") != experiments.end());
EXPECT_TRUE(experiments.find("non_existing_experiment") == experiments.end());
}
TEST(DatasetUtilsTest, CountBytes) {
std::vector<Tensor> uncompressed = {
CreateTensor<int64_t>(TensorShape{128, 2}),
CreateTensor<int64_t>(TensorShape{64, 4})};
EXPECT_EQ(GetAllocatedBytes(uncompressed), 4096);
EXPECT_EQ(GetTotalBytes(uncompressed), 4096);
CompressedElement compressed_element;
TF_ASSERT_OK(CompressElement(uncompressed, &compressed_element));
std::vector<Tensor> compressed{{DT_VARIANT, TensorShape({})}};
compressed.front().scalar<Variant>()() = compressed_element;
EXPECT_EQ(GetAllocatedBytes(compressed), compressed_element.ByteSizeLong());
EXPECT_EQ(GetTotalBytes(compressed), compressed_element.ByteSizeLong());
}
TEST_F(DatasetOpsTestBase, TestVariantEqualityChecking) {
Tensor scalar_0{DT_VARIANT, TensorShape({})};
scalar_0.scalar<Variant>()() = TestVariant({CreateTensor<int64_t>({}, {0})});
TF_EXPECT_OK(ExpectEqual(scalar_0, scalar_0));
Tensor scalar_1{DT_VARIANT, TensorShape({})};
scalar_1.scalar<Variant>()() = TestVariant({CreateTensor<int64_t>({}, {1})});
EXPECT_THAT(
ExpectEqual(scalar_0, scalar_1),
absl_testing::StatusIs(tsl::error::INTERNAL, HasSubstr("aren't equal")));
Tensor nonscalar{DT_VARIANT, TensorShape({2})};
EXPECT_THAT(ExpectEqual(nonscalar, nonscalar),
absl_testing::StatusIs(tsl::error::INTERNAL,
HasSubstr("must be scalars")));
Tensor unsupported{DT_VARIANT, TensorShape({})};
unsupported.scalar<Variant>()() = 0;
EXPECT_THAT(
ExpectEqual(unsupported, unsupported),
absl_testing::StatusIs(tsl::error::INTERNAL, HasSubstr("types must be")));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/finalization_utils.cc
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2021 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/core/data/finalization_utils.h"
#include "absl/status/statusor.h"
#include "tensorflow/core/data/root_dataset.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/refcount.h"
namespace tensorflow {
namespace data {
absl::StatusOr<DatasetBase*> GetFinalizedDataset(OpKernelContext* ctx,
const DatasetBase* dataset) {
return dataset->Finalize(
ctx, [ctx, dataset]() -> absl::StatusOr<core::RefCountPtr<DatasetBase>> {
core::RefCountPtr<DatasetBase> dataset_ref_ptr;
DatasetBase* raw_ptr;
TF_RETURN_IF_ERROR(data::FinalizeDataset(ctx, dataset, &raw_ptr));
dataset_ref_ptr.reset(raw_ptr);
return dataset_ref_ptr;
});
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/finalization_utils.h
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2021 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_CORE_DATA_FINALIZATION_UTILS_H_
#define TENSORFLOW_CORE_DATA_FINALIZATION_UTILS_H_
#include <functional>
#include "absl/status/statusor.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/op_kernel.h"
namespace tensorflow {
namespace data {
// Returns the finalized version of the dataset. The returned DatasetBase is
// unowned and lives for as long as this dataset.
absl::StatusOr<DatasetBase*> GetFinalizedDataset(OpKernelContext* ctx,
const DatasetBase* dataset);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_FINALIZATION_UTILS_H_
tensorflow/core/data/flat_map_utils.cc
+230
View File
@@ -0,0 +1,230 @@
/* Copyright 2024 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/core/data/flat_map_utils.h"
#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <deque>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/synchronization/mutex.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/data/captured_function.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function_handle_cache.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
namespace tensorflow {
namespace data {
FlatMapRandomAccessHandler::FlatMapRandomAccessHandler(
OpKernelContext* const ctx, const DatasetBase* input_dataset,
CapturedFunction& captured_map_func)
: input_dataset_(input_dataset),
captured_map_func_(captured_map_func),
unbounded_thread_pool_(ctx->env(),
"tf_data_flat_map_random_access_handler") {
absl::Status status =
ctx->function_library()->Clone(&flib_def_, &pflr_, &flr_, true);
if (!status.ok()) {
cumulative_cardinalities_ = std::move(status);
return;
}
function_handle_cache_ =
std::make_unique<FunctionHandleCache>(pflr_->GetFLR("/device:CPU:0"));
IteratorContext::Params params(ctx);
params.cancellation_manager = &cancellation_manager_;
params.env = ctx->env();
params.flr = flr_;
params.function_handle_cache = function_handle_cache_.get();
params.resource_mgr = &resource_mgr_;
params.thread_factory = unbounded_thread_pool_.get_thread_factory();
params.thread_pool = &unbounded_thread_pool_;
ctx_ = std::make_unique<IteratorContext>(std::move(params));
}
FlatMapRandomAccessHandler::~FlatMapRandomAccessHandler() {
for (DatasetBase* dataset : input_datasets_) {
dataset->Unref();
}
input_datasets_.clear();
}
absl::StatusOr<int64_t> FlatMapRandomAccessHandler::Cardinality() {
TF_RETURN_IF_ERROR(cumulative_cardinalities_.status());
if (cumulative_cardinalities_->empty()) {
cumulative_cardinalities_ = ComputeCardinalities();
}
TF_RETURN_IF_ERROR(cumulative_cardinalities_.status());
return cumulative_cardinalities_->back();
}
absl::StatusOr<int64_t> FlatMapRandomAccessHandler::CumulativeCardinality(
size_t index) {
TF_RETURN_IF_ERROR(cumulative_cardinalities_.status());
if (index >= cumulative_cardinalities_->size()) {
return absl::OutOfRangeError(absl::StrCat(
"Dataset index exceeds the number of input datasets. Got index: ",
index, ", number of input datasets: ",
cumulative_cardinalities_->size(), "."));
}
return (*cumulative_cardinalities_)[index];
}
absl::StatusOr<std::vector<int64_t>>
FlatMapRandomAccessHandler::ComputeCardinalities() {
if (input_datasets_.empty()) {
TF_ASSIGN_OR_RETURN(input_datasets_, MakeInputDatasets());
}
std::vector<int64_t> cumulative_cardinalities;
cumulative_cardinalities.reserve(input_datasets_.size());
for (size_t i = 0; i < input_datasets_.size(); ++i) {
int64_t input_cardinality = input_datasets_[i]->Cardinality();
if (input_cardinality == kInfiniteCardinality ||
input_cardinality == kUnknownCardinality) {
cumulative_cardinalities.push_back(input_cardinality);
return cumulative_cardinalities;
}
int64_t cumulative_cardinality = input_cardinality;
if (i > 0) {
cumulative_cardinality += cumulative_cardinalities.back();
}
cumulative_cardinalities.push_back(cumulative_cardinality);
}
if (cumulative_cardinalities.empty()) {
cumulative_cardinalities.push_back(0);
}
return cumulative_cardinalities;
}
absl::StatusOr<int64_t> FlatMapRandomAccessHandler::GetDatasetIndex(
size_t element_position) {
TF_ASSIGN_OR_RETURN(int64_t cardinality, Cardinality());
if (cardinality < 0) {
return absl::InvalidArgumentError(absl::StrCat(
"Failed to globally shuffle flat map dataset. Global shuffling "
"requires finite cardinality. Got ",
cardinality, "."));
}
if (element_position >= cardinality) {
return absl::OutOfRangeError(absl::StrCat(
"Element index exceeds the flat map dataset cardinality. Got index: ",
element_position, ", cardinality: ", cardinality, "."));
}
return std::upper_bound(cumulative_cardinalities_->begin(),
cumulative_cardinalities_->end(), element_position) -
cumulative_cardinalities_->begin();
}
absl::StatusOr<std::vector<std::unique_ptr<IteratorBase>>>
FlatMapRandomAccessHandler::MakeInputIterators(
IteratorContext* ctx, const DatasetBaseIterator* parent,
const std::string& prefix) {
if (input_datasets_.empty()) {
TF_ASSIGN_OR_RETURN(input_datasets_, MakeInputDatasets());
}
std::vector<std::unique_ptr<IteratorBase>> result;
if (input_datasets_.empty()) {
return result;
}
result.resize(input_datasets_.size());
for (size_t i = 0; i < input_datasets_.size(); ++i) {
TF_RETURN_IF_ERROR(input_datasets_[i]->MakeIterator(
ctx, parent, absl::StrCat(prefix, "[", i, "]"), &result[i]));
}
return result;
}
absl::StatusOr<std::deque<DatasetBase*>>
FlatMapRandomAccessHandler::MakeInputDatasets() const {
std::unique_ptr<IteratorBase> iterator;
TF_RETURN_IF_ERROR(input_dataset_->MakeIterator(
ctx_.get(), /*parent=*/nullptr, "Iterator", &iterator));
std::unique_ptr<InstantiatedCapturedFunction> map_func;
TF_RETURN_IF_ERROR(captured_map_func_.Instantiate(ctx_.get(), &map_func));
absl::Mutex mu;
std::deque<DatasetBase*> input_datasets;
absl::Status status; // Guarded by `mu`.
std::vector<std::unique_ptr<tsl::Thread>> threads;
while (true) {
std::vector<Tensor> input_tensors;
bool end_of_sequence = false;
TF_RETURN_IF_ERROR(
iterator->GetNext(ctx_.get(), &input_tensors, &end_of_sequence));
if (end_of_sequence) {
break;
}
input_datasets.push_back(nullptr);
DatasetBase*& input_dataset = input_datasets.back();
threads.push_back(ctx_->StartThread(
"flat_map_random_access_iterator",
[this, input_tensors = std::move(input_tensors), &input_dataset,
&map_func, &status, &mu]() {
absl::StatusOr<DatasetBase*> dataset =
MakeInputDataset(std::move(input_tensors), *map_func);
if (!dataset.ok()) {
absl::MutexLock l(mu);
status.Update(dataset.status());
return;
}
input_dataset = *dataset;
}));
}
threads.clear();
TF_RETURN_IF_ERROR(std::move(status));
return input_datasets;
}
absl::StatusOr<DatasetBase*> FlatMapRandomAccessHandler::MakeInputDataset(
std::vector<Tensor> input_tensors,
const InstantiatedCapturedFunction& map_func) const {
std::vector<Tensor> mapped_tensors;
TF_RETURN_IF_ERROR(
map_func.Run(ctx_.get(), std::move(input_tensors), &mapped_tensors));
if (!(mapped_tensors.size() == 1 && mapped_tensors[0].dtype() == DT_VARIANT &&
TensorShapeUtils::IsScalar(mapped_tensors[0].shape()))) {
return absl::InvalidArgumentError(
"Flat map function must return a single scalar of dtype DT_VARIANT "
"representing a dataset.");
}
DatasetBase* mapped_dataset = nullptr;
TF_RETURN_IF_ERROR(
GetDatasetFromVariantTensor(mapped_tensors[0], &mapped_dataset));
mapped_dataset->Ref();
return mapped_dataset;
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/flat_map_utils.h
+112
View File
@@ -0,0 +1,112 @@
/* Copyright 2024 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_CORE_DATA_FLAT_MAP_UTILS_H_
#define TENSORFLOW_CORE_DATA_FLAT_MAP_UTILS_H_
#include <cstddef>
#include <cstdint>
#include <deque>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "xla/tsl/platform/threadpool.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/data/captured_function.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function_handle_cache.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/data/iterator_ops.h"
#include "tsl/platform/refcount.h"
namespace tensorflow {
namespace data {
// Utility class for computing the cardinality of a flat map dataset.
class FlatMapRandomAccessHandler {
public:
// Initializes the counter. This will save necessary information from `ctx`.
// `input_dataset` is the input dataset passed to `flat_map` (not the flat_map
// dataset). `captured_map_func` is the captured map function.
FlatMapRandomAccessHandler(OpKernelContext* ctx,
const DatasetBase* input_dataset,
CapturedFunction& captured_map_func);
virtual ~FlatMapRandomAccessHandler();
FlatMapRandomAccessHandler(const FlatMapRandomAccessHandler&) = delete;
FlatMapRandomAccessHandler& operator=(const FlatMapRandomAccessHandler&) =
delete;
// Returns the dataset cardinality.
absl::StatusOr<int64_t> Cardinality();
// Returns the cumulative cardinality at the index-th dataset.
absl::StatusOr<int64_t> CumulativeCardinality(size_t index);
// Given the flattened element position `element_position`, returns the index
// of the dataset to which the element belongs.
absl::StatusOr<int64_t> GetDatasetIndex(size_t element_position);
// Creates the dataset iterators.
absl::StatusOr<std::vector<std::unique_ptr<IteratorBase>>> MakeInputIterators(
IteratorContext* ctx, const DatasetBaseIterator* parent,
const std::string& prefix);
private:
// Computes the cumulative cardinalities.
absl::StatusOr<std::vector<int64_t>> ComputeCardinalities();
// Creates the input datasets. Each dataset is the result of applying the map
// function to one element from the input iterator.
absl::StatusOr<std::deque<DatasetBase*>> MakeInputDatasets() const;
absl::StatusOr<DatasetBase*> MakeInputDataset(
std::vector<Tensor> input_tensors,
const InstantiatedCapturedFunction& map_func) const;
const DatasetBase* input_dataset_;
CapturedFunction& captured_map_func_;
// The iterator context which bundles together the necessary runtime support
// to create and get elements from the input dataset.
std::unique_ptr<IteratorContext> ctx_;
FunctionLibraryRuntime* flr_;
std::unique_ptr<FunctionLibraryDefinition> flib_def_;
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr_;
std::unique_ptr<thread::ThreadPool> interop_threadpool_;
std::unique_ptr<FunctionHandleCache> function_handle_cache_;
std::function<void(std::function<void()>)> runner_;
ResourceMgr resource_mgr_;
CancellationManager cancellation_manager_;
UnboundedThreadPool unbounded_thread_pool_;
// Input datasets generated by running the map function. Each dataset is the
// result of applying the map function to one element from the input iterator.
std::deque<DatasetBase*> input_datasets_;
// Cumulative cardinalities. Before `ComputeCardinalities` is called, this is
// an empty vector. After `ComputeCardinalities` is called, the last element
// is the dataset cardinality.
absl::StatusOr<std::vector<int64_t>> cumulative_cardinalities_ =
std::vector<int64_t>{};
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_FLAT_MAP_UTILS_H_
tensorflow/core/data/global_shuffle_utils.cc
+124
View File
@@ -0,0 +1,124 @@
/* Copyright 2024 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/core/data/global_shuffle_utils.h"
#include <cstdint>
#include <optional>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
namespace data {
namespace {
constexpr absl::string_view kGlobalShuffleIteratorNextIndex =
"global_shuffle_iterator_next_index";
}
IteratorContextWithIndexMapper::IteratorContextWithIndexMapper(
IteratorContext* ctx, const IteratorBase* iterator)
: ctx_(ctx) {
if (ctx_->index_mapper()) {
IteratorContext::Params params(ctx_);
params.index_mapper = iterator->GetIndexMapper(ctx_->index_mapper());
ctx_with_index_mapper_.emplace(params);
}
}
IteratorContext* IteratorContextWithIndexMapper::Get() {
return ctx_with_index_mapper_.has_value() ? &ctx_with_index_mapper_.value()
: ctx_;
}
void IteratorContextWithIndexMapper::MergeCheckpoint() {
if (ctx_with_index_mapper_.has_value()) {
ctx_->MergeCheckpoint(ctx_with_index_mapper_->checkpoint());
}
}
absl::Status GlobalShuffleIterator::GetNext(IteratorContext* ctx,
std::vector<Tensor>* out_tensors,
bool* end_of_sequence) {
if (!ctx->index_mapper()) {
return absl::FailedPreconditionError(absl::StrCat(
"Trying to get a random element from dataset ", dataset_->DebugString(),
" which is not globally shuffled."));
}
absl::MutexLock l(mu_);
absl::StatusOr<int64_t> shuffled_index =
absl::NotFoundError("Default not found");
while (absl::IsNotFound(shuffled_index.status())) {
shuffled_index = ctx->index_mapper()(element_count_++);
}
if (absl::IsOutOfRange(shuffled_index.status())) {
*end_of_sequence = true;
return absl::OkStatus();
}
TF_RETURN_IF_ERROR(shuffled_index.status());
absl::Status status =
dataset_->Get(AnyContext(ctx), shuffled_index.value(), out_tensors);
if (absl::IsOutOfRange(status)) {
*end_of_sequence = true;
return absl::OkStatus();
}
TF_RETURN_IF_ERROR(status);
*end_of_sequence = false;
return absl::OkStatus();
}
absl::Status GlobalShuffleIterator::Save(
const std::string& parent_iterator_prefix, SerializationContext* ctx,
IteratorStateWriter* writer) {
absl::MutexLock l(mu_);
TF_RETURN_IF_ERROR(writer->WriteScalar(
parent_iterator_prefix, kGlobalShuffleIteratorNextIndex, element_count_));
return absl::OkStatus();
}
absl::Status GlobalShuffleIterator::Restore(
const std::string& parent_iterator_prefix, IteratorContext* ctx,
IteratorStateReader* reader) {
if (!ctx->restored_element_count().has_value()) {
return absl::FailedPreconditionError(absl::StrCat(
"Trying to restore random element count for dataset ",
dataset_->DebugString(), " which is not globally shuffled."));
}
absl::MutexLock l(mu_);
TF_RETURN_IF_ERROR(reader->ReadScalar(parent_iterator_prefix,
kGlobalShuffleIteratorNextIndex,
&element_count_));
return absl::OkStatus();
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/global_shuffle_utils.h
+100
View File
@@ -0,0 +1,100 @@
/* Copyright 2024 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_CORE_DATA_GLOBAL_SHUFFLE_UTILS_H_
#define TENSORFLOW_CORE_DATA_GLOBAL_SHUFFLE_UTILS_H_
#include <cstdint>
#include <optional>
#include <string>
#include <vector>
#include "absl/base/thread_annotations.h"
#include "absl/status/status.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
namespace data {
// Builds and selects the `IteratorContext` to use based on whether the dataset
// is globally shuffled.
//
// Example usage in `Iterator::GetNextInternal`:
//
// ```
// IteratorContextWithIndexMapper ctx_with_index_mapper(ctx, this);
// TF_RETURN_IF_ERROR(input_impl_->GetNext(
// ctx_with_index_mapper.Get(), out_tensors, end_of_sequence));
// ctx_with_index_mapper.MergeCheckpoint();
// ```
//
// The iterator should also implement `GetIndexMapper` if it needs to customize
// the index mapping behavior.
class IteratorContextWithIndexMapper {
public:
// Caller keeps ownership of both pointers.
explicit IteratorContextWithIndexMapper(IteratorContext* ctx,
const IteratorBase* iterator);
virtual ~IteratorContextWithIndexMapper() = default;
IteratorContextWithIndexMapper(const IteratorContextWithIndexMapper&) =
delete;
IteratorContextWithIndexMapper& operator=(
const IteratorContextWithIndexMapper&) = delete;
IteratorContext* Get();
void MergeCheckpoint();
private:
IteratorContext* ctx_;
std::optional<IteratorContext> ctx_with_index_mapper_;
};
// For source datasets that support random access, this class adapts the dataset
// random access API to support globally shuffled iterators.
class GlobalShuffleIterator {
public:
// The dataset is expected to support random access by implementing the
// absl::Status Get(int64_t index, std::vector<Tensor>* out_tensors) const.
explicit GlobalShuffleIterator(const DatasetBase* dataset)
: dataset_(dataset) {}
// Returns the next shuffled element.
// REQUIRES: ctx->index_mapper() != nullptr.
absl::Status GetNext(IteratorContext* ctx, std::vector<Tensor>* out_tensors,
bool* end_of_sequence);
absl::Status Save(const std::string& parent_iterator_prefix,
SerializationContext* ctx, IteratorStateWriter* writer);
// Restores the element count.
// REQUIRES: ctx->restored_element_count() != nullopt.
absl::Status Restore(const std::string& parent_iterator_prefix,
IteratorContext* ctx, IteratorStateReader* reader);
private:
const DatasetBase* const dataset_;
mutable absl::Mutex mu_;
// Count of elements produced by this iterator when it runs in the random
// access mode.
int64_t element_count_ ABSL_GUARDED_BY(mu_) = 0;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_GLOBAL_SHUFFLE_UTILS_H_
tensorflow/core/data/hash_utils.cc
+817
View File
@@ -0,0 +1,817 @@
/* Copyright 2020 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/core/data/hash_utils.h"
#include <array>
#include <cstddef>
#include <memory>
#include <queue>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/hash/hash.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/data/dataset_utils.h"
#include "tensorflow/core/data/serialization_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/op_def_builder.h"
#include "tensorflow/core/framework/op_def_util.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/regexp.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/work_sharder.h"
namespace tensorflow {
namespace data {
namespace {
// clang-format off
constexpr std::array<const char*, 3> kOpsWithSeed = {
"AnonymousRandomSeedGenerator",
"ShuffleDataset",
"ShuffleAndRepeatDataset"
};
// clang-format on
constexpr char kSeedInputName[] = "seed";
constexpr char kSeed2InputName[] = "seed2";
constexpr char kSeedGeneratorInputName[] = "seed_generator";
template <std::size_t SIZE>
bool IsNodeOfType(const NodeDef& node,
const std::array<const char*, SIZE>& op_types) {
for (const auto& type : op_types) {
if (MatchesAnyVersion(type, node.op())) {
return true;
}
}
return false;
}
absl::Status GetSink(const GraphDef& graph_def, const NodeDef** sink) {
for (auto& node : graph_def.node()) {
if (node.op() == kRetvalOp) {
*sink = &node;
break;
}
}
if (sink == nullptr) {
return absl::InternalError("Cannot find sink node for dataset graph.");
}
return absl::OkStatus();
}
absl::Status ShouldIgnoreInput(const NodeDef& node, int i, bool* result) {
*result = false;
if (IsNodeOfType(node, kOpsWithSeed)) {
const OpRegistrationData* reg;
auto status = OpRegistry::Global()->LookUp(node.op(), &reg);
if (status.ok()) {
if (reg->op_def.input_arg_size() > i) {
const std::string input_arg_name = reg->op_def.input_arg(i).name();
if (input_arg_name == kSeedInputName ||
input_arg_name == kSeed2InputName ||
input_arg_name == kSeedGeneratorInputName) {
VLOG(2) << "Ignoring arg: " << input_arg_name
<< " from node: " << node.name();
*result = true;
return absl::OkStatus();
}
}
} else if (absl::IsNotFound(status)) {
LOG(WARNING) << "Cannot find " << node.op()
<< " in global op registry, so cannot determine which "
"inputs are seeds.";
} else {
return status;
}
}
return absl::OkStatus();
}
absl::Status ParseInputNodeName(absl::string_view input_name,
absl::string_view* node_name,
absl::string_view* suffix,
bool* is_control_input) {
if (input_name[0] == '^') {
*node_name = input_name.substr(1);
*is_control_input = true;
return absl::OkStatus();
}
std::pair<absl::string_view, absl::string_view> node_spec =
absl::StrSplit(input_name, absl::MaxSplits(':', 1));
*node_name = node_spec.first;
*suffix = node_spec.second;
*is_control_input = false;
return absl::OkStatus();
}
// Given a graph_def and a root_node, this class computes a fingerprint that
// tries to capture the structure of the graph rooted at the provided node.
// It does not at any point rely on the names of the nodes in the graph and
// just relies on the connections between different nodes. In the presence of
// multiple cycles in the graph, there is a non-zero possibility that two
// graphs with different structure might end up with the same fingerprint
// as in order to break cycles we prune away some edges (in a deterministic
// fashion though). Idea for this algorithm was borrowed from:
// https://stackoverflow.com/questions/11338746/directed-graphs-with-a-given-root-node-match-another-directed-graph-for-equali
class GraphHasher {
using NodeCache = absl::flat_hash_map<const NodeDef*, uint64_t>;
using FunctionCache = absl::flat_hash_map<const FunctionDef*, uint64_t>;
using AttrCache =
absl::flat_hash_map<std::pair<const NodeDef*, bool>, uint64_t>;
public:
// `GraphHasher` does not take ownership of `graph_def`, `root_node`, or
// `flib_def`.
explicit GraphHasher(const GraphDef* graph, const NodeDef* root,
const FunctionLibraryDefinition* flib)
: graph_(graph), root_(root), flib_(flib) {
node_cache_ = std::make_shared<NodeCache>();
function_cache_ = std::make_shared<FunctionCache>();
attr_cache_ = std::make_shared<AttrCache>();
}
explicit GraphHasher(const GraphDef* graph, const NodeDef* root,
const FunctionLibraryDefinition* flib,
std::shared_ptr<NodeCache> node_cache,
std::shared_ptr<FunctionCache> function_cache,
std::shared_ptr<AttrCache> attr_cache)
: graph_(graph),
root_(root),
flib_(flib),
node_cache_(node_cache),
function_cache_(function_cache),
attr_cache_(attr_cache) {}
absl::Status Init() {
// Construct a map of name -> NodeDef to avoid repeated linear searches.
absl::flat_hash_map<absl::string_view, const NodeDef*> node_def_by_name;
node_def_by_name.reserve(graph_->node_size());
for (const auto& node : graph_->node()) {
auto result = node_def_by_name.emplace(node.name(), &node);
if (TF_PREDICT_FALSE(!result.second)) {
auto node_name_formatter =
[](std::string* out,
const decltype(node_def_by_name)::value_type& item) {
absl::StrAppend(out, "'", item.first, "'");
};
return absl::InternalError(absl::StrCat(
"Encountered graph with duplicate node name '", node.name(),
"' in [", absl::StrJoin(node_def_by_name, ",", node_name_formatter),
"]"));
}
}
// Pre-process the graph to do a BFS and prune away cycles that might cause
// problems.
absl::flat_hash_set<absl::string_view> visited;
std::queue<const NodeDef*> bfs_queue;
bfs_queue.push(root_);
while (!bfs_queue.empty()) {
const NodeDef* node = bfs_queue.front();
bfs_queue.pop();
if (visited.contains(node->name())) {
continue;
}
visited.insert(node->name());
NodeRep node_rep;
for (int i = 0; i < node->input_size(); ++i) {
DCHECK_GT(node->input(i).length(), 0);
// We skip trying to take the hash of the seeds of any ops, as they
// are irrelevant to the hash of the graph and may vary from run to run.
bool should_ignore_input = false;
TF_RETURN_IF_ERROR(ShouldIgnoreInput(*node, i, &should_ignore_input));
if (should_ignore_input) continue;
absl::string_view node_name, suffix;
bool is_control_input;
TF_RETURN_IF_ERROR(ParseInputNodeName(node->input(i), &node_name,
&suffix, &is_control_input));
auto* input_node = gtl::FindPtrOrNull(node_def_by_name, node_name);
if (input_node == nullptr) {
return absl::InternalError(
absl::StrCat("Graph node [", node->name(), "] has input [",
node_name, "] that doesn't exist in graph"));
}
// If we've already seen this node before, skip it and don't add it to
// the queue.
if (visited.contains(node_name)) {
EdgeRep cycle_edge(node, input_node);
cycle_forming_edges_.insert(cycle_edge.GetHash());
continue;
}
if (is_control_input) {
node_rep.node_control_inputs.push_back(input_node);
} else {
node_rep.node_inputs.push_back(std::make_pair(input_node, suffix));
bfs_queue.push(input_node);
}
}
nodes_[node] = node_rep;
}
return absl::OkStatus();
}
absl::Status HashRoot(uint64_t* hash) { return HashNode(root_, hash); }
absl::Status CheckEqual(GraphHasher* that) {
return CheckNodesEqual(root_, that, that->root_);
}
private:
absl::Status HashNode(const NodeDef* node, uint64_t* hash) {
auto it = node_cache_->find(node);
if (it != node_cache_->end()) {
*hash = it->second;
return absl::OkStatus();
}
NodeRep* node_rep = gtl::FindOrNull(nodes_, node);
if (node_rep == nullptr) {
return absl::InvalidArgumentError(
absl::StrCat("Could not find node: ", node->name()));
}
uint64_t non_input_hash;
TF_RETURN_IF_ERROR(
HashNodeNonInput(node, /*hash_functions=*/true, &non_input_hash));
uint64_t control_inputs_hash;
TF_RETURN_IF_ERROR(
HashControlInputs(node_rep->node_control_inputs, &control_inputs_hash));
// Hash regular inputs. We combine them in an ordered fashion.
uint64_t inputs_hash = 0;
for (const auto& input : node_rep->node_inputs) {
uint64_t node_hash = 0;
EdgeRep edge(node, input.first);
// If the edge was pruned we get the non input node hash to avoid cycles.
if (cycle_forming_edges_.contains(edge.GetHash())) {
TF_RETURN_IF_ERROR(
HashNodeNonInput(input.first, /*hash_functions=*/true, &node_hash));
} else {
TF_RETURN_IF_ERROR(HashNode(input.first, &node_hash));
}
inputs_hash = Hash64Combine(
inputs_hash, Hash64Combine(node_hash, Hash64(input.second.data(),
input.second.size())));
}
*hash = Hash64Combine(non_input_hash,
Hash64Combine(control_inputs_hash, inputs_hash));
auto result = node_cache_->emplace(node, *hash);
if (!result.second) {
return absl::InternalError(absl::StrCat("Computed the hash for node ",
node->DebugString(), " twice!"));
}
return absl::OkStatus();
}
absl::Status CheckNodesEqual(const NodeDef* this_node, GraphHasher* that,
const NodeDef* that_node) {
absl::Status s = CheckNodesEqualHelper(this_node, that, that_node);
if (!s.ok()) {
return absl::FailedPreconditionError(
absl::StrCat("Nodes ", this_node->name(), " and ", that_node->name(),
" are not the same:\n", s));
}
return s;
}
absl::Status CheckNodesEqualHelper(const NodeDef* this_node,
GraphHasher* that,
const NodeDef* that_node) {
TF_RETURN_IF_ERROR(CheckNodesEqualNonInput(this_node, that, that_node,
/*compare_functions=*/true));
TF_RETURN_IF_ERROR(
CheckControlInputsEqual(nodes_[this_node].node_control_inputs, that,
that->nodes_[that_node].node_control_inputs));
auto& this_node_inputs = nodes_[this_node].node_inputs;
auto& that_node_inputs = that->nodes_[that_node].node_inputs;
if (this_node_inputs.size() != that_node_inputs.size()) {
return absl::FailedPreconditionError(absl::StrCat(
"Nodes have different numbers of node inputs: ",
this_node_inputs.size(), " vs ", that_node_inputs.size()));
}
for (int i = 0; i < this_node_inputs.size(); ++i) {
const NodeDef* this_input = this_node_inputs[i].first;
const NodeDef* that_input = that_node_inputs[i].first;
if (is_cycle_forming_edge(this_node, this_input)) {
TF_RETURN_IF_ERROR(CheckNodesEqualNonInput(this_input, that, that_input,
/*compare_functions=*/true));
} else {
TF_RETURN_IF_ERROR(CheckNodesEqual(this_input, that, that_input));
}
absl::string_view this_input_suffix = this_node_inputs[i].second;
absl::string_view that_input_suffix = that_node_inputs[i].second;
if (this_input_suffix != that_input_suffix) {
return absl::FailedPreconditionError(absl::StrCat(
"Node inputs ", this_input->name(), " and ", that_input->name(),
" have different suffixes: ", this_input_suffix, " vs ",
that_input_suffix));
}
}
return absl::OkStatus();
}
absl::Status HashNodeNonInput(const NodeDef* node, bool hash_functions,
uint64_t* hash) {
auto iter = attr_cache_->find(std::make_pair(node, hash_functions));
if (iter != attr_cache_->end()) {
*hash = iter->second;
return absl::OkStatus();
}
// Hash Attrs. We get the list of attrs from the op registry and then look
// up their values in the NodeDef attr map. This avoids looping over
// a map which is non-deterministic.
uint64_t attrs_hash = 0;
const OpRegistrationData* reg;
TF_RETURN_IF_ERROR(flib_->LookUp(node->op(), &reg));
uint64_t op_hash = 0;
if (reg->is_function_op) {
if (hash_functions) {
TF_RETURN_IF_ERROR(HashFunction(node->op(), node->attr(), &op_hash));
}
} else {
op_hash = Hash64(node->op());
}
for (const auto& attr : reg->op_def.attr()) {
const auto& attr_key = attr.name();
// Ignore "metadata" attribute of tf.data operations.
if (DatasetOpKernel::IsDatasetOp(reg->op_def) && attr_key == "metadata")
continue;
auto node_attr_iter = node->attr().find(attr_key);
if (node_attr_iter == node->attr().end()) {
continue;
}
const auto& attr_value = node_attr_iter->second;
if (attr_key == kColocationAttrName ||
attr_key == kColocationGroupPrefix) {
continue;
}
uint64_t attr_hash = 0;
TF_RETURN_IF_ERROR(
HashAttr(attr_key, attr_value, hash_functions, &attr_hash));
attrs_hash = Hash64Combine(attrs_hash, attr_hash);
}
// Hash Device.
uint64_t device_hash = Hash64(node->device());
*hash = Hash64Combine(op_hash, Hash64Combine(attrs_hash, device_hash));
auto result =
attr_cache_->emplace(std::make_pair(node, hash_functions), *hash);
if (!result.second) {
return absl::InternalError(absl::StrCat(
"Computed the hash for non-input node: ", node->DebugString(),
" and hash function bool: ", hash_functions, "twice!"));
}
return absl::OkStatus();
}
absl::Status CheckNodesEqualNonInput(const NodeDef* this_node,
GraphHasher* that,
const NodeDef* that_node,
bool compare_functions) {
// We get the list of attrs from the op registry and then look
// up their values in the NodeDef attr map. This avoids looping over
// a map which is non-deterministic.
const OpRegistrationData* reg;
TF_RETURN_IF_ERROR(flib_->LookUp(this_node->op(), &reg));
if (reg->is_function_op) {
if (compare_functions) {
TF_RETURN_IF_ERROR(
CheckFunctionsEqual(this_node->op(), this_node->attr(), that,
that_node->op(), that_node->attr()));
}
} else {
if (this_node->op() != that_node->op()) {
return absl::FailedPreconditionError(absl::StrCat(
"ops for nodes ", this_node->name(), " and ", that_node->name(),
" are different: ", this_node->op(), " != ", that_node->op()));
}
}
for (const auto& attr : reg->op_def.attr()) {
const auto& attr_key = attr.name();
const bool this_has_attr = this_node->attr().contains(attr_key);
const bool that_has_attr = that_node->attr().contains(attr_key);
if (this_has_attr != that_has_attr) {
return absl::FailedPreconditionError(
absl::StrCat("attr with key ", attr_key, " is different for nodes ",
this_node->name(), " and ", that_node->name(),
". Present in former: ", this_has_attr,
". Present in latter: ", that_has_attr));
}
if (!this_has_attr) {
continue;
}
if (attr_key == kColocationAttrName ||
attr_key == kColocationGroupPrefix) {
continue;
}
const auto& this_attr = this_node->attr().at(attr_key);
const auto& that_attr = that_node->attr().at(attr_key);
TF_RETURN_IF_ERROR(CheckAttrsEqual(attr_key, this_attr, that, that_attr,
compare_functions));
}
if (this_node->device() != that_node->device()) {
return absl::FailedPreconditionError(
absl::StrCat("Devices are different for nodes ", this_node->name(),
" and ", that_node->name(), ": ", this_node->device(),
" vs ", that_node->device()));
}
return absl::OkStatus();
}
absl::Status HashAttr(const std::string& attr_name,
const AttrValue& attr_value, bool hash_functions,
uint64_t* hash) {
uint64_t value_hash = 0;
if (attr_value.has_func()) {
if (hash_functions) {
TF_RETURN_IF_ERROR(HashFunction(attr_value.func(), &value_hash));
}
} else if (attr_value.has_list() && attr_value.list().func_size() > 0) {
if (hash_functions) {
for (auto& func : attr_value.list().func()) {
uint64_t func_hash;
TF_RETURN_IF_ERROR(HashFunction(func, &func_hash));
value_hash = Hash64Combine(value_hash, func_hash);
}
}
} else {
value_hash = DeterministicProtoHash64(attr_value);
}
*hash = Hash64Combine(Hash64(attr_name), value_hash);
return absl::OkStatus();
}
absl::Status CheckAttrsEqual(const std::string& attr_name,
const AttrValue& this_attr, GraphHasher* that,
const AttrValue& that_attr,
bool compare_functions) {
if (this_attr.has_func() != that_attr.has_func()) {
return absl::FailedPreconditionError(absl::StrCat(
"AttrValues are of different types: ", this_attr.DebugString(),
" vs ", that_attr.DebugString()));
}
if (this_attr.has_func()) {
if (compare_functions) {
TF_RETURN_IF_ERROR(
CheckFunctionsEqual(this_attr.func(), that, that_attr.func()));
}
return absl::OkStatus();
}
if (this_attr.has_list() != that_attr.has_list()) {
return absl::FailedPreconditionError(absl::StrCat(
"AttrValues are of different types: ", this_attr.DebugString(),
" vs ", that_attr.DebugString()));
}
if (this_attr.has_list()) {
if (this_attr.list().func_size() != that_attr.list().func_size()) {
return absl::FailedPreconditionError(absl::StrCat(
"AttrValues have func lists of different sizes: ",
this_attr.DebugString(), " vs ", that_attr.DebugString()));
}
if (compare_functions) {
for (int i = 0; i < this_attr.list().func_size(); ++i) {
TF_RETURN_IF_ERROR(CheckFunctionsEqual(this_attr.list().func(i), that,
that_attr.list().func(i)));
}
}
return absl::OkStatus();
}
uint64_t this_hash, that_hash;
TF_RETURN_IF_ERROR(
HashAttr(attr_name, this_attr, /*hash_functions=*/true, &this_hash));
TF_RETURN_IF_ERROR(that->HashAttr(attr_name, that_attr,
/*hash_functions=*/true, &that_hash));
if (this_hash != that_hash) {
return absl::FailedPreconditionError(
absl::StrCat("AttrValues are different: ", this_attr.DebugString(),
" vs ", that_attr.DebugString()));
}
return absl::OkStatus();
}
absl::Status HashFunction(const NameAttrList& func, uint64_t* hash) {
return HashFunction(func.name(), func.attr(), hash);
}
absl::Status HashFunction(const std::string& name, const AttrValueMap& attrs,
uint64_t* hash) {
const FunctionDef* fdef = flib_->Find(name);
auto it = function_cache_->find(fdef);
if (it != function_cache_->end()) {
*hash = it->second;
return absl::OkStatus();
}
// Convert to a GraphDef.
std::unique_ptr<FunctionBody> fbody;
TF_RETURN_IF_ERROR(
FunctionDefToBodyHelper(*fdef, AttrSlice(&attrs), flib_, &fbody));
GraphDef graph_def = fbody->graph->ToGraphDefDebug();
// For each return node, we create a new GraphHasher to compute a hash.
// We then combine these hashes to produce the hash ordered.
uint64_t ret_nodes_hash = 0;
for (const auto& ret_node : fbody->ret_nodes) {
uint64_t ret_node_hash = 0;
GraphHasher hasher(&graph_def, &ret_node->def(), flib_, node_cache_,
function_cache_, attr_cache_);
TF_RETURN_IF_ERROR(hasher.Init());
TF_RETURN_IF_ERROR(hasher.HashRoot(&ret_node_hash));
ret_nodes_hash = Hash64Combine(ret_nodes_hash, ret_node_hash);
}
std::vector<const NodeDef*> control_rets;
control_rets.reserve(fbody->control_ret_nodes.size());
for (const auto& control_ret_node : fbody->control_ret_nodes) {
control_rets.push_back(&control_ret_node->def());
}
uint64_t control_ret_nodes_hash = 0;
TF_RETURN_IF_ERROR(
HashControlInputs(control_rets, &control_ret_nodes_hash));
*hash = Hash64Combine(ret_nodes_hash, control_ret_nodes_hash);
auto result = function_cache_->emplace(fdef, *hash);
if (!result.second) {
return absl::InternalError(
absl::StrCat("Computed the hash for function ", name, " twice!"));
}
return absl::OkStatus();
}
absl::Status CheckFunctionsEqual(const NameAttrList& this_func,
GraphHasher* that,
const NameAttrList& that_func) {
return CheckFunctionsEqual(this_func.name(), this_func.attr(), that,
that_func.name(), that_func.attr());
}
absl::Status CheckFunctionsEqual(const std::string& this_name,
const AttrValueMap& this_attrs,
GraphHasher* that,
const std::string& that_name,
const AttrValueMap& that_attrs) {
absl::Status s = CheckFunctionsEqualHelper(this_name, this_attrs, that,
that_name, that_attrs);
if (!s.ok()) {
return absl::FailedPreconditionError(
absl::StrCat("Functions ", this_name, " and ", that_name,
" are not the same:\n", s));
}
return s;
}
absl::Status CheckFunctionsEqualHelper(const std::string& this_name,
const AttrValueMap& this_attrs,
GraphHasher* that,
const std::string& that_name,
const AttrValueMap& that_attrs) {
const FunctionDef* this_fdef = flib_->Find(this_name);
const FunctionDef* that_fdef = that->flib_->Find(that_name);
// Convert to GraphDefs.
std::unique_ptr<FunctionBody> this_fbody;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
*this_fdef, AttrSlice(&this_attrs), flib_, &this_fbody));
GraphDef this_graph_def = this_fbody->graph->ToGraphDefDebug();
std::unique_ptr<FunctionBody> that_fbody;
TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
*that_fdef, AttrSlice(&that_attrs), that->flib_, &that_fbody));
GraphDef that_graph_def = that_fbody->graph->ToGraphDefDebug();
if (this_fbody->ret_nodes.size() != that_fbody->ret_nodes.size()) {
return absl::FailedPreconditionError(absl::StrCat(
"Different numbers of ret nodes for functions ", this_name, " and ",
that_name, ": ", this_fbody->ret_nodes.size(), " vs ",
that_fbody->ret_nodes.size()));
}
for (int i = 0; i < this_fbody->ret_nodes.size(); ++i) {
const NodeDef* this_root = &this_fbody->ret_nodes[i]->def();
const NodeDef* that_root = &that_fbody->ret_nodes[i]->def();
GraphHasher this_hasher(&this_graph_def, this_root, flib_, node_cache_,
function_cache_, attr_cache_);
TF_RETURN_IF_ERROR(this_hasher.Init());
GraphHasher that_hasher(&that_graph_def, that_root, that->flib_,
node_cache_, function_cache_, attr_cache_);
TF_RETURN_IF_ERROR(that_hasher.Init());
TF_RETURN_IF_ERROR(this_hasher.CheckEqual(&that_hasher));
}
std::vector<const NodeDef*> this_control_rets;
this_control_rets.reserve(this_fbody->control_ret_nodes.size());
for (const auto& control_ret_node : this_fbody->control_ret_nodes) {
this_control_rets.push_back(&control_ret_node->def());
}
std::vector<const NodeDef*> that_control_rets;
that_control_rets.reserve(that_fbody->control_ret_nodes.size());
for (const auto& control_ret_node : that_fbody->control_ret_nodes) {
that_control_rets.push_back(&control_ret_node->def());
}
TF_RETURN_IF_ERROR(
CheckControlInputsEqual(this_control_rets, that, that_control_rets));
return absl::OkStatus();
}
absl::Status HashControlInputs(const std::vector<const NodeDef*>& inputs,
uint64_t* hash) {
*hash = 0;
for (const NodeDef* input : inputs) {
uint64_t node_hash = 0;
TF_RETURN_IF_ERROR(
HashNodeNonInput(input, /*hash_functions=*/false, &node_hash));
*hash = Hash64CombineUnordered(*hash, node_hash);
}
return absl::OkStatus();
}
absl::Status CheckControlInputsEqual(
const std::vector<const NodeDef*>& this_inputs, GraphHasher* that,
const std::vector<const NodeDef*>& that_inputs) {
absl::flat_hash_map<uint64_t, const NodeDef*> this_hashes;
for (const NodeDef* input : this_inputs) {
uint64_t node_hash = 0;
TF_RETURN_IF_ERROR(
HashNodeNonInput(input, /*hash_functions=*/false, &node_hash));
this_hashes[node_hash] = input;
}
absl::flat_hash_map<uint64_t, const NodeDef*> that_hashes;
for (const NodeDef* input : that_inputs) {
uint64_t node_hash = 0;
TF_RETURN_IF_ERROR(
HashNodeNonInput(input, /*hash_functions=*/false, &node_hash));
auto this_iter = this_hashes.find(node_hash);
if (this_iter != this_hashes.end()) {
this_hashes.erase(this_iter);
} else {
that_hashes[node_hash] = input;
}
}
if (!this_hashes.empty()) {
auto formatter = [](std::string* out,
const decltype(this_hashes)::value_type& item) {
out->append(item.second->name());
};
return absl::FailedPreconditionError(absl::StrCat(
"Control dependencies are different. One node has dependencies [",
absl::StrJoin(this_hashes, ", ", formatter),
"], which don't match any of the other node's dependencies [",
absl::StrJoin(that_hashes, ", ", formatter), "]"));
}
return absl::OkStatus();
}
private:
bool is_cycle_forming_edge(const NodeDef* start, const NodeDef* end) {
EdgeRep edge(start, end);
return cycle_forming_edges_.contains(edge.GetHash());
}
struct NodeRep {
std::vector<const NodeDef*> node_control_inputs;
std::vector<std::pair<const NodeDef*, absl::string_view>> node_inputs;
};
struct EdgeRep {
const NodeDef* start_node;
const NodeDef* end_node;
EdgeRep(const NodeDef* start, const NodeDef* end)
: start_node(start), end_node(end) {}
uint64_t GetHash() {
return Hash64Combine(absl::Hash<const NodeDef*>()(start_node),
absl::Hash<const NodeDef*>()(end_node));
}
};
const GraphDef* const graph_; // Not owned.
const NodeDef* const root_; // Not owned.
const FunctionLibraryDefinition* const flib_; // Not owned.
// Edges that need to be pruned as their presence will cause cycles.
absl::flat_hash_set<uint64_t> cycle_forming_edges_;
absl::flat_hash_map<const NodeDef*, NodeRep> nodes_;
std::shared_ptr<NodeCache> node_cache_;
std::shared_ptr<FunctionCache> function_cache_;
std::shared_ptr<AttrCache> attr_cache_;
};
} // anonymous namespace
absl::Status HashTensor(const Tensor& tensor, uint64_t* hash) {
const tstring* s = nullptr;
// Hash tensor type.
*hash = Hash64Combine(0, tensor.dtype());
// Hash tensor shape.
for (int i = 0; i < tensor.shape().dims(); ++i) {
*hash = Hash64Combine(*hash, tensor.shape().dim_size(i));
}
// Hash tensor data.
switch (tensor.dtype()) {
case DT_RESOURCE:
case DT_VARIANT:
return absl::UnimplementedError(absl::StrCat(
"Hashing ", DataTypeString(tensor.dtype()), " is not supported."));
case DT_STRING:
s = tensor.flat<tstring>().data();
for (int i = 0; i < tensor.NumElements(); ++i, ++s) {
*hash = Hash64Combine(*hash, Hash64(s->data(), s->size()));
}
break;
default:
*hash = Hash64(tensor.tensor_data().data(), tensor.tensor_data().size());
}
return absl::OkStatus();
}
absl::Status HashNode(const GraphDef& graph, const NodeDef& node,
uint64_t* hash) {
const FunctionLibraryDefinition flib_def(OpRegistry::Global(),
graph.library());
return HashNode(graph, node, flib_def, hash);
}
absl::Status HashNode(const GraphDef& graph, const NodeDef& node,
const FunctionLibraryDefinition& flib_def,
uint64_t* hash) {
GraphHasher hasher(&graph, &node, &flib_def);
TF_RETURN_IF_ERROR(hasher.Init());
return hasher.HashRoot(hash);
}
absl::Status HashGraph(const GraphDef& graph_def, uint64_t* hash) {
const NodeDef* sink = nullptr;
TF_RETURN_IF_ERROR(GetSink(graph_def, &sink));
return HashNode(graph_def, *sink, hash);
}
absl::Status CheckGraphsEqual(const GraphDef& a, const GraphDef& b) {
const NodeDef* sink_a;
TF_RETURN_IF_ERROR(GetSink(a, &sink_a));
const NodeDef* sink_b;
TF_RETURN_IF_ERROR(GetSink(b, &sink_b));
return CheckSubgraphsEqual(a, sink_a, b, sink_b);
}
absl::Status CheckSubgraphsEqual(const GraphDef& a, const NodeDef* node_a,
const GraphDef& b, const NodeDef* node_b) {
const FunctionLibraryDefinition flib_def_a(OpRegistry::Global(), a.library());
GraphHasher hasher_a(&a, node_a, &flib_def_a);
TF_RETURN_IF_ERROR(hasher_a.Init());
const FunctionLibraryDefinition flib_def_b(OpRegistry::Global(), b.library());
GraphHasher hasher_b(&b, node_b, &flib_def_b);
TF_RETURN_IF_ERROR(hasher_b.Init());
return hasher_a.CheckEqual(&hasher_b);
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/hash_utils.h
+66
View File
@@ -0,0 +1,66 @@
/* Copyright 2020 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_CORE_DATA_HASH_UTILS_H_
#define TENSORFLOW_CORE_DATA_HASH_UTILS_H_
#include "absl/status/status.h"
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
namespace data {
// Returns a stable hash of the subgraph rooted at the given node.
//
// NOTE: There is currently no guarantee that the hash of a subgraph will stay
// the same between TensorFlow builds.
absl::Status HashNode(const GraphDef& graph, const NodeDef& node,
uint64_t* hash);
absl::Status HashNode(const GraphDef& graph, const NodeDef& node,
const FunctionLibraryDefinition& flib_def,
uint64_t* hash);
// Returns a stable hash of the given tensor.
//
// NOTE: There is currently no guarantee that the hash of a subgraph will stay
// the same between TensorFlow builds.
absl::Status HashTensor(const Tensor& tensor, uint64_t* hash);
// Returns a stable hash of the given graph.
//
// NOTE: There is currently no guarantee that the hash of a subgraph will stay
// the same between TensorFlow builds.
absl::Status HashGraph(const GraphDef& graph, uint64_t* hash);
// Determines whether the given graphs are equal, following the same logic used
// for HashGraph. Returns OK if the graphs can be determined to be equal,
// otherwise returns an error message explaining why the graphs couldn't be
// determined to be equal.
absl::Status CheckGraphsEqual(const GraphDef& a, const GraphDef& b);
// Determines whether the subgraphs rooted at the given nodes are equal
// following the same logic used for HashGraph. Returns OK if the graphs can be
// determined to be equal, otherwise returns an error message explaining why the
// graphs couldn't be determined to be equal.
absl::Status CheckSubgraphsEqual(const GraphDef& a, const NodeDef* node_a,
const GraphDef& b, const NodeDef* node_b);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_HASH_UTILS_H_
tensorflow/core/data/hash_utils_test.cc
+1364
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/metric_utils.cc
+94
View File
@@ -0,0 +1,94 @@
/* Copyright 2022 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/core/data/metric_utils.h"
#include <algorithm>
#include <cstdint>
#include <string>
#include <vector>
#include "absl/time/time.h"
#include "tensorflow/core/data/utils.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mutex.h"
namespace tensorflow {
namespace data {
namespace {
// Safely subtracts `y` from `x` avoiding underflow.
uint64_t safe_sub(uint64_t x, uint64_t y) { return x >= y ? x - y : 0; }
} // namespace
IteratorMetricsCollector::IteratorMetricsCollector(
const std::string& device_type, const Env& env)
: device_type_(device_type), env_(env) {}
absl::Time IteratorMetricsCollector::RecordStart() {
const uint64_t start_time_us = env_.NowMicros();
if (!ShouldCollectMetrics()) {
return absl::FromUnixMicros(start_time_us);
}
mutex_lock l(mu_);
if (end_time_us_ == 0) {
// We initialize `end_time_us_` to the start time of the first request to
// make it possible to use the delta between `end_time_us_` and subsequent
// `GetNext()` end time to incrementally collect the duration of the
// iterator's lifetime.
end_time_us_ = start_time_us;
}
uint64_t gap_time_us = 0;
if (num_active_calls_ == 0) {
first_start_time_us_ = start_time_us;
gap_time_us = safe_sub(start_time_us, end_time_us_);
}
metrics::RecordTFDataIteratorGap(gap_time_us);
num_active_calls_++;
return absl::FromUnixMicros(start_time_us);
}
void IteratorMetricsCollector::RecordStop(absl::Time start_time,
const std::vector<Tensor>& output) {
if (!ShouldCollectMetrics()) {
return;
}
const uint64_t end_time_us = env_.NowMicros();
const int64_t latency_micros =
safe_sub(end_time_us, absl::ToUnixMicros(start_time));
AddLatencySample(latency_micros);
IncrementThroughput(GetTotalBytes(output));
mutex_lock l(mu_);
metrics::RecordTFDataIteratorLifetime(safe_sub(end_time_us, end_time_us_));
end_time_us_ = std::max(end_time_us_, end_time_us);
num_active_calls_--;
if (num_active_calls_ == 0) {
metrics::RecordTFDataIteratorBusy(
safe_sub(end_time_us_, first_start_time_us_));
}
}
bool IteratorMetricsCollector::ShouldCollectMetrics() const {
return device_type_ == DEVICE_CPU;
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/metric_utils.h
+87
View File
@@ -0,0 +1,87 @@
/* Copyright 2022 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_CORE_DATA_METRIC_UTILS_H_
#define TENSORFLOW_CORE_DATA_METRIC_UTILS_H_
#include <cstdint>
#include <string>
#include <vector>
#include "absl/time/time.h"
#include "tensorflow/core/data/tfdataz_metrics.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
namespace data {
// Exports the metrics for `GetNext` calls by tf.data iterators. When the user
// calls `RecordStart` and `RecordStop`, it will export a latency sample. It
// also exports throughput, tf.data iterator life time, etc. This class is
// thread-safe. Example usage:
//
// ```
// IteratorMetricsCollector metrics_collector(DEVICE_CPU, env);
// absl::Time start_time = metrics_collector.RecordStart();
// auto status = iterator_->GetNext(IteratorContext(std::move(params)),
// out_tensors, end_of_sequence);
// metrics_collector.RecordStop(start_time, *out_tensors);
// ```
class IteratorMetricsCollector {
public:
// Constructs a `IteratorMetricsCollector`. `device_type` is one of the
// devices defined in `types.h` (DEVICE_CPU, DEVICE_GPU, DEVICE_TPU, etc).
// We only collect metrics for CPU devices. This is a heuristic to avoid
// collecting metrics for device-side iterators created by the multi-device
// iterator mechanism.
IteratorMetricsCollector(const std::string& device_type, const Env& env);
// Starts the timer for the next `GetNext` call. Returns the start time.
absl::Time RecordStart();
// Records metrics for the most recent `GetNext` call, including the latency,
// bytes fetched, iterator life time, etc. `start_time` is the start time
// returned by `RecordStart`. `output` is the output of the `GetNext` call.
void RecordStop(absl::Time start_time, const std::vector<Tensor>& output);
private:
// We only collect metrics for CPU devices.
bool ShouldCollectMetrics() const;
// One of the devices defined in `types.h`
// (DEVICE_CPU, DEVICE_GPU, DEVICE_TPU, etc).
const std::string device_type_;
const Env& env_;
mutex mu_;
// Records the number of currently active `GetNext` calls.
uint64_t num_active_calls_ TF_GUARDED_BY(mu_) = 0;
// Records the start time (in microseconds) of the first `RecordStart()` call
// that followed the last period of inactivity.
uint64_t first_start_time_us_ TF_GUARDED_BY(mu_) = 0;
// Records the end time (in microseconds) of the most recent `RecordStop()`
// call.
uint64_t end_time_us_ TF_GUARDED_BY(mu_) = 0;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_METRIC_UTILS_H_
tensorflow/core/data/metric_utils_test.cc
+170
View File
@@ -0,0 +1,170 @@
/* Copyright 2022 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/core/data/metric_utils.h"
#include <cstdint>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/lib/monitoring/test_utils.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace data {
namespace {
using tensorflow::monitoring::testing::CellReader;
using tensorflow::monitoring::testing::Histogram;
TEST(MetricUtilsTest, CollectMetrics) {
CellReader<Histogram> latency("/tensorflow/data/getnext_duration");
CellReader<int64_t> iterator_lifetime("/tensorflow/data/iterator_lifetime");
CellReader<int64_t> iterator_busy("/tensorflow/data/iterator_busy");
EXPECT_FLOAT_EQ(latency.Delta().num(), 0.0);
EXPECT_EQ(iterator_lifetime.Delta(), 0);
EXPECT_EQ(iterator_busy.Delta(), 0);
IteratorMetricsCollector metrics_collector(DEVICE_CPU, *Env::Default());
absl::Time start_time = metrics_collector.RecordStart();
absl::SleepFor(absl::Seconds(1));
metrics_collector.RecordStop(start_time, /*output=*/{});
Histogram latency_histogram = latency.Delta();
EXPECT_FLOAT_EQ(latency_histogram.num(), 1.0);
EXPECT_GT(latency_histogram.sum(), 0.0);
EXPECT_GT(iterator_lifetime.Delta(), 0);
EXPECT_GT(iterator_busy.Delta(), 0);
}
TEST(MetricUtilsTest, ShouldNotCollectMetrics) {
CellReader<Histogram> latency("/tensorflow/data/getnext_duration");
CellReader<int64_t> iterator_lifetime("/tensorflow/data/iterator_lifetime");
CellReader<int64_t> iterator_busy("/tensorflow/data/iterator_busy");
EXPECT_FLOAT_EQ(latency.Delta().num(), 0.0);
EXPECT_EQ(iterator_lifetime.Delta(), 0);
EXPECT_EQ(iterator_busy.Delta(), 0);
IteratorMetricsCollector metrics_collector(DEVICE_TPU, *Env::Default());
absl::Time start_time = metrics_collector.RecordStart();
absl::SleepFor(absl::Seconds(1));
metrics_collector.RecordStop(start_time, /*output=*/{});
EXPECT_FLOAT_EQ(latency.Delta().num(), 0.0);
EXPECT_EQ(iterator_lifetime.Delta(), 0);
EXPECT_EQ(iterator_busy.Delta(), 0);
}
TEST(MetricUtilsTest, ConcurrentThreads) {
CellReader<Histogram> latency("/tensorflow/data/getnext_duration");
CellReader<int64_t> iterator_lifetime("/tensorflow/data/iterator_lifetime");
CellReader<int64_t> iterator_busy("/tensorflow/data/iterator_busy");
EXPECT_FLOAT_EQ(latency.Delta().num(), 0.0);
EXPECT_EQ(iterator_lifetime.Delta(), 0);
EXPECT_EQ(iterator_busy.Delta(), 0);
IteratorMetricsCollector metrics_collector(DEVICE_CPU, *Env::Default());
absl::Time start_time = metrics_collector.RecordStart();
auto thread = absl::WrapUnique(Env::Default()->StartThread(
/*thread_options=*/{}, /*name=*/"Concurrent metric collection thread",
[&metrics_collector]() {
absl::Time concurrent_start_time = metrics_collector.RecordStart();
absl::SleepFor(absl::Seconds(1));
metrics_collector.RecordStop(concurrent_start_time, /*output=*/{});
}));
absl::SleepFor(absl::Seconds(1));
metrics_collector.RecordStop(start_time, /*output=*/{});
thread.reset();
Histogram latency_histogram = latency.Delta();
EXPECT_FLOAT_EQ(latency_histogram.num(), 2.0);
EXPECT_GT(latency_histogram.sum(),
absl::ToInt64Microseconds(absl::Seconds(2)));
// The iterator busy time and lifetime do not count the latency twice.
EXPECT_GE(iterator_lifetime.Delta(),
absl::ToInt64Microseconds(absl::Seconds(1)));
EXPECT_LT(iterator_lifetime.Delta(),
absl::ToInt64Microseconds(absl::Seconds(1.5)));
EXPECT_GE(iterator_busy.Delta(), absl::ToInt64Microseconds(absl::Seconds(1)));
EXPECT_LT(iterator_busy.Delta(),
absl::ToInt64Microseconds(absl::Seconds(1.5)));
}
TEST(MetricUtilsTest, OverlappingThreads) {
CellReader<Histogram> latency("/tensorflow/data/getnext_duration");
CellReader<int64_t> iterator_lifetime("/tensorflow/data/iterator_lifetime");
CellReader<int64_t> iterator_busy("/tensorflow/data/iterator_busy");
EXPECT_FLOAT_EQ(latency.Delta().num(), 0.0);
EXPECT_EQ(iterator_lifetime.Delta(), 0);
EXPECT_EQ(iterator_busy.Delta(), 0);
// Two overlapping threads collect metrics:
// Thread 1 (end - start = 1 sec): |---------|
// Thread 2 (end - start = 2 sec): |------------------|
// Overlap: 0.5 sec.
// Iterator busy time: 2.5 sec.
IteratorMetricsCollector metrics_collector(DEVICE_CPU, *Env::Default());
absl::Time start_time = metrics_collector.RecordStart();
absl::SleepFor(absl::Seconds(0.5));
auto thread = absl::WrapUnique(Env::Default()->StartThread(
/*thread_options=*/{}, /*name=*/"Concurrent metric collection thread",
[&metrics_collector]() {
absl::Time concurrent_start_time = metrics_collector.RecordStart();
absl::SleepFor(absl::Seconds(2));
metrics_collector.RecordStop(concurrent_start_time, /*output=*/{});
}));
absl::SleepFor(absl::Seconds(0.5));
metrics_collector.RecordStop(start_time, /*output=*/{});
absl::SleepFor(absl::Seconds(1.5));
thread.reset();
Histogram latency_histogram = latency.Delta();
EXPECT_FLOAT_EQ(latency_histogram.num(), 2.0);
EXPECT_GT(latency_histogram.sum(),
absl::ToInt64Microseconds(absl::Seconds(3)));
// The iterator busy time and lifetime should not count the overlap twice.
EXPECT_GE(iterator_lifetime.Delta(),
absl::ToInt64Microseconds(absl::Seconds(2.5)));
EXPECT_LT(iterator_lifetime.Delta(),
absl::ToInt64Microseconds(absl::Seconds(2.9)));
EXPECT_GE(iterator_busy.Delta(),
absl::ToInt64Microseconds(absl::Seconds(2.5)));
EXPECT_LT(iterator_busy.Delta(),
absl::ToInt64Microseconds(absl::Seconds(2.9)));
}
TEST(MetricUtilsTest, RecordBytesFetched) {
CellReader<int64_t> bytes_fetched("/tensorflow/data/bytes_fetched");
IteratorMetricsCollector metrics_collector(DEVICE_CPU, *Env::Default());
absl::Time start_time = metrics_collector.RecordStart();
std::vector<Tensor> output;
output.push_back(Tensor(DT_FLOAT, TensorShape({10})));
output.push_back(Tensor(DT_INT32, TensorShape({5, 2})));
metrics_collector.RecordStop(start_time, output);
EXPECT_EQ(bytes_fetched.Delta(), 80);
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/name_utils.cc
+84
View File
@@ -0,0 +1,84 @@
/* 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 "tensorflow/core/data/name_utils.h"
#include <vector>
#include "absl/base/attributes.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
namespace tensorflow {
namespace data {
namespace name_utils {
ABSL_CONST_INIT const char kDelimiter[] = "::";
ABSL_CONST_INIT const char kDefaultDatasetDebugStringPrefix[] = "";
constexpr char kDataset[] = "Dataset";
constexpr char kOp[] = "Op";
constexpr char kVersion[] = "V";
std::string OpName(const std::string& dataset_type) {
return OpName(dataset_type, OpNameParams());
}
std::string OpName(const std::string& dataset_type,
const OpNameParams& params) {
if (params.op_version == 1) {
return absl::StrCat(dataset_type, kDataset);
}
return absl::StrCat(dataset_type, kDataset, kVersion, params.op_version);
}
std::string ArgsToString(const std::vector<std::string>& args) {
if (args.empty()) {
return "";
}
return absl::StrCat("(", absl::StrJoin(args, ", "), ")");
}
std::string DatasetDebugString(const std::string& dataset_type) {
return DatasetDebugString(dataset_type, DatasetDebugStringParams());
}
std::string DatasetDebugString(const std::string& dataset_type,
const DatasetDebugStringParams& params) {
OpNameParams op_name_params;
op_name_params.op_version = params.op_version;
std::string op_name = OpName(dataset_type, op_name_params);
return strings::StrCat(op_name, kOp, ArgsToString(params.args), kDelimiter,
params.dataset_prefix, kDataset);
}
std::string IteratorPrefix(const std::string& dataset_type,
const std::string& prefix) {
return IteratorPrefix(dataset_type, prefix, IteratorPrefixParams());
}
std::string IteratorPrefix(const std::string& dataset_type,
const std::string& prefix,
const IteratorPrefixParams& params) {
if (params.op_version == 1) {
return absl::StrCat(prefix, kDelimiter, params.dataset_prefix,
dataset_type);
}
return strings::StrCat(prefix, kDelimiter, params.dataset_prefix,
dataset_type, kVersion, params.op_version);
}
} // namespace name_utils
} // namespace data
} // namespace tensorflow
tensorflow/core/data/name_utils.h
+111
View File
@@ -0,0 +1,111 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_NAME_UTILS_H_
#define TENSORFLOW_CORE_DATA_NAME_UTILS_H_
#include <vector>
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace data {
namespace name_utils {
extern const char kDelimiter[];
extern const char kDefaultDatasetDebugStringPrefix[];
struct OpNameParams {
int op_version = 1;
};
struct DatasetDebugStringParams {
template <typename... T>
void set_args(T... input_args) {
args = {static_cast<const strings::AlphaNum&>(input_args).data()...};
}
int op_version = 1;
std::string dataset_prefix = "";
std::vector<std::string> args;
};
struct IteratorPrefixParams {
int op_version = 1;
std::string dataset_prefix = "";
};
// Merge the given args in the format of "(arg1, arg2, ..., argn)".
//
// e.g. ArgsToString({"1", "2", "3"}) -> "(1, 2, 3)"; ArgsToString({}) -> "".
std::string ArgsToString(const std::vector<std::string>& args);
// Returns the dataset op name.
//
// e.g. OpName("Map") -> "MapDataset".
std::string OpName(const std::string& dataset_type);
// Returns the dataset op names.
//
// e.g. OpName(ConcatenateDatasetOp::kDatasetType, OpNameParams())
// -> "ConcatenateDataset"
//
// OpNameParams params;
// params.op_version = 2;
// OpName(ParallelInterleaveDatasetOp::kDatasetType, params)
// -> "ParallelInterleaveDatasetV2"
std::string OpName(const std::string& dataset_type, const OpNameParams& params);
// Returns a human-readable debug string for this dataset in the format of
// "FooDatasetOp(arg1, arg2, ...)::Dataset".
//
// e.g. DatasetDebugString("Map") -> "MapDatasetOp::Dataset";
std::string DatasetDebugString(const std::string& dataset_type);
// Returns a human-readable debug string for this dataset in the format of
// "FooDatasetOp(arg1, arg2, ...)::Dataset".
//
// e.g.
// DatasetDebugStringParams range_params;
// range_params.set_args(0, 10, 3);
// DatasetDebugString(RangeDatasetOp::kDatasetType, range_params)
// -> "RangeDatasetOp(0, 10, 3)::Dataset");
std::string DatasetDebugString(const std::string& dataset_type,
const DatasetDebugStringParams& params);
// Returns a string that identifies the sequence of iterators leading up to
// the iterator of this dataset.
//
// e.g. IteratorPrefix("Map", "Iterator::Range") -> "Iterator::Range::Map".
std::string IteratorPrefix(const std::string& dataset_type,
const std::string& prefix);
// Returns a string that identifies the sequence of iterators leading up to
// the iterator of this dataset.
//
// e.g.
// IteratorPrefixParams params;
// params.op_version = 2;
// IteratorPrefix(BatchDatasetOp::KDatasetType, "Iterator::Range", params) ->
// "Iterator::Range::BatchV2".
std::string IteratorPrefix(const std::string& dataset_type,
const std::string& prefix,
const IteratorPrefixParams& params);
} // namespace name_utils
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_NAME_UTILS_H_
tensorflow/core/data/name_utils_test.cc
+62
View File
@@ -0,0 +1,62 @@
/* 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 "tensorflow/core/data/name_utils.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace data {
namespace {
TEST(DeviceNameUtils, ArgsToString) {
EXPECT_EQ(name_utils::ArgsToString({}), "");
EXPECT_EQ(name_utils::ArgsToString({"a"}), "(a)");
EXPECT_EQ(name_utils::ArgsToString({"1", "2", "3"}), "(1, 2, 3)");
}
TEST(NameUtilsTest, DatasetDebugString) {
EXPECT_EQ(name_utils::DatasetDebugString("Concatenate"),
"ConcatenateDatasetOp::Dataset");
name_utils::DatasetDebugStringParams range_params;
range_params.set_args(0, 10, 3);
EXPECT_EQ(name_utils::DatasetDebugString("Range", range_params),
"RangeDatasetOp(0, 10, 3)::Dataset");
name_utils::DatasetDebugStringParams shuffle_params;
shuffle_params.dataset_prefix = "FixedSeed";
shuffle_params.set_args(10, 1, 2);
EXPECT_EQ(name_utils::DatasetDebugString("Shuffle", shuffle_params),
"ShuffleDatasetOp(10, 1, 2)::FixedSeedDataset");
name_utils::DatasetDebugStringParams parallel_interleave_params;
parallel_interleave_params.op_version = 2;
EXPECT_EQ(name_utils::DatasetDebugString("ParallelInterleave",
parallel_interleave_params),
"ParallelInterleaveDatasetV2Op::Dataset");
}
TEST(NameUtilsTest, OpName) {
EXPECT_EQ(name_utils::OpName("Range"), "RangeDataset");
EXPECT_EQ(name_utils::OpName("Concatenate", name_utils::OpNameParams()),
"ConcatenateDataset");
name_utils::OpNameParams params;
params.op_version = 2;
EXPECT_EQ(name_utils::OpName("ParallelInterleave", params),
"ParallelInterleaveDatasetV2");
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/rewrite_utils.cc
+362
View File
@@ -0,0 +1,362 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/data/rewrite_utils.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/refcount.h"
// On mobile we do not provide this functionality because not all of its
// dependencies are available there.
#if !defined(IS_MOBILE_PLATFORM)
#include <algorithm>
#include <functional>
#include <map>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/substitute.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_runner.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/data/dataset_utils.h"
#include "tensorflow/core/data/hash_utils.h"
#include "tensorflow/core/data/serialization_utils.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def_util.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/grappler/optimizers/data/function_utils.h"
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/grappler/optimizers/meta_optimizer.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace data {
namespace {
constexpr char kOptimizerName[] = "tf_data_meta_optimizer";
constexpr char kOptimizers[] = "optimizers";
constexpr char kOptimizerConfigs[] = "optimizer_configs";
void AddFakeSinks(FunctionDef* function_def) {
int counter = 0;
for (const auto& output : function_def->signature().output_arg()) {
NodeDef* node = function_def->add_node_def();
tensorflow::grappler::function_utils::SetUniqueFunctionNodeName(
absl::StrCat("FakeSink", counter++), function_def, node);
node->set_op("Identity");
node->add_input(function_def->ret().at(output.name()));
(*node->mutable_attr())["T"].set_type(output.type());
(*function_def->mutable_ret())[output.name()] =
absl::StrCat(node->name(), ":output:0");
}
}
void RemoveFakeSinks(FunctionDef* function_def) {
// Map from identity node names to their input tensor strings
std::map<std::string, std::string> identity_map;
for (const auto& node : function_def->node_def()) {
if (node.op() == "Identity" && node.input_size() == 1) {
identity_map[node.name()] = node.input(0);
}
}
for (const auto& output_arg : function_def->signature().output_arg()) {
const std::string& tensor = function_def->ret().at(output_arg.name());
const std::string& output_node = tensor.substr(0, tensor.find(':'));
if (identity_map.find(output_node) != identity_map.end()) {
(*function_def->mutable_ret())[output_arg.name()] =
identity_map.at(output_node);
}
}
}
absl::Status ApplyRewrites(
OpKernelContext* ctx,
const std::function<RewriterConfig(void)> config_factory,
GraphDef* graph_def, std::string* dataset_node) {
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
GetGrapplerItem(graph_def, dataset_node, /*add_fake_sinks=*/true);
std::unordered_map<std::string, tensorflow::DeviceProperties> device_map;
tensorflow::grappler::VirtualCluster cluster(device_map);
// Run data optimizer using grappler's meta optimizer.
tensorflow::ConfigProto config;
*config.mutable_graph_options()->mutable_rewrite_options() = config_factory();
TF_RETURN_IF_ERROR(tensorflow::grappler::RunMetaOptimizer(
std::move(*grappler_item), config, ctx->device(), &cluster, graph_def));
// Remove fake sinks after optimizations are done.
//
// TODO(b/118820916): When MetaOptimizer adds provisions for function retvals
// to be optimizable, we will no longer need this.
for (auto& function_def : *graph_def->mutable_library()->mutable_function()) {
RemoveFakeSinks(&function_def);
}
return absl::OkStatus();
}
} // anonymous namespace
RewriterConfig CreateRewriterConfig(
const absl::flat_hash_set<tstring>& optimizations,
const absl::flat_hash_set<tstring>& optimizations_configs) {
RewriterConfig rewriter_config;
rewriter_config.add_optimizers(kOptimizerName);
rewriter_config.set_meta_optimizer_iterations(RewriterConfig::ONE);
rewriter_config.set_fail_on_optimizer_errors(true);
auto custom_optimizer = rewriter_config.add_custom_optimizers();
custom_optimizer->set_name(kOptimizerName);
auto* custom_optimizations_list =
(*custom_optimizer->mutable_parameter_map())[kOptimizers].mutable_list();
const auto& registered_optimizers =
grappler::CustomGraphOptimizerRegistry::GetRegisteredOptimizers();
for (const auto& optimization : optimizations) {
if (std::find(registered_optimizers.begin(), registered_optimizers.end(),
optimization) != registered_optimizers.end()) {
custom_optimizations_list->add_s(optimization.data(),
optimization.size());
} else {
VLOG(1) << "Optimization " << optimization << " is not registered.";
}
}
auto* config_list =
(*custom_optimizer->mutable_parameter_map())[kOptimizerConfigs]
.mutable_list();
for (const auto& config : optimizations_configs) {
config_list->add_s(config.data(), config.size());
}
return rewriter_config;
}
absl::Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input,
std::function<RewriterConfig(void)> config_factory,
bool record_fingerprint,
core::RefCountPtr<DatasetBase>* rewritten_input) {
std::vector<std::pair<std::string, Tensor>> input_list;
GraphDef graph_def;
std::string output_node;
TF_RETURN_IF_ERROR(
AsGraphDefForRewrite(ctx, input, &input_list, &graph_def, &output_node));
VLOG(3) << "Before graph rewrites: " << graph_def.DebugString();
TF_RETURN_IF_ERROR(
ApplyRewrites(ctx, config_factory, &graph_def, &output_node));
VLOG(3) << "After graph rewrites: " << graph_def.DebugString();
// Instantiate the optimized input pipeline by running the optimized graph
// using the optimized function library.
FunctionLibraryRuntime* flr = nullptr;
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr = nullptr;
std::unique_ptr<FunctionLibraryDefinition> lib_def = nullptr;
TF_RETURN_IF_ERROR(
ctx->function_library()->Clone(&lib_def, &pflr, &flr, true));
// Some functions may have been modified without having their names changed
// (for example, nested dataset graphs from FlatMap or Interleave).
TF_RETURN_IF_ERROR(AddToFunctionLibrary(lib_def.get(), graph_def.library()));
Graph graph(OpRegistry::Global());
TF_RETURN_IF_ERROR(ImportGraphDef({}, graph_def, &graph, nullptr));
std::vector<Tensor> outputs;
GraphRunner graph_runner(flr->device());
TF_RETURN_IF_ERROR(
graph_runner.Run(&graph, flr, input_list, {output_node}, &outputs));
DatasetBase* rewritten_dataset;
TF_RETURN_IF_ERROR(
GetDatasetFromVariantTensor(outputs[0], &rewritten_dataset));
rewritten_dataset->Ref();
rewritten_input->reset(rewritten_dataset);
if (record_fingerprint) {
(*ctx->runner())([graph_def = std::move(graph_def),
lib_def = lib_def.release(),
input_list = std::move(input_list),
output_node = std::move(output_node)]() {
std::unique_ptr<FunctionLibraryDefinition> lib_def_owner(lib_def);
const NodeDef* node_def = nullptr;
for (const auto& node : graph_def.node()) {
if (node.name() == output_node) {
node_def = &node;
break;
}
}
if (node_def == nullptr) {
VLOG(3) << "Failed to find node: " << output_node;
return;
}
uint64_t hash = 0;
absl::Status s = HashNode(graph_def, *node_def, *lib_def, &hash);
if (!s.ok()) {
VLOG(3) << "Failed to hash graph: " << s;
return;
}
for (const auto& pair : input_list) {
hash = Hash64CombineUnordered(hash, Hash64(pair.first));
uint64_t tensor_hash = 0;
absl::Status s = HashTensor(pair.second, &tensor_hash);
if (s.ok()) {
hash = Hash64CombineUnordered(hash, tensor_hash);
} else {
VLOG(3) << "Failed to hash tensor: " << s;
}
}
std::string graph_hash = absl::StrCat(absl::Hex(hash, absl::kZeroPad16));
metrics::RecordTFDataFingerprint(graph_hash);
});
}
return absl::OkStatus();
}
std::unique_ptr<tensorflow::grappler::GrapplerItem> GetGrapplerItem(
GraphDef* graph_def, std::string* dataset_node, bool add_fake_sinks,
bool apply_optimizations) {
// Add an identity node as the fetch node, otherwise we might get 'placeholder
// is both fed and fetched' errors in some cases when using input list with
// placeholder dataset nodes.
NodeDef* node = graph_def->mutable_node()->Add();
tensorflow::grappler::graph_utils::SetUniqueGraphNodeName("Sink", graph_def,
node);
node->set_op("Identity");
node->add_input(*dataset_node);
(*node->mutable_attr())["T"].set_type(DT_VARIANT);
*dataset_node = node->name();
if (add_fake_sinks) {
// Add fake sink node to graph and functions to allow rewriting the actual
// sink nodes.
//
// TODO(b/118820916): When MetaOptimizer adds provisions for function
// retvals to be optimizable, we will no longer need this.
for (auto& function_def :
*graph_def->mutable_library()->mutable_function()) {
AddFakeSinks(&function_def);
}
}
// Create metagraph.
MetaGraphDef meta_graph_def;
(*meta_graph_def.mutable_graph_def()) = *graph_def;
// Grappler determines fetch ops from collection 'train_op'.
CollectionDef collection_def;
auto node_list = collection_def.mutable_node_list();
node_list->add_value(*dataset_node);
(*meta_graph_def.mutable_collection_def())["train_op"] = collection_def;
// Create Grappler item.
tensorflow::grappler::ItemConfig item_config;
item_config.apply_optimizations = apply_optimizations;
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
tensorflow::grappler::GrapplerItemFromMetaGraphDef(
"graph", meta_graph_def, item_config);
// Grappler should not optimize function library of tf.data graphs. The
// tf.data meta optimizer takes care of optimizing tf.data functions.
grappler_item->optimization_options().optimize_function_library = false;
return grappler_item;
}
absl::flat_hash_set<tstring> SelectOptimizations(
const absl::flat_hash_set<std::string>& experiments,
const absl::flat_hash_set<tstring>& optimizations_enabled,
const absl::flat_hash_set<tstring>& optimizations_disabled,
const absl::flat_hash_set<tstring>& optimizations_default) {
absl::flat_hash_set<tstring> optimizations;
// Add the enabled optimizations.
optimizations.insert(optimizations_enabled.begin(),
optimizations_enabled.end());
// Add all default optimization that are not disabled.
for (const auto& optimization : optimizations_default) {
if (!optimizations_disabled.contains(optimization)) {
optimizations.insert(optimization);
}
}
// Add experiments that correspond to an optimization unless the optimization
// is disabled.
const auto& registered_optimizers =
grappler::CustomGraphOptimizerRegistry::GetRegisteredOptimizers();
for (const auto& experiment : experiments) {
if (std::find(registered_optimizers.begin(), registered_optimizers.end(),
experiment) != registered_optimizers.end() &&
!optimizations_disabled.contains(experiment)) {
optimizations.insert(experiment);
}
}
return optimizations;
}
absl::StatusOr<std::string> GetDatasetNode(const GraphDef& graph_def) {
// Symbolic `_Retval` node indicates which node corresponds to the dataset.
for (const auto& node : graph_def.node()) {
if (node.op() == kRetvalOp) {
return node.input(0);
}
}
return absl::NotFoundError(
absl::Substitute("Dataset node for graph is not found:\n$0",
graph_def.ShortDebugString()));
}
absl::StatusOr<NodeDef> GetDatasetNodeDef(const GraphDef& graph_def) {
TF_ASSIGN_OR_RETURN(std::string dataset_node_name, GetDatasetNode(graph_def));
for (const auto& node : graph_def.node()) {
if (node.name() == dataset_node_name) {
return node;
}
}
return absl::NotFoundError(
absl::Substitute("Dataset node for graph is not found:\n$0",
graph_def.ShortDebugString()));
}
} // namespace data
} // namespace tensorflow
#endif // !IS_MOBILE_PLATFORM
tensorflow/core/data/rewrite_utils.h
+93
View File
@@ -0,0 +1,93 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_REWRITE_UTILS_H_
#define TENSORFLOW_CORE_DATA_REWRITE_UTILS_H_
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/core/platform/platform.h"
// On mobile we do not provide this functionality because not all of its
// dependencies are available there.
#if !defined(IS_MOBILE_PLATFORM)
#include <functional>
#include <memory>
#include <string>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace data {
RewriterConfig CreateRewriterConfig(
const absl::flat_hash_set<tstring>& optimizations,
const absl::flat_hash_set<tstring>& optimizations_configs);
// Rewrites the input dataset using the given config. The rewritten_input
// stored in the core::RefCountPtr<DatasetBase>* output parameter is owned.
absl::Status RewriteDataset(OpKernelContext* ctx, const DatasetBase* input,
std::function<RewriterConfig(void)> config_factory,
bool record_fingerprint,
core::RefCountPtr<DatasetBase>* rewritten_input);
// Creates a grappler item for `graph_def`, which is required for graph
// optimization.
// `dataset_node` is the name of the node corresponding to the dataset.
// If `add_fake_sinks` is true, it adds fake sink node to graph and functions to
// allow rewriting the actual sink nodes.
// If `apply_optimizations` is true, general grappler optimizations at level
// `tensorflow::OptimizerOptions::L1` are applied to the graph.
// TODO(b/118820916): When MetaOptimizer adds provisions for function retvals to
// be optimizable, we will no longer need to add fake nodes.
std::unique_ptr<tensorflow::grappler::GrapplerItem> GetGrapplerItem(
GraphDef* graph_def, std::string* dataset_node, bool add_fake_sinks,
bool apply_optimizations = true);
// Returns the name of the node corresponding to the dataset. It is indicated by
// the symbolic `_Retval` node.
absl::StatusOr<std::string> GetDatasetNode(const GraphDef& graph_def);
// Like `GetDatasetNode` above, but returns the entire node object.
absl::StatusOr<NodeDef> GetDatasetNodeDef(const GraphDef& graph_def);
// Determines which optimizations should be applied.
//
// The result will contain any optimizations that are explicitly enabled, any
// default optimization that are not explicitly disabled, and any experiment
// that corresponds to an optimization as long as the optimization is not
// explicitly disabled.
absl::flat_hash_set<tstring> SelectOptimizations(
const absl::flat_hash_set<std::string>& experiments,
const absl::flat_hash_set<tstring>& optimizations_enabled,
const absl::flat_hash_set<tstring>& optimizations_disabled,
const absl::flat_hash_set<tstring>& optimizations_default);
} // namespace data
} // namespace tensorflow
#endif // !IS_MOBILE_PLATFORM
#endif // TENSORFLOW_CORE_DATA_REWRITE_UTILS_H_
tensorflow/core/data/rewrite_utils_test.cc
+153
View File
@@ -0,0 +1,153 @@
/* Copyright 2021 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/core/data/rewrite_utils.h"
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::test::AsScalar;
using ::tensorflow::test::function::GDef;
using ::tensorflow::test::function::NDef;
using ::testing::ElementsAre;
NodeDef GetMapNode(absl::string_view name, absl::string_view input_node_name,
absl::string_view function_name) {
return NDef(
name, /*op=*/"MapDataset", {std::string(input_node_name)},
{{"f", FunctionDefHelper::FunctionRef(std::string(function_name))},
{"Targuments", {}},
{"output_shapes", absl::Span<const TensorShape>{TensorShape()}},
{"output_types", absl::Span<const DataType>{DT_INT64}}});
}
FunctionDef XTimesX() {
return FunctionDefHelper::Create(
/*function_name=*/"XTimesX",
/*in_def=*/{"x: int64"},
/*out_def=*/{"y: int64"},
/*attr_def=*/{},
/*node_def=*/{{{"y"}, "Mul", {"x", "x"}, {{"T", DT_INT64}}}},
/*ret_def=*/{{"y", "y:z:0"}});
}
GraphDef GetRangeSquareDatasetDef(const int64_t range) {
return GDef(
{NDef("start", "Const", /*inputs=*/{},
{{"value", AsScalar<int64_t>(0)}, {"dtype", DT_INT64}}),
NDef("stop", "Const", /*inputs=*/{},
{{"value", AsScalar<int64_t>(range)}, {"dtype", DT_INT64}}),
NDef("step", "Const", /*inputs=*/{},
{{"value", AsScalar<int64_t>(1)}, {"dtype", DT_INT64}}),
NDef("range", "RangeDataset", /*inputs=*/{"start", "stop", "step"},
{{"output_shapes", absl::Span<const TensorShape>{TensorShape()}},
{"output_types", absl::Span<const DataType>{DT_INT64}}}),
GetMapNode("map", "range", "XTimesX"),
NDef("dataset", "_Retval", /*inputs=*/{"map"},
{{"T", DT_VARIANT}, {"index", 0}})},
{XTimesX()});
}
TEST(GraphUtilTest, GetFetchNode) {
GraphDef graph = GetRangeSquareDatasetDef(10);
TF_ASSERT_OK_AND_ASSIGN(std::string dataset_node, GetDatasetNode(graph));
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
GetGrapplerItem(&graph, &dataset_node, /*add_fake_sinks=*/false);
EXPECT_THAT(grappler_item->fetch, ElementsAre("Sink"));
}
TEST(GraphUtilTest, GetFetchNodeDef) {
GraphDef graph = GetRangeSquareDatasetDef(10);
TF_ASSERT_OK_AND_ASSIGN(NodeDef dataset_nodedef, GetDatasetNodeDef(graph));
std::string dataset_node = dataset_nodedef.name();
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
GetGrapplerItem(&graph, &dataset_node, /*add_fake_sinks=*/false);
EXPECT_THAT(grappler_item->fetch, ElementsAre("Sink"));
}
struct SelectOptimizationsTestCase {
absl::flat_hash_set<std::string> experiments;
absl::flat_hash_set<tstring> optimizations_enabled;
absl::flat_hash_set<tstring> optimizations_disabled;
absl::flat_hash_set<tstring> optimizations_default;
std::vector<std::string> expected;
};
class SelectOptimizationsTest
: public ::testing::TestWithParam<SelectOptimizationsTestCase> {};
TEST_P(SelectOptimizationsTest, DatasetUtils) {
const SelectOptimizationsTestCase test_case = GetParam();
auto optimizations = SelectOptimizations(
test_case.experiments, test_case.optimizations_enabled,
test_case.optimizations_disabled, test_case.optimizations_default);
EXPECT_THAT(
std::vector<std::string>(optimizations.begin(), optimizations.end()),
::testing::UnorderedElementsAreArray(test_case.expected));
}
INSTANTIATE_TEST_SUITE_P(
Test, SelectOptimizationsTest,
::testing::Values(
SelectOptimizationsTestCase{
/*experiments=*/{}, /*optimizations_enabled=*/{},
/*optimizations_disabled=*/{}, /*optimizations_default=*/{},
/*expected=*/{}},
SelectOptimizationsTestCase{
/*experiments=*/{"map_and_batch_fusion"},
/*optimizations_enabled=*/{"bar"},
/*optimizations_disabled=*/{}, /*optimizations_default=*/{"baz"},
/*expected=*/{"map_and_batch_fusion", "bar", "baz"}},
SelectOptimizationsTestCase{
/*experiments=*/{"this_is_not_an_optimization"},
/*optimizations_enabled=*/{"bar"},
/*optimizations_disabled=*/{}, /*optimizations_default=*/{"baz"},
/*expected=*/{"bar", "baz"}},
SelectOptimizationsTestCase{/*experiments=*/{},
/*optimizations_enabled=*/{"foo"},
/*optimizations_disabled=*/{"baz"},
/*optimizations_default=*/{"bar", "baz"},
/*expected=*/{"foo", "bar"}},
SelectOptimizationsTestCase{
/*experiments=*/{"foo"}, /*optimizations_enabled=*/{"bar"},
/*optimizations_disabled=*/{"foo"},
/*optimizations_default=*/{"baz"}, /*expected=*/{"bar", "baz"}}));
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/root_dataset.cc
+530
View File
@@ -0,0 +1,530 @@
/* Copyright 2021 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/core/data/root_dataset.h"
#include <algorithm>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "tensorflow/core/data/dataset_utils.h"
#include "tensorflow/core/data/name_utils.h"
#include "tensorflow/core/data/rewrite_utils.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/model.h"
#include "tensorflow/core/framework/model.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/host_info.h"
#include "tensorflow/core/platform/mem.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringprintf.h"
#include "tsl/platform/host_info.h"
namespace tensorflow {
namespace data {
namespace {
constexpr char kDatasetType[] = "Root";
constexpr char kAlgorithm[] = "algorithm";
constexpr char kCpuBudget[] = "cpu_budget";
constexpr char kExperiments[] = "experiments";
constexpr char kReadRoundtripLatency[] = "read_latency_usec";
constexpr char kReadResponseBytes[] = "read_bytes";
constexpr char kIntraOpParallelism[] = "intra_op_parallelism";
constexpr char kMemBandwidth[] = "mem_bw_used_megabytes_per_sec";
constexpr char kPrivateThreadpoolSize[] = "threadpool_size";
constexpr char kRamBudget[] = "ram_budget_megabytes";
constexpr char kRamUsage[] = "ram_usage_megabytes";
constexpr char kMaxBufferBytes[] = "max_buffered_megabytes";
constexpr char kWarmStart[] = "warm_start";
// If value `x` matches `y`, returns default value `z`. Otherwise, return `x`.
inline int64_t value_or_default(int64_t x, int64_t y, int64_t z) {
return x == y ? z : x;
}
void SetRootDatasetParams(const Options& options, RootDataset::Params* params) {
if (ShouldConfigureMaxIntraOpParallelism(options)) {
params->max_intra_op_parallelism =
options.threading_options().max_intra_op_parallelism();
}
if (ShouldUsePrivateThreadPool(options)) {
params->private_threadpool_size =
options.threading_options().private_threadpool_size();
}
params->autotune = ShouldUseAutotuning(options);
params->autotune_algorithm = model::AutotuneAlgorithm::DEFAULT;
auto experiments = GetExperiments();
if (experiments.contains("stage_based_autotune") ||
experiments.contains("stage_based_autotune_v2")) {
params->autotune_algorithm = model::AutotuneAlgorithm::STAGE_BASED;
}
if (options.autotune_options().optional_autotune_algorithm_case() ==
AutotuneOptions::kAutotuneAlgorithm) {
params->autotune_algorithm =
options.autotune_options().autotune_algorithm();
}
int64_t cpu_budget_from_options = options.autotune_options().cpu_budget();
if (cpu_budget_from_options == 0) {
params->autotune_cpu_budget_func = [] { return GetCpuBudget(); };
} else {
params->autotune_cpu_budget_func = [cpu_budget_from_options] {
return cpu_budget_from_options;
};
}
params->autotune_ram_budget_from_options =
options.autotune_options().ram_budget();
double ram_budget_share;
if (experiments.contains("autotune_buffer_optimization")) {
// When running this experiment, increase the ram_budget since it already
// takes into account the ram usage in buffer sizing, which is not the
// case for prefetch autotuner. Without this, we see degradation in some
// jobs for lack of buffers while ram usage is low.
ram_budget_share = 0.9;
} else {
ram_budget_share = model::kRamBudgetShare;
}
params->ram_budget_share = ram_budget_share;
}
void AddTraceMetadata(const RootDataset::Params& params, const Options& options,
TraceMeMetadata* trace_metadata) {
if (params.autotune) {
trace_metadata->push_back(std::make_pair(
kAlgorithm, model::AutotuneAlgorithm_Name(params.autotune_algorithm)));
trace_metadata->push_back(std::make_pair(
kCpuBudget,
absl::StrFormat("%lld", static_cast<long long>(
params.autotune_cpu_budget_func()))));
int64_t ram_budget = params.ComputeInitialAutotuneRamBudget();
trace_metadata->push_back(std::make_pair(
kRamBudget,
absl::StrFormat("%lld", static_cast<long long>(ram_budget / 1.0e6))));
}
if (params.max_intra_op_parallelism >= 0) {
trace_metadata->push_back(std::make_pair(
kIntraOpParallelism,
absl::StrFormat("%lld", static_cast<long long>(value_or_default(
params.max_intra_op_parallelism, 0,
port::MaxParallelism())))));
}
if (params.private_threadpool_size >= 0) {
trace_metadata->push_back(std::make_pair(
kPrivateThreadpoolSize,
absl::StrFormat("%lld", static_cast<long long>(value_or_default(
params.private_threadpool_size, 0,
port::MaxParallelism())))));
}
auto experiments = GetExperiments();
if (!experiments.empty()) {
trace_metadata->push_back(
std::make_pair(kExperiments, absl::StrJoin(experiments, " ")));
}
trace_metadata->push_back(
std::make_pair(kWarmStart, options.warm_start() ? "true" : "false"));
}
} // namespace
// static
absl::Status RootDataset::FromOptions(const DatasetBase* input,
DatasetBase** output) {
Params params;
SetRootDatasetParams(input->options(), &params);
*output = new RootDataset(input, params);
(*output)->Initialize(input->metadata());
for (const auto& framework : input->options().framework_type()) {
metrics::RecordTFDataFrameworkType(framework);
}
return absl::OkStatus();
}
absl::Status RootDataset::FromOptions(core::RefCountPtr<DatasetBase> input,
DatasetBase** output) {
Params params;
for (const auto& framework : input->options().framework_type()) {
metrics::RecordTFDataFrameworkType(framework);
}
SetRootDatasetParams(input->options(), &params);
Metadata metadata = input->metadata();
*output = new RootDataset(std::move(input), params);
(*output)->Initialize(metadata);
return absl::OkStatus();
}
class RootDataset::Iterator : public DatasetIterator<RootDataset> {
public:
explicit Iterator(const Params& params)
: DatasetIterator<RootDataset>(params) {
if (dataset()->params_.max_intra_op_parallelism >= 0) {
max_intra_op_parallelism_ =
value_or_default(dataset()->params_.max_intra_op_parallelism, 0,
port::MaxParallelism());
}
if (dataset()->params_.private_threadpool_size >= 0) {
threadpool_size_ =
value_or_default(dataset()->params_.private_threadpool_size, 0,
port::MaxParallelism());
thread_pool_ = std::make_unique<thread::ThreadPool>(
Env::Default(), ThreadOptions{}, "data_private_threadpool",
threadpool_size_);
}
cancellation_manager_ = std::make_unique<CancellationManager>();
}
~Iterator() override { cancellation_manager_->StartCancel(); }
bool SymbolicCheckpointCompatible() const override { return true; }
absl::Status Initialize(IteratorContext* ctx) override {
// prefetch_autotuner.h currently disregards `autotune` parameter
// so no matter whether dataset()->params_.autotune is on or not
// we need to pass ram_budget_manager_ to the downstream dataset operations
ram_budget_manager_ = std::make_shared<model::RamBudgetManager>(
dataset()->params_.ComputeInitialAutotuneRamBudget());
if (dataset()->params_.autotune) {
if (ctx->model() != nullptr) {
model_ = ctx->model();
} else {
model_ = std::make_shared<model::Model>();
ctx->SetModel(model_);
}
absl::flat_hash_set<std::string> experiments = GetExperiments();
if (experiments.contains("stage_based_autotune_v2")) {
model_->AddExperiment("stage_based_autotune_v2");
}
if (experiments.contains("autotune_buffer_optimization")) {
model_->AddExperiment("autotune_buffer_optimization");
}
}
IteratorContext iter_ctx(CreateParams(ctx));
if (model_) {
auto factory = [&iter_ctx, this](model::Node::Args args) {
return CreateNode(&iter_ctx, std::move(args));
};
model_->AddNode(std::move(factory), prefix(), nullptr, &node_);
}
TF_RETURN_IF_ERROR(dataset()->input_->MakeIterator(&iter_ctx, this,
prefix(), &input_impl_));
ctx->MergeCheckpoint(iter_ctx.checkpoint());
return absl::OkStatus();
}
absl::Status GetNextInternal(IteratorContext* ctx,
std::vector<Tensor>* out_tensors,
bool* end_of_sequence) override {
{
tf_shared_lock l(mu_);
if (model_ != nullptr && end_time_usec_ > 0) {
model_->RecordIteratorGapTime(ctx->env()->NowMicros() - end_time_usec_);
}
}
if (dataset()->params_.autotune) {
TF_RETURN_IF_ERROR(EnsureModelThreadStarted(ctx));
}
IteratorContext iter_ctx(CreateParams(ctx));
TF_RETURN_IF_ERROR(
input_impl_->GetNext(&iter_ctx, out_tensors, end_of_sequence));
ctx->MergeCheckpoint(iter_ctx.checkpoint());
{
mutex_lock l(mu_);
end_time_usec_ = std::max(ctx->env()->NowMicros(), end_time_usec_);
}
return absl::OkStatus();
}
protected:
std::shared_ptr<model::Node> CreateNode(
IteratorContext* ctx, model::Node::Args args) const override {
return model::MakeKnownRatioNode(std::move(args), /*ratio=*/1);
}
absl::Status SaveInternal(SerializationContext* ctx,
IteratorStateWriter* writer) override {
TF_RETURN_IF_ERROR(SaveInput(ctx, writer, input_impl_));
return absl::OkStatus();
}
absl::Status RestoreInternal(IteratorContext* ctx,
IteratorStateReader* reader) override {
IteratorContext iter_ctx(CreateParams(ctx));
TF_RETURN_IF_ERROR(RestoreInput(&iter_ctx, reader, input_impl_));
ctx->MergeCheckpoint(iter_ctx.checkpoint());
return absl::OkStatus();
}
TraceMeMetadata GetTraceMeMetadata() const override {
tensorflow::data::TraceMeMetadata traceme_metadata =
dataset()->traceme_metadata_;
const int64_t mem_bw = port::GetMemoryBandwidthInfo().bw_used;
if (mem_bw != INT64_MAX) {
traceme_metadata.push_back(std::make_pair(
kMemBandwidth,
absl::StrFormat("%lld", static_cast<long long>(mem_bw))));
}
const auto memory_info = port::GetMemoryInfo();
const auto memory_usage = memory_info.total - memory_info.free;
traceme_metadata.push_back(std::make_pair(
kRamUsage,
absl::StrFormat("%lld out of %lld (%.2f%%)",
static_cast<long long>(memory_usage / 1.0e6),
static_cast<long long>(memory_info.total / 1.0e6),
static_cast<double>(100 * memory_usage) /
static_cast<double>(memory_info.total))));
const auto io_statistics = tsl::port::GetIOStatistics();
if (io_statistics.roundtrip_latency_usec.count > 0) {
traceme_metadata.push_back(std::make_pair(
kReadRoundtripLatency,
absl::StrFormat(
"(count: %lld, mean: %lld, std dev: %lld)",
static_cast<long long>(
io_statistics.roundtrip_latency_usec.count),
static_cast<long long>(io_statistics.roundtrip_latency_usec.mean),
static_cast<long long>(
io_statistics.roundtrip_latency_usec.std_dev))));
}
if (io_statistics.response_bytes.count > 0) {
traceme_metadata.push_back(std::make_pair(
kReadResponseBytes,
absl::StrFormat(
"(count: %lld, mean: %lld, std dev: %lld)",
static_cast<long long>(io_statistics.response_bytes.count),
static_cast<long long>(io_statistics.response_bytes.mean),
static_cast<long long>(io_statistics.response_bytes.std_dev))));
}
return traceme_metadata;
}
private:
IteratorContext::Params CreateParams(IteratorContext* ctx) {
IteratorContext::Params params(ctx);
// prefetch_autotuner.h currently disregards `autotune` parameter
// so no matter whether dataset()->params_.autotune is on or not
// we need to pass ram_budget_manager_ to the downstream dataset operations
params.ram_budget_manager = ram_budget_manager_;
// After cl/548201925, `ctx->model()` may not be `nullptr` even when
// autotuning is off, but `model_` should be `nullptr` and it should have
// been set to a valid model in `Initialize()` if autotuning is on. We
// should simply set `params.model` to `model_` here.
params.model = model_;
if (dataset()->params_.private_threadpool_size >= 0) {
params.runner = [pool = thread_pool_.get()](std::function<void()> c) {
pool->Schedule(std::move(c));
};
params.runner_threadpool_size = threadpool_size_;
}
if (dataset()->params_.max_intra_op_parallelism >= 0) {
params.runner =
RunnerWithMaxParallelism(params.runner, max_intra_op_parallelism_);
}
params.options = &dataset()->options();
return params;
}
absl::Status EnsureModelThreadStarted(IteratorContext* ctx) {
mutex_lock l(mu_);
if (!model_thread_) {
RunMode run_mode = ctx->run_mode();
model_thread_ = ctx->StartThread("tf_data_model", [this, run_mode]() {
RootDataset::Params params = dataset()->params_;
std::function<int64_t(int64_t)> ram_budget_func;
std::optional<int64_t> raw_ram_budget;
if (params.autotune_ram_budget_from_options > 0) {
raw_ram_budget = params.autotune_ram_budget_from_options;
} else if (run_mode != RunMode::STANDALONE) {
// Dynamic RAM budget should only apply to tf.data service.
raw_ram_budget = params.ComputeInitialAutotuneRamBudget();
}
absl::Status status = model_->OptimizeLoop(
params.autotune_algorithm, params.autotune_cpu_budget_func,
params.ram_budget_share, raw_ram_budget, *ram_budget_manager_,
cancellation_manager_.get());
if (!status.ok()) {
LOG(WARNING) << "Optimization loop failed: " << status;
}
});
}
return absl::OkStatus();
}
std::shared_ptr<model::Model> model_ = nullptr;
// `ram_budget_manager_` coordinates the memory budget and allocation
// between prefetch legacy autotune and `tensorflow::data::model::Model`
std::shared_ptr<model::RamBudgetManager> ram_budget_manager_ = nullptr;
// Controls cancellation of `model_thread_`. Must be ordered before
// `model_thread_` so that `model_thread_` is destroyed first.
std::unique_ptr<CancellationManager> cancellation_manager_;
mutex mu_;
std::unique_ptr<Thread> model_thread_ TF_GUARDED_BY(mu_);
int64_t max_intra_op_parallelism_;
int64_t threadpool_size_;
std::unique_ptr<thread::ThreadPool> thread_pool_;
// The end time of the previous `GetNextInternal` call.
uint64_t end_time_usec_ TF_GUARDED_BY(mu_) = 0;
// Must be ordered last as its execution may depend on other members.
std::unique_ptr<IteratorBase> input_impl_;
};
RootDataset::RootDataset(const DatasetBase* input, const Params& params)
: DatasetBase(DatasetContext({name_utils::OpName(kDatasetType),
name_utils::OpName(kDatasetType)})),
input_(input),
params_(std::move(params)) {
AddTraceMetadata(params_, input_->options(), &traceme_metadata_);
}
RootDataset::RootDataset(core::RefCountPtr<DatasetBase> input,
const Params& params)
: DatasetBase(DatasetContext({name_utils::OpName(kDatasetType),
name_utils::OpName(kDatasetType)})),
params_(std::move(params)) {
owned_input_ = std::move(input);
input_ = owned_input_.get();
random_indexing_compatible_ = absl::OkStatus();
if (input_ != nullptr) {
random_indexing_compatible_ = input_->RandomIndexingCompatible();
}
AddTraceMetadata(params_, input_->options(), &traceme_metadata_);
}
RootDataset::~RootDataset() = default;
std::unique_ptr<IteratorBase> RootDataset::MakeIteratorInternal(
const std::string& prefix) const {
return std::make_unique<Iterator>(
Iterator::Params{this, name_utils::IteratorPrefix(kDatasetType, prefix)});
}
const DataTypeVector& RootDataset::output_dtypes() const {
return input_->output_dtypes();
}
const std::vector<PartialTensorShape>& RootDataset::output_shapes() const {
return input_->output_shapes();
}
std::string RootDataset::DebugString() const {
return name_utils::DatasetDebugString(kDatasetType);
}
int64_t RootDataset::CardinalityInternal(CardinalityOptions options) const {
return input_->Cardinality(options);
}
absl::Status RootDataset::Get(OpKernelContext* ctx, int64_t index,
std::vector<Tensor>* out_tensors) const {
std::vector<const DatasetBase*> inputs;
TF_RETURN_IF_ERROR(this->InputDatasets(&inputs));
return inputs[0]->Get(ctx, index, out_tensors);
}
absl::Status RootDataset::InputDatasets(
std::vector<const DatasetBase*>* inputs) const {
inputs->push_back(input_);
return absl::OkStatus();
}
absl::Status RootDataset::CheckExternalState() const {
return input_->CheckExternalState();
}
absl::Status RootDataset::AsGraphDefInternal(SerializationContext* ctx,
DatasetGraphDefBuilder* b,
Node** output) const {
return absl::UnimplementedError(
"RootDataset does not support serialization.");
}
#if !defined(IS_MOBILE_PLATFORM)
namespace {
// Initialize the rewritten output dataset with the metadata from the input
// dataset. Note: Do not override the `name` field in the metadata.
void InitializeRewrittenDatasetMetadata(const DatasetBase* input,
DatasetBase* rewritten_output) {
Metadata rewritten_output_metadata = rewritten_output->metadata();
rewritten_output_metadata.set_data_service_address(
input->metadata().data_service_address());
rewritten_output->Initialize(rewritten_output_metadata);
}
} // namespace
absl::Status FinalizeDataset(OpKernelContext* ctx, const DatasetBase* input,
DatasetBase** output) {
const Options& options = input->options();
absl::flat_hash_set<tstring> optimizations_enabled;
absl::flat_hash_set<tstring> optimizations_disabled;
absl::flat_hash_set<tstring> optimizations_default;
GetOptimizations(options, &optimizations_enabled, &optimizations_disabled,
&optimizations_default);
// Disable `enable_gradient_descent` as it assumes presence of ModelDatasetOp.
optimizations_disabled.insert("enable_gradient_descent");
auto experiments = GetExperiments();
LogAndRecordExperiments(experiments);
auto optimizations =
SelectOptimizations(experiments, optimizations_enabled,
optimizations_disabled, optimizations_default);
if (optimizations.empty()) {
return RootDataset::FromOptions(input, output);
}
auto optimization_configs = CreateGraphRewriteConfigs(options);
auto config_factory = [&optimizations, &optimization_configs]() {
return CreateRewriterConfig(optimizations, optimization_configs);
};
core::RefCountPtr<DatasetBase> rewritten_output;
absl::Status s =
RewriteDataset(ctx, input, std::move(config_factory),
/*record_fingerprint=*/false, &rewritten_output);
*output = rewritten_output.get();
bool rewritten = (*output != input);
if (absl::IsDeadlineExceeded(s)) {
// Ignore DeadlineExceeded as it implies that the attempted rewrite took too
// long which should not prevent further computation.
LOG(WARNING) << s;
} else if (!s.ok()) {
return s;
}
if (!rewritten) {
return RootDataset::FromOptions(input, output);
} else {
InitializeRewrittenDatasetMetadata(input, rewritten_output.get());
return RootDataset::FromOptions(std::move(rewritten_output), output);
}
return absl::OkStatus();
}
#else // !IS_MOBILE_PLATFORM
Status FinalizeDataset(OpKernelContext* ctx, const DatasetBase* input,
DatasetBase** output) {
return RootDataset::FromOptions(input, output);
}
#endif // !IS_MOBILE_PLATFORM
} // namespace data
} // namespace tensorflow
tensorflow/core/data/root_dataset.h
+108
View File
@@ -0,0 +1,108 @@
/* Copyright 2021 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_CORE_DATA_ROOT_DATASET_H_
#define TENSORFLOW_CORE_DATA_ROOT_DATASET_H_
#include <cstdint>
#include <functional>
#include <memory>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/model.pb.h"
#include "tensorflow/core/platform/mem.h"
#include "tensorflow/core/platform/refcount.h"
namespace tensorflow {
namespace data {
// Dataset transformation responsible for internal tf.data logic such as
// autotuning, applying threading configuration.
class RootDataset : public DatasetBase {
public:
struct Params {
bool autotune = true;
model::AutotuneAlgorithm autotune_algorithm;
std::function<int64_t()> autotune_cpu_budget_func;
double ram_budget_share;
int64_t autotune_ram_budget_from_options;
int64_t max_intra_op_parallelism = 1;
int64_t private_threadpool_size = 0;
int64_t ComputeInitialAutotuneRamBudget() const {
if (autotune_ram_budget_from_options > 0) {
return autotune_ram_budget_from_options;
} else {
return ram_budget_share * port::AvailableRam();
}
}
};
static absl::Status FromOptions(const DatasetBase* input,
DatasetBase** output);
static absl::Status FromOptions(core::RefCountPtr<DatasetBase> input,
DatasetBase** output);
~RootDataset() override;
const DataTypeVector& output_dtypes() const override;
const std::vector<PartialTensorShape>& output_shapes() const override;
int64_t CardinalityInternal(CardinalityOptions options) const override;
absl::Status Get(OpKernelContext* ctx, int64_t index,
std::vector<Tensor>* out_tensors) const override;
absl::Status CheckExternalState() const override;
std::string DebugString() const override;
absl::Status InputDatasets(
std::vector<const DatasetBase*>* inputs) const override;
std::unique_ptr<IteratorBase> MakeIteratorInternal(
const std::string& prefix) const override;
absl::Status RandomIndexingCompatible() const override {
return random_indexing_compatible_;
}
protected:
absl::Status AsGraphDefInternal(SerializationContext* ctx,
DatasetGraphDefBuilder* b,
Node** output) const override;
private:
class Iterator;
RootDataset(const DatasetBase* input, const Params& params);
RootDataset(core::RefCountPtr<DatasetBase> input, const Params& params);
const DatasetBase* input_;
core::RefCountPtr<DatasetBase> owned_input_;
const Params params_;
TraceMeMetadata traceme_metadata_;
absl::Status random_indexing_compatible_;
};
// Finalizes the `input` dataset, which is expected to be called before the
// dataset is about to be iterated. This can for instance apply static graph
// optimizations or inject internal tf.data transformations responsible for
// autotuning or threading configuration. The caller must ensure that the
// input dataset to be finalized outlives the output.
absl::Status FinalizeDataset(OpKernelContext* ctx, const DatasetBase* input,
DatasetBase** output);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_ROOT_DATASET_H_
tensorflow/core/data/serialization_utils.cc
+669
View File
@@ -0,0 +1,669 @@
/* 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 "tensorflow/core/data/serialization_utils.h"
#include <cstddef>
#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/graph_runner.h"
#include "tensorflow/core/data/compression_utils.h"
#include "tensorflow/core/data/dataset_utils.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/framework/variant_op_registry.h"
#include "tensorflow/core/framework/variant_tensor_data.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/platform/stringpiece.h"
namespace tensorflow {
namespace data {
namespace {
constexpr char kDelimiter[] = "@@";
constexpr char kComponent[] = "component";
constexpr char kNumComponents[] = "num_components";
constexpr char kNumElements[] = "num_elements";
constexpr char kIsDataset[] = ".is_dataset";
constexpr char kIteratorVariantTypeName[] = "tensorflow::Iterator";
constexpr char kOutputNode[] = ".output_node";
absl::Status FromGraphDef(
FunctionLibraryRuntime* flr, const GraphDef& graph_def,
const std::vector<std::pair<std::string, Tensor>>& input_list,
const std::string& output_node, Tensor* result) {
FunctionLibraryRuntime* cloned_flr = nullptr;
std::unique_ptr<ProcessFunctionLibraryRuntime> pflr = nullptr;
std::unique_ptr<FunctionLibraryDefinition> lib_def = nullptr;
TF_RETURN_IF_ERROR(flr->Clone(&lib_def, &pflr, &cloned_flr, true));
TF_RETURN_IF_ERROR(AddToFunctionLibrary(lib_def.get(), graph_def.library()));
Graph graph(OpRegistry::Global());
TF_RETURN_IF_ERROR(ImportGraphDef({}, graph_def, &graph, nullptr));
std::vector<Tensor> outputs;
GraphRunner graph_runner(cloned_flr->device());
TF_RETURN_IF_ERROR(graph_runner.Run(&graph, cloned_flr, input_list,
{output_node}, &outputs));
*result = outputs[0];
return absl::OkStatus();
}
// FindStatefulOps searches `graph_def` for all of its stateful ops storing
// their names in `stateful_op_names`.
absl::Status FindStatefulOps(const GraphDef& graph_def,
std::vector<std::string>* stateful_op_names) {
FunctionLibraryDefinition lib_def(OpRegistry::Global(), graph_def.library());
// Iterate over all nodes in the graph.
for (const auto& node : graph_def.node()) {
// Each Dataset graph has a _Retval op in the end which is marked stateful
if (node.op() == FunctionLibraryDefinition::kRetOp) continue;
if (!IsNodeStateful(lib_def, node).ok()) {
stateful_op_names->push_back(node.op());
}
}
// Iterate over all functions.
for (const auto& fdef : graph_def.library().function()) {
if (!fdef.signature().is_stateful()) continue;
for (const auto& node : fdef.node_def()) {
if (!IsNodeStateful(lib_def, node).ok()) {
stateful_op_names->push_back(
absl::StrCat(node.op(), " in function: ", fdef.signature().name()));
}
}
}
return absl::OkStatus();
}
} // namespace
absl::Status ReadElementsFromCheckpoint(
IteratorContext* ctx, IteratorStateReader* reader,
absl::string_view key_prefix, std::vector<std::vector<Tensor>>* elements) {
int64_t num_elements;
TF_RETURN_IF_ERROR(
reader->ReadScalar(key_prefix, kNumElements, &num_elements));
if (num_elements < 0) {
return absl::InternalError(
absl::StrCat("Num_elements in tf.data checkpoint must be >= 0, got: ",
num_elements));
}
DCHECK(elements->empty());
elements->reserve(num_elements);
for (int i = 0; i < num_elements; ++i) {
std::string element_prefix = absl::StrCat(key_prefix, "::", i);
int64_t num_components;
TF_RETURN_IF_ERROR(
reader->ReadScalar(element_prefix, kNumComponents, &num_components));
if (num_components < 0) {
return absl::InternalError(
absl::StrCat("Num of Tensor size in tf.data checkpoint must be >= 0, "
"got: ",
num_components));
}
elements->emplace_back();
std::vector<Tensor>& element = elements->at(i);
element.reserve(num_components);
for (int j = 0; j < num_components; ++j) {
element.emplace_back();
TF_RETURN_IF_ERROR(reader->ReadTensor(
ctx->flr(), element_prefix, absl::StrCat(kComponent, "[", j, "]"),
&element.back()));
}
}
return absl::OkStatus();
}
absl::Status WriteElement(IteratorStateWriter* writer,
absl::string_view key_prefix,
const std::vector<std::vector<Tensor>>& elements,
int64_t index) {
const std::vector<Tensor>& element = elements[index];
std::string element_prefix = absl::StrCat(key_prefix, "::", index);
TF_RETURN_IF_ERROR(
writer->WriteScalar(element_prefix, kNumComponents, element.size()));
for (int j = 0; j < element.size(); ++j) {
TF_RETURN_IF_ERROR(writer->WriteTensor(
element_prefix, absl::StrCat(kComponent, "[", j, "]"), element[j]));
}
return absl::OkStatus();
}
absl::Status WriteElementsToCheckpoint(
IteratorStateWriter* writer, absl::string_view key_prefix,
const std::vector<std::vector<Tensor>>& elements) {
TF_RETURN_IF_ERROR(
writer->WriteScalar(key_prefix, kNumElements, elements.size()));
for (int i = 0; i < elements.size(); ++i) {
TF_RETURN_IF_ERROR(WriteElement(writer, key_prefix, elements, i));
}
return absl::OkStatus();
}
absl::Status UpdateCheckpointElements(
IteratorStateWriter* writer, absl::string_view key_prefix,
const std::vector<std::vector<Tensor>>& elements,
const absl::flat_hash_set<int64_t>& checkpoint_indices) {
TF_RETURN_IF_ERROR(
writer->WriteScalar(key_prefix, kNumElements, elements.size()));
for (int64_t i : checkpoint_indices) {
TF_RETURN_IF_ERROR(WriteElement(writer, key_prefix, elements, i));
}
return absl::OkStatus();
}
VariantTensorDataReader::VariantTensorDataReader(
const std::vector<const tensorflow::VariantTensorData*>& data) {
for (const auto& d : data) {
std::string metadata;
d->get_metadata(&metadata);
auto keys = str_util::Split(metadata, kDelimiter, str_util::SkipEmpty());
const std::string name = keys[0];
data_[name] = d;
map_[name] = std::map<std::string, size_t>();
for (size_t i = 1; i < keys.size(); ++i) {
map_[name][keys[i]] = i - 1;
}
}
}
absl::Status VariantTensorDataReader::ReadScalar(absl::string_view key,
int64_t* val) const {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return ReadScalar(prefix, key, val);
}
absl::Status VariantTensorDataReader::ReadScalar(absl::string_view name,
absl::string_view key,
int64_t* val) const {
return ReadScalarInternal(name, key, val);
}
absl::Status VariantTensorDataReader::ReadScalar(absl::string_view key,
tstring* val) const {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return ReadScalar(prefix, key, val);
}
absl::Status VariantTensorDataReader::ReadScalar(absl::string_view name,
absl::string_view key,
tstring* val) const {
return ReadScalarInternal(name, key, val);
}
absl::Status VariantTensorDataReader::ReadTensor(absl::string_view key,
Tensor* val) const {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return ReadTensor(prefix, key, val);
}
absl::Status VariantTensorDataReader::ReadTensor(FunctionLibraryRuntime* flr,
absl::string_view key,
Tensor* val) const {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return ReadTensorInternal(flr, prefix, key, val);
}
absl::Status VariantTensorDataReader::ReadTensor(absl::string_view name,
absl::string_view key,
Tensor* val) const {
return ReadTensor(/*flr=*/nullptr, name, key, val);
}
absl::Status VariantTensorDataReader::ReadTensor(FunctionLibraryRuntime* flr,
absl::string_view name,
absl::string_view key,
Tensor* val) const {
return ReadTensorInternal(flr, name, key, val);
}
bool VariantTensorDataReader::Contains(absl::string_view key) const {
std::string prefix;
if (!ExtractIteratorPrefix(key, &prefix).ok()) {
return false;
}
return Contains(prefix, key);
}
bool VariantTensorDataReader::Contains(absl::string_view n,
absl::string_view key) const {
std::string name(n);
auto it = map_.find(name);
if (it == map_.end()) {
return false;
}
const auto& bucket = it->second;
return bucket.find(std::string(key)) != bucket.end();
}
template <typename T>
absl::Status VariantTensorDataReader::ReadScalarInternal(absl::string_view n,
absl::string_view key,
T* val) const {
std::string name(n);
auto it = map_.find(name);
if (it == map_.end()) {
return absl::NotFoundError(name);
}
const auto& bucket = it->second;
auto key_it = bucket.find(std::string(key));
if (key_it == bucket.end()) {
return absl::NotFoundError(key);
}
if (key_it->second >= data_.at(name)->tensors().size()) {
return absl::OutOfRangeError(
"VariantTensorDataReader tensor index out of range.");
}
const Tensor& t = data_.at(name)->tensors(key_it->second);
if (t.NumElements() != 1) {
return absl::InvalidArgumentError(
absl::StrCat("Expected a scalar, but found a tensor with ",
t.NumElements(), " elements."));
}
if (t.dtype() != DataTypeToEnum<T>::v()) {
return absl::InvalidArgumentError(absl::StrCat(
"Expected scalar of type ", DataTypeString(DataTypeToEnum<T>::v()),
", but found tensor of type ", DataTypeString(t.dtype())));
}
*val = t.scalar<T>()();
return absl::OkStatus();
}
absl::Status VariantTensorDataReader::ReadTensorInternal(
FunctionLibraryRuntime* flr, absl::string_view n, absl::string_view key,
Tensor* val) const {
if (Contains(n, absl::StrCat(key, kIsDataset))) {
return ReadDatasetInternal(flr, n, key, val);
}
std::string name(n);
auto it = map_.find(name);
if (it == map_.end()) {
return absl::NotFoundError(name);
}
const auto& bucket = it->second;
auto key_it = bucket.find(std::string(key));
if (key_it == bucket.end()) {
return absl::NotFoundError(key);
}
if (key_it->second >= data_.at(name)->tensors().size()) {
return absl::OutOfRangeError(
"VariantTensorDataReader tensor index out of range.");
}
*val = data_.at(name)->tensors(key_it->second);
return absl::OkStatus();
}
absl::Status VariantTensorDataReader::ReadDatasetInternal(
FunctionLibraryRuntime* flr, absl::string_view n, absl::string_view key,
Tensor* val) const {
if (flr == nullptr) {
return absl::InternalError(
"Function library runtime is needed to restore a dataset.");
}
tstring output_node, serialized_graph_def;
TF_RETURN_IF_ERROR(
ReadScalar(n, absl::StrCat(key, kOutputNode), &output_node));
TF_RETURN_IF_ERROR(ReadScalar(n, key, &serialized_graph_def));
GraphDef graph_def;
graph_def.ParseFromString(serialized_graph_def);
TF_RETURN_IF_ERROR(FromGraphDef(flr, graph_def, {}, output_node, val));
return absl::OkStatus();
}
std::map<std::string, Tensor> VariantTensorDataReader::ReadAllTensors() {
std::map<std::string, Tensor> result;
for (const auto& entry : map_) {
std::string key1 = entry.first;
for (const auto& inner : entry.second) {
std::string key2 = inner.first;
size_t index = inner.second;
result[absl::StrCat(key1, kDelimiter, key2)] =
data_[key1]->tensors(index);
}
}
return result;
}
absl::Status VariantTensorDataWriter::WriteScalar(absl::string_view key,
const int64_t val) {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return WriteScalar(prefix, key, val);
}
absl::Status VariantTensorDataWriter::WriteScalar(absl::string_view name,
absl::string_view key,
const int64_t val) {
return WriteScalarInternal(name, key, val);
}
absl::Status VariantTensorDataWriter::WriteScalar(absl::string_view key,
const tstring& val) {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return WriteScalar(prefix, key, val);
}
absl::Status VariantTensorDataWriter::WriteScalar(absl::string_view name,
absl::string_view key,
const tstring& val) {
return WriteScalarInternal(name, key, val);
}
absl::Status VariantTensorDataWriter::WriteTensor(absl::string_view key,
const Tensor& val) {
std::string prefix;
TF_RETURN_IF_ERROR(ExtractIteratorPrefix(key, &prefix));
return WriteTensor(prefix, key, val);
}
absl::Status VariantTensorDataWriter::WriteTensor(absl::string_view name,
absl::string_view key,
const Tensor& val) {
return WriteTensorInternal(name, key, val);
}
void VariantTensorDataWriter::MaybeFlush() {
if (is_flushed_) return;
for (auto& keys : keys_) {
const std::string name = keys.first;
std::string metadata = name;
for (size_t i = 0; i < keys_[name].size(); ++i) {
absl::StrAppend(&metadata, kDelimiter, keys_[name][i]);
}
data_[name]->set_metadata(metadata);
}
is_flushed_ = true;
}
void VariantTensorDataWriter::Reset() {
is_flushed_ = false;
data_.clear();
keys_.clear();
}
void VariantTensorDataWriter::ReleaseData(
std::vector<std::unique_ptr<VariantTensorData>>* variants) {
MaybeFlush();
for (auto& it : data_) {
variants->push_back(std::move(it.second));
}
Reset();
}
void VariantTensorDataWriter::GetData(
std::vector<const VariantTensorData*>* variants) {
MaybeFlush();
for (auto& it : data_) {
variants->push_back(it.second.get());
}
}
template <typename T>
absl::Status VariantTensorDataWriter::WriteScalarInternal(
absl::string_view name, absl::string_view key, const T& val) {
if (is_flushed_) {
return absl::FailedPreconditionError(
"Cannot call WriteScalar after GetData or ReleaseData is called");
}
Tensor val_t = Tensor(DataTypeToEnum<T>::v(), TensorShape({}));
val_t.scalar<T>()() = val;
return WriteTensorInternal(name, key, val_t);
}
absl::Status VariantTensorDataWriter::WriteTensorInternal(absl::string_view n,
absl::string_view key,
const Tensor& val) {
DatasetBase* dataset;
if (GetDatasetFromVariantTensor(val, &dataset).ok()) {
return WriteDatasetInternal(n, key, dataset);
}
if (is_flushed_) {
return absl::FailedPreconditionError(
"Cannot call WriteTensor after GetData or ReleaseData is called");
}
DCHECK_EQ(key.find(kDelimiter), std::string::npos);
std::string name(n);
if (keys_.count(name) == 0) {
keys_[name] = std::vector<std::string>();
}
keys_[name].push_back(std::string(key));
if (data_.count(name) == 0) {
data_[name] = std::make_unique<VariantTensorData>();
data_[name]->set_type_name("tensorflow::Iterator");
}
*(data_[name]->add_tensors()) = val;
return absl::OkStatus();
}
absl::Status VariantTensorDataWriter::WriteDatasetInternal(
absl::string_view n, absl::string_view key, const DatasetBase* dataset) {
GraphDef graph_def;
SerializationContext ctx((SerializationContext::Params()));
TF_RETURN_IF_ERROR(AsGraphDef(dataset, std::move(ctx), &graph_def));
std::string output_node;
for (const auto& node : graph_def.node()) {
if (node.op() == kRetvalOp) {
output_node = node.input(0);
break;
}
}
std::string result;
graph_def.SerializeToString(&result);
TF_RETURN_IF_ERROR(WriteScalar(n, absl::StrCat(key, kIsDataset), ""));
TF_RETURN_IF_ERROR(
WriteScalar(n, absl::StrCat(key, kOutputNode), output_node));
TF_RETURN_IF_ERROR(WriteScalar(n, key, result));
return absl::OkStatus();
}
std::string IteratorStateVariant::TypeName() {
return kIteratorVariantTypeName;
}
IteratorStateVariant::IteratorStateVariant(const IteratorStateVariant& other) {
if (other.data_) {
data_ = std::make_unique<VariantTensorData>(*other.data_);
}
}
absl::Status IteratorStateVariant::InitializeFromVariantData(
std::unique_ptr<VariantTensorData> data) {
data_ = std::move(data);
return absl::OkStatus();
}
void IteratorStateVariant::Encode(VariantTensorData* data) const {
CompressedElement compressed_tensors;
absl::Status s = CompressElement(data_->tensors(), &compressed_tensors);
if (!s.ok()) {
LOG(WARNING) << "Failed to compress iterator state variant: " << s;
*data = *data_;
return;
}
data->set_type_name(TypeName());
data->set_metadata(data_->metadata_string());
Tensor tensor(DT_VARIANT, TensorShape({}));
tensor.scalar<Variant>()() = std::move(compressed_tensors);
*data->add_tensors() = std::move(tensor);
}
bool IteratorStateVariant::Decode(VariantTensorData data) {
if (data.type_name() != TypeName()) {
return false;
}
const CompressedElement* compressed = GetCompressedElement(data);
if (!compressed) {
data_ = std::make_unique<VariantTensorData>(std::move(data));
return true;
}
std::vector<Tensor> tensors;
absl::Status s = UncompressElement(*compressed, &tensors);
if (!s.ok()) {
LOG(WARNING) << "Failed to uncompress iterator state variant: " << s;
data_ = std::make_unique<VariantTensorData>(std::move(data));
return true;
}
data_ = std::make_unique<VariantTensorData>();
data_->set_type_name(TypeName());
data_->set_metadata(std::move(data.metadata_string()));
for (auto& tensor : tensors) {
*data_->add_tensors() = std::move(tensor);
}
return true;
}
const CompressedElement* IteratorStateVariant::GetCompressedElement(
const VariantTensorData& data) {
bool should_uncompress =
data.tensors_size() == 1 &&
TensorShapeUtils::IsScalar(data.tensors(0).shape()) &&
data.tensors(0).dtype() == DT_VARIANT;
if (!should_uncompress) {
return nullptr;
}
const Variant& variant = data.tensors(0).scalar<Variant>()();
return variant.get<CompressedElement>();
}
std::string IteratorStateVariant::DebugString() const {
if (data_) {
return absl::StrCat("IteratorStateVariant<", data_->DebugString(), ">");
} else {
return absl::StrCat("IteratorStateVariant<empty>");
}
}
// Register the reader class in the global variant decode_fn registry
// so that a Variant containing a serialized representation of iterator state
// can be decoded using DecodeUnaryVariant. If we don't do this we will need
// to manually decode the returned Variant using MaybeDecodeAndCopy in
// DeserializeIteratorOp which is not recommended.
REGISTER_UNARY_VARIANT_DECODE_FUNCTION(IteratorStateVariant,
kIteratorVariantTypeName);
absl::Status AsGraphDefForRewrite(
OpKernelContext* ctx, const DatasetBase* input,
std::vector<std::pair<std::string, Tensor>>* input_list, GraphDef* result,
std::string* dataset_node) {
SerializationContext::Params params(ctx);
params.input_list = input_list;
params.external_state_policy = ExternalStatePolicy::POLICY_IGNORE;
params.is_graph_rewrite = true;
SerializationContext serialization_ctx(params);
TF_RETURN_IF_ERROR(AsGraphDef(input, std::move(serialization_ctx), result));
// Symbolic `_Retval` node indicates which node corresponds to the dataset.
for (const auto& node : result->node()) {
if (node.op() == kRetvalOp) {
*dataset_node = node.input(0);
}
}
return absl::OkStatus();
}
absl::Status AsGraphDef(const DatasetBase* dataset,
SerializationContext&& serialization_ctx,
GraphDef* graph_def) {
if (serialization_ctx.external_state_policy() ==
ExternalStatePolicy::POLICY_FAIL) {
TF_RETURN_IF_ERROR(dataset->CheckExternalState());
}
if (serialization_ctx.external_state_policy() ==
ExternalStatePolicy::POLICY_WARN) {
std::vector<std::string> stateful_op_names;
TF_RETURN_IF_ERROR(FindStatefulOps(*graph_def, &stateful_op_names));
if (!stateful_op_names.empty()) {
LOG(WARNING) << "We found the following stateful ops in the dataset "
"construction graph whose state would not be "
"serialized and might "
"cause subtle bugs: "
<< absl::StrJoin(stateful_op_names, ", ");
}
}
GraphDefBuilder b;
DatasetBase::DatasetGraphDefBuilder db(&b);
Node* output_node = nullptr;
TF_RETURN_IF_ERROR(
db.AddInputDataset(&serialization_ctx, dataset, &output_node));
// Insert a purely symbolic _Retval node to indicate to consumers which node
// represents `dataset`.
ops::UnaryOp(std::string(kRetvalOp), output_node,
b.opts()
.WithName("dataset")
.WithAttr("T", DT_VARIANT)
.WithAttr("index", 0));
TF_RETURN_IF_ERROR(b.ToGraphDef(graph_def));
return absl::OkStatus();
}
absl::StatusOr<absl::flat_hash_map<std::string, int64_t>> CheckpointStats(
const std::string& checkpoint_bytes) {
TensorProto proto;
if (!ParseProtoUnlimited(&proto, checkpoint_bytes)) {
return absl::InvalidArgumentError(
"Failed to parse checkpoint bytes into proto.");
}
Tensor t;
if (!t.FromProto(proto)) {
return absl::InvalidArgumentError(
"Failed to parse checkpoint tensor from proto.");
}
auto variant = t.scalar<Variant>()();
auto* w = variant.get<IteratorStateVariant>();
if (!w) {
return absl::InvalidArgumentError(
"Failed to access IteratorStateVariant inside checkpoint tensor");
}
const VariantTensorData* data = w->GetData();
auto reader = std::make_unique<VariantTensorDataReader>(
std::vector<const VariantTensorData*>{data});
absl::flat_hash_map<std::string, int64_t> stats;
for (const auto& [key, tensor] : reader->ReadAllTensors()) {
stats[key] = tensor.TotalBytes();
}
return stats;
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/serialization_utils.h
+244
View File
@@ -0,0 +1,244 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_SERIALIZATION_UTILS_H_
#define TENSORFLOW_CORE_DATA_SERIALIZATION_UTILS_H_
#include <cstddef>
#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/variant_tensor_data.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace data {
inline constexpr absl::string_view kRetvalOp = "_Retval";
// Reads dataset elements from the checkpoint reader using the given key prefix.
absl::Status ReadElementsFromCheckpoint(
IteratorContext* ctx, IteratorStateReader* reader,
absl::string_view key_prefix, std::vector<std::vector<Tensor>>* elements);
// Writes dataset elements to the checkpoint writer using the given key prefix.
// The elements can be read back by passing the same key prefix to
// ReadElementsFromCheckpoint. Only one list of elements can be written under
// the same key_prefix.
absl::Status WriteElementsToCheckpoint(
IteratorStateWriter* writer, absl::string_view key_prefix,
const std::vector<std::vector<Tensor>>& elements);
// Updates the dataset elements in the checkpoint for given `checkpoint_indices`
// using the given key prefix, assuming that vector of elements have
// checkpointed these before. The elements can be read back by passing the same
// key prefix to ReadElementsFromCheckpoint.
absl::Status UpdateCheckpointElements(
IteratorStateWriter* writer, absl::string_view key_prefix,
const std::vector<std::vector<Tensor>>& elements,
const absl::flat_hash_set<int64_t>& checkpoint_indices);
// Helper class for reading data from a vector of VariantTensorData objects.
class VariantTensorDataReader : public IteratorStateReader {
public:
explicit VariantTensorDataReader(
const std::vector<const VariantTensorData*>& data);
bool Contains(absl::string_view key) const override;
bool Contains(absl::string_view name, absl::string_view key) const override;
absl::Status ReadScalar(absl::string_view key, int64_t* val) const override;
absl::Status ReadScalar(absl::string_view name, absl::string_view key,
int64_t* val) const override;
absl::Status ReadScalar(absl::string_view key, tstring* val) const override;
absl::Status ReadScalar(absl::string_view name, absl::string_view key,
tstring* val) const override;
absl::Status ReadTensor(absl::string_view key, Tensor* val) const override;
absl::Status ReadTensor(FunctionLibraryRuntime* flr, absl::string_view key,
Tensor* val) const override;
absl::Status ReadTensor(absl::string_view name, absl::string_view key,
Tensor* val) const override;
absl::Status ReadTensor(FunctionLibraryRuntime* flr, absl::string_view name,
absl::string_view key, Tensor* val) const override;
private:
template <typename T>
absl::Status ReadScalarInternal(absl::string_view name, absl::string_view key,
T* val) const;
absl::Status ReadTensorInternal(FunctionLibraryRuntime* flr,
absl::string_view name, absl::string_view key,
Tensor* val) const;
absl::Status ReadDatasetInternal(FunctionLibraryRuntime* flr,
absl::string_view name,
absl::string_view key, Tensor* val) const;
// Produces all key/value pairs stored in this reader. Useful for debugging.
std::map<std::string, Tensor> ReadAllTensors();
// For access to ReadAllTensors()
friend absl::StatusOr<absl::flat_hash_map<std::string, int64_t>>
CheckpointStats(const std::string& checkpoint_bytes);
std::map<std::string, std::map<std::string, size_t>> map_;
std::map<std::string, const VariantTensorData*> data_; // Not owned.
};
// Helper class used to build a list of VariantTensorData objects, one for each
// iterator which is determined from the key supplied from the Write* calls.
// Sample usage:
// VariantTensorDataWriter writer;
// writer.WriteScalar(full_name("buffer_size"), buffer_.size());
// writer.WriteScalar(full_name("num_threads"), threadpool_.size());
// ....
// std::vector<std::unique_ptr<VariantTensorData>> variants;
// writer.ReleaseData(&variants);
// Now the VariantTensorData objects can be used to serialize.
class VariantTensorDataWriter : public IteratorStateWriter {
public:
absl::Status WriteScalar(absl::string_view key, int64_t val) override;
absl::Status WriteScalar(absl::string_view name, absl::string_view key,
int64_t val) override;
absl::Status WriteScalar(absl::string_view key, const tstring& val) override;
absl::Status WriteScalar(absl::string_view name, absl::string_view key,
const tstring& val) override;
absl::Status WriteTensor(absl::string_view key, const Tensor& val) override;
absl::Status WriteTensor(absl::string_view name, absl::string_view key,
const Tensor& val) override;
// Releases the built VariantTensorData's to `variants`. Clears out all
// class state.
void ReleaseData(std::vector<std::unique_ptr<VariantTensorData>>* variants);
// Obtains a read-only version of the VariantTensorData's built.
void GetData(std::vector<const VariantTensorData*>* variants);
private:
void MaybeFlush();
void Reset();
template <typename T>
absl::Status WriteScalarInternal(absl::string_view name,
absl::string_view key, const T& val);
absl::Status WriteTensorInternal(absl::string_view name,
absl::string_view key, const Tensor& val);
absl::Status WriteDatasetInternal(absl::string_view name,
absl::string_view key,
const DatasetBase* dataset);
bool is_flushed_ = false;
std::map<std::string, std::unique_ptr<VariantTensorData>> data_;
std::map<std::string, std::vector<std::string>> keys_;
};
// Wrapper for encoding/decoding the iterator state stored in a Variant tensor.
// The `GetData()` method returns an VariantTensorData object which contains all
// the state needed to restore a single iterator.
//
// Usage example:
//
// Encoding:
//
// Tensor t(DT_VARIANT, TensorShape({}));
// t->scalar<Variant>()() = IteratorStateVariant();
//
// Encode() sets the type_name of the VariantTensorData object to
// IteratorStateVariant::TypeName().
//
// Decoding:
//
// Variant v = <VariantTensorDataProto object>;
// DecodeUnaryVariant(&v);
// IteratorStateVariant* wrapper = v.get<IteratorStateVariant>();
// IteratorStateReader reader({wrapper->GetData()});
// iterator_resource->Restore(ctx, &reader);
//
// The type_name of the VariantTensorData object to be decoded must match
// IteratorStateVariant::TypeName().
class IteratorStateVariant {
public:
IteratorStateVariant() = default;
IteratorStateVariant(const IteratorStateVariant& other);
IteratorStateVariant& operator=(IteratorStateVariant&& other) = default;
IteratorStateVariant& operator=(const IteratorStateVariant& other) = delete;
static std::string TypeName();
// Initializes `this` from a VariantTensorData object.
absl::Status InitializeFromVariantData(
std::unique_ptr<VariantTensorData> data);
// Returns a borrowed pointer to the underlying VariantTensorData.
const VariantTensorData* GetData() const { return data_.get(); }
// Encodes this `IteratorStateVariant` into `*data`. Data will be compressed
// and stored as a scalar `CompressedElement` tensor, or left uncompressed if
// compression fails.
void Encode(VariantTensorData* data) const;
// Decodes from `data`. If `data` contains a single scalar `CompressedElement`
// tensor, it is assumed to be compressed by `Encode`, and will be
// uncompressed as part of `Decode`.
bool Decode(VariantTensorData data);
std::string DebugString() const;
private:
// Returns the compressed element in `data`. If `data` does not contain a
// compressed element, returns nullptr.
static const CompressedElement* GetCompressedElement(
const VariantTensorData& data);
std::unique_ptr<VariantTensorData> data_;
};
// Returns a GraphDef representation of the given dataset.
absl::Status AsGraphDef(const DatasetBase* dataset,
SerializationContext&& serialization_ctx,
GraphDef* graph_def);
// Returns a GraphDef representation of the given dataset suitable for
// optimization rewrites. It sets serialization parameters to export a minimum
// graph with additional information for optimization (i.e. ignoring external
// state, not serializing data tensors, not failing if there are datasets which
// do not have AsGraphDef implemented). Sets the `dataset_node` parameter to the
// dataset's node name in the resulting GraphDef.
absl::Status AsGraphDefForRewrite(
OpKernelContext* ctx, const DatasetBase* input,
std::vector<std::pair<std::string, Tensor>>* input_list, GraphDef* result,
std::string* dataset_node);
// Analyzes the bytes of a tf.data iterator checkpoint to identify all of the
// keys in the checkpoint along with their sizes in bytes.
absl::StatusOr<absl::flat_hash_map<std::string, int64_t>> CheckpointStats(
const std::string& checkpoint_bytes);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_KERNELS_SERIALIZATION_UTILS_H_
tensorflow/core/data/serialization_utils_test.cc
+332
View File
@@ -0,0 +1,332 @@
/* Copyright 2021 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/core/data/serialization_utils.h"
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/data/dataset_test_base.h"
#include "tensorflow/core/data/dataset_utils.h"
#include "tensorflow/core/data/test_utils.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/public/version.h"
#include "tensorflow/core/util/work_sharder.h"
namespace tensorflow {
namespace data {
namespace {
std::string full_name(std::string key) { return FullName("Iterator:", key); }
TEST(SerializationUtilsTest, CheckpointElementsRoundTrip) {
std::vector<std::vector<Tensor>> elements;
elements.push_back(CreateTensors<int32_t>(TensorShape({3}), {{1, 2, 3}}));
elements.push_back(CreateTensors<int32_t>(TensorShape({2}), {{4, 5}}));
VariantTensorDataWriter writer;
tstring test_prefix = full_name("test_prefix");
TF_ASSERT_OK(WriteElementsToCheckpoint(&writer, test_prefix, elements));
std::vector<const VariantTensorData*> data;
writer.GetData(&data);
VariantTensorDataReader reader(data);
std::vector<std::vector<Tensor>> read_elements;
TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestContext> ctx,
TestContext::Create());
TF_ASSERT_OK(ReadElementsFromCheckpoint(ctx->iter_ctx(), &reader, test_prefix,
&read_elements));
ASSERT_EQ(elements.size(), read_elements.size());
for (int i = 0; i < elements.size(); ++i) {
std::vector<Tensor>& original = elements[i];
std::vector<Tensor>& read = read_elements[i];
ASSERT_EQ(original.size(), read.size());
for (int j = 0; j < original.size(); ++j) {
EXPECT_EQ(original[j].NumElements(), read[j].NumElements());
EXPECT_EQ(original[j].flat<int32_t>()(0), read[j].flat<int32_t>()(0));
}
}
}
TEST(SerializationUtilsTest, VariantTensorDataRoundtrip) {
VariantTensorDataWriter writer;
TF_ASSERT_OK(writer.WriteScalar(full_name("Int64"), 24));
Tensor input_tensor(DT_FLOAT, {1});
input_tensor.flat<float>()(0) = 2.0f;
TF_ASSERT_OK(writer.WriteTensor(full_name("Tensor"), input_tensor));
std::vector<const VariantTensorData*> data;
writer.GetData(&data);
VariantTensorDataReader reader(data);
int64_t val_int64;
TF_ASSERT_OK(reader.ReadScalar(full_name("Int64"), &val_int64));
EXPECT_EQ(val_int64, 24);
Tensor val_tensor;
TF_ASSERT_OK(reader.ReadTensor(full_name("Tensor"), &val_tensor));
EXPECT_EQ(input_tensor.NumElements(), val_tensor.NumElements());
EXPECT_EQ(input_tensor.flat<float>()(0), val_tensor.flat<float>()(0));
}
TEST(SerializationUtilsTest, VariantTensorDataNonExistentKey) {
VariantTensorData data;
absl::StrAppend(&data.metadata_, "key1", "@@");
data.tensors_.push_back(Tensor(DT_INT64, {1}));
std::vector<const VariantTensorData*> reader_data;
reader_data.push_back(&data);
VariantTensorDataReader reader(reader_data);
int64_t val_int64;
tstring val_string;
Tensor val_tensor;
EXPECT_EQ(error::NOT_FOUND,
reader.ReadScalar(full_name("NonExistentKey"), &val_int64).code());
EXPECT_EQ(error::NOT_FOUND,
reader.ReadScalar(full_name("NonExistentKey"), &val_string).code());
EXPECT_EQ(error::NOT_FOUND,
reader.ReadTensor(full_name("NonExistentKey"), &val_tensor).code());
}
TEST(SerializationUtilsTest, VariantTensorDataRoundtripIteratorName) {
VariantTensorDataWriter writer;
TF_ASSERT_OK(writer.WriteScalar("Iterator", "Int64", 24));
Tensor input_tensor(DT_FLOAT, {1});
input_tensor.flat<float>()(0) = 2.0f;
TF_ASSERT_OK(writer.WriteTensor("Iterator", "Tensor", input_tensor));
std::vector<const VariantTensorData*> data;
writer.GetData(&data);
VariantTensorDataReader reader(data);
int64_t val_int64;
TF_ASSERT_OK(reader.ReadScalar("Iterator", "Int64", &val_int64));
EXPECT_EQ(val_int64, 24);
Tensor val_tensor;
TF_ASSERT_OK(reader.ReadTensor("Iterator", "Tensor", &val_tensor));
EXPECT_EQ(input_tensor.NumElements(), val_tensor.NumElements());
EXPECT_EQ(input_tensor.flat<float>()(0), val_tensor.flat<float>()(0));
}
TEST(SerializationUtilsTest, VariantTensorDataNonExistentKeyIteratorName) {
VariantTensorData data;
absl::StrAppend(&data.metadata_, "key1", "@@");
data.tensors_.push_back(Tensor(DT_INT64, {1}));
std::vector<const VariantTensorData*> reader_data;
reader_data.push_back(&data);
VariantTensorDataReader reader(reader_data);
int64_t val_int64;
tstring val_string;
Tensor val_tensor;
EXPECT_EQ(error::NOT_FOUND,
reader.ReadScalar("Iterator", "NonExistentKey", &val_int64).code());
EXPECT_EQ(
error::NOT_FOUND,
reader.ReadScalar("Iterator", "NonExistentKey", &val_string).code());
EXPECT_EQ(
error::NOT_FOUND,
reader.ReadTensor("Iterator", "NonExistentKey", &val_tensor).code());
}
TEST(SerializationUtilsTest, VariantTensorDataWriteAfterFlushing) {
VariantTensorDataWriter writer;
TF_ASSERT_OK(writer.WriteScalar(full_name("Int64"), 24));
std::vector<const VariantTensorData*> data;
writer.GetData(&data);
Tensor input_tensor(DT_FLOAT, {1});
input_tensor.flat<float>()(0) = 2.0f;
EXPECT_EQ(error::FAILED_PRECONDITION,
writer.WriteTensor(full_name("Tensor"), input_tensor).code());
}
TEST(SerializationUtilsTest, VariantTensorDataReaderOOB) {
VariantTensorData data;
data.metadata_ = "Iterator@@key1";
std::vector<const VariantTensorData*> reader_data;
reader_data.push_back(&data);
VariantTensorDataReader reader(reader_data);
int64_t val_int64;
EXPECT_EQ(reader.ReadScalar("Iterator", "key1", &val_int64).code(),
error::OUT_OF_RANGE);
Tensor val_tensor;
EXPECT_EQ(reader.ReadTensor("Iterator", "key1", &val_tensor).code(),
error::OUT_OF_RANGE);
}
class ParameterizedIteratorStateVariantTest
: public DatasetOpsTestBase,
public ::testing::WithParamInterface<std::vector<Tensor>> {
protected:
VariantTensorData GetVariantTensorData() const {
std::vector<Tensor> tensors = GetParam();
VariantTensorData data;
data.set_type_name(IteratorStateVariant::TypeName());
for (Tensor& tensor : tensors) {
*data.add_tensors() = std::move(tensor);
}
return data;
}
absl::StatusOr<VariantTensorData> EncodeAndDecode(
const VariantTensorData& data) const {
IteratorStateVariant encoder;
TF_RETURN_IF_ERROR(encoder.InitializeFromVariantData(
std::make_unique<VariantTensorData>(data)));
VariantTensorData encoded_data;
encoder.Encode(&encoded_data);
IteratorStateVariant decoder;
decoder.Decode(encoded_data);
return *decoder.GetData();
}
absl::StatusOr<VariantTensorData> DecodeUncompressed(
const VariantTensorData& data) const {
IteratorStateVariant decoder;
decoder.Decode(data);
return *decoder.GetData();
}
};
class ParemeterizedCheckpointIndicesTest
: public DatasetOpsTestBase,
public ::testing::WithParamInterface<absl::flat_hash_set<int64_t>> {
protected:
absl::flat_hash_set<int64_t> GetCheckpointIndices() const {
absl::flat_hash_set<int64_t> checkpoint_indices = GetParam();
return checkpoint_indices;
}
};
std::vector<std::vector<Tensor>> TestCases() {
return {
CreateTensors<int64_t>(TensorShape{1}, {{1}}), // int64
CreateTensors<int64_t>(TensorShape{1}, {{1}, {2}}), // multiple int64
CreateTensors<tstring>(TensorShape{1}, {{"a"}, {"b"}}), // tstring
{CreateTensor<tstring>(TensorShape{1}, {"a"}),
CreateTensor<int64_t>(TensorShape{1}, {1})}, // mixed tstring/int64
{}, // empty
{CreateTensor<int64_t>(TensorShape{128, 128}),
CreateTensor<int64_t>(TensorShape{64, 2})}, // larger components
};
}
std::vector<absl::flat_hash_set<int64_t>> CheckpointIndicesTestCases() {
return {
{/*checkpoint_indices*/},
{/*checkpoint_indices*/ 0},
{/*checkpoint_indices*/ 0, 1},
{/*checkpoint_indices*/ 0, 1, 2},
{/*checkpoint_indices*/ 1, 3, 4},
{/*checkpoint_indices*/ 1, 2, 3, 4},
{/*checkpoint_indices*/ 0, 1, 2, 3, 4},
};
}
TEST_P(ParameterizedIteratorStateVariantTest, EncodeAndDecode) {
VariantTensorData data = GetVariantTensorData();
TF_ASSERT_OK_AND_ASSIGN(VariantTensorData result, EncodeAndDecode(data));
EXPECT_EQ(result.type_name(), data.type_name());
for (int i = 0; i < result.tensors_size(); ++i) {
test::ExpectEqual(result.tensors(i), data.tensors(i));
}
}
TEST_P(ParameterizedIteratorStateVariantTest, DecodeUncompressed) {
VariantTensorData data = GetVariantTensorData();
TF_ASSERT_OK_AND_ASSIGN(VariantTensorData result, DecodeUncompressed(data));
EXPECT_EQ(result.type_name(), data.type_name());
for (int i = 0; i < result.tensors_size(); ++i) {
test::ExpectEqual(result.tensors(i), data.tensors(i));
}
}
TEST_P(ParemeterizedCheckpointIndicesTest,
CheckpointElementsRoundTripUsingIndices) {
std::vector<std::vector<Tensor>> elements;
elements.push_back(CreateTensors<int32_t>(TensorShape({3}), {{1, 2, 3}}));
elements.push_back(CreateTensors<int32_t>(TensorShape({2}), {{4, 5}}));
elements.push_back(
CreateTensors<int32_t>(TensorShape({5}), {{6, 7, 8, 9, 10}}));
elements.push_back(
CreateTensors<int32_t>(TensorShape({4}), {{11, 12, 13, 14}}));
elements.push_back(CreateTensors<int32_t>(TensorShape({2}), {{15, 16}}));
VariantTensorDataWriter writer;
tstring test_prefix = full_name("test_prefix");
// Generate checkpoint for entire buffer
absl::flat_hash_set<int64_t> checkpoint_indices_write = {0, 1, 2, 3, 4};
TF_ASSERT_OK(WriteElementsToCheckpoint(&writer, test_prefix, elements));
// Update the elements at checkpoint indices
for (auto index : GetCheckpointIndices()) {
elements.at(index) = CreateTensors<int32_t>(TensorShape({1}), {{1}});
}
TF_ASSERT_OK(UpdateCheckpointElements(&writer, test_prefix, elements,
GetCheckpointIndices()));
std::vector<const VariantTensorData*> data;
writer.GetData(&data);
VariantTensorDataReader reader(data);
std::vector<std::vector<Tensor>> read_elements;
TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<TestContext> ctx,
TestContext::Create());
TF_ASSERT_OK(ReadElementsFromCheckpoint(ctx->iter_ctx(), &reader, test_prefix,
&read_elements));
ASSERT_EQ(elements.size(), read_elements.size());
// Check if checkpoint state of entire buffer is as expected
for (int index = 0; index < elements.size(); ++index) {
std::vector<Tensor>& original = elements[index];
std::vector<Tensor>& read = read_elements[index];
ASSERT_EQ(original.size(), read.size());
for (int j = 0; j < original.size(); ++j) {
EXPECT_EQ(original[j].NumElements(), read[j].NumElements());
EXPECT_EQ(original[j].flat<int32_t>()(0), read[j].flat<int32_t>()(0));
}
}
}
INSTANTIATE_TEST_SUITE_P(Instantiation, ParameterizedIteratorStateVariantTest,
::testing::ValuesIn(TestCases()));
INSTANTIATE_TEST_SUITE_P(Instantiation, ParemeterizedCheckpointIndicesTest,
::testing::ValuesIn(CheckpointIndicesTestCases()));
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/BUILD
+1198
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/service/auto_scaler.cc
+316
View File
@@ -0,0 +1,316 @@
/* 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.
==============================================================================*/
#include "tensorflow/core/data/service/auto_scaler.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <memory>
#include <numeric>
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/time/time.h"
#include "tensorflow/core/framework/metrics.h"
#include "tsl/platform/mutex.h"
#include "tsl/platform/thread_annotations.h"
namespace tensorflow {
namespace data {
constexpr double kAutoScalerOutlierSigmas = 1.0;
template <typename T>
double GetMedian(const absl::flat_hash_map<T, double>& rates) {
std::vector<double> sorted_rates;
for (const auto& [id, rate] : rates) {
sorted_rates.push_back(rate);
}
std::sort(sorted_rates.begin(), sorted_rates.end());
return sorted_rates[sorted_rates.size() / 2];
}
template <typename T>
double GetMean(const absl::flat_hash_map<T, double>& rates) {
double rates_sum = 0.0;
for (const auto& [id, rate] : rates) {
rates_sum += rate;
}
if (rates_sum == 0.0) return 0.0;
return rates_sum / static_cast<double>(rates.size());
}
template <typename T>
double GetStandardDeviation(const absl::flat_hash_map<T, double>& rates,
double mean) {
double squared_distances_sum = 0.0;
for (const auto& [id, rate] : rates) {
squared_distances_sum += (rate - mean) * (rate - mean);
}
if (squared_distances_sum == 0.0 || rates.size() <= 1) return 0.0;
return std::sqrt(squared_distances_sum /
static_cast<double>(rates.size() - 1));
}
// Discards rates that are more than (std_dev * outlier_sigmas) far from the
// mean, and replaces them with the median. Puts the result in
// `rates_without_outliers`.
template <typename T>
void ReplaceOutliers(const absl::flat_hash_map<T, double>& rates,
std::vector<double>& rates_without_outliers,
double outlier_sigmas) {
if (rates.empty()) return;
double mean = GetMean(rates);
double median = GetMedian(rates);
double standard_deviation = GetStandardDeviation(rates, mean);
double lower_threshold = mean - standard_deviation * outlier_sigmas;
double upper_threshold = mean + standard_deviation * outlier_sigmas;
for (const auto& [id, rate] : rates) {
if (rate >= lower_threshold && rate <= upper_threshold) {
rates_without_outliers.push_back(rate);
} else {
rates_without_outliers.push_back(median);
}
}
}
std::optional<int64_t> AutoScaler::GetOptimalNumberOfWorkers() const
TF_LOCKS_EXCLUDED(mu_) {
tsl::mutex_lock l(mu_);
if (worker_throughputs_.empty() || consumption_rates_.empty())
return std::nullopt;
std::vector<double> consumption_rates_without_outliers;
// TODO(armandouv): Discard outlier replacement when we ensure reported time
// values are correct.
// Outliers can make the estimate have an unfeasible value (very high or very
// low).
ReplaceOutliers(consumption_rates_, consumption_rates_without_outliers,
kAutoScalerOutlierSigmas);
double consumption_rates_sum_ =
std::accumulate(consumption_rates_without_outliers.begin(),
consumption_rates_without_outliers.end(), 0.0);
std::vector<double> worker_throughputs_without_outliers;
ReplaceOutliers(worker_throughputs_, worker_throughputs_without_outliers,
kAutoScalerOutlierSigmas);
double worker_throughputs_sum_ =
std::accumulate(worker_throughputs_without_outliers.begin(),
worker_throughputs_without_outliers.end(), 0.0);
double average_worker_throughput =
worker_throughputs_sum_ / static_cast<double>(worker_throughputs_.size());
int64_t optimal_number_of_workers =
ceil(consumption_rates_sum_ / average_worker_throughput);
return std::max(int64_t{1}, optimal_number_of_workers);
}
absl::Status AutoScaler::ReportProcessingTime(const std::string& worker_address,
absl::Duration processing_time)
TF_LOCKS_EXCLUDED(mu_) {
if (processing_time <= absl::ZeroDuration()) {
return absl::InvalidArgumentError(absl::StrCat(
"Cannot update processing_time with a ZeroDuration or negative value: ",
absl::FormatDuration(processing_time)));
}
double worker_throughput = 1.0 / absl::ToDoubleSeconds(processing_time);
tsl::mutex_lock l(mu_);
worker_throughputs_[worker_address] = worker_throughput;
return absl::OkStatus();
}
absl::Status AutoScaler::ReportTargetProcessingTime(
int64_t consumer_id, absl::Duration target_processing_time)
TF_LOCKS_EXCLUDED(mu_) {
if (target_processing_time <= absl::ZeroDuration()) {
return absl::InvalidArgumentError(
absl::StrCat("Cannot update target_processing_time with a ZeroDuration "
"or negative value: ",
absl::FormatDuration(target_processing_time)));
}
double consumption_rate = 1.0 / absl::ToDoubleSeconds(target_processing_time);
tsl::mutex_lock l(mu_);
consumption_rates_[consumer_id] = consumption_rate;
return absl::OkStatus();
}
absl::Status AutoScaler::RemoveWorker(const std::string& worker_address)
TF_LOCKS_EXCLUDED(mu_) {
tsl::mutex_lock l(mu_);
if (!worker_throughputs_.contains(worker_address))
return absl::NotFoundError(
absl::StrCat("Worker with address ", worker_address, " not found"));
worker_throughputs_.erase(worker_address);
return absl::OkStatus();
}
absl::Status AutoScaler::RemoveConsumer(int64_t consumer_id)
TF_LOCKS_EXCLUDED(mu_) {
tsl::mutex_lock l(mu_);
if (!consumption_rates_.contains(consumer_id))
return absl::NotFoundError(
absl::StrCat("Consumer with ID ", consumer_id, " not found"));
consumption_rates_.erase(consumer_id);
return absl::OkStatus();
}
void MultipleIterationsAutoScaler::EnsureIterationIsRegistered(
int64_t iteration_id) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (!auto_scalers_.contains(iteration_id)) {
auto_scalers_[iteration_id] = std::make_unique<AutoScaler>();
}
}
absl::Status MultipleIterationsAutoScaler::UnregisterIteration(
int64_t iteration_id) TF_LOCKS_EXCLUDED(mu_) {
tsl::mutex_lock l(mu_);
if (!auto_scalers_.contains(iteration_id))
return absl::NotFoundError(absl::StrCat("AutoScaler for iteration_id ",
iteration_id, " does not exist"));
auto_scalers_.erase(iteration_id);
return absl::OkStatus();
}
absl::Status MultipleIterationsAutoScaler::UpdateOptimalNumberOfWorkersMetric(
int64_t current_number_of_workers) TF_LOCKS_EXCLUDED(mu_) {
if (current_number_of_workers <= 0)
return absl::InvalidArgumentError(
"The current number of workers must be positive");
std::optional<int64_t> optimal_number_of_workers =
GetOptimalNumberOfWorkers();
if (!optimal_number_of_workers)
return absl::UnavailableError(
"Cannot update the optimal number of workers metric because there are "
"no reported processing and target processing times for at least one "
"iteration");
VLOG(3) << "Estimated optimal number of workers: "
<< optimal_number_of_workers.value();
// Limit the estimate to wait for target processing times to converge to a
// feasible value. First, start increasing exponentially by 4x. Once
// increases are greater than 500, scale linearly.
int64_t bound_optimal_number_of_workers = optimal_number_of_workers.value();
if (bound_optimal_number_of_workers > current_number_of_workers * 4 ||
bound_optimal_number_of_workers > current_number_of_workers + 500) {
bound_optimal_number_of_workers = std::min(current_number_of_workers * 4,
current_number_of_workers + 500);
}
// Limit the estimate to at most 100k workers.
bound_optimal_number_of_workers =
std::min(bound_optimal_number_of_workers, int64_t{100000});
VLOG(3) << "Bound optimal number of workers: "
<< bound_optimal_number_of_workers;
metrics::RecordTFDataServiceOptimalNumberOfWorkers(
bound_optimal_number_of_workers);
return absl::OkStatus();
}
std::optional<int64_t> MultipleIterationsAutoScaler::GetOptimalNumberOfWorkers()
const TF_LOCKS_EXCLUDED(mu_) {
int64_t optimal_number_of_workers = 0;
{
tsl::tf_shared_lock l(mu_);
for (const auto& [iteration_id, auto_scaler] : auto_scalers_) {
std::optional<int64_t> current_optimal_number_of_workers =
auto_scaler->GetOptimalNumberOfWorkers();
if (!current_optimal_number_of_workers.has_value()) continue;
optimal_number_of_workers = std::max(
optimal_number_of_workers, current_optimal_number_of_workers.value());
}
}
if (optimal_number_of_workers == 0)
return std::nullopt;
else
return optimal_number_of_workers;
}
absl::Status MultipleIterationsAutoScaler::ReportProcessingTime(
int64_t iteration_id, const std::string& worker_address,
absl::Duration processing_time) TF_LOCKS_EXCLUDED(mu_) {
tsl::mutex_lock l(mu_);
EnsureIterationIsRegistered(iteration_id);
absl::Status status = auto_scalers_[iteration_id]->ReportProcessingTime(
worker_address, processing_time);
return status;
}
absl::Status MultipleIterationsAutoScaler::ReportTargetProcessingTime(
int64_t iteration_id, int64_t consumer_id,
absl::Duration target_processing_time) TF_LOCKS_EXCLUDED(mu_) {
tsl::mutex_lock l(mu_);
EnsureIterationIsRegistered(iteration_id);
absl::Status status = auto_scalers_[iteration_id]->ReportTargetProcessingTime(
consumer_id, target_processing_time);
return status;
}
absl::Status MultipleIterationsAutoScaler::RemoveWorker(
int64_t iteration_id, const std::string& worker_address)
TF_LOCKS_EXCLUDED(mu_) {
tsl::tf_shared_lock l(mu_);
if (!auto_scalers_.contains(iteration_id))
return absl::NotFoundError(absl::StrCat(
"There are no reported times for iteration_id ", iteration_id));
absl::Status status =
auto_scalers_[iteration_id]->RemoveWorker(worker_address);
return status;
}
absl::Status MultipleIterationsAutoScaler::RemoveConsumer(int64_t iteration_id,
int64_t consumer_id)
TF_LOCKS_EXCLUDED(mu_) {
tsl::tf_shared_lock l(mu_);
if (!auto_scalers_.contains(iteration_id))
return absl::NotFoundError(absl::StrCat(
"There are no reported times for iteration_id ", iteration_id));
absl::Status status =
auto_scalers_[iteration_id]->RemoveConsumer(consumer_id);
return status;
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/auto_scaler.h
+180
View File
@@ -0,0 +1,180 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_SERVICE_AUTO_SCALER_H_
#define TENSORFLOW_CORE_DATA_SERVICE_AUTO_SCALER_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/time/time.h"
#include "xla/tsl/platform/status.h"
#include "tsl/platform/mutex.h"
#include "tsl/platform/thread_annotations.h"
namespace tensorflow {
namespace data {
// Estimates the optimal number of tf.data service workers for an Iteration
// based on the current workload.
// Note: It is assumed that all reported times correspond to the same Iteration.
//
// Glossary:
// * Consumer: A client that consumes elements from tf.data service.
// * Worker: A tf.data service worker.
// * Processing time (PT): The estimated time it takes a worker to process and
// produce an element.
// * Target processing time (TPT): From the perspective of a consumer,
// it is the maximum time a tf.data input pipeline can take to produce an
// element such that the downstream processor wait time is 0. In other words,
// this is the ideal time the tf.data pipeline should take to produce an element
// so that training doesn't slow down due to waiting for elements. This means
// that we want processing time <= target processing time, so that when an
// element is requested, the pipeline has processed it already.
// * Worker throughput (WT): It is the multiplicative inverse of processing time
// (1 / PT). This refers to the number of elements produced by a worker per
// second.
// * Consumption rate (CR): It is the multiplicative inverse of target
// processing time (1 / TPT). This refers to the number of elements requested by
// a consumer per second.
//
// **AutoScaler overview**
//
// 1. It keeps track of the most recent worker throughputs reported by each
// worker in the data service cluster, as well as the most recent consumption
// rates reported by each consumer. WTs and CRs are derived from reporting PTs
// and TPTs, respectively.
// 2. Having this information, it estimates the optimal number of workers N as
// follows:
// N = (Sum of CRs reported by all consumers) /
// (Average of WTs reported by all workers)
//
// AutoScaler is thread-safe.
class AutoScaler {
public:
AutoScaler() = default;
// Returns the estimated optimal number of workers according to the current
// observed workload. If there are no previously reported processing and
// target processing times, returns nullopt.
std::optional<int64_t> GetOptimalNumberOfWorkers() const
TF_LOCKS_EXCLUDED(mu_);
// Reports the latest observed processing time from the worker with
// `worker_address`. Returns an error if `processing_time` is ZeroDuration or
// negative.
absl::Status ReportProcessingTime(const std::string& worker_address,
absl::Duration processing_time)
TF_LOCKS_EXCLUDED(mu_);
// Reports the latest observed target processing time from the consumer
// identified by `consumer_id`. Returns an error if `target_processing_time`
// is ZeroDuration or negative.
absl::Status ReportTargetProcessingTime(int64_t consumer_id,
absl::Duration target_processing_time)
TF_LOCKS_EXCLUDED(mu_);
// Unregisters the worker with `worker_address`, removing its reported
// processing time from consideration of the current workload estimation.
// Returns an error if the specified worker does not exist.
absl::Status RemoveWorker(const std::string& worker_address)
TF_LOCKS_EXCLUDED(mu_);
// Unregisters the consumer identified by `consumer_id`, removing its reported
// target processing time from consideration of the current workload
// estimation. Returns an error if the specified consumer does not exist.
absl::Status RemoveConsumer(int64_t consumer_id) TF_LOCKS_EXCLUDED(mu_);
private:
mutable tsl::mutex mu_;
// Map from worker address to worker throughput.
absl::flat_hash_map<std::string, double> worker_throughputs_
TF_GUARDED_BY(mu_);
// Map from consumer id to consumption rate.
absl::flat_hash_map<int64_t, double> consumption_rates_ TF_GUARDED_BY(mu_);
};
// Exports a metric (/tensorflow/data/service/optimal_number_of_workers) with
// the estimated optimal number of tf.data service workers, according to
// the observed cluster workload.
//
// It estimates the number of workers as the maximum of the estimated optimal
// number of workers for all Iterations running in the tf.data service cluster.
//
// MultipleIterationsAutoScaler is thread-safe.
class MultipleIterationsAutoScaler {
public:
MultipleIterationsAutoScaler() = default;
// Unregisters iteration with `iteration_id`, removing its reported
// times from consideration of the current workload estimation.
// Returns an error if the specified iteration does not exist.
absl::Status UnregisterIteration(int64_t iteration_id) TF_LOCKS_EXCLUDED(mu_);
// Updates the metric value with the current estimated optimal number of
// workers. The estimate is limited to min(4 * `current_number_of_workers`,
// `current_number_of_workers` + 500). Returns an error if there are no
// previously reported processing and target processing times for at least one
// iteration, or `current_number_of_workers` is not positive.
absl::Status UpdateOptimalNumberOfWorkersMetric(
int64_t current_number_of_workers) TF_LOCKS_EXCLUDED(mu_);
// Returns the estimated optimal number of workers according to the current
// observed workload. If there are no previously reported processing and
// target processing times for at least one iteration, returns nullopt.
std::optional<int64_t> GetOptimalNumberOfWorkers() const
TF_LOCKS_EXCLUDED(mu_);
// Reports the latest observed processing time from the worker with
// `worker_address` for iteration with `iteration_id`. Returns an error if
// `processing_time` is ZeroDuration or negative.
absl::Status ReportProcessingTime(int64_t iteration_id,
const std::string& worker_address,
absl::Duration processing_time)
TF_LOCKS_EXCLUDED(mu_);
// Reports the latest observed target processing time from the consumer
// identified by `consumer_id` for iteration with `iteration_id`. Returns an
// error if `target_processing_time` is ZeroDuration or negative.
absl::Status ReportTargetProcessingTime(int64_t iteration_id,
int64_t consumer_id,
absl::Duration target_processing_time)
TF_LOCKS_EXCLUDED(mu_);
// Unregisters the worker with `worker_address` for iteration with
// `iteration_id`, removing its reported processing time from consideration of
// the current workload estimation. Returns an error if there are no
// previously reported processing times for iteration with `iteration_id` and
// the specified worker.
absl::Status RemoveWorker(int64_t iteration_id,
const std::string& worker_address)
TF_LOCKS_EXCLUDED(mu_);
// Unregisters the consumer identified by `consumer_id` for iteration with
// `iteration_id`, removing its reported target processing time from
// consideration of the current workload estimation. Returns an error if there
// are no previously reported processing times for iteration with
// `iteration_id` and the specified consumer.
absl::Status RemoveConsumer(int64_t iteration_id, int64_t consumer_id)
TF_LOCKS_EXCLUDED(mu_);
private:
// Registers iteration with `iteration_id` if it does not exist already,
// allowing its future reported times to be considered for the current
// workload estimation.
void EnsureIterationIsRegistered(int64_t iteration_id)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
mutable tsl::mutex mu_;
// Map from iteration id to AutoScaler.
absl::flat_hash_map<int64_t, std::unique_ptr<AutoScaler>> auto_scalers_
TF_GUARDED_BY(mu_);
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_AUTO_SCALER_H_
tensorflow/core/data/service/auto_scaler_test.cc
+763
View File
@@ -0,0 +1,763 @@
/* 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.
==============================================================================*/
#include "tensorflow/core/data/service/auto_scaler.h"
#include <cstdint>
#include <optional>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/time/time.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status_matchers.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
namespace tensorflow {
namespace data {
namespace {
using ::tsl::testing::StatusIs;
TEST(AutoScalerTest, GetOptimalNumberOfWorkersInitialState) {
AutoScaler auto_scaler;
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), std::nullopt);
}
TEST(AutoScalerTest, GetOptimalNumberOfWorkersNoRegisteredWorkers) {
AutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(10)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), std::nullopt);
}
TEST(AutoScalerTest, GetOptimalNumberOfWorkersNoRegisteredConsumers) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(10)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), std::nullopt);
}
// Worker 0:
// - Processing time = 0.2 [s] -> Throughput = 5 [elements/s]
// Consumer 0:
// - Target processing time = 0.025 [s] -> Consumption rate = 40 [elements/s]
//
// Average throughput = 5 [elements/s]
// Sum of consumption rates = 40 [elements/s]
// Estimated number of workers = 40 / 5 = 8
TEST(AutoScalerTest, GetOptimalNumberOfWorkersExpectedEstimate1) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(auto_scaler.ReportTargetProcessingTime(0, absl::Seconds(0.025)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 8);
}
// Worker 0:
// - Processing time = 0.2 [s] -> Throughput = 5 [elements/s]
// Worker 1:
// - Processing time = 0.15 [s] -> Throughput = 6.6666 [elements/s]
// Consumer 0:
// - Target processing time = 0.025 [s] -> Consumption rate = 40 [elements/s]
// Consumer 1:
// - Target processing time = 0.05 [s] -> Consumption rate = 20 [elements/s]
//
// Average throughput = 5.833 [elements/s]
// Sum of consumption rates = 60 [elements/s]
// Estimated number of workers = 60 / 5.833 = ⌈10.28⌉ = 11
TEST(AutoScalerTest, GetOptimalNumberOfWorkersExpectedEstimate2) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/1:20000",
absl::Seconds(0.15)));
TF_ASSERT_OK(auto_scaler.ReportTargetProcessingTime(0, absl::Seconds(0.025)));
TF_ASSERT_OK(auto_scaler.ReportTargetProcessingTime(1, absl::Seconds(0.05)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 11);
}
// Worker 0:
// - Processing time = 0.1 [s] -> Throughput = 10 [elements/s]
// Worker 1:
// - Processing time = 0.2 [s] -> Throughput = 5 [elements/s]
// Consumer 0:
// - Target processing time = 0.01 [s] -> Consumption rate = 100 [elements/s]
// Consumer 1:
// - Target processing time = 0.02 [s] -> Consumption rate = 50 [elements/s]
//
// Average throughput = 7.5 [elements/s]
// Sum of consumption rates = 150 [elements/s]
// Estimated number of workers = 150 / 7.5 = 20
TEST(AutoScalerTest, GetOptimalNumberOfWorkersExpectedEstimate3) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Seconds(0.1)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/1:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(auto_scaler.ReportTargetProcessingTime(0, absl::Seconds(0.01)));
TF_ASSERT_OK(auto_scaler.ReportTargetProcessingTime(1, absl::Seconds(0.02)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 20);
}
// TODO(armandouv): Delete when we ensure reported time values are correct.
// If outliers are not discarded, the number of workers will be
// unrealistically high (e.g. ~100k workers).
// Worker 0:
// - Processing time = 0.08 [s] -> Throughput = 12.5 [elements/s]
// Consumer 0:
// - Target processing time = 0.0000005 [s] -> Consumption rate = 2000000
// [elements/s] (DISCARDED, replaced by median = 500)
// Consumer 1:
// - Target processing time = 0.003 [s] -> Consumption rate = 333.333
// [elements/s]
// Consumer 2:
// - Target processing time = 0.002 [s] -> Consumption rate = 500 [elements/s]
//
// Average throughput = 12.5 [elements/s]
// Sum of consumption rates = 1333.33 [elements/s]
// Estimated number of workers = 1333.33 / 12.5 = ⌈106.66⌉ = 107
TEST(AutoScalerTest, GetOptimalNumberOfWorkersRemoveOutliersTPT) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Nanoseconds(80000000)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Nanoseconds(500)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, absl::Nanoseconds(3000000)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(2, absl::Nanoseconds(2000000)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 107);
}
// For workers, there can be very small PTs that cause the estimation to be very
// low.
// Worker 0:
// - Processing time = 0.08 [s] -> Throughput = 12.5 [elements/s]
// Worker 1:
// - Processing time = 0.07 [s] -> Throughput = 14.285 [elements/s]
// Worker 3:
// - Processing time = 0.000001 [s] -> Throughput = 1000000 [elements/s]
// (DISCARDED, replaced by median = 14.285)
// Consumer 0:
// - Target processing time = 0.0003 [s] -> Consumption rate = 3333.33
// [elements/s]
//
// Average throughput = 13.69 [elements/s]
// Sum of consumption rates = 3333.33 [elements/s]
// Estimated number of workers = 3333.33 / 13.69 = ⌈243.48⌉ = 244
TEST(AutoScalerTest, GetOptimalNumberOfWorkersRemoveOutliersPT) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Nanoseconds(80000000)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/1:20000",
absl::Nanoseconds(70000000)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/2:20000",
absl::Nanoseconds(1000)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Nanoseconds(300000)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 244);
}
TEST(AutoScalerTest, ReportProcessingTimeNewWorker) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(10)));
}
TEST(AutoScalerTest, ReportProcessingTimeExistingWorker) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(20)));
}
TEST(AutoScalerTest, ReportProcessingTimeNewAndExisting) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/1:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/2:20000",
absl::Microseconds(30)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(30)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/1:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/2:20000",
absl::Microseconds(10)));
}
TEST(AutoScalerTest, ReportProcessingTimeZeroDuration) {
AutoScaler auto_scaler;
absl::Status result = auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::ZeroDuration());
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(AutoScalerTest, ReportProcessingTimeNegativeDuration) {
AutoScaler auto_scaler;
absl::Status result = auto_scaler.ReportProcessingTime(
"/worker/task/0:20000", absl::Microseconds(-10));
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(AutoScalerTest, ReportTargetProcessingTimeNewConsumer) {
AutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(10)));
}
TEST(AutoScalerTest, ReportTargetProcessingTimeExistingConsumer) {
AutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(20)));
}
TEST(AutoScalerTest, ReportTargetProcessingTimeNewAndExisting) {
AutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, absl::Microseconds(20)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(2, absl::Microseconds(30)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(30)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, absl::Microseconds(20)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(2, absl::Microseconds(10)));
}
TEST(AutoScalerTest, ReportTargetProcessingTimeZeroDuration) {
AutoScaler auto_scaler;
absl::Status result =
auto_scaler.ReportTargetProcessingTime(0, absl::ZeroDuration());
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(AutoScalerTest, ReportTargetProcessingTimeNegativeDuration) {
AutoScaler auto_scaler;
absl::Status result =
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(-10));
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(AutoScalerTest, RemoveWorkerSuccessful) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/1:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveWorker("/worker/task/0:20000"));
TF_ASSERT_OK(auto_scaler.RemoveWorker("/worker/task/1:20000"));
}
TEST(AutoScalerTest, RemoveNonexistentWorker) {
AutoScaler auto_scaler;
EXPECT_THAT(auto_scaler.RemoveWorker("/worker/task/0:20000"),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(AutoScalerTest, RemoveWorkerAfterNewPTReported) {
AutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime("/worker/task/0:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveWorker("/worker/task/0:20000"));
}
TEST(AutoScalerTest, RemoveConsumerSuccessful) {
AutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(30)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, absl::Microseconds(30)));
TF_ASSERT_OK(auto_scaler.RemoveConsumer(0));
TF_ASSERT_OK(auto_scaler.RemoveConsumer(1));
}
TEST(AutoScalerTest, RemoveNonexistentConsumer) {
AutoScaler auto_scaler;
EXPECT_THAT(auto_scaler.RemoveConsumer(0),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(AutoScalerTest, RemoveConsumerAfterNewTPTReported) {
AutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(30)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveConsumer(0));
}
TEST(MultipleIterationsAutoScalerTest, UnregisterExistingIteration) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(5)));
TF_ASSERT_OK(auto_scaler.UnregisterIteration(0));
}
TEST(MultipleIterationsAutoScalerTest, UnregisterNonexistentIteration) {
MultipleIterationsAutoScaler auto_scaler;
EXPECT_THAT(auto_scaler.UnregisterIteration(0),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricInvalidCurrentWorkers) {
MultipleIterationsAutoScaler auto_scaler;
absl::Status status = auto_scaler.UpdateOptimalNumberOfWorkersMetric(0);
EXPECT_THAT(status,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
status = auto_scaler.UpdateOptimalNumberOfWorkersMetric(-1);
EXPECT_THAT(status,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricNoReportedTimes) {
MultipleIterationsAutoScaler auto_scaler;
absl::Status status = auto_scaler.UpdateOptimalNumberOfWorkersMetric(1);
EXPECT_THAT(status, absl_testing::StatusIs(absl::StatusCode::kUnavailable));
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricNoReportedPTs) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(5)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(5)));
absl::Status status = auto_scaler.UpdateOptimalNumberOfWorkersMetric(1);
EXPECT_THAT(status, absl_testing::StatusIs(absl::StatusCode::kUnavailable));
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricNoReportedTPTs) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(10)));
absl::Status status = auto_scaler.UpdateOptimalNumberOfWorkersMetric(1);
EXPECT_THAT(status, absl_testing::StatusIs(absl::StatusCode::kUnavailable));
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricWithReportedTimes) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(5)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(5)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.UpdateOptimalNumberOfWorkersMetric(1));
monitoring::testing::CellReader<int64_t> cell_reader(
"/tensorflow/data/service/optimal_number_of_workers");
EXPECT_GT(cell_reader.Read(), 0);
metrics::RecordTFDataServiceOptimalNumberOfWorkers(0);
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricIncreaseWithinLimit) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(500)));
// Estimated workers = 50. Current workers = 15.
// 50 <= 15 * 4 = 60, so the estimate is not modified.
TF_ASSERT_OK(auto_scaler.UpdateOptimalNumberOfWorkersMetric(15));
monitoring::testing::CellReader<int64_t> cell_reader(
"/tensorflow/data/service/optimal_number_of_workers");
EXPECT_EQ(cell_reader.Read(), 50);
metrics::RecordTFDataServiceOptimalNumberOfWorkers(0);
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetric4xIncreaseLimit) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(1)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
// Estimated workers = 10. Current workers = 2.
// 10 > 4 * 2 = 8, so the estimate is limited to 8.
TF_ASSERT_OK(auto_scaler.UpdateOptimalNumberOfWorkersMetric(2));
monitoring::testing::CellReader<int64_t> cell_reader(
"/tensorflow/data/service/optimal_number_of_workers");
EXPECT_EQ(cell_reader.Read(), 8);
metrics::RecordTFDataServiceOptimalNumberOfWorkers(0);
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetric500IncreaseLimit) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(1)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10000)));
// Estimated workers = 10000. Current workers = 1000.
// 10000 > 1000 * 4 = 4000, and 10000 > 1000 + 500 = 1500, so the estimate is
// limited to min(4000, 1500) = 1500.
TF_ASSERT_OK(auto_scaler.UpdateOptimalNumberOfWorkersMetric(1000));
monitoring::testing::CellReader<int64_t> cell_reader(
"/tensorflow/data/service/optimal_number_of_workers");
EXPECT_EQ(cell_reader.Read(), 1500);
metrics::RecordTFDataServiceOptimalNumberOfWorkers(0);
}
TEST(MultipleIterationsAutoScalerTest,
UpdateOptimalNumberOfWorkersMetricMaxLimit) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(1)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(200000)));
// Estimated workers = 200000. Current workers = 99700.
// The estimate is limited to 100k workers.
TF_ASSERT_OK(auto_scaler.UpdateOptimalNumberOfWorkersMetric(99700));
monitoring::testing::CellReader<int64_t> cell_reader(
"/tensorflow/data/service/optimal_number_of_workers");
EXPECT_EQ(cell_reader.Read(), 100000);
metrics::RecordTFDataServiceOptimalNumberOfWorkers(0);
}
TEST(MultipleIterationsAutoScalerTest, GetOptimalNumberOfWorkersInitialState) {
MultipleIterationsAutoScaler auto_scaler;
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), std::nullopt);
}
TEST(MultipleIterationsAutoScalerTest,
GetOptimalNumberOfWorkersNoRegisteredWorkers) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(5)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(5)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), std::nullopt);
}
TEST(MultipleIterationsAutoScalerTest,
GetOptimalNumberOfWorkersNoRegisteredConsumers) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(10)));
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), std::nullopt);
}
TEST(MultipleIterationsAutoScalerTest,
GetOptimalNumberOfWorkersExpectedEstimate1) {
MultipleIterationsAutoScaler auto_scaler;
// Estimated number of workers for iteration 0 = 8
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Seconds(0.025)));
// Estimated number of workers for iteration 1 = 11
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/1:20000",
absl::Seconds(0.15)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Seconds(0.025)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 1, absl::Seconds(0.05)));
// max(8, 11) = 11 workers
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 11);
}
TEST(MultipleIterationsAutoScalerTest,
GetOptimalNumberOfWorkersExpectedEstimate2) {
MultipleIterationsAutoScaler auto_scaler;
// Estimated number of workers for iteration 0 = 8
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Seconds(0.025)));
// Estimated number of workers for iteration 1 = 11
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/1:20000",
absl::Seconds(0.15)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Seconds(0.025)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 1, absl::Seconds(0.05)));
// Estimated number of workers for iteration 2 = 20
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(2, "/worker/task/0:20000",
absl::Seconds(0.1)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(2, "/worker/task/1:20000",
absl::Seconds(0.2)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(2, 0, absl::Seconds(0.01)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(2, 1, absl::Seconds(0.02)));
// max(8, 11, 20) = 20 workers
EXPECT_EQ(auto_scaler.GetOptimalNumberOfWorkers(), 20);
}
TEST(MultipleIterationsAutoScalerTest, ReportProcessingTimeNewIteration) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
}
TEST(MultipleIterationsAutoScalerTest, ReportProcessingTimeNewWorker) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/1:20000",
absl::Microseconds(10)));
}
TEST(MultipleIterationsAutoScalerTest, ReportProcessingTimeExistingWorker) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(10)));
}
TEST(MultipleIterationsAutoScalerTest, ReportProcessingTimeNewAndExisting) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/1:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/1:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/1:20000",
absl::Microseconds(30)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/1:20000",
absl::Microseconds(30)));
}
TEST(MultipleIterationsAutoScalerTest, ReportProcessingTimeZeroDuration) {
MultipleIterationsAutoScaler auto_scaler;
absl::Status result = auto_scaler.ReportProcessingTime(
0, "/worker/task/0:20000", absl::ZeroDuration());
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(MultipleIterationsAutoScalerTest, ReportProcessingTimeNegativeDuration) {
MultipleIterationsAutoScaler auto_scaler;
absl::Status result = auto_scaler.ReportProcessingTime(
0, "/worker/task/0:20000", absl::Microseconds(-10));
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(MultipleIterationsAutoScalerTest, ReportTargetProcessingTimeNewIteration) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
}
TEST(MultipleIterationsAutoScalerTest, ReportTargetProcessingTimeNewConsumer) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 1, absl::Microseconds(10)));
}
TEST(MultipleIterationsAutoScalerTest,
ReportTargetProcessingTimeExistingWorker) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(10)));
}
TEST(MultipleIterationsAutoScalerTest,
ReportTargetProcessingTimeNewAndExisting) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 1, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 1, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(20)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 1, absl::Microseconds(30)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(20)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 1, absl::Microseconds(30)));
}
TEST(MultipleIterationsAutoScalerTest, ReportTargetProcessingTimeZeroDuration) {
MultipleIterationsAutoScaler auto_scaler;
absl::Status result =
auto_scaler.ReportTargetProcessingTime(0, 0, absl::ZeroDuration());
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(MultipleIterationsAutoScalerTest,
ReportTargetProcessingTimeNegativeDuration) {
MultipleIterationsAutoScaler auto_scaler;
absl::Status result =
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(-10));
EXPECT_THAT(result,
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(MultipleIterationsAutoScalerTest, RemoveWorkerUnregisteredIteration) {
MultipleIterationsAutoScaler auto_scaler;
EXPECT_THAT(auto_scaler.RemoveWorker(0, "/worker/task/1:20000"),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
EXPECT_THAT(auto_scaler.RemoveWorker(1, "/worker/task/1:20000"),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(MultipleIterationsAutoScalerTest, RemoveWorkerSuccessful) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(1, "/worker/task/0:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveWorker(0, "/worker/task/0:20000"));
TF_ASSERT_OK(auto_scaler.RemoveWorker(1, "/worker/task/0:20000"));
}
TEST(MultipleIterationsAutoScalerTest, RemoveNonexistentWorker) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
EXPECT_THAT(auto_scaler.RemoveWorker(0, "/worker/task/1:20000"),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(MultipleIterationsAutoScalerTest, RemoveWorkerAfterNewPTReported) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(10)));
TF_ASSERT_OK(auto_scaler.ReportProcessingTime(0, "/worker/task/0:20000",
absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveWorker(0, "/worker/task/0:20000"));
}
TEST(MultipleIterationsAutoScalerTest, RemoveConsumerUnregisteredIteration) {
MultipleIterationsAutoScaler auto_scaler;
EXPECT_THAT(auto_scaler.RemoveConsumer(0, 0),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
EXPECT_THAT(auto_scaler.RemoveConsumer(1, 0),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(MultipleIterationsAutoScalerTest, RemoveConsumerSuccessful) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(1, 0, absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveConsumer(0, 0));
TF_ASSERT_OK(auto_scaler.RemoveConsumer(1, 0));
}
TEST(MultipleIterationsAutoScalerTest, RemoveNonexistentConsumer) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
EXPECT_THAT(auto_scaler.RemoveConsumer(0, 1),
absl_testing::StatusIs(absl::StatusCode::kNotFound));
}
TEST(MultipleIterationsAutoScalerTest, RemoveConsumerAfterNewTPTReported) {
MultipleIterationsAutoScaler auto_scaler;
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(10)));
TF_ASSERT_OK(
auto_scaler.ReportTargetProcessingTime(0, 0, absl::Microseconds(20)));
TF_ASSERT_OK(auto_scaler.RemoveConsumer(0, 0));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/byte_size.cc
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2024 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/core/data/service/byte_size.h"
#include <cstddef>
#include <string>
#include "absl/strings/str_cat.h"
namespace tensorflow {
namespace data {
size_t ByteSize::ToUnsignedBytes() const { return bytes_; }
double ByteSize::ToDoubleBytes() const { return static_cast<double>(bytes_); }
double ByteSize::ToDoubleKB() const { return *this / ByteSize::KB(1); }
double ByteSize::ToDoubleMB() const { return *this / ByteSize::MB(1); }
double ByteSize::ToDoubleGB() const { return *this / ByteSize::GB(1); }
double ByteSize::ToDoubleTB() const { return *this / ByteSize::TB(1); }
std::string ByteSize::DebugString() const {
if (*this < ByteSize::KB(1)) {
return absl::StrCat(ToUnsignedBytes(), "B");
}
if (*this < ByteSize::MB(1)) {
return absl::StrCat(ToDoubleKB(), "KB");
}
if (*this < ByteSize::GB(1)) {
return absl::StrCat(ToDoubleMB(), "MB");
}
if (*this < ByteSize::TB(1)) {
return absl::StrCat(ToDoubleGB(), "GB");
}
return absl::StrCat(ToDoubleTB(), "TB");
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/byte_size.h
+198
View File
@@ -0,0 +1,198 @@
/* Copyright 2024 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_CORE_DATA_SERVICE_BYTE_SIZE_H_
#define TENSORFLOW_CORE_DATA_SERVICE_BYTE_SIZE_H_
#include <cstddef>
#include <ostream>
#include <string>
namespace tensorflow {
namespace data {
// A `ByteSize` represents data space usage measured in bytes. It is constructed
// using Bytes, KB, MB, GB, or TB. Supports common arithmetic operations. Uses
// `size_t` in its internal representation. Thus, it only supports non-negative
// sizes, and the maximum byte size is std::numeric_limits<size_t>::max().
//
// Usage example:
//
// constexpr ByteSize kAllocatedMemoryLimit = ByteSize::MB(64);
//
// Tensor data = ...
// ByteSize tensor_size = ByteSize::Bytes(data.AllocatedBytes());
// if (tensor_size > 0.95 * kAllocatedMemoryLimit) {
// LOG(WARNING) << "Tensor memory usage is " << tensor_size << ". This is "
// << "close to the limit " << kAllocatedMemoryLimit << ".";
// }
class ByteSize final {
public:
// The default is 0 bytes.
constexpr ByteSize() = default;
constexpr ByteSize(const ByteSize&) = default;
ByteSize& operator=(const ByteSize&) = default;
// Constructs byte sizes of bytes, KB, MB, GB, and TB.
constexpr static ByteSize Bytes(size_t n);
// In this and following templates, `T` should be a numeric type,
// e.g.: size_t, double, etc.
template <class T>
constexpr static ByteSize KB(T n);
template <class T>
constexpr static ByteSize MB(T n);
template <class T>
constexpr static ByteSize GB(T n);
template <class T>
constexpr static ByteSize TB(T n);
// Compound assignment operators.
ByteSize& operator+=(ByteSize rhs);
// Does not support negative bytes. If *this < rhs, returns 0 bytes.
ByteSize& operator-=(ByteSize rhs);
template <class T>
ByteSize& operator*=(T rhs);
template <class T>
ByteSize& operator/=(T rhs);
// Converts the measurement into the specified unit.
size_t ToUnsignedBytes() const;
double ToDoubleBytes() const;
double ToDoubleKB() const;
double ToDoubleMB() const;
double ToDoubleGB() const;
double ToDoubleTB() const;
// Returns a human-readable string of the byte size. For example, "5KB",
// "1GB", etc.
std::string DebugString() const;
private:
constexpr explicit ByteSize(double bytes) : bytes_(bytes) {}
size_t bytes_ = 0;
};
constexpr ByteSize ByteSize::Bytes(size_t n) { return ByteSize(n); };
template <class T>
constexpr ByteSize ByteSize::KB(T n) {
return ByteSize::Bytes(n * (size_t{1} << 10));
}
template <class T>
constexpr ByteSize ByteSize::MB(T n) {
return ByteSize::Bytes(n * (size_t{1} << 20));
}
template <class T>
constexpr ByteSize ByteSize::GB(T n) {
return ByteSize::Bytes(n * (size_t{1} << 30));
}
template <class T>
constexpr ByteSize ByteSize::TB(T n) {
return ByteSize::Bytes(n * (size_t{1} << 40));
}
// Compound assignments.
inline ByteSize& ByteSize::operator+=(ByteSize rhs) {
bytes_ += rhs.ToUnsignedBytes();
return *this;
}
inline ByteSize& ByteSize::operator-=(ByteSize rhs) {
if (bytes_ < rhs.ToUnsignedBytes()) {
bytes_ = 0;
return *this;
}
bytes_ -= rhs.ToUnsignedBytes();
return *this;
}
template <class T>
inline ByteSize& ByteSize::operator*=(T rhs) {
bytes_ *= rhs;
return *this;
}
template <class T>
inline ByteSize& ByteSize::operator/=(T rhs) {
bytes_ /= rhs;
return *this;
}
// Binary arithmetic operators.
inline ByteSize operator+(ByteSize lhs, ByteSize rhs) {
return lhs += rhs;
}
inline ByteSize operator-(ByteSize lhs, ByteSize rhs) {
return lhs -= rhs;
}
template <class T>
inline ByteSize operator*(ByteSize lhs, T rhs) { return lhs *= rhs; }
template <class T>
inline ByteSize operator*(T lhs, ByteSize rhs) { return rhs *= lhs; }
template <class T>
inline ByteSize operator/(ByteSize lhs, T rhs) { return lhs /= rhs; }
inline double operator/(ByteSize lhs, ByteSize rhs) {
return lhs.ToDoubleBytes() / rhs.ToDoubleBytes();
}
// Comparison operators.
inline bool operator<(ByteSize lhs, ByteSize rhs) {
return lhs.ToUnsignedBytes() < rhs.ToUnsignedBytes();
}
inline bool operator>(ByteSize lhs, ByteSize rhs) {
return rhs < lhs;
}
inline bool operator>=(ByteSize lhs, ByteSize rhs) {
return !(lhs < rhs);
}
inline bool operator<=(ByteSize lhs, ByteSize rhs) {
return !(rhs < lhs);
}
inline bool operator==(ByteSize lhs, ByteSize rhs) {
return lhs.ToUnsignedBytes() == rhs.ToUnsignedBytes();
}
inline bool operator!=(ByteSize lhs, ByteSize rhs) {
return !(lhs == rhs);
}
// Output operator, which supports logging with LOG(*).
inline std::ostream& operator<<(std::ostream& os, ByteSize byte_size) {
return os << byte_size.DebugString();
}
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_BYTE_SIZE_H_
tensorflow/core/data/service/byte_size_test.cc
+385
View File
@@ -0,0 +1,385 @@
/* Copyright 2024 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/core/data/service/byte_size.h"
#include <cstddef>
#include <string>
#include <vector>
#include "absl/algorithm/container.h"
#include "xla/tsl/platform/test.h"
namespace tensorflow {
namespace data {
namespace {
using ::testing::Eq;
using ::testing::Not;
TEST(ByteSizeTest, Constructors) {
EXPECT_EQ(ByteSize::Bytes(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::Bytes(1), ByteSize::Bytes(1));
EXPECT_EQ(ByteSize::Bytes(1024), ByteSize::Bytes(1024));
EXPECT_EQ(ByteSize::Bytes(1024), ByteSize::KB(1));
EXPECT_EQ(ByteSize::Bytes(size_t{1} << 63), ByteSize::TB(size_t{1} << 23));
EXPECT_EQ(ByteSize::KB(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::KB(1), ByteSize::Bytes(size_t{1} << 10));
EXPECT_EQ(ByteSize::KB(0.9), ByteSize::Bytes(1024 * 0.9));
EXPECT_EQ(ByteSize::KB(1.5), ByteSize::Bytes(1024 * 1.5));
EXPECT_EQ(ByteSize::KB(1.5), ByteSize::KB(1.5));
EXPECT_EQ(ByteSize::KB(1024), ByteSize::MB(1));
EXPECT_EQ(ByteSize::MB(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::MB(1), ByteSize::Bytes(size_t{1} << 20));
EXPECT_EQ(ByteSize::MB(0.9), ByteSize::Bytes(size_t{1} << 20) * 0.9);
EXPECT_EQ(ByteSize::MB(1.5), ByteSize::Bytes(size_t{1} << 20) * 1.5);
EXPECT_EQ(ByteSize::MB(1.5), ByteSize::MB(1.5));
EXPECT_EQ(ByteSize::MB(1024), ByteSize::GB(1));
EXPECT_EQ(ByteSize::GB(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::GB(1), ByteSize::Bytes(size_t{1} << 30));
EXPECT_EQ(ByteSize::GB(0.9), ByteSize::Bytes(size_t{1} << 30) * 0.9);
EXPECT_EQ(ByteSize::GB(1.5), ByteSize::Bytes(size_t{1} << 30) * 1.5);
EXPECT_EQ(ByteSize::GB(1.5), ByteSize::GB(1.5));
EXPECT_EQ(ByteSize::GB(1024), ByteSize::TB(1));
EXPECT_EQ(ByteSize::TB(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::TB(1), ByteSize::Bytes(size_t{1} << 40));
EXPECT_EQ(ByteSize::TB(0.9), ByteSize::Bytes(size_t{1} << 40) * 0.9);
EXPECT_EQ(ByteSize::TB(1.5), ByteSize::Bytes(size_t{1} << 40) * 1.5);
EXPECT_EQ(ByteSize::TB(1.5), ByteSize::TB(1.5));
EXPECT_EQ(ByteSize::TB(1024), ByteSize::TB(1024));
EXPECT_EQ(ByteSize::TB(size_t{1} << 23), ByteSize::TB(size_t{1} << 23));
EXPECT_THAT(ByteSize::Bytes(0), Not(Eq(ByteSize::Bytes(1))));
EXPECT_THAT(ByteSize::Bytes(1025), Not(Eq(ByteSize::KB(1))));
EXPECT_THAT(ByteSize::KB(1), Not(Eq(ByteSize::MB(1))));
EXPECT_THAT(ByteSize::MB(1), Not(Eq(ByteSize::GB(1))));
EXPECT_THAT(ByteSize::GB(1), Not(Eq(ByteSize::TB(1))));
EXPECT_THAT(ByteSize::TB(1), Not(Eq(ByteSize::TB(2))));
}
TEST(ByteSizeTest, ConstexprConstruction) {
constexpr ByteSize default_byte_size;
EXPECT_EQ(default_byte_size, ByteSize::Bytes(0));
constexpr ByteSize bytes = ByteSize::Bytes(1);
EXPECT_EQ(bytes, ByteSize::Bytes(1));
constexpr ByteSize kb = ByteSize::KB(1);
EXPECT_EQ(kb, ByteSize::KB(1));
constexpr ByteSize mb = ByteSize::MB(1);
EXPECT_EQ(mb, ByteSize::MB(1));
constexpr ByteSize gb = ByteSize::GB(1);
EXPECT_EQ(gb, ByteSize::GB(1));
constexpr ByteSize tb = ByteSize::TB(1);
EXPECT_EQ(tb, ByteSize::TB(1));
constexpr ByteSize tb_copy(tb);
EXPECT_EQ(tb_copy, tb);
}
TEST(ByteSizeTest, ConvertToBytes) {
EXPECT_EQ(ByteSize::Bytes(0).ToUnsignedBytes(), 0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(0).ToDoubleBytes(), 0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(0).ToDoubleKB(), 0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(0).ToDoubleMB(), 0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(0).ToDoubleGB(), 0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(0).ToDoubleTB(), 0);
EXPECT_EQ(ByteSize::Bytes(1).ToUnsignedBytes(), 1);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(1).ToDoubleBytes(), 1.0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(1).ToDoubleKB(), 1.0 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(1).ToDoubleMB(), 1.0 / 1024 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(1).ToDoubleGB(), 1.0 / 1024 / 1024 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(1).ToDoubleTB(),
1.0 / 1024 / 1024 / 1024 / 1024);
EXPECT_EQ(ByteSize::KB(0.25).ToUnsignedBytes(), 0.25 * (size_t{1} << 10));
EXPECT_DOUBLE_EQ(ByteSize::KB(0.25).ToDoubleBytes(), 0.25 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::KB(0.25).ToDoubleKB(), 0.25);
EXPECT_DOUBLE_EQ(ByteSize::KB(0.25).ToDoubleMB(), 0.25 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::KB(0.25).ToDoubleGB(), 0.25 / 1024 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::KB(0.25).ToDoubleTB(), 0.25 / 1024 / 1024 / 1024);
EXPECT_EQ(ByteSize::MB(0.5).ToUnsignedBytes(), 0.5 * (size_t{1} << 20));
EXPECT_DOUBLE_EQ(ByteSize::MB(0.5).ToDoubleBytes(), 0.5 * 1024 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::MB(0.5).ToDoubleKB(), 0.5 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::MB(0.5).ToDoubleMB(), 0.5);
EXPECT_DOUBLE_EQ(ByteSize::MB(0.5).ToDoubleGB(), 0.5 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::MB(0.5).ToDoubleTB(), 0.5 / 1024 / 1024);
EXPECT_EQ(ByteSize::GB(10).ToUnsignedBytes(), 10.0 * (size_t{1} << 30));
EXPECT_DOUBLE_EQ(ByteSize::GB(10).ToDoubleBytes(), 10.0 * 1024 * 1024 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::GB(10).ToDoubleKB(), 10.0 * 1024 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::GB(10).ToDoubleMB(), 10.0 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::GB(10).ToDoubleGB(), 10.0);
EXPECT_DOUBLE_EQ(ByteSize::GB(10).ToDoubleTB(), 10.0 / 1024);
EXPECT_EQ(ByteSize::TB(1024).ToUnsignedBytes(), 1024 * (size_t{1} << 40));
EXPECT_DOUBLE_EQ(ByteSize::TB(1024).ToDoubleBytes(),
1024.0 * 1024 * 1024 * 1024 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::TB(1024).ToDoubleKB(),
1024.0 * 1024 * 1024 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::TB(1024).ToDoubleMB(), 1024.0 * 1024 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::TB(1024).ToDoubleGB(), 1024.0 * 1024);
EXPECT_DOUBLE_EQ(ByteSize::TB(1024).ToDoubleTB(), 1024.0);
}
TEST(ByteSizeTest, Arithmetics) {
// Add.
EXPECT_EQ(ByteSize::Bytes(0) + ByteSize::Bytes(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::Bytes(0) + ByteSize::Bytes(1), ByteSize::Bytes(1));
EXPECT_EQ(ByteSize::Bytes(512) + ByteSize::Bytes(512), ByteSize::KB(1));
EXPECT_EQ(ByteSize::Bytes(512) + ByteSize::KB(1), ByteSize::KB(1.5));
EXPECT_EQ(ByteSize::KB(0.5) + ByteSize::KB(1), ByteSize::KB(1.5));
EXPECT_EQ(ByteSize::MB(1) + ByteSize::KB(512), ByteSize::MB(1.5));
EXPECT_EQ(ByteSize::MB(1) + ByteSize::Bytes(512), ByteSize::Bytes(1049088));
EXPECT_EQ(ByteSize::GB(0.5) + ByteSize::MB(256) + ByteSize::MB(256),
ByteSize::GB(1));
std::vector<ByteSize> GBs(1024, ByteSize::GB(1));
EXPECT_EQ(absl::c_accumulate(GBs, ByteSize::Bytes(0)), ByteSize::TB(1));
EXPECT_EQ(ByteSize::TB(1) + ByteSize::TB(0.5) + ByteSize::GB(512),
ByteSize::TB(2));
// Substract.
EXPECT_EQ(ByteSize::Bytes(0) - ByteSize::Bytes(0), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::KB(1) - ByteSize::Bytes(512), ByteSize::KB(0.5));
EXPECT_EQ(ByteSize::MB(1) - ByteSize::KB(512) - ByteSize::KB(512),
ByteSize::MB(0));
EXPECT_EQ(ByteSize::GB(1) - ByteSize::MB(512), ByteSize::GB(0.5));
EXPECT_EQ(ByteSize::GB(0.5) - ByteSize::MB(512), ByteSize::GB(0));
EXPECT_EQ(ByteSize::GB(1) - ByteSize::MB(512) - ByteSize::MB(512),
ByteSize::GB(0));
EXPECT_EQ(ByteSize::TB(1) - ByteSize::GB(512) - ByteSize::GB(512),
ByteSize::GB(0));
// No negative bytes.
EXPECT_EQ(ByteSize::Bytes(0) - ByteSize::Bytes(1), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::Bytes(0) - ByteSize::GB(1), ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::MB(1) - ByteSize::GB(1), ByteSize::Bytes(0));
// Multiply.
EXPECT_EQ(ByteSize::Bytes(0) * 0, ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::KB(1) * 0, ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::MB(1) * 0, ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::GB(1) * 0, ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::TB(1) * 0, ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::Bytes(1) * 1024, ByteSize::KB(1));
EXPECT_EQ(ByteSize::KB(1) * 1024, ByteSize::MB(1));
EXPECT_EQ(ByteSize::MB(1) * 1024, ByteSize::GB(1));
EXPECT_EQ(ByteSize::GB(1) * 1024, ByteSize::TB(1));
EXPECT_EQ(ByteSize::Bytes(1) * 1.1, ByteSize::Bytes(1));
EXPECT_EQ(ByteSize::KB(1) * 1.2, ByteSize::KB(1.2));
EXPECT_EQ(ByteSize::MB(1) * 1.3, ByteSize::MB(1.3));
EXPECT_EQ(ByteSize::GB(1) * 1.4, ByteSize::GB(1.4));
EXPECT_EQ(ByteSize::TB(1) * 1.5, ByteSize::TB(1.5));
EXPECT_EQ(ByteSize::KB(1) * 0.5, ByteSize::Bytes(512));
EXPECT_EQ(ByteSize::MB(1) * 0.5, ByteSize::KB(512));
EXPECT_EQ(ByteSize::GB(1) * 0.5, ByteSize::MB(512));
EXPECT_EQ(ByteSize::TB(1) * 0.25, ByteSize::GB(256));
EXPECT_EQ(1024 * ByteSize::Bytes(1), ByteSize::KB(1));
EXPECT_EQ(1024 * ByteSize::KB(1), ByteSize::MB(1));
EXPECT_EQ(1024 * ByteSize::MB(1), ByteSize::GB(1));
EXPECT_EQ(1024 * ByteSize::GB(1), ByteSize::TB(1));
EXPECT_EQ(0.9 * ByteSize::TB(1), ByteSize::GB(921.6));
EXPECT_EQ(0 * ByteSize::TB(1), ByteSize::Bytes(0));
// Divide.
EXPECT_EQ(ByteSize::Bytes(0) / 1, ByteSize::Bytes(0));
EXPECT_EQ(ByteSize::KB(1) / 2, ByteSize::KB(0.5));
EXPECT_EQ(ByteSize::MB(1) / 2, ByteSize::KB(512));
EXPECT_EQ(ByteSize::GB(1) / 2, ByteSize::MB(512));
EXPECT_EQ(ByteSize::TB(1.5) / 2, ByteSize::GB(768));
EXPECT_EQ(ByteSize::KB(1) / 0.5, ByteSize::KB(2));
EXPECT_EQ(ByteSize::MB(1) / 0.5, ByteSize::MB(2));
EXPECT_EQ(ByteSize::GB(1) / 0.5, ByteSize::GB(2));
EXPECT_EQ(ByteSize::TB(1) / 0.25, ByteSize::TB(4));
// Ratio.
EXPECT_DOUBLE_EQ(ByteSize::Bytes(0) / ByteSize::KB(1), 0.0);
EXPECT_DOUBLE_EQ(ByteSize::Bytes(1) / ByteSize::TB(1),
1.0 / 1024 / 1024 / 1024 / 1024);
EXPECT_DOUBLE_EQ(ByteSize::KB(1) / ByteSize::KB(2), 0.5);
EXPECT_DOUBLE_EQ(ByteSize::KB(512) / ByteSize::MB(1), 0.5);
EXPECT_DOUBLE_EQ(ByteSize::KB(1) / ByteSize::MB(1), 1.0 / 1024.0);
EXPECT_DOUBLE_EQ(ByteSize::MB(1) / ByteSize::GB(1), 1.0 / 1024.0);
EXPECT_DOUBLE_EQ(ByteSize::GB(1) / ByteSize::TB(1), 1.0 / 1024.0);
}
TEST(ByteSizeTest, Assignments) {
ByteSize byte_size;
EXPECT_EQ(byte_size, ByteSize::Bytes(0));
byte_size = ByteSize::Bytes(1);
EXPECT_EQ(byte_size, ByteSize::Bytes(1));
for (size_t i = 0; i < 1023; ++i) {
byte_size += ByteSize::Bytes(1);
}
EXPECT_EQ(byte_size, ByteSize::KB(1));
for (size_t i = 0; i < 10; ++i) {
byte_size *= 2;
}
EXPECT_EQ(byte_size, ByteSize::MB(1));
byte_size *= 1024 * 1024;
EXPECT_EQ(byte_size, ByteSize::TB(1));
for (size_t i = 0; i < 10; ++i) {
byte_size /= 2;
}
EXPECT_EQ(byte_size, ByteSize::GB(1));
for (size_t i = 0; i < 4; ++i) {
byte_size -= ByteSize::MB(256);
}
EXPECT_EQ(byte_size, ByteSize::Bytes(0));
// No negative bytes. The result will be 0 bytes.
byte_size -= ByteSize::Bytes(1);
EXPECT_EQ(byte_size, ByteSize::Bytes(0));
}
TEST(ByteSizeTest, Comparisons) {
EXPECT_LE(ByteSize::Bytes(0), ByteSize::Bytes(0));
EXPECT_LT(ByteSize::Bytes(0), ByteSize::Bytes(1));
EXPECT_LE(ByteSize::Bytes(0), ByteSize::Bytes(1));
EXPECT_LT(ByteSize::Bytes(1), ByteSize::Bytes(1024));
EXPECT_LE(ByteSize::Bytes(1), ByteSize::Bytes(1024));
EXPECT_LT(ByteSize::Bytes(1024), ByteSize::Bytes(1024 * 1024));
EXPECT_LE(ByteSize::Bytes(1024), ByteSize::Bytes(1024 * 1024));
EXPECT_LT(ByteSize::Bytes(1024), ByteSize::KB(1.1));
EXPECT_LE(ByteSize::Bytes(1024), ByteSize::KB(1.1));
EXPECT_LE(ByteSize::KB(0), ByteSize::Bytes(0));
EXPECT_LE(ByteSize::KB(1), ByteSize::Bytes(1024));
EXPECT_LT(ByteSize::KB(0), ByteSize::Bytes(1));
EXPECT_LE(ByteSize::KB(0), ByteSize::Bytes(1));
EXPECT_LT(ByteSize::KB(0.9), ByteSize::Bytes(1024));
EXPECT_LE(ByteSize::KB(0.9), ByteSize::Bytes(1024));
EXPECT_LT(ByteSize::KB(1), ByteSize::KB(1024));
EXPECT_LE(ByteSize::KB(1), ByteSize::KB(1024));
EXPECT_LT(ByteSize::KB(1), ByteSize::MB(1));
EXPECT_LE(ByteSize::KB(1), ByteSize::MB(1));
EXPECT_LT(ByteSize::KB(1024), ByteSize::MB(1.1));
EXPECT_LE(ByteSize::KB(1024), ByteSize::MB(1.1));
EXPECT_LE(ByteSize::MB(0), ByteSize::Bytes(0));
EXPECT_LT(ByteSize::MB(0), ByteSize::Bytes(1));
EXPECT_LE(ByteSize::MB(0), ByteSize::Bytes(1));
EXPECT_LT(ByteSize::MB(0.9), ByteSize::KB(1024));
EXPECT_LE(ByteSize::MB(0.9), ByteSize::KB(1024));
EXPECT_LT(ByteSize::MB(1), ByteSize::MB(1024));
EXPECT_LE(ByteSize::MB(1), ByteSize::MB(1024));
EXPECT_LT(ByteSize::MB(1), ByteSize::GB(1));
EXPECT_LE(ByteSize::MB(1), ByteSize::GB(1));
EXPECT_LT(ByteSize::MB(1024), ByteSize::GB(1.1));
EXPECT_LE(ByteSize::MB(1024), ByteSize::GB(1.1));
EXPECT_LE(ByteSize::GB(0), ByteSize::Bytes(0));
EXPECT_LT(ByteSize::GB(0), ByteSize::Bytes(1));
EXPECT_LE(ByteSize::GB(0), ByteSize::Bytes(1));
EXPECT_LT(ByteSize::GB(0.9), ByteSize::MB(1024));
EXPECT_LE(ByteSize::GB(0.9), ByteSize::MB(1024));
EXPECT_LT(ByteSize::GB(1), ByteSize::GB(1024));
EXPECT_LE(ByteSize::GB(1), ByteSize::GB(1024));
EXPECT_LT(ByteSize::GB(1), ByteSize::TB(1));
EXPECT_LE(ByteSize::GB(1), ByteSize::TB(1));
EXPECT_LT(ByteSize::GB(1024), ByteSize::TB(1.1));
EXPECT_LE(ByteSize::GB(1024), ByteSize::TB(1.1));
EXPECT_LE(ByteSize::TB(0), ByteSize::Bytes(0));
EXPECT_LT(ByteSize::TB(0), ByteSize::Bytes(1));
EXPECT_LE(ByteSize::TB(0), ByteSize::Bytes(1));
EXPECT_LT(ByteSize::TB(0.9), ByteSize::GB(1024));
EXPECT_LE(ByteSize::TB(0.9), ByteSize::GB(1024));
EXPECT_LT(ByteSize::TB(1), ByteSize::TB(1024));
EXPECT_LE(ByteSize::TB(1), ByteSize::TB(1024));
EXPECT_LT(ByteSize::TB(1024), ByteSize::TB(1025));
EXPECT_LE(ByteSize::TB(1024), ByteSize::TB(1025));
EXPECT_GT(ByteSize::TB(1), ByteSize::GB(1));
EXPECT_GT(ByteSize::GB(1), ByteSize::MB(1));
EXPECT_GT(ByteSize::MB(1), ByteSize::KB(1));
EXPECT_GT(ByteSize::KB(1), ByteSize::Bytes(1));
EXPECT_GT(ByteSize::Bytes(1), ByteSize::Bytes(0));
EXPECT_GT(ByteSize::TB(1), ByteSize::GB(1));
EXPECT_GT(ByteSize::TB(1), ByteSize::GB(1) + ByteSize::MB(1) +
ByteSize::KB(1) + ByteSize::Bytes(1));
EXPECT_GT(ByteSize::GB(1), 0.0000001 * ByteSize::TB(1));
EXPECT_GT(ByteSize::MB(1), ByteSize::KB(1) * 1023);
EXPECT_GT(ByteSize::KB(1), ByteSize::KB(3) / 4);
EXPECT_GT(ByteSize::Bytes(1), ByteSize::TB(0));
EXPECT_GE(ByteSize::TB(0.5), ByteSize::GB(0.5));
EXPECT_GE(ByteSize::GB(0.5), ByteSize::MB(0.5));
EXPECT_GE(ByteSize::MB(0.5), ByteSize::KB(0.5));
EXPECT_GE(ByteSize::KB(0.5), ByteSize::Bytes(1));
EXPECT_GE(ByteSize::Bytes(1), ByteSize::Bytes(0));
EXPECT_GE(ByteSize::TB(0), ByteSize::Bytes(0));
EXPECT_GE(ByteSize::GB(0), ByteSize::Bytes(0));
EXPECT_GE(ByteSize::MB(0), ByteSize::Bytes(0));
EXPECT_GE(ByteSize::KB(0), ByteSize::Bytes(0));
EXPECT_GE(ByteSize::Bytes(0), ByteSize::Bytes(0));
}
TEST(ByteSizeTest, DebugString) {
EXPECT_EQ(ByteSize::Bytes(0).DebugString(), "0B");
EXPECT_EQ(ByteSize::Bytes(1).DebugString(), "1B");
EXPECT_EQ(ByteSize::Bytes(size_t{1} << 10).DebugString(), "1KB");
EXPECT_EQ(ByteSize::Bytes(size_t{1} << 20).DebugString(), "1MB");
EXPECT_EQ(ByteSize::Bytes(size_t{1} << 30).DebugString(), "1GB");
EXPECT_EQ(ByteSize::Bytes(size_t{1} << 40).DebugString(), "1TB");
EXPECT_EQ(ByteSize::KB(0.5).DebugString(), "512B");
EXPECT_EQ(ByteSize::KB(1).DebugString(), "1KB");
EXPECT_EQ(ByteSize::KB(1.5).DebugString(), "1.5KB");
EXPECT_EQ(ByteSize::KB(1024).DebugString(), "1MB");
EXPECT_EQ(ByteSize::KB(1024 * 1024).DebugString(), "1GB");
EXPECT_EQ(ByteSize::KB(1024 * 1024 * 1024).DebugString(), "1TB");
EXPECT_EQ(ByteSize::MB(0.5).DebugString(), "512KB");
EXPECT_EQ(ByteSize::MB(1).DebugString(), "1MB");
EXPECT_EQ(ByteSize::MB(1.5).DebugString(), "1.5MB");
EXPECT_EQ(ByteSize::MB(1024).DebugString(), "1GB");
EXPECT_EQ(ByteSize::MB(1024 * 1024).DebugString(), "1TB");
EXPECT_EQ(ByteSize::GB(0.5).DebugString(), "512MB");
EXPECT_EQ(ByteSize::GB(1).DebugString(), "1GB");
EXPECT_EQ(ByteSize::GB(1.5).DebugString(), "1.5GB");
EXPECT_EQ(ByteSize::GB(1024).DebugString(), "1TB");
EXPECT_EQ(ByteSize::TB(0.5).DebugString(), "512GB");
EXPECT_EQ(ByteSize::TB(1).DebugString(), "1TB");
EXPECT_EQ(ByteSize::TB(1.5).DebugString(), "1.5TB");
EXPECT_EQ(ByteSize::TB(1024).DebugString(), "1024TB");
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/client/BUILD
+177
View File
@@ -0,0 +1,177 @@
# tf.data service client library.
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "tf_grpc_cc_dependencies")
load("//tensorflow/core/platform:build_config.bzl", "tf_protos_profiler_service")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
cc_library(
name = "common",
hdrs = [
"common.h",
],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data/service:common_proto_cc",
"@com_google_absl//absl/time",
],
)
cc_library(
name = "data_service_client",
srcs = ["data_service_client.cc"],
hdrs = ["data_service_client.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":common",
":validate_utils",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core/data:utils",
"//tensorflow/core/data/service:common",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:dispatcher_proto_cc",
"//tensorflow/core/data/service:grpc_util",
"//tensorflow/core/data/service:worker_client",
"//tensorflow/core/data/service:worker_impl",
"//tensorflow/core/data/service:worker_proto_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:statusor",
"//tensorflow/core/profiler/lib:traceme",
"//tensorflow/core/profiler/lib:traceme_encode",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:random",
"@tsl//tsl/platform:retrying_utils",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tf_cc_test(
name = "data_service_client_test",
srcs = ["data_service_client_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":common",
":data_service_client",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:common",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:test_cluster",
"//tensorflow/core/data/service:test_util",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:status_matchers",
"//tensorflow/core/platform:statusor",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/time",
"@com_google_googletest//:gtest",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
] + tf_grpc_cc_dependencies() + tf_protos_profiler_service(),
)
cc_library(
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:dispatcher_proto_cc",
"//tensorflow/core/data/service:grpc_util",
"//tensorflow/core/platform:env_time",
"//tensorflow/core/platform:statusor",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:errors",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
tf_cc_test(
name = "utils_test",
srcs = ["utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:test_cluster",
"//tensorflow/core/data/service:test_util",
"@com_google_googletest//:gtest",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:status_matchers",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/protobuf:protos_all_cc",
] + tf_grpc_cc_dependencies() + tf_protos_profiler_service(),
)
cc_library(
name = "validate_utils",
srcs = ["validate_utils.cc"],
hdrs = ["validate_utils.h"],
# copybara:uncomment copts = ["-Wthread-safety-analysis"],
deps = [
":common",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data/service:common",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:worker_impl",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "validate_utils_test",
srcs = ["validate_utils_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":common",
":validate_utils",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:worker_impl",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:status_matchers",
"@com_google_googletest//:gtest_main",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tensorflow/core/data/service/client/common.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2022 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_CORE_DATA_SERVICE_CLIENT_COMMON_H_
#define TENSORFLOW_CORE_DATA_SERVICE_CLIENT_COMMON_H_
#include <cstdint>
#include <optional>
#include <string>
#include "absl/time/time.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
// tf.data service parameters.
struct DataServiceParams final {
std::string dataset_id;
ProcessingModeDef processing_mode;
std::string address;
std::string protocol;
std::string data_transfer_protocol;
std::string job_name;
int64_t repetition = 0;
std::optional<int64_t> num_consumers;
std::optional<int64_t> consumer_index;
int64_t max_outstanding_requests = 0;
absl::Duration task_refresh_interval;
TargetWorkers target_workers = TargetWorkers::TARGET_WORKERS_UNSPECIFIED;
DataServiceMetadata metadata;
std::optional<CrossTrainerCacheOptions> cross_trainer_cache_options;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_CLIENT_COMMON_H_
tensorflow/core/data/service/client/data_service_client.cc
+997
View File
@@ -0,0 +1,997 @@
/* Copyright 2022 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/core/data/service/client/data_service_client.h"
#include <algorithm>
#include <chrono>
#include <cstdint>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <random>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_map.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/ascii.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/substitute.h"
#include "absl/time/time.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/service/client/common.h"
#include "tensorflow/core/data/service/client/validate_utils.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/grpc_util.h"
#include "tensorflow/core/data/service/worker.pb.h"
#include "tensorflow/core/data/service/worker_client.h"
#include "tensorflow/core/data/service/worker_impl.h"
#include "tensorflow/core/data/utils.h"
#include "tensorflow/core/framework/allocator.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/framework/model.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#include "tensorflow/core/profiler/lib/traceme_encode.h"
#include "tsl/platform/random.h"
#include "tsl/platform/retrying_utils.h"
namespace tensorflow {
namespace data {
namespace {
bool IsColocatedTask(const TaskInfo& task) {
return absl::c_any_of(task.worker_tags(), [](std::string_view worker_tag) {
return absl::AsciiStrToUpper(worker_tag) == kColocatedWorkerTag;
});
}
absl::StatusOr<DataTransferServerInfo> GetTransferServer(
const std::string& protocol, const TaskInfo& task_info) {
for (const auto& transfer_server : task_info.transfer_servers()) {
if (transfer_server.protocol() == protocol) {
return transfer_server;
}
}
return absl::NotFoundError(absl::StrCat("Protocol '", protocol,
"' is not available for worker '",
task_info.worker_address(), "'."));
}
} // namespace
DataServiceClient::DataServiceClient(const DataServiceParams& params)
: params_(params),
max_outstanding_requests_(params.max_outstanding_requests) {}
DataServiceClient::~DataServiceClient() {
VLOG(2) << "Destroying data service client for iteration id "
<< iteration_client_id_;
task_thread_manager_.reset();
if (initialized_) {
absl::Status s = dispatcher_->ReleaseIterationClient(iteration_client_id_);
if (!s.ok()) {
LOG(WARNING) << "Failed to release iteration client id: " << s;
}
}
for (auto& worker_thread : worker_threads_) {
worker_thread.reset();
}
DeleteLocalWorkerTasks();
VLOG(2) << "Destroyed data service dataset iterator for iteration id "
<< iteration_client_id_;
}
absl::Status DataServiceClient::Initialize(
const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
Allocator* allocator) {
accelerator_device_info_ = accelerator_device_info;
allocator_ = allocator;
TF_RETURN_IF_ERROR(ValidateDataServiceParams(params_));
VLOG(3) << "Connecting to " << params_.address
<< " in tf.data service client.";
dispatcher_ = std::make_unique<DataServiceDispatcherClient>(params_.address,
params_.protocol);
int64_t deadline_micros = std::numeric_limits<int64_t>::max();
std::optional<std::string> job_name;
if (!params_.job_name.empty()) {
job_name = params_.job_name;
}
TF_RETURN_IF_ERROR(grpc_util::Retry(
[&]() {
return dispatcher_->GetOrCreateJob(
params_.dataset_id, params_.processing_mode, job_name,
params_.num_consumers,
params_.cross_trainer_cache_options.has_value(),
params_.target_workers, job_id_);
},
/*description=*/
absl::StrCat("get or create job with dispatcher at ", params_.address),
deadline_micros));
TF_RETURN_IF_ERROR(grpc_util::Retry(
[&]() {
return dispatcher_->GetOrCreateIteration(job_id_, params_.repetition,
iteration_client_id_);
},
/*description=*/
absl::StrCat("get or create iteration with dispatcher at ",
params_.address),
deadline_micros));
initialized_ = true;
return absl::OkStatus();
}
absl::StatusOr<GetNextResult> DataServiceClient::GetNext(
DataServiceContextFactory context_factory) TF_LOCKS_EXCLUDED(mu_) {
VLOG(3) << "Getting the next element from tf.data service client.";
mutex_lock l(mu_);
if (ctx_ == nullptr) {
ctx_ = context_factory();
}
EnsureThreadsStarted();
std::shared_ptr<Result> result;
do {
while (!ResultReady() && !Finished() && !cancelled_ && status_.ok()) {
VLOG(3) << "Blocking in GetNext: " << DebugString();
get_next_cv_.wait(l);
}
if (cancelled_) {
VLOG(3) << "Returning from GetNext due to cancellation";
return absl::CancelledError("Data service iterator was cancelled");
}
if (!status_.ok()) {
VLOG(3) << "Returning from GetNext with error " << status_;
return status_;
}
if (results_.empty()) {
VLOG(3) << "Returning from GetNext with end_of_sequence";
return GetNextResult::EndOfSequence();
}
if (!ResultReady()) {
VLOG(3) << "Returning from GetNext with internal error";
return absl::InternalError(
"Expected a result to be ready, but none were.");
}
result = PopNextResult();
worker_thread_cv_.notify_one();
if (result->skip) {
VLOG(3) << "Skipping result from task " << result->task_id;
}
} while (result->skip);
GetNextResult next;
next.end_of_sequence = result->end_of_sequence;
if (next.end_of_sequence) {
VLOG(1) << "Returning end_of_sequence";
return next;
}
VLOG(1) << "Returning the next element from data service dataset's "
<< "Iterator: task " << result->task_id << ", element "
<< result->element_index;
if (IsCoordinatedRead()) {
VLOG(1) << "Consumer " << *params_.consumer_index << ": Result "
<< get_next_index_++;
}
next.tensors.swap(result->element);
return next;
}
void DataServiceClient::Cancel() TF_LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
for (const auto& task : tasks_) {
task->worker->TryCancel();
}
cancelled_ = true;
worker_thread_cv_.notify_all();
manager_thread_cv_.notify_all();
get_next_cv_.notify_all();
}
TraceMeMetadata DataServiceClient::GetTraceMeMetadata() const {
TraceMeMetadata result;
int64_t num_tasks = -1;
int64_t autotuned_max_outstanding_requests = model::kAutotune;
if (mu_.try_lock()) {
num_tasks = tasks_.size() - finished_tasks_;
autotuned_max_outstanding_requests = max_outstanding_requests_;
mu_.unlock();
}
result.push_back(std::make_pair(
"num_tasks",
num_tasks == -1
? kTraceInfoUnavailable
: absl::StrFormat("%lld", static_cast<long long>(num_tasks))));
result.push_back(std::make_pair("job_name", params_.job_name));
result.push_back(std::make_pair(
"max_outstanding_requests",
absl::StrFormat(
"%lld", static_cast<long long>(params_.max_outstanding_requests))));
if (params_.max_outstanding_requests == model::kAutotune) {
result.push_back(std::make_pair(
"autotuned_max_outstanding_requests",
absl::StrFormat("%lld", static_cast<long long>(
autotuned_max_outstanding_requests))));
}
return result;
}
void DataServiceClient::EnsureThreadsStarted()
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (!task_thread_manager_ && !cancelled_) {
task_thread_manager_ = ctx_->StartThread("task-thread-manager",
[this]() { TaskThreadManager(); });
}
}
bool DataServiceClient::Finished() const TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
return num_running_worker_threads_ == 0 && !ShouldWaitForNext();
}
bool DataServiceClient::ShouldWaitForNext() const
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (should_finish_iteration_) {
return !iteration_finished_;
}
return tasks_.empty() || finished_tasks_ < tasks_.size();
}
void DataServiceClient::DeleteLocalWorkerTasks() TF_LOCKS_EXCLUDED(mu_) {
std::vector<std::shared_ptr<Task>> tasks;
{
mutex_lock l(mu_);
tasks = tasks_;
}
for (const std::shared_ptr<Task>& task : tasks) {
std::shared_ptr<DataServiceWorkerImpl> worker =
LocalWorkers::Get(task->info.worker_address());
if (worker && ShouldDeleteLocalTask(task->info)) {
worker->DeleteLocalTask(task->info);
}
}
}
// Deletes the task if it is only read by the local client.
bool DataServiceClient::ShouldDeleteLocalTask(const TaskInfo& task) const
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (IsCoordinatedRead()) {
return false;
}
if (params_.target_workers == TARGET_WORKERS_LOCAL) {
return true;
}
return params_.target_workers == TARGET_WORKERS_AUTO && IsColocatedTask(task);
}
void DataServiceClient::TaskThreadManager() TF_LOCKS_EXCLUDED(mu_) {
auto cleanup =
gtl::MakeCleanup([] { VLOG(1) << "Task thread manager exiting"; });
VLOG(1) << "Starting task thread manager";
uint64_t next_check = Env::Default()->NowMicros();
while (true) {
{
mutex_lock l(mu_);
// All units are microseconds.
while (!cancelled_ && Env::Default()->NowMicros() < next_check) {
int64_t remaining_time = next_check - Env::Default()->NowMicros();
VLOG(4) << "Task thread manager waiting for " << remaining_time << "us";
manager_thread_cv_.wait_for(l,
std::chrono::microseconds(remaining_time));
}
if (cancelled_) {
VLOG(3) << "Task thread manager finished";
return;
}
}
Heartbeat();
UpdateBufferSize();
UpdateWorkerThreads();
next_check = Env::Default()->NowMicros() +
absl::ToInt64Microseconds(params_.task_refresh_interval);
}
}
void DataServiceClient::TryBlockRound(int64_t round)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (round_robin_round_limit_.has_value() &&
round_robin_round_limit_.value() == round) {
return;
}
if (current_round_ >= round) {
// In the next heartbeat, notify the dispatcher that we failed to add
// the task.
VLOG(1) << "Rejecting request to block round " << round
<< ", because processing has already begun for round "
<< current_round_;
return;
}
VLOG(1) << "Accepting request to block round " << round;
round_robin_round_limit_ = round;
}
void DataServiceClient::UpdateIterationFinished(bool iteration_finished)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (!iteration_finished) {
return;
}
iteration_finished_ = true;
get_next_cv_.notify_all();
worker_thread_cv_.notify_all();
}
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
DataServiceClient::CreateWorkerClient(const std::string& protocol,
const TaskInfo& task_info) {
TF_ASSIGN_OR_RETURN(DataTransferServerInfo transfer_server,
GetTransferServer(protocol, task_info));
return CreateDataServiceWorkerClient(params_.protocol, transfer_server,
accelerator_device_info_, allocator_);
}
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
DataServiceClient::CreateGrpcWorkerClient(const TaskInfo& task_info) {
return CreateWorkerClient(kGrpcTransferProtocol, task_info);
}
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
DataServiceClient::CreateAlternativeWorkerClientMaybeWithGrpcFallback(
const DataTransferServerInfo& transfer_server, const TaskInfo& task_info) {
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>> worker =
CreateDataServiceWorkerClient(params_.protocol, transfer_server,
accelerator_device_info_, allocator_);
if (worker.ok()) {
LOG(INFO) << "Successfully started client for data transfer protocol '"
<< transfer_server.protocol() << "' for worker '"
<< task_info.worker_address() << "'.";
return worker;
}
std::string client_creation_error_message =
absl::StrCat("Failed to start client for data transfer protocol '",
transfer_server.protocol(), "' for worker '",
task_info.worker_address(), "'.");
if (!transfer_server.fall_back_to_grpc_at_client_creation_time()) {
return absl::InternalError(
absl::StrCat(client_creation_error_message,
" Original error: ", worker.status().message()));
}
LOG(INFO) << client_creation_error_message
<< "; falling back to gRPC. Original error: " << worker.status();
metrics::RecordTFDataServiceDataTransferProtocolFallback(
transfer_server.protocol(),
static_cast<error::Code>(worker.status().raw_code()),
std::string(worker.status().message()));
return CreateGrpcWorkerClient(task_info);
}
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
DataServiceClient::CreateWorkerClient(const TaskInfo& task_info) {
if (params_.data_transfer_protocol == kLocalTransferProtocol ||
ForceLocalProtocol(task_info.worker_address())) {
DataTransferServerInfo info;
info.set_protocol(kLocalTransferProtocol);
info.set_address(task_info.worker_address());
return CreateDataServiceWorkerClient(params_.protocol, info,
accelerator_device_info_, allocator_);
}
if (!params_.data_transfer_protocol.empty()) {
TF_ASSIGN_OR_RETURN(
DataTransferServerInfo transfer_server,
GetTransferServer(params_.data_transfer_protocol, task_info));
return CreateAlternativeWorkerClientMaybeWithGrpcFallback(transfer_server,
task_info);
}
if (std::string default_protocol = DefaultDataTransferProtocol();
default_protocol != kGrpcTransferProtocol) {
absl::StatusOr<DataTransferServerInfo> transfer_server =
GetTransferServer(default_protocol, task_info);
if (transfer_server.ok()) {
return CreateAlternativeWorkerClientMaybeWithGrpcFallback(
*transfer_server, task_info);
}
VLOG(1) << "Failed to find transfer server for default data transfer "
"protocol '"
<< default_protocol << "' for worker '"
<< task_info.worker_address()
<< "'; falling back to grpc. Original error: "
<< transfer_server.status();
metrics::RecordTFDataServiceDataTransferProtocolFallback(
default_protocol, error::Code::NOT_FOUND,
"Failed to find transfer server for default protocol");
}
return CreateGrpcWorkerClient(task_info);
}
absl::Status DataServiceClient::AddTask(const TaskInfo& task_info)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<DataServiceWorkerClient> worker,
CreateWorkerClient(task_info));
metrics::RecordTFDataServiceDataTransferProtocolUsed(
worker->GetDataTransferProtocol(),
/*user_specified=*/!params_.data_transfer_protocol.empty());
tasks_.push_back(std::make_shared<Task>(task_info, std::move(worker)));
worker_thread_cv_.notify_one();
if (IsCoordinatedRead()) {
VLOG(1) << "Consumer " << params_.consumer_index.value() << " adding task "
<< task_info.task_id() << " to read from worker "
<< task_info.worker_address()
<< ". Task starting round: " << task_info.starting_round();
DCHECK_LE(current_round_, task_info.starting_round());
if (current_round_ == task_info.starting_round()) {
DCHECK_EQ(next_task_index_, 0);
}
}
return absl::OkStatus();
}
void DataServiceClient::Heartbeat() TF_LOCKS_EXCLUDED(mu_) {
ClientHeartbeatRequest req;
req.set_iteration_client_id(iteration_client_id_);
if (IsCoordinatedRead()) {
mutex_lock l(mu_);
req.set_current_round(current_round_);
if (round_robin_round_limit_.has_value()) {
req.set_blocked_round(round_robin_round_limit_.value());
}
}
{
mutex_lock l(mu_);
double target_processing_time_nsec = ctx_->GetTargetProcessingTimeNsec();
req.set_target_processing_time_nsec(target_processing_time_nsec);
}
ClientHeartbeatResponse resp;
absl::Status s = dispatcher_->ClientHeartbeat(req, resp);
if (!s.ok()) {
if (IsPreemptedError(s)) {
LOG(WARNING)
<< "Failed to heartbeat to dispatcher from iteration client id "
<< iteration_client_id_ << ". Dispatcher address: " << params_.address
<< ". Error: " << s;
return;
}
mutex_lock l(mu_);
status_ = s;
get_next_cv_.notify_all();
}
mutex_lock l(mu_);
UpdateIterationFinished(resp.iteration_finished());
if (resp.optional_block_round_case() ==
ClientHeartbeatResponse::kBlockRound) {
TryBlockRound(resp.block_round());
} else {
round_robin_round_limit_ = std::nullopt;
worker_thread_cv_.notify_all();
}
UpdateTasks(resp);
RecordTFMetrics(resp);
}
void DataServiceClient::UpdateTasks(const ClientHeartbeatResponse& resp)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
absl::flat_hash_map<int64_t, TaskInfo> task_id_to_task;
for (auto& task : resp.task_info()) {
task_id_to_task[task.task_id()] = task;
}
if (iteration_finished_) {
return;
}
int index = 0;
while (index < tasks_.size()) {
std::shared_ptr<Task> task = tasks_[index];
if (task_id_to_task.contains(task->info.task_id())) {
// Remove already-known tasks from `task_id_to_task`, so that at the
// end of the loop, only new tasks remain.
task_id_to_task.erase(task->info.task_id());
++index;
} else {
// Task has been removed.
if (task->end_of_sequence) {
finished_tasks_--;
}
tasks_.erase(tasks_.begin() + index);
if (index < next_task_index_) {
next_task_index_--;
}
if (!tasks_.empty() && next_task_index_ >= tasks_.size()) {
AdvanceTaskIndex();
}
}
}
for (auto& task : resp.task_info()) {
auto it = task_id_to_task.find(task.task_id());
if (it == task_id_to_task.end()) {
continue;
}
if (!ShouldReadFromTask(task)) {
VLOG(3) << "Skipping untargeted worker task " << task.task_id();
should_finish_iteration_ = false;
continue;
}
absl::Status s = AddTask(it->second);
if (!s.ok()) {
status_ = s;
get_next_cv_.notify_all();
break;
}
}
if (!IsCoordinatedRead()) {
// Shuffle task order within each client to avoid thundering herd effect.
std::mt19937 rng(tsl::random::New64());
std::shuffle(tasks_.begin(), tasks_.end(), rng);
}
}
bool DataServiceClient::ShouldReadFromTask(const TaskInfo& task) const
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (IsCoordinatedRead()) {
return true;
}
const bool is_local_task =
(LocalWorkers::Get(task.worker_address()) != nullptr);
if (params_.target_workers == TARGET_WORKERS_LOCAL && !is_local_task) {
return false;
}
// Cross-TF/TPU host reads may cause resource contention on the TF/TPU
// hosts. tf.data service avoids reading from non-local TF-hosted workers.
const bool is_cross_tf_host_read = !is_local_task && IsColocatedTask(task);
if (params_.target_workers == TARGET_WORKERS_AUTO && is_cross_tf_host_read) {
return false;
}
return true;
}
void DataServiceClient::RecordTFMetrics(const ClientHeartbeatResponse& resp)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
for (const auto& task : resp.task_info()) {
if (worker_uids_.contains(task.worker_uid())) {
continue;
}
metrics::RecordTFDataServiceClientIterators(
task.worker_uid(), resp.deployment_mode(), params_.processing_mode,
IsCoordinatedRead());
worker_uids_.insert(task.worker_uid());
}
}
void DataServiceClient::UpdateBufferSize() TF_LOCKS_EXCLUDED(mu_) {
if (params_.max_outstanding_requests == model::kAutotune) {
// Adjust `max_outstanding_requests_` to account for newly added tasks.
// `tasks_` includes the local tasks, so we subtract one from the
// configured local task buffer size.
mutex_lock l(mu_);
int64_t max_outstanding_requests = ctx_->UpdateMaxOutstandingRequests(
max_outstanding_requests_, tasks_.size());
if (max_outstanding_requests > max_outstanding_requests_) {
worker_thread_cv_.notify_all();
}
VLOG(3) << "Updated `max_outstanding_requests` from "
<< max_outstanding_requests_ << " to " << max_outstanding_requests
<< " with " << tasks_.size() << " tasks.";
max_outstanding_requests_ = max_outstanding_requests;
}
}
void DataServiceClient::UpdateWorkerThreads() TF_LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
const int64_t max_num_threads =
std::min<int64_t>(tasks_.size(), max_outstanding_requests_);
while (num_running_worker_threads_ < max_num_threads && !cancelled_ &&
status_.ok()) {
num_running_worker_threads_++;
auto done = [this]() {
mutex_lock l(mu_);
num_running_worker_threads_--;
get_next_cv_.notify_all();
};
int64_t thread_index = worker_threads_.size();
worker_threads_.push_back(
ctx_->StartThread("tf-data-service-task_thread",
[this, thread_index, done = std::move(done)]() {
RunWorkerThread(thread_index, std::move(done));
}));
}
}
void DataServiceClient::RunWorkerThread(int64_t thread_index,
std::function<void()> done)
TF_LOCKS_EXCLUDED(mu_) {
auto cleanup = gtl::MakeCleanup([done = std::move(done)]() {
done();
VLOG(1) << "Worker thread exiting";
});
VLOG(1) << "Starting worker thread";
std::shared_ptr<Task> task_to_process;
int64_t num_consecutive_skipped = 0;
constexpr int64_t MAX_ROUND_FALLBACK_TO_BLOCKING = 5;
bool allow_skip = true;
while (true) {
std::shared_ptr<Result> result;
{
mutex_lock l(mu_);
if (task_to_process) {
task_to_process->in_use = false;
--outstanding_requests_;
task_to_process = nullptr;
worker_thread_cv_.notify_one();
}
while (true) {
if (cancelled_ || !ShouldWaitForNext()) {
return;
}
task_to_process = GetTaskToProcess();
if (task_to_process) {
VLOG(3) << "Selected a task to process: "
<< task_to_process->info.ShortDebugString();
break;
}
worker_thread_cv_.wait(l);
}
DCHECK(task_to_process != nullptr);
task_to_process->in_use = true;
++outstanding_requests_;
if (IsCoordinatedRead()) {
// Reserve a spot in the results_ queue.
results_.push(std::make_shared<Result>());
ctx_->RecordBufferEnqueue(results_.back()->element);
result = results_.back();
} else {
// The result will be added to results_ when it's ready.
result = std::make_shared<Result>();
}
VLOG(3) << "Processing task " << task_to_process->info.task_id();
}
int64_t deadline_micros = std::numeric_limits<int64_t>::max();
absl::Status s = GetElementTraced(task_to_process.get(), deadline_micros,
/*enqueue_result=*/!IsCoordinatedRead(),
allow_skip, result, thread_index);
if (!s.ok()) {
mutex_lock l(mu_);
VLOG(1) << "Failed to get element from worker "
<< task_to_process->info.worker_address() << ": " << s;
task_to_process->in_use = false;
--outstanding_requests_;
status_ = errors::CreateWithUpdatedMessage(
s, absl::StrCat("Failed to get element from worker ",
task_to_process->info.worker_address(), ": ",
s.message()));
get_next_cv_.notify_all();
return;
}
if (!IsCoordinatedRead()) {
if (mutex_lock l(mu_); result->skip) {
num_consecutive_skipped++;
if (num_consecutive_skipped >=
MAX_ROUND_FALLBACK_TO_BLOCKING * tasks_.size()) {
// Switches to blocking call when we already skip
// all workers enough rounds.
// This is to ensures we do not spam the network traffic.
allow_skip = false;
VLOG(1) << "`allow_skip` is turned off. Switching to blocking "
"get element calls to the workers.";
}
} else {
num_consecutive_skipped = 0;
allow_skip = true;
}
}
}
}
// Reports whether we can request another element without violating
// `max_outstanding_requests_`.
bool DataServiceClient::ShouldProcessTask() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
// When doing round-robin reads, outstanding requests pre-allocate a
// result in `results_`, so we only need to check the size of `results_`.
if (IsCoordinatedRead()) {
return results_.size() < max_outstanding_requests_;
}
// Otherwise, results aren't added to `results_` until the data has been
// successfully retrieved. We need to count requests already added to
// `results_` as well as in-progress requests.
return results_.size() + outstanding_requests_ < max_outstanding_requests_;
}
// Searches for a task to process, visiting tasks in-order and giving every
// task a chance to proceed.
std::shared_ptr<DataServiceClient::Task> DataServiceClient::GetTaskToProcess()
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
if (!ShouldProcessTask()) {
return nullptr;
}
for (int i = 0; i < tasks_.size(); ++i) {
std::shared_ptr<Task>& task = tasks_[next_task_index_];
if (IsCoordinatedRead() &&
(task->in_use ||
current_round_ >= round_robin_round_limit_.value_or(
std::numeric_limits<int64_t>::max()))) {
VLOG(4) << "No round robin task found. in_use: " << task->in_use
<< ". current_round: " << current_round_
<< ". round_robin_round_limit: "
<< round_robin_round_limit_.value_or(-1);
return nullptr;
}
if (current_round_ < task->info.starting_round() || task->in_use ||
task->end_of_sequence || task->removed) {
VLOG(3) << "Skipping task " << next_task_index_
<< ". starting round: " << task->info.starting_round()
<< ". current round: " << current_round_
<< ". task->in_use: " << task->in_use
<< ". end_of_sequence: " << task->end_of_sequence
<< ". task->removed: " << task->removed;
AdvanceTaskIndex();
continue;
}
task->round = current_round_;
AdvanceTaskIndex();
return task;
}
return nullptr;
}
// Increments the next task index, starting over if all tasks have been
// processed.
void DataServiceClient::AdvanceTaskIndex() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
next_task_index_++;
if (next_task_index_ >= tasks_.size()) {
current_round_++;
next_task_index_ = 0;
}
}
absl::Status DataServiceClient::TryGetElement(const Task& task, bool allow_skip,
GetElementResult& result) {
GetElementRequest req;
req.set_task_id(task.info.task_id());
req.set_skipped_previous_round(task.skipped_previous_round);
if (IsCoordinatedRead()) {
req.set_consumer_index(params_.consumer_index.value());
req.set_round_index(task.round);
req.set_allow_skip(true);
} else {
req.set_allow_skip(allow_skip);
}
if (params_.cross_trainer_cache_options) {
req.set_trainer_id(params_.cross_trainer_cache_options->trainer_id());
}
return task.worker->GetElement(req, result);
}
void DataServiceClient::ProcessGetElementResponse(
bool enqueue_result, GetElementResult& get_element_result,
std::shared_ptr<Result> result, Task& task) TF_LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
result->ready = true;
result->end_of_sequence = get_element_result.end_of_sequence;
result->skip = get_element_result.skip;
if (!get_element_result.end_of_sequence && !get_element_result.skip) {
task.skipped_previous_round = false;
result->element = std::move(get_element_result.components);
result->element_index = get_element_result.element_index;
result->task_id = task.info.task_id();
} else if (get_element_result.skip) {
task.skipped_previous_round = true;
} else {
task.end_of_sequence = true;
finished_tasks_++;
}
if (enqueue_result && !result->end_of_sequence && !result->skip) {
ctx_->RecordBufferEnqueue(result->element);
results_.push(std::move(result));
}
get_next_cv_.notify_all();
}
absl::Status DataServiceClient::GetElementTraced(
Task* task, int64_t deadline_micros, bool enqueue_result, bool allow_skip,
std::shared_ptr<Result> result, int64_t thread_index) {
VLOG(3) << "Getting an element for task id " << task->info.task_id();
tsl::profiler::TraceMe activity("GetDataServiceElement",
tsl::profiler::TraceMeLevel::kInfo);
activity.AppendMetadata([&]() {
return tsl::profiler::TraceMeEncode(
{{"address", task->info.worker_address()}});
});
if (IsCoordinatedRead()) {
VLOG(3) << "Requesting element from consumer index "
<< params_.consumer_index.value() << ", round " << task->round;
activity.AppendMetadata([&]() {
return tsl::profiler::TraceMeEncode(
{{"consumer_index", params_.consumer_index.value()},
{"round_index", task->round}});
});
}
absl::Status s = GetElement(task, deadline_micros, enqueue_result, allow_skip,
result, thread_index);
VLOG(3) << "Got an element for task id " << task->info.task_id();
return s;
}
absl::Status DataServiceClient::MaybeRemoveTask(Task& task,
int64_t deadline_micros,
Result& result)
TF_LOCKS_EXCLUDED(mu_) {
bool removed;
VLOG(1) << "Requesting task removal for worker " << task.info.worker_address()
<< " in round " << task.round;
TF_RETURN_IF_ERROR(grpc_util::Retry(
[&] {
return dispatcher_->MaybeRemoveTask(task.info.task_id(),
params_.consumer_index.value(),
task.round, removed);
},
/*should_retry=*/
[&] {
mutex_lock l(mu_);
return !cancelled_;
},
/*description=*/"request task removal ", deadline_micros));
if (removed) {
mutex_lock l(mu_);
task.removed = true;
result.ready = true;
result.skip = true;
get_next_cv_.notify_all();
return absl::OkStatus();
}
VLOG(1) << "Failed to remove task for worker " << task.info.worker_address();
return absl::OkStatus();
}
absl::Status DataServiceClient::GetElement(Task* task, int64_t deadline_micros,
bool enqueue_result, bool allow_skip,
std::shared_ptr<Result> result,
int64_t thread_index)
TF_LOCKS_EXCLUDED(mu_) {
GetElementResult get_element_result;
while (true) {
absl::Status s = TryGetElement(*task, allow_skip, get_element_result);
if (s.ok()) {
task->num_retries = 0;
if (get_element_result.skip) {
metrics::RecordTFDataClientGetElementAction(
"skip_empty_buffer", absl::StrCat(iteration_client_id_),
task->info.worker_address(), absl::StrCat(thread_index));
} else if (!get_element_result.end_of_sequence) {
metrics::RecordTFDataClientGetElementAction(
"success", absl::StrCat(iteration_client_id_),
task->info.worker_address(), absl::StrCat(thread_index));
}
break;
}
if (!IsPreemptedError(s)) {
if (task->worker->GetDataTransferProtocol() == kGrpcTransferProtocol ||
task->worker->GetDataTransferProtocol() == kLocalTransferProtocol) {
return absl::Status(
s.code(),
absl::StrCat(
"Failed to get an element, with a nonretryable error: ",
s.message()));
}
if (!task->worker->FallBackToGrpcAtGetElementTime()) {
return absl::Status(
s.code(),
absl::StrCat("Failed to get an element over data "
"transfer protocol '",
task->worker->GetDataTransferProtocol(),
"', with a nonretryable error: ", s.message()));
}
LOG(ERROR) << "Failed to get an element over data transfer protocol '"
<< task->worker->GetDataTransferProtocol()
<< "', with a nonretryable error; falling back to grpc. "
"Original error: "
<< s;
metrics::RecordTFDataServiceDataTransferProtocolError(
task->worker->GetDataTransferProtocol(),
static_cast<error::Code>(s.raw_code()), std::string(s.message()));
mutex_lock l(mu_);
TF_ASSIGN_OR_RETURN(std::unique_ptr<DataServiceWorkerClient> worker,
CreateGrpcWorkerClient(task->info));
task->worker = std::move(worker);
continue;
}
{
mutex_lock l(mu_);
if (cancelled_) {
return absl::CancelledError("DataServiceDataset iterator cancelled");
}
}
int64_t now_micros = Env::Default()->NowMicros();
if (now_micros > deadline_micros) {
return s;
}
if (IsCoordinatedRead() && task->num_retries > 0) {
TF_RETURN_IF_ERROR(MaybeRemoveTask(*task, deadline_micros, *result));
mutex_lock l(mu_);
if (result->skip) {
return absl::OkStatus();
}
}
int64_t backoff_until = std::min(
deadline_micros,
now_micros + absl::ToInt64Microseconds(
tsl::ComputeRetryBackoff(task->num_retries++)));
VLOG(1) << "Failed to get an element from worker "
<< task->info.worker_address() << ": " << s << ". Will retry in "
<< (backoff_until - now_micros) << " microseconds";
Env::Default()->SleepForMicroseconds(backoff_until - now_micros);
if (!IsCoordinatedRead()) {
mutex_lock l(mu_);
// Mark the result as skipped so that we try reading from a different
// task before returning to this one.
result->ready = true;
result->skip = true;
metrics::RecordTFDataClientGetElementAction(
"skip_error", absl::StrCat(iteration_client_id_),
task->info.worker_address(), absl::StrCat(thread_index));
return absl::OkStatus();
}
}
ProcessGetElementResponse(enqueue_result, get_element_result, result, *task);
return absl::OkStatus();
}
bool DataServiceClient::ResultReady() const TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
return !results_.empty() && results_.front()->ready;
}
std::shared_ptr<DataServiceClient::Result> DataServiceClient::PopNextResult()
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
std::shared_ptr<Result> result = results_.front();
results_.pop();
ctx_->RecordBufferDequeue(result->element);
return result;
}
bool DataServiceClient::IsCoordinatedRead() const {
return params_.num_consumers.has_value();
}
std::string DataServiceClient::DebugString() const
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
return absl::Substitute(
"results_ { size: $0 front.ready: $1 } iteration_finished_: $2 "
"tasks { size: $3 } finished_tasks_: $4 "
"num_running_worker_threads_: $5",
results_.size(), !results_.empty() && results_.front()->ready,
iteration_finished_, tasks_.size(), finished_tasks_,
num_running_worker_threads_);
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/client/data_service_client.h
+278
View File
@@ -0,0 +1,278 @@
/* Copyright 2022 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_CORE_DATA_SERVICE_CLIENT_DATA_SERVICE_CLIENT_H_
#define TENSORFLOW_CORE_DATA_SERVICE_CLIENT_DATA_SERVICE_CLIENT_H_
#include <cstdint>
#include <functional>
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "tensorflow/core/data/service/client/common.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/worker_client.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
namespace data {
// Interface for interacting with the tf.data service iterator context.
class DataServiceContext {
public:
virtual ~DataServiceContext() = default;
virtual std::unique_ptr<Thread> StartThread(const std::string& name,
std::function<void()> fn) = 0;
virtual void RecordBufferEnqueue(const std::vector<Tensor>& element) = 0;
virtual void RecordBufferDequeue(const std::vector<Tensor>& element) = 0;
// Returns the time in nanoseconds a tf.data input pipeline can take to
// produce an element such that the downstream processor wait time is 0.
// Returns 0 if there are not sufficient recorded iterator gap times to
// produce a good estimate, or the tf.data Model instance is null.
virtual double GetTargetProcessingTimeNsec() const = 0;
// Updates the `max_outstanding_requests` with
// `requested_outstanding_requests`.
// Returns the new max outstanding requests which may be different from the
// requested one depending on available ram.
virtual int64_t UpdateMaxOutstandingRequests(
int64_t max_outstanding_requests,
int64_t requested_outstanding_requests) = 0;
};
using DataServiceContextFactory =
std::function<std::unique_ptr<DataServiceContext>()>;
// API for reading data from tf.data service.
//
// The client works by reading from tf.data workers in parallel and interleaving
// the dataset elements. It periodically queries the dispatcher to decide which
// workers to read from (in case workers are added or removed). The data reading
// is non-deterministic. This class is thread-safe.
class DataServiceClient {
public:
explicit DataServiceClient(const DataServiceParams& params);
virtual ~DataServiceClient();
DataServiceClient(const DataServiceClient&) = delete;
DataServiceClient& operator=(const DataServiceClient&) = delete;
// Initializes the client.
absl::Status Initialize(
const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info,
Allocator* allocator);
// Reads the next element from tf.data workers. Blocks if the next element is
// not ready.
virtual absl::StatusOr<GetNextResult> GetNext(
DataServiceContextFactory context_factory);
// Cancels the client.
void Cancel();
TraceMeMetadata GetTraceMeMetadata() const;
private:
struct Task {
Task(const TaskInfo& info, std::unique_ptr<DataServiceWorkerClient> worker)
: info(info), worker(std::move(worker)) {}
const TaskInfo info;
// Client for fetching task elements from the tf.data service worker.
std::unique_ptr<DataServiceWorkerClient> worker;
// The next round to read from the task.
int64_t round = 0;
// Whether the task has been removed. The task will eventually be
// deleted from `tasks_` on the next dispatcher heartbeat.
bool removed = false;
bool skipped_previous_round = false;
// Indicates whether a worker thread is currently processing the task.
bool in_use TF_GUARDED_BY(&DataServiceClient::mu_) = false;
// Indicates whether the worker has returned end_of_sequence for the task.
bool end_of_sequence TF_GUARDED_BY(&DataServiceClient::mu_) = false;
// Number of retries. The more it is retried, the longer it should wait
// before the next retry.
int64_t num_retries = 0;
};
struct Result {
Result() = default;
Result(Result&&) = default;
Result& operator=(Result&&) = default;
Result(const Result&) = delete;
Result& operator=(const Result&) = delete;
// Whether the result has been computed yet. GetNext needs to block
// until the next result is ready.
bool ready TF_GUARDED_BY(&DataServiceClient::mu_) = false;
std::vector<Tensor> element TF_GUARDED_BY(&DataServiceClient::mu_);
// The element's index within the tf.data worker it came from. Used for
// debugging.
int64_t element_index TF_GUARDED_BY(&DataServiceClient::mu_) = -1;
// The id of the task that generated the result.
int64_t task_id TF_GUARDED_BY(&DataServiceClient::mu_) = -1;
bool end_of_sequence TF_GUARDED_BY(&DataServiceClient::mu_) = false;
bool skip TF_GUARDED_BY(&DataServiceClient::mu_) = false;
};
void EnsureThreadsStarted();
void CancelThreads();
// Returns whether the client has finished and should return.
bool Finished() const;
// Returns whether the job has more data.
bool ShouldWaitForNext() const;
void DeleteLocalWorkerTasks();
bool ShouldDeleteLocalTask(const TaskInfo& task) const;
// Periodically refresh the task list.
// Maintain one thread fetching elements for each task.
// TODO(aaudibert): Instead of polling, have dispatcher send updates when
// the list of tasks changes.
void TaskThreadManager();
void TryBlockRound(int64_t round) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
void UpdateIterationFinished(bool iteration_finished);
absl::Status AddTask(const TaskInfo& task_info);
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>> CreateWorkerClient(
const TaskInfo& task_info);
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>> CreateWorkerClient(
const std::string& protocol, const TaskInfo& task_info);
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
CreateGrpcWorkerClient(const TaskInfo& task_info);
absl::StatusOr<std::unique_ptr<DataServiceWorkerClient>>
CreateAlternativeWorkerClientMaybeWithGrpcFallback(
const DataTransferServerInfo& transfer_server, const TaskInfo& task_info);
void Heartbeat();
void UpdateTasks(const ClientHeartbeatResponse& resp);
bool ShouldReadFromTask(const TaskInfo& task) const;
void RecordTFMetrics(const ClientHeartbeatResponse& resp);
void UpdateBufferSize();
void UpdateWorkerThreads();
// `thread_index` is the index of the running worker thread in
// `worker_threads_`, used for observability.
void RunWorkerThread(int64_t thread_index, std::function<void()> done);
// Reports whether we can request another element without violating
// `max_outstanding_requests_`.
bool ShouldProcessTask();
// Searches for a task to process, visiting tasks in-order and giving every
// task a chance to proceed.
std::shared_ptr<Task> GetTaskToProcess();
void AdvanceTaskIndex();
absl::Status TryGetElement(const Task& task, bool allow_skip,
GetElementResult& result);
void ProcessGetElementResponse(bool enqueue_result,
GetElementResult& get_element_result,
std::shared_ptr<Result> result, Task& task);
absl::Status GetElementTraced(Task* task, int64_t deadline_micros,
bool enqueue_result, bool allow_skip,
std::shared_ptr<Result> result,
int64_t thread_index = -1);
absl::Status MaybeRemoveTask(Task& task, int64_t deadline_micros,
Result& result);
absl::Status GetElement(Task* task, int64_t deadline_micros,
bool enqueue_result, bool allow_skip,
std::shared_ptr<Result> result,
int64_t thread_index = -1);
bool ResultReady() const;
std::shared_ptr<Result> PopNextResult();
bool IsCoordinatedRead() const;
std::string DebugString() const;
const DataServiceParams params_;
mutable mutex mu_;
condition_variable get_next_cv_ TF_GUARDED_BY(mu_);
condition_variable worker_thread_cv_ TF_GUARDED_BY(mu_);
condition_variable manager_thread_cv_ TF_GUARDED_BY(mu_);
bool cancelled_ TF_GUARDED_BY(mu_) = false;
// Number of outstanding requests.
int64_t outstanding_requests_ TF_GUARDED_BY(mu_) = 0;
// max_outstanding_requests controls how many elements may be held in memory
// at the same time. This count includes both in-progress requests for
// elements as well as completed requests which haven't yet been produced.
int64_t max_outstanding_requests_ TF_GUARDED_BY(mu_);
// The number of threads in `worker_threads_` which are still running.
int64_t num_running_worker_threads_ TF_GUARDED_BY(mu_) = 0;
// The index of the next task in `tasks_` to read from.
int64_t next_task_index_ TF_GUARDED_BY(mu_) = 0;
// The number tasks in the `tasks_` list that have reached end_of_sequence.
int64_t finished_tasks_ TF_GUARDED_BY(mu_) = 0;
// List of tasks to read from.
std::vector<std::shared_ptr<Task>> tasks_ TF_GUARDED_BY(mu_);
// The current round robin round we are engaged in. A round involves reading
// from each task once.
int64_t current_round_ TF_GUARDED_BY(mu_) = 0;
// Maximum round robin round to read up to before blocking, not inclusive.
// INVARIANT: current_round_ <= round_robin_round_limit_.
// If current_round_ == round_robin_round_limit_,
// next_task_index_ must be 0.
std::optional<int64_t> round_robin_round_limit_ TF_GUARDED_BY(mu_);
// A status to be returned from the next call to `GetNext`. This is set by
// asynchronous threads when they encounter errors.
absl::Status status_ TF_GUARDED_BY(mu_) = absl::OkStatus();
// A queue of results for `GetElement` requests to read from. When doing
// strict round robin reads, the queue will contain placeholder results with
// their `Result::ready` field false until their data has been retrieved
// from a worker. When not doing round-robin reads, results are only added
// to the queue after they are ready, to avoid head-of-line blocking.
std::queue<std::shared_ptr<Result>> results_ TF_GUARDED_BY(mu_);
bool initialized_ = false;
std::unique_ptr<DataServiceContext> ctx_ TF_GUARDED_BY(mu_);
// Set once in Initialize().
int64_t job_id_;
int64_t iteration_client_id_;
std::unique_ptr<DataServiceDispatcherClient> dispatcher_;
const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info_;
Allocator* allocator_;
int64_t get_next_index_ TF_GUARDED_BY(mu_) = 0;
bool iteration_finished_ TF_GUARDED_BY(mu_) = false;
bool should_finish_iteration_ TF_GUARDED_BY(mu_) = true;
// The set of worker UIDs that we have already recorded metrics for.
absl::flat_hash_set<int64_t> worker_uids_ TF_GUARDED_BY(mu_);
std::vector<std::unique_ptr<Thread>> worker_threads_ TF_GUARDED_BY(mu_);
std::unique_ptr<Thread> task_thread_manager_ TF_GUARDED_BY(mu_);
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_CLIENT_DATA_SERVICE_CLIENT_H_
tensorflow/core/data/service/client/data_service_client_test.cc
+241
View File
@@ -0,0 +1,241 @@
/* Copyright 2022 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/core/data/service/client/data_service_client.h"
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include "absl/memory/memory.h"
#include "absl/time/time.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/service/client/common.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/test_cluster.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::data::testing::RangeDataset;
using ::tensorflow::testing::IsOkAndHolds;
using ::tensorflow::testing::StatusIs;
using ::testing::_;
using ::testing::AtLeast;
using ::testing::ElementsAreArray;
using ::testing::HasSubstr;
using ::testing::UnorderedElementsAreArray;
DataServiceParams GetDataServiceParams(
const std::string& dataset_id, const std::string& data_service_address,
const ProcessingModeDef::ShardingPolicy sharding_policy) {
DataServiceParams params;
params.dataset_id = dataset_id;
params.processing_mode.set_sharding_policy(sharding_policy);
params.address = data_service_address;
params.protocol = "grpc";
params.data_transfer_protocol = "grpc";
params.job_name = "test_job";
params.repetition = 0;
params.max_outstanding_requests = 100;
params.task_refresh_interval = absl::Milliseconds(100);
return params;
}
std::vector<int64_t> Range(const int64_t range) {
std::vector<int64_t> result;
for (int64_t i = 0; i < range; ++i) {
result.push_back(i);
}
return result;
}
class TestDataServiceContext : public DataServiceContext {
public:
TestDataServiceContext() = default;
~TestDataServiceContext() override = default;
std::unique_ptr<Thread> StartThread(const std::string& name,
std::function<void()> fn) override {
return absl::WrapUnique(
Env::Default()->StartThread({}, name, std::move(fn)));
}
MOCK_METHOD(void, RecordBufferEnqueue, (const std::vector<Tensor>& element),
(override));
MOCK_METHOD(void, RecordBufferDequeue, (const std::vector<Tensor>& element),
(override));
double GetTargetProcessingTimeNsec() const override { return 1.0e6; }
int64_t UpdateMaxOutstandingRequests(int64_t max_outstanding_requests,
int64_t new_size) override {
return new_size;
}
};
std::unique_ptr<TestDataServiceContext> GetTestDataServiceContext() {
return std::make_unique<TestDataServiceContext>();
}
template <class T>
StatusOr<std::vector<T>> GetResults(DataServiceClient& client) {
std::vector<T> results;
while (true) {
TF_ASSIGN_OR_RETURN(GetNextResult next,
client.GetNext(GetTestDataServiceContext));
if (next.end_of_sequence) {
return results;
}
results.push_back(next.tensors[0].unaligned_flat<T>().data()[0]);
}
return results;
}
template <class T>
StatusOr<T> GetNext(DataServiceClient& client) {
TF_ASSIGN_OR_RETURN(GetNextResult next,
client.GetNext(GetTestDataServiceContext));
if (next.end_of_sequence) {
return errors::OutOfRange(
"The tf.data service has reached the end of sequence");
}
return next.tensors[0].unaligned_flat<T>().data()[0];
}
TEST(DataServiceClientTest, NoSharding) {
TestCluster test_cluster(/*num_workers=*/1);
TF_ASSERT_OK(test_cluster.Initialize());
DatasetClient<int64_t> test_dataset(test_cluster);
TF_ASSERT_OK_AND_ASSIGN(std::string dataset_id,
test_dataset.RegisterDataset(RangeDataset(10)));
DataServiceParams params = GetDataServiceParams(
dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::OFF);
DataServiceClient client(params);
TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
/*allocator=*/nullptr));
EXPECT_THAT(GetResults<int64_t>(client),
absl_testing::IsOkAndHolds(ElementsAreArray(Range(10))));
client.Cancel();
}
TEST(DataServiceClientTest, DynamicSharding) {
TestCluster test_cluster(/*num_workers=*/3);
TF_ASSERT_OK(test_cluster.Initialize());
DatasetClient<int64_t> test_dataset(test_cluster);
TF_ASSERT_OK_AND_ASSIGN(std::string dataset_id,
test_dataset.RegisterDataset(RangeDataset(10)));
DataServiceParams params = GetDataServiceParams(
dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::DYNAMIC);
DataServiceClient client(params);
TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
/*allocator=*/nullptr));
EXPECT_THAT(GetResults<int64_t>(client),
absl_testing::IsOkAndHolds(UnorderedElementsAreArray(Range(10))));
client.Cancel();
}
TEST(DataServiceClientTest, StaticSharding) {
TestCluster test_cluster(/*num_workers=*/3);
TF_ASSERT_OK(test_cluster.Initialize());
DatasetClient<int64_t> dataset_client(test_cluster);
TF_ASSERT_OK_AND_ASSIGN(std::string dataset_id,
dataset_client.RegisterDataset(RangeDataset(10)));
DataServiceParams params =
GetDataServiceParams(dataset_id, test_cluster.DispatcherAddress(),
ProcessingModeDef::FILE_OR_DATA);
DataServiceClient client(params);
TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
/*allocator=*/nullptr));
EXPECT_THAT(GetResults<int64_t>(client),
absl_testing::IsOkAndHolds(UnorderedElementsAreArray(Range(10))));
client.Cancel();
}
TEST(DataServiceClientTest, RecordBufferEvents) {
TestCluster test_cluster(/*num_workers=*/1);
TF_ASSERT_OK(test_cluster.Initialize());
DatasetClient<int64_t> test_dataset(test_cluster);
TF_ASSERT_OK_AND_ASSIGN(std::string dataset_id,
test_dataset.RegisterDataset(RangeDataset(10)));
DataServiceParams params = GetDataServiceParams(
dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::OFF);
DataServiceClient client(params);
TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
/*allocator=*/nullptr));
auto mock_context = std::make_unique<TestDataServiceContext>();
TestDataServiceContext* ctx = mock_context.get();
EXPECT_CALL(*ctx, RecordBufferEnqueue(_)).Times(AtLeast(1));
EXPECT_CALL(*ctx, RecordBufferDequeue(_)).Times(AtLeast(1));
TF_ASSERT_OK_AND_ASSIGN(GetNextResult next, client.GetNext([&mock_context]() {
return std::move(mock_context);
}));
client.Cancel();
}
TEST(DataServiceClientTest, Cancel) {
TestCluster test_cluster(/*num_workers=*/1);
TF_ASSERT_OK(test_cluster.Initialize());
DatasetClient<int64_t> dataset_client(test_cluster);
TF_ASSERT_OK_AND_ASSIGN(std::string dataset_id,
dataset_client.RegisterDataset(RangeDataset(10)));
DataServiceParams params = GetDataServiceParams(
dataset_id, test_cluster.DispatcherAddress(), ProcessingModeDef::OFF);
DataServiceClient client(params);
TF_ASSERT_OK(client.Initialize(/*accelerator_device_info=*/nullptr,
/*allocator=*/nullptr));
client.Cancel();
EXPECT_THAT(client.GetNext(GetTestDataServiceContext),
absl_testing::StatusIs(error::CANCELLED));
}
TEST(DataServiceClientTest, ValidationError) {
DataServiceParams params = GetDataServiceParams(
"dataset_id", "tf_data_service_address", ProcessingModeDef::OFF);
params.target_workers = TARGET_WORKERS_LOCAL;
DataServiceClient client(params);
EXPECT_THAT(
client.Initialize(/*accelerator_device_info=*/nullptr,
/*allocator=*/nullptr),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr(
"Local reads require local tf.data workers, but no local worker "
"is found.")));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/client/utils.cc
+142
View File
@@ -0,0 +1,142 @@
/* Copyright 2022 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/core/data/service/client/utils.h"
#include <cstdint>
#include <string>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/substitute.h"
#include "absl/time/time.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/grpc_util.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/env_time.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace data {
namespace {
// Same timeout used by the RegisterDatasetOp.
constexpr absl::Duration kGetMetadataRetryTimeout = absl::Hours(1);
} // namespace
absl::StatusOr<DataServiceMetadata> GetDataServiceMetadata(
const std::string& dataset_id, const std::string& address,
const std::string& protocol) {
DataServiceDispatcherClient client(address, protocol);
DataServiceMetadata metadata;
absl::Time deadline =
absl::FromUnixMicros(EnvTime::NowMicros()) + kGetMetadataRetryTimeout;
absl::Status status = grpc_util::Retry(
[&]() { return client.GetDataServiceMetadata(dataset_id, metadata); },
absl::Substitute("Get data service metadata for dataset $0, "
"with dispatcher at $1.",
dataset_id, address),
absl::ToUnixMicros(deadline));
if (absl::IsNotFound(status)) {
return absl::NotFoundError(absl::StrCat(
"Dataset id ", dataset_id,
" not found. It must be registered with `register_dataset` before "
"calling `from_dataset_id`."));
}
TF_RETURN_IF_ERROR(status);
return metadata;
}
absl::StatusOr<bool> CompressionDisabledAtRuntime(
const std::string& dataset_id, const std::string& address,
const std::string& protocol, bool disable_compression_at_runtime) {
DataServiceDispatcherClient client(address, protocol);
DisableCompressionAtRuntimeResponse response;
absl::Time deadline =
absl::FromUnixMicros(EnvTime::NowMicros()) + kGetMetadataRetryTimeout;
TF_RETURN_IF_ERROR(grpc_util::Retry(
[&]() {
return client.DisableCompressionAtRuntime(
dataset_id, disable_compression_at_runtime, response);
},
absl::Substitute(
"Get compression disabled at runtime with dispatcher at $0.",
address),
absl::ToUnixMicros(deadline)));
return response.compression_disabled_at_runtime();
}
absl::StatusOr<DataServiceConfig> GetDataServiceConfig(
const std::string& address, const std::string& protocol) {
DataServiceDispatcherClient client(address, protocol);
DataServiceConfig config;
absl::Time deadline =
absl::FromUnixMicros(EnvTime::NowMicros()) + kGetMetadataRetryTimeout;
TF_RETURN_IF_ERROR(grpc_util::Retry(
[&]() { return client.GetDataServiceConfig(config); },
absl::Substitute("Get data service config with dispatcher at $0.",
address),
absl::ToUnixMicros(deadline)));
return config;
}
absl::StatusOr<DataServiceMetadata::Compression> GetValidatedCompression(
const std::string& dataset_id, const DataServiceMetadata& metadata) {
if (metadata.compression() == DataServiceMetadata::COMPRESSION_UNSPECIFIED) {
return absl::InternalError(absl::Substitute(
"Got invalid compression $0 for dataset $1. A proper compression "
"should be registered in `register_dataset`.",
DataServiceMetadata::Compression_Name(metadata.compression()),
dataset_id));
}
return metadata.compression();
}
int64_t EstimateCardinality(const ProcessingModeDef& processing_mode,
const DataServiceMetadata& metadata,
bool is_coordinated_read) {
if (is_coordinated_read) {
// Coordinated reads require the dataset to be infinite.
return kInfiniteCardinality;
}
if (metadata.cardinality() == 0) {
return 0;
}
if (metadata.cardinality() == kInfiniteCardinality) {
// Sharding may cause an infinite dataset to be empty. For example, in
// `range(10).batch(10, drop_remainder=True).repeat()`, inserting `shard`
// before `batch` will cause the dataset to be empty.
// This case is rare, and there is significant performance hit for dynamic
// sharding if it reports unknown cardinality, so it is reasonable to
// report infinite cardinality. For DATA sharding, it is ok to report
// infinite cardinality since it inserts `shard` after `repeat`.
if (processing_mode.sharding_policy() == ProcessingModeDef::OFF ||
processing_mode.sharding_policy() == ProcessingModeDef::DYNAMIC ||
processing_mode.sharding_policy() == ProcessingModeDef::DATA) {
return kInfiniteCardinality;
}
}
return kUnknownCardinality;
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/client/utils.h
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2022 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_CORE_DATA_SERVICE_CLIENT_UTILS_H_
#define TENSORFLOW_CORE_DATA_SERVICE_CLIENT_UTILS_H_
#include <cstdint>
#include <optional>
#include <string>
#include "absl/status/statusor.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
// Gets the `DataServiceMetadata` for `dataset_id`.
absl::StatusOr<DataServiceMetadata> GetDataServiceMetadata(
const std::string& dataset_id, const std::string& address,
const std::string& protocol);
// Gets the `DisableCompressAtRuntimeResponse.compression_disabled_at_runtime`
// for the given dataset.
absl::StatusOr<bool> CompressionDisabledAtRuntime(
const std::string& dataset_id, const std::string& address,
const std::string& protocol, bool disable_compression_at_runtime);
// Gets the `DataServiceConfig` for the data service running at `address`.
absl::StatusOr<DataServiceConfig> GetDataServiceConfig(
const std::string& address, const std::string& protocol);
// Gets the compression from `metadata`. If `metadata` specifies no valid
// compression, returns an internal error.
absl::StatusOr<DataServiceMetadata::Compression> GetValidatedCompression(
const std::string& dataset_id, const DataServiceMetadata& metadata);
// Estimates the cardinality of a data service dataset.
int64_t EstimateCardinality(const ProcessingModeDef& processing_mode,
const DataServiceMetadata& metadata,
bool is_coordinated_read);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_CLIENT_UTILS_H_
tensorflow/core/data/service/client/utils_test.cc
+134
View File
@@ -0,0 +1,134 @@
/* Copyright 2022 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/core/data/service/client/utils.h"
#include <optional>
#include <string>
#include <gmock/gmock.h>
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/test_cluster.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/status_matchers.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::data::testing::EqualsProto;
using ::tensorflow::data::testing::RangeDataset;
using ::tsl::testing::IsOkAndHolds;
using ::tsl::testing::StatusIs;
TEST(UtilsTest, GetDataServiceMetadata) {
TestCluster test_cluster(/*num_workers=*/1);
TF_ASSERT_OK(test_cluster.Initialize());
DataServiceDispatcherClient dispatcher_client(
test_cluster.DispatcherAddress(), "grpc");
DataServiceMetadata metadata;
metadata.set_compression(DataServiceMetadata::COMPRESSION_SNAPPY);
metadata.set_cardinality(100);
std::string dataset_id;
TF_ASSERT_OK(dispatcher_client.RegisterDataset(
RangeDataset(10), metadata, /*requested_dataset_id=*/std::nullopt,
dataset_id));
EXPECT_THAT(GetDataServiceMetadata(dataset_id,
test_cluster.DispatcherAddress(), "grpc"),
absl_testing::IsOkAndHolds(EqualsProto(metadata)));
}
TEST(UtilsTest, GetDataServiceMetadataNotFound) {
TestCluster test_cluster(/*num_workers=*/1);
TF_ASSERT_OK(test_cluster.Initialize());
EXPECT_THAT(GetDataServiceMetadata(/*dataset_id=*/"not found",
test_cluster.DispatcherAddress(), "grpc"),
absl_testing::StatusIs(error::NOT_FOUND));
}
TEST(UtilsTest, GetDataServiceConfig) {
TestCluster test_cluster(/*num_workers=*/1);
TF_ASSERT_OK(test_cluster.Initialize());
TF_ASSERT_OK_AND_ASSIGN(
DataServiceConfig service_config,
GetDataServiceConfig(test_cluster.DispatcherAddress(), "grpc"));
EXPECT_EQ(service_config.deployment_mode(), DEPLOYMENT_MODE_COLOCATED);
}
TEST(UtilsTest, GetValidatedCompression) {
DataServiceMetadata metadata;
metadata.set_compression(DataServiceMetadata::COMPRESSION_SNAPPY);
EXPECT_THAT(
GetValidatedCompression("dataset_id", metadata),
absl_testing::IsOkAndHolds(DataServiceMetadata::COMPRESSION_SNAPPY));
}
TEST(UtilsTest, InvalidCompression) {
DataServiceMetadata metadata;
EXPECT_THAT(GetValidatedCompression("dataset_id", metadata),
absl_testing::StatusIs(error::INTERNAL));
}
TEST(UtilsTest, EstimateCardinalityEmptyDataset) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
DataServiceMetadata metadata;
metadata.set_cardinality(0);
EXPECT_EQ(EstimateCardinality(processing_mode, metadata,
/*is_coordinated_read=*/false),
0);
}
TEST(UtilsTest, EstimateCardinalityInfiniteDataset) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
DataServiceMetadata metadata;
metadata.set_cardinality(kInfiniteCardinality);
EXPECT_EQ(EstimateCardinality(processing_mode, metadata,
/*is_coordinated_read=*/false),
kInfiniteCardinality);
processing_mode.set_sharding_policy(ProcessingModeDef::DYNAMIC);
EXPECT_EQ(EstimateCardinality(processing_mode, metadata,
/*is_coordinated_read=*/false),
kInfiniteCardinality);
}
TEST(UtilsTest, EstimateCardinalityCoordinatedRead) {
ProcessingModeDef processing_mode;
DataServiceMetadata metadata;
EXPECT_EQ(EstimateCardinality(processing_mode, metadata,
/*is_coordinated_read=*/true),
kInfiniteCardinality);
}
TEST(UtilsTest, EstimateCardinalityUnknownCardinality) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
DataServiceMetadata metadata;
metadata.set_cardinality(10);
EXPECT_EQ(EstimateCardinality(processing_mode, metadata,
/*is_coordinated_read=*/false),
kUnknownCardinality);
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/client/validate_utils.cc
+103
View File
@@ -0,0 +1,103 @@
/* Copyright 2022 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/core/data/service/client/validate_utils.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/data/service/client/common.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/worker_impl.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
namespace {
// Validates local worker related parameters.
absl::Status ValidateLocalWorkers(
const DataServiceParams& data_service_params) {
if (data_service_params.target_workers != TARGET_WORKERS_LOCAL) {
return absl::OkStatus();
}
if (LocalWorkers::Empty()) {
if (IsStaticShard(data_service_params.processing_mode)) {
return absl::InvalidArgumentError(absl::StrCat(
"Static sharding policy <",
ProcessingModeDef::ShardingPolicy_Name(
data_service_params.processing_mode.sharding_policy()),
"> requires local tf.data workers, but no local worker is found. "
"You need to run local tf.data service workers in your training "
"workers. Static sharding also requires a fixed worker pool and "
"a list of worker addresses in the DispatcherConfig. See the "
"\"Processing Modes\" section in the module doc for details."));
}
return absl::InvalidArgumentError(
"Local reads require local tf.data workers, but no local worker "
"is found. You need to run local tf.data service workers in your "
"training workers.");
}
if (data_service_params.num_consumers.has_value()) {
return absl::InvalidArgumentError(
"Coordinated reads require non-local workers, but `target_workers` "
"is \"LOCAL\".");
}
return absl::OkStatus();
}
// Validates cross-trainer cache related parameters.
absl::Status ValidateCrossTrainerCache(
const DataServiceParams& data_service_params) {
if (!data_service_params.cross_trainer_cache_options.has_value()) {
return absl::OkStatus();
}
if (data_service_params.job_name.empty()) {
return absl::InvalidArgumentError(
"Cross-trainer caching requires named jobs. Got empty `job_name`.");
}
if (data_service_params.metadata.cardinality() >= 0) {
return absl::InvalidArgumentError(absl::StrCat(
"Cross-trainer caching requires the input dataset to be infinite. "
"Got input with cardinality ",
data_service_params.metadata.cardinality()));
}
if (data_service_params.repetition > 1) {
return absl::InvalidArgumentError(absl::StrCat(
"Cross-trainer caching requires infinite datasets and disallows "
"multiple repetitions of the same dataset. Got repetition ",
data_service_params.repetition));
}
if (data_service_params.num_consumers.has_value()) {
return absl::InvalidArgumentError(absl::StrCat(
"Cross-trainer caching does not support coordinated reads. "
"Got number of coordinated consumers: ",
data_service_params.num_consumers.value()));
}
return absl::OkStatus();
}
} // namespace
absl::Status ValidateDataServiceParams(
const DataServiceParams& data_service_params) {
TF_RETURN_IF_ERROR(ValidateLocalWorkers(data_service_params));
TF_RETURN_IF_ERROR(ValidateCrossTrainerCache(data_service_params));
return absl::OkStatus();
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/client/validate_utils.h
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2022 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_CORE_DATA_SERVICE_CLIENT_VALIDATE_UTILS_H_
#define TENSORFLOW_CORE_DATA_SERVICE_CLIENT_VALIDATE_UTILS_H_
#include "absl/status/status.h"
#include "tensorflow/core/data/service/client/common.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace data {
// Validates data service dataset parameters.
absl::Status ValidateDataServiceParams(
const DataServiceParams& data_service_params);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_CLIENT_VALIDATE_UTILS_H_
tensorflow/core/data/service/client/validate_utils_test.cc
+181
View File
@@ -0,0 +1,181 @@
/* Copyright 2022 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/core/data/service/client/validate_utils.h"
#include <memory>
#include <gmock/gmock.h>
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/service/client/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/worker_impl.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::testing::StatusIs;
using ::testing::HasSubstr;
DataServiceParams GetDefaultParams() {
DataServiceParams params;
params.dataset_id = "dataset_id";
params.processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
params.address = "localhost";
params.protocol = "grpc";
params.data_transfer_protocol = "grpc";
params.metadata.set_cardinality(kUnknownCardinality);
return params;
}
std::shared_ptr<DataServiceWorkerImpl> GetLocalWorker() {
experimental::WorkerConfig config;
config.set_protocol("grpc");
config.set_dispatcher_address("localhost");
config.set_worker_address("localhost");
return std::make_shared<DataServiceWorkerImpl>(config);
}
TEST(ValidateUtilsTest, DefaultParams) {
TF_EXPECT_OK(ValidateDataServiceParams(GetDefaultParams()));
}
TEST(ValidateUtilsTest, LocalWorkerSuccess) {
DataServiceParams params = GetDefaultParams();
LocalWorkers::Add("localhost", GetLocalWorker());
params.target_workers = TARGET_WORKERS_LOCAL;
TF_EXPECT_OK(ValidateDataServiceParams(params));
LocalWorkers::Remove("localhost");
}
TEST(ValidateUtilsTest, NoLocalWorker) {
DataServiceParams params = GetDefaultParams();
params.target_workers = TARGET_WORKERS_LOCAL;
EXPECT_THAT(
ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr(
"Local reads require local tf.data workers, but no local worker "
"is found.")));
}
TEST(ValidateUtilsTest, NoLocalWorkerStaticSharding) {
DataServiceParams params = GetDefaultParams();
params.processing_mode.set_sharding_policy(ProcessingModeDef::FILE_OR_DATA);
params.target_workers = TARGET_WORKERS_LOCAL;
EXPECT_THAT(
ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr(
"Static sharding policy <FILE_OR_DATA> requires local tf.data "
"workers, but no local worker is found.")));
}
TEST(ValidateUtilsTest, LocalReadDisallowsCoordinatedRead) {
DataServiceParams params = GetDefaultParams();
LocalWorkers::Add("localhost", GetLocalWorker());
params.num_consumers = 1;
params.consumer_index = 0;
params.target_workers = TARGET_WORKERS_LOCAL;
EXPECT_THAT(ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr("Coordinated reads require non-local workers, but "
"`target_workers` is \"LOCAL\".")));
LocalWorkers::Remove("localhost");
}
TEST(ValidateUtilsTest, CrossTrainerCacheSuccess) {
DataServiceParams params = GetDefaultParams();
params.job_name = "job_name";
params.repetition = 1;
params.metadata.set_cardinality(kInfiniteCardinality);
params.cross_trainer_cache_options.emplace();
params.cross_trainer_cache_options->set_trainer_id("trainer ID");
TF_EXPECT_OK(ValidateDataServiceParams(params));
}
TEST(ValidateUtilsTest, CrossTrainerCacheRequiresJobName) {
DataServiceParams params = GetDefaultParams();
params.repetition = 1;
params.metadata.set_cardinality(kInfiniteCardinality);
params.cross_trainer_cache_options.emplace();
params.cross_trainer_cache_options->set_trainer_id("trainer ID");
EXPECT_THAT(
ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
"Cross-trainer caching requires named jobs. Got empty `job_name`."));
}
TEST(ValidateUtilsTest, CrossTrainerCacheRequiresInfiniteDataset) {
DataServiceParams params = GetDefaultParams();
params.job_name = "job_name";
params.repetition = 1;
params.metadata.set_cardinality(10);
params.cross_trainer_cache_options.emplace();
params.cross_trainer_cache_options->set_trainer_id("trainer ID");
EXPECT_THAT(ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr("Cross-trainer caching requires the input "
"dataset to be infinite.")));
}
TEST(ValidateUtilsTest, CrossTrainerCacheDisallowsRepetition) {
DataServiceParams params = GetDefaultParams();
params.job_name = "job_name";
params.repetition = 5;
params.metadata.set_cardinality(kInfiniteCardinality);
params.cross_trainer_cache_options.emplace();
params.cross_trainer_cache_options->set_trainer_id("trainer ID");
EXPECT_THAT(
ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr(
"Cross-trainer caching requires infinite datasets and disallows "
"multiple repetitions of the same dataset.")));
}
TEST(ValidateUtilsTest, CrossTrainerCacheDisallowsCoordinatedRead) {
DataServiceParams params = GetDefaultParams();
params.job_name = "job_name";
params.repetition = 1;
params.num_consumers = 1;
params.consumer_index = 0;
params.metadata.set_cardinality(kInfiniteCardinality);
params.cross_trainer_cache_options.emplace();
params.cross_trainer_cache_options->set_trainer_id("trainer ID");
EXPECT_THAT(
ValidateDataServiceParams(params),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr(
"Cross-trainer caching does not support coordinated reads.")));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/common.cc
+141
View File
@@ -0,0 +1,141 @@
/* Copyright 2021 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/core/data/service/common.h"
#include <string>
#include "absl/strings/string_view.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
namespace {
constexpr const char kAuto[] = "AUTO";
constexpr const char kAny[] = "ANY";
constexpr const char kLocal[] = "LOCAL";
constexpr const char kColocated[] = "COLOCATED";
constexpr const char kRemote[] = "REMOTE";
constexpr const char kHybrid[] = "HYBRID";
} // namespace
bool IsNoShard(const ProcessingModeDef& processing_mode) {
return processing_mode.sharding_policy() == ProcessingModeDef::OFF;
}
bool IsDynamicShard(const ProcessingModeDef& processing_mode) {
return processing_mode.sharding_policy() == ProcessingModeDef::DYNAMIC;
}
bool IsStaticShard(const ProcessingModeDef& processing_mode) {
return processing_mode.sharding_policy() == ProcessingModeDef::FILE ||
processing_mode.sharding_policy() == ProcessingModeDef::DATA ||
processing_mode.sharding_policy() == ProcessingModeDef::FILE_OR_DATA ||
processing_mode.sharding_policy() == ProcessingModeDef::HINT;
}
absl::Status ValidateProcessingMode(const ProcessingModeDef& processing_mode) {
if (!IsNoShard(processing_mode) && !IsDynamicShard(processing_mode) &&
!IsStaticShard(processing_mode)) {
return absl::InternalError(absl::StrCat(
"ProcessingMode ", processing_mode.ShortDebugString(),
" does not "
"specify a valid sharding policy. Please add the policy to either "
"`IsDynamicShard` or `IsStaticShard` (i.e., auto-shard)."));
}
return absl::OkStatus();
}
absl::StatusOr<AutoShardPolicy> ToAutoShardPolicy(
const ProcessingModeDef::ShardingPolicy sharding_policy) {
switch (sharding_policy) {
case ProcessingModeDef::FILE:
return AutoShardPolicy::FILE;
case ProcessingModeDef::DATA:
return AutoShardPolicy::DATA;
case ProcessingModeDef::FILE_OR_DATA:
return AutoShardPolicy::AUTO;
case ProcessingModeDef::HINT:
return AutoShardPolicy::HINT;
case ProcessingModeDef::DYNAMIC:
case ProcessingModeDef::OFF:
return AutoShardPolicy::OFF;
default:
return absl::InternalError(absl::StrCat(
"tf.data service sharding policy ",
ProcessingModeDef::ShardingPolicy_Name(sharding_policy),
" is not convertible to a valid auto-shard policy. If you're "
"defining a new sharding policy, please update the policy mapping."));
}
}
absl::StatusOr<TargetWorkers> ParseTargetWorkers(absl::string_view s) {
std::string str_upper = absl::AsciiStrToUpper(s);
if (str_upper.empty() || str_upper == kAuto) {
return TARGET_WORKERS_AUTO;
}
if (str_upper == kAny) {
return TARGET_WORKERS_ANY;
}
if (str_upper == kLocal) {
return TARGET_WORKERS_LOCAL;
}
return absl::InvalidArgumentError(
absl::StrCat("Unrecognized target workers: ", s));
}
std::string TargetWorkersToString(TargetWorkers target_workers) {
switch (target_workers) {
case TARGET_WORKERS_AUTO:
return kAuto;
case TARGET_WORKERS_ANY:
return kAny;
case TARGET_WORKERS_LOCAL:
return kLocal;
default:
DCHECK(false);
return "UNKNOWN";
}
}
absl::StatusOr<DeploymentMode> ParseDeploymentMode(absl::string_view s) {
std::string str_upper = absl::AsciiStrToUpper(s);
if (str_upper == kColocated) {
return DEPLOYMENT_MODE_COLOCATED;
}
if (str_upper == kRemote) {
return DEPLOYMENT_MODE_REMOTE;
}
if (str_upper == kHybrid) {
return DEPLOYMENT_MODE_HYBRID;
}
return absl::InvalidArgumentError(
absl::StrCat("Invalid tf.data service deployment mode: ", s,
". Supported modes are "
"COLOCATED, REMOTE, and HYBRID."));
}
bool IsPreemptedError(const absl::Status& status) {
return absl::IsAborted(status) || absl::IsCancelled(status) ||
absl::IsUnavailable(status);
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/common.h
+120
View File
@@ -0,0 +1,120 @@
/* Copyright 2021 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_CORE_DATA_SERVICE_COMMON_H_
#define TENSORFLOW_CORE_DATA_SERVICE_COMMON_H_
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
// Increment this when making backwards-incompatible changes to communication
// between tf.data clients and servers.
constexpr int kDataServiceVersion = 9;
// If the user starts a colocated tf.data worker on each TF host, the worker
// will be applied a "COLOCATED" tag. This is used to avoid reading from tf.data
// workers on other TF hosts when the host runs a local tf.data service worker.
constexpr absl::string_view kColocatedWorkerTag = "COLOCATED";
// Container to hold the result of a `GetNext` call.
struct GetNextResult final {
explicit GetNextResult() = default;
GetNextResult(const GetNextResult&) = delete;
GetNextResult& operator=(const GetNextResult&) = delete;
GetNextResult(GetNextResult&&) = default;
GetNextResult& operator=(GetNextResult&&) = delete;
static GetNextResult EndOfSequence() {
GetNextResult result;
result.end_of_sequence = true;
return result;
}
std::vector<Tensor> tensors;
bool end_of_sequence = false;
};
// Returns true if `processing_mode` specifies no sharding policy.
bool IsNoShard(const ProcessingModeDef& processing_mode);
// Returns true if `processing_mode` is dynamic sharding.
bool IsDynamicShard(const ProcessingModeDef& processing_mode);
// Returns true if `processing_mode` is static sharding.
bool IsStaticShard(const ProcessingModeDef& processing_mode);
// Returns an internal error if `processing_mode` is invalid.
absl::Status ValidateProcessingMode(const ProcessingModeDef& processing_mode);
// Converts tf.data service `sharding_policy` to `AutoShardPolicy`. Returns an
// internal error if `sharding_policy` is not supported.
absl::StatusOr<AutoShardPolicy> ToAutoShardPolicy(
ProcessingModeDef::ShardingPolicy sharding_policy);
// Parses a string representing a `TargetWorkers` (case-insensitive).
// Returns InvalidArgument if the string is not recognized.
absl::StatusOr<TargetWorkers> ParseTargetWorkers(absl::string_view s);
// Converts a `TargetWorkers` enum to string.
std::string TargetWorkersToString(TargetWorkers target_workers);
// Parses a string representing a `DeploymentMode` (case-insensitive).
// Returns InvalidArgument if the string is not recognized.
absl::StatusOr<DeploymentMode> ParseDeploymentMode(absl::string_view s);
// Returns true if `status` is a retriable error that indicates preemption.
bool IsPreemptedError(const absl::Status& status);
// Base class for data service clients. Data service clients are
// threadsafe.
class DataServiceClientBase {
public:
DataServiceClientBase(const std::string& address, const std::string& protocol)
: address_(address), protocol_(protocol) {}
virtual ~DataServiceClientBase() = default;
// Not copyable or movable.
DataServiceClientBase(const DataServiceClientBase&) = delete;
DataServiceClientBase& operator=(const DataServiceClientBase&) = delete;
// Initializes the client. Calling `Initialize()` is not required since the
// first RPC will perform any necessary initialization. However, it can be
// useful to call `Initialize()` proactively so that any errors that happen
// during initialization can be surfaced earlier.
virtual absl::Status Initialize() { return EnsureInitialized(); }
protected:
// Initializes the client if it isn't already initialized.
virtual absl::Status EnsureInitialized() = 0;
const std::string address_;
const std::string protocol_;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_COMMON_H_
tensorflow/core/data/service/common.proto
+162
View File
@@ -0,0 +1,162 @@
// Copyright 2026 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.
// ==============================================================================
syntax = "proto3";
package tensorflow.data;
import "xla/tsl/protobuf/status.proto";
import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/protobuf/data_service.proto";
import "tensorflow/core/protobuf/snapshot.proto";
// Next tag: 2
message DatasetDef {
// We represent datasets as tensorflow GraphDefs which define the operations
// needed to create a tf.data dataset.
GraphDef graph = 1;
}
// Next tag: 3
message IterationKeyDef {
string name = 1;
int64 iteration = 2;
}
// Next tag: 14
message TaskDef {
reserved 6;
// The dataset to iterate over.
oneof dataset {
DatasetDef dataset_def = 1;
string path = 2;
}
string dataset_id = 3;
int64 task_id = 4;
int64 iteration_id = 5;
// In distributed epoch processing mode, we use one split provider for each
// source that feeds into the dataset. In parallel_epochs mode,
// `num_split_providers` is always zero.
int64 num_split_providers = 9;
// Address of the worker that the task is assigned to.
string worker_address = 8;
ProcessingModeDef processing_mode_def = 10;
// Optional number of consumers. If set, the results of the task will be
// provided to consumers round-robin.
oneof optional_num_consumers {
int64 num_consumers = 7;
}
// Number of workers and the worker index. These are only populated when the
// `processing_mode_def` specifies a static sharding policy.
int64 num_workers = 11;
int64 worker_index = 12;
// True if cross-trainer cache is enabled.
bool use_cross_trainer_cache = 13;
}
// Next tag: 9
message TaskInfo {
// The address of the worker processing the task.
string worker_address = 1;
// The data transfer servers of the worker processing the task.
repeated DataTransferServerInfo transfer_servers = 8;
// Tags attached to the worker. This allows reading from selected workers.
// For example, by applying a "COLOCATED" tag, tf.data service is able to read
// from the local tf.data worker if one exists, then from off-TF-host workers,
// to avoid cross-TF-host reads.
repeated string worker_tags = 6;
// The task id.
int64 task_id = 2;
// The id of the iteration that the task is part of.
int64 iteration_id = 3;
// The UID of the worker Borg job, used for telemetry.
int64 worker_uid = 7;
// The round to start reading from the task in. For non-round-robin reads,
// this is always 0.
int64 starting_round = 5;
reserved 4;
}
// Next tag: 5
message SnapshotTaskDef {
// The base directory at which the snapshot is being materialized.
string base_path = 1;
// The index of the stream that the worker has been assigned to process.
int64 stream_index = 2;
// The number of source datasets (split providers).
int64 num_sources = 3;
// Snapshot metadata including the element spec and compression method.
experimental.DistributedSnapshotMetadata metadata = 4;
}
// Next tag: 5
message SnapshotTaskProgress {
SnapshotTaskDef snapshot_task = 1;
// True if the snapshot is complete successfully. Unset if the snapshot is not
// complete or an error has occurred.
bool completed = 2;
// If any error occurs during the snapshot processing, `status` will be filled
// with the error status.
StatusProto status = 3;
}
message SnapshotStreamInfo {
// The index of the stream being processed or having been processed.
int64 index = 1;
enum State {
// Unspecified. Invalid state.
UNSPECIFIED = 0;
// The dispatcher thinks the stream has a live worker.
ASSIGNED = 1;
// The dispatcher doesn't think the stream has a live worker.
ORPHAN = 2;
// The dispatcher doesn't know whether or not the stream has a live worker.
UNKNOWN = 3;
// The dispatcher thinks the stream has no more splits left to be processed.
DONE = 4;
}
State state = 2;
}
// Specifies which tf.data service workers to read from.
enum TargetWorkers {
TARGET_WORKERS_UNSPECIFIED = 0;
// tf.data service runtime decides which workers to read from.
TARGET_WORKERS_AUTO = 1;
// Reads from any available worker.
TARGET_WORKERS_ANY = 2;
// Only reads from local workers. If no local worker is found, it is an error.
TARGET_WORKERS_LOCAL = 3;
}
// Information about one of a worker server's data transfer servers.
// Next tag: 6
message DataTransferServerInfo {
string protocol = 1;
string address = 2;
// If provided, properties of the server used to determine compatibility with
// a client.
bytes compatibility_info = 3;
// If `true`, data service clients should fall back to gRPC for this server if
// they fail to create a data transfer client for it.
bool fall_back_to_grpc_at_client_creation_time = 4;
// If `true`, data service clients should fall back to gRPC for this server if
// it nonretryably fails to transfer an element.
bool fall_back_to_grpc_at_get_element_time = 5;
}
tensorflow/core/data/service/common_test.cc
+207
View File
@@ -0,0 +1,207 @@
/* Copyright 2021 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/core/data/service/common.h"
#include <vector>
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::testing::IsOkAndHolds;
using ::tensorflow::testing::StatusIs;
using ::testing::HasSubstr;
std::vector<ProcessingModeDef::ShardingPolicy> EnumerateShardingPolicies() {
std::vector<ProcessingModeDef::ShardingPolicy> result;
const ::tensorflow::protobuf::EnumDescriptor* enum_descriptor =
::tensorflow::protobuf::GetEnumDescriptor<
ProcessingModeDef::ShardingPolicy>();
for (int i = 0; i < enum_descriptor->value_count(); ++i) {
result.push_back(static_cast<ProcessingModeDef::ShardingPolicy>(
enum_descriptor->value(i)->number()));
}
return result;
}
TEST(CommonTest, NoShard) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
EXPECT_TRUE(IsNoShard(processing_mode));
EXPECT_FALSE(IsDynamicShard(processing_mode));
EXPECT_FALSE(IsStaticShard(processing_mode));
}
TEST(CommonTest, DynamicShard) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::DYNAMIC);
EXPECT_FALSE(IsNoShard(processing_mode));
EXPECT_TRUE(IsDynamicShard(processing_mode));
EXPECT_FALSE(IsStaticShard(processing_mode));
}
TEST(CommonTest, StaticShard) {
ProcessingModeDef processing_mode;
std::vector<ProcessingModeDef::ShardingPolicy> policies = {
ProcessingModeDef::FILE, ProcessingModeDef::DATA,
ProcessingModeDef::FILE_OR_DATA, ProcessingModeDef::HINT};
for (const ProcessingModeDef::ShardingPolicy policy : policies) {
processing_mode.set_sharding_policy(policy);
EXPECT_FALSE(IsNoShard(processing_mode));
EXPECT_FALSE(IsDynamicShard(processing_mode));
EXPECT_TRUE(IsStaticShard(processing_mode));
}
}
TEST(CommonTest, DefaultShardingPolicyIsNoShard) {
ProcessingModeDef processing_mode;
EXPECT_TRUE(IsNoShard(processing_mode));
EXPECT_FALSE(IsDynamicShard(processing_mode));
EXPECT_FALSE(IsStaticShard(processing_mode));
}
TEST(CommonTest, ToAutoShardPolicy) {
EXPECT_THAT(ToAutoShardPolicy(ProcessingModeDef::FILE_OR_DATA),
absl_testing::IsOkAndHolds(AutoShardPolicy::AUTO));
EXPECT_THAT(ToAutoShardPolicy(ProcessingModeDef::HINT),
absl_testing::IsOkAndHolds(AutoShardPolicy::HINT));
EXPECT_THAT(ToAutoShardPolicy(ProcessingModeDef::OFF),
absl_testing::IsOkAndHolds(AutoShardPolicy::OFF));
EXPECT_THAT(ToAutoShardPolicy(ProcessingModeDef::DYNAMIC),
absl_testing::IsOkAndHolds(AutoShardPolicy::OFF));
}
TEST(CommonTest, ConvertValidShardingPolicyToAutoShardPolicy) {
for (const ProcessingModeDef::ShardingPolicy sharding_policy :
EnumerateShardingPolicies()) {
TF_EXPECT_OK(ToAutoShardPolicy(sharding_policy).status());
}
}
TEST(CommonTest, ConvertInvalidShardingPolicyToAutoShardPolicy) {
const ProcessingModeDef::ShardingPolicy sharding_policy =
static_cast<ProcessingModeDef::ShardingPolicy>(-100);
EXPECT_THAT(
ToAutoShardPolicy(sharding_policy),
absl_testing::StatusIs(error::INTERNAL,
HasSubstr("please update the policy mapping.")));
}
TEST(CommonTest, ValidateProcessingMode) {
for (const ProcessingModeDef::ShardingPolicy policy :
EnumerateShardingPolicies()) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(policy);
TF_EXPECT_OK(ValidateProcessingMode(processing_mode));
}
}
TEST(CommonTest, InvalidProcessingMode) {
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(
static_cast<ProcessingModeDef::ShardingPolicy>(100));
EXPECT_THAT(ValidateProcessingMode(processing_mode),
absl_testing::StatusIs(
error::INTERNAL,
HasSubstr("does not specify a valid sharding policy.")));
}
TEST(CommonTest, ParseTargetWorkers) {
EXPECT_THAT(ParseTargetWorkers("AUTO"),
absl_testing::IsOkAndHolds(TARGET_WORKERS_AUTO));
EXPECT_THAT(ParseTargetWorkers("Auto"),
absl_testing::IsOkAndHolds(TARGET_WORKERS_AUTO));
EXPECT_THAT(ParseTargetWorkers("ANY"),
absl_testing::IsOkAndHolds(TARGET_WORKERS_ANY));
EXPECT_THAT(ParseTargetWorkers("any"),
absl_testing::IsOkAndHolds(TARGET_WORKERS_ANY));
EXPECT_THAT(ParseTargetWorkers("LOCAL"),
absl_testing::IsOkAndHolds(TARGET_WORKERS_LOCAL));
EXPECT_THAT(ParseTargetWorkers("local"),
absl_testing::IsOkAndHolds(TARGET_WORKERS_LOCAL));
EXPECT_THAT(ParseTargetWorkers(""),
absl_testing::IsOkAndHolds(TARGET_WORKERS_AUTO));
}
TEST(CommonTest, ParseInvalidTargetWorkers) {
EXPECT_THAT(ParseTargetWorkers("TARGET_WORKERS_UNSPECIFIED"),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
EXPECT_THAT(ParseTargetWorkers("UNSET"),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST(CommonTest, TargetWorkersToString) {
EXPECT_EQ(TargetWorkersToString(TARGET_WORKERS_AUTO), "AUTO");
EXPECT_EQ(TargetWorkersToString(TARGET_WORKERS_ANY), "ANY");
EXPECT_EQ(TargetWorkersToString(TARGET_WORKERS_LOCAL), "LOCAL");
}
TEST(CommonTest, ParseDeploymentMode) {
EXPECT_THAT(
ParseDeploymentMode("COLOCATED"),
absl_testing::IsOkAndHolds(DeploymentMode::DEPLOYMENT_MODE_COLOCATED));
EXPECT_THAT(
ParseDeploymentMode("Colocated"),
absl_testing::IsOkAndHolds(DeploymentMode::DEPLOYMENT_MODE_COLOCATED));
EXPECT_THAT(
ParseDeploymentMode("REMOTE"),
absl_testing::IsOkAndHolds(DeploymentMode::DEPLOYMENT_MODE_REMOTE));
EXPECT_THAT(
ParseDeploymentMode("remote"),
absl_testing::IsOkAndHolds(DeploymentMode::DEPLOYMENT_MODE_REMOTE));
EXPECT_THAT(
ParseDeploymentMode("HYBRID"),
absl_testing::IsOkAndHolds(DeploymentMode::DEPLOYMENT_MODE_HYBRID));
EXPECT_THAT(
ParseDeploymentMode("hybrid"),
absl_testing::IsOkAndHolds(DeploymentMode::DEPLOYMENT_MODE_HYBRID));
}
TEST(CommonTest, ParseInvalidDeploymentMode) {
EXPECT_THAT(ParseDeploymentMode("DEPLOYMENT_MODE_UNSPECIFIED"),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
TEST(CommonTest, IsPreemptedError) {
EXPECT_TRUE(IsPreemptedError(absl::AbortedError("Aborted")));
EXPECT_TRUE(IsPreemptedError(absl::CancelledError("Cancelled")));
EXPECT_TRUE(IsPreemptedError(absl::UnavailableError("Unavailable")));
EXPECT_FALSE(IsPreemptedError(absl::OkStatus()));
}
TEST(CommonTest, IsPermanentError) {
EXPECT_FALSE(
IsPreemptedError(absl::FailedPreconditionError("Failed precondition")));
EXPECT_FALSE(IsPreemptedError(absl::InternalError("Internal")));
EXPECT_FALSE(
IsPreemptedError(absl::InvalidArgumentError("Invalid argument")));
EXPECT_FALSE(IsPreemptedError(absl::NotFoundError("Not found")));
EXPECT_FALSE(IsPreemptedError(absl::OutOfRangeError("Out of range")));
EXPECT_FALSE(IsPreemptedError(absl::UnknownError("Unknown")));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/credentials_factory.cc
+123
View File
@@ -0,0 +1,123 @@
/* Copyright 2020 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/core/data/service/credentials_factory.h"
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/mutex.h"
namespace tensorflow {
namespace data {
namespace {
mutex* get_lock() {
static mutex lock(LINKER_INITIALIZED);
return &lock;
}
using CredentialsFactories =
std::unordered_map<std::string, CredentialsFactory*>;
CredentialsFactories& credentials_factories() {
static auto& factories = *new CredentialsFactories();
return factories;
}
} // namespace
void CredentialsFactory::Register(CredentialsFactory* factory) {
mutex_lock l(*get_lock());
if (!credentials_factories().insert({factory->Protocol(), factory}).second) {
LOG(ERROR)
<< "Two credentials factories are being registered with protocol "
<< factory->Protocol() << ". Which one gets used is undefined.";
}
}
absl::Status CredentialsFactory::Get(absl::string_view protocol,
CredentialsFactory** out) {
mutex_lock l(*get_lock());
auto it = credentials_factories().find(std::string(protocol));
if (it != credentials_factories().end()) {
*out = it->second;
return absl::OkStatus();
}
std::vector<std::string> available_types;
for (const auto& factory : credentials_factories()) {
available_types.push_back(factory.first);
}
return absl::NotFoundError(
absl::StrCat("No credentials factory has been registered for ",
"protocol ", protocol, ". The available types are: [ ",
absl::StrJoin(available_types, ", "), " ]"));
}
absl::Status CredentialsFactory::CreateServerCredentials(
absl::string_view protocol,
std::shared_ptr<::grpc::ServerCredentials>* out) {
CredentialsFactory* factory;
TF_RETURN_IF_ERROR(CredentialsFactory::Get(protocol, &factory));
TF_RETURN_IF_ERROR(factory->CreateServerCredentials(out));
return absl::OkStatus();
}
absl::Status CredentialsFactory::CreateClientCredentials(
absl::string_view protocol,
std::shared_ptr<::grpc::ChannelCredentials>* out) {
CredentialsFactory* factory;
TF_RETURN_IF_ERROR(CredentialsFactory::Get(protocol, &factory));
TF_RETURN_IF_ERROR(factory->CreateClientCredentials(out));
return absl::OkStatus();
}
bool CredentialsFactory::Exists(absl::string_view protocol) {
mutex_lock l(*get_lock());
return credentials_factories().find(std::string(protocol)) !=
credentials_factories().end();
}
class InsecureCredentialsFactory : public CredentialsFactory {
public:
std::string Protocol() override { return "grpc"; }
absl::Status CreateServerCredentials(
std::shared_ptr<::grpc::ServerCredentials>* out) override {
*out = ::grpc::InsecureServerCredentials();
return absl::OkStatus();
}
absl::Status CreateClientCredentials(
std::shared_ptr<::grpc::ChannelCredentials>* out) override {
*out = ::grpc::InsecureChannelCredentials();
return absl::OkStatus();
}
};
class InsecureCredentialsRegistrar {
public:
InsecureCredentialsRegistrar() {
auto factory = new InsecureCredentialsFactory();
CredentialsFactory::Register(factory);
}
};
static InsecureCredentialsRegistrar registrar;
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/credentials_factory.h
+77
View File
@@ -0,0 +1,77 @@
/* Copyright 2020 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_CORE_DATA_SERVICE_CREDENTIALS_FACTORY_H_
#define TENSORFLOW_CORE_DATA_SERVICE_CREDENTIALS_FACTORY_H_
#include <memory>
#include <string>
#include "grpcpp/grpcpp.h"
#include "grpcpp/security/credentials.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace data {
// Credential factory implementations should be threadsafe since all callers
// to `GetCredentials` will get the same instance of `CredentialsFactory`.
class CredentialsFactory {
public:
virtual ~CredentialsFactory() = default;
// Returns a protocol name for the credentials factory. This is the string to
// look up with `GetCredentials` to find the registered credentials factory.
virtual std::string Protocol() = 0;
// Stores server credentials to `*out`.
virtual absl::Status CreateServerCredentials(
std::shared_ptr<::grpc::ServerCredentials>* out) = 0;
// Stores client credentials to `*out`.
virtual absl::Status CreateClientCredentials(
std::shared_ptr<::grpc::ChannelCredentials>* out) = 0;
// Registers a credentials factory.
static void Register(CredentialsFactory* factory);
// Creates server credentials using the credentials factory registered as
// `protocol`, and stores them to `*out`.
static absl::Status CreateServerCredentials(
absl::string_view protocol,
std::shared_ptr<::grpc::ServerCredentials>* out);
// Creates client credentials using the credentials factory registered as
// `protocol`, and stores them to `*out`.
static absl::Status CreateClientCredentials(
absl::string_view protocol,
std::shared_ptr<::grpc::ChannelCredentials>* out);
// Returns whether a factory has been registered under the given protocol
// name.
static bool Exists(absl::string_view protocol);
private:
// Gets the credentials factory registered via `Register` for the specified
// protocol, and stores it to `*out`.
static absl::Status Get(const absl::string_view protocol,
CredentialsFactory** out);
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_CREDENTIALS_FACTORY_H_
tensorflow/core/data/service/credentials_factory_test.cc
+94
View File
@@ -0,0 +1,94 @@
/* Copyright 2020 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/core/data/service/credentials_factory.h"
#include <memory>
#include <string>
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace data {
namespace {
constexpr char kFailedToCreateServerCredentials[] =
"Failed to create server credentials.";
constexpr char kFailedToCreateClientCredentials[] =
"Failed to create client credentials.";
class TestCredentialsFactory : public CredentialsFactory {
public:
std::string Protocol() override { return "test"; }
absl::Status CreateServerCredentials(
std::shared_ptr<grpc::ServerCredentials>* out) override {
return absl::InternalError(kFailedToCreateServerCredentials);
}
absl::Status CreateClientCredentials(
std::shared_ptr<grpc::ChannelCredentials>* out) override {
return absl::InternalError(kFailedToCreateClientCredentials);
}
};
} // namespace
TEST(CredentialsFactory, Register) {
TestCredentialsFactory test_factory;
CredentialsFactory::Register(&test_factory);
std::shared_ptr<grpc::ServerCredentials> server_credentials;
ASSERT_EQ(absl::InternalError(kFailedToCreateServerCredentials),
CredentialsFactory::CreateServerCredentials(test_factory.Protocol(),
&server_credentials));
std::shared_ptr<grpc::ChannelCredentials> client_credentials;
ASSERT_EQ(absl::InternalError(kFailedToCreateClientCredentials),
CredentialsFactory::CreateClientCredentials(test_factory.Protocol(),
&client_credentials));
}
TEST(CredentialsFactory, DefaultGrpcProtocol) {
std::shared_ptr<grpc::ServerCredentials> server_credentials;
TF_ASSERT_OK(
CredentialsFactory::CreateServerCredentials("grpc", &server_credentials));
std::shared_ptr<grpc::ChannelCredentials> client_credentials;
TF_ASSERT_OK(
CredentialsFactory::CreateClientCredentials("grpc", &client_credentials));
}
TEST(CredentialsFactory, MissingServerProtocol) {
std::shared_ptr<grpc::ServerCredentials> server_credentials;
absl::Status s = CredentialsFactory::CreateServerCredentials(
"unknown_protocol", &server_credentials);
ASSERT_EQ(error::Code::NOT_FOUND, s.code());
ASSERT_TRUE(
absl::StrContains(s.ToString(),
"No credentials factory has been registered for "
"protocol unknown_protocol"));
}
TEST(CredentialsFactory, MissingClientProtocol) {
std::shared_ptr<grpc::ChannelCredentials> client_credentials;
absl::Status s = CredentialsFactory::CreateClientCredentials(
"unknown_protocol", &client_credentials);
ASSERT_EQ(error::Code::NOT_FOUND, s.code());
ASSERT_TRUE(
absl::StrContains(s.ToString(),
"No credentials factory has been registered for "
"protocol unknown_protocol"));
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/cross_trainer_cache.h
+355
View File
@@ -0,0 +1,355 @@
/* Copyright 2022 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_CORE_DATA_SERVICE_CROSS_TRAINER_CACHE_H_
#define TENSORFLOW_CORE_DATA_SERVICE_CROSS_TRAINER_CACHE_H_
#include <cstddef>
#include <deque>
#include <functional>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/data/service/byte_size.h"
#include "tensorflow/core/framework/metrics.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
namespace data {
// Sliding-window cache shared across concurrent trainers. Readers call `Get` to
// read elements they haven't read. After a trainer reads an element, it remains
// in the cache and the data is shared with other trainers. This is useful for
// datasets involving expensive computation, and multiple models use the same
// data for training. For example, for hyperparameter tuning.
//
// The cache progresses when a trainer that has consumed all elements in the
// cache requests additional data. It has a bounded size. Elements are garbage
// collected when the cache becomes full. Consequently, trainers read from a
// sliding window through the dataset and may not read the full dataset.
//
// The `CrossTrainerCache` class is thread-safe.
//
// Example usage:
//
// // `InfiniteRange` returns 1, 2, 3, ... in the `GetNext` calls.
// class InfiniteRange : public CachableSequence<int64_t> {
// public:
// StatusOr<int64_t> GetNext() override {
// return next_++;
// }
//
// size_t GetElementSizeBytes(const int64_t& element) const override {
// return sizeof(element);
// }
//
// private:
// int64_t next_ = 1;
// };
//
// CrossTrainerCache<int64_t> cache(
// /*max_cache_size_bytes=*/10 * (size_t{1} << 30), // 10GB
// std::make_unique<InfiniteRange>());
//
// std::shared_ptr<int64_t> next;
// TF_ASSIGN_OR_RETURN(next, cache.Get("Trainer 1")); // Returns 1
// TF_ASSIGN_OR_RETURN(next, cache.Get("Trainer 2")); // Returns 1
// TF_ASSIGN_OR_RETURN(next, cache.Get("Trainer 1")); // Returns 2
// TF_ASSIGN_OR_RETURN(next, cache.Get("Trainer 2")); // Returns 2
// To use the cache, the user needs to define a `CachableSequence` to generate
// an infinite sequence of data. It should implement a `GetNext` method to
// produce elements, and a `GetElementSizeBytes` method to estimate the element
// size in bytes.
template <class ElementType>
class CachableSequence {
public:
virtual ~CachableSequence() = default;
// Returns the next element to be cached.
virtual StatusOr<ElementType> GetNext() = 0;
// Returns the estimated size of the element in bytes.
virtual size_t GetElementSizeBytes(const ElementType&) const = 0;
};
// Sliding-window cache shared across concurrent trainers.
template <class ElementType>
class CrossTrainerCache {
public:
// Creates a `CrossTrainerCache` with `max_cache_size_bytes` of memory budget.
// The cache should be able to hold at least one element, i.e.:
// REQUIRES: `max_cache_size_bytes >= max(GetElementSizeBytes(*))`
explicit CrossTrainerCache(
size_t max_cache_size_bytes,
std::unique_ptr<CachableSequence<ElementType>> cachable_sequence);
virtual ~CrossTrainerCache() = default;
CrossTrainerCache(const CrossTrainerCache&) = delete;
CrossTrainerCache& operator=(const CrossTrainerCache&) = delete;
// Gets the next element for a trainer. A `trainer_id` identifies the trainer
// reading from the cache. A trainer reads the next element it hasn't read
// before. After a trainer reads data, the data is cached and reused by other
// trainers.
StatusOr<std::shared_ptr<const ElementType>> Get(
const std::string& trainer_id);
// Cancels the cache with `status` and notifies the readers. After cancelling,
// all `Get` calls will return `status`.
// REQUIRES: !status.ok()
void Cancel(absl::Status status);
// Returns true if the cache has been cancelled.
bool IsCancelled() const;
private:
struct CacheQueryResult {
std::shared_ptr<const ElementType> element;
bool cache_hit;
};
// Returns the next element and metrics about this query.
StatusOr<CacheQueryResult> GetCacheQueryResult(const std::string& trainer_id);
// Returns true if element is ready for `trainer_id`. An element is ready if
// other trainers have read the data and the data remains in the cache. If the
// data is not ready, one of the trainers need to extend the cache.
bool IsElementReady(const std::string& trainer_id);
// Returns the absolute element index relative to the dataset (not relative to
// the cached elements).
size_t GetElementIndex(const std::string& trainer_id);
// Returns the next element for `trainer_id`.
StatusOr<std::shared_ptr<const ElementType>> GetElement(
const std::string& trainer_id);
// Reads a new element and writes it into the cache.
absl::Status ExtendCache();
// Frees old elements to keep the cache size below `max_cache_size_bytes_`.
// `new_element_size_bytes` is the size of the new element being inserted.
void FreeSpace(size_t new_element_size_bytes);
// Records the cache hit rate and cache size.
void RecordMetrics(const CacheQueryResult& result);
// Maximum cache size in bytes.
const size_t max_cache_size_bytes_;
// The element sequence over which the sliding window cache operates.
std::unique_ptr<CachableSequence<ElementType>> cachable_sequence_;
mutable mutex mu_;
mutable condition_variable cv_;
// If `status_` is non-OK, the cache is cancelled, and all method calls will
// return this status.
absl::Status status_ TF_GUARDED_BY(mu_) = absl::OkStatus();
// `cache_` stores the cached elements.
std::deque<std::shared_ptr<const ElementType>> cache_ TF_GUARDED_BY(mu_);
size_t cache_size_bytes_ TF_GUARDED_BY(mu_) = 0;
size_t cache_start_index_ TF_GUARDED_BY(mu_) = 0;
// True if one thread is extending the cache.
bool extending_cache_ TF_GUARDED_BY(mu_) = false;
// Maps trainer IDs to element indices. The indices are absolute indices
// within the dataset. The actual index to use with `cache_` would be
// `trainer_to_element_index_map_[trainer_id] - cache_start_index_`.
absl::flat_hash_map<std::string, size_t> trainer_to_element_index_map_
TF_GUARDED_BY(mu_);
};
template <class ElementType>
CrossTrainerCache<ElementType>::CrossTrainerCache(
size_t max_cache_size_bytes,
std::unique_ptr<CachableSequence<ElementType>> cachable_sequence)
: max_cache_size_bytes_(max_cache_size_bytes),
cachable_sequence_(std::move(cachable_sequence)) {
DCHECK_GT(max_cache_size_bytes, 0)
<< "CrossTrainerCache size must be greater than 0.";
VLOG(2) << "Initialized tf.data service cross-trainer cache with "
<< ByteSize::Bytes(max_cache_size_bytes) << " of memory.";
}
template <class ElementType>
StatusOr<std::shared_ptr<const ElementType>>
CrossTrainerCache<ElementType>::Get(const std::string& trainer_id)
TF_LOCKS_EXCLUDED(mu_) {
if (trainer_id.empty()) {
return absl::InvalidArgumentError(
"tf.data service cross-trainer cache requires a non-empty trainer ID.");
}
TF_ASSIGN_OR_RETURN(CacheQueryResult result, GetCacheQueryResult(trainer_id));
RecordMetrics(result);
return result.element;
}
template <class ElementType>
StatusOr<typename CrossTrainerCache<ElementType>::CacheQueryResult>
CrossTrainerCache<ElementType>::GetCacheQueryResult(
const std::string& trainer_id) {
bool should_extend_cache = false;
while (true) {
{
mutex_lock l(mu_);
TF_RETURN_IF_ERROR(status_);
if (IsElementReady(trainer_id)) {
TF_ASSIGN_OR_RETURN(std::shared_ptr<const ElementType> element,
GetElement(trainer_id));
return CacheQueryResult{element,
/*is_cache_hit=*/!should_extend_cache};
}
// Extends the cache or waits for another thread to extend the cache. When
// concurrent trainers wait for the next element, only one of them should
// extend the cache.
if (extending_cache_) {
should_extend_cache = false;
cv_.wait(l);
} else {
should_extend_cache = true;
extending_cache_ = true;
}
}
if (should_extend_cache) {
absl::Status s = ExtendCache();
mutex_lock l(mu_);
extending_cache_ = false;
cv_.notify_all();
TF_RETURN_IF_ERROR(s);
}
}
}
template <class ElementType>
bool CrossTrainerCache<ElementType>::IsElementReady(
const std::string& trainer_id) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
return GetElementIndex(trainer_id) < cache_start_index_ + cache_.size();
}
template <class ElementType>
StatusOr<std::shared_ptr<const ElementType>>
CrossTrainerCache<ElementType>::GetElement(const std::string& trainer_id)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
size_t element_index = GetElementIndex(trainer_id);
if (element_index >= std::numeric_limits<size_t>::max()) {
return absl::InternalError(absl::StrCat(
"tf.data service caching element index exceeds integer limit. Got ",
element_index));
}
std::shared_ptr<const ElementType> result =
cache_[element_index - cache_start_index_];
trainer_to_element_index_map_[trainer_id] = element_index + 1;
return result;
}
template <class ElementType>
size_t CrossTrainerCache<ElementType>::GetElementIndex(
const std::string& trainer_id) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
size_t element_index = trainer_to_element_index_map_[trainer_id];
if (element_index < cache_start_index_) {
element_index = cache_start_index_;
}
return element_index;
}
template <class ElementType>
absl::Status CrossTrainerCache<ElementType>::ExtendCache()
TF_LOCKS_EXCLUDED(mu_) {
TF_ASSIGN_OR_RETURN(ElementType element, cachable_sequence_->GetNext());
size_t new_element_size_bytes =
cachable_sequence_->GetElementSizeBytes(element);
if (new_element_size_bytes > max_cache_size_bytes_) {
return absl::InvalidArgumentError(absl::StrCat(
"tf.data service element size is larger than cache size in bytes. Got ",
"element size: ", new_element_size_bytes,
" and cache size: ", max_cache_size_bytes_));
}
mutex_lock l(mu_);
TF_RETURN_IF_ERROR(status_);
FreeSpace(new_element_size_bytes);
cache_.push_back(std::make_shared<ElementType>(std::move(element)));
cache_size_bytes_ += new_element_size_bytes;
return absl::OkStatus();
}
template <class ElementType>
void CrossTrainerCache<ElementType>::FreeSpace(size_t new_element_size_bytes)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
size_t num_elements_discarded = 0;
while (!cache_.empty() &&
cache_size_bytes_ + new_element_size_bytes > max_cache_size_bytes_) {
size_t free_bytes =
cachable_sequence_->GetElementSizeBytes(*cache_.front());
cache_.pop_front();
cache_size_bytes_ -= free_bytes;
++cache_start_index_;
++num_elements_discarded;
}
VLOG(3) << "Freed " << num_elements_discarded << " element(s) from "
<< "tf.data service cross-trainer cache. Memory usage: "
<< ByteSize::Bytes(cache_size_bytes_) << ".";
}
template <class ElementType>
void CrossTrainerCache<ElementType>::Cancel(absl::Status status)
TF_LOCKS_EXCLUDED(mu_) {
DCHECK(!status.ok())
<< "Cancelling CrossTrainerCache requires a non-OK status. Got "
<< status;
VLOG(2) << "Cancel tf.data service cross-trainer cache with status "
<< status;
mutex_lock l(mu_);
status_ = std::move(status);
cv_.notify_all();
}
template <class ElementType>
bool CrossTrainerCache<ElementType>::IsCancelled() const
TF_LOCKS_EXCLUDED(mu_) {
mutex_lock l(mu_);
return !status_.ok();
}
template <class ElementType>
void CrossTrainerCache<ElementType>::RecordMetrics(
const CacheQueryResult& result) {
metrics::RecordTFDataServiceCrossTrainerCacheQuery(result.cache_hit);
size_t cache_size_bytes = 0;
{
mutex_lock l(mu_);
cache_size_bytes = cache_size_bytes_;
}
metrics::RecordTFDataServiceCrossTrainerCacheSizeBytes(cache_size_bytes);
}
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_CROSS_CLIENT_CACHE_H_
tensorflow/core/data/service/cross_trainer_cache_test.cc
+487
View File
@@ -0,0 +1,487 @@
/* Copyright 2022 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/core/data/service/cross_trainer_cache.h"
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "absl/time/time.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/monitoring/cell_reader.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/random.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::monitoring::testing::CellReader;
using ::tensorflow::testing::IsOkAndHolds;
using ::tensorflow::testing::StatusIs;
using ::testing::Gt;
using ::testing::HasSubstr;
using ::testing::Pointee;
using ::testing::UnorderedElementsAreArray;
class InfiniteRange : public CachableSequence<int64_t> {
public:
absl::StatusOr<int64_t> GetNext() override { return next_++; }
size_t GetElementSizeBytes(const int64_t& element) const override {
return sizeof(element);
}
private:
// No need to guard this variable because only one thread can write the cache.
int64_t next_ = 0;
};
class TensorDataset : public CachableSequence<Tensor> {
public:
absl::StatusOr<Tensor> GetNext() override { return Tensor("Test Tensor"); }
size_t GetElementSizeBytes(const Tensor& element) const override {
return element.TotalBytes();
}
};
class SlowDataset : public CachableSequence<Tensor> {
public:
explicit SlowDataset(absl::Duration delay) : delay_(delay) {}
absl::StatusOr<Tensor> GetNext() override {
Env::Default()->SleepForMicroseconds(absl::ToInt64Microseconds(delay_));
return Tensor("Test Tensor");
}
size_t GetElementSizeBytes(const Tensor& element) const override {
return element.TotalBytes();
}
private:
absl::Duration delay_;
};
template <class T>
class ElementOrErrorDataset : public CachableSequence<T> {
public:
explicit ElementOrErrorDataset(const std::vector<StatusOr<T>>& elements)
: elements_(elements) {}
StatusOr<T> GetNext() override {
if (next_ >= elements_.size()) {
return absl::OutOfRangeError("Out of range.");
}
return elements_[next_++];
}
size_t GetElementSizeBytes(const T& element) const override {
return sizeof(element);
}
private:
const std::vector<StatusOr<T>> elements_;
int64_t next_ = 0;
};
template <>
size_t ElementOrErrorDataset<std::string>::GetElementSizeBytes(
const std::string& element) const {
return element.size();
}
template <>
size_t ElementOrErrorDataset<Tensor>::GetElementSizeBytes(
const Tensor& element) const {
return element.TotalBytes();
}
std::vector<int64_t> GetRange(const size_t range) {
std::vector<int64_t> result;
for (int64_t i = 0; i < range; ++i) {
result.push_back(i);
}
return result;
}
bool SequenceIsIncreasing(const std::vector<int64_t> sequence) {
for (int i = 1; i < sequence.size(); ++i) {
if (sequence[i - 1] > sequence[i - 1]) {
return false;
}
}
return true;
}
TEST(CrossTrainerCacheTest, GetFromOneTrainer) {
const size_t num_elements = 10;
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/1024, std::make_unique<InfiniteRange>());
for (size_t i = 0; i < num_elements; ++i) {
EXPECT_THAT(cache.Get("Trainer ID"),
absl_testing::IsOkAndHolds(Pointee(i)));
}
}
TEST(CrossTrainerCacheTest, GetFromMultipleTrainers) {
const size_t num_elements = 10;
const size_t num_trainers = 10;
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/1024, std::make_unique<InfiniteRange>());
for (size_t i = 0; i < num_elements; ++i) {
// All the readers get the same element in one step.
for (size_t j = 0; j < num_trainers; ++j) {
const std::string trainer_id = absl::StrCat("Trainer ", j);
EXPECT_THAT(cache.Get(trainer_id),
absl_testing::IsOkAndHolds(Pointee(i)));
}
}
}
TEST(CrossTrainerCacheTest, SlowTrainersSkipData) {
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/5 * sizeof(int64_t),
std::make_unique<InfiniteRange>());
EXPECT_THAT(cache.Get("Fast trainer 1"),
absl_testing::IsOkAndHolds(Pointee(0)));
EXPECT_THAT(cache.Get("Fast trainer 2"),
absl_testing::IsOkAndHolds(Pointee(0)));
EXPECT_THAT(cache.Get("Slow trainer 1"),
absl_testing::IsOkAndHolds(Pointee(0)));
EXPECT_THAT(cache.Get("Slow trainer 2"),
absl_testing::IsOkAndHolds(Pointee(0)));
for (int i = 1; i < 20; ++i) {
EXPECT_THAT(cache.Get("Fast trainer 1"),
absl_testing::IsOkAndHolds(Pointee(i)));
EXPECT_THAT(cache.Get("Fast trainer 2"),
absl_testing::IsOkAndHolds(Pointee(i)));
}
// When 19 is cached, 14 must have been discarded.
EXPECT_THAT(cache.Get("Slow trainer 1"),
absl_testing::IsOkAndHolds(Pointee(Gt(14))));
EXPECT_THAT(cache.Get("Slow trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(14))));
for (int i = 20; i < 100; ++i) {
EXPECT_THAT(cache.Get("Fast trainer 1"),
absl_testing::IsOkAndHolds(Pointee(i)));
EXPECT_THAT(cache.Get("Fast trainer 2"),
absl_testing::IsOkAndHolds(Pointee(i)));
}
// When 99 is cached, 94 must have been discarded.
EXPECT_THAT(cache.Get("Slow trainer 1"),
absl_testing::IsOkAndHolds(Pointee(Gt(94))));
EXPECT_THAT(cache.Get("Slow trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(94))));
}
TEST(CrossTrainerCacheTest, NewTrainersStartLate) {
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/5 * sizeof(int64_t),
std::make_unique<InfiniteRange>());
for (int i = 0; i < 100; ++i) {
EXPECT_THAT(cache.Get("Old trainer"),
absl_testing::IsOkAndHolds(Pointee(i)));
}
// New trainers start to read after the first trainer has finished.
for (int j = 0; j < 100; ++j) {
EXPECT_THAT(cache.Get(absl::StrCat("New trainer ", j)),
absl_testing::IsOkAndHolds(Pointee(Gt(94))));
}
}
TEST(CrossTrainerCacheTest, AlternateTrainerExtendsCache) {
// The cache size is smaller than one int64_t.
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/sizeof(int64_t),
std::make_unique<InfiniteRange>());
EXPECT_THAT(cache.Get("Trainer 1"), absl_testing::IsOkAndHolds(Pointee(0)));
EXPECT_THAT(cache.Get("Trainer 1"), absl_testing::IsOkAndHolds(Pointee(1)));
EXPECT_THAT(cache.Get("Trainer 1"), absl_testing::IsOkAndHolds(Pointee(2)));
// When 2 is cached, 0 must have been discarded.
EXPECT_THAT(cache.Get("Trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(0))));
EXPECT_THAT(cache.Get("Trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(1))));
EXPECT_THAT(cache.Get("Trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(2))));
// When 3 is cached, 1 must have been discarded.
EXPECT_THAT(cache.Get("Trainer 1"),
absl_testing::IsOkAndHolds(Pointee(Gt(1))));
EXPECT_THAT(cache.Get("Trainer 1"),
absl_testing::IsOkAndHolds(Pointee(Gt(2))));
EXPECT_THAT(cache.Get("Trainer 1"),
absl_testing::IsOkAndHolds(Pointee(Gt(3))));
// When 4 is cached, 2 must have been discarded.
EXPECT_THAT(cache.Get("Trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(2))));
EXPECT_THAT(cache.Get("Trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(3))));
EXPECT_THAT(cache.Get("Trainer 2"),
absl_testing::IsOkAndHolds(Pointee(Gt(4))));
// When 5 is cached, 3 must have been discarded.
EXPECT_THAT(cache.Get("Trainer 3"),
absl_testing::IsOkAndHolds(Pointee(Gt(3))));
EXPECT_THAT(cache.Get("Trainer 3"),
absl_testing::IsOkAndHolds(Pointee(Gt(4))));
EXPECT_THAT(cache.Get("Trainer 3"),
absl_testing::IsOkAndHolds(Pointee(Gt(5))));
}
TEST(CrossTrainerCacheTest, CacheHitMetrics) {
CellReader<int64_t> cell_reader(
"/tensorflow/data/service/cross_trainer_cache_queries");
EXPECT_EQ(cell_reader.Delta("true"), 0);
EXPECT_EQ(cell_reader.Delta("false"), 0);
EXPECT_EQ(cell_reader.Read("true"), 0);
EXPECT_EQ(cell_reader.Read("false"), 0);
const size_t num_elements = 10;
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/1024, std::make_unique<InfiniteRange>());
for (size_t i = 0; i < num_elements; ++i) {
EXPECT_THAT(cache.Get("Trainer 1"), absl_testing::IsOkAndHolds(Pointee(i)));
}
EXPECT_EQ(cell_reader.Delta("true"), 0);
EXPECT_EQ(cell_reader.Delta("false"), 10);
EXPECT_EQ(cell_reader.Read("true"), 0);
EXPECT_EQ(cell_reader.Read("false"), 10);
for (size_t i = 0; i < num_elements; ++i) {
EXPECT_THAT(cache.Get("Trainer 2"), absl_testing::IsOkAndHolds(Pointee(i)));
}
EXPECT_EQ(cell_reader.Delta("true"), 10);
EXPECT_EQ(cell_reader.Delta("false"), 0);
EXPECT_EQ(cell_reader.Read("true"), 10);
EXPECT_EQ(cell_reader.Read("false"), 10);
}
TEST(CrossTrainerCacheTest, CacheSizeMetrics) {
CellReader<int64_t> cell_reader(
"/tensorflow/data/service/cross_trainer_cache_size_bytes");
const size_t num_elements = 5;
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/num_elements * sizeof(int64_t),
std::make_unique<InfiniteRange>());
for (size_t i = 0; i < num_elements; ++i) {
EXPECT_THAT(cache.Get("Trainer 1"), absl_testing::IsOkAndHolds(Pointee(i)));
EXPECT_EQ(cell_reader.Read(), (i + 1) * sizeof(int64_t));
}
// The cache size does not increase after reaching `num_elements`.
for (size_t i = 0; i < 100; ++i) {
EXPECT_THAT(cache.Get("Trainer 1"),
absl_testing::IsOkAndHolds(Pointee(num_elements + i)));
EXPECT_EQ(cell_reader.Read(), 5 * sizeof(int64_t));
}
}
TEST(CrossTrainerCacheTest, ConcurrentReaders) {
size_t num_trainers = 10;
size_t num_elements_to_read = 200;
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/3 * sizeof(int64_t),
std::make_unique<InfiniteRange>());
std::vector<std::vector<int64_t>> results;
std::vector<std::unique_ptr<Thread>> reader_threads;
results.reserve(num_trainers);
for (size_t i = 0; i < num_trainers; ++i) {
results.emplace_back();
std::vector<int64_t>& result = results.back();
reader_threads.push_back(absl::WrapUnique(Env::Default()->StartThread(
/*thread_options=*/{}, /*name=*/absl::StrCat("Trainer_", i),
[&cache, num_elements_to_read, &result]() {
for (size_t i = 0; i < num_elements_to_read; ++i) {
// Randomly slows down some trainers.
if (random::New64() % 5 == 0) {
Env::Default()->SleepForMicroseconds(2000);
}
TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<const int64_t> next,
cache.Get(absl::StrCat("Trainer_", i)));
result.push_back(*next);
}
})));
}
reader_threads.clear();
// Verifies all trainers can read `num_elements_to_read` elements.
EXPECT_EQ(results.size(), num_trainers);
for (const std::vector<int64_t>& result : results) {
EXPECT_EQ(result.size(), num_elements_to_read);
EXPECT_TRUE(SequenceIsIncreasing(result));
}
}
TEST(CrossTrainerCacheTest, ConcurrentReadersFromOneTrainer) {
size_t num_trainers = 10;
size_t num_elements_to_read = 100;
CrossTrainerCache<int64_t> cache(
/*max_cache_size_bytes=*/3 * sizeof(int64_t),
std::make_unique<InfiniteRange>());
mutex mu;
std::vector<int64_t> results; // Guarded by `mu`.
std::vector<std::unique_ptr<Thread>> reader_threads;
for (size_t i = 0; i < num_trainers; ++i) {
reader_threads.push_back(absl::WrapUnique(Env::Default()->StartThread(
/*thread_options=*/{}, /*name=*/absl::StrCat("Thread_", i),
[&cache, num_elements_to_read, &results, &mu]() {
for (size_t i = 0; i < num_elements_to_read; ++i) {
// Randomly slows down some trainers.
if (random::New64() % 5 == 0) {
Env::Default()->SleepForMicroseconds(1000);
}
TF_ASSERT_OK_AND_ASSIGN(std::shared_ptr<const int64_t> next,
cache.Get("Trainer ID"));
mutex_lock l(mu);
results.push_back(*next);
}
})));
}
reader_threads.clear();
// Verifies the readers have read all elements because they have the same
// trainer ID.
EXPECT_THAT(results, UnorderedElementsAreArray(GetRange(1000)));
}
TEST(CrossTrainerCacheTest, Cancel) {
size_t num_trainers = 10;
CrossTrainerCache<Tensor> cache(
/*max_cache_size_bytes=*/1000, std::make_unique<TensorDataset>());
EXPECT_FALSE(cache.IsCancelled());
mutex mu;
absl::Status status; // Guarded by `mu`.
std::vector<std::unique_ptr<Thread>> reader_threads;
for (size_t i = 0; i < num_trainers; ++i) {
reader_threads.push_back(absl::WrapUnique(Env::Default()->StartThread(
/*thread_options=*/{}, /*name=*/absl::StrCat("Trainer_", i),
[&cache, &status, &mu]() {
for (int j = 0; true; ++j) {
absl::StatusOr<std::shared_ptr<const Tensor>> tensor =
cache.Get(absl::StrCat("Trainer_", j % 1000));
{
mutex_lock l(mu);
status = tensor.status();
}
if (!tensor.status().ok()) {
return;
}
test::ExpectEqual(*tensor.value(), Tensor("Test Tensor"));
}
})));
}
Env::Default()->SleepForMicroseconds(1000000);
cache.Cancel(absl::CancelledError("Cancelled"));
reader_threads.clear();
mutex_lock l(mu);
EXPECT_THAT(status, absl_testing::StatusIs(error::CANCELLED));
EXPECT_THAT(cache.Get("New trainer"),
absl_testing::StatusIs(error::CANCELLED));
EXPECT_TRUE(cache.IsCancelled());
}
TEST(CrossTrainerCacheTest, Errors) {
auto elements = std::make_unique<ElementOrErrorDataset<std::string>>(
std::vector<absl::StatusOr<std::string>>{
std::string("First element"),
absl::CancelledError("Cancelled"),
std::string("Second element"),
absl::InvalidArgumentError("InvalidArgument"),
std::string("Third element"),
absl::UnavailableError("Unavailable"),
});
CrossTrainerCache<std::string> cache(
/*max_cache_size_bytes=*/1000, std::move(elements));
EXPECT_THAT(
cache.Get("Trainer ID"),
absl_testing::IsOkAndHolds(Pointee(std::string("First element"))));
EXPECT_THAT(cache.Get("Trainer ID"),
absl_testing::StatusIs(error::CANCELLED));
EXPECT_THAT(
cache.Get("Trainer ID"),
absl_testing::IsOkAndHolds(Pointee(std::string("Second element"))));
EXPECT_THAT(cache.Get("Trainer ID"),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
EXPECT_THAT(
cache.Get("Trainer ID"),
absl_testing::IsOkAndHolds(Pointee(std::string("Third element"))));
EXPECT_THAT(cache.Get("Trainer ID"),
absl_testing::StatusIs(error::UNAVAILABLE));
// Errors are not stored in the cache.
EXPECT_THAT(
cache.Get("New Trainer"),
absl_testing::IsOkAndHolds(Pointee(std::string("First element"))));
EXPECT_THAT(
cache.Get("New Trainer"),
absl_testing::IsOkAndHolds(Pointee(std::string("Second element"))));
EXPECT_THAT(
cache.Get("New Trainer"),
absl_testing::IsOkAndHolds(Pointee(std::string("Third element"))));
}
TEST(CrossTrainerCacheTest, CacheSizeIsTooSmall) {
// The cache size is smaller than one int64_t.
CrossTrainerCache<Tensor> cache(
/*max_cache_size_bytes=*/1, std::make_unique<TensorDataset>());
EXPECT_THAT(cache.Get("Trainer ID"),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr("tf.data service element size is larger than "
"cache size in bytes.")));
}
TEST(CrossTrainerCacheTest, TrainerIDMustBeNonEmpty) {
CrossTrainerCache<Tensor> cache(
/*max_cache_size_bytes=*/1000, std::make_unique<TensorDataset>());
EXPECT_THAT(cache.Get(""),
absl_testing::StatusIs(error::INVALID_ARGUMENT,
"tf.data service cross-trainer cache "
"requires a non-empty trainer ID."));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/data_service_test.cc
+360
View File
@@ -0,0 +1,360 @@
/* Copyright 2021 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 <vector>
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/data/service/test_cluster.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::data::testing::InterleaveTextlineDataset;
using ::tensorflow::data::testing::RangeDataset;
using ::tensorflow::data::testing::RangeDatasetWithShardHint;
using ::tensorflow::data::testing::WaitWhile;
using ::tensorflow::testing::IsOkAndHolds;
using ::testing::ElementsAre;
using ::testing::ElementsAreArray;
using ::testing::HasSubstr;
using ::testing::Pair;
using ::testing::SizeIs;
using ::testing::TestWithParam;
using ::testing::UnorderedElementsAre;
using ::testing::UnorderedElementsAreArray;
using ::testing::Values;
tstring LocalTempFilename() {
std::string path;
CHECK(Env::Default()->LocalTempFilename(&path));
return tstring(path);
}
std::vector<int64_t> Range(const int64_t range) {
std::vector<int64_t> result;
for (int64_t i = 0; i < range; ++i) {
result.push_back(i);
}
return result;
}
TEST(DataServiceTest, RangeDataset_NoShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
EXPECT_THAT(
dataset_client.Read(RangeDataset(20), ProcessingModeDef::OFF,
TARGET_WORKERS_AUTO),
absl_testing::IsOkAndHolds(UnorderedElementsAre(
Pair(cluster.WorkerAddress(0), ElementsAreArray(Range(20))),
Pair(cluster.WorkerAddress(1), ElementsAreArray(Range(20))),
Pair(cluster.WorkerAddress(2), ElementsAreArray(Range(20))),
Pair(cluster.WorkerAddress(3), ElementsAreArray(Range(20))),
Pair(cluster.WorkerAddress(4), ElementsAreArray(Range(20))))));
}
TEST(DataServiceTest, RangeDataset_DynamicShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
TF_ASSERT_OK_AND_ASSIGN(
DatasetClient<int64_t>::WorkerResultMap worker_results,
dataset_client.Read(RangeDataset(20), ProcessingModeDef::DYNAMIC,
TARGET_WORKERS_AUTO));
std::vector<int64_t> result;
for (const auto& worker_result : worker_results) {
result.insert(result.end(), worker_result.second.begin(),
worker_result.second.end());
}
EXPECT_THAT(result, UnorderedElementsAreArray(Range(20)));
}
using DataServiceTest_DataShard =
::testing::TestWithParam<ProcessingModeDef::ShardingPolicy>;
TEST_P(DataServiceTest_DataShard, RangeDataset_DataShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
EXPECT_THAT(
dataset_client.Read(RangeDataset(20), GetParam(), TARGET_WORKERS_LOCAL),
absl_testing::IsOkAndHolds(UnorderedElementsAre(
Pair(cluster.WorkerAddress(0), ElementsAre(0, 5, 10, 15)),
Pair(cluster.WorkerAddress(1), ElementsAre(1, 6, 11, 16)),
Pair(cluster.WorkerAddress(2), ElementsAre(2, 7, 12, 17)),
Pair(cluster.WorkerAddress(3), ElementsAre(3, 8, 13, 18)),
Pair(cluster.WorkerAddress(4), ElementsAre(4, 9, 14, 19)))));
}
INSTANTIATE_TEST_SUITE_P(ShardingPolicy, DataServiceTest_DataShard,
Values(ProcessingModeDef::FILE_OR_DATA,
ProcessingModeDef::DATA));
TEST(DataServiceTest, RangeDataset_HintShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
EXPECT_THAT(
dataset_client.Read(RangeDatasetWithShardHint(20),
ProcessingModeDef::HINT, TARGET_WORKERS_LOCAL),
absl_testing::IsOkAndHolds(UnorderedElementsAre(
Pair(cluster.WorkerAddress(0), ElementsAre(0, 5, 10, 15)),
Pair(cluster.WorkerAddress(1), ElementsAre(1, 6, 11, 16)),
Pair(cluster.WorkerAddress(2), ElementsAre(2, 7, 12, 17)),
Pair(cluster.WorkerAddress(3), ElementsAre(3, 8, 13, 18)),
Pair(cluster.WorkerAddress(4), ElementsAre(4, 9, 14, 19)))));
}
TEST(DataServiceTest, TextlineDataset_NoShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<tstring> dataset_client(cluster);
std::vector<tstring> filenames = {LocalTempFilename(), LocalTempFilename(),
LocalTempFilename(), LocalTempFilename(),
LocalTempFilename()};
TF_ASSERT_OK_AND_ASSIGN(
const DatasetDef dataset,
InterleaveTextlineDataset(
filenames, {"0", "1\n1", "2\n2\n2", "3\n3\n3\n3", "4\n4\n4\n4\n4"}));
std::vector<tstring> expected = {"0", "1", "2", "3", "4", "1", "2", "3",
"4", "2", "3", "4", "3", "4", "4"};
EXPECT_THAT(
dataset_client.Read(dataset, ProcessingModeDef::OFF, TARGET_WORKERS_ANY),
absl_testing::IsOkAndHolds(UnorderedElementsAre(
Pair(cluster.WorkerAddress(0), ElementsAreArray(expected)),
Pair(cluster.WorkerAddress(1), ElementsAreArray(expected)),
Pair(cluster.WorkerAddress(2), ElementsAreArray(expected)),
Pair(cluster.WorkerAddress(3), ElementsAreArray(expected)),
Pair(cluster.WorkerAddress(4), ElementsAreArray(expected)))));
}
TEST(DataServiceTest, TextlineDataset_DataShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<tstring> dataset_client(cluster);
std::vector<tstring> filenames = {LocalTempFilename(), LocalTempFilename(),
LocalTempFilename(), LocalTempFilename(),
LocalTempFilename()};
TF_ASSERT_OK_AND_ASSIGN(
const DatasetDef dataset,
InterleaveTextlineDataset(
filenames, {"0", "1\n1", "2\n2\n2", "3\n3\n3\n3", "4\n4\n4\n4\n4"}));
EXPECT_THAT(dataset_client.Read(dataset, ProcessingModeDef::DATA,
TARGET_WORKERS_LOCAL),
absl_testing::IsOkAndHolds(UnorderedElementsAre(
Pair(cluster.WorkerAddress(0), ElementsAre("0", "1", "3")),
Pair(cluster.WorkerAddress(1), ElementsAre("1", "2", "4")),
Pair(cluster.WorkerAddress(2), ElementsAre("2", "3", "3")),
Pair(cluster.WorkerAddress(3), ElementsAre("3", "4", "4")),
Pair(cluster.WorkerAddress(4), ElementsAre("4", "2", "4")))));
}
using DataServiceTest_FileShard =
::testing::TestWithParam<ProcessingModeDef::ShardingPolicy>;
TEST_P(DataServiceTest_FileShard, TextlineDataset_FileShard) {
TestCluster cluster(/*num_workers=*/5);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<tstring> dataset_client(cluster);
std::vector<tstring> filenames = {LocalTempFilename(), LocalTempFilename(),
LocalTempFilename(), LocalTempFilename(),
LocalTempFilename()};
TF_ASSERT_OK_AND_ASSIGN(
const DatasetDef dataset,
InterleaveTextlineDataset(
filenames, {"0", "1\n1", "2\n2\n2", "3\n3\n3\n3", "4\n4\n4\n4\n4"}));
EXPECT_THAT(
dataset_client.Read(dataset, ProcessingModeDef::FILE_OR_DATA,
TARGET_WORKERS_LOCAL),
absl_testing::IsOkAndHolds(UnorderedElementsAre(
Pair(cluster.WorkerAddress(0), ElementsAre("0")),
Pair(cluster.WorkerAddress(1), ElementsAre("1", "1")),
Pair(cluster.WorkerAddress(2), ElementsAre("2", "2", "2")),
Pair(cluster.WorkerAddress(3), ElementsAre("3", "3", "3", "3")),
Pair(cluster.WorkerAddress(4),
ElementsAre("4", "4", "4", "4", "4")))));
}
INSTANTIATE_TEST_SUITE_P(ShardingPolicy, DataServiceTest_FileShard,
Values(ProcessingModeDef::FILE_OR_DATA,
ProcessingModeDef::FILE));
TEST(DataServiceTest, GcMissingClientsWithSmallTimeout) {
TestCluster::Config config;
config.num_workers = 5;
config.job_gc_check_interval_ms = 10;
config.job_gc_timeout_ms = 10;
config.client_timeout_ms = 10;
TestCluster cluster(config);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
TF_ASSERT_OK_AND_ASSIGN(int64_t iteration_client_id,
dataset_client.CreateIteration(RangeDataset(10)));
Env::Default()->SleepForMicroseconds(1000 * 1000); // 1 second.
// Iteration should not be garbage collected before the client has started
// reading.
EXPECT_THAT(cluster.NumActiveIterations(), absl_testing::IsOkAndHolds(1));
TF_ASSERT_OK(dataset_client.GetTasks(iteration_client_id).status());
// Iteration should be garbage collected within 10 seconds.
absl::Time wait_start = absl::Now();
TF_ASSERT_OK(WaitWhile([&]() -> absl::StatusOr<bool> {
TF_ASSIGN_OR_RETURN(size_t num_iterations, cluster.NumActiveIterations());
return num_iterations > 0;
}));
EXPECT_LT(absl::Now(), wait_start + absl::Seconds(10));
}
TEST(DataServiceTest, DontGcMissingClientsWithLargeTimeout) {
TestCluster::Config config;
config.num_workers = 5;
config.job_gc_check_interval_ms = 10;
config.job_gc_timeout_ms = 10;
config.client_timeout_ms = 10000000000;
TestCluster cluster(config);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
TF_ASSERT_OK(dataset_client.CreateIteration(RangeDataset(10)).status());
Env::Default()->SleepForMicroseconds(1000 * 1000); // 1 second.
// Iteration should not be garbage collected, since the client hasn't timed
// out.
EXPECT_THAT(cluster.NumActiveIterations(), absl_testing::IsOkAndHolds(1));
}
TEST(DataServiceTest, GetWorkers) {
TestCluster cluster(1);
TF_ASSERT_OK(cluster.Initialize());
DataServiceDispatcherClient dispatcher(cluster.DispatcherAddress(), "grpc");
std::vector<WorkerInfo> workers;
TF_EXPECT_OK(dispatcher.GetWorkers(workers));
EXPECT_EQ(1, workers.size());
}
TEST(DataServiceTest, DispatcherStateExport) {
TestCluster cluster(1);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
TF_ASSERT_OK(dataset_client.CreateIteration(RangeDataset(10)).status());
ServerStateExport server_state_export = cluster.ExportDispatcherState();
EXPECT_THAT(server_state_export.dispatcher_state_export().worker_addresses(),
ElementsAre(HasSubstr("localhost")));
ASSERT_THAT(server_state_export.dispatcher_state_export().iterations(),
SizeIs(1));
EXPECT_EQ(
server_state_export.dispatcher_state_export().iterations(0).dataset_id(),
"1000");
EXPECT_THAT(server_state_export.dispatcher_state_export()
.iterations(0)
.iteration_key()
.name(),
HasSubstr("anonymous_job"));
EXPECT_EQ(
server_state_export.dispatcher_state_export().iterations(0).num_clients(),
1);
EXPECT_FALSE(
server_state_export.dispatcher_state_export().iterations(0).finished());
}
TEST(DataServiceTest, WorkerStateExport) {
TestCluster::Config config;
config.num_workers = 1;
config.worker_heartbeat_interval_ms = 300;
TestCluster cluster(config);
TF_ASSERT_OK(cluster.Initialize());
DatasetClient<int64_t> dataset_client(cluster);
TF_ASSERT_OK(dataset_client.CreateIteration(RangeDataset(10)).status());
ServerStateExport server_state_export = cluster.ExportWorkerState(0);
EXPECT_THAT(server_state_export.worker_state_export()
.worker_config()
.dispatcher_address(),
HasSubstr("localhost"));
ASSERT_THAT(server_state_export.worker_state_export().tasks(), SizeIs(1));
EXPECT_THAT(server_state_export.worker_state_export().tasks(0).path(),
HasSubstr("In-memory dataset graphs are omitted for brevity."));
TF_ASSERT_OK_AND_ASSIGN(
auto result, dataset_client.Read(RangeDataset(10), ProcessingModeDef::OFF,
TARGET_WORKERS_AUTO));
absl::SleepFor(absl::Seconds(3));
server_state_export = cluster.ExportWorkerState(0);
ASSERT_THAT(server_state_export.worker_state_export().finished_task_ids(),
SizeIs(1));
}
TEST(DataServiceTest, RejectInvalidDatasetId) {
TestCluster cluster(/*num_workers=*/1);
TF_ASSERT_OK(cluster.Initialize());
DataServiceDispatcherClient dispatcher(cluster.DispatcherAddress(), "grpc");
std::string dataset_id;
EXPECT_THAT(dispatcher.RegisterDataset(
RangeDataset(10), DataServiceMetadata(),
/*requested_dataset_id=*/"nested/path", dataset_id),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
EXPECT_THAT(dispatcher.RegisterDataset(
RangeDataset(10), DataServiceMetadata(),
/*requested_dataset_id=*/"some/nested/path", dataset_id),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
#if defined(_WIN32)
EXPECT_THAT(dispatcher.RegisterDataset(
RangeDataset(10), DataServiceMetadata(),
/*requested_dataset_id=*/"invalid\\path", dataset_id),
::tensorflow::testing::StatusIs(error::INVALID_ARGUMENT));
EXPECT_THAT(
dispatcher.RegisterDataset(RangeDataset(10), DataServiceMetadata(),
/*requested_dataset_id=*/"c:path", dataset_id),
::tensorflow::testing::StatusIs(error::INVALID_ARGUMENT));
#endif
EXPECT_THAT(
dispatcher.RegisterDataset(RangeDataset(10), DataServiceMetadata(),
/*requested_dataset_id=*/"..", dataset_id),
absl_testing::StatusIs(error::INVALID_ARGUMENT));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/data_transfer.cc
+143
View File
@@ -0,0 +1,143 @@
/* Copyright 2021 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/core/data/service/data_transfer.h"
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/strings/str_join.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/framework/variant.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace data {
namespace {
mutex* get_lock() {
static mutex lock(LINKER_INITIALIZED);
return &lock;
}
using DataTransferServerFactories =
std::unordered_map<std::string, DataTransferServer::ServerFactoryT>;
DataTransferServerFactories& transfer_server_factories() {
static auto& factories = *new DataTransferServerFactories();
return factories;
}
using DataTransferClientFactories =
std::unordered_map<std::string, DataTransferClient::ClientFactoryT>;
DataTransferClientFactories& transfer_client_factories() {
static auto& factories = *new DataTransferClientFactories();
return factories;
}
} // namespace
GetElementResult GetElementResult::Copy() const {
GetElementResult copy;
copy.components = components;
copy.element_index = element_index;
copy.end_of_sequence = end_of_sequence;
copy.skip = skip;
return copy;
}
size_t GetElementResult::EstimatedMemoryUsageBytes() const {
size_t size_bytes = components.size() * sizeof(Tensor) +
sizeof(element_index) + sizeof(end_of_sequence) +
sizeof(skip);
for (const Tensor& tensor : components) {
size_bytes += tensor.TotalBytes();
if (tensor.dtype() != DT_VARIANT) {
continue;
}
// Estimates the memory usage of a compressed element.
const Variant& variant = tensor.scalar<Variant>()();
const CompressedElement* compressed = variant.get<CompressedElement>();
if (compressed) {
size_bytes += compressed->SpaceUsedLong();
}
}
return size_bytes;
}
void DataTransferServer::Register(std::string name, ServerFactoryT factory) {
mutex_lock l(*get_lock());
if (!transfer_server_factories().insert({name, factory}).second) {
LOG(ERROR)
<< "Two data transfer server factories are being registered with name "
<< name << ". Which one gets used is undefined.";
}
}
absl::Status DataTransferServer::Build(
std::string name, GetElementT get_element,
std::shared_ptr<DataTransferServer>* out) {
mutex_lock l(*get_lock());
auto it = transfer_server_factories().find(name);
if (it != transfer_server_factories().end()) {
return it->second(get_element, out);
}
std::vector<std::string> available_names;
for (const auto& factory : transfer_server_factories()) {
available_names.push_back(factory.first);
}
return absl::NotFoundError(absl::StrCat(
"No data transfer server factory has been registered for name ", name,
". The available names are: [ ", absl::StrJoin(available_names, ", "),
" ]"));
}
void DataTransferClient::Register(std::string name, ClientFactoryT factory) {
mutex_lock l(*get_lock());
if (!transfer_client_factories().insert({name, factory}).second) {
LOG(ERROR)
<< "Two data transfer client factories are being registered with name "
<< name << ". Which one gets used is undefined.";
}
}
absl::Status DataTransferClient::Build(
std::string name, Config config, std::unique_ptr<DataTransferClient>* out) {
mutex_lock l(*get_lock());
auto it = transfer_client_factories().find(name);
if (it != transfer_client_factories().end()) {
return it->second(config, out);
}
std::vector<std::string> available_names;
for (const auto& factory : transfer_client_factories()) {
available_names.push_back(factory.first);
}
return absl::NotFoundError(absl::StrCat(
"No data transfer client factory has been registered for name ", name,
". The available names are: [ ", absl::StrJoin(available_names, ", "),
" ]"));
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/data_transfer.h
+152
View File
@@ -0,0 +1,152 @@
/* Copyright 2021 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_CORE_DATA_SERVICE_DATA_TRANSFER_H_
#define TENSORFLOW_CORE_DATA_SERVICE_DATA_TRANSFER_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/core/data/service/worker.pb.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/framework/dataset.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
// The result of a GetElement request. Exactly one of the following will be
// true: (1) `components` is nonempty (2) `end_of_sequence` is true (3) `skip`
// is true.
struct GetElementResult {
GetElementResult() = default;
GetElementResult(const GetElementResult&) = delete;
GetElementResult& operator=(const GetElementResult&) = delete;
GetElementResult(GetElementResult&&) = default;
GetElementResult& operator=(GetElementResult&&) = default;
// Creates a copy of this result. This is used to create multiple copies of
// the same cached value.
GetElementResult Copy() const;
// Estimated memory used by this object, measured in bytes.
size_t EstimatedMemoryUsageBytes() const;
// A dataset element produced by a GetElement request.
std::vector<Tensor> components;
// The element's index within the task it came from.
int64_t element_index = 0;
// If true, indicates that there is no more data to read.
bool end_of_sequence = false;
// If true, indicates that there is still data, but the caller should skip
// reading from the worker. This is used for load balancing when doing round
// robin reads.
bool skip = false;
};
// Client for communicating with the tf.data service transfer server.
class DataTransferClient {
public:
struct Config {
absl::string_view protocol;
std::string address;
const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info;
Allocator* allocator;
};
using ClientFactoryT =
std::function<absl::Status(Config, std::unique_ptr<DataTransferClient>*)>;
virtual ~DataTransferClient() = default;
// Fetches the next element.
virtual absl::Status GetElement(const GetElementRequest& req,
GetElementResult& result) = 0;
// Makes a best effort to cancel all outstanding calls in progress for the
// client, and causes further calls to return Cancelled status.
virtual void TryCancel() = 0;
// Registers a DataTransferClient factory under `name`.
static void Register(std::string name, ClientFactoryT factory);
// Builds a DataTransferClient from the factory registered under `name`.
static absl::Status Build(std::string name, Config config,
std::unique_ptr<DataTransferClient>* out);
// Returns a string describing properties of the client relevant for checking
// compatibility with a server for a given protocol.
virtual absl::StatusOr<std::string> GetCompatibilityInfo() const {
return std::string();
}
// Returns an error if the client is incompatible with a server which has the
// properties described in `server_compatibility_info`.
virtual absl::Status CheckCompatibility(
const std::string& server_compatibility_info) const {
return absl::OkStatus();
}
protected:
Env* const env_ = Env::Default();
};
// Server for communicating with the tf.data service transfer client.
class DataTransferServer {
public:
using GetElementT =
std::function<absl::Status(const GetElementRequest*, GetElementResult*)>;
using ServerFactoryT = std::function<absl::Status(
GetElementT, std::shared_ptr<DataTransferServer>*)>;
virtual ~DataTransferServer() = default;
// Starts DataTransferServer, it should be available for requests afterwards.
virtual absl::Status Start(const experimental::WorkerConfig& config) = 0;
// Return the port that this server is listening on.
virtual int Port() const = 0;
// Register a DataTransferServer factory under `name`.
static void Register(std::string name, ServerFactoryT factory);
// Builds a DataTransferServer from the factory registered with `name`.
static absl::Status Build(std::string name, GetElementT get_element,
std::shared_ptr<DataTransferServer>* out);
// Returns a string describing properties of the server relevant for checking
// compatibility with a client for a given protocol.
virtual absl::StatusOr<std::string> GetCompatibilityInfo() const {
return std::string();
}
// If `true`, data service clients should fall back to gRPC for this server if
// they fail to create a data transfer client for it.
virtual bool FallBackToGrpcAtClientCreationTime() const { return true; }
// If `true`, data service clients should fall back to gRPC for this server if
// it nonretryably fails to transfer an element.
virtual bool FallBackToGrpcAtGetElementTime() const { return true; }
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_DATA_TRANSFER_H_
tensorflow/core/data/service/data_transfer_test.cc
+117
View File
@@ -0,0 +1,117 @@
/* Copyright 2021 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/core/data/service/data_transfer.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace data {
namespace {
class TestDataTransferServer : public DataTransferServer {
public:
explicit TestDataTransferServer(bool* called) : called_(called) {}
absl::Status Start(const experimental::WorkerConfig& unused_config) override {
*called_ = true;
return absl::OkStatus();
}
int Port() const override { return 0; }
private:
bool* called_;
};
template <class T>
GetElementResult MakeElementResult(T value) {
GetElementResult result;
result.components.push_back(Tensor(std::move(value)));
result.element_index = 0;
result.end_of_sequence = false;
return result;
}
TEST(DataTransferTest, RegisterDataTransferServerBuilder) {
bool called = false;
DataTransferServer::Register("test", [&called](auto ignore, auto* server) {
*server = std::make_shared<TestDataTransferServer>(&called);
return absl::OkStatus();
});
std::shared_ptr<DataTransferServer> server;
TF_ASSERT_OK(DataTransferServer::Build("test", {}, &server));
EXPECT_FALSE(called);
TF_ASSERT_OK(server->Start(/*config=*/{}));
EXPECT_TRUE(called);
}
TEST(DataTransferTest, EstimateMemoryUsageBytes) {
GetElementResult empty;
EXPECT_GT(empty.EstimatedMemoryUsageBytes(), 0);
Tensor tensor(DT_INT64, TensorShape({10, 100}));
GetElementResult int64_result = MakeElementResult(tensor);
EXPECT_GT(int64_result.EstimatedMemoryUsageBytes(), 1000 * sizeof(int64_t));
EXPECT_GT(int64_result.EstimatedMemoryUsageBytes(),
int64_result.components[0].AllocatedBytes());
EXPECT_GE(int64_result.EstimatedMemoryUsageBytes(), sizeof(int64_result));
}
TEST(DataTransferTest, EstimateVariantMemoryUsageBytes) {
const size_t data_size = 1000;
std::unique_ptr<CompressedElement> compressed{
protobuf::Arena::Create<CompressedElement>(nullptr)};
compressed->set_data(std::string(data_size, 'a'));
Tensor tensor(DT_VARIANT, TensorShape({}));
tensor.scalar<Variant>()() = *compressed;
GetElementResult variant_result = MakeElementResult(tensor);
EXPECT_GT(variant_result.EstimatedMemoryUsageBytes(), data_size);
EXPECT_GT(variant_result.EstimatedMemoryUsageBytes(),
compressed->ByteSizeLong());
EXPECT_GT(variant_result.EstimatedMemoryUsageBytes(),
compressed->SpaceUsedLong());
}
TEST(DataTransferTest, CopyGetElementResult) {
std::string hello_world = "hello, world!";
GetElementResult result = MakeElementResult(hello_world);
ASSERT_EQ(result.components.size(), 1);
EXPECT_GT(result.EstimatedMemoryUsageBytes(), hello_world.size());
GetElementResult copy = result.Copy();
ASSERT_EQ(copy.components.size(), 1);
test::ExpectEqual(result.components[0], copy.components[0]);
EXPECT_EQ(copy.EstimatedMemoryUsageBytes(),
result.EstimatedMemoryUsageBytes());
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/dataset_store.cc
+73
View File
@@ -0,0 +1,73 @@
/* Copyright 2020 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/core/data/service/dataset_store.h"
#include <memory>
#include <string>
#include "absl/memory/memory.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/utils.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/path.h"
namespace tensorflow {
namespace data {
FileSystemDatasetStore::FileSystemDatasetStore(const std::string& datasets_dir)
: datasets_dir_(datasets_dir) {}
absl::Status FileSystemDatasetStore::Put(const std::string& key,
const DatasetDef& dataset) {
std::string path_to_write = io::JoinPath(datasets_dir_, key);
TF_RETURN_IF_ERROR(WriteDatasetDef(path_to_write, dataset));
return absl::OkStatus();
}
absl::Status FileSystemDatasetStore::Get(
const std::string& key, std::shared_ptr<const DatasetDef>& dataset_def) {
std::string path = io::JoinPath(datasets_dir_, key);
TF_RETURN_IF_ERROR(Env::Default()->FileExists(path));
DatasetDef def;
TF_RETURN_IF_ERROR(ReadDatasetDef(path, def));
dataset_def = std::make_shared<const DatasetDef>(def);
return absl::OkStatus();
}
absl::Status MemoryDatasetStore::Put(const std::string& key,
const DatasetDef& dataset) {
auto& stored_dataset = datasets_[key];
stored_dataset = std::make_shared<const DatasetDef>(dataset);
return absl::OkStatus();
}
absl::Status MemoryDatasetStore::Get(
const std::string& key, std::shared_ptr<const DatasetDef>& dataset_def) {
auto& stored_dataset = datasets_[key];
if (!stored_dataset) {
return absl::NotFoundError(
absl::StrCat("Dataset with key ", key, " not found"));
}
dataset_def = stored_dataset;
return absl::OkStatus();
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/dataset_store.h
+81
View File
@@ -0,0 +1,81 @@
/* Copyright 2020 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_CORE_DATA_SERVICE_DATASET_STORE_H_
#define TENSORFLOW_CORE_DATA_SERVICE_DATASET_STORE_H_
#include <memory>
#include <string>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/core/data/service/dispatcher_state.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/env.h"
namespace tensorflow {
namespace data {
// An interface for storing and getting dataset definitions.
class DatasetStore {
public:
virtual ~DatasetStore() = default;
// Stores the given dataset under the given key. Overwrites a dataset if it
// already exists.
virtual absl::Status Put(const std::string& key,
const DatasetDef& dataset) = 0;
// Gets the dataset for the given key, storing the dataset in `dataset_def`.
virtual absl::Status Get(const std::string& key,
std::shared_ptr<const DatasetDef>& dataset_def) = 0;
};
// Dataset store which reads and writes datasets within a directory.
// The dataset with key `key` is stored at the path "datasets_dir/key".
class FileSystemDatasetStore : public DatasetStore {
public:
explicit FileSystemDatasetStore(const std::string& datasets_dir);
FileSystemDatasetStore(const FileSystemDatasetStore&) = delete;
FileSystemDatasetStore& operator=(const FileSystemDatasetStore&) = delete;
absl::Status Put(const std::string& key, const DatasetDef& dataset) override;
absl::Status Get(const std::string& key,
std::shared_ptr<const DatasetDef>& dataset_def) override;
private:
const std::string datasets_dir_;
};
// DatasetStore which stores all datasets in memory. This is useful when the
// dispatcher doesn't have a work directory configured.
class MemoryDatasetStore : public DatasetStore {
public:
MemoryDatasetStore() = default;
MemoryDatasetStore(const MemoryDatasetStore&) = delete;
MemoryDatasetStore& operator=(const MemoryDatasetStore&) = delete;
absl::Status Put(const std::string& key, const DatasetDef& dataset) override;
absl::Status Get(const std::string& key,
std::shared_ptr<const DatasetDef>& dataset_def) override;
private:
// Mapping from key to dataset definition.
absl::flat_hash_map<std::string, std::shared_ptr<const DatasetDef>> datasets_;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_DATASET_STORE_H_
tensorflow/core/data/service/dataset_store_test.cc
+121
View File
@@ -0,0 +1,121 @@
/* Copyright 2020 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/core/data/service/dataset_store.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/log/check.h"
#include "absl/memory/memory.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace data {
namespace {
const char kFileSystem[] = "file_system";
const char kMemory[] = "memory";
std::string NewDatasetsDir() {
std::string dir = io::JoinPath(testing::TmpDir(), "datasets");
if (Env::Default()->FileExists(dir).ok()) {
int64_t undeleted_files;
int64_t undeleted_dirs;
CHECK_OK(Env::Default()->DeleteRecursively(dir, &undeleted_files,
&undeleted_dirs));
}
CHECK_OK(Env::Default()->RecursivelyCreateDir(dir));
return dir;
}
std::unique_ptr<DatasetStore> MakeStore(const std::string& type) {
if (type == kFileSystem) {
return std::make_unique<FileSystemDatasetStore>(NewDatasetsDir());
} else if (type == kMemory) {
return std::make_unique<MemoryDatasetStore>();
} else {
CHECK(false) << "unexpected type: " << type;
}
}
DatasetDef DatasetDefWithVersion(int32_t version) {
DatasetDef def;
def.mutable_graph()->set_version(version);
return def;
}
} // namespace
class DatasetStoreTest : public ::testing::Test,
public ::testing::WithParamInterface<std::string> {};
TEST_P(DatasetStoreTest, StoreAndGet) {
std::unique_ptr<DatasetStore> store = MakeStore(GetParam());
std::string key = "key";
DatasetDef dataset_def = DatasetDefWithVersion(1);
TF_ASSERT_OK(store->Put(key, dataset_def));
std::shared_ptr<const DatasetDef> result;
TF_ASSERT_OK(store->Get(key, result));
EXPECT_EQ(result->graph().version(), dataset_def.graph().version());
}
TEST_P(DatasetStoreTest, StoreAndGetMultiple) {
std::unique_ptr<DatasetStore> store = MakeStore(GetParam());
int64_t num_datasets = 10;
std::vector<std::string> keys;
for (int i = 0; i < num_datasets; ++i) {
std::string key = absl::StrCat("key", i);
DatasetDef dataset_def = DatasetDefWithVersion(i);
TF_ASSERT_OK(store->Put(key, dataset_def));
keys.push_back(key);
}
for (int i = 0; i < num_datasets; ++i) {
std::shared_ptr<const DatasetDef> result;
TF_ASSERT_OK(store->Get(keys[i], result));
EXPECT_EQ(result->graph().version(), i);
}
}
TEST_P(DatasetStoreTest, StoreAlreadyExists) {
std::unique_ptr<DatasetStore> store = MakeStore(GetParam());
int32_t version = 1;
DatasetDef dataset_def = DatasetDefWithVersion(version);
std::string key = "key";
TF_ASSERT_OK(store->Put(key, dataset_def));
TF_EXPECT_OK(store->Put(key, dataset_def));
std::shared_ptr<const DatasetDef> result;
TF_ASSERT_OK(store->Get(key, result));
EXPECT_EQ(result->graph().version(), version);
}
TEST_P(DatasetStoreTest, GetMissing) {
std::unique_ptr<DatasetStore> store = MakeStore(GetParam());
std::shared_ptr<const DatasetDef> result;
absl::Status s = store->Get("missing", result);
EXPECT_EQ(s.code(), error::NOT_FOUND);
}
INSTANTIATE_TEST_SUITE_P(DatasetStoreTests, DatasetStoreTest,
::testing::Values(kFileSystem, kMemory));
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/dispatcher.proto
+418
View File
@@ -0,0 +1,418 @@
// Copyright 2026 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.
// ==============================================================================
syntax = "proto3";
package tensorflow.data;
import "tensorflow/core/data/service/common.proto";
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/protobuf/data_service.proto";
import "tensorflow/core/protobuf/snapshot.proto";
// Next tag: 3
message ActiveTask {
int64 task_id = 1;
// Estimated time it takes this Task to produce an element, in nanoseconds.
double processing_time_nsec = 2;
}
// Next tag: 9
message WorkerHeartbeatRequest {
string worker_address = 1;
repeated DataTransferServerInfo transfer_servers = 7;
repeated string worker_tags = 4;
// The UID of the worker Borg job, used for telemetry.
int64 worker_uid = 5;
repeated int64 current_tasks = 2;
// The status of any active snapshot tasks, keyed by snapshot path.
map<string, SnapshotTaskProgress> snapshot_task_progress = 6;
reserved 3;
// TODO(armandouv): Deprecate current_tasks and extract task ids from here.
repeated ActiveTask active_tasks = 8;
}
// Next tag: 4
message WorkerHeartbeatResponse {
repeated TaskDef new_tasks = 1;
repeated int64 tasks_to_delete = 2;
// Snapshots to process.
repeated SnapshotTaskDef snapshot_tasks = 3;
}
// Next tag: 3
message TaskProgress {
// The task that this message is about.
int64 task_id = 1;
// Whether the task has completed.
bool completed = 2;
}
// Next tag: 3
message WorkerUpdateRequest {
string worker_address = 1;
repeated TaskProgress updates = 2;
}
// Next tag: 1
message WorkerUpdateResponse {}
// Next tag: 2
message GetDatasetDefRequest {
string dataset_id = 1;
}
// Next tag: 2
message GetDatasetDefResponse {
DatasetDef dataset_def = 1;
}
// Next tag: 4
message GetSplitRequest {
int64 iteration_id = 1;
int64 repetition = 2;
int64 split_provider_index = 3;
}
// Next tag: 3
message GetSplitResponse {
TensorProto split = 1;
bool end_of_splits = 2;
}
// Next tag: 1
message GetVersionRequest {}
// Next tag: 2
message GetVersionResponse {
int64 version = 1;
}
// Next tag: 5
message GetOrRegisterDatasetRequest {
// The dataset to register.
DatasetDef dataset = 1;
// Metadata related to tf.data service.
DataServiceMetadata metadata = 3;
oneof optional_dataset_id {
// If provided, tf.data service will register the dataset with the specified
// ID. Otherwise, it will generate a unique dataset ID.
string dataset_id = 4;
}
reserved 2;
}
// Next tag: 2
message GetOrRegisterDatasetResponse {
// The id for the registered dataset.
string dataset_id = 1;
}
// Next tag: 2
message GetDataServiceMetadataRequest {
// The dataset id to get the data service dataset metadata.
string dataset_id = 1;
}
// Next tag: 2
message GetDataServiceMetadataResponse {
// The retrieved data service dataset metadata.
DataServiceMetadata metadata = 1;
}
// Next tag: 1
message GetDataServiceConfigRequest {}
// Next tag: 2
message GetDataServiceConfigResponse {
DataServiceConfig config = 1;
}
// Next tag: 7
message GetOrCreateJobRequest {
// The id of the dataset to create a job for.
string dataset_id = 1;
// A mode controlling how the tf.data service produces data for the job.
ProcessingModeDef processing_mode_def = 2;
// Optional job name identifying a shared job. If not set, the RPC will always
// create a new job.
oneof optional_job_name {
string job_name = 3;
}
// Optional number of consumers. If set, the job's tasks will provide
// their elements to consumers round-robin.
oneof optional_num_consumers {
int64 num_consumers = 4;
}
// True if cross-trainer cache is enabled.
bool use_cross_trainer_cache = 5;
// Specifies which workers the client of this job reads from.
TargetWorkers target_workers = 6;
}
// Next tag: 2
message GetOrCreateJobResponse {
int64 job_id = 1;
}
// Next tag: 3
message GetOrCreateIterationRequest {
// The job to create an iteration for.
int64 job_id = 1;
// Which repetition of the job to read from.
int64 repetition = 2;
}
// Next tag: 2
message GetOrCreateIterationResponse {
// An id for the client that will read from the iteration. When the client is
// done with the iteration, they should call ReleaseIterationClient with this
// id.
int64 iteration_client_id = 1;
}
// Next tag: 4
message MaybeRemoveTaskRequest {
int64 task_id = 1;
int64 consumer_index = 2;
int64 round = 3;
}
// Next tag: 2
message MaybeRemoveTaskResponse {
bool removed = 1;
}
// Next tag: 2
message ReleaseIterationClientRequest {
int64 iteration_client_id = 1;
}
// Next tag: 1
message ReleaseIterationClientResponse {}
// Next tag: 6
message ClientHeartbeatRequest {
reserved 3;
// The iteration client id to heartbeat for.
int64 iteration_client_id = 1;
// Reports which round the client is currently reading from when doing
// round-robin reads.
oneof optional_current_round {
int64 current_round = 2;
}
// Reports whether the client has successfully blocked the indicated round
// from starting. This enables the dispatcher to add a new task in the
// blocked round or later.
oneof optional_blocked_round {
int64 blocked_round = 4;
}
// Target processing time in nanoseconds observed by the client.
double target_processing_time_nsec = 5;
}
// Next tag: 5
message ClientHeartbeatResponse {
// A list of all tasks that the client should read from.
repeated TaskInfo task_info = 1;
// Tells the client not to start the given round if possible.
oneof optional_block_round {
int64 block_round = 3;
}
// Whether the iteration has finished.
bool iteration_finished = 2;
// tf.data service deployment mode. Supported values are "REMOTE",
// "COLOCATED", and "HYBRID". If unspecified, it is assumed to be "REMOTE".
DeploymentMode deployment_mode = 4;
}
// Next tag: 3
message WorkerInfo {
string address = 1;
reserved 2;
}
// Next tag: 1
message GetWorkersRequest {}
// Next tag: 2
message GetWorkersResponse {
// A list of all workers.
repeated WorkerInfo workers = 1;
}
// Next tag: 4
message SnapshotRequest {
// The dataset to snapshot.
DatasetDef dataset = 1;
// The path to which to materialize the snapshot.
string path = 2;
// The metadata for the snapshot.
experimental.DistributedSnapshotMetadata metadata = 3;
}
// Next tag: 1
message SnapshotResponse {}
// Next tag: 6
message GetSnapshotSplitRequest {
// The address of the worker requesting the split.
string worker_address = 4;
// The base path of the snapshot materialization.
string base_path = 1;
// The index of the snapshot stream from which to get the split.
int64 stream_index = 2;
// The index of the dataset source from which to get the split.
int64 source_index = 3;
// The repetition of the dataset from which to get the split.
int64 repetition_index = 5;
}
// Next tag: 4
message GetSnapshotSplitResponse {
oneof response {
// The split to process.
TensorProto split = 1;
// If true, there are no splits left to be processed for this stream.
bool end_of_splits = 2;
}
// The local split index within the stream source, starting at zero and
// incrementing by one for each split assigned to the worker. This local index
// is used by the worker to keep track of which split it has read up to. If
// `end_of_splits` is true, this equals the total number of splits.
int64 local_split_index = 3;
}
// Next tag: 2
message GetSnapshotStreamsRequest {
// The path at which the snapshot is being materialized.
string path = 1;
}
// Next tag: 2
message GetSnapshotStreamsResponse {
// Information about all streams for the snapshot.
repeated SnapshotStreamInfo streams = 1;
}
// Next tag: 5
message DisableCompressionAtRuntimeRequest {
string dataset_id = 1;
// If `true`, compression should be disabled at runtime.
bool disable_compression_at_runtime = 4;
reserved 2, 3;
}
// Next tag: 4
message DisableCompressionAtRuntimeResponse {
oneof decision {
// If `true`, the dispatcher has decided to disable compression, and workers
// and trainers will ignore the compression specified at registration time.
// If `false`, the dispatcher has decided to not disable compression, and
// workers and trainers will adhere to the compression specified at
// registration time.
bool compression_disabled_at_runtime = 1;
// If `true`, compression was not specified at registration time, so there
// is no runtime compression disabling decision to make.
bool no_compression_to_disable = 3;
}
reserved 2;
}
service DispatcherService {
// Performs a periodic worker heartbeat.
rpc WorkerHeartbeat(WorkerHeartbeatRequest) returns (WorkerHeartbeatResponse);
// Updates the dispatcher with information about the worker's state.
rpc WorkerUpdate(WorkerUpdateRequest) returns (WorkerUpdateResponse);
// Gets a dataset definition.
rpc GetDatasetDef(GetDatasetDefRequest) returns (GetDatasetDefResponse);
// Gets the next split for a given iteration.
rpc GetSplit(GetSplitRequest) returns (GetSplitResponse);
// Returns the API version of the server.
rpc GetVersion(GetVersionRequest) returns (GetVersionResponse);
// Registers a dataset with the server, or returns its id if it is already
// registered.
//
// The dataset is constructed in a new graph, so it must not refer to
// external resources or variables.
rpc GetOrRegisterDataset(GetOrRegisterDatasetRequest)
returns (GetOrRegisterDatasetResponse);
// Gets a job if it already exists, otherwise creates it.
rpc GetOrCreateJob(GetOrCreateJobRequest) returns (GetOrCreateJobResponse);
// Gets an iteration if it already exists, otherwise creates it.
rpc GetOrCreateIteration(GetOrCreateIterationRequest)
returns (GetOrCreateIterationResponse);
// Attempts to remove a task from a round-robin read iteration.
rpc MaybeRemoveTask(MaybeRemoveTaskRequest) returns (MaybeRemoveTaskResponse);
// Releases an iteration client so that an iteration may eventually be cleaned
// up.
rpc ReleaseIterationClient(ReleaseIterationClientRequest)
returns (ReleaseIterationClientResponse);
// Heartbeats from the client. This lets the dispatcher know that the client
// is still active, and gives the dispatcher a chance to notify the client
// of new tasks.
rpc ClientHeartbeat(ClientHeartbeatRequest) returns (ClientHeartbeatResponse);
// Reports a list of all workers registered with the dispatcher.
rpc GetWorkers(GetWorkersRequest) returns (GetWorkersResponse);
// Returns the data service metadata for the registered dataset.
rpc GetDataServiceMetadata(GetDataServiceMetadataRequest)
returns (GetDataServiceMetadataResponse);
// Returns the config of a data service cluster.
rpc GetDataServiceConfig(GetDataServiceConfigRequest)
returns (GetDataServiceConfigResponse);
// Initiates the process of materializing a dataset's output to disk.
rpc Snapshot(SnapshotRequest) returns (SnapshotResponse);
// Gets the next split for the given stream of the given snapshot. Returns an
// error if there has been some miscommunication between the worker and
// dispatcher regarding stream assignment and the worker should stop (though
// due to stream leases this case should never happen).
rpc GetSnapshotSplit(GetSnapshotSplitRequest)
returns (GetSnapshotSplitResponse);
// Returns information about all streams for the given snapshot.
rpc GetSnapshotStreams(GetSnapshotStreamsRequest)
returns (GetSnapshotStreamsResponse);
// Returns information about the decision to disable compression at runtime
// for the given dataset.
rpc DisableCompressionAtRuntime(DisableCompressionAtRuntimeRequest)
returns (DisableCompressionAtRuntimeResponse);
}
tensorflow/core/data/service/dispatcher_client.cc
+388
View File
@@ -0,0 +1,388 @@
/* Copyright 2021 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/core/data/service/dispatcher_client.h"
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <vector>
#include "grpcpp/client_context.h"
#include "grpcpp/create_channel.h"
#include "grpcpp/security/credentials.h"
#include "grpcpp/support/channel_arguments.h"
#include "grpcpp/support/status.h"
#include "absl/strings/str_cat.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/credentials_factory.h"
#include "tensorflow/core/data/service/dispatcher.grpc.pb.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/grpc_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
absl::Status DataServiceDispatcherClient::Initialize() {
mutex_lock l(mu_);
if (stub_) {
return absl::OkStatus();
}
std::shared_ptr<grpc::ChannelCredentials> credentials;
TF_RETURN_IF_ERROR(
CredentialsFactory::CreateClientCredentials(protocol_, &credentials));
grpc::ChannelArguments args;
args.SetMaxReceiveMessageSize(std::numeric_limits<int32_t>::max());
args.SetInt(GRPC_ARG_USE_LOCAL_SUBCHANNEL_POOL, true);
auto channel = grpc::CreateCustomChannel(address_, credentials, args);
stub_ = DispatcherService::NewStub(channel);
GetVersionRequest req;
GetVersionResponse resp;
grpc::ClientContext ctx;
grpc::Status s = stub_->GetVersion(&ctx, req, &resp);
if (!s.ok()) {
return grpc_util::WrapError(
absl::StrCat("Failed to get dispatcher version from dispatcher "
"running at ",
address_),
s);
}
if (resp.version() != kDataServiceVersion) {
return absl::FailedPreconditionError(absl::StrCat(
"Version mismatch with tf.data service server. The server is running "
"version ",
resp.version(), ", while the client is running version ",
kDataServiceVersion,
". Please ensure that the client and server side are running the "
"same version of TensorFlow. If you're running an MPM binary, make "
"sure the server is running an up-to-date MPM."));
}
return absl::OkStatus();
}
absl::StatusOr<WorkerHeartbeatResponse>
DataServiceDispatcherClient::WorkerHeartbeat(
const WorkerHeartbeatRequest& request) {
WorkerHeartbeatResponse response;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->WorkerHeartbeat(&client_ctx, request, &response);
if (!status.ok()) {
return grpc_util::WrapError("Failed to perform worker heartbeat", status);
}
return response;
}
absl::Status DataServiceDispatcherClient::WorkerUpdate(
const std::string& worker_address,
std::vector<TaskProgress>& task_progress) {
WorkerUpdateRequest req;
req.set_worker_address(worker_address);
for (const auto& update : task_progress) {
*(req.add_updates()) = update;
}
WorkerUpdateResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->WorkerUpdate(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to send worker update", status);
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetDatasetDef(
const std::string& dataset_id, DatasetDef& dataset_def) {
GetDatasetDefRequest req;
req.set_dataset_id(dataset_id);
GetDatasetDefResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->GetDatasetDef(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to get dataset def", status);
}
dataset_def = resp.dataset_def();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetSplit(int64_t iteration_id,
int64_t repetition,
int64_t split_provider_index,
Tensor& split,
bool& end_of_splits) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetSplitRequest req;
req.set_iteration_id(iteration_id);
req.set_repetition(repetition);
req.set_split_provider_index(split_provider_index);
GetSplitResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->GetSplit(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to get split", status);
}
end_of_splits = resp.end_of_splits();
if (!end_of_splits) {
if (!split.FromProto(resp.split())) {
return absl::InternalError("Failed to parse split tensor proto");
}
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::Snapshot(
const DatasetDef& dataset, const std::string& path,
const experimental::DistributedSnapshotMetadata& metadata) {
TF_RETURN_IF_ERROR(EnsureInitialized());
SnapshotRequest req;
*req.mutable_dataset() = dataset;
req.set_path(path);
*req.mutable_metadata() = metadata;
SnapshotResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->Snapshot(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to snapshot", status);
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetSnapshotSplit(
const std::string& worker_address, const std::string& base_path,
int64_t stream_index, int64_t source_index, int64_t repetition_index,
Tensor& split, int64_t& local_split_index, bool& end_of_splits) {
GetSnapshotSplitRequest req;
req.set_worker_address(worker_address);
req.set_base_path(base_path);
req.set_stream_index(stream_index);
req.set_source_index(source_index);
req.set_repetition_index(repetition_index);
GetSnapshotSplitResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->GetSnapshotSplit(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to get snapshot split", status);
}
local_split_index = resp.local_split_index();
end_of_splits = resp.end_of_splits();
if (end_of_splits) {
return absl::OkStatus();
}
if (!split.FromProto(resp.split())) {
return absl::InternalError(absl::StrCat(
"Failed to parse split tensor proto: ", resp.split().DebugString()));
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::RegisterDataset(
const DatasetDef& dataset, const DataServiceMetadata& metadata,
const std::optional<std::string>& requested_dataset_id,
std::string& dataset_id) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetOrRegisterDatasetRequest req;
*req.mutable_dataset() = dataset;
*req.mutable_metadata() = metadata;
if (requested_dataset_id.has_value()) {
req.set_dataset_id(*requested_dataset_id);
}
GetOrRegisterDatasetResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->GetOrRegisterDataset(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to register dataset", status);
}
dataset_id = resp.dataset_id();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetOrCreateJob(
const std::string& dataset_id, const ProcessingModeDef& processing_mode,
const std::optional<std::string>& job_name,
std::optional<int64_t> num_consumers, bool use_cross_trainer_cache,
TargetWorkers target_workers, int64_t& job_id) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetOrCreateJobRequest req;
req.set_dataset_id(dataset_id);
*req.mutable_processing_mode_def() = processing_mode;
if (job_name.has_value()) {
req.set_job_name(job_name.value());
}
if (num_consumers.has_value()) {
req.set_num_consumers(num_consumers.value());
}
req.set_target_workers(target_workers);
req.set_use_cross_trainer_cache(use_cross_trainer_cache);
GetOrCreateJobResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->GetOrCreateJob(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError(
absl::StrCat("Failed to get or create job for dataset with id ",
dataset_id),
status);
}
job_id = resp.job_id();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetOrCreateIteration(
int64_t job_id, int64_t repetition, int64_t& iteration_client_id) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetOrCreateIterationRequest req;
req.set_job_id(job_id);
req.set_repetition(repetition);
GetOrCreateIterationResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->GetOrCreateIteration(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError(
absl::StrCat("Failed to get or create iteration for job with id ",
job_id),
status);
}
iteration_client_id = resp.iteration_client_id();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::ReleaseIterationClient(
int64_t iteration_client_id) {
TF_RETURN_IF_ERROR(EnsureInitialized());
ReleaseIterationClientRequest req;
req.set_iteration_client_id(iteration_client_id);
ReleaseIterationClientResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->ReleaseIterationClient(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError(
absl::StrCat("Failed to release iteration client with id ",
iteration_client_id),
status);
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::MaybeRemoveTask(
int64_t task_id, int64_t consumer_index, int64_t round, bool& removed) {
TF_RETURN_IF_ERROR(EnsureInitialized());
MaybeRemoveTaskRequest req;
req.set_task_id(task_id);
req.set_consumer_index(consumer_index);
req.set_round(round);
MaybeRemoveTaskResponse resp;
grpc::ClientContext client_ctx;
grpc::Status status = stub_->MaybeRemoveTask(&client_ctx, req, &resp);
if (!status.ok()) {
return grpc_util::WrapError("Failed to call MaybeRemoveTask", status);
}
removed = resp.removed();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::ClientHeartbeat(
ClientHeartbeatRequest& req, ClientHeartbeatResponse& resp) {
TF_RETURN_IF_ERROR(EnsureInitialized());
grpc::ClientContext ctx;
grpc::Status s = stub_->ClientHeartbeat(&ctx, req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get tasks", s);
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetWorkers(
std::vector<WorkerInfo>& workers) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetWorkersRequest req;
GetWorkersResponse resp;
grpc::ClientContext ctx;
grpc::Status s = stub_->GetWorkers(&ctx, req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get workers", s);
}
workers.clear();
for (auto& worker : resp.workers()) {
workers.push_back(worker);
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetDataServiceMetadata(
const std::string& dataset_id, DataServiceMetadata& metadata) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetDataServiceMetadataRequest req;
req.set_dataset_id(dataset_id);
GetDataServiceMetadataResponse resp;
grpc::ClientContext ctx;
grpc::Status s = stub_->GetDataServiceMetadata(&ctx, req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get data service metadata", s);
}
metadata = resp.metadata();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::GetDataServiceConfig(
DataServiceConfig& config) {
TF_RETURN_IF_ERROR(EnsureInitialized());
GetDataServiceConfigRequest request;
GetDataServiceConfigResponse response;
grpc::ClientContext ctx;
grpc::Status s = stub_->GetDataServiceConfig(&ctx, request, &response);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get data service config", s);
}
config = response.config();
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::DisableCompressionAtRuntime(
const std::string& dataset_id, bool disable_compression_at_runtime,
DisableCompressionAtRuntimeResponse& response) {
TF_RETURN_IF_ERROR(EnsureInitialized());
grpc::ClientContext ctx;
DisableCompressionAtRuntimeRequest request;
request.set_dataset_id(dataset_id);
request.set_disable_compression_at_runtime(disable_compression_at_runtime);
grpc::Status s = stub_->DisableCompressionAtRuntime(&ctx, request, &response);
if (!s.ok()) {
return grpc_util::WrapError(
"Failed to get runtime compression disabling decision", s);
}
return absl::OkStatus();
}
absl::Status DataServiceDispatcherClient::EnsureInitialized() {
return grpc_util::Retry(
[this] { return Initialize(); }, "Initialize dispatcher client",
/*deadline_micros=*/std::numeric_limits<int64_t>::max());
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/dispatcher_client.h
+153
View File
@@ -0,0 +1,153 @@
/* Copyright 2021 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_CORE_DATA_SERVICE_DISPATCHER_CLIENT_H_
#define TENSORFLOW_CORE_DATA_SERVICE_DISPATCHER_CLIENT_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dispatcher.grpc.pb.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
#include "tensorflow/core/protobuf/snapshot.pb.h"
namespace tensorflow {
namespace data {
// Client for communicating with the tf.data service dispatcher.
class DataServiceDispatcherClient : public DataServiceClientBase {
public:
DataServiceDispatcherClient(const std::string& address,
const std::string& protocol)
: DataServiceClientBase(address, protocol) {}
absl::Status Initialize() override;
// Sends a heartbeat to the dispatcher. If the worker wasn't already
// registered with the dispatcher, this will register the worker. The
// dispatcher will report which new tasks the worker should run, and which
// tasks it should delete.
absl::StatusOr<WorkerHeartbeatResponse> WorkerHeartbeat(
const WorkerHeartbeatRequest& request);
// Updates the dispatcher with information about the worker's state.
absl::Status WorkerUpdate(const std::string& worker_address,
std::vector<TaskProgress>& task_progress);
// Gets a dataset definition for the given dataset id, and stores the
// definition in `dataset_def`.
absl::Status GetDatasetDef(const std::string& dataset_id,
DatasetDef& dataset_def);
// Gets the next split for the specified iteration id, repetition, and split
// provider index.
absl::Status GetSplit(int64_t iteration_id, int64_t repetition,
int64_t split_provider_index, Tensor& split,
bool& end_of_splits);
// Gets the next split for the specified source of a stream of the snapshot in
// `base_path`. If `end_of_splits` returns true, then there are no more splits
// to be processed for the specified stream source.
virtual absl::Status GetSnapshotSplit(
const std::string& worker_address, const std::string& base_path,
int64_t stream_index, int64_t source_index, int64_t repetition_index,
Tensor& split, int64_t& local_split_index, bool& end_of_splits);
// Initiates the process of materializing `dataset`'s output to `path`.
absl::Status Snapshot(
const DatasetDef& dataset, const std::string& path,
const experimental::DistributedSnapshotMetadata& metadata);
// Registers a dataset with the tf.data service, and stores the generated
// dataset id in `dataset_id`.
absl::Status RegisterDataset(
const DatasetDef& dataset, const DataServiceMetadata& metadata,
const std::optional<std::string>& requested_dataset_id,
std::string& dataset_id);
// If `job_name` is set, looks up a job matching `job_name`.
// If `job_name` is absent or no matching job is found, creates a
// new job. The resulting job id is stored in `job_id`.
absl::Status GetOrCreateJob(const std::string& dataset_id,
const ProcessingModeDef& processing_mode,
const std::optional<std::string>& job_name,
std::optional<int64_t> num_consumers,
bool use_cross_trainer_cache,
TargetWorkers target_workers, int64_t& job_id);
// Looks up an iteration of a job, creating an iteration if one doesn't
// already exist. The returned `iteration_client_id` can be used to query
// information about the iteration. The client should call
// `ReleaseIterationClient` when finished with the iteration, so that
// resources can be reclaimed.
absl::Status GetOrCreateIteration(int64_t job_id, int64_t repetition,
int64_t& iteration_client_id);
// Releases a iteration client id, indicating that the id will no longer be
// used to read from the iteration.
absl::Status ReleaseIterationClient(int64_t iteration_client_id);
// Attempts to remove a task. The task is removed if all consumers try to
// remove the task in the same round.
absl::Status MaybeRemoveTask(int64_t task_id, int64_t consumer_index,
int64_t round, bool& removed);
// Heartbeats to the dispatcher, getting back the tasks that should be
// running, and whether the iteration is finished.
absl::Status ClientHeartbeat(ClientHeartbeatRequest& req,
ClientHeartbeatResponse& resp);
// Queries the dispatcher for its registered workers. The worker info will be
// stored in `workers`.
absl::Status GetWorkers(std::vector<WorkerInfo>& workers);
// Returns data service metadata for the registered dataset.
absl::Status GetDataServiceMetadata(const std::string& dataset_id,
DataServiceMetadata& metadata);
// Returns data service config of the data service cluster.
absl::Status GetDataServiceConfig(DataServiceConfig& config);
// Returns information about the decision to disable compression at runtime
// for a given dataset.
absl::Status DisableCompressionAtRuntime(
const std::string& dataset_id, bool disable_compression_at_runtime,
DisableCompressionAtRuntimeResponse& response);
protected:
absl::Status EnsureInitialized() override;
private:
mutex mu_;
// Initialization is guarded by `mu_`, but using the stub does not require
// holding `mu_`
std::unique_ptr<DispatcherService::Stub> stub_;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_DISPATCHER_CLIENT_H_
tensorflow/core/data/service/dispatcher_client_test.cc
+439
View File
@@ -0,0 +1,439 @@
/* Copyright 2021 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/core/data/service/dispatcher_client.h"
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <optional>
#include <string>
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/test.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/data_transfer.h"
#include "tensorflow/core/data/service/dataset_store.h"
#include "tensorflow/core/data/service/snapshot/path_utils.h"
#include "tensorflow/core/data/service/test_cluster.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/protobuf/snapshot.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
#include "tsl/platform/path.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::data::experimental::DistributedSnapshotMetadata;
using ::tensorflow::data::testing::CreateDummyDistributedSnapshotMetadata;
using ::tensorflow::data::testing::EqualsProto;
using ::tensorflow::data::testing::InfiniteDataset;
using ::tensorflow::data::testing::LocalTempFilename;
using ::tensorflow::data::testing::RangeDataset;
using ::tensorflow::testing::StatusIs;
using ::testing::AllOf;
using ::testing::HasSubstr;
constexpr const char kProtocol[] = "grpc";
DataServiceMetadata GetDefaultMetadata() {
StructuredValue decoded_spec;
TensorShapeProto::Dim* dim =
decoded_spec.mutable_tensor_shape_value()->add_dim();
dim->set_size(1);
dim->set_name(absl::StrCat("dim"));
DataServiceMetadata metadata;
metadata.set_element_spec(decoded_spec.SerializeAsString());
metadata.set_compression(DataServiceMetadata::COMPRESSION_SNAPPY);
metadata.set_cardinality(kUnknownCardinality);
return metadata;
}
class DispatcherClientTest : public ::testing::Test {
protected:
absl::Status SetUpTfDataService(int64_t num_workers,
int64_t worker_max_concurrent_snapshots = 0) {
TestCluster::Config config;
config.num_workers = num_workers;
config.work_dir = tsl::io::JoinPath(tsl::testing::TmpDir(), "work_dir");
config.worker_max_concurrent_snapshots = worker_max_concurrent_snapshots;
test_cluster_ = std::make_unique<TestCluster>(config);
TF_RETURN_IF_ERROR(test_cluster_->Initialize());
dispatcher_client_ = std::make_unique<DataServiceDispatcherClient>(
test_cluster_->DispatcherAddress(), kProtocol);
return absl::OkStatus();
}
// Creates a dataset and returns the dataset ID.
absl::StatusOr<std::string> RegisterDataset(
const DatasetDef& dataset, const DataServiceMetadata& metadata,
const std::optional<std::string>& requested_dataset_id = std::nullopt) {
std::string dataset_id;
TF_RETURN_IF_ERROR(dispatcher_client_->RegisterDataset(
dataset, metadata, requested_dataset_id, dataset_id));
return dataset_id;
}
// Starts snapshots and returns the directories.
absl::StatusOr<absl::flat_hash_set<std::string>> StartDummySnapshots(
int64_t num_snapshots) {
DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
// Create a set of local file paths to which snapshots will be materialized.
absl::flat_hash_set<std::string> directories;
for (int64_t i = 0; i < num_snapshots; ++i) {
directories.insert(LocalTempFilename());
}
for (const auto& directory : directories) {
TF_RETURN_IF_ERROR(
dispatcher_client_->Snapshot(RangeDataset(10), directory, metadata));
}
return directories;
}
std::unique_ptr<TestCluster> test_cluster_;
std::unique_ptr<DataServiceDispatcherClient> dispatcher_client_;
};
TEST_F(DispatcherClientTest, GetDataServiceMetadata) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
metadata.set_cardinality(10);
TF_ASSERT_OK_AND_ASSIGN(const std::string dataset_id,
RegisterDataset(RangeDataset(10), metadata));
DataServiceMetadata result;
TF_ASSERT_OK(dispatcher_client_->GetDataServiceMetadata(dataset_id, result));
EXPECT_THAT(result, EqualsProto(metadata));
}
TEST_F(DispatcherClientTest, DatasetDoesNotExist) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
EXPECT_THAT(
dispatcher_client_->GetDataServiceMetadata(
/*dataset_id=*/"not-found", metadata),
absl_testing::StatusIs(error::NOT_FOUND,
HasSubstr("Dataset id not-found not found")));
}
TEST_F(DispatcherClientTest, SnapshotAlreadyStarted) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
std::string directory = LocalTempFilename();
TF_ASSERT_OK(
dispatcher_client_->Snapshot(RangeDataset(10), directory, metadata));
EXPECT_THAT(
dispatcher_client_->Snapshot(RangeDataset(10), directory, metadata),
absl_testing::StatusIs(error::ALREADY_EXISTS,
HasSubstr("already started")));
}
TEST_F(DispatcherClientTest, GetDataServiceConfig) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceConfig config;
TF_ASSERT_OK(dispatcher_client_->GetDataServiceConfig(config));
EXPECT_EQ(config.deployment_mode(), DEPLOYMENT_MODE_COLOCATED);
}
TEST_F(DispatcherClientTest, SnapshotSkeletonWritten) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
TF_ASSERT_OK_AND_ASSIGN(absl::flat_hash_set<std::string> paths,
StartDummySnapshots(/*num_snapshots=*/3));
for (const auto& path : paths) {
TF_ASSERT_OK(Env::Default()->FileExists(CommittedChunksDirectory(path)));
TF_ASSERT_OK(Env::Default()->FileExists(StreamsDirectory(path)));
}
}
TEST_F(DispatcherClientTest, SnapshotMetadataAndDatasetDefWritten) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
TF_ASSERT_OK_AND_ASSIGN(absl::flat_hash_set<std::string> paths,
StartDummySnapshots(/*num_snapshots=*/3));
for (const auto& path : paths) {
TF_ASSERT_OK(
Env::Default()->FileExists(io::JoinPath(path, "snapshot.metadata")));
TF_ASSERT_OK(
Env::Default()->FileExists(io::JoinPath(path, "dataset_def.proto")));
}
}
TEST_F(DispatcherClientTest, SnapshotsInHeartbeat) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1,
/*worker_max_concurrent_snapshots=*/3));
TF_ASSERT_OK_AND_ASSIGN(absl::flat_hash_set<std::string> paths,
StartDummySnapshots(/*num_snapshots=*/3));
WorkerHeartbeatRequest worker_heartbeat_request;
worker_heartbeat_request.set_worker_address(test_cluster_->WorkerAddress(0));
for (int64_t i = 1; i <= 3; ++i) {
TF_ASSERT_OK_AND_ASSIGN(
WorkerHeartbeatResponse worker_heartbeat_response,
dispatcher_client_->WorkerHeartbeat(worker_heartbeat_request));
ASSERT_EQ(worker_heartbeat_response.snapshot_tasks_size(), i);
for (const auto& snapshot_task :
worker_heartbeat_response.snapshot_tasks()) {
ASSERT_TRUE(paths.count(snapshot_task.base_path()));
ASSERT_EQ(snapshot_task.stream_index(), 0);
}
}
}
TEST_F(DispatcherClientTest, GetSnapshotSplit) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
TF_ASSERT_OK_AND_ASSIGN(absl::flat_hash_set<std::string> paths,
StartDummySnapshots(/*num_snapshots=*/3));
WorkerHeartbeatRequest worker_heartbeat_request;
worker_heartbeat_request.set_worker_address(test_cluster_->WorkerAddress(0));
TF_ASSERT_OK_AND_ASSIGN(
WorkerHeartbeatResponse worker_heartbeat_response,
dispatcher_client_->WorkerHeartbeat(worker_heartbeat_request));
for (int64_t i = 0; i < 5; ++i) {
for (const auto& snapshot_task :
worker_heartbeat_response.snapshot_tasks()) {
GetSnapshotSplitRequest get_snapshot_split_request;
Tensor split;
int64_t local_split_index = 0;
bool end_of_splits = false;
TF_ASSERT_OK(dispatcher_client_->GetSnapshotSplit(
test_cluster_->WorkerAddress(0), snapshot_task.base_path(),
snapshot_task.stream_index(),
/*source_index=*/0, /*repetition_index=*/0, split, local_split_index,
end_of_splits));
EXPECT_EQ(local_split_index, i);
EXPECT_FALSE(end_of_splits);
}
}
}
TEST_F(DispatcherClientTest, GetSnapshotSplitMultipleStreams) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/3,
/*worker_max_concurrent_snapshots=*/1));
TF_ASSERT_OK_AND_ASSIGN(absl::flat_hash_set<std::string> paths,
StartDummySnapshots(/*num_snapshots=*/3));
absl::flat_hash_set<std::string> snapshots_in_progress;
for (int64_t i = 0; i < 3; ++i) {
WorkerHeartbeatRequest worker_heartbeat_request;
worker_heartbeat_request.set_worker_address(
test_cluster_->WorkerAddress(i));
TF_ASSERT_OK_AND_ASSIGN(
WorkerHeartbeatResponse worker_heartbeat_response,
dispatcher_client_->WorkerHeartbeat(worker_heartbeat_request));
EXPECT_EQ(worker_heartbeat_response.snapshot_tasks().size(), 1);
for (const auto& snapshot_task :
worker_heartbeat_response.snapshot_tasks()) {
snapshots_in_progress.insert(snapshot_task.base_path());
GetSnapshotSplitRequest get_snapshot_split_request;
Tensor split;
int64_t local_split_index = 0;
bool end_of_splits = false;
TF_ASSERT_OK(dispatcher_client_->GetSnapshotSplit(
test_cluster_->WorkerAddress(i), snapshot_task.base_path(),
snapshot_task.stream_index(),
/*source_index=*/0, /*repetition_index=*/0, split, local_split_index,
end_of_splits));
EXPECT_EQ(local_split_index, 0);
EXPECT_FALSE(end_of_splits);
}
}
// Each worker writes one snapshot; each snapshot has been assigned a worker.
EXPECT_EQ(snapshots_in_progress, paths);
}
TEST_F(DispatcherClientTest, RegisterDatasetWithExplicitId) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
metadata.set_cardinality(10);
TF_ASSERT_OK_AND_ASSIGN(
const std::string dataset_id1,
RegisterDataset(RangeDataset(10), metadata,
/*requested_dataset_id=*/"dataset_id"));
EXPECT_EQ(dataset_id1, "dataset_id");
// Registers a dataset with the same dataset ID.
TF_ASSERT_OK_AND_ASSIGN(
const std::string dataset_id2,
RegisterDataset(RangeDataset(10), metadata,
/*requested_dataset_id=*/"dataset_id"));
EXPECT_EQ(dataset_id1, dataset_id2);
}
TEST_F(DispatcherClientTest, DatasetsDoNotMatch) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
metadata.set_cardinality(10);
TF_ASSERT_OK_AND_ASSIGN(
const std::string dataset_id1,
RegisterDataset(RangeDataset(10), metadata,
/*requested_dataset_id=*/"dataset_id"));
EXPECT_EQ(dataset_id1, "dataset_id");
// Registers a dataset with the same dataset ID but different metadata.
metadata.set_cardinality(kInfiniteCardinality);
EXPECT_THAT(
RegisterDataset(InfiniteDataset(), metadata,
/*requested_dataset_id=*/"dataset_id"),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
HasSubstr(
"Datasets with the same ID should have the same structure")));
}
TEST_F(DispatcherClientTest, EnableCrossTrainerCache) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
metadata.set_cardinality(kInfiniteCardinality);
TF_ASSERT_OK_AND_ASSIGN(const std::string dataset_id,
RegisterDataset(InfiniteDataset(), metadata));
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
std::string job_name = "job";
int64_t job_id;
TF_ASSERT_OK(dispatcher_client_->GetOrCreateJob(
dataset_id, processing_mode, job_name,
/*num_consumers=*/std::nullopt,
/*use_cross_trainer_cache=*/true, TARGET_WORKERS_AUTO, job_id));
int64_t iteration_client_id;
TF_ASSERT_OK(dispatcher_client_->GetOrCreateIteration(
job_id, /*repetition=*/0, iteration_client_id));
WorkerHeartbeatRequest worker_heartbeat_request;
worker_heartbeat_request.set_worker_address(test_cluster_->WorkerAddress(0));
TF_ASSERT_OK_AND_ASSIGN(
WorkerHeartbeatResponse worker_heartbeat_response,
dispatcher_client_->WorkerHeartbeat(worker_heartbeat_request));
ASSERT_EQ(worker_heartbeat_response.new_tasks_size(), 1);
EXPECT_TRUE(worker_heartbeat_response.new_tasks(0).use_cross_trainer_cache());
}
TEST_F(DispatcherClientTest, CreateNamedJob) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
metadata.set_cardinality(10);
TF_ASSERT_OK_AND_ASSIGN(const std::string dataset_id,
RegisterDataset(RangeDataset(10), metadata));
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
std::string job_name = "job";
int64_t job_id_1 = -1;
TF_ASSERT_OK(dispatcher_client_->GetOrCreateJob(
dataset_id, processing_mode, job_name,
/*num_consumers=*/std::nullopt,
/*use_cross_trainer_cache=*/true, TARGET_WORKERS_AUTO, job_id_1));
int64_t job_id_2 = -2;
// Creating the same job should succeed and receive the same job id.
TF_ASSERT_OK(dispatcher_client_->GetOrCreateJob(
dataset_id, processing_mode, job_name,
/*num_consumers=*/std::nullopt,
/*use_cross_trainer_cache=*/true, TARGET_WORKERS_AUTO, job_id_2));
ASSERT_EQ(job_id_1, job_id_2);
}
TEST_F(DispatcherClientTest, NamedJobsDoNotMatch) {
TF_ASSERT_OK(SetUpTfDataService(/*num_workers=*/1));
DataServiceMetadata metadata = GetDefaultMetadata();
metadata.set_cardinality(10);
TF_ASSERT_OK_AND_ASSIGN(const std::string dataset_id,
RegisterDataset(RangeDataset(10), metadata));
int64_t job_id = 0;
ProcessingModeDef processing_mode;
processing_mode.set_sharding_policy(ProcessingModeDef::OFF);
std::string job_name = "job";
TF_ASSERT_OK(dispatcher_client_->GetOrCreateJob(
dataset_id, processing_mode, job_name,
/*num_consumers=*/std::nullopt,
/*use_cross_trainer_cache=*/false, TARGET_WORKERS_AUTO, job_id));
// Creating the same iteration with a different argument should fail.
processing_mode.set_sharding_policy(ProcessingModeDef::DYNAMIC);
EXPECT_THAT(
dispatcher_client_->GetOrCreateJob(dataset_id, processing_mode, job_name,
/*num_consumers=*/std::nullopt,
/*use_cross_trainer_cache=*/true,
TARGET_WORKERS_AUTO, job_id),
absl_testing::StatusIs(
error::INVALID_ARGUMENT,
AllOf(HasSubstr("but found an existing job with different "
"parameters: "),
HasSubstr("Existing processing mode: <"),
HasSubstr("Existing cross-trainer cache: <disabled>"))));
}
class DispatcherClientTest_DatasetId
: public DispatcherClientTest,
public ::testing::WithParamInterface<std::optional<std::string>> {};
TEST_P(DispatcherClientTest_DatasetId, SyncDatasetStoreWithDispatcherState) {
TestCluster::Config config;
config.num_workers = 1;
config.work_dir = tsl::io::JoinPath(tsl::testing::TmpDir(), "work_dir");
test_cluster_ = std::make_unique<TestCluster>(config);
TF_ASSERT_OK(test_cluster_->Initialize());
dispatcher_client_ = std::make_unique<DataServiceDispatcherClient>(
test_cluster_->DispatcherAddress(), kProtocol);
DatasetDef dataset_def = RangeDataset(10);
std::optional<std::string> requested_dataset_id = GetParam();
std::string dataset_id;
TF_ASSERT_OK(dispatcher_client_->RegisterDataset(
dataset_def, GetDefaultMetadata(),
/*requested_dataset_id=*/std::nullopt, dataset_id));
EXPECT_EQ(dataset_id, "1000");
// Writes an inconsistent dataset file. It should be discarded when the user
// registers a new dataset.
std::string datasets_dir = tsl::io::JoinPath(config.work_dir, "datasets");
FileSystemDatasetStore dataset_store(datasets_dir);
TF_ASSERT_OK(dataset_store.Put("1001", dataset_def));
if (requested_dataset_id.has_value()) {
TF_ASSERT_OK(dataset_store.Put(*requested_dataset_id, dataset_def));
}
TF_ASSERT_OK(dispatcher_client_->RegisterDataset(
dataset_def, GetDefaultMetadata(),
/*requested_dataset_id=*/requested_dataset_id, dataset_id));
if (requested_dataset_id.has_value()) {
EXPECT_EQ(dataset_id, *requested_dataset_id);
} else {
EXPECT_EQ(dataset_id, "1001");
}
}
INSTANTIATE_TEST_SUITE_P(DatasetId, DispatcherClientTest_DatasetId,
::testing::Values(std::nullopt, "dataset_id"));
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/dispatcher_impl.cc
+1643
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/data/service/dispatcher_impl.h
+412
View File
@@ -0,0 +1,412 @@
/* Copyright 2020 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_CORE_DATA_SERVICE_DISPATCHER_IMPL_H_
#define TENSORFLOW_CORE_DATA_SERVICE_DISPATCHER_IMPL_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/time/time.h"
#include "tensorflow/core/data/service/auto_scaler.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dataset_store.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/dispatcher_state.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/data/service/snapshot/snapshot_manager.h"
#include "tensorflow/core/data/service/task_remover.h"
#include "tensorflow/core/data/service/worker.grpc.pb.h"
#include "tensorflow/core/framework/dataset.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
// A service which coordinates a pool of workers to serve dataset elements over
// RPC.
//
// Glossary:
// * Dataset: A definition of how to generate a potentially large collection of
// elements.
// * Iteration: A coordinated phase of reading from the tf.data service. An
// iteration produces some amount of data, and (potentially multiple)
// consumers consume the data from the iteration until there is no data left.
// Each iteration has a ProcessingModeDef which determines what data it
// produces.
// * Task: An iteration is broken into multiple tasks, which each represent
// iterating over all of or part of the dataset. Workers process tasks.
// * Consumer: A process reading from the tf.data service.
//
// **Adding workers**
//
// tf.data service supports adding workers mid-iteration. When a new worker
// connects to the dispatcher, the dispatcher creates a new task for the worker,
// one task for each outstanding iteration. Consumers periodically heartbeat to
// the dispatcher to learn about new tasks.
//
// For non-round-robin-reads, there is no coordination among consumers. Each
// consumer will start reading from the new task as soon as it learns about the
// task from its heartbeat. Round robin reads, on the other hand, require
// consumers to read from the same task at each step. This requires coordination
// to ensure that all consumers start reading from the new task in the same
// round.
//
// The protocol for adding round robin tasks works as follows:
//
// - The dispatcher keeps track of which round each round-robin iteration is on.
// This
// information is reported by consumers in their heartbeats.
// - When a new worker joins and there is an outstanding round-robin iteration,
// we create a new task for the iteration and assign it to the worker.
// However, we don't yet report the task in consumer heartbeats.
// We call the task a "pending task" and add it to its iteration's "pending
// tasks" queue.
// - When we create a pending task, we choose a "target round" to try adding
// the task to. The target round is chosen by adding a "target round delta" to
// the latest reported round for the iteration.
// - When a consumer heartbeats for an iteration and there is a pending task for
// that iteration, the dispatcher sends a heartbeat response telling the
// consumer to block before reading from the target round.
// - When a consumer receives a heartbeat response telling it to block
// (before reading) a round, the consumer try to block the round. If the
// consumer has already started the round, it will too late to block the
// round.
// - When consumers heartbeat, they tell the dispatcher their current round and
// whether they have blocked themselves from reading past a certain round. If
// a consumer reports a current round exceeding the target round, the target
// round has failed and needs to be increased. We choose a new target round by
// doubling the previous target round delta. If the consumer reports that it
// has blocked before the target round, we record that the consumer is ready
// to add the new task. Once all consumers are ready to add the new task, we
// remove the task from the pending tasks list and begin reporting the task to
// consumers. We set the "starting_round" field of the task to indicate the
// target round where all consumers should start reading from the task.
// - If a new worker joins while there are already pending tasks, a pending
// task for the new worker is created and queued behind the existing tasks.
// The new task won't be considered until all previous pending tasks have been
// successfully added.
//
// An example of executing this protocol with two consumers could go as follows:
// 1. Consumers read up to round 50 and heartbeat that they are on round 50.
// 2. A new worker joins. Dispatcher chooses round 51 as the target round.
// 3. Consumer 1 heartbeats that its current round is 50. Dispatcher tells it to
// block round 51.
// 4. Consumer 2 heartbeats that its current round is 51. Dispatcher realizes
// that it is too late to block round 51 and chooses round 53 as the new
// target round. Dispatcher tells consumer 2 to block round 53.
// 5. Consumer 1 heartbeats that its current round is 50 and that it has blocked
// round 51. Dispatcher tells it to block round 53 instead. Dispatcher
// records that consumer 1 is ready to add a task in round 53.
// 6. Consumer 2 heartbeats that its current round is 52 and it has blocked
// round 53. Dispatcher realizes that all consumers are blocked on round 53
// or earlier and promotes the task from pending to regular. Dispatcher sends
// consumer 2 a task list containing the new task, and tells consumer 2 that
// it no longer needs to block.
// 7. Consumer 1 heartbeats. Dispatcher sends consumer 1 the task list
// containing the new task, and tells it that it no longer needs to block.
//
class DataServiceDispatcherImpl {
public:
explicit DataServiceDispatcherImpl(
const experimental::DispatcherConfig& config);
~DataServiceDispatcherImpl();
// Starts the dispatcher. If there is a journal, this will read from the
// journal to restore the dispatcher's state.
absl::Status Start();
// Stops the dispatcher. After stopping, RPCs should return without blocking.
void Stop();
// Returns the number of active iterations.
size_t NumActiveIterations() TF_LOCKS_EXCLUDED(mu_);
// See dispatcher.proto for API documentation.
/// Worker-facing API.
absl::Status WorkerHeartbeat(const WorkerHeartbeatRequest* request,
WorkerHeartbeatResponse* response);
absl::Status WorkerUpdate(const WorkerUpdateRequest* request,
WorkerUpdateResponse* response);
absl::Status GetDatasetDef(const GetDatasetDefRequest* request,
GetDatasetDefResponse* response);
absl::Status GetSplit(const GetSplitRequest* request,
GetSplitResponse* response);
/// Client-facing API.
absl::Status GetVersion(const GetVersionRequest* request,
GetVersionResponse* response);
absl::Status GetOrRegisterDataset(const GetOrRegisterDatasetRequest* request,
GetOrRegisterDatasetResponse* response);
absl::Status GetDataServiceMetadata(
const GetDataServiceMetadataRequest* request,
GetDataServiceMetadataResponse* response);
absl::Status GetDataServiceConfig(const GetDataServiceConfigRequest* request,
GetDataServiceConfigResponse* response);
absl::Status GetOrCreateJob(const GetOrCreateJobRequest* request,
GetOrCreateJobResponse* response);
absl::Status GetOrCreateIteration(const GetOrCreateIterationRequest* request,
GetOrCreateIterationResponse* response);
absl::Status ReleaseIterationClient(
const ReleaseIterationClientRequest* request,
ReleaseIterationClientResponse* response);
absl::Status MaybeRemoveTask(const MaybeRemoveTaskRequest* request,
MaybeRemoveTaskResponse* response);
absl::Status ClientHeartbeat(const ClientHeartbeatRequest* request,
ClientHeartbeatResponse* response);
absl::Status GetWorkers(const GetWorkersRequest* request,
GetWorkersResponse* response);
absl::Status Snapshot(const SnapshotRequest* request,
SnapshotResponse* response);
absl::Status GetSnapshotSplit(const GetSnapshotSplitRequest* request,
GetSnapshotSplitResponse* response);
absl::Status GetSnapshotStreams(const GetSnapshotStreamsRequest* request,
GetSnapshotStreamsResponse* response);
absl::Status DisableCompressionAtRuntime(
const DisableCompressionAtRuntimeRequest* request,
DisableCompressionAtRuntimeResponse* response);
// Exports the dispatcher state for debugging.
DispatcherStateExport ExportState() const;
private:
// A thread which periodically checks for iterations to clean up, clients to
// release, workers to consider missing, and snapshot streams to reassign.
void MaintenanceThread();
// Restores split providers from the state in `iteration` and stores them in
// `restored`.
absl::Status RestoreSplitProviders(
const DispatcherState::Iteration& iteration,
std::vector<std::unique_ptr<SplitProvider>>& restored)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Makes split providers for the specified `dataset_id`, and stores them in
// `split_providers`.
absl::Status MakeSplitProviders(
const std::string& dataset_id,
std::vector<std::unique_ptr<SplitProvider>>& split_providers)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Registers a dataset, storing the new dataset's id in `dataset_id`.
absl::Status RegisterDataset(const DatasetDef& dataset,
const DataServiceMetadata& metadata,
const std::string& requested_dataset_id,
std::string& dataset_id)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Finds the dataset ID with the requested dataset ID.
// Returns nullptr if no such dataset exists.
absl::StatusOr<std::optional<std::string>> FindDataset(
const GetOrRegisterDatasetRequest& request);
// Gets a worker's stub from `worker_stubs_`, or if none exists, creates a
// stub and stores it in `worker_stubs_`. A borrowed pointer to the stub is
// stored in `out_stub`.
absl::Status GetOrCreateWorkerStub(const std::string& worker_address,
WorkerService::Stub*& out_stub)
TF_LOCKS_EXCLUDED(mu_);
// Creates a job and stores it in `job`.
absl::Status CreateJob(const std::string& job_name,
const GetOrCreateJobRequest& request,
std::shared_ptr<const DispatcherState::Job>& job)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Creates an iteration and stores it in `iteration`. This method updates the
// dispatcher state with the new iteration, but does not assign tasks to
// workers.
absl::Status CreateIteration(
const GetOrCreateIterationRequest& request,
std::shared_ptr<const DispatcherState::Iteration>& iteration)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Creates tasks for the specified worker, one task for every unfinished
// iteration.
absl::Status CreateTasksForWorker(const std::string& worker_address);
// Finds tasks that should be deleted from a worker, updating the heartbeat
// response.
absl::Status FindTasksToDelete(
const absl::flat_hash_set<int64_t>& current_tasks,
const std::vector<std::shared_ptr<const DispatcherState::Task>>&
assigned_tasks,
WorkerHeartbeatResponse* response);
// Finds new tasks that should be assigned to a worker and adds them to
// the heartbeat response.
absl::Status FindNewTasks(
const std::string& worker_address,
const absl::flat_hash_set<int64_t>& current_tasks,
std::vector<std::shared_ptr<const DispatcherState::Task>>& assigned_tasks,
WorkerHeartbeatResponse* response);
// Reports the processing time of each active task to `auto_scaler_`.
void ReportProcessingTimesFromActiveTasks(
const std::vector<ActiveTask>& active_tasks,
const std::string& worker_address) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Acquires an iteration client id to read from the given iteration and sets
// `iteration_client_id`.
absl::Status AcquireIterationClientId(
const std::shared_ptr<const DispatcherState::Iteration>& iteration,
int64_t& iteration_client_id) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Creates one task for each worker, for the given iteration. The created
// tasks are stored in `tasks`. This method only updates dispatcher metadata
// with the new tasks, but doesn't assign the tasks to the workers.
absl::Status CreateTasksForIteration(
std::shared_ptr<const DispatcherState::Iteration> iteration,
std::vector<std::shared_ptr<const DispatcherState::Task>>& tasks)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Creates a new task for an iteration. The created task may be either
// pending or active.
absl::Status CreateTask(
std::shared_ptr<const DispatcherState::Iteration> iteration,
const std::string& worker_address,
std::shared_ptr<const DispatcherState::Task>& task)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Creates a pending task for a round robin iteration. All consumers need to
// agree on which round to add the task in before the pending task can be
// promoted to a regular task.
absl::Status CreatePendingTask(
std::shared_ptr<const DispatcherState::Iteration> iteration,
const std::string& worker_address) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Creates a new active task for an iteration, storing the created task in
// `task`.
absl::Status CreateActiveTask(
std::shared_ptr<const DispatcherState::Iteration> iteration,
const std::string& worker_address,
std::shared_ptr<const DispatcherState::Task>& task);
// Assigns the list of tasks to the workers indicated by their
// `worker_address` fields.
absl::Status AssignTasks(
std::vector<std::shared_ptr<const DispatcherState::Task>> tasks)
TF_LOCKS_EXCLUDED(mu_);
// Assigns a task to the worker indicated by its `worker_address` field.
absl::Status AssignTask(std::shared_ptr<const DispatcherState::Task> task)
TF_LOCKS_EXCLUDED(mu_);
// Validates that an existing job matches a given request.
// Returns an error status describing any difference.
absl::Status ValidateMatchingJob(
std::shared_ptr<const DispatcherState::Job> job,
const GetOrCreateJobRequest& request) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Fills out a TaskDef with information about a task.
absl::Status PopulateTaskDef(
std::shared_ptr<const DispatcherState::Task> task,
TaskDef* task_def) const TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Checks that the dispatcher has started, returning UNAVAILABLE if it hasn't.
absl::Status CheckStarted() TF_LOCKS_EXCLUDED(mu_);
// Restores ongoing tf.data snapshots.
absl::Status RestoreSnapshots();
// Records that a split was produced by a call to `GetSplit`.
absl::Status RecordSplitProduced(int64_t iteration_id, int64_t repetition,
int64_t split_provider_index, bool finished)
TF_LOCKS_EXCLUDED(mu_);
// Applies a state update, updating both the journal and the in-memory state.
absl::Status Apply(const Update& update) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Applies a state update, but doesn't update the journal. Only meant to be
// used when recovering state when the dispatcher starts.
absl::Status ApplyWithoutJournaling(const Update& update)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Removes the client with `client_id` from `auto_scaler_`
void RemoveClientFromAutoScaler(int64_t client_id)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Releases iteration clients that haven't heartbeated recently.
absl::Status ReleaseMissingClients() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Removes the worker with `worker_address` from `auto_scaler_`, which is
// potentially associated with multiple iterations.
void RemoveWorkerFromAutoScaler(const std::string& worker_address)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Checks for workers that haven't heartbeated recently and alerts the
// snapshot managers.
void DetectMissingWorkers() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Scans for old iterations and marks them as finished.
absl::Status GcOldIterations() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Returns true if an iteration should be garbage collected.
bool ShouldGcIteration(const DispatcherState::Iteration& iteration,
int64_t now_us) const;
// Gets a `DatasetDef` from `dataset_store_` for the given dataset id, and
// stores it in `dataset_def`.
absl::Status GetDatasetDef(const std::string& dataset_id,
std::shared_ptr<const DatasetDef>& dataset_def)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
// Gets a `DatasetDef` from `dataset_store_` for the given dataset, and
// stores it in `dataset_def`.
absl::Status GetDatasetDef(const DispatcherState::Dataset& dataset,
std::shared_ptr<const DatasetDef>& dataset_def)
TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
const experimental::DispatcherConfig config_;
Env* env_;
mutable mutex mu_;
// Uses a separate mutex for `GetSplit` requests. `GetSplit` may be blocking.
// Locking `mu_` in `GetSplit` could block all other RPCs.
mutable mutex get_split_mu_;
bool started_ TF_GUARDED_BY(mu_) = false;
bool cancelled_ TF_GUARDED_BY(mu_) = false;
// Cached worker stubs for communicating with workers.
absl::flat_hash_map<std::string, std::unique_ptr<WorkerService::Stub>>
worker_stubs_ TF_GUARDED_BY(mu_);
// Store of dataset definitions.
std::unique_ptr<DatasetStore> dataset_store_ TF_GUARDED_BY(mu_);
// Mapping from iteration id to the split providers for the iteration.
absl::flat_hash_map<int64_t, std::vector<std::unique_ptr<SplitProvider>>>
split_providers_ TF_GUARDED_BY(mu_);
// Mapping from round robin iteration id to the round the iteration is
// currently on. This is based on the data provided by client heartbeats,
// and may be stale.
absl::flat_hash_map<int64_t, int64_t> round_robin_rounds_ TF_GUARDED_BY(mu_);
// Map from task id to a TaskRemover which determines when to remove the task.
absl::flat_hash_map<int64_t, std::shared_ptr<TaskRemover>>
remove_task_requests_ TF_GUARDED_BY(mu_);
// Map from client id to the time of the client's last heartbeat.
absl::flat_hash_map<int64_t, absl::Time> latest_client_heartbeats_time_
TF_GUARDED_BY(mu_);
// Map from worker address to the time of the worker's last heartbeat.
absl::flat_hash_map<std::string, absl::Time> latest_worker_heartbeats_time_
TF_GUARDED_BY(mu_);
// A manager for each snapshot resumed or started during the lifetime of this
// dispatcher instance. Note that these are *not* garbage collected; managers
// for completed snapshots will remain here for the lifetime of the dispatcher
// instance. They will even be recovered if the dispatcher is restarted.
absl::flat_hash_map<std::string, std::unique_ptr<SnapshotManager>> snapshots_
TF_GUARDED_BY(mu_);
// A single stream assignment manager shared by all managers in `snapshots_`.
SnapshotAssignmentManager snapshot_assignment_manager_;
std::optional<std::unique_ptr<JournalWriter>> journal_writer_
TF_GUARDED_BY(mu_);
DispatcherState state_ TF_GUARDED_BY(mu_);
// Condition variable for waking up the gc thread.
condition_variable maintenance_thread_cv_;
std::unique_ptr<Thread> maintenance_thread_;
MultipleIterationsAutoScaler auto_scaler_;
DataServiceDispatcherImpl(const DataServiceDispatcherImpl&) = delete;
void operator=(const DataServiceDispatcherImpl&) = delete;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_DISPATCHER_IMPL_H_
tensorflow/core/data/service/dispatcher_state.cc
+503
View File
@@ -0,0 +1,503 @@
/* Copyright 2020 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/core/data/service/dispatcher_state.h"
#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/journal.h"
#include "tensorflow/core/data/service/journal.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
DispatcherState::DispatcherState()
: worker_index_resolver_(std::vector<std::string>{}) {}
DispatcherState::DispatcherState(
const experimental::DispatcherConfig& dispatcher_config)
: worker_index_resolver_(dispatcher_config.worker_addresses()) {}
absl::Status DispatcherState::Apply(const Update& update) {
switch (update.update_type_case()) {
case Update::kRegisterDataset:
RegisterDataset(update.register_dataset());
break;
case Update::kRegisterWorker:
RegisterWorker(update.register_worker());
break;
case Update::kCreateJob:
CreateJob(update.create_job());
break;
case Update::kCreateIteration:
CreateIteration(update.create_iteration());
break;
case Update::kProduceSplit:
ProduceSplit(update.produce_split());
break;
case Update::kAcquireIterationClient:
AcquireIterationClient(update.acquire_iteration_client());
break;
case Update::kReleaseIterationClient:
ReleaseIterationClient(update.release_iteration_client());
break;
case Update::kGarbageCollectIteration:
GarbageCollectIteration(update.garbage_collect_iteration());
break;
case Update::kRemoveTask:
RemoveTask(update.remove_task());
break;
case Update::kCreatePendingTask:
CreatePendingTask(update.create_pending_task());
break;
case Update::kClientHeartbeat:
ClientHeartbeat(update.client_heartbeat());
break;
case Update::kCreateTask:
CreateTask(update.create_task());
break;
case Update::kFinishTask:
FinishTask(update.finish_task());
break;
case Update::kSnapshot:
Snapshot(update.snapshot());
break;
case Update::kCompressionDisabledAtRuntime:
CompressionDisabledAtRuntime(update.compression_disabled_at_runtime());
break;
case Update::UPDATE_TYPE_NOT_SET:
return absl::InternalError("Update type not set.");
}
return absl::OkStatus();
}
void DispatcherState::RegisterDataset(
const RegisterDatasetUpdate& register_dataset) {
std::string dataset_id = register_dataset.dataset_id();
auto dataset =
std::make_shared<Dataset>(dataset_id, register_dataset.metadata());
DCHECK(!datasets_by_id_.contains(dataset_id));
datasets_by_id_[dataset_id] = dataset;
UpdateNextAvailableDatasetId();
}
void DispatcherState::RegisterWorker(
const RegisterWorkerUpdate& register_worker) {
std::string address = register_worker.worker_address();
DCHECK(!workers_.contains(address));
workers_[address] = std::make_shared<Worker>(register_worker);
tasks_by_worker_[address] =
absl::flat_hash_map<int64_t, std::shared_ptr<Task>>();
worker_index_resolver_.AddWorker(address);
}
void DispatcherState::CreateJob(const CreateJobUpdate& create_job) {
int64_t job_id = create_job.job_id();
std::string job_name = create_job.job_name();
std::optional<int64_t> num_consumers;
if (create_job.optional_num_consumers_case() ==
CreateJobUpdate::kNumConsumers) {
num_consumers = create_job.num_consumers();
}
auto job = std::make_shared<Job>(
job_id, create_job.dataset_id(), create_job.processing_mode_def(),
job_name, num_consumers, create_job.use_cross_trainer_cache(),
create_job.target_workers());
DCHECK(!jobs_by_id_.contains(job_id));
jobs_by_id_[job_id] = job;
DCHECK(!jobs_by_name_.contains(job_name));
jobs_by_name_[job_name] = job;
next_available_job_id_ = std::max(next_available_job_id_, job_id + 1);
}
absl::Status DispatcherState::JobFromId(int64_t job_id,
std::shared_ptr<const Job>& job) const {
auto it = jobs_by_id_.find(job_id);
if (it == jobs_by_id_.end()) {
return absl::NotFoundError(
absl::StrCat("Job with id ", job_id, " not found"));
}
job = it->second;
return absl::OkStatus();
}
absl::Status DispatcherState::JobByName(const std::string& job_name,
std::shared_ptr<const Job>& job) const {
auto it = jobs_by_name_.find(job_name);
if (it == jobs_by_name_.end()) {
return absl::NotFoundError(
absl::StrCat("Job with name ", job_name, " not found"));
}
job = it->second;
return absl::OkStatus();
}
void DispatcherState::CreateIteration(
const CreateIterationUpdate& create_iteration) {
int64_t iteration_id = create_iteration.iteration_id();
int64_t job_id = create_iteration.job_id();
DCHECK(jobs_by_id_.contains(job_id));
auto& job = jobs_by_id_[job_id];
DCHECK(job);
IterationKey iteration_key(job->job_name, create_iteration.repetition());
auto iteration = std::make_shared<Iteration>(
iteration_id, iteration_key, create_iteration.num_split_providers(), job);
DCHECK(!iterations_.contains(iteration_id));
iterations_[iteration_id] = iteration;
tasks_by_iteration_[iteration_id] = std::vector<std::shared_ptr<Task>>();
DCHECK(!iterations_by_key_.contains(iteration_key) ||
iterations_by_key_[iteration_key]->garbage_collected);
iterations_by_key_[iteration_key] = iteration;
next_available_iteration_id_ =
std::max(next_available_iteration_id_, iteration_id + 1);
}
void DispatcherState::ProduceSplit(const ProduceSplitUpdate& produce_split) {
std::shared_ptr<Iteration> iteration =
iterations_[produce_split.iteration_id()];
DCHECK(iteration->distributed_epoch_state.has_value());
DistributedEpochState& state = iteration->distributed_epoch_state.value();
int64_t provider_index = produce_split.split_provider_index();
DCHECK_GE(produce_split.repetition(), state.repetitions[provider_index]);
state.repetitions[provider_index] = produce_split.repetition();
if (produce_split.finished()) {
state.repetitions[provider_index]++;
state.indices[provider_index] = 0;
return;
}
state.indices[provider_index]++;
}
void DispatcherState::AcquireIterationClient(
const AcquireIterationClientUpdate& acquire_iteration_client) {
int64_t iteration_client_id = acquire_iteration_client.iteration_client_id();
std::shared_ptr<Iteration>& iteration =
iterations_for_client_ids_[iteration_client_id];
DCHECK(!iteration);
iteration = iterations_[acquire_iteration_client.iteration_id()];
DCHECK(iteration);
iteration->num_clients++;
next_available_iteration_client_id_ =
std::max(next_available_iteration_client_id_, iteration_client_id + 1);
}
void DispatcherState::ReleaseIterationClient(
const ReleaseIterationClientUpdate& release_iteration_client) {
int64_t iteration_client_id = release_iteration_client.iteration_client_id();
std::shared_ptr<Iteration>& iteration =
iterations_for_client_ids_[iteration_client_id];
DCHECK(iteration);
iteration->num_clients--;
DCHECK_GE(iteration->num_clients, 0);
iteration->last_client_released_micros =
release_iteration_client.time_micros();
iterations_for_client_ids_.erase(iteration_client_id);
}
void DispatcherState::GarbageCollectIteration(
const GarbageCollectIterationUpdate& garbage_collect_iteration) {
int64_t iteration_id = garbage_collect_iteration.iteration_id();
for (auto& task : tasks_by_iteration_[iteration_id]) {
task->finished = true;
tasks_by_worker_[task->worker_address].erase(task->task_id);
}
iterations_[iteration_id]->finished = true;
iterations_[iteration_id]->garbage_collected = true;
}
void DispatcherState::RemoveTask(const RemoveTaskUpdate& remove_task) {
std::shared_ptr<Task>& task = tasks_[remove_task.task_id()];
DCHECK(task);
task->removed = true;
auto& tasks_for_iteration =
tasks_by_iteration_[task->iteration->iteration_id];
for (auto it = tasks_for_iteration.begin(); it != tasks_for_iteration.end();
++it) {
if ((*it)->task_id == task->task_id) {
tasks_for_iteration.erase(it);
break;
}
}
tasks_by_worker_[task->worker_address].erase(task->task_id);
tasks_.erase(task->task_id);
VLOG(1) << "Removed task " << remove_task.task_id() << " from worker "
<< task->worker_address;
}
void DispatcherState::CreatePendingTask(
const CreatePendingTaskUpdate& create_pending_task) {
int64_t task_id = create_pending_task.task_id();
auto& task = tasks_[task_id];
DCHECK_EQ(task, nullptr);
auto& iteration = iterations_[create_pending_task.iteration_id()];
DCHECK_NE(iteration, nullptr);
task = std::make_shared<Task>(create_pending_task, iteration);
iteration->pending_tasks.emplace(task, create_pending_task.starting_round());
tasks_by_worker_[create_pending_task.worker_address()][task->task_id] = task;
next_available_task_id_ = std::max(next_available_task_id_, task_id + 1);
}
void DispatcherState::ClientHeartbeat(
const ClientHeartbeatUpdate& client_heartbeat) {
int64_t iteration_client_id = client_heartbeat.iteration_client_id();
auto& iteration = iterations_for_client_ids_[iteration_client_id];
DCHECK(!iteration->pending_tasks.empty());
auto& task = iteration->pending_tasks.front();
if (client_heartbeat.has_task_rejected()) {
task.failures++;
task.ready_consumers.clear();
task.target_round = client_heartbeat.task_rejected().new_target_round();
}
if (client_heartbeat.task_accepted()) {
task.ready_consumers.insert(iteration_client_id);
if (task.ready_consumers.size() == iteration->job->num_consumers.value()) {
VLOG(1) << "Promoting task " << task.task->task_id
<< " from pending to active";
task.task->starting_round = task.target_round;
tasks_by_iteration_[iteration->iteration_id].push_back(task.task);
iteration->pending_tasks.pop();
}
}
}
void DispatcherState::CreateTask(const CreateTaskUpdate& create_task) {
int64_t task_id = create_task.task_id();
auto& task = tasks_[task_id];
DCHECK_EQ(task, nullptr);
auto& iteration = iterations_[create_task.iteration_id()];
DCHECK_NE(iteration, nullptr);
task = std::make_shared<Task>(create_task, iteration);
tasks_by_iteration_[create_task.iteration_id()].push_back(task);
tasks_by_worker_[create_task.worker_address()][task->task_id] = task;
next_available_task_id_ = std::max(next_available_task_id_, task_id + 1);
}
void DispatcherState::FinishTask(const FinishTaskUpdate& finish_task) {
VLOG(2) << "Marking task " << finish_task.task_id() << " as finished";
int64_t task_id = finish_task.task_id();
auto& task = tasks_[task_id];
DCHECK(task != nullptr);
task->finished = true;
tasks_by_worker_[task->worker_address].erase(task->task_id);
bool all_finished = true;
for (const auto& task_for_iteration :
tasks_by_iteration_[task->iteration->iteration_id]) {
if (!task_for_iteration->finished) {
all_finished = false;
}
}
VLOG(3) << "Iteration " << task->iteration->iteration_id
<< " finished: " << all_finished;
iterations_[task->iteration->iteration_id]->finished = all_finished;
}
std::string DispatcherState::NextAvailableDatasetId() const {
return absl::StrCat(next_available_dataset_id_);
}
void DispatcherState::UpdateNextAvailableDatasetId() {
while (datasets_by_id_.contains(absl::StrCat(next_available_dataset_id_))) {
++next_available_dataset_id_;
}
}
absl::Status DispatcherState::DatasetFromId(
const std::string& id, std::shared_ptr<const Dataset>& dataset) const {
auto it = datasets_by_id_.find(id);
if (it == datasets_by_id_.end()) {
return absl::NotFoundError(absl::StrCat("Dataset id ", id, " not found"));
}
dataset = it->second;
return absl::OkStatus();
}
absl::Status DispatcherState::WorkerFromAddress(
const std::string& address, std::shared_ptr<const Worker>& worker) const {
auto it = workers_.find(address);
if (it == workers_.end()) {
return absl::NotFoundError(
absl::StrCat("Worker with address ", address, " not found."));
}
worker = it->second;
return absl::OkStatus();
}
std::vector<std::shared_ptr<const DispatcherState::Worker>>
DispatcherState::ListWorkers() const {
std::vector<std::shared_ptr<const Worker>> workers;
workers.reserve(workers_.size());
for (const auto& it : workers_) {
workers.push_back(it.second);
}
return workers;
}
std::vector<std::shared_ptr<const DispatcherState::Iteration>>
DispatcherState::ListIterations() const {
std::vector<std::shared_ptr<const DispatcherState::Iteration>> iterations;
iterations.reserve(iterations_.size());
for (const auto& it : iterations_) {
iterations.push_back(it.second);
}
return iterations;
}
absl::Status DispatcherState::IterationFromId(
int64_t id, std::shared_ptr<const Iteration>& iteration) const {
auto it = iterations_.find(id);
if (it == iterations_.end()) {
return absl::NotFoundError(absl::StrCat("Iteration id ", id, " not found"));
}
iteration = it->second;
return absl::OkStatus();
}
absl::Status DispatcherState::IterationByKey(
IterationKey iteration_key,
std::shared_ptr<const Iteration>& iteration) const {
auto it = iterations_by_key_.find(iteration_key);
if (it == iterations_by_key_.end()) {
return absl::NotFoundError(absl::StrCat(
"Iteration key ", iteration_key.DebugString(), " not found"));
}
iteration = it->second;
return absl::OkStatus();
}
int64_t DispatcherState::NextAvailableJobId() const {
return next_available_job_id_;
}
int64_t DispatcherState::NextAvailableIterationId() const {
return next_available_iteration_id_;
}
absl::Status DispatcherState::IterationForIterationClientId(
int64_t iteration_client_id, std::shared_ptr<const Iteration>& iteration) {
iteration = iterations_for_client_ids_[iteration_client_id];
if (!iteration) {
return absl::NotFoundError(
absl::StrCat("Iteration client id not found: ", iteration_client_id));
}
return absl::OkStatus();
}
std::vector<int64_t> DispatcherState::ListActiveClientIds() {
std::vector<int64_t> ids;
for (const auto& it : iterations_for_client_ids_) {
if (it.second && !it.second->finished) {
ids.push_back(it.first);
}
}
return ids;
}
int64_t DispatcherState::NextAvailableIterationClientId() const {
return next_available_iteration_client_id_;
}
absl::Status DispatcherState::TaskFromId(
int64_t id, std::shared_ptr<const Task>& task) const {
auto it = tasks_.find(id);
if (it == tasks_.end()) {
return absl::NotFoundError(absl::StrCat("Task ", id, " not found"));
}
task = it->second;
return absl::OkStatus();
}
absl::Status DispatcherState::TasksForIteration(
int64_t iteration_id,
std::vector<std::shared_ptr<const Task>>& tasks) const {
auto it = tasks_by_iteration_.find(iteration_id);
if (it == tasks_by_iteration_.end()) {
return absl::NotFoundError(
absl::StrCat("Iteration ", iteration_id, " not found"));
}
tasks.clear();
tasks.reserve(it->second.size());
for (const auto& task : it->second) {
tasks.push_back(task);
}
return absl::OkStatus();
}
absl::Status DispatcherState::TasksForWorker(
absl::string_view worker_address,
std::vector<std::shared_ptr<const Task>>& tasks) const {
tasks.clear();
auto it = tasks_by_worker_.find(worker_address);
if (it == tasks_by_worker_.end()) {
return absl::NotFoundError(
absl::StrCat("Worker ", worker_address, " not found"));
}
const absl::flat_hash_map<int64_t, std::shared_ptr<Task>>& worker_tasks =
it->second;
tasks.reserve(worker_tasks.size());
for (const auto& task : worker_tasks) {
tasks.push_back(task.second);
}
return absl::OkStatus();
}
int64_t DispatcherState::NextAvailableTaskId() const {
return next_available_task_id_;
}
absl::Status DispatcherState::ValidateWorker(
absl::string_view worker_address) const {
return worker_index_resolver_.ValidateWorker(worker_address);
}
absl::StatusOr<int64_t> DispatcherState::GetWorkerIndex(
absl::string_view worker_address) const {
return worker_index_resolver_.GetWorkerIndex(worker_address);
}
void DispatcherState::Snapshot(const SnapshotUpdate& snapshot) {
snapshot_paths_.insert(snapshot.path());
}
void DispatcherState::CompressionDisabledAtRuntime(
const CompressionDisabledAtRuntimeUpdate& compression_disabled_at_runtime) {
compression_disabled_at_runtime_.insert({
compression_disabled_at_runtime.dataset_id(),
compression_disabled_at_runtime.compression_disabled(),
});
}
std::optional<bool> DispatcherState::CompressionDisabledAtRuntime(
const std::string& dataset_id) const {
if (auto it = compression_disabled_at_runtime_.find(dataset_id);
it != compression_disabled_at_runtime_.end()) {
return it->second;
}
return std::nullopt;
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/dispatcher_state.h
+381
View File
@@ -0,0 +1,381 @@
/* Copyright 2020 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_CORE_DATA_SERVICE_DISPATCHER_STATE_H_
#define TENSORFLOW_CORE_DATA_SERVICE_DISPATCHER_STATE_H_
#include <cstdint>
#include <memory>
#include <optional>
#include <queue>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/graph_rewriters.h"
#include "tensorflow/core/data/service/journal.h"
#include "tensorflow/core/data/service/journal.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
// A class encapsulating the journaled state of the dispatcher. All state
// modifications must be done via `Apply`. This helps to ensure that
// replaying the journal will allow us to restore the exact same state.
//
// The following usage pattern will keep the journal in sync with the state of
// the dispatcher:
// {
// mutex_lock l(mu_);
// Update update = ... // create an update
// dispatcher_state.Apply(update);
// journal_writer.write(Update);
// // Unlock mu_
// }
//
// The division of functionality between DispatcherImpl and DispatcherState is
// as follows:
// - DispatcherImpl is responsible for handling RPC requests, reading from
// DispatcherState, and deciding what updates to apply to DispatcherState.
// DispatcherImpl handles all synchronization.
// - DispatcherState is responsible for making the state changes requested by
// DispatcherImpl and for providing DispatcherImpl with read-only access to
// the state.
//
// DispatcherState is thread-compatible but not thread-safe.
class DispatcherState {
public:
DispatcherState();
explicit DispatcherState(
const experimental::DispatcherConfig& dispatcher_config);
DispatcherState(const DispatcherState&) = delete;
DispatcherState& operator=(const DispatcherState&) = delete;
// Applies the given update to the dispatcher's state.
absl::Status Apply(const Update& update);
// A dataset registered with the dispatcher.
struct Dataset {
explicit Dataset(const std::string& dataset_id,
const DataServiceMetadata& metadata)
: dataset_id(dataset_id), metadata(metadata) {}
const std::string dataset_id;
const DataServiceMetadata metadata;
};
// A worker registered with the dispatcher.
struct Worker {
explicit Worker(const RegisterWorkerUpdate& register_worker)
: address(register_worker.worker_address()),
transfer_servers({register_worker.transfer_servers().begin(),
register_worker.transfer_servers().end()}),
tags(register_worker.worker_tags().begin(),
register_worker.worker_tags().end()),
uid(register_worker.worker_uid()) {}
const std::string address;
const std::vector<DataTransferServerInfo> transfer_servers;
const std::vector<std::string> tags;
const int64_t uid;
};
// A key for identifying an iteration. The key contains a job name,
// as well as a repetition number describing which repetition of the job
// we are on.
struct IterationKey {
explicit IterationKey(absl::string_view name, int64_t repetition)
: name(name), repetition(repetition) {}
friend bool operator==(const IterationKey& lhs, const IterationKey& rhs) {
return lhs.name == rhs.name && lhs.repetition == rhs.repetition;
}
template <typename H>
friend H AbslHashValue(H h, const IterationKey& k) {
return H::combine(std::move(h), k.name, k.repetition);
}
std::string DebugString() const {
return absl::StrCat(name, "/", repetition);
}
const std::string name;
const int64_t repetition;
};
struct DistributedEpochState {
explicit DistributedEpochState(int64_t num_split_providers)
: repetitions(num_split_providers), indices(num_split_providers) {}
// The current repetition for each split provider.
std::vector<int64_t> repetitions;
// Number of splits produced so far by each split provider.
std::vector<int64_t> indices;
};
struct Task;
struct PendingTask {
explicit PendingTask(std::shared_ptr<Task> task, int64_t target_round)
: task(std::move(task)), target_round(target_round) {}
std::shared_ptr<Task> task;
// The target round where we want to insert the task.
int64_t target_round;
// Which consumers have responded that they have successfully blocked
// before the target round.
absl::flat_hash_set<int64_t> ready_consumers;
// How many times we have failed to add the task.
int64_t failures = 0;
};
struct Job {
explicit Job(int64_t id, const std::string& dataset_id,
const ProcessingModeDef& processing_mode, std::string job_name,
std::optional<int64_t> num_consumers,
bool use_cross_trainer_cache, TargetWorkers target_workers)
: id(id),
dataset_id(dataset_id),
processing_mode(processing_mode),
job_name(job_name),
num_consumers(num_consumers),
use_cross_trainer_cache(use_cross_trainer_cache),
target_workers(target_workers) {}
const int64_t id;
const std::string dataset_id;
const ProcessingModeDef processing_mode;
const std::string job_name;
const std::optional<int64_t> num_consumers;
const bool use_cross_trainer_cache;
const TargetWorkers target_workers;
};
// An iteration for processing a dataset.
struct Iteration {
explicit Iteration(int64_t iteration_id, IterationKey iteration_key,
int64_t num_split_providers, std::shared_ptr<Job> job)
: iteration_id(iteration_id), iteration_key(iteration_key), job(job) {
if (IsDynamicShard(job->processing_mode)) {
distributed_epoch_state = DistributedEpochState(num_split_providers);
}
}
bool IsRoundRobin() const { return job->num_consumers.has_value(); }
std::string DebugString() const {
return absl::StrCat(iteration_key.name, "_", iteration_key.repetition);
}
const int64_t iteration_id;
const IterationKey iteration_key;
const std::shared_ptr<Job> job;
std::optional<DistributedEpochState> distributed_epoch_state;
std::queue<PendingTask> pending_tasks;
int64_t num_clients = 0;
int64_t last_client_released_micros = -1;
bool finished = false;
// Indicates whether the iteration was garbage collected.
bool garbage_collected = false;
};
struct Task {
template <class T>
explicit Task(const T& create_task_update,
const std::shared_ptr<Iteration>& iteration)
: task_id(create_task_update.task_id()),
iteration(iteration),
worker_address(create_task_update.worker_address()),
transfer_servers(create_task_update.transfer_servers().begin(),
create_task_update.transfer_servers().end()),
worker_tags(create_task_update.worker_tags().begin(),
create_task_update.worker_tags().end()),
worker_uid(create_task_update.worker_uid()) {}
const int64_t task_id;
const std::shared_ptr<Iteration> iteration;
const std::string worker_address;
const std::vector<DataTransferServerInfo> transfer_servers;
const std::vector<std::string> worker_tags;
const int64_t worker_uid;
int64_t starting_round = 0;
bool finished = false;
bool removed = false;
};
using TasksById = absl::flat_hash_map<int64_t, std::shared_ptr<Task>>;
// Returns the next available dataset ID.
std::string NextAvailableDatasetId() const;
// Gets a dataset by id. Returns NOT_FOUND if there is no such dataset.
absl::Status DatasetFromId(const std::string& id,
std::shared_ptr<const Dataset>& dataset) const;
// Gets a worker by address. Returns NOT_FOUND if there is no such worker.
absl::Status WorkerFromAddress(const std::string& address,
std::shared_ptr<const Worker>& worker) const;
// Lists all workers registered with the dispatcher.
std::vector<std::shared_ptr<const Worker>> ListWorkers() const;
// Returns the next available job id.
int64_t NextAvailableJobId() const;
// Gets a job by id. Returns NOT_FOUND if there is no such job.
absl::Status JobFromId(int64_t job_id, std::shared_ptr<const Job>& job) const;
// Gets a job by name. Returns NOT_FOUND if there is no such job.
absl::Status JobByName(const std::string& job_name,
std::shared_ptr<const Job>& job) const;
// Returns the next available iteration id.
int64_t NextAvailableIterationId() const;
// Returns a list of all iterations.
std::vector<std::shared_ptr<const Iteration>> ListIterations() const;
// Gets an iteration by id. Returns NOT_FOUND if there is no such iteration.
absl::Status IterationFromId(
int64_t id, std::shared_ptr<const Iteration>& iteration) const;
// Gets an iteration by key. Returns NOT_FOUND if there is no such iteration.
absl::Status IterationByKey(
IterationKey key, std::shared_ptr<const Iteration>& iteration) const;
// Returns the iteration associated with the given iteration client id.
// Returns NOT_FOUND if the iteration_client_id is unknown or has been
// released.
absl::Status IterationForIterationClientId(
int64_t iteration_client_id, std::shared_ptr<const Iteration>& iteration);
// Returns a list of all active client ids.
std::vector<int64_t> ListActiveClientIds();
// Returns the next available iteration client id.
int64_t NextAvailableIterationClientId() const;
// Returns the next available task id.
int64_t NextAvailableTaskId() const;
// Gets a task by id. Returns NOT_FOUND if there is no such task.
absl::Status TaskFromId(int64_t id, std::shared_ptr<const Task>& task) const;
// Stores a list of all tasks for the given iteration to `tasks`. Returns
// NOT_FOUND if there is no such iteration.
absl::Status TasksForIteration(
int64_t iteration_id,
std::vector<std::shared_ptr<const Task>>& tasks) const;
// Stores a list of all tasks for the given worker to `tasks`. Returns
// NOT_FOUND if there is no such worker.
absl::Status TasksForWorker(
const absl::string_view worker_address,
std::vector<std::shared_ptr<const Task>>& tasks) const;
// If the dispatcher config explicitly specifies a list of workers, validates
// `worker_address` is in the list.
absl::Status ValidateWorker(absl::string_view worker_address) const;
// If the dispatcher config specifies worker addresses, `GetWorkerIndex`
// returns the worker index according to the list. This is useful for
// deterministically sharding a dataset among a fixed set of workers.
absl::StatusOr<int64_t> GetWorkerIndex(
absl::string_view worker_address) const;
// Returns the paths of all snapshots initiated during the lifetime of this
// journal.
const absl::flat_hash_set<std::string>& ListSnapshotPaths() const {
return snapshot_paths_;
}
// Returns a bool describing whether or not compression was disabled at
// runtime for the given dataset, if such a decision has been made.
std::optional<bool> CompressionDisabledAtRuntime(
const std::string& dataset_id) const;
// Returns the current number of registered workers.
int64_t GetNumberOfRegisteredWorkers() const { return workers_.size(); }
private:
void RegisterDataset(const RegisterDatasetUpdate& register_dataset);
void RegisterWorker(const RegisterWorkerUpdate& register_worker);
void CreateJob(const CreateJobUpdate& create_job);
void CreateIteration(const CreateIterationUpdate& create_iteration);
void ProduceSplit(const ProduceSplitUpdate& produce_split);
void AcquireIterationClient(
const AcquireIterationClientUpdate& acquire_iteration_client);
void ReleaseIterationClient(
const ReleaseIterationClientUpdate& release_iteration_client);
void GarbageCollectIteration(
const GarbageCollectIterationUpdate& garbage_collect_iteration);
void RemoveTask(const RemoveTaskUpdate& remove_task);
void CreatePendingTask(const CreatePendingTaskUpdate& create_pending_task);
void ClientHeartbeat(const ClientHeartbeatUpdate& client_heartbeat);
void CreateTask(const CreateTaskUpdate& create_task);
void FinishTask(const FinishTaskUpdate& finish_task);
void Snapshot(const SnapshotUpdate& snapshot);
void CompressionDisabledAtRuntime(const CompressionDisabledAtRuntimeUpdate&
compression_disabled_at_runtime);
// Updates the next available dataset ID.
void UpdateNextAvailableDatasetId();
int64_t next_available_dataset_id_ = 1000;
// Registered datasets, keyed by dataset ids.
absl::flat_hash_map<std::string, std::shared_ptr<Dataset>> datasets_by_id_;
// Registered workers, keyed by address.
absl::flat_hash_map<std::string, std::shared_ptr<Worker>> workers_;
// Assigns an index to each worker according to worker addresses list
// specified in the dispatcher config.
WorkerIndexResolver worker_index_resolver_;
int64_t next_available_job_id_ = 5000;
// Jobs, keyed by job ids.
absl::flat_hash_map<int64_t, std::shared_ptr<Job>> jobs_by_id_;
// Jobs, keyed by job names.
absl::flat_hash_map<std::string, std::shared_ptr<Job>> jobs_by_name_;
int64_t next_available_iteration_id_ = 2000;
// Iterations, keyed by iteration ids.
absl::flat_hash_map<int64_t, std::shared_ptr<Iteration>> iterations_;
// Iterations, keyed by their iteration keys.
absl::flat_hash_map<IterationKey, std::shared_ptr<Iteration>>
iterations_by_key_;
int64_t next_available_iteration_client_id_ = 3000;
// Mapping from client ids to the iterations they are associated with.
absl::flat_hash_map<int64_t, std::shared_ptr<Iteration>>
iterations_for_client_ids_;
int64_t next_available_task_id_ = 4000;
// Tasks, keyed by task ids.
TasksById tasks_;
// List of tasks associated with each iteration.
absl::flat_hash_map<int64_t, std::vector<std::shared_ptr<Task>>>
tasks_by_iteration_;
// Tasks, keyed by worker addresses. The values are a map from task id to
// task.
absl::flat_hash_map<std::string, TasksById> tasks_by_worker_;
// Paths for all snapshots initiated during the lifetime of this journal.
absl::flat_hash_set<std::string> snapshot_paths_;
// A mapping of dataset id to a boolean describing whether or not compression
// was disabled at runtime for that dataset.
absl::flat_hash_map<std::string, bool> compression_disabled_at_runtime_;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_DISPATCHER_STATE_H_
tensorflow/core/data/service/dispatcher_state_test.cc
+664
View File
@@ -0,0 +1,664 @@
/* Copyright 2020 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/core/data/service/dispatcher_state.h"
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status_matchers.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/journal.pb.h"
#include "tensorflow/core/platform/random.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
namespace {
using Dataset = DispatcherState::Dataset;
using Worker = DispatcherState::Worker;
using IterationKey = DispatcherState::IterationKey;
using Job = DispatcherState::Job;
using Iteration = DispatcherState::Iteration;
using Task = DispatcherState::Task;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::SizeIs;
using ::testing::UnorderedElementsAre;
using ::tsl::testing::StatusIs;
absl::Status RegisterDataset(const std::string& dataset_id,
DispatcherState& state) {
Update update;
RegisterDatasetUpdate* register_dataset = update.mutable_register_dataset();
register_dataset->set_dataset_id(dataset_id);
return state.Apply(update);
}
absl::Status RegisterWorker(std::string worker_address,
DispatcherState& state) {
Update update;
update.mutable_register_worker()->set_worker_address(worker_address);
return state.Apply(update);
}
absl::Status CreateJob(int64_t job_id, const std::string& dataset_id,
const std::string& job_name, DispatcherState& state) {
Update update;
CreateJobUpdate* create_job = update.mutable_create_job();
create_job->set_job_id(job_id);
create_job->set_dataset_id(dataset_id);
create_job->set_job_name(job_name);
return state.Apply(update);
}
absl::Status CreateIteration(int64_t iteration_id,
const std::string& dataset_id,
const IterationKey& named_iteration_key,
DispatcherState& state) {
int64_t job_id = state.NextAvailableJobId();
TF_RETURN_IF_ERROR(
CreateJob(job_id, dataset_id, named_iteration_key.name, state));
Update update;
CreateIterationUpdate* create_iteration = update.mutable_create_iteration();
create_iteration->set_job_id(job_id);
create_iteration->set_iteration_id(iteration_id);
create_iteration->set_repetition(named_iteration_key.repetition);
return state.Apply(update);
}
absl::Status CreateIteration(int64_t iteration_id,
const std::string& dataset_id,
DispatcherState& state) {
IterationKey key(/*name=*/absl::StrCat(random::New64()), /*repetition=*/0);
return CreateIteration(iteration_id, dataset_id, key, state);
}
absl::Status AcquireIterationClientId(int64_t iteration_id,
int64_t iteration_client_id,
DispatcherState& state) {
Update update;
AcquireIterationClientUpdate* acquire_iteration_client =
update.mutable_acquire_iteration_client();
acquire_iteration_client->set_iteration_id(iteration_id);
acquire_iteration_client->set_iteration_client_id(iteration_client_id);
return state.Apply(update);
}
absl::Status ReleaseIterationClientId(int64_t iteration_client_id,
int64_t release_time,
DispatcherState& state) {
Update update;
ReleaseIterationClientUpdate* release_iteration_client =
update.mutable_release_iteration_client();
release_iteration_client->set_iteration_client_id(iteration_client_id);
release_iteration_client->set_time_micros(release_time);
return state.Apply(update);
}
absl::Status CreateTask(int64_t task_id, int64_t iteration_id,
const std::string& worker_address,
DispatcherState& state) {
Update update;
CreateTaskUpdate* create_task = update.mutable_create_task();
create_task->set_task_id(task_id);
create_task->set_iteration_id(iteration_id);
create_task->set_worker_address(worker_address);
return state.Apply(update);
}
absl::Status FinishTask(int64_t task_id, DispatcherState& state) {
Update update;
FinishTaskUpdate* finish_task = update.mutable_finish_task();
finish_task->set_task_id(task_id);
return state.Apply(update);
}
absl::Status Snapshot(const std::string& path, DispatcherState& state) {
Update update;
SnapshotUpdate* snapshot = update.mutable_snapshot();
snapshot->set_path(path);
return state.Apply(update);
}
} // namespace
TEST(DispatcherState, RegisterDataset) {
DispatcherState state;
std::string dataset_id = state.NextAvailableDatasetId();
int64_t dataset_id_int;
ASSERT_TRUE(absl::SimpleAtoi(dataset_id, &dataset_id_int));
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
EXPECT_EQ(state.NextAvailableDatasetId(), absl::StrCat(dataset_id_int + 1));
std::shared_ptr<const Dataset> dataset;
TF_EXPECT_OK(state.DatasetFromId(dataset_id, dataset));
EXPECT_TRUE(dataset->metadata.element_spec().empty());
EXPECT_EQ(dataset->metadata.compression(),
DataServiceMetadata::COMPRESSION_UNSPECIFIED);
}
TEST(DispatcherState, RegisterDatasetWithExplicitID) {
DispatcherState state;
TF_EXPECT_OK(RegisterDataset("dataset_id", state));
std::shared_ptr<const Dataset> dataset;
TF_EXPECT_OK(state.DatasetFromId("dataset_id", dataset));
EXPECT_EQ(dataset->dataset_id, "dataset_id");
}
TEST(DispatcherState, RegisterDatasetsWithDifferentIDs) {
DispatcherState state;
TF_EXPECT_OK(RegisterDataset("dataset_id1", state));
TF_EXPECT_OK(RegisterDataset("dataset_id2", state));
std::shared_ptr<const Dataset> dataset;
TF_EXPECT_OK(state.DatasetFromId("dataset_id1", dataset));
EXPECT_EQ(dataset->dataset_id, "dataset_id1");
TF_EXPECT_OK(state.DatasetFromId("dataset_id2", dataset));
EXPECT_EQ(dataset->dataset_id, "dataset_id2");
}
TEST(DispatcherState, RegisterDatasetCompression) {
DispatcherState state;
const std::string dataset_id = state.NextAvailableDatasetId();
Update update;
RegisterDatasetUpdate* register_dataset = update.mutable_register_dataset();
register_dataset->set_dataset_id(dataset_id);
register_dataset->mutable_metadata()->set_compression(
DataServiceMetadata::COMPRESSION_SNAPPY);
TF_ASSERT_OK(state.Apply(update));
{
std::shared_ptr<const Dataset> dataset;
TF_EXPECT_OK(state.DatasetFromId(dataset_id, dataset));
EXPECT_EQ(dataset->metadata.compression(),
DataServiceMetadata::COMPRESSION_SNAPPY);
}
}
TEST(DispatcherState, RegisterDatasetElementSpec) {
DispatcherState state;
const std::string dataset_id = state.NextAvailableDatasetId();
Update update;
RegisterDatasetUpdate* register_dataset = update.mutable_register_dataset();
register_dataset->set_dataset_id(dataset_id);
register_dataset->mutable_metadata()->set_element_spec(
"encoded_element_spec");
TF_ASSERT_OK(state.Apply(update));
{
std::shared_ptr<const Dataset> dataset;
TF_EXPECT_OK(state.DatasetFromId(dataset_id, dataset));
EXPECT_EQ(dataset->metadata.element_spec(), "encoded_element_spec");
}
}
TEST(DispatcherState, MissingDatasetId) {
DispatcherState state;
std::shared_ptr<const Dataset> dataset;
absl::Status s = state.DatasetFromId("missing_dataset_id", dataset);
EXPECT_EQ(s.code(), error::NOT_FOUND);
}
TEST(DispatcherState, NextAvailableDatasetId) {
DispatcherState state;
std::string dataset_id = state.NextAvailableDatasetId();
int64_t dataset_id_int;
ASSERT_TRUE(absl::SimpleAtoi(dataset_id, &dataset_id_int));
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
EXPECT_NE(state.NextAvailableDatasetId(), dataset_id);
EXPECT_EQ(state.NextAvailableDatasetId(), absl::StrCat(dataset_id_int + 1));
EXPECT_EQ(state.NextAvailableDatasetId(), state.NextAvailableDatasetId());
}
TEST(DispatcherState, RegisterWorker) {
DispatcherState state;
std::string address = "test_worker_address";
TF_EXPECT_OK(RegisterWorker(address, state));
std::shared_ptr<const Worker> worker;
TF_EXPECT_OK(state.WorkerFromAddress(address, worker));
EXPECT_EQ(worker->address, address);
}
TEST(DispatcherState, RegisterWorkerInFixedWorkerSet) {
experimental::DispatcherConfig config;
config.add_worker_addresses("/worker/task/0");
config.add_worker_addresses("/worker/task/1");
config.add_worker_addresses("/worker/task/2");
DispatcherState state(config);
TF_EXPECT_OK(state.ValidateWorker("/worker/task/0:20000"));
TF_EXPECT_OK(state.ValidateWorker("/worker/task/1:20000"));
TF_EXPECT_OK(state.ValidateWorker("/worker/task/2:20000"));
TF_EXPECT_OK(RegisterWorker("/worker/task/0:20000", state));
TF_EXPECT_OK(RegisterWorker("/worker/task/1:20000", state));
TF_EXPECT_OK(RegisterWorker("/worker/task/2:20000", state));
std::shared_ptr<const Worker> worker;
TF_EXPECT_OK(state.WorkerFromAddress("/worker/task/0:20000", worker));
EXPECT_EQ(worker->address, "/worker/task/0:20000");
}
TEST(DispatcherState, RegisterInvalidWorkerInFixedWorkerSet) {
experimental::DispatcherConfig config;
config.add_worker_addresses("/worker/task/0");
config.add_worker_addresses("/worker/task/1");
config.add_worker_addresses("/worker/task/2");
DispatcherState state(config);
EXPECT_THAT(state.ValidateWorker("localhost:20000"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("The worker's address is not configured")));
// Tests that `RegisterWorker` always returns OK, and ignores errors. This is
// because the journal records are supposed to be valid. If there is an error,
// it should be caught by `ValidateWorker` and not written to the journal.
TF_EXPECT_OK(RegisterWorker("localhost:20000", state));
std::shared_ptr<const Worker> worker;
EXPECT_THAT(state.WorkerFromAddress("/worker/task/0:20000", worker),
absl_testing::StatusIs(
error::NOT_FOUND,
"Worker with address /worker/task/0:20000 not found."));
}
TEST(DispatcherState, ListWorkers) {
DispatcherState state;
std::string address_1 = "address_1";
std::string address_2 = "address_2";
{
std::vector<std::shared_ptr<const Worker>> workers = state.ListWorkers();
EXPECT_THAT(workers, IsEmpty());
}
TF_EXPECT_OK(RegisterWorker(address_1, state));
{
std::vector<std::shared_ptr<const Worker>> workers = state.ListWorkers();
EXPECT_THAT(workers, SizeIs(1));
}
TF_EXPECT_OK(RegisterWorker(address_2, state));
{
std::vector<std::shared_ptr<const Worker>> workers = state.ListWorkers();
EXPECT_THAT(workers, SizeIs(2));
}
}
TEST(DispatcherState, MissingWorker) {
DispatcherState state;
std::shared_ptr<const Worker> worker;
absl::Status s = state.WorkerFromAddress("test_worker_address", worker);
EXPECT_EQ(s.code(), error::NOT_FOUND);
}
TEST(DispatcherState, UnknownUpdate) {
DispatcherState state;
Update update;
absl::Status s = state.Apply(update);
EXPECT_EQ(s.code(), error::INTERNAL);
}
TEST(DispatcherState, JobName) {
DispatcherState state;
std::string dataset_id = state.NextAvailableDatasetId();
int64_t job_id = state.NextAvailableJobId();
std::string job_name = "test_name";
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateJob(job_id, dataset_id, job_name, state));
std::shared_ptr<const Job> job;
TF_EXPECT_OK(state.JobByName(job_name, job));
EXPECT_EQ(state.NextAvailableJobId(), job_id + 1);
EXPECT_EQ(job->dataset_id, dataset_id);
EXPECT_FALSE(job->use_cross_trainer_cache);
}
TEST(DispatcherState, JobData) {
DispatcherState state;
std::string dataset_id = state.NextAvailableDatasetId();
int64_t job_id = state.NextAvailableJobId();
int64_t num_consumers = 8;
bool use_cross_trainer_cache = true;
TF_ASSERT_OK(RegisterDataset(dataset_id, state));
Update update;
CreateJobUpdate* create_job = update.mutable_create_job();
create_job->set_job_id(job_id);
create_job->set_dataset_id(dataset_id);
create_job->set_num_consumers(num_consumers);
create_job->set_use_cross_trainer_cache(use_cross_trainer_cache);
TF_ASSERT_OK(state.Apply(update));
std::shared_ptr<const Job> job;
TF_ASSERT_OK(state.JobFromId(job_id, job));
EXPECT_EQ(job->num_consumers, num_consumers);
EXPECT_EQ(job->use_cross_trainer_cache, use_cross_trainer_cache);
}
TEST(DispatcherState, CrossTrainerCacheTask) {
DispatcherState state;
std::string dataset_id = state.NextAvailableDatasetId();
std::string worker_address = "test_worker_address";
TF_ASSERT_OK(RegisterDataset(dataset_id, state));
int64_t job_id = state.NextAvailableJobId();
Update job_update;
CreateJobUpdate* create_job = job_update.mutable_create_job();
create_job->set_job_id(job_id);
create_job->set_dataset_id(dataset_id);
create_job->set_use_cross_trainer_cache(true);
TF_ASSERT_OK(state.Apply(job_update));
int64_t iteration_id = state.NextAvailableIterationId();
Update iteration_update;
CreateIterationUpdate* create_iteration =
iteration_update.mutable_create_iteration();
create_iteration->set_job_id(job_id);
create_iteration->set_iteration_id(iteration_id);
TF_ASSERT_OK(state.Apply(iteration_update));
int64_t task_id = state.NextAvailableTaskId();
TF_EXPECT_OK(CreateTask(task_id, iteration_id, worker_address, state));
std::shared_ptr<const Task> task;
TF_EXPECT_OK(state.TaskFromId(task_id, task));
EXPECT_EQ(task->iteration->iteration_id, iteration_id);
EXPECT_EQ(task->task_id, task_id);
EXPECT_EQ(task->worker_address, worker_address);
EXPECT_TRUE(task->iteration->job->use_cross_trainer_cache);
}
TEST(DispatcherState, CreateTask) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
std::string worker_address = "test_worker_address";
DispatcherState state;
int64_t task_id = state.NextAvailableTaskId();
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id, iteration_id, worker_address, state));
EXPECT_EQ(state.NextAvailableTaskId(), task_id + 1);
{
std::shared_ptr<const Task> task;
TF_EXPECT_OK(state.TaskFromId(task_id, task));
EXPECT_EQ(task->iteration->iteration_id, iteration_id);
EXPECT_EQ(task->task_id, task_id);
EXPECT_EQ(task->worker_address, worker_address);
EXPECT_FALSE(task->iteration->job->use_cross_trainer_cache);
}
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForIteration(iteration_id, tasks));
EXPECT_THAT(tasks, SizeIs(1));
}
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForWorker(worker_address, tasks));
EXPECT_EQ(1, tasks.size());
}
}
TEST(DispatcherState, CreateTasksForSameIteration) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t task_id_1 = 8;
int64_t task_id_2 = 9;
std::string worker_address = "test_worker_address";
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id_1, iteration_id, worker_address, state));
TF_EXPECT_OK(CreateTask(task_id_2, iteration_id, worker_address, state));
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForIteration(iteration_id, tasks));
EXPECT_THAT(tasks, SizeIs(2));
}
}
TEST(DispatcherState, CreateTasksForDifferentIterations) {
std::string dataset_id = "dataset_id";
int64_t iteration_id_1 = 3;
int64_t iteration_id_2 = 4;
int64_t task_id_1 = 8;
int64_t task_id_2 = 9;
std::string worker_address = "test_worker_address";
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id_1, dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id_2, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id_1, iteration_id_1, worker_address, state));
TF_EXPECT_OK(CreateTask(task_id_2, iteration_id_2, worker_address, state));
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForIteration(iteration_id_1, tasks));
EXPECT_THAT(tasks, SizeIs(1));
}
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForIteration(iteration_id_2, tasks));
EXPECT_THAT(tasks, SizeIs(1));
}
}
TEST(DispatcherState, CreateTasksForSameWorker) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t task_id_1 = 8;
int64_t task_id_2 = 9;
std::string worker_address = "test_worker_address";
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id_1, iteration_id, worker_address, state));
TF_EXPECT_OK(CreateTask(task_id_2, iteration_id, worker_address, state));
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForWorker(worker_address, tasks));
EXPECT_EQ(2, tasks.size());
}
}
TEST(DispatcherState, CreateTasksForDifferentWorkers) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t task_id_1 = 8;
int64_t task_id_2 = 9;
std::string worker_address_1 = "test_worker_address_1";
std::string worker_address_2 = "test_worker_address_2";
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id_1, iteration_id, worker_address_1, state));
TF_EXPECT_OK(CreateTask(task_id_2, iteration_id, worker_address_2, state));
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForWorker(worker_address_1, tasks));
EXPECT_EQ(1, tasks.size());
}
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForWorker(worker_address_2, tasks));
EXPECT_EQ(1, tasks.size());
}
}
TEST(DispatcherState, GetTasksForWorkerEmpty) {
std::string worker_address = "test_worker_address";
DispatcherState state;
TF_EXPECT_OK(RegisterWorker(worker_address, state));
{
std::vector<std::shared_ptr<const Task>> tasks;
TF_EXPECT_OK(state.TasksForWorker(worker_address, tasks));
EXPECT_EQ(0, tasks.size());
}
}
TEST(DispatcherState, FinishTask) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t task_id = 4;
std::string worker_address = "test_worker_address";
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id, iteration_id, worker_address, state));
TF_EXPECT_OK(FinishTask(task_id, state));
std::shared_ptr<const Task> task;
TF_EXPECT_OK(state.TaskFromId(task_id, task));
EXPECT_TRUE(task->finished);
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(state.IterationFromId(iteration_id, iteration));
EXPECT_TRUE(iteration->finished);
}
TEST(DispatcherState, FinishMultiTaskIteration) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t task_id_1 = 4;
int64_t task_id_2 = 5;
std::string worker_address = "test_worker_address";
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(CreateTask(task_id_1, iteration_id, worker_address, state));
TF_EXPECT_OK(CreateTask(task_id_2, iteration_id, worker_address, state));
TF_EXPECT_OK(FinishTask(task_id_1, state));
{
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(state.IterationFromId(iteration_id, iteration));
EXPECT_FALSE(iteration->finished);
}
TF_EXPECT_OK(FinishTask(task_id_2, state));
{
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(state.IterationFromId(iteration_id, iteration));
EXPECT_TRUE(iteration->finished);
}
}
TEST(DispatcherState, AcquireIterationClientId) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t iteration_client_id_1 = 1;
int64_t iteration_client_id_2 = 2;
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id_1, state));
{
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(state.IterationFromId(iteration_id, iteration));
EXPECT_EQ(iteration->num_clients, 1);
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id_2, state));
EXPECT_EQ(iteration->num_clients, 2);
}
{
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(
state.IterationForIterationClientId(iteration_client_id_1, iteration));
EXPECT_EQ(iteration->iteration_id, iteration_id);
}
{
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(
state.IterationForIterationClientId(iteration_client_id_2, iteration));
EXPECT_EQ(iteration->iteration_id, iteration_id);
}
}
TEST(DispatcherState, ReleaseIterationClientId) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t iteration_client_id = 6;
int64_t release_time = 100;
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id, state));
TF_EXPECT_OK(
ReleaseIterationClientId(iteration_client_id, release_time, state));
std::shared_ptr<const Iteration> iteration;
TF_EXPECT_OK(state.IterationFromId(iteration_id, iteration));
EXPECT_EQ(iteration->num_clients, 0);
absl::Status s =
state.IterationForIterationClientId(iteration_client_id, iteration);
EXPECT_EQ(s.code(), error::NOT_FOUND);
}
TEST(DispatcherState, ListActiveClientsEmpty) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t iteration_client_id = 6;
int64_t release_time = 100;
DispatcherState state;
EXPECT_THAT(state.ListActiveClientIds(), IsEmpty());
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id, state));
TF_EXPECT_OK(
ReleaseIterationClientId(iteration_client_id, release_time, state));
EXPECT_THAT(state.ListActiveClientIds(), IsEmpty());
}
TEST(DispatcherState, ListActiveClients) {
std::string dataset_id = "dataset_id";
int64_t iteration_id = 3;
int64_t iteration_client_id_1 = 6;
int64_t iteration_client_id_2 = 7;
int64_t iteration_client_id_3 = 8;
int64_t release_time = 100;
DispatcherState state;
TF_EXPECT_OK(RegisterDataset(dataset_id, state));
TF_EXPECT_OK(CreateIteration(iteration_id, dataset_id, state));
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id_1, state));
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id_2, state));
TF_EXPECT_OK(
ReleaseIterationClientId(iteration_client_id_2, release_time, state));
TF_EXPECT_OK(
AcquireIterationClientId(iteration_id, iteration_client_id_3, state));
EXPECT_THAT(state.ListActiveClientIds(), UnorderedElementsAre(6, 8));
}
TEST(DispatcherState, ListSnapshotPaths) {
DispatcherState state;
absl::flat_hash_set<std::string> snapshot_paths = {"p1", "p2"};
for (const auto& snapshot_path : snapshot_paths) {
TF_EXPECT_OK(Snapshot(snapshot_path, state));
}
EXPECT_EQ(state.ListSnapshotPaths(), snapshot_paths);
}
TEST(DispatcherState, GetNumberOfRegisteredWorkers) {
DispatcherState state;
std::string address_1 = "address_1";
std::string address_2 = "address_2";
EXPECT_EQ(state.GetNumberOfRegisteredWorkers(), 0);
TF_EXPECT_OK(RegisterWorker(address_1, state));
EXPECT_EQ(state.GetNumberOfRegisteredWorkers(), 1);
TF_EXPECT_OK(RegisterWorker(address_2, state));
EXPECT_EQ(state.GetNumberOfRegisteredWorkers(), 2);
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/export.proto
+58
View File
@@ -0,0 +1,58 @@
// Copyright 2026 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.
// ==============================================================================
syntax = "proto3";
package tensorflow.data;
import "tensorflow/core/data/service/common.proto";
import "tensorflow/core/protobuf/data_service.proto";
import "tensorflow/core/protobuf/service_config.proto";
// State of the dispatcher server, exported to improve debuggability.
// Next tag: 4
message DispatcherStateExport {
message Iteration {
string dataset_id = 1;
int64 iteration_id = 2;
IterationKeyDef iteration_key = 3;
ProcessingModeDef processing_mode = 4;
int64 num_consumers = 6;
int64 num_clients = 8;
bool finished = 10;
bool garbage_collected = 11;
}
experimental.DispatcherConfig dispatcher_config = 1;
repeated string worker_addresses = 2;
repeated Iteration iterations = 3;
}
// State of the worker server, exported to improve debuggability.
// Next tag: 5
message WorkerStateExport {
experimental.WorkerConfig worker_config = 1;
repeated TaskDef tasks = 2;
repeated int64 finished_task_ids = 3;
repeated int64 deleted_task_ids = 4;
}
// State of the tf.data service server, exported to improve debuggability.
// The dispatcher and worker servers will populate the corresponding fields.
// Next tag: 3
message ServerStateExport {
DispatcherStateExport dispatcher_state_export = 1;
WorkerStateExport worker_state_export = 2;
}
tensorflow/core/data/service/graph_rewriters.cc
+231
View File
@@ -0,0 +1,231 @@
/* Copyright 2021 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/core/data/service/graph_rewriters.h"
#include <cstdlib>
#include <iterator>
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include "absl/algorithm/container.h"
#include "absl/strings/match.h"
#include "absl/strings/str_join.h"
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "absl/types/optional.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/data/rewrite_utils.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/url.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
#include "tensorflow/core/grappler/optimizers/data/auto_shard.h"
#include "tensorflow/core/grappler/optimizers/data/graph_utils.h"
#include "tensorflow/core/grappler/optimizers/data/optimizer_base.h"
#include "tensorflow/core/grappler/optimizers/data/remove_compression_map.h"
#include "tensorflow/core/kernels/data/experimental/auto_shard_dataset_op.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::data::experimental::AutoShardDatasetOp;
// Don't apply general grappler optimizations when performing these rewrites.
// Sometimes there is a conflict among multiple applications of these general
// optimizations to the same graph (see b/303524867).
constexpr bool kApplyGeneralGrapplerOptimizations = false;
// A dynamic port has form %port% or %port_foo% that is to be replaced with the
// actual port.
bool HasDynamicPort(absl::string_view address) {
URL url(address);
return url.has_port() && absl::StartsWith(url.port(), "%port") &&
absl::EndsWith(url.port(), "%");
}
// Returns true if `config_address` has no port or a dynamic port (e.g.: %port%)
// and `worker_address` has an actual port (number of named port).
//
// For example, it returns true for the following cases:
//
// config_address worker_address
// ----------------------------------------------------------
// /worker/task/0 /worker/task/0:worker
// /worker/task/0:%port% /worker/task/0:10000
// /worker/task/0:%port_worker% /worker/task/0:worker
// /worker/task/0:%port_worker% /worker/task/0:10000
// localhost localhost:10000
// localhost:%port% localhost:10000
bool ShouldReplaceDynamicPort(absl::string_view config_address,
absl::string_view worker_address) {
URL config_url(config_address), worker_url(worker_address);
return (!config_url.has_port() || HasDynamicPort(config_address)) &&
worker_url.has_port() && config_url.host() == worker_url.host();
}
} // namespace
absl::StatusOr<GraphDef>
RemoveCompressionMapRewriter::ApplyRemoveCompressionMapRewrite(
const GraphDef& graph_def) {
grappler::RemoveCompressionMap remove_compression_map;
tensorflow::RewriterConfig::CustomGraphOptimizer config = GetRewriteConfig();
TF_RETURN_IF_ERROR(remove_compression_map.Init(&config));
GraphDef input_graph = graph_def;
TF_ASSIGN_OR_RETURN(std::string dataset_node, GetDatasetNode(input_graph));
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
GetGrapplerItem(&input_graph, &dataset_node, /*add_fake_sinks=*/false,
kApplyGeneralGrapplerOptimizations);
GraphDef rewritten_graph;
std::unordered_map<std::string, tensorflow::DeviceProperties> device_map;
tensorflow::grappler::VirtualCluster cluster(device_map);
grappler::AutoShard::OptimizationStats stats;
TF_RETURN_IF_ERROR(remove_compression_map.OptimizeAndCollectStats(
&cluster, *grappler_item, &rewritten_graph, &stats));
return rewritten_graph;
}
tensorflow::RewriterConfig::CustomGraphOptimizer
RemoveCompressionMapRewriter::GetRewriteConfig() const {
tensorflow::RewriterConfig::CustomGraphOptimizer config;
config.set_name("tf-data-service-remove-compression-map");
return config;
}
absl::StatusOr<AutoShardRewriter> AutoShardRewriter::Create(
const TaskDef& task_def) {
TF_ASSIGN_OR_RETURN(
AutoShardPolicy auto_shard_policy,
ToAutoShardPolicy(task_def.processing_mode_def().sharding_policy()));
return AutoShardRewriter(auto_shard_policy, task_def.num_workers(),
task_def.worker_index());
}
absl::StatusOr<GraphDef> AutoShardRewriter::ApplyAutoShardRewrite(
const GraphDef& graph_def) {
if (auto_shard_policy_ == AutoShardPolicy::OFF) {
return graph_def;
}
VLOG(2) << "Applying auto-shard policy "
<< AutoShardPolicy_Name(auto_shard_policy_)
<< ". Number of workers: " << num_workers_
<< "; worker index: " << worker_index_ << ".";
grappler::AutoShard autoshard;
tensorflow::RewriterConfig::CustomGraphOptimizer config = GetRewriteConfig();
TF_RETURN_IF_ERROR(autoshard.Init(&config));
GraphDef input_graph = graph_def;
TF_ASSIGN_OR_RETURN(std::string dataset_node, GetDatasetNode(input_graph));
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
GetGrapplerItem(&input_graph, &dataset_node, /*add_fake_sinks=*/false,
kApplyGeneralGrapplerOptimizations);
GraphDef rewritten_graph;
std::unordered_map<std::string, tensorflow::DeviceProperties> device_map;
tensorflow::grappler::VirtualCluster cluster(device_map);
grappler::AutoShard::OptimizationStats stats;
TF_RETURN_IF_ERROR(autoshard.OptimizeAndCollectStats(
&cluster, *grappler_item, &rewritten_graph, &stats));
return rewritten_graph;
}
AutoShardRewriter::AutoShardRewriter(AutoShardPolicy auto_shard_policy,
int64_t num_workers, int64_t worker_index)
: auto_shard_policy_(auto_shard_policy),
num_workers_(num_workers),
worker_index_(worker_index) {}
tensorflow::RewriterConfig::CustomGraphOptimizer
AutoShardRewriter::GetRewriteConfig() const {
tensorflow::RewriterConfig::CustomGraphOptimizer config;
config.set_name("tf-data-service-auto-shard");
(*config.mutable_parameter_map())[AutoShardDatasetOp::kNumWorkers].set_i(
num_workers_);
(*config.mutable_parameter_map())[AutoShardDatasetOp::kIndex].set_i(
worker_index_);
(*config.mutable_parameter_map())[AutoShardDatasetOp::kAutoShardPolicy].set_i(
auto_shard_policy_);
// This parameter is used internally by tf.distribute to rebatch the dataset.
// It is not used outside the context of `experimental_distribute_dataset`.
(*config.mutable_parameter_map())[AutoShardDatasetOp::kNumReplicas].set_i(1);
return config;
}
absl::Status WorkerIndexResolver::ValidateWorker(
absl::string_view worker_address) const {
if (worker_addresses_.empty()) {
return absl::OkStatus();
}
for (absl::string_view config_address : worker_addresses_) {
if (config_address == worker_address ||
ShouldReplaceDynamicPort(config_address, worker_address)) {
return absl::OkStatus();
}
}
return absl::FailedPreconditionError(absl::Substitute(
"Failed to assign an index for worker $0. Configured workers list: [$1]. "
"The worker's address is not configured, or other workers are already "
"running at the configured host. If your worker has restarted, make sure "
"it runs at the same address and port.",
worker_address, absl::StrJoin(worker_addresses_, ", ")));
}
void WorkerIndexResolver::AddWorker(absl::string_view worker_address) {
for (std::string& config_address : worker_addresses_) {
if (config_address == worker_address) {
return;
}
if (ShouldReplaceDynamicPort(config_address, worker_address)) {
config_address = std::string(worker_address);
return;
}
}
}
absl::StatusOr<int64_t> WorkerIndexResolver::GetWorkerIndex(
absl::string_view worker_address) const {
const auto it = absl::c_find(worker_addresses_, worker_address);
if (it == worker_addresses_.cend()) {
return absl::NotFoundError(absl::Substitute(
"Failed to shard dataset in tf.data service: Worker $0 is not in the "
"workers list. Got workers list $1.",
worker_address, absl::StrJoin(worker_addresses_, ",")));
}
return std::distance(worker_addresses_.cbegin(), it);
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/graph_rewriters.h
+108
View File
@@ -0,0 +1,108 @@
/* Copyright 2021 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_CORE_DATA_SERVICE_GRAPH_REWRITERS_H_
#define TENSORFLOW_CORE_DATA_SERVICE_GRAPH_REWRITERS_H_
#include <cstdint>
#include <string>
#include <vector>
#include "absl/strings/string_view.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace data {
// Rewrites the dataset graph by removing the compression map.
class RemoveCompressionMapRewriter {
public:
// Returns `graph_def` with the compression map removed.
absl::StatusOr<GraphDef> ApplyRemoveCompressionMapRewrite(
const GraphDef& graph_def);
private:
tensorflow::RewriterConfig::CustomGraphOptimizer GetRewriteConfig() const;
};
// Rewrites the dataset graph by applying an auto-shard policy.
class AutoShardRewriter {
public:
// Creates an `AutoShardRewriter` according to `task_def`. Returns an error if
// the sharding policy is not a valid auto-shard policy.
static absl::StatusOr<AutoShardRewriter> Create(const TaskDef& task_def);
// Applies auto-sharding to `graph_def`. If auto-shard policy is OFF, returns
// the same graph as `graph_def`. Otherwise, returns the re-written graph.
absl::StatusOr<GraphDef> ApplyAutoShardRewrite(const GraphDef& graph_def);
private:
AutoShardRewriter(AutoShardPolicy auto_shard_policy, int64_t num_workers,
int64_t worker_index);
// Creates a rewrite config based on the auto-shard policy.
tensorflow::RewriterConfig::CustomGraphOptimizer GetRewriteConfig() const;
const AutoShardPolicy auto_shard_policy_;
const int64_t num_workers_;
const int64_t worker_index_;
};
// Maps a worker to its index, given a list of workers. For example, suppose
// `worker_addresses` contains
// /worker/task/0:worker, /worker/task/1:worker, /worker/task/2:worker,
// then
// /worker/task/0:worker maps to index 0,
// /worker/task/1:worker maps to index 1,
// /worker/task/2:worker maps to index 2.
// This is useful for deterministically sharding a dataset among a fixed set of
// tf.data service workers.
class WorkerIndexResolver {
public:
// Constructs a `WorkerIndexResolver` to generate worker indexes according to
// the specified worker addresses. The worker addresses can be "host" or
// "host:port", where "port" is a number, named port, or "%port%" to be
// replaced with the actual port.
template <class T>
explicit WorkerIndexResolver(const T& worker_addresses)
: worker_addresses_(worker_addresses.cbegin(), worker_addresses.cend()) {}
// Validates `worker_address`. Returns an error if the `worker_addresses` list
// is non-empty and `worker_address` is not specified in the worker addresses
// list (with optional port replacement).
absl::Status ValidateWorker(absl::string_view worker_address) const;
// Processes a worker at address `worker_address`. Its index can be retrieved
// by calling `GetWorkerIndex`.
void AddWorker(absl::string_view worker_address);
// Returns the worker index for the worker at `worker_address`. Returns a
// NotFound error if the worker is not registered.
absl::StatusOr<int64_t> GetWorkerIndex(
absl::string_view worker_address) const;
private:
std::vector<std::string> worker_addresses_;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_GRAPH_REWRITERS_H_
tensorflow/core/data/service/graph_rewriters_test.cc
+533
View File
@@ -0,0 +1,533 @@
/* Copyright 2021 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/core/data/service/graph_rewriters.h"
#include <string>
#include "absl/strings/string_view.h"
#include "absl/strings/substitute.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/framework/dataset_options.pb.h"
#include "tensorflow/core/framework/function_testlib.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::tensorflow::data::testing::EqualsProto;
using ::tensorflow::data::testing::RangeDatasetWithShardHint;
using ::tensorflow::data::testing::RangeSquareDataset;
using ::tensorflow::testing::IsOkAndHolds;
using ::tensorflow::testing::StatusIs;
using ::testing::HasSubstr;
using ::testing::SizeIs;
absl::StatusOr<NodeDef> GetNode(const GraphDef& graph_def,
absl::string_view name) {
for (const NodeDef& node : graph_def.node()) {
if (node.name() == name) {
return node;
}
}
return absl::NotFoundError(absl::Substitute(
"Node $0 not found in graph $1.", name, graph_def.ShortDebugString()));
}
absl::StatusOr<int64_t> GetValue(const GraphDef& graph_def,
absl::string_view name) {
for (const NodeDef& node : graph_def.node()) {
if (node.name() == name) {
return node.attr().at("value").tensor().int64_val()[0];
}
}
return absl::NotFoundError(absl::Substitute(
"Node $0 not found in graph $1.", name, graph_def.ShortDebugString()));
}
TaskDef GetTaskDef(const ProcessingModeDef::ShardingPolicy sharding_policy,
const int64_t num_workers, const int64_t worker_index) {
TaskDef task_def;
task_def.mutable_processing_mode_def()->set_sharding_policy(sharding_policy);
task_def.set_num_workers(num_workers);
task_def.set_worker_index(worker_index);
return task_def;
}
TEST(AutoShardRewriterTest, AutoShard) {
TaskDef task_def = GetTaskDef(ProcessingModeDef::FILE_OR_DATA,
/*num_workers=*/3, /*worker_index=*/1);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
TF_ASSERT_OK_AND_ASSIGN(GraphDef rewritten_graph,
rewriter.ApplyAutoShardRewrite(dataset.graph()));
TF_ASSERT_OK_AND_ASSIGN(NodeDef shard_node,
GetNode(rewritten_graph, "ShardDataset"));
ASSERT_THAT(shard_node.input(), SizeIs(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(1)),
absl_testing::IsOkAndHolds(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(2)),
absl_testing::IsOkAndHolds(1));
}
TEST(AutoShardRewriterTest, ShardByData) {
TaskDef task_def = GetTaskDef(ProcessingModeDef::DATA, /*num_workers=*/3,
/*worker_index=*/1);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
TF_ASSERT_OK_AND_ASSIGN(GraphDef rewritten_graph,
rewriter.ApplyAutoShardRewrite(dataset.graph()));
TF_ASSERT_OK_AND_ASSIGN(NodeDef shard_node,
GetNode(rewritten_graph, "ShardDataset"));
ASSERT_THAT(shard_node.input(), SizeIs(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(1)),
absl_testing::IsOkAndHolds(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(2)),
absl_testing::IsOkAndHolds(1));
}
TEST(AutoShardRewriterTest, ShardByFile) {
TaskDef task_def = GetTaskDef(ProcessingModeDef::FILE, /*num_workers=*/3,
/*worker_index=*/1);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
EXPECT_THAT(
rewriter.ApplyAutoShardRewrite(dataset.graph()),
absl_testing::StatusIs(error::NOT_FOUND,
HasSubstr("Found an unshardable source dataset")));
}
TEST(AutoShardRewriterTest, ShardByHint) {
TaskDef task_def = GetTaskDef(ProcessingModeDef::HINT, /*num_workers=*/3,
/*worker_index=*/1);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeDatasetWithShardHint(10);
TF_ASSERT_OK_AND_ASSIGN(GraphDef rewritten_graph,
rewriter.ApplyAutoShardRewrite(dataset.graph()));
TF_ASSERT_OK_AND_ASSIGN(NodeDef shard_node,
GetNode(rewritten_graph, "ShardDataset"));
ASSERT_THAT(shard_node.input(), SizeIs(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(1)),
absl_testing::IsOkAndHolds(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(2)),
absl_testing::IsOkAndHolds(1));
}
TEST(AutoShardRewriterTest, NoShard) {
TaskDef task_def =
GetTaskDef(ProcessingModeDef::OFF, /*num_workers=*/3, /*worker_index=*/1);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
EXPECT_THAT(rewriter.ApplyAutoShardRewrite(dataset.graph()),
absl_testing::IsOkAndHolds(EqualsProto(dataset.graph())));
}
TEST(AutoShardRewriterTest, EmptyDataset) {
TaskDef task_def =
GetTaskDef(ProcessingModeDef::FILE_OR_DATA, /*num_workers=*/3,
/*worker_index=*/1);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(0);
TF_ASSERT_OK_AND_ASSIGN(GraphDef rewritten_graph,
rewriter.ApplyAutoShardRewrite(dataset.graph()));
TF_ASSERT_OK_AND_ASSIGN(NodeDef shard_node,
GetNode(rewritten_graph, "ShardDataset"));
ASSERT_THAT(shard_node.input(), SizeIs(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(1)),
absl_testing::IsOkAndHolds(3));
EXPECT_THAT(GetValue(rewritten_graph, shard_node.input(2)),
absl_testing::IsOkAndHolds(1));
}
TEST(AutoShardRewriterTest, NoWorkers) {
TaskDef task_def =
GetTaskDef(ProcessingModeDef::FILE_OR_DATA, /*num_workers=*/0,
/*worker_index=*/0);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
EXPECT_THAT(
rewriter.ApplyAutoShardRewrite(dataset.graph()),
absl_testing::StatusIs(error::INVALID_ARGUMENT,
"num_workers should be >= 1, currently 0"));
}
TEST(AutoShardRewriterTest, NoWorkersWhenShardIsOff) {
TaskDef task_def =
GetTaskDef(ProcessingModeDef::OFF, /*num_workers=*/0, /*worker_index=*/0);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
EXPECT_THAT(rewriter.ApplyAutoShardRewrite(dataset.graph()),
absl_testing::IsOkAndHolds(EqualsProto(dataset.graph())));
}
TEST(AutoShardRewriterTest, WorkerIndexOutOfRange) {
TaskDef task_def =
GetTaskDef(ProcessingModeDef::FILE_OR_DATA, /*num_workers=*/2,
/*worker_index=*/5);
TF_ASSERT_OK_AND_ASSIGN(AutoShardRewriter rewriter,
AutoShardRewriter::Create(task_def));
DatasetDef dataset = RangeSquareDataset(10);
EXPECT_THAT(
rewriter.ApplyAutoShardRewrite(dataset.graph()),
absl_testing::StatusIs(error::INVALID_ARGUMENT,
"index should be >= 0 and < 2, currently 5"));
}
TEST(WorkerIndexResolverTest, AddOneWorker) {
WorkerIndexResolver resolver(std::vector<std::string>{"localhost"});
EXPECT_THAT(resolver.GetWorkerIndex("localhost:12345"),
absl_testing::StatusIs(error::NOT_FOUND));
TF_EXPECT_OK(resolver.ValidateWorker("localhost:12345"));
resolver.AddWorker("localhost:12345");
EXPECT_THAT(resolver.GetWorkerIndex("localhost:12345"),
absl_testing::IsOkAndHolds(0));
}
TEST(WorkerIndexResolverTest, AddMultipleWorkers) {
WorkerIndexResolver resolver(std::vector<std::string>{
"/worker/task/0", "/worker/task/1", "/worker/task/2"});
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:34567"));
resolver.AddWorker("/worker/task/2:12345");
resolver.AddWorker("/worker/task/1:23456");
resolver.AddWorker("/worker/task/0:34567");
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:34567"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:12345"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, NamedPorts) {
WorkerIndexResolver resolver(
std::vector<std::string>{"/worker/task/0:worker", "/worker/task/1:worker",
"/worker/task/2:worker"});
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:worker"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:worker"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:worker"));
resolver.AddWorker("/worker/task/2:worker");
resolver.AddWorker("/worker/task/1:worker");
resolver.AddWorker("/worker/task/0:worker");
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:worker"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:worker"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:worker"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, DynamicPorts) {
WorkerIndexResolver resolver(std::vector<std::string>{
"/worker/task/0:%port_worker%", "/worker/task/1:%port_worker%",
"/worker/task/2:%port_worker%"});
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:worker"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:worker"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:worker"));
resolver.AddWorker("/worker/task/2:worker");
resolver.AddWorker("/worker/task/1:worker");
resolver.AddWorker("/worker/task/0:worker");
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:worker"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:worker"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:worker"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, AnonymousPorts) {
WorkerIndexResolver resolver(
std::vector<std::string>{"/worker/task/0:%port%", "/worker/task/1:%port%",
"/worker/task/2:%port%"});
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:10000"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:10001"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:10002"));
resolver.AddWorker("/worker/task/2:10000");
resolver.AddWorker("/worker/task/1:10001");
resolver.AddWorker("/worker/task/0:10002");
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:10002"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:10001"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:10000"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, NumericPorts) {
WorkerIndexResolver resolver(std::vector<std::string>{
"/worker/task/0:12345", "/worker/task/1:23456", "/worker/task/2:34567"});
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:34567"),
absl_testing::IsOkAndHolds(2));
// Adding duplicate workers is a no-op.
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:34567"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:12345"));
resolver.AddWorker("/worker/task/2:34567");
resolver.AddWorker("/worker/task/1:23456");
resolver.AddWorker("/worker/task/0:12345");
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, IPv6Addresses) {
WorkerIndexResolver resolver(std::vector<std::string>{
"[1080:0:0:0:8:800:200C:417A]", "[1080:0:0:0:8:800:200C:417B]",
"[1080:0:0:0:8:800:200C:417C]"});
TF_EXPECT_OK(resolver.ValidateWorker("[1080:0:0:0:8:800:200C:417A]:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("[1080:0:0:0:8:800:200C:417B]:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("[1080:0:0:0:8:800:200C:417C]:34567"));
resolver.AddWorker("[1080:0:0:0:8:800:200C:417A]:12345");
resolver.AddWorker("[1080:0:0:0:8:800:200C:417B]:23456");
resolver.AddWorker("[1080:0:0:0:8:800:200C:417C]:34567");
EXPECT_THAT(resolver.GetWorkerIndex("[1080:0:0:0:8:800:200C:417A]:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("[1080:0:0:0:8:800:200C:417B]:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("[1080:0:0:0:8:800:200C:417C]:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, IPv6AddressesWithDynamicPort) {
WorkerIndexResolver resolver(
std::vector<std::string>{"[1080:0:0:0:8:800:200C:417A]:%port%",
"[1080:0:0:0:8:800:200C:417B]:%port%",
"[1080:0:0:0:8:800:200C:417C]:%port%"});
TF_EXPECT_OK(resolver.ValidateWorker("[1080:0:0:0:8:800:200C:417A]:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("[1080:0:0:0:8:800:200C:417B]:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("[1080:0:0:0:8:800:200C:417C]:34567"));
resolver.AddWorker("[1080:0:0:0:8:800:200C:417A]:12345");
resolver.AddWorker("[1080:0:0:0:8:800:200C:417B]:23456");
resolver.AddWorker("[1080:0:0:0:8:800:200C:417C]:34567");
EXPECT_THAT(resolver.GetWorkerIndex("[1080:0:0:0:8:800:200C:417A]:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("[1080:0:0:0:8:800:200C:417B]:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("[1080:0:0:0:8:800:200C:417C]:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, AddressesWithProtocols) {
WorkerIndexResolver resolver(std::vector<std::string>{
"http://127.0.0.1", "http://127.0.0.1", "http://127.0.0.1"});
TF_EXPECT_OK(resolver.ValidateWorker("http://127.0.0.1:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("http://127.0.0.1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("http://127.0.0.1:34567"));
resolver.AddWorker("http://127.0.0.1:12345");
resolver.AddWorker("http://127.0.0.1:23456");
resolver.AddWorker("http://127.0.0.1:34567");
EXPECT_THAT(resolver.GetWorkerIndex("http://127.0.0.1:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("http://127.0.0.1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("http://127.0.0.1:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, AddressesWithProtocolsAndDynamicPorts) {
WorkerIndexResolver resolver(std::vector<std::string>{
"http://127.0.0.1:%port_name%", "http://127.0.0.1:%port_name%",
"http://127.0.0.1:%port_name%"});
TF_EXPECT_OK(resolver.ValidateWorker("http://127.0.0.1:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("http://127.0.0.1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("http://127.0.0.1:34567"));
resolver.AddWorker("http://127.0.0.1:12345");
resolver.AddWorker("http://127.0.0.1:23456");
resolver.AddWorker("http://127.0.0.1:34567");
EXPECT_THAT(resolver.GetWorkerIndex("http://127.0.0.1:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("http://127.0.0.1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("http://127.0.0.1:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, HostNameHasColons) {
WorkerIndexResolver resolver(
std::vector<std::string>{":worker:task:0:%port%", ":worker:task:1:%port%",
":worker:task:2:34567"});
TF_EXPECT_OK(resolver.ValidateWorker(":worker:task:0:12345"));
TF_EXPECT_OK(resolver.ValidateWorker(":worker:task:1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker(":worker:task:2:34567"));
resolver.AddWorker(":worker:task:0:12345");
resolver.AddWorker(":worker:task:1:23456");
resolver.AddWorker(":worker:task:2:34567");
EXPECT_THAT(resolver.GetWorkerIndex(":worker:task:0:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex(":worker:task:1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex(":worker:task:2:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, ChangeWorkerPort) {
WorkerIndexResolver resolver(std::vector<std::string>{
"/worker/task/0", "/worker/task/1", "/worker/task/2"});
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:34567"));
resolver.AddWorker("/worker/task/2:12345");
resolver.AddWorker("/worker/task/1:23456");
resolver.AddWorker("/worker/task/0:34567");
EXPECT_THAT(resolver.ValidateWorker("/worker/task/0:99999"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("already running at the configured host")));
EXPECT_THAT(resolver.ValidateWorker("/worker/task/1:99999"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("already running at the configured host")));
EXPECT_THAT(resolver.ValidateWorker("/worker/task/2:99999"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("already running at the configured host")));
}
TEST(WorkerIndexResolverTest, WorkerNotFound) {
WorkerIndexResolver resolver(std::vector<std::string>{
"/worker/task/0", "/worker/task/1", "/worker/task/2"});
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:34567"),
absl_testing::StatusIs(error::NOT_FOUND));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:23456"),
absl_testing::StatusIs(error::NOT_FOUND));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:12345"),
absl_testing::StatusIs(error::NOT_FOUND));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/3:45678"),
absl_testing::StatusIs(error::NOT_FOUND));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/2:12345"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/1:23456"));
TF_EXPECT_OK(resolver.ValidateWorker("/worker/task/0:34567"));
EXPECT_THAT(resolver.ValidateWorker("/worker/task/3:45678"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("The worker's address is not configured")));
resolver.AddWorker("/worker/task/3:45678");
resolver.AddWorker("/worker/task/2:12345");
resolver.AddWorker("/worker/task/1:23456");
resolver.AddWorker("/worker/task/0:34567");
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/0:34567"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/1:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("/worker/task/2:12345"),
absl_testing::IsOkAndHolds(2));
EXPECT_THAT(
resolver.GetWorkerIndex("/worker/task/3:45678"),
absl_testing::StatusIs(
error::NOT_FOUND,
HasSubstr(
"Worker /worker/task/3:45678 is not in the workers list.")));
}
TEST(WorkerIndexResolverTest, MultipleWorkersInOneHost) {
WorkerIndexResolver resolver(
std::vector<std::string>{"localhost", "localhost", "localhost"});
TF_EXPECT_OK(resolver.ValidateWorker("localhost:12345"));
resolver.AddWorker("localhost:12345");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:23456"));
resolver.AddWorker("localhost:23456");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:34567"));
resolver.AddWorker("localhost:34567");
EXPECT_THAT(resolver.GetWorkerIndex("localhost:12345"),
absl_testing::IsOkAndHolds(0));
EXPECT_THAT(resolver.GetWorkerIndex("localhost:23456"),
absl_testing::IsOkAndHolds(1));
EXPECT_THAT(resolver.GetWorkerIndex("localhost:34567"),
absl_testing::IsOkAndHolds(2));
}
TEST(WorkerIndexResolverTest, MoreWorkersThanConfigured) {
WorkerIndexResolver resolver(std::vector<std::string>{
"localhost:%port%", "localhost:%port%", "localhost:%port%"});
TF_EXPECT_OK(resolver.ValidateWorker("localhost:12345"));
resolver.AddWorker("localhost:12345");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:23456"));
resolver.AddWorker("localhost:23456");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:34567"));
resolver.AddWorker("localhost:34567");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:12345"));
resolver.AddWorker("localhost:12345");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:23456"));
resolver.AddWorker("localhost:23456");
TF_EXPECT_OK(resolver.ValidateWorker("localhost:34567"));
resolver.AddWorker("localhost:34567");
EXPECT_THAT(resolver.ValidateWorker("localhost:45678"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("already running at the configured host")));
EXPECT_THAT(resolver.ValidateWorker("localhost:56789"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("already running at the configured host")));
}
TEST(WorkerIndexResolverTest, WorkerNotConfigured) {
WorkerIndexResolver resolver(std::vector<std::string>{""});
EXPECT_THAT(resolver.GetWorkerIndex("localhost:12345"),
absl_testing::StatusIs(error::NOT_FOUND));
EXPECT_THAT(resolver.ValidateWorker("localhost:12345"),
absl_testing::StatusIs(
error::FAILED_PRECONDITION,
HasSubstr("The worker's address is not configured")));
resolver.AddWorker("localhost:12345");
EXPECT_THAT(resolver.GetWorkerIndex("localhost:12345"),
absl_testing::StatusIs(error::NOT_FOUND));
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/grpc_dispatcher_impl.cc
+75
View File
@@ -0,0 +1,75 @@
/* 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 "tensorflow/core/data/service/grpc_dispatcher_impl.h"
#include "grpcpp/server_context.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
using ::grpc::ServerBuilder;
using ::grpc::ServerContext;
GrpcDispatcherImpl::GrpcDispatcherImpl(
const experimental::DispatcherConfig& config, ServerBuilder& server_builder)
: impl_(config) {
server_builder.RegisterService(this);
VLOG(1) << "Registered data service dispatcher";
}
absl::Status GrpcDispatcherImpl::Start() { return impl_.Start(); }
void GrpcDispatcherImpl::Stop() { impl_.Stop(); }
size_t GrpcDispatcherImpl::NumActiveIterations() {
return impl_.NumActiveIterations();
}
DispatcherStateExport GrpcDispatcherImpl::ExportState() const {
return impl_.ExportState();
}
#define HANDLER(method) \
grpc::Status GrpcDispatcherImpl::method(ServerContext* context, \
const method##Request* request, \
method##Response* response) { \
return ToGrpcStatus(impl_.method(request, response)); \
}
HANDLER(WorkerHeartbeat);
HANDLER(WorkerUpdate);
HANDLER(GetDatasetDef);
HANDLER(GetSplit);
HANDLER(GetVersion);
HANDLER(GetOrRegisterDataset);
HANDLER(ReleaseIterationClient);
HANDLER(MaybeRemoveTask);
HANDLER(GetOrCreateJob);
HANDLER(GetOrCreateIteration);
HANDLER(ClientHeartbeat);
HANDLER(GetWorkers);
HANDLER(GetDataServiceMetadata);
HANDLER(GetDataServiceConfig);
HANDLER(Snapshot);
HANDLER(GetSnapshotSplit);
HANDLER(GetSnapshotStreams);
HANDLER(DisableCompressionAtRuntime);
#undef HANDLER
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/grpc_dispatcher_impl.h
+78
View File
@@ -0,0 +1,78 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_SERVICE_GRPC_DISPATCHER_IMPL_H_
#define TENSORFLOW_CORE_DATA_SERVICE_GRPC_DISPATCHER_IMPL_H_
#include "grpcpp/server_builder.h"
#include "tensorflow/core/data/service/dispatcher.grpc.pb.h"
#include "tensorflow/core/data/service/dispatcher_impl.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
// This class is a wrapper that handles communication for gRPC.
class GrpcDispatcherImpl : public DispatcherService::Service {
public:
// Constructs a GrpcDispatcherImpl with the given config, and registers it
// with `server_builder`.
explicit GrpcDispatcherImpl(const experimental::DispatcherConfig& config,
::grpc::ServerBuilder& server_builder);
~GrpcDispatcherImpl() override { Stop(); }
absl::Status Start();
void Stop();
size_t NumActiveIterations();
DispatcherStateExport ExportState() const;
#define HANDLER(method) \
::grpc::Status method(::grpc::ServerContext* context, \
const method##Request* request, \
method##Response* response) override;
HANDLER(WorkerHeartbeat);
HANDLER(WorkerUpdate);
HANDLER(GetDatasetDef);
HANDLER(GetSplit);
HANDLER(GetVersion);
HANDLER(GetOrRegisterDataset);
HANDLER(ReleaseIterationClient);
HANDLER(MaybeRemoveTask);
HANDLER(GetOrCreateJob);
HANDLER(GetOrCreateIteration);
HANDLER(ClientHeartbeat);
HANDLER(GetWorkers);
HANDLER(GetDataServiceMetadata);
HANDLER(GetDataServiceConfig);
HANDLER(Snapshot);
HANDLER(GetSnapshotSplit);
HANDLER(GetSnapshotStreams);
HANDLER(DisableCompressionAtRuntime);
#undef HANDLER
private:
DataServiceDispatcherImpl impl_;
GrpcDispatcherImpl(const GrpcDispatcherImpl&) = delete;
void operator=(const GrpcDispatcherImpl&) = delete;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_GRPC_DISPATCHER_IMPL_H_
tensorflow/core/data/service/grpc_dispatcher_impl_test.cc
+180
View File
@@ -0,0 +1,180 @@
/* Copyright 2021 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/core/data/service/grpc_dispatcher_impl.h"
#include <cstdint>
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include "grpcpp/channel.h"
#include "grpcpp/client_context.h"
#include "grpcpp/create_channel.h"
#include "grpcpp/security/credentials.h"
#include "grpcpp/support/channel_arguments.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/credentials_factory.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/server_lib.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::grpc::Channel;
using ::grpc::ChannelArguments;
using ::grpc::ChannelCredentials;
using ::grpc::ClientContext;
constexpr const char kHostAddress[] = "localhost";
constexpr const char kProtocol[] = "grpc";
class GrpcDispatcherImplTest : public ::testing::Test {
protected:
void SetUp() override {
TF_ASSERT_OK(SetUpDispatcherServer());
TF_ASSERT_OK(SetUpDispatcherClientStub());
}
absl::Status SetUpDispatcherServer() {
experimental::DispatcherConfig config;
config.set_protocol(kProtocol);
TF_RETURN_IF_ERROR(NewDispatchServer(config, dispatcher_server_));
return dispatcher_server_->Start();
}
absl::Status SetUpDispatcherClientStub() {
std::shared_ptr<ChannelCredentials> credentials;
TF_RETURN_IF_ERROR(
CredentialsFactory::CreateClientCredentials(kProtocol, &credentials));
ChannelArguments args;
args.SetMaxReceiveMessageSize(std::numeric_limits<int32_t>::max());
args.SetInt(GRPC_ARG_USE_LOCAL_SUBCHANNEL_POOL, true);
std::shared_ptr<Channel> channel =
::grpc::CreateCustomChannel(GetDispatcherAddress(), credentials, args);
dispatcher_client_stub_ = DispatcherService::NewStub(channel);
return absl::OkStatus();
}
std::string GetDispatcherAddress() const {
return absl::StrCat(kHostAddress, ":", dispatcher_server_->BoundPort());
}
std::unique_ptr<DispatchGrpcDataServer> dispatcher_server_;
std::unique_ptr<DispatcherService::Stub> dispatcher_client_stub_;
};
TEST_F(GrpcDispatcherImplTest, GrpcTest) {
ClientContext ctx;
GetVersionRequest req;
GetVersionResponse resp;
TF_ASSERT_OK(
FromGrpcStatus(dispatcher_client_stub_->GetVersion(&ctx, req, &resp)));
EXPECT_EQ(resp.version(), kDataServiceVersion);
}
TEST_F(GrpcDispatcherImplTest, GetSplitInvalidProviderIndex) {
// Register a dataset.
std::string dataset_id;
{
ClientContext ctx;
GetOrRegisterDatasetRequest req;
*req.mutable_dataset()->mutable_graph() = testing::RangeDataset(10).graph();
GetOrRegisterDatasetResponse resp;
TF_ASSERT_OK(FromGrpcStatus(
dispatcher_client_stub_->GetOrRegisterDataset(&ctx, req, &resp)));
dataset_id = resp.dataset_id();
}
// Create an iteration.
int64_t iteration_id = -1;
{
ClientContext ctx;
GetOrCreateJobRequest job_req;
job_req.set_dataset_id(dataset_id);
job_req.mutable_processing_mode_def()->set_sharding_policy(
ProcessingModeDef::DYNAMIC);
GetOrCreateJobResponse job_resp;
TF_ASSERT_OK(FromGrpcStatus(
dispatcher_client_stub_->GetOrCreateJob(&ctx, job_req, &job_resp)));
ClientContext iter_ctx;
GetOrCreateIterationRequest req;
req.set_job_id(job_resp.job_id());
req.set_repetition(0);
GetOrCreateIterationResponse resp;
TF_ASSERT_OK(FromGrpcStatus(
dispatcher_client_stub_->GetOrCreateIteration(&iter_ctx, req, &resp)));
// Dynamically find iteration_id
for (int i = 0; i < 10000; ++i) {
ClientContext ctx;
GetSplitRequest split_req;
split_req.set_iteration_id(i);
split_req.set_repetition(0);
split_req.set_split_provider_index(0);
GetSplitResponse split_resp;
::grpc::Status status =
dispatcher_client_stub_->GetSplit(&ctx, split_req, &split_resp);
if (status.error_code() != ::grpc::StatusCode::NOT_FOUND) {
iteration_id = i;
break;
}
}
ASSERT_NE(iteration_id, -1)
<< "Could not find iteration_id. Iteration was not created correctly?";
}
// GetSplit with invalid provider index (too small).
{
ClientContext ctx;
GetSplitRequest req;
req.set_iteration_id(iteration_id);
req.set_repetition(0);
req.set_split_provider_index(-1);
GetSplitResponse resp;
::grpc::Status status = dispatcher_client_stub_->GetSplit(&ctx, req, &resp);
EXPECT_EQ(status.error_code(), ::grpc::StatusCode::INVALID_ARGUMENT);
}
// GetSplit with invalid provider index (too large).
{
ClientContext ctx;
GetSplitRequest req;
req.set_iteration_id(iteration_id);
req.set_repetition(0);
req.set_split_provider_index(1000000);
GetSplitResponse resp;
::grpc::Status status = dispatcher_client_stub_->GetSplit(&ctx, req, &resp);
EXPECT_EQ(status.error_code(), ::grpc::StatusCode::INVALID_ARGUMENT);
}
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/grpc_util.cc
+97
View File
@@ -0,0 +1,97 @@
/* Copyright 2020 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/core/data/service/grpc_util.h"
#include <algorithm>
#include <functional>
#include <string>
#include "absl/time/time.h"
#include "tensorflow/core/data/service/common.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/env_time.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/retrying_utils.h"
namespace tensorflow {
namespace data {
namespace grpc_util {
constexpr char kStreamRemovedMessage[] = "Stream removed";
absl::Status WrapError(const std::string& message,
const ::grpc::Status& status) {
if (status.ok()) {
return absl::InternalError(absl::StrCat(
"Expected a non-ok grpc status. Wrapping message: ", message));
} else {
// FromGrpcStatus checks for "Stream removed" as well, but only when the
// status code is "Unknown". We have observed that sometimes stream removed
// errors use other status codes (b/258285154).
// TODO(aaudibert): Upstream this to FromGrpcStatus.
if (status.error_message() == kStreamRemovedMessage) {
return absl::Status(absl::StatusCode::kUnavailable,
kStreamRemovedMessage);
}
absl::Status s = FromGrpcStatus(status);
return absl::Status(s.code(),
absl::StrCat(message, ": ", status.error_message()));
}
}
absl::Status Retry(const std::function<absl::Status()>& f,
const std::string& description, int64_t deadline_micros) {
return Retry(
f, [] { return true; }, description, deadline_micros);
}
absl::Status Retry(const std::function<absl::Status()>& f,
const std::function<bool()>& should_retry,
const std::string& description, int64_t deadline_micros) {
absl::Status s = f();
for (int num_retries = 0;; ++num_retries) {
if (!IsPreemptedError(s)) {
return s;
}
int64_t now_micros = EnvTime::NowMicros();
if (now_micros > deadline_micros || !should_retry()) {
return s;
}
int64_t deadline_with_backoff_micros =
now_micros +
absl::ToInt64Microseconds(tsl::ComputeRetryBackoff(num_retries));
// Wait for a short period of time before retrying. If our backoff would put
// us past the deadline, we truncate it to ensure our attempt starts before
// the deadline.
int64_t backoff_until =
std::min(deadline_with_backoff_micros, deadline_micros);
int64_t wait_time_micros = backoff_until - now_micros;
if (wait_time_micros > 100 * 1000) {
LOG(INFO) << "Failed to " << description << ": " << s
<< ". Will retry in " << wait_time_micros / 1000 << "ms.";
}
Env::Default()->SleepForMicroseconds(wait_time_micros);
s = f();
}
return s;
}
} // namespace grpc_util
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/grpc_util.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2020 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_CORE_DATA_SERVICE_GRPC_UTIL_H_
#define TENSORFLOW_CORE_DATA_SERVICE_GRPC_UTIL_H_
#include <functional>
#include <string>
#include "grpcpp/grpcpp.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace data {
namespace grpc_util {
// Wraps a grpc::Status in a tensorflow::Status with the given message.
absl::Status WrapError(const std::string& message,
const ::grpc::Status& status);
// Retries the given function if the function produces UNAVAILABLE, ABORTED, or
// CANCELLED status codes. We retry these codes because they can all indicate
// preemption of a server. The retries continue until the deadline is exceeded
// or the `should_retry` callback returns false. `description` may be used to
// log that retries are happening. It should contain a description of the action
// being retried, e.g. "register dataset" The retry loop uses exponential
// backoff between retries. `deadline_micros` is interpreted as microseconds
// since the epoch.
absl::Status Retry(const std::function<absl::Status()>& f,
const std::function<bool()>& should_retry,
const std::string& description, int64_t deadline_micros);
// Same as `Retry` above, but with a `should_retry` callback that always returns
// `true`.
absl::Status Retry(const std::function<absl::Status()>& f,
const std::string& description, int64_t deadline_micros);
} // namespace grpc_util
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_GRPC_UTIL_H_
tensorflow/core/data/service/grpc_util_test.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2020 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/core/data/service/grpc_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace data {
namespace grpc_util {
TEST(GrpcUtil, WrapInvalidArgument) {
grpc::Status s(grpc::StatusCode::INVALID_ARGUMENT, "test message");
absl::Status wrapped = WrapError("wrapping message", s);
ASSERT_EQ(wrapped,
absl::InvalidArgumentError("wrapping message: test message"));
}
TEST(GrpcUtil, WrapOk) {
grpc::Status s;
absl::Status wrapped = WrapError("wrapping message", s);
ASSERT_EQ(wrapped,
absl::InternalError("Expected a non-ok grpc status. Wrapping "
"message: wrapping message"));
}
} // namespace grpc_util
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/grpc_worker_impl.cc
+75
View File
@@ -0,0 +1,75 @@
/* 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 "tensorflow/core/data/service/grpc_worker_impl.h"
#include <memory>
#include <string>
#include <vector>
#include "grpcpp/server_builder.h"
#include "grpcpp/server_context.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/data/service/worker_impl.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
using ::grpc::ServerBuilder;
using ::grpc::ServerContext;
GrpcWorkerImpl::GrpcWorkerImpl(const experimental::WorkerConfig& config,
ServerBuilder& server_builder)
: impl_(std::make_shared<DataServiceWorkerImpl>(config)) {
server_builder.RegisterService(this);
VLOG(1) << "Registered data service worker";
}
absl::Status GrpcWorkerImpl::Start(
const std::string& worker_address,
const std::vector<DataTransferServerInfo>& transfer_servers) {
worker_address_ = worker_address;
TF_RETURN_IF_ERROR(impl_->Start(worker_address, transfer_servers));
LocalWorkers::Add(worker_address, impl_);
return absl::OkStatus();
}
void GrpcWorkerImpl::Stop() {
LocalWorkers::Remove(worker_address_);
impl_->Stop();
}
WorkerStateExport GrpcWorkerImpl::ExportState() const {
return impl_->ExportState();
}
#define HANDLER(method) \
::grpc::Status GrpcWorkerImpl::method(ServerContext* context, \
const method##Request* request, \
method##Response* response) { \
return ToGrpcStatus(impl_->method(request, response)); \
}
HANDLER(ProcessTask);
HANDLER(GetElement);
HANDLER(GetWorkerTasks);
HANDLER(GetSnapshotTaskProgresses);
#undef HANDLER
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/grpc_worker_impl.h
+81
View File
@@ -0,0 +1,81 @@
/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_SERVICE_GRPC_WORKER_IMPL_H_
#define TENSORFLOW_CORE_DATA_SERVICE_GRPC_WORKER_IMPL_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "grpcpp/server_builder.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/data/service/worker.grpc.pb.h"
#include "tensorflow/core/data/service/worker.pb.h"
#include "tensorflow/core/data/service/worker_impl.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
// This class is a wrapper that handles communication for gRPC.
class GrpcWorkerImpl : public WorkerService::Service {
public:
// Constructs a GrpcWorkerImpl with the given config, and registers it with
// `server_builder`.
explicit GrpcWorkerImpl(const experimental::WorkerConfig& config,
::grpc::ServerBuilder& server_builder);
~GrpcWorkerImpl() override { Stop(); }
absl::Status Start(
const std::string& worker_address,
const std::vector<DataTransferServerInfo>& transfer_servers);
void Stop();
std::function<absl::Status(const GetElementRequest*, GetElementResult*)>
get_element_getter() {
return [this](const GetElementRequest* request, GetElementResult* result) {
return impl_->GetElementResult(request, result);
};
}
WorkerStateExport ExportState() const;
#define HANDLER(method) \
::grpc::Status method(::grpc::ServerContext* context, \
const method##Request* request, \
method##Response* response) override;
HANDLER(ProcessTask);
HANDLER(GetElement);
HANDLER(GetWorkerTasks);
HANDLER(GetSnapshotTaskProgresses);
#undef HANDLER
private:
std::string worker_address_;
// A std::shared_ptr allows clients to access local servers and directly call
// the servers' methods to avoid RPC calls and data copy.
std::shared_ptr<DataServiceWorkerImpl> impl_;
GrpcWorkerImpl(const GrpcWorkerImpl&) = delete;
void operator=(const GrpcWorkerImpl&) = delete;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_GRPC_WORKER_IMPL_H_
tensorflow/core/data/service/grpc_worker_impl_test.cc
+118
View File
@@ -0,0 +1,118 @@
/* Copyright 2021 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/core/data/service/grpc_worker_impl.h"
#include <limits>
#include <memory>
#include <string>
#include <utility>
#include "grpcpp/channel.h"
#include "grpcpp/client_context.h"
#include "grpcpp/create_channel.h"
#include "grpcpp/security/credentials.h"
#include "grpcpp/support/channel_arguments.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/credentials_factory.h"
#include "tensorflow/core/data/service/dispatcher.pb.h"
#include "tensorflow/core/data/service/server_lib.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/data/service/worker.pb.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::grpc::Channel;
using ::grpc::ChannelArguments;
using ::grpc::ChannelCredentials;
using ::grpc::ClientContext;
constexpr const char kHostAddress[] = "localhost";
constexpr const char kProtocol[] = "grpc";
class GrpcWorkerImplTest : public ::testing::Test {
protected:
void SetUp() override {
TF_ASSERT_OK(SetUpDispatcherServer());
TF_ASSERT_OK(SetUpWorkerServer());
TF_ASSERT_OK(SetUpWorkerClientStub());
}
absl::Status SetUpDispatcherServer() {
experimental::DispatcherConfig config;
config.set_protocol(kProtocol);
TF_RETURN_IF_ERROR(NewDispatchServer(config, dispatcher_server_));
return dispatcher_server_->Start();
}
absl::Status SetUpWorkerServer() {
experimental::WorkerConfig config;
config.set_protocol(kProtocol);
config.set_dispatcher_address(GetDispatcherAddress());
config.set_worker_address(absl::StrCat(kHostAddress, ":%port%"));
TF_RETURN_IF_ERROR(NewWorkerServer(config, worker_server_));
return worker_server_->Start();
}
absl::Status SetUpWorkerClientStub() {
std::shared_ptr<ChannelCredentials> credentials;
TF_RETURN_IF_ERROR(
CredentialsFactory::CreateClientCredentials(kProtocol, &credentials));
ChannelArguments args;
args.SetMaxReceiveMessageSize(std::numeric_limits<int32_t>::max());
args.SetInt(GRPC_ARG_USE_LOCAL_SUBCHANNEL_POOL, true);
std::shared_ptr<Channel> channel =
::grpc::CreateCustomChannel(GetWorkerAddress(), credentials, args);
worker_client_stub_ = WorkerService::NewStub(channel);
return absl::OkStatus();
}
std::string GetDispatcherAddress() const {
return absl::StrCat(kHostAddress, ":", dispatcher_server_->BoundPort());
}
std::string GetWorkerAddress() const {
return absl::StrCat(kHostAddress, ":", worker_server_->BoundPort());
}
std::unique_ptr<DispatchGrpcDataServer> dispatcher_server_;
std::unique_ptr<WorkerGrpcDataServer> worker_server_;
std::unique_ptr<WorkerService::Stub> worker_client_stub_;
};
TEST_F(GrpcWorkerImplTest, GetWorkerTasks) {
ClientContext ctx;
GetWorkerTasksRequest req;
GetWorkerTasksResponse resp;
TF_ASSERT_OK(
FromGrpcStatus(worker_client_stub_->GetWorkerTasks(&ctx, req, &resp)));
EXPECT_EQ(resp.tasks_size(), 0);
}
} // namespace
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/journal.cc
+145
View File
@@ -0,0 +1,145 @@
/* Copyright 2020 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/core/data/service/journal.h"
#include <algorithm>
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "tensorflow/core/data/service/journal.pb.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/regexp.h"
namespace tensorflow {
namespace data {
namespace {
constexpr absl::string_view kJournal = "journal";
absl::Status ParseSequenceNumber(const std::string& journal_file,
int64_t* sequence_number) {
if (!RE2::FullMatch(journal_file, ".*_(\\d+)", sequence_number)) {
return absl::InvalidArgumentError(
absl::StrCat("Failed to parse journal file name: ", journal_file));
}
return absl::OkStatus();
}
} // namespace
std::string DataServiceJournalFile(const std::string& journal_dir,
int64_t sequence_number) {
return io::JoinPath(journal_dir,
absl::StrCat(kJournal, "_", sequence_number));
}
FileJournalWriter::FileJournalWriter(Env* env, const std::string& journal_dir)
: env_(env), journal_dir_(journal_dir) {}
absl::Status FileJournalWriter::EnsureInitialized() {
if (writer_) {
return absl::OkStatus();
}
std::vector<std::string> journal_files;
TF_RETURN_IF_ERROR(env_->RecursivelyCreateDir(journal_dir_));
TF_RETURN_IF_ERROR(env_->GetChildren(journal_dir_, &journal_files));
int64_t latest_sequence_number = -1;
for (const auto& file : journal_files) {
int64_t sequence_number;
TF_RETURN_IF_ERROR(ParseSequenceNumber(file, &sequence_number));
latest_sequence_number = std::max(latest_sequence_number, sequence_number);
}
std::string journal_file =
DataServiceJournalFile(journal_dir_, latest_sequence_number + 1);
TF_RETURN_IF_ERROR(env_->NewAppendableFile(journal_file, &file_));
writer_ = std::make_unique<io::RecordWriter>(file_.get());
VLOG(1) << "Created journal writer to write to " << journal_file;
return absl::OkStatus();
}
absl::Status FileJournalWriter::Write(const Update& update) {
TF_RETURN_IF_ERROR(EnsureInitialized());
std::string s = update.SerializeAsString();
if (s.empty()) {
return absl::InternalError(absl::StrCat(
"Failed to serialize update ", update.DebugString(), " to string"));
}
TF_RETURN_IF_ERROR(writer_->WriteRecord(s));
TF_RETURN_IF_ERROR(writer_->Flush());
TF_RETURN_IF_ERROR(file_->Sync());
if (VLOG_IS_ON(4)) {
VLOG(4) << "Wrote journal entry: " << update.DebugString();
}
return absl::OkStatus();
}
FileJournalReader::FileJournalReader(Env* env, absl::string_view journal_dir)
: env_(env), journal_dir_(journal_dir) {}
absl::Status FileJournalReader::EnsureInitialized() {
if (reader_) {
return absl::OkStatus();
}
return UpdateFile(DataServiceJournalFile(journal_dir_, 0));
}
absl::Status FileJournalReader::Read(Update& update, bool& end_of_journal) {
TF_RETURN_IF_ERROR(EnsureInitialized());
while (true) {
tstring record;
absl::Status s = reader_->ReadRecord(&record);
if (absl::IsOutOfRange(s)) {
sequence_number_++;
std::string next_journal_file =
DataServiceJournalFile(journal_dir_, sequence_number_);
if (absl::IsNotFound(env_->FileExists(next_journal_file))) {
VLOG(3) << "Next journal file " << next_journal_file
<< " does not exist. End of journal reached.";
end_of_journal = true;
return absl::OkStatus();
}
TF_RETURN_IF_ERROR(UpdateFile(next_journal_file));
continue;
}
TF_RETURN_IF_ERROR(s);
if (!update.ParseFromString(record)) {
return absl::DataLossError("Failed to parse journal record.");
}
if (VLOG_IS_ON(4)) {
VLOG(4) << "Read journal entry: " << update.DebugString();
}
end_of_journal = false;
return absl::OkStatus();
}
}
absl::Status FileJournalReader::UpdateFile(const std::string& filename) {
VLOG(1) << "Reading from journal file " << filename;
TF_RETURN_IF_ERROR(env_->NewRandomAccessFile(filename, &file_));
io::RecordReaderOptions opts;
opts.buffer_size = 2 << 20; // 2MB
reader_ = std::make_unique<io::SequentialRecordReader>(file_.get(), opts);
return absl::OkStatus();
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/journal.h
+118
View File
@@ -0,0 +1,118 @@
/* Copyright 2020 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_CORE_DATA_SERVICE_JOURNAL_H_
#define TENSORFLOW_CORE_DATA_SERVICE_JOURNAL_H_
#include <memory>
#include <string>
#include "tensorflow/core/data/service/journal.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/env.h"
namespace tensorflow {
namespace data {
// Returns the location of the journal file within the journal directory.
std::string DataServiceJournalFile(const std::string& journal_dir,
int64_t sequence_number);
// Interface for writing to a journal.
class JournalWriter {
public:
virtual ~JournalWriter() = default;
// Writes and syncs an update to the journal.
virtual absl::Status Write(const Update& update) = 0;
// Initializes the writer if it is not yet initialized.
virtual absl::Status EnsureInitialized() = 0;
};
// FileJournalWriter is not thread-safe, requiring external synchronization when
// used by multiple threads.
//
// FileJournalWriter writes journal files to a configured journal directory. The
// directory is laid out in the following format:
//
// journal_dir/
// journal_0
// journal_1
// ...
//
// When the writer is created, it lists the directory to find the next available
// journal file name. For example, if the journal directory contains
// "journal_0", "journal_1", and "journal_2", the writer will write to
// "journal_3". The writer will flush updates as they are written, so that they
// can be stored durably in case of machine failure.
class FileJournalWriter : public JournalWriter {
public:
// Creates a journal writer to write to the given journal directory.
// If there is already journal data there, the journal writer will append to
// the existing journal.
explicit FileJournalWriter(Env* env, const std::string& journal_dir);
FileJournalWriter(const FileJournalWriter&) = delete;
FileJournalWriter& operator=(const FileJournalWriter&) = delete;
absl::Status Write(const Update& update) override;
absl::Status EnsureInitialized() override;
private:
Env* env_;
const std::string journal_dir_;
std::unique_ptr<WritableFile> file_;
std::unique_ptr<io::RecordWriter> writer_;
};
// Interface for reading from a journal.
class JournalReader {
public:
virtual ~JournalReader() = default;
// Reads the next update from the journal. Sets `end_of_journal=true` if
// there are no more updates left in the journal.
virtual absl::Status Read(Update& update, bool& end_of_journal) = 0;
};
// JournalReader is not thread-safe, requiring external synchronization when
// used by multiple threads.
//
// The journal reader reads through all journal files in the configured journal
// directory, in order of their sequence numbers. See FileJournalWriter above.
class FileJournalReader : public JournalReader {
public:
explicit FileJournalReader(Env* env, absl::string_view journal_dir);
FileJournalReader(const FileJournalReader&) = delete;
FileJournalReader& operator=(const FileJournalReader&) = delete;
absl::Status Read(Update& update, bool& end_of_journal) override;
private:
// Initializes the reader if it is not yet initialized.
absl::Status EnsureInitialized();
// Updates the `FileJournalReader` to read from a new file.
absl::Status UpdateFile(const std::string& filename);
Env* env_;
const std::string journal_dir_;
// Sequence number of current journal file.
int64_t sequence_number_ = 0;
std::unique_ptr<RandomAccessFile> file_;
std::unique_ptr<io::SequentialRecordReader> reader_;
};
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_JOURNAL_H_
tensorflow/core/data/service/journal.proto
+176
View File
@@ -0,0 +1,176 @@
// Copyright 2026 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.
// ==============================================================================
syntax = "proto3";
package tensorflow.data;
import "tensorflow/core/data/service/common.proto";
import "tensorflow/core/protobuf/data_service.proto";
// Message representing journaled dispatcher metadata updates. When we apply
// one of these changes to the dispatcher's in-memory state, we also write an
// Update message to the journal.
// Next tag: 17
message Update {
oneof update_type {
RegisterDatasetUpdate register_dataset = 1;
RegisterWorkerUpdate register_worker = 5;
CreateJobUpdate create_job = 14;
CreateIterationUpdate create_iteration = 2;
ProduceSplitUpdate produce_split = 8;
AcquireIterationClientUpdate acquire_iteration_client = 6;
ReleaseIterationClientUpdate release_iteration_client = 7;
GarbageCollectIterationUpdate garbage_collect_iteration = 12;
RemoveTaskUpdate remove_task = 11;
CreatePendingTaskUpdate create_pending_task = 9;
ClientHeartbeatUpdate client_heartbeat = 10;
CreateTaskUpdate create_task = 3;
FinishTaskUpdate finish_task = 4;
SnapshotUpdate snapshot = 15;
CompressionDisabledAtRuntimeUpdate compression_disabled_at_runtime = 16;
}
reserved 13;
}
// Next tag: 5
message RegisterDatasetUpdate {
string dataset_id = 1;
uint64 fingerprint = 2;
DataServiceMetadata metadata = 3;
bool dedupe_by_dataset_id = 4;
}
// Next tag: 6
message RegisterWorkerUpdate {
string worker_address = 1;
repeated DataTransferServerInfo transfer_servers = 5;
repeated string worker_tags = 3;
int64 worker_uid = 4;
reserved 2;
}
// Next tag: 9
message CreateJobUpdate {
int64 job_id = 1;
string job_name = 2;
string dataset_id = 3;
ProcessingModeDef processing_mode_def = 4;
// Optional number of consumers. If set, the iteration's tasks will provide
// their elements to consumers round-robin.
oneof optional_num_consumers {
int64 num_consumers = 6;
}
// Specifies which workers the client of this iteration reads from.
TargetWorkers target_workers = 7;
// True if cross-trainer cache is enabled.
bool use_cross_trainer_cache = 8;
}
// Next tag: 5
message CreateIterationUpdate {
int64 iteration_id = 1;
int64 job_id = 2;
int64 repetition = 3;
int64 num_split_providers = 4;
}
// Next tag: 5
message ProduceSplitUpdate {
int64 iteration_id = 1;
int64 repetition = 2;
int64 split_provider_index = 4;
// Whether the split provider reached its end.
bool finished = 3;
}
// Next tag: 3
message AcquireIterationClientUpdate {
int64 iteration_id = 1;
int64 iteration_client_id = 2;
}
// Next tag: 3
message ReleaseIterationClientUpdate {
int64 iteration_client_id = 1;
// The time when the client was released, measured in microseconds since the
// epoch.
int64 time_micros = 2;
}
// Next tag: 2
message GarbageCollectIterationUpdate {
int64 iteration_id = 1;
}
// Next tag: 2
message RemoveTaskUpdate {
int64 task_id = 1;
}
// Indicates that a client failed to block before reaching the target round.
// Next tag: 2
message TaskRejected {
// A new target round to try adding the task in.
int64 new_target_round = 1;
}
// Updates dispatcher state based on a client heartbeat.
// Next tag: 4
message ClientHeartbeatUpdate {
int64 iteration_client_id = 1;
bool task_accepted = 2;
TaskRejected task_rejected = 3;
}
// Next tag: 9
message CreatePendingTaskUpdate {
int64 task_id = 1;
int64 iteration_id = 2;
string worker_address = 3;
repeated DataTransferServerInfo transfer_servers = 8;
repeated string worker_tags = 6;
int64 worker_uid = 7;
int64 starting_round = 5;
reserved 4;
}
// Next tag: 10
message CreateTaskUpdate {
reserved 3, 5;
int64 task_id = 1;
int64 iteration_id = 2;
string worker_address = 4;
repeated DataTransferServerInfo transfer_servers = 9;
repeated string worker_tags = 7;
int64 worker_uid = 8;
reserved 6;
}
// Next tag: 2
message FinishTaskUpdate {
int64 task_id = 1;
}
// Next tag: 2
message SnapshotUpdate {
string path = 1;
}
// Next tag: 3
message CompressionDisabledAtRuntimeUpdate {
string dataset_id = 1;
bool compression_disabled = 2;
}
tensorflow/core/data/service/journal_test.cc
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2020 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/core/data/service/journal.h"
#include <memory>
#include <string>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/journal.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/data_service.pb.h"
namespace tensorflow {
namespace data {
namespace {
using ::testing::HasSubstr;
bool NewJournalDir(std::string& journal_dir) {
std::string filename = testing::TmpDir();
if (!Env::Default()->CreateUniqueFileName(&filename, "journal_dir")) {
return false;
}
journal_dir = filename;
return true;
}
Update MakeCreateIterationUpdate() {
Update update;
CreateIterationUpdate* create_iteration = update.mutable_create_iteration();
create_iteration->set_job_id(3);
create_iteration->set_iteration_id(8);
create_iteration->set_repetition(5);
return update;
}
Update MakeFinishTaskUpdate() {
Update update;
FinishTaskUpdate* finish_task = update.mutable_finish_task();
finish_task->set_task_id(8);
return update;
}
Update MakeRegisterDatasetUpdate() {
Update update;
RegisterDatasetUpdate* register_dataset = update.mutable_register_dataset();
register_dataset->set_dataset_id("dataset_id");
register_dataset->set_fingerprint(3);
return update;
}
absl::Status CheckJournalContent(absl::string_view journal_dir,
const std::vector<Update>& expected) {
FileJournalReader reader(Env::Default(), journal_dir);
for (const auto& update : expected) {
Update result;
bool end_of_journal = true;
TF_RETURN_IF_ERROR(reader.Read(result, end_of_journal));
EXPECT_FALSE(end_of_journal);
// We can't use the testing::EqualsProto matcher because it is not available
// in OSS.
EXPECT_EQ(result.SerializeAsString(), update.SerializeAsString());
}
Update result;
bool end_of_journal = false;
TF_RETURN_IF_ERROR(reader.Read(result, end_of_journal));
EXPECT_TRUE(end_of_journal);
return absl::OkStatus();
}
} // namespace
TEST(Journal, RoundTripMultiple) {
std::string journal_dir;
EXPECT_TRUE(NewJournalDir(journal_dir));
std::vector<Update> updates = {MakeCreateIterationUpdate(),
MakeRegisterDatasetUpdate(),
MakeFinishTaskUpdate()};
FileJournalWriter writer(Env::Default(), journal_dir);
for (const auto& update : updates) {
TF_EXPECT_OK(writer.Write(update));
}
TF_EXPECT_OK(CheckJournalContent(journal_dir, updates));
}
TEST(Journal, AppendExistingJournal) {
std::string journal_dir;
EXPECT_TRUE(NewJournalDir(journal_dir));
std::vector<Update> updates = {MakeCreateIterationUpdate(),
MakeRegisterDatasetUpdate(),
MakeFinishTaskUpdate()};
for (const auto& update : updates) {
FileJournalWriter writer(Env::Default(), journal_dir);
TF_EXPECT_OK(writer.Write(update));
}
TF_EXPECT_OK(CheckJournalContent(journal_dir, updates));
}
TEST(Journal, MissingFile) {
std::string journal_dir;
EXPECT_TRUE(NewJournalDir(journal_dir));
FileJournalReader reader(Env::Default(), journal_dir);
Update result;
bool end_of_journal = true;
absl::Status s = reader.Read(result, end_of_journal);
EXPECT_TRUE(absl::IsNotFound(s));
}
TEST(Journal, NonRecordData) {
std::string journal_dir;
EXPECT_TRUE(NewJournalDir(journal_dir));
TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(journal_dir));
{
std::unique_ptr<WritableFile> file;
TF_ASSERT_OK(Env::Default()->NewAppendableFile(
DataServiceJournalFile(journal_dir, /*sequence_number=*/0), &file));
TF_ASSERT_OK(file->Append("not record data"));
}
FileJournalReader reader(Env::Default(), journal_dir);
Update result;
bool end_of_journal = true;
absl::Status s = reader.Read(result, end_of_journal);
EXPECT_THAT(s.message(), HasSubstr("corrupted record"));
EXPECT_EQ(s.code(), error::DATA_LOSS);
}
TEST(Journal, InvalidRecordData) {
std::string journal_dir;
EXPECT_TRUE(NewJournalDir(journal_dir));
TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(journal_dir));
{
std::unique_ptr<WritableFile> file;
TF_ASSERT_OK(Env::Default()->NewAppendableFile(
DataServiceJournalFile(journal_dir, /*sequence_number=*/0), &file));
auto writer = std::make_unique<io::RecordWriter>(file.get());
TF_ASSERT_OK(writer->WriteRecord("not serialized proto"));
}
FileJournalReader reader(Env::Default(), journal_dir);
Update result;
bool end_of_journal = true;
absl::Status s = reader.Read(result, end_of_journal);
EXPECT_THAT(s.message(), HasSubstr("Failed to parse journal record"));
EXPECT_EQ(s.code(), error::DATA_LOSS);
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/py_utils.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2021 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/core/data/service/py_utils.h"
#include <string>
#include "tensorflow/core/data/service/credentials_factory.h"
namespace tensorflow {
namespace data {
std::string DefaultProtocol() {
#if defined(PLATFORM_GOOGLE)
if (CredentialsFactory::Exists("grpc+loas")) {
return "grpc+loas";
}
static absl::once_flag log_once;
absl::call_once(log_once, [] {
LOG(WARNING)
<< "loas credentials factory is not available, falling back to "
"using insecure credentials.";
});
#endif // PLATFORM_GOOGLE
return "grpc";
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/py_utils.h
+33
View File
@@ -0,0 +1,33 @@
/* Copyright 2021 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_CORE_DATA_SERVICE_PY_UTILS_H_
#define TENSORFLOW_CORE_DATA_SERVICE_PY_UTILS_H_
#include <string>
// Utilities called from the Python API through pybind. We define this file
// separately from other utils to keep the transitive closure of dependencies
// minimal, avoiding linking conflicts.
namespace tensorflow {
namespace data {
// Returns the default protocol to use for tf.data service control flow.
std::string DefaultProtocol();
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_PY_UTILS_H_
tensorflow/core/data/service/server_lib.cc
+306
View File
@@ -0,0 +1,306 @@
/* 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/core/data/service/server_lib.h"
#include <iterator>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "grpcpp/server.h"
#include "grpcpp/server_builder.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/credentials_factory.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/data/service/grpc_dispatcher_impl.h"
#include "tensorflow/core/data/service/grpc_util.h"
#include "tensorflow/core/data/service/grpc_worker_impl.h"
#include "tensorflow/core/data/service/worker_client.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/str_util.h"
namespace tensorflow {
namespace data {
namespace {
// See `WorkerConfig` docs.
constexpr char kPortPlaceholder[] = "%port%";
constexpr char kDataTransferPortPlaceholder[] = "%dts_port%";
} // namespace
GrpcDataServerBase::GrpcDataServerBase(
int port, const std::string& protocol, const std::string& server_type,
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> options)
: requested_port_(port),
protocol_(protocol),
server_type_(server_type),
bound_port_(port),
server_options_(std::move(options)) {}
absl::Status GrpcDataServerBase::Start() {
if (stopped_) {
return absl::FailedPreconditionError(
"Server cannot be started after it has been stopped.");
}
if (started_) {
return absl::OkStatus();
}
::grpc::ServerBuilder builder;
for (std::unique_ptr<::grpc::ServerBuilderOption>& option : server_options_) {
builder.SetOption(std::move(option));
}
server_options_.clear();
std::shared_ptr<::grpc::ServerCredentials> credentials;
TF_RETURN_IF_ERROR(
CredentialsFactory::CreateServerCredentials(protocol_, &credentials));
builder.AddListeningPort(absl::StrCat("0.0.0.0:", requested_port_),
credentials, &bound_port_);
builder.SetMaxReceiveMessageSize(-1);
AddDataServiceToBuilder(builder);
AddProfilerServiceToBuilder(builder);
server_ = builder.BuildAndStart();
if (!server_) {
return absl::InternalError("Could not start gRPC server");
}
TF_RETURN_IF_ERROR(StartServiceInternal());
started_ = true;
LOG(INFO) << "Started tf.data " << server_type_
<< " running at 0.0.0.0:" << BoundPort();
return absl::OkStatus();
}
void GrpcDataServerBase::Stop() {
if (stopped_) {
return;
}
if (server_) {
StopServiceInternal();
server_->Shutdown();
LOG(INFO) << "Shut down " << server_type_ << " server running at port "
<< BoundPort();
}
stopped_ = true;
}
void GrpcDataServerBase::Join() { server_->Wait(); }
int GrpcDataServerBase::BoundPort() { return bound_port(); }
void GrpcDataServerBase::AddProfilerServiceToBuilder(
::grpc::ServerBuilder& builder) {
profiler_service_ = tsl::profiler::CreateProfilerService();
builder.RegisterService(profiler_service_.get());
}
DispatchGrpcDataServer::DispatchGrpcDataServer(
const experimental::DispatcherConfig& config,
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> options)
: GrpcDataServerBase(config.port(), config.protocol(), "DispatchServer",
std::move(options)),
config_(config) {}
DispatchGrpcDataServer::~DispatchGrpcDataServer() { delete service_; }
void DispatchGrpcDataServer::AddDataServiceToBuilder(
::grpc::ServerBuilder& builder) {
service_ = std::make_unique<GrpcDispatcherImpl>(config_, builder).release();
}
absl::Status DispatchGrpcDataServer::StartServiceInternal() {
return service_->Start();
}
void DispatchGrpcDataServer::StopServiceInternal() { service_->Stop(); }
absl::Status DispatchGrpcDataServer::NumWorkers(int* num_workers) {
GetWorkersRequest req;
GetWorkersResponse resp;
::grpc::ServerContext ctx;
::grpc::Status s = service_->GetWorkers(&ctx, &req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get workers", s);
}
*num_workers = resp.workers_size();
return absl::OkStatus();
}
absl::Status DispatchGrpcDataServer::SnapshotStreams(
const std::string& path, std::vector<SnapshotStreamInfoWrapper>* streams) {
GetSnapshotStreamsRequest req;
req.set_path(path);
GetSnapshotStreamsResponse resp;
::grpc::ServerContext ctx;
::grpc::Status s = service_->GetSnapshotStreams(&ctx, &req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get snapshot streams", s);
}
for (const auto& stream : resp.streams()) {
streams->push_back(SnapshotStreamInfoWrapper(stream));
}
return absl::OkStatus();
}
size_t DispatchGrpcDataServer::NumActiveIterations() {
return service_->NumActiveIterations();
}
ServerStateExport DispatchGrpcDataServer::ExportState() const {
ServerStateExport server_state_export;
*server_state_export.mutable_dispatcher_state_export() =
service_->ExportState();
return server_state_export;
}
WorkerGrpcDataServer::WorkerGrpcDataServer(
const experimental::WorkerConfig& config,
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> options)
: GrpcDataServerBase(config.port(), config.protocol(), "WorkerServer",
std::move(options)),
config_(config) {}
WorkerGrpcDataServer::~WorkerGrpcDataServer() { delete service_; }
void WorkerGrpcDataServer::AddDataServiceToBuilder(
::grpc::ServerBuilder& builder) {
service_ = std::make_unique<GrpcWorkerImpl>(config_, builder).release();
}
void WorkerGrpcDataServer::MaybeStartAlternativeDataTransferServer(
std::vector<DataTransferServerInfo>& transfer_servers) {
if (config_.data_transfer_protocol().empty() ||
config_.data_transfer_protocol() == kGrpcTransferProtocol) {
return;
}
absl::Status s = DataTransferServer::Build(config_.data_transfer_protocol(),
service_->get_element_getter(),
&transfer_server_);
if (!s.ok()) {
LOG(ERROR) << "failed to build " << config_.data_transfer_protocol()
<< " server for worker " << config_.worker_address() << ": "
<< s;
return;
}
s = transfer_server_->Start(config_);
if (!s.ok()) {
LOG(ERROR) << "failed to start " << config_.data_transfer_protocol()
<< " server for worker " << config_.worker_address() << ": "
<< s;
return;
}
LOG(INFO) << "Data transfer server started at 0.0.0.0:"
<< transfer_server_->Port() << " for protocol "
<< config_.data_transfer_protocol() << " for worker "
<< config_.worker_address();
DataTransferServerInfo alternative_transfer_server;
alternative_transfer_server.set_protocol(config_.data_transfer_protocol());
alternative_transfer_server.set_address(str_util::StringReplace(
config_.data_transfer_address(), kDataTransferPortPlaceholder,
absl::StrCat(transfer_server_->Port()),
/*replace_all=*/false));
absl::StatusOr<std::string> compatibility_info =
transfer_server_->GetCompatibilityInfo();
if (!compatibility_info.ok()) {
LOG(ERROR)
<< "failed to populate compatibility information for worker server "
<< config_.worker_address() << " for protocol "
<< config_.data_transfer_protocol() << ": "
<< compatibility_info.status();
return;
}
alternative_transfer_server.set_compatibility_info(*compatibility_info);
alternative_transfer_server.set_fall_back_to_grpc_at_client_creation_time(
transfer_server_->FallBackToGrpcAtClientCreationTime());
alternative_transfer_server.set_fall_back_to_grpc_at_get_element_time(
transfer_server_->FallBackToGrpcAtGetElementTime());
transfer_servers.push_back(alternative_transfer_server);
}
absl::Status WorkerGrpcDataServer::StartServiceInternal() {
std::string base_address = config_.worker_address();
if (base_address.empty()) {
base_address = absl::StrCat("localhost:", kPortPlaceholder);
}
std::string worker_address = str_util::StringReplace(
base_address, kPortPlaceholder, absl::StrCat(bound_port()),
/*replace_all=*/false);
DataTransferServerInfo grpc_transfer_server;
grpc_transfer_server.set_protocol(kGrpcTransferProtocol);
grpc_transfer_server.set_address(worker_address);
std::vector<DataTransferServerInfo> transfer_servers = {grpc_transfer_server};
MaybeStartAlternativeDataTransferServer(transfer_servers);
TF_RETURN_IF_ERROR(service_->Start(worker_address, transfer_servers));
return absl::OkStatus();
}
void WorkerGrpcDataServer::StopServiceInternal() { service_->Stop(); }
absl::Status WorkerGrpcDataServer::NumTasks(int* num_tasks) {
GetWorkerTasksRequest req;
GetWorkerTasksResponse resp;
::grpc::ServerContext ctx;
::grpc::Status s = service_->GetWorkerTasks(&ctx, &req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get tasks", s);
}
*num_tasks = resp.tasks_size();
return absl::OkStatus();
}
absl::Status WorkerGrpcDataServer::SnapshotTaskProgresses(
std::vector<SnapshotTaskProgressWrapper>* snapshot_task_progresses) {
GetSnapshotTaskProgressesRequest req;
GetSnapshotTaskProgressesResponse resp;
::grpc::ServerContext ctx;
::grpc::Status s = service_->GetSnapshotTaskProgresses(&ctx, &req, &resp);
if (!s.ok()) {
return grpc_util::WrapError("Failed to get tasks", s);
}
for (const auto& progress : resp.snapshot_task_progresses()) {
snapshot_task_progresses->push_back(SnapshotTaskProgressWrapper(progress));
}
return absl::OkStatus();
}
ServerStateExport WorkerGrpcDataServer::ExportState() const {
ServerStateExport server_state_export;
*server_state_export.mutable_worker_state_export() = service_->ExportState();
return server_state_export;
}
absl::Status NewDispatchServer(
const experimental::DispatcherConfig& config,
std::unique_ptr<DispatchGrpcDataServer>& out_server) {
out_server = std::make_unique<DispatchGrpcDataServer>(config);
return absl::OkStatus();
}
absl::Status NewWorkerServer(
const experimental::WorkerConfig& config,
std::unique_ptr<WorkerGrpcDataServer>& out_server) {
out_server = std::make_unique<WorkerGrpcDataServer>(config);
return absl::OkStatus();
}
} // namespace data
} // namespace tensorflow
tensorflow/core/data/service/server_lib.h
+189
View File
@@ -0,0 +1,189 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef TENSORFLOW_CORE_DATA_SERVICE_SERVER_LIB_H_
#define TENSORFLOW_CORE_DATA_SERVICE_SERVER_LIB_H_
#include <memory>
#include <string>
#include <vector>
#include "grpcpp/server.h"
#include "grpcpp/server_builder.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/data_transfer.h"
#include "tensorflow/core/data/service/export.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/profiler/rpc/profiler_service_impl.h"
#include "tensorflow/core/protobuf/service_config.pb.h"
namespace tensorflow {
namespace data {
// Forward declared because transitively depending on .grpc.pb.h files causes
// issues in the pywrap build.
class GrpcDispatcherImpl;
class GrpcWorkerImpl;
// A grpc server for the tf.data service.
class GrpcDataServerBase {
public:
// Constructs a tf.data server with the specified port. If the port is 0, the
// server will find an available port in `Start()`. The chosen port can be
// found by calling `BoundPort()`.
GrpcDataServerBase(
int requested_port, const std::string& protocol,
const std::string& server_type,
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> options = {});
virtual ~GrpcDataServerBase() = default;
// Starts the server running asynchronously.
absl::Status Start();
// Stops the server. This will block until all outstanding requests complete.
void Stop();
// Blocks until the server stops.
void Join();
// Returns the port bound by the server. Only valid after calling Start().
int BoundPort();
// Exports the server state to improve debuggability.
virtual ServerStateExport ExportState() const = 0;
protected:
virtual void AddDataServiceToBuilder(::grpc::ServerBuilder& builder) = 0;
void AddProfilerServiceToBuilder(::grpc::ServerBuilder& builder);
// Starts the service. This will be called after building the service, so
// bound_port() will return the actual bound port.
virtual absl::Status StartServiceInternal() = 0;
virtual void StopServiceInternal() {}
int bound_port() { return bound_port_; }
const int requested_port_;
const std::string protocol_;
const std::string server_type_;
private:
int bound_port_;
bool started_ = false;
bool stopped_ = false;
std::unique_ptr<::grpc::Server> server_;
// TensorFlow profiler service implementation.
std::unique_ptr<grpc::ProfilerService::Service> profiler_service_ = nullptr;
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> server_options_;
};
// A wrapper for `SnapshotStreamInfo` for use with pybind.
struct SnapshotStreamInfoWrapper {
SnapshotStreamInfoWrapper() = default;
explicit SnapshotStreamInfoWrapper(const SnapshotStreamInfo& info)
: index(info.index()), state(info.state()) {}
int64_t index;
int64_t state;
};
class DispatchGrpcDataServer : public GrpcDataServerBase {
public:
explicit DispatchGrpcDataServer(
const experimental::DispatcherConfig& config,
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> options = {});
~DispatchGrpcDataServer() override;
// Returns the number of workers registered with the dispatcher.
absl::Status NumWorkers(int* num_workers);
// Returns the number of active (non-finished) iterations running on the
// dispatcher.
size_t NumActiveIterations();
// Returns information about all the streams for the snapshot at `path`.
absl::Status SnapshotStreams(const std::string& path,
std::vector<SnapshotStreamInfoWrapper>* streams);
ServerStateExport ExportState() const override;
protected:
void AddDataServiceToBuilder(::grpc::ServerBuilder& builder) override;
absl::Status StartServiceInternal() override;
void StopServiceInternal() override;
private:
const experimental::DispatcherConfig config_;
// Owned. We use a raw pointer because GrpcDispatcherImpl is forward-declared.
GrpcDispatcherImpl* service_;
};
// A wrapper for `SnapshotTaskProgress` for use with pybind.
struct SnapshotTaskProgressWrapper {
SnapshotTaskProgressWrapper() = default;
explicit SnapshotTaskProgressWrapper(const SnapshotTaskProgress& progress)
: snapshot_task_base_path(progress.snapshot_task().base_path()),
snapshot_task_stream_index(progress.snapshot_task().stream_index()),
completed(progress.completed()) {}
std::string snapshot_task_base_path;
int64_t snapshot_task_stream_index;
bool completed;
};
class WorkerGrpcDataServer : public GrpcDataServerBase {
public:
explicit WorkerGrpcDataServer(
const experimental::WorkerConfig& config,
std::vector<std::unique_ptr<::grpc::ServerBuilderOption>> options = {});
~WorkerGrpcDataServer() override;
// Returns the number of tasks currently being executed by the worker.
absl::Status NumTasks(int* num_tasks);
// Returns the progresses of the snapshot tasks currently being executed by
// the worker.
absl::Status SnapshotTaskProgresses(
std::vector<SnapshotTaskProgressWrapper>* snapshot_task_progresses);
ServerStateExport ExportState() const override;
protected:
void AddDataServiceToBuilder(::grpc::ServerBuilder& builder) override;
absl::Status StartServiceInternal() override;
void StopServiceInternal() override;
private:
// If an alternative data transfer protocol is configured, tries to start a
// transfer server for it, adding an entry to `transfer_servers` if
// successful.
void MaybeStartAlternativeDataTransferServer(
std::vector<DataTransferServerInfo>& transfer_servers);
const experimental::WorkerConfig config_;
// Owned. We use a raw pointer because GrpcWorkerImpl is forward-declared.
GrpcWorkerImpl* service_;
std::shared_ptr<DataTransferServer> transfer_server_;
};
// Creates a dispatch tf.data server and stores it in `out_server`.
absl::Status NewDispatchServer(
const experimental::DispatcherConfig& config,
std::unique_ptr<DispatchGrpcDataServer>& out_server);
// Creates a worker tf.data server and stores it in `out_server`.
absl::Status NewWorkerServer(const experimental::WorkerConfig& config,
std::unique_ptr<WorkerGrpcDataServer>& out_server);
} // namespace data
} // namespace tensorflow
#endif // TENSORFLOW_CORE_DATA_SERVICE_SERVER_LIB_H_
tensorflow/core/data/service/snapshot/BUILD
+598
View File
@@ -0,0 +1,598 @@
# Distributed snapshot library.
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@xla//xla/tsl:tsl.default.bzl", "get_compatible_with_portable")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "tf_grpc_cc_dependencies", "tf_kernel_library")
load("//tensorflow/core/platform:build_config.bzl", "tf_protos_profiler_service")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
tf_cc_test(
name = "distributed_snapshot_test",
srcs = ["distributed_snapshot_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":path_utils",
":test_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:test_cluster",
"//tensorflow/core/data/service:test_util",
"//tensorflow/core/framework:tensor_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:tstring",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
] + tf_grpc_cc_dependencies() + tf_protos_profiler_service(),
)
cc_library(
name = "file_utils",
srcs = ["file_utils.cc"],
hdrs = ["file_utils.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":path_utils",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data:snapshot_utils",
"@com_google_absl//absl/functional:function_ref",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:random",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_to_from_proto",
],
)
tf_cc_test(
name = "file_utils_test",
srcs = ["file_utils_test.cc"],
deps = [
":file_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:dataset_test_base",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:test_util",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
tf_kernel_library(
name = "list_snapshot_chunks_dataset_op",
srcs = ["list_snapshot_chunks_dataset_op.cc"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
":snapshot_chunk_provider",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core/data:name_utils",
"//tensorflow/core/data:split_utils",
"//tensorflow/core/framework:dataset_options_proto_cc",
"//tensorflow/core/framework:op_requires",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@tsl//tsl/platform:tstring",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
cc_library(
name = "parallel_tfrecord_writer",
srcs = ["parallel_tfrecord_writer.cc"],
hdrs = ["parallel_tfrecord_writer.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":utils",
"//tensorflow/core:framework",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data/service:byte_size",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:random",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "parallel_tfrecord_writer_test",
srcs = ["parallel_tfrecord_writer_test.cc"],
deps = [
":parallel_tfrecord_writer",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/platform:test",
],
)
cc_library(
name = "path_utils",
srcs = ["path_utils.cc"],
hdrs = ["path_utils.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:path",
],
)
tf_cc_test(
name = "path_utils_test",
srcs = ["path_utils_test.cc"],
deps = [
":path_utils",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/status:status_matchers",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
cc_library(
name = "prefetched_split_provider",
srcs = ["prefetched_split_provider.cc"],
hdrs = ["prefetched_split_provider.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "prefetched_split_provider_test",
srcs = ["prefetched_split_provider_test.cc"],
deps = [
":prefetched_split_provider",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:split_provider",
"//tensorflow/core/data/service:test_util",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/platform:test",
],
)
cc_library(
name = "snapshot_split_provider",
srcs = ["snapshot_split_provider.cc"],
hdrs = ["snapshot_split_provider.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
":path_utils",
"//tensorflow/core:framework",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:dispatcher_proto_cc",
"//tensorflow/core/data/service:grpc_util",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:mutex",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "snapshot_split_provider_test",
srcs = ["snapshot_split_provider_test.cc"],
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
deps = [
":file_utils",
":path_utils",
":snapshot_split_provider",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:serialization_utils",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:dispatcher_client",
"//tensorflow/core/data/service:test_util",
"//tensorflow/core/framework:tensor_testutil",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
cc_library(
name = "snapshot_manager",
srcs = ["snapshot_manager.cc"],
hdrs = ["snapshot_manager.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
":path_utils",
":prefetched_split_provider",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:dispatcher_proto_cc",
"//tensorflow/core/data/service:split_provider",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:mutex",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_to_from_proto",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
tf_cc_test(
name = "snapshot_manager_test",
size = "small",
srcs = ["snapshot_manager_test.cc"],
deps = [
":path_utils",
":snapshot_manager",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:dispatcher_proto_cc",
"//tensorflow/core/data/service:test_util",
"//tensorflow/core/framework:tensor_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:status",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:status_to_from_proto",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
tf_kernel_library(
name = "snapshot_chunk_dataset_op",
srcs = ["snapshot_chunk_dataset_op.cc"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data:name_utils",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data:utils",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:tstring",
"@xla//xla/tsl/platform:env",
],
)
cc_library(
name = "snapshot_chunk_provider",
srcs = ["snapshot_chunk_provider.cc"],
hdrs = ["snapshot_chunk_provider.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
":path_utils",
"//tensorflow/core:framework",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:retrying_utils",
"@tsl//tsl/platform:tstring",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_to_from_proto",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:status_proto_cc",
],
)
tf_cc_test(
name = "snapshot_chunk_provider_test",
size = "small",
srcs = ["snapshot_chunk_provider_test.cc"],
deps = [
":file_utils",
":path_utils",
":snapshot_chunk_provider",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:serialization_utils",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:tstring",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:status_to_from_proto",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:status_proto_cc",
],
)
cc_library(
name = "snapshot_stream_writer",
srcs = ["snapshot_stream_writer.cc"],
hdrs = ["snapshot_stream_writer.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
":parallel_tfrecord_writer",
":path_utils",
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data:utils",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/data/service:common",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:task_runner",
"//tensorflow/core/data/service:worker_proto_cc",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:mutex",
"@tsl//tsl/platform:path",
"@tsl//tsl/platform:thread_annotations",
"@tsl//tsl/profiler/lib:traceme",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "snapshot_stream_writer_checkpoint_test",
size = "medium",
srcs = ["snapshot_stream_writer_checkpoint_test.cc"],
deps = [
":path_utils",
":snapshot_stream_writer",
":test_utils",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:task_runner",
"//tensorflow/core/data/service:test_util",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:random",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:status_matchers",
],
)
tf_cc_test(
name = "snapshot_stream_writer_test",
size = "small",
srcs = ["snapshot_stream_writer_test.cc"],
deps = [
":file_utils",
":path_utils",
":snapshot_stream_writer",
":test_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data:standalone",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:task_runner",
"//tensorflow/core/data/service:test_util",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/lib/monitoring:cell_reader",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/protobuf:protos_all_cc",
],
)
cc_library(
name = "test_utils",
srcs = ["test_utils.cc"],
hdrs = ["test_utils.h"],
compatible_with = get_compatible_with_portable(),
deps = [
":file_utils",
":path_utils",
":snapshot_stream_writer",
"//tensorflow/core:framework",
"//tensorflow/core/data:snapshot_utils",
"//tensorflow/core/data:standalone",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/data/service:common_proto_cc",
"//tensorflow/core/data/service:task_runner",
"//tensorflow/core/framework:types_proto_cc",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:path",
"@xla//xla/tsl/platform:env",
],
)
cc_library(
name = "utils",
srcs = ["utils.cc"],
hdrs = ["utils.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/core:framework",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/framework:tensor_proto_cc",
"@com_google_absl//absl/strings",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
],
)
tf_cc_test(
name = "utils_test",
size = "small",
srcs = ["utils_test.cc"],
deps = [
":utils",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/data/service:byte_size",
"//tensorflow/core/framework:dataset_proto_cc",
"//tensorflow/core/framework:types_proto_cc",
"@tsl//tsl/platform:protobuf",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:status",
],
)
tensorflow/core/data/service/snapshot/distributed_snapshot_test.cc
+217
View File
@@ -0,0 +1,217 @@
/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <cstdint>
#include <memory>
#include <string>
#include <tuple>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/time/time.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/lib/io/compression.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/platform/test.h"
#include "tensorflow/core/data/service/common.pb.h"
#include "tensorflow/core/data/service/dispatcher_client.h"
#include "tensorflow/core/data/service/snapshot/path_utils.h"
#include "tensorflow/core/data/service/snapshot/test_utils.h"
#include "tensorflow/core/data/service/test_cluster.h"
#include "tensorflow/core/data/service/test_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/snapshot.pb.h"
#include "tsl/platform/path.h"
#include "tsl/platform/tstring.h"
namespace tensorflow {
namespace data {
namespace {
using testing::CreateDummyDistributedSnapshotMetadata;
using ::testing::IsEmpty;
using testing::LocalTempFilename;
using testing::RangeDataset;
using ::testing::UnorderedElementsAre;
using tsl::testing::IsOkAndHolds;
constexpr const char kProtocol[] = "grpc";
class TestSnapshotCluster {
public:
explicit TestSnapshotCluster(int64_t num_workers) {
TestCluster::Config config;
config.num_workers = num_workers;
config.worker_heartbeat_interval_ms = 100;
config.work_dir = tsl::io::JoinPath(tsl::testing::TmpDir(), "work_dir");
test_cluster_ = std::make_unique<TestCluster>(config);
TF_CHECK_OK(test_cluster_->Initialize());
dispatcher_client_ = std::make_unique<DataServiceDispatcherClient>(
test_cluster_->DispatcherAddress(), kProtocol);
}
DataServiceDispatcherClient& dispatcher() const {
return *dispatcher_client_;
}
private:
std::unique_ptr<TestCluster> test_cluster_;
std::unique_ptr<DataServiceDispatcherClient> dispatcher_client_;
};
absl::Status WaitForFileExists(const std::string& file_path) {
while (true) {
absl::Status status = Env::Default()->FileExists(file_path);
if (!absl::IsNotFound(status)) {
TF_RETURN_IF_ERROR(status);
}
if (status.ok()) {
return absl::OkStatus();
}
Env::Default()->SleepForMicroseconds(
absl::ToInt64Microseconds(absl::Seconds(1)));
}
return absl::OkStatus();
}
absl::Status WaitForSnapshotComplete(const std::string& base_path) {
return WaitForFileExists(SnapshotDoneFilePath(base_path));
}
class DistributedSnapshotTest : public ::testing::TestWithParam<int64_t> {
protected:
int64_t NumWorkers() const { return GetParam(); }
};
TEST_P(DistributedSnapshotTest, WriteSnapshot) {
TestSnapshotCluster data_service(NumWorkers());
DatasetDef dataset = RangeDataset(10);
experimental::DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
std::string snapshot_path = LocalTempFilename();
TF_ASSERT_OK(
data_service.dispatcher().Snapshot(dataset, snapshot_path, metadata));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshot_path));
EXPECT_THAT(testing::ReadSnapshot<int64_t>(snapshot_path,
tsl::io::compression::kNone),
absl_testing::IsOkAndHolds(
UnorderedElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)));
}
TEST_P(DistributedSnapshotTest, WriteMultipleSnapshots) {
TestSnapshotCluster data_service(NumWorkers());
experimental::DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
std::vector<std::string> snapshots = {
LocalTempFilename(), LocalTempFilename(), LocalTempFilename()};
TF_ASSERT_OK(data_service.dispatcher().Snapshot(RangeDataset(0), snapshots[0],
metadata));
TF_ASSERT_OK(data_service.dispatcher().Snapshot(RangeDataset(10),
snapshots[1], metadata));
TF_ASSERT_OK(data_service.dispatcher().Snapshot(RangeDataset(20),
snapshots[2], metadata));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshots[0]));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshots[1]));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshots[2]));
EXPECT_THAT(
testing::ReadSnapshot<int64_t>(snapshots[0], tsl::io::compression::kNone),
absl_testing::IsOkAndHolds(IsEmpty()));
EXPECT_THAT(
testing::ReadSnapshot<int64_t>(snapshots[1], tsl::io::compression::kNone),
absl_testing::IsOkAndHolds(
UnorderedElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)));
EXPECT_THAT(
testing::ReadSnapshot<int64_t>(snapshots[2], tsl::io::compression::kNone),
absl_testing::IsOkAndHolds(UnorderedElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8,
9, 10, 11, 12, 13, 14, 15,
16, 17, 18, 19)));
}
TEST_P(DistributedSnapshotTest, ChooseFromDatasets) {
// Tests snapshotting this dataset:
// datasets = [tf.data.Dataset.from_tensor_slices(["a", "a", "a", "a", "a"]),
// tf.data.Dataset.from_tensor_slices(["b", "b", "b", "b", "b"]),
// tf.data.Dataset.from_tensor_slices(["c", "c", "c", "c", "c"])]
// choice_dataset = tf.data.Dataset.range(3).repeat()
// dataset = tf.data.Dataset.choose_from_datasets(datasets, choice_dataset)
TestSnapshotCluster data_service(NumWorkers());
TF_ASSERT_OK_AND_ASSIGN(DatasetDef dataset,
testing::GetTestDataset("choose_from_datasets"));
experimental::DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
std::string snapshot_path = LocalTempFilename();
TF_ASSERT_OK(
data_service.dispatcher().Snapshot(dataset, snapshot_path, metadata));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshot_path));
EXPECT_THAT(testing::ReadSnapshot<tsl::tstring>(snapshot_path,
tsl::io::compression::kNone),
absl_testing::IsOkAndHolds(
UnorderedElementsAre("a", "b", "c", "a", "b", "c", "a", "b",
"c", "a", "b", "c", "a", "b", "c")));
}
TEST_P(DistributedSnapshotTest, EmptyDataset) {
TestSnapshotCluster data_service(NumWorkers());
DatasetDef dataset = RangeDataset(0);
experimental::DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
std::string snapshot_path = LocalTempFilename();
TF_ASSERT_OK(
data_service.dispatcher().Snapshot(dataset, snapshot_path, metadata));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshot_path));
EXPECT_THAT(testing::ReadSnapshot<int64_t>(snapshot_path,
tsl::io::compression::kNone),
absl_testing::IsOkAndHolds(IsEmpty()));
}
INSTANTIATE_TEST_SUITE_P(NumWorkers, DistributedSnapshotTest,
::testing::Values(1, 5));
class DistributedSnapshotCompressionTest
: public ::testing::TestWithParam<std::tuple<int64_t, std::string>> {
protected:
int64_t NumWorkers() const { return std::get<0>(GetParam()); }
std::string Compression() const { return std::get<1>(GetParam()); }
};
TEST_P(DistributedSnapshotCompressionTest, Compression) {
TestSnapshotCluster data_service(NumWorkers());
DatasetDef dataset = RangeDataset(20);
experimental::DistributedSnapshotMetadata metadata =
CreateDummyDistributedSnapshotMetadata();
metadata.set_compression(Compression());
std::string snapshot_path = LocalTempFilename();
TF_ASSERT_OK(
data_service.dispatcher().Snapshot(dataset, snapshot_path, metadata));
TF_ASSERT_OK(WaitForSnapshotComplete(snapshot_path));
EXPECT_THAT(testing::ReadSnapshot<int64_t>(snapshot_path, Compression()),
absl_testing::IsOkAndHolds(
UnorderedElementsAre(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
13, 14, 15, 16, 17, 18, 19)));
}
INSTANTIATE_TEST_SUITE_P(
NumWorkersAndCompression, DistributedSnapshotCompressionTest,
::testing::Combine(::testing::Values(1, 5),
::testing::Values(tsl::io::compression::kGzip,
tsl::io::compression::kSnappy,
tsl::io::compression::kZlib)));
} // namespace
} // namespace data
} // namespace tensorflow

Some files were not shown because too many files have changed in this diff Show More