Adds support for cbor2 serialization for function calling [SDK-584] (#3474)

Deployed functions will now support both pickle and cbor inputs, and try to mirror the output format with the input format when allowed.
In addition, functions will now advertise which input and output formats they support to allow a lookup to introspect this.

This enables future versions of modal clients (in particular libmodal clients in other languages) to use cbor instead of pickle for payloads.

Author: freider

.vscode/settings.json

@@ -1,4 +1,7 @@
 {
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,
+    "python.testing.pytestArgs": [
+        "test"
+    ],
 }

modal/_container_entrypoint.py

@@ -32,14 +32,14 @@
     _PartialFunctionFlags,
 )
 from modal._serialization import deserialize, deserialize_params
-from modal._utils.async_utils import TaskContext, synchronizer
+from modal._utils.async_utils import TaskContext, aclosing, synchronizer
 from modal._utils.function_utils import (
     callable_has_non_self_params,
 )
 from modal.app import App, _App
 from modal.client import Client, _Client
 from modal.config import logger
-from modal.exception import ExecutionError, InputCancellation, InvalidError
+from modal.exception import ExecutionError, InputCancellation
 from modal.running_app import RunningApp, running_app_from_layout
 from modal_proto import api_pb2
 
@@ -184,17 +184,13 @@ def call_function(
     batch_wait_ms: int,
 ):
     async def run_input_async(io_context: IOContext) -> None:
-        started_at = time.time()
         reset_context = execution_context._set_current_context_ids(
             io_context.input_ids, io_context.function_call_ids, io_context.attempt_tokens
         )
+        started_at = time.time()
         async with container_io_manager.handle_input_exception.aio(io_context, started_at):
-            res = io_context.call_finalized_function()
             # TODO(erikbern): any exception below shouldn't be considered a user exception
             if io_context.finalized_function.is_generator:
-                if not inspect.isasyncgen(res):
-                    raise InvalidError(f"Async generator function returned value of type {type(res)}")
-
                 # Send up to this many outputs at a time.
                 current_function_call_id = execution_context.current_function_call_id()
                 assert current_function_call_id is not None  # Set above.
@@ -204,33 +200,24 @@ async def run_input_async(io_context: IOContext) -> None:
                 async with container_io_manager.generator_output_sender(
                     current_function_call_id,
                     current_attempt_token,
-                    io_context.finalized_function.data_format,
+                    io_context._generator_output_format(),
                     generator_queue,
                 ):
                     item_count = 0
-                    async for value in res:
-                        await container_io_manager._queue_put.aio(generator_queue, value)
-                        item_count += 1
+                    async with aclosing(io_context.call_generator_async()) as gen:
+                        async for value in gen:
+                            await container_io_manager._queue_put.aio(generator_queue, value)
+                            item_count += 1
 
-                message = api_pb2.GeneratorDone(items_total=item_count)
-                await container_io_manager.push_outputs.aio(
-                    io_context,
-                    started_at,
-                    message,
-                    api_pb2.DATA_FORMAT_GENERATOR_DONE,
+                await container_io_manager._send_outputs.aio(
+                    started_at, io_context.output_items_generator_done(started_at, item_count)
                 )
             else:
-                if not inspect.iscoroutine(res) or inspect.isgenerator(res) or inspect.isasyncgen(res):
-                    raise InvalidError(
-                        f"Async (non-generator) function returned value of type {type(res)}"
-                        " You might need to use @app.function(..., is_generator=True)."
-                    )
-                value = await res
+                value = await io_context.call_function_async()
                 await container_io_manager.push_outputs.aio(
                     io_context,
                     started_at,
                     value,
-                    io_context.finalized_function.data_format,
                 )
         reset_context()
 
@@ -240,13 +227,9 @@ def run_input_sync(io_context: IOContext) -> None:
             io_context.input_ids, io_context.function_call_ids, io_context.attempt_tokens
         )
         with container_io_manager.handle_input_exception(io_context, started_at):
-            res = io_context.call_finalized_function()
-
             # TODO(erikbern): any exception below shouldn't be considered a user exception
             if io_context.finalized_function.is_generator:
-                if not inspect.isgenerator(res):
-                    raise InvalidError(f"Generator function returned value of type {type(res)}")
-
+                gen = io_context.call_generator_sync()
                 # Send up to this many outputs at a time.
                 current_function_call_id = execution_context.current_function_call_id()
                 assert current_function_call_id is not None  # Set above.
