8000 Internal change · tensorflow/tensorflow@88a220c · GitHub
[go: up one dir, main page]

Skip to content

Commit 88a220c

Browse files
Internal change
PiperOrigin-RevId: 759800545
1 parent ddbe3bf commit 88a220c

File tree

10 files changed

+77
-9
lines changed

10 files changed

+77
-9
lines changed

tensorflow/core/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -417,6 +417,7 @@ tf_cuda_library(
417417
"//tensorflow/core/framework:tensor_slice.h",
418418
"//tensorflow/core/framework:tensor_types.h",
419419
"//tensorflow/core/framework:tensor_util.h",
420+
"//tensorflow/core/framework:tf_data_file_logger_options.h",
420421
"//tensorflow/core/framework:thread_factory.h",
421422
"//tensorflow/core/framework:tracking_allocator.h",
422423
"//tensorflow/core/framework:type_index.h",

tensorflow/core/data/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -757,6 +757,7 @@ tf_cc_test(
757757
# copybara:uncomment extra_copts = ["-Wthread-safety-analysis"],
758758
deps = [
759759
":utils",
760+
"//tensorflow/core:framework",
760761
"@com_google_googletest//:gtest_main",
761762
],
762763
)

tensorflow/core/data/utils.cc

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ limitations under the License.
2222
#include "absl/container/flat_hash_map.h"
2323
#include "absl/status/statusor.h"
2424
#include "tensorflow/core/framework/metrics.h"
25+
#include "tensorflow/core/framework/tf_data_file_logger_options.h"
2526
#include "tensorflow/core/protobuf/data_service.pb.h"
2627

2728
namespace tensorflow {
@@ -47,7 +48,7 @@ absl::StatusOr<bool> DisableCompressionAtRuntime(
4748
return false;
4849
}
4950

50-
void LogFilenames(const std::vector<std::string>& files) {}
51+
void LogFilenames(const LogFilenamesOptions& options) {}
5152

5253
} // namespace data
5354
} // namespace tensorflow

tensorflow/core/data/utils.h

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@ limitations under the License.
1818
#include <memory>
1919
#include <optional>
2020
#include <string>
21+
#include <vector>
2122

2223
#include "absl/container/flat_hash_map.h"
2324
#include "absl/status/statusor.h"
25+
#include "tensorflow/core/framework/tf_data_file_logger_options.h"
2426
#include "tensorflow/core/protobuf/data_service.pb.h"
2527

2628
namespace tensorflow {
@@ -54,7 +56,7 @@ absl::StatusOr<bool> DisableCompressionAtRuntime(
5456
// every call. Thread safe.
5557
// TODO (shushanik) Implement streamz error reporting in case the logging is not
5658
// successful
57-
void LogFilenames(const std::vector<std::string>& files);
59+
void LogFilenames(const LogFilenamesOptions& options);
5860

5961
} // namespace data
6062
} // namespace tensorflow

tensorflow/core/data/utils_test.cc

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ limitations under the License.
1717
#include <vector>
1818

1919
#include <gtest/gtest.h>
20+
#include "tensorflow/core/framework/tf_data_file_logger_options.h"
2021

