Run only one ML worker

app.py

@@ -1,12 +1,33 @@
 import gradio as gr
-import asyncio
+
 from giskard.ml_worker.ml_worker import MLWorker
 from pydantic import AnyHttpUrl
 from giskard.settings import settings
 from urllib.parse import urlparse
 
-def run_ml_worker(url, api_key, hf_token):
+import asyncio
+import threading
+
+previous_url = ""
+ml_worker = None
+
+
+def run_ml_worker(ml_worker: MLWorker):
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    loop.run_until_complete(ml_worker.start())
+    loop.close()
+
+
+def start_ml_worker(url, api_key, hf_token):
+    global ml_worker, previous_url
     # Always run an external ML worker
+    if ml_worker is not None:
+        print(f"Stopping ML worker for {previous_url}")
+        ml_worker.stop()
+        print("ML worker stopped")
+
     parsed_url = urlparse(url)
     backend_url = AnyHttpUrl(
         url=f"{parsed_url.scheme if parsed_url.scheme else 'http'}://{parsed_url.hostname}"
@@ -17,8 +38,10 @@ def run_ml_worker(url, api_key, hf_token):
     )
     print(f"Starting ML worker for {backend_url}")
     ml_worker = MLWorker(False, backend_url, api_key, hf_token)
-    asyncio.run(ml_worker.start())
-    return "ML worker finished its work"
+    previous_url = backend_url
+    thread = threading.Thread(target=run_ml_worker, args=(ml_worker,))
+    thread.start()
+    return f"ML worker running for {backend_url}"
 
-iface = gr.Interface(fn=run_ml_worker, inputs=["text", "text", "text"], outputs="text")
+iface = gr.Interface(fn=start_ml_worker, inputs=["text", "text", "text"], outputs="text")
 iface.launch()