Use websockets for json communication

codalab/bin/ws_server.py

@@ -1,75 +1,204 @@
+# This is the real ws-server, basically.
 # Main entry point for CodaLab cl-ws-server.
 import argparse
+import json
 import asyncio
+from collections import defaultdict
 import logging
 import re
-from typing import Any, Dict
+from typing import Any, Dict, List
 import websockets
+from dataclasses import dataclass
+import threading
+import time
+
+ACK=b'a'
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.WARNING)
+logger.setLevel(logging.INFO)
 logging.basicConfig(format='%(asctime)s %(message)s %(pathname)s %(lineno)d')
 
-worker_to_ws: Dict[str, Any] = {}
 
+"""
+TODO!!! WE NEED TO ADD A SECRET OR SOME SORT OF AUTH FOR THE WS SERVER ENDPOITNS
+Otherwise, people could create a custom local worker build and wreak havoc on the ws server...
+Note that this was already an issue, though; people could've just hit the checkpoint endpoint...
+"""
 
-async def rest_server_handler(websocket):
-    """Handles routes of the form: /main. This route is called by the rest-server
-    whenever a worker needs to be pinged (to ask it to check in). The body of the
-    message is the worker id to ping. This function sends a message to the worker
-    with that worker id through an appropriate websocket.
+@dataclass
+class WS:
     """
