Feat: Add user data to worker pools (#724)

Resolve BE-2035
beam-cloud · Nov 21, 2024 · e211a3b · e211a3b
1 parent 18ce8fb
commit e211a3b
Show file tree

Hide file tree

Showing 6 changed files with 368 additions and 351 deletions.
diff --git a/pkg/gateway/gateway.proto b/pkg/gateway/gateway.proto
@@ -356,6 +356,7 @@ message Machine {
   string created = 13;
   string agent_version = 14;
   MachineMetrics machine_metrics = 15;
+  string user_data = 16;
 }
 
 message MachineMetrics {

diff --git a/pkg/gateway/services/machine.go b/pkg/gateway/services/machine.go
@@ -190,6 +190,7 @@ func (gws *GatewayService) CreateMachine(ctx context.Context, in *pb.CreateMachi
 			ProviderName:      string(*pool.Provider),
 			TailscaleUrl:      gws.appConfig.Tailscale.ControlURL,
 			TailscaleAuth:     gws.appConfig.Tailscale.AuthKey,
+			UserData:          pool.UserData,
 		},
 	}, nil
 }

diff --git a/pkg/types/config.go b/pkg/types/config.go
@@ -220,6 +220,7 @@ type WorkerPoolConfig struct {
 	RequiresPoolSelector bool                              `key:"requiresPoolSelector" json:"requires_pool_selector"`
 	Priority             int32                             `key:"priority" json:"priority"`
 	Preemptable          bool                              `key:"preemptable" json:"preemptable"`
+	UserData             string                            `key:"userData" json:"user_data"`
 }
 
 type WorkerPoolJobSpecConfig struct {