Testing pytest.

PiperOrigin-RevId: 886439501

Author: iindyk

.github/workflows/tests.yml

@@ -16,6 +16,6 @@ jobs:
       id-token: write
     with:
       pypi_project_url: ''
-      run_tests_with_bazel: true
+      run_tests_with_bazel: false
       is_nightly: false
       is_ci_tests: true

grain/_src/python/checkpoint/handler_test.py

@@ -29,10 +29,6 @@ py_library(
     name = "tfrecord_dataset",
     srcs = ["tfrecord_dataset.py"],
     srcs_version = "PY3",
-    visibility = [
-        "//third_party/py/grain:internal",
-        "//third_party/py/maxtext:__pkg__",
-    ],
     deps = ["//grain/_src/python/dataset"],
 )
 

grain/_src/python/checkpoint/orbax_integration_test.py

@@ -21,7 +21,6 @@ main() {
   # Enable host OS specific configs. For instance, "build:linux" will be used
   # automatically when building on Linux.
   write_to_bazelrc "build --enable_platform_specific_config"
-  write_to_bazelrc "build --verbose_failures"
   # Bazel 7.0.0 no longer supports dynamic symbol lookup on macOS. To resolve
   # undefined symbol errors in macOS arm64 builds, explicitly add the necessary
   # linker flags until dependencies are well defined. See
@@ -133,8 +132,10 @@ main() {
     $PYTHON_BIN -m pip install --force ${OUTPUT_DIR}/all_dist/grain*.whl
     $PYTHON_BIN -m pip install -r test_requirements.in  --only-binary pyarrow
     # TF does not have Python 3.14 support yet.
+    # Once majority of our users are on TF >= 2.21, we should update the
+    # protobuf dependency and bump the smoke test to TF >= 2.21.
     if (( "${PYTHON_MINOR_VERSION}" < 14 )); then
-      $PYTHON_BIN -m pip install tensorflow  --only-binary h5py
+      $PYTHON_BIN -m pip install "tensorflow<2.21"  --only-binary h5py
     fi
 
     pushd "${OUTPUT_DIR}/all_dist"

grain/_src/python/dataset/sources/BUILD

@@ -16,7 +16,7 @@ dependencies = [
     "cloudpickle",
     "etils[epath,epy]",
     "numpy",
-    "protobuf>=5.28.3",
+    "protobuf>=5.28.3,<7",
 ]
 readme = "README.md"
 license = { file = "LICENSE" }

grain/oss/build_whl.sh

@@ -20,10 +20,10 @@ array-record
 absl-py
 etils[epath,epy]
 cloudpickle
-jax
+jax<0.9
 numpy
 attrs
-pyarrow
+pyarrow<23.0  # This should be update with protobuf dependency update.
 pytest
 parameterized
 typing_extensions