beam-thunder integration - #1814
Conversation
added gateway-side changes
tailscale ip formatting
beam thunder integration
luke-lombardi
left a comment
luke-lombardi
left a comment
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, |
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) |
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) |
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
integrated thunder sdk with gateway
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
All reported issues were addressed across 13 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
added startup hook functionality
| gpuCount = 1 | ||
| } | ||
|
|
||
| log.Info(). |
There was a problem hiding this comment.
Is this log required? we log this in the scheduler too I think
| return err | ||
| } | ||
|
|
||
| func (r *RedisRepository) DeleteZone(ctx context.Context, workspaceID, poolName string) error { |
There was a problem hiding this comment.
I think it may make sense to push this into the repository layer. repository/thunder_redis.go or something
| "github.com/opencontainers/runtime-spec/specs-go" | ||
| ) | ||
|
|
||
| const startupHookShutdownTimeout = 5 * time.Second |
There was a problem hiding this comment.
This code can probably also live inside startup_hook.go maybe?
re-organized hooks and thunder redis
Cleaned up beam worker
removed ttl param from thunder repo
cleanup for beta9 gateway package
cleaned up agent code
re-ordered comment
post testing
2 participants
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 virtual GPUs with pool-level control and startup safety. Previously only physical GPUs were used and containers published “started” immediately; now pools can enable virtualization, workers choose Thunder per request, and “started” waits until the Thunder installer completes.
Pool control: adds
WorkerPoolConfig.gpuVirtualized(surfaced in OpenAPI and Python SDK) and propagates to workers viaWORKER_GPU_VIRTUALIZED.Agent flow: queries pool mode with
GetAgentPoolVirtualization, discovers a cache-local privatethunderdIP, installs viaCreateNodeEnrollment, and uninstalls on shutdown (ExecStopPost); increasesTimeoutStopSecto 60s and soft-fails with logs if unavailable.Worker behavior: selects Thunder vs physical GPUs per request, injects Thunder
LD_PRELOAD, mountsnvidia-smi,libcuda.so.1,libnvidia-ml.so.1, and the Thunder library, and uses pre-init/startup hooks so “started” is withheld until install finishes; creates/revokes client enrollments on start/teardown.Gateway: registers a new
ThunderService(Redis-backed) usinggithub.com/Thunder-Compute/thunder-sdkfor client/node enrollments and per-pool zones; exposesCreateNodeEnrollment/DeleteNodeEnrollmentvia agent-token auth and returns a clear error if Thunder is disabled; adds consistent Redis keys and a pool lock.Tooling: generates
pkg/gateway/services/thunder/thunder.proto, addsNewThunderServiceClient, and updatesbin/gen_proto.sh; updates GPU manager interfaces to accept context/request.Migration
THUNDER_API_URLandTHUNDER_API_TOKENin the gateway environment.thunderd; node enrollment runs automatically.WorkerPoolConfig.gpuVirtualized.Written for commit 5ab900b. Summary will update on new commits.