diff --git a/MODULE.bazel b/MODULE.bazel index f91bc78..e3b51eb 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -27,7 +27,7 @@ bazel_dep(name = "platforms", version = "0.0.11") # We force the protobuf version to be 28.3 because 29.x (that is required by rules_python) causes # segfault when importing TensorFlow due to a conflict with TF's protobuf version. -bazel_dep(name = "protobuf", version = "29.0-rc3") +bazel_dep(name = "protobuf", version = "28.3") single_version_override( module_name = "protobuf", version = "28.3", diff --git a/cpp/array_record_writer.cc b/cpp/array_record_writer.cc index 0196439..f1542c0 100644 --- a/cpp/array_record_writer.cc +++ b/cpp/array_record_writer.cc @@ -247,7 +247,7 @@ class ArrayRecordWriterBase::SubmitChunkCallback // return false if we can't schedule the callback. // return true and inc num_concurrent_chunk_writers if we can add a new one. bool TrackConcurrentChunkWriters() { - absl::MutexLock l(mu_); + absl::MutexLock l(&mu_); if (num_concurrent_chunk_writers_ >= max_parallelism_) { return false; } @@ -490,7 +490,7 @@ void ArrayRecordWriterBase::SubmitChunkCallback::operator()( footer.set_num_records(num_records); array_footer_.push_back(std::move(footer)); - absl::MutexLock l(mu_); + absl::MutexLock l(&mu_); num_concurrent_chunk_writers_--; } diff --git a/cpp/sequenced_chunk_writer.cc b/cpp/sequenced_chunk_writer.cc index 20c3250..450dc37 100644 --- a/cpp/sequenced_chunk_writer.cc +++ b/cpp/sequenced_chunk_writer.cc @@ -34,7 +34,7 @@ namespace array_record { bool SequencedChunkWriterBase::CommitFutureChunk( std::future>&& future_chunk) { - absl::MutexLock l(mu_); + absl::MutexLock l(&mu_); if (!ok()) { return false; } @@ -64,13 +64,13 @@ bool SequencedChunkWriterBase::SubmitFutureChunks(bool block) { if (block) { // When blocking, we block both on mutex acquisition and on future // completion. - absl::MutexLock lock(mu_); + absl::MutexLock lock(&mu_); riegeli::ChunkWriter* writer = get_writer(); while (!queue_.empty()) { TrySubmitFirstFutureChunk(writer); } return ok(); - } else if (mu_.try_lock()) { + } else if (mu_.TryLock()) { // When non-blocking, we only proceed if we can lock the mutex without // blocking, and we only process those futures that are ready. We need // to unlock the mutex manually in this case, and take care to call ok() @@ -82,7 +82,7 @@ bool SequencedChunkWriterBase::SubmitFutureChunks(bool block) { TrySubmitFirstFutureChunk(writer); } bool result = ok(); - mu_.unlock(); + mu_.Unlock(); return result; } else { return true; @@ -118,7 +118,7 @@ void SequencedChunkWriterBase::TrySubmitFirstFutureChunk( return; } if (pad_to_block_boundary_) { - if (!chunk_writer->PadToBlockBoundary()) { + if (!chunk_writer->PadToBlockBoundary()) { Fail(riegeli::Annotate( chunk_writer->status(), diff --git a/oss/build_whl.sh b/oss/build_whl.sh index 3819881..70f14fa 100755 --- a/oss/build_whl.sh +++ b/oss/build_whl.sh @@ -81,8 +81,10 @@ function main() { echo $(date) : "=== Auditing wheel" auditwheel repair --plat ${AUDITWHEEL_PLATFORM} -w dist dist/*.whl cp dist/*manylinux*.whl "${DEST}" + WHEEL_FILE="${DEST}/$(basename "$(ls -1 dist/*manylinux*.whl | head -n 1)")" else cp dist/*.whl "${DEST}" + WHEEL_FILE="${DEST}/$(basename "$(ls -1 dist/*.whl | head -n 1)")" fi echo $(date) : "=== Listing wheel" @@ -92,12 +94,13 @@ function main() { echo $(date) : "=== Output wheel file is in: ${DEST}" # Install ArrayRecord from the wheel and run smoke tests. - $PYTHON_BIN -m pip install ${OUTPUT_DIR}/all_dist/array_record*.whl + $PYTHON_BIN -m pip install "${WHEEL_FILE}" $PYTHON_BIN -c 'import array_record' $PYTHON_BIN -c 'from array_record.python import array_record_data_source' # TF does not have a Python 3.14 wheel yet. if (( "${PYTHON_MINOR_VERSION}" < 14 )); then - $PYTHON_BIN -m pip install jax tensorflow>=2.20.0 grain --only-binary h5py + $PYTHON_BIN -m pip install jax "tensorflow>=2.20.0" grain --only-binary :all: + $PYTHON_BIN -m pip install --force-reinstall --no-deps "${WHEEL_FILE}" $PYTHON_BIN oss/test_import_tensorflow.py $PYTHON_BIN oss/test_import_grain.py fi diff --git a/python/BUILD b/python/BUILD index d247e61..7579231 100644 --- a/python/BUILD +++ b/python/BUILD @@ -10,6 +10,21 @@ licenses(["notice"]) pybind_extension( name = "array_record_module", srcs = ["array_record_module.cc"], + # Restrict exported symbols of array_record_module.so in OSS builds to only + # PyInit_array_record_module. By default in OSS Bazel builds, pybind11_bazel + # exports all statically linked C++ symbols (including Protobuf and Abseil) + # globally. When imported in the same process as TensorFlow + # (libtensorflow_framework), duplicate Protobuf/Abseil symbols collide in + # the dynamic linker and cause a segmentation fault on macOS and Linux. + additional_linker_inputs = select({ + "@platforms//os:linux": [":array_record_module.lds"], + "//conditions:default": [], + }), + linkopts = select({ + "@platforms//os:osx": ["-Wl,-exported_symbol,_PyInit_array_record_module"], + "@platforms//os:linux": ["-Wl,--version-script=$(location :array_record_module.lds)"], + "//conditions:default": [], + }), deps = [ "@abseil-cpp//absl/status", "@abseil-cpp//absl/strings", diff --git a/python/array_record_module.lds b/python/array_record_module.lds new file mode 100644 index 0000000..17f3719 --- /dev/null +++ b/python/array_record_module.lds @@ -0,0 +1,29 @@ +/* Copyright 2026 Google LLC. 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. +==============================================================================*/ + +/* Linker version script for array_record_module.so in Linux OSS builds. + * + * Exports only the Python module initialization entry point + * (PyInit_array_record_module) and hides all other statically linked C++ + * symbols (such as Protobuf and Abseil). Without this script, importing + * array_record alongside TensorFlow (libtensorflow_framework) causes C++ + * Protobuf/Abseil symbol collisions and segmentation faults at runtime. + */ +{ + global: + PyInit_array_record_module; + local: + *; +};