feat(gpud): clean up --endpoint flag parsing, persist in systemd service file (#681)

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

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

Author: gyuho

.gitignore

@@ -26,3 +26,5 @@ go.work.sum
 /bin/
 
 .DS_Store
+
+coverage.html

cmd/gpud/command/join.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -30,6 +31,7 @@ import (
 func cmdJoin(cliContext *cli.Context) (retErr error) {
 	rootCtx, rootCancel := context.WithTimeout(context.Background(), 3*time.Minute)
 	defer rootCancel()
+
 	endpoint := cliContext.String("endpoint")
 	clusterName := cliContext.String("cluster-name")
 	provider := cliContext.String("provider")
@@ -174,7 +176,7 @@ func cmdJoin(cliContext *cli.Context) (retErr error) {
 		}
 	}
 	fmt.Println("Please wait while control plane is initializing basic setup for your machine, this may take up to one minute...")
-	response, err := http.Post(fmt.Sprintf("%s/api/v1/join", endpoint), "application/json", bytes.NewBuffer(rawPayload))
+	response, err := http.Post(createJoinURL(endpoint), "application/json", bytes.NewBuffer(rawPayload))
 	if err != nil {
 		return err
 	}
@@ -291,3 +293,13 @@ func runCommand(ctx context.Context, script string, result *string) error {
 	}
 	return nil
 }
+
+// createJoinURL creates a URL for the join endpoint
+func createJoinURL(endpoint string) string {
+	host := endpoint
+	url, _ := url.Parse(endpoint)
+	if url.Host != "" {
+		host = url.Host
+	}
+	return fmt.Sprintf("https://%s/api/v1/join", host)
+}

cmd/gpud/command/notify.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"time"
 
 	"github.com/urfave/cli"
@@ -60,7 +61,7 @@ func notification(ctx context.Context, endpoint string, req payload) error {
 		Status string `json:"status"`
 	}
 	rawPayload, _ := json.Marshal(&req)
-	response, err := http.Post(fmt.Sprintf("https://%s/api/v1/notification", endpoint), "application/json", bytes.NewBuffer(rawPayload))
+	response, err := http.Post(createNotificationURL(endpoint), "application/json", bytes.NewBuffer(rawPayload))
 	if err != nil {
 		return err
 	}
@@ -79,3 +80,13 @@ func notification(ctx context.Context, endpoint string, req payload) error {
 	}
 	return nil
 }
+
+// createNotificationURL creates a URL for the notification endpoint
+func createNotificationURL(endpoint string) string {
+	host := endpoint
+	url, _ := url.Parse(endpoint)
+	if url.Host != "" {
+		host = url.Host
+	}
+	return fmt.Sprintf("https://%s/api/v1/notification", host)
+}

cmd/gpud/command/up.go

@@ -43,7 +43,8 @@ func cmdUp(cliContext *cli.Context) (retErr error) {
 		return fmt.Errorf("gpud binary not found at %s (you may run 'cp %s %s' to fix the installation)", systemd.DefaultBinPath, bin, systemd.DefaultBinPath)
 	}
 
