From 67d0dd7f9c76c8f3bb1a665a44567fe09b7b8ef1 Mon Sep 17 00:00:00 2001 From: pytorchbot Date: Mon, 22 Apr 2024 14:51:44 -0700 Subject: [PATCH] Add a pure python wrapper to pybindings.portable_lib (#3137) (#3218) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/3137 When installed as a pip wheel, we must import `torch` before trying to import the pybindings shared library extension. This will load libtorch.so and related libs, ensuring that the pybindings lib can resolve those runtime dependencies. So, add a pure python wrapper that lets us do this when users say `import executorch.extension.pybindings.portable_lib` We only need this for OSS, so don't bother doing this for other pybindings targets. Reviewed By: orionr, mikekgfb Differential Revision: D56317150 fbshipit-source-id: 920382636732aa276c25a76163afb7d28b1846d0 (cherry picked from commit 969aa96904c436ee43bca0dd98bb431d298c1b72) Co-authored-by: Dave Bort --- CMakeLists.txt | 5 +++- extension/pybindings/TARGETS | 16 +++++++++---- extension/pybindings/portable_lib.py | 34 ++++++++++++++++++++++++++++ extension/pybindings/pybindings.pyi | 12 ++++++---- setup.py | 2 +- 5 files changed, 59 insertions(+), 10 deletions(-) create mode 100644 extension/pybindings/portable_lib.py diff --git a/CMakeLists.txt b/CMakeLists.txt index fa19d5bcf2..d3eb9041a2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -548,8 +548,11 @@ if(EXECUTORCH_BUILD_PYBIND) # pybind portable_lib pybind11_add_module(portable_lib extension/pybindings/pybindings.cpp) + # The actual output file needs a leading underscore so it can coexist with + # portable_lib.py in the same python package. + set_target_properties(portable_lib PROPERTIES OUTPUT_NAME "_portable_lib") target_compile_definitions(portable_lib - PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=portable_lib) + PUBLIC EXECUTORCH_PYTHON_MODULE_NAME=_portable_lib) target_include_directories(portable_lib PRIVATE ${TORCH_INCLUDE_DIRS}) target_compile_options(portable_lib PUBLIC ${_pybind_compile_options}) target_link_libraries( diff --git a/extension/pybindings/TARGETS b/extension/pybindings/TARGETS index 0b4e9ef304..9dee0e208b 100644 --- a/extension/pybindings/TARGETS +++ b/extension/pybindings/TARGETS @@ -30,9 +30,9 @@ runtime.genrule( srcs = [":pybinding_types"], outs = { "aten_lib.pyi": ["aten_lib.pyi"], - "portable_lib.pyi": ["portable_lib.pyi"], + "_portable_lib.pyi": ["_portable_lib.pyi"], }, - cmd = "cp $(location :pybinding_types)/* $OUT/portable_lib.pyi && cp $(location :pybinding_types)/* $OUT/aten_lib.pyi", + cmd = "cp $(location :pybinding_types)/* $OUT/_portable_lib.pyi && cp $(location :pybinding_types)/* $OUT/aten_lib.pyi", visibility = ["//executorch/extension/pybindings/..."], ) @@ -46,8 +46,9 @@ executorch_pybindings( executorch_pybindings( compiler_flags = ["-std=c++17"], cppdeps = PORTABLE_MODULE_DEPS + MODELS_ATEN_OPS_LEAN_MODE_GENERATED_LIB, - python_module_name = "portable_lib", - types = ["//executorch/extension/pybindings:pybindings_types_gen[portable_lib.pyi]"], + # Give this an underscore prefix because it has a pure python wrapper. + python_module_name = "_portable_lib", + types = ["//executorch/extension/pybindings:pybindings_types_gen[_portable_lib.pyi]"], visibility = ["PUBLIC"], ) @@ -58,3 +59,10 @@ executorch_pybindings( types = ["//executorch/extension/pybindings:pybindings_types_gen[aten_lib.pyi]"], visibility = ["PUBLIC"], ) + +runtime.python_library( + name = "portable_lib", + srcs = ["portable_lib.py"], + visibility = ["@EXECUTORCH_CLIENTS"], + deps = [":_portable_lib"], +) diff --git a/extension/pybindings/portable_lib.py b/extension/pybindings/portable_lib.py new file mode 100644 index 0000000000..b9ed089f91 --- /dev/null +++ b/extension/pybindings/portable_lib.py @@ -0,0 +1,34 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +# When installed as a pip wheel, we must import `torch` before trying to import +# the pybindings shared library extension. This will load libtorch.so and +# related libs, ensuring that the pybindings lib can resolve those runtime +# dependencies. +import torch as _torch + +# Let users import everything from the C++ _portable_lib extension as if this +# python file defined them. Although we could import these dynamically, it +# wouldn't preserve the static type annotations. +from executorch.extension.pybindings._portable_lib import ( # noqa: F401 + # Disable "imported but unused" (F401) checks. + _create_profile_block, # noqa: F401 + _dump_profile_results, # noqa: F401 + _get_operator_names, # noqa: F401 + _load_bundled_program_from_buffer, # noqa: F401 + _load_for_executorch, # noqa: F401 + _load_for_executorch_from_buffer, # noqa: F401 + _load_for_executorch_from_bundled_program, # noqa: F401 + _reset_profile_results, # noqa: F401 + BundledModule, # noqa: F401 + ExecuTorchModule, # noqa: F401 +) + +# Clean up so that `dir(portable_lib)` is the same as `dir(_portable_lib)` +# (apart from some __dunder__ names). +del _torch diff --git a/extension/pybindings/pybindings.pyi b/extension/pybindings/pybindings.pyi index 0392b5b7ce..8f2b43cac9 100644 --- a/extension/pybindings/pybindings.pyi +++ b/extension/pybindings/pybindings.pyi @@ -7,9 +7,12 @@ # pyre-strict from typing import Any, Dict, List, Sequence, Tuple -class ExecutorchModule: +class ExecuTorchModule: + # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def __call__(self, inputs: Any) -> List[Any]: ... + # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def run_method(self, method_name: str, inputs: Sequence[Any]) -> List[Any]: ... + # pyre-ignore[2, 3]: "Any" in parameter and return type annotations. def forward(self, inputs: Sequence[Any]) -> List[Any]: ... # Bundled program methods. def load_bundled_input( @@ -30,16 +33,17 @@ class BundledModule: ... def _load_for_executorch( path: str, enable_etdump: bool = False -) -> ExecutorchModule: ... +) -> ExecuTorchModule: ... def _load_for_executorch_from_buffer( buffer: bytes, enable_etdump: bool = False -) -> ExecutorchModule: ... +) -> ExecuTorchModule: ... def _load_for_executorch_from_bundled_program( module: BundledModule, enable_etdump: bool = False -) -> ExecutorchModule: ... +) -> ExecuTorchModule: ... def _load_bundled_program_from_buffer( buffer: bytes, non_const_pool_size: int = ... ) -> BundledModule: ... +def _get_operator_names() -> List[str]: ... def _create_profile_block(name: str) -> None: ... def _dump_profile_results() -> bytes: ... def _reset_profile_results() -> None: ... diff --git a/setup.py b/setup.py index 6458fdc10c..fa67bebf0b 100644 --- a/setup.py +++ b/setup.py @@ -435,7 +435,7 @@ def get_ext_modules() -> list[Extension]: # portable kernels, and a selection of backends. This lets users # load and execute .pte files from python. BuiltExtension( - "portable_lib.*", "executorch.extension.pybindings.portable_lib" + "_portable_lib.*", "executorch.extension.pybindings._portable_lib" ) )