beam thunder integration

bin/gen_proto.sh

@@ -109,6 +109,9 @@ protoc -I ./googleapis -I ./pkg/types -I ./pkg/gateway/ --python_betterproto_bet
 # Internal cache services
 protoc -I ./pkg/cache/ --go_out=./proto --go_opt=paths=source_relative --go-grpc_out=./proto --go-grpc_opt=paths=source_relative ./pkg/cache/cache.proto
 
+# Thunder service
+protoc -I $PROTOC_INCLUDE_PATH -I ./googleapis -I ./pkg/types -I ./pkg/gateway/services/thunder/ --go_out=./proto --go_opt=paths=source_relative --go-grpc_out=./proto --go-grpc_opt=paths=source_relative ./pkg/gateway/services/thunder/thunder.proto
+
 # Repository services
 protoc -I $PROTOC_INCLUDE_PATH -I ./googleapis -I ./pkg/types -I ./pkg/gateway/services/repository/ --go_out=./proto --go_opt=paths=source_relative --go-grpc_out=./proto --go-grpc_opt=paths=source_relative ./pkg/gateway/services/repository/container_repo.proto
 protoc -I $PROTOC_INCLUDE_PATH -I ./googleapis -I ./pkg/types -I ./pkg/gateway/services/repository/ --go_out=./proto --go_opt=paths=source_relative --go-grpc_out=./proto --go-grpc_opt=paths=source_relative ./pkg/gateway/services/repository/worker_repo.proto

docs/openapi/gateway.swagger.json

@@ -3520,6 +3520,9 @@
         },
         "checkpointTrigger": {
           "$ref": "#/definitions/typesCheckpointTrigger"
+        },
+        "gpuVirtualized": {
+          "type": "boolean"
         }
       }
     },

pkg/abstractions/common/instance.go