2122
namespace tensorflow::data {
2223
namespace {
@@ -56,9 +57,9 @@ TEST(LocalityOptimizedPath, TfDataPath) {
5657
}
5758

5859
TEST(LocalityOptimizedPath, LogFilenames) {
59-
EXPECT_NO_FATAL_FAILURE(LogFilenames(
60-
std::vector<std::string>({"/path/file1", "file2.txt", "a"})));
61-
EXPECT_NO_FATAL_FAILURE(LogFilenames(std::vector<std::string>({})));
60+
LogFilenamesOptions options({"/path/file1", "file2.txt", "a"}, "");
61+
EXPECT_NO_FATAL_FAILURE(LogFilenames(options));
62+
EXPECT_NO_FATAL_FAILURE(LogFilenames(options));
6263
}
6364

6465
} // namespace

tensorflow/core/framework/BUILD

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ exports_files(
106106
"tensor_reference.h",
107107
"tensor_slice.h",
108108
"tensor_util.h",
109+
"tf_data_file_logger_options.h",
109110
"thread_factory.h",
110111
"tracking_allocator.h",
111112
"versions.h",
@@ -260,6 +261,7 @@ filegroup(
260261
"tensor_slice.h",
261262
"tensor_types.h",
262263
"tensor_util.h",
264+
"tf_data_file_logger_options.h",
263265
"thread_factory.h",
264266
"tracking_allocator.h",
265267
"type_index.h",
@@ -342,6 +344,7 @@ filegroup(
342344
"tensor_shape.h",
343345
"tensor_types.h",
344346
"tensor_util.h",
347+
"tf_data_file_logger_options.h",
345348
"tracking_allocator.h",
346349
"type_index.h",
347350
"type_traits.h",
Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
#ifndef TENSORFLOW_CORE_FRAMEWORK_TF_DATA_FILE_LOGGER_OPTIONS_H_
16+
#define TENSORFLOW_CORE_FRAMEWORK_TF_DATA_FILE_LOGGER_OPTIONS_H_
17+
18+
#include <string>
19+
#include <vector>
20+
21+
namespace tensorflow {
22+
namespace data {
23+
24+
struct LogFilenamesOptions {
25+
std::vector<std::string> files;
26+
std::string data_service_address;
27+
28+
LogFilenamesOptions() : data_service_address("") {}
29+
LogFilenamesOptions(const std::vector<std::string>& files,
30+
const std::string& data_service_address)
31+
: files(files), data_service_address(data_service_address) {}
32+
};
33+
34+
inline bool operator==(const LogFilenamesOptions& lhs,
35+
const LogFilenamesOptions& rhs) {
36+
return lhs.files == rhs.files &&
37+
lhs.data_service_address == rhs.data_service_address;
38+
}
39+
40+
inline bool operator!=(const LogFilenamesOptions& lhs,
41+
const LogFilenamesOptions& rhs) {
42+
return !(lhs == rhs);
43+
}
44+
45+
} // namespace data
46+
} // namespace tensorflow
47+
48+
#endif // TENSORFLOW_CORE_FRAMEWORK_TF_DATA_FILE_LOGGER_OPTIONS_H_

tensorflow/core/kernels/data/fixed_length_record_dataset_op.cc

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ limitations under the License.
2020
#include "tensorflow/core/framework/metrics.h"
2121
#include "tensorflow/core/framework/partial_tensor_shape.h"
2222
#include "tensorflow/core/framework/tensor.h"
23+
#include "tensorflow/core/framework/tf_data_file_logger_options.h"
2324
#include "tensorflow/core/lib/io/buffered_inputstream.h"
2425
#include "tensorflow/core/lib/io/inputbuffer.h"
2526
#include "tensorflow/core/lib/io/random_inputstream.h"
@@ -133,7 +134,9 @@ class FixedLengthRecordDatasetOp::Dataset : public DatasetBase {
133134
: DatasetIterator<Dataset>(params) {}
134135

135136
absl::Status Initialize(IteratorContext* ctx) override {
136-
LogFilenames(dataset()->filenames_);
137+
LogFilenamesOptions log_filenames_options(dataset()->filenames_,
138+
ctx->data_service_address());
139+
LogFilenames(log_filenames_options);
137140
return absl::OkStatus();
138141
}
139142

@@ -264,7 +267,9 @@ class FixedLengthRecordDatasetOp::Dataset : public DatasetBase {
264267
: DatasetIterator<Dataset>(params) {}
265268

266269
absl::Status Initialize(IteratorContext* ctx) override {
267-
LogFilenames(dataset()->filenames_);
270+
LogFilenamesOptions log_filenames_options(dataset()->filenames_,
271+
ctx->data_service_address());
272+
LogFilenames(log_filenames_options);
268273
return absl::OkStatus();
269274
}
270275

tensorflow/core/kernels/data/text_line_dataset_op.cc

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ limitations under the License.
2020
#include "tensorflow/core/framework/metrics.h"
2121
#include "tensorflow/core/framework/partial_tensor_shape.h"
2222
#include "tensorflow/core/framework/tensor.h"
23+
#include "tensorflow/core/framework/tf_data_file_logger_options.h"
2324
#include "tensorflow/core/lib/io/buffered_inputstream.h"
2425
#include "tensorflow/core/lib/io/inputbuffer.h"
2526
#include "tensorflow/core/lib/io/random_inputstream.h"
@@ -101,7 +102,9 @@ class TextLineDatasetOp::Dataset : public DatasetBase {
101102
: DatasetIterator<Dataset>(params) {}
102103

103104
absl::Status Initialize(IteratorContext* ctx) override {
104-
LogFilenames(dataset()->filenames_);
105+
LogFilenamesOptions log_filenames_options(dataset()->filenames_,
106+
ctx->data_service_address());
107+
LogFilenames(log_filenames_options);
105108
return absl::OkStatus();
106109
}
107110

tensorflow/core/kernels/data/tf_record_dataset_op.cc

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ limitations under the License.
2222
#include "tensorflow/core/framework/metrics.h"
2323
#include "tensorflow/core/framework/partial_tensor_shape.h"
2424
#include "tensorflow/core/framework/tensor.h"
25+
#include "tensorflow/core/framework/tf_data_file_logger_options.h"
2526
#include "tensorflow/core/lib/io/buffered_inputstream.h"
2627
#include "tensorflow/core/lib/io/inputbuffer.h"
2728
#include "tensorflow/core/lib/io/random_inputstream.h"
@@ -134,7 +135,9 @@ class TFRecordDatasetOp::Dataset : public DatasetBase {
134135
: DatasetIterator<Dataset>(params) {}
135136

136137
absl::Status Initialize(IteratorContext* ctx) override {
137-
LogFilenames(dataset()->filenames_);
138+
LogFilenamesOptions log_filenames_options(dataset()->filenames_,
139+
ctx->data_service_address());
140+
LogFilenames(log_filenames_options);
138141
return absl::OkStatus();
139142
}
140143

0 commit comments

Comments
 (0)
0