First step of having multiprocess in worker

Author: Hartorn

.gitignore

@@ -1,3 +1,4 @@
+enron_with_categories
 .history
 .mypy_cache
 docker-stack.yml

giskard/client/giskard_client.py

@@ -75,6 +75,8 @@ def __call__(self, r):
 class GiskardClient:
     def __init__(self, url: str, key: str, hf_token: str = None):
         self.host_url = url
+        self.key = key
+        self.hf_token = hf_token
         base_url = urljoin(url, "/api/v2/")
         self._session = sessions.BaseUrlSession(base_url=base_url)
         self._session.mount(base_url, ErrorHandlingAdapter())

giskard/commands/cli_worker.py

@@ -1,10 +1,11 @@
+from typing import Optional
+
 import asyncio
 import functools
 import logging
 import os
 import platform
 import sys
-from typing import Optional
 
 import click
 import lockfile
@@ -13,19 +14,19 @@
 from lockfile.pidlockfile import PIDLockFile, read_pid_from_pidfile, remove_existing_pidfile
 from pydantic import AnyHttpUrl
 
-from giskard.cli_utils import common_options
 from giskard.cli_utils import (
+    common_options,
     create_pid_file_path,
+    follow_file,
+    get_log_path,
     remove_stale_pid_file,
     run_daemon,
-    get_log_path,
     tail,
-    follow_file,
     validate_url,
 )
 from giskard.path_utils import run_dir
 from giskard.settings import settings
-from giskard.utils.analytics_collector import anonymize, analytics
+from giskard.utils.analytics_collector import analytics, anonymize
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +86,13 @@ def wrapper(*args, **kwargs):
     envvar="GSK_HF_TOKEN",
     help="Access token for Giskard hosted in a private Hugging Face Spaces",
 )
-def start_command(url: AnyHttpUrl, is_server, api_key, is_daemon, hf_token):
+@click.option(
+    "--parallelism",
+    "nb_workers",
+    default=None,
+    help="Number of processes to use for parallelism (None for number of cpu)",
+)
+def start_command(url: AnyHttpUrl, is_server, api_key, is_daemon, hf_token, nb_workers):
     """\b
     Start ML Worker.
 
@@ -102,7 +109,7 @@ def start_command(url: AnyHttpUrl, is_server, api_key, is_daemon, hf_token):
     )
     api_key = initialize_api_key(api_key, is_server)
     hf_token = initialize_hf_token(hf_token, is_server)
-    _start_command(is_server, url, api_key, is_daemon, hf_token)
+    _start_command(is_server, url, api_key, is_daemon, hf_token, nb_workers)
 
 
 def initialize_api_key(api_key, is_server):
@@ -126,7 +133,7 @@ def initialize_hf_token(hf_token, is_server):
     return hf_token
 
 
-def _start_command(is_server, url: AnyHttpUrl, api_key, is_daemon, hf_token=None):
+def _start_command(is_server, url: AnyHttpUrl, api_key, is_daemon, hf_token=None, nb_workers=None):
     from giskard.ml_worker.ml_worker import MLWorker
 
     start_msg = "Starting ML Worker"
@@ -154,7 +161,7 @@ def _start_command(is_server, url: AnyHttpUrl, api_key, is_daemon, hf_token=None
             run_daemon(is_server, url, api_key, hf_token)
         else:
             ml_worker = MLWorker(is_server, url, api_key, hf_token)
-            asyncio.get_event_loop().run_until_complete(ml_worker.start())
+            asyncio.get_event_loop().run_until_complete(ml_worker.start(nb_workers))
     except KeyboardInterrupt:
         logger.info("Exiting")
         if ml_worker:

giskard/ml_worker/ml_worker.py

@@ -1,16 +1,20 @@
+from typing import Optional
+
 import logging
 import random
 import secrets
-import stomp
 import time
+
+import stomp
 from pydantic import AnyHttpUrl
-from websocket._exceptions import WebSocketException, WebSocketBadStatusException
+from websocket._exceptions import WebSocketBadStatusException, WebSocketException
 
 import giskard
+from giskard.cli_utils import validate_url
 from giskard.client.giskard_client import GiskardClient
 from giskard.ml_worker.testing.registry.registry import load_plugins
 from giskard.settings import settings
-from giskard.cli_utils import validate_url
+from giskard.utils import shutdown_pool, start_pool
 
 logger = logging.getLogger(__name__)
 
@@ -140,9 +144,9 @@ def _connect_websocket_client(self, is_server=False):
     def is_remote_worker(self):
         return self.ml_worker_id is not INTERNAL_WORKER_ID
 
-    async def start(self):
+    async def start(self, nb_workers: Optional[int] = None):
         load_plugins()
-
+        start_pool(nb_workers)
         if self.ws_conn:
             self.ws_stopping = False
             self.connect_websocket_client()
@@ -162,3 +166,4 @@ def stop(self):
         if self.ws_conn:
             self.ws_stopping = True
             self.ws_conn.disconnect()
+        shutdown_pool()

giskard/ml_worker/websocket/__init__.py

@@ -72,7 +72,8 @@ class Catalog(WorkerReply):
 
 
 class DataRow(BaseModel):
-    columns: Dict[str, str]
+    columns: Dict[str, str] = Field(..., repr=False)
+
 
 
 class DataFrame(BaseModel):