Browse Source

Merge pull request #4372 from JTHesse/main

build: Adding ability to install requirements from frontmatter for tools and functions
Timothy Jaeryang Baek 8 months ago
parent
commit
9682806476
2 changed files with 13 additions and 1 deletions
  1. 13 0
      backend/apps/webui/utils.py
  2. 0 1
      backend/start.sh

+ 13 - 0
backend/apps/webui/utils.py

@@ -1,6 +1,8 @@
 from importlib import util
 import os
 import re
+import sys
+import subprocess
 
 from config import TOOLS_DIR, FUNCTIONS_DIR
 
@@ -52,6 +54,7 @@ def load_toolkit_module_by_id(toolkit_id):
     frontmatter = extract_frontmatter(toolkit_path)
 
     try:
+        install_frontmatter_requirements(frontmatter.get("requirements", ""))
         spec.loader.exec_module(module)
         print(f"Loaded module: {module.__name__}")
         if hasattr(module, "Tools"):
@@ -73,6 +76,7 @@ def load_function_module_by_id(function_id):
     frontmatter = extract_frontmatter(function_path)
 
     try:
+        install_frontmatter_requirements(frontmatter.get("requirements", ""))
         spec.loader.exec_module(module)
         print(f"Loaded module: {module.__name__}")
         if hasattr(module, "Pipe"):
@@ -88,3 +92,12 @@ def load_function_module_by_id(function_id):
         # Move the file to the error folder
         os.rename(function_path, f"{function_path}.error")
         raise e
+
+def install_frontmatter_requirements(requirements):
+    if requirements:
+        req_list = [req.strip() for req in requirements.split(',')]
+        for req in req_list:
+            print(f"Installing requirement: {req}")
+            subprocess.check_call([sys.executable, "-m", "pip", "install", req])
+    else:
+        print("No requirements found in frontmatter.")

+ 0 - 1
backend/start.sh

@@ -30,7 +30,6 @@ if [[ "${USE_CUDA_DOCKER,,}" == "true" ]]; then
   export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/python3.11/site-packages/torch/lib:/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib"
 fi
 
-
 # Check if SPACE_ID is set, if so, configure for space
 if [ -n "$SPACE_ID" ]; then
   echo "Configuring for HuggingFace Space deployment"