Merge pull request #32 from yuchen0cc/main

avoid client init race in threading

Author: liulanzheng

oss-torch-connector/osstorchconnector/_oss_client.py

@@ -1,6 +1,7 @@
 import os
 from typing import Iterator, Iterable, Any
 import logging
+import threading
 
 log = logging.getLogger(__name__)
 
@@ -30,15 +31,18 @@ def __init__(self, endpoint: str, cred_path: str = "", config_path: str = "", uu
         self._total = total
         self._cred_provider = cred_provider
 
+    _lock = threading.Lock()
     @property
     def _client(self) -> DataSet:
-        if self._client_pid is None or self._client_pid != os.getpid():
-            # does OSS client survive forking ? NO
-            if self._client_pid != os.getpid() and self._real_client is not None:
-                log.info("OssClient delete dataset")
-                # del self._real_client
-            self._client_pid = os.getpid()
-            self._real_client = self._client_builder()
+        with OssClient._lock:
+            if self._client_pid is None or self._client_pid != os.getpid() :
+                # does OSS client survive forking ? NO
+                if self._client_pid != os.getpid() and self._real_client is not None:
+                    log.info("OssClient delete dataset")
+                    # del self._real_client
+                self._client_pid = os.getpid()
+                self._real_client = self._client_builder()
+
         return self._real_client
 
     def _client_builder(self) -> DataSet:

oss-torch-connector/osstorchconnector/oss_filesystem.py

@@ -173,8 +173,8 @@ def validate_checkpoint_id(cls, checkpoint_id: Union[str, os.PathLike]) -> bool:
             return False
         return True
 
-    def writer(self, path : str) -> OssStorageWriter:
-        return OssStorageWriter(self, path)
+    def writer(self, path : str, **kwargs) -> OssStorageWriter:
+        return OssStorageWriter(self, path, **kwargs)
 
     def reader(self, path : str) -> OssStorageReader:
         return OssStorageReader(self, path)