Introduce a task server which runs on local processes

colmena/task_server/local.py

@@ -0,0 +1,46 @@
+"""Use Python's :class:`~concurrent.futures.Executor` to run workers on a local system"""
+from typing import Optional, Union, Callable, Collection
+from concurrent.futures import Future, Executor, ThreadPoolExecutor, ProcessPoolExecutor
+
+from colmena.models.results import Result
+
+from .base import FutureBasedTaskServer, convert_to_colmena_method
+from colmena.queue.base import ColmenaQueues
+from colmena.models.methods import ColmenaMethod
+
+
+class LocalTaskServer(FutureBasedTaskServer):
+    """Use Python's native concurrent libraries to execute tasks"""
+
+    def __init__(self,
+                 queues: ColmenaQueues,
+                 methods: Collection[Union[Callable, ColmenaMethod]],
+                 threads: bool = True,
+                 num_workers: Optional[int] = None):
+        """
+        Args:
+            methods: Methods to be served
+            queues: Queues used to commmunicate with thinker
+            threads: Use threads instead of workers
+            num_workers: Number of workers to deploy.
+        """
+        self._methods = dict(
+            (m.name, m) for m in map(convert_to_colmena_method, methods)
+        )
+        self._executor: Optional[Executor] = None
+        self.num_workers = num_workers
+        self.threads = threads
+        super().__init__(queues=queues, method_names=list(self._methods.keys()))
+
+    def _submit(self, task: Result, topic: str) -> Future | None:
+        return self._executor.submit(self._methods[task.method], task)
+
+    def _setup(self):
+        self._executor = (
+            ThreadPoolExecutor(self.num_workers)
+            if self.threads else
+            ProcessPoolExecutor(self.num_workers)
+        )
+
+    def _cleanup(self):
+        self._executor.shutdown()

colmena/task_server/tests/test_local.py

@@ -0,0 +1,32 @@
+"""Test local task server"""
+
+from pytest import fixture
+
+from colmena.queue.python import PipeQueues
+from colmena.task_server.local import LocalTaskServer
+
+
+def f(x):
+    return x
+
+
+@fixture
+def queues():
+    return PipeQueues()
+
+
+@fixture
+def server(queues):
+    server = LocalTaskServer(queues, methods=[f], num_workers=2)
+    server.start()
+    yield server
+    queues.send_kill_signal()
+    server.join()
+
+
+def test_local_service(server, queues):
+    queues.send_inputs(1, method='f')
+
+    result = queues.get_result(timeout=4)
+    assert result.success, result.failure_info.traceback
+    assert result.value == 1