diff --git a/spinifex/config/config.go b/spinifex/config/config.go index c140f63ec..fb60eb464 100644 --- a/spinifex/config/config.go +++ b/spinifex/config/config.go @@ -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 diff --git a/spinifex/daemon/bedrock_deps.go b/spinifex/daemon/bedrock_deps.go new file mode 100644 index 000000000..0ef2820ec --- /dev/null +++ b/spinifex/daemon/bedrock_deps.go @@ -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() +} diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index 76186faff..a95293ea8 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -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 diff --git a/spinifex/gateway/bedrock/catalog.go b/spinifex/gateway/bedrock/catalog.go index 3f88c947a..98d58d70c 100644 --- a/spinifex/gateway/bedrock/catalog.go +++ b/spinifex/gateway/bedrock/catalog.go @@ -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, diff --git a/spinifex/handlers/bedrock/ami.go b/spinifex/handlers/bedrock/ami.go new file mode 100644 index 000000000..a186f3f03 --- /dev/null +++ b/spinifex/handlers/bedrock/ami.go @@ -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 +} diff --git a/spinifex/handlers/bedrock/capacity.go b/spinifex/handlers/bedrock/capacity.go new file mode 100644 index 000000000..aca100f4c --- /dev/null +++ b/spinifex/handlers/bedrock/capacity.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/capacity_test.go b/spinifex/handlers/bedrock/capacity_test.go new file mode 100644 index 000000000..0d579428b --- /dev/null +++ b/spinifex/handlers/bedrock/capacity_test.go @@ -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)) +} diff --git a/spinifex/handlers/bedrock/fakes_test.go b/spinifex/handlers/bedrock/fakes_test.go new file mode 100644 index 000000000..0a7cbd331 --- /dev/null +++ b/spinifex/handlers/bedrock/fakes_test.go @@ -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"}, + } +} diff --git a/spinifex/handlers/bedrock/launch.go b/spinifex/handlers/bedrock/launch.go new file mode 100644 index 000000000..98a18e901 --- /dev/null +++ b/spinifex/handlers/bedrock/launch.go @@ -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 +} diff --git a/spinifex/handlers/bedrock/launch_test.go b/spinifex/handlers/bedrock/launch_test.go new file mode 100644 index 000000000..70bce85a2 --- /dev/null +++ b/spinifex/handlers/bedrock/launch_test.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/readiness.go b/spinifex/handlers/bedrock/readiness.go new file mode 100644 index 000000000..d039162a9 --- /dev/null +++ b/spinifex/handlers/bedrock/readiness.go @@ -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 +} diff --git a/spinifex/handlers/bedrock/readiness_test.go b/spinifex/handlers/bedrock/readiness_test.go new file mode 100644 index 000000000..78ab5fcf3 --- /dev/null +++ b/spinifex/handlers/bedrock/readiness_test.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/service.go b/spinifex/handlers/bedrock/service.go new file mode 100644 index 000000000..d699b8b69 --- /dev/null +++ b/spinifex/handlers/bedrock/service.go @@ -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 +} diff --git a/spinifex/handlers/bedrock/service_nats.go b/spinifex/handlers/bedrock/service_nats.go new file mode 100644 index 000000000..11ebcaaf3 --- /dev/null +++ b/spinifex/handlers/bedrock/service_nats.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/service_nats_test.go b/spinifex/handlers/bedrock/service_nats_test.go new file mode 100644 index 000000000..fe28570c2 --- /dev/null +++ b/spinifex/handlers/bedrock/service_nats_test.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/service_test.go b/spinifex/handlers/bedrock/service_test.go new file mode 100644 index 000000000..7bd7e4dae --- /dev/null +++ b/spinifex/handlers/bedrock/service_test.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/state.go b/spinifex/handlers/bedrock/state.go new file mode 100644 index 000000000..88eba9f0e --- /dev/null +++ b/spinifex/handlers/bedrock/state.go @@ -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"` +} diff --git a/spinifex/handlers/bedrock/state_test.go b/spinifex/handlers/bedrock/state_test.go new file mode 100644 index 000000000..20b9a6599 --- /dev/null +++ b/spinifex/handlers/bedrock/state_test.go @@ -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) + }) + } +} diff --git a/spinifex/handlers/bedrock/store.go b/spinifex/handlers/bedrock/store.go new file mode 100644 index 000000000..6dbcadfe0 --- /dev/null +++ b/spinifex/handlers/bedrock/store.go @@ -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 +} diff --git a/spinifex/handlers/bedrock/store_test.go b/spinifex/handlers/bedrock/store_test.go new file mode 100644 index 000000000..5f0c331cb --- /dev/null +++ b/spinifex/handlers/bedrock/store_test.go @@ -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) +} diff --git a/spinifex/handlers/bedrock/types.go b/spinifex/handlers/bedrock/types.go new file mode 100644 index 000000000..e37b82ce6 --- /dev/null +++ b/spinifex/handlers/bedrock/types.go @@ -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{} diff --git a/spinifex/handlers/bedrock/userdata.go b/spinifex/handlers/bedrock/userdata.go new file mode 100644 index 000000000..8969670e7 --- /dev/null +++ b/spinifex/handlers/bedrock/userdata.go @@ -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() +} diff --git a/spinifex/handlers/bedrock/vpc.go b/spinifex/handlers/bedrock/vpc.go new file mode 100644 index 000000000..5a469317d --- /dev/null +++ b/spinifex/handlers/bedrock/vpc.go @@ -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 +} diff --git a/spinifex/tags/tags.go b/spinifex/tags/tags.go index 647e14686..fb75f630a 100644 --- a/spinifex/tags/tags.go +++ b/spinifex/tags/tags.go @@ -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 }