[Typing] 添加 libpaddle 的 stub 文件生成脚本 (#66357)

megemini · web-flow · commit 478d7cbca349 · 2024-09-05T11:10:08.000+08:00
* [init] pybind11 stub

* [init] pybind11 stub

* [patch] pybind11-stubgen

* [Change] move libpaddle stub to _typing

* [Add] genenrate module or file

* [Add] setup config

* [Change] req

* [Change] req

* [Change] python/requirements.txt

* [Add] core.pyi _C_ops.pyi

* [Change] make libpaddle a module

* [Change] remove base __init__ import libpaddle

* [Add] parse ops yaml

* [Update] some comments

* [Fix] setup with _C_ops.pyi

* [Update] PYBIND11_MAPPING

* [Fix] Tensor as paddle.Tensor

* [Fix] NO GPUPlace

* [Fix] NO GPUPinnedPlace

* [Change] import libpaddle from root

* [Fix] var mapping

* [Update] refactor ops yaml base api class

* [Change] remove patch missing import

* [Add] patch pybind11 annotations

* [Change] convert bad exp to typing.Any

* [Update] pybind11 mapping

* [Fix] mapping invalid to Any

* [Change] desc for setup

* [Change] filter [build] out of setup_requires

* [Update] desc

* [Fix] codestyle

* [Update] package_data format

* [Update] clear pybind11 mapping

* [Fix] typing
diff --git a/.gitignore b/.gitignore
@@ -110,3 +110,4 @@ paddle/phi/kernels/fusion/cutlass/conv2d/build
 paddle/phi/kernels/fusion/cutlass/conv2d/cutlass
 paddle/phi/kernels/fusion/cutlass/gemm_epilogue/build
 paddle/phi/kernels/fusion/cutlass/gemm_epilogue/cutlass
+python/paddle/_typing/libs/**/*.pyi
diff --git a/python/paddle/_C_ops.pyi b/python/paddle/_C_ops.pyi
@@ -0,0 +1,16 @@
+# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle._typing.libs.libpaddle.eager.ops import *  # noqa: F403
+from paddle._typing.libs.libpaddle.pir.ops import *  # noqa: F403
diff --git a/python/paddle/_typing/libs/README.md b/python/paddle/_typing/libs/README.md
@@ -0,0 +1,5 @@
+# Stub files for pybind11 C++ APIs
+
+The files in this directory (`python/paddle/_typing/libs`) are auto generated from pybind11 C++ APIs.
+
+Please do NOT edit these files, or the changes may be lost.
diff --git a/python/paddle/base/core.pyi b/python/paddle/base/core.pyi
@@ -12,93 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
-
-class Place:
-    def custom_device_id(self) -> int: ...
-    def custom_device_type(self) -> str: ...
-    def gpu_device_id(self) -> int: ...
-    def ipu_device_id(self) -> int: ...
-    def is_cpu_place(self) -> bool: ...
-    def is_cuda_pinned_place(self) -> bool: ...
-    def is_custom_place(self) -> bool: ...
-    def is_gpu_place(self) -> bool: ...
-    def is_ipu_place(self) -> bool: ...
-    def is_xpu_place(self) -> bool: ...
-    def set_place(self, place: Place) -> None: ...
-    def xpu_device_id(self) -> int: ...
-
-class CPUPlace:
-    def is_float16_supported(self) -> bool: ...
-    def is_bfloat16_supported(self) -> bool: ...
-
-class CUDAPlace:
-    def __init__(self, id: int, /) -> None: ...
-    def get_device_id(self) -> int: ...
-    def is_float16_supported(self) -> bool: ...
-    def is_bfloat16_supported(self) -> bool: ...
-
-class CUDAPinnedPlace: ...
-
-class NPUPlace:
-    def __init__(self, id: int, /) -> None: ...
-
-class IPUPlace: ...
-
-class CustomPlace:
-    def __init__(self, name: str, id: int, /) -> None: ...
-    def get_device_id(self) -> int: ...
-    def get_device_type(self) -> str: ...
-
-class MLUPlace:
-    def __init__(self, id: int, /) -> None: ...
-
-class XPUPlace:
-    def __init__(self, id: int, /) -> None: ...
-    def get_xpu_device_count(self) -> int: ...
-    def set_xpu_debug_level(self, level: int) -> None: ...
-    def get_xpu_device_version(self, device_id: int) -> int: ...
-    def get_xpu_device_op_support_types(
-        self, op_name: str, version: int
-    ) -> list[str]: ...
-    def get_xpu_device_op_list(self, version: int) -> dict[str, Any]: ...
-    def is_float16_supported(self) -> bool: ...
-    def is_bfloat16_supported(self) -> bool: ...
-
-class Placement:
-    def __init__(self) -> None: ...
-    def __eq__(self, arg0: Placement) -> bool: ...
-    def __hash__(self) -> int: ...
-    def __ne__(self, arg0: Placement) -> bool: ...
-    def is_partial(self) -> bool: ...
-    def is_replicated(self) -> bool: ...
-    def is_shard(self, dim: int | None = None) -> bool: ...
-
-class Shard(Placement):
-    def __init__(self, arg0: int) -> None: ...
-    def get_dim(self) -> int: ...
-
-class finfo:
-    min: float
-    max: float
-    eps: float
-    resolution: float
-    smallest_normal: float
-    tiny: float
-    bits: int
-    dtype: str
-
-class iinfo:
-    min: int
-    max: int
-    bits: int
-    dtype: str
-
-def is_compiled_with_cuda() -> bool: ...
-def set_nan_inf_debug_path(arg0: str) -> None: ...
-
-class task:
-    def is_completed(self) -> bool: ...
-    def is_sync(self) -> bool: ...
-    def synchronize(self) -> None: ...
-    def wait(self, timeout: int = 0) -> bool: ...
+from paddle._typing.libs.libpaddle import *  # noqa: F403
diff --git a/python/paddle/base/libpaddle.pyi b/python/paddle/base/libpaddle.pyi
@@ -0,0 +1,15 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from paddle._typing.libs.libpaddle import *  # noqa: F403
diff --git a/python/paddle/distributed/spawn.py b/python/paddle/distributed/spawn.py
@@ -530,7 +530,7 @@ def spawn(
             ...     process_group = group.process_group if group else None
             ...     # 2. create data parallel layer & optimizer
             ...     layer = LinearNet()
-            ...     dp_layer = paddle.DataParallel(layer, group = process_group)
+            ...     dp_layer = paddle.DataParallel(layer, group = process_group)  # type: ignore[arg-type]
             ...     loss_fn = nn.MSELoss()
             ...     adam = opt.Adam(
             ...         learning_rate=0.001, parameters=dp_layer.parameters())
diff --git a/python/requirements.txt b/python/requirements.txt
@@ -7,3 +7,4 @@ astor
 opt_einsum==3.3.0
 networkx
 typing_extensions
+pybind11-stubgen==2.5.1 # generate stub file from pybind11 C++ APIs [build]
diff --git a/python/setup.py.in b/python/setup.py.in
@@ -13,6 +13,7 @@ import shlex
 import ctypes
 
 from contextlib import contextmanager
+from pathlib import Path
 from setuptools import Command
 from setuptools import setup, Distribution, Extension
 from setuptools.command.install import install as InstallCommandBase
@@ -769,7 +770,12 @@ packages=['paddle',
           'paddle.pir',
           'paddle.decomposition',
           'paddle._typing',
-          ]
+          'paddle._typing.libs',
+          'paddle._typing.libs.libpaddle',
+          'paddle._typing.libs.libpaddle.pir',
+          'paddle._typing.libs.libpaddle.eager',
+          'paddle._typing.libs.libpaddle.eager.ops',
+        ]
 
 with open('@PADDLE_SOURCE_DIR@/python/requirements.txt') as f:
     setup_requires = f.read().splitlines()
@@ -788,6 +794,9 @@ if sys.version_info >= (3,8):
             or "<\"3.7\"" in setup_requires_i
             or "<=\"3.7\"" in setup_requires_i
             or "<\"3.8\"" in setup_requires_i
+            or setup_requires_i.strip().endswith(
+                '[build]'
+            )  # remove `[build]` requirements
         ):
             continue
         setup_requires_tmp+=[setup_requires_i]
