From 0cec0649291ee5961477acbe7ca38d67e6d366ca Mon Sep 17 00:00:00 2001
From: Zach Langbert <zach@geode.io>
Date: Mon, 18 Aug 2025 16:54:17 -0400
Subject: [PATCH] add basic gpu support

---
 internal/commands/deploy.go          | 11 +++++++++++
 internal/commands/deployment.go      | 23 +++++++++++++++++++++++
 internal/commands/deployment_test.go |  1 +
 3 files changed, 35 insertions(+)

diff --git a/internal/commands/deploy.go b/internal/commands/deploy.go
index 336ba31..2fd25b2 100644
--- a/internal/commands/deploy.go
+++ b/internal/commands/deploy.go
@@ -27,6 +27,7 @@ var Deploy = &cli.Command{
 		containerPortFlag,
 		requestedMemoryFlag,
 		requestedCPUFlag,
+		requestedGPUFlag,
 		additionalContainerPortsFlag,
 		envVarsFlag,
 		idleTimeoutFlag,
@@ -67,6 +68,8 @@ var Deploy = &cli.Command{
 			deploymentTag = &deploy.DeploymentTag
 		}
 
+		gpu := float64(deploy.RequestedGPU)
+
 		res, err := deploy.SDK.DeploymentsV3.CreateDeployment(
 			ctx,
 			components.DeploymentConfigV3{
@@ -77,6 +80,7 @@ var Deploy = &cli.Command{
 				ContainerPort:            deploy.ContainerPort,
 				RequestedMemoryMB:        deploy.RequestedMemoryMB,
 				RequestedCPU:             deploy.RequestedCPU,
+				ExperimentalRequestedGPU: &gpu, // TODO: add support for final API field when it exists
 				AdditionalContainerPorts: deploy.AdditionalContainerPorts,
 				Env:                      deploy.Env,
 				DeploymentTag:            deploymentTag,
@@ -149,6 +153,13 @@ func (c *DeployConfig) Merge(latest *components.DeploymentV3, isIdleTimeoutDefau
 		c.RequestedCPU = latest.RequestedCPU
 	}
 
+	if c.RequestedGPU == 0 {
+		// TODO: add support for final API field when it exists
+		if latest.ExperimentalRequestedGPU != nil {
+			c.RequestedGPU = int64(*latest.ExperimentalRequestedGPU)
+		}
+	}
+
 	if len(c.AdditionalContainerPorts) == 0 {
 		c.AdditionalContainerPorts = latest.AdditionalContainerPorts
 	}
diff --git a/internal/commands/deployment.go b/internal/commands/deployment.go
index 0bcfb65..66b6ab7 100644
--- a/internal/commands/deployment.go
+++ b/internal/commands/deployment.go
@@ -171,6 +171,8 @@ var Deployment = &cli.Command{
 					deploymentTag = &deployment.DeploymentTag
 				}
 
+				gpu := float64(deployment.RequestedGPU)
+
 				res, err := deployment.SDK.DeploymentsV3.CreateDeployment(
 					ctx,
 					components.DeploymentConfigV3{
@@ -181,6 +183,7 @@ var Deployment = &cli.Command{
 						ContainerPort:            deployment.ContainerPort,
 						RequestedMemoryMB:        deployment.RequestedMemoryMB,
 						RequestedCPU:             deployment.RequestedCPU,
+						ExperimentalRequestedGPU: &gpu, // TODO: add support for final API field when it exists
 						AdditionalContainerPorts: deployment.AdditionalContainerPorts,
 						Env:                      deployment.Env,
 						DeploymentTag:            deploymentTag,
@@ -297,6 +300,17 @@ var (
 		Category:   "Deployment:",
 	}
 
+	requestedGPUFlag = &workaround.FloatFlag{
+		Name: "requested-gpu",
+		Sources: cli.NewValueSourceChain(
+			cli.EnvVar(deploymentEnvVar("REQUESTED_GPU")),
+			altsrc.ConfigFile(configFlag.Name, "deployment.requested-gpu"),
+		),
+		Usage:      "`<gpus>` to allocate to your process",
+		Persistent: true,
+		Category:   "Deployment:",
+	}
+
 	fromLatestFlag = &cli.BoolFlag{
 		Name:     "from-latest",
 		Sources:  cli.EnvVars(deploymentEnvVar("FROM_LATEST")),
@@ -457,6 +471,7 @@ type CreateDeploymentConfig struct {
 	ContainerPort            int
 	RequestedMemoryMB        float64
 	RequestedCPU             float64
+	RequestedGPU             int64
 	AdditionalContainerPorts []components.ContainerPort
 	Env                      []components.DeploymentConfigV3Env
 	DeploymentTag            string
@@ -490,6 +505,7 @@ func (c *CreateDeploymentConfig) Load(cmd *cli.Command) error {
 	c.ContainerPort = int(cmd.Int(containerPortFlag.Name))
 	c.RequestedMemoryMB = cmd.Float(requestedMemoryFlag.Name)
 	c.RequestedCPU = cmd.Float(requestedCPUFlag.Name)
+	c.RequestedGPU = cmd.Int(requestedGPUFlag.Name)
 	c.DeploymentTag = cmd.String(deploymentTagFlag.Name)
 
 	addlPorts := cmd.StringSlice(additionalContainerPortsFlag.Name)
@@ -542,6 +558,13 @@ func (c *CreateDeploymentConfig) Merge(latest *components.DeploymentV3) {
 		c.RequestedCPU = latest.RequestedCPU
 	}
 
+	if c.RequestedGPU == 0 {
+		// TODO: add support for final API field when it exists
+		if latest.ExperimentalRequestedGPU != nil {
+			c.RequestedGPU = int64(*latest.ExperimentalRequestedGPU)
+		}
+	}
+
 	if len(c.AdditionalContainerPorts) == 0 {
 		c.AdditionalContainerPorts = latest.AdditionalContainerPorts
 	}
diff --git a/internal/commands/deployment_test.go b/internal/commands/deployment_test.go
index a89bc76..c66515a 100644
--- a/internal/commands/deployment_test.go
+++ b/internal/commands/deployment_test.go
@@ -196,6 +196,7 @@ func Test_Integration_DeploymentCommands_Happy(t *testing.T) {
 					"containerPort": 8000,
 					"requestedMemoryMB": 1024,
 					"requestedCPU": 0.5,
+					"experimentalRequestedGPU": 0,
 					"additionalContainerPorts": [
 						{
 							"transportType": "tcp",