Merge pull request #515 from latchbio/aidan/add-l40s-gpu-type

feat: add l40s gpu type task decorator

CHANGELOG.md

@@ -16,6 +16,16 @@ Types of changes
 
 # Latch SDK Changelog
 
+## 2.55.1 - 2024-12-14
+
+### Added
+
+* Add `g6e_xlarge_task`, `g6e_2xlarge_task`, `g6e_4xlarge_task`, `g6e_8xlarge_task`, `g6e_12xlarge_task`, `g6e_16xlarge_task`, and `g6e_24xlarge_task` to request G6e GPU instances.
+
+### Updated
+
+* Update latch base image to 'cb01-main'
+
 ## 2.55.0 - 2024-12-13
 
 ### Changed

pyproject.toml

@@ -12,7 +12,7 @@ include = ["src/**/*.py", "src/latch_cli/services/init/*"]
 
 [project]
 name = "latch"
-version = "2.55.0"
+version = "2.55.1"
 description = "The Latch SDK"
 authors = [{ name = "Kenny Workman", email = "[email protected]" }]
 maintainers = [

src/latch/__init__.py

@@ -20,6 +20,13 @@
 from latch.resources.tasks import (
     custom_memory_optimized_task,
     custom_task,
+    g6e_xlarge_task,
+    g6e_2xlarge_task,
+    g6e_4xlarge_task,
+    g6e_8xlarge_task,
+    g6e_12xlarge_task,
+    g6e_16xlarge_task,
+    g6e_24xlarge_task,
     large_gpu_task,
     large_task,
     medium_task,

src/latch/resources/tasks.py

@@ -627,3 +627,94 @@ def nextflow_runtime_task(cpu: int, memory: int, storage_gib: int = 50):
     ]
 
     return functools.partial(task, task_config=task_config)
+
+
+def _get_l40s_pod(instance_type: str, cpu: int, memory_gib: int, gpus: int) -> Pod:
+    """Helper function to create L40s GPU pod configurations."""
+    primary_container = V1Container(name="primary")
+
+    backoff_cpu = cpu - 2
+    backoff_memory = min(memory_gib - 4, int(memory_gib * 0.9))
+
+    resources = V1ResourceRequirements(
+        requests={
+            "cpu": str(backoff_cpu),
+            "memory": f"{backoff_memory}Gi",
+            "nvidia.com/gpu": str(gpus),
+            "ephemeral-storage": "4500Gi",
+        },
+        limits={
+            "cpu": str(cpu),
+            "memory": f"{memory_gib}Gi",
+            "nvidia.com/gpu": str(gpus),
+            "ephemeral-storage": "5000Gi",
+        },
+    )
+    primary_container.resources = resources
+
+    return Pod(
+        pod_spec=V1PodSpec(
+            containers=[primary_container],
+            tolerations=[
+                V1Toleration(
+                    effect="NoSchedule",
+                    key="ng",
+                    value=instance_type
+                )
+            ],
+        ),
+        primary_container_name="primary",
+        annotations={
+            "cluster-autoscaler.kubernetes.io/safe-to-evict": "false",
+        },
+    )
+
+
+g6e_xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-xlarge", cpu=4, memory_gib=32, gpus=1)
+)
+"""4 vCPUs, 32 GiB RAM, 1 L40s GPU"""
+
+g6e_2xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-2xlarge", cpu=8, memory_gib=64, gpus=1)
+)
+"""8 vCPUs, 64 GiB RAM, 1 L40s GPU"""
+
+g6e_4xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-4xlarge", cpu=16, memory_gib=128, gpus=1)
+)
+"""16 vCPUs, 128 GiB RAM, 1 L40s GPU"""
+
+g6e_8xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-8xlarge", cpu=32, memory_gib=256, gpus=1)
+)
+"""32 vCPUs, 256 GiB RAM, 1 L40s GPU"""
+
+g6e_12xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-12xlarge", cpu=48, memory_gib=384, gpus=1)
+)
+"""48 vCPUs, 384 GiB RAM, 1 L40s GPU"""
+
+g6e_16xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-16xlarge", cpu=64, memory_gib=512, gpus=1)
+)
+"""64 vCPUs, 512 GiB RAM, 1 L40s GPUs"""
+
+g6e_24xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-24xlarge", cpu=96, memory_gib=768, gpus=4)
+)
+"""96 vCPUs, 768 GiB RAM, 4 L40s GPUs"""
+
+"""
+g6e_48xlarge_task = functools.partial(
+    task,
+    task_config=_get_l40s_pod("g6e-48xlarge", cpu=192, memory_gib=1536, gpus=8)
+)
+192 vCPUs, 1536 GiB RAM, 8 L40s GPUs"""

src/latch_cli/constants.py

@@ -20,7 +20,7 @@ class Units(int, Enum):
 @dataclass(frozen=True)
 class LatchConstants:
     base_image: str = (
-        "812206152185.dkr.ecr.us-west-2.amazonaws.com/latch-base:fe0b-main"
+        "812206152185.dkr.ecr.us-west-2.amazonaws.com/latch-base:cb01-main"
     )
     nextflow_latest_version: str = "v2.3.0"