@@ -337,6 +337,7 @@ func (i *AutoscaledInstance) CheckConcurrencyLimit() error {
 	request := &types.ContainerRequest{
 		Cpu:              i.StubConfig.Runtime.Cpu,
 		GpuCount:         uint32(gpuCount),
+		GpuVirtualized:   i.StubConfig.Runtime.GpuVirtualized,
 		WorkspaceId:      i.Workspace.ExternalId,
 		Workspace:        *i.Workspace,
 		AllowMarketplace: i.StubConfig.AllowMarketplace,

pkg/abstractions/endpoint/instance.go

@@ -112,6 +112,7 @@ func (i *endpointInstance) startContainers(containersToRun int) error {
 			Memory:            i.StubConfig.Runtime.Memory,
 			GpuRequest:        gpuRequest,
 			GpuCount:          uint32(gpuCount),
+			GpuVirtualized:    i.StubConfig.Runtime.GpuVirtualized,
 			AppId:             i.Stub.App.ExternalId,
 			ImageId:           i.StubConfig.Runtime.ImageId,
 			StubId:            i.Stub.ExternalId,

pkg/abstractions/function/task.go

@@ -47,6 +47,7 @@ func (t *FunctionTask) Execute(ctx context.Context, options ...interface{}) erro
 		if err := t.fs.scheduler.CheckConcurrencyLimit(&types.ContainerRequest{
 			Cpu:              cpu,
 			GpuCount:         uint32(gpuCount),
+			GpuVirtualized:   stubConfig.Runtime.GpuVirtualized,
 			WorkspaceId:      stub.Workspace.ExternalId,
 			Workspace:        stub.Workspace,
 			StubId:           stub.ExternalId,
@@ -204,22 +205,23 @@ func (t *FunctionTask) run(ctx context.Context, stub *types.StubWithRelated, tas
 	}
 
 	runRequest := &types.ContainerRequest{
-		ContainerId:  t.containerId,
-		Env:          env,
-		Cpu:          stubConfig.Runtime.Cpu,
-		Memory:       stubConfig.Runtime.Memory,
-		GpuRequest:   gpuRequest,
-		GpuCount:     uint32(gpuCount),
-		ImageId:      stubConfig.Runtime.ImageId,
-		StubId:       stub.ExternalId,
-		TaskId:       task.ExternalId,
-		AppId:        stub.App.ExternalId,
-		WorkspaceId:  stub.Workspace.ExternalId,
-		Workspace:    stub.Workspace,
-		EntryPoint:   []string{stubConfig.PythonVersion, "-m", "beta9.runner.function"},
-		Mounts:       mounts,
-		Stub:         *stub,
-		PoolSelector: stubConfig.PoolSelector(),
+		ContainerId:    t.containerId,
+		Env:            env,
+		Cpu:            stubConfig.Runtime.Cpu,
+		Memory:         stubConfig.Runtime.Memory,
+		GpuRequest:     gpuRequest,
+		GpuCount:       uint32(gpuCount),
+		GpuVirtualized: stubConfig.Runtime.GpuVirtualized,
+		ImageId:        stubConfig.Runtime.ImageId,
+		StubId:         stub.ExternalId,
+		TaskId:         task.ExternalId,
+		AppId:          stub.App.ExternalId,
+		WorkspaceId:    stub.Workspace.ExternalId,
+		Workspace:      stub.Workspace,
+		EntryPoint:     []string{stubConfig.PythonVersion, "-m", "beta9.runner.function"},
+		Mounts:         mounts,
+		Stub:           *stub,
+		PoolSelector:   stubConfig.PoolSelector(),
 	}
 	if err := abstractions.ConfigureContainerRequestNetwork(runRequest, stubConfig); err != nil {
 		return err

pkg/abstractions/pod/instance.go

@@ -144,6 +144,7 @@ func (i *podInstance) startContainers(containersToRun int) error {
 			Memory:            i.StubConfig.Runtime.Memory,
 			GpuRequest:        gpuRequest,
 			GpuCount:          uint32(gpuCount),
+			GpuVirtualized:    i.StubConfig.Runtime.GpuVirtualized,
 			ImageId:           i.StubConfig.Runtime.ImageId,
 			StubId:            i.Stub.ExternalId,
 			AppId:             i.Stub.App.ExternalId,

pkg/abstractions/pod/pod.go

@@ -470,6 +470,15 @@ func (s *GenericPodService) run(ctx context.Context, authInfo *auth.AuthInfo, st
 		gpuCount = 1
 	}
 
+	log.Info().
+		Str("stub_id", stub.ExternalId).
+		Str("container_id", containerId).
+		Str("gpu", stubConfig.Runtime.Gpu.String()).
+		Strs("gpu_request", gpuRequest).
+		Uint32("gpu_count", uint32(gpuCount)).
+		Bool("gpu_virtualized", stubConfig.Runtime.GpuVirtualized).
+		Msg("pod service built sandbox run request from stub config")
+
 	checkpointEnabled := stubConfig.CheckpointEnabled
 	if gpuCount > 1 {
 		checkpointEnabled = false
@@ -507,6 +516,7 @@ func (s *GenericPodService) run(ctx context.Context, authInfo *auth.AuthInfo, st
 		Memory:            stubConfig.Runtime.Memory,
 		GpuRequest:        gpuRequest,
 		GpuCount:          uint32(gpuCount),
+		GpuVirtualized:    stubConfig.Runtime.GpuVirtualized,
 		Mounts:            mounts,
 		Stub:              *stub,
 		ImageId:           *imageId,

pkg/abstractions/shell/shell.go

@@ -562,6 +562,7 @@ func (ss *SSHShellService) CreateStandaloneShell(ctx context.Context, in *pb.Cre
 		Memory:           stubConfig.Runtime.Memory,
 		GpuRequest:       gpuRequest,
 		GpuCount:         uint32(gpuCount),
+		GpuVirtualized:   stubConfig.Runtime.GpuVirtualized,
 		ImageId:          stubConfig.Runtime.ImageId,
 		StubId:           stub.ExternalId,
 		AppId:            "",

pkg/abstractions/taskqueue/instance.go

@@ -93,6 +93,7 @@ func (i *taskQueueInstance) startContainers(containersToRun int) error {
 			Memory:            i.StubConfig.Runtime.Memory,
 			GpuRequest:        gpuRequest,
 			GpuCount:          uint32(gpuCount),
+			GpuVirtualized:    i.StubConfig.Runtime.GpuVirtualized,
 			ImageId:           i.StubConfig.Runtime.ImageId,
 			StubId:            i.Stub.ExternalId,
 			AppId:             i.Stub.App.ExternalId,

pkg/agent/thunder.go

@@ -0,0 +1,136 @@
+package agent
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net"
+	"net/netip"
+	"os/exec"
+	"strings"
+	"time"
+
+	"github.com/beam-cloud/beta9/pkg/common"
+	pb "github.com/beam-cloud/beta9/proto"
+	"google.golang.org/grpc"
+)
+
+const thunderNodeInstallTimeout = 2 * time.Minute
+
+var runThunderNodeInstallCommand = defaultRunThunderNodeInstallCommand
+
+type nodeEnrollmentGatewayClient interface {
+	CreateNodeEnrollment(context.Context, *pb.CreateNodeEnrollmentRequest, ...grpc.CallOption) (*pb.CreateNodeEnrollmentResponse, error)
+	DeleteNodeEnrollment(context.Context, *pb.DeleteNodeEnrollmentRequest, ...grpc.CallOption) (*pb.DeleteNodeEnrollmentResponse, error)
+}
+
+func setupThunderNode(ctx context.Context, client nodeEnrollmentGatewayClient, agentToken string, tailscaleIPs []netip.Addr, stdout, stderr io.Writer) error {
+	if client == nil {
+		return fmt.Errorf("gateway client is required for Thunder node enrollment")
+	}
+
+	reachableIP, err := thunderReachableNodeIP(tailscaleIPs)
+	if err != nil {
+		return err
+	}
+
+	res, err := client.CreateNodeEnrollment(ctx, &pb.CreateNodeEnrollmentRequest{AgentToken: agentToken})
+	if err != nil {
+		return fmt.Errorf("create Thunder node enrollment: %w", err)
+	}
+	if res == nil || !res.Ok {
+		return fmt.Errorf("create Thunder node enrollment: %s", firstNonEmpty(res.GetErrorMsg(), "gateway rejected request"))
+	}
+
+	enrollmentToken := strings.TrimSpace(res.EnrollmentToken)
+	if enrollmentToken == "" {
+		verbosef(stdout, "Thunder node enrollment already exists\n")
+		return nil
+	}
+
+	cmd := thunderNodeInstallCommand(reachableIP, enrollmentToken)
+	installCtx, cancel := context.WithTimeout(ctx, thunderNodeInstallTimeout)
+	defer cancel()
+
+	statusf(stdout, "Installing Thunder node")
+	if err := runThunderNodeInstallCommand(installCtx, cmd); err != nil {
+		deleteThunderNodeEnrollment(context.Background(), client, agentToken, stderr)
+		return fmt.Errorf("install Thunder node: %w", err)
+	}
+	statusf(stdout, "Thunder node ready")
+	return nil
+}
+
+func deleteThunderNodeEnrollment(ctx context.Context, client nodeEnrollmentGatewayClient, agentToken string, stderr io.Writer) {
+	if client == nil {
+		return
+	}
+	if stderr == nil {
+		stderr = io.Discard
+	}
+	deleteCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+
+	res, err := client.DeleteNodeEnrollment(deleteCtx, &pb.DeleteNodeEnrollmentRequest{AgentToken: agentToken})
+	if err != nil {
+		fmt.Fprintf(stderr, "failed to delete Thunder node enrollment: %v\n", err)
+		return
+	}
+	if res == nil || !res.Ok {
+		fmt.Fprintf(stderr, "failed to delete Thunder node enrollment: %s\n", firstNonEmpty(res.GetErrorMsg(), "gateway rejected request"))
+	}
+}
+
+func thunderReachableNodeIP(ips []netip.Addr) (string, error) {
+	for _, ip := range ips {
+		if ip.IsValid() && ip.Is4() {
+			return ip.String(), nil
+		}
+	}
+	for _, ip := range ips {
+		if ip.IsValid() {
+			return ip.String(), nil
+		}
+	}
+	return "", fmt.Errorf("tailscale IP address is required for Thunder node enrollment")
+}
+
+func hostTailscaleIPs() []netip.Addr {
+	iface, err := net.InterfaceByName("tailscale0")
+	if err != nil {
+		return nil
+	}
+	addrs, err := iface.Addrs()
+	if err != nil {
+		return nil
+	}
+
+	ips := make([]netip.Addr, 0, len(addrs))
+	for _, addr := range addrs {
+		prefix, err := netip.ParsePrefix(addr.String())
+		if err == nil {
+			ips = append(ips, prefix.Addr())
+			continue
+		}
+		if ip, err := netip.ParseAddr(addr.String()); err == nil {
+			ips = append(ips, ip)
+		}
+	}
+	return ips
+}
+
+func thunderNodeInstallCommand(reachableIP, enrollmentToken string) string {
+	return "curl -fsSL https://get.thundercompute.com/install.sh | sudo THUNDER_INSTALL_MODE=thunderd THUNDERD_IP=" + common.ShellQuote(reachableIP) + " THUNDER_ENROLLMENT_TOKEN=" + common.ShellQuote(enrollmentToken) + " sh"
+}
+
+func defaultRunThunderNodeInstallCommand(ctx context.Context, command string) error {
+	out, err := exec.CommandContext(ctx, "sh", "-c", command).CombinedOutput()
+	if err != nil {
+		output := strings.TrimSpace(string(out))
+		if output != "" {
+			return fmt.Errorf("%w: %s", err, output)
+		}
+		return err
+	}
+	return nil
+}