modal 0.73.116__py3-none-any.whl → 0.73.128__py3-none-any.whl

modal/__init__.py

@@ -27,6 +27,7 @@ try:
         asgi_app,
         batched,
         build,
+        concurrent,
         enter,
         exit,
         fastapi_endpoint,
@@ -82,6 +83,7 @@ __all__ = [
     "asgi_app",
     "batched",
     "build",
+    "concurrent",
     "current_function_call_id",
     "current_input_id",
     "enable_output",

modal/_functions.py

@@ -1,4 +1,5 @@
 # Copyright Modal Labs 2023
+import asyncio
 import dataclasses
 import inspect
 import textwrap
@@ -24,7 +25,7 @@ from ._pty import get_pty_info
 from ._resolver import Resolver
 from ._resources import convert_fn_config_to_resources_config
 from ._runtime.execution_context import current_input_id, is_local
-from ._serialization import serialize, serialize_proto_params
+from ._serialization import apply_defaults, serialize, serialize_proto_params, validate_params
 from ._traceback import print_server_warnings
 from ._utils.async_utils import (
     TaskContext,
@@ -174,7 +175,7 @@ class _Invocation:
             input_jwt=input.input_jwt,
             input_id=input.input_id,
             item=item,
-            sync_client_retries_enabled=response.sync_client_retries_enabled
+            sync_client_retries_enabled=response.sync_client_retries_enabled,
         )
         return _Invocation(client.stub, function_call_id, client, retry_context)
 
@@ -256,9 +257,13 @@ class _Invocation:
             try:
                 return await self._get_single_output(ctx.input_jwt)
             except (UserCodeException, FunctionTimeoutError) as exc:
-                await user_retry_manager.raise_or_sleep(exc)
+                delay_ms = user_retry_manager.get_delay_ms()
+                if delay_ms is None:
+                    raise exc
+                await asyncio.sleep(delay_ms / 1000)
             except InternalFailure:
-                # For system failures on the server, we retry immediately.
+                # For system failures on the server, we retry immediately,
+                # and the failure does not count towards the retry policy.
                 pass
             await self._retry_input()
 
@@ -430,7 +435,8 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
         max_containers: Optional[int] = None,
         buffer_containers: Optional[int] = None,
         scaledown_window: Optional[int] = None,
-        allow_concurrent_inputs: Optional[int] = None,
+        max_concurrent_inputs: Optional[int] = None,
+        target_concurrent_inputs: Optional[int] = None,
         batch_max_size: Optional[int] = None,
         batch_wait_ms: Optional[int] = None,
         cloud: Optional[str] = None,
@@ -781,7 +787,8 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                     runtime_perf_record=config.get("runtime_perf_record"),
                     app_name=app_name,
                     is_builder_function=is_builder_function,
-                    target_concurrent_inputs=allow_concurrent_inputs or 0,
+                    max_concurrent_inputs=max_concurrent_inputs or 0,
+                    target_concurrent_inputs=target_concurrent_inputs or 0,
                     batch_max_size=batch_max_size or 0,
                     batch_linger_ms=batch_wait_ms or 0,
                     worker_id=config.get("worker_id"),
@@ -968,8 +975,11 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                         "Can't use positional arguments with modal.parameter-based synthetic constructors.\n"
                         "Use (<parameter_name>=value) keyword arguments when constructing classes instead."
                     )
-                serialized_params = serialize_proto_params(kwargs, parent._class_parameter_info.schema)
-                can_use_parent = len(parent._class_parameter_info.schema) == 0
+                schema = parent._class_parameter_info.schema
+                kwargs_with_defaults = apply_defaults(kwargs, schema)
+                validate_params(kwargs_with_defaults, schema)
+                serialized_params = serialize_proto_params(kwargs_with_defaults)
+                can_use_parent = len(parent._class_parameter_info.schema) == 0  # no parameters
             else:
                 can_use_parent = len(args) + len(kwargs) == 0 and options is None
                 serialized_params = serialize((args, kwargs))
@@ -1304,6 +1314,7 @@ class _Function(typing.Generic[P, ReturnType, OriginalReturnType], _Object, type
                 order_outputs,
                 return_exceptions,
                 count_update_callback,
+                api_pb2.FUNCTION_CALL_INVOCATION_TYPE_SYNC
             )
         ) as stream:
             async for item in stream:

modal/_partial_function.py

@@ -59,6 +59,8 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
     force_build: bool
     cluster_size: Optional[int]  # Experimental: Clustered functions
     build_timeout: Optional[int]
+    max_concurrent_inputs: Optional[int]
+    target_concurrent_inputs: Optional[int]
 
     def __init__(
         self,
@@ -72,6 +74,8 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
         cluster_size: Optional[int] = None,  # Experimental: Clustered functions
         force_build: bool = False,
         build_timeout: Optional[int] = None,
+        max_concurrent_inputs: Optional[int] = None,
+        target_concurrent_inputs: Optional[int] = None,
     ):
         self.raw_f = raw_f
         self.flags = flags
@@ -89,6 +93,8 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
         self.cluster_size = cluster_size  # Experimental: Clustered functions
         self.force_build = force_build
         self.build_timeout = build_timeout
+        self.max_concurrent_inputs = max_concurrent_inputs
+        self.target_concurrent_inputs = target_concurrent_inputs
 
     def _get_raw_f(self) -> Callable[P, ReturnType]:
         return self.raw_f
@@ -143,6 +149,8 @@ class _PartialFunction(typing.Generic[P, ReturnType, OriginalReturnType]):
             batch_wait_ms=self.batch_wait_ms,
             force_build=self.force_build,
             build_timeout=self.build_timeout,
+            max_concurrent_inputs=self.max_concurrent_inputs,
+            target_concurrent_inputs=self.target_concurrent_inputs,
         )
 
 
@@ -722,3 +730,49 @@ def _batched(
         )
 
     return wrapper
+
+
+def _concurrent(
+    _warn_parentheses_missing=None,
+    *,
+    max_inputs: int,  # Hard limit on each container's input concurrency
+    target_inputs: Optional[int] = None,  # Input concurrency that Modal's autoscaler should target
+) -> Callable[[Union[Callable[..., Any], _PartialFunction]], _PartialFunction]:
+    """Decorator that allows individual containers to handle multiple inputs concurrently.
+
+    The concurrency mechanism depends on whether the function is async or not:
+    - Async functions will run inputs on a single thread as asyncio tasks.
+    - Synchronous functions will use multi-threading. The code must be thread-safe.
+
+    Input concurrency will be most useful for workflows that are IO-bound
+    (e.g., making network requests) or when running an inference server that supports
+    dynamic batching.
+
+    When `target_inputs` is set, Modal's autoscaler will try to provision resources such
+    that each container is running that many inputs concurrently. Containers may burst up to
+    up to `max_inputs` if resources are insufficient to remain at the target concurrency.
+    """
+    if _warn_parentheses_missing is not None:
+        raise InvalidError(
+            "Positional arguments are not allowed. Did you forget parentheses? Suggestion: `@modal.concurrent()`."
+        )
+
+    if target_inputs and target_inputs > max_inputs:
+        raise InvalidError("`target_inputs` parameter cannot be greater than `max_inputs`.")
+
+    def wrapper(obj: Union[Callable[..., Any], _PartialFunction]) -> _PartialFunction:
+        if isinstance(obj, _PartialFunction):
+            # Risky that we need to mutate the parameters here; should make this safer
+            obj.max_concurrent_inputs = max_inputs
+            obj.target_concurrent_inputs = target_inputs
+            obj.add_flags(_PartialFunctionFlags.FUNCTION)
+            return obj
+
+        return _PartialFunction(
+            obj,
+            _PartialFunctionFlags.FUNCTION,
+            max_concurrent_inputs=max_inputs,
+            target_concurrent_inputs=target_inputs,
+        )
+
+    return wrapper

modal/_runtime/container_io_manager.py

@@ -63,7 +63,9 @@ class IOContext:
     """
 
     input_ids: list[str]
+    retry_counts: list[int]
     function_call_ids: list[str]
+    function_inputs: list[api_pb2.FunctionInput]
     finalized_function: "modal._runtime.user_code_imports.FinalizedFunction"
 
     _cancel_issued: bool = False
@@ -72,6 +74,7 @@ class IOContext:
     def __init__(
         self,
         input_ids: list[str],
+        retry_counts: list[int],
         function_call_ids: list[str],
         finalized_function: "modal._runtime.user_code_imports.FinalizedFunction",
         function_inputs: list[api_pb2.FunctionInput],
@@ -79,9 +82,10 @@ class IOContext:
         client: _Client,
     ):
         self.input_ids = input_ids
+        self.retry_counts = retry_counts
         self.function_call_ids = function_call_ids
         self.finalized_function = finalized_function
-        self._function_inputs = function_inputs
+        self.function_inputs = function_inputs
         self._is_batched = is_batched
         self._client = client
 
@@ -90,11 +94,11 @@ class IOContext:
         cls,
         client: _Client,
         finalized_functions: dict[str, "modal._runtime.user_code_imports.FinalizedFunction"],
-        inputs: list[tuple[str, str, api_pb2.FunctionInput]],
+        inputs: list[tuple[str, int, str, api_pb2.FunctionInput]],
         is_batched: bool,
     ) -> "IOContext":
         assert len(inputs) >= 1 if is_batched else len(inputs) == 1
-        input_ids, function_call_ids, function_inputs = zip(*inputs)
+        input_ids, retry_counts, function_call_ids, function_inputs = zip(*inputs)
 
         async def _populate_input_blobs(client: _Client, input: api_pb2.FunctionInput) -> api_pb2.FunctionInput:
             # If we got a pointer to a blob, download it from S3.
@@ -111,7 +115,7 @@ class IOContext:
         method_name = function_inputs[0].method_name
         assert all(method_name == input.method_name for input in function_inputs)
         finalized_function = finalized_functions[method_name]
-        return cls(input_ids, function_call_ids, finalized_function, function_inputs, is_batched, client)
+        return cls(input_ids, retry_counts, function_call_ids, finalized_function, function_inputs, is_batched, client)
 
     def set_cancel_callback(self, cb: Callable[[], None]):
         self._cancel_callback = cb
@@ -135,7 +139,7 @@ class IOContext:
         # to make sure we handle user exceptions properly
         # and don't retry
         deserialized_args = [
-            deserialize(input.args, self._client) if input.args else ((), {}) for input in self._function_inputs
+            deserialize(input.args, self._client) if input.args else ((), {}) for input in self.function_inputs
         ]
         if not self._is_batched:
             return deserialized_args[0]
@@ -551,7 +555,7 @@ class _ContainerIOManager:
         self,
         batch_max_size: int,
         batch_wait_ms: int,
-    ) -> AsyncIterator[list[tuple[str, str, api_pb2.FunctionInput]]]:
+    ) -> AsyncIterator[list[tuple[str, int, str, api_pb2.FunctionInput]]]:
         request = api_pb2.FunctionGetInputsRequest(function_id=self.function_id)
         iteration = 0
         while self._fetching_inputs:
@@ -586,8 +590,7 @@ class _ContainerIOManager:
                         if item.kill_switch:
                             logger.debug(f"Task {self.task_id} input kill signal input.")
                             return
-
-                        inputs.append((item.input_id, item.function_call_id, item.input))
+                        inputs.append((item.input_id, item.retry_count, item.function_call_id, item.input))
                         if item.input.final_input:
                             if request.batch_max_size > 0:
                                 logger.debug(f"Task {self.task_id} Final input not expected in batch input stream")
@@ -648,8 +651,9 @@ class _ContainerIOManager:
                 output_created_at=output_created_at,
                 result=result,
                 data_format=data_format,
+                retry_count=retry_count,
             )
-            for input_id, result in zip(io_context.input_ids, results)
+            for input_id, retry_count, result in zip(io_context.input_ids, io_context.retry_counts, results)
         ]
         await retry_transient_errors(
             self._client.stub.FunctionPutOutputs,

modal/_runtime/container_io_manager.pyi

@@ -14,7 +14,9 @@ class Sentinel: ...
 
 class IOContext:
     input_ids: list[str]
+    retry_counts: list[int]
     function_call_ids: list[str]
+    function_inputs: list[modal_proto.api_pb2.FunctionInput]
     finalized_function: modal._runtime.user_code_imports.FinalizedFunction
     _cancel_issued: bool
     _cancel_callback: typing.Optional[collections.abc.Callable[[], None]]
@@ -22,6 +24,7 @@ class IOContext:
     def __init__(
         self,
         input_ids: list[str],
+        retry_counts: list[int],
         function_call_ids: list[str],
         finalized_function: modal._runtime.user_code_imports.FinalizedFunction,
         function_inputs: list[modal_proto.api_pb2.FunctionInput],
@@ -33,7 +36,7 @@ class IOContext:
         cls,
         client: modal.client._Client,
         finalized_functions: dict[str, modal._runtime.user_code_imports.FinalizedFunction],
-        inputs: list[tuple[str, str, modal_proto.api_pb2.FunctionInput]],
+        inputs: list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]],
         is_batched: bool,
     ) -> IOContext: ...
     def set_cancel_callback(self, cb: collections.abc.Callable[[], None]): ...
@@ -116,7 +119,7 @@ class _ContainerIOManager:
     def get_max_inputs_to_fetch(self): ...
     def _generate_inputs(
         self, batch_max_size: int, batch_wait_ms: int
-    ) -> collections.abc.AsyncIterator[list[tuple[str, str, modal_proto.api_pb2.FunctionInput]]]: ...
+    ) -> collections.abc.AsyncIterator[list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]]]: ...
     def run_inputs_outputs(
         self,
         finalized_functions: dict[str, modal._runtime.user_code_imports.FinalizedFunction],
@@ -287,10 +290,10 @@ class ContainerIOManager:
     class ___generate_inputs_spec(typing_extensions.Protocol[SUPERSELF]):
         def __call__(
             self, batch_max_size: int, batch_wait_ms: int
-        ) -> typing.Iterator[list[tuple[str, str, modal_proto.api_pb2.FunctionInput]]]: ...
+        ) -> typing.Iterator[list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]]]: ...
         def aio(
             self, batch_max_size: int, batch_wait_ms: int
-        ) -> collections.abc.AsyncIterator[list[tuple[str, str, modal_proto.api_pb2.FunctionInput]]]: ...
+        ) -> collections.abc.AsyncIterator[list[tuple[str, int, str, modal_proto.api_pb2.FunctionInput]]]: ...
 
     _generate_inputs: ___generate_inputs_spec[typing_extensions.Self]
 

modal/_serialization.py

@@ -400,6 +400,7 @@ class ParamTypeInfo:
     default_field: str
     proto_field: str
     converter: typing.Callable[[str], typing.Any]
+    type: type
 
 
 PYTHON_TO_PROTO_TYPE: dict[type, "api_pb2.ParameterType.ValueType"] = {
@@ -411,75 +412,112 @@ PYTHON_TO_PROTO_TYPE: dict[type, "api_pb2.ParameterType.ValueType"] = {
 
 PROTO_TYPE_INFO = {
     # Protobuf type enum -> encode/decode helper metadata
-    api_pb2.PARAM_TYPE_STRING: ParamTypeInfo(default_field="string_default", proto_field="string_value", converter=str),
-    api_pb2.PARAM_TYPE_INT: ParamTypeInfo(default_field="int_default", proto_field="int_value", converter=int),
+    api_pb2.PARAM_TYPE_STRING: ParamTypeInfo(
+        default_field="string_default",
+        proto_field="string_value",
+        converter=str,
+        type=str,
+    ),
+    api_pb2.PARAM_TYPE_INT: ParamTypeInfo(
+        default_field="int_default",
+        proto_field="int_value",
+        converter=int,
+        type=int,
+    ),
     api_pb2.PARAM_TYPE_BYTES: ParamTypeInfo(
-        default_field="bytes_default", proto_field="bytes_value", converter=assert_bytes
+        default_field="bytes_default",
+        proto_field="bytes_value",
+        converter=assert_bytes,
+        type=bytes,
     ),
 }
 
 
-def serialize_proto_params(python_params: dict[str, Any], schema: typing.Sequence[api_pb2.ClassParameterSpec]) -> bytes:
-    proto_params: list[api_pb2.ClassParameterValue] = []
+def apply_defaults(
+    python_params: typing.Mapping[str, Any], schema: typing.Sequence[api_pb2.ClassParameterSpec]
+) -> dict[str, Any]:
+    """Apply any declared defaults from the provided schema, if values aren't provided in python_params
+
+    Conceptually similar to inspect.BoundArguments.apply_defaults.
+
+    Note: Apply this before serializing parameters in order to get consistent parameter
+        pools regardless if a value is explicitly provided or not.
+    """
+    result = {**python_params}
     for schema_param in schema:
-        type_info = PROTO_TYPE_INFO.get(schema_param.type)
-        if not type_info:
-            raise ValueError(f"Unsupported parameter type: {schema_param.type}")
+        if schema_param.has_default and schema_param.name not in python_params:
+            default_field_name = schema_param.WhichOneof("default_oneof")
+            if default_field_name is None:
+                raise InvalidError(f"{schema_param.name} declared as having a default, but has no default value")
+            result[schema_param.name] = getattr(schema_param, default_field_name)
+    return result
+
+
+def serialize_proto_params(python_params: dict[str, Any]) -> bytes:
+    proto_params: list[api_pb2.ClassParameterValue] = []
+    for param_name, python_value in python_params.items():
+        python_type = type(python_value)
+        protobuf_type = get_proto_parameter_type(python_type)
+        type_info = PROTO_TYPE_INFO.get(protobuf_type)
         proto_param = api_pb2.ClassParameterValue(
-            name=schema_param.name,
-            type=schema_param.type,
+            name=param_name,
+            type=protobuf_type,
         )
-        python_value = python_params.get(schema_param.name)
-        if python_value is None:
-            if schema_param.has_default:
-                python_value = getattr(schema_param, type_info.default_field)
-            else:
-                raise ValueError(f"Missing required parameter: {schema_param.name}")
         try:
             converted_value = type_info.converter(python_value)
         except ValueError as exc:
-            raise ValueError(f"Invalid type for parameter {schema_param.name}: {exc}")
+            raise ValueError(f"Invalid type for parameter {param_name}: {exc}")
         setattr(proto_param, type_info.proto_field, converted_value)
         proto_params.append(proto_param)
     proto_bytes = api_pb2.ClassParameterSet(parameters=proto_params).SerializeToString(deterministic=True)
     return proto_bytes
 
 
-def deserialize_proto_params(serialized_params: bytes, schema: list[api_pb2.ClassParameterSpec]) -> dict[str, Any]:
-    # TODO: this currently requires the schema to decode a payload, but we should make the validation
-    #       distinct from the deserialization
+def deserialize_proto_params(serialized_params: bytes) -> dict[str, Any]:
     proto_struct = api_pb2.ClassParameterSet()
     proto_struct.ParseFromString(serialized_params)
-    value_by_name = {p.name: p for p in proto_struct.parameters}
     python_params = {}
-    for schema_param in schema:
-        if schema_param.name not in value_by_name:
-            # TODO: handle default values? Could just be a flag on the FunctionParameter schema spec,
-            #  allowing it to not be supplied in the FunctionParameterSet?
-            raise AttributeError(f"Constructor arguments don't match declared parameters (missing {schema_param.name})")
-        param_value = value_by_name[schema_param.name]
-        if schema_param.type != param_value.type:
-            raise ValueError(
-                "Constructor arguments types don't match declared parameters "
-                f"({schema_param.name}: type {schema_param.type} != type {param_value.type})"
-            )
+    for param in proto_struct.parameters:
         python_value: Any
-        if schema_param.type == api_pb2.PARAM_TYPE_STRING:
-            python_value = param_value.string_value
-        elif schema_param.type == api_pb2.PARAM_TYPE_INT:
-            python_value = param_value.int_value
-        elif schema_param.type == api_pb2.PARAM_TYPE_BYTES:
-            python_value = param_value.bytes_value
+        if param.type == api_pb2.PARAM_TYPE_STRING:
+            python_value = param.string_value
+        elif param.type == api_pb2.PARAM_TYPE_INT:
+            python_value = param.int_value
+        elif param.type == api_pb2.PARAM_TYPE_BYTES:
+            python_value = param.bytes_value
         else:
-            # TODO(elias): based on `parameters` declared types, we could add support for
-            #  custom non proto types encoded as bytes in the proto, e.g. PARAM_TYPE_PYTHON_PICKLE
-            raise NotImplementedError("Only strings and ints are supported parameter value types at the moment")
+            raise NotImplementedError(f"Unimplemented parameter type: {param.type}.")
 
-        python_params[schema_param.name] = python_value
+        python_params[param.name] = python_value
 
     return python_params
 
 
+def validate_params(params: dict[str, Any], schema: typing.Sequence[api_pb2.ClassParameterSpec]):
+    # first check that all declared values are provided
+    for schema_param in schema:
+        if schema_param.name not in params:
+            # we expect all values to be present - even defaulted ones (defaults are applied on payload construction)
+            raise InvalidError(f"Missing required parameter: {schema_param.name}")
+        python_value = params[schema_param.name]
+        python_type = type(python_value)
+        param_protobuf_type = get_proto_parameter_type(python_type)
+        if schema_param.type != param_protobuf_type:
+            expected_python_type = PROTO_TYPE_INFO[schema_param.type].type
+            raise TypeError(
+                f"Parameter '{schema_param.name}' type error: expected {expected_python_type.__name__}, "
+                f"got {python_type.__name__}"
+            )
+
+    schema_fields = {p.name for p in schema}
+    # then check that no extra values are provided
+    non_declared_fields = params.keys() - schema_fields
+    if non_declared_fields:
+        raise InvalidError(
+            f"The following parameter names were provided but are not present in the schema: {non_declared_fields}"
+        )
+
+
 def deserialize_params(serialized_params: bytes, function_def: api_pb2.Function, _client: "modal.client._Client"):
     if function_def.class_parameter_info.format in (
         api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_UNSPECIFIED,
@@ -488,11 +526,21 @@ def deserialize_params(serialized_params: bytes, function_def: api_pb2.Function,
         # legacy serialization format - pickle of `(args, kwargs)` w/ support for modal object arguments
         param_args, param_kwargs = deserialize(serialized_params, _client)
     elif function_def.class_parameter_info.format == api_pb2.ClassParameterInfo.PARAM_SERIALIZATION_FORMAT_PROTO:
-        param_args = ()
-        param_kwargs = deserialize_proto_params(serialized_params, list(function_def.class_parameter_info.schema))
+        param_args = ()  # we use kwargs only for our implicit constructors
+        param_kwargs = deserialize_proto_params(serialized_params)
+        # TODO: We can probably remove the validation below since we do validation in the caller?
+        validate_params(param_kwargs, list(function_def.class_parameter_info.schema))
     else:
         raise ExecutionError(
             f"Unknown class parameter serialization format: {function_def.class_parameter_info.format}"
         )
 
     return param_args, param_kwargs
+
+
+def get_proto_parameter_type(parameter_type: type) -> "api_pb2.ParameterType.ValueType":
+    if parameter_type not in PYTHON_TO_PROTO_TYPE:
+        type_name = getattr(parameter_type, "__name__", repr(parameter_type))
+        supported = ", ".join(parameter_type.__name__ for parameter_type in PYTHON_TO_PROTO_TYPE.keys())
+        raise InvalidError(f"{type_name} is not a supported parameter type. Use one of: {supported}")
+    return PYTHON_TO_PROTO_TYPE[parameter_type]

modal/_utils/async_utils.py

@@ -12,6 +12,7 @@ from dataclasses import dataclass
 from typing import (
     Any,
     Callable,
+    Generic,
     Optional,
     TypeVar,
     Union,
@@ -26,6 +27,10 @@ from typing_extensions import ParamSpec, assert_type
 from ..exception import InvalidError
 from .logger import logger
 
+T = TypeVar("T")
+P = ParamSpec("P")
+V = TypeVar("V")
+
 synchronizer = synchronicity.Synchronizer()
 
 
@@ -260,7 +265,72 @@ def run_coro_blocking(coro):
         return fut.result()
 
 
-async def queue_batch_iterator(q: asyncio.Queue, max_batch_size=100, debounce_time=0.015):
+class TimestampPriorityQueue(Generic[T]):
+    """
+    A priority queue that schedules items to be processed at specific timestamps.
+    """
+
+    _MAX_PRIORITY = float("inf")
+
+    def __init__(self, maxsize: int = 0):
+        self.condition = asyncio.Condition()
+        self._queue: asyncio.PriorityQueue[tuple[float, Union[T, None]]] = asyncio.PriorityQueue(maxsize=maxsize)
+
+    async def close(self):
+        await self.put(self._MAX_PRIORITY, None)
+
+    async def put(self, timestamp: float, item: Union[T, None]):
+        """
+        Add an item to the queue to be processed at a specific timestamp.
+        """
+        await self._queue.put((timestamp, item))
+        async with self.condition:
+            self.condition.notify_all()  # notify any waiting coroutines
+
+    async def get(self) -> Union[T, None]:
+        """
+        Get the next item from the queue that is ready to be processed.
+        """
+        while True:
+            async with self.condition:
+                while self.empty():
+                    await self.condition.wait()
+                # peek at the next item
+                timestamp, item = await self._queue.get()
+                now = time.time()
+                if timestamp < now:
+                    return item
+                if timestamp == self._MAX_PRIORITY:
+                    return None
+                # not ready yet, calculate sleep time
+                sleep_time = timestamp - now
+                self._queue.put_nowait((timestamp, 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)
+                except asyncio.TimeoutError:
+                    continue
+
+    def empty(self) -> bool:
+        return self._queue.empty()
+
+    def qsize(self) -> int:
+        return self._queue.qsize()
+
+    async def clear(self):
+        """
+        Clear the retry queue. Used for testing to simulate reading all elements from queue using queue_batch_iterator.
+        """
+        while not self.empty():
+            await self.get()
+
+    def __len__(self):
+        return self._queue.qsize()
+
+
+async def queue_batch_iterator(
+    q: Union[asyncio.Queue, TimestampPriorityQueue], max_batch_size=100, debounce_time=0.015
+):
     """
     Read from a queue but return lists of items when queue is large
 
@@ -405,11 +475,6 @@ def on_shutdown(coro):
     _shutdown_tasks.append(asyncio.create_task(wrapper()))
 
 
-T = TypeVar("T")
-P = ParamSpec("P")
-V = TypeVar("V")
-
-
 def asyncify(f: Callable[P, T]) -> Callable[P, typing.Coroutine[None, None, T]]:
     """Convert a blocking function into one that runs in the current loop's executor."""