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
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
2 issues found and verified against the latest diff
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="pkg/gateway/services/thunder/db.go">
<violation number="1" location="pkg/gateway/services/thunder/db.go:85">
P1: Custom agent: **Prevent Redundant Code Duplication**
The three CRUD flows (client enrollment, node enrollment, zone) repeat nearly identical `Get`, `Save`, `Delete`, and `List` code. For example, each `SaveX` validates required fields, marshals JSON, and runs `TxPipelined(Set + SAdd)`; each `DeleteX` runs `TxPipelined(Del + SRem)`; each `ListX` does `SMembers`, `sort.Strings`, key building, and `listJSON`. Since the file already uses Go generics via `listJSON[T]`, you can extract these into shared generic helpers that take key builders and validation callbacks. This would reduce ~12 duplicated method bodies to ~4 generic helpers and prevent inconsistencies when the persistence pattern evolves.</violation>
</file>
<file name="pkg/gateway/services/thunder/service_test.go">
<violation number="1" location="pkg/gateway/services/thunder/service_test.go:21">
P2: Custom agent: **Prevent Redundant Code Duplication**
This new 854-line test file repeats two substantial blocks that could be extracted into shared test helpers. The client-enrollment setup boilerplate (newThunderRedisClient + container/worker test repositories + AddWorker for worker-1/pool-1/machine-1 + SetContainerState for container-1/workspace-1/worker-1) is duplicated nearly verbatim across at least five tests (TestServiceCreateAndDeleteClientEnrollment, TestServiceCreateClientEnrollmentReusesExistingZone, TestServiceCreateClientEnrollmentReplacesExistingToken, TestServiceCreateClientEnrollmentSucceedsWhenPreviousTokenRevokeFails, TestServicePrivateWorkerTokenIsWorkspaceScoped), differing only in Gpu/GpuCount. Likewise, the httptest mock Thunder API handler (Authorization check, POST zones, POST enrollment-token, DELETE enrollment-token revocation) is repeated with high similarity across the client and node enrollment tests, differing only in role/gpuType/gpuCount assertions. Consolidating these into helpers (e.g., a setupEnrollmentTest(t, gpu, gpuCount) helper and a reusable Thunder API stub factory) would reduce maintenance burden and reduce the risk that the tests drift out of sync with each other when the enrollment flow next changes.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| @@ -0,0 +1,262 @@ | |||
| package thunder | |||
There was a problem hiding this comment.
P1: Custom agent: Prevent Redundant Code Duplication
The three CRUD flows (client enrollment, node enrollment, zone) repeat nearly identical Get, Save, Delete, and List code. For example, each SaveX validates required fields, marshals JSON, and runs TxPipelined(Set + SAdd); each DeleteX runs TxPipelined(Del + SRem); each ListX does SMembers, sort.Strings, key building, and listJSON. Since the file already uses Go generics via listJSON[T], you can extract these into shared generic helpers that take key builders and validation callbacks. This would reduce ~12 duplicated method bodies to ~4 generic helpers and prevent inconsistencies when the persistence pattern evolves.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At pkg/gateway/services/thunder/db.go, line 85:
<comment>The three CRUD flows (client enrollment, node enrollment, zone) repeat nearly identical `Get`, `Save`, `Delete`, and `List` code. For example, each `SaveX` validates required fields, marshals JSON, and runs `TxPipelined(Set + SAdd)`; each `DeleteX` runs `TxPipelined(Del + SRem)`; each `ListX` does `SMembers`, `sort.Strings`, key building, and `listJSON`. Since the file already uses Go generics via `listJSON[T]`, you can extract these into shared generic helpers that take key builders and validation callbacks. This would reduce ~12 duplicated method bodies to ~4 generic helpers and prevent inconsistencies when the persistence pattern evolves.</comment>
<file context>
@@ -0,0 +1,262 @@
+ return &state, true, nil
+}
+
+func (r *RedisRepository) SaveClientEnrollment(ctx context.Context, state *ClientEnrollmentState, ttl time.Duration) error {
+ if state == nil || state.ContainerID == "" {
+ return errThunderStateRequired
</file context>
There was a problem hiding this comment.
2 issues found across 24 files (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="pkg/types/types.proto">
<violation number="1">
P3: When removing a field from a proto message, reserve its number (and name) so it can never be silently reused later. In the same PR you correctly did this in gateway.proto with `reserved 47;`, but here in types.proto the `gpu_virtualized = 37;` field was removed from `ContainerRequest` without a matching `reserved 37;` declaration. Since field 37 is now completely free, a future addition could accidentally take over tag 37, and any older SDK or serialized payload still carrying on-the-wire field 37 would then be decoded as that unrelated new field, silently corrupting requests. Because this is prototype wiring removed consistently on both sides the immediate impact is low, but adding the reservation now costs nothing and matches the pattern you already used in gateway.proto. Consider adding `reserved 37; reserved "gpu_virtualized";` after `task_id = 36;` and regenerating.</violation>
</file>
<file name="pkg/types/scheduler.go">
<violation number="1">
P2: Custom agent: **Prevent Redundant Code Duplication**
Removing `RequiresPhysicalGPU()` from `ContainerRequest` without adding an equivalent helper on the `Worker` duplicates the physical-vs-virtual GPU decision inline across multiple call sites. The expression `request.RequiresGPU() && !s.gpuVirtualizedForRequest(request)` now appears in `pkg/worker/lifecycle.go` at lines 298 and 1271. Extracting a shared `Worker`-level method such as `requiresPhysicalGPU(request)` would restore the centralized intent and avoid repeated inline logic.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
integrated thunder sdk with gateway
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 end-to-end with pool-level control, gateway-managed enrollments, and startup hooks so containers aren’t marked started until Thunder is installed. Adds an agent systemd stop hook to uninstall Thunder, improves scheduling/logs, and removes the TTL parameter from the Thunder Redis repository.
New Features
gpuVirtualizedto agent/worker and OpenAPI; addGetAgentPoolVirtualizationRPC; add GPU scheduling request logs.LD_PRELOAD; mountnvidia-smi,libcuda.so.1,libnvidia-ml.so.1, and Thunder lib; send metrics through the selected GPU manager; revoke enrollments and unassign on teardown.CreateNodeEnrollment/DeleteNodeEnrollmentusing cache-locality private IP; pass poolgpuVirtualized; soft-fail with logs if unavailable; add systemdExecStopPostto uninstall Thunder on shutdown.ThunderServicewith Redis-backed repository (client/node enrollments, per-pool locks, zone mappings) usinggithub.com/Thunder-Compute/thunder-sdk; register gRPC and enable agent-token RPCs in the auth interceptor.pkg/gateway/services/thunder/thunder.proto; update OpenAPI withgpuVirtualized.Migration
THUNDER_API_URLandTHUNDER_API_TOKENin the gateway environment.thunderd; node enrollment runs automatically.WorkerPoolConfig.gpuVirtualized(propagated viaWORKER_GPU_VIRTUALIZED).Written for commit f4ce2b9. Summary will update on new commits.