diff --git a/tensorflow/cc/saved_model/fingerprinting_utils.cc b/tensorflow/cc/saved_model/fingerprinting_utils.cc index 218128bd60b4b6..92e61671c9132d 100644 --- a/tensorflow/cc/saved_model/fingerprinting_utils.cc +++ b/tensorflow/cc/saved_model/fingerprinting_utils.cc @@ -219,6 +219,8 @@ absl::StatusOr HashFields( if (chunked_message.has_chunk_index() && matches == field_tags.size()) { // chunked_field_tags are an exact match with field_tags. Hash referenced // chunk. + TF_RETURN_IF_ERROR(tools::proto_splitter::ValidateChunkIndex( + chunked_message.chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( std::string chunk, ReadChunk(reader, chunks_info[chunked_message.chunk_index()])); @@ -244,6 +246,8 @@ absl::StatusOr HashFields( merged_message = mfr.parent->GetReflection()->MutableMessage(mfr.parent, mfr.field); } + TF_RETURN_IF_ERROR(tools::proto_splitter::ValidateChunkIndex( + chunked_message.chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( std::string chunk, ReadChunk(reader, chunks_info[chunked_message.chunk_index()])); diff --git a/tensorflow/cc/saved_model/fingerprinting_utils_test.cc b/tensorflow/cc/saved_model/fingerprinting_utils_test.cc index e457ef128dcd3f..d51c01682c5fec 100644 --- a/tensorflow/cc/saved_model/fingerprinting_utils_test.cc +++ b/tensorflow/cc/saved_model/fingerprinting_utils_test.cc @@ -416,6 +416,41 @@ TEST(FingerprintingTest, CreateFingerprintDefNonExistentDirectoryReturnsError) { EXPECT_FALSE(result.ok()); } +TEST(FingerprintingTest, TestHashFieldsReturnsErrorOnInvalidChunkIndex) { + std::string cpb_file = io::JoinPath( + TensorFlowSrcRoot(), "tools/proto_splitter/testdata", "many-field.cpb"); + TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file)); + + auto read_metadata = GetChunkMetadata(reader); + if (!read_metadata.ok()) { + reader.Close(); + TF_ASSERT_OK(read_metadata.status()); + } + ChunkMetadata chunk_metadata = read_metadata.value(); + + ChunkedMessage invalid_chunked_message; + invalid_chunked_message.set_chunk_index(0); + auto* invalid_field = invalid_chunked_message.add_chunked_fields(); + invalid_field->add_field_tag()->set_field(1); + + std::vector chunks_info = std::vector( + chunk_metadata.chunks().begin(), chunk_metadata.chunks().end()); + invalid_field->mutable_message()->set_chunk_index(chunks_info.size()); + + RepeatedPtrField target_tags; + target_tags.Add()->set_field(1); + + ManyFields merged_message; + + auto statusor = HashFields(invalid_chunked_message, reader, chunks_info, + target_tags, &merged_message); + + EXPECT_FALSE(statusor.ok()); + EXPECT_EQ(statusor.status().code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(statusor.status().message()), + ::testing::HasSubstr("is out of range")); +} + } // namespace } // namespace tensorflow::saved_model::fingerprinting diff --git a/tensorflow/python/autograph/converters/call_trees.py b/tensorflow/python/autograph/converters/call_trees.py index 3d694d45a17b0b..ca9c1e3d057d97 100644 --- a/tensorflow/python/autograph/converters/call_trees.py +++ b/tensorflow/python/autograph/converters/call_trees.py @@ -32,6 +32,31 @@ # TODO(mdan): Rename to FunctionCallsTransformer. +# Python random module functions that produce values traced as constants. +# These cause issues when used inside tf.function because the random value +# is computed during tracing and becomes a static constant in the graph. +_PYTHON_RANDOM_FUNCTIONS = frozenset([ + 'random.random', + 'random.randint', + 'random.randrange', + 'random.uniform', + 'random.choice', + 'random.choices', + 'random.sample', + 'random.shuffle', + 'random.gauss', + 'random.normalvariate', + 'random.betavariate', + 'random.expovariate', + 'random.gammavariate', + 'random.lognormvariate', + 'random.vonmisesvariate', + 'random.paretovariate', + 'random.weibullvariate', + 'random.triangular', + 'random.getrandbits', +]) + class _Function(object): @@ -42,6 +67,7 @@ def __init__(self): set_trace_warned = False +python_random_warned = False class _ArgTemplateBuilder(object): @@ -187,6 +213,26 @@ def visit_Call(self, node): set_trace_warned = True return node + # Warn when Python random module functions are used inside tf.function. + # These values are computed at trace time and become constants in the graph, + # which can cause shape mismatches or unexpected behavior at runtime, + # especially with XLA compilation. + if full_name in _PYTHON_RANDOM_FUNCTIONS: + global python_random_warned + if not python_random_warned: + ag_logging.warning( + "Detected use of Python's `%s()` inside a tf.function. " + 'The random value is computed during tracing and becomes a ' + 'constant in the graph, which may cause shape mismatches or ' + 'unexpected behavior, especially with XLA compilation. ' + 'Use `tf.random` functions instead for dynamic random values. ' + 'For example, replace `random.randint(a, b)` with ' + '`tf.random.uniform([], a, b + 1, dtype=tf.int32)`. ' + 'See https://www.tensorflow.org/guide/function#executing_python_side_effects', + full_name, + ) + python_random_warned = True + if (full_name == 'print' and not self.ctx.user.options.uses(converter.Feature.BUILTIN_FUNCTIONS)): return node diff --git a/tensorflow/python/autograph/converters/call_trees_test.py b/tensorflow/python/autograph/converters/call_trees_test.py index eb69d718a4c846..9e78f62021fef6 100644 --- a/tensorflow/python/autograph/converters/call_trees_test.py +++ b/tensorflow/python/autograph/converters/call_trees_test.py @@ -14,6 +14,7 @@ # ============================================================================== """Tests for call_trees module.""" +import random import types from tensorflow.python.autograph.converters import call_trees @@ -257,6 +258,73 @@ def test_method(self, a): self.assertEqual(321, tr(tc, 1)) self.assertListEqual(mock.calls, [((1,), None)]) + def test_python_random_warning(self): + """Test that using Python random module triggers a warning.""" + # Reset the warning flag for testing + call_trees.python_random_warned = False + + def f(): + return random.randint(1, 10) + + with self.assertLogs(level='WARNING') as logs: + tr, mock = self._transform_with_mock(f) + + self.assertLen(logs.output, 1) + self.assertIn( + "Detected use of Python's `random.randint()` inside a tf.function.", + logs.output[0], + ) + + # The function should still be callable + result = tr() + self.assertIsInstance(result, int) + self.assertGreaterEqual(result, 1) + self.assertLessEqual(result, 10) + + def test_python_random_randrange_warning(self): + """Test that using Python random.randrange triggers a warning.""" + # Reset the warning flag for testing + call_trees.python_random_warned = False + + def f(): + return random.randrange(0, 100) + + with self.assertLogs(level='WARNING') as logs: + tr, mock = self._transform_with_mock(f) + + self.assertLen(logs.output, 1) + self.assertIn( + "Detected use of Python's `random.randrange()` inside a tf.function.", + logs.output[0], + ) + + # The function should still be callable + result = tr() + self.assertIsInstance(result, int) + self.assertGreaterEqual(result, 0) + self.assertLess(result, 100) + + def test_python_random_choice_warning(self): + """Test that using Python random.choice triggers a warning.""" + # Reset the warning flag for testing + call_trees.python_random_warned = False + + def f(): + return random.choice([1, 2, 3]) + + with self.assertLogs(level='WARNING') as logs: + tr, mock = self._transform_with_mock(f) + + self.assertLen(logs.output, 1) + self.assertIn( + "Detected use of Python's `random.choice()` inside a tf.function.", + logs.output[0], + ) + + # The function should still be callable + result = tr() + self.assertIn(result, [1, 2, 3]) + if __name__ == '__main__': test.main() diff --git a/tensorflow/tools/proto_splitter/cc/util.cc b/tensorflow/tools/proto_splitter/cc/util.cc index b38cdea137d444..095e6342d64879 100644 --- a/tensorflow/tools/proto_splitter/cc/util.cc +++ b/tensorflow/tools/proto_splitter/cc/util.cc @@ -801,5 +801,12 @@ absl::StatusOr OnlyContainsPb(absl::string_view prefix) { return false; } +absl::Status ValidateChunkIndex(uint64_t chunk_index, size_t chunks_size) { + if (chunk_index < chunks_size) return absl::OkStatus(); + return absl::FailedPreconditionError(absl::StrCat("Chunk index ", chunk_index, + " is out of range for ", + chunks_size, " chunks.")); +} + } // namespace tools::proto_splitter } // namespace tensorflow diff --git a/tensorflow/tools/proto_splitter/cc/util.h b/tensorflow/tools/proto_splitter/cc/util.h index ee525d752acf6c..54eff47aa047a4 100644 --- a/tensorflow/tools/proto_splitter/cc/util.h +++ b/tensorflow/tools/proto_splitter/cc/util.h @@ -15,6 +15,7 @@ limitations under the License. #ifndef TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_ #define TENSORFLOW_TOOLS_PROTO_SPLITTER_CC_UTIL_H_ +#include #include #include #include @@ -164,6 +165,9 @@ absl::StatusOr ReadChunk( // file exists. Returns an error if neither .pb nor .cpb exist. absl::StatusOr OnlyContainsPb(absl::string_view prefix); +// Validates that a chunk index is within the bounds of a chunks vector. +absl::Status ValidateChunkIndex(uint64_t chunk_index, size_t chunks_size); + } // namespace tools::proto_splitter } // namespace tensorflow diff --git a/tensorflow/tools/proto_splitter/merge.cc b/tensorflow/tools/proto_splitter/merge.cc index d2243312769ce9..514e3f33caf810 100644 --- a/tensorflow/tools/proto_splitter/merge.cc +++ b/tensorflow/tools/proto_splitter/merge.cc @@ -59,6 +59,8 @@ absl::Status Merger::Merge(const std::vector>& chunks, if (chunked_message.has_chunk_index()) { // Chunks referenced by fields should be merged into the parent chunk. + TF_RETURN_IF_ERROR( + ValidateChunkIndex(chunked_message.chunk_index(), chunks.size())); merged_message->MergeFrom(*chunks[chunked_message.chunk_index()].get()); } @@ -237,6 +239,8 @@ absl::Status Merger::ReadFields(const ChunkedMessage& chunked_message, tsl::protobuf::Message* merged_message) { if (chunked_message.has_chunk_index()) { // Chunks referenced by fields should be merged into the parent chunk. + TF_RETURN_IF_ERROR( + ValidateChunkIndex(chunked_message.chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( std::string chunk, ReadChunk(reader, chunks_info[chunked_message.chunk_index()])); @@ -305,12 +309,16 @@ absl::Status Merger::ProcessField( std::string chunk; switch (op) { case MergerOp::READ: { + TF_RETURN_IF_ERROR(ValidateChunkIndex( + chunked_field.message().chunk_index(), chunks_info.size())); TF_ASSIGN_OR_RETURN( chunk, ReadChunk(reader, chunks_info[chunked_field.message().chunk_index()])); break; } case MergerOp::MERGE: { + TF_RETURN_IF_ERROR(ValidateChunkIndex( + chunked_field.message().chunk_index(), chunks.size())); chunk = chunks[chunked_field.message().chunk_index()]->SerializeAsString(); break; diff --git a/tensorflow/tools/proto_splitter/merge_test.cc b/tensorflow/tools/proto_splitter/merge_test.cc index b49c901c7f9745..f0f2c0c73f7828 100644 --- a/tensorflow/tools/proto_splitter/merge_test.cc +++ b/tensorflow/tools/proto_splitter/merge_test.cc @@ -312,6 +312,72 @@ TEST(MergeTest, TestProcessFieldReturnsErrorOnInvalidFieldNumber) { ::testing::HasSubstr("not found in message descriptor")); } +TEST(MergeTest, TestMergeReturnsErrorOnInvalidRootChunkIndex) { + ::tensorflow::proto_splitter::ChunkedMessage chunked_message; + chunked_message.set_chunk_index(1); + + std::vector> chunks; + chunks.push_back( + std::make_unique<::tensorflow::proto_splitter_testdata::ManyFields>()); + + ::tensorflow::proto_splitter_testdata::ManyFields merged; + absl::Status status = Merger::Merge(chunks, chunked_message, &merged); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("Chunk index 1 is out of range")); +} + +TEST(MergeTest, TestMergeReturnsErrorOnInvalidNestedChunkIndex) { + ::tensorflow::proto_splitter::ChunkedMessage chunked_message; + auto* chunk_field = chunked_message.add_chunked_fields(); + auto* tag = chunk_field->add_field_tag(); + tag->set_field(1); + chunk_field->mutable_message()->set_chunk_index(1); + + std::vector> chunks; + chunks.push_back( + std::make_unique<::tensorflow::proto_splitter_testdata::ManyFields>()); + + ::tensorflow::proto_splitter_testdata::ManyFields merged; + absl::Status status = Merger::Merge(chunks, chunked_message, &merged); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("Chunk index 1 is out of range")); +} + +TEST(MergeTest, TestReadPartialReturnsErrorOnInvalidRootChunkIndex) { + const std::string path = + io::JoinPath(testing::TensorFlowSrcRoot(), + "tools/proto_splitter/testdata", "many-field"); + TF_ASSERT_OK_AND_ASSIGN(auto reader, tools::proto_splitter::GetRiegeliReader( + absl::StrCat(path, ".cpb"))); + + auto read_metadata = GetChunkMetadata(reader); + if (!read_metadata.ok()) { + reader.Close(); + TF_ASSERT_OK(read_metadata.status()); + } + reader.Close(); + + ::tensorflow::proto_splitter::ChunkMetadata chunk_metadata = + read_metadata.value(); + chunk_metadata.mutable_message()->set_chunk_index( + chunk_metadata.chunks_size()); + + proto_splitter_testdata::ManyFields merged_many_fields; + absl::Status status = + Merger::ReadPartial(path, chunk_metadata, &merged_many_fields); + + EXPECT_FALSE(status.ok()); + EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition); + EXPECT_THAT(std::string(status.message()), + ::testing::HasSubstr("is out of range")); +} + } // namespace } // namespace tensorflow::tools::proto_splitter