Segregate the user-API and support in python runtime packages. (#14671)

This keeps us from loading the Python native library while resolving console script executables. Prior to this, it would trip up Tracy because technically, the python interpreter would attach first and then the console binary would attach. By segregating the packages like this, we make such things impossible.
iree-org · Aug 14, 2023 · 48f7394 · 48f7394
1 parent 9376e5f
commit 48f7394
Show file tree

Hide file tree

Showing 12 changed files with 51 additions and 31 deletions.
diff --git a/runtime/bindings/python/CMakeLists.txt b/runtime/bindings/python/CMakeLists.txt
@@ -114,7 +114,6 @@ iree_py_library(
   SRCS
     "iree/runtime/__init__.py"
     "iree/runtime/_binding.py"
-    "iree/runtime/_runtime_libs.py"
     "iree/runtime/array_interop.py"
     "iree/runtime/benchmark.py"
     "iree/runtime/flags.py"
@@ -123,12 +122,13 @@ iree_py_library(
     "iree/runtime/system_setup.py"
     "iree/runtime/tracing.py"
     "iree/runtime/version.py"
-    "iree/runtime/scripts/iree_benchmark_trace/__main__.py"
-    "iree/runtime/scripts/iree_benchmark_module/__main__.py"
-    "iree/runtime/scripts/iree_cpuinfo/__main__.py"
-    "iree/runtime/scripts/iree_run_trace/__main__.py"
-    "iree/runtime/scripts/iree_run_module/__main__.py"
-    "iree/runtime/scripts/iree_tracy_capture/__main__.py"
+    "iree/_runtime/libs.py"
+    "iree/_runtime/scripts/iree_benchmark_trace/__main__.py"
+    "iree/_runtime/scripts/iree_benchmark_module/__main__.py"
+    "iree/_runtime/scripts/iree_cpuinfo/__main__.py"
+    "iree/_runtime/scripts/iree_run_trace/__main__.py"
+    "iree/_runtime/scripts/iree_run_module/__main__.py"
+    "iree/_runtime/scripts/iree_tracy_capture/__main__.py"
   PYEXT_DEPS
     iree_runtime_bindings_python_PyExtRt
 )

diff --git a/runtime/bindings/python/iree/_runtime/README.md b/runtime/bindings/python/iree/_runtime/README.md
@@ -0,0 +1,14 @@
+# Non user-API runtime support.
+
+This package contains elements of the runtime which are not user
+visible parts of the API. It is a pure namespace package that does
+no native initialization by default (unlike the `runtime` package
+which initializes the native library as part of exporting its API).
+
+This distinction is important because some runtime libraries (like
+`tracy` will attempt to connect to a profiler on load and there are
+cases where we want to do things without triggering such load-time
+behavior).
+
+As peers to this package, there are native `_runtime_libs*` packages
+for different variations of native libraries and tools.
diff --git a/...ings/python/iree/runtime/_runtime_libs.py → ...ime/bindings/python/iree/_runtime/libs.py b/...ings/python/iree/runtime/_runtime_libs.py → ...ime/bindings/python/iree/_runtime/libs.py
@@ -3,7 +3,7 @@
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-"""Trampoline to the right iree._runtime_libs* module."""
+"""Trampoline to the right iree._runtime_libs*._runtime module."""
 
 import os
 import sys
@@ -12,7 +12,7 @@
 variant = os.getenv("IREE_PY_RUNTIME", "default")
 if variant == "tracy":
     try:
-        import iree._runtime_libs_tracy as _libs
+        import iree._runtime_libs_tracy as libs
     except ModuleNotFoundError as e:
         raise ModuleNotFoundError(
             "IREE Tracy runtime requested via IREE_PY_RUNTIME but it is not "
@@ -25,8 +25,8 @@
             f"Unknown value for IREE_PY_RUNTIME env var ({variant}): " f"Using default"
         )
         variant = "default"
-    import iree._runtime_libs as _libs
+    import iree._runtime_libs as libs
 
-_libs.name = variant
-_libs.library_path = _libs.__path__[0]
-sys.modules[__name__] = _libs
+libs.name = variant
+libs.library_path = libs.__path__[0]
+sys.modules[__name__] = libs
diff --git a/...scripts/iree_benchmark_module/__main__.py → ...scripts/iree_benchmark_module/__main__.py b/...scripts/iree_benchmark_module/__main__.py → ...scripts/iree_benchmark_module/__main__.py
@@ -7,13 +7,13 @@
 import os
 import subprocess
 import sys
-from ... import _runtime_libs
+from ... import libs
 
 
 def main(args=None):
     if args is None:
         args = sys.argv[1:]
-    exe = os.path.join(_runtime_libs.library_path, "iree-benchmark-module")
+    exe = os.path.join(libs.library_path, "iree-benchmark-module")
     return subprocess.call(args=[exe] + args)
 
 

diff --git a/...untime/scripts/iree_run_trace/__main__.py → .../scripts/iree_benchmark_trace/__main__.py b/...untime/scripts/iree_run_trace/__main__.py → .../scripts/iree_benchmark_trace/__main__.py
@@ -7,13 +7,13 @@
 import os
 import subprocess
 import sys
-from ... import _runtime_libs
+from ... import libs
 
 
 def main(args=None):
     if args is None:
         args = sys.argv[1:]
-    exe = os.path.join(_runtime_libs.library_path, "iree-run-trace")
+    exe = os.path.join(libs.library_path, "iree-benchmark-trace")
     return subprocess.call(args=[exe] + args)
 
 

diff --git a/.../runtime/scripts/iree_cpuinfo/__main__.py → ..._runtime/scripts/iree_cpuinfo/__main__.py b/.../runtime/scripts/iree_cpuinfo/__main__.py → ..._runtime/scripts/iree_cpuinfo/__main__.py
diff --git a/...time/scripts/iree_dump_module/__main__.py → ...time/scripts/iree_dump_module/__main__.py b/...time/scripts/iree_dump_module/__main__.py → ...time/scripts/iree_dump_module/__main__.py
@@ -15,7 +15,7 @@
 def main(args=None):
     if args is None:
         args = sys.argv[1:]
-    exe = os.path.join(_runtime_libs.library_path, "iree-dump-module")
+    exe = os.path.join(_runtime_libs.__path__[0], "iree-dump-module")
     return subprocess.call(args=[exe] + args)
 
 

diff --git a/...ntime/scripts/iree_run_module/__main__.py → ...ntime/scripts/iree_run_module/__main__.py b/...ntime/scripts/iree_run_module/__main__.py → ...ntime/scripts/iree_run_module/__main__.py
@@ -7,13 +7,13 @@
 import os
 import subprocess
 import sys
-from ... import _runtime_libs
+from ... import libs
 
 
 def main(args=None):
     if args is None:
         args = sys.argv[1:]
-    exe = os.path.join(_runtime_libs.library_path, "iree-run-module")
+    exe = os.path.join(libs.library_path, "iree-run-module")
     return subprocess.call(args=[exe] + args)
 
 

diff --git a/.../scripts/iree_benchmark_trace/__main__.py → ...untime/scripts/iree_run_trace/__main__.py b/.../scripts/iree_benchmark_trace/__main__.py → ...untime/scripts/iree_run_trace/__main__.py
@@ -7,13 +7,13 @@
 import os
 import subprocess
 import sys
-from ... import _runtime_libs
+from ... import libs
 
 
 def main(args=None):
     if args is None:
         args = sys.argv[1:]
-    exe = os.path.join(_runtime_libs.library_path, "iree-benchmark-trace")
+    exe = os.path.join(libs.library_path, "iree-run-trace")
     return subprocess.call(args=[exe] + args)
 
 

diff --git a/...me/scripts/iree_tracy_capture/__main__.py → ...me/scripts/iree_tracy_capture/__main__.py b/...me/scripts/iree_tracy_capture/__main__.py → ...me/scripts/iree_tracy_capture/__main__.py
diff --git a/runtime/bindings/python/iree/runtime/_binding.py b/runtime/bindings/python/iree/runtime/_binding.py
@@ -9,10 +9,13 @@
 locates the actual _runtime module based on environment configuration.
 
 TODO: We could rename this to _runtime since it a trampoline for
-the _runtime module we load from elsehwhere.
+the _runtime native extension module we load from elsehwhere.
 """
 
 import sys
-from ._runtime_libs import _runtime
+
+# "_runtime" is the native extension within the _runtime_libs package
+# selected.
+from .._runtime.libs import _runtime
 
 sys.modules[__name__] = _runtime
diff --git a/runtime/setup.py b/runtime/setup.py
@@ -435,6 +435,8 @@ def generate_version_py():
         include=[
             "iree.runtime",
             "iree.runtime.*",
+            "iree._runtime",
+            "iree._runtime.*",
         ],
     )
     + [
@@ -521,6 +523,7 @@ def generate_version_py():
             # Note: Must be relative path, so we line this up with the absolute
             # path built above. Note that this must exist prior to the call.
             "iree.runtime": "bindings/python/iree/runtime",
+            "iree._runtime": "bindings/python/iree/_runtime",
             "iree._runtime_libs": f"{CMAKE_INSTALL_DIR_REL}/python_packages/iree_runtime/iree/_runtime_libs",
         },
         {
@@ -563,16 +566,16 @@ def generate_version_py():
     ),
     entry_points={
         "console_scripts": [
-            "iree-run-module = iree.runtime.scripts.iree_run_module.__main__:main",
-            "iree-run-trace = iree.runtime.scripts.iree_run_trace.__main__:main",
-            "iree-benchmark-module = iree.runtime.scripts.iree_benchmark_module.__main__:main",
-            "iree-benchmark-trace = iree.runtime.scripts.iree_benchmark_trace.__main__:main",
-            "iree-dump-module = iree.runtime.scripts.iree_dump_module.__main__:main",
-            "iree-cpuinfo = iree.runtime.scripts.iree_cpuinfo.__main__:main",
+            "iree-run-module = iree._runtime.scripts.iree_run_module.__main__:main",
+            "iree-run-trace = iree._runtime.scripts.iree_run_trace.__main__:main",
+            "iree-benchmark-module = iree._runtime.scripts.iree_benchmark_module.__main__:main",
+            "iree-benchmark-trace = iree._runtime.scripts.iree_benchmark_trace.__main__:main",
+            "iree-dump-module = iree._runtime.scripts.iree_dump_module.__main__:main",
+            "iree-cpuinfo = iree._runtime.scripts.iree_cpuinfo.__main__:main",
         ]
         + (
             [
-                "iree-tracy-capture = iree.runtime.scripts.iree_tracy_capture.__main__:main"
+                "iree-tracy-capture = iree._runtime.scripts.iree_tracy_capture.__main__:main"
             ]
             if ENABLE_TRACY_TOOLS
             else []