-    # Got a message from the rest server.
-    worker_id = await websocket.recv()
-    logger.warning(f"Got a message from the rest server, to ping worker: {worker_id}.")
-
-    try:
-        worker_ws = worker_to_ws[worker_id]
-        await worker_ws.send(worker_id)
-    except KeyError:
-        logger.error(f"Websocket not found for worker: {worker_id}")
+    Stores websocket object and whether or not the websocket is available.
+    TODO: give this a better, less confusing name.
+    """
+    _ws: Any = None
+    _is_available: bool = True
+    _lock: threading.Lock = threading.Lock()
+    _timeout: float = 86400
+    _last_use: float = None
+
+    @property
+    def ws(self):
+        return self._ws
+
+    @property
+    def lock(self):
+        return self._lock
+
+    @property
+    def is_available(self):
+        return self._is_available
+
+    @is_available.setter
+    def is_available(self, value):
+        self._is_available = value
+
+    @property
+    def timeout(self):
+        return self._timeout
+
+    @property
+    def last_use(self):
+        return self._last_use
+
+    @last_use.setter
+    def last_use(self, value):
+        self._last_use = value
+
+worker_to_ws: Dict[str, Dict[str, WS]] = defaultdict(dict)  # Maps worker to list of its websockets (since each worker has a pool of connections)
+server_worker_to_ws:  Dict[str, Dict[str, WS]] = defaultdict(dict)   # Map the rest-server websocket connection to the corresponding worker socket connection.
+
+async def connection_handler(websocket, worker_id):
+    """
+    Handles routes of the form: /server/connect/worker_id. This route is called by the rest-server
+    (or bundle-manager) when they want a connection to a worker.
+    Returns the id of the socket to connect to, which will be used in later requests.
+    """
+    logger.debug(f"Got a message from the rest server, to connect to worker: {worker_id}.")
+    socket_id = None
+    #logger.error([ws.is_available for _,ws in worker_to_ws[worker_id].items()])
+    for s_id, ws in worker_to_ws[worker_id].items():
+        logger.debug("about to lock")
+        with ws.lock:
+            logger.debug("locked")
+            if ws.is_available or time.time() - ws.last_use >= ws.timeout:
+                logger.debug("available")
+                ws.last_use = time.time()
+                socket_id = s_id
+                worker_to_ws[worker_id][socket_id].is_available = False
+                logger.debug("breaking")
+                break
+
+    logger.info(f"For worker {worker_id}, sending server socket ID {socket_id}")
+    if not socket_id:
+        logger.error(f"No socket ids available for worker {worker_id}")
+    #import pdb; pdb.set_trace()
+    await websocket.send(json.dumps({'socket_id': socket_id}))
+    logger.debug("Sent.")
+
+
+async def disconnection_handler(websocket, worker_id, socket_id):
+    """
+    Handles routes of the form: /server/connect/worker_id/socket_id. This route is called by the rest-server
+    (or bundle-manager) when they want a connection to a worker.
+    Returns the id of the socket to connect to, which will be used in later requests.
+    """
+    with worker_to_ws[worker_id][socket_id].lock:
+        logger.info(f"For worker {worker_id}, disconnecting socket ID {socket_id}")
+        if worker_to_ws[worker_id][socket_id].is_available:
+            # For now, just log the error.
+            logging.error("Available socket set for disconnection")
+        worker_to_ws[worker_id][socket_id].is_available = True
+
+
+async def exchange(from_ws, to_ws, worker_id, socket_id):    
+    # Send and receive all data until connection closes.
+    # (We may be streaming a file, in which case we need to receive and then send
+    # lots of chunks)
+    while True:
+        # Don't use async for so we can avoid two couroutines waiting on same socket.
+        try:
+            data = await from_ws.recv()
+            await to_ws.send(data)
+        except websockets.exceptions.ConnectionClosed:
+            break
+        # ahhh... this won't work. When the message is sent, it's actually buffered at the client, so it being sent doesn't mean it was received...
+        # Shoot... I wonder if tehre's a way around htis...
+        # In fact, seee here: https://stackoverflow.com/questions/46549892/does-websocket-send-guarantee-consumption
+        # we need to receive an ACK from the to_ws lol and then send it to the from_ws. yuck lol
+        # But that's not too hard; we can do that pretty easily.
+        # this will create lot sof extra traffic, but that's OK
+
+        # Now, we're getting the "recv() called by two coroutines for same websocket"
+        # It makes sense: the thread sending to ws=websocket(worker_id, socket_id) is calling data = await ws.recv()
+        # and the recv caller is calling async for data in ws: and so both are calling recv().
+        # That's an issue that's not solvable, I don't think... shoot.
+
+        # The answer in this case might be to keep two versions of websockets per worker.
+        # So, one for the worker and one for the server (for each worker_id, socket_id combination).
+        # Why do this? So that we can wait and send properly... It'd be very annoying, though, for sure
+        # I don't like it. Is there any better way to do this?
+        # I don't think so... I think this works better, unfortunately. It's kind of gross, but that's alright.
+        # We'll need a separate send and recv handler now for server and worker... Kind of annoying.
+        # Might be able to get by it with some clever instantiation... oh well
+
+        # No, I think we can avoid that.
+
+async def send_handler(websocket, worker_id, socket_id):
+    """Handles routes of the form: /send/{worker_id}/{socket_id}. This route is called by
+    the rest-server or bundle-manager when either wants to send a message/stream to the worker.
+    """
+    with worker_to_ws[worker_id][socket_id].lock:
+        worker_to_ws[worker_id][socket_id].ws.last_use = time.time()
+    await exchange(websocket, worker_to_ws[worker_id][socket_id].ws, worker_id, socket_id)
 
+async def recv_handler(websocket, worker_id, socket_id):
+    """Handles routes of the form: /recv/{worker_id}/{socket_id}. This route is called by
+    the rest-server or bundle-manager when either wants to receive a message/stream from the worker.
+    """
+    with worker_to_ws[worker_id][socket_id].lock:
+        worker_to_ws[worker_id][socket_id].ws.last_use = time.time()
+    await exchange(worker_to_ws[worker_id][socket_id].ws, websocket, worker_id, socket_id)
 
-async def worker_handler(websocket, worker_id):
-    """Handles routes of the form: /worker/{id}. This route is called when
+async def worker_handler(websocket, worker_id, socket_id):
+    """Handles routes of the form: /worker/{worker_id}/{socket_id}. This route is called when
     a worker first connects to the ws-server, creating a connection that can
     be used to ask the worker to check-in later.
     """
     # runs on worker connect
-    worker_to_ws[worker_id] = websocket
-    logger.warning(f"Connected to worker {worker_id}!")
+    worker_to_ws[worker_id][socket_id] = WS(websocket)
+    logger.warning(f"Worker {worker_id} connected; has {len(worker_to_ws[worker_id])} connections")
 
+    # keep connection alive.
     while True:
         try:
