Use websockets for json communication #4490

AndrewJGaut · 2023-06-14T21:15:11Z

This PR addresses #4431. It modifies JSON messaging between the REST server / bundle manager and workers.

An in-depth document describing these changes and their motivation is given here.

Summary

Previously, messages were sent to workers in the HTTP response to worker checkins. However, this was inefficient; a server thread would try to send the message through an AF Unix socket corresponding to the worker using TCP, and so if the sending and receiving threads were not synchronized, the message would not be sent. The need for synchronization made server to worker communication highly inefficient. To make this slightly more efficient, a websocket server was added and would send a ping to the worker meant to receive the message to cause it to checkin immediately so that the receiving server thread would start listening as soon as possible on the AF Unix socket to receive the message. However, desynchronization and corresponding inefficiency persisted.

After this PR, JSON messages will be sent to workers immediately through the websocket server. This means that:
(1) Workers will no longer receive any messages through HTTP responses to the checkin.
(2) No synchronization is required between server threads for JSON messaging.
(3) The websocket server no longer uses the ping functionality.

We also add in authentication of workers and server to the websocket server. (The need for this is described here.)

Speedup

When looking at the time tests, the only test that uses worker communication is the run test. For that test, before the websockets change, it would take over 11 seconds (11.86 seconds here, 11.44 seconds here). After the websockets change, we get 5.57 seconds. That's a speedup of over 2x! (To see the time in those links, click on Run tests using Docker runtime and then search cl info -f name and it will be the number right above the line that comes up.)

Future Work / TODO

I might do these in a future PR:

Only release lock once worker is done with whatever it's doing
More fine-grained locking for the worker (lock per bundle uuid)
Dynamic resizing of number of websockets for worker depending on 'demand'

codalab/bin/ws_server.py

epicfaace · 2023-06-28T18:34:54Z

"working quite well now" 👀 👀

…ake it secure. Not yet tested b/c on plane and internet isn't good enough to run CodaLab (and download images and such)

…lane)

…codes so that we have proper tests

This reverts commit 4eb3d7b.

This reverts commit cbd1505.

codalab/bin/ws_server.py

…th exponential backoff, it's still not very aggressive at all

…have any sockets open with ws-server

codalab/model/worker_model.py

codalab/bin/ws_server.py

epicfaace · 2023-08-09T02:05:32Z

codalab/bin/ws_server.py

+            return
+
+    logger.warning(f"All websockets for worker {worker_id} are currently busy.")
+    await server_websocket.close(1013, f"All websockets for worker {worker_id} are currently busy.")


what is 1013?

Hmm... weird. I see this listed in several other places (e.g. here) as Try again later, but it's not in the RFC.

Perhaps I should change this to a different status code, but none of the ones listed in the RFC are fitting... what do you think?

codalab/bin/ws_server.py

codalab/lib/codalab_manager.py

epicfaace · 2023-08-09T02:07:33Z

codalab/lib/download_manager.py

            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)


rename to get_json_message?

I do want to keep the with_sock portion (I changed it to with_unix_socket in this newest version) since it helps make clear that it's using AF_Unix sockets. However, I could change recv back to get -- I thought recv was clearer since that's typically the API for getting data from sockets.

Let me know. I can definitely change this.

codalab/lib/download_manager.py

codalab/model/bundle_model.py

codalab/worker/main.py

codalab/bin/ws_server.py

codalab/model/bundle_model.py

codalab/model/worker_model.py

…or if the client tries to connect to an invalid path.

epicfaace

Looks good to me. This new architecture is all documented in your doc, correct? Can you link to it? (maybe we also add a link to it as a comment in the code in ws_server.py so we never lose it)

epicfaace · 2023-08-31T01:45:54Z

codalab/worker/worker.py

@@ -836,7 +852,7 @@ def netcat_fn():
                        break
                    total_data.append(data)
                s.close()
-                reply(None, {}, b''.join(total_data))
+                reply(None, {}, io.BytesIO(b''.join(total_data)))


Do we need this change?

epicfaace · 2023-08-31T01:46:25Z

codalab/worker/worker.py

@@ -136,6 +140,11 @@ def __init__(

        self.ws_server = ws_server

+        assert (
+            num_coroutines > 0 and type(num_coroutines) is int


Do we need this?

epicfaace · 2023-09-06T22:25:11Z

@AndrewJGaut is this ready to merge?

AndrewJGaut · 2023-09-10T17:59:32Z

Looks good to me. This new architecture is all documented in your doc, correct? Can you link to it? (maybe we also add a link to it as a comment in the code in ws_server.py so we never lose it)

I figured it was enough to link the doc in the PR description, since it's stated pretty clearly. Do you think we should in the code as well? I can do that.

AndrewJGaut · 2023-09-10T17:59:47Z

@AndrewJGaut is this ready to merge?

Let me test it out in dev one more time to be sure.

…se-websockets-for-json-communication

AndrewJGaut · 2023-09-10T19:24:39Z

@AndrewJGaut is this ready to merge?

Let me test it out in dev one more time to be sure.

Confirmed it's good to merge.

This reverts commit 7c460c1.

AndrewJGaut and others added 5 commits June 14, 2023 12:19

first commit

3221c90

Fix bugs

9c0ec8d

Fix more bugsz

a139498

slight change

b226f40

minor bug fixes

93f3e62

epicfaace reviewed Jun 14, 2023

View reviewed changes

codalab/bin/ws_server.py Outdated Show resolved Hide resolved

epicfaace reviewed Jun 14, 2023

View reviewed changes

codalab/bin/ws_server.py Outdated Show resolved Hide resolved

AndrewJGaut added 2 commits June 28, 2023 11:08

It's working quite well now

0db9264

merge conflict

1a65356

AndrewJGaut and others added 18 commits June 28, 2023 15:01

Update to solve that last bug. I think it should work now

ae3e3be

bug fix that was causing messages to be send to tall worker threads

85774a4

Fix unittests issue and fix formatting

81064a7

Minor change to increase robustness of message sending

5326420

Some more formatting changes

34550f0

Add in worker auth and use wss rather than ws for websocket URLs to m…

57aad3d

…ake it secure. Not yet tested b/c on plane and internet isn't good enough to run CodaLab (and download images and such)

Added in server auth with secret. Aslo still need to test (still on p…

42127e5

…lane)

Fixed issues and got auth working. Now, I'll work on returning error …

aa37b0e

…codes so that we have proper tests

Add in tests for authentication functionality (for worker and server)

3e71d67

Slight cleanup to data sending code

7a0bd70

Adding in ssl certification

cbd1505

Add in SSL stuff for worker; still testing on dev

4eb3d7b

Revert "Add in SSL stuff for worker; still testing on dev"

aa74981

This reverts commit 4eb3d7b.

Revert "Adding in ssl certification"

123fc7e

This reverts commit cbd1505.

Fixed formatting

6efc0cf

Very minor formatting change to ignore one line for MyPy

35bdab0

Another minor formatting change

9f5417c

Merge branch 'master' into use-websockets-for-json-communication

f0ec0f1

AndrewJGaut marked this pull request as ready for review July 16, 2023 16:37

AndrewJGaut commented Jul 16, 2023

View reviewed changes

codalab/bin/ws_server.py Outdated Show resolved Hide resolved

AndrewJGaut marked this pull request as ready for review July 17, 2023 01:49

AndrewJGaut added 3 commits July 17, 2023 00:11

Make websocket locks more robust and improve error messaging

e34fb72

More permissible retries in case of other errors (e.g. like 1013). Wi…

42819a6

…th exponential backoff, it's still not very aggressive at all

minor change to have a different error message if worker doesn't yet …

8e868a4

…have any sockets open with ws-server

AndrewJGaut mentioned this pull request Jul 22, 2023

Fix/4431 migrate communication to use websockets only #4451

Closed

epicfaace reviewed Aug 2, 2023

View reviewed changes

codalab/model/worker_model.py Outdated Show resolved Hide resolved

epicfaace reviewed Aug 9, 2023

View reviewed changes

codalab/model/worker_model.py Outdated Show resolved Hide resolved

epicfaace reviewed Aug 9, 2023

View reviewed changes

codalab/bin/ws_server.py Outdated Show resolved Hide resolved

codalab/bin/ws_server.py Outdated Show resolved Hide resolved

codalab/model/bundle_model.py Show resolved Hide resolved

epicfaace reviewed Aug 9, 2023

View reviewed changes

codalab/model/worker_model.py Outdated Show resolved Hide resolved

AndrewJGaut added 5 commits August 15, 2023 22:10

Rename send_json and send_json_message_with_sock

f36ed3a

Rearrange worker_model to minimize diff

4ed30a2

Final changes

d115a96

Fix formatting

5ac4853

Minor changes to get auth working again and to robustly return an err…

8acedff

…or if the client tries to connect to an invalid path.

AndrewJGaut requested a review from epicfaace August 23, 2023 05:18

epicfaace approved these changes Aug 31, 2023

View reviewed changes

AndrewJGaut added 2 commits September 10, 2023 11:09

Merge branch 'master' of github.com:codalab/codalab-worksheets into u…

3d25b77

…se-websockets-for-json-communication

Merge in master and make some minor changes

bbe9402

AndrewJGaut merged commit 7c460c1 into master Sep 10, 2023
63 checks passed

AndrewJGaut deleted the use-websockets-for-json-communication branch September 10, 2023 19:24

epicfaace mentioned this pull request Sep 26, 2023

Bump version to 1.7.2 #4534

Closed

AndrewJGaut added a commit that referenced this pull request Oct 7, 2023

Revert "Use websockets for json communication (#4490)"

c471585

This reverts commit 7c460c1.

AndrewJGaut mentioned this pull request Oct 7, 2023

Revert "Use websockets for json communication" #4548

Merged

epicfaace pushed a commit that referenced this pull request Oct 10, 2023

Revert "Use websockets for json communication (#4490)" (#4548)

bf1dff7

This reverts commit 7c460c1.

AndrewJGaut restored the use-websockets-for-json-communication branch October 11, 2023 03:00

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use websockets for json communication #4490

Use websockets for json communication #4490

AndrewJGaut commented Jun 14, 2023 •

edited

Loading

epicfaace commented Jun 28, 2023

epicfaace Aug 9, 2023

AndrewJGaut Aug 16, 2023 •

edited

Loading

epicfaace Aug 9, 2023

AndrewJGaut Aug 16, 2023

epicfaace left a comment

epicfaace Aug 31, 2023

epicfaace Aug 31, 2023

epicfaace commented Sep 6, 2023

AndrewJGaut commented Sep 10, 2023

AndrewJGaut commented Sep 10, 2023

AndrewJGaut commented Sep 10, 2023

Use websockets for json communication #4490

Use websockets for json communication #4490

Conversation

AndrewJGaut commented Jun 14, 2023 • edited Loading

Summary

Speedup

Future Work / TODO

epicfaace commented Jun 28, 2023

epicfaace Aug 9, 2023

Choose a reason for hiding this comment

AndrewJGaut Aug 16, 2023 • edited Loading

Choose a reason for hiding this comment

epicfaace Aug 9, 2023

Choose a reason for hiding this comment

AndrewJGaut Aug 16, 2023

Choose a reason for hiding this comment

epicfaace left a comment

Choose a reason for hiding this comment

epicfaace Aug 31, 2023

Choose a reason for hiding this comment

epicfaace Aug 31, 2023

Choose a reason for hiding this comment

epicfaace commented Sep 6, 2023

AndrewJGaut commented Sep 10, 2023

AndrewJGaut commented Sep 10, 2023

AndrewJGaut commented Sep 10, 2023

AndrewJGaut commented Jun 14, 2023 •

edited

Loading

AndrewJGaut Aug 16, 2023 •

edited

Loading