feat(gossip, login): separate gossip pkg, implement new login/gossip operations (#641)

@eahydra @cardyok Now implementing very basic login + gossip.

Needs a lot of cleanups later.

Signed-off-by: Gyuho Lee <[email protected]>

Author: gyuho

api/v1/gossip.go

@@ -0,0 +1,14 @@
+package v1
+
+// GossipRequest is the request for the gossip request.
+type GossipRequest struct {
+	MachineID     string   `json:"machineID"`
+	DaemonVersion string   `json:"daemonVersion"`
+	Components    []string `json:"components"`
+}
+
+// GossipResponse is the response for the gossip request.
+type GossipResponse struct {
+	Error  string `json:"error"`
+	Status string `json:"status"`
+}

cmd/gpud/command/command.go

@@ -1,6 +1,7 @@
 package command
 
 import (
+	"errors"
 	"fmt"
 	"time"
 
@@ -23,7 +24,6 @@ var (
 	logFile  string
 
 	statusWatch bool
-	uid         string
 
 	annotations   string
 	listenAddress string
@@ -56,6 +56,10 @@ const (
 	warningSign = "\033[31m✘\033[0m"
 )
 
+var (
+	ErrEmptyToken = errors.New("token is empty")
+)
+
 func App() *cli.App {
 	app := cli.NewApp()
 
@@ -183,11 +187,6 @@ sudo rm /etc/systemd/system/gpud.service
 					Usage:       "set the annotations",
 					Destination: &annotations,
 				},
-				cli.StringFlag{
-					Name:        "uid",
-					Usage:       "uid for this machine",
-					Destination: &uid,
-				},
 				&cli.BoolFlag{
 					Name:        "pprof",
 					Usage:       "enable pprof (default: false)",
@@ -388,11 +387,6 @@ sudo rm /etc/systemd/system/gpud.service
 					Usage: "endpoint for control plane",
 					Value: "mothership-machine.app.lepton.ai",
 				},
-				cli.StringFlag{
-					Name:        "uid",
-					Usage:       "uid for this machine",
-					Destination: &uid,
-				},
 			},
 			Subcommands: []cli.Command{
 				{

cmd/gpud/command/login.go

@@ -3,20 +3,31 @@ package command
 import (
 	"context"
 	"fmt"
-	"os"
 	"time"
 
+	"github.com/urfave/cli"
+
+	apiv1 "github.com/leptonai/gpud/api/v1"
 	client "github.com/leptonai/gpud/client/v1"
 	"github.com/leptonai/gpud/pkg/config"
-	gpud_state "github.com/leptonai/gpud/pkg/gpud-state"
+	gpudstate "github.com/leptonai/gpud/pkg/gpud-state"
 	"github.com/leptonai/gpud/pkg/login"
 	"github.com/leptonai/gpud/pkg/server"
 	"github.com/leptonai/gpud/pkg/sqlite"
-
-	"github.com/urfave/cli"
 )
 
 func cmdLogin(cliContext *cli.Context) error {
+	token := cliContext.String("token")
+	if token == "" {
+		fmt.Print("Please visit https://dashboard.lepton.ai/ under Settings/Tokens to fetch your token\nPlease enter your token:")
+		if _, err := fmt.Scanln(&token); err != nil && err.Error() != "unexpected newline" {
+			return fmt.Errorf("failed reading input: %w", err)
+		}
+	}
+	if token == "" {
+		return ErrEmptyToken
+	}
+
 	rootCtx, rootCancel := context.WithTimeout(context.Background(), 2*time.Minute)
 	defer rootCancel()
 
@@ -31,75 +42,54 @@ func cmdLogin(cliContext *cli.Context) error {
 	if err != nil {
 		return fmt.Errorf("failed to get state file: %w", err)
 	}
-
-	dbRW, err := sqlite.Open(stateFile)
-	if err != nil {
-		return fmt.Errorf("failed to open state file: %w", err)
-	}
-	defer dbRW.Close()
-
 	dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
 	if err != nil {
 		return fmt.Errorf("failed to open state file: %w", err)
 	}
 	defer dbRO.Close()
 
-	uid, err := gpud_state.CreateMachineIDIfNotExist(rootCtx, dbRW, dbRO, "")
+	machineID, err := gpudstate.ReadMachineID(rootCtx, dbRO)
 	if err != nil {
-		return fmt.Errorf("failed to get machine uid: %w", err)
+		return err
 	}
-
-	components, err := gpud_state.GetComponents(rootCtx, dbRO, uid)
-	if err != nil {
-		return fmt.Errorf("failed to get components: %w", err)
+	if machineID != "" {
+		fmt.Printf("machine ID %s already assigned (skipping login)\n", machineID)
+		return nil
 	}
 
-	cliToken := cliContext.String("token")
 	endpoint := cliContext.String("endpoint")
 
-	dbToken, _ := gpud_state.GetLoginInfo(rootCtx, dbRO, uid)
-	token := dbToken
-	if cliToken != "" {
-		token = cliToken
-	} else {
-		fmt.Println("trying token from local store, if you want to override, use --token flag")
+	// machine ID has not been assigned yet
+	// thus request one and blocks until the login request is processed
+	loginResp, err := login.SendRequest(rootCtx, endpoint, apiv1.LoginRequest{Token: token})
+	if err != nil {
+		return err
 	}
+	machineID = loginResp.MachineID
 
-	if token == "" {
-		fmt.Print("Please visit https://dashboard.lepton.ai/ under Settings/Tokens to fetch your token\nPlease enter your token:")
-		if _, err := fmt.Scanln(&token); err != nil && err.Error() != "unexpected newline" {
-			return fmt.Errorf("failed reading input: %w", err)
-		}
+	// consume the login response to persist the machine ID
+	dbRW, err := sqlite.Open(stateFile)
+	if err != nil {
+		return fmt.Errorf("failed to open state file: %w", err)
+	}
+	defer dbRW.Close()
+	if err := gpudstate.RecordMachineID(rootCtx, dbRW, dbRO, machineID); err != nil {
+		return fmt.Errorf("failed to record machine ID: %w", err)
 	}
 
 	fifoFile, err := config.DefaultFifoFile()
 	if err != nil {
 		return fmt.Errorf("failed to get fifo file: %w", err)
 	}
 
-	if token != "" && endpoint != "" {
-		hostname, err := os.Hostname()
-		if err != nil {
-			hostname = "UnknownName"
-		}
-		if err := login.Login(hostname, token, endpoint, components, uid); err != nil {
-			return err
-		}
-	} else {
-		fmt.Println("login skipped since token or endpoint not provided...")
-		return nil
-	}
-
 	if err := server.WriteToken(token, fifoFile); err != nil {
 		return fmt.Errorf("failed to write token: %v", err)
 	}
 
-	if token != dbToken {
-		if err = gpud_state.UpdateLoginInfo(rootCtx, dbRW, uid, token); err != nil {
-			fmt.Println("machine logged in but failed to update token:", err)
-		}
+	if err = gpudstate.UpdateLoginInfo(rootCtx, dbRW, machineID, token); err != nil {
+		fmt.Println("machine logged in but failed to update token:", err)
 	}
 
-	fmt.Printf("%s successfully logged into lepton.ai\n", checkMark)
+	fmt.Printf("%s successfully logged in with machine id %s\n", checkMark, loginResp.MachineID)
 	return nil
 }

cmd/gpud/command/run.go

@@ -16,8 +16,10 @@ import (
 
 	"github.com/leptonai/gpud/pkg/config"
 	gpud_manager "github.com/leptonai/gpud/pkg/gpud-manager"
+	gpudstate "github.com/leptonai/gpud/pkg/gpud-state"
 	"github.com/leptonai/gpud/pkg/log"
 	lepServer "github.com/leptonai/gpud/pkg/server"
+	"github.com/leptonai/gpud/pkg/sqlite"
 	pkd_systemd "github.com/leptonai/gpud/pkg/systemd"
 	"github.com/leptonai/gpud/version"
 )
@@ -97,6 +99,23 @@ func cmdRun(cliContext *cli.Context) error {
 	}
 	m.Start(rootCtx)
 
+	stateFile, err := config.DefaultStateFile()
+	if err != nil {
+		return fmt.Errorf("failed to get state file: %w", err)
+	}
+	dbRO, err := sqlite.Open(stateFile, sqlite.WithReadOnly(true))
+	if err != nil {
+		return fmt.Errorf("failed to open state file: %w", err)
+	}
+	defer dbRO.Close()
+	uid, err := gpudstate.ReadMachineID(rootCtx, dbRO)
+	if err != nil {
+		return fmt.Errorf("failed to read machine ID: %w", err)
+	}
+	if uid == "" {
+		log.Logger.Warnw("machine ID not found, running in local mode not connected to any control plane")
+	}
+
 	server, err := lepServer.New(rootCtx, cfg, cliContext.String("endpoint"), uid, m)
 	if err != nil {
 		return err

cmd/gpud/command/utils.go

@@ -24,5 +24,5 @@ func GetUID(ctx context.Context) (string, error) {
 	}
 	defer dbRO.Close()
 
-	return gpud_state.GetMachineID(ctx, dbRO)
+	return gpud_state.ReadMachineID(ctx, dbRO)
 }

pkg/config/default.go

@@ -244,14 +244,6 @@ func setupDefaultDir() (string, error) {
 	return d, nil
 }
 
-func DefaultConfigFile() (string, error) {
-	dir, err := setupDefaultDir()
-	if err != nil {
-		return "", err
-	}
-	return filepath.Join(dir, "gpud.yaml"), nil
-}
-
 func DefaultStateFile() (string, error) {
 	dir, err := setupDefaultDir()
 	if err != nil {

pkg/gossip/client.go

@@ -0,0 +1,69 @@
+package gossip
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+
+	apiv1 "github.com/leptonai/gpud/api/v1"
+	"github.com/leptonai/gpud/pkg/log"
+)
+
+// SendRequest sends a gossip request.
+func SendRequest(ctx context.Context, endpoint string, req apiv1.GossipRequest) (*apiv1.GossipResponse, error) {
+	url := createURL(endpoint)
+	return sendRequest(ctx, url, req)
+}
+
+func sendRequest(ctx context.Context, url string, req apiv1.GossipRequest) (*apiv1.GossipResponse, error) {
+	if os.Getenv("GPUD_NO_USAGE_STATS") == "true" {
+		log.Logger.Debug("gossip skipped since GPUD_NO_USAGE_STATS=true specified")
+		return nil, nil
+	}
+
+	log.Logger.Infow("sending gossip request", "url", url)
+
+	b, err := json.Marshal(req)
+	if err != nil {
+		return nil, fmt.Errorf("error marshaling gossip request: %w", err)
+	}
+
+	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(b))
+	if err != nil {
+		return nil, fmt.Errorf("error creating request: %w", err)
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+
+	client := &http.Client{}
+	httpResp, err := client.Do(httpReq)
+	if err != nil {
+		return nil, err
+	}
+	defer httpResp.Body.Close()
+
+	body, err := io.ReadAll(httpResp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("error reading response body: %w", err)
+	}
+
+	var resp apiv1.GossipResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return nil, fmt.Errorf("error unmarshalling gossip response: %w", err)
+	}
+
+	if httpResp.StatusCode != http.StatusOK {
+		return &resp, fmt.Errorf("login request failed with status code %d (%+v)", httpResp.StatusCode, resp)
+	}
+
+	log.Logger.Infow("gossip request processed", "url", url)
+	return &resp, nil
+}
+
+// createURL creates a URL for the gossip endpoint
+func createURL(endpoint string) string {
+	return fmt.Sprintf("https://%s/api/v1/gossip", endpoint)
+}