-	if err := systemdInit(); err != nil {
+	endpoint := cliContext.String("endpoint")
+	if err := systemdInit(endpoint); err != nil {
 		fmt.Printf("%s failed to initialize systemd files\n", warningSign)
 		return err
 	}
@@ -62,8 +63,8 @@ func cmdUp(cliContext *cli.Context) (retErr error) {
 	return nil
 }
 
-func systemdInit() error {
-	if err := systemd.CreateDefaultEnvFile(); err != nil {
+func systemdInit(endpoint string) error {
+	if err := systemd.CreateDefaultEnvFile(endpoint); err != nil {
 		return err
 	}
 	systemdUnitFileData := systemd.GPUDService

pkg/gossip/client.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/url"
 	"os"
 
 	apiv1 "github.com/leptonai/gpud/api/v1"
@@ -65,5 +66,10 @@ func sendRequest(ctx context.Context, url string, req apiv1.GossipRequest) (*api
 
 // createURL creates a URL for the gossip endpoint
 func createURL(endpoint string) string {
-	return fmt.Sprintf("https://%s/api/v1/gossip", endpoint)
+	host := endpoint
+	url, _ := url.Parse(endpoint)
+	if url.Host != "" {
+		host = url.Host
+	}
+	return fmt.Sprintf("https://%s/api/v1/gossip", host)
 }

pkg/gossip/client_test.go

@@ -19,6 +19,7 @@ func TestCreateURL(t *testing.T) {
 		endpoint string
 		expected string
 	}{
+		{"https://example.com", "https://example.com/api/v1/gossip"},
 		{"example.com", "https://example.com/api/v1/gossip"},
 		{"api.leptonai.com", "https://api.leptonai.com/api/v1/gossip"},
 	}

pkg/gpud-manager/systemd/systemd.go

@@ -27,23 +27,32 @@ func DefaultBinExists() bool {
 
 // CreateDefaultEnvFile creates the default environment file for gpud systemd service.
 // Assume systemdctl is already installed, and runs on the linux system.
-func CreateDefaultEnvFile() error {
-	return writeEnvFile(DefaultEnvFile)
+func CreateDefaultEnvFile(endpoint string) error {
+	return writeEnvFile(DefaultEnvFile, endpoint)
 }
 
 const defaultEnvFileContent = `# gpud environment variables are set here
 FLAGS="--log-level=info --log-file=/var/log/gpud.log"
 `
 
-func writeEnvFile(file string) error {
+func createDefaultEnvFileContent(endpoint string) string {
+	if endpoint == "" {
+		return defaultEnvFileContent
+	}
+	return fmt.Sprintf(`# gpud environment variables are set here
+FLAGS="--log-level=info --log-file=/var/log/gpud.log --endpoint=%s"
+`, endpoint)
+}
+
+func writeEnvFile(file string, endpoint string) error {
 	if _, err := os.Stat(file); err == nil {
-		return addLogFileFlagIfExists(file)
+		return updateFlagsFromExistingEnvFile(file, endpoint)
 	}
-	return atomicfile.WriteFile(file, []byte(defaultEnvFileContent), 0644)
+	return atomicfile.WriteFile(file, []byte(createDefaultEnvFileContent(endpoint)), 0644)
 }
 
-func addLogFileFlagIfExists(file string) error {
-	lines, err := processEnvFileLines(file)
+func updateFlagsFromExistingEnvFile(file string, endpoint string) error {
+	lines, err := processEnvFileLines(file, endpoint)
 	if err != nil {
 		return err
 	}
@@ -52,7 +61,7 @@ func addLogFileFlagIfExists(file string) error {
 
 // processEnvFileLines reads all lines from the environment file and processes each line,
 // adding the log-file flag to the FLAGS variable if it doesn't already exist.
-func processEnvFileLines(file string) ([]string, error) {
+func processEnvFileLines(file string, endpoint string) ([]string, error) {
 	readFile, err := os.OpenFile(file, os.O_RDONLY, 0644)
 	if err != nil {
 		return nil, err
@@ -70,15 +79,22 @@ func processEnvFileLines(file string) ([]string, error) {
 			continue
 		}
 
-		// FLAGS already contains --log-file flag
-		if strings.Contains(line, "--log-file=") {
+		// FLAGS already contains --log-file flag and --endpoint flag
+		if strings.Contains(line, "--log-file=") && (endpoint != "" && strings.Contains(line, "--endpoint=")) {
 			lines = append(lines, line)
 			continue
 		}
 
 		// remove the trailing " character
 		line = strings.TrimSuffix(line, "\"")
-		line = fmt.Sprintf("%s --log-file=/var/log/gpud.log\"", line)
+
+		if !strings.Contains(line, "--log-file=") {
+			line = fmt.Sprintf("%s --log-file=/var/log/gpud.log\"", line)
+		}
+
+		if endpoint != "" && !strings.Contains(line, "--endpoint=") {
+			line = fmt.Sprintf("%s --endpoint=%s\"", line, endpoint)
+		}
 
 		lines = append(lines, line)
 	}