support s3's fallback

Author: littleEast7

alluxiofs/client/core.py

@@ -234,8 +234,9 @@ def _check_response(self, response):
         self.logger.debug(
             f"Alluxio server returned {response.status_code}: {message}"
         )
-
-        if response.status_code == 404:
+        if response.status_code == 104:
+            raise ConnectionResetError("Connection reset by peer")
+        elif response.status_code == 404:
             raise FileNotFoundError(message)
         elif response.status_code == 400:
             raise ValueError(message)
@@ -826,23 +827,14 @@ def get_ufs_info_from_worker(self):
             worker_host = self.config.load_balance_domain
         else:
             worker_host = self.loadbalancer.get_worker().host
-        try:
-            url = GET_UFS_SECRET_INFO.format(
-                domain=worker_host,
-                http_port=self.config.worker_http_port,
-            )
-            response = requests.get(url)
-            response.raise_for_status()
-            info = response.content
-            return info
-        except Exception as e:
-            raise Exception(
-                EXCEPTION_CONTENT.format(
-                    worker_host=worker_host,
-                    http_port=self.config.worker_http_port,
-                    error=f"Error when get ufsInfo, {e}",
-                )
-            )
+        url = GET_UFS_SECRET_INFO.format(
+            domain=worker_host,
+            http_port=self.config.worker_http_port,
+        )
+        response = requests.get(url)
+        self._check_response(response)
+        info = response.content
+        return info
 
     def _all_chunk_generator(
         self, worker_host, worker_http_port, path_id, file_path, chunk_size
@@ -881,10 +873,7 @@ def _all_chunk_generator(
                     url_chunk, headers=headers, stream=True
                 ) as response:
                     # Check for connection reset error (status code 104)
-                    if response.status_code == 104:
-                        raise ConnectionResetError("Connection reset by peer")
-
-                    response.raise_for_status()
+                    self._check_response(response)
                     for chunk in response.iter_content(chunk_size=chunk_size):
                         if chunk:
                             out.write(chunk)
@@ -1727,7 +1716,7 @@ def _get_preferred_worker_address(self, full_ufs_path):
             )
         try:
             response = requests.get(url)
-            response.raise_for_status()
+            self._check_response(response)
             data = json.loads(response.content)
             ip = data["Host"]
             port = data["HttpServerPort"]
@@ -2228,7 +2217,7 @@ def _get_preferred_worker_host(self, full_ufs_path):
             )
         try:
             response = requests.get(url)
-            response.raise_for_status()
+            self._check_response(response)
             data = json.loads(response.content)
             ip = data["Host"]
         except HTTPError:

alluxiofs/client/ufs_manager.py

@@ -223,10 +223,10 @@ def must_get_ufs_from_path(self, path: str):
         ufs = self.get_ufs_from_cache(path)
         if ufs is None:
             self.logger.error(
-                f"No registered UFS found in alluxio for path: {path}"
+                f"No registered UFS found in alluxio python sdk for path: {path}"
             )
             raise ValueError(
-                f"No registered UFS found in alluxio for path: {path}"
+                f"No registered UFS found in alluxio python sdk for path: {path}"
             )
         return ufs
 

alluxiofs/client/utils.py

@@ -24,10 +24,22 @@
     "endpoint": "endpoint",
 }
 
+S3_SETUP_OPTIONS_MAP = {
+    "access_key": "key",
+    "secret_key": "secret",
+    "endpoint": "endpoint_url",
+}
+
 
 def convert_ufs_info_to(ufs, info):
     if ufs == "oss":
         res = {OSS_SETUP_OPTIONS_MAP[k]: info[k] for k in info}
+    elif ufs == "s3" or ufs == "s3a":
+        res = {
+            S3_SETUP_OPTIONS_MAP[k]: info[k]
+            for k in info
+            if k in S3_SETUP_OPTIONS_MAP
+        }
     else:
         res = info
     return res
@@ -37,13 +49,14 @@ def parameters_adapter(fs, fs_method, final_kwargs):
     """Adapts parameters for different filesystems and methods."""
     adapted_kwargs = final_kwargs.copy()
     method_name = getattr(fs_method, "__name__", str(fs_method))
+    protocols = (fs.protocol,) if isinstance(fs.protocol, str) else fs.protocol
 
-    if fs.protocol == "bos":
+    if "bos" in protocols:
         if method_name in ["cp_file"]:
             adapted_kwargs["src_path"] = final_kwargs.pop("path1")
             adapted_kwargs["target_path"] = final_kwargs.pop("path2")
 
-    if fs.protocol == "oss":
+    if "oss" in protocols:
         if method_name in ["pipe_file"]:
             headers = {}
             mode = adapted_kwargs.pop("mode", "overwrite")
@@ -54,6 +67,10 @@ def parameters_adapter(fs, fs_method, final_kwargs):
                 # This will cause the upload to fail if the file exists
                 headers["x-oss-forbid-overwrite"] = "true"
 
+    if any(p in ["s3", "s3a"] for p in protocols):
+        if method_name in ["_pipe_file", "pipe_file"]:
+            adapted_kwargs["data"] = final_kwargs.pop("value")
+
     return adapted_kwargs