Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[libshortfin] Fix python stub generation. #195

Merged
merged 2 commits into from
Sep 18, 2024
Merged

[libshortfin] Fix python stub generation. #195

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0ff14b7
[libshortfin] Fix python stub generation.
stellaraccident Sep 18, 2024
8f8b6cc
Do not generate across build types
stellaraccident Sep 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 32 additions & 7 deletions libshortfin/python/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,18 +35,43 @@ target_link_libraries(shortfin_python_extension
PRIVATE ${SHORTFIN_LINK_LIBRARY_NAME}
)

if (SHORTFIN_ENABLE_TRACING)
function(shortfin_python_stubs build_type)
nanobind_add_stub(
shortfin_python_extension_stub
MODULE _shortfin_tracy.lib
OUTPUT _shortfin_tracy/lib.pyi
MODULE _shortfin_${build_type}.lib
OUTPUT _shortfin_${build_type}/lib.pyi
DEPENDS shortfin_python_extension
)
else()

endfunction()

function(shortfin_python_stubs build_variant)
set(output_root "${CMAKE_CURRENT_BINARY_DIR}/_shortfin_${build_variant}")
file(MAKE_DIRECTORY ${output_root})
nanobind_add_stub(
shortfin_python_extension_stub
MODULE _shortfin_default.lib
OUTPUT _shortfin_default/lib.pyi
shortfin_python_extension_stub_lib_${build_variant}
MODULE _shortfin_${build_variant}.lib
OUTPUT ${output_root}/lib/__init__.pyi
DEPENDS shortfin_python_extension
)

nanobind_add_stub(
shortfin_python_extension_stub_array_${build_variant}
MODULE _shortfin_${build_variant}.lib.array
OUTPUT ${output_root}/lib/array.pyi
DEPENDS shortfin_python_extension
)

nanobind_add_stub(
shortfin_python_extension_stub_local_${build_variant}
MODULE _shortfin_${build_variant}.lib.local
OUTPUT ${output_root}/lib/local.pyi
DEPENDS shortfin_python_extension
)
endfunction()

if (SHORTFIN_ENABLE_TRACING)
shortfin_python_stubs(tracy)
else()
shortfin_python_stubs(default)
endif()
35 changes: 20 additions & 15 deletions libshortfin/python/_shortfin/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,23 +7,28 @@
# The proper way to import this package is via:
# from _shortfin import lib as sfl

from typing import TYPE_CHECKING

import os
import sys
import warnings

variant = os.getenv("SHORTFIN_PY_RUNTIME", "default")

if variant == "tracy":
try:
from _shortfin_tracy import lib
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
"Shortfin Tracy runtime requested via SHORTFIN_PY_RUNTIME but it is not enabled in this build"
)
print("-- Using Tracy runtime (SHORTFIN_PY_RUNTIME=tracy)", file=sys.stderr)
else:
if variant != "default":
warnings.warn(
f"Unknown value for SHORTFIN_PY_RUNTIME env var ({variant}): Using default"
)
if TYPE_CHECKING:
from _shortfin_default import lib
else:
variant = os.getenv("SHORTFIN_PY_RUNTIME", "default")

if variant == "tracy":
try:
from _shortfin_tracy import lib
except ModuleNotFoundError as e:
raise ModuleNotFoundError(
"Shortfin Tracy runtime requested via SHORTFIN_PY_RUNTIME but it is not enabled in this build"
)
print("-- Using Tracy runtime (SHORTFIN_PY_RUNTIME=tracy)", file=sys.stderr)
else:
if variant != "default":
warnings.warn(
f"Unknown value for SHORTFIN_PY_RUNTIME env var ({variant}): Using default"
)
from _shortfin_default import lib
Loading