From ad1a298486490066ec831ff5031ad7db8068a25d Mon Sep 17 00:00:00 2001
From: Deven Navani <deven.navani@gmail.com>
Date: Mon, 2 Sep 2024 17:43:53 +0000
Subject: [PATCH] Add _experimental_buffer_containers

---
 modal/app.py          | 2 ++
 modal/functions.py    | 2 ++
 modal_proto/api.proto | 1 +
 3 files changed, 5 insertions(+)

diff --git a/modal/app.py b/modal/app.py
index 057911675..b0f5d1393 100644
--- a/modal/app.py
+++ b/modal/app.py
@@ -580,6 +580,7 @@ def function(
             SchedulerPlacement
         ] = None,  # Experimental controls over fine-grained scheduling (alpha).
         _experimental_gpus: Sequence[GPU_T] = [],  # Experimental controls over GPU fallbacks (alpha).
+        _experimental_buffer_containers: Optional[int] = None,  # Number of additional, idle containers to keep around.
     ) -> Callable[[Union[Callable[P, R], _PartialFunction[P, R]]], _Function[P, R]]:
         """Decorator to register a new Modal function with this app."""
         if isinstance(_warn_parentheses_missing, _Image):
@@ -711,6 +712,7 @@ def f(self, x):
                 scheduler_placement=scheduler_placement,
                 _experimental_boost=_experimental_boost,
                 _experimental_gpus=_experimental_gpus,
+                _experimental_buffer_containers=_experimental_buffer_containers,
             )
 
             self._add_function(function, webhook_config is not None)
diff --git a/modal/functions.py b/modal/functions.py
index 9a00e635b..4174e6b41 100644
--- a/modal/functions.py
+++ b/modal/functions.py
@@ -508,6 +508,7 @@ def from_args(
         max_inputs: Optional[int] = None,
         ephemeral_disk: Optional[int] = None,
         _experimental_gpus: Sequence[GPU_T] = [],
+        _experimental_buffer_containers: Optional[int] = None,
     ) -> None:
         """mdmd:hidden"""
         tag = info.get_tag()
@@ -828,6 +829,7 @@ async def _load(self: _Function, resolver: Resolver, existing_object_id: Optiona
                     ],
                     i6pn_enabled=config.get("i6pn_enabled"),
                     _experimental_concurrent_cancellations=True,
+                    _experimental_buffer_containers=_experimental_buffer_containers,
                 )
                 assert resolver.app_id
                 request = api_pb2.FunctionCreateRequest(
diff --git a/modal_proto/api.proto b/modal_proto/api.proto
index fef96710e..08f9eac59 100644
--- a/modal_proto/api.proto
+++ b/modal_proto/api.proto
@@ -1054,6 +1054,7 @@ message Function {
   bool i6pn_enabled = 62;
   bool _experimental_concurrent_cancellations = 63;
   uint32 max_concurrent_inputs = 64;
+  uint32 _experimental_buffer_containers = 65;
 }
 
 message FunctionBindParamsRequest {