Merge pull request #25 from learning-at-home/rename

tesseract -> hivemind

CONTRIBUTING.md

@@ -1,7 +1,7 @@
 
 
 ### Rules for collaborating:
-Tesseract is still in the early stage of development, we expect only a handful of collaborators with individual roles.
+Hivemind is still in the early stage of development, we expect only a handful of collaborators with individual roles.
 
 1. Before you write any code, please contact us to avoid duplicate work:
    * Report bugs and propose new features via issues. We don't have templates at this point;
@@ -20,5 +20,5 @@ Tesseract is still in the early stage of development, we expect only a handful o
 
 ### Tips & tricks
 * You can find a wealth of pytorch debugging tricks at [their contributing page](https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md).
-* Tesseract is optimized for development in pycharm CE 2019.3 or newer.
+* Hivemind is optimized for development in pycharm CE 2019.3 or newer.
   * When working on tests, please mark "tests" as sources root.

README.md

@@ -1,4 +1,4 @@
-## Tesseract
+## Hivemind
 
 [![Build status](https://circleci.com/gh/learning-at-home/tesseract.svg?style=shield)](https://circleci.com/gh/learning-at-home/tesseract)
 [![Documentation Status](https://readthedocs.org/projects/learning-at-home/badge/?version=latest)](https://learning-at-home.readthedocs.io/en/latest/?badge=latest)

docs/conf.py

@@ -23,8 +23,8 @@ from recommonmark.parser import CommonMarkParser
 
 # -- Project information -----------------------------------------------------
 sys.path.insert(0, '..')
-src_path = '../tesseract'
-project = 'tesseract'
+src_path = '../hivemind'
+project = 'hivemind'
 copyright = '2020, Learning@home & contributors'
 author = 'Learning@home & contributors'
 
@@ -119,7 +119,7 @@ html_static_path = ['_static']
 # -- Options for HTMLHelp output ---------------------------------------------
 
 # Output file base name for HTML help builder.
-htmlhelp_basename = 'tesseractdoc'
+htmlhelp_basename = 'hiveminddoc'
 
 
 # -- Options for LaTeX output ------------------------------------------------
@@ -146,7 +146,7 @@ latex_elements = {
 # (source start file, target name, title,
 #  author, documentclass [howto, manual, or own class]).
 latex_documents = [
-    (master_doc, 'tesseract.tex', 'tesseract Documentation',
+    (master_doc, 'hivemind.tex', 'hivemind Documentation',
      'Learning@home \\& contributors', 'manual'),
 ]
 
@@ -156,7 +156,7 @@ latex_documents = [
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
 man_pages = [
-    (master_doc, 'tesseract', 'tesseract Documentation',
+    (master_doc, 'hivemind', 'hivemind Documentation',
      [author], 1)
 ]
 
@@ -167,8 +167,8 @@ man_pages = [
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
-    (master_doc, 'tesseract', 'tesseract Documentation',
-     author, 'tesseract', 'One line description of project.',
+    (master_doc, 'hivemind', 'hivemind Documentation',
+     author, 'hivemind', 'One line description of project.',
      'Miscellaneous'),
 ]
 
@@ -240,7 +240,7 @@ def linkcode_resolve(domain, info):
     if domain != 'py' or not info['module']:
         return None
     try:
-        filename = 'tesseract/%s#L%d-L%d' % find_source()
+        filename = 'hivemind/%s#L%d-L%d' % find_source()
     except Exception:
         filename = info['module'].replace('.', '/') + '.py'
     return "https://github.com/learning-at-home/tesseract/blob/%s/%s" % (branch, filename)

docs/index.rst

@@ -1,7 +1,7 @@
-``learning@home::tesseract``
+``learning@home::hivemind``
 ====================================
 
-Tesseract lets you train huge neural networks on computers provided by volunteers. Powered by pytorch.
+Hivemind lets you train huge neural networks on computers provided by volunteers. Powered by pytorch.
 
 .. image:: _static/bug.gif
 
@@ -28,4 +28,4 @@ Indices and tables
 * :ref:`modindex`
 * :ref:`search`
 
-.. _GitHub: https://github.com/learning-at-home/tesseract
+.. _GitHub: https://github.com/learning-at-home/hivemind

docs/modules/client.rst

@@ -1,9 +1,9 @@
-``tesseract.client``
+``hidemind.client``
 ====================
 
-.. automodule:: tesseract.client
+.. automodule:: hivemind.client
 
-.. currentmodule:: tesseract.client
+.. currentmodule:: hivemind.client
 
 .. raw:: html
 

docs/modules/server.rst

@@ -1,17 +1,17 @@
-``tesseract.server & runtime``
+``hivemind.server & runtime``
 ========================================
 
-.. automodule:: tesseract.server
+.. automodule:: hivemind.server
 
-.. currentmodule:: tesseract.server
+.. currentmodule:: hivemind.server
 
-.. autoclass:: TesseractServer
+.. autoclass:: Server
    :members:
    :member-order: bysource
 
-.. currentmodule:: tesseract.runtime
+.. currentmodule:: hivemind.runtime
 
-.. autoclass:: TesseractRuntime
+.. autoclass:: Runtime
     :members:
     :member-order: bysource
 

docs/user/quickstart.md

@@ -1,6 +1,6 @@
 # Quick start [nothing here yet]
 
-This will eventually become a tutorial on how to host a tesseract node or connect to an existing node.
+This will eventually become a tutorial on how to host a hivemind node or connect to an existing node.
 
 ![img](https://media.giphy.com/media/3oz8xtBx06mcZWoNJm/giphy.gif)
 
@@ -16,30 +16,30 @@ This will eventually become a tutorial on how to host a tesseract node or connec
 
 ## How do I run it?
 
-Currently, there is no way to do it easily. There are some tests (you can check [`./tests/benchmark_throughput.py`](https://github.com/learning-at-home/tesseract/blob/master/tests/benchmark_throughput.py)
+Currently, there is no way to do it easily. There are some tests (you can check [`./tests/benchmark_throughput.py`](https://github.com/learning-at-home/hivemind/blob/master/tests/benchmark_throughput.py)
  or look into CI logs) and we want to expand them. If you want to
 do something complex with it, please contact us by opening an issue (less preferred: [telegram](https://t.me/justheuristic)).
 
-## `tesseract` quick tour
+## `hivemind` quick tour
 
 **Trainer process:**
 
-- **`RemoteExpert`**(`tesseract/client/remote_expert.py`) behaves like a pytorch
+- **`RemoteExpert`**(`hivemind/client/remote_expert.py`) behaves like a pytorch
   module with autograd support but actually sends request to a remote runtime.
-- **`RemoteMixtureOfExperts`**(`tesseract/client/remote_moe.py`) finds best experts
+- **`RemoteMixtureOfExperts`**(`hivemind/client/remote_moe.py`) finds best experts
   for a given input and either returns them as `RemoteExpert` or applies them
   right away.
 
 **Runtime process:**
 
-- **`TesseractRuntime`** (`tesseract/runtime/__init__.py`) aggregates batches
+- **`Runtime`** (`hivemind/runtime/__init__.py`) aggregates batches
   and performs inference/training of experts according to their priority.
-- **`TesseractServer`** (`tesseract/server/__init__.py`) wraps runtime and
+- **`Server`** (`hivemind/server/__init__.py`) wraps runtime and
   periodically uploads experts into `DHTNode`.
 
 **DHT:**
 
-- **`DHTNode`**(`tesseract/dht/__init__.py`) is a node of
+- **`DHTNode`**(`hivemind/dht/__init__.py`) is a node of
   Kademlia-based DHT that stores metadata used by trainer and runtime.
 
 ## Limitations
@@ -52,6 +52,6 @@ do something complex with it, please contact us by opening an issue (less prefer
 
 **Runtime**:
 * You can achieve 4x less network load by passing quantized uint8 activations across experts.
-    Implement your own quantization or wait for tesseract v0.8.
+    Implement your own quantization or wait for hivemind v0.8.
 * Currently runtime can form batches that exceed maximal batch_size by task_size - 1. 
     We will fix that in the nearest patch.

tesseract/client/expert.py → hivemind/client/expert.py

@@ -16,8 +16,8 @@ class RemoteExpert(nn.Module):
     Sending wrong input shapes can cause RemoteExpert to freeze indefinitely due to error in runtime.
 
     :param uid: unique expert identifier
-    :param host: hostname where TesseractServer operates
-    :param port: port to which TesseractServer listens
+    :param host: hostname where Server operates
+    :param port: port to which Server listens
     """
 
     def __init__(self, uid, host='127.0.0.1', port=8080):

tesseract/client/moe.py → hivemind/client/moe.py

@@ -23,7 +23,7 @@ class RemoteMixtureOfExperts(nn.Module):
      the missing experts
 
     :param in_features: common input size for experts and gating function
-    :param grid_size: tesseract dimensions that form expert uid (see below)
+    :param grid_size: hivemind dimensions that form expert uid (see below)
     :param uid_prefix: common prefix for all expert uids
      expert uid follows the pattern {uid_prefix}.{0...grid_size[0]}.{0...grid_size[1]}...{0...grid_size[-1]}
     :param dht: DHTNode where the experts reside

tesseract/dht/__init__.py → hivemind/dht/__init__.py

@@ -6,8 +6,8 @@ from typing import Tuple, List, Optional
 
 from kademlia.network import Server
 
-from tesseract.client import RemoteExpert
-from tesseract.utils import run_forever, SharedFuture, PickleSerializer
+from hivemind.client import RemoteExpert
+from hivemind.utils import run_forever, SharedFuture, PickleSerializer
 
 
 class DHTNode(mp.Process):
@@ -44,7 +44,7 @@ class DHTNode(mp.Process):
         """
         self.start()
         if await_ready and not self.ready.wait(timeout=timeout):
-            raise TimeoutError("TesseractServer didn't notify .ready in {timeout} seconds")
+            raise TimeoutError("Server didn't notify .ready in {timeout} seconds")
 
     def shutdown(self) -> None:
         """ Shuts down the dht process """

tesseract/runtime/__init__.py → hivemind/runtime/__init__.py

@@ -12,15 +12,15 @@ from .expert_backend import ExpertBackend
 from .task_pool import TaskPool, TaskPoolBase
 
 
-class TesseractRuntime(threading.Thread):
+class Runtime(threading.Thread):
     """
     A group of processes that processes incoming requests for multiple experts on a shared device.
-    TesseractRuntime is usually created and managed by TesseractServer, humans need not apply.
+    Runtime is usually created and managed by Server, humans need not apply.
 
     For debugging, you can start runtime manually with .start() or .run()
 
     >>> expert_backends = {'expert_name': ExpertBackend(**kwargs)}
-    >>> runtime = TesseractRuntime(expert_backends)
+    >>> runtime = Runtime(expert_backends)
     >>> runtime.start()  # start runtime in background thread. To start in current thread, use runtime.run()
     >>> runtime.ready.wait()  # await for runtime to load all experts on device and create request pools
     >>> future = runtime.expert_backends['expert_name'].forward_pool.submit_task(*expert_inputs)

tesseract/runtime/expert_backend.py → hivemind/runtime/expert_backend.py

@@ -9,7 +9,7 @@ from ..utils import nested_flatten, nested_pack, nested_compare, BatchTensorProt
 
 class ExpertBackend(nn.Module):
     """
-    ExpertBackend is a wrapper around torch module that allows it to run tasks asynchronously with TesseractRuntime
+    ExpertBackend is a wrapper around torch module that allows it to run tasks asynchronously with Runtime
     By default, ExpertBackend handles three types of requests:
 
      - forward - receive inputs and compute outputs. Concurrent requests will be batched for better GPU utilization.
@@ -57,7 +57,7 @@ class ExpertBackend(nn.Module):
 
     def forward(self, *inputs: torch.Tensor) -> Tuple[torch.Tensor, ...]:
         """
-        Apply forward pass to an aggregated batch of requests. Used by TesseractRuntime, do not call this manually;
+        Apply forward pass to an aggregated batch of requests. Used by Runtime, do not call this manually;
         To submit a request for asynchronous processing, please use ``ExpertBackend.forward_pool.submit_task``.
 
         Subclassing:
@@ -78,7 +78,7 @@ class ExpertBackend(nn.Module):
 
     def backward(self, *inputs: torch.Tensor) -> Tuple[torch.Tensor, ...]:
         """
-        Apply backward pass to an aggregated batch of requests. Used by TesseractRuntime, do not call this manually
+        Apply backward pass to an aggregated batch of requests. Used by Runtime, do not call this manually
         To submit a request for asynchronous processing, please use ``ExpertBackend.backward_pool.submit_task``.
 
         Subclassing:
@@ -86,7 +86,7 @@ class ExpertBackend(nn.Module):
 
            It should return gradients w.r.t. inputs that follow ``nested_flatten(self.forward_schema)``;
 
-           TesseractRuntime doesn't guarantee that backward will be performed in the same order and for the same data
+           Runtime doesn't guarantee that backward will be performed in the same order and for the same data
            as forward, so we recommend stateless backward pass that re-runs expert forward pass inside backward.
 
            .. todo state, randomness, etc
@@ -127,6 +127,6 @@ class ExpertBackend(nn.Module):
                     keyword_names=tuple(self.kwargs_schema.keys()))
 
     def get_pools(self) -> Sequence[TaskPool]:
-        """ return all pools that should be processed by ``TesseractRuntime`` """
+        """ return all pools that should be processed by ``Runtime`` """
         return self.forward_pool, self.backward_pool
 

tesseract/runtime/task_pool.py → hivemind/runtime/task_pool.py

@@ -20,7 +20,7 @@ Task = namedtuple("Task", ("future", "args"))
 
 
 class TaskPoolBase(mp.Process):
-    """ A pool that accepts tasks and forms batches for parallel processing, interacts with TesseractRuntime """
+    """ A pool that accepts tasks and forms batches for parallel processing, interacts with Runtime """
 
     def __init__(self, process_func: callable):
         super().__init__()
@@ -55,10 +55,10 @@ class TaskPoolBase(mp.Process):
 
 class TaskPool(TaskPoolBase):
     """
-    Request aggregator that accepts processing requests, groups them into batches, waits for TesseractRuntime
+    Request aggregator that accepts processing requests, groups them into batches, waits for Runtime
     to process these batches and dispatches results back to request sources. Operates as a background process.
 
-    :param process_func: function to be applied to every formed batch; called by TesseractRuntime
+    :param process_func: function to be applied to every formed batch; called by Runtime
         Note that process_func should accept only \*args Tensors and return a flat tuple of Tensors
     :param max_batch_size: process at most this many inputs in a batch (task contains have one or several inputs)
     :param min_batch_size: process at least this many inputs in a batch, otherwise wait for more
@@ -81,7 +81,7 @@ class TaskPool(TaskPoolBase):
         self.tasks = mp.Queue(maxsize=pool_size or 0)
         self.undispatched_task_timestamps = mp.SimpleQueue()
 
-        # interaction with TesseractRuntime
+        # interaction with Runtime
         self.batch_receiver, self.batch_sender = mp.Pipe(duplex=False)  # send/recv arrays that contain batch inputs
         self.batch_received = mp.Event()  # runtime can notify pool that it can send next batch
         self.outputs_receiver, self.outputs_sender = mp.Pipe(duplex=False)  # send/recv arrays that contain outputs

tesseract/server/__init__.py → hivemind/server/__init__.py

@@ -7,18 +7,18 @@ from typing import Dict
 from .connection_handler import handle_connection
 from .dht_handler import DHTHandlerThread
 from ..dht import DHTNode
-from ..runtime import TesseractRuntime, ExpertBackend
+from ..runtime import Runtime, ExpertBackend
 
 
-class TesseractServer(threading.Thread):
+class Server(threading.Thread):
     """
-    TesseractServer allows you to host "experts" - pytorch sub-networks used by Decentralized Mixture of Experts.
-    After creation, a server should be started: see TesseractServer.run or TesseractServer.run_in_background.
+    Server allows you to host "experts" - pytorch sub-networks used by Decentralized Mixture of Experts.
+    After creation, a server should be started: see Server.run or Server.run_in_background.
 
     A working server does 3 things:
-     - processes incoming forward/backward requests via TesseractRuntime (created by the server)
+     - processes incoming forward/backward requests via Runtime (created by the server)
      - publishes updates to expert status every :update_period: seconds
-     - follows orders from TesseractController - if it exists
+     - follows orders from HivemindController - if it exists
 
     :type dht: DHTNode or None. Server with dht=None will NOT be visible from DHT,
      but it will still support accessing experts directly with RemoteExpert(uid=UID, host=IPADDR, port=PORT).
@@ -40,15 +40,15 @@ class TesseractServer(threading.Thread):
         self.dht, self.experts, self.update_period = dht, expert_backends, update_period
         self.addr, self.port = addr, port
         self.conn_handlers = self._create_connection_handlers(conn_handler_processes)
-        self.runtime = TesseractRuntime(self.experts, **kwargs)
+        self.runtime = Runtime(self.experts, **kwargs)
 
         if start:
             self.run_in_background(await_ready=True)
 
     def run(self):
         """
-        Starts TesseractServer in the current thread. Initializes dht if necessary, starts connection handlers,
-        runs TesseractRuntime (self.runtime) to process incoming requests.
+        Starts Server in the current thread. Initializes dht if necessary, starts connection handlers,
+        runs Runtime (self.runtime) to process incoming requests.
         """
         if self.dht:
             if not self.dht.is_alive():
@@ -71,12 +71,12 @@ class TesseractServer(threading.Thread):
 
     def run_in_background(self, await_ready=True, timeout=None):
         """
-        Starts TesseractServer in a background thread. if await_ready, this method will wait until background server
+        Starts Server in a background thread. if await_ready, this method will wait until background server
         is ready to process incoming requests or for :timeout: seconds max.
         """
         self.start()
         if await_ready and not self.ready.wait(timeout=timeout):
-            raise TimeoutError("TesseractServer didn't notify .ready in {timeout} seconds")
+            raise TimeoutError("Server didn't notify .ready in {timeout} seconds")
 
     @property
     def ready(self) -> mp.synchronize.Event:
@@ -104,7 +104,7 @@ class TesseractServer(threading.Thread):
 
     def shutdown(self):
         """
-        Gracefully terminate a tesseract server, process-safe.
+        Gracefully terminate a hivemind server, process-safe.
         Please note that terminating server otherwise (e.g. by killing processes) may result in zombie processes.
         If you did already cause a zombie outbreak, your only option is to kill them with -9 (SIGKILL).
         """

tesseract/server/connection_handler.py → hivemind/server/connection_handler.py

@@ -1,8 +1,8 @@
 from socket import socket
 from typing import Tuple, Dict
 
-from tesseract.runtime.expert_backend import ExpertBackend
-from tesseract.utils import PytorchSerializer, Connection
+from hivemind.runtime.expert_backend import ExpertBackend
+from hivemind.utils import PytorchSerializer, Connection
 
 
 def handle_connection(connection_tuple: Tuple[socket, str], experts: Dict[str, ExpertBackend]):

tesseract/utils/threading.py → hivemind/utils/threading.py

@@ -3,7 +3,7 @@ from concurrent.futures import Future, ThreadPoolExecutor, as_completed, Timeout
 import time
 from typing import Optional, List
 
-GLOBAL_EXECUTOR = ThreadPoolExecutor(max_workers=os.environ.get("TESSERACT_THREADS", float('inf')))
+GLOBAL_EXECUTOR = ThreadPoolExecutor(max_workers=os.environ.get("HIVEMIND_THREADS", float('inf')))
 
 
 def run_in_background(func: callable, *args, **kwargs) -> Future:

setup.py

@@ -10,18 +10,18 @@ with open('requirements.txt') as requirements_file:
     install_requires = [str(requirement) for requirement in parse_requirements(requirements_file)]
 
 # loading version from setup.py
-with codecs.open(os.path.join(here, 'tesseract/__init__.py'), encoding='utf-8') as init_file:
+with codecs.open(os.path.join(here, 'hivemind/__init__.py'), encoding='utf-8') as init_file:
     version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", init_file.read(), re.M)
     version_string = version_match.group(1)
 
 setup(
-    name='tesseract',
+    name='hivemind',
     version=version_string,
     description='',
     long_description='',
     author='Learning@home authors',
     author_email='mryabinin@hse.ru',
-    packages=['tesseract'],
+    packages=['hivemind'],
     license='MIT',
     install_requires=install_requires,
     classifiers=[

tests/benchmark_throughput.py

@@ -7,14 +7,14 @@ import time
 
 import torch
 from test_utils import layers, print_device_info
-from tesseract import find_open_port
+from hivemind import find_open_port
 
-import tesseract
+import hivemind
 
 
 def client_process(can_start, benchmarking_failed, port, num_experts, batch_size, hid_dim, num_batches, backprop=True):
     can_start.wait()
-    experts = [tesseract.RemoteExpert(f"expert{i}", port=port) for i in range(num_experts)]
+    experts = [hivemind.RemoteExpert(f"expert{i}", port=port) for i in range(num_experts)]
 
     try:
         dummy_batch = torch.randn(batch_size, hid_dim)
@@ -62,14 +62,14 @@ def benchmark_throughput(num_experts=16, num_handlers=None, num_clients=128, num
         experts = {}
         for i in range(num_experts):
             expert = torch.jit.script(layers.name_to_block[expert_cls](hid_dim))
-            experts[f'expert{i}'] = tesseract.ExpertBackend(name=f'expert{i}',
-                                                            expert=expert, opt=torch.optim.Adam(expert.parameters()),
-                                                            args_schema=(tesseract.BatchTensorProto(hid_dim),),
-                                                            outputs_schema=tesseract.BatchTensorProto(hid_dim),
-                                                            max_batch_size=max_batch_size,
-                                                            )
+            experts[f'expert{i}'] = hivemind.ExpertBackend(name=f'expert{i}',
+                                                           expert=expert, opt=torch.optim.Adam(expert.parameters()),
+                                                           args_schema=(hivemind.BatchTensorProto(hid_dim),),
+                                                           outputs_schema=hivemind.BatchTensorProto(hid_dim),
+                                                           max_batch_size=max_batch_size,
+                                                           )
         timestamps['created_experts'] = time.perf_counter()
-        server = tesseract.TesseractServer(None, experts, port=port, conn_handler_processes=num_handlers, device=device)
+        server = hivemind.Server(None, experts, port=port, conn_handler_processes=num_handlers, device=device)
         server.start()
         server.ready.wait()
         timestamps['server_ready'] = time.perf_counter()

tests/test_moe.py

@@ -1,5 +1,5 @@
 import torch
-import tesseract
+import hivemind
 from test_utils.run_server import background_server
 
 
@@ -20,8 +20,8 @@ def test_remote_module_call():
 
     with background_server(num_experts=num_experts, device='cpu',
                            no_optimizer=True, no_dht=True) as (localhost, server_port, dht_port):
-        experts = [tesseract.RemoteExpert(uid=f'expert.{i}', port=server_port) for i in range(num_experts)]
-        moe_output, = tesseract.client.moe._RemoteMoECall.apply(
+        experts = [hivemind.RemoteExpert(uid=f'expert.{i}', port=server_port) for i in range(num_experts)]
+        moe_output, = hivemind.client.moe._RemoteMoECall.apply(
             logits, experts[:len(logits)], k_min, timeout_after_k_min, backward_k_min, timeout_total, backward_timeout,
             [(None,), {}], xx)
 
@@ -45,15 +45,15 @@ def test_remote_module_call():
 
 def test_compute_expert_scores():
     try:
-        dht = tesseract.DHTNode(port=tesseract.find_open_port(), start=True)
-        moe = tesseract.client.moe.RemoteMixtureOfExperts(
+        dht = hivemind.DHTNode(port=hivemind.find_open_port(), start=True)
+        moe = hivemind.client.moe.RemoteMixtureOfExperts(
             dht=dht, in_features=1024, grid_size=[40], k_best=4, k_min=1, timeout_after_k_min=1,
             uid_prefix='expert')
         gx, gy = torch.randn(4, 5, requires_grad=True), torch.torch.randn(4, 3, requires_grad=True)
         ii = [[4, 0, 2], [3, 1, 1, 1, 3], [0], [3, 2]]
         jj = [[2, 2, 1], [0, 1, 2, 0, 1], [0], [1, 2]]
         batch_experts = [
-            [tesseract.RemoteExpert(uid=f'expert.{ii[b][e]}.{jj[b][e]}') for e in range(len(ii[b]))]
+            [hivemind.RemoteExpert(uid=f'expert.{ii[b][e]}.{jj[b][e]}') for e in range(len(ii[b]))]
             for b in range(len(ii))
         ]  # note: these experts do not exists on server, we use them only to test moe compute_expert_scores
         logits = moe.compute_expert_scores([gx, gy], batch_experts)

tests/test_utils/run_server.py

@@ -4,14 +4,14 @@ import multiprocessing as mp
 import argparse
 
 import torch
-import tesseract
+import hivemind
 from .layers import name_to_block
 
 
 def make_dummy_server(host='0.0.0.0', port=None, num_experts=1, expert_cls='ffn', hidden_dim=1024, num_handlers=None,
                       expert_prefix='expert', expert_offset=0, max_batch_size=16384, device=None, no_optimizer=False,
                       no_dht=False, initial_peers=(), dht_port=None, root_port=None, verbose=True, start=False,
-                      UID_DELIMETER=tesseract.DHTNode.UID_DELIMETER, **kwargs) -> tesseract.TesseractServer:
+                      UID_DELIMETER=hivemind.DHTNode.UID_DELIMETER, **kwargs) -> hivemind.Server:
     """ A context manager that creates server in a background thread, awaits .ready on entry and shutdowns on exit """
     if verbose and len(kwargs) != 0:
         print("Ignored kwargs:", kwargs)
@@ -24,8 +24,8 @@ def make_dummy_server(host='0.0.0.0', port=None, num_experts=1, expert_cls='ffn'
     if not no_dht:
         if not len(initial_peers):
             print("No initial peers provided. Starting additional dht as an initial peer.")
-            dht_root = tesseract.DHTNode(
-                *initial_peers, port=root_port or tesseract.find_open_port(), start=True)
+            dht_root = hivemind.DHTNode(
+                *initial_peers, port=root_port or hivemind.find_open_port(), start=True)
             print(f"Initializing DHT with port {dht_root.port}")
             initial_peers = (('localhost', dht_root.port), )
         else:
@@ -33,8 +33,8 @@ def make_dummy_server(host='0.0.0.0', port=None, num_experts=1, expert_cls='ffn'
             if root_port is not None:
                 print(f"Warning: root_port={root_port} will not be used since we already have peers.")
 
-        dht = tesseract.DHTNode(
-            *initial_peers, port=dht_port or tesseract.find_open_port(), start=True)
+        dht = hivemind.DHTNode(
+            *initial_peers, port=dht_port or hivemind.find_open_port(), start=True)
         if verbose:
             print(f"Running dht node on port {dht.port}")
 
@@ -44,14 +44,14 @@ def make_dummy_server(host='0.0.0.0', port=None, num_experts=1, expert_cls='ffn'
         expert = torch.jit.script(name_to_block[expert_cls](hidden_dim))
         opt = torch.optim.SGD(expert.parameters(), 0.0) if no_optimizer else torch.optim.Adam(expert.parameters())
         expert_uid = f'{expert_prefix}{UID_DELIMETER}{i + expert_offset}'
-        experts[expert_uid] = tesseract.ExpertBackend(name=expert_uid, expert=expert, opt=opt,
-                                                      args_schema=(tesseract.BatchTensorProto(hidden_dim),),
-                                                      outputs_schema=tesseract.BatchTensorProto(hidden_dim),
-                                                      max_batch_size=max_batch_size,
-                                                      )
+        experts[expert_uid] = hivemind.ExpertBackend(name=expert_uid, expert=expert, opt=opt,
+                                                     args_schema=(hivemind.BatchTensorProto(hidden_dim),),
+                                                     outputs_schema=hivemind.BatchTensorProto(hidden_dim),
+                                                     max_batch_size=max_batch_size,
+                                                     )
     # actually start server
-    server = tesseract.TesseractServer(
-        dht, experts, addr=host, port=port or tesseract.find_open_port(),
+    server = hivemind.Server(
+        dht, experts, addr=host, port=port or hivemind.find_open_port(),
         conn_handler_processes=num_handlers, device=device)
 
     if start: