Serialize the ASGI protocol directly with Modal functions #905

ekzhang · 2023-09-21T22:42:28Z

First step to making Python pickle-independent web endpoints.

This adds a DataFormat to function inputs and outputs, and fills in Pickle by default. We then use the new ASGI-based serialization format (via protobuf) for all web endpoint outputs.

This also adds support on the client for taking ASGI-encoded inputs, which will be useful in the future, since data_format has been filled in for FunctionInput as well.

ekzhang · 2023-09-21T23:24:52Z

Also after our discussion today and some Googling, I saw there's an open issue for WebTransport support in ASGI:

django/asgiref#280

Not something we want to implement right now, but that gives me more confidence that we can probably just rely on pseudo-ASGI-over-protobuf as a internal format for Modal's web endpoint handling, regardless of language in the future.

erikbern · 2023-09-22T22:23:37Z

modal_proto/api.proto

@@ -90,6 +97,68 @@ enum WebhookAsyncMode {
  WEBHOOK_ASYNC_MODE_AUTO = 4; // redirect to polling endpoint if execution time nears the http timeout
 }

+// A web endpoint connection-related message.


would it make sense to put this in a separate proto file?

Hmm I don't have strong opinions either way. I think it's possible, but also everything is contained in a single message right now.

erikbern · 2023-09-22T22:32:24Z

modal/_container_entrypoint.py

            status=api_pb2.GenericResult.GENERIC_STATUS_SUCCESS,
-            data=self.serialize(data),
+            data=self.serialize_data_format(data, data_format),


this would start writing ASGI serialized format before we added support for it on the other side? i'm not quite following

(nvm, just saw the comment that this will be broken up)

Yep! The rest of this file is just plumbing for serialize_data_format

erikbern · 2023-09-22T22:34:56Z

I was confused about the deployment strategy since this current PR starts to write ASGI already but I just saw you mentioned it in the server-side PR that this will be merged in two parts. That makes more sense!

Is the plan to break out the proto changes only or will this be broken up into multiple parts? Just thinking what's easiest to review

ekzhang · 2023-09-23T01:29:11Z

Is the plan to break out the proto changes only or will this be broken up into multiple parts? Just thinking what's easiest to review

It's reviewable right now in full, with the server changes. There's some data dependencies.

The deployment plan is in https://github.com/modal-labs/modal/pull/7745 — first just split out the api.proto and _serialization.py changes of this PR first. Then the server changes are merged, then we cycle worker versions, then we merge this PR.

It's hard to just review one section of the code since they depend on each other. The api.proto / _serialization.py changes won't make sense without looking at how they're used.

ekzhang

Description of changes

ekzhang · 2023-09-23T01:34:34Z

modal/_asgi.py

@@ -23,8 +23,6 @@ async def receive():

        # Run the ASGI app, while draining the send message queue at the same time,
        # and yielding results.
-        # TODO(gongy): we currently create one ASGI instance per concurrent input,
-        # but it would be sufficient to share this background ASGI context.


I removed this comment because I think it's outdated — we only call the function once in import_function() to produce the ASGI app right now even for multiple inputs and regardless of input concurrency. Then this gets passed into asgi_app_wrapper. (cc @gongy)

ekzhang · 2023-09-23T01:39:45Z

modal/_container_entrypoint.py

@@ -161,6 +161,12 @@ def serialize(self, obj: Any) -> bytes:
    def deserialize(self, data: bytes) -> Any:
        return deserialize(data, self._client)

+    def serialize_data_format(self, obj: Any, data_format: int) -> bytes:
+        return serialize_data_format(obj, data_format)


This file contains the main client change in the PR. We call serialize_data_format with FunctionPutOutputs, which adds data_format=api_pb2.DATA_FORMAT_PICKLE to all FunctionPutOutputsRequest from the client.

Except for web endpoints, which we serialize with api_pb2.DATA_FORMAT_ASGI.

The reason why we need to cycle worker versions with the new protobuf changes before we can deploy this is that workers proxy the FunctionPutOutputs RPC, and prost deletes any unknown Protobuf fields.

ekzhang · 2023-09-23T01:40:51Z

modal/_container_entrypoint.py

            status=api_pb2.GenericResult.GENERIC_STATUS_SUCCESS,
-            data=self.serialize(data),
+            data=self.serialize_data_format(data, data_format),


Yep! The rest of this file is just plumbing for serialize_data_format

modal/_serialization.py

ekzhang · 2023-09-23T01:42:53Z

modal/_serialization.py

+    elif msg_type == "http.disconnect":
+        return api_pb2.Asgi(http_disconnect=api_pb2.Asgi.HttpDisconnect())
+    else:
+        logger.debug("skipping serialization of unknown ASGI message type %r", msg_type)


If we get a message like a WebSocket response, or an unknown ASGI extension which somehow got enabled despite us not advertising support for it, we skip the serialization.

would it be better to raise an exception? this will just ignore it effectively (people won't see the debug log unless they override the verbosity level)

I think there's some kinds of messages we want to ignore as unimplemented without throwing (like the lifespan protocol). But yeah for websocket.* events might be nice to throw explicitly

I think we should ignore it, because there are many messages that we can't predict. The idea is "accept as many servers as possible, but require compliance on our side by users."

WebSockets just don't work on Modal, and that's fine documented as-is I think. We don't need to crash the user's app after a WebSocket request.

could potentially also warn. I'm a bit nervous being completely silent will cause confusion to users

Hmm fair enough — on second thought I don't think we actually forward WebSocket connection events on our side though, so it's a bit of a moot point?

ekzhang · 2023-09-23T01:43:45Z

modal/_serialization.py

+def deserialize_data_format(s: bytes, data_format: int, client) -> Any:
+    if data_format == api_pb2.DATA_FORMAT_UNSPECIFIED:
+        # TODO: Remove this after Modal client version 0.52, when the data_format field is always set.
+        return deserialize(s, client)


We need to keep this for compatibility with clients before today, since they won't be aware to set data_format. I made sure to set data_format in all locations and tested the code.

ekzhang · 2023-09-23T01:45:08Z

modal/functions.py

@@ -140,7 +141,7 @@ async def _process_result(result, stub, client=None):
        raise RemoteError(result.exception)

    try:
-        return deserialize(data, client)
+        return deserialize_data_format(data, data_format, client)


This is internal to the _Invocation class, which I updated to be data format-aware. So FunctionGetOutputs can read ASGI-serialized messages too.

modal/_serialization.py

erikbern

Looks good! I've been trying to think of ways this could break but there's nothing I can think of.

aksh-at · 2023-09-23T20:13:56Z

modal/_container_entrypoint.py

@@ -161,6 +161,12 @@ def serialize(self, obj: Any) -> bytes:
    def deserialize(self, data: bytes) -> Any:
        return deserialize(data, self._client)

+    def serialize_data_format(self, obj: Any, data_format: int) -> bytes:


tiny nit, should this be

Suggested change

def serialize_data_format(self, obj: Any, data_format: int) -> bytes:

def serialize_data_format(self, obj: Any, data_format: api_pb2.DataFormat.ValueType) -> bytes:

Honestly don't know how this affects mypy, can I defer it until later?

I've had a bunch of mypy issues with ValueType

aksh-at · 2023-09-23T20:19:02Z

modal/_serialization.py

+    elif msg_type == "http.disconnect":
+        return api_pb2.Asgi(http_disconnect=api_pb2.Asgi.HttpDisconnect())
+    else:
+        logger.debug("skipping serialization of unknown ASGI message type %r", msg_type)


I think there's some kinds of messages we want to ignore as unimplemented without throwing (like the lifespan protocol). But yeah for websocket.* events might be nice to throw explicitly

* Just the protobuf parts of #905 * Fix mypy and copyright * Add skip_old_by to my new test * Fix docstring typo

ekzhang marked this pull request as draft September 21, 2023 22:42

ekzhang added 2 commits September 22, 2023 16:33

Add AsgiProtocol message to the client protobuf

cd0c8a6

Update Asgi message docs

c1b15cb

ekzhang force-pushed the ekzhang/better-asgi branch from 3f30fd3 to c1b15cb Compare September 22, 2023 16:33

ekzhang added 4 commits September 22, 2023 19:43

Add DATA_FORMAT_PICKLE

aa2488c

Finish client support for ASGI serialization

614769f

Fix mypy issue

31431bc

Refactor and fix tests

302773c

ekzhang requested a review from erikbern September 22, 2023 21:48

ekzhang marked this pull request as ready for review September 22, 2023 21:48

ekzhang added 2 commits September 22, 2023 21:50

Remove an error that was used for strictness

0f82be7

Fix mypy

b4b5f57

ekzhang changed the title ~~Simulate the ASGI protocol directly with Modal functions~~ Serialize the ASGI protocol directly with Modal functions Sep 22, 2023

erikbern reviewed Sep 22, 2023

View reviewed changes

ekzhang commented Sep 23, 2023

View reviewed changes

erikbern reviewed Sep 23, 2023

View reviewed changes

modal/_serialization.py Outdated Show resolved Hide resolved

erikbern approved these changes Sep 23, 2023

View reviewed changes

aksh-at approved these changes Sep 23, 2023

View reviewed changes

ekzhang added 2 commits September 25, 2023 15:30

Merge branch 'main' into ekzhang/better-asgi

be20780

Fix typo on trailers

eee505d

ekzhang added a commit that referenced this pull request Sep 25, 2023

Just the protobuf parts of #905

8eddb99

erikbern mentioned this pull request Sep 25, 2023

Just the protobuf parts of #905 #918

Merged

Add serialization_test to ensure ASGI roundtrips

48db306

ekzhang added a commit that referenced this pull request Sep 25, 2023

Just the protobuf parts of #905

ac3292d

ekzhang added a commit that referenced this pull request Sep 25, 2023

Just the protobuf parts of #905 (#918)

99019a9

* Just the protobuf parts of #905 * Fix mypy and copyright * Add skip_old_by to my new test * Fix docstring typo

Merge branch 'main' into ekzhang/better-asgi

06e13a7

ekzhang merged commit e0e8735 into main Sep 25, 2023

ekzhang deleted the ekzhang/better-asgi branch September 25, 2023 21:06

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Serialize the ASGI protocol directly with Modal functions #905

Serialize the ASGI protocol directly with Modal functions #905

ekzhang commented Sep 21, 2023 •

edited

Loading

ekzhang commented Sep 21, 2023

erikbern Sep 22, 2023

ekzhang Sep 23, 2023

erikbern Sep 22, 2023

erikbern Sep 22, 2023

ekzhang Sep 23, 2023

erikbern commented Sep 22, 2023 •

edited

Loading

ekzhang commented Sep 23, 2023 •

edited

Loading

ekzhang left a comment

ekzhang Sep 23, 2023

ekzhang Sep 23, 2023

ekzhang Sep 23, 2023

ekzhang Sep 23, 2023

erikbern Sep 23, 2023

aksh-at Sep 23, 2023

ekzhang Sep 25, 2023

erikbern Sep 25, 2023

ekzhang Sep 25, 2023

ekzhang Sep 23, 2023

ekzhang Sep 23, 2023

erikbern left a comment

aksh-at Sep 23, 2023

ekzhang Sep 25, 2023

erikbern Sep 25, 2023

aksh-at Sep 23, 2023

	def serialize_data_format(self, obj: Any, data_format: int) -> bytes:
	def serialize_data_format(self, obj: Any, data_format: api_pb2.DataFormat.ValueType) -> bytes:

Serialize the ASGI protocol directly with Modal functions #905

Serialize the ASGI protocol directly with Modal functions #905

Conversation

ekzhang commented Sep 21, 2023 • edited Loading

ekzhang commented Sep 21, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

erikbern commented Sep 22, 2023 • edited Loading

ekzhang commented Sep 23, 2023 • edited Loading

ekzhang left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

erikbern left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ekzhang commented Sep 21, 2023 •

edited

Loading

erikbern commented Sep 22, 2023 •

edited

Loading

ekzhang commented Sep 23, 2023 •

edited

Loading