@@ -256,25 +239,20 @@ def run_input_sync(io_context: IOContext) -> None:
                 with container_io_manager.generator_output_sender(
                     current_function_call_id,
                     current_attempt_token,
-                    io_context.finalized_function.data_format,
+                    io_context._generator_output_format(),
                     generator_queue,
                 ):
                     item_count = 0
-                    for value in res:
+                    for value in gen:
                         container_io_manager._queue_put(generator_queue, value)
                         item_count += 1
 
-                message = api_pb2.GeneratorDone(items_total=item_count)
-                container_io_manager.push_outputs(io_context, started_at, message, api_pb2.DATA_FORMAT_GENERATOR_DONE)
-            else:
-                if inspect.iscoroutine(res) or inspect.isgenerator(res) or inspect.isasyncgen(res):
-                    raise InvalidError(
-                        f"Sync (non-generator) function return value of type {type(res)}."
-                        " You might need to use @app.function(..., is_generator=True)."
-                    )
-                container_io_manager.push_outputs(
-                    io_context, started_at, res, io_context.finalized_function.data_format
+                container_io_manager._send_outputs(
+                    started_at, io_context.output_items_generator_done(started_at, item_count)
                 )
+            else:
+                values = io_context.call_function_sync()
+                container_io_manager.push_outputs(io_context, started_at, values)
         reset_context()
 
     if container_io_manager.input_concurrency_enabled:

modal/_functions.py

@@ -150,8 +150,7 @@ async def create(
             args,
             kwargs,
             stub,
-            max_object_size_bytes=function._max_object_size_bytes,
-            method_name=function._use_method_name,
+            function=function,
             function_call_invocation_type=function_call_invocation_type,
         )
 
@@ -439,8 +438,7 @@ async def create(
             args,
             kwargs,
             control_plane_stub,
-            max_object_size_bytes=function._max_object_size_bytes,
-            method_name=function._use_method_name,
+            function=function,
         )
 
         request = api_pb2.AttemptStartRequest(
@@ -698,6 +696,7 @@ def from_local(
         experimental_options: Optional[dict[str, str]] = None,
         _experimental_proxy_ip: Optional[str] = None,
         _experimental_custom_scaling_factor: Optional[float] = None,
+        restrict_output: bool = False,
     ) -> "_Function":
         """mdmd:hidden
 
@@ -834,17 +833,23 @@ def from_local(
                     is_web_endpoint=is_web_endpoint,
                     ignore_first_argument=True,
                 )
+                if is_web_endpoint:
+                    method_input_formats = [api_pb2.DATA_FORMAT_ASGI]
+                    method_output_formats = [api_pb2.DATA_FORMAT_ASGI]
+                else:
+                    method_input_formats = [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]
+                    if restrict_output:
+                        method_output_formats = [api_pb2.DATA_FORMAT_CBOR]
+                    else:
+                        method_output_formats = [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]
+
                 method_definition = api_pb2.MethodDefinition(
                     webhook_config=partial_function.params.webhook_config,
                     function_type=function_type,
                     function_name=function_name,
                     function_schema=method_schema,
-                    supported_input_formats=[api_pb2.DATA_FORMAT_ASGI]
-                    if is_web_endpoint
-                    else [api_pb2.DATA_FORMAT_PICKLE],
-                    supported_output_formats=[api_pb2.DATA_FORMAT_ASGI]
-                    if is_web_endpoint
-                    else [api_pb2.DATA_FORMAT_PICKLE],
+                    supported_input_formats=method_input_formats,
+                    supported_output_formats=method_output_formats,
                 )
                 method_definitions[method_name] = method_definition
 
@@ -869,16 +874,18 @@ def _deps(only_explicit_mounts=False) -> list[_Object]:
             return deps
 
         if info.is_service_class():
-            # classes don't have data formats themselves - methods do
+            # classes don't have data formats themselves - input/output formats are set per method above
             supported_input_formats = []
             supported_output_formats = []
         elif webhook_config is not None:
             supported_input_formats = [api_pb2.DATA_FORMAT_ASGI]
             supported_output_formats = [api_pb2.DATA_FORMAT_ASGI]
         else:
-            # TODO: add CBOR support
-            supported_input_formats = [api_pb2.DATA_FORMAT_PICKLE]
-            supported_output_formats = [api_pb2.DATA_FORMAT_PICKLE]
+            supported_input_formats = [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]
+            if restrict_output:
+                supported_output_formats = [api_pb2.DATA_FORMAT_CBOR]
+            else:
+                supported_output_formats = [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]
 
         async def _preload(self: _Function, resolver: Resolver, existing_object_id: Optional[str]):
             assert resolver.client and resolver.client.stub