[Typing] Add type annotations for python/paddle/incubate/distributed/utils/io/{dist_save,save_for_auto}.py (#67679)

megemini · web-flow · commit 0b9c75291592 · 2024-08-26T10:19:18.000+08:00
diff --git a/python/paddle/incubate/distributed/utils/io/dist_save.py b/python/paddle/incubate/distributed/utils/io/dist_save.py
@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 import re
 import sys
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
 import paddle
 import paddle.distributed as dist
@@ -24,11 +27,33 @@
 
 from .save_for_auto import save_for_auto_inference
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from io import BytesIO
+
+    from typing_extensions import Unpack
+
+    from paddle import Tensor
+    from paddle._typing import NestedStructure
+    from paddle.nn.layer.layers import _StateDict
+    from paddle.static import Program
+
+    class _SaveConfig(TypedDict, total=False):
+        use_binary_format: bool
+        gather_to: int | Sequence[int] | None
+        state_type: Literal['params', 'opt']
+        max_grouped_size: str | int
+
+
 __all__ = ["save", "save_for_auto_inference"]
 
 
 @dygraph_only
-def save(state_dict, path, **configs):
+def save(
+    state_dict: dict[str, Any] | _StateDict | NestedStructure[Tensor] | Program,
+    path: str | BytesIO,
+    **configs: Unpack[_SaveConfig],
+) -> None:
     '''
     Save a state dict to the specified path in both distributed and single-card environment.
 
@@ -73,6 +98,7 @@ def save(state_dict, path, **configs):
         .. code-block:: python
 
             >>> # doctest: +SKIP('TODO: the error will be fix in the feature')
+            >>> # type: ignore
             >>> import paddle
             >>> paddle.distributed.init_process_group(backend='nccl')
             >>> paddle.distributed.fleet.init(is_collective=True)
diff --git a/python/paddle/incubate/distributed/utils/io/save_for_auto.py b/python/paddle/incubate/distributed/utils/io/save_for_auto.py
@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 import os
 import pickle
 import re
+from typing import TYPE_CHECKING
 
 import numpy as np
 
@@ -28,11 +31,16 @@
 )
 from paddle.distributed.fleet.utils.log_util import logger
 
+if TYPE_CHECKING:
+    from paddle.nn import Layer
+
 __all__ = ["save_for_auto_inference"]
 
 
 @dygraph_only
-def save_for_auto_inference(path_prefix, dist_model, cvt2cpu=False):
+def save_for_auto_inference(
+    path_prefix: str, dist_model: Layer, cvt2cpu: bool = False
+) -> None:
     """
     Description：
         Save model parameters for auto parallel inference.
@@ -56,9 +64,10 @@ def save_for_auto_inference(path_prefix, dist_model, cvt2cpu=False):
         .. code-block:: python
 
             >>> # doctest: +SKIP('model not exist')
-            >>> dist_model = build_distributed_model()
+            >>> from paddle.incubate.distributed.utils.io import save_for_auto_inference
+            >>> dist_model = build_distributed_model()  # type: ignore[name-defined]
             >>> path_prefix = "path/to/save_infer"
-            >>> save_for_auto_inference(path_prefix, dist_model=dist_model, original_model=single_model, cvt2cpu=False)
+            >>> save_for_auto_inference(path_prefix, dist_model=dist_model, cvt2cpu=False)
 
     Outputs:
         path/to/save_infer_dist0.pdparams path/to/save_infer_dist1.pdparams path/to/save_infer_dist2.pdparams ...
