Add some multi-threading tests to tensorstore python

The new tests pass when built with the python GIL.
The new tests fail when built in free-threaded mode and TSAN.

PiperOrigin-RevId: 824918753
Change-Id: I4d006d41660fcecc0e35c4c7050ec3214ad0f084

Author: laramiel

python/tensorstore/BUILD

@@ -143,6 +143,7 @@ tensorstore_pytest_test(
     name = "dim_test",
     size = "small",
     srcs = ["tests/dim_test.py"],
+    tags = ["cpu:2"],
     deps = [
         ":conftest",
         ":tensorstore",
@@ -752,6 +753,7 @@ tensorstore_pytest_test(
     name = "spec_test",
     size = "small",
     srcs = ["tests/spec_test.py"],
+    tags = ["cpu:2"],
     deps = [
         ":conftest",
         ":tensorstore",
@@ -774,6 +776,7 @@ tensorstore_pytest_test(
     name = "tensorstore_test",
     size = "small",
     srcs = ["tests/tensorstore_test.py"],
+    tags = ["cpu:2"],
     deps = [
         ":conftest",
         ":tensorstore",
@@ -794,8 +797,9 @@ tensorstore_pytest_test(
 
 tensorstore_pytest_test(
     name = "future_test",
-    size = "small",
+    size = "medium",
     srcs = ["tests/future_test.py"],
+    tags = ["cpu:2"],
     deps = [
         ":conftest",
         ":tensorstore",
@@ -1100,6 +1104,7 @@ tensorstore_pytest_test(
     name = "kvstore_test",
     size = "small",
     srcs = ["tests/kvstore_test.py"],
+    tags = ["cpu:2"],
     deps = [
         ":conftest",
         ":tensorstore",

python/tensorstore/generate_type_stubs.py

@@ -122,7 +122,13 @@ def transform_init_ast(
   #
   # These are internal implementation details that aren't useful in the type
   # stubs.
-  excluded_symbols = {"_unpickle", "__reduce__", "__getstate__", "__setstate__"}
+  excluded_symbols = {
+      "_unpickle",
+      "__reduce__",
+      "__getstate__",
+      "__setstate__",
+      "__conditional_annotations__",
+  }
 
   class _InitPyVisitor(ast.NodeTransformer):
 

python/tensorstore/tests/dim_test.py

@@ -14,6 +14,8 @@
 """Tests for tensorstore.Dim"""
 
 import pickle
+import threading
+import time
 
 import pytest
 import tensorstore as ts
@@ -117,3 +119,48 @@ def test_hull() -> None:
 def test_pickle() -> None:
   x = ts.Dim(inclusive_min=3, size=10)
   assert pickle.loads(pickle.dumps(x)) == x
+
+
+def test_dim_concurrent() -> None:
+  """Tests concurrent access to Dim properties."""
+  dim = ts.Dim()
+
+  stop = threading.Event()
+
+  def read_props() -> None:
+    while not stop.is_set():
+      _ = dim.inclusive_min
+      _ = dim.implicit_lower
+      _ = dim.implicit_upper
+      _ = dim.label
+      _ = dim == ts.Dim()
+      _ = f"{dim}"
+      _ = repr(dim)
+
+  def update_props() -> None:
+    time.sleep(0.01)
+    i = 0
+    while not stop.is_set():
+      if (i % 2) == 0:
+        dim.implicit_lower = True
+        dim.implicit_upper = False
+        dim.label = "x"
+      else:
+        dim.implicit_lower = False
+        dim.implicit_upper = True
+        dim.label = ""
+      i += 1
+
+  threads = []
+  for _ in range(4):
+    threads.append(threading.Thread(target=read_props))
+    threads.append(threading.Thread(target=update_props))
+
+  for t in threads:
+    t.start()
+
+  time.sleep(0.3)
+  stop.set()
+
+  for t in threads:
+    t.join()

python/tensorstore/tests/future_test.py

@@ -15,6 +15,7 @@
 import asyncio
 import os
 import pickle
+import random
 import signal
 import threading
 import time
@@ -233,3 +234,147 @@ def callback(f: ts.Future) -> None:
   del p
   assert f.done()
   assert exc is not None
+
+
+def _get_delay() -> float:
+  return random.random() * 0.01
+
+
+def _run_threads(
+    promise: ts.Promise,
+    i: int,
+    stop: threading.Event,
+    threads: list[threading.Thread],
+) -> None:
+  """Runs a list of threads concurrently with setting the promise result."""
+
+  delay = _get_delay()
+  for t in threads:
+    t.start()
+
+  time.sleep(delay)
+  try:
+    promise.set_result(i)
+  except:  # pylint: disable=bare-except
+    pass
+
+  time.sleep(0.01)
+  stop.set()
+
+  for t in threads:
+    t.join()
+
+
+def test_future_concurrent_set_cancel_callback() -> None:
+  """Test multi-treaded races between adding callbacks and cancellation."""
+
+  def _concurrent_set_cancel_callback(i: int) -> None:
+    callback_called = threading.Event()
+    callback_result = [None]
+    callback_error = [None]
+
+    def callback(f: ts.Future) -> None:
+      nonlocal callback_called
+      nonlocal callback_result
+      nonlocal callback_error
+      try:
+        callback_result[0] = f.result()
+      except asyncio.CancelledError as e:
+        callback_error[0] = e
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        callback_error[0] = e
+      finally:
+        callback_called.set()
+
+    def do_add_callback(delay: float, f: ts.Future) -> None:
+      time.sleep(delay)
+      f.add_done_callback(callback)
+
+    def do_cancel(delay: float, f: ts.Future) -> None:
+      time.sleep(delay)
+      f.cancel()
+
+    promise, future = ts.Promise.new()
+    _run_threads(
+        promise,
+        i,
+        threading.Event(),  # unused
+        [
+            threading.Thread(target=do_cancel, args=(_get_delay(), future)),
+            threading.Thread(
+                target=do_add_callback, args=(_get_delay(), future)
+            ),
+        ],
+    )
+
+    assert future.done()
+    assert callback_called.wait(timeout=5)
+
+    if future.cancelled():
+      assert isinstance(callback_error[0], asyncio.CancelledError)
+      with pytest.raises(asyncio.CancelledError):
+        future.result()
+    else:
+      # If not cancelled, result must have been set.
+      assert callback_error[0] is None
+      assert future.result() == i
+      assert callback_result[0] == i
+
+  for i in range(20):
+    _concurrent_set_cancel_callback(i)
+
+
+def test_future_concurrent_ops() -> None:
+  """Test multi-treaded races between adding callbacks and setting result."""
+
+  def _concurrent_ops() -> None:
+    events = [threading.Event() for _ in range(8)]
+    results = {}
+    lock = threading.Lock()
+    stop = threading.Event()
+
+    def make_callback(
+        idx: int, event: threading.Event
+    ) -> Callable[[ts.Future], None]:
+      def callback(f: ts.Future) -> None:
+        try:
+          result = f.result()
+          with lock:
+            results[idx] = result
+        finally:
+          event.set()
+
+      return callback
+
+    promise, future = ts.Promise.new()
+
+    def do_add_callback(
+        delay: float, callback: Callable[[ts.Future], None]
+    ) -> None:
+      time.sleep(delay)
+      future.add_done_callback(callback)
+
+    def read_props() -> None:
+      while not stop.is_set():
+        _ = future.done()
+        _ = future.cancelled()
+
+    threads = []
+    for i in range(len(events)):
+      threads.append(
+          threading.Thread(
+              target=do_add_callback,
+              
+              args=(_get_delay(), make_callback(i, events[i])),
+          )
+      )
+      threads.append(threading.Thread(target=read_props))
+
+    _run_threads(promise, 42, stop, threads)
+
+    for i in range(len(events)):
+      assert events[i].wait(timeout=5)
+      assert results[i] == 42
+
+  for _ in range(20):
+    _concurrent_ops()

python/tensorstore/tests/kvstore_test.py

@@ -17,6 +17,9 @@
 import pathlib
 import pickle
 import tempfile
+import threading
+import time
+from typing import Callable
 
 import pytest
 import tensorstore as ts
@@ -157,3 +160,78 @@ def test_copy_range_to_ocdbt_memory_bad_path() -> None:
   ).result()
   with pytest.raises(NotImplementedError):
     child.experimental_copy_range_to(parent).result()
+
+
+def _run_threads(
+    stop: threading.Event,
+    read_props: Callable[[], None],
+    update_props: Callable[[], None],
+) -> None:
+  threads = []
+  for _ in range(4):
+    threads.append(threading.Thread(target=read_props))
+    threads.append(threading.Thread(target=update_props))
+
+  for t in threads:
+    t.start()
+
+  time.sleep(0.3)
+  stop.set()
+
+  for t in threads:
+    t.join()
+
+
+def test_kvstore_spec_concurrent_update_and_read() -> None:
+  """Validates that concurrent updates and reads do not crash."""
+  s = ts.KvStore.Spec('memory://')
+
+  stop = threading.Event()
+
+  def read_props() -> None:
+    while not stop.is_set():
+      _ = s.path
+      _ = s.url
+      _ = s.base
+      _ = s == ts.KvStore.Spec('memory://')
+      _ = s.to_json()
+      _ = f'{s}'
+      _ = repr(s)
+
+  def update_props() -> None:
+    time.sleep(0.01)
+    i = 0
+    while not stop.is_set():
+      if (i % 2) == 0:
+        s.path = 'abc/'
+      else:
+        s.path = 'def/'
+      i += 1
+
+  _run_threads(stop, read_props, update_props)
+
+
+def test_kvstore_keyrange_concurrent() -> None:
+  """Tests concurrent access to KvStore.KeyRange properties."""
+  kr = ts.KvStore.KeyRange('a', 'z')
+
+  stop = threading.Event()
+
+  def read_props() -> None:
+    while not stop.is_set():
+      _ = kr.inclusive_min
+      _ = kr.exclusive_max
+      _ = kr.empty
+      _ = kr == ts.KvStore.KeyRange('a', 'z')
+      _ = f'{kr}'
+      _ = repr(kr)
+
+  def update_props() -> None:
+    time.sleep(0.01)
+    while not stop.is_set():
+      kr.inclusive_min = 'b'
+      kr.exclusive_max = 'y'
+      kr.inclusive_min = 'a'
+      kr.exclusive_max = 'z'
+
+  _run_threads(stop, read_props, update_props)