Added fix for assertion error during parallel access of mountpoint client. Added integration tests for parallel access of mountpoint client. (#237)

* Added integration test to test parallel saving of checkpoints

* Added integration test for parallel access of mountpoint client

* Added fix to prevent assert error in multi threaded environment

* Modified integration test

* Modified integration test and lock approach

* Added comments

* Added some more comments

Author: rajdchak

s3torchconnector/src/s3torchconnector/_s3client/_s3client.py

@@ -3,6 +3,7 @@
 
 import logging
 import os
+import threading
 from functools import partial
 from typing import Optional, Any
 
@@ -32,6 +33,9 @@ def _identity(obj: Any) -> Any:
     return obj
 
 
+_client_lock = threading.Lock()
+
+
 class S3Client:
     def __init__(
         self,
@@ -51,10 +55,15 @@ def __init__(
 
     @property
     def _client(self) -> MountpointS3Client:
+        # This is a fast check to avoid acquiring the lock unnecessarily.
         if self._client_pid is None or self._client_pid != os.getpid():
-            self._client_pid = os.getpid()
-            # `MountpointS3Client` does not survive forking, so re-create it if the PID has changed.
-            self._real_client = self._client_builder()
+            # Acquire the lock to ensure thread-safety when creating the client.
+            with _client_lock:
+                # This double-check ensures that the client is only created once.
+                if self._client_pid is None or self._client_pid != os.getpid():
+                    # `MountpointS3Client` does not survive forking, so re-create it if the PID has changed.
+                    self._real_client = self._client_builder()
+                    self._client_pid = os.getpid()
         assert self._real_client is not None
         return self._real_client
 

s3torchconnector/tst/e2e/test_mountpoint_client_parallel_access.py

@@ -0,0 +1,81 @@
+import os
+import random
+import time
+import threading
+import pytest
+from s3torchconnector._s3client import S3Client
+from s3torchconnectorclient._mountpoint_s3_client import MountpointS3Client
+
+
+class S3ClientWithoutLock(S3Client):
+    @property
+    def _client(self) -> MountpointS3Client:
+        if self._client_pid is None or self._client_pid != os.getpid():
+            self._client_pid = os.getpid()
+            # `MountpointS3Client` does not survive forking, so re-create it if the PID has changed.
+            self._real_client = self._client_builder()
+        assert self._real_client is not None
+        return self._real_client
+
+    def _client_builder(self):
+        time.sleep(1)
+        return super()._client_builder()
+
+
+class S3ClientWithLock(S3Client):
+    def _client_builder(self):
+        time.sleep(1)
+        return super()._client_builder()
+
+
+def access_client(client, error_event):
+    try:
+        if not error_event.is_set():
+            client._client
+            print(f"Successfully accessed by thread {threading.current_thread().name}")
+    except AssertionError as e:
+        print(f"AssertionError in thread {threading.current_thread().name}: {e}")
+        error_event.set()
+
+
+def test_multiple_thread_accessing_mountpoint_client_in_parallel_without_lock():
+    print("Running test without lock...")
+    client = S3ClientWithoutLock("us-west-2")
+    if not access_mountpoint_client_in_parallel(client):
+        pytest.fail(
+            "Test failed as AssertionError did not happen in one of the threads."
+        )
+
+
+def test_multiple_thread_accessing_mountpoint_client_in_parallel_with_lock():
+    print("Running test with lock...")
+    client = S3ClientWithLock("us-west-2")
+    if access_mountpoint_client_in_parallel(client):
+        pytest.fail("Test failed as AssertionError happened in one of the threads.")
+
+
+def access_mountpoint_client_in_parallel(client):
+
+    error_event = threading.Event()
+    # Create and start multiple threads
+    accessor_threads = []
+    num_accessor_threads = 10
+
+    for i in range(num_accessor_threads):
+        if error_event.is_set():
+            break
+        accessor_thread = threading.Thread(
+            target=access_client,
+            args=(
+                client,
+                error_event,
+            ),
+            name=f"Accessor-{i + 1}",
+        )
+        accessor_threads.append(accessor_thread)
+        accessor_thread.start()
+
+    for thread in accessor_threads:
+        thread.join()
+
+    return error_event.is_set()