@@ -1283,20 +1292,34 @@ if '${WITH_CPP_DIST}' == 'ON':
     install_cpp_dist_and_build_test(paddle_install_dir, paddle_lib_test_dir)
 
 # type hints
-package_data['paddle'] = package_data.get('paddle', []) + ['py.typed']
-package_data['paddle.framework'] = package_data.get(
-    'paddle.framework', []
-) + ['*.pyi']
-package_data['paddle.base'] = package_data.get('paddle.base', []) + [
-    '*.pyi'
-]
-package_data['paddle.tensor'] = package_data.get('paddle.tensor', []) + ['tensor.pyi']
-
-
-def generate_tensor_stub(paddle_binary_dir, paddle_source_dir):
-    print('-'*2, 'Generate stub file tensor.pyi ... ')
+def extend_type_hints_package_data(package_data):
+    type_hints_files = {
+        'paddle': ['py.typed','*.pyi'],
+        'paddle.framework': ['*.pyi'],
+        'paddle.base': ['*.pyi'],
+        'paddle.tensor': ['tensor.pyi'],
+        'paddle._typing': ['*.pyi'],
+        'paddle._typing.libs': ['*.pyi','*.md'],
+        'paddle._typing.libs.libpaddle': ['*.pyi'],
+        'paddle._typing.libs.libpaddle.pir': ['*.pyi'],
+        'paddle._typing.libs.libpaddle.eager': ['*.pyi'],
+        'paddle._typing.libs.libpaddle.eager.ops': ['*.pyi'],
+    }
+    for pkg, files in type_hints_files.items():
+        if pkg not in package_data:
+            package_data[pkg] = []
+        package_data[pkg] += files
+
+    return package_data
+
+package_data = extend_type_hints_package_data(package_data)
+
+
+def generate_stub_files(paddle_binary_dir, paddle_source_dir):
     script_path = paddle_source_dir + '/tools/'
     sys.path.append(script_path)
