From 6ab403c3d5a3670975d4b3de3289c8119803f675 Mon Sep 17 00:00:00 2001
From: Scott Todd <scott.todd0@gmail.com>
Date: Thu, 26 Sep 2024 14:21:55 -0700
Subject: [PATCH] [libshortfin] Re-land "Fix standalone Python package build to
 include subpackages." (#225)

Progress on https://github.com/nod-ai/SHARK-Platform/issues/130.
This reverts https://github.com/nod-ai/SHARK-Platform/pull/224 to
re-land https://github.com/nod-ai/SHARK-Platform/pull/223.

This now uses `SOURCE_DIR` instead of `SETUPPY_DIR` to fix package
discovery when running in pre-built mode, which will fix the errors
reported on CI:
```
 ImportError while loading conftest '/home/runner/work/SHARK-Platform/SHARK-Platform/libshortfin/tests/conftest.py'.
tests/conftest.py:10: in <module>
    import shortfin as sf
E   ModuleNotFoundError: No module named 'shortfin'
```

Logs from a failed build:
```
   setup.py running in pre-built mode:
    SOURCE_DIR = /home/runner/work/SHARK-Platform/SHARK-Platform/libshortfin
    BINARY_DIR = /home/runner/work/SHARK-Platform/SHARK-Platform/libshortfin/build
  Found libshortfin packages: ['_shortfin_default']
```

Logs from a successful build:
```
   setup.py running in pre-built mode:
    SOURCE_DIR = /home/runner/work/SHARK-Platform/SHARK-Platform/libshortfin
    BINARY_DIR = /home/runner/work/SHARK-Platform/SHARK-Platform/libshortfin/build
  Found libshortfin packages: ['shortfin_apps', 'shortfin', '_shortfin', 'shortfin_apps.llm', 'shortfin_apps.llm.components', 'shortfin.interop', 'shortfin.support', 'shortfin.interop.fastapi', 'shortfin.interop.support']
```

---------

Co-authored-by: Marius Brehler <marius.brehler@gmail.com>
---
 libshortfin/README.md | 35 +++++++++++++++++++++++++++--------
 libshortfin/setup.py  | 25 ++++++++++++++++---------
 2 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/libshortfin/README.md b/libshortfin/README.md
index c61a7e38e..697662ac0 100644
--- a/libshortfin/README.md
+++ b/libshortfin/README.md
@@ -3,12 +3,18 @@
 Build options
 
 1. Native C++ build
-2. Python release build
-3. Python dev build
+2. Local Python release build
+3. Package Python release build
+4. Python dev build
+
+Prerequisites
+
+* A modern C/C++ compiler, such as clang 18 or gcc 12
+* A modern Python, such as Python 3.12
 
 ## Native C++ Builds
 
-```
+```bash
 cmake -GNinja -S. -Bbuild \
     -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \
     -DCMAKE_LINKER_TYPE=LLD
@@ -18,15 +24,28 @@ If Python bindings are enabled in this mode (`-DSHORTFIN_BUILD_PYTHON_BINDINGS=O
 then `pip install -e build/` will install from the build dir (and support
 build/continue).
 
-## Python Release Builds
+## Local Python Release Builds
 
-```
+```bash
 pip install -v -e .
 ```
 
-## Python Dev Builds
+## Package Python Release Builds
+
+```bash
+# Build shortfin.*.whl into the dist/ directory
+#   e.g. `shortfin-0.9-cp312-cp312-linux_x86_64.whl`
+python3 -m pip wheel -v -w dist .
 
+# Install the built wheel.
+python3 -m pip install dist/*.whl
 ```
+
+TODO: helper script to build under manylinux using Docker
+
+## Python Dev Builds
+
+```bash
 # Install build system pre-reqs (since we are building in dev mode, this
 # is not done for us). See source of truth in pyproject.toml:
 pip install setuptools wheel
@@ -81,13 +100,13 @@ this is not considered to be a significant issue.
 
 Run platform independent tests only:
 
-```
+```bash
 pytest tests/
 ```
 
 Run tests including for a specific platform:
 
-```
+```bash
 pytest tests/ --system amdgpu
 ```
 
diff --git a/libshortfin/setup.py b/libshortfin/setup.py
index a98d2fb00..4b6dd7ac8 100644
--- a/libshortfin/setup.py
+++ b/libshortfin/setup.py
@@ -11,6 +11,7 @@
 import os
 from pathlib import Path
 from distutils.command.build import build as _build
+from setuptools import find_namespace_packages
 from setuptools.command.build_ext import build_ext as _build_ext
 from setuptools.command.build_py import build_py as _build_py
 
@@ -307,20 +308,26 @@ def populate_built_package(abs_dir):
 if ENABLE_TRACY:
     populate_built_package(os.path.join(PYTHON_TRACY_BINARY_DIR / "_shortfin_tracy"))
 
+packages = find_namespace_packages(
+    where=os.path.join(SOURCE_DIR, "python"),
+    include=[
+        "_shortfin",
+        "_shortfin_default",
+        "shortfin",
+        "shortfin.*",
+        "shortfin_apps",
+        "shortfin_apps.*",
+    ]
+    + (["_shortfin_tracy"] if ENABLE_TRACY else []),
+)
+print(f"Found libshortfin packages: {packages}")
+
 setup(
     name="shortfin",
     version="0.9",
     description="Shortfin native library implementation",
     author="SHARK Authors",
-    packages=(
-        [
-            "_shortfin",
-            "_shortfin_default",
-            "shortfin",
-            "shortfin_apps",
-        ]
-        + (["_shortfin_tracy"] if ENABLE_TRACY else [])
-    ),
+    packages=packages,
     zip_safe=False,
     package_dir=combine_dicts(
         {