Hotfix to allow re-registering of ttld machines (#810)

beam-cloud · Dec 26, 2024 · c789b7e · c789b7e
1 parent b36b742
commit c789b7e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/pkg/api/v1/machine.go b/pkg/api/v1/machine.go
@@ -83,14 +83,19 @@ func (g *MachineGroup) RegisterMachine(ctx echo.Context) error {
 		hostName = fmt.Sprintf("%s.%s.%s", request.HostName, g.config.Tailscale.User, g.config.Tailscale.HostName)
 	}
 
+	poolConfig, ok := g.config.Worker.Pools[request.PoolName]
+	if !ok {
+		return HTTPInternalServerError("Invalid pool name")
+	}
+
 	err = g.providerRepo.RegisterMachine(request.ProviderName, request.PoolName, request.MachineID, &types.ProviderMachineState{
 		MachineId: request.MachineID,
 		Token:     request.Token,
 		HostName:  hostName,
 		Cpu:       cpu,
 		Memory:    memory,
 		GpuCount:  uint32(gpuCount),
-	})
+	}, &poolConfig)
 	if err != nil {
 		return HTTPInternalServerError("Failed to register machine")
 	}

diff --git a/pkg/repository/base.go b/pkg/repository/base.go
@@ -158,7 +158,7 @@ type ProviderRepository interface {
 	RemoveMachine(providerName, poolName, machineId string) error
 	SetMachineKeepAlive(providerName, poolName, machineId, agentVersion string, metrics *types.ProviderMachineMetrics) error
 	SetLastWorkerSeen(providerName, poolName, machineId string) error
-	RegisterMachine(providerName, poolName, machineId string, newMachineInfo *types.ProviderMachineState) error
+	RegisterMachine(providerName, poolName, machineId string, newMachineInfo *types.ProviderMachineState, poolConfig *types.WorkerPoolConfig) error
 	WaitForMachineRegistration(providerName, poolName, machineId string) (*types.ProviderMachineState, error)
 	ListAllMachines(providerName, poolName string, useLock bool) ([]*types.ProviderMachine, error)
 	SetMachineLock(providerName, poolName, machineId string) error

diff --git a/pkg/repository/provider_redis.go b/pkg/repository/provider_redis.go
@@ -299,12 +299,20 @@ func (r *ProviderRedisRepository) RemoveMachine(providerName, poolName, machineI
 	return nil
 }
 
-func (r *ProviderRedisRepository) RegisterMachine(providerName, poolName, machineId string, newMachineInfo *types.ProviderMachineState) error {
+func (r *ProviderRedisRepository) RegisterMachine(providerName, poolName, machineId string, newMachineInfo *types.ProviderMachineState, poolConfig *types.WorkerPoolConfig) error {
 	stateKey := common.RedisKeys.ProviderMachineState(providerName, poolName, machineId)
 
 	machineInfo, err := r.getMachineStateFromKey(stateKey)
 	if err != nil {
-		return fmt.Errorf("failed to get machine state <%v>: %w", stateKey, err)
+		// TODO: This is a temporary fix to allow the machine to be registered
+		// without having to update the machine state, in the future we should tie
+		// registration token to machine ID and store that somewhere else persistently
+		machineInfo = &types.ProviderMachineState{}
+		machineInfo.Gpu = poolConfig.GPUType
+		machineInfo.Created = fmt.Sprintf("%d", time.Now().UTC().Unix())
+		machineInfo.LastKeepalive = fmt.Sprintf("%d", time.Now().UTC().Unix())
+		machineInfo.PoolName = newMachineInfo.PoolName
+		machineInfo.MachineId = newMachineInfo.MachineId
 	}
 
 	machineInfo.HostName = newMachineInfo.HostName