+
+    print('-' * 2, 'Generate stub file tensor.pyi ... ')
     import gen_tensor_stub
 
     gen_tensor_stub.generate_stub_file(
@@ -1309,7 +1332,42 @@ def generate_tensor_stub(paddle_binary_dir, paddle_source_dir):
         paddle_binary_dir + '/python/paddle/tensor/tensor.pyi',
         paddle_source_dir + '/python/paddle/tensor/tensor.pyi',
     )
-    print('-'*2, 'End Generate stub file tensor.pyi ... ')
+    print('-' * 2, 'End Generate stub file tensor.pyi ... ')
+
+    print('-' * 2, 'Generate stub file for python binding APIs ... ')
+    import gen_pybind11_stub
+
+    gen_pybind11_stub.generate_stub_file(
+        output_dir=str(Path(paddle_binary_dir) / 'python/paddle/_typing/libs/'),
+        module_name='paddle.base.libpaddle',
+        ignore_all_errors=True,
+        ops_yaml=[
+            paddle_source_dir
+            + "/paddle/phi/ops/yaml/ops.yaml;paddle.base.libpaddle.eager.ops",
+            paddle_source_dir
+            + "/paddle/phi/ops/yaml/ops.yaml;paddle.base.libpaddle.pir.ops",
+            paddle_source_dir
+            + "/paddle/phi/ops/yaml/sparse_ops.yaml;paddle.base.libpaddle.eager.ops;sparse",
+            paddle_source_dir
+            + "/paddle/phi/ops/yaml/sparse_ops.yaml;paddle.base.libpaddle.pir.ops;sparse",
+            paddle_source_dir
+            + "/paddle/phi/ops/yaml/strings_ops.yaml;paddle.base.libpaddle.eager.ops;strings",
+            paddle_source_dir
+            + "/paddle/phi/ops/yaml/strings_ops.yaml;paddle.base.libpaddle.pir.ops;strings",
+        ],
+    )
+
+    libpaddle_dst = paddle_source_dir + '/python/paddle/_typing/libs/libpaddle'
+    if Path(libpaddle_dst).exists():
+        shutil.rmtree(libpaddle_dst)
+
+    shutil.copytree(
+        paddle_binary_dir + '/python/paddle/_typing/libs/libpaddle',
+        libpaddle_dst,
+    )
+
+    print('-' * 2, 'End Generate stub for python binding APIs ... ')
+
 
 # generate stub file `tensor.pyi`
 if os.getenv("SKIP_STUB_GEN", '').lower() not in [
@@ -1320,7 +1378,7 @@ if os.getenv("SKIP_STUB_GEN", '').lower() not in [
     'on',
     '1',
 ]:
-    generate_tensor_stub('${PADDLE_BINARY_DIR}', '${PADDLE_SOURCE_DIR}')
+    generate_stub_files('${PADDLE_BINARY_DIR}', '${PADDLE_SOURCE_DIR}')
 
 
 with redirect_stdout():
diff --git a/setup.py b/setup.py
diff --git a/tools/gen_pybind11_stub.py b/tools/gen_pybind11_stub.py
