[Typing][C-47] Add type annotations for python/paddle/distributed/fleet/utils/fs.py (#66766)

megemini · web-flow · commit 37b59a5a4e9c · 2024-07-31T09:22:01.000+08:00
diff --git a/python/paddle/distributed/fleet/utils/fs.py b/python/paddle/distributed/fleet/utils/fs.py
@@ -14,6 +14,8 @@
 
 # type: ignore[abstract]
 
+from __future__ import annotations
+
 import abc
 import functools
 import multiprocessing
@@ -22,12 +24,28 @@
 import shutil
 import subprocess
 import time
+from typing import TYPE_CHECKING, Callable, Literal, TypedDict, TypeVar
 
 # (TODO: GhostScreaming) It will be removed later.
 from paddle.base import core
 
 from .log_util import logger
 
+if TYPE_CHECKING:
+    from typing_extensions import ParamSpec
+
+    _InputT = ParamSpec("_InputT")
+    _RetT = TypeVar("_RetT")
+
+    _HDFSClientConfig = TypedDict(
+        '_HDFSClientConfig', {'fs.default.name': str, 'hadoop.job.ugi': str}
+    )
+
+    class _FileInfo(TypedDict):
+        path: str
+        size: int
+
+
 __all__ = []
 
 
@@ -128,7 +146,7 @@ class LocalFS(FS):
 
     """
 
-    def ls_dir(self, fs_path):
+    def ls_dir(self, fs_path: str) -> tuple[list[str], list[str]]:
         """
         List directories and files under `fs_path` .
 
@@ -162,7 +180,7 @@ def ls_dir(self, fs_path):
 
         return dirs, files
 
-    def mkdirs(self, fs_path):
+    def mkdirs(self, fs_path: str) -> None:
         """
         Create a local directory.
 
@@ -183,7 +201,7 @@ def mkdirs(self, fs_path):
         assert not os.path.isfile(fs_path), f"{fs_path} is already a file"
         os.makedirs(fs_path, exist_ok=True)
 
-    def rename(self, fs_src_path, fs_dst_path):
+    def rename(self, fs_src_path: str, fs_dst_path: str) -> None:
         """
         Rename the file.
 
@@ -217,7 +235,7 @@ def _rmr(self, fs_path):
     def _rm(self, fs_path):
         os.remove(fs_path)
 
-    def delete(self, fs_path):
+    def delete(self, fs_path: str) -> None:
         """
         Delete the local file path, whether it's a file or directory.
 
@@ -243,10 +261,10 @@ def delete(self, fs_path):
 
         return self._rmr(fs_path)
 
-    def need_upload_download(self):
+    def need_upload_download(self) -> Literal[False]:
         return False
 
-    def is_file(self, fs_path):
+    def is_file(self, fs_path: str) -> bool:
         """
         Whether the local file path is a file.
 
@@ -271,7 +289,7 @@ def is_file(self, fs_path):
         """
         return os.path.isfile(fs_path)
 
-    def is_dir(self, fs_path):
+    def is_dir(self, fs_path: str) -> bool:
         """
         Whether the local file path is a directory.
 
@@ -296,7 +314,7 @@ def is_dir(self, fs_path):
         """
         return os.path.isdir(fs_path)
 
-    def is_exist(self, fs_path):
+    def is_exist(self, fs_path: str) -> bool:
         """
         Whether the local file path exists.
 
@@ -319,7 +337,7 @@ def is_exist(self, fs_path):
         """
         return os.path.exists(fs_path)
 
-    def touch(self, fs_path, exist_ok=True):
+    def touch(self, fs_path: str, exist_ok: bool = True) -> None:
         """
         Create a local file.
 
@@ -346,7 +364,13 @@ def touch(self, fs_path, exist_ok=True):
         with open(fs_path, 'a'):
             pass
 
-    def mv(self, src_path, dst_path, overwrite=False, test_exists=False):
+    def mv(
+        self,
+        src_path: str,
+        dst_path: str,
+        overwrite: bool = False,
+        test_exists: bool = False,
+    ) -> None:
         """
         Move a local file or directory from `src_path` to `dst_path` .
 
@@ -378,7 +402,7 @@ def mv(self, src_path, dst_path, overwrite=False, test_exists=False):
 
         return self.rename(src_path, dst_path)
 
-    def list_dirs(self, fs_path):
+    def list_dirs(self, fs_path: str) -> list[str]:
         """
         Only list directories under `fs_path` .
 
@@ -408,10 +432,12 @@ def list_dirs(self, fs_path):
         return dirs
 
 
-def _handle_errors(max_time_out=None):
-    def decorator(f):
+def _handle_errors(
+    max_time_out: float | None = None,
+) -> Callable[[Callable[_InputT, _RetT]], Callable[_InputT, _RetT]]:
+    def decorator(f: Callable[_InputT, _RetT]) -> Callable[_InputT, _RetT]:
         @functools.wraps(f)
-        def handler(*args, **kwargs):
+        def handler(*args: _InputT.args, **kwargs: _InputT.kwargs) -> _RetT:
             o = args[0]
             time_out = max_time_out
             if time_out is None:
@@ -473,13 +499,15 @@ class HDFSClient(FS):
 
     """
 
+    pre_commands: list[str]
+
     def __init__(
         self,
-        hadoop_home,
-        configs,
-        time_out=5 * 60 * 1000,  # ms
-        sleep_inter=1000,
-    ):  # ms
+        hadoop_home: str,
+        configs: _HDFSClientConfig,
+        time_out: int = 5 * 60 * 1000,  # ms
+        sleep_inter: int = 1000,
+    ) -> None:  # ms
         self.pre_commands = []
         hadoop_bin = f'{hadoop_home}/bin/hadoop'
         self.pre_commands.append(hadoop_bin)
@@ -525,9 +553,11 @@ def _run_safe_cmd(self, cmd, redirect_stderr=False, retry_times=5):
                     exe_cmd,
                     check=True,
                     stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT
-                    if redirect_stderr
-                    else subprocess.PIPE,
+                    stderr=(
+                        subprocess.STDOUT
+                        if redirect_stderr
+                        else subprocess.PIPE
+                    ),
                     text=True,
                 )
                 output = process.stdout
@@ -543,7 +573,7 @@ def _run_safe_cmd(self, cmd, redirect_stderr=False, retry_times=5):
             raise FSShellCmdAborted(cmd)
 
     @_handle_errors()
-    def list_dirs(self, fs_path):
+    def list_dirs(self, fs_path: str) -> list[str]:
         """
         Only list directories under `fs_path` .
 
@@ -577,7 +607,7 @@ def list_dirs(self, fs_path):
         return dirs
 
     @_handle_errors()
-    def ls_dir(self, fs_path):
+    def ls_dir(self, fs_path: str) -> tuple[list[str], list[str]]:
         """
         List directories and files under `fs_path` .
 
@@ -641,7 +671,7 @@ def _test_match(self, lines):
         return None
 
     @_handle_errors()
-    def is_dir(self, fs_path):
+    def is_dir(self, fs_path: str) -> bool:
         """
         Whether the remote HDFS path is a directory.
 
@@ -687,7 +717,7 @@ def _is_dir(self, fs_path):
 
         return True
 
-    def is_file(self, fs_path):
+    def is_file(self, fs_path: str) -> bool:
         """
         Whether the remote HDFS path is a file.
 
@@ -720,7 +750,7 @@ def is_file(self, fs_path):
         return not self._is_dir(fs_path)
 
     @_handle_errors()
-    def is_exist(self, fs_path):
+    def is_exist(self, fs_path: str) -> bool:
         """
         Whether the remote HDFS path exists.
 
@@ -755,7 +785,9 @@ def is_exist(self, fs_path):
 
         return True
 
-    def upload_dir(self, local_dir, dest_dir, overwrite=False):
+    def upload_dir(
+        self, local_dir: str, dest_dir: str, overwrite: bool = False
+    ) -> None:
         """
         upload dir to hdfs
         Args:
@@ -775,7 +807,13 @@ def upload_dir(self, local_dir, dest_dir, overwrite=False):
         self._try_upload(local_dir, dest_dir)
 
     # can't retry
-    def upload(self, local_path, fs_path, multi_processes=5, overwrite=False):
+    def upload(
+        self,
+        local_path: str,
+        fs_path: str,
+        multi_processes: int = 5,
+        overwrite: bool = False,
+    ) -> None:
         """
         Upload the local path to remote HDFS.
 
@@ -867,7 +905,13 @@ def _try_upload(self, local_path, fs_path):
             raise e
 
     # can't retry
-    def download(self, fs_path, local_path, multi_processes=5, overwrite=False):
+    def download(
+        self,
+        fs_path: str,
+        local_path: str,
+        multi_processes: int = 5,
+        overwrite: bool = False,
+    ) -> None:
         """
         Download remote HDFS path to the local.
 
@@ -941,7 +985,7 @@ def _try_download(self, fs_path, local_path):
             raise e
 
     @_handle_errors()
-    def mkdirs(self, fs_path):
+    def mkdirs(self, fs_path: str) -> None:
         """
         Create a remote HDFS directory.
 
@@ -986,7 +1030,13 @@ def mkdirs(self, fs_path):
             if ret != 0:
                 raise ExecuteError(cmd)
 
-    def mv(self, fs_src_path, fs_dst_path, overwrite=False, test_exists=True):
+    def mv(
+        self,
+        fs_src_path: str,
+        fs_dst_path: str,
+        overwrite: bool = False,
+        test_exists: bool = True,
+    ) -> None:
         """
         Move a remote HDFS file or directory from `fs_src_path` to `fs_dst_path` .
 
@@ -1051,7 +1101,7 @@ def _rm(self, fs_path):
             raise ExecuteError(cmd)
 
     @_handle_errors()
-    def delete(self, fs_path):
+    def delete(self, fs_path: str) -> None:
         """
         Delete a remote HDFS path, whether it's a file or directory.
 
@@ -1084,7 +1134,7 @@ def delete(self, fs_path):
 
         return self._rm(fs_path)
 
-    def touch(self, fs_path, exist_ok=True):
+    def touch(self, fs_path: str, exist_ok: bool = True) -> None:
         """
         Create a remote HDFS file.
 
@@ -1124,15 +1174,15 @@ def _touchz(self, fs_path):
         if ret != 0:
             raise ExecuteError(cmd)
 
-    def need_upload_download(self):
+    def need_upload_download(self) -> Literal[True]:
         return True
 
-    def cat(self, fs_path=None):
+    def cat(self, fs_path: str | None = None) -> str:
         """
         Cat a remote HDFS file.
 
         Args:
-            fs_path(str): The HDFS file path.
+            fs_path(str|None): The HDFS file path.
 
         Returns:
             file content
@@ -1194,7 +1244,7 @@ def _split_files(self, files, trainer_id, trainers):
 
         return trainer_files[trainer_id]
 
-    def list_files_info(self, path_list):
+    def list_files_info(self, path_list: list[str]) -> list[_FileInfo]:
         """
         list_files return file path and size
         Args:
