Add proto definitions for GPU fallbacks

modal-labs · Sep 3, 2024 · d5c16a8 · d5c16a8
1 parent 89d4a00
commit d5c16a8
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 9 deletions.
diff --git a/modal/functions.py b/modal/functions.py
@@ -820,14 +820,17 @@ async def _load(self: _Function, resolver: Resolver, existing_object_id: Optiona
                     scheduler_placement=scheduler_placement.proto if scheduler_placement else None,
                     is_class=info.is_service_class(),
                     class_parameter_info=info.class_parameter_info(),
-                    _experimental_resources=[
-                        convert_fn_config_to_resources_config(
-                            cpu=cpu, memory=memory, gpu=_experimental_gpu, ephemeral_disk=ephemeral_disk
+                    i6pn_enabled=config.get("i6pn_enabled"),
+                    _experimental_concurrent_cancellations=True,
+                    _experimental_task_templates=[
+                        api_pb2.TaskTemplate(
+                            priority=1,
+                            resource=convert_fn_config_to_resources_config(
+                                cpu=cpu, memory=memory, gpu=_experimental_gpu, ephemeral_disk=ephemeral_disk
+                            ),
                         )
                         for _experimental_gpu in _experimental_gpus
                     ],
-                    i6pn_enabled=config.get("i6pn_enabled"),
-                    _experimental_concurrent_cancellations=True,
                 )
                 assert resolver.app_id
                 request = api_pb2.FunctionCreateRequest(

diff --git a/modal_proto/api.proto b/modal_proto/api.proto
@@ -943,7 +943,6 @@ message FileEntry {
   uint64 size = 4;
 }
 
-
 message Function {
   string module_name = 1;
   string function_name = 2;
@@ -1045,15 +1044,17 @@ message Function {
 
   ClassParameterInfo class_parameter_info = 56;
 
-  repeated Resources _experimental_resources = 57; // overrides `resources` field above
-
+  reserved 57; // _experimental_resources
   reserved 58;
   reserved 59;
   uint32 batch_max_size = 60; // Maximum number of inputs to fetch at once
   uint64 batch_linger_ms = 61; // Miliseconds to block before a response is needed
   bool i6pn_enabled = 62;
   bool _experimental_concurrent_cancellations = 63;
   uint32 max_concurrent_inputs = 64;
+
+  bool _experimental_task_templates_enabled = 65;  // forces going through the new gpu-fallbacks integration path, even if no fallback options are specified
+  repeated TaskTemplate _experimental_task_templates = 66;  // for fallback options, where the first/most-preferred "template" is derived from fields above
 }
 
 message FunctionBindParamsRequest {
@@ -1762,7 +1763,7 @@ message Sandbox {
   // to look at fine-grained placement constraints.
   reserved 16; // _experimental_scheduler
   optional SchedulerPlacement scheduler_placement = 17;
-  repeated Resources _experimental_resources = 18; // overrides `resources` field above
+  reserved 18;  // _experimental_resources
 
   string worker_id = 19; // for internal debugging use only
   oneof open_ports_oneof {
@@ -2072,6 +2073,11 @@ message TaskStats {
   double started_at = 4;
 }
 
+message TaskTemplate {
+  uint32 priority = 1;
+  Resources resources = 2;
+  uint32 concurrent_inputs = 3;
+}
 
 message TokenFlowCreateRequest {
   string utm_source = 3;