modal 1.1.5.dev43__py3-none-any.whl → 1.1.5.dev45__py3-none-any.whl

modal/_container_entrypoint.py

@@ -32,14 +32,14 @@ from modal._partial_function import (
     _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 @@ def call_function(
                 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 call_function(
             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 call_function(
                 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 @@ class _Invocation:
             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 @@ class _InputPlaneInvocation:
             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 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         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 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     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 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
             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