Rename TesseractNetwork to DHTNode

docs/user/quickstart.md

@@ -35,11 +35,11 @@ do something complex with it, please contact us by opening an issue (less prefer
 - **`TesseractRuntime`** (`tesseract/runtime/__init__.py`) aggregates batches
   and performs inference/training of experts according to their priority.
 - **`TesseractServer`** (`tesseract/server/__init__.py`) wraps runtime and
-  periodically uploads experts into `TesseractNetwork`.
+  periodically uploads experts into `DHTNode`.
 
 **DHT:**
 
-- **`TesseractNetwork`**(`tesseract/network/__init__.py`) is a node of
+- **`DHTNode`**(`tesseract/dht/__init__.py`) is a node of
   Kademlia-based DHT that stores metadata used by trainer and runtime.
 
 ## Limitations

tesseract/__init__.py

@@ -1,6 +1,6 @@
 from .client import *
-from .network import *
+from .dht import *
 from .server import *
 from .utils import *
 
-__version__ = '0.7.1'
+__version__ = '0.7.1'

tesseract/client/moe.py

@@ -26,8 +26,8 @@ class RemoteMixtureOfExperts(nn.Module):
     :param grid_size: tesseract 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 network: TesseractNetwork where the experts reside
-    :param num_workers: number of threads for parallel network operation
+    :param dht: DHTNode where the experts reside
+    :param num_workers: number of threads for parallel dht operation
     :param k_best: queries this many experts with highest scores
     :param k_min: makes sure at least this many experts returned output
     :param timeout_after_k_min: waits for this many seconds after k_min experts returned results.
@@ -37,11 +37,11 @@ class RemoteMixtureOfExperts(nn.Module):
      allow_broadcasting=True will flatten first d-1 input dimensions, apply RemoteMixtureOfExperts and un-flatten again
      allow_broadcasting=False will raise an error
     """
-    def __init__(self, *, in_features, grid_size: Tuple[int], network, k_best, k_min=1,
+    def __init__(self, *, in_features, grid_size: Tuple[int], dht, k_best, k_min=1,
                  forward_timeout=None, timeout_after_k_min=1.0, backward_k_min=1, backward_timeout=None,
                  uid_prefix='', expert_padding=None, allow_broadcasting=True):
         super().__init__()
-        self.network, self.grid_size = network, grid_size
+        self.dht, self.grid_size = dht, grid_size
         self.uid_prefix, self.expert_padding = uid_prefix, expert_padding
         self.k_best, self.k_min, self.backward_k_min = k_best, k_min, backward_k_min
         self.forward_timeout, self.timeout_after_k_min, self.backward_timeout = forward_timeout, timeout_after_k_min, backward_timeout
@@ -94,7 +94,7 @@ class RemoteMixtureOfExperts(nn.Module):
         :param grid_scores: scores predicted for each dimension in the grid,
         :type grid_scores: a sequence of tensors of shape[batch_size, self.grid_size[i]]
         :param k_best: how many of the top experts participate in the computation
-        :param kwargs: extra keyword parameters passed to self.network.first_k_active
+        :param kwargs: extra keyword parameters passed to self.dht.first_k_active
         :returns: a list of *batch_size* lists that contain chosen experts for one sample each inner list contains \
          RemoteExpert instances for *up to* k_best experts
         """
@@ -104,7 +104,7 @@ class RemoteMixtureOfExperts(nn.Module):
         beam = np.array([[self.uid_prefix]] * batch_size, dtype=object)  # [batch_size, up_to_beam_size]
         scores = np.zeros([batch_size, 1], dtype=np.float64)
 
-        delimeters = np.array(self.network.UID_DELIMETER)[None, None, None]  # pre-compute numpy array for fast concat
+        delimeters = np.array(self.dht.UID_DELIMETER)[None, None, None]  # pre-compute numpy array for fast concat
 
         for dim_index, dim_scores in enumerate(grid_scores):
             dim_scores = check_numpy(dim_scores)
@@ -121,7 +121,7 @@ class RemoteMixtureOfExperts(nn.Module):
             # select k best candidates according to scores but only those that are still active
             new_order = np.argsort(- new_scores, axis=-1)
             top_alive_lookups = [
-                run_in_background(self.network.first_k_active, cands[order], k_best, **kwargs)
+                run_in_background(self.dht.first_k_active, cands[order], k_best, **kwargs)
                 for cands, order in zip(new_candidates, new_order)]
 
             batch_cand_to_score = [
@@ -137,7 +137,7 @@ class RemoteMixtureOfExperts(nn.Module):
             scores = np.array([row + [-float('inf')] * (k_best - len(row))
                                for row in top_alive_scores], dtype='float32')
 
-        unique_experts = self.network.get_experts(list(set(
+        unique_experts = self.dht.get_experts(list(set(
             uid for row in beam for uid in row if uid != self.expert_padding)))
         if self._outputs_schema is None:
             self._outputs_schema = next(iter(unique_experts)).info['outputs_schema']
@@ -160,8 +160,8 @@ class RemoteMixtureOfExperts(nn.Module):
 
         grid_indices = np.zeros([len(flat_experts), len(grid_scores)], dtype=np.int64)
         for i, expert in enumerate(flat_experts):
-            expert_indices = expert.uid[len(self.uid_prefix) + len(self.network.UID_DELIMETER):]
-            expert_indices = list(map(int, expert_indices.split(self.network.UID_DELIMETER)))
+            expert_indices = expert.uid[len(self.uid_prefix) + len(self.dht.UID_DELIMETER):]
+            expert_indices = list(map(int, expert_indices.split(self.dht.UID_DELIMETER)))
             grid_indices[i] = expert_indices
 
         scores_per_dim = [

tesseract/network/__init__.py → tesseract/dht/__init__.py

@@ -10,7 +10,7 @@ from tesseract.client import RemoteExpert
 from tesseract.utils import run_forever, SharedFuture, PickleSerializer
 
 
-class TesseractNetwork(mp.Process):
+class DHTNode(mp.Process):
     UID_DELIMETER = '.'  # splits expert uids over this delimeter
     HEARTBEAT_EXPIRATION = 120  # expert is inactive iff it fails to post timestamp for *this many seconds*
     make_key = "{}::{}".format
@@ -39,7 +39,7 @@ class TesseractNetwork(mp.Process):
 
     def run_in_background(self, await_ready=True, timeout=None):
         """
-        Starts TesseractNetwork in a background process. if await_ready, this method will wait until background network
+        Starts DHTNode in a background process. if await_ready, this method will wait until background dht
         is ready to process incoming requests or for :timeout: seconds max.
         """
         self.start()
@@ -47,11 +47,11 @@ class TesseractNetwork(mp.Process):
             raise TimeoutError("TesseractServer didn't notify .ready in {timeout} seconds")
 
     def shutdown(self) -> None:
-        """ Shuts down the network process """
+        """ Shuts down the dht process """
         if self.is_alive():
             self.kill()
         else:
-            warnings.warn("Network shutdown has no effect: network process is already not alive")
+            warnings.warn("DHT shutdown has no effect: dht process is already not alive")
 
     def get_experts(self, uids: List[str], heartbeat_expiration=HEARTBEAT_EXPIRATION) -> List[Optional[RemoteExpert]]:
         """ Find experts across DHT using their ids; Return a list of [RemoteExpert if found else None]"""

tesseract/server/__init__.py

@@ -5,8 +5,8 @@ from socket import socket, AF_INET, SOCK_STREAM, SO_REUSEADDR, SOL_SOCKET, timeo
 from typing import Dict
 
 from .connection_handler import handle_connection
-from .network_handler import NetworkHandlerThread
-from ..network import TesseractNetwork
+from .dht_handler import DHTHandlerThread
+from ..dht import DHTNode
 from ..runtime import TesseractRuntime, ExpertBackend
 
 
@@ -20,24 +20,24 @@ class TesseractServer(threading.Thread):
      - publishes updates to expert status every :update_period: seconds
      - follows orders from TesseractController - if it exists
 
-    :type network: TesseractNetwork or None. Server with network=None will NOT be visible from DHT,
+    :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).
     :param expert_backends: dict{expert uid (str) : ExpertBackend} for all expert hosted by this server.
-    :param addr: server's network address that determines how it can be accessed. Default is local connections only.
+    :param addr: server's dht address that determines how it can be accessed. Default is local connections only.
     :param port: port to which server listens for requests such as expert forward or backward pass.
     :param conn_handler_processes: maximum number of simultaneous requests. Please note that the default value of 1
         if too small for normal functioning, we recommend 4 handlers per expert backend.
     :param update_period: how often will server attempt to publish its state (i.e. experts) to the DHT;
-        if network is None, this parameter is ignored.
+        if dht is None, this parameter is ignored.
     :param start: if True, the server will immediately start as a background thread and returns control after server
         is ready (see .ready below)
     """
 
-    def __init__(self, network: TesseractNetwork, expert_backends: Dict[str, ExpertBackend], addr='127.0.0.1',
+    def __init__(self, dht: DHTNode, expert_backends: Dict[str, ExpertBackend], addr='127.0.0.1',
                  port: int = 8080, conn_handler_processes: int = 1, update_period: int = 30, start=False,
                  **kwargs):
         super().__init__()
-        self.network, self.experts, self.update_period = network, expert_backends, update_period
+        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)
@@ -47,16 +47,16 @@ class TesseractServer(threading.Thread):
 
     def run(self):
         """
-        Starts TesseractServer in the current thread. Initializes network if necessary, starts connection handlers,
+        Starts TesseractServer in the current thread. Initializes dht if necessary, starts connection handlers,
         runs TesseractRuntime (self.runtime) to process incoming requests.
         """
-        if self.network:
-            if not self.network.is_alive():
-                self.network.run_in_background(await_ready=True)
+        if self.dht:
+            if not self.dht.is_alive():
+                self.dht.run_in_background(await_ready=True)
 
-            network_thread = NetworkHandlerThread(experts=self.experts, network=self.network,
+            dht_handler_thread = DHTHandlerThread(experts=self.experts, dht=self.dht,
                                                   addr=self.addr, port=self.port, update_period=self.update_period)
-            network_thread.start()
+            dht_handler_thread.start()
 
         for process in self.conn_handlers:
             if not process.is_alive():
@@ -66,8 +66,8 @@ class TesseractServer(threading.Thread):
 
         for process in self.conn_handlers:
             process.join()
-        if self.network:
-            network_thread.join()
+        if self.dht:
+            dht_handler_thread.join()
 
     def run_in_background(self, await_ready=True, timeout=None):
         """
@@ -112,8 +112,8 @@ class TesseractServer(threading.Thread):
         for process in self.conn_handlers:
             process.terminate()
 
-        if self.network is not None:
-            self.network.shutdown()
+        if self.dht is not None:
+            self.dht.shutdown()
 
         self.runtime.shutdown()
 

tesseract/server/network_handler.py → tesseract/server/dht_handler.py

@@ -1,20 +1,20 @@
 import threading
 import time
 
-from ..network import TesseractNetwork
+from ..dht import DHTNode
 
 
-class NetworkHandlerThread(threading.Thread):
-    def __init__(self, experts, network: TesseractNetwork,
+class DHTHandlerThread(threading.Thread):
+    def __init__(self, experts, dht: DHTNode,
                  update_period: int = 5, addr: str = '127.0.0.1', port: int = 8080):
-        super(NetworkHandlerThread, self).__init__()
+        super(DHTHandlerThread, self).__init__()
         self.port = port
         self.addr = addr
         self.experts = experts
-        self.network = network
+        self.dht = dht
         self.update_period = update_period
 
     def run(self) -> None:
         while True:
-            self.network.declare_experts(self.experts.keys(), self.addr, self.port)
+            self.dht.declare_experts(self.experts.keys(), self.addr, self.port)
             time.sleep(self.update_period)

tests/test_moe.py

@@ -18,8 +18,8 @@ def test_remote_module_call():
     logits = torch.randn(3, requires_grad=True)
     random_proj = torch.randn_like(xx)
 
-    with background_server(num_experts=num_experts,  device='cpu',
-                           no_optimizer=True, no_network=True) as (localhost, server_port, network_port):
+    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(
             logits, experts[:len(logits)], k_min, timeout_after_k_min, backward_k_min, timeout_total, backward_timeout,
@@ -45,9 +45,9 @@ def test_remote_module_call():
 
 def test_compute_expert_scores():
     try:
-        dht = tesseract.TesseractNetwork(port=tesseract.find_open_port(), start=True)
+        dht = tesseract.DHTNode(port=tesseract.find_open_port(), start=True)
         moe = tesseract.client.moe.RemoteMixtureOfExperts(
-            network=dht, in_features=1024, grid_size=[40], k_best=4, k_min=1, timeout_after_k_min=1,
+            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]]

tests/test_utils/run_server.py

@@ -10,8 +10,8 @@ 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_network=False, initial_peers=(), network_port=None, root_port=None, verbose=True, start=False,
-                      UID_DELIMETER=tesseract.TesseractNetwork.UID_DELIMETER, **kwargs) -> tesseract.TesseractServer:
+                      no_dht=False, initial_peers=(), dht_port=None, root_port=None, verbose=True, start=False,
+                      UID_DELIMETER=tesseract.DHTNode.UID_DELIMETER, **kwargs) -> tesseract.TesseractServer:
     """ 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)
@@ -19,12 +19,12 @@ def make_dummy_server(host='0.0.0.0', port=None, num_experts=1, expert_cls='ffn'
     num_handlers = num_handlers if num_handlers is not None else num_experts * 8
     device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
 
-    # initialize network
-    network = None
-    if not no_network:
+    # initialize dht
+    dht = None
+    if not no_dht:
         if not len(initial_peers):
-            print("No initial peers provided. Starting additional network as an initial peer.")
-            dht_root = tesseract.TesseractNetwork(
+            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)
             print(f"Initializing DHT with port {dht_root.port}")
             initial_peers = (('localhost', dht_root.port), )
@@ -33,10 +33,10 @@ 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.")
 
-        network = tesseract.TesseractNetwork(
-            *initial_peers, port=network_port or tesseract.find_open_port(), start=True)
+        dht = tesseract.DHTNode(
+            *initial_peers, port=dht_port or tesseract.find_open_port(), start=True)
         if verbose:
-            print(f"Running network node on port {network.port}")
+            print(f"Running dht node on port {dht.port}")
 
     # initialize experts
     experts = {}
@@ -51,7 +51,7 @@ def make_dummy_server(host='0.0.0.0', port=None, num_experts=1, expert_cls='ffn'
                                                       )
     # actually start server
     server = tesseract.TesseractServer(
-        network, experts, addr=host, port=port or tesseract.find_open_port(),
+        dht, experts, addr=host, port=port or tesseract.find_open_port(),
         conn_handler_processes=num_handlers, device=device)
 
     if start:
@@ -71,8 +71,8 @@ def background_server(*args, verbose=True, **kwargs):
     def server_runner():
         try:
             server = make_dummy_server(*args, verbose=verbose, start=True, **kwargs)
-            network_port = server.network.port if server.network is not None else None
-            send_addr.send((server.addr, server.port, network_port))
+            dht_port = server.dht.port if server.dht is not None else None
+            send_addr.send((server.addr, server.port, dht_port))
             trigger_shutdown.wait()
         finally:
             if verbose:
@@ -106,9 +106,9 @@ if __name__ == '__main__':
     parser.add_argument('--max_batch_size', type=int, default=16384, required=False)
     parser.add_argument('--device', type=str, default=None, required=False)
     parser.add_argument('--no_optimizer', action='store_true')
-    parser.add_argument('--no_network', action='store_true')
+    parser.add_argument('--no_dht', action='store_true')
     parser.add_argument('--initial_peers', type=str, default="[]", required=False)
-    parser.add_argument('--network_port', type=int, default=None, required=False)
+    parser.add_argument('--dht_port', type=int, default=None, required=False)
     parser.add_argument('--root_port', type=int, default=None, required=False)
 
     parser.add_argument('--increase_file_limit', action='store_true')