Merge pull request #120 from depot/feat/gpu

feat: enable gpus
depot · Aug 6, 2024 · 2f0c622 · 2f0c622
2 parents 5646bbd + b622c1d
commit 2f0c622
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 0 deletions.
diff --git a/proto/depot/cloud/v3/machine.proto b/proto/depot/cloud/v3/machine.proto
@@ -95,6 +95,7 @@ message RegisterMachineResponse {
     optional string vector_config = 18;
     optional bool enable_cni = 19;
     optional bool use_buildkit_private = 20;
+    optional bool enable_gpu = 21;
   }
 
   // EngineTask represents an instruction to start an engine daemon

diff --git a/src/gen/ts/depot/cloud/v3/machine_pb.ts b/src/gen/ts/depot/cloud/v3/machine_pb.ts
@@ -569,6 +569,11 @@ export class RegisterMachineResponse_BuildKitTask extends Message<RegisterMachin
    */
   useBuildkitPrivate?: boolean
 
+  /**
+   * @generated from field: optional bool enable_gpu = 21;
+   */
+  enableGpu?: boolean
+
   constructor(data?: PartialMessage<RegisterMachineResponse_BuildKitTask>) {
     super()
     proto3.util.initPartial(data, this)
@@ -596,6 +601,7 @@ export class RegisterMachineResponse_BuildKitTask extends Message<RegisterMachin
     {no: 18, name: 'vector_config', kind: 'scalar', T: 9 /* ScalarType.STRING */, opt: true},
     {no: 19, name: 'enable_cni', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
     {no: 20, name: 'use_buildkit_private', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
+    {no: 21, name: 'enable_gpu', kind: 'scalar', T: 8 /* ScalarType.BOOL */, opt: true},
   ])
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): RegisterMachineResponse_BuildKitTask {

diff --git a/src/tasks/buildkit.ts b/src/tasks/buildkit.ts
@@ -193,6 +193,11 @@ keepBytes = ${cacheSizeBytes}
     env.DEPOT_RESOLVER_CONCURRENCY = task.resolverConcurrency.toString()
   }
 
+  if (task.enableGpu) {
+    console.log('Enabling GPU')
+    env.DEPOT_ENABLE_GPU = 'true'
+  }
+
   const args = task.enableDebugLogging ? ['--debug'] : []
 
   async function runBuildKit() {