client: Fix GatewayClientSync loop handling issues

Revamp the design of GatewayClientSync to ensure it works in any
environment. If an event loop is already running in the current thread,
it should create a background thread with its own event loop to run the
synchronous code. If no event loop is running, it should create and use
its own in the current thread, taking care of cleaning up afterwards.
Add docstrings and arguments to sync client methods to match their async
counterparts and help with development.

Signed-off-by: Phoevos Kalemkeris <[email protected]>

Author: phoevos

client/cogstack_model_gateway_client/client.py

@@ -1,5 +1,7 @@
 import asyncio
 import json
+import threading
+import time
 from collections.abc import Iterable
 from functools import wraps
 
@@ -128,7 +130,7 @@ async def process(
         wait_for_completion: bool = True,
         return_result: bool = True,
     ):
-        """Generate annotations for the the provided text."""
+        """Generate annotations for the provided text."""
         return await self.submit_task(
             model_name=model_name,
             task="process",
@@ -290,9 +292,10 @@ async def deploy_model(
 class GatewayClientSync:
     def __init__(self, *args, **kwargs):
         self._client = GatewayClient(*args, **kwargs)
-        self._loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(self._loop)
-        self._loop.run_until_complete(self._client.__aenter__())
+        self._loop = None
+        self._own_loop = False
+        self._thread = None
+        self._setup_event_loop()
 
     @property
     def base_url(self):
@@ -322,40 +325,204 @@ def timeout(self):
     def timeout(self, value: float):
         self._client.timeout = value
 
+    def _setup_event_loop(self):
+        """Set up event loop handling for sync operations."""
+        try:
+            existing_loop = asyncio.get_running_loop()
+            self._loop = existing_loop
+            self._own_loop = False
+            self._setup_background_loop()
+        except RuntimeError:
+            self._loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self._loop)
+            self._own_loop = True
+            self._loop.run_until_complete(self._client.__aenter__())
+
+    def _setup_background_loop(self):
+        """Set up a background thread with persistent event loop and client."""
+
+        def run_background_loop():
+            self._background_loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self._background_loop)
+
+            self._background_loop.run_until_complete(self._client.__aenter__())
+            self._background_loop.run_forever()
+
+        self._thread = threading.Thread(target=run_background_loop, daemon=True)
+        self._thread.start()
+
+        start_time, timeout = time.time(), 5.0
+        while not hasattr(self, "_background_loop") and time.time() - start_time < timeout:
+            time.sleep(0.01)
+
+        if not hasattr(self, "_background_loop"):
+            raise RuntimeError("Failed to start background event loop")
+
+    def _run_async(self, coro):
+        """Run an async coroutine, handling different event loop scenarios."""
+        if self._own_loop:
+            return self._loop.run_until_complete(coro)
+        else:
+            future = asyncio.run_coroutine_threadsafe(coro, self._background_loop)
+            return future.result(timeout=60)
+
     def __del__(self):
         try:
-            if hasattr(self, "_client") and self._client and self._client._client is not None:
-                self._loop.run_until_complete(self._client.__aexit__(None, None, None))
-                self._loop.close()
+            if not hasattr(self, "_client") or not self._client:
+                return
+
+            if self._own_loop and self._client._client is not None:
+                try:
+                    self._loop.run_until_complete(self._client.__aexit__(None, None, None))
+                    self._loop.close()
+                except Exception:
+                    pass
+            elif hasattr(self, "_background_loop") and hasattr(self, "_thread"):
+                if self._background_loop and not self._background_loop.is_closed():
+                    try:
+                        future = asyncio.run_coroutine_threadsafe(
+                            self._client.__aexit__(None, None, None), self._background_loop
+                        )
+                        future.result(timeout=1.0)
+                    except Exception:
+                        # If cleanup fails, we can't do much about it in __del__
+                        pass
+                    try:
+                        self._background_loop.call_soon_threadsafe(self._background_loop.stop)
+                    except Exception:
+                        pass
+                if self._thread and self._thread.is_alive():
+                    self._thread.join(timeout=1.0)
         except Exception:
             pass
 
