Prespecified Task Resource

The Latch SDK currently supports a set of prespecified task resource requests represented as decorators:

  • small_task: 2 cpus, 4 gigs of memory, 0 gpus
  • medium_task: 32 cpus, 128 gigs of memory, 0 gpus
  • large_task: 96 cpus, 192 gig sof memory, 0 gpus
  • small_gpu_task: 8 cpus, 32 gigs of memory, 1 gpu (24 gigs of VRAM, 9,216 cuda cores)
  • large_gpu_task: 31 cpus, 120 gigs of memory, 1 gpu (24 gigs of VRAM, 9,216 cuda cores)

We use the tasks as follows:

from latch import small_task, large_gpu_task

@small_task
def my_task(
    ...
):
    ...

@large_gpu_task
def inference(
    ...
):
    ...

Custom Task Resource

You can also arbitrarily specify task resources using @custom_task:

from latch import custom_task

@custom_task(cpu, memory) # cpu: int, memory: int
def my_task(
    ...
):
    ...

Dynamic Task Resource

You can dynamically define task resources based on the tasks’ input parameters by passing functions as arguments for the custom_task decorator. The provided functions will execute at runtime, and the task will launch with the resulting resource values:

...

from latch import custom_task
from latch.types.file import LatchFile

def allocate_cpu(files: List[LatchFile], **kwargs) -> int: # number of cores to allocate
    return min(8, len(files))

def allocate_storage(files: List[LatchFile], **kwargs) -> int: # GiBs of storage to allocate
    return sum([1.5 * file.size() for file in files]) // 1024**3

@custom_task(cpu=allocate_cpu, memory=8, storage=allocate_storage)
def my_task(files: List[LatchFile], count: int):
    ...

In the provided example, the allocate_cpu function is designed to process the input parameters files. Upon execution, the function returns an integer representing the total number of CPU cores that should be allocated to the task based on the input file size.

The parameters passed to the resource functions at runtime are the same as those passed to the task function. Therefore, the resource functions must only accept parameters that exist in the task function signature. See how the my_task function and the allocate_cpu function both accept a parameter named files of type List[LatchFile].