tensorflow · copybara-service · May 20, 2025 · May 20, 2025
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
@@ -417,6 +417,7 @@ tf_cuda_library(
         "//tensorflow/core/framework:tensor_slice.h",
         "//tensorflow/core/framework:tensor_types.h",
         "//tensorflow/core/framework:tensor_util.h",
+        "//tensorflow/core/framework:tf_data_file_logger_options.h",
         "//tensorflow/core/framework:thread_factory.h",
         "//tensorflow/core/framework:tracking_allocator.h",
         "//tensorflow/core/framework:type_index.h",

diff --git a/tensorflow/core/data/BUILD b/tensorflow/core/data/BUILD
@@ -757,6 +757,7 @@ tf_cc_test(
     # copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
     deps = [
         ":utils",
+        "//tensorflow/core:framework",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/tensorflow/core/data/utils.cc b/tensorflow/core/data/utils.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/statusor.h"
 #include "tensorflow/core/framework/metrics.h"
+#include "tensorflow/core/framework/tf_data_file_logger_options.h"
 #include "tensorflow/core/protobuf/data_service.pb.h"
 
 namespace tensorflow {
@@ -47,7 +48,7 @@ absl::StatusOr<bool> DisableCompressionAtRuntime(
   return false;
 }
 
-void LogFilenames(const std::vector<std::string>& files) {}
+void LogFilenames(const LogFilenamesOptions& options) {}
 
 }  // namespace data
 }  // namespace tensorflow
diff --git a/tensorflow/core/data/utils.h b/tensorflow/core/data/utils.h
@@ -18,9 +18,11 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <string>
+#include <vector>
 
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/statusor.h"
+#include "tensorflow/core/framework/tf_data_file_logger_options.h"
 #include "tensorflow/core/protobuf/data_service.pb.h"
 
 namespace tensorflow {
@@ -54,7 +56,7 @@ absl::StatusOr<bool> DisableCompressionAtRuntime(
 // every call. Thread safe.
 // TODO (shushanik) Implement streamz error reporting in case the logging is not
 // successful
-void LogFilenames(const std::vector<std::string>& files);
+void LogFilenames(const LogFilenamesOptions& options);
 
 }  // namespace data
 }  // namespace tensorflow

diff --git a/tensorflow/core/data/utils_test.cc b/tensorflow/core/data/utils_test.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include <vector>
 
 #include <gtest/gtest.h>
+#include "tensorflow/core/framework/tf_data_file_logger_options.h"
 
 namespace tensorflow::data {
 namespace {
@@ -56,9 +57,12 @@ TEST(LocalityOptimizedPath, TfDataPath) {
 }
 
 TEST(LocalityOptimizedPath, LogFilenames) {
-  EXPECT_NO_FATAL_FAILURE(LogFilenames(
-      std::vector<std::string>({"/path/file1", "file2.txt", "a"})));
-  EXPECT_NO_FATAL_FAILURE(LogFilenames(std::vector<std::string>({})));
+  LogFilenamesOptions options = {.files = {"/path/file1", "file2.txt", "a"},
+                                 .data_service_address = ""};
+  EXPECT_NO_FATAL_FAILURE(LogFilenames(options));
+
+  LogFilenamesOptions options_default_address = {.files = {}};
+  EXPECT_NO_FATAL_FAILURE(LogFilenames(options_default_address));
 }
 
 }  // namespace

diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
@@ -106,6 +106,7 @@ exports_files(
         "tensor_reference.h",
         "tensor_slice.h",
         "tensor_util.h",
+        "tf_data_file_logger_options.h",
         "thread_factory.h",
         "tracking_allocator.h",
         "versions.h",
@@ -260,6 +261,7 @@ filegroup(
         "tensor_slice.h",
         "tensor_types.h",
         "tensor_util.h",
+        "tf_data_file_logger_options.h",
         "thread_factory.h",
         "tracking_allocator.h",
         "type_index.h",
@@ -342,6 +344,7 @@ filegroup(
         "tensor_shape.h",
         "tensor_types.h",
         "tensor_util.h",
+        "tf_data_file_logger_options.h",
         "tracking_allocator.h",
         "type_index.h",
         "type_traits.h",

diff --git a/tensorflow/core/framework/tf_data_file_logger_options.h b/tensorflow/core/framework/tf_data_file_logger_options.h
@@ -0,0 +1,43 @@
+/* Copyright 2025 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_FRAMEWORK_TF_DATA_FILE_LOGGER_OPTIONS_H_
+#define TENSORFLOW_CORE_FRAMEWORK_TF_DATA_FILE_LOGGER_OPTIONS_H_
+
+#include <string>
+#include <vector>
+
+namespace tensorflow {
+namespace data {
+
+struct LogFilenamesOptions {
+  std::vector<std::string> files;
+  std::string data_service_address = "";
+};
+
+inline bool operator==(const LogFilenamesOptions& lhs,
+                       const LogFilenamesOptions& rhs) {
+  return lhs.files == rhs.files &&
+         lhs.data_service_address == rhs.data_service_address;
+}
+
+inline bool operator!=(const LogFilenamesOptions& lhs,
+                       const LogFilenamesOptions& rhs) {
+  return !(lhs == rhs);
+}
+
+}  // namespace data
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_FRAMEWORK_TF_DATA_FILE_LOGGER_OPTIONS_H_
diff --git a/tensorflow/core/kernels/data/fixed_length_record_dataset_op.cc b/tensorflow/core/kernels/data/fixed_length_record_dataset_op.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/partial_tensor_shape.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tf_data_file_logger_options.h"
 #include "tensorflow/core/lib/io/buffered_inputstream.h"
 #include "tensorflow/core/lib/io/inputbuffer.h"
 #include "tensorflow/core/lib/io/random_inputstream.h"
@@ -133,7 +134,10 @@ class FixedLengthRecordDatasetOp::Dataset : public DatasetBase {
         : DatasetIterator<Dataset>(params) {}
 
     absl::Status Initialize(IteratorContext* ctx) override {
-      LogFilenames(dataset()->filenames_);
+      LogFilenamesOptions log_filenames_options = {
+          .files = dataset()->filenames_,
+          .data_service_address = ctx->data_service_address()};
+      LogFilenames(log_filenames_options);
       return absl::OkStatus();
     }
 
@@ -264,7 +268,10 @@ class FixedLengthRecordDatasetOp::Dataset : public DatasetBase {
         : DatasetIterator<Dataset>(params) {}
 
     absl::Status Initialize(IteratorContext* ctx) override {
-      LogFilenames(dataset()->filenames_);
+      LogFilenamesOptions log_filenames_options = {
+          .files = dataset()->filenames_,
+          .data_service_address = ctx->data_service_address()};
+      LogFilenames(log_filenames_options);
       return absl::OkStatus();
     }
 

diff --git a/tensorflow/core/kernels/data/text_line_dataset_op.cc b/tensorflow/core/kernels/data/text_line_dataset_op.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/partial_tensor_shape.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tf_data_file_logger_options.h"
 #include "tensorflow/core/lib/io/buffered_inputstream.h"
 #include "tensorflow/core/lib/io/inputbuffer.h"
 #include "tensorflow/core/lib/io/random_inputstream.h"
@@ -101,7 +102,10 @@ class TextLineDatasetOp::Dataset : public DatasetBase {
         : DatasetIterator<Dataset>(params) {}
 
     absl::Status Initialize(IteratorContext* ctx) override {
-      LogFilenames(dataset()->filenames_);
+      LogFilenamesOptions log_filenames_options = {
+          .files = dataset()->filenames_,
+          .data_service_address = ctx->data_service_address()};
+      LogFilenames(log_filenames_options);
       return absl::OkStatus();
     }
 

diff --git a/tensorflow/core/kernels/data/tf_record_dataset_op.cc b/tensorflow/core/kernels/data/tf_record_dataset_op.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/core/framework/metrics.h"
 #include "tensorflow/core/framework/partial_tensor_shape.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tf_data_file_logger_options.h"
 #include "tensorflow/core/lib/io/buffered_inputstream.h"
 #include "tensorflow/core/lib/io/inputbuffer.h"
 #include "tensorflow/core/lib/io/random_inputstream.h"
@@ -134,7 +135,10 @@ class TFRecordDatasetOp::Dataset : public DatasetBase {
         : DatasetIterator<Dataset>(params) {}
 
     absl::Status Initialize(IteratorContext* ctx) override {
-      LogFilenames(dataset()->filenames_);
+      LogFilenamesOptions log_filenames_options = {
+          .files = dataset()->filenames_,
+          .data_service_address = ctx->data_service_address()};
+      LogFilenames(log_filenames_options);
       return absl::OkStatus();
     }