beam-thunder integration - #1814
Open
rohanphadnis-thunder wants to merge 28 commits into
Open
Conversation
added gateway-side changes
tailscale ip formatting
beam thunder integration
luke-lombardi
requested changes
Aug 3, 2026
luke-lombardi
left a comment
luke-lombardi
left a comment
Contributor
There was a problem hiding this comment.
High level this looks good, mostly just some nit-picks around code organization
| request := &types.ContainerRequest{ | ||
| Cpu: i.StubConfig.Runtime.Cpu, | ||
| GpuCount: uint32(gpuCount), | ||
| GpuVirtualized: i.StubConfig.Runtime.GpuVirtualized, |
Contributor
There was a problem hiding this comment.
Is this still required? I know we talked about removing from the client, but I feel like it should also be removed from the container request as well.
| } | ||
|
|
||
| func (s *Worker) deleteContainer(containerId string) { | ||
| s.thunderSetupTracker.Delete(containerId) |
Contributor
There was a problem hiding this comment.
feels like the thundersetuptracker calls could be pushed down into the runtime layer potentially? feels like it could get confusing having those calls up here next to the core container lifecycle pieces
| }() | ||
|
|
||
| exitCode, _ = s.runContainer(ctx, request, outputLogger, outputWriter, startedChan, checkpointPIDChan, opts.StartupStartedAt, opts.StartupPortBindings, opts.CheckpointFilesystemRestore) | ||
| exitCode, _ = s.runContainer(ctx, request, outputLogger, outputWriter, startedChan, checkpointPIDChan, thunderInstallResult, opts.StartupStartedAt, opts.StartupPortBindings, opts.CheckpointFilesystemRestore) |
Contributor
There was a problem hiding this comment.
same comment as above, if theres a way to get this to be jammed into a runtime (under pkg/runtime) I think that would be cleaner
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
2 participants
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.
Beam-Thunder Integration
This PR is for integrating beta9 with Thunder Compute. Key changes:
Phase 1:
gpu_virtualizedgpu_virtualized, a boolean value is propagated from the Python beam client SDK to the beam gateway to each individual worker. Note that this is still a temporary change, preserved for its usefulness in prototyping. Eventually, GPU virtualization is controlled at the level of the pool. If a pool is virtualized, all its containers would be virtualized as well.Phase 2: Beam Worker
The beam worker implements
beta9/pkg/worker/thunder.go. This contains theContainerThunderManager, a struct which implements theGPUManagerinterface. It mountsnvidia-smi,libcuda.so, andlibnvidia-ml.sointo the container from the host, but blocks the mounting of any physical GPU. For assigning and unassigning GPUs, it will use the gateway's thunder service and mint/revoke enrollment tokens as needed. Finally, when the container starts, the worker will run the curl installer in the container.Phase 3: Beam Agent
The beam agent change is very easy.
tailscale0interface.CreateNodeEnrollmentto get a node enrollment token and run the node curl installer.Phase 4: Beam Gateway
The beam gateway provides RPC calls to the beam agent and the beam worker in order to mint these enrollment tokens. It acts as a central client for the Thunder Compute API. This design ensures that the global API token isn't propagated to each beam agent/worker.
To make the Thunder information stateful, 3 new redis maps are added:
Summary by cubic
Integrates Thunder Compute to enable virtualized GPU support via a new
gpu_virtualizedflag, with centralized Thunder enrollment across gateway, agent, and worker. This lets you schedule GPUs without exposing physical devices to containers.New Features
gpu_virtualizedtotypes.ContainerRequestand propagate through all abstractions and SDK; add scheduler helper to distinguish physical vs virtual GPUs.ContainerThunderManagerto mountnvidia-smi,libcuda.so.1,libnvidia-ml.so.1, and Thunder lib while blocking physical GPU devices when virtualized.ThunderService; wait for Thunder setup before exec and clean up on teardown.tailscale0IP via new gateway RPCs:CreateNodeEnrollmentandDeleteNodeEnrollment; soft-fails if unavailable.ThunderService(pkg/gateway/services/thunder) with Redis-backed state and HTTP client; centralizeTHUNDER_API_TOKEN; mint/revoke client and node enrollment tokens; map pools to Thunder zones.gpu_virtualizedoption to Pods, Sandboxes, and generic runners; forward toGetOrCreateStub.bin/gen_proto.sh.Migration
THUNDER_API_URLandTHUNDER_API_TOKENin the gateway environment.gpu_virtualized=Truein the beta9 Python SDK when launching Pods/Sandboxes/Functions.Written for commit 9a6a9bb. Summary will update on new commits.