Skip to content

Commit 857052c

Browse files
committed
Add per-job GPU support via workflow labels
Parse gpu:T4/gpu:A100/etc from job labels and pass to modal.Sandbox.create(). Supports T4, L4, A100, A100-80GB, H100. GPU defaults to None (CPU-only) if no gpu: label is present.
1 parent ac67438 commit 857052c

1 file changed

Lines changed: 27 additions & 1 deletion

File tree

app.py

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,14 @@ def format(self, record: logging.LogRecord) -> str:
9393
if MAX_CONCURRENT_PER_REPO <= 0:
9494
MAX_CONCURRENT_PER_REPO = None
9595

96+
GPU_MAP = {
97+
"t4": modal.gpu.T4(),
98+
"l4": modal.gpu.L4(),
99+
"a100": modal.gpu.A100(),
100+
"a100-80gb": modal.gpu.A100_80GB(),
101+
"h100": modal.gpu.H100(),
102+
}
103+
96104
# =============================================================================
97105
# TRUST MODEL
98106
# =============================================================================
@@ -476,6 +484,23 @@ async def github_webhook(request: Request):
476484
)
477485
return {"status": "ignored", "reason": "no modal label"}
478486

487+
gpu_config = None
488+
for label in job_labels:
489+
if label.startswith("gpu:"):
490+
gpu_key = label.split(":", 1)[1].lower()
491+
gpu_config = GPU_MAP.get(gpu_key)
492+
if not gpu_config:
493+
logger.warning(
494+
"Unknown GPU type requested",
495+
extra={"job_id": job_id, "gpu": gpu_key},
496+
)
497+
else:
498+
logger.info(
499+
"GPU requested for job",
500+
extra={"job_id": job_id, "gpu": gpu_key},
501+
)
502+
break
503+
479504
# Repository allowlist validation
480505
if ALLOWED_REPOS and repo_full_name not in ALLOWED_REPOS:
481506
logger.warning(
@@ -618,7 +643,7 @@ async def github_webhook(request: Request):
618643
)
619644
raise HTTPException(status_code=500, detail="Internal server error")
620645

621-
logger.info("Spawning sandbox", extra={"job_id": job_id, "repo": repo_full_name})
646+
logger.info("Spawning sandbox", extra={"job_id": job_id, "repo": repo_full_name, "gpu": gpu_config is not None})
622647

623648
cmd = (
624649
"bash -c '/start-dockerd.sh &'"
@@ -637,6 +662,7 @@ async def github_webhook(request: Request):
637662
app=app,
638663
timeout=TIMEOUT_SECONDS,
639664
env={"GHA_JIT_CONFIG": jit_config},
665+
gpu=gpu_config,
640666
experimental_options={"enable_docker": True},
641667
)
642668

0 commit comments

Comments
 (0)