pass kwargs to custom_task

latchbio · Feb 6, 2024 · b547413 · b547413
1 parent f8fcd3b
commit b547413
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/latch/resources/tasks.py b/latch/resources/tasks.py
@@ -361,7 +361,7 @@ def custom_memory_optimized_task(cpu: int, memory: int):
     return functools.partial(task, task_config=task_config)
 
 
-def custom_task(cpu: int, memory: int, *, storage_gib: int = 500):
+def custom_task(cpu: int, memory: int, *, storage_gib: int = 500, **kwargs):
     """Returns a custom task configuration requesting
     the specified CPU/RAM allocations
 
@@ -461,4 +461,4 @@ def custom_task(cpu: int, memory: int, *, storage_gib: int = 500):
                 " 4949 GiB)"
             )
 
-    return functools.partial(task, task_config=task_config)
+    return functools.partial(task, task_config=task_config, **kwargs)