-    def submit_task(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.submit_task(*args, **kwargs))
+    def submit_task(
+        self,
+        model_name: str = None,
+        task: str = None,
+        data=None,
+        json=None,
+        files=None,
+        params=None,
+        headers=None,
+        wait_for_completion: bool = False,
+        return_result: bool = True,
+    ):
+        """Submit a task to the Gateway and return the task info."""
+        return self._run_async(
+            self._client.submit_task(
+                model_name=model_name,
+                task=task,
+                data=data,
+                json=json,
+                files=files,
+                params=params,
+                headers=headers,
+                wait_for_completion=wait_for_completion,
+                return_result=return_result,
+            )
+        )
+
+    def process(
+        self,
+        text: str,
+        model_name: str = None,
+        wait_for_completion: bool = True,
+        return_result: bool = True,
+    ):
+        """Generate annotations for the provided text."""
+        return self._run_async(
+            self._client.process(
+                text=text,
+                model_name=model_name,
+                wait_for_completion=wait_for_completion,
+                return_result=return_result,
+            )
+        )
 
-    def process(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.process(*args, **kwargs))
+    def process_bulk(
+        self,
+        texts: list[str],
+        model_name: str = None,
+        wait_for_completion: bool = True,
+        return_result: bool = True,
+    ):
+        """Generate annotations for a list of texts."""
+        return self._run_async(
+            self._client.process_bulk(
+                texts=texts,
+                model_name=model_name,
+                wait_for_completion=wait_for_completion,
+                return_result=return_result,
+            )
+        )
 
-    def process_bulk(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.process_bulk(*args, **kwargs))
+    def redact(
+        self,
+        text: str,
+        concepts_to_keep: Iterable[str] = None,
+        warn_on_no_redaction: bool = None,
+        mask: str = None,
+        hash: bool = None,
+        model_name: str = None,
+        wait_for_completion: bool = True,
+        return_result: bool = True,
+    ):
+        """Redact sensitive information from the provided text."""
+        return self._run_async(
+            self._client.redact(
+                text=text,
+                concepts_to_keep=concepts_to_keep,
+                warn_on_no_redaction=warn_on_no_redaction,
+                mask=mask,
+                hash=hash,
+                model_name=model_name,
+                wait_for_completion=wait_for_completion,
+                return_result=return_result,
+            )
+        )
 
-    def redact(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.redact(*args, **kwargs))
+    def get_task(self, task_uuid: str, detail: bool = True):
+        """Get a Gateway task details by its UUID."""
+        return self._run_async(self._client.get_task(task_uuid=task_uuid, detail=detail))
 
-    def get_task(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.get_task(*args, **kwargs))
+    def get_task_result(self, task_uuid: str, parse: bool = True):
+        """Get the result of a Gateway task by its UUID.
 
-    def get_task_result(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.get_task_result(*args, **kwargs))
+        If parse is True, try to infer and parse the result as JSON, JSONL, or text.
+        Otherwise, return raw bytes.
+        """
+        return self._run_async(self._client.get_task_result(task_uuid=task_uuid, parse=parse))
 
-    def wait_for_task(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.wait_for_task(*args, **kwargs))
+    def wait_for_task(self, task_uuid: str, detail: bool = True, raise_on_error: bool = False):
+        """Poll Gateway until the task reaches a final state."""
+        return self._run_async(
+            self._client.wait_for_task(
+                task_uuid=task_uuid, detail=detail, raise_on_error=raise_on_error
+            )
+        )
 
-    def get_models(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.get_models(*args, **kwargs))
+    def get_models(self, verbose: bool = False):
+        """Get the list of available models from the Gateway."""
+        return self._run_async(self._client.get_models(verbose=verbose))
 
-    def get_model(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.get_model(*args, **kwargs))
+    def get_model(self, model_name: str = None):
+        """Get details of a specific model."""
+        return self._run_async(self._client.get_model(model_name=model_name))
 
-    def deploy_model(self, *args, **kwargs):
-        return self._loop.run_until_complete(self._client.deploy_model(*args, **kwargs))
+    def deploy_model(
+        self,
+        model_name: str = None,
+        tracking_id: str = None,
+        model_uri: str = None,
+        ttl: int = None,
+    ):
+        """Deploy a CogStack Model Serve model through the Gateway."""
+        return self._run_async(
+            self._client.deploy_model(
+                model_name=model_name,
+                tracking_id=tracking_id,
+                model_uri=model_uri,
+                ttl=ttl,
+            )
+        )