Add support for multiple uvicorn workers via -w

Author: mpagni12

README.md

@@ -24,7 +24,7 @@ Alternatively, you can dynamically fetch the compounds data on startup using a S
 mol-search-sparql-service -s fetch_rhea.rq -e https://sparql.rhea-db.org/sparql -p 8000
 ```
 
-Other available optionally flags include `-d` (`--daemon`) to run the server in the background and write stdout/stderr to `server.log`. Port defaults to `8010` if `-p` is omitted.
+Other available optional flags include `-w` (`--workers`) to deploy multiple Uvicorn worker processes (default 1) and `-d` (`--daemon`) to run the server in the background and write stdout/stderr to `server.log`. Port defaults to `8010` if `-p` is omitted.
 
 With 3 columns:
 

src/mol_search_sparql_service/main.py

@@ -16,6 +16,7 @@ def main():
 
     parser.add_argument('-e', '--endpoint', help='SPARQL endpoint URL (required if using -s)')
     parser.add_argument('-p', '--port', type=int, default=8010, help='Port to run the server on (default: 8010)')
+    parser.add_argument('-w', '--workers', type=int, default=1, help='Number of Uvicorn workers (default: 1)')
     parser.add_argument('-d', '--daemon', action='store_true', help='Daemonize the server process')
     args = parser.parse_args()
 
@@ -37,15 +38,24 @@ def main():
             resp.raise_for_status()
 
             fd, temp_path = tempfile.mkstemp(suffix='.tsv')
+            # Important: Keep the temp file around for workers by not deleting it immediately
             with os.fdopen(fd, 'wb') as f:
                 f.write(resp.content)
-                
+            
+            # Pass to environment so workers can use it
+            os.environ['COMPOUNDS_FILE'] = temp_path
+            os.environ['DELETE_COMPOUNDS_FILE'] = '1'
+            
             engine.load_and_compile(temp_path)
-            os.remove(temp_path)
+            # We will rely on sparql_service or OS to clean this up, or clean it up after uvicorn exits.
         except Exception as e:
             print(f"Failed to fetch data from endpoint: {e}")
             sys.exit(1)
 
+    # Store explicit file path in environment for workers
+    if args.file:
+        os.environ['COMPOUNDS_FILE'] = args.file
+    
     # Daemonize if requested
     if args.daemon:
         # Double fork to detach from terminal completely
@@ -65,12 +75,15 @@ def main():
             os.dup2(f.fileno(), sys.stderr.fileno())
 
     # 2. Start Server
-    print(f"Starting SPARQL endpoint on port {args.port}...")
-    
+    print(f"Starting SPARQL endpoint on port {args.port} with {args.workers} worker(s)...")
     import rdflib.plugins.sparql
     print(f"DEBUG: CUSTOM_EVALS keys: {list(rdflib.plugins.sparql.CUSTOM_EVALS.keys())}")
 
-    uvicorn.run(app, host="0.0.0.0", port=args.port)
+    # Uvicorn requires an import string when using multiple workers
+    if args.workers > 1:
+        uvicorn.run("mol_search_sparql_service.sparql_service:app", host="0.0.0.0", port=args.port, workers=args.workers)
+    else:
+        uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 if __name__ == "__main__":
     main()

src/mol_search_sparql_service/sparql_service.py

@@ -200,5 +200,29 @@ def sparql_assistant() -> str:
     # Functions are already registered via decorators on 'g'
 )
 
+# Startup event to load data in worker processes
+import os
+import atexit
+
+@app.on_event("startup")
+async def startup_event():
+    compounds_file = os.environ.get('COMPOUNDS_FILE')
+    # If the engine has no datasets (wasn't loaded by main.py due to workers>1), load it now
+    if compounds_file and not engine.datasets:
+        print(f"Worker initializing engine from: {compounds_file}")
+        engine.load_and_compile(compounds_file)
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    # Attempt to cleanup temp files if this is the "main" process 
+    # (or just let the OS handle temp files).
+    delete_file = os.environ.get('DELETE_COMPOUNDS_FILE')
+    compounds_file = os.environ.get('COMPOUNDS_FILE')
+    if delete_file == '1' and compounds_file and os.path.exists(compounds_file):
+        try:
+            os.remove(compounds_file)
+        except:
+            pass
+
 # Mount MCP Server (FastMCP's internal app)
 app.mount("/mcp", mcp.sse_app())