beam thunder integration - #1
Merged
Merged
Conversation
rohanphadnis-thunder
marked this pull request as ready for review
July 17, 2026 17:05
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 3 potential issues.
Autofix Details
Bugbot Autofix prepared fixes for all 3 issues found in the latest run.
- ✅ Fixed: Central token exposed in sandbox
- The sandbox config now receives the scoped per-client token instead of the privileged worker API token.
- ✅ Fixed: Cleanup lock blocks on Thunder HTTP
- GPU unassignment now completes before acquiring the worker-wide container cleanup lock.
- ✅ Fixed: Failed unregister orphans Thunder clients
- Failed Thunder unregister requests now retain the local allocation token so cleanup can be retried.
Or push these changes by commenting:
@cursor push 02f9b59231
Preview (02f9b59231)
diff --git a/pkg/worker/lifecycle.go b/pkg/worker/lifecycle.go
--- a/pkg/worker/lifecycle.go
+++ b/pkg/worker/lifecycle.go
@@ -163,13 +163,13 @@
s.containerInstances.Set(containerId, instance)
}
- s.containerLock.Lock()
-
// De-allocate GPU devices so they are available for new containers
if request != nil && request.RequiresGPU() {
s.gpuManagerForRequest(request).UnassignGPUDevices(containerId)
}
+ s.containerLock.Lock()
+
// Tear down container network components - best effort
if err := s.containerNetworkManager.TearDown(request.ContainerId); err != nil {
log.Warn().Str("container_id", request.ContainerId).Err(err).Msg("failed to clean up container network")
diff --git a/pkg/worker/lifecycle_test.go b/pkg/worker/lifecycle_test.go
--- a/pkg/worker/lifecycle_test.go
+++ b/pkg/worker/lifecycle_test.go
@@ -4,6 +4,8 @@
"context"
"io"
"log/slog"
+ "net/http"
+ "net/http/httptest"
"os"
"path/filepath"
"strings"
@@ -55,6 +57,63 @@
require.False(t, handled)
}
+func TestClearContainerDoesNotHoldContainerLockDuringThunderUnassign(t *testing.T) {
+ deleteStarted := make(chan struct{})
+ releaseDelete := make(chan struct{})
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ require.Equal(t, thunderDeleteClientPath, r.URL.Path)
+ close(deleteStarted)
+ <-releaseDelete
+ w.WriteHeader(http.StatusOK)
+ }))
+ defer server.Close()
+
+ manager := NewContainerThunderManager(server.URL, "central-token", server.Client())
+ manager.allocations.Set("container-1", thunderAllocation{Token: "client-token"})
+ worker := &Worker{
+ ctx: context.Background(),
+ containerInstances: common.NewSafeMap[*ContainerInstance](),
+ containerThunderManager: manager,
+ containerNetworkManager: &fakeContainerNetworkController{},
+ completedRequests: make(chan *types.ContainerRequest, 1),
+ config: types.AppConfig{Worker: types.WorkerConfig{
+ TerminationGracePeriod: 3600,
+ }},
+ }
+ request := &types.ContainerRequest{
+ ContainerId: "container-1",
+ Gpu: "H100",
+ GpuCount: 1,
+ GpuVirtualized: true,
+ }
+
+ clearDone := make(chan struct{})
+ go func() {
+ worker.clearContainer(request.ContainerId, request, 0)
+ close(clearDone)
+ }()
+ <-deleteStarted
+
+ lockAcquired := make(chan struct{})
+ go func() {
+ worker.containerLock.Lock()
+ worker.containerLock.Unlock()
+ close(lockAcquired)
+ }()
+ select {
+ case <-lockAcquired:
+ case <-time.After(time.Second):
+ t.Fatal("container lock was held during Thunder unregister")
+ }
+
+ close(releaseDelete)
+ select {
+ case <-clearDone:
+ case <-time.After(time.Second):
+ t.Fatal("container cleanup did not complete")
+ }
+}
+
func TestContainerResolvConfSourceFallsBackForLoopbackHostResolver(t *testing.T) {
hostResolv := filepath.Join(t.TempDir(), "resolv.conf")
require.NoError(t, os.WriteFile(hostResolv, []byte("nameserver 127.0.0.53\noptions edns0\n"), 0o644))
diff --git a/pkg/worker/thunder.go b/pkg/worker/thunder.go
--- a/pkg/worker/thunder.go
+++ b/pkg/worker/thunder.go
@@ -124,6 +124,7 @@
payload := thunderDeleteClientRequest{DeviceID: containerId, Token: allocation.Token}
if err := c.doThunderRequest(http.MethodPost, thunderDeleteClientPath, payload, nil); err != nil {
log.Error().Str("container_id", containerId).Err(err).Msg("failed to unregister Thunder virtual GPU client")
+ return
}
c.allocations.Delete(containerId)
}
@@ -245,7 +246,7 @@
if err := createThunderToken(rootPath, allocation.Token); err != nil {
return err
}
- if err := c.createThunderConfig(rootPath, request); err != nil {
+ if err := c.createThunderConfig(rootPath, request, allocation.Token); err != nil {
return err
}
return nil
@@ -279,7 +280,7 @@
return writeThunderFile(rootPath, thunderTokenPath, []byte(token), 0644)
}
-func (c *ContainerThunderManager) createThunderConfig(rootPath string, request *types.ContainerRequest) error {
+func (c *ContainerThunderManager) createThunderConfig(rootPath string, request *types.ContainerRequest, clientToken string) error {
config := thunderConfigFile{
DeviceID: request.ContainerId,
GPUCount: int(thunderGPUCount(request)),
@@ -287,7 +288,7 @@
EnableGRPCTLS: false,
CentralApiUrl: c.apiURL,
CentralZoneId: "thunder-beam",
- CentralApiToken: c.apiToken,
+ CentralApiToken: clientToken,
}
contents, err := json.Marshal(config)
diff --git a/pkg/worker/thunder_test.go b/pkg/worker/thunder_test.go
--- a/pkg/worker/thunder_test.go
+++ b/pkg/worker/thunder_test.go
@@ -95,6 +95,42 @@
assert.Equal(t, thunderDeleteClientRequest{DeviceID: "container-123", Token: "client-token"}, deletePayload)
}
+func TestThunderUnassignRetainsAllocationAfterDeleteFailure(t *testing.T) {
+ var deleteAttempts int32
+ server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ switch r.URL.Path {
+ case thunderRegisterClientPath:
+ _ = json.NewEncoder(w).Encode(thunderRegisterClientResponse{Token: "client-token"})
+ case thunderDeleteClientPath:
+ if atomic.AddInt32(&deleteAttempts, 1) <= thunderRequestAttempts {
+ w.WriteHeader(http.StatusInternalServerError)
+ return
+ }
+ _ = json.NewEncoder(w).Encode(map[string]bool{"success": true})
+ default:
+ t.Fatalf("unexpected path %s", r.URL.Path)
+ }
+ }))
+ defer server.Close()
+
+ manager := NewContainerThunderManager(server.URL, "central-token", server.Client())
+ request := &types.ContainerRequest{ContainerId: "container-123", Gpu: "H100", GpuCount: 1, GpuVirtualized: true}
+ if _, err := manager.AssignGPUDevices(request); err != nil {
+ t.Fatal(err)
+ }
+
+ manager.UnassignGPUDevices(request.ContainerId)
+ if _, ok := manager.allocations.Get(request.ContainerId); !ok {
+ t.Fatal("allocation was deleted after failed unregister")
+ }
+
+ manager.UnassignGPUDevices(request.ContainerId)
+ if _, ok := manager.allocations.Get(request.ContainerId); ok {
+ t.Fatal("allocation was retained after successful unregister")
+ }
+ assert.Equal(t, int32(thunderRequestAttempts+1), atomic.LoadInt32(&deleteAttempts))
+}
+
func TestThunderAssignRetriesUnsuccessfulStatus(t *testing.T) {
var attempts int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -196,6 +232,6 @@
EnableGRPCTLS: false,
CentralApiUrl: server.URL,
CentralZoneId: "thunder-beam",
- CentralApiToken: "central-token",
+ CentralApiToken: "client-token",
}, config)
}You can send follow-ups to the cloud agent here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 3 total unresolved issues (including 2 from previous reviews).
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Entrypoint runs before Thunder install
- Wrapped virtualized-GPU entrypoints with a readiness barrier that is released only after the Thunder installer succeeds.
Or push these changes by commenting:
@cursor push caed28bddc
Preview (caed28bddc)
diff --git a/pkg/worker/lifecycle.go b/pkg/worker/lifecycle.go
--- a/pkg/worker/lifecycle.go
+++ b/pkg/worker/lifecycle.go
@@ -1180,6 +1180,13 @@
})
}
+ if request.GpuVirtualized {
+ if err := s.configureThunderClientInstall(request, spec); err != nil {
+ log.Error().Err(err).Str("container_id", request.ContainerId).Msg("failed to configure Thunder client install")
+ return
+ }
+ }
+
// Add Docker capabilities if enabled for sandbox containers.
if request.DockerEnabled && request.Stub.Type.Kind() == types.StubTypeSandbox {
runtime.AddDockerInDockerCapabilities(spec)
diff --git a/pkg/worker/thunder.go b/pkg/worker/thunder.go
--- a/pkg/worker/thunder.go
+++ b/pkg/worker/thunder.go
@@ -285,6 +285,28 @@
return "", false
}
+func (s *Worker) configureThunderClientInstall(request *types.ContainerRequest, spec *specs.Spec) error {
+ if s == nil || request == nil || !request.GpuVirtualized {
+ return nil
+ }
+ manager, ok := s.containerThunderManager.(*ContainerThunderManager)
+ if !ok || manager == nil {
+ return fmt.Errorf("thunder manager unavailable")
+ }
+ allocation, ok := manager.allocations.Get(request.ContainerId)
+ if !ok || strings.TrimSpace(allocation.EnrollmentToken) == "" {
+ return fmt.Errorf("missing Thunder enrollment token for container %s", request.ContainerId)
+ }
+ if spec == nil || spec.Process == nil || len(spec.Process.Args) == 0 {
+ return fmt.Errorf("container process unavailable for Thunder install")
+ }
+
+ readyPath := thunderInstallReadyPath(request.ContainerId)
+ cmd := "while [ ! -f " + common.ShellQuote(readyPath) + " ]; do sleep 0.1; done; exec \"$@\""
+ spec.Process.Args = append([]string{"sh", "-c", cmd, "--"}, spec.Process.Args...)
+ return nil
+}
+
func (s *Worker) installThunderClient(ctx context.Context, request *types.ContainerRequest) error {
if s == nil || request == nil || !request.GpuVirtualized {
return nil
@@ -312,7 +334,7 @@
cwd = instance.Spec.Process.Cwd
}
}
- cmd := "curl -fsSL https://get.thundercompute.com/install.sh | sudo THUNDER_INSTALL_MODE=client THUNDER_AUTH_TOKEN=" + common.ShellQuote(allocation.EnrollmentToken) + " sh"
+ cmd := "curl -fsSL https://get.thundercompute.com/install.sh | sudo THUNDER_INSTALL_MODE=client THUNDER_AUTH_TOKEN=" + common.ShellQuote(allocation.EnrollmentToken) + " sh && : > " + common.ShellQuote(thunderInstallReadyPath(request.ContainerId))
proc := specs.Process{
Args: []string{"sh", "-c", cmd},
Cwd: cwd,
@@ -326,6 +348,10 @@
return nil
}
+func thunderInstallReadyPath(containerID string) string {
+ return "/tmp/.beta9-thunder-client-ready-" + containerID
+}
+
func containsEnvKey(env []string, key string) bool {
prefix := key + "="
for _, item := range env {
diff --git a/pkg/worker/thunder_test.go b/pkg/worker/thunder_test.go
--- a/pkg/worker/thunder_test.go
+++ b/pkg/worker/thunder_test.go
@@ -92,7 +92,7 @@
assert.Equal(t, fmt.Sprintf(thunderEnrollmentTokenNodePath, "token-id"), deletePath)
}
-func TestInstallThunderClientExecutesInstaller(t *testing.T) {
+func TestThunderClientInstallBlocksEntrypointUntilReady(t *testing.T) {
rt := &mockRuntime{name: "runc"}
manager := NewContainerThunderManager("https://worker-default.example", "worker-default-token", nil)
manager.allocations.Set("container-123", thunderAllocation{
@@ -100,27 +100,40 @@
APIURL: "https://worker-default.example",
APIToken: "worker-default-token",
})
+ spec := &specs.Spec{
+ Process: &specs.Process{
+ Args: []string{"python", "app.py"},
+ Cwd: "/workspace",
+ Env: []string{"A=1"},
+ },
+ }
instances := common.NewSafeMap[*ContainerInstance]()
- instances.Set("container-123", &ContainerInstance{
- Runtime: rt,
- Spec: &specs.Spec{Process: &specs.Process{Cwd: "/workspace", Env: []string{"A=1"}}},
- })
+ instances.Set("container-123", &ContainerInstance{Runtime: rt, Spec: spec})
worker := &Worker{
containerThunderManager: manager,
containerInstances: instances,
}
+ request := &types.ContainerRequest{ContainerId: "container-123", GpuVirtualized: true}
- err := worker.installThunderClient(context.Background(), &types.ContainerRequest{ContainerId: "container-123", GpuVirtualized: true})
+ err := worker.configureThunderClientInstall(request, spec)
if err != nil {
t.Fatal(err)
}
+ assert.Equal(t, []string{"sh", "-c", "while [ ! -f '/tmp/.beta9-thunder-client-ready-container-123' ]; do sleep 0.1; done; exec \"$@\"", "--", "python", "app.py"}, spec.Process.Args)
+ assert.Equal(t, "/workspace", spec.Process.Cwd)
+ assert.Equal(t, []string{"A=1"}, spec.Process.Env)
+
+ err = worker.installThunderClient(context.Background(), request)
+ if err != nil {
+ t.Fatal(err)
+ }
if len(rt.execCalls) != 1 {
t.Fatalf("execCalls = %d, want 1", len(rt.execCalls))
}
call := rt.execCalls[0]
assert.Equal(t, "container-123", call.containerID)
assert.Equal(t, "/workspace", call.proc.Cwd)
- assert.Equal(t, []string{"sh", "-c", "curl -fsSL https://get.thundercompute.com/install.sh | sudo THUNDER_INSTALL_MODE=client THUNDER_AUTH_TOKEN='enroll-token' sh"}, call.proc.Args)
+ assert.Equal(t, []string{"sh", "-c", "curl -fsSL https://get.thundercompute.com/install.sh | sudo THUNDER_INSTALL_MODE=client THUNDER_AUTH_TOKEN='enroll-token' sh && : > '/tmp/.beta9-thunder-client-ready-container-123'"}, call.proc.Args)
assert.Contains(t, call.proc.Env, "A=1")
assert.Contains(t, call.proc.Env, "PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin")
}You can send follow-ups to the cloud agent here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 4 total unresolved issues (including 3 from previous reviews).
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Unsafe Spec access before nil check
- Added Spec and Process nil guards before environment access so Thunder installation returns a controlled error instead of panicking.
Or push these changes by commenting:
@cursor push 97648a5d4e
Preview (97648a5d4e)
diff --git a/pkg/worker/thunder.go b/pkg/worker/thunder.go
--- a/pkg/worker/thunder.go
+++ b/pkg/worker/thunder.go
@@ -330,7 +330,7 @@
return fmt.Errorf("missing Thunder enrollment token for container %s", request.ContainerId)
}
instance, ok := s.containerInstances.Get(request.ContainerId)
- if !ok || instance == nil || instance.Runtime == nil {
+ if !ok || instance == nil || instance.Runtime == nil || instance.Spec == nil || instance.Spec.Process == nil {
return fmt.Errorf("container runtime unavailable for Thunder install")
}You can send follow-ups to the cloud agent here.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
There are 3 total unresolved issues (including 2 from previous reviews).
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Thunder wait can deadlock cleanup
- The Thunder result channel now closes on every startup goroutine exit, and the waiter treats closure without a result as an aborted install instead of blocking.
Or push these changes by commenting:
@cursor push f2a5f85e81
Preview (f2a5f85e81)
diff --git a/pkg/worker/lifecycle.go b/pkg/worker/lifecycle.go
--- a/pkg/worker/lifecycle.go
+++ b/pkg/worker/lifecycle.go
@@ -1273,6 +1273,8 @@
}()
go func() {
+ defer close(thunderInstallResult)
+
// When the process starts monitor it and potentially checkpoint it
pid, ok := <-startedChan
if !ok {
@@ -1550,8 +1552,8 @@
if request.Stub.Type.Kind() == types.StubTypeSandbox && request.GpuVirtualized {
select {
- case err := <-thunderInstallResult:
- if err != nil {
+ case err, ok := <-thunderInstallResult:
+ if !ok || err != nil {
return
}
case <-ctx.Done():You can send follow-ups to the cloud agent here.
added gateway-side changes
tailscale ip formatting
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: Wrong Tailscale IP source
- Thunder enrollment now reads Tailscale IPs from the userspace tsnet LocalClient status instead of the missing kernel tailscale0 interface.
Or push these changes by commenting:
@cursor push 664252a45a
Preview (664252a45a)
diff --git a/pkg/agent/thunder.go b/pkg/agent/thunder.go
--- a/pkg/agent/thunder.go
+++ b/pkg/agent/thunder.go
@@ -4,7 +4,6 @@
"context"
"fmt"
"io"
- "net"
"net/netip"
"os/exec"
"strings"
@@ -95,30 +94,6 @@
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"
}
diff --git a/pkg/agent/transport.go b/pkg/agent/transport.go
--- a/pkg/agent/transport.go
+++ b/pkg/agent/transport.go
@@ -6,6 +6,7 @@
"fmt"
"io"
"net"
+ "net/netip"
"sort"
"strconv"
"strings"
@@ -79,14 +80,16 @@
}
hostname := credential.Hostname
+ var tailscaleIPs []netip.Addr
if localClient, err := server.LocalClient(); err == nil {
if status, err := localClient.Status(ctx); err == nil {
if status.Self != nil && status.Self.DNSName != "" {
hostname = strings.TrimSuffix(status.Self.DNSName, ".")
}
+ tailscaleIPs = status.TailscaleIPs
}
}
- if err := setupThunderNode(ctx, client, agentToken, hostTailscaleIPs(), stdout, stderr); err != nil {
+ if err := setupThunderNode(ctx, client, agentToken, tailscaleIPs, stdout, stderr); err != nil {
fmt.Fprintf(stderr, "Thunder node enrollment skipped: %v\n", err)
} else {
defer deleteThunderNodeEnrollment(context.Background(), client, agentToken, stderr)You can send follow-ups to the cloud agent here.
Reviewed by Cursor Bugbot for commit 532ee76. Configure here.
| fmt.Fprintf(stderr, "Thunder node enrollment skipped: %v\n", err) | ||
| } else { | ||
| defer deleteThunderNodeEnrollment(context.Background(), client, agentToken, stderr) | ||
| } |
There was a problem hiding this comment.
Wrong Tailscale IP source
High Severity
Thunder node setup reads IPs from the kernel tailscale0 interface, but the agent uses userspace tsnet and does not create that interface. Enrollment is skipped on every agent start, so Thunder server nodes never come up.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 532ee76. Configure here.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR achieves the following:
gpu_virtualizedexposed to client SDK and propagated through workflow (client -> gateway -> worker)gpu_virtualized == true. NewContainerThunderManagerimplements theGPUManagerinterface.Next Steps:
Note
High Risk
Changes core GPU scheduling and sandbox startup, adds external Thunder API/enrollment flows, and new unauthenticated gateway RPCs that rely on agent-token validation in handlers.
Overview
Introduces Thunder integration for virtualized GPU workloads: clients can set
gpu_virtualizedon stubs; the gateway persists it and every abstraction path copies it ontoContainerRequest.On the worker, virtualized requests use a new
ContainerThunderManager(sameGPUManagerinterface) instead of NVIDIA CDI/device pinning. Sandboxes run the Thunder client installer after the process manager is ready, blockContainerSandboxExecuntil setup finishes, and bind Thunder-related mounts/LD_PRELOADwithout assigning physical GPUs.RequiresPhysicalGPU()gates runtime GPU capability and CDI injection.The gateway registers a
ThunderService(gRPC + Redis for zones and client/node enrollment tokens) and talks to the Thunder HTTP API. Agent-token RPCsCreateNodeEnrollment/DeleteNodeEnrollmentenroll compute agents as Thunder nodes (installer via Tailscale IP); the agent transport runs that on startup and cleans up on shutdown. Workers callCreateClientEnrollment/DeleteClientEnrollmentwith worker auth for per-container clients.Reviewed by Cursor Bugbot for commit 532ee76. Bugbot is set up for automated code reviews on this repo. Configure here.