[shortfin] Add extras_require for optional deps.

nod-ai · Sep 30, 2024 · 4c2f3ec · 4c2f3ec
1 parent 72665fc
commit 4c2f3ec
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/shortfin/setup.py b/shortfin/setup.py
@@ -381,6 +381,27 @@ def populate_built_package(abs_dir):
         [CMakeExtension("_shortfin_default.lib")]
         + ([CMakeExtension("_shortfin_tracy.lib")] if ENABLE_TRACY else [])
     ),
+    extras_require={
+        "testing": [
+            "pytest",
+            "requests",
+        ],
+        "apps": [
+            "dataclasses-json",
+            "tokenizers",
+        ],
+        "onnx": [
+            "onnx",
+        ],
+        "torch": [
+            # Note that this pulls in `torch>=2.3.0`, which then pulls in all of CUDA.
+            "iree-turbine",
+        ],
+        "server": [
+            "fastapi",
+            "uvicorn",
+        ],
+    },
     cmdclass={
         "build": CustomBuild,
         "build_ext": NoopBuildExtension,