modal 1.0.6.dev58__py3-none-any.whl → 1.2.3.dev7__py3-none-any.whl

modal/_runtime/user_code_imports.py

@@ -29,7 +29,7 @@ class FinalizedFunction:
     callable: Callable[..., Any]
     is_async: bool
     is_generator: bool
-    data_format: int  # api_pb2.DataFormat
+    supported_output_formats: Sequence["api_pb2.DataFormat.ValueType"]
     lifespan_manager: Optional["LifespanManager"] = None
 
 
@@ -93,9 +93,9 @@ def construct_webhook_callable(
 
 @dataclass
 class ImportedFunction(Service):
-    user_cls_instance: Any
     app: modal.app._App
     service_deps: Optional[Sequence["modal._object._Object"]]
+    user_cls_instance = None
 
     _user_defined_callable: Callable[..., Any]
 
@@ -108,6 +108,7 @@ class ImportedFunction(Service):
         is_generator = fun_def.function_type == api_pb2.Function.FUNCTION_TYPE_GENERATOR
 
         webhook_config = fun_def.webhook_config
+
         if not webhook_config.type:
             # for non-webhooks, the runnable is straight forward:
             return {
@@ -115,7 +116,10 @@ class ImportedFunction(Service):
                     callable=self._user_defined_callable,
                     is_async=is_async,
                     is_generator=is_generator,
-                    data_format=api_pb2.DATA_FORMAT_PICKLE,
+                    supported_output_formats=fun_def.supported_output_formats
+                    # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]` is only
+                    # needed for tests
+                    or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR],
                 )
             }
 
@@ -129,7 +133,8 @@ class ImportedFunction(Service):
                 lifespan_manager=lifespan_manager,
                 is_async=True,
                 is_generator=True,
-                data_format=api_pb2.DATA_FORMAT_ASGI,
+                # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_ASGI]` is only needed for tests
+                supported_output_formats=fun_def.supported_output_formats or [api_pb2.DATA_FORMAT_ASGI],
             )
         }
 
@@ -154,6 +159,7 @@ class ImportedClass(Service):
             # Use the function definition for whether this is a generator (overriden by webhooks)
             is_generator = _partial.params.is_generator
             webhook_config = _partial.params.webhook_config
+            method_def = fun_def.method_definitions[method_name]
 
             bound_func = user_func.__get__(self.user_cls_instance)
 
@@ -163,7 +169,10 @@ class ImportedClass(Service):
                     callable=bound_func,
                     is_async=is_async,
                     is_generator=bool(is_generator),
-                    data_format=api_pb2.DATA_FORMAT_PICKLE,
+                    # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR]` is only
+                    # needed for tests
+                    supported_output_formats=method_def.supported_output_formats
+                    or [api_pb2.DATA_FORMAT_PICKLE, api_pb2.DATA_FORMAT_CBOR],
                 )
             else:
                 web_callable, lifespan_manager = construct_webhook_callable(
@@ -174,7 +183,8 @@ class ImportedClass(Service):
                     lifespan_manager=lifespan_manager,
                     is_async=True,
                     is_generator=True,
-                    data_format=api_pb2.DATA_FORMAT_ASGI,
+                    # FIXME (elias): the following `or [api_pb2.DATA_FORMAT_ASGI]` is only needed for tests
+                    supported_output_formats=method_def.supported_output_formats or [api_pb2.DATA_FORMAT_ASGI],
                 )
             finalized_functions[method_name] = finalized_function
         return finalized_functions
@@ -199,7 +209,6 @@ def get_user_class_instance(_cls: modal.cls._Cls, args: tuple[Any, ...], kwargs:
 
 def import_single_function_service(
     function_def: api_pb2.Function,
-    ser_cls: Optional[type],  # used only for @build functions
     ser_fun: Optional[Callable[..., Any]],
 ) -> Service:
     """Imports a function dynamically, and locates the app.
@@ -228,12 +237,9 @@ def import_single_function_service(
     service_deps: Optional[Sequence["modal._object._Object"]] = None
     active_app: modal.app._App
 
-    user_cls_or_cls: typing.Union[None, type, modal.cls.Cls]
-    user_cls_instance = None
-
     if ser_fun is not None:
         # This is a serialized function we already fetched from the server
-        user_cls_or_cls, user_defined_callable = ser_cls, ser_fun
+        user_defined_callable = ser_fun
         active_app = get_active_app_fallback(function_def)
     else:
         # Load the module dynamically
@@ -244,58 +250,22 @@ def import_single_function_service(
             raise LocalFunctionError("Attempted to load a function defined in a function scope")
 
         parts = qual_name.split(".")
-        if len(parts) == 1:
-            # This is a function
-            user_cls_or_cls = None
-            f = getattr(module, qual_name)
-            if isinstance(f, Function):
-                _function: modal._functions._Function[Any, Any, Any] = synchronizer._translate_in(f)  # type: ignore
-                service_deps = _function.deps(only_explicit_mounts=True)
-                user_defined_callable = _function.get_raw_f()
-                assert _function._app  # app should always be set on a decorated function
-                active_app = _function._app
-            else:
-                user_defined_callable = f
-                active_app = get_active_app_fallback(function_def)
-
-        elif len(parts) == 2:
-            # This path should only be triggered by @build class builder methods and can be removed
-            # once @build is deprecated.
-            assert not function_def.use_method_name  # new "placeholder methods" should not be invoked directly!
-            assert function_def.is_builder_function
-            cls_name, fun_name = parts
-            user_cls_or_cls = getattr(module, cls_name)
-            if isinstance(user_cls_or_cls, modal.cls.Cls):
-                # The cls decorator is in global scope
-                _cls = typing.cast(modal.cls._Cls, synchronizer._translate_in(user_cls_or_cls))
-                user_defined_callable = _cls._callables[fun_name]
-                # Intentionally not including these, since @build functions don't actually
-                # forward the information from their parent class.
-                # service_deps = _cls._get_class_service_function().deps(only_explicit_mounts=True)
-                assert _cls._app
-                active_app = _cls._app
-            else:
-                # This is non-decorated class
-                user_defined_callable = getattr(user_cls_or_cls, fun_name)  # unbound method
-                active_app = get_active_app_fallback(function_def)
-        else:
+        if len(parts) != 1:
             raise InvalidError(f"Invalid function qualname {qual_name}")
 
-    # Instantiate the class if it's defined
-    if user_cls_or_cls:
-        if isinstance(user_cls_or_cls, modal.cls.Cls):
-            # This code is only used for @build methods on classes
-            _cls = typing.cast(modal.cls._Cls, user_cls_or_cls)
-            user_cls_instance = get_user_class_instance(_cls, (), {})
-            # Bind the unbound method to the instance as self (using the descriptor protocol!)
+        f = getattr(module, qual_name)
+        if isinstance(f, Function):
+            _function: modal._functions._Function[Any, Any, Any] = synchronizer._translate_in(f)  # type: ignore
+            service_deps = _function.deps(only_explicit_mounts=True)
+            user_defined_callable = _function.get_raw_f()
+            assert _function._app  # app should always be set on a decorated function
+            active_app = _function._app
         else:
-            # serialized=True or "undecorated"
-            user_cls_instance = user_cls_or_cls()
-
-        user_defined_callable = user_defined_callable.__get__(user_cls_instance)
+            # function isn't decorated in global scope
+            user_defined_callable = f
+            active_app = get_active_app_fallback(function_def)
 
     return ImportedFunction(
-        user_cls_instance,
         active_app,
         service_deps,
         user_defined_callable,

modal/_serialization.py

@@ -6,6 +6,14 @@ import typing
 from inspect import Parameter
 from typing import Any
 
+from modal._traceback import extract_traceback
+from modal.config import config
+
+try:
+    import cbor2  # type: ignore
+except ImportError:  # pragma: no cover - optional dependency
+    cbor2 = None
+
 import google.protobuf.message
 
 from modal._utils.async_utils import synchronizer
@@ -15,7 +23,7 @@ from ._object import _Object
 from ._type_manager import parameter_serde_registry, schema_registry
 from ._vendor import cloudpickle
 from .config import logger
-from .exception import DeserializationError, ExecutionError, InvalidError
+from .exception import DeserializationError, ExecutionError, InvalidError, SerializationError
 from .object import Object
 
 if typing.TYPE_CHECKING:
@@ -346,6 +354,12 @@ def _deserialize_asgi(asgi: api_pb2.Asgi) -> Any:
         return None
 
 
+def get_preferred_payload_format() -> "api_pb2.DataFormat.ValueType":
+    payload_format = (config.get("payload_format") or "pickle").lower()
+    data_format = api_pb2.DATA_FORMAT_CBOR if payload_format == "cbor" else api_pb2.DATA_FORMAT_PICKLE
+    return data_format
+
+
 def serialize_data_format(obj: Any, data_format: int) -> bytes:
     """Similar to serialize(), but supports other data formats."""
     if data_format == api_pb2.DATA_FORMAT_PICKLE:
@@ -355,6 +369,21 @@ def serialize_data_format(obj: Any, data_format: int) -> bytes:
     elif data_format == api_pb2.DATA_FORMAT_GENERATOR_DONE:
         assert isinstance(obj, api_pb2.GeneratorDone)
         return obj.SerializeToString(deterministic=True)
+    elif data_format == api_pb2.DATA_FORMAT_CBOR:
+        if cbor2 is None:
+            raise InvalidError("CBOR support requires the 'cbor2' package to be installed.")
+        try:
+            return cbor2.dumps(obj)
+        except cbor2.CBOREncodeTypeError:
+            try:
+                typename = f"{type(obj).__module__}.{type(obj).__name__}"
+            except Exception:
+                typename = str(type(obj))
+            raise SerializationError(
+                # TODO (elias): add documentation link for more information on this
+                f"Can not serialize type {typename} as cbor. If you need to use a custom data type, "
+                "try to serialize it yourself e.g. by using pickle.dumps(my_data)"
+            )
     else:
         raise InvalidError(f"Unknown data format {data_format!r}")
 
@@ -366,6 +395,10 @@ def deserialize_data_format(s: bytes, data_format: int, client) -> Any:
         return _deserialize_asgi(api_pb2.Asgi.FromString(s))
     elif data_format == api_pb2.DATA_FORMAT_GENERATOR_DONE:
         return api_pb2.GeneratorDone.FromString(s)
+    elif data_format == api_pb2.DATA_FORMAT_CBOR:
+        if cbor2 is None:
+            raise InvalidError("CBOR support requires the 'cbor2' package to be installed.")
+        return cbor2.loads(s)
     else:
         raise InvalidError(f"Unknown data format {data_format!r}")
 
@@ -579,3 +612,26 @@ def get_callable_schema(
         arguments=arguments,
         return_type=return_type_proto,
     )
+
+
+def pickle_exception(exc: BaseException) -> bytes:
+    try:
+        return serialize(exc)
+    except Exception as serialization_exc:
+        # We can't always serialize exceptions.
+        err = f"Failed to serialize exception {exc} of type {type(exc)}: {serialization_exc}"
+        logger.info(err)
+        return serialize(SerializationError(err))
+
+
+def pickle_traceback(exc: BaseException, task_id: str) -> tuple[bytes, bytes]:
+    serialized_tb, tb_line_cache = b"", b""
+
+    try:
+        tb_dict, line_cache = extract_traceback(exc, task_id)
+        serialized_tb = serialize(tb_dict)
+        tb_line_cache = serialize(line_cache)
+    except Exception:
+        logger.info("Failed to serialize exception traceback.")
+
+    return serialized_tb, tb_line_cache

modal/_utils/async_utils.py

@@ -1,12 +1,14 @@
 # Copyright Modal Labs 2022
 import asyncio
 import concurrent.futures
+import contextlib
 import functools
 import inspect
 import itertools
 import sys
 import time
 import typing
+import warnings
 from collections.abc import AsyncGenerator, AsyncIterable, Awaitable, Iterable, Iterator
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
@@ -51,6 +53,10 @@ def synchronize_api(obj, target_module=None):
     return synchronizer.create_blocking(obj, blocking_name, target_module=target_module)
 
 
+# Used for testing to configure the `n_attempts` that `retry` will use.
+RETRY_N_ATTEMPTS_OVERRIDE: Optional[int] = None
+
+
 def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout=90):
     """Decorator that calls an async function multiple times, with a given timeout.
 
@@ -75,8 +81,13 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
     def decorator(fn):
         @functools.wraps(fn)
         async def f_wrapped(*args, **kwargs):
+            if RETRY_N_ATTEMPTS_OVERRIDE is not None:
+                local_n_attempts = RETRY_N_ATTEMPTS_OVERRIDE
+            else:
+                local_n_attempts = n_attempts
+
             delay = base_delay
-            for i in range(n_attempts):
+            for i in range(local_n_attempts):
                 t0 = time.time()
                 try:
                     return await asyncio.wait_for(fn(*args, **kwargs), timeout=timeout)
@@ -84,12 +95,12 @@ def retry(direct_fn=None, *, n_attempts=3, base_delay=0, delay_factor=2, timeout
                     logger.debug(f"Function {fn} was cancelled")
                     raise
                 except Exception as e:
-                    if i >= n_attempts - 1:
+                    if i >= local_n_attempts - 1:
                         raise
                     logger.debug(
                         f"Failed invoking function {fn}: {e}"
                         f" (took {time.time() - t0}s, sleeping {delay}s"
-                        f" and trying {n_attempts - i - 1} more times)"
+                        f" and trying {local_n_attempts - i - 1} more times)"
                     )
                 await asyncio.sleep(delay)
                 delay *= delay_factor
@@ -125,7 +136,8 @@ class TaskContext:
     _loops: set[asyncio.Task]
 
     def __init__(self, grace: Optional[float] = None):
-        self._grace = grace
+        self._grace = grace  # grace is the time we want for tasks to finish before cancelling them
+        self._cancellation_grace: float = 1.0  # extra graceperiod for the cancellation itself to "bubble up"
         self._loops = set()
 
     async def start(self):
@@ -157,22 +169,29 @@ class TaskContext:
             # still needs to be handled
             # (https://stackoverflow.com/a/63356323/2475114)
             if gather_future:
-                try:
+                with contextlib.suppress(asyncio.CancelledError):
                     await gather_future
-                except asyncio.CancelledError:
-                    pass
 
+            cancelled_tasks: list[asyncio.Task] = []
             for task in self._tasks:
                 if task.done() and not task.cancelled():
                     # Raise any exceptions if they happened.
                     # Only tasks without a done_callback will still be present in self._tasks
                     task.result()
 
-                if task.done() or task in self._loops:  # Note: Legacy code, we can probably cancel loops.
+                if task.done():
                     continue
 
                 # Cancel any remaining unfinished tasks.
                 task.cancel()
+                cancelled_tasks.append(task)
+
+            cancellation_gather = asyncio.gather(*cancelled_tasks, return_exceptions=True)
+            try:
+                await asyncio.wait_for(cancellation_gather, timeout=self._cancellation_grace)
+            except asyncio.TimeoutError:
+                warnings.warn(f"Internal warning: Tasks did not cancel in a timely manner: {cancelled_tasks}")
+
             await asyncio.sleep(0)  # wake up coroutines waiting for cancellations
 
     async def __aexit__(self, exc_type, value, tb):
@@ -279,7 +298,9 @@ class TimestampPriorityQueue(Generic[T]):
 
     def __init__(self, maxsize: int = 0):
         self.condition = asyncio.Condition()
-        self._queue: asyncio.PriorityQueue[tuple[float, Union[T, None]]] = asyncio.PriorityQueue(maxsize=maxsize)
+        self._queue: asyncio.PriorityQueue[tuple[float, int, Union[T, None]]] = asyncio.PriorityQueue(maxsize=maxsize)
+        # Used to tiebreak items with the same timestamp that are not comparable. (eg. protos)
+        self._counter = itertools.count()
 
     async def close(self):
         await self.put(self._MAX_PRIORITY, None)
@@ -288,7 +309,7 @@ class TimestampPriorityQueue(Generic[T]):
         """
         Add an item to the queue to be processed at a specific timestamp.
         """
-        await self._queue.put((timestamp, item))
+        await self._queue.put((timestamp, next(self._counter), item))
         async with self.condition:
             self.condition.notify_all()  # notify any waiting coroutines
 
@@ -301,7 +322,7 @@ class TimestampPriorityQueue(Generic[T]):
                 while self.empty():
                     await self.condition.wait()
                 # peek at the next item
-                timestamp, item = await self._queue.get()
+                timestamp, counter, item = await self._queue.get()
                 now = time.time()
                 if timestamp < now:
                     return item
@@ -309,7 +330,7 @@ class TimestampPriorityQueue(Generic[T]):
                     return None
                 # not ready yet, calculate sleep time
                 sleep_time = timestamp - now
-                self._queue.put_nowait((timestamp, item))  # put it back
+                self._queue.put_nowait((timestamp, counter, item))  # put it back
                 # wait until either the timeout or a new item is added
                 try:
                     await asyncio.wait_for(self.condition.wait(), timeout=sleep_time)

modal/_utils/auth_token_manager.py

@@ -9,7 +9,6 @@ from typing import Any
 from modal.exception import ExecutionError
 from modal_proto import api_pb2, modal_api_grpc
 
-from .grpc_utils import retry_transient_errors
 from .logger import logger
 
 
@@ -27,7 +26,7 @@ class _AuthTokenManager:
         self._expiry = 0.0
         self._lock: typing.Union[asyncio.Lock, None] = None
 
-    async def get_token(self):
+    async def get_token(self) -> str:
         """
         When called, the AuthTokenManager can be in one of three states:
         1. Has a valid cached token. It is returned to the caller.
@@ -66,9 +65,7 @@ class _AuthTokenManager:
             # new token. Once we have a new token, the other coroutines will unblock and return from here.
             if self._token and not self._needs_refresh():
                 return
-            resp: api_pb2.AuthTokenGetResponse = await retry_transient_errors(
-                self._stub.AuthTokenGet, api_pb2.AuthTokenGetRequest()
-            )
+            resp: api_pb2.AuthTokenGetResponse = await self._stub.AuthTokenGet(api_pb2.AuthTokenGetRequest())
             if not resp.token:
                 # Not expected
                 raise ExecutionError(