OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2026-01-15 12:15:51 +00:00
Code Issues Actions 52 Packages Projects Releases Wiki Activity

Fix flaky tests related to multithreaded stringstream access (#170586)

Browse Source 
Authored with claude code

The flaky test is PyTorchStreamWriterAndReader.LoadWithMultiThreads

Signed-off-by: Edward Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/170586
Approved by: https://github.com/Skylion007
This commit is contained in:
Edward Yang
2025-12-16 14:36:00 -05:00
committed by PyTorch MergeBot
parent 3e5e0a3e4e
commit e7ab9015e7
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
caffe2/serialize/inline_container_test.cc
View File

@@ -175,10 +175,15 @@ TEST(PyTorchStreamWriterAndReader, LoadWithMultiThreads) {
// Inplace multi-threading getRecord(name, dst, n, additional_readers) test
additionalReader.clear();
// Each IStreamAdapter needs its own independent stream to avoid data races.
// IStreamAdapter does not synchronize access to the underlying istream.
std::vector<std::unique_ptr<std::istringstream>> additional_streams;
std::vector<uint8_t> dst1(size1), dst2(size2);
for (int i = 0; i < 10; ++i) {
// Test various sizes of read threads
additionalReader.push_back(std::make_unique<IStreamAdapter>(&iss));
additional_streams.push_back(std::make_unique<std::istringstream>(the_file));
additionalReader.push_back(
std::make_unique<IStreamAdapter>(additional_streams.back().get()));
ret = reader.getRecord("key1", dst1.data(), size1, additionalReader);
ASSERT_EQ(ret, size1);
@@ -611,10 +616,15 @@ TEST(PyTorchStreamWriterAndReader, LoadWithMultiThreadsWithAllocator) {
// Inplace multi-threading getRecord(name, dst, n, additional_readers) test
additionalReader.clear();
// Each IStreamAdapter needs its own independent stream to avoid data races.
// IStreamAdapter does not synchronize access to the underlying istream.
std::vector<std::unique_ptr<std::istringstream>> additional_streams;
std::vector<uint8_t> dst1(size1), dst2(size2);
for (int i = 0; i < 10; ++i) {
// Test various sizes of read threads
additionalReader.push_back(std::make_unique<IStreamAdapter>(&iss));
additional_streams.push_back(std::make_unique<std::istringstream>(the_file));
additionalReader.push_back(
std::make_unique<IStreamAdapter>(additional_streams.back().get()));
ret = reader.getRecord("key1", dst1.data(), size1, additionalReader);
ASSERT_EQ(ret, size1);