feat: add daemon-side lifecycle management for Bedrock serving VMs (#676)

Introduces handlers/bedrock, the daemon-side owner of vLLM serving VM
placement and lifecycle. State lives in a new JetStream KV bucket
(bedrock-endpoints) keyed {accountID}/{modelID}, cycling through
ABSENT -> STARTING -> READY -> DRAINING -> ABSENT. The daemon is the
sole KV writer; the gateway (or any other caller) requests transitions
over NATS via bedrock.endpoint.{ensure,describe,list,delete}, wired
into subscribeAll exactly as handlers/eks registers its own subjects.

Ensure is idempotent and single-flighted at two layers: a local
per-key mutex collapses concurrent same-process callers, and a
JetStream KV Create/CAS on the record's generation collapses races
across daemon replicas, so a burst of concurrent Ensure calls for one
model launches exactly one VM either way. Capacity is admission-checked
against gpu.Manager's free devices before a launch is claimed, so a
refusal never leaves a stuck STARTING record. Readiness is an explicit
poll of the serving VM's /v1/models endpoint until 200 or
startupTimeout, never a sleep.

Serving VMs boot via the system-instance path against a dedicated
system VPC, and follow the Phase-0 weights resolution: a snapshot
staged out-of-band is COW-cloned per endpoint into a data volume the
VM attaches post-boot, decoupled from the boot AMI.
This commit is contained in:
Josh Siviour
2026-08-03 15:10:55 +10:00
committed by GitHub
parent e6ac1ad9b5
commit ad8b2e5aad
24 changed files with 2680 additions and 1 deletions
+19
View File
@@ -105,6 +105,7 @@ type Config struct {
Northstar NorthstarConfig `json:"Northstar" mapstructure:"northstar"`
RDS RDSConfig `json:"RDS" mapstructure:"rds"`
ACM ACMConfig `json:"ACM" mapstructure:"acm"`
Bedrock BedrockConfig `json:"Bedrock" mapstructure:"bedrock"`
BaseDir string `json:"BaseDir" mapstructure:"base_dir"`
WalDir string `json:"WalDir" mapstructure:"wal_dir"`
@@ -201,6 +202,24 @@ type RDSConfig struct {
// supernet, so a name-hash collision can never place an RDS subnet in EKS space.
const RDSDefaultSystemVPCSupernet = "10.248.0.0/14"
// Every self-hosted vLLM serving VM's primary NIC lives in the shared Bedrock
// system VPC, mirroring RDS's DB-VM VPC: one shared VPC per region rather
// than one per endpoint, since a serving VM has no customer ENI to isolate.
type BedrockConfig struct {
// The IPv4 /14 the system VPC's /22 is carved from. It must not overlap the
// RDS or EKS control-plane supernets or any customer VPC CIDR.
SystemVPCSupernet string `json:"SystemVPCSupernet" mapstructure:"system_vpc_supernet"`
// Clamped to 1..3. Zero defaults to one, which is all a single-AZ platform
// can place across.
SystemVPCPrivateSubnets int `json:"SystemVPCPrivateSubnets" mapstructure:"system_vpc_private_subnets"`
}
// BedrockDefaultSystemVPCSupernet anchors the Bedrock system VPC address
// space at 10.244.0.0/14, immediately below and disjoint from both the RDS
// (10.248.0.0/14) and EKS control-plane (10.252.0.0/14) supernets.
const BedrockDefaultSystemVPCSupernet = "10.244.0.0/14"
// ACMConfig holds the operator-level ACM certificate-issuance configuration.
// Deliberately small: four keys, nothing deployment-specific and nothing
// derivable. In particular there is no allowed-domains list (public modes are
+82
View File
@@ -0,0 +1,82 @@
package daemon
import (
"log/slog"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
"github.com/mulgadc/spinifex/spinifex/gpu"
handlers_bedrock "github.com/mulgadc/spinifex/spinifex/handlers/bedrock"
handlers_rds "github.com/mulgadc/spinifex/spinifex/handlers/rds"
handlers_systemvpc "github.com/mulgadc/spinifex/spinifex/handlers/systemvpc"
"github.com/nats-io/nats.go/jetstream"
)
// buildBedrockLaunchDeps assembles the launch-time collaborators: VPC/volume
// plumbing mirrors buildRDSLaunchDeps exactly (same underlying EC2 handler
// set), plus the weights snapshot resolver the RDS launcher has no equivalent
// of.
func (d *Daemon) buildBedrockLaunchDeps() handlers_bedrock.LaunchDeps {
js, err := jetstream.New(d.natsConn)
var weights = gateway_bedrock.NoopWeightsResolver
if err != nil {
slog.Warn("Bedrock: jetstream init failed; serving VMs cannot resolve staged weights", "err", err)
} else {
clusterSize := 1
if d.clusterConfig != nil {
clusterSize = len(d.clusterConfig.Nodes)
}
weights = gateway_bedrock.NewWeightsStore(js, clusterSize)
}
return handlers_bedrock.LaunchDeps{
Config: d.config,
SystemVPC: handlers_systemvpc.Deps{
VPC: d.vpcService,
SG: d.vpcService,
IGW: d.igwService,
RT: d.routeTableService,
NGW: d.natGatewayService,
EIP: d.eipService,
NATSConn: d.natsConn,
},
VPC: d.vpcService,
Instance: d,
Image: d.imageService,
Volume: d.volumeService,
Attacher: handlers_rds.NewNATSVolumeAttacher(d.natsConn),
Weights: weights,
}
}
// buildBedrockServiceDeps assembles the ServiceDeps the Bedrock endpoint
// lifecycle Service needs. GPU is left at its zero value (a true nil
// interface) when no GPU manager is present: assigning a typed nil
// *daemonGPUSnapshotter instead would produce a non-nil interface wrapping a
// nil pointer, defeating admitCapacity's own "snapshotter == nil" check.
func (d *Daemon) buildBedrockServiceDeps() handlers_bedrock.ServiceDeps {
clusterSize := 1
if d.clusterConfig != nil {
clusterSize = len(d.clusterConfig.Nodes)
}
deps := handlers_bedrock.ServiceDeps{
Config: d.config,
Launch: d.buildBedrockLaunchDeps(),
NodeID: d.node,
Replicas: clusterSize,
}
if d.gpuManager != nil {
deps.GPU = &daemonGPUSnapshotter{mgr: d.gpuManager}
}
return deps
}
// daemonGPUSnapshotter adapts *gpu.Manager to handlers_bedrock's Snapshot-only
// capacity-check surface, keeping the gpu package out of that handler package
// the same way daemonGPUClaimer keeps it out of handlers_ec2_instance.
type daemonGPUSnapshotter struct {
mgr *gpu.Manager
}
func (g *daemonGPUSnapshotter) Snapshot() []gpu.PoolEntry {
return g.mgr.Snapshot()
}
+20
View File
@@ -34,6 +34,7 @@ import (
"github.com/mulgadc/spinifex/spinifex/config"
"github.com/mulgadc/spinifex/spinifex/gpu"
handlers_acm "github.com/mulgadc/spinifex/spinifex/handlers/acm"
handlers_bedrock "github.com/mulgadc/spinifex/spinifex/handlers/bedrock"
handlers_dns "github.com/mulgadc/spinifex/spinifex/handlers/dns"
handlers_ec2_account "github.com/mulgadc/spinifex/spinifex/handlers/ec2/account"
handlers_ec2_eigw "github.com/mulgadc/spinifex/spinifex/handlers/ec2/eigw"
@@ -151,6 +152,7 @@ type Daemon struct {
ecsScheduler *handlers_ecs.Scheduler
rdsService *handlers_rds.Service
rdsReconciler *handlers_rds.Reconciler
bedrockService *handlers_bedrock.Service
acmService *handlers_acm.ACMServiceImpl
acmRenewalWorker *handlers_acm.Worker
ecrMetaService *handlers_ecr.MetaServiceImpl
@@ -1084,6 +1086,18 @@ func (d *Daemon) subscribeAll() error {
)
}
// Bedrock serving-endpoint lifecycle (daemon is the sole KV writer; the
// gateway reads/requests transitions over these subjects instead of
// touching JetStream directly).
if d.bedrockService != nil {
subs = append(subs,
natsSub{handlers_bedrock.SubjectEnsureEndpoint, handleNATSRequest(d.bedrockService.Ensure), "spinifex-workers"},
natsSub{handlers_bedrock.SubjectDescribeEndpoint, handleNATSRequest(d.bedrockService.Describe), "spinifex-workers"},
natsSub{handlers_bedrock.SubjectListEndpoints, handleNATSRequest(d.bedrockService.List), "spinifex-workers"},
natsSub{handlers_bedrock.SubjectDeleteEndpoint, handleNATSRequest(d.bedrockService.Delete), "spinifex-workers"},
)
}
// ACM gateway → daemon subscriptions (certificate import and managed issuance).
if d.acmService != nil {
subs = append(subs,
@@ -1662,6 +1676,12 @@ func (d *Daemon) startCluster() error {
d.rdsReconciler.Run(d.ctx)
})
// Bedrock serving-endpoint lifecycle: no reconciler yet (no idle-reclaim or
// health sweep — out of this bead's scope), just the request-driven
// ensure/describe/list/delete surface RDS-style, constructed synchronously
// since it touches no JetStream KV until its first request.
d.bedrockService = handlers_bedrock.NewService(d.natsConn, d.buildBedrockServiceDeps())
// ACM certificates hold private keys, so unlike EKS/ECS's IAM dependency
// (which degrades to "feature disabled" without a master key) there is no
// safe degraded mode here: a missing key must fail daemon startup, not
+24
View File
@@ -187,6 +187,30 @@ func ListFoundationModels(ctx context.Context, accountID string, resolver Creden
return &bedrock.ListFoundationModelsOutput{ModelSummaries: summaries}, nil
}
// ServingSpec is the subset of a self-host catalog entry the daemon-side
// launcher needs to place and boot a serving VM. A separate type (rather than
// exporting catalogEntry) keeps the launcher from depending on the AWS-shaped
// summary/details fields this package alone renders.
type ServingSpec struct {
MinVRAMMiB int
InstanceType string
VLLMArgs []string
}
// LookupServingSpec returns modelID's serving spec. ok is false for an
// unknown model or a provider (non-self-host) entry, which has none.
func LookupServingSpec(modelID string) (spec ServingSpec, ok bool) {
entry, found := lookupCatalogEntry(modelID)
if !found || entry.Provider != tierSelfHost {
return ServingSpec{}, false
}
return ServingSpec{
MinVRAMMiB: entry.MinVRAMMiB,
InstanceType: entry.InstanceType,
VLLMArgs: entry.VLLMArgs,
}, true
}
// GetFoundationModel looks up a single model by exact modelId. Unknown
// models, and self-host models with no resolvable weights snapshot, return
// ResourceNotFoundException. A weights resolve error is an internal fault,
+67
View File
@@ -0,0 +1,67 @@
package handlers_bedrock
import (
"context"
"errors"
"fmt"
"log/slog"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/mulgadc/spinifex/spinifex/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
)
// servingAMIRoleTag identifies the single vllm-serving runtime AMI, the same
// way engineTagKey/engineVersionTagKey identify an RDS engine build. There is
// only one serving runtime today (vLLM), so no version filter is needed —
// every self-host catalog entry boots the same guest OS/runtime and differs
// only by the weights volume it is handed and its VLLMArgs.
const servingAMIRoleTag = "spinifex:bedrock-role"
// ErrServingAMINotFound is returned when no vllm-serving AMI is registered.
var ErrServingAMINotFound = errors.New("bedrock: no vllm-serving AMI found")
type amiResolver interface {
DescribeImages(ctx context.Context, input *ec2.DescribeImagesInput, accountID string) (*ec2.DescribeImagesOutput, error)
}
// resolveServingAMI finds the vllm-serving AMI, taking the most recently
// imported one if more than one is registered (a rebuild leaves the old image
// in place until an operator prunes it).
func resolveServingAMI(ctx context.Context, imgSvc amiResolver) (string, error) {
out, err := imgSvc.DescribeImages(ctx, &ec2.DescribeImagesInput{
Filters: []*ec2.Filter{
{Name: aws.String("tag:" + tags.ManagedByKey), Values: aws.StringSlice([]string{tags.ManagedByBedrock})},
{Name: aws.String("tag:" + servingAMIRoleTag), Values: aws.StringSlice([]string{"vllm-serving"})},
},
}, utils.GlobalAccountID)
if err != nil {
return "", fmt.Errorf("bedrock: describe vllm-serving AMI: %w", err)
}
if out == nil {
return "", ErrServingAMINotFound
}
var newestID, newestCreated string
matches := 0
for _, image := range out.Images {
if image == nil || aws.StringValue(image.ImageId) == "" {
continue
}
matches++
created := aws.StringValue(image.CreationDate)
if newestID == "" || created > newestCreated {
newestID = aws.StringValue(image.ImageId)
newestCreated = created
}
}
if newestID == "" {
return "", ErrServingAMINotFound
}
if matches > 1 {
slog.WarnContext(ctx, "bedrock: multiple vllm-serving AMIs registered; using newest",
"imageId", newestID, "matches", matches)
}
return newestID, nil
}
+51
View File
@@ -0,0 +1,51 @@
package handlers_bedrock
import (
"fmt"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/gpu"
)
// gpuSnapshotter reports the current GPU/MIG pool. Narrowed to Snapshot so
// this package depends on neither Claim nor Release — those happen inside the
// shared system-instance launch path, not here.
type gpuSnapshotter interface {
Snapshot() []gpu.PoolEntry
}
// checkCapacity reports whether the free pool (whole devices or MIG slices
// with Available=true) has at least one entry with minVRAMMiB free. It is a
// fast pre-admission check only: the deeper claim/allocate decision still
// happens inside the shared RunInstances path when the VM actually launches,
// so a device can still be lost to a racing launch between this check and
// that claim — the guest fails to boot in that rare case, not silently OOMs.
func checkCapacity(snapshot []gpu.PoolEntry, minVRAMMiB int) error {
if minVRAMMiB <= 0 {
return fmt.Errorf("bedrock: invalid MinVRAMMiB %d", minVRAMMiB)
}
for _, entry := range snapshot {
if !entry.Available {
continue
}
memMiB := entry.Device.MemoryMiB
if entry.MIGInstance != nil {
memMiB = entry.MIGInstance.Profile.MemoryMiB
}
if memMiB >= int64(minVRAMMiB) {
return nil
}
}
return awserrors.Errorf(awserrors.ErrorModelNotReadyException,
"bedrock: no free GPU device has %d MiB of VRAM free", minVRAMMiB)
}
// admitCapacity checks minVRAMMiB against snapshotter's current free pool.
// A nil snapshotter (no GPU manager: passthrough disabled or no GPUs present)
// admits nothing, since every self-host model in the catalog requires a GPU.
func admitCapacity(snapshotter gpuSnapshotter, minVRAMMiB int) error {
if snapshotter == nil {
return awserrors.Errorf(awserrors.ErrorModelNotReadyException, "bedrock: no GPU manager on this node")
}
return checkCapacity(snapshotter.Snapshot(), minVRAMMiB)
}
+102
View File
@@ -0,0 +1,102 @@
package handlers_bedrock
import (
"testing"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/gpu"
"github.com/stretchr/testify/assert"
)
func TestCheckCapacity(t *testing.T) {
tests := []struct {
name string
snapshot []gpu.PoolEntry
minVRAM int
wantErr bool
}{
{
name: "whole device has enough free VRAM",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: true},
},
minVRAM: 5120,
},
{
name: "MIG slice has enough free VRAM",
snapshot: []gpu.PoolEntry{
{MIGInstance: &gpu.MIGInstance{Profile: gpu.MIGProfile{MemoryMiB: 10240}}, Available: true},
},
minVRAM: 5120,
},
{
name: "device exists but not available (already claimed)",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: false},
},
minVRAM: 5120,
wantErr: true,
},
{
name: "available device too small",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 2048}, Available: true},
},
minVRAM: 5120,
wantErr: true,
},
{
name: "empty pool",
snapshot: nil,
minVRAM: 5120,
wantErr: true,
},
{
name: "picks the sufficient device among several",
snapshot: []gpu.PoolEntry{
{Device: gpu.GPUDevice{MemoryMiB: 2048}, Available: true},
{Device: gpu.GPUDevice{MemoryMiB: 16384}, Available: true},
},
minVRAM: 5120,
},
{
name: "invalid MinVRAMMiB",
snapshot: []gpu.PoolEntry{{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: true}},
minVRAM: 0,
wantErr: true,
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
err := checkCapacity(tc.snapshot, tc.minVRAM)
if tc.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}
func TestCheckCapacity_ErrorCode(t *testing.T) {
err := checkCapacity(nil, 5120)
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
}
type stubSnapshotter struct {
entries []gpu.PoolEntry
}
func (s *stubSnapshotter) Snapshot() []gpu.PoolEntry { return s.entries }
func TestAdmitCapacity_NilSnapshotter(t *testing.T) {
err := admitCapacity(nil, 5120)
assert.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
}
func TestAdmitCapacity_DelegatesToCheckCapacity(t *testing.T) {
snap := &stubSnapshotter{entries: []gpu.PoolEntry{{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: true}}}
assert.NoError(t, admitCapacity(snap, 5120))
assert.Error(t, admitCapacity(snap, 16384))
}
+359
View File
@@ -0,0 +1,359 @@
package handlers_bedrock
import (
"context"
"fmt"
"net/http"
"sync"
"sync/atomic"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/mulgadc/spinifex/spinifex/config"
"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
handlers_systemvpc "github.com/mulgadc/spinifex/spinifex/handlers/systemvpc"
)
// fakeSystemVPC is the whole EC2 VPC family as the system-VPC builder sees
// it: every describe is empty and every create succeeds, mirroring RDS's own
// launch_test.go fixture of the same name (the systemvpc package's own rules
// are covered in its package tests, not re-verified here).
type fakeSystemVPC struct {
seq int
igwCreated bool
}
var (
_ handlers_systemvpc.VPCProvisioner = (*fakeSystemVPC)(nil)
_ handlers_systemvpc.RouteTableProvisioner = (*fakeSystemVPC)(nil)
_ handlers_systemvpc.NATGatewayProvisioner = (*fakeSystemVPC)(nil)
_ handlers_systemvpc.EIPProvisioner = (*fakeSystemVPC)(nil)
_ handlers_systemvpc.IGWProvisioner = (*fakeSystemVPC)(nil)
)
func (f *fakeSystemVPC) id(prefix string) *string {
f.seq++
return aws.String(fmt.Sprintf("%s-%04d", prefix, f.seq))
}
func (f *fakeSystemVPC) deps() handlers_systemvpc.Deps {
return handlers_systemvpc.Deps{VPC: f, IGW: f, RT: f, NGW: f, EIP: f}
}
func (f *fakeSystemVPC) CreateVpc(context.Context, *ec2.CreateVpcInput, string) (*ec2.CreateVpcOutput, error) {
return &ec2.CreateVpcOutput{Vpc: &ec2.Vpc{VpcId: f.id("vpc")}}, nil
}
func (f *fakeSystemVPC) DeleteVpc(context.Context, *ec2.DeleteVpcInput, string) (*ec2.DeleteVpcOutput, error) {
return &ec2.DeleteVpcOutput{}, nil
}
func (f *fakeSystemVPC) DescribeVpcs(context.Context, *ec2.DescribeVpcsInput, string) (*ec2.DescribeVpcsOutput, error) {
return &ec2.DescribeVpcsOutput{}, nil
}
func (f *fakeSystemVPC) CreateSubnet(_ context.Context, in *ec2.CreateSubnetInput, _ string) (*ec2.CreateSubnetOutput, error) {
return &ec2.CreateSubnetOutput{Subnet: &ec2.Subnet{SubnetId: f.id("subnet-bedrocksys"), CidrBlock: in.CidrBlock}}, nil
}
func (f *fakeSystemVPC) DeleteSubnet(context.Context, *ec2.DeleteSubnetInput, string) (*ec2.DeleteSubnetOutput, error) {
return &ec2.DeleteSubnetOutput{}, nil
}
func (f *fakeSystemVPC) DescribeSubnets(context.Context, *ec2.DescribeSubnetsInput, string) (*ec2.DescribeSubnetsOutput, error) {
return &ec2.DescribeSubnetsOutput{}, nil
}
func (f *fakeSystemVPC) CreateRouteTable(context.Context, *ec2.CreateRouteTableInput, string) (*ec2.CreateRouteTableOutput, error) {
return &ec2.CreateRouteTableOutput{RouteTable: &ec2.RouteTable{RouteTableId: f.id("rtb")}}, nil
}
func (f *fakeSystemVPC) DeleteRouteTable(context.Context, *ec2.DeleteRouteTableInput, string) (*ec2.DeleteRouteTableOutput, error) {
return &ec2.DeleteRouteTableOutput{}, nil
}
func (f *fakeSystemVPC) DescribeRouteTables(context.Context, *ec2.DescribeRouteTablesInput, string) (*ec2.DescribeRouteTablesOutput, error) {
return &ec2.DescribeRouteTablesOutput{}, nil
}
func (f *fakeSystemVPC) CreateRoute(context.Context, *ec2.CreateRouteInput, string) (*ec2.CreateRouteOutput, error) {
return &ec2.CreateRouteOutput{}, nil
}
func (f *fakeSystemVPC) AssociateRouteTable(context.Context, *ec2.AssociateRouteTableInput, string) (*ec2.AssociateRouteTableOutput, error) {
return &ec2.AssociateRouteTableOutput{AssociationId: f.id("rtbassoc")}, nil
}
func (f *fakeSystemVPC) DisassociateRouteTable(context.Context, *ec2.DisassociateRouteTableInput, string) (*ec2.DisassociateRouteTableOutput, error) {
return &ec2.DisassociateRouteTableOutput{}, nil
}
func (f *fakeSystemVPC) CreateNatGateway(context.Context, *ec2.CreateNatGatewayInput, string) (*ec2.CreateNatGatewayOutput, error) {
return &ec2.CreateNatGatewayOutput{NatGateway: &ec2.NatGateway{NatGatewayId: f.id("nat")}}, nil
}
func (f *fakeSystemVPC) DeleteNatGateway(context.Context, *ec2.DeleteNatGatewayInput, string) (*ec2.DeleteNatGatewayOutput, error) {
return &ec2.DeleteNatGatewayOutput{}, nil
}
func (f *fakeSystemVPC) DescribeNatGateways(context.Context, *ec2.DescribeNatGatewaysInput, string) (*ec2.DescribeNatGatewaysOutput, error) {
return &ec2.DescribeNatGatewaysOutput{}, nil
}
func (f *fakeSystemVPC) AllocateAddress(context.Context, *ec2.AllocateAddressInput, string) (*ec2.AllocateAddressOutput, error) {
return &ec2.AllocateAddressOutput{AllocationId: f.id("eipalloc"), PublicIp: aws.String("198.51.100.9")}, nil
}
func (f *fakeSystemVPC) ReleaseAddress(context.Context, *ec2.ReleaseAddressInput, string) (*ec2.ReleaseAddressOutput, error) {
return &ec2.ReleaseAddressOutput{}, nil
}
func (f *fakeSystemVPC) CreateInternetGateway(context.Context, *ec2.CreateInternetGatewayInput, string) (*ec2.CreateInternetGatewayOutput, error) {
f.igwCreated = true
return &ec2.CreateInternetGatewayOutput{InternetGateway: &ec2.InternetGateway{InternetGatewayId: aws.String("igw-bedrocksys")}}, nil
}
func (f *fakeSystemVPC) AttachInternetGateway(context.Context, *ec2.AttachInternetGatewayInput, string) (*ec2.AttachInternetGatewayOutput, error) {
return &ec2.AttachInternetGatewayOutput{}, nil
}
func (f *fakeSystemVPC) DetachInternetGateway(context.Context, *ec2.DetachInternetGatewayInput, string) (*ec2.DetachInternetGatewayOutput, error) {
return &ec2.DetachInternetGatewayOutput{}, nil
}
func (f *fakeSystemVPC) DeleteInternetGateway(context.Context, *ec2.DeleteInternetGatewayInput, string) (*ec2.DeleteInternetGatewayOutput, error) {
return &ec2.DeleteInternetGatewayOutput{}, nil
}
// DescribeInternetGateways answers the post-attach lookup: the builder
// attaches an IGW and then re-reads it, so the pre-create lookup must find
// nothing (or no IGW is ever created) and only the post-attach one answered.
func (f *fakeSystemVPC) DescribeInternetGateways(_ context.Context, in *ec2.DescribeInternetGatewaysInput, _ string) (*ec2.DescribeInternetGatewaysOutput, error) {
if !f.igwCreated {
return &ec2.DescribeInternetGatewaysOutput{}, nil
}
return &ec2.DescribeInternetGatewaysOutput{InternetGateways: []*ec2.InternetGateway{{
InternetGatewayId: aws.String("igw-bedrocksys"),
Attachments: []*ec2.InternetGatewayAttachment{{VpcId: in.Filters[0].Values[0]}},
}}}, nil
}
// testModelID is the real self-host catalog entry (MinVRAMMiB=5120), used
// rather than a fabricated one so tests exercise LookupServingSpec against
// the actual catalog data instead of a parallel test fixture.
const testModelID = "meta.llama3-2-1b-instruct-v1:0"
// fakeVPC is an in-memory launchVPCProvisioner: every CreateNetworkInterface
// call returns a fresh ENI ID/IP so concurrent launches never collide.
type fakeVPC struct {
mu sync.Mutex
nextID int
created []string
deleted []string
detached []string
}
func (f *fakeVPC) CreateNetworkInterface(_ context.Context, _ *ec2.CreateNetworkInterfaceInput, _ string) (*ec2.CreateNetworkInterfaceOutput, error) {
f.mu.Lock()
f.nextID++
id := f.nextID
eniID := fmt.Sprintf("eni-%d", id)
f.created = append(f.created, eniID)
f.mu.Unlock()
return &ec2.CreateNetworkInterfaceOutput{NetworkInterface: &ec2.NetworkInterface{
NetworkInterfaceId: aws.String(fmt.Sprintf("eni-%d", id)),
PrivateIpAddress: aws.String(fmt.Sprintf("10.244.1.%d", id%250+1)),
MacAddress: aws.String("02:00:00:00:00:01"),
}}, nil
}
func (f *fakeVPC) DeleteNetworkInterface(_ context.Context, in *ec2.DeleteNetworkInterfaceInput, _ string) (*ec2.DeleteNetworkInterfaceOutput, error) {
f.mu.Lock()
f.deleted = append(f.deleted, aws.StringValue(in.NetworkInterfaceId))
f.mu.Unlock()
return &ec2.DeleteNetworkInterfaceOutput{}, nil
}
func (f *fakeVPC) DetachENI(_ context.Context, _, eniID string) error {
f.mu.Lock()
f.detached = append(f.detached, eniID)
f.mu.Unlock()
return nil
}
// fakeVolume is an in-memory launchVolumeProvisioner.
type fakeVolume struct {
mu sync.Mutex
nextID int
created []*ec2.CreateVolumeInput
deleted []string
// failCreate, when set, makes CreateVolume return this error instead.
failCreate error
}
func (f *fakeVolume) CreateVolume(_ context.Context, in *ec2.CreateVolumeInput, _ string) (*ec2.Volume, error) {
f.mu.Lock()
f.created = append(f.created, in)
f.mu.Unlock()
if f.failCreate != nil {
return nil, f.failCreate
}
f.mu.Lock()
f.nextID++
id := f.nextID
f.mu.Unlock()
return &ec2.Volume{VolumeId: aws.String(fmt.Sprintf("vol-%d", id))}, nil
}
func (f *fakeVolume) DeleteVolume(_ context.Context, in *ec2.DeleteVolumeInput, _ string) (*ec2.DeleteVolumeOutput, error) {
f.mu.Lock()
f.deleted = append(f.deleted, aws.StringValue(in.VolumeId))
f.mu.Unlock()
return &ec2.DeleteVolumeOutput{}, nil
}
// fakeAttacher is an in-memory volumeAttacher. failErr, when set, makes
// AttachVolume fail so launch_test.go can exercise the post-boot rollback path.
type fakeAttacher struct {
failErr error
}
func (f *fakeAttacher) AttachVolume(_ context.Context, _, _, _, device string) (string, error) {
if f.failErr != nil {
return "", f.failErr
}
return device, nil
}
// fakeInstanceLauncher counts launches, so the concurrency test can assert
// exactly one VM was launched regardless of how many Ensure calls raced.
type fakeInstanceLauncher struct {
mu sync.Mutex
launchCount atomic.Int32
requests []*sysinstance.SystemInstanceInput
terminated []string
nextID int
failLaunch error
}
var _ sysinstance.SystemInstanceLauncher = (*fakeInstanceLauncher)(nil)
func (f *fakeInstanceLauncher) LaunchSystemInstance(in *sysinstance.SystemInstanceInput) (*sysinstance.SystemInstanceOutput, error) {
f.launchCount.Add(1)
f.mu.Lock()
f.requests = append(f.requests, in)
f.mu.Unlock()
if f.failLaunch != nil {
return nil, f.failLaunch
}
f.mu.Lock()
f.nextID++
id := f.nextID
f.mu.Unlock()
return &sysinstance.SystemInstanceOutput{InstanceID: fmt.Sprintf("i-%d", id)}, nil
}
// lastInput returns the most recent LaunchSystemInstance request, for tests
// that assert on how a launch was wired.
func (f *fakeInstanceLauncher) lastInput() *sysinstance.SystemInstanceInput {
f.mu.Lock()
defer f.mu.Unlock()
if len(f.requests) == 0 {
return nil
}
return f.requests[len(f.requests)-1]
}
func (f *fakeInstanceLauncher) launches() []*sysinstance.SystemInstanceInput {
f.mu.Lock()
defer f.mu.Unlock()
return f.requests
}
func (f *fakeInstanceLauncher) TerminateSystemInstance(instanceID string) error {
f.mu.Lock()
f.terminated = append(f.terminated, instanceID)
f.mu.Unlock()
return nil
}
// fakeAMIResolver always resolves the same serving AMI.
type fakeAMIResolver struct{}
func (fakeAMIResolver) DescribeImages(_ context.Context, _ *ec2.DescribeImagesInput, _ string) (*ec2.DescribeImagesOutput, error) {
return &ec2.DescribeImagesOutput{Images: []*ec2.Image{{
ImageId: aws.String("ami-vllm-serving"),
CreationDate: aws.String("2026-01-01T00:00:00.000Z"),
}}}, nil
}
// fakeWeightsResolver resolves a fixed snapshot ID for every model unless
// resolvable is false, mimicking "no weights staged".
type fakeWeightsResolver struct {
snapshotID string
resolvable bool
err error
}
func (f fakeWeightsResolver) Resolve(_ context.Context, _ string) (string, bool, error) {
return f.snapshotID, f.resolvable, f.err
}
// stubRoundTripper returns statusCode for every request, letting readiness
// tests skip real networking entirely.
type stubRoundTripper struct {
statusCode int32 // atomic: TestEnsure tests flip this to simulate boot delay
}
func (s *stubRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
code := int(atomic.LoadInt32(&s.statusCode))
return &http.Response{
StatusCode: code,
Body: http.NoBody,
Header: make(http.Header),
}, nil
}
// launchHarness bundles every LaunchDeps fake so a test can reach into any of
// them, mirroring handlers_rds's launchHarness.
type launchHarness struct {
sysvpc *fakeSystemVPC
vpc *fakeVPC
launcher *fakeInstanceLauncher
images fakeAMIResolver
volumes *fakeVolume
attacher *fakeAttacher
weights fakeWeightsResolver
}
func newLaunchHarness() *launchHarness {
return &launchHarness{
sysvpc: &fakeSystemVPC{},
vpc: &fakeVPC{},
launcher: &fakeInstanceLauncher{},
volumes: &fakeVolume{},
attacher: &fakeAttacher{},
weights: fakeWeightsResolver{snapshotID: "snap-llama32-1b", resolvable: true},
}
}
func (h *launchHarness) deps() LaunchDeps {
return LaunchDeps{
Config: &config.Config{Region: "us-east-1", AZ: "us-east-1a"},
SystemVPC: h.sysvpc.deps(),
VPC: h.vpc,
Instance: h.launcher,
Image: h.images,
Volume: h.volumes,
Attacher: h.attacher,
Weights: h.weights,
}
}
func testLaunchInput() LaunchInput {
return LaunchInput{
ModelID: testModelID,
InstanceType: "g5.xlarge",
VLLMArgs: []string{"--dtype=bfloat16"},
}
}
+336
View File
@@ -0,0 +1,336 @@
package handlers_bedrock
import (
"context"
"errors"
"fmt"
"log/slog"
"net"
"slices"
"strconv"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/config"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
"github.com/mulgadc/spinifex/spinifex/handlers/sysinstance"
handlers_systemvpc "github.com/mulgadc/spinifex/spinifex/handlers/systemvpc"
"github.com/mulgadc/spinifex/spinifex/tags"
"github.com/mulgadc/spinifex/spinifex/utils"
)
// vllmServePort is the vLLM OpenAI-compatible server's listen port, baked
// into the vllm-serving AMI's systemd unit. Both the readiness probe and
// real inference traffic (from the awsgw process) dial this port on the
// serving VM's system-VPC ENI private IP.
const vllmServePort = 8000
// weightsVolumeDevice is where a serving VM's COW-cloned weights volume is
// attached. Fixed like RDS's dataVolumeDevice: the guest locates it by
// device, not by name, and one VM never carries more than one weights volume.
const weightsVolumeDevice = "/dev/sdf"
// bedrockInstanceTagKey links an ENI or volume back to its serving VM, so a
// teardown or leak sweep can find them without the KV record.
const bedrockInstanceTagKey = "spinifex:bedrock-endpoint"
// A fresh budget rather than the caller's remainder: the unwind is often
// triggered by that deadline expiring, and a dead context deletes nothing.
const rollbackTimeout = 60 * time.Second
// ErrNoWeightsStaged is returned when no weights snapshot has been staged for
// a model that otherwise has a serving spec — an operator must PutWeights
// before the model is servable.
var ErrNoWeightsStaged = errors.New("bedrock: no weights snapshot staged for this model")
type launchVPCProvisioner interface {
CreateNetworkInterface(ctx context.Context, input *ec2.CreateNetworkInterfaceInput, accountID string) (*ec2.CreateNetworkInterfaceOutput, error)
DeleteNetworkInterface(ctx context.Context, input *ec2.DeleteNetworkInterfaceInput, accountID string) (*ec2.DeleteNetworkInterfaceOutput, error)
DetachENI(ctx context.Context, accountID, eniID string) error
}
type launchVolumeProvisioner interface {
CreateVolume(ctx context.Context, input *ec2.CreateVolumeInput, accountID string) (*ec2.Volume, error)
DeleteVolume(ctx context.Context, input *ec2.DeleteVolumeInput, accountID string) (*ec2.DeleteVolumeOutput, error)
}
// Split from launchVolumeProvisioner because attach is routed to the node
// owning the VM, not answered by whichever node picks the request up.
type volumeAttacher interface {
AttachVolume(ctx context.Context, accountID, instanceID, volumeID, device string) (string, error)
}
type LaunchDeps struct {
Config *config.Config
SystemVPC handlers_systemvpc.Deps
VPC launchVPCProvisioner
Instance sysinstance.SystemInstanceLauncher
Image amiResolver
Volume launchVolumeProvisioner
Attacher volumeAttacher
Weights gateway_bedrock.WeightsResolver
}
// LaunchInput describes the serving VM to launch. Everything here is already
// validated and admission-checked by the caller (Service.Ensure): the launch
// helper resolves and wires, it does not re-police capacity or catalog rules.
type LaunchInput struct {
ModelID string
InstanceType string
VLLMArgs []string
}
type LaunchOutput struct {
InstanceID string
ENIID string
PrivateIP string
WeightsVolumeID string
// BaseURL is the vLLM server's base address, derived from PrivateIP.
BaseURL string
// Tears down everything this launch created, for a caller whose readiness
// probe times out after it returned. The launch runs it itself on its own
// failures, so it is only ever invoked once.
Unwind func(context.Context)
}
// LaunchServingVM boots one self-hosted vLLM serving VM for in.ModelID. On any
// failure every resource this call created is torn down, so a retried Ensure
// does not accumulate orphan ENIs, volumes and VMs.
func LaunchServingVM(ctx context.Context, deps LaunchDeps, in LaunchInput) (out *LaunchOutput, err error) {
if err := validateLaunchInput(in); err != nil {
return nil, err
}
region, az := deps.Config.Region, deps.Config.AZ
amiID, err := resolveServingAMI(ctx, deps.Image)
if err != nil {
return nil, err
}
snapshotID, resolvable, err := deps.Weights.Resolve(ctx, in.ModelID)
if err != nil {
return nil, fmt.Errorf("bedrock: resolve weights for %s: %w", in.ModelID, err)
}
if !resolvable {
return nil, fmt.Errorf("%w: %s", ErrNoWeightsStaged, in.ModelID)
}
sysRefs, err := EnsureSystemVPC(ctx, deps.SystemVPC, &deps.Config.Bedrock, utils.GlobalAccountID, region)
if err != nil {
return nil, err
}
systemSubnetID := sysRefs.PrivateSubnetIDs[0]
// Unwind in reverse creation order on any failure below. Each step appends
// its own undo as soon as the resource exists.
var rollback []func(context.Context)
// The VM is torn down first regardless of creation order: the ENI and
// weights volume are attached while it runs, and deleting those is
// rejected as InUse.
var terminateVM func(context.Context)
// A context detached from the caller's, so a failure caused by the
// caller's own deadline expiring can still clean up.
unwind := func(ctx context.Context) {
rbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), rollbackTimeout)
defer cancel()
if terminateVM != nil {
terminateVM(rbCtx)
}
for _, v := range slices.Backward(rollback) {
v(rbCtx)
}
}
defer func() {
if err != nil {
unwind(ctx)
}
}()
eni, err := createLaunchENI(ctx, deps.VPC, systemSubnetID, in.ModelID)
if err != nil {
return nil, err
}
rollback = append(rollback, func(ctx context.Context) {
deleteLaunchENI(ctx, deps.VPC, eni.id)
})
weightsVolumeID, err := createWeightsVolume(ctx, deps.Volume, az, snapshotID, in.ModelID)
if err != nil {
return nil, err
}
rollback = append(rollback, func(ctx context.Context) {
if _, delErr := deps.Volume.DeleteVolume(ctx, &ec2.DeleteVolumeInput{VolumeId: aws.String(weightsVolumeID)}, utils.GlobalAccountID); delErr != nil {
slog.WarnContext(ctx, "bedrock: rollback delete of orphaned weights volume failed",
"model", in.ModelID, "volumeId", weightsVolumeID, "err", delErr)
}
})
userData := buildServeUserData(serveUserDataInput{
ModelID: in.ModelID,
VLLMArgs: in.VLLMArgs,
WeightsDevice: weightsVolumeDevice,
ServePort: vllmServePort,
})
sysOut, err := deps.Instance.LaunchSystemInstance(&sysinstance.SystemInstanceInput{
BootMode: sysinstance.BootAMI,
ManagedBy: tags.ManagedByBedrock,
InstanceType: in.InstanceType,
ImageID: amiID,
// The VM and its primary NIC live in the system account: a serving VM
// has no customer owner of its own, only the accountID the endpoint
// record is keyed under (a separate, KV-level concern from VM ownership).
AccountID: utils.GlobalAccountID,
SubnetID: systemSubnetID,
ENIID: eni.id,
ENIMac: eni.mac,
ENIIP: eni.ip,
UserData: userData,
})
if err != nil {
return nil, fmt.Errorf("bedrock: launch serving VM for %s: %w", in.ModelID, err)
}
if sysOut == nil || sysOut.InstanceID == "" {
return nil, fmt.Errorf("bedrock: launch serving VM for %s: launcher returned no instance", in.ModelID)
}
instanceID := sysOut.InstanceID
terminateVM = func(ctx context.Context) {
if termErr := deps.Instance.TerminateSystemInstance(instanceID); termErr != nil &&
!errors.Is(termErr, sysinstance.ErrSystemInstanceNotFound) {
slog.WarnContext(ctx, "bedrock: rollback terminate of failed serving VM failed",
"model", in.ModelID, "instanceId", instanceID, "err", termErr)
}
}
device, err := deps.Attacher.AttachVolume(ctx, utils.GlobalAccountID, instanceID, weightsVolumeID, weightsVolumeDevice)
if err != nil {
return nil, fmt.Errorf("bedrock: attach weights volume %s to %s: %w", weightsVolumeID, instanceID, err)
}
slog.InfoContext(ctx, "bedrock: serving VM launched",
"model", in.ModelID, "instanceId", instanceID, "ami", amiID,
"eni", eni.id, "eniIp", eni.ip, "weightsVolume", weightsVolumeID, "device", device)
return &LaunchOutput{
InstanceID: instanceID,
ENIID: eni.id,
PrivateIP: eni.ip,
WeightsVolumeID: weightsVolumeID,
BaseURL: "http://" + net.JoinHostPort(eni.ip, strconv.Itoa(vllmServePort)),
Unwind: unwind,
}, nil
}
// TerminateServingVM tears down a DRAINING endpoint's VM, ENI and weights
// volume. The VM is terminated first (ENI/volume deletes would otherwise be
// rejected as InUse), and every step is best-effort and logged rather than
// aborted on the first failure, so one stuck resource does not block the
// others from being reclaimed. The caller decides whether the returned error
// should block the KV record's DRAINING->ABSENT transition.
func TerminateServingVM(ctx context.Context, deps LaunchDeps, rec EndpointRecord) error {
var errs []error
if rec.InstanceID != "" {
if err := deps.Instance.TerminateSystemInstance(rec.InstanceID); err != nil &&
!errors.Is(err, sysinstance.ErrSystemInstanceNotFound) {
errs = append(errs, fmt.Errorf("terminate instance %s: %w", rec.InstanceID, err))
}
}
if rec.ENIID != "" {
deleteLaunchENI(ctx, deps.VPC, rec.ENIID)
}
if rec.WeightsVolumeID != "" {
if _, err := deps.Volume.DeleteVolume(ctx, &ec2.DeleteVolumeInput{VolumeId: aws.String(rec.WeightsVolumeID)}, utils.GlobalAccountID); err != nil &&
!awserrors.IsNotFound(err) {
errs = append(errs, fmt.Errorf("delete weights volume %s: %w", rec.WeightsVolumeID, err))
}
}
return errors.Join(errs...)
}
func validateLaunchInput(in LaunchInput) error {
switch {
case in.ModelID == "":
return errors.New("bedrock: LaunchServingVM empty model id")
case in.InstanceType == "":
return errors.New("bedrock: LaunchServingVM empty instance type")
}
return nil
}
type launchENI struct {
id string
ip string
mac string
}
// The system NIC is unreachable from any customer VPC — a serving VM has no
// customer-facing surface — so it needs no security group of its own.
func createLaunchENI(ctx context.Context, vpcSvc launchVPCProvisioner, subnetID, modelID string) (*launchENI, error) {
out, err := vpcSvc.CreateNetworkInterface(ctx, &ec2.CreateNetworkInterfaceInput{
SubnetId: aws.String(subnetID),
Description: aws.String("Bedrock serving NIC for " + modelID),
TagSpecifications: []*ec2.TagSpecification{{
ResourceType: aws.String("network-interface"),
Tags: []*ec2.Tag{
{Key: aws.String(tags.ManagedByKey), Value: aws.String(tags.ManagedByBedrock)},
{Key: aws.String(bedrockInstanceTagKey), Value: aws.String(modelID)},
},
}},
}, utils.GlobalAccountID)
if err != nil {
return nil, fmt.Errorf("bedrock: create ENI in subnet %s: %w", subnetID, err)
}
if out == nil || out.NetworkInterface == nil ||
aws.StringValue(out.NetworkInterface.NetworkInterfaceId) == "" ||
aws.StringValue(out.NetworkInterface.PrivateIpAddress) == "" {
return nil, fmt.Errorf("bedrock: create ENI in subnet %s: incomplete interface returned", subnetID)
}
ni := out.NetworkInterface
return &launchENI{
id: aws.StringValue(ni.NetworkInterfaceId),
ip: aws.StringValue(ni.PrivateIpAddress),
mac: aws.StringValue(ni.MacAddress),
}, nil
}
// Detach comes first because a terminated VM can leave the attachment record
// behind, which a plain delete rejects as InUse.
func deleteLaunchENI(ctx context.Context, vpcSvc launchVPCProvisioner, eniID string) {
if err := vpcSvc.DetachENI(ctx, utils.GlobalAccountID, eniID); err != nil && !awserrors.IsNotFound(err) {
slog.DebugContext(ctx, "bedrock: rollback ENI detach failed", "eniId", eniID, "err", err)
}
if _, err := vpcSvc.DeleteNetworkInterface(ctx, &ec2.DeleteNetworkInterfaceInput{
NetworkInterfaceId: aws.String(eniID),
}, utils.GlobalAccountID); err != nil && !awserrors.IsNotFound(err) {
slog.WarnContext(ctx, "bedrock: rollback delete of orphaned ENI failed", "eniId", eniID, "err", err)
}
}
// createWeightsVolume COW-clones modelID's staged weights snapshot into a
// fresh volume this endpoint's VM attaches — Phase-0's bake-into-a-volume,
// CreateSnapshot, COW-clone-per-endpoint resolution.
func createWeightsVolume(ctx context.Context, volSvc launchVolumeProvisioner, az, snapshotID, modelID string) (string, error) {
volume, err := volSvc.CreateVolume(ctx, &ec2.CreateVolumeInput{
AvailabilityZone: aws.String(az),
SnapshotId: aws.String(snapshotID),
VolumeType: aws.String("gp3"),
TagSpecifications: []*ec2.TagSpecification{{
ResourceType: aws.String("volume"),
Tags: []*ec2.Tag{
{Key: aws.String(tags.ManagedByKey), Value: aws.String(tags.ManagedByBedrock)},
{Key: aws.String(bedrockInstanceTagKey), Value: aws.String(modelID)},
},
}},
}, utils.GlobalAccountID)
if err != nil {
return "", fmt.Errorf("bedrock: clone weights volume for %s from snapshot %s: %w", modelID, snapshotID, err)
}
if volume == nil || aws.StringValue(volume.VolumeId) == "" {
return "", fmt.Errorf("bedrock: clone weights volume for %s: empty volume id", modelID)
}
return aws.StringValue(volume.VolumeId), nil
}
+123
View File
@@ -0,0 +1,123 @@
package handlers_bedrock
import (
"errors"
"testing"
"github.com/aws/aws-sdk-go/aws"
"github.com/mulgadc/spinifex/spinifex/tags"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLaunchServingVM_WiresLaunchInput(t *testing.T) {
h := newLaunchHarness()
out, err := LaunchServingVM(t.Context(), h.deps(), testLaunchInput())
require.NoError(t, err)
require.NotNil(t, out)
in := h.launcher.lastInput()
require.NotNil(t, in)
assert.Equal(t, tags.ManagedByBedrock, in.ManagedBy)
assert.Equal(t, "ami-vllm-serving", in.ImageID)
assert.Equal(t, "g5.xlarge", in.InstanceType)
assert.NotEmpty(t, in.ENIID)
assert.NotEmpty(t, in.ENIIP)
assert.Contains(t, in.UserData, "VLLM_MODEL_ID="+testModelID)
assert.Contains(t, in.UserData, "VLLM_ARGS=--dtype=bfloat16")
assert.Equal(t, out.ENIID, in.ENIID)
assert.Equal(t, out.PrivateIP, in.ENIIP)
assert.Equal(t, "http://"+out.PrivateIP+":8000", out.BaseURL)
assert.NotEmpty(t, out.WeightsVolumeID)
assert.NotEmpty(t, out.InstanceID)
}
func TestLaunchServingVM_ClonesWeightsSnapshotIntoAVolume(t *testing.T) {
h := newLaunchHarness()
h.weights = fakeWeightsResolver{snapshotID: "snap-abc123", resolvable: true}
out, err := LaunchServingVM(t.Context(), h.deps(), testLaunchInput())
require.NoError(t, err)
require.Len(t, h.volumes.created, 1)
assert.Equal(t, "snap-abc123", aws.StringValue(h.volumes.created[0].SnapshotId))
assert.NotContains(t, h.volumes.deleted, out.WeightsVolumeID, "must not have been rolled back on success")
}
func TestLaunchServingVM_NoWeightsStaged(t *testing.T) {
h := newLaunchHarness()
h.weights = fakeWeightsResolver{resolvable: false}
_, err := LaunchServingVM(t.Context(), h.deps(), testLaunchInput())
require.Error(t, err)
assert.ErrorIs(t, err, ErrNoWeightsStaged)
assert.Empty(t, h.launcher.launches(), "no VM may be launched without staged weights")
}
func TestLaunchServingVM_RejectsIncompleteInput(t *testing.T) {
for name, mutate := range map[string]func(*LaunchInput){
"no model id": func(in *LaunchInput) { in.ModelID = "" },
"no instance type": func(in *LaunchInput) { in.InstanceType = "" },
} {
t.Run(name, func(t *testing.T) {
h := newLaunchHarness()
in := testLaunchInput()
mutate(&in)
_, err := LaunchServingVM(t.Context(), h.deps(), in)
require.Error(t, err)
assert.Zero(t, h.vpc.nextID, "must fail before touching any provisioner")
})
}
}
func TestLaunchServingVM_RollsBackOnVMLaunchFailure(t *testing.T) {
h := newLaunchHarness()
h.launcher.failLaunch = errors.New("no capacity")
_, err := LaunchServingVM(t.Context(), h.deps(), testLaunchInput())
require.Error(t, err)
// The ENI and the cloned weights volume must not outlive the failed launch.
assert.NotEmpty(t, h.vpc.deleted)
assert.NotEmpty(t, h.volumes.deleted)
}
func TestLaunchServingVM_RollsBackOnAttachFailure(t *testing.T) {
h := newLaunchHarness()
h.attacher = &fakeAttacher{failErr: errors.New("no free hot-plug port")}
_, err := LaunchServingVM(t.Context(), h.deps(), testLaunchInput())
require.Error(t, err)
assert.NotEmpty(t, h.launcher.terminated, "a VM with nothing to serve must not be left running")
assert.NotEmpty(t, h.vpc.deleted)
assert.NotEmpty(t, h.volumes.deleted)
}
func TestTerminateServingVM_TerminatesDetachesAndDeletes(t *testing.T) {
h := newLaunchHarness()
rec := EndpointRecord{
InstanceID: "i-42",
ENIID: "eni-42",
WeightsVolumeID: "vol-42",
}
err := TerminateServingVM(t.Context(), h.deps(), rec)
require.NoError(t, err)
assert.Equal(t, []string{"i-42"}, h.launcher.terminated)
assert.Contains(t, h.vpc.detached, "eni-42")
assert.Contains(t, h.vpc.deleted, "eni-42")
assert.Contains(t, h.volumes.deleted, "vol-42")
}
func TestTerminateServingVM_EmptyRecordIsANoop(t *testing.T) {
h := newLaunchHarness()
err := TerminateServingVM(t.Context(), h.deps(), EndpointRecord{})
require.NoError(t, err)
assert.Empty(t, h.launcher.terminated)
assert.Empty(t, h.vpc.deleted)
assert.Empty(t, h.volumes.deleted)
}
+62
View File
@@ -0,0 +1,62 @@
package handlers_bedrock
import (
"context"
"fmt"
"net/http"
"time"
)
// defaultStartupTimeout bounds how long Ensure's launch goroutine polls a
// freshly-launched serving VM's /v1/models endpoint before giving up and
// reverting the record to ABSENT. PROVISIONAL: chosen conservatively for a
// small model (Llama 3.2 1B) cold-booting a microVM plus vLLM's own weight
// load; needs correcting against a real measurement on wattle once the model
// actually runs there.
const defaultStartupTimeout = 5 * time.Minute
// readinessPollInterval is the spacing between /v1/models probes. Short
// enough that a fast-booting VM isn't held back by the poll cadence, long
// enough not to hammer a guest still bringing its network stack up.
const readinessPollInterval = 2 * time.Second
// waitReady polls baseURL + "/v1/models" every pollInterval until it returns
// HTTP 200 or ctx's deadline (set by the caller from startupTimeout) expires.
// Never a sleep in place of a real check: every non-200 (including connection
// refused, while the guest is still booting) is swallowed and retried, since
// readiness is the very thing being waited for. pollInterval is a parameter
// (not just the readinessPollInterval constant) so tests can poll fast
// without waiting out the production cadence.
func waitReady(ctx context.Context, client *http.Client, baseURL string, pollInterval time.Duration) error {
url := baseURL + "/v1/models"
ticker := time.NewTicker(pollInterval)
defer ticker.Stop()
for {
if probeOnce(ctx, client, url) {
return nil
}
select {
case <-ctx.Done():
return fmt.Errorf("bedrock: readiness probe of %s: %w", url, ctx.Err())
case <-ticker.C:
}
}
}
// probeOnce issues a single GET and reports whether it returned HTTP 200.
// Any transport error (guest not listening yet) or non-200 status counts as
// not-yet-ready, not a failure — the caller's context deadline is what
// eventually turns "not yet" into an error.
func probeOnce(ctx context.Context, client *http.Client, url string) bool {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return false
}
resp, err := client.Do(req)
if err != nil {
return false
}
defer resp.Body.Close()
return resp.StatusCode == http.StatusOK
}
@@ -0,0 +1,71 @@
package handlers_bedrock
import (
"context"
"net/http"
"net/http/httptest"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestWaitReady_ImmediateSuccess(t *testing.T) {
var requests atomic.Int32
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
assert.Equal(t, "/v1/models", r.URL.Path)
w.WriteHeader(http.StatusOK)
}))
defer ts.Close()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
err := waitReady(ctx, ts.Client(), ts.URL, 10*time.Millisecond)
require.NoError(t, err)
assert.GreaterOrEqual(t, requests.Load(), int32(1))
}
func TestWaitReady_NonOKThenOK(t *testing.T) {
var requests atomic.Int32
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
n := requests.Add(1)
if n < 3 {
w.WriteHeader(http.StatusServiceUnavailable)
return
}
w.WriteHeader(http.StatusOK)
}))
defer ts.Close()
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
err := waitReady(ctx, ts.Client(), ts.URL, 10*time.Millisecond)
require.NoError(t, err)
assert.GreaterOrEqual(t, requests.Load(), int32(3))
}
func TestWaitReady_TimeoutNeverReady(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusServiceUnavailable)
}))
defer ts.Close()
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Millisecond)
defer cancel()
err := waitReady(ctx, ts.Client(), ts.URL, 10*time.Millisecond)
require.Error(t, err)
assert.ErrorIs(t, err, context.DeadlineExceeded)
}
func TestWaitReady_ConnectionRefusedRetriedNotFailed(t *testing.T) {
// Nothing listens on this URL: every attempt is a transport error, which
// waitReady must treat as "not yet ready" and keep polling, not fail fast.
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Millisecond)
defer cancel()
err := waitReady(ctx, http.DefaultClient, "http://127.0.0.1:1", 10*time.Millisecond)
require.Error(t, err)
assert.ErrorIs(t, err, context.DeadlineExceeded)
}
+412
View File
@@ -0,0 +1,412 @@
package handlers_bedrock
import (
"context"
"errors"
"fmt"
"log/slog"
"net/http"
"sync"
"time"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/config"
gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock"
"github.com/mulgadc/spinifex/spinifex/utils"
"github.com/nats-io/nats.go"
"github.com/nats-io/nats.go/jetstream"
)
// EndpointService is the serving-endpoint lifecycle surface, implemented both
// by the daemon-side Service (the sole KV writer, D2) and by
// NATSEndpointService (the NATS-forwarding client gateway-side code and tests
// use instead of importing this package's JetStream/launch dependencies).
type EndpointService interface {
Ensure(ctx context.Context, in *EnsureEndpointInput, accountID string) (*EnsureEndpointOutput, error)
Describe(ctx context.Context, in *DescribeEndpointInput, accountID string) (*DescribeEndpointOutput, error)
List(ctx context.Context, in *ListEndpointsInput, accountID string) (*ListEndpointsOutput, error)
Delete(ctx context.Context, in *DeleteEndpointInput, accountID string) (*DeleteEndpointOutput, error)
}
// ServiceDeps are the daemon-supplied collaborators the Service needs beyond
// NATS: the launch/VPC/volume plumbing and the GPU capacity view.
type ServiceDeps struct {
Config *config.Config
Launch LaunchDeps
GPU gpuSnapshotter
// NodeID identifies which daemon replica's node a launched VM lands on
// (stamped into EndpointRecord.NodeID).
NodeID string
// StartupTimeout bounds the readiness probe; zero takes defaultStartupTimeout.
StartupTimeout time.Duration
// HTTPClient is the readiness-probe client; nil takes a zero-value
// *http.Client (context-deadline-bound, no extra per-request timeout needed).
HTTPClient *http.Client
// PollInterval overrides readinessPollInterval; zero takes the default.
// Test-only knob so a readiness test does not wait out the production cadence.
PollInterval time.Duration
// Replicas sizes the bedrock-endpoints KV bucket on first create.
Replicas int
}
// Service is the daemon-side, KV-backed handler set for the serving-endpoint
// lifecycle. One per daemon; every replica subscribes the same NATS subjects
// on a shared queue group, so exactly one replica handles a given request, but
// any replica may read/write any key — hence the KV CAS layer beneath the
// local per-key mutex (see ensureMu).
type Service struct {
nc *nats.Conn
deps ServiceDeps
// ensureMu collapses concurrent Ensure calls FOR THE SAME KEY on THIS
// replica into one KV round trip: without it, N concurrent local callers
// all hit kv.Create and only the loser count changes, wasting N-1 network
// round trips on every burst. It does nothing for concurrent Ensures
// landing on different replicas — that race is only ever resolved by the
// KV Create/CAS below, which is why both layers are required.
ensureMu keyMutex
// launchWG tracks in-flight async launches for test determinism (WaitLaunches).
launchWG sync.WaitGroup
}
var _ EndpointService = (*Service)(nil)
// NewService constructs a Service bound to nc, using deps for launch/capacity.
func NewService(nc *nats.Conn, deps ServiceDeps) *Service {
return &Service{nc: nc, deps: deps}
}
// WaitLaunches blocks until every async launch this Service has started so
// far has finished. Test-only.
func (s *Service) WaitLaunches() { s.launchWG.Wait() }
func (s *Service) js() (jetstream.JetStream, error) {
if s.nc == nil {
return nil, errors.New("bedrock service: nil nats connection")
}
return jetstream.New(s.nc)
}
func (s *Service) bucket(ctx context.Context) (jetstream.KeyValue, error) {
js, err := s.js()
if err != nil {
return nil, err
}
return GetOrCreateEndpointsBucket(ctx, js, s.deps.Replicas)
}
func (s *Service) startupTimeout() time.Duration {
if s.deps.StartupTimeout > 0 {
return s.deps.StartupTimeout
}
return defaultStartupTimeout
}
func (s *Service) pollInterval() time.Duration {
if s.deps.PollInterval > 0 {
return s.deps.PollInterval
}
return readinessPollInterval
}
func (s *Service) httpClient() *http.Client {
if s.deps.HTTPClient != nil {
return s.deps.HTTPClient
}
return &http.Client{}
}
// Ensure idempotently guarantees modelID has a running or starting endpoint.
// A model with no serving spec (unknown, or a provider entry) is rejected
// before any KV touch. Capacity is admission-checked before the claim so a
// refusal never leaves a STARTING record for a concurrent caller to observe.
func (s *Service) Ensure(ctx context.Context, in *EnsureEndpointInput, _ string) (*EnsureEndpointOutput, error) {
if in == nil || in.ModelID == "" {
return nil, errors.New(awserrors.ErrorInvalidParameterValue)
}
spec, ok := gateway_bedrock.LookupServingSpec(in.ModelID)
if !ok {
return nil, errors.New(awserrors.ErrorResourceNotFoundException)
}
// Serving endpoints are shared platform infra today, not per-tenant: every
// account's Ensure for the same model converges on one shared VM, stored
// under the system account. The {accountID}/{modelID} key shape stays
// ready for per-tenant serving pools without a schema change later.
storeAccountID := utils.GlobalAccountID
key := EndpointKey(storeAccountID, in.ModelID)
unlock := s.ensureMu.lock(key)
defer unlock()
kv, err := s.bucket(ctx)
if err != nil {
return nil, err
}
var existing EndpointRecord
found, err := getJSON(ctx, kv, key, &existing)
if err != nil {
return nil, fmt.Errorf("bedrock: read endpoint %s: %w", in.ModelID, err)
}
if found {
switch existing.State {
case StateStarting, StateReady:
return &EnsureEndpointOutput{Endpoint: existing}, nil
case StateDraining:
return nil, awserrors.Errorf(awserrors.ErrorModelNotReadyException,
"bedrock: endpoint for %s is draining", in.ModelID)
}
}
if err := admitCapacity(s.deps.GPU, spec.MinVRAMMiB); err != nil {
return nil, err
}
if err := validateTransition(StateAbsent, StateStarting); err != nil {
return nil, err
}
rec := EndpointRecord{
AccountID: storeAccountID,
ModelID: in.ModelID,
State: StateStarting,
NodeID: s.deps.NodeID,
CreatedAt: time.Now().UTC(),
Generation: 1,
}
if _, err := createJSONRevision(ctx, kv, key, rec); err != nil {
if errors.Is(err, jetstream.ErrKeyExists) {
// Lost the cross-replica race after the read above: some other
// replica's Ensure claimed the key in between. Its record is the
// answer, not an error.
var winner EndpointRecord
if ok, gerr := getJSON(ctx, kv, key, &winner); gerr == nil && ok {
return &EnsureEndpointOutput{Endpoint: winner}, nil
}
}
return nil, fmt.Errorf("bedrock: claim endpoint %s: %w", in.ModelID, err)
}
// The launch outlives the request, so it runs on its own background
// context rather than the caller's, which is cancelled once this reply is
// written.
launchCtx := context.Background()
s.launchWG.Go(func() {
defer func() {
if r := recover(); r != nil {
slog.ErrorContext(launchCtx, "bedrock: launch panic", "model", in.ModelID, "panic", r)
s.abortLaunch(launchCtx, key, in.ModelID)
}
}()
s.runLaunch(launchCtx, key, rec, spec)
})
return &EnsureEndpointOutput{Endpoint: rec}, nil
}
// runLaunch performs the slow launch + readiness probe and writes the
// terminal state. Only this goroutine writes key while the record is
// STARTING (guaranteed by Ensure returning early for any concurrent caller
// that observes STARTING), so no CAS-conflict retry loop is needed here.
func (s *Service) runLaunch(ctx context.Context, key string, rec EndpointRecord, spec gateway_bedrock.ServingSpec) {
out, err := LaunchServingVM(ctx, s.deps.Launch, LaunchInput{
ModelID: rec.ModelID,
InstanceType: spec.InstanceType,
VLLMArgs: spec.VLLMArgs,
})
if err != nil {
slog.ErrorContext(ctx, "bedrock: launch serving VM failed", "model", rec.ModelID, "err", err)
s.abortLaunch(ctx, key, rec.ModelID)
return
}
rec.InstanceID = out.InstanceID
rec.ENIID = out.ENIID
rec.WeightsVolumeID = out.WeightsVolumeID
rec.BaseURL = out.BaseURL
timeoutCtx, cancel := context.WithTimeout(ctx, s.startupTimeout())
defer cancel()
if err := waitReady(timeoutCtx, s.httpClient(), rec.BaseURL, s.pollInterval()); err != nil {
slog.ErrorContext(ctx, "bedrock: readiness probe timed out; unwinding launch",
"model", rec.ModelID, "instanceId", rec.InstanceID, "err", err)
out.Unwind(ctx)
s.abortLaunch(ctx, key, rec.ModelID)
return
}
if err := validateTransition(StateStarting, StateReady); err != nil {
slog.ErrorContext(ctx, "bedrock: illegal transition to READY", "model", rec.ModelID, "err", err)
return
}
rec.State = StateReady
rec.ReadyAt = time.Now().UTC()
rec.Generation++
kv, err := s.bucket(ctx)
if err != nil {
slog.ErrorContext(ctx, "bedrock: bucket unavailable to record READY", "model", rec.ModelID, "err", err)
return
}
_, rev, gerr := readCurrent(ctx, kv, key)
if gerr != nil {
slog.ErrorContext(ctx, "bedrock: re-read endpoint before READY write failed", "model", rec.ModelID, "err", gerr)
return
}
if err := updateJSON(ctx, kv, key, rev, rec); err != nil {
slog.ErrorContext(ctx, "bedrock: CAS write of READY state failed", "model", rec.ModelID, "err", err)
}
}
// abortLaunch reverts a STARTING record to ABSENT (deletes the key) so a
// retried Ensure gets a clean claim instead of a dead record. Runs on the
// launch's own background context: the reader that triggered this failure
// may already be gone, but the revert must still happen.
func (s *Service) abortLaunch(ctx context.Context, key, modelID string) {
if err := validateTransition(StateStarting, StateAbsent); err != nil {
slog.ErrorContext(ctx, "bedrock: illegal abort transition", "model", modelID, "err", err)
return
}
kv, err := s.bucket(ctx)
if err != nil {
slog.ErrorContext(ctx, "bedrock: bucket unavailable to abort launch", "model", modelID, "err", err)
return
}
if err := deleteJSON(ctx, kv, key); err != nil {
slog.ErrorContext(ctx, "bedrock: revert of failed launch failed; record may be stuck STARTING",
"model", modelID, "err", err)
}
}
// readCurrent re-reads key's current record and revision.
func readCurrent(ctx context.Context, kv jetstream.KeyValue, key string) (EndpointRecord, uint64, error) {
var rec EndpointRecord
rev, found, err := getJSONRevision(ctx, kv, key, &rec)
if err != nil {
return EndpointRecord{}, 0, err
}
if !found {
return EndpointRecord{}, 0, fmt.Errorf("bedrock: endpoint key %s vanished mid-launch", key)
}
return rec, rev, nil
}
// Describe returns modelID's current endpoint record, or a synthetic ABSENT
// record when none exists — never an error, since "not provisioned" is a
// normal answer, not a fault.
func (s *Service) Describe(ctx context.Context, in *DescribeEndpointInput, _ string) (*DescribeEndpointOutput, error) {
if in == nil || in.ModelID == "" {
return nil, errors.New(awserrors.ErrorInvalidParameterValue)
}
storeAccountID := utils.GlobalAccountID
kv, err := s.bucket(ctx)
if err != nil {
return nil, err
}
var rec EndpointRecord
found, err := getJSON(ctx, kv, EndpointKey(storeAccountID, in.ModelID), &rec)
if err != nil {
return nil, fmt.Errorf("bedrock: describe endpoint %s: %w", in.ModelID, err)
}
if !found {
rec = EndpointRecord{AccountID: storeAccountID, ModelID: in.ModelID, State: StateAbsent}
}
return &DescribeEndpointOutput{Endpoint: rec}, nil
}
// List returns every endpoint record in the shared endpoints bucket.
func (s *Service) List(ctx context.Context, _ *ListEndpointsInput, _ string) (*ListEndpointsOutput, error) {
kv, err := s.bucket(ctx)
if err != nil {
return nil, err
}
recs, err := ListEndpoints(ctx, kv, utils.GlobalAccountID)
if err != nil {
return nil, err
}
return &ListEndpointsOutput{Endpoints: recs}, nil
}
// Delete moves a READY endpoint to DRAINING, tears down its VM/ENI/weights
// volume, then deletes the record. Runs synchronously: unlike Ensure this is
// an explicit, infrequent operator action, so there is no client-latency
// reason to background it. An already-ABSENT endpoint is a no-op success.
func (s *Service) Delete(ctx context.Context, in *DeleteEndpointInput, _ string) (*DeleteEndpointOutput, error) {
if in == nil || in.ModelID == "" {
return nil, errors.New(awserrors.ErrorInvalidParameterValue)
}
storeAccountID := utils.GlobalAccountID
key := EndpointKey(storeAccountID, in.ModelID)
unlock := s.ensureMu.lock(key)
defer unlock()
kv, err := s.bucket(ctx)
if err != nil {
return nil, err
}
rec, rev, found, err := getFullJSON(ctx, kv, key)
if err != nil {
return nil, fmt.Errorf("bedrock: read endpoint %s: %w", in.ModelID, err)
}
if !found {
return &DeleteEndpointOutput{}, nil
}
if rec.State != StateReady {
return nil, awserrors.Errorf(awserrors.ErrorModelNotReadyException,
"bedrock: endpoint for %s is not READY (state=%s)", in.ModelID, rec.State)
}
if err := validateTransition(rec.State, StateDraining); err != nil {
return nil, err
}
rec.State = StateDraining
rec.Generation++
if err := updateJSON(ctx, kv, key, rev, rec); err != nil {
return nil, fmt.Errorf("bedrock: mark endpoint %s draining: %w", in.ModelID, err)
}
if err := TerminateServingVM(ctx, s.deps.Launch, rec); err != nil {
return nil, fmt.Errorf("bedrock: delete endpoint %s: %w", in.ModelID, err)
}
if err := validateTransition(StateDraining, StateAbsent); err != nil {
return nil, err
}
if err := deleteJSON(ctx, kv, key); err != nil {
return nil, fmt.Errorf("bedrock: remove endpoint %s record: %w", in.ModelID, err)
}
return &DeleteEndpointOutput{}, nil
}
// getFullJSON is getJSONRevision with the found flag surfaced alongside the
// value and revision, for callers (Delete) that branch on all three.
func getFullJSON(ctx context.Context, kv jetstream.KeyValue, key string) (EndpointRecord, uint64, bool, error) {
var rec EndpointRecord
rev, found, err := getJSONRevision(ctx, kv, key, &rec)
return rec, rev, found, err
}
// keyMutex hands out a per-key *sync.Mutex, creating it on first use. Entries
// are never removed: the key space is the model catalog, which is small and
// static, so the map cannot grow unbounded.
type keyMutex struct {
mu sync.Mutex
locks map[string]*sync.Mutex
}
// lock blocks until key's mutex is held and returns the matching unlock func.
func (k *keyMutex) lock(key string) func() {
k.mu.Lock()
m, ok := k.locks[key]
if !ok {
m = &sync.Mutex{}
if k.locks == nil {
k.locks = make(map[string]*sync.Mutex)
}
k.locks[key] = m
}
k.mu.Unlock()
m.Lock()
return m.Unlock
}
+46
View File
@@ -0,0 +1,46 @@
package handlers_bedrock
import (
"context"
"time"
"github.com/mulgadc/spinifex/spinifex/utils"
"github.com/nats-io/nats.go"
)
// defaultNATSTimeout bounds a request-reply round trip. Ensure's own slow
// work (VM launch + readiness probe) happens after the daemon has already
// replied STARTING, so this timeout only needs to cover the claim, not the
// launch.
const defaultNATSTimeout = 30 * time.Second
// NATSEndpointService is the NATS-forwarding EndpointService client: the
// gateway process (or any other caller off the daemon's own node) uses this
// instead of importing handlers_bedrock's JetStream/launch dependencies.
type NATSEndpointService struct {
nc *nats.Conn
timeout time.Duration
}
var _ EndpointService = (*NATSEndpointService)(nil)
// NewNATSEndpointService constructs a client bound to nc.
func NewNATSEndpointService(nc *nats.Conn) *NATSEndpointService {
return &NATSEndpointService{nc: nc, timeout: defaultNATSTimeout}
}
func (s *NATSEndpointService) Ensure(ctx context.Context, in *EnsureEndpointInput, accountID string) (*EnsureEndpointOutput, error) {
return utils.NATSRequest[EnsureEndpointOutput](ctx, s.nc, SubjectEnsureEndpoint, in, s.timeout, accountID)
}
func (s *NATSEndpointService) Describe(ctx context.Context, in *DescribeEndpointInput, accountID string) (*DescribeEndpointOutput, error) {
return utils.NATSRequest[DescribeEndpointOutput](ctx, s.nc, SubjectDescribeEndpoint, in, s.timeout, accountID)
}
func (s *NATSEndpointService) List(ctx context.Context, in *ListEndpointsInput, accountID string) (*ListEndpointsOutput, error) {
return utils.NATSRequest[ListEndpointsOutput](ctx, s.nc, SubjectListEndpoints, in, s.timeout, accountID)
}
func (s *NATSEndpointService) Delete(ctx context.Context, in *DeleteEndpointInput, accountID string) (*DeleteEndpointOutput, error) {
return utils.NATSRequest[DeleteEndpointOutput](ctx, s.nc, SubjectDeleteEndpoint, in, s.timeout, accountID)
}
@@ -0,0 +1,123 @@
package handlers_bedrock
import (
"context"
"encoding/json"
"net/http"
"testing"
"time"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/mulgadc/spinifex/spinifex/utils"
"github.com/nats-io/nats.go"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// stubNATSHandler mirrors daemon.handleNATSRequest's unmarshal -> service ->
// marshal/error -> respond wire contract, without importing the daemon
// package (which would pull in the whole cluster runtime for one handler
// shape). It exists purely so this test exercises the exact bytes
// NATSEndpointService and the real daemon subscription put on the wire.
func stubNATSHandler[I any, O any](serviceFn func(context.Context, *I, string) (*O, error)) nats.MsgHandler {
return func(msg *nats.Msg) {
accountID := utils.AccountIDFromMsg(msg)
input := new(I)
if errResp := utils.UnmarshalJsonPayload(input, msg.Data); errResp != nil {
_ = msg.Respond(errResp)
return
}
output, err := serviceFn(context.Background(), input, accountID)
if err != nil {
payload := utils.GenerateErrorPayloadWithMessage(awserrors.ValidErrorCodeFromError(err), err.Error())
_ = msg.Respond(payload)
return
}
data, err := json.Marshal(output)
if err != nil {
_ = msg.Respond(utils.GenerateErrorPayloadWithMessage(awserrors.ErrorServerInternal, err.Error()))
return
}
_ = msg.Respond(data)
}
}
// subscribeService registers every EndpointService subject on the
// "spinifex-workers" queue group, exactly as daemon.subscribeAll does for the
// real Service, so the client-side round trip in these tests exercises the
// production subject names and queue group.
func subscribeService(t *testing.T, nc *nats.Conn, svc EndpointService) {
t.Helper()
subs := []struct {
subject string
handler nats.MsgHandler
}{
{SubjectEnsureEndpoint, stubNATSHandler(svc.Ensure)},
{SubjectDescribeEndpoint, stubNATSHandler(svc.Describe)},
{SubjectListEndpoints, stubNATSHandler(svc.List)},
{SubjectDeleteEndpoint, stubNATSHandler(svc.Delete)},
}
for _, s := range subs {
sub, err := nc.QueueSubscribe(s.subject, "spinifex-workers", s.handler)
require.NoError(t, err)
t.Cleanup(func() { _ = sub.Unsubscribe() })
}
}
func TestNATSEndpointService_EnsureDescribeListDelete_RoundTrip(t *testing.T) {
h := newLaunchHarness()
svc, nc := newTestService(t, h, http.StatusOK, sufficientGPU())
subscribeService(t, nc, svc)
client := NewNATSEndpointService(nc)
ensureOut, err := client.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "000000000000")
require.NoError(t, err)
assert.Equal(t, StateStarting, ensureOut.Endpoint.State)
assert.Equal(t, testModelID, ensureOut.Endpoint.ModelID)
svc.WaitLaunches()
descOut, err := client.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "000000000000")
require.NoError(t, err)
assert.Equal(t, StateReady, descOut.Endpoint.State)
assert.NotEmpty(t, descOut.Endpoint.BaseURL)
listOut, err := client.List(t.Context(), &ListEndpointsInput{}, "000000000000")
require.NoError(t, err)
require.Len(t, listOut.Endpoints, 1)
assert.Equal(t, testModelID, listOut.Endpoints[0].ModelID)
_, err = client.Delete(t.Context(), &DeleteEndpointInput{ModelID: testModelID}, "000000000000")
require.NoError(t, err)
descOut, err = client.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "000000000000")
require.NoError(t, err)
assert.Equal(t, StateAbsent, descOut.Endpoint.State)
}
// TestNATSEndpointService_EnsureError_CarriesCodeAndMessage confirms a
// service-side awserrors.Errorf refusal survives the NATS round trip with
// both its wire code and its human message intact, not just a bare code.
func TestNATSEndpointService_EnsureError_CarriesCodeAndMessage(t *testing.T) {
h := newLaunchHarness()
svc, nc := newTestService(t, h, http.StatusOK, &stubSnapshotter{}) // no free capacity
subscribeService(t, nc, svc)
client := NewNATSEndpointService(nc)
_, err := client.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "000000000000")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
assert.Contains(t, err.Error(), "VRAM", "the refusal's human message must survive the NATS round trip, not just its code")
}
func TestNATSEndpointService_NoResponder(t *testing.T) {
_, nc, _ := testutil.StartTestJetStream(t)
client := NewNATSEndpointService(nc)
ctx, cancel := context.WithTimeout(t.Context(), 2*time.Second)
defer cancel()
_, err := client.Describe(ctx, &DescribeEndpointInput{ModelID: testModelID}, "000000000000")
require.Error(t, err)
}
+235
View File
@@ -0,0 +1,235 @@
package handlers_bedrock
import (
"net/http"
"sync"
"testing"
"time"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/config"
"github.com/mulgadc/spinifex/spinifex/gpu"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/mulgadc/spinifex/spinifex/utils"
"github.com/nats-io/nats.go"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// sufficientGPU is a GPU snapshot with plenty of room for testModelID's
// 5120 MiB MinVRAMMiB floor.
func sufficientGPU() *stubSnapshotter {
return &stubSnapshotter{entries: []gpu.PoolEntry{{Device: gpu.GPUDevice{MemoryMiB: 8192}, Available: true}}}
}
// newTestService wires a Service against an embedded JetStream server and h's
// fakes, with a short startup timeout/poll interval so readiness tests never
// wait out anything close to production timing. statusCode drives every
// readiness probe response via a stubRoundTripper.
func newTestService(t *testing.T, h *launchHarness, statusCode int32, gpuSnap gpuSnapshotter) (*Service, *nats.Conn) {
t.Helper()
_, nc, _ := testutil.StartTestJetStream(t)
deps := ServiceDeps{
Config: &config.Config{Region: "us-east-1", AZ: "us-east-1a"},
Launch: h.deps(),
GPU: gpuSnap,
NodeID: "node-1",
StartupTimeout: 200 * time.Millisecond,
HTTPClient: &http.Client{Transport: &stubRoundTripper{statusCode: statusCode}},
PollInterval: 5 * time.Millisecond,
Replicas: 1,
}
return NewService(nc, deps), nc
}
func TestEnsure_ConcurrentSameProcessLaunchesExactlyOneVM(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, sufficientGPU())
const n = 20
var wg sync.WaitGroup
outs := make([]*EnsureEndpointOutput, n)
errs := make([]error, n)
for i := range n {
wg.Add(1)
go func(i int) {
defer wg.Done()
out, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
outs[i], errs[i] = out, err
}(i)
}
wg.Wait()
s.WaitLaunches()
for i := range n {
require.NoError(t, errs[i])
require.NotNil(t, outs[i])
}
assert.EqualValues(t, 1, h.launcher.launchCount.Load())
desc, err := s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.Equal(t, StateReady, desc.Endpoint.State)
}
// TestEnsure_CrossReplicaCASCollapsesToOneLaunch exercises the layer the
// per-process ensureMu cannot: two independent Service instances (standing in
// for two daemon replicas) sharing one JetStream backend. Only the KV
// Create/CAS in store.go decides the winner here.
func TestEnsure_CrossReplicaCASCollapsesToOneLaunch(t *testing.T) {
h := newLaunchHarness()
_, nc, _ := testutil.StartTestJetStream(t)
newDeps := func() ServiceDeps {
return ServiceDeps{
Launch: h.deps(),
GPU: sufficientGPU(),
StartupTimeout: 200 * time.Millisecond,
HTTPClient: &http.Client{Transport: &stubRoundTripper{statusCode: http.StatusOK}},
PollInterval: 5 * time.Millisecond,
Replicas: 1,
}
}
s1 := NewService(nc, newDeps())
s2 := NewService(nc, newDeps())
var wg sync.WaitGroup
var err1, err2 error
wg.Add(2)
go func() {
defer wg.Done()
_, err1 = s1.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
}()
go func() {
defer wg.Done()
_, err2 = s2.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
}()
wg.Wait()
s1.WaitLaunches()
s2.WaitLaunches()
require.NoError(t, err1)
require.NoError(t, err2)
assert.EqualValues(t, 1, h.launcher.launchCount.Load())
}
func TestEnsure_CapacityRefusal(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, &stubSnapshotter{})
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
assert.Zero(t, h.launcher.launchCount.Load())
desc, err := s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.Equal(t, StateAbsent, desc.Endpoint.State, "a refused admission must never leave a STARTING record")
}
func TestEnsure_UnknownModelRejected(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, sufficientGPU())
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: "no-such-model"}, "")
require.Error(t, err)
assert.Zero(t, h.launcher.launchCount.Load())
}
func TestEnsure_ReachesReady(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, sufficientGPU())
out, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.Equal(t, StateStarting, out.Endpoint.State)
s.WaitLaunches()
desc, err := s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.Equal(t, StateReady, desc.Endpoint.State)
assert.NotEmpty(t, desc.Endpoint.BaseURL)
assert.False(t, desc.Endpoint.ReadyAt.IsZero())
assert.EqualValues(t, 2, desc.Endpoint.Generation)
}
func TestEnsure_ReadinessTimeoutAbortsToAbsent(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusServiceUnavailable, sufficientGPU())
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
s.WaitLaunches()
desc, err := s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.Equal(t, StateAbsent, desc.Endpoint.State, "a launch that never becomes ready must revert, not stick STARTING")
assert.NotEmpty(t, h.launcher.terminated, "the VM that never became ready must be unwound")
assert.NotEmpty(t, h.vpc.deleted)
assert.NotEmpty(t, h.volumes.deleted)
}
func TestDelete_ReadyToAbsent(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, sufficientGPU())
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
s.WaitLaunches()
desc, err := s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
require.Equal(t, StateReady, desc.Endpoint.State)
instanceID := desc.Endpoint.InstanceID
_, err = s.Delete(t.Context(), &DeleteEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
desc, err = s.Describe(t.Context(), &DescribeEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.Equal(t, StateAbsent, desc.Endpoint.State)
assert.Contains(t, h.launcher.terminated, instanceID)
}
func TestDelete_IdempotentOnAbsent(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, sufficientGPU())
out, err := s.Delete(t.Context(), &DeleteEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
assert.NotNil(t, out)
assert.Empty(t, h.launcher.terminated)
}
func TestDelete_RefusesNonReadyState(t *testing.T) {
h := newLaunchHarness()
s, nc := newTestService(t, h, http.StatusOK, sufficientGPU())
js := testutil.NewJetStream(t, nc)
kv, err := GetOrCreateEndpointsBucket(t.Context(), js, 1)
require.NoError(t, err)
key := EndpointKey(utils.GlobalAccountID, testModelID)
_, err = createJSONRevision(t.Context(), kv, key, EndpointRecord{
AccountID: utils.GlobalAccountID, ModelID: testModelID, State: StateStarting, Generation: 1,
})
require.NoError(t, err)
_, err = s.Delete(t.Context(), &DeleteEndpointInput{ModelID: testModelID}, "")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorModelNotReadyException, awserrors.ValidErrorCodeFromError(err))
}
func TestList_ReturnsAllEnsuredEndpoints(t *testing.T) {
h := newLaunchHarness()
s, _ := newTestService(t, h, http.StatusOK, sufficientGPU())
_, err := s.Ensure(t.Context(), &EnsureEndpointInput{ModelID: testModelID}, "")
require.NoError(t, err)
s.WaitLaunches()
out, err := s.List(t.Context(), &ListEndpointsInput{}, "")
require.NoError(t, err)
require.Len(t, out.Endpoints, 1)
assert.Equal(t, testModelID, out.Endpoints[0].ModelID)
}
+84
View File
@@ -0,0 +1,84 @@
// Package handlers_bedrock owns placement and lifecycle of self-hosted vLLM
// serving VMs: it decides where a model's endpoint runs, launches and probes
// the VM, and persists the result as the single source of truth other
// components (the gateway's inference path) read over NATS.
package handlers_bedrock
import (
"errors"
"fmt"
"slices"
"time"
)
// EndpointState is the lifecycle state of one model's serving endpoint.
type EndpointState string
const (
// StateAbsent is the zero value: no record exists (or it was deleted). Not
// itself stored — its "record" is simply a missing KV key.
StateAbsent EndpointState = "ABSENT"
// StateStarting covers admission through VM launch and the readiness probe.
StateStarting EndpointState = "STARTING"
// StateReady means the readiness probe observed a 200 from /v1/models;
// BaseURL is safe to route inference traffic to.
StateReady EndpointState = "READY"
// StateDraining means a delete was requested; the VM is being torn down.
StateDraining EndpointState = "DRAINING"
)
// ErrIllegalTransition is returned by validateTransition for any state change
// not in the table below.
var ErrIllegalTransition = errors.New("bedrock: illegal endpoint state transition")
// legalTransitions enumerates every allowed from->to move. A STARTING that
// fails (launch error or readiness timeout) goes back to ABSENT — the record
// is deleted rather than parked in a terminal FAILED state — so a retried
// Ensure gets a clean claim instead of piling up dead records.
var legalTransitions = map[EndpointState][]EndpointState{
StateAbsent: {StateStarting},
StateStarting: {StateReady, StateAbsent},
StateReady: {StateDraining},
StateDraining: {StateAbsent},
}
// validateTransition reports whether moving from -> to is legal.
func validateTransition(from, to EndpointState) error {
if slices.Contains(legalTransitions[from], to) {
return nil
}
return fmt.Errorf("%w: %s -> %s", ErrIllegalTransition, from, to)
}
// EndpointRecord is the KV-persisted state of one model's serving endpoint.
// JSON tags are explicit as this struct crosses NATS subjects (DescribeEndpoint).
type EndpointRecord struct {
AccountID string `json:"account_id"`
ModelID string `json:"model_id"`
State EndpointState `json:"state"`
InstanceID string `json:"instance_id,omitempty"`
// NodeID is the cluster node the VM landed on, so a caller can route a
// node-scoped follow-up (e.g. a terminate) without a fleet-wide fan-out.
NodeID string `json:"node_id,omitempty"`
// BaseURL is the vLLM OpenAI-compatible server's base address
// ("http://{eniIP}:{port}"), set once STARTING reaches READY.
BaseURL string `json:"base_url,omitempty"`
// ENIID and WeightsVolumeID are not in the bead's minimum field list but
// are required to tear the VM down cleanly on DRAINING->ABSENT: without
// them a delete has no way to find and release these resources again.
ENIID string `json:"eni_id,omitempty"`
WeightsVolumeID string `json:"weights_volume_id,omitempty"`
CreatedAt time.Time `json:"created_at,omitzero"`
ReadyAt time.Time `json:"ready_at,omitzero"`
// Pinned exempts this endpoint from the (not-yet-built) idle-reclaim sweep.
Pinned bool `json:"pinned,omitempty"`
// Generation increments on every write and is the CAS single-flight token:
// a concurrent Ensure that lost the local-mutex race re-reads the record and
// checks Generation before deciding whether a launch is still needed.
Generation uint64 `json:"generation"`
}
+55
View File
@@ -0,0 +1,55 @@
package handlers_bedrock
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestValidateTransition(t *testing.T) {
allStates := []EndpointState{StateAbsent, StateStarting, StateReady, StateDraining}
legal := map[EndpointState]map[EndpointState]bool{
StateAbsent: {StateStarting: true},
StateStarting: {StateReady: true, StateAbsent: true},
StateReady: {StateDraining: true},
StateDraining: {StateAbsent: true},
}
for _, from := range allStates {
for _, to := range allStates {
t.Run(string(from)+"->"+string(to), func(t *testing.T) {
err := validateTransition(from, to)
if legal[from][to] {
assert.NoError(t, err)
} else {
assert.Error(t, err)
assert.ErrorIs(t, err, ErrIllegalTransition)
}
})
}
}
}
// A few explicitly illegal transitions worth naming, so a future edit to the
// table's shape (e.g. adding a state) breaks a readable test, not just the
// generated matrix above.
func TestValidateTransition_IllegalCases(t *testing.T) {
cases := []struct {
name string
from EndpointState
to EndpointState
}{
{"ready cannot restart", StateReady, StateStarting},
{"absent cannot jump to ready", StateAbsent, StateReady},
{"draining cannot go ready", StateDraining, StateReady},
{"absent cannot self-transition", StateAbsent, StateAbsent},
{"starting cannot re-enter starting", StateStarting, StateStarting},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
err := validateTransition(tc.from, tc.to)
assert.ErrorIs(t, err, ErrIllegalTransition)
})
}
}
+143
View File
@@ -0,0 +1,143 @@
package handlers_bedrock
import (
"context"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"strings"
"github.com/mulgadc/spinifex/spinifex/kvutil"
"github.com/nats-io/nats.go/jetstream"
)
// KVBucketEndpoints is the single JetStream KV bucket holding every serving
// endpoint's state, across all accounts. Unlike EKS/RDS there is no
// per-account bucket: self-host serving VMs are shared platform infra, not
// per-tenant resources, so one bucket keyed {accountID}/{modelID} is enough
// and avoids growing a bucket per account that never runs a serving endpoint.
const KVBucketEndpoints = "bedrock-endpoints"
// KVBucketEndpointsHistory pins one revision per key: a superseded state is
// not meaningful to retain, only the current one CAS-guards against races.
const KVBucketEndpointsHistory = 1
// EndpointKey returns the KV key for accountID's modelID endpoint record.
// Model IDs contain ':' (e.g. "meta.llama3-2-1b-instruct-v1:0"), which NATS
// rejects in a KV key, so the segment is base64url-encoded, mirroring
// gateway_bedrock's weightsKey.
func EndpointKey(accountID, modelID string) string {
return accountID + "/" + base64.RawURLEncoding.EncodeToString([]byte(modelID))
}
// endpointsPrefix returns the KV key prefix under which accountID's endpoint
// records live, for ListEndpoints to enumerate.
func endpointsPrefix(accountID string) string {
return accountID + "/"
}
// GetOrCreateEndpointsBucket returns the shared bedrock-endpoints KV bucket,
// creating it on first use at the given replica count (clamped to a minimum
// of 1).
func GetOrCreateEndpointsBucket(ctx context.Context, js jetstream.JetStream, replicas int) (jetstream.KeyValue, error) {
kv, err := kvutil.GetOrCreateBucketWithReplicas(ctx, js, KVBucketEndpoints, KVBucketEndpointsHistory, replicas)
if err != nil {
return nil, fmt.Errorf("bedrock: create endpoints KV bucket: %w", err)
}
return kv, nil
}
// Returns (false, nil) when the key is absent.
func getJSON(ctx context.Context, kv jetstream.KeyValue, key string, out any) (bool, error) {
entry, err := kv.Get(ctx, key)
if errors.Is(err, jetstream.ErrKeyNotFound) {
return false, nil
}
if err != nil {
return false, err
}
if err := json.Unmarshal(entry.Value(), out); err != nil {
return false, fmt.Errorf("unmarshal %s: %w", key, err)
}
return true, nil
}
// getJSON plus the entry revision, for callers that follow with a CAS update.
func getJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, out any) (uint64, bool, error) {
entry, err := kv.Get(ctx, key)
if errors.Is(err, jetstream.ErrKeyNotFound) {
return 0, false, nil
}
if err != nil {
return 0, false, err
}
if err := json.Unmarshal(entry.Value(), out); err != nil {
return 0, false, fmt.Errorf("unmarshal %s: %w", key, err)
}
return entry.Revision(), true, nil
}
// createJSONRevision writes v at key only if nothing is stored there,
// returning the created entry's revision. Returns jetstream.ErrKeyExists when
// the key is already taken — the caller's cue to fall back to the read+CAS path.
func createJSONRevision(ctx context.Context, kv jetstream.KeyValue, key string, v any) (uint64, error) {
data, err := json.Marshal(v)
if err != nil {
return 0, err
}
rev, err := kv.Create(ctx, key, data)
if err != nil {
return 0, err
}
return rev, nil
}
// updateJSON writes v at key only if the stored entry is still at rev — the
// CAS primitive single-flight and state transitions build on.
func updateJSON(ctx context.Context, kv jetstream.KeyValue, key string, rev uint64, v any) error {
data, err := json.Marshal(v)
if err != nil {
return err
}
if _, err := kv.Update(ctx, key, data, rev); err != nil {
return fmt.Errorf("update %s: %w", key, err)
}
return nil
}
// deleteJSON removes key unconditionally. Purge (not just Delete) so no
// stale revision survives to confuse a subsequent createJSONRevision's
// "does the key exist" check under History=1.
func deleteJSON(ctx context.Context, kv jetstream.KeyValue, key string) error {
if err := kv.Purge(ctx, key); err != nil && !errors.Is(err, jetstream.ErrKeyNotFound) {
return fmt.Errorf("purge %s: %w", key, err)
}
return nil
}
// ListEndpoints returns every endpoint record under accountID's key prefix.
func ListEndpoints(ctx context.Context, kv jetstream.KeyValue, accountID string) ([]EndpointRecord, error) {
keys, err := kvutil.Keys(ctx, kv)
if err != nil {
return nil, fmt.Errorf("bedrock: list endpoint keys: %w", err)
}
prefix := endpointsPrefix(accountID)
out := make([]EndpointRecord, 0, len(keys))
for _, key := range keys {
if !strings.HasPrefix(key, prefix) {
continue
}
var rec EndpointRecord
ok, getErr := getJSON(ctx, kv, key, &rec)
if getErr != nil {
return nil, fmt.Errorf("bedrock: read endpoint %s: %w", key, getErr)
}
if !ok {
// Deleted between Keys() and Get() — not an error, just gone.
continue
}
out = append(out, rec)
}
return out, nil
}
+96
View File
@@ -0,0 +1,96 @@
package handlers_bedrock
import (
"testing"
"github.com/mulgadc/spinifex/spinifex/testutil"
"github.com/nats-io/nats.go/jetstream"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestEndpointKey_Base64EncodesColonBearingModelID(t *testing.T) {
key := EndpointKey("000000000000", "meta.llama3-2-1b-instruct-v1:0")
assert.Equal(t, "000000000000/"+"bWV0YS5sbGFtYTMtMi0xYi1pbnN0cnVjdC12MTow", key)
}
func newTestBucket(t *testing.T) jetstream.KeyValue {
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateEndpointsBucket(t.Context(), js, 1)
require.NoError(t, err)
return kv
}
func TestStore_CreateGetUpdateDelete(t *testing.T) {
kv := newTestBucket(t)
key := EndpointKey("000000000000", "test-model")
// Absent key reads as not-found, not an error.
var rec EndpointRecord
found, err := getJSON(t.Context(), kv, key, &rec)
require.NoError(t, err)
assert.False(t, found)
rec = EndpointRecord{AccountID: "000000000000", ModelID: "test-model", State: StateStarting, Generation: 1}
rev, err := createJSONRevision(t.Context(), kv, key, rec)
require.NoError(t, err)
assert.Positive(t, rev)
// A second create at the same key is rejected: this is the claim
// primitive Ensure relies on for cross-replica single-flight.
_, err = createJSONRevision(t.Context(), kv, key, rec)
assert.ErrorIs(t, err, jetstream.ErrKeyExists)
var got EndpointRecord
gotRev, found, err := getJSONRevision(t.Context(), kv, key, &got)
require.NoError(t, err)
require.True(t, found)
assert.Equal(t, rev, gotRev)
assert.Equal(t, StateStarting, got.State)
got.State = StateReady
got.Generation = 2
require.NoError(t, updateJSON(t.Context(), kv, key, gotRev, got))
// The CAS update must fail against the now-stale revision — this is what
// stops a launch goroutine and a concurrent writer stomping each other.
err = updateJSON(t.Context(), kv, key, gotRev, got)
assert.Error(t, err)
require.NoError(t, deleteJSON(t.Context(), kv, key))
found, err = getJSON(t.Context(), kv, key, &rec)
require.NoError(t, err)
assert.False(t, found)
// Deleting an already-absent key is idempotent.
require.NoError(t, deleteJSON(t.Context(), kv, key))
}
func TestListEndpoints_FiltersByAccountPrefix(t *testing.T) {
kv := newTestBucket(t)
seed := func(accountID, modelID string, state EndpointState) {
rec := EndpointRecord{AccountID: accountID, ModelID: modelID, State: state, Generation: 1}
_, err := createJSONRevision(t.Context(), kv, EndpointKey(accountID, modelID), rec)
require.NoError(t, err)
}
seed("000000000000", "model-a", StateReady)
seed("000000000000", "model-b", StateStarting)
seed("111111111111", "model-c", StateReady)
recs, err := ListEndpoints(t.Context(), kv, "000000000000")
require.NoError(t, err)
require.Len(t, recs, 2)
modelIDs := []string{recs[0].ModelID, recs[1].ModelID}
assert.ElementsMatch(t, []string{"model-a", "model-b"}, modelIDs)
recs, err = ListEndpoints(t.Context(), kv, "111111111111")
require.NoError(t, err)
require.Len(t, recs, 1)
assert.Equal(t, "model-c", recs[0].ModelID)
recs, err = ListEndpoints(t.Context(), kv, "222222222222")
require.NoError(t, err)
assert.Empty(t, recs)
}
+47
View File
@@ -0,0 +1,47 @@
package handlers_bedrock
// NATS subjects the daemon subscribes and NATSEndpointService forwards to,
// following the "eks.MethodName" naming EKS uses for its own internal
// (non-AWS-SDK-shaped) subjects.
const (
SubjectEnsureEndpoint = "bedrock.endpoint.ensure"
SubjectDescribeEndpoint = "bedrock.endpoint.describe"
SubjectListEndpoints = "bedrock.endpoint.list"
SubjectDeleteEndpoint = "bedrock.endpoint.delete"
)
// EnsureEndpointInput requests that modelID have a running (or starting)
// serving endpoint. Idempotent: a model already STARTING or READY returns its
// current record without launching a second VM.
type EnsureEndpointInput struct {
ModelID string `json:"model_id"`
}
type EnsureEndpointOutput struct {
Endpoint EndpointRecord `json:"endpoint"`
}
// DescribeEndpointInput looks up modelID's current endpoint record.
type DescribeEndpointInput struct {
ModelID string `json:"model_id"`
}
type DescribeEndpointOutput struct {
Endpoint EndpointRecord `json:"endpoint"`
}
// ListEndpointsInput has no fields today; accountID comes from the NATS
// message header like every other handler, not the payload.
type ListEndpointsInput struct{}
type ListEndpointsOutput struct {
Endpoints []EndpointRecord `json:"endpoints"`
}
// DeleteEndpointInput requests a READY endpoint move to DRAINING and its VM
// be torn down. Idempotent: an already-ABSENT endpoint returns success.
type DeleteEndpointInput struct {
ModelID string `json:"model_id"`
}
type DeleteEndpointOutput struct{}
+44
View File
@@ -0,0 +1,44 @@
package handlers_bedrock
import (
"fmt"
"strings"
)
// servingAgentEnvPath is where the vllm-serving AMI's baked systemd unit
// reads its launch parameters from, mirroring RDS's agentEnvPath convention.
// The unit itself (and this path) lives in the AMI build, out of this bead's
// scope — this is the contract the launcher writes to it.
const servingAgentEnvPath = "/etc/conf.d/vllm-serve"
type serveUserDataInput struct {
ModelID string
VLLMArgs []string
WeightsDevice string
ServePort int
}
// buildServeUserData renders the cloud-config write_files blob that hands the
// vllm-serving guest its model ID, extra vLLM args and the device its cloned
// weights volume landed on, following buildAgentUserData's write_files shape.
func buildServeUserData(in serveUserDataInput) string {
body := strings.Join([]string{
"VLLM_MODEL_ID=" + in.ModelID,
"VLLM_ARGS=" + strings.Join(in.VLLMArgs, " "),
"VLLM_WEIGHTS_DEVICE=" + in.WeightsDevice,
fmt.Sprintf("VLLM_SERVE_PORT=%d", in.ServePort),
}, "\n")
var buf strings.Builder
buf.WriteString("#cloud-config\n")
buf.WriteString("write_files:\n")
fmt.Fprintf(&buf, " - path: %s\n", servingAgentEnvPath)
buf.WriteString(" permissions: '0600'\n")
buf.WriteString(" content: |\n")
for line := range strings.SplitSeq(body, "\n") {
buf.WriteString(" ")
buf.WriteString(line)
buf.WriteString("\n")
}
return buf.String()
}
+74
View File
@@ -0,0 +1,74 @@
package handlers_bedrock
import (
"context"
"errors"
"fmt"
"github.com/mulgadc/spinifex/spinifex/config"
handlers_systemvpc "github.com/mulgadc/spinifex/spinifex/handlers/systemvpc"
"github.com/mulgadc/spinifex/spinifex/tags"
)
// One shared Bedrock system VPC per region holds every serving VM's primary
// NIC. A serving VM has no customer ENI to isolate — inference traffic and
// the readiness probe both dial the ENI's private IP directly — so a VPC per
// endpoint would add a NAT gateway and EIP without adding isolation.
//
// Own tag keys and address space, so no EKS/RDS teardown or orphan reaper can
// reach it even though it shares their VPC builder and topology.
const (
bedrockSystemVPCTagKey = "spinifex:bedrock-system-vpc"
bedrockSystemVPCRoleTagKey = "spinifex:bedrock-role"
bedrockSystemVPCRolePrefix = "bedrock"
)
// SystemVPCName is stable across releases: it seeds both the owner tag and
// the deterministic address hash.
func SystemVPCName(region string) string {
return "bedrock-system-" + region
}
// cfg supplies the operator-overridable address space and subnet count; a nil
// or unset cfg falls back to the defaults.
func SystemVPCSpec(cfg *config.BedrockConfig, region string) handlers_systemvpc.Spec {
supernet := config.BedrockDefaultSystemVPCSupernet
privateSubnets := 1
if cfg != nil {
if cfg.SystemVPCSupernet != "" {
supernet = cfg.SystemVPCSupernet
}
if cfg.SystemVPCPrivateSubnets > 0 {
privateSubnets = cfg.SystemVPCPrivateSubnets
}
}
return handlers_systemvpc.Spec{
Owner: handlers_systemvpc.Owner{
Name: SystemVPCName(region),
ManagedBy: tags.ManagedByBedrock,
OwnerTagKey: bedrockSystemVPCTagKey,
RoleTagKey: bedrockSystemVPCRoleTagKey,
},
Region: region,
RolePrefix: bedrockSystemVPCRolePrefix,
Supernet: supernet,
PrivateSubnets: privateSubnets,
}
}
// EnsureSystemVPC is idempotent. The private subnet the serving VMs sit in
// routes 0.0.0.0/0 to the VPC's NAT gateway, which is the guest's only egress
// to fetch anything the baked AMI doesn't already carry.
func EnsureSystemVPC(ctx context.Context, deps handlers_systemvpc.Deps, cfg *config.BedrockConfig, accountID, region string) (*handlers_systemvpc.Refs, error) {
if region == "" {
return nil, errors.New("bedrock: EnsureSystemVPC empty region")
}
refs, err := handlers_systemvpc.Ensure(ctx, deps, SystemVPCSpec(cfg, region), accountID)
if err != nil {
return nil, fmt.Errorf("bedrock: ensure system VPC for %s: %w", region, err)
}
if len(refs.PrivateSubnetIDs) == 0 {
return nil, fmt.Errorf("bedrock: system VPC %s has no private subnet to place serving VMs in", refs.VpcID)
}
return refs, nil
}
+5 -1
View File
@@ -26,6 +26,10 @@ const (
// UI's listings and the platform's own sweeps.
ManagedByRDS = "rds"
// ManagedByBedrock identifies self-hosted vLLM serving VMs, their ENIs, the
// vllm-serving AMI, and the shared Bedrock system VPC.
ManagedByBedrock = "bedrock"
// LBARNKey stores the parent LB ARN on ELBv2-managed ENIs.
LBARNKey = "spinifex:lb-arn"
@@ -45,5 +49,5 @@ const (
// VMs bind a system.TerminateInstance.{id} subject so a cluster-wide teardown
// invoked on any node can route a terminate to the owning node.
func IsSystemManaged(managedBy string) bool {
return managedBy == ManagedByELBv2 || managedBy == ManagedByEKS || managedBy == ManagedByRDS
return managedBy == ManagedByELBv2 || managedBy == ManagedByEKS || managedBy == ManagedByRDS || managedBy == ManagedByBedrock
}