-            await asyncio.wait_for(websocket.recv(), timeout=60)
+            await asyncio.sleep(60)
         except asyncio.futures.TimeoutError:
             pass
         except websockets.exceptions.ConnectionClosed:
             logger.error(f"Socket connection closed with worker {worker_id}.")
             break
 
-
 ROUTES = (
-    (r'^.*/main$', rest_server_handler),
-    (r'^.*/worker/(.+)$', worker_handler),
+    (r'^.*/send/(.+)/(.+)$', send_handler),
+    (r'^.*/recv/(.+)/(.+)$', recv_handler),
+    (r'^.*/server/connect/(.+)$', connection_handler),
+    (r'^.*/server/disconnect/(.+)/(.+)$', disconnection_handler),
+    (r'^.*/worker/(.+)/(.+)$', worker_handler),
 )
-
-
 async def ws_handler(websocket, *args):
     """Handler for websocket connections. Routes websockets to the appropriate
     route handler defined in ROUTES."""
-    logger.warning(f"websocket handler, path: {websocket.path}.")
+    logger.debug(f"websocket handler, path: {websocket.path}.")
     for (pattern, handler) in ROUTES:
         match = re.match(pattern, websocket.path)
         if match:
             return await handler(websocket, *match.groups())
-    assert False
-
 
 async def async_main():
     """Main function that runs the websocket server."""
     parser = argparse.ArgumentParser()
-    parser.add_argument('--port', help='Port to run the server on.', type=int, required=True)
+    parser.add_argument('--port', help='Port to run the server on.', type=int, required=False, default=2901)
     args = parser.parse_args()
     logging.debug(f"Running ws-server on 0.0.0.0:{args.port}")
     async with websockets.serve(ws_handler, "0.0.0.0", args.port):
@@ -84,4 +213,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()

codalab/lib/download_manager.py

@@ -137,7 +137,7 @@ def _get_target_info_within_bundle(self, target, depth):
                 read_args = {'type': 'get_target_info', 'depth': depth}
                 self._send_read_message(worker, response_socket_id, target, read_args)
                 with closing(self._worker_model.start_listening(response_socket_id)) as sock:
-                    result = self._worker_model.get_json_message(sock, 60)
+                    result = self._worker_model.recv_json_message_with_sock(sock, 60)
                 if result is None:  # dead workers are a fact of life now
                     logging.info('Unable to reach worker, bundle state {}'.format(bundle_state))
                     raise NotFoundError(
@@ -365,8 +365,8 @@ def _send_read_message(self, worker, response_socket_id, target, read_args):
             'path': target.subpath,
             'read_args': read_args,
         }
-        if not self._worker_model.send_json_message(
-            worker['socket_id'], worker['worker_id'], message, 60
+        if not self._worker_model.connect_and_send_json(
+            message, worker['worker_id']
         ):  # dead workers are a fact of life now
             logging.info('Unable to reach worker')
 
@@ -378,21 +378,21 @@ def _send_netcat_message(self, worker, response_socket_id, uuid, port, message):
             'port': port,
             'message': message,
         }
-        if not self._worker_model.send_json_message(
-            worker['socket_id'], worker['worker_id'], message, 60
+        if not self._worker_model.connect_and_send_json(
+            message, worker['worker_id']
         ):  # dead workers are a fact of life now
             logging.info('Unable to reach worker')
 
     def _get_read_response_stream(self, response_socket_id):
         with closing(self._worker_model.start_listening(response_socket_id)) as sock:
-            header_message = self._worker_model.get_json_message(sock, 60)
+            header_message = self._worker_model.recv_json_message_with_sock(sock, 60)
             precondition(header_message is not None, 'Unable to reach worker')
             if 'error_code' in header_message:
                 raise http_error_to_exception(
                     header_message['error_code'], header_message['error_message']
                 )
 
-            fileobj = self._worker_model.get_stream(sock, 60)
+            fileobj = self._worker_model.recv_stream(sock, 60)
             precondition(fileobj is not None, 'Unable to reach worker')
             return fileobj
 
@@ -416,4 +416,4 @@ def __getattr__(self, attr):
 
     def close(self):
         self._fileobj.close()
-        self._worker_model.deallocate_socket(self._socket_id)
+        self._worker_model.deallocate_socket(self._socket_id)