# Distributed snapshot library.

load("@local_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/strings",
        "@com_google_absl//absl/time",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:tstring",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_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",
        "@local_tsl//tsl/platform:protobuf",
        "@local_tsl//tsl/platform:random",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_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:statusor",
        "@local_tsl//tsl/platform:path",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_xla//xla/tsl/platform:statusor",
        "@local_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",
        "@local_tsl//tsl/platform:tstring",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_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",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:random",
        "@local_tsl//tsl/profiler/lib:traceme",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_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:statusor",
        "@com_google_absl//absl/strings",
        "@local_tsl//tsl/platform:path",
        "@local_xla//xla/tsl/lib/core:status_test_util",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_xla//xla/tsl/platform:statusor",
        "@local_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",
        "@local_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",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_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",
        "@local_tsl//tsl/platform:path",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_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:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@local_tsl//tsl/platform:path",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_xla//xla/tsl/platform:statusor",
        "@local_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",
        "@local_tsl//tsl/platform:mutex",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:thread_annotations",
        "@local_xla//xla/tsl/platform:env",
        "@local_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/strings",
        "@com_google_absl//absl/time",
        "@local_tsl//tsl/platform:path",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_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",
        "@local_tsl//tsl/platform:mutex",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:thread_annotations",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_to_from_proto",
        "@local_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",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:status",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_xla//xla/tsl/platform:status_to_from_proto",
        "@local_xla//xla/tsl/platform:statusor",
        "@local_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",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:tstring",
        "@local_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",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:retrying_utils",
        "@local_tsl//tsl/platform:tstring",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_to_from_proto",
        "@local_xla//xla/tsl/platform:statusor",
        "@local_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:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/synchronization",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:tstring",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_xla//xla/tsl/platform:status_to_from_proto",
        "@local_xla//xla/tsl/platform:statusor",
        "@local_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",
        "@local_tsl//tsl/platform:mutex",
        "@local_tsl//tsl/platform:path",
        "@local_tsl//tsl/platform:thread_annotations",
        "@local_tsl//tsl/profiler/lib:traceme",
        "@local_xla//xla/tsl/platform:env",
        "@local_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:statusor",
        "@com_google_absl//absl/time",
        "@local_tsl//tsl/platform:random",
        "@local_xla//xla/tsl/platform:env",
        "@local_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:statusor",
        "@com_google_absl//absl/strings",
        "@local_tsl//tsl/platform:path",
        "@local_xla//xla/tsl/lib/monitoring:cell_reader",
        "@local_xla//xla/tsl/platform:env",
        "@local_xla//xla/tsl/platform:status_matchers",
        "@local_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",
        "@local_tsl//tsl/platform:path",
        "@local_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",
        "@local_xla//xla/tsl/platform:errors",
        "@local_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",
        "@local_tsl//tsl/platform:protobuf",
        "@local_xla//xla/tsl/platform:errors",
        "@local_xla//xla/tsl/platform:status",
    ],
)
