diff --git a/pkgs/development/python-modules/orbax-checkpoint/default.nix b/pkgs/development/python-modules/orbax-checkpoint/default.nix
index cf58f457ff23..279b88e79b97 100644
--- a/pkgs/development/python-modules/orbax-checkpoint/default.nix
+++ b/pkgs/development/python-modules/orbax-checkpoint/default.nix
@@ -15,6 +15,7 @@
   jax,
   msgpack,
   numpy,
+  prometheus-client,
   protobuf,
   psutil,
   pyyaml,
@@ -24,21 +25,24 @@
   uvloop,
 
   # tests
+  aiosqlite,
   chex,
   fastapi,
   google-cloud-logging,
+  greenlet,
   httpx,
   mock,
   optax,
   portpicker,
   pytestCheckHook,
   safetensors,
+  sqlalchemy,
   torch,
 }:
 
 buildPythonPackage (finalAttrs: {
   pname = "orbax-checkpoint";
-  version = "0.11.39";
+  version = "0.11.40";
   pyproject = true;
   __structuredAttrs = true;
 
@@ -46,7 +50,7 @@ buildPythonPackage (finalAttrs: {
     owner = "google";
     repo = "orbax";
     tag = "v${finalAttrs.version}";
-    hash = "sha256-KrehggcpKqMd51SdEo3uzYyvH8M15tmECHzvGLBhD/4=";
+    hash = "sha256-Z1T1mt12kdY6EMY+95m12kW9nHcGj77f87i4PY9ibBU=";
   };
 
   sourceRoot = "${finalAttrs.src.name}/checkpoint";
@@ -61,6 +65,7 @@ buildPythonPackage (finalAttrs: {
     jax
     msgpack
     numpy
+    prometheus-client
     protobuf
     psutil
     pyyaml
@@ -73,19 +78,28 @@ buildPythonPackage (finalAttrs: {
   ++ etils.optional-dependencies.epy;
 
   nativeCheckInputs = [
+    aiosqlite
     chex
     fastapi
     google-cloud-logging
+    greenlet
     httpx
     mock
     optax
     portpicker
     pytestCheckHook
     safetensors
+    sqlalchemy
     torch
   ];
 
   disabledTests = [
+    # ValueError: Distributed system is not available; please initialize it via `jax.distributed.initialize()` at the start of your program.
+    "NumpyHandlerTest"
+    "SerializationTest"
+    "SingleReplicaArrayHandlerTest"
+    "UtilsTest"
+
     # Flaky
     # AssertionError: 2 not greater than 2.0046136379241943
     "test_async_mkdir_parallel"
@@ -108,6 +122,7 @@ buildPythonPackage (finalAttrs: {
 
     # ValueError: cannot reshape array of size 1 into shape (0,2)
     "test_get_leaf_memory_per_device"
+    "test_memory_size"
     "test_number_of_broadcasts"
     "test_tree_memory_per_device"
   ]
@@ -144,6 +159,7 @@ buildPythonPackage (finalAttrs: {
 
     # ValueError: Distributed system is not available; please initialize it via `jax.distributed.initialize()` at the start of your program.
     "orbax/checkpoint/_src/handlers/array_checkpoint_handler_test.py"
+    "orbax/checkpoint/experimental/v1/_src/layout/safetensors_layout_multiprocess_test.py"
 
     # import file mismatch:
     # imported module 'registry_test' has this __file__ attribute:
@@ -153,6 +169,13 @@ buildPythonPackage (finalAttrs: {
     # HINT: remove __pycache__ / .pyc files and/or use a unique basename for your test file module
     "orbax/checkpoint/experimental/v1/_src/serialization/registry_test.py"
 
+    # import file mismatch:
+    # imported module 'serialization_test' has this __file__ attribute:
+    #   /build/source/checkpoint/orbax/checkpoint/_src/serialization/serialization_test.py
+    # which is not the same as the test file we want to collect:
+    #   /build/source/checkpoint/orbax/checkpoint/experimental/v1/_src/metadata/serialization_test.py
+    "orbax/checkpoint/experimental/v1/_src/metadata/serialization_test.py"
+
     # E   FileNotFoundError: [Errno 2] No such file or directory:
     # '/build/absl_testing/DefaultSnapshotTest/runTest/root/path/to/source/data.txt'
     "orbax/checkpoint/_src/path/snapshot/snapshot_test.py"
@@ -161,15 +184,19 @@ buildPythonPackage (finalAttrs: {
     "orbax/checkpoint/_src/multihost/multihost_test.py"
 
     # Circular dependency flax
+    "orbax/checkpoint/_src/checkpointers/async_checkpointer_test.py"
+    "orbax/checkpoint/_src/checkpointers/checkpointer_test.py"
     "orbax/checkpoint/_src/handlers/pytree_checkpoint_handler_test.py"
     "orbax/checkpoint/_src/metadata/empty_values_test.py"
     "orbax/checkpoint/_src/metadata/tree_rich_types_test.py"
     "orbax/checkpoint/_src/metadata/tree_test.py"
+    "orbax/checkpoint/_src/serialization/local_type_handlers_test.py"
     "orbax/checkpoint/_src/testing/test_tree_utils.py"
     "orbax/checkpoint/_src/tree/parts_of_test.py"
     "orbax/checkpoint/_src/tree/structure_utils_test.py"
     "orbax/checkpoint/_src/tree/utils_test.py"
     "orbax/checkpoint/checkpoint_manager_test.py"
+    "orbax/checkpoint/experimental/v1/_src/handlers/pytree_handler_test.py"
     "orbax/checkpoint/single_host_test.py"
     "orbax/checkpoint/transform_utils_test.py"
     "orbax/checkpoint/_src/handlers/standard_checkpoint_handler_test.py"