mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 17:11:30 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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"},
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"`
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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{}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user