diff --git a/.dockerignore b/.dockerignore index a3aab7a..d1f1a7a 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,3 +1,11 @@ # More info: https://docs.docker.com/engine/reference/builder/#dockerignore-file # Ignore build and test binaries. bin/ + +# Secrets must never enter the build context, so they can never be baked into an +# image layer even if the Dockerfile later switches to a broad `COPY . .`. The +# manager reads credentials from the environment at runtime (per-provider +# Secrets), never from a file in the image. +.env +.env.* +!.env.example diff --git a/.env.example b/.env.example index 388c2fb..83df3b8 100644 --- a/.env.example +++ b/.env.example @@ -16,6 +16,19 @@ MODAL_TOKEN_ID= MODAL_TOKEN_SECRET= +# --- AWS provider ---------------------------------------------------------- +# SECRETS ONLY. In production prefer IRSA / instance role and leave these blank +# (the SDK's default credential chain finds the role) — the AWS secret is then +# skipped, which is fine. Set them only for local/dev without a role. Both keys +# are required together; leave both blank to skip. +AWS_ACCESS_KEY_ID= +AWS_SECRET_ACCESS_KEY= +# OPTIONAL. Required ONLY for temporary STS/SSO credentials (ASIA... access keys), +# which are invalid without it; leave blank for long-lived IAM user keys +# (AKIA...). Note temporary creds expire (minutes–hours) — fine for a quick test, +# not a long-running deployment. +AWS_SESSION_TOKEN= + # --- Additional providers (add as adapters land) --------------------------- # Each provider gets its OWN secret (see hack/deploy.sh PROVIDER_SECRETS), e.g.: # RUNPOD_API_KEY= diff --git a/api/v1alpha1/groupversion_info.go b/api/v1alpha1/groupversion_info.go index dab8323..5abce62 100644 --- a/api/v1alpha1/groupversion_info.go +++ b/api/v1alpha1/groupversion_info.go @@ -82,6 +82,24 @@ const ( // default" (e.g. Modal is OnDemand-only and ignores it). CapacityTypeAnnotation = "nebula.inftyai.com/capacity-type" + // RegionAnnotation carries the optimizer-chosen provider region. Like + // CapacityTypeAnnotation it is a provisioning input absent from the Pod's own + // spec, so the placement controller writes it when it ungates the Pod and the + // virtual kubelet reads it back on CreatePod into ProvisionRequest.Region. + // Empty/absent means "the provider's configured default region" — region-simple + // providers (Modal, RunPod) ignore it. + RegionAnnotation = "nebula.inftyai.com/region" + + // BlocklistTTLAnnotation carries the pool's FailoverPolicy.BlocklistTTL down to + // the virtual kubelet. Like the two annotations above it is a provisioning-time + // input the Pod cannot otherwise express: the TTL is a NodePool policy, but the + // VK handler (which provisions per-Pod and never sees the pool) needs it to know + // how long to blocklist a placement that just failed. The placement controller + // stamps it when it ungates the Pod; the handler reads it on a Provision failure + // to bound the block it records. Absent/unparseable means the handler's built-in + // default TTL. + BlocklistTTLAnnotation = "nebula.inftyai.com/blocklist-ttl" + // TerminateInstanceFinalizer is held by every NodeClaim to guarantee teardown. // The virtual kubelet owns the happy path (DeletePod → provider.Terminate, // keyed on the Pod-derived claim name), but its teardown is edge-triggered and diff --git a/api/v1alpha1/nodeclaim_types.go b/api/v1alpha1/nodeclaim_types.go index 902f473..a76ca9e 100644 --- a/api/v1alpha1/nodeclaim_types.go +++ b/api/v1alpha1/nodeclaim_types.go @@ -31,6 +31,16 @@ type NodeClaimSpec struct { // +optional CapacityType CapacityType `json:"capacityType,omitempty"` + // Region is the provider region the placement optimizer selected, in the + // provider's own vocabulary (e.g. AWS "us-east-1"). Stored durably alongside + // Provider/CapacityType because it is a provisioning input that cannot be read + // off the Pod, and Provision needs it to re-issue the request in the same + // region after a controller restart. Immutable, like Provider. Empty means + // "the provider's configured default region" — region-simple providers (Modal, + // RunPod) leave it empty. + // +optional + Region string `json:"region,omitempty"` + // PoolRef is the NodePool whose policy produced this claim, for reporting. // +optional PoolRef string `json:"poolRef,omitempty"` diff --git a/api/v1alpha1/nodepool_types.go b/api/v1alpha1/nodepool_types.go index 080d132..ba2696c 100644 --- a/api/v1alpha1/nodepool_types.go +++ b/api/v1alpha1/nodepool_types.go @@ -12,12 +12,18 @@ import ( // fixed: capacity type first, provider second. // // FOR each capacityType in CapacityTypes (in listed order): // outer: hard tier -// candidates = Providers x {this capacityType}, available now, minus blocklist +// candidates = Providers x (each provider's Regions) x {this capacityType}, +// available now, minus blocklist // region nests per provider // IF candidates non-empty: -// pick one via Strategy (LowestPrice | Ordered | Weighted) // inner: rank providers +// pick one via Strategy (LowestPrice | Ordered | Weighted) // inner: rank candidates // DONE // // else fall through to the next capacity tier // +// Region is a per-provider axis (see ProviderSpec.Regions), nested under each +// provider because a region name only means something to one provider. It +// widens the candidate key to {provider, region, accelerator, capacityType} +// without changing the tier-first ordering above. +// // So CapacityTypes is a hard preference: every provider's Spot is tried before // ANY provider's OnDemand. This is deliberate — "spot everywhere before any // on-demand" is the least surprising behaviour, even if a different provider's @@ -28,13 +34,14 @@ import ( // static property of the spec, so it is enforced at admission by the CEL rule // below rather than surfaced as a status condition after the fact. // +kubebuilder:validation:XValidation:rule="self.strategy != 'Weighted' || self.providers.all(p, has(p.weight))",message="strategy Weighted requires a weight on every provider" +// +kubebuilder:validation:XValidation:rule="self.providers.all(p, p.name != 'aws' || (has(p.regions) && size(p.regions) > 0))",message="provider aws requires at least one region" type NodePoolSpec struct { // Providers is the ordered set of NeoClouds this pool is allowed to use. // A Pod bound to this pool can only ever be placed on a provider in this // list. Order is significant only for the Ordered strategy (it is the // inner, provider-ranking axis). // +kubebuilder:validation:MinItems=1 - Providers []ProviderRef `json:"providers"` + Providers []ProviderSpec `json:"providers"` // CapacityTypes is the OUTER axis: the purchase models to try, in fallback // order. e.g. [Spot, OnDemand] means "use spot on any provider first; only @@ -57,8 +64,10 @@ type NodePoolSpec struct { Failover *FailoverPolicy `json:"failover,omitempty"` } -// ProviderRef names a provider and, for the Weighted strategy, its share. -type ProviderRef struct { +// ProviderSpec is one provider's entry in a pool: which provider, and the +// per-provider placement policy (Weighted share, allowed regions). It is not a +// mere reference — it carries config — so it is a Spec, not a Ref. +type ProviderSpec struct { // Name is the provider identifier, matching the ProviderLabel value on that // provider's virtual node (e.g. "runpod", "modal", "kubernetes"). Name string `json:"name"` @@ -68,6 +77,29 @@ type ProviderRef struct { // +kubebuilder:validation:Minimum=1 // +optional Weight *int32 `json:"weight,omitempty"` + + // Regions constrains this provider to a subset of its regions, in the + // provider's OWN vocabulary (e.g. ["us-east-1","eu-west-2"] for AWS). Region + // is provider-namespaced — there is no cross-provider region vocabulary — so + // it lives here per provider, not on the pool. Two cases: + // - omitted/empty => the provider's configured default region (the region + // its client resolved from env/config/instance metadata at startup). This + // is the no-surprise default for region-simple providers (Modal, RunPod), + // which have a single region and ignore this field. + // - explicit list => exactly those regions. + // AWS is the exception: it is region-aware with no meaningful single default, + // so a `- name: aws` entry MUST list at least one region. That is enforced at + // admission by the CEL rule on NodePoolSpec (a per-provider requirement, so it + // belongs on the spec where all provider entries are visible, not as a blanket + // MinItems that would burden region-simple providers). An "all regions" + // wildcard is intentionally NOT supported yet: it only makes sense once the + // price-ranking optimizer can expand it against the provider's catalog and + // choose among the results, so it is reserved for then. At most 8 regions; + // maxLength bounds each entry. + // +optional + // +kubebuilder:validation:MaxItems=8 + // +kubebuilder:validation:items:MaxLength=32 + Regions []string `json:"regions,omitempty"` } // PlacementStrategy ranks providers WITHIN a capacity tier (the inner axis). diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index e619508..ffb03a9 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -195,7 +195,7 @@ func (in *NodePoolSpec) DeepCopyInto(out *NodePoolSpec) { *out = *in if in.Providers != nil { in, out := &in.Providers, &out.Providers - *out = make([]ProviderRef, len(*in)) + *out = make([]ProviderSpec, len(*in)) for i := range *in { (*in)[i].DeepCopyInto(&(*out)[i]) } @@ -267,21 +267,26 @@ func (in *PodReference) DeepCopy() *PodReference { } // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. -func (in *ProviderRef) DeepCopyInto(out *ProviderRef) { +func (in *ProviderSpec) DeepCopyInto(out *ProviderSpec) { *out = *in if in.Weight != nil { in, out := &in.Weight, &out.Weight *out = new(int32) **out = **in } + if in.Regions != nil { + in, out := &in.Regions, &out.Regions + *out = make([]string, len(*in)) + copy(*out, *in) + } } -// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ProviderRef. -func (in *ProviderRef) DeepCopy() *ProviderRef { +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new ProviderSpec. +func (in *ProviderSpec) DeepCopy() *ProviderSpec { if in == nil { return nil } - out := new(ProviderRef) + out := new(ProviderSpec) in.DeepCopyInto(out) return out } diff --git a/cmd/main.go b/cmd/main.go index 2e530a2..13b5bd0 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -45,7 +45,9 @@ import ( nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" "github.com/InftyAI/Nebula/internal/controller" webhookv1 "github.com/InftyAI/Nebula/internal/webhook/v1" + "github.com/InftyAI/Nebula/pkg/failover" "github.com/InftyAI/Nebula/pkg/provider" + awsprovider "github.com/InftyAI/Nebula/pkg/provider/aws" "github.com/InftyAI/Nebula/pkg/provider/fake" "github.com/InftyAI/Nebula/pkg/provider/modal" "github.com/InftyAI/Nebula/pkg/vnode" @@ -219,6 +221,13 @@ func main() { // startup already sees its provider. registerProviders(context.Background()) + // One shared failover blocklist, written by the virtual kubelet handlers on a + // Provision failure and read by the placement controller to skip a candidate + // that just failed (zone → region → tier). It is in-memory, process-wide state + // (empty on restart, self-refilling), so a single instance is threaded into + // both sides rather than persisted. + blocklist := failover.New() + if err := (&controller.NodePoolReconciler{ Client: mgr.GetClient(), Scheme: mgr.GetScheme(), @@ -234,8 +243,9 @@ func main() { os.Exit(1) } if err := (&controller.PodPlacementReconciler{ - Client: mgr.GetClient(), - Scheme: mgr.GetScheme(), + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + Blocklist: blocklist, }).SetupWithManager(mgr); err != nil { setupLog.Error(err, "unable to create controller", "controller", "PodPlacement") os.Exit(1) @@ -246,7 +256,7 @@ func main() { // provider.Terminate on DeletePod, so an ungated Pod bound to a provider's // virtual node materializes an external instance. Each Runner is a // manager.Runnable, so it shares the manager's lifecycle and leader election. - if err := setupVirtualNodes(mgr); err != nil { + if err := setupVirtualNodes(mgr, blocklist); err != nil { setupLog.Error(err, "unable to set up virtual nodes") os.Exit(1) } @@ -295,7 +305,7 @@ func main() { // provider. The Runner needs a typed clientset (the virtual kubelet's node/pod // controllers use client-go directly, not the controller-runtime client), built // from the same rest.Config the manager uses. -func setupVirtualNodes(mgr ctrl.Manager) error { +func setupVirtualNodes(mgr ctrl.Manager, blocklist vnode.Blocklister) error { clientset, err := kubernetes.NewForConfig(mgr.GetConfig()) if err != nil { return err @@ -305,7 +315,7 @@ func setupVirtualNodes(mgr ctrl.Manager) error { if !ok { continue } - if err := mgr.Add(vnode.NewRunner(prov, clientset)); err != nil { + if err := mgr.Add(vnode.NewRunner(prov, clientset, blocklist)); err != nil { return err } setupLog.Info("registered virtual node", "provider", name, "node", vnode.NodeName(name)) @@ -327,6 +337,21 @@ func registerProviders(ctx context.Context) { setupLog.Info("registered provider", "provider", p.Name()) } + // AWS. Region is the only NON-SECRET config, read from AWS_REGION (empty => the + // SDK resolves its default from shared config / IMDS). The adapter is otherwise + // self-configuring: it resolves the GPU AMI and default-VPC subnets itself, so + // no launch template or pre-created infra is needed. Credentials are secrets and + // are NEVER read here: the SDK client uses the default credential chain (IRSA / + // instance-role / AWS_ACCESS_KEY_ID delivered via a Secret), so nothing + // sensitive touches a flag or this call. When AWS is not configured (no + // resolvable region, no GPU AMI in the region), registration is a non-fatal skip. + if p, err := awsprovider.NewSDKClient(ctx, os.Getenv("AWS_REGION")); err != nil { + setupLog.Info("skipping AWS provider registration", "reason", err.Error()) + } else { + provider.Register(p) + setupLog.Info("registered provider", "provider", p.Name()) + } + // The fake provider is an in-memory backend used only by the e2e suite to // exercise the full control-plane loop without cloud credentials. It ships in // the binary but registers ONLY when explicitly enabled, so it can never place diff --git a/config/catalog/kustomization.yaml b/config/catalog/kustomization.yaml index efc7dd7..39d2f64 100644 --- a/config/catalog/kustomization.yaml +++ b/config/catalog/kustomization.yaml @@ -16,6 +16,7 @@ configMapGenerator: - name: catalog files: - modal.csv=../../pkg/provider/catalog/data/modal.csv + - aws.csv=../../pkg/provider/catalog/data/aws.csv generatorOptions: # Stable name (no content-hash suffix) so `kubectl edit` and the volume diff --git a/config/crd/bases/nebula.inftyai.com_nodeclaims.yaml b/config/crd/bases/nebula.inftyai.com_nodeclaims.yaml index 7e2943e..d3bc033 100644 --- a/config/crd/bases/nebula.inftyai.com_nodeclaims.yaml +++ b/config/crd/bases/nebula.inftyai.com_nodeclaims.yaml @@ -104,6 +104,16 @@ spec: Recovery from preemption is delete-and-recreate. Held durably so teardown knows which provider API to call even after status is lost. type: string + region: + description: |- + Region is the provider region the placement optimizer selected, in the + provider's own vocabulary (e.g. AWS "us-east-1"). Stored durably alongside + Provider/CapacityType because it is a provisioning input that cannot be read + off the Pod, and Provision needs it to re-issue the request in the same + region after a controller restart. Immutable, like Provider. Empty means + "the provider's configured default region" — region-simple providers (Modal, + RunPod) leave it empty. + type: string required: - podRef - provider diff --git a/config/crd/bases/nebula.inftyai.com_nodepools.yaml b/config/crd/bases/nebula.inftyai.com_nodepools.yaml index 5599385..19aa2a7 100644 --- a/config/crd/bases/nebula.inftyai.com_nodepools.yaml +++ b/config/crd/bases/nebula.inftyai.com_nodepools.yaml @@ -55,18 +55,23 @@ spec: resolves along two orthogonal axes, and the ORDER between them is\nfixed: capacity type first, provider second.\n\n\tFOR each capacityType in CapacityTypes (in listed order): // outer: hard tier\n\t candidates - = Providers x {this capacityType}, available now, minus blocklist\n\t + = Providers x (each provider's Regions) x {this capacityType},\n\t available + now, minus blocklist // region nests per provider\n\t \ IF candidates non-empty:\n\t pick one via Strategy (LowestPrice - | Ordered | Weighted) // inner: rank providers\n\t DONE\n\t // - else fall through to the next capacity tier\n\nSo CapacityTypes is a - hard preference: every provider's Spot is tried before\nANY provider's - OnDemand. This is deliberate — \"spot everywhere before any\non-demand\" - is the least surprising behaviour, even if a different provider's\non-demand - were momentarily cheaper. Strategy only ranks providers *within*\nthe - active capacity tier; it never crosses tiers.\n\nThe Weighted strategy - requires a weight on every provider ref. This is a\nstatic property - of the spec, so it is enforced at admission by the CEL rule\nbelow rather - than surfaced as a status condition after the fact." + | Ordered | Weighted) // inner: rank candidates\n\t DONE\n\t + \ // else fall through to the next capacity tier\n\nRegion is a per-provider + axis (see ProviderSpec.Regions), nested under each\nprovider because + a region name only means something to one provider. It\nwidens the candidate + key to {provider, region, accelerator, capacityType}\nwithout changing + the tier-first ordering above.\n\nSo CapacityTypes is a hard preference: + every provider's Spot is tried before\nANY provider's OnDemand. This + is deliberate — \"spot everywhere before any\non-demand\" is the least + surprising behaviour, even if a different provider's\non-demand were + momentarily cheaper. Strategy only ranks providers *within*\nthe active + capacity tier; it never crosses tiers.\n\nThe Weighted strategy requires + a weight on every provider ref. This is a\nstatic property of the spec, + so it is enforced at admission by the CEL rule\nbelow rather than surfaced + as a status condition after the fact." properties: capacityTypes: default: @@ -109,14 +114,41 @@ spec: list. Order is significant only for the Ordered strategy (it is the inner, provider-ranking axis). items: - description: ProviderRef names a provider and, for the Weighted - strategy, its share. + description: |- + ProviderSpec is one provider's entry in a pool: which provider, and the + per-provider placement policy (Weighted share, allowed regions). It is not a + mere reference — it carries config — so it is a Spec, not a Ref. properties: name: description: |- Name is the provider identifier, matching the ProviderLabel value on that provider's virtual node (e.g. "runpod", "modal", "kubernetes"). type: string + regions: + description: |- + Regions constrains this provider to a subset of its regions, in the + provider's OWN vocabulary (e.g. ["us-east-1","eu-west-2"] for AWS). Region + is provider-namespaced — there is no cross-provider region vocabulary — so + it lives here per provider, not on the pool. Two cases: + - omitted/empty => the provider's configured default region (the region + its client resolved from env/config/instance metadata at startup). This + is the no-surprise default for region-simple providers (Modal, RunPod), + which have a single region and ignore this field. + - explicit list => exactly those regions. + AWS is the exception: it is region-aware with no meaningful single default, + so a `- name: aws` entry MUST list at least one region. That is enforced at + admission by the CEL rule on NodePoolSpec (a per-provider requirement, so it + belongs on the spec where all provider entries are visible, not as a blanket + MinItems that would burden region-simple providers). An "all regions" + wildcard is intentionally NOT supported yet: it only makes sense once the + price-ranking optimizer can expand it against the provider's catalog and + choose among the results, so it is reserved for then. At most 8 regions; + maxLength bounds each entry. + items: + maxLength: 32 + type: string + maxItems: 8 + type: array weight: description: |- Weight is the relative share of new placements for the Weighted strategy. @@ -145,6 +177,9 @@ spec: x-kubernetes-validations: - message: strategy Weighted requires a weight on every provider rule: self.strategy != 'Weighted' || self.providers.all(p, has(p.weight)) + - message: provider aws requires at least one region + rule: self.providers.all(p, p.name != 'aws' || (has(p.regions) && size(p.regions) + > 0)) status: description: NodePoolStatus surfaces the current placement picture for observability. diff --git a/config/manager/manager.yaml b/config/manager/manager.yaml index 895d52b..0a16800 100644 --- a/config/manager/manager.yaml +++ b/config/manager/manager.yaml @@ -72,6 +72,14 @@ spec: # embedded in the binary. - name: NEBULA_CATALOG_DIR value: /etc/nebula/catalog + # AWS NON-SECRET config (never a credential): the region to operate in. + # The adapter is otherwise self-configuring — it resolves the GPU AMI and + # the default VPC's subnets itself — so no launch template or other infra + # config is needed. This is a plain value here, NOT in a Secret. Empty => + # the SDK resolves its default region (shared config / IMDS); with no + # resolvable region AWS registration is a non-fatal skip. + - name: AWS_REGION + value: "us-east-1" envFrom: # Provider credentials live in a per-provider Secret, one secretRef per # provider — NOT a single shared secret. This matches the "creds-absent → @@ -85,6 +93,14 @@ spec: - secretRef: name: nebula-modal-credentials optional: true + # AWS credentials (AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY), read by + # the SDK's default chain. optional=true so an absent Secret is fine — + # in production this is the norm: prefer IRSA / instance role and skip + # this Secret entirely. Non-secret region/launch-template config is set + # as plain env above, never here. + - secretRef: + name: nebula-aws-credentials + optional: true # Add one secretRef per provider as adapters land, e.g.: # - secretRef: # name: nebula-runpod-credentials diff --git a/config/samples/deployment.yaml b/config/samples/deployment.yaml index 826652a..30efdcc 100644 --- a/config/samples/deployment.yaml +++ b/config/samples/deployment.yaml @@ -36,15 +36,9 @@ spec: metadata: labels: app: gpu-workload-sample - # --- opt in (must be on the Pod template) --- nebula.inftyai.com/enabled: "true" - # Place against this pool (see nebula_v1alpha1_nodepool.yaml). - nebula.inftyai.com/nodepool: nodepool-sample - # Requested accelerator TYPE. Matched case-insensitively against the - # provider catalog (e.g. Modal offers T4/L4/A10G/A100-40GB/A100-80GB/ - # H100/H200). The COUNT is the nvidia.com/gpu resource below — omit it - # for a single GPU. - nebula.inftyai.com/accelerator-type: a100-40gb + nebula.inftyai.com/nodepool: aws + nebula.inftyai.com/accelerator-type: l4 spec: # Do NOT set nodeName or a provider nodeSelector yourself — the placement # controller fills the nodeSelector in when it ungates the Pod. Setting diff --git a/config/samples/nebula_v1alpha1_nodepool_aws.yaml b/config/samples/nebula_v1alpha1_nodepool_aws.yaml new file mode 100644 index 0000000..46740ea --- /dev/null +++ b/config/samples/nebula_v1alpha1_nodepool_aws.yaml @@ -0,0 +1,35 @@ +apiVersion: nebula.inftyai.com/v1alpha1 +kind: NodePool +metadata: + labels: + app.kubernetes.io/managed-by: nebula + name: aws +spec: + # Place onto AWS EC2. The provider name must match the ProviderLabel value on + # the AWS virtual node ("aws"); the adapter registers whenever the manager has a + # resolvable region (AWS_REGION or the SDK default chain) and credentials (IRSA / + # instance role / keys) — it self-configures the GPU AMI and subnets, so no other + # setup is needed. Otherwise the pool reports Ready=False/UnknownProvider. + providers: + - name: aws + # regions is per-provider and in AWS's OWN vocabulary. Unlike region-simple + # providers (Modal), AWS is region-aware with no meaningful default, so at + # least one region is REQUIRED here (enforced at admission). List more to + # widen the candidate set; prices are similar across regions, so which types + # are actually offered is resolved by the live probe. + regions: + - us-east-1 + # - us-west-2 + # Outer axis: prefer cheap interruptible Spot, fall back to OnDemand when Spot + # capacity is exhausted. (EC2 has a real Spot tier with a ~2-minute notice; a + # Spot no-capacity failure blocks only Spot in that region, so OnDemand is still + # tried.) Reserved is omitted — EC2 Reserved is a billing construct, not a + # distinct provisioning path here. + capacityTypes: + - Spot + - OnDemand + # Inner axis: within the active capacity tier. A single-provider pool makes this + # a no-op, but LowestPrice is the natural choice once more AWS regions are listed. + strategy: LowestPrice + failover: + blocklistTTL: 10m diff --git a/config/samples/nebula_v1alpha1_nodepool.yaml b/config/samples/nebula_v1alpha1_nodepool_neocloud.yaml similarity index 95% rename from config/samples/nebula_v1alpha1_nodepool.yaml rename to config/samples/nebula_v1alpha1_nodepool_neocloud.yaml index 55ba7a1..e9aafd5 100644 --- a/config/samples/nebula_v1alpha1_nodepool.yaml +++ b/config/samples/nebula_v1alpha1_nodepool_neocloud.yaml @@ -3,7 +3,7 @@ kind: NodePool metadata: labels: app.kubernetes.io/managed-by: nebula - name: nodepool-sample + name: neocloud spec: # The NeoClouds this pool may place onto. Order matters only for the Ordered # strategy (the inner, provider-ranking axis). diff --git a/go.mod b/go.mod index e4ec53d..6b222b3 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,10 @@ module github.com/InftyAI/Nebula go 1.24.0 require ( + github.com/aws/aws-sdk-go-v2 v1.43.0 + github.com/aws/aws-sdk-go-v2/config v1.32.31 + github.com/aws/aws-sdk-go-v2/service/ec2 v1.317.0 + github.com/aws/smithy-go v1.27.3 github.com/modal-labs/modal-client/go v0.9.0 github.com/onsi/ginkgo/v2 v2.27.2 github.com/onsi/gomega v1.38.2 @@ -11,6 +15,7 @@ require ( k8s.io/api v0.33.3 k8s.io/apimachinery v0.33.3 k8s.io/client-go v0.33.3 + k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738 sigs.k8s.io/controller-runtime v0.21.0 ) @@ -19,6 +24,17 @@ require ( github.com/Masterminds/semver/v3 v3.4.0 // indirect github.com/antlr4-go/antlr/v4 v4.13.0 // indirect github.com/aristanetworks/gomap v0.0.0-20230726210543-f4e41046dced // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.19.30 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 // indirect + github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 // indirect + github.com/aws/aws-sdk-go-v2/service/signin v1.5.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.45.0 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/blang/semver/v4 v4.0.0 // indirect github.com/cenkalti/backoff/v4 v4.3.0 // indirect @@ -104,7 +120,6 @@ require ( k8s.io/component-base v0.33.0 // indirect k8s.io/klog/v2 v2.130.1 // indirect k8s.io/kube-openapi v0.0.0-20250318190949-c8a335a9a2ff // indirect - k8s.io/utils v0.0.0-20241104100929-3ea5e8cea738 // indirect sigs.k8s.io/apiserver-network-proxy/konnectivity-client v0.31.2 // indirect sigs.k8s.io/json v0.0.0-20241010143419-9aa6b5e7a4b3 // indirect sigs.k8s.io/randfill v1.0.0 // indirect diff --git a/go.sum b/go.sum index c783216..83be65a 100644 --- a/go.sum +++ b/go.sum @@ -10,6 +10,36 @@ github.com/aristanetworks/gomap v0.0.0-20230726210543-f4e41046dced h1:HxlRMDx/Ve github.com/aristanetworks/gomap v0.0.0-20230726210543-f4e41046dced/go.mod h1:p7lmI+ecoe1RTyD11SPXWsSQ3H+pJ4cp5y7vtKW4QdM= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs= +github.com/aws/aws-sdk-go-v2 v1.43.0 h1:fharf/WhbRAVZ1du0QL7roNFxZ6T/sWr+4Ni617bwSI= +github.com/aws/aws-sdk-go-v2 v1.43.0/go.mod h1:5pKeft2eJj+gElQ38Jqg4ibCqh+/AK33/0X3hip7IjM= +github.com/aws/aws-sdk-go-v2/config v1.32.31 h1:n4nY9O3QKoHIkL85EX+V8RcMFtOhlpTFhGArg915PXk= +github.com/aws/aws-sdk-go-v2/config v1.32.31/go.mod h1:PN0NYDCCoOpGGsZ2+elDUidmHfQBPyYzN2GCgl8HEBs= +github.com/aws/aws-sdk-go-v2/credentials v1.19.30 h1:TTCvvzFU6gXa4iJecNG/0F/B0oYTiazoRECr2XyLHrY= +github.com/aws/aws-sdk-go-v2/credentials v1.19.30/go.mod h1:jKxAp2AEncnliinzpgOSZDFv6+VjvWhjw/AtbfsWT9U= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31 h1:kfVL5wAunCJycL6MOQ6aNh6PlAYEymflcjuKmrWUA0o= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.31/go.mod h1:nWfRNDAppujCQgOUd43lKT4yeLv9z3nJ3bw1G3BgQKo= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31 h1:Z8F3hfCY33IGpJjFAnv0wvtv1FIKj1GHmRDEYqy64tw= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.31/go.mod h1:aVyUoytEyOViR6jhq6jula0xkc5NfBE2hgeF6BvOrao= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31 h1:hyOxUyXdh3AyjE93gBgsfziJag9ACwcs+ZpDBLzi8mw= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.31/go.mod h1:OERqI9k0draSLB8O8woxY3q25ZWTELRK4RRoLMuMZFo= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32 h1:0MrUL35H/Y4kdFfItoR5jCgtDQ4Z/8LudAoIHRfA4hE= +github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.32/go.mod h1:2tNZkuWz54arj8mHVf+8Y7cKkcD8Wr/fBpENgEXpjLc= +github.com/aws/aws-sdk-go-v2/service/ec2 v1.317.0 h1:IkqA16g2hkQntk/K5+srT65TueoTDa7vGhZwqG9w6T4= +github.com/aws/aws-sdk-go-v2/service/ec2 v1.317.0/go.mod h1:dmz3SHr11/hwUijR6xfE/xDRNHcjJwJWZ9ASZdkjGeg= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13 h1:mbRIur/BiHK6SKPjoBIXSE/hJ6g6JGRLuxQy1jGjlN4= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.13/go.mod h1:ITg9em2KbJx1s0y4aqRX5OYWG6HBZ5TVR//OdpEZ2CQ= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31 h1:w2SIhW92DZPFrSL4ksVCr8IYff5OZwIcxg8+95tzvAI= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.31/go.mod h1:wAhpCQbkov+IcvjozJbd2xRCoZybUEHNkcFunssNACg= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.0 h1:OHH5iTQvVGmfHjX/5Q+vFuA/Rf2x6/95aJ/75QCQSm4= +github.com/aws/aws-sdk-go-v2/service/signin v1.5.0/go.mod h1:mCF3AK9PpL49oOrhniUXWAfhVBVQ/XbytoE5eccZUIs= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.0 h1:CaJyYhxBE0M/HJX/YvSaSmQlsI91VHB0lKU8LtLxL3A= +github.com/aws/aws-sdk-go-v2/service/sso v1.33.0/go.mod h1:+e6BMRMPjBQoCw/WovYR9GLy2IU0z4Q77smOB1DraSg= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0 h1:tC323YV77QdafeBr6LUhLDTsboyuyHLNRwAyCP44kGU= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.38.0/go.mod h1:SfLK1sgviHmbI+MozR9iDwDjL4cdCVZtahsjoR+z7wg= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.0 h1:Pd6PNlp4t8PTXxqzstICl52Wsy78vpjFZ7PRUj44mJc= +github.com/aws/aws-sdk-go-v2/service/sts v1.45.0/go.mod h1:rmQ0TnHzuLPmabgjPcsywhsSOmaBDgzR4zvDxSPsGdg= +github.com/aws/smithy-go v1.27.3 h1:F3Zb497UhhskkfpJmfkXswyo+t0sh9OTBnIHjogWbVY= +github.com/aws/smithy-go v1.27.3/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/blang/semver/v4 v4.0.0 h1:1PFHFE6yCCTv8C1TeyNNarDzntLi7wMI5i/pzqYIsAM= diff --git a/hack/deploy.sh b/hack/deploy.sh index 05a92e2..e56804a 100755 --- a/hack/deploy.sh +++ b/hack/deploy.sh @@ -62,6 +62,14 @@ fi PROVIDER_SECRETS=( # Only secrets belong here. "nebula-modal-credentials|MODAL_TOKEN_ID MODAL_TOKEN_SECRET|" + # AWS: creds are the only secret. The access key + secret are required together + # (a lone key is a misconfig), so a blank pair skips the Secret — the SDK then + # relies on IRSA / instance role, which is the preferred path. AWS_SESSION_TOKEN + # is OPTIONAL: it is REQUIRED for temporary STS/SSO creds (ASIA... keys) and + # unused for long-lived IAM user keys (AKIA...), so it is only added to the + # Secret when set. The region is NON-SECRET and lives on the manager Deployment, + # not in this Secret; the adapter self-configures the rest (GPU AMI + subnets). + "nebula-aws-credentials|AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY|AWS_SESSION_TOKEN" # "nebula-runpod-credentials|RUNPOD_API_KEY|" ) diff --git a/internal/controller/nodeclaim_controller_test.go b/internal/controller/nodeclaim_controller_test.go index 8a1b9c1..9dbe953 100644 --- a/internal/controller/nodeclaim_controller_test.go +++ b/internal/controller/nodeclaim_controller_test.go @@ -73,7 +73,7 @@ func (f *fakeProvider) MapAccelerator(c string) (string, bool) { } return "", false } -func (f *fakeProvider) ClassifyProvisionError(error) provider.BlockScope { +func (f *fakeProvider) ClassifyProvisionError(error, string) provider.BlockScope { return provider.BlockScope{} } diff --git a/internal/controller/nodepool_controller_test.go b/internal/controller/nodepool_controller_test.go index 21684ef..2598b2e 100644 --- a/internal/controller/nodepool_controller_test.go +++ b/internal/controller/nodepool_controller_test.go @@ -58,7 +58,7 @@ func newPoolReconciler(t *testing.T, known []string, objs ...client.Object) (*No return r, c } -func newPool(name string, strategy nebulav1alpha1.PlacementStrategy, providers ...nebulav1alpha1.ProviderRef) *nebulav1alpha1.NodePool { +func newPool(name string, strategy nebulav1alpha1.PlacementStrategy, providers ...nebulav1alpha1.ProviderSpec) *nebulav1alpha1.NodePool { return &nebulav1alpha1.NodePool{ ObjectMeta: metav1.ObjectMeta{Name: name}, Spec: nebulav1alpha1.NodePoolSpec{ @@ -106,7 +106,7 @@ func poolReadyCond(p *nebulav1alpha1.NodePool) *metav1.Condition { func TestPool_ValidPoolBecomesReady(t *testing.T) { pool := newPool("gpu", nebulav1alpha1.StrategyOrdered, - nebulav1alpha1.ProviderRef{Name: "modal"}) + nebulav1alpha1.ProviderSpec{Name: "modal"}) r, c := newPoolReconciler(t, []string{"modal"}, pool) reconcilePool(t, r, "gpu") @@ -122,8 +122,8 @@ func TestPool_ValidPoolBecomesReady(t *testing.T) { func TestPool_UnknownProviderNotReady(t *testing.T) { pool := newPool("gpu", nebulav1alpha1.StrategyOrdered, - nebulav1alpha1.ProviderRef{Name: "modal"}, - nebulav1alpha1.ProviderRef{Name: "ghost"}) + nebulav1alpha1.ProviderSpec{Name: "modal"}, + nebulav1alpha1.ProviderSpec{Name: "ghost"}) r, c := newPoolReconciler(t, []string{"modal"}, pool) // ghost not registered reconcilePool(t, r, "gpu") @@ -139,8 +139,8 @@ func TestPool_UnknownProviderNotReady(t *testing.T) { func TestPool_PlacedCountsBoundClaims(t *testing.T) { pool := newPool("gpu", nebulav1alpha1.StrategyOrdered, - nebulav1alpha1.ProviderRef{Name: "modal"}, - nebulav1alpha1.ProviderRef{Name: "runpod"}) + nebulav1alpha1.ProviderSpec{Name: "modal"}, + nebulav1alpha1.ProviderSpec{Name: "runpod"}) // 2 bound on modal, 1 on runpod for this pool; 1 bound for another pool // (ignored); 1 pending on modal for this pool (ignored — not Bound). r, c := newPoolReconciler(t, []string{"modal", "runpod"}, @@ -168,7 +168,7 @@ func TestPool_PlacedCountsBoundClaims(t *testing.T) { func TestPool_PlacedNilWhenNothingRunning(t *testing.T) { pool := newPool("gpu", nebulav1alpha1.StrategyOrdered, - nebulav1alpha1.ProviderRef{Name: "modal"}) + nebulav1alpha1.ProviderSpec{Name: "modal"}) r, c := newPoolReconciler(t, []string{"modal"}, pool, poolClaim("pend", "gpu", "modal", nebulav1alpha1.NodeClaimProvisioning), diff --git a/internal/controller/nodepool_validation_test.go b/internal/controller/nodepool_validation_test.go index 85fb718..f27c388 100644 --- a/internal/controller/nodepool_validation_test.go +++ b/internal/controller/nodepool_validation_test.go @@ -30,7 +30,7 @@ import ( // run CRD x-kubernetes-validations). They require envtest binaries; when those // are absent the whole suite is skipped in BeforeSuite. var _ = Describe("NodePool spec validation (CEL)", func() { - newWeightedPool := func(name string, refs ...nebulav1alpha1.ProviderRef) *nebulav1alpha1.NodePool { + newWeightedPool := func(name string, refs ...nebulav1alpha1.ProviderSpec) *nebulav1alpha1.NodePool { return &nebulav1alpha1.NodePool{ ObjectMeta: metav1.ObjectMeta{Name: name}, Spec: nebulav1alpha1.NodePoolSpec{ @@ -43,8 +43,8 @@ var _ = Describe("NodePool spec validation (CEL)", func() { It("rejects a Weighted pool with a provider missing a weight", func() { pool := newWeightedPool("weighted-missing", - nebulav1alpha1.ProviderRef{Name: "modal", Weight: weight(3)}, - nebulav1alpha1.ProviderRef{Name: "runpod"}, // no weight + nebulav1alpha1.ProviderSpec{Name: "modal", Weight: weight(3)}, + nebulav1alpha1.ProviderSpec{Name: "runpod"}, // no weight ) err := k8sClient.Create(ctx, pool) Expect(err).To(HaveOccurred()) @@ -53,8 +53,8 @@ var _ = Describe("NodePool spec validation (CEL)", func() { It("admits a Weighted pool with a weight on every provider", func() { pool := newWeightedPool("weighted-ok", - nebulav1alpha1.ProviderRef{Name: "modal", Weight: weight(3)}, - nebulav1alpha1.ProviderRef{Name: "runpod", Weight: weight(1)}, + nebulav1alpha1.ProviderSpec{Name: "modal", Weight: weight(3)}, + nebulav1alpha1.ProviderSpec{Name: "runpod", Weight: weight(1)}, ) Expect(k8sClient.Create(ctx, pool)).To(Succeed()) Expect(k8sClient.Delete(ctx, pool)).To(Succeed()) @@ -64,7 +64,7 @@ var _ = Describe("NodePool spec validation (CEL)", func() { pool := &nebulav1alpha1.NodePool{ ObjectMeta: metav1.ObjectMeta{Name: "ordered-noweights"}, Spec: nebulav1alpha1.NodePoolSpec{ - Providers: []nebulav1alpha1.ProviderRef{{Name: "modal"}}, + Providers: []nebulav1alpha1.ProviderSpec{{Name: "modal"}}, Strategy: nebulav1alpha1.StrategyOrdered, }, } diff --git a/internal/controller/pod_placement_controller.go b/internal/controller/pod_placement_controller.go index 38c923b..5033dcb 100644 --- a/internal/controller/pod_placement_controller.go +++ b/internal/controller/pod_placement_controller.go @@ -27,6 +27,7 @@ import ( logf "sigs.k8s.io/controller-runtime/pkg/log" nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/failover" "github.com/InftyAI/Nebula/pkg/provider" "github.com/InftyAI/Nebula/pkg/util" ) @@ -57,6 +58,20 @@ type PodPlacementReconciler struct { // Providers resolves a provider name to its backend; defaults to the // process-wide registry. Overridable in tests. Providers func(name string) (provider.Provider, bool) + + // Blocklist is the shared failover blocklist the VK handlers write to on a + // Provision failure; selectPlacement reads it to skip a candidate that just + // failed. May be nil (no candidate is ever considered blocked), keeping tests + // and blocklist-less wiring simple. + Blocklist Blocklister +} + +// Blocklister is the read side of pkg/failover.Blocklist the placement controller +// depends on: it asks whether a candidate placement is currently excluded. The +// concrete type is injected from main; this narrow interface keeps the controller +// decoupled and a nil value a no-op. +type Blocklister interface { + Blocked(c failover.Candidate) bool } // +kubebuilder:rbac:groups="",resources=pods,verbs=get;list;watch;update;patch;delete @@ -130,10 +145,11 @@ func (r *PodPlacementReconciler) Reconcile(ctx context.Context, req ctrl.Request } // Stamp the routing decision and release the Pod to the scheduler. - if err := r.place(ctx, &pod, placement); err != nil { + if err := r.place(ctx, &pod, pool, placement); err != nil { return ctrl.Result{}, err } - log.Info("placed Pod", "pod", pod.Name, "provider", placement.provider, "capacityType", placement.capacityType) + log.Info("placed Pod", "pod", pod.Name, "provider", placement.provider, + "capacityType", placement.capacityType, "region", placement.region) return ctrl.Result{}, nil } @@ -141,6 +157,10 @@ func (r *PodPlacementReconciler) Reconcile(ctx context.Context, req ctrl.Request type placement struct { provider string capacityType nebulav1alpha1.CapacityType + // region is the provider region to provision in (provider's own vocabulary). + // Empty means the provider's configured default region; region-simple + // providers leave it empty. + region string } // needsPlacement reports whether the Pod is an opted-in workload still held by diff --git a/internal/controller/pod_placement_controller_test.go b/internal/controller/pod_placement_controller_test.go index d6d65f5..20552b1 100644 --- a/internal/controller/pod_placement_controller_test.go +++ b/internal/controller/pod_placement_controller_test.go @@ -19,6 +19,7 @@ package controller import ( "context" "testing" + "time" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -29,9 +30,36 @@ import ( "sigs.k8s.io/controller-runtime/pkg/reconcile" nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/failover" "github.com/InftyAI/Nebula/pkg/provider" ) +// fakeBlocklist is a test Blocklister that reports a fixed set of candidates as +// blocked. A candidate is blocked when it matches an entry on every NON-empty +// field (empty fields are wildcards), mirroring the real blocklist's coverage. +type fakeBlocklist struct { + blocked []failover.Candidate +} + +func (b *fakeBlocklist) Blocked(c failover.Candidate) bool { + for _, e := range b.blocked { + if e.Provider != "" && e.Provider != c.Provider { + continue + } + if e.AcceleratorType != "" && e.AcceleratorType != c.AcceleratorType { + continue + } + if e.CapacityType != "" && e.CapacityType != c.CapacityType { + continue + } + if e.Region != "" && e.Region != c.Region { + continue + } + return true + } + return false +} + // newPlacementReconciler wires a PodPlacementReconciler over a fake client. func newPlacementReconciler(t *testing.T, objs []client.Object, provs ...*fakeProvider) (*PodPlacementReconciler, client.Client) { t.Helper() @@ -74,9 +102,9 @@ func gatedPod(name, ns, uid, pool, gpu string) *corev1.Pod { } func poolWith(name string, capTypes []nebulav1alpha1.CapacityType, providers ...string) *nebulav1alpha1.NodePool { - refs := make([]nebulav1alpha1.ProviderRef, 0, len(providers)) + refs := make([]nebulav1alpha1.ProviderSpec, 0, len(providers)) for _, p := range providers { - refs = append(refs, nebulav1alpha1.ProviderRef{Name: p}) + refs = append(refs, nebulav1alpha1.ProviderSpec{Name: p}) } return &nebulav1alpha1.NodePool{ ObjectMeta: metav1.ObjectMeta{Name: name}, @@ -87,6 +115,18 @@ func poolWith(name string, capTypes []nebulav1alpha1.CapacityType, providers ... } } +// poolWithRegions builds a single-provider pool whose provider ref lists the given +// regions, with the given capacity-type fallback order. +func poolWithRegions(name string, capTypes []nebulav1alpha1.CapacityType, providerName string, regions ...string) *nebulav1alpha1.NodePool { + return &nebulav1alpha1.NodePool{ + ObjectMeta: metav1.ObjectMeta{Name: name}, + Spec: nebulav1alpha1.NodePoolSpec{ + Providers: []nebulav1alpha1.ProviderSpec{{Name: providerName, Regions: regions}}, + CapacityTypes: capTypes, + }, + } +} + func reconcilePod(t *testing.T, r *PodPlacementReconciler, ns, name string) { t.Helper() if _, err := r.Reconcile(context.Background(), reconcile.Request{ @@ -337,6 +377,97 @@ func TestPlacement_IdempotentOnRetry(t *testing.T) { } } +func TestPlacement_FailsOverToNextRegionWhenBlocked(t *testing.T) { + // The pool lists two regions on one provider. us-east-1 is blocklisted, so the + // inner region walk advances to us-west-2 within the SAME capacity tier. + pod := gatedPod("p1", "default", "uid-1", "pool-a", "H100") + pool := poolWithRegions("pool-a", []nebulav1alpha1.CapacityType{nebulav1alpha1.CapacityOnDemand}, + provider.ProviderModal, "us-east-1", "us-west-2") + prov := &fakeProvider{name: provider.ProviderModal, gpus: []string{"H100"}} + r, c := newPlacementReconciler(t, []client.Object{pod, pool}, prov) + r.Blocklist = &fakeBlocklist{blocked: []failover.Candidate{ + {Provider: provider.ProviderModal, AcceleratorType: "H100", CapacityType: nebulav1alpha1.CapacityOnDemand, Region: "us-east-1"}, + }} + + reconcilePod(t, r, "default", "p1") + + got := getPod(t, c, "default", "p1") + if got.Annotations[nebulav1alpha1.RegionAnnotation] != "us-west-2" { + t.Fatalf("expected failover to us-west-2, got region %q", got.Annotations[nebulav1alpha1.RegionAnnotation]) + } +} + +func TestPlacement_CapacityIsOuterAxis(t *testing.T) { + // Two providers, tiers [Spot, OnDemand]. Spot is blocked on BOTH providers but + // OnDemand is free. Capacity is the OUTER axis, so the walk must exhaust every + // provider's Spot before dropping any provider to OnDemand — landing on the + // first provider's OnDemand, not the second provider's (still-Spot) offer. + pod := gatedPod("p1", "default", "uid-1", "pool-a", "H100") + pool := poolWith("pool-a", []nebulav1alpha1.CapacityType{nebulav1alpha1.CapacitySpot, nebulav1alpha1.CapacityOnDemand}, + "runpod", provider.ProviderModal) + runpod := &fakeProvider{name: "runpod", gpus: []string{"H100"}} + modal := &fakeProvider{name: provider.ProviderModal, gpus: []string{"H100"}} + r, c := newPlacementReconciler(t, []client.Object{pod, pool}, runpod, modal) + r.Blocklist = &fakeBlocklist{blocked: []failover.Candidate{ + // Both providers' Spot is exhausted (wildcard provider, Spot only). + {CapacityType: nebulav1alpha1.CapacitySpot}, + }} + + reconcilePod(t, r, "default", "p1") + + got := getPod(t, c, "default", "p1") + if got.Annotations[nebulav1alpha1.CapacityTypeAnnotation] != "OnDemand" { + t.Fatalf("expected the walk to drop to OnDemand, got %q", got.Annotations[nebulav1alpha1.CapacityTypeAnnotation]) + } + // ...and to the FIRST provider (runpod), since OnDemand is walked provider-first. + if got.Spec.NodeSelector[nebulav1alpha1.ProviderLabel] != "runpod" { + t.Fatalf("expected first provider runpod at the OnDemand tier, got %v", got.Spec.NodeSelector) + } +} + +func TestPlacement_AllCandidatesBlockedLeavesPodGated(t *testing.T) { + // Every (tier, provider, region) candidate is blocked (DenyAll on the provider). + // Placement has nowhere to go, so it leaves the Pod gated for a later retry + // (when a block expires) rather than placing onto a failing candidate. + pod := gatedPod("p1", "default", "uid-1", "pool-a", "H100") + pool := poolWith("pool-a", []nebulav1alpha1.CapacityType{nebulav1alpha1.CapacitySpot, nebulav1alpha1.CapacityOnDemand}, + provider.ProviderModal) + prov := &fakeProvider{name: provider.ProviderModal, gpus: []string{"H100"}} + r, c := newPlacementReconciler(t, []client.Object{pod, pool}, prov) + r.Blocklist = &fakeBlocklist{blocked: []failover.Candidate{ + {Provider: provider.ProviderModal}, // whole-provider block (auth/quota) + }} + + reconcilePod(t, r, "default", "p1") + + got := getPod(t, c, "default", "p1") + if !hasGateNamed(got) { + t.Fatal("expected the Pod to stay gated when every candidate is blocked") + } + // No claim should have been created for an unplaced Pod. + var nc nebulav1alpha1.NodeClaim + if err := c.Get(context.Background(), types.NamespacedName{Name: "default-p1"}, &nc); err == nil { + t.Fatal("expected no claim when placement is blocked everywhere") + } +} + +func TestPlacement_StampsBlocklistTTLAnnotation(t *testing.T) { + // The pool's FailoverPolicy.BlocklistTTL must reach the Pod so the VK handler + // knows how long to blocklist a placement that fails. + pod := gatedPod("p1", "default", "uid-1", "pool-a", "H100") + pool := poolWith("pool-a", []nebulav1alpha1.CapacityType{nebulav1alpha1.CapacityOnDemand}, provider.ProviderModal) + pool.Spec.Failover = &nebulav1alpha1.FailoverPolicy{BlocklistTTL: metav1.Duration{Duration: 7 * time.Minute}} + prov := &fakeProvider{name: provider.ProviderModal, gpus: []string{"H100"}} + r, c := newPlacementReconciler(t, []client.Object{pod, pool}, prov) + + reconcilePod(t, r, "default", "p1") + + got := getPod(t, c, "default", "p1") + if got.Annotations[nebulav1alpha1.BlocklistTTLAnnotation] != "7m0s" { + t.Fatalf("expected blocklist-ttl annotation 7m0s, got %q", got.Annotations[nebulav1alpha1.BlocklistTTLAnnotation]) + } +} + // terminalOwnedPod builds an opted-in Pod in a terminal phase. When ownedByRS is // true it carries a controlling ReplicaSet ownerReference (so a controller would // recreate it); otherwise it is a bare Pod. diff --git a/internal/controller/pod_placement_helpers.go b/internal/controller/pod_placement_helpers.go index 444dab1..ee99931 100644 --- a/internal/controller/pod_placement_helpers.go +++ b/internal/controller/pod_placement_helpers.go @@ -28,6 +28,7 @@ import ( "sigs.k8s.io/controller-runtime/pkg/reconcile" nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/failover" "github.com/InftyAI/Nebula/pkg/util" ) @@ -49,15 +50,33 @@ func (r *PodPlacementReconciler) poolFor(ctx context.Context, pod *corev1.Pod) ( return &pool, nil } -// selectPlacement implements the v1 "first matching provider" policy: walk the -// pool's providers in listed order and pick the first whose adapter offers the -// Pod's requested GPU type. The capacity tier is the pool's first preference -// (the outer axis); provider quirks (e.g. Modal being OnDemand-only) are handled -// at Provision time, not here. Returns ok=false when nothing matches. +// selectPlacement resolves the pool's policy along the two orthogonal axes, in the +// fixed order the NodePoolSpec doc mandates: capacity tier is the OUTER axis and +// region (nested per provider) is the INNER one. It walks // -// This is the seam the richer optimizer replaces later (price ranking, capacity -// fallback, blocklist); the caller's flow does not depend on how the choice is -// made, only that a (provider, capacityType) comes back. +// FOR each capacityType in CapacityTypes (listed order): // outer: hard tier +// FOR each provider (listed order = Ordered strategy): +// FOR each of that provider's regions: // inner +// skip if the candidate is blocklisted; else place here +// +// so every provider's Spot is tried (in every one of its regions) before ANY +// provider's OnDemand — "capacity is the outer axis". Skipping blocklisted +// candidates is what turns a single provision failure into failover: a region- +// scoped block (a Spot limit in us-east-1) advances to the next region, then the +// next tier, without hot-looping against the placement that just failed. The +// adapter already handled the finer zone axis (sweeping a region's AZs) before the +// failure ever reached the blocklist, so this loop only walks regions, not zones. +// +// Returns ok=false when every (tier, provider, region) candidate is either +// unservable (provider unregistered or does not offer the accelerator) or blocked; +// the caller then leaves the Pod gated for a later retry (a pool edit, a provider +// registering, or a block expiring). Provider quirks (e.g. Modal being +// OnDemand-only) are still handled at Provision time, not here. +// +// Strategy (LowestPrice/Weighted price-ranking) is a later swap-in for the inner +// ordering; today the inner walk is listed order (Ordered), which the caller's +// flow does not depend on — only that a (provider, capacityType, region) comes +// back. func (r *PodPlacementReconciler) selectPlacement(pod *corev1.Pod, pool *nebulav1alpha1.NodePool) (placement, bool) { // Only the accelerator type is needed for matching; the count is a provisioning // detail the adapter reads. A malformed request is treated as "no accelerator" @@ -65,33 +84,71 @@ func (r *PodPlacementReconciler) selectPlacement(pod *corev1.Pod, pool *nebulav1 // would surface the real error. accel, _, _ := util.AcceleratorRequest(pod) - for _, ref := range pool.Spec.Providers { - prov, ok := r.provider(ref.Name) - if !ok { - continue // unregistered; NodePool status surfaces this separately - } - // A CPU-only Pod (no accelerator) matches any provider; an accelerator Pod - // only matches a provider whose catalog maps that accelerator. - if accel != "" { - if _, offered := prov.MapAccelerator(accel); !offered { - continue + for _, tier := range capacityTiers(pool) { // outer: capacity + for _, ref := range pool.Spec.Providers { // provider (Ordered = listed order) + prov, ok := r.provider(ref.Name) + if !ok { + continue // unregistered; NodePool status surfaces this separately + } + // A CPU-only Pod (no accelerator) matches any provider; an accelerator + // Pod only matches a provider whose catalog maps that accelerator. + if accel != "" { + if _, offered := prov.MapAccelerator(accel); !offered { + continue + } + } + for _, region := range regionsFor(ref) { // inner: region + if r.blocked(ref.Name, accel, tier, region) { + continue // failed recently; try the next region, then the next tier + } + return placement{ + provider: ref.Name, + capacityType: tier, + region: region, + }, true } } - return placement{ - provider: ref.Name, - capacityType: preferredCapacityType(pool), - }, true } return placement{}, false } -// preferredCapacityType is the pool's first-choice tier (the outer axis's head). -// v1 does not do cross-tier fallback here; empty means "provider default". -func preferredCapacityType(pool *nebulav1alpha1.NodePool) nebulav1alpha1.CapacityType { +// capacityTiers is the outer axis to walk: the pool's CapacityTypes in fallback +// order. An empty list means "the provider default tier" — a single unnamed +// candidate ("") so the walk still runs once. (Admission defaults the field, so +// this only guards a hand-built pool.) +func capacityTiers(pool *nebulav1alpha1.NodePool) []nebulav1alpha1.CapacityType { if len(pool.Spec.CapacityTypes) == 0 { - return "" + return []nebulav1alpha1.CapacityType{""} + } + return pool.Spec.CapacityTypes +} + +// regionsFor is the inner axis for one provider ref: the regions to try, in listed +// order. An empty/omitted list means "the provider's configured default region", +// represented as a single empty-string candidate so the walk runs once for +// region-simple providers (Modal, RunPod). AWS is required by admission to list at +// least one region (the CEL rule on NodePoolSpec). An "all regions" wildcard is not +// supported yet (see ProviderSpec.Regions), so there is nothing to expand here. +func regionsFor(ref nebulav1alpha1.ProviderSpec) []string { + if len(ref.Regions) == 0 { + return []string{""} // omitted => provider default region (region-simple providers) } - return pool.Spec.CapacityTypes[0] + return ref.Regions +} + +// blocked reports whether the (provider, accelerator, tier, region) candidate is +// currently excluded by the failover blocklist. It is nil-safe: with no blocklist +// wired (tests, or a blocklist-less build) nothing is ever blocked. +func (r *PodPlacementReconciler) blocked(provName, accel string, tier nebulav1alpha1.CapacityType, region string) bool { + if r.Blocklist == nil { + return false + } + return r.Blocklist.Blocked(failover.Candidate{ + Provider: provName, + AcceleratorType: accel, + CapacityType: tier, + Region: region, + }) } // ensureClaim creates the NodeClaim for this placement if it does not already @@ -127,6 +184,7 @@ func (r *PodPlacementReconciler) ensureClaim(ctx context.Context, pod *corev1.Po }, Provider: p.provider, CapacityType: p.capacityType, + Region: p.region, PoolRef: pool.Name, }, } @@ -156,20 +214,26 @@ func (r *PodPlacementReconciler) ensureClaim(ctx context.Context, pod *corev1.Po // place stamps the routing decision onto the Pod and removes the gate, atomically // from the Pod's perspective (one Update). After this, the scheduler is free to // bind the Pod to the chosen provider's virtual node. -func (r *PodPlacementReconciler) place(ctx context.Context, pod *corev1.Pod, p placement) error { +func (r *PodPlacementReconciler) place(ctx context.Context, pod *corev1.Pod, pool *nebulav1alpha1.NodePool, p placement) error { // Route to the provider's virtual node. if pod.Spec.NodeSelector == nil { pod.Spec.NodeSelector = map[string]string{} } pod.Spec.NodeSelector[nebulav1alpha1.ProviderLabel] = p.provider - // Carry the capacity tier the VK handler reads on CreatePod (the one input - // that is not otherwise on the Pod). Skip when empty (provider default). + // Carry the capacity tier and region the VK handler reads on CreatePod (inputs + // that are not otherwise on the Pod). Skip each when empty (provider default). if p.capacityType != "" { - if pod.Annotations == nil { - pod.Annotations = map[string]string{} - } - pod.Annotations[nebulav1alpha1.CapacityTypeAnnotation] = string(p.capacityType) + setAnnotation(pod, nebulav1alpha1.CapacityTypeAnnotation, string(p.capacityType)) + } + if p.region != "" { + setAnnotation(pod, nebulav1alpha1.RegionAnnotation, p.region) + } + // Carry the pool's blocklist TTL so the VK handler (which never sees the pool) + // knows how long to exclude a placement that fails. Only stamp an explicit + // policy value; an unset policy leaves the handler on its own default. + if pool.Spec.Failover != nil && pool.Spec.Failover.BlocklistTTL.Duration > 0 { + setAnnotation(pod, nebulav1alpha1.BlocklistTTLAnnotation, pool.Spec.Failover.BlocklistTTL.Duration.String()) } // Remove our gate, releasing the Pod to the scheduler. Preserve any other @@ -179,6 +243,14 @@ func (r *PodPlacementReconciler) place(ctx context.Context, pod *corev1.Pod, p p return r.Update(ctx, pod) } +// setAnnotation sets one annotation on the Pod, allocating the map on first use. +func setAnnotation(pod *corev1.Pod, key, value string) { + if pod.Annotations == nil { + pod.Annotations = map[string]string{} + } + pod.Annotations[key] = value +} + // removeGate returns gates with the named gate removed, preserving order. func removeGate(gates []corev1.PodSchedulingGate, name string) []corev1.PodSchedulingGate { out := gates[:0] diff --git a/pkg/failover/blocklist.go b/pkg/failover/blocklist.go new file mode 100644 index 0000000..8281da7 --- /dev/null +++ b/pkg/failover/blocklist.go @@ -0,0 +1,181 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package failover holds the in-memory, TTL-bounded blocklist that turns a single +// provision failure into a temporary exclusion, so placement can fail over to the +// next candidate (zone → region → tier) instead of hot-looping against a provider +// that just rejected the same request. +// +// The design mirrors SkyPilot's blocklist-granularity rule: a failure is recorded +// at the exact granularity that failed (a provider.BlockScope), and a candidate is +// blocked only if some live entry's scope MATCHES it. Empty scope fields are +// wildcards, so "no Spot capacity for H100 in us-east-1" never disqualifies an +// OnDemand request, a different accelerator, or another region — but an auth +// failure (DenyAll) blocks the whole provider until its TTL lapses. +// +// It is deliberately in-memory and per-process: the blocklist is a hint that +// bounds churn, not durable state. On a manager restart it is empty and the next +// attempt re-probes the provider — the worst case is one wasted attempt, which the +// provider itself rejects and re-records. This keeps the control plane free of a +// blocklist CRD and the reconcile hot-path free of an API round-trip. +package failover + +import ( + "sync" + "time" + + "k8s.io/utils/clock" + + nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/provider" +) + +// entry is one recorded block: the provider it applies to, the scope that failed, +// and when the block lapses. It is compared against a query candidate by Blocked. +type entry struct { + provider string + scope provider.BlockScope + expiresAt time.Time +} + +// Candidate is a placement being considered — the (provider, accelerator, tier, +// region) tuple the blocklist is queried against. It is the query counterpart to a +// recorded BlockScope: Blocked reports whether any live entry's scope covers it. +type Candidate struct { + Provider string + AcceleratorType string + CapacityType nebulav1alpha1.CapacityType + Region string +} + +// Blocklist is a concurrency-safe, TTL-bounded set of provider blocks. The zero +// value is not usable; construct with New. It is safe for concurrent use by the +// vnode handlers (which Record failures) and the placement controller (which +// queries Blocked) sharing one instance. +type Blocklist struct { + mu sync.Mutex + entries []entry + clock clock.Clock +} + +// New returns an empty Blocklist backed by the real clock. +func New() *Blocklist { + return &Blocklist{clock: clock.RealClock{}} +} + +// newWithClock is the test seam: an injectable clock so TTL expiry is exercised +// without sleeping. +func newWithClock(c clock.Clock) *Blocklist { + return &Blocklist{clock: c} +} + +// Record adds a block for prov at scope, lapsing after ttl. A DenyAll scope blocks +// the whole provider; a scoped block confines it to the matching wildcard fields. +// A non-positive ttl or an empty provider is a no-op (nothing to bound, or nothing +// to key on) — the caller is expected to skip zero scopes, but Record is defensive +// so a misfire cannot install a permanent or provider-less block. Recording also +// opportunistically drops expired entries so the slice cannot grow without bound. +func (b *Blocklist) Record(prov string, scope provider.BlockScope, ttl time.Duration) { + if prov == "" || ttl <= 0 { + return + } + b.mu.Lock() + defer b.mu.Unlock() + + now := b.clock.Now() + b.gcLocked(now) + b.entries = append(b.entries, entry{ + provider: prov, + scope: scope, + expiresAt: now.Add(ttl), + }) +} + +// Blocked reports whether c is currently excluded by any live entry. It is the +// query the placement loop makes before selecting a candidate: a true result means +// "skip this (provider, accelerator, tier, region) and try the next one". +func (b *Blocklist) Blocked(c Candidate) bool { + b.mu.Lock() + defer b.mu.Unlock() + + now := b.clock.Now() + b.gcLocked(now) + for i := range b.entries { + if b.entries[i].provider == c.Provider && scopeCovers(b.entries[i].scope, c) { + return true + } + } + return false +} + +// gcLocked drops entries that have expired as of now. Caller holds b.mu. +func (b *Blocklist) gcLocked(now time.Time) { + if len(b.entries) == 0 { + return + } + kept := b.entries[:0] + for _, e := range b.entries { + if e.expiresAt.After(now) { + kept = append(kept, e) + } + } + // Zero the tail so dropped entries are not retained by the backing array. + for i := len(kept); i < len(b.entries); i++ { + b.entries[i] = entry{} + } + b.entries = kept +} + +// scopeCovers reports whether scope (a recorded block) applies to candidate c. +// DenyAll covers everything on the provider. Otherwise each field must match the +// candidate's, per the three-state semantics of BlockScope's pointer fields: +// +// - nil => the axis is not applicable: it matches ONLY a candidate whose field +// is empty. A CPU-only Pod (no accelerator) or a region-simple provider (no +// region) blocks without widening across an axis it never had. +// - &"" => wildcard: matches any value on that axis. +// - &"v" => exact: matches only a candidate whose field equals "v". +// +// This is the match half of the blocklist-granularity rule: the block is as broad +// as the failure was, no broader. +func scopeCovers(scope provider.BlockScope, c Candidate) bool { + if scope.DenyAll { + return true + } + if !ptrMatches(scope.AcceleratorType, c.AcceleratorType) { + return false + } + if scope.CapacityType != "" && scope.CapacityType != c.CapacityType { + return false + } + if !ptrMatches(scope.Region, c.Region) { + return false + } + return true +} + +// ptrMatches applies BlockScope's three-state rule to one axis: nil matches only an +// empty candidate value (the axis is not applicable), &"" is a wildcard that matches +// anything, and &"v" matches only the exact value. +func ptrMatches(pattern *string, value string) bool { + if pattern == nil { + return value == "" + } + if *pattern == "" { + return true + } + return *pattern == value +} diff --git a/pkg/failover/blocklist_test.go b/pkg/failover/blocklist_test.go new file mode 100644 index 0000000..422e628 --- /dev/null +++ b/pkg/failover/blocklist_test.go @@ -0,0 +1,216 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package failover + +import ( + "testing" + "time" + + testclock "k8s.io/utils/clock/testing" + + nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/provider" +) + +// baseTime is an arbitrary fixed instant; the fake clock starts here so tests are +// deterministic (package scripts forbid wall-clock reads anyway). +var baseTime = time.Date(2026, 7, 24, 0, 0, 0, 0, time.UTC) + +// ptr is a compact helper for the three-state BlockScope pointer fields. +func ptr(s string) *string { return &s } + +func spotH100EastScope() provider.BlockScope { + return provider.BlockScope{ + AcceleratorType: ptr("H100"), + CapacityType: nebulav1alpha1.CapacitySpot, + Region: ptr("us-east-1"), + } +} + +// cand builds a query Candidate compactly, keeping the table rows within the +// line-length limit. +func cand(prov, accel string, tier nebulav1alpha1.CapacityType, region string) Candidate { + return Candidate{Provider: prov, AcceleratorType: accel, CapacityType: tier, Region: region} +} + +const ( + spot = nebulav1alpha1.CapacitySpot + onDemand = nebulav1alpha1.CapacityOnDemand +) + +func TestBlocked_ScopeMatchIsPrecise(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + bl.Record("aws", spotH100EastScope(), 10*time.Minute) + + tests := []struct { + name string + c Candidate + want bool + }{ + { + name: "exact match is blocked", + c: cand("aws", "H100", spot, "us-east-1"), + want: true, + }, + { + name: "different tier (OnDemand) survives", + c: cand("aws", "H100", onDemand, "us-east-1"), + want: false, + }, + { + name: "different accelerator survives", + c: cand("aws", "A100", spot, "us-east-1"), + want: false, + }, + { + name: "different region survives", + c: cand("aws", "H100", spot, "us-west-2"), + want: false, + }, + { + name: "different provider survives", + c: cand("modal", "H100", spot, "us-east-1"), + want: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := bl.Blocked(tt.c); got != tt.want { + t.Errorf("Blocked(%+v) = %v, want %v", tt.c, got, tt.want) + } + }) + } +} + +func TestBlocked_WildcardFieldsMatchAnyValue(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + // Wildcard AcceleratorType + wildcard Region (&"") => blocks Spot on aws + // everywhere, any GPU. The pointers are deliberately non-nil-empty: nil would + // mean "no such axis" and match only an empty candidate field. + bl.Record("aws", provider.BlockScope{ + AcceleratorType: ptr(""), + CapacityType: spot, + Region: ptr(""), + }, 10*time.Minute) + + if !bl.Blocked(cand("aws", "H100", spot, "us-east-1")) { + t.Error("H100/us-east-1 Spot should be blocked by the wildcard Spot scope") + } + if !bl.Blocked(cand("aws", "A100", spot, "eu-west-1")) { + t.Error("A100/eu-west-1 Spot should also be blocked by the wildcard Spot scope") + } + if bl.Blocked(cand("aws", "H100", onDemand, "us-east-1")) { + t.Error("OnDemand must survive a Spot-only wildcard block") + } +} + +// A nil axis is "not applicable": it matches only a candidate whose field is empty, +// never a populated one. This is how a region-simple provider (Modal: nil Region) +// or a CPU-only Pod (nil AcceleratorType) blocks without widening across an axis it +// never had. Contrast &"" (wildcard), which matches any value — see the test above. +func TestBlocked_NilFieldMatchesOnlyEmptyCandidate(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + // A region-simple provider's block: nil Region, nil AcceleratorType, Spot only. + bl.Record("modal", provider.BlockScope{CapacityType: spot}, 10*time.Minute) + + // A region-simple candidate carries an empty region and (CPU Pod) empty + // accelerator, so the nil axes match it. + if !bl.Blocked(cand("modal", "", spot, "")) { + t.Error("nil-axis Spot block should match the empty-field candidate") + } + // A populated candidate field is NOT matched by a nil axis: it is out of scope, + // not wildcarded in. + if bl.Blocked(cand("modal", "H100", spot, "")) { + t.Error("nil AcceleratorType must not match a populated accelerator") + } + if bl.Blocked(cand("modal", "", spot, "us-east-1")) { + t.Error("nil Region must not match a populated region") + } +} + +func TestBlocked_DenyAllBlocksWholeProvider(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + bl.Record("aws", provider.BlockScope{DenyAll: true}, 10*time.Minute) + + // Every accelerator/tier/region on aws is blocked... + if !bl.Blocked(cand("aws", "H100", onDemand, "us-west-2")) { + t.Error("DenyAll must block any aws candidate") + } + // ...but a different provider is untouched (a block never spans providers). + if bl.Blocked(Candidate{Provider: "modal", AcceleratorType: "H100"}) { + t.Error("DenyAll on aws must not block modal") + } +} + +func TestBlocked_ExpiresAfterTTL(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + bl.Record("aws", spotH100EastScope(), 10*time.Minute) + + c := cand("aws", "H100", spot, "us-east-1") + if !bl.Blocked(c) { + t.Fatal("should be blocked immediately after Record") + } + + // Just before expiry: still blocked. + fc.Step(10*time.Minute - time.Second) + if !bl.Blocked(c) { + t.Error("should still be blocked one second before TTL") + } + + // Past expiry: cleared. + fc.Step(2 * time.Second) + if bl.Blocked(c) { + t.Error("should be unblocked after TTL lapses") + } +} + +func TestRecord_NoOpOnInvalidInput(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + + // Empty provider and non-positive TTL must not install a block. + bl.Record("", provider.BlockScope{DenyAll: true}, 10*time.Minute) + bl.Record("aws", provider.BlockScope{DenyAll: true}, 0) + bl.Record("aws", provider.BlockScope{DenyAll: true}, -time.Minute) + + if bl.Blocked(Candidate{Provider: "aws", AcceleratorType: "H100"}) { + t.Error("invalid Record inputs must not block anything") + } + if bl.Blocked(Candidate{Provider: "", AcceleratorType: "H100"}) { + t.Error("empty-provider block must not exist") + } +} + +func TestRecord_ExpiredEntriesAreGarbageCollected(t *testing.T) { + fc := testclock.NewFakeClock(baseTime) + bl := newWithClock(fc) + bl.Record("aws", spotH100EastScope(), time.Minute) + + // Let the first entry expire, then record another: the GC on Record should drop + // the stale one so the slice does not grow unbounded. + fc.Step(2 * time.Minute) + bl.Record("aws", provider.BlockScope{DenyAll: true}, time.Minute) + + if n := len(bl.entries); n != 1 { + t.Fatalf("len(entries) = %d, want 1 (stale entry should be GC'd on Record)", n) + } +} diff --git a/pkg/provider/aws/aws.go b/pkg/provider/aws/aws.go new file mode 100644 index 0000000..983958b --- /dev/null +++ b/pkg/provider/aws/aws.go @@ -0,0 +1,480 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package aws implements the provider.Provider interface for Amazon EC2, the +// first hyperscaler (region-aware) backend. +// +// AWS's shape drives several adapter decisions: +// - You do not request a GPU by accelerator type; you request an INSTANCE TYPE +// (p5.48xlarge = 8x H100). The canonical-accelerator -> instance-type mapping +// lives in the catalog's accelerator_id column, so the shared +// catalog.Base.MapAccelerator resolves it and this adapter needs no override. +// - EC2 is region-aware. Every ProvisionRequest carries the optimizer-chosen +// Region (empty => the region this adapter's client was configured with), and +// List/Get report the region an instance actually runs in. Capacity failures +// are per-region, so ClassifyProvisionError scopes the block to one region. +// - EC2 has a real Spot tier with a ~2-minute interruption notice, and instances +// can be stopped/started — so SupportsSpot and SupportsStop are true, unlike +// the NeoCloud adapters. +// - EC2 has native tags, so the ClaimName rides in a tag (no name-encoding hack). +// +// The concrete EC2 API lives behind the Client seam so this package holds only +// provider-agnostic translation and is unit-testable without AWS credentials or +// the SDK. A real Client wrapping the AWS SDK is wired in via NewSDKClient. +package aws + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + corev1 "k8s.io/api/core/v1" + + nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/provider" + "github.com/InftyAI/Nebula/pkg/provider/catalog" + "github.com/InftyAI/Nebula/pkg/util" +) + +// preemptionNotice is EC2's Spot interruption warning lead time (the "2-minute +// warning"). Declared so the control plane can poll Spot claims faster and drain +// ahead of a reclaim. +const preemptionNotice = 2 * time.Minute + +// spotPollInterval is how often to re-list on this adapter. Because EC2 Spot +// interruptions are common and abrupt (the 2-minute notice is not pushed to us), +// we poll faster than the vnode default so a reclaim is noticed promptly. +const spotPollInterval = 10 * time.Second + +// provisionTimeout is this adapter's override of the vnode handler's generic +// Provision deadline (Capabilities.ProvisionTimeout). EC2 capacity is per-zone, +// so RunInstance fails over across the region's availability zones on a capacity +// error; AWS raises the deadline above the generic default to leave room for that +// inner zone sweep, while still capping it so a capacity-starved region yields +// promptly to the outer region-level failover rather than stalling a Pod. It +// bounds the launch attempts, not "the container became healthy" (that is the +// poll loop's job). +const provisionTimeout = 2 * time.Minute + +// ErrSpotCapacity is a marker the Client wraps onto a Spot-tier capacity failure +// (alongside provider.ErrNoCapacity) so ClassifyProvisionError — which the +// interface hands only the error, not the request — can recover that the failing +// tier was Spot and block only Spot, leaving OnDemand serviceable. A real client +// wraps it as fmt.Errorf("...: %w: %w", provider.ErrNoCapacity, aws.ErrSpotCapacity). +var ErrSpotCapacity = errors.New("aws: spot capacity") + +// compile-time assertion that Provider satisfies the interface. +var _ provider.Provider = (*Provider)(nil) + +// ClaimTagKey is the EC2 tag under which the NodeClaim name is stored, so +// List/Get can recover Nebula identity. EC2 has native tags, so no name-encoding +// hack is needed. +const ClaimTagKey = "nebula.inftyai.com/claim" + +// Client is the narrow seam over EC2's API, expressed in provider-agnostic terms +// so a real SDK-backed implementation and a test fake are interchangeable. Only +// the operations the adapter needs are exposed. +type Client interface { + // RunInstance launches exactly one instance from spec and returns its EC2 + // instance id. + RunInstance(ctx context.Context, spec InstanceSpec) (id string, err error) + // TerminateInstance terminates an instance by id. Must be idempotent: + // terminating an already-gone instance returns nil. + TerminateInstance(ctx context.Context, id string) error + // DescribeInstance returns one instance, or (nil, nil) if it no longer exists. + DescribeInstance(ctx context.Context, id string) (*EC2Instance, error) + // ListInstances returns every Nebula-owned instance (filtered by the + // ClaimTagKey tag) across the region, in as few calls as possible. + ListInstances(ctx context.Context) ([]EC2Instance, error) + // AvailableInstanceTypes returns the set of EC2 instance types the client's + // region actually offers, as a set keyed by instance type. It backs the + // per-region availability filter in Offerings: a static catalog row whose + // instance type is absent here is not offered in this region. Which types a + // region offers is AWS-authoritative and changes over time, so it is queried + // live rather than hand-maintained in the catalog. + AvailableInstanceTypes(ctx context.Context) (map[string]bool, error) +} + +// InstanceSpec is the resolved, EC2-shaped request the Client turns into a +// RunInstances call. The adapter builds it from the Pod (source of truth) plus +// the resolved instance type and capacity tier. +type InstanceSpec struct { + // InstanceType is the EC2 instance type resolved from the accelerator (e.g. + // "p5.48xlarge"), via the catalog's accelerator_id column. + InstanceType string + // Image is the container image, from the Pod's first container. The Client is + // responsible for launching it (e.g. via a GPU AMI + user-data, or ECS/EKS). + Image string + // Command is the container command+args, from the Pod. + Command []string + // Env is the environment, flattened from the Pod's container env. + Env map[string]string + // Spot requests interruptible capacity when true (OnDemand otherwise). + Spot bool + // Region is where to launch, in EC2's vocabulary. Empty => the Client's + // configured default region. + Region string + // Tags carry Nebula identity; ClaimTagKey holds the NodeClaim name. + Tags map[string]string +} + +// EC2Instance is the adapter-level view of one EC2 instance as observed. +type EC2Instance struct { + ID string + Tags map[string]string + State string // EC2's own state name, normalized by toState. + Region string + // PublicEndpoint is the reachable address once running (e.g. public DNS/IP). + PublicEndpoint string + // Spot is true when the instance was launched as Spot capacity. + Spot bool + // StatusChecksPassed is true once BOTH EC2 reachability checks (system and + // instance, the "2/2 checks passed") report ok. An instance enters the running + // state a minute or two before its checks pass, so toState holds a running-but- + // unchecked instance at Pending: "running" is not yet "reachable". + StatusChecksPassed bool +} + +// Provider is the EC2 implementation of provider.Provider. It embeds catalog.Base +// for the generic catalog methods (Name, Offerings, and MapAccelerator — which +// resolves the accelerator_id/instance-type mapping straight from the catalog, so +// no override is needed here) and implements the EC2-specific lifecycle. +type Provider struct { + catalog.Base + client Client + // region is the adapter's configured default region, used when a + // ProvisionRequest does not pin one and stamped onto observed instances that + // do not report their own. + region string +} + +// New returns an EC2 Provider backed by client and price catalog. region is the +// default region the client was configured with (used when a request omits one); +// cat is the catalog.Lookup seam so tests can inject a fake. +func New(client Client, cat catalog.Lookup, region string) *Provider { + return &Provider{ + Base: catalog.Base{ProviderName: provider.ProviderAWS, Catalog: cat}, + client: client, + region: region, + } +} + +// Capabilities implements provider.Provider. See the package doc for why each +// trait is set the way it is. +func (p *Provider) Capabilities() provider.Capabilities { + return provider.Capabilities{ + SupportsStop: true, // EC2 instances stop/start + SupportsSpot: true, // real interruptible tier + NativeTags: true, // EC2 tags carry identity + PreemptionNotice: preemptionNotice, // Spot 2-minute warning + PollInterval: spotPollInterval, // Spot reclaims are abrupt; poll faster than default + ProvisionTimeout: provisionTimeout, // caps the per-zone capacity failover loop + } +} + +// Offerings implements provider.Provider, overriding the generic +// catalog.Base.Offerings. It combines the two halves of AWS's price/availability +// truth: +// +// - The catalog CSV holds the DURABLE facts — the (accelerator, count) -> +// instance-type mapping and a seed price — with a BLANK region, since the same +// instance type serves the same pair in every region. +// - AWS itself holds the PER-REGION truth: which instance types the configured +// region actually offers, queried live via the AvailableInstanceTypes probe. +// +// So each static row is stamped with this adapter's region and its Available flag +// is AND-ed with the live probe: a row survives as available only if the catalog +// seeded it available AND the region currently offers its instance type. Rows for +// types the region does not offer are still returned (so the optimizer can see the +// price) but marked unavailable. A probe failure is surfaced as an error rather +// than silently reporting the stale seed as truth. +func (p *Provider) Offerings(ctx context.Context) ([]provider.Offering, error) { + rows := p.Catalog.Offerings(p.ProviderName) + if len(rows) == 0 { + return nil, nil + } + avail, err := p.client.AvailableInstanceTypes(ctx) + if err != nil { + return nil, fmt.Errorf("aws: probe instance-type availability in %s: %w", p.region, err) + } + out := make([]provider.Offering, len(rows)) + for i, o := range rows { + o.Region = p.region + o.Available = o.Available && avail[o.AcceleratorID] + out[i] = o + } + return out, nil +} + +// Provision implements provider.Provider. The Pod is the source of truth for the +// workload; req carries the claim identity, capacity tier, and region. +func (p *Provider) Provision(ctx context.Context, pod *corev1.Pod, req provider.ProvisionRequest) (string, error) { + if pod == nil { + return "", errors.New("aws: nil pod") + } + if req.ClaimName == "" { + return "", errors.New("aws: empty ClaimName in ProvisionRequest") + } + + // Idempotency: if an instance already carries this claim tag, return it rather + // than launching a second (guards a retry after a partial create). + if existing, err := p.findByClaim(ctx, req.ClaimName); err != nil { + return "", err + } else if existing != nil { + return existing.ID, nil + } + + spec, err := p.instanceSpecFromPod(pod, req) + if err != nil { + return "", err + } + // The Provision deadline is enforced generically by the vnode handler (from + // Capabilities.ProvisionTimeout), so RunInstance simply honors ctx as it fails + // over across zones — no adapter-local WithTimeout here. + return p.client.RunInstance(ctx, spec) +} + +// Terminate implements provider.Provider. Idempotent by the Client contract. +func (p *Provider) Terminate(ctx context.Context, instanceID string) error { + if instanceID == "" { + return nil // nothing provisioned yet; treat as already gone + } + return p.client.TerminateInstance(ctx, instanceID) +} + +// Get implements provider.Provider. +func (p *Provider) Get(ctx context.Context, instanceID string) (*provider.Instance, error) { + ec2, err := p.client.DescribeInstance(ctx, instanceID) + if err != nil { + return nil, err + } + if ec2 == nil { + return nil, nil // absent => terminated, per interface contract + } + inst := p.toInstance(*ec2) + return &inst, nil +} + +// List implements provider.Provider. +func (p *Provider) List(ctx context.Context) ([]provider.Instance, error) { + instances, err := p.client.ListInstances(ctx) + if err != nil { + return nil, err + } + out := make([]provider.Instance, 0, len(instances)) + for _, ec2 := range instances { + out = append(out, p.toInstance(ec2)) + } + return out, nil +} + +// ClassifyProvisionError implements provider.Provider. The failure CATEGORIES +// and scope-derivation rule are shared (provider.ClassifyError / the ErrX +// sentinels), so this method supplies only what is EC2-specific. Two things set +// AWS apart from the OnDemand-only NeoClouds: +// +// - Both tiers exist, so the tier to stamp on an accelerator-scoped block is +// read from the wrapped sentinel: a Spot no-capacity error must block only +// Spot, leaving OnDemand serviceable. ClassifyError does not see the request, +// so we detect the Spot case here (ErrNoCapacity wrapped with the spot marker) +// and pass the right tier. +// - EC2 capacity is per-region and this adapter is bound to one region (its +// client's), so an accelerator/capacity block is confined to p.region — a +// "no capacity in us-east-1" failure does not disqualify the same request +// against the us-west-2 adapter. Whole-provider blocks (auth/quota) are left +// region-wide, since bad credentials fail in every region. +func (p *Provider) ClassifyProvisionError(err error, accelerator string) provider.BlockScope { + if err == nil { + return provider.BlockScope{} + } + tier := nebulav1alpha1.CapacityOnDemand + if errors.Is(err, ErrSpotCapacity) { + tier = nebulav1alpha1.CapacitySpot + } + scope := provider.ClassifyError(err, tier, accelerator) + // Accelerator/capacity blocks are per-region (EC2 capacity is regional and + // this adapter is bound to one region); a DenyAll (auth/quota) fails in every + // region, so it stays region-wide (Region left nil). The exact-region pointer + // confines the block to p.region so the same request survives in another region. + // + // p.region is correct ONLY because this adapter is single-region: the client is + // pinned to cfg.Region at construction and runInSubnet ignores spec.Region, so + // every launch lands in p.region regardless of the request's Region — the block + // region and the launch region are the same. If the adapter is ever made + // multi-region (runInSubnet honoring spec.Region, re-resolving AMI/subnets per + // region), this MUST become the request's region (regionOrDefault(req.Region)) — + // which means threading it through ClassifyProvisionError — or failover would + // block the wrong region and keep retrying the one that actually failed. + if !scope.DenyAll { + region := p.region + scope.Region = ®ion + } + return scope +} + +// findByClaim returns the instance tagged with claimName, or nil if none. +func (p *Provider) findByClaim(ctx context.Context, claimName string) (*provider.Instance, error) { + instances, err := p.client.ListInstances(ctx) + if err != nil { + return nil, err + } + for _, ec2 := range instances { + if ec2.Tags[ClaimTagKey] == claimName { + inst := p.toInstance(ec2) + return &inst, nil + } + } + return nil, nil +} + +// instanceSpecFromPod reads the workload off the Pod (source of truth) and the +// accelerator type (from the AcceleratorTypeLabel), maps it to an EC2 instance +// type via the catalog, and stamps the claim tag, capacity tier, and region. +func (p *Provider) instanceSpecFromPod(pod *corev1.Pod, req provider.ProvisionRequest) (InstanceSpec, error) { + if len(pod.Spec.Containers) == 0 { + return InstanceSpec{}, errors.New("aws: pod has no containers") + } + c := pod.Spec.Containers[0] + + env := make(map[string]string, len(c.Env)) + for _, e := range c.Env { + // ValueFrom (secrets/configmaps) is not resolved here; the real Client + // wiring must project those. Plain values are copied through. + if e.ValueFrom == nil { + env[e.Name] = e.Value + } + } + + // Accelerator type comes from the AcceleratorTypeLabel; the count rides on the + // nvidia.com/gpu resource. On EC2 both are lookup keys: the instance type is the + // one whose (accelerator_type, gpu_count) pair matches, since the GPU count is + // baked into the instance type (T4x1 = g4dn.xlarge, T4x8 = g4dn.metal) rather + // than a free knob. So — unlike the identity-mapped NeoClouds — this does NOT + // use the type-only MapAccelerator; it resolves by type AND count. + canonical, count, err := util.AcceleratorRequest(pod) + if err != nil { + return InstanceSpec{}, fmt.Errorf("aws: %w", err) + } + if canonical == "" || count <= 0 { + return InstanceSpec{}, errors.New( + "aws: pod requests no accelerator; EC2 GPU provisioning needs an accelerator type and count") + } + instanceType, ok := p.instanceTypeFor(canonical, count) + if !ok { + return InstanceSpec{}, fmt.Errorf("aws: no EC2 instance type for %s x%d", canonical, count) + } + + return InstanceSpec{ + InstanceType: instanceType, + Image: c.Image, + Command: append(append([]string{}, c.Command...), c.Args...), + Env: env, + Spot: req.CapacityType == nebulav1alpha1.CapacitySpot, + Region: p.regionOrDefault(req.Region), + Tags: map[string]string{ClaimTagKey: req.ClaimName}, + }, nil +} + +// instanceTypeFor resolves the EC2 instance type serving canonical accelerators +// at the requested count, by matching a catalog offering on BOTH the accelerator +// type (case-insensitively) and gpu_count. This is the AWS-specific replacement +// for the type-only MapAccelerator: because an EC2 instance type has a fixed GPU +// count, (type, count) — not type alone — selects it. Region and capacity tier +// are NOT part of the key here: the same (type, count) maps to the same instance +// type in every region/tier, so the first matching row's AcceleratorID is +// authoritative. Returns ok=false when no offering serves that pair (e.g. an +// unsupported accelerator, or a count with no single-instance shape — there is no +// T4x2 instance type). +func (p *Provider) instanceTypeFor(canonical string, count int32) (instanceType string, ok bool) { + for _, o := range p.Catalog.Offerings(p.ProviderName) { + if o.GPUCount == count && strings.EqualFold(o.AcceleratorType, canonical) && o.AcceleratorID != "" { + return o.AcceleratorID, true + } + } + return "", false +} + +// regionOrDefault returns the requested region, or the adapter's configured +// default when the request does not pin one. +func (p *Provider) regionOrDefault(region string) string { + if region != "" { + return region + } + return p.region +} + +// toInstance normalizes an observed EC2 instance into the provider-agnostic +// Instance. +func (p *Provider) toInstance(ec2 EC2Instance) provider.Instance { + tier := nebulav1alpha1.CapacityOnDemand + if ec2.Spot { + tier = nebulav1alpha1.CapacitySpot + } + region := ec2.Region + if region == "" { + region = p.region + } + return provider.Instance{ + ID: ec2.ID, + ClaimName: ec2.Tags[ClaimTagKey], + State: toState(ec2.State, ec2.StatusChecksPassed), + Endpoint: ec2.PublicEndpoint, + CapacityType: tier, + Region: region, + } +} + +// EC2 instance state names (a subset; see the EC2 API InstanceState). toState +// normalizes these to the provider-agnostic lifecycle state. +const ( + stateRunning = "running" + statePending = "pending" + stateStopping = "stopping" + stateStopped = "stopped" + stateShuttingDown = "shutting-down" + stateTerminated = "terminated" +) + +// toState maps EC2's instance-state name to the provider-agnostic lifecycle +// state. "running" is up but not necessarily reachable; "pending" is coming up; +// "terminated"/"shutting-down" are gone; "stopping"/"stopped" are treated as +// Terminated for scheduling purposes (a stopped instance is not serving the +// workload, and the NodeClaim ledger's recovery model is delete-and-recreate). +// An unrecognized state maps to Pending so the poll loop keeps watching rather +// than declaring a premature terminal state. +// +// A running instance is only reported Running once its reachability checks pass +// (statusChecksPassed). EC2 flips an instance to "running" a minute or two before +// its 2/2 checks clear; reporting Running that early would advance the Pod (and +// the owning Deployment) before the instance can actually be reached, so a +// running-but-unchecked instance is held at Pending until the checks pass. +func toState(ec2State string, statusChecksPassed bool) provider.InstanceState { + switch ec2State { + case stateRunning: + if !statusChecksPassed { + return provider.InstancePending + } + return provider.InstanceRunning + case statePending: + return provider.InstancePending + case stateTerminated, stateShuttingDown, stateStopping, stateStopped: + return provider.InstanceTerminated + default: + return provider.InstancePending + } +} diff --git a/pkg/provider/aws/aws_test.go b/pkg/provider/aws/aws_test.go new file mode 100644 index 0000000..aca4168 --- /dev/null +++ b/pkg/provider/aws/aws_test.go @@ -0,0 +1,584 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package aws + +import ( + "context" + "errors" + "fmt" + "reflect" + "testing" + "time" + + awssdk "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/ec2" + ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/api/resource" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + + nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" + "github.com/InftyAI/Nebula/pkg/provider" + "github.com/InftyAI/Nebula/pkg/util" +) + +// fakeClient is an in-memory Client for tests. It records the last RunInstance +// spec and lets tests seed existing instances and inject errors. +type fakeClient struct { + instances []EC2Instance + lastSpec InstanceSpec + runCnt int + runErr error + runID string + terminated []string + // available is the set the AvailableInstanceTypes probe returns; availErr + // injects a probe failure. Both default to zero (empty set / no error), so an + // Offerings call on a bare fake reports every row unavailable but still present. + available map[string]bool + availErr error +} + +func (f *fakeClient) RunInstance(_ context.Context, spec InstanceSpec) (string, error) { + f.runCnt++ + f.lastSpec = spec + if f.runErr != nil { + return "", f.runErr + } + id := f.runID + if id == "" { + id = "i-new" + } + f.instances = append(f.instances, EC2Instance{ + ID: id, + Tags: spec.Tags, + State: statePending, + Region: spec.Region, + Spot: spec.Spot, + }) + return id, nil +} + +func (f *fakeClient) TerminateInstance(_ context.Context, id string) error { + f.terminated = append(f.terminated, id) + return nil +} + +func (f *fakeClient) DescribeInstance(_ context.Context, id string) (*EC2Instance, error) { + for i := range f.instances { + if f.instances[i].ID == id { + inst := f.instances[i] + return &inst, nil + } + } + return nil, nil +} + +func (f *fakeClient) ListInstances(_ context.Context) ([]EC2Instance, error) { + return f.instances, nil +} + +func (f *fakeClient) AvailableInstanceTypes(_ context.Context) (map[string]bool, error) { + if f.availErr != nil { + return nil, f.availErr + } + return f.available, nil +} + +// fakeCatalog is a trivial catalog.Lookup for tests. Its rows carry the +// AcceleratorID (EC2 instance type) so MapAccelerator resolves the way AWS +// requires — by instance type, not accelerator name. +type fakeCatalog struct{ rows []provider.Offering } + +func (c fakeCatalog) Offerings(_ string) []provider.Offering { return c.rows } + +const testRegion = "us-east-1" + +// offering is a compact constructor for a region-aware test catalog row (the +// full struct literal is too wide to read inline). count is the gpu_count, the +// AWS lookup key that (with the accelerator type) selects the instance type. +func offering(accel, id string, count int32, tier nebulav1alpha1.CapacityType, price float64) provider.Offering { + return provider.Offering{ + AcceleratorType: accel, + AcceleratorID: id, + GPUCount: count, + CapacityType: tier, + PricePerHour: price, + Available: true, + Region: testRegion, + } +} + +// newTestProvider builds a Provider with a fake client and a small region-aware +// catalog whose (accelerator_type, gpu_count) pairs map to EC2 instance types. +// T4 appears at two counts so the count-aware resolution is exercised. +func newTestProvider(f *fakeClient) *Provider { + return New(f, fakeCatalog{rows: []provider.Offering{ + offering("H100", "p5.48xlarge", 8, nebulav1alpha1.CapacityOnDemand, 98.32), + offering("H100", "p5.48xlarge", 8, nebulav1alpha1.CapacitySpot, 34.41), + offering("A100-80GB", "p4de.24xlarge", 8, nebulav1alpha1.CapacityOnDemand, 40.96), + offering("T4", "g4dn.xlarge", 1, nebulav1alpha1.CapacityOnDemand, 0.526), + offering("T4", "g4dn.metal", 8, nebulav1alpha1.CapacityOnDemand, 7.824), + }}, testRegion) +} + +// gpuPod builds a Pod whose accelerator type rides on the accelerator-type label +// and whose count rides on the container's nvidia.com/gpu resource; count<=0 +// means CPU-only (no label, no GPU resource). +func gpuPod(accel string, count int64) *corev1.Pod { + c := corev1.Container{ + Name: "main", + Image: "myimg:latest", + Command: []string{"run"}, + Args: []string{"--flag"}, + Env: []corev1.EnvVar{{Name: "FOO", Value: "bar"}}, + } + pod := &corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{Name: "p", Namespace: "default"}, + Spec: corev1.PodSpec{Containers: []corev1.Container{c}}, + } + if accel != "" && count > 0 { + pod.Labels = map[string]string{nebulav1alpha1.AcceleratorTypeLabel: accel} + pod.Spec.Containers[0].Resources.Limits = corev1.ResourceList{ + util.NvidiaGPUResource: *resource.NewQuantity(count, resource.DecimalSI), + } + } + return pod +} + +func TestProvision_MapsAcceleratorToInstanceType(t *testing.T) { + f := &fakeClient{runID: "i-1"} + p := newTestProvider(f) + + id, err := p.Provision(context.Background(), gpuPod("H100", 8), provider.ProvisionRequest{ + ClaimName: "claim-a", + CapacityType: nebulav1alpha1.CapacityOnDemand, + Region: "us-west-2", + }) + if err != nil { + t.Fatalf("Provision: %v", err) + } + if id != "i-1" { + t.Fatalf("id = %q, want i-1", id) + } + // AWS requests by instance type: H100 must resolve to its accelerator_id. + if f.lastSpec.InstanceType != "p5.48xlarge" { + t.Fatalf("InstanceType = %q, want p5.48xlarge", f.lastSpec.InstanceType) + } + if f.lastSpec.Image != "myimg:latest" { + t.Fatalf("image = %q", f.lastSpec.Image) + } + if f.lastSpec.Spot { + t.Fatalf("Spot = true, want false for an OnDemand request") + } + // The request pinned a region; it must ride through verbatim. + if f.lastSpec.Region != "us-west-2" { + t.Fatalf("Region = %q, want us-west-2 (from request)", f.lastSpec.Region) + } + if got := f.lastSpec.Tags[ClaimTagKey]; got != "claim-a" { + t.Fatalf("claim tag = %q, want claim-a", got) + } + if len(f.lastSpec.Command) != 2 || f.lastSpec.Command[0] != "run" { + t.Fatalf("command = %v", f.lastSpec.Command) + } +} + +func TestProvision_LowercaseAcceleratorLabel(t *testing.T) { + f := &fakeClient{runID: "i-lc"} + p := newTestProvider(f) + + // A user may write the accelerator-type label in any case; it must resolve to + // the canonical catalog row (and thus the right instance type). + if _, err := p.Provision(context.Background(), gpuPod("h100", 8), provider.ProvisionRequest{ + ClaimName: "claim-lc", + }); err != nil { + t.Fatalf("Provision: %v", err) + } + if f.lastSpec.InstanceType != "p5.48xlarge" { + t.Fatalf("InstanceType = %q, want p5.48xlarge from lowercase label", f.lastSpec.InstanceType) + } +} + +func TestProvision_CountSelectsInstanceType(t *testing.T) { + // The GPU count is a lookup key on AWS: the same accelerator at different + // counts must resolve to DIFFERENT instance types (T4x1 = g4dn.xlarge, + // T4x8 = g4dn.metal), because the count is baked into the instance type. + cases := []struct { + count int64 + wantType string + }{ + {1, "g4dn.xlarge"}, + {8, "g4dn.metal"}, + } + for _, tc := range cases { + f := &fakeClient{runID: "i-t4"} + p := newTestProvider(f) + req := provider.ProvisionRequest{ClaimName: "claim-t4"} + if _, err := p.Provision(context.Background(), gpuPod("T4", tc.count), req); err != nil { + t.Fatalf("Provision(T4 x%d): %v", tc.count, err) + } + if f.lastSpec.InstanceType != tc.wantType { + t.Fatalf("T4 x%d -> InstanceType %q, want %q", tc.count, f.lastSpec.InstanceType, tc.wantType) + } + } +} + +func TestProvision_UnsupportedCountIsError(t *testing.T) { + f := &fakeClient{} + p := newTestProvider(f) + // T4 x2 has no instance type (there is no 2-GPU T4 shape): must error rather + // than silently picking the x1 or x8 row. + req := provider.ProvisionRequest{ClaimName: "claim-t4x2"} + if _, err := p.Provision(context.Background(), gpuPod("T4", 2), req); err == nil { + t.Fatal("expected an error for an unsupported (accelerator, count) pair") + } + if f.runCnt != 0 { + t.Fatalf("RunInstance called %d times, want 0", f.runCnt) + } +} + +func TestProvision_SpotSetsMarketOption(t *testing.T) { + f := &fakeClient{runID: "i-spot"} + p := newTestProvider(f) + + if _, err := p.Provision(context.Background(), gpuPod("H100", 8), provider.ProvisionRequest{ + ClaimName: "claim-spot", + CapacityType: nebulav1alpha1.CapacitySpot, + }); err != nil { + t.Fatalf("Provision: %v", err) + } + if !f.lastSpec.Spot { + t.Fatalf("Spot = false, want true for a Spot request") + } +} + +func TestProvision_EmptyRegionFallsBackToAdapterDefault(t *testing.T) { + f := &fakeClient{runID: "i-def"} + p := newTestProvider(f) + + // An empty request Region means "use the adapter's configured default region". + if _, err := p.Provision(context.Background(), gpuPod("H100", 8), provider.ProvisionRequest{ + ClaimName: "claim-def", + }); err != nil { + t.Fatalf("Provision: %v", err) + } + if f.lastSpec.Region != testRegion { + t.Fatalf("Region = %q, want adapter default %q", f.lastSpec.Region, testRegion) + } +} + +func TestProvision_NoAcceleratorIsError(t *testing.T) { + f := &fakeClient{} + p := newTestProvider(f) + + // EC2 GPU provisioning is by instance type; a Pod with no accelerator has no + // instance type to launch, so it must error rather than silently guessing. + req := provider.ProvisionRequest{ClaimName: "claim-cpu"} + if _, err := p.Provision(context.Background(), gpuPod("", 0), req); err == nil { + t.Fatal("expected an error for a Pod requesting no accelerator") + } + if f.runCnt != 0 { + t.Fatalf("RunInstance called %d times, want 0", f.runCnt) + } +} + +func TestProvision_Idempotent(t *testing.T) { + f := &fakeClient{ + instances: []EC2Instance{{ + ID: "i-existing", + Tags: map[string]string{ClaimTagKey: "claim-a"}, + State: stateRunning, + Region: testRegion, + }}, + } + p := newTestProvider(f) + + id, err := p.Provision(context.Background(), gpuPod("H100", 8), provider.ProvisionRequest{ClaimName: "claim-a"}) + if err != nil { + t.Fatalf("Provision: %v", err) + } + if id != "i-existing" { + t.Fatalf("id = %q, want i-existing (idempotent reuse)", id) + } + if f.runCnt != 0 { + t.Fatalf("RunInstance called %d times, want 0 (idempotent)", f.runCnt) + } +} + +func TestProvision_UnsupportedAccelerator(t *testing.T) { + f := &fakeClient{} + p := newTestProvider(f) + req := provider.ProvisionRequest{ClaimName: "claim-x"} + if _, err := p.Provision(context.Background(), gpuPod("TPU-v4", 1), req); err == nil { + t.Fatal("expected error for unsupported accelerator") + } +} + +func TestGetAndList_NormalizeInstance(t *testing.T) { + f := &fakeClient{ + instances: []EC2Instance{{ + ID: "i-1", + Tags: map[string]string{ClaimTagKey: "claim-a"}, + State: stateRunning, + Region: "eu-west-1", + PublicEndpoint: "ec2-1-2-3-4.compute.amazonaws.com", + Spot: true, + StatusChecksPassed: true, // 2/2 checks passed => reported Running + }}, + } + p := newTestProvider(f) + + got, err := p.Get(context.Background(), "i-1") + if err != nil { + t.Fatalf("Get: %v", err) + } + if got == nil { + t.Fatal("Get returned nil for an existing instance") + } + if got.ClaimName != "claim-a" || got.State != provider.InstanceRunning { + t.Fatalf("Get = %+v", got) + } + if got.Region != "eu-west-1" { + t.Fatalf("Region = %q, want eu-west-1 (reported by the instance)", got.Region) + } + if got.CapacityType != nebulav1alpha1.CapacitySpot { + t.Fatalf("CapacityType = %q, want Spot", got.CapacityType) + } + if got.Endpoint != "ec2-1-2-3-4.compute.amazonaws.com" { + t.Fatalf("Endpoint = %q", got.Endpoint) + } + + // A missing instance is (nil, nil): absence == terminated per the contract. + missing, err := p.Get(context.Background(), "i-gone") + if err != nil { + t.Fatalf("Get(missing): %v", err) + } + if missing != nil { + t.Fatalf("Get(missing) = %+v, want nil", missing) + } + + list, err := p.List(context.Background()) + if err != nil { + t.Fatalf("List: %v", err) + } + if len(list) != 1 || list[0].ID != "i-1" { + t.Fatalf("List = %+v", list) + } +} + +func TestToInstance_RegionFallsBackToAdapterDefault(t *testing.T) { + p := newTestProvider(&fakeClient{}) + // An instance that does not report its own region is stamped with the + // adapter's configured region. + inst := p.toInstance(EC2Instance{ID: "i-x", State: stateRunning}) + if inst.Region != testRegion { + t.Fatalf("Region = %q, want adapter default %q", inst.Region, testRegion) + } +} + +func TestTerminate_Idempotent(t *testing.T) { + f := &fakeClient{} + p := newTestProvider(f) + + // Empty id: nothing was provisioned; treat as already gone (no client call). + if err := p.Terminate(context.Background(), ""); err != nil { + t.Fatalf("Terminate(\"\"): %v", err) + } + if len(f.terminated) != 0 { + t.Fatalf("Terminate(\"\") called client, want no-op") + } + if err := p.Terminate(context.Background(), "i-1"); err != nil { + t.Fatalf("Terminate: %v", err) + } + if len(f.terminated) != 1 || f.terminated[0] != "i-1" { + t.Fatalf("terminated = %v, want [i-1]", f.terminated) + } +} + +func TestClassifyProvisionError(t *testing.T) { + p := newTestProvider(&fakeClient{}) + const accel = "H100" + denyAll := provider.BlockScope{DenyAll: true} + // EC2 capacity/accelerator blocks are confined to the adapter's region (an + // exact-match Region pointer) AND to the requested accelerator (passed in by the + // caller — the provider owns the whole scope). Auth/quota (DenyAll) ignores both. + region, wantAccel := testRegion, accel + onDemandRegional := provider.BlockScope{ + AcceleratorType: &wantAccel, CapacityType: nebulav1alpha1.CapacityOnDemand, Region: ®ion, + } + spotRegional := provider.BlockScope{ + AcceleratorType: &wantAccel, CapacityType: nebulav1alpha1.CapacitySpot, Region: ®ion, + } + + // A Spot capacity failure the Client wraps with both sentinels: no-capacity + // (category) + spot (tier), so the block confines to Spot in this region. + spotNoCapacity := fmt.Errorf("%w: %w", provider.ErrNoCapacity, ErrSpotCapacity) + stringNoCapacity := fmt.Errorf("InsufficientInstanceCapacity: no capacity") + + tests := []struct { + name string + err error + want provider.BlockScope + }{ + {"auth is provider-wide (no region)", provider.ErrAuth, denyAll}, + {"quota is provider-wide (no region)", provider.ErrQuota, denyAll}, + {"capacity is regional OnDemand", provider.ErrNoCapacity, onDemandRegional}, + {"wrapped capacity is regional", fmt.Errorf("run: %w", provider.ErrNoCapacity), onDemandRegional}, + {"spot capacity blocks only Spot in region", spotNoCapacity, spotRegional}, + {"string no-capacity is regional OnDemand", stringNoCapacity, onDemandRegional}, + {"unknown is provider-wide", fmt.Errorf("weird transient blip"), denyAll}, + {"nil", nil, provider.BlockScope{}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := p.ClassifyProvisionError(tt.err, accel) + if !reflect.DeepEqual(got, tt.want) { + t.Fatalf("got %+v, want %+v", got, tt.want) + } + }) + } +} + +func TestCapabilities(t *testing.T) { + p := newTestProvider(&fakeClient{}) + caps := p.Capabilities() + if !caps.SupportsStop || !caps.SupportsSpot || !caps.NativeTags { + t.Fatalf("unexpected caps: %+v", caps) + } + if caps.PreemptionNotice != 2*time.Minute { + t.Fatalf("PreemptionNotice = %v, want 2m", caps.PreemptionNotice) + } + if p.Name() != provider.ProviderAWS { + t.Fatalf("name = %q, want %q", p.Name(), provider.ProviderAWS) + } +} + +func TestOfferings_StampsRegionAndFiltersByLiveProbe(t *testing.T) { + // The region offers only g4dn.xlarge and p5.48xlarge (not g4dn.metal or + // p4de.24xlarge), so those two rows stay available and the rest, though still + // returned so the optimizer sees their price, are marked unavailable. + f := &fakeClient{available: map[string]bool{ + "g4dn.xlarge": true, + "p5.48xlarge": true, + }} + p := newTestProvider(f) + + offs, err := p.Offerings(context.Background()) + if err != nil { + t.Fatalf("Offerings: %v", err) + } + // Every catalog row is returned (none dropped), each stamped with the region. + if len(offs) != 5 { + t.Fatalf("got %d offerings, want all 5 catalog rows", len(offs)) + } + availByType := map[string]bool{} + for _, o := range offs { + if o.Region != testRegion { + t.Fatalf("offering %q region = %q, want %q", o.AcceleratorID, o.Region, testRegion) + } + availByType[o.AcceleratorID] = o.Available + } + if !availByType["g4dn.xlarge"] || !availByType["p5.48xlarge"] { + t.Fatalf("types the region offers must stay available: %+v", availByType) + } + if availByType["g4dn.metal"] || availByType["p4de.24xlarge"] { + t.Fatalf("types the region does not offer must be unavailable: %+v", availByType) + } +} + +func TestOfferings_ProbeErrorPropagates(t *testing.T) { + // A probe failure must surface as an error, not silently report the stale seed + // availability as truth. + f := &fakeClient{availErr: errors.New("throttled")} + p := newTestProvider(f) + if _, err := p.Offerings(context.Background()); err == nil { + t.Fatal("expected the probe error to propagate") + } +} + +func TestToState(t *testing.T) { + // The EC2-state → provider-state mapping is load-bearing: an unknown state + // must map to Pending so the poll loop keeps watching rather than declaring a + // premature terminal state; stop/stopping fold into Terminated. "running" is + // gated on the reachability checks — running-but-unchecked stays Pending. + cases := []struct { + state string + checksPassed bool + want provider.InstanceState + }{ + {stateRunning, true, provider.InstanceRunning}, + // Running but the 2/2 checks have not cleared yet: still Pending. + {stateRunning, false, provider.InstancePending}, + {statePending, true, provider.InstancePending}, + {statePending, false, provider.InstancePending}, + {stateStopping, false, provider.InstanceTerminated}, + {stateStopped, false, provider.InstanceTerminated}, + {stateShuttingDown, false, provider.InstanceTerminated}, + {stateTerminated, false, provider.InstanceTerminated}, + {"", false, provider.InstancePending}, + {"rebooting", true, provider.InstancePending}, + } + for _, tc := range cases { + if got := toState(tc.state, tc.checksPassed); got != tc.want { + t.Fatalf("toState(%q, checks=%v) = %q, want %q", tc.state, tc.checksPassed, got, tc.want) + } + } +} + +func TestResolveGPUAMI_PicksNewestAndErrsWhenAbsent(t *testing.T) { + // The newest CreationDate wins so a driver/runtime refresh is picked up. + f := &fakeEC2{imagesOut: &ec2.DescribeImagesOutput{Images: []ec2types.Image{ + {ImageId: awssdk.String("ami-old"), CreationDate: awssdk.String("2024-01-01T00:00:00.000Z")}, + {ImageId: awssdk.String("ami-new"), CreationDate: awssdk.String("2026-01-01T00:00:00.000Z")}, + }}} + c := newSDKClient(f) + got, err := c.resolveGPUAMI(context.Background()) + if err != nil { + t.Fatalf("resolveGPUAMI: %v", err) + } + if got != "ami-new" { + t.Fatalf("resolveGPUAMI = %q, want ami-new (newest)", got) + } + + // No matching image => ErrConfig (AWS unusable in the region, non-fatal skip). + f2 := &fakeEC2{imagesOut: &ec2.DescribeImagesOutput{}} + c2 := newSDKClient(f2) + if _, err := c2.resolveGPUAMI(context.Background()); !errors.Is(err, ErrConfig) { + t.Fatalf("resolveGPUAMI(no images) err = %v, want ErrConfig", err) + } +} + +func TestDiscoverDefaultSubnets_ReturnsPerAZTargets(t *testing.T) { + f := &fakeEC2{subnetsOut: &ec2.DescribeSubnetsOutput{Subnets: []ec2types.Subnet{ + {SubnetId: awssdk.String("subnet-a"), AvailabilityZone: awssdk.String("us-east-1a")}, + {SubnetId: awssdk.String("subnet-b"), AvailabilityZone: awssdk.String("us-east-1b")}, + }}} + c := newSDKClient(f) + got, err := c.discoverDefaultSubnets(context.Background()) + if err != nil { + t.Fatalf("discoverDefaultSubnets: %v", err) + } + if len(got) != 2 || got[0].id != "subnet-a" || got[1].az != "us-east-1b" { + t.Fatalf("subnets = %+v, want the two default-VPC subnets", got) + } + + // No default VPC (empty result) is not an error: launch just skips zone failover. + f2 := &fakeEC2{subnetsOut: &ec2.DescribeSubnetsOutput{}} + c2 := newSDKClient(f2) + if got, err := c2.discoverDefaultSubnets(context.Background()); err != nil || len(got) != 0 { + t.Fatalf("discoverDefaultSubnets(no default VPC) = (%+v, %v), want (nil, nil)", got, err) + } +} diff --git a/pkg/provider/aws/client.go b/pkg/provider/aws/client.go new file mode 100644 index 0000000..86d274d --- /dev/null +++ b/pkg/provider/aws/client.go @@ -0,0 +1,576 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package aws + +import ( + "context" + "errors" + "fmt" + + awssdk "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/config" + "github.com/aws/aws-sdk-go-v2/service/ec2" + ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types" + smithy "github.com/aws/smithy-go" + + "github.com/InftyAI/Nebula/pkg/provider/catalog" +) + +// ec2API is the subset of the EC2 SDK client the sdkClient depends on. Narrowing +// it to an interface lets client.go's translation be tested against a fake SDK +// without a live account (the real *ec2.Client satisfies it). +type ec2API interface { + RunInstances( + ctx context.Context, in *ec2.RunInstancesInput, optFns ...func(*ec2.Options), + ) (*ec2.RunInstancesOutput, error) + TerminateInstances( + ctx context.Context, in *ec2.TerminateInstancesInput, optFns ...func(*ec2.Options), + ) (*ec2.TerminateInstancesOutput, error) + DescribeInstances( + ctx context.Context, in *ec2.DescribeInstancesInput, optFns ...func(*ec2.Options), + ) (*ec2.DescribeInstancesOutput, error) + // DescribeInstanceStatus reports the system+instance reachability checks (the + // "2/2 checks passed" a launched instance clears a minute or two AFTER it + // enters the running state), used to gate the Running lifecycle state so a Pod + // is not reported Running before its instance is actually reachable. + DescribeInstanceStatus( + ctx context.Context, in *ec2.DescribeInstanceStatusInput, optFns ...func(*ec2.Options), + ) (*ec2.DescribeInstanceStatusOutput, error) + DescribeInstanceTypeOfferings( + ctx context.Context, in *ec2.DescribeInstanceTypeOfferingsInput, optFns ...func(*ec2.Options), + ) (*ec2.DescribeInstanceTypeOfferingsOutput, error) + // DescribeImages resolves the region's GPU AMI at construction (see + // resolveGPUAMI); DescribeSubnets discovers the default VPC's per-AZ subnets + // RunInstance fails over across (see discoverDefaultSubnets). + DescribeImages( + ctx context.Context, in *ec2.DescribeImagesInput, optFns ...func(*ec2.Options), + ) (*ec2.DescribeImagesOutput, error) + DescribeSubnets( + ctx context.Context, in *ec2.DescribeSubnetsInput, optFns ...func(*ec2.Options), + ) (*ec2.DescribeSubnetsOutput, error) +} + +// sdkClient is the real Client, backed by aws-sdk-go-v2's EC2 API. All +// EC2-specific SDK calls live here so the adapter (aws.go) and its tests stay +// SDK-free; the pure translation (user-data, error mapping) lives in translate.go. +// +// Execution model (see the package doc): a NodeClaim maps to one EC2 instance +// launched DIRECTLY from RunInstances — no Launch Template, no pre-created infra. +// Everything the launch needs is SELF-CONFIGURED from credentials alone at +// construction, so registering AWS in a new region requires nothing but access: +// +// - amiID: the region's GPU AMI (NVIDIA driver + container runtime baked in), +// resolved live via DescribeImages against the Amazon-owned ECS GPU-optimized +// AMI (see resolveGPUAMI). AMI ids are per-region, so this cannot be a constant. +// - subnets: the default VPC's subnets, one per availability zone, discovered +// via DescribeSubnets (default-for-az=true). EC2 capacity is per-AZ, so trying +// another zone's subnet is the cheapest recovery from InsufficientInstance +// capacity before condemning the whole region. +// - security group: left UNSET. Launching into a default-VPC subnet with no +// SecurityGroupIds makes EC2 attach that VPC's default SG (outbound-open, +// inbound-from-itself) — exactly what a fire-and-forget `docker pull && run` +// needs, with no inbound exposure. +// +// The adapter overrides only what is workload-specific: instance type, user-data +// (the container bootstrap), Spot market option, and the identity tags. +type sdkClient struct { + ec2 ec2API + // region is the resolved region this client operates in, echoed onto observed + // instances that do not otherwise report one. + region string + // amiID is the region's GPU AMI, resolved at construction. Every instance + // launches from it; it is NON-SECRET, AWS-published config, not a credential. + amiID string + // subnets are the default VPC's per-AZ subnets RunInstance fails over across on + // a capacity error, discovered at construction. Empty when the region has no + // default VPC: RunInstance then makes a single attempt letting EC2 pick the + // subnet (still valid, just no zone failover). + subnets []subnet +} + +// subnet is one candidate launch target: its id and the AZ it lives in (for +// logging/ordering). RunInstance launches into one of these per attempt. +type subnet struct { + id string + az string +} + +// compile-time assertion that sdkClient satisfies the adapter's Client seam. +var _ Client = (*sdkClient)(nil) + +// NewSDKClient builds an EC2-backed Provider for region, ready to register. +// +// It is fully self-configuring: given only credentials and a region, it resolves +// the GPU AMI and the default VPC's subnets itself, so registering AWS in a new +// region needs no pre-created Launch Template, tagged subnets, or any other infra +// — matching the creds-only UX of the NeoCloud adapters. +// +// Security contract (fixed regardless of backing): +// - region is the ONLY non-secret config, passed as a plain argument (flag/env). +// Empty means "resolve the SDK's default region" from the standard chain +// (AWS_REGION / shared config / IMDS); the resolved value is echoed on +// observed instances and stamped onto region-scoped blocks. A region that +// resolves to empty is a config skip (ErrConfig). +// - Credentials are SECRETS and are NEVER accepted here. The SDK's default +// credential chain (config.LoadDefaultConfig) is used, preferring IRSA / +// instance-role and falling back to AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY +// from the environment (a Kubernetes Secret / .env, not a flag), keeping +// secrets out of process arguments and logs. +// +// The catalog is loaded via catalog.Load() (embedded CSV / mounted ConfigMap), +// identical to the other adapters. +func NewSDKClient(ctx context.Context, region string) (*Provider, error) { + cat, err := catalog.Load() + if err != nil { + return nil, fmt.Errorf("aws: load price catalog: %w", err) + } + + opts := []func(*config.LoadOptions) error{} + if region != "" { + opts = append(opts, config.WithRegion(region)) + } + cfg, err := config.LoadDefaultConfig(ctx, opts...) + if err != nil { + return nil, fmt.Errorf("aws: load SDK config: %w", err) + } + if cfg.Region == "" { + // Neither the argument nor the SDK chain resolved a region; without one the + // EC2 endpoint is undefined. Treat as unconfigured (non-fatal skip). + return nil, fmt.Errorf("aws: no region resolved: %w", ErrConfig) + } + + c := &sdkClient{ + ec2: ec2.NewFromConfig(cfg), + region: cfg.Region, + } + + // Resolve the region's GPU AMI (required — no AMI, nothing to launch) and the + // default VPC's per-AZ subnets (best-effort — no default VPC leaves c.subnets + // empty and RunInstance lets EC2 pick the subnet, just without zone failover). + amiID, err := c.resolveGPUAMI(ctx) + if err != nil { + return nil, fmt.Errorf("aws: resolve GPU AMI in %s: %w", cfg.Region, err) + } + c.amiID = amiID + + subnets, err := c.discoverDefaultSubnets(ctx) + if err != nil { + return nil, fmt.Errorf("aws: discover default-VPC subnets in %s: %w", cfg.Region, err) + } + c.subnets = subnets + + return New(c, cat, cfg.Region), nil +} + +// gpuAMINameFilter matches the Amazon-owned Amazon-Linux-2 ECS GPU-optimized AMI. +// DescribeImages is filtered to this pattern and owner "amazon", then the newest +// by creation date is chosen — so the AMI stays current without a hand-maintained +// per-region id table. This AMI ships the NVIDIA driver and a Docker runtime, +// which is all buildUserData's `docker pull && docker run --gpus all` needs. +const gpuAMINameFilter = "amzn2-ami-ecs-gpu-hvm-*-x86_64-ebs" + +// ErrConfig marks a missing/invalid non-secret AWS configuration (no resolvable +// region, no GPU AMI in the region). It is a distinct sentinel so cmd/main.go can +// treat "AWS not configured" as a non-fatal skip (errors.Is) — AWS simply is not +// registered — rather than crashing the manager, mirroring a creds-absent skip +// for the other providers. +var ErrConfig = errors.New("aws: not configured") + +// RunInstance implements Client. It launches exactly one instance directly (from +// the resolved GPU AMI into a default-VPC subnet), overriding only the +// workload-specific fields. +// +// EC2 capacity is per-availability-zone, so a launch is not a single shot: when a +// zone reports InsufficientInstanceCapacity (or a Spot-capacity variant), +// RunInstance retries in the next discovered zone before giving up. Only when +// every candidate zone is capacity-starved does it return the wrapped +// ErrNoCapacity, so the outer region-level failover (ClassifyProvisionError + +// NodePool regions) fires only after this cheaper inner failover is exhausted. A +// non-capacity failure (auth, quota, bad config) is terminal and returned +// immediately — retrying other zones would not help and only burns the deadline. +// The ctx deadline (set by the vnode handler from Capabilities.ProvisionTimeout) +// bounds the whole sweep. +// +// When no subnets were discovered (no default VPC), the loop runs once with a +// zero-value target, letting EC2 pick the subnet — a valid single-shot launch, +// just without zone failover. +func (c *sdkClient) RunInstance(ctx context.Context, spec InstanceSpec) (string, error) { + userData, err := buildUserData(spec) + if err != nil { + return "", err + } + + targets := c.subnets + if len(targets) == 0 { + targets = []subnet{{}} // single attempt, letting EC2 pick the subnet + } + + var lastErr error + for _, sn := range targets { + // Honor the deadline between attempts: a slow first zone must not let the + // sweep overrun the provision timeout. + if err := ctx.Err(); err != nil { + if lastErr != nil { + return "", lastErr + } + return "", err + } + + id, err := c.runInSubnet(ctx, spec, userData, sn) + if err == nil { + return id, nil + } + classified := classifyEC2Error(err, spec.Spot) + // Only a ZONE-LOCAL capacity shortfall justifies trying the next zone: a + // sibling AZ's default subnet may still satisfy the launch. Everything else + // is terminal for this sweep and surfaced immediately — auth/quota/unknown + // (retrying other zones would not help), and crucially a region/account- + // scoped capacity error (a Spot price ceiling or per-region Spot limit), for + // which every AZ would fail identically. Stopping now hands such an error + // straight to region/tier failover instead of burning the deadline on a + // futile sweep. It still carries ErrNoCapacity, so ClassifyProvisionError + // blocks the region correctly. + if !errors.Is(classified, errZoneLocal) { + return "", classified + } + lastErr = classified + } + // Every zone was capacity-starved; return the last capacity error so + // ClassifyProvisionError can block the region and region-failover can proceed. + return "", lastErr +} + +// runInSubnet issues one RunInstances attempt. sn.id empty => let EC2 pick the +// subnet (no default VPC discovered); otherwise launch into that subnet (and thus +// its AZ). SecurityGroupIds is deliberately unset: launching into a default-VPC +// subnet attaches that VPC's default SG (outbound-open, no inbound exposure), +// which is all a fire-and-forget container launch needs. +func (c *sdkClient) runInSubnet(ctx context.Context, spec InstanceSpec, userData string, sn subnet) (string, error) { + in := &ec2.RunInstancesInput{ + MinCount: awssdk.Int32(1), + MaxCount: awssdk.Int32(1), + ImageId: awssdk.String(c.amiID), + InstanceType: ec2types.InstanceType(spec.InstanceType), + UserData: awssdk.String(userData), + TagSpecifications: []ec2types.TagSpecification{{ + ResourceType: ec2types.ResourceTypeInstance, + Tags: ec2Tags(spec.Tags), + }}, + } + if sn.id != "" { + in.SubnetId = awssdk.String(sn.id) + } + if spec.Spot { + in.InstanceMarketOptions = &ec2types.InstanceMarketOptionsRequest{ + MarketType: ec2types.MarketTypeSpot, + } + } + + out, err := c.ec2.RunInstances(ctx, in) + if err != nil { + return "", err // caller classifies (needs the raw error to detect capacity) + } + if len(out.Instances) == 0 || out.Instances[0].InstanceId == nil { + return "", errors.New("aws: RunInstances returned no instance id") + } + return *out.Instances[0].InstanceId, nil +} + +// resolveGPUAMI finds the region's GPU AMI: the newest Amazon-owned image whose +// name matches gpuAMINameFilter (the ECS GPU-optimized AMI). AMI ids are +// per-region, so this is queried live rather than hardcoded — the whole point of +// the self-configuring model. A region that returns no matching image is a config +// error (ErrConfig): AWS is effectively not usable there, and the caller skips it +// non-fatally rather than launching from a missing AMI. +func (c *sdkClient) resolveGPUAMI(ctx context.Context) (string, error) { + out, err := c.ec2.DescribeImages(ctx, &ec2.DescribeImagesInput{ + Owners: []string{"amazon"}, + Filters: []ec2types.Filter{ + {Name: awssdk.String("name"), Values: []string{gpuAMINameFilter}}, + {Name: awssdk.String("state"), Values: []string{"available"}}, + }, + }) + if err != nil { + return "", err + } + // Pick the newest by CreationDate (RFC3339 strings sort lexicographically in + // chronological order), so a driver/runtime refresh is picked up automatically. + var newest ec2types.Image + for _, img := range out.Images { + if img.ImageId == nil || img.CreationDate == nil { + continue + } + if newest.CreationDate == nil || *img.CreationDate > *newest.CreationDate { + newest = img + } + } + if newest.ImageId == nil { + return "", fmt.Errorf("no GPU AMI (%s) offered: %w", gpuAMINameFilter, ErrConfig) + } + return *newest.ImageId, nil +} + +// discoverDefaultSubnets lists the default VPC's subnets — one default subnet per +// availability zone (default-for-az=true) — as RunInstance's per-zone capacity +// failover targets. It pages through all matches. An empty result (the region has +// no default VPC) is NOT an error: the caller launches without pinning a subnet, +// giving up zone failover but still functioning. Using only the default subnets +// keeps Nebula from placing instances on networks an operator did not intend. +func (c *sdkClient) discoverDefaultSubnets(ctx context.Context) ([]subnet, error) { + in := &ec2.DescribeSubnetsInput{ + Filters: []ec2types.Filter{{ + Name: awssdk.String("default-for-az"), + Values: []string{"true"}, + }}, + } + var out []subnet + for { + page, err := c.ec2.DescribeSubnets(ctx, in) + if err != nil { + return nil, err + } + for _, s := range page.Subnets { + if s.SubnetId == nil { + continue + } + sn := subnet{id: *s.SubnetId} + if s.AvailabilityZone != nil { + sn.az = *s.AvailabilityZone + } + out = append(out, sn) + } + if page.NextToken == nil { + break + } + in.NextToken = page.NextToken + } + return out, nil +} + +// TerminateInstance implements Client. Idempotent: an already-gone instance +// (InvalidInstanceID.NotFound) returns nil so the finalizer can retry safely. +func (c *sdkClient) TerminateInstance(ctx context.Context, id string) error { + _, err := c.ec2.TerminateInstances(ctx, &ec2.TerminateInstancesInput{ + InstanceIds: []string{id}, + }) + if err != nil && !isNotFound(err) { + return err + } + return nil +} + +// DescribeInstance implements Client. Returns (nil, nil) when the instance no +// longer exists (absent == terminated per the interface contract). +func (c *sdkClient) DescribeInstance(ctx context.Context, id string) (*EC2Instance, error) { + out, err := c.ec2.DescribeInstances(ctx, &ec2.DescribeInstancesInput{ + InstanceIds: []string{id}, + }) + if err != nil { + if isNotFound(err) { + return nil, nil + } + return nil, err + } + for _, r := range out.Reservations { + if len(r.Instances) == 0 { + continue + } + inst := c.observe(r.Instances[0]) + // Fold in reachability checks so a single-instance Get reports Running only + // once the 2/2 checks pass, matching ListInstances. A status probe failure + // is non-fatal: the instance is still returned (checks stay false). + if inst.ID != "" { + if passed, err := c.statusChecksByID(ctx, []string{inst.ID}); err == nil { + inst.StatusChecksPassed = passed[inst.ID] + } + } + return &inst, nil + } + return nil, nil +} + +// ListInstances implements Client. It scopes the list server-side to instances +// carrying the ClaimTagKey tag — the tag every Nebula instance is launched with — +// so instances Nebula does not own are never returned, and pages through all +// results. This is the engine of the poll loop. +func (c *sdkClient) ListInstances(ctx context.Context) ([]EC2Instance, error) { + in := &ec2.DescribeInstancesInput{ + Filters: []ec2types.Filter{{ + Name: awssdk.String("tag-key"), + Values: []string{ClaimTagKey}, + }}, + } + var out []EC2Instance + var ids []string + for { + page, err := c.ec2.DescribeInstances(ctx, in) + if err != nil { + return nil, err + } + for _, r := range page.Reservations { + for i := range r.Instances { + inst := c.observe(r.Instances[i]) + out = append(out, inst) + if inst.ID != "" { + ids = append(ids, inst.ID) + } + } + } + if page.NextToken == nil { + break + } + in.NextToken = page.NextToken + } + + // Fold in reachability checks: DescribeInstances does not report them, so a + // separate DescribeInstanceStatus call decides which running instances have + // actually passed their 2/2 checks. A failure here is not fatal — the instances + // are still returned (StatusChecksPassed stays false), so the poll loop simply + // holds them at Pending one more tick rather than the whole List erroring out. + passed, err := c.statusChecksByID(ctx, ids) + if err != nil { + return out, nil + } + for i := range out { + out[i].StatusChecksPassed = passed[out[i].ID] + } + return out, nil +} + +// statusChecksByID returns the subset of ids whose system AND instance +// reachability checks both report "ok" (the "2/2 checks passed"). It pages +// through DescribeInstanceStatus; an id absent from the result (or not ok) maps +// to false. Called with no ids it returns an empty map without hitting the API. +func (c *sdkClient) statusChecksByID(ctx context.Context, ids []string) (map[string]bool, error) { + passed := make(map[string]bool, len(ids)) + if len(ids) == 0 { + return passed, nil + } + in := &ec2.DescribeInstanceStatusInput{InstanceIds: ids} + for { + page, err := c.ec2.DescribeInstanceStatus(ctx, in) + if err != nil { + return nil, err + } + for _, s := range page.InstanceStatuses { + if s.InstanceId == nil { + continue + } + passed[*s.InstanceId] = statusChecksOK(s) + } + if page.NextToken == nil { + break + } + in.NextToken = page.NextToken + } + return passed, nil +} + +// statusChecksOK reports whether both reachability checks (system and instance) +// are "ok". EC2 exposes them as an overall summary plus per-check details; the +// summary is "ok" only when both underlying checks pass, so it is the single +// field to read. +func statusChecksOK(s ec2types.InstanceStatus) bool { + return s.SystemStatus != nil && s.SystemStatus.Status == ec2types.SummaryStatusOk && + s.InstanceStatus != nil && s.InstanceStatus.Status == ec2types.SummaryStatusOk +} + +// AvailableInstanceTypes implements Client via DescribeInstanceTypeOfferings, +// scoped to this client's region (LocationType=region), paging through all +// results. The returned set is keyed by instance type; a type present here is +// offered in the region. +func (c *sdkClient) AvailableInstanceTypes(ctx context.Context) (map[string]bool, error) { + in := &ec2.DescribeInstanceTypeOfferingsInput{ + LocationType: ec2types.LocationTypeRegion, + Filters: []ec2types.Filter{{ + Name: awssdk.String("location"), + Values: []string{c.region}, + }}, + } + out := map[string]bool{} + for { + page, err := c.ec2.DescribeInstanceTypeOfferings(ctx, in) + if err != nil { + return nil, err + } + for _, o := range page.InstanceTypeOfferings { + out[string(o.InstanceType)] = true + } + if page.NextToken == nil { + break + } + in.NextToken = page.NextToken + } + return out, nil +} + +// observe normalizes a live EC2 SDK instance into the adapter-level EC2Instance +// view: id, tags, state name, region (from the placement AZ), a best-effort +// public endpoint, and whether it is Spot. It is a point-in-time read and never +// blocks. +func (c *sdkClient) observe(inst ec2types.Instance) EC2Instance { + out := EC2Instance{ + Region: c.region, + Tags: map[string]string{}, + } + if inst.InstanceId != nil { + out.ID = *inst.InstanceId + } + if inst.State != nil { + out.State = string(inst.State.Name) + } + for _, t := range inst.Tags { + if t.Key != nil && t.Value != nil { + out.Tags[*t.Key] = *t.Value + } + } + // Region is a coarser fact than the AZ; the AZ (e.g. "us-east-1a") is prefixed + // by the region, but we already know the region this client operates in, so we + // keep c.region rather than trimming the AZ. + if inst.PublicDnsName != nil && *inst.PublicDnsName != "" { + out.PublicEndpoint = *inst.PublicDnsName + } else if inst.PublicIpAddress != nil { + out.PublicEndpoint = *inst.PublicIpAddress + } + out.Spot = inst.InstanceLifecycle == ec2types.InstanceLifecycleTypeSpot + return out +} + +// ec2Tags converts the adapter's tag map into EC2's tag slice. +func ec2Tags(tags map[string]string) []ec2types.Tag { + out := make([]ec2types.Tag, 0, len(tags)) + for _, k := range sortedKeys(tags) { + v := tags[k] + out = append(out, ec2types.Tag{Key: awssdk.String(k), Value: awssdk.String(v)}) + } + return out +} + +// isNotFound reports whether err is EC2's "instance does not exist" condition, so +// Terminate/Describe can treat it as already gone (idempotency). +func isNotFound(err error) bool { + var apiErr smithy.APIError + if !errors.As(err, &apiErr) { + return false + } + switch apiErr.ErrorCode() { + case "InvalidInstanceID.NotFound", "InvalidInstanceID.Malformed": + return true + default: + return false + } +} diff --git a/pkg/provider/aws/client_test.go b/pkg/provider/aws/client_test.go new file mode 100644 index 0000000..3a27cb3 --- /dev/null +++ b/pkg/provider/aws/client_test.go @@ -0,0 +1,608 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package aws + +import ( + "context" + "encoding/base64" + "errors" + "strings" + "testing" + + awssdk "github.com/aws/aws-sdk-go-v2/aws" + "github.com/aws/aws-sdk-go-v2/service/ec2" + ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types" + smithy "github.com/aws/smithy-go" + + "github.com/InftyAI/Nebula/pkg/provider" +) + +// fakeEC2 is an in-memory ec2API for testing the sdkClient translation without a +// live AWS account. It records the last RunInstances input and lets tests seed +// describe results and inject errors. +type fakeEC2 struct { + lastRun *ec2.RunInstancesInput + runOut *ec2.RunInstancesOutput + runErr error + runCalls int + terminated []string + terminateErr error + describePages []*ec2.DescribeInstancesOutput + describeErr error + describeIdx int + + // statusPages seeds DescribeInstanceStatus (reachability checks); statusErr + // injects a failure. Left nil, the fake reports no status rows — every instance + // then reads as checks-not-passed, matching a freshly launched instance. + statusPages []*ec2.DescribeInstanceStatusOutput + statusErr error + statusIdx int + // lastStatusIn records the last DescribeInstanceStatus input so a test can + // assert which instance ids were queried. + lastStatusIn *ec2.DescribeInstanceStatusInput + + offeringPages []*ec2.DescribeInstanceTypeOfferingsOutput + offeringErr error + offeringIdx int + // lastOfferingIn records the last DescribeInstanceTypeOfferings input so a test + // can assert the region filter was applied. + lastOfferingIn *ec2.DescribeInstanceTypeOfferingsInput + + imagesOut *ec2.DescribeImagesOutput + imagesErr error + subnetsOut *ec2.DescribeSubnetsOutput + subnetsErr error +} + +func (f *fakeEC2) DescribeImages( + _ context.Context, _ *ec2.DescribeImagesInput, _ ...func(*ec2.Options), +) (*ec2.DescribeImagesOutput, error) { + if f.imagesErr != nil { + return nil, f.imagesErr + } + if f.imagesOut != nil { + return f.imagesOut, nil + } + return &ec2.DescribeImagesOutput{}, nil +} + +func (f *fakeEC2) DescribeSubnets( + _ context.Context, _ *ec2.DescribeSubnetsInput, _ ...func(*ec2.Options), +) (*ec2.DescribeSubnetsOutput, error) { + if f.subnetsErr != nil { + return nil, f.subnetsErr + } + if f.subnetsOut != nil { + return f.subnetsOut, nil + } + return &ec2.DescribeSubnetsOutput{}, nil +} + +func (f *fakeEC2) RunInstances( + _ context.Context, in *ec2.RunInstancesInput, _ ...func(*ec2.Options), +) (*ec2.RunInstancesOutput, error) { + f.lastRun = in + f.runCalls++ + if f.runErr != nil { + return nil, f.runErr + } + return f.runOut, nil +} + +func (f *fakeEC2) TerminateInstances( + _ context.Context, in *ec2.TerminateInstancesInput, _ ...func(*ec2.Options), +) (*ec2.TerminateInstancesOutput, error) { + if f.terminateErr != nil { + return nil, f.terminateErr + } + f.terminated = append(f.terminated, in.InstanceIds...) + return &ec2.TerminateInstancesOutput{}, nil +} + +func (f *fakeEC2) DescribeInstances( + _ context.Context, _ *ec2.DescribeInstancesInput, _ ...func(*ec2.Options), +) (*ec2.DescribeInstancesOutput, error) { + if f.describeErr != nil { + return nil, f.describeErr + } + if f.describeIdx >= len(f.describePages) { + return &ec2.DescribeInstancesOutput{}, nil + } + page := f.describePages[f.describeIdx] + f.describeIdx++ + return page, nil +} + +func (f *fakeEC2) DescribeInstanceStatus( + _ context.Context, in *ec2.DescribeInstanceStatusInput, _ ...func(*ec2.Options), +) (*ec2.DescribeInstanceStatusOutput, error) { + f.lastStatusIn = in + if f.statusErr != nil { + return nil, f.statusErr + } + if f.statusIdx >= len(f.statusPages) { + return &ec2.DescribeInstanceStatusOutput{}, nil + } + page := f.statusPages[f.statusIdx] + f.statusIdx++ + return page, nil +} + +func (f *fakeEC2) DescribeInstanceTypeOfferings( + _ context.Context, in *ec2.DescribeInstanceTypeOfferingsInput, _ ...func(*ec2.Options), +) (*ec2.DescribeInstanceTypeOfferingsOutput, error) { + f.lastOfferingIn = in + if f.offeringErr != nil { + return nil, f.offeringErr + } + if f.offeringIdx >= len(f.offeringPages) { + return &ec2.DescribeInstanceTypeOfferingsOutput{}, nil + } + page := f.offeringPages[f.offeringIdx] + f.offeringIdx++ + return page, nil +} + +func apiErr(code string) error { + return &smithy.GenericAPIError{Code: code, Message: code} +} + +func newSDKClient(f *fakeEC2) *sdkClient { + return &sdkClient{ec2: f, region: testRegion, amiID: "ami-123"} +} + +func TestBuildUserData(t *testing.T) { + spec := InstanceSpec{ + Image: "myrepo/train:v1", + Command: []string{"python", "train.py"}, + Env: map[string]string{"B": "2", "A": "1"}, + } + encoded, err := buildUserData(spec) + if err != nil { + t.Fatalf("buildUserData: %v", err) + } + raw, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + t.Fatalf("user-data is not valid base64: %v", err) + } + script := string(raw) + + if !strings.Contains(script, "docker pull 'myrepo/train:v1'") { + t.Fatalf("script missing image pull:\n%s", script) + } + if !strings.Contains(script, "--gpus all") { + t.Fatalf("script missing GPU attach:\n%s", script) + } + // Env is rendered in sorted key order for determinism (A before B). + if idxA, idxB := strings.Index(script, "'A=1'"), strings.Index(script, "'B=2'"); idxA < 0 || idxB < 0 || idxA > idxB { + t.Fatalf("env not rendered in sorted order:\n%s", script) + } + if !strings.Contains(script, "'python' 'train.py'") { + t.Fatalf("script missing command:\n%s", script) + } +} + +func TestBuildUserData_EmptyImageErrors(t *testing.T) { + if _, err := buildUserData(InstanceSpec{}); err == nil { + t.Fatal("expected an error for an empty image") + } +} + +func TestBuildUserData_QuotesHostileValues(t *testing.T) { + // A value with a single quote must not break out of its shell argument. + spec := InstanceSpec{ + Image: "img", + Env: map[string]string{"X": "a'; rm -rf /; echo '"}, + } + encoded, err := buildUserData(spec) + if err != nil { + t.Fatalf("buildUserData: %v", err) + } + raw, _ := base64.StdEncoding.DecodeString(encoded) + script := string(raw) + // The dangerous substring must appear only inside a quoted argument, never as + // a bare `rm -rf /` command line. + if strings.Contains(script, "\nrm -rf /") { + t.Fatalf("hostile env value escaped quoting:\n%s", script) + } +} + +func TestClassifyEC2Error(t *testing.T) { + tests := []struct { + name string + code string + spot bool + wantNoCap bool + wantSpot bool + wantQuota bool + wantAuth bool + wantZone bool + }{ + {"capacity ondemand is zone-local", "InsufficientInstanceCapacity", false, true, false, false, false, true}, + {"capacity spot is zone-local", "InsufficientInstanceCapacity", true, true, true, false, false, true}, + {"host capacity is zone-local", "InsufficientHostCapacity", false, true, false, false, false, true}, + {"type unsupported in AZ is zone-local", "Unsupported", false, true, false, false, false, true}, + // Region/account-scoped Spot limits: capacity + Spot, but NOT zone-local, so + // the adapter's per-zone sweep must stop rather than iterate every AZ. Spot + // applies even when the caller flagged the request OnDemand — these codes only + // arise on Spot requests. + {"spot price too low is region-scoped", "SpotMaxPriceTooLow", true, true, true, false, false, false}, + {"max spot count is region-scoped", "MaxSpotInstanceCountExceeded", true, true, true, false, false, false}, + {"instance limit is quota", "InstanceLimitExceeded", false, false, false, true, false, false}, + {"unauthorized is auth", "UnauthorizedOperation", false, false, false, false, true, false}, + {"unknown code passes through", "SomethingElse", false, false, false, false, false, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := classifyEC2Error(apiErr(tt.code), tt.spot) + if errors.Is(got, provider.ErrNoCapacity) != tt.wantNoCap { + t.Errorf("ErrNoCapacity: got %v, want %v (%v)", errors.Is(got, provider.ErrNoCapacity), tt.wantNoCap, got) + } + if errors.Is(got, ErrSpotCapacity) != tt.wantSpot { + t.Errorf("ErrSpotCapacity: got %v, want %v", errors.Is(got, ErrSpotCapacity), tt.wantSpot) + } + if errors.Is(got, provider.ErrQuota) != tt.wantQuota { + t.Errorf("ErrQuota: got %v, want %v", errors.Is(got, provider.ErrQuota), tt.wantQuota) + } + if errors.Is(got, provider.ErrAuth) != tt.wantAuth { + t.Errorf("ErrAuth: got %v, want %v", errors.Is(got, provider.ErrAuth), tt.wantAuth) + } + if errors.Is(got, errZoneLocal) != tt.wantZone { + t.Errorf("errZoneLocal: got %v, want %v", errors.Is(got, errZoneLocal), tt.wantZone) + } + }) + } + + if classifyEC2Error(nil, false) != nil { + t.Fatal("classifyEC2Error(nil) must be nil") + } + // A non-API error is returned unchanged so string heuristics can still apply. + plain := errors.New("dial tcp: timeout") + if got := classifyEC2Error(plain, false); got != plain { + t.Fatalf("plain error should pass through unchanged, got %v", got) + } +} + +func TestSDKRunInstance_OnDemand(t *testing.T) { + f := &fakeEC2{runOut: &ec2.RunInstancesOutput{ + Instances: []ec2types.Instance{{InstanceId: awssdk.String("i-42")}}, + }} + c := newSDKClient(f) + + id, err := c.RunInstance(context.Background(), InstanceSpec{ + InstanceType: "p5.48xlarge", + Image: "img", + Tags: map[string]string{ClaimTagKey: "claim-a"}, + }) + if err != nil { + t.Fatalf("RunInstance: %v", err) + } + if id != "i-42" { + t.Fatalf("id = %q, want i-42", id) + } + if f.lastRun.InstanceType != ec2types.InstanceType("p5.48xlarge") { + t.Fatalf("InstanceType = %q", f.lastRun.InstanceType) + } + if f.lastRun.ImageId == nil || *f.lastRun.ImageId != "ami-123" { + t.Fatalf("AMI not set from resolved GPU AMI: %+v", f.lastRun.ImageId) + } + // Security groups are deliberately unset: the default-VPC subnet supplies the + // default SG (outbound-open, no inbound exposure). + if len(f.lastRun.SecurityGroupIds) != 0 { + t.Fatalf("SecurityGroupIds must be unset, got %+v", f.lastRun.SecurityGroupIds) + } + // OnDemand: no spot market option. + if f.lastRun.InstanceMarketOptions != nil { + t.Fatalf("OnDemand request must not set market options") + } + // Claim tag rides on the instance tag spec. + if len(f.lastRun.TagSpecifications) == 0 || *f.lastRun.TagSpecifications[0].Tags[0].Value != "claim-a" { + t.Fatalf("claim tag not set: %+v", f.lastRun.TagSpecifications) + } +} + +func TestSDKRunInstance_SpotSetsMarketOption(t *testing.T) { + f := &fakeEC2{runOut: &ec2.RunInstancesOutput{ + Instances: []ec2types.Instance{{InstanceId: awssdk.String("i-spot")}}, + }} + c := newSDKClient(f) + + if _, err := c.RunInstance(context.Background(), InstanceSpec{ + InstanceType: "p5.48xlarge", Image: "img", Spot: true, + Tags: map[string]string{ClaimTagKey: "c"}, + }); err != nil { + t.Fatalf("RunInstance: %v", err) + } + if f.lastRun.InstanceMarketOptions == nil || + f.lastRun.InstanceMarketOptions.MarketType != ec2types.MarketTypeSpot { + t.Fatalf("Spot request must set spot market option: %+v", f.lastRun.InstanceMarketOptions) + } +} + +func TestSDKRunInstance_CapacityErrorWrapped(t *testing.T) { + f := &fakeEC2{runErr: apiErr("InsufficientInstanceCapacity")} + c := newSDKClient(f) + + _, err := c.RunInstance(context.Background(), InstanceSpec{ + InstanceType: "p5.48xlarge", Image: "img", Spot: true, + Tags: map[string]string{ClaimTagKey: "c"}, + }) + if !errors.Is(err, provider.ErrNoCapacity) || !errors.Is(err, ErrSpotCapacity) { + t.Fatalf("err = %v, want wrapped ErrNoCapacity + ErrSpotCapacity", err) + } +} + +// A zone-local capacity shortfall must make RunInstance sweep EVERY discovered +// subnet before giving up — a sibling AZ may still have capacity. +func TestSDKRunInstance_ZoneLocalSweepsAllSubnets(t *testing.T) { + f := &fakeEC2{runErr: apiErr("InsufficientInstanceCapacity")} + c := &sdkClient{ec2: f, region: testRegion, amiID: "ami-123", subnets: []subnet{ + {id: "subnet-a", az: "us-east-1a"}, + {id: "subnet-b", az: "us-east-1b"}, + {id: "subnet-c", az: "us-east-1c"}, + }} + + _, err := c.RunInstance(context.Background(), InstanceSpec{ + InstanceType: "p5.48xlarge", Image: "img", + Tags: map[string]string{ClaimTagKey: "c"}, + }) + if !errors.Is(err, provider.ErrNoCapacity) { + t.Fatalf("err = %v, want ErrNoCapacity after exhausting all zones", err) + } + if f.runCalls != 3 { + t.Fatalf("runCalls = %d, want 3 (one attempt per subnet)", f.runCalls) + } +} + +// A region/account-scoped Spot limit must STOP the sweep at the first subnet: +// every AZ would fail identically, so iterating them wastes the deadline. It still +// bubbles up ErrNoCapacity so region-level failover fires. +func TestSDKRunInstance_RegionScopedStopsSweep(t *testing.T) { + f := &fakeEC2{runErr: apiErr("MaxSpotInstanceCountExceeded")} + c := &sdkClient{ec2: f, region: testRegion, amiID: "ami-123", subnets: []subnet{ + {id: "subnet-a", az: "us-east-1a"}, + {id: "subnet-b", az: "us-east-1b"}, + {id: "subnet-c", az: "us-east-1c"}, + }} + + _, err := c.RunInstance(context.Background(), InstanceSpec{ + InstanceType: "p5.48xlarge", Image: "img", Spot: true, + Tags: map[string]string{ClaimTagKey: "c"}, + }) + if !errors.Is(err, provider.ErrNoCapacity) || !errors.Is(err, ErrSpotCapacity) { + t.Fatalf("err = %v, want ErrNoCapacity + ErrSpotCapacity", err) + } + if f.runCalls != 1 { + t.Fatalf("runCalls = %d, want 1 (region-scoped error must not sweep)", f.runCalls) + } +} + +func TestSDKTerminate_IdempotentOnNotFound(t *testing.T) { + f := &fakeEC2{terminateErr: apiErr("InvalidInstanceID.NotFound")} + c := newSDKClient(f) + if err := c.TerminateInstance(context.Background(), "i-gone"); err != nil { + t.Fatalf("Terminate should swallow NotFound, got %v", err) + } + + f2 := &fakeEC2{} + c2 := newSDKClient(f2) + if err := c2.TerminateInstance(context.Background(), "i-1"); err != nil { + t.Fatalf("Terminate: %v", err) + } + if len(f2.terminated) != 1 || f2.terminated[0] != "i-1" { + t.Fatalf("terminated = %v, want [i-1]", f2.terminated) + } +} + +func TestSDKDescribe_AbsentIsNil(t *testing.T) { + // NotFound => (nil, nil): absence == terminated. + f := &fakeEC2{describeErr: apiErr("InvalidInstanceID.NotFound")} + c := newSDKClient(f) + got, err := c.DescribeInstance(context.Background(), "i-gone") + if err != nil || got != nil { + t.Fatalf("DescribeInstance(absent) = (%+v, %v), want (nil, nil)", got, err) + } + + // Empty reservations also => nil. + f2 := &fakeEC2{describePages: []*ec2.DescribeInstancesOutput{{}}} + c2 := newSDKClient(f2) + got, err = c2.DescribeInstance(context.Background(), "i-x") + if err != nil || got != nil { + t.Fatalf("DescribeInstance(empty) = (%+v, %v), want (nil, nil)", got, err) + } +} + +func TestSDKDescribe_ObservesInstance(t *testing.T) { + f := &fakeEC2{describePages: []*ec2.DescribeInstancesOutput{{ + Reservations: []ec2types.Reservation{{Instances: []ec2types.Instance{{ + InstanceId: awssdk.String("i-1"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNameRunning}, + PublicDnsName: awssdk.String("ec2-1-2-3-4.compute.amazonaws.com"), + InstanceLifecycle: ec2types.InstanceLifecycleTypeSpot, + Tags: []ec2types.Tag{{Key: awssdk.String(ClaimTagKey), Value: awssdk.String("claim-a")}}, + }}}}, + }}} + c := newSDKClient(f) + + got, err := c.DescribeInstance(context.Background(), "i-1") + if err != nil || got == nil { + t.Fatalf("DescribeInstance = (%+v, %v)", got, err) + } + if got.ID != "i-1" || got.State != string(ec2types.InstanceStateNameRunning) { + t.Fatalf("observed = %+v", got) + } + if got.PublicEndpoint != "ec2-1-2-3-4.compute.amazonaws.com" { + t.Fatalf("endpoint = %q", got.PublicEndpoint) + } + if !got.Spot { + t.Fatalf("Spot lifecycle not detected") + } + if got.Tags[ClaimTagKey] != "claim-a" { + t.Fatalf("claim tag = %q", got.Tags[ClaimTagKey]) + } + if got.Region != testRegion { + t.Fatalf("region = %q, want %q", got.Region, testRegion) + } +} + +func TestSDKAvailableInstanceTypes_PagesAndScopesToRegion(t *testing.T) { + f := &fakeEC2{offeringPages: []*ec2.DescribeInstanceTypeOfferingsOutput{ + { + InstanceTypeOfferings: []ec2types.InstanceTypeOffering{ + {InstanceType: ec2types.InstanceType("g4dn.xlarge")}, + }, + NextToken: awssdk.String("page2"), + }, + { + InstanceTypeOfferings: []ec2types.InstanceTypeOffering{ + {InstanceType: ec2types.InstanceType("p5.48xlarge")}, + }, + }, + }} + c := newSDKClient(f) + + got, err := c.AvailableInstanceTypes(context.Background()) + if err != nil { + t.Fatalf("AvailableInstanceTypes: %v", err) + } + if len(got) != 2 || !got["g4dn.xlarge"] || !got["p5.48xlarge"] { + t.Fatalf("available = %+v, want both paged types", got) + } + // The probe must be scoped to the client's region (LocationType=region + + // location filter) so it reports THIS region's availability, not global. + if f.lastOfferingIn.LocationType != ec2types.LocationTypeRegion { + t.Fatalf("LocationType = %q, want region", f.lastOfferingIn.LocationType) + } + if len(f.lastOfferingIn.Filters) == 0 || + len(f.lastOfferingIn.Filters[0].Values) == 0 || + f.lastOfferingIn.Filters[0].Values[0] != testRegion { + t.Fatalf("location filter = %+v, want %q", f.lastOfferingIn.Filters, testRegion) + } +} + +func TestSDKAvailableInstanceTypes_ErrorPropagates(t *testing.T) { + f := &fakeEC2{offeringErr: errors.New("boom")} + c := newSDKClient(f) + if _, err := c.AvailableInstanceTypes(context.Background()); err == nil { + t.Fatal("expected the probe error to propagate") + } +} + +func TestSDKList_PagesAndFiltersByClaimTag(t *testing.T) { + f := &fakeEC2{describePages: []*ec2.DescribeInstancesOutput{ + { + Reservations: []ec2types.Reservation{{Instances: []ec2types.Instance{{ + InstanceId: awssdk.String("i-1"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNameRunning}, + }}}}, + NextToken: awssdk.String("page2"), + }, + { + Reservations: []ec2types.Reservation{{Instances: []ec2types.Instance{{ + InstanceId: awssdk.String("i-2"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNamePending}, + }}}}, + }, + }} + c := newSDKClient(f) + + list, err := c.ListInstances(context.Background()) + if err != nil { + t.Fatalf("ListInstances: %v", err) + } + if len(list) != 2 || list[0].ID != "i-1" || list[1].ID != "i-2" { + t.Fatalf("list = %+v, want both pages", list) + } +} + +func TestSDKList_FoldsInStatusChecks(t *testing.T) { + // Two running instances; only i-1 has passed both reachability checks. List must + // report i-1 as checks-passed and i-2 as not, so the poll loop holds i-2 at + // Pending until its checks clear. + f := &fakeEC2{ + describePages: []*ec2.DescribeInstancesOutput{{ + Reservations: []ec2types.Reservation{{Instances: []ec2types.Instance{ + { + InstanceId: awssdk.String("i-1"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNameRunning}, + }, + { + InstanceId: awssdk.String("i-2"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNameRunning}, + }, + }}}}, + }, + statusPages: []*ec2.DescribeInstanceStatusOutput{{ + InstanceStatuses: []ec2types.InstanceStatus{ + { + InstanceId: awssdk.String("i-1"), + SystemStatus: &ec2types.InstanceStatusSummary{Status: ec2types.SummaryStatusOk}, + InstanceStatus: &ec2types.InstanceStatusSummary{Status: ec2types.SummaryStatusOk}, + }, + { + InstanceId: awssdk.String("i-2"), + SystemStatus: &ec2types.InstanceStatusSummary{Status: ec2types.SummaryStatusOk}, + InstanceStatus: &ec2types.InstanceStatusSummary{Status: ec2types.SummaryStatusInitializing}, + }, + }, + }}, + } + c := newSDKClient(f) + + list, err := c.ListInstances(context.Background()) + if err != nil { + t.Fatalf("ListInstances: %v", err) + } + byID := map[string]EC2Instance{} + for _, inst := range list { + byID[inst.ID] = inst + } + if !byID["i-1"].StatusChecksPassed { + t.Error("i-1 has 2/2 checks ok; want StatusChecksPassed=true") + } + if byID["i-2"].StatusChecksPassed { + t.Error("i-2 instance check is still initializing; want StatusChecksPassed=false") + } + // The status probe must be scoped to exactly the listed ids. + if got := f.lastStatusIn.InstanceIds; len(got) != 2 { + t.Fatalf("status probe ids = %v, want the two listed instances", got) + } +} + +func TestSDKList_StatusProbeFailureIsNonFatal(t *testing.T) { + // A DescribeInstanceStatus failure must NOT fail the whole List: the instances + // are still returned (checks default to not-passed), so a transient status-API + // hiccup costs at most one extra Pending tick rather than stranding the loop. + f := &fakeEC2{ + describePages: []*ec2.DescribeInstancesOutput{{ + Reservations: []ec2types.Reservation{{Instances: []ec2types.Instance{{ + InstanceId: awssdk.String("i-1"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNameRunning}, + }}}}, + }}, + statusErr: errors.New("throttled"), + } + c := newSDKClient(f) + + list, err := c.ListInstances(context.Background()) + if err != nil { + t.Fatalf("ListInstances must not propagate a status-probe error: %v", err) + } + if len(list) != 1 || list[0].StatusChecksPassed { + t.Fatalf("list = %+v, want the instance returned with checks not passed", list) + } +} diff --git a/pkg/provider/aws/translate.go b/pkg/provider/aws/translate.go new file mode 100644 index 0000000..0e1d7b7 --- /dev/null +++ b/pkg/provider/aws/translate.go @@ -0,0 +1,162 @@ +/* +Copyright 2026 The InftyAI Team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package aws + +import ( + "encoding/base64" + "errors" + "fmt" + "sort" + "strings" + + smithy "github.com/aws/smithy-go" + + "github.com/InftyAI/Nebula/pkg/provider" +) + +// These are the pure (AWS-SDK-free) translations the SDK client relies on: +// building the cloud-init user-data that starts the container, and mapping raw +// EC2 API errors onto the shared provider sentinels. They live apart from +// client.go so they can be unit-tested without any AWS calls. + +// buildUserData renders the base64 cloud-init user-data that boots the workload +// container on a GPU AMI. The chosen execution model (see the package doc) is a +// prebuilt GPU AMI — NVIDIA driver + a container runtime already installed — plus +// this user-data that pulls spec.Image and runs it with the Pod's command/env and +// all GPUs attached. The script is intentionally minimal and idempotent-ish: it +// is the seam where a richer bootstrap (systemd unit, log shipping, health +// reporting) would later grow. +// +// spec.Image is required; a spec with no image is a programmer error (the adapter +// reads it off the Pod's first container) and yields an error rather than a +// container-less instance. +func buildUserData(spec InstanceSpec) (string, error) { + if spec.Image == "" { + return "", errors.New("aws: empty image in instance spec") + } + + var b strings.Builder + b.WriteString("#!/bin/bash\n") + b.WriteString("set -euo pipefail\n") + // Pull first so an image error surfaces before we try to run. + fmt.Fprintf(&b, "docker pull %s\n", shellQuote(spec.Image)) + + // docker run --gpus all -d, with env and command appended. Env keys are sorted + // so the rendered script is deterministic (stable across reconciles and easy to + // assert in tests). + b.WriteString("docker run --rm --gpus all") + for _, k := range sortedKeys(spec.Env) { + fmt.Fprintf(&b, " -e %s", shellQuote(k+"="+spec.Env[k])) + } + b.WriteString(" " + shellQuote(spec.Image)) + for _, arg := range spec.Command { + b.WriteString(" " + shellQuote(arg)) + } + b.WriteString("\n") + + return base64.StdEncoding.EncodeToString([]byte(b.String())), nil +} + +// sortedKeys returns m's keys in sorted order, for deterministic rendering. +func sortedKeys(m map[string]string) []string { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +// shellQuote wraps s in single quotes, escaping any embedded single quote, so an +// image/env/command value with spaces or shell metacharacters cannot break out of +// its argument or inject commands into the bootstrap script. +func shellQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" +} + +// errZoneLocal is an INTERNAL marker wrapped onto a capacity failure that is +// specific to the availability zone of the attempted subnet, not the whole region +// — a sibling AZ's default subnet may still satisfy the launch. RunInstance reads +// it to decide whether to keep sweeping subnets (zone-local) or stop immediately +// and bubble a region-scoped block up to control-plane failover. It always rides +// ALONGSIDE provider.ErrNoCapacity, so the control plane's ClassifyProvisionError +// derives the same BlockScope regardless of scope; only the adapter's inner sweep +// distinguishes the two. +var errZoneLocal = errors.New("aws: zone-local capacity shortfall") + +// classifyEC2Error maps a raw EC2 API error onto the shared provider sentinels so +// the adapter's ClassifyProvisionError can derive a BlockScope uniformly. It +// recognizes EC2's error codes (smithy.APIError) first, falling back to the +// original error otherwise. spot indicates whether the failing request was Spot, +// so a capacity failure can additionally carry ErrSpotCapacity and confine the +// block to the Spot tier. +// +// Capacity failures split by SCOPE, which drives the adapter's per-zone sweep: +// - ZONE-LOCAL (InsufficientInstanceCapacity/InsufficientHostCapacity, and +// "Unsupported" — how EC2 reports the instance type is not offered in the AZ of +// the chosen subnet): a sibling AZ may still satisfy the launch, so these carry +// errZoneLocal and RunInstance advances to the next subnet. If every AZ fails +// it collapses to a region no-capacity that hands off to region-level failover. +// - REGION/ACCOUNT-SCOPED (SpotMaxPriceTooLow, MaxSpotInstanceCountExceeded): a +// Spot price ceiling or per-region Spot limit that applies across the whole +// region, not one AZ. Sweeping sibling AZs is futile, so these OMIT errZoneLocal +// — RunInstance stops immediately and lets region/tier failover take over. They +// still carry ErrSpotCapacity (they are Spot-only), confining the block to Spot. +// +// Returning the wrapped sentinel (not a BlockScope) keeps the +// error-category → scope rule in one place (provider.ClassifyError, via the +// adapter), so this function only has to recognize AWS's vocabulary. +func classifyEC2Error(err error, spot bool) error { + if err == nil { + return nil + } + var apiErr smithy.APIError + if !errors.As(err, &apiErr) { + return err // not an API error; let string heuristics in ClassifyError handle it + } + + switch apiErr.ErrorCode() { + case "InsufficientInstanceCapacity", "InsufficientHostCapacity", "Unsupported": + return wrapNoCapacity(err, spot, true /* zoneLocal */) + case "SpotMaxPriceTooLow", "MaxSpotInstanceCountExceeded": + // Region/account-scoped Spot limits: NOT zone-local, so the sweep stops. + // These only arise on Spot requests, so ErrSpotCapacity always applies. + return wrapNoCapacity(err, true /* spot */, false /* zoneLocal */) + case "InstanceLimitExceeded", "VcpuLimitExceeded", "RequestLimitExceeded": + return fmt.Errorf("%w: %w", err, provider.ErrQuota) + case "UnauthorizedOperation", "AuthFailure", "Blocked", + "OptInRequired", "PendingVerification": + return fmt.Errorf("%w: %w", err, provider.ErrAuth) + default: + return err + } +} + +// wrapNoCapacity wraps err with provider.ErrNoCapacity plus, conditionally, the +// Spot-tier marker (ErrSpotCapacity) and the zone-local sweep marker +// (errZoneLocal). errors.Join lets errors.Is find every marker regardless of how +// many apply. +func wrapNoCapacity(err error, spot, zoneLocal bool) error { + errs := []error{err, provider.ErrNoCapacity} + if spot { + errs = append(errs, ErrSpotCapacity) + } + if zoneLocal { + errs = append(errs, errZoneLocal) + } + return errors.Join(errs...) +} diff --git a/pkg/provider/catalog/base.go b/pkg/provider/catalog/base.go index 046c049..106c5dc 100644 --- a/pkg/provider/catalog/base.go +++ b/pkg/provider/catalog/base.go @@ -76,14 +76,21 @@ func (b Base) Offerings(context.Context) ([]provider.Offering, error) { return b.Catalog.Offerings(b.ProviderName), nil } -// MapAccelerator is the default identity translation: it confirms the canonical -// accelerator is in this provider's catalog (case-insensitively) and returns the -// canonical name unchanged. Providers whose identifiers differ from the -// canonical names override this method. Returns ok=false when the provider does -// not offer the accelerator. +// MapAccelerator translates a canonical accelerator type into this provider's +// own accelerator id using the catalog as the mapping table: it finds the row +// whose AcceleratorType matches (case-insensitively) and returns that row's +// AcceleratorID. When AcceleratorID is blank (a provider whose id equals the +// canonical name) it falls back to the canonical name, so an identity-mapped +// provider needs no per-name data. Because the mapping lives entirely in the +// CSV, a provider whose ids diverge (AWS instance types) does NOT need to +// override this — it just populates the accelerator_id column. Returns ok=false +// when the provider does not offer the accelerator. func (b Base) MapAccelerator(canonical string) (providerAcceleratorID string, ok bool) { for _, o := range b.Catalog.Offerings(b.ProviderName) { if strings.EqualFold(o.AcceleratorType, canonical) { + if o.AcceleratorID != "" { + return o.AcceleratorID, true + } return o.AcceleratorType, true } } diff --git a/pkg/provider/catalog/catalog.go b/pkg/provider/catalog/catalog.go index 40529e8..e39afe0 100644 --- a/pkg/provider/catalog/catalog.go +++ b/pkg/provider/catalog/catalog.go @@ -119,48 +119,98 @@ func (c *Catalog) Offerings(providerName string) []provider.Offering { return out } +// Required and optional CSV column names. Parsing is header-driven (by name, not +// position) so a provider can add the optional columns — or reorder — without a +// parser change, and NeoCloud CSVs stay in the original minimal form. +const ( + colAccelerator = "accelerator_type" // required: canonical Nebula accelerator type + colAcceleratorID = "accelerator_id" // optional: provider's own accelerator id (e.g. AWS instance type) + colCapacityType = "capacity_type" // required: Spot | OnDemand | Reserved + colPrice = "price_per_hour" // required: approximate USD/GPU-hour + colAvailable = "available" // required: whether the provider offers it + colGPUCount = "gpu_count" // optional: accelerators the id provides (AWS: baked into instance type) + colRegion = "region" // optional: provider region this row prices + colUpdated = "updated" // optional: documentation only, ignored +) + // parseCSV reads offering rows. Lines beginning with '#' are comments; the first -// non-comment line is the header. Column order is fixed: -// accelerator,capacity_type,price_per_hour,available,updated. The trailing -// updated column is documentation only and ignored here. -// -// There is deliberately no provider_id column yet: every provider we currently -// support (Modal) names its accelerators identically to Nebula's canonical -// names, so MapAccelerator is pure identity. When a provider whose identifiers -// diverge lands (e.g. RunPod's "NVIDIA H100 80GB HBM3"), add a provider_id -// column here and have MapAccelerator return it — no per-provider Go table. +// non-comment line is the header, and columns are resolved BY NAME from it, so +// order does not matter and optional columns may be absent. Required columns: +// accelerator_type, capacity_type, price_per_hour, available. Optional columns: +// accelerator_id (a provider's own accelerator id, e.g. AWS "p5.48xlarge", for a +// provider whose MapAccelerator override translates canonical names — for Modal, +// whose mapping is identity, it is simply left blank), gpu_count (how many +// accelerators the id provides, a lookup dimension only where the count is baked +// into the offering — AWS instance types — and blank/0 for providers that take +// the count as a runtime parameter, like Modal), region (the provider region a +// row prices, for region-aware providers), and updated (documentation only). func parseCSV(r io.Reader) ([]provider.Offering, error) { cr := csv.NewReader(r) cr.Comment = '#' - cr.FieldsPerRecord = -1 // tolerate a trailing "updated" column we ignore + cr.FieldsPerRecord = -1 // rows may carry a different column count than legacy fixed-order cr.TrimLeadingSpace = true records, err := cr.ReadAll() if err != nil { return nil, err } + if len(records) == 0 { + return nil, nil + } - offerings := make([]provider.Offering, 0, len(records)) + // Map header names to column indices. + idx := make(map[string]int, len(records[0])) + for i, h := range records[0] { + idx[strings.TrimSpace(h)] = i + } + for _, req := range []string{colAccelerator, colCapacityType, colPrice, colAvailable} { + if _, ok := idx[req]; !ok { + return nil, fmt.Errorf("missing required column %q in header", req) + } + } + + // field returns the named column's value for a row, or "" if the column is + // absent (optional) or the row is short. + field := func(rec []string, name string) string { + i, ok := idx[name] + if !ok || i >= len(rec) { + return "" + } + return strings.TrimSpace(rec[i]) + } + + offerings := make([]provider.Offering, 0, len(records)-1) for i, rec := range records { if i == 0 { continue // header } - if len(rec) < 4 { - return nil, fmt.Errorf("row %d: expected >=4 columns, got %d", i, len(rec)) - } - price, err := strconv.ParseFloat(strings.TrimSpace(rec[2]), 64) + price, err := strconv.ParseFloat(field(rec, colPrice), 64) if err != nil { - return nil, fmt.Errorf("row %d: price_per_hour %q: %w", i, rec[2], err) + return nil, fmt.Errorf("row %d: price_per_hour %q: %w", i, field(rec, colPrice), err) } - available, err := strconv.ParseBool(strings.TrimSpace(rec[3])) + available, err := strconv.ParseBool(field(rec, colAvailable)) if err != nil { - return nil, fmt.Errorf("row %d: available %q: %w", i, rec[3], err) + return nil, fmt.Errorf("row %d: available %q: %w", i, field(rec, colAvailable), err) + } + // gpu_count is optional; blank => 0 (not a lookup dimension for this + // provider, e.g. Modal, which takes the count as a runtime parameter). A + // present-but-unparseable value is an error, not a silent 0. + var gpuCount int32 + if v := field(rec, colGPUCount); v != "" { + n, err := strconv.ParseInt(v, 10, 32) + if err != nil { + return nil, fmt.Errorf("row %d: gpu_count %q: %w", i, v, err) + } + gpuCount = int32(n) } offerings = append(offerings, provider.Offering{ - AcceleratorType: strings.TrimSpace(rec[0]), - CapacityType: nebulav1alpha1.CapacityType(strings.TrimSpace(rec[1])), + AcceleratorType: field(rec, colAccelerator), + CapacityType: nebulav1alpha1.CapacityType(field(rec, colCapacityType)), PricePerHour: price, Available: available, + Region: field(rec, colRegion), + AcceleratorID: field(rec, colAcceleratorID), + GPUCount: gpuCount, }) } return offerings, nil diff --git a/pkg/provider/catalog/catalog_test.go b/pkg/provider/catalog/catalog_test.go index 6234435..44bbcaa 100644 --- a/pkg/provider/catalog/catalog_test.go +++ b/pkg/provider/catalog/catalog_test.go @@ -59,9 +59,93 @@ func TestLoad_EmbeddedModal(t *testing.T) { } } +func TestLoad_EmbeddedAWS(t *testing.T) { + c, err := Load() + if err != nil { + t.Fatalf("Load: %v", err) + } + offs := c.Offerings("aws") + if len(offs) == 0 { + t.Fatal("expected embedded aws offerings") + } + + // AWS requests by instance type, so every row must carry an accelerator_id (the + // EC2 instance type). The embedded catalog is REGION-LESS — region is stamped by + // the adapter's live probe, not hand-maintained here — so rows have no region. + // Both capacity tiers appear. + var haveSpot, haveOnDemand bool + for _, o := range offs { + if o.Region != "" { + t.Fatalf("embedded aws catalog must be region-less, got region %q for %q", o.Region, o.AcceleratorType) + } + if o.AcceleratorID == "" { + t.Fatalf("aws offering %q has no accelerator_id (EC2 instance type)", o.AcceleratorType) + } + switch o.CapacityType { + case nebulav1alpha1.CapacitySpot: + haveSpot = true + case nebulav1alpha1.CapacityOnDemand: + haveOnDemand = true + } + } + if !haveSpot || !haveOnDemand { + t.Fatalf("expected both Spot and OnDemand aws rows, got spot=%v ondemand=%v", haveSpot, haveOnDemand) + } + + // gpu_count is a lookup dimension for AWS: the same accelerator must appear at + // more than one count (e.g. T4 x1 and T4 x8 are different instance types). + counts := map[int32]string{} + for _, o := range offs { + if o.AcceleratorType == "T4" && o.CapacityType == nebulav1alpha1.CapacityOnDemand { + counts[o.GPUCount] = o.AcceleratorID + } + } + if counts[1] == counts[8] || counts[1] == "" || counts[8] == "" { + t.Fatalf("T4 must map to distinct instance types per count, got %+v", counts) + } +} + +func TestLoadFrom_OptionalColumnsAndComments(t *testing.T) { + dir := t.TempDir() + // Header-driven parsing: columns are located by name, so a file may include + // accelerator_id/region and interleave comment/blank lines. A blank + // accelerator_id is allowed (identity mapping) and a blank region is allowed + // (region-simple provider). + csv := "# a comment line\n" + + "accelerator_type,accelerator_id,capacity_type,price_per_hour,available,region,updated\n" + + "\n" + + "H100,p5.48xlarge,Spot,34.41,true,us-east-1,2026-07-27\n" + + "# another comment\n" + + "T4,,OnDemand,0.50,true,,2026-07-27\n" + if err := os.WriteFile(filepath.Join(dir, "modal.csv"), []byte(csv), 0o644); err != nil { + t.Fatal(err) + } + + c, err := LoadFrom(dir) + if err != nil { + t.Fatalf("LoadFrom: %v", err) + } + offs := c.Offerings("modal") + if len(offs) != 2 { + t.Fatalf("expected 2 rows (comments/blanks skipped), got %d: %+v", len(offs), offs) + } + + byType := map[string]provider.Offering{} + for _, o := range offs { + byType[o.AcceleratorType] = o + } + h := byType["H100"] + if h.AcceleratorID != "p5.48xlarge" || h.Region != "us-east-1" || h.CapacityType != nebulav1alpha1.CapacitySpot { + t.Fatalf("H100 row parsed wrong: %+v", h) + } + if t4 := byType["T4"]; t4.AcceleratorID != "" || t4.Region != "" { + t.Fatalf("T4 row: blank accelerator_id/region must stay blank, got %+v", t4) + } +} + func TestLoadFrom_OverrideDir(t *testing.T) { dir := t.TempDir() - csv := "accelerator,capacity_type,price_per_hour,available,updated\n" + + csv := "accelerator_type,capacity_type,price_per_hour,available,updated\n" + "H100,OnDemand,9.99,true,2026-07-25\n" if err := os.WriteFile(filepath.Join(dir, "modal.csv"), []byte(csv), 0o644); err != nil { t.Fatal(err) @@ -79,7 +163,7 @@ func TestLoadFrom_OverrideDir(t *testing.T) { func TestLoad_OverrideEnvWins(t *testing.T) { dir := t.TempDir() - csv := "accelerator,capacity_type,price_per_hour,available,updated\n" + + csv := "accelerator_type,capacity_type,price_per_hour,available,updated\n" + "T4,OnDemand,0.11,true,2026-07-25\n" if err := os.WriteFile(filepath.Join(dir, "modal.csv"), []byte(csv), 0o644); err != nil { t.Fatal(err) diff --git a/pkg/provider/catalog/data/aws.csv b/pkg/provider/catalog/data/aws.csv new file mode 100644 index 0000000..458e5be --- /dev/null +++ b/pkg/provider/catalog/data/aws.csv @@ -0,0 +1,56 @@ +# AWS EC2 GPU price/availability catalog — community-maintained. +# +# Unlike Modal, AWS does not name GPUs by accelerator type, and the GPU count is +# NOT a free runtime knob: you request an INSTANCE TYPE whose accelerator count is +# fixed (T4x1 = g4dn.xlarge, T4x8 = g4dn.metal). So the mapping key is +# (accelerator_type, gpu_count): accelerator_id holds the EC2 instance type that +# serves that pair, and the AWS adapter selects it by matching BOTH the canonical +# accelerator and the Pod's requested GPU count. +# +# The `region` column is left BLANK here on purpose. This file holds only the +# durable facts — the (accelerator, count) -> instance-type mapping and a seed +# on-demand price — one row per {accelerator, count, capacity_type}. Region truth +# (which types a region actually offers, and live spot prices) is NOT +# hand-maintained: AWS exposes it at runtime, so the adapter's Offerings(ctx) +# probes DescribeInstanceTypeOfferings per configured region and stamps/filters +# these rows. That keeps this file small and stops per-region availability from +# drifting out of date. +# +# On-demand prices are approximate USD per INSTANCE-hour transcribed by hand from +# https://aws.amazon.com/ec2/pricing/on-demand/ (us-east-1 rates as the seed); +# treat them as a starting point, not a billing source. Spot prices vary by +# region and minute and are left to the live probe (the seed spot rows carry a +# recent us-east-1 typical rate only as a fallback). +# +# Columns (shared header across all provider CSVs; unused cells left blank): +# accelerator_type canonical Nebula accelerator type +# accelerator_id EC2 instance type that serves {accelerator_type, gpu_count} +# gpu_count accelerators the instance type provides (the lookup key) +# capacity_type Spot | OnDemand | Reserved +# price_per_hour approximate USD/instance-hour (seed; live probe may override) +# available seed availability (the live probe is authoritative per region) +# region BLANK — stamped per configured region by the live probe +# updated YYYY-MM-DD the row was last verified +accelerator_type,accelerator_id,gpu_count,capacity_type,price_per_hour,available,region,updated +T4,g4dn.xlarge,1,OnDemand,0.526,true,,2026-07-27 +T4,g4dn.xlarge,1,Spot,0.158,true,,2026-07-27 +T4,g4dn.12xlarge,4,OnDemand,3.912,true,,2026-07-27 +T4,g4dn.12xlarge,4,Spot,1.174,true,,2026-07-27 +T4,g4dn.metal,8,OnDemand,7.824,true,,2026-07-27 +T4,g4dn.metal,8,Spot,2.347,true,,2026-07-27 +A10G,g5.xlarge,1,OnDemand,1.006,true,,2026-07-27 +A10G,g5.xlarge,1,Spot,0.352,true,,2026-07-27 +A10G,g5.12xlarge,4,OnDemand,5.672,true,,2026-07-27 +A10G,g5.12xlarge,4,Spot,1.985,true,,2026-07-27 +A10G,g5.48xlarge,8,OnDemand,16.288,true,,2026-07-27 +A10G,g5.48xlarge,8,Spot,5.700,true,,2026-07-27 +L4,g6.xlarge,1,OnDemand,0.805,true,,2026-07-27 +L4,g6.xlarge,1,Spot,0.282,true,,2026-07-27 +L4,g6.12xlarge,4,OnDemand,4.602,true,,2026-07-27 +L4,g6.12xlarge,4,Spot,1.610,true,,2026-07-27 +L4,g6.48xlarge,8,OnDemand,13.350,true,,2026-07-27 +L4,g6.48xlarge,8,Spot,4.672,true,,2026-07-27 +A100-80GB,p4de.24xlarge,8,OnDemand,40.966,true,,2026-07-27 +A100-80GB,p4de.24xlarge,8,Spot,14.338,true,,2026-07-27 +H100,p5.48xlarge,8,OnDemand,98.320,true,,2026-07-27 +H100,p5.48xlarge,8,Spot,34.412,true,,2026-07-27 diff --git a/pkg/provider/catalog/data/modal.csv b/pkg/provider/catalog/data/modal.csv index 2596277..caaae26 100644 --- a/pkg/provider/catalog/data/modal.csv +++ b/pkg/provider/catalog/data/modal.csv @@ -10,21 +10,29 @@ # accelerator column defines both the supported set + price (for the optimizer) # and, since Modal names its GPUs identically to Nebula's canonical names, the # accelerator vocabulary MapAccelerator resolves against. No accelerator table -# lives in Go anymore. (If a future provider's identifiers diverge from the -# canonical names, add a provider_id column here rather than a Go map.) +# lives in Go anymore. # -# Columns: -# accelerator canonical Nebula accelerator type -# capacity_type Spot | OnDemand | Reserved -# price_per_hour approximate on-demand USD/GPU-hour -# available whether Modal currently offers it (a live probe may override) -# updated YYYY-MM-DD the row was last verified -accelerator,capacity_type,price_per_hour,available,updated -T4,OnDemand,0.59,true,2026-07-25 -L4,OnDemand,0.80,true,2026-07-25 -A10G,OnDemand,1.10,true,2026-07-25 -L40S,OnDemand,1.95,true,2026-07-25 -A100-40GB,OnDemand,2.10,true,2026-07-25 -A100-80GB,OnDemand,2.50,true,2026-07-25 -H100,OnDemand,3.95,true,2026-07-25 -H200,OnDemand,4.55,true,2026-07-25 +# Columns (shared header across all provider CSVs; unused cells left blank): +# accelerator_type canonical Nebula accelerator type +# accelerator_id BLANK — Modal names its GPUs exactly like the canonical types, +# so the mapping is identity and MapAccelerator falls back to +# accelerator_type. A provider whose ids diverge (AWS instance +# types) fills it in. +# gpu_count BLANK — Modal takes the accelerator count as a runtime +# parameter (off the Pod's nvidia.com/gpu resource), so it is +# not a catalog lookup dimension. AWS, where the count is baked +# into the instance type, fills it in. +# capacity_type Spot | OnDemand | Reserved +# price_per_hour approximate on-demand USD/GPU-hour +# available whether Modal currently offers it (a live probe may override) +# region BLANK — Modal is region-simple +# updated YYYY-MM-DD the row was last verified +accelerator_type,accelerator_id,gpu_count,capacity_type,price_per_hour,available,region,updated +T4,,,OnDemand,0.59,true,,2026-07-25 +L4,,,OnDemand,0.80,true,,2026-07-25 +A10G,,,OnDemand,1.10,true,,2026-07-25 +L40S,,,OnDemand,1.95,true,,2026-07-25 +A100-40GB,,,OnDemand,2.10,true,,2026-07-25 +A100-80GB,,,OnDemand,2.50,true,,2026-07-25 +H100,,,OnDemand,3.95,true,,2026-07-25 +H200,,,OnDemand,4.55,true,,2026-07-25 diff --git a/pkg/provider/errors.go b/pkg/provider/errors.go index 0f5ee13..dfd6452 100644 --- a/pkg/provider/errors.go +++ b/pkg/provider/errors.go @@ -52,20 +52,41 @@ var ( // provider; an unrecognized error is treated conservatively as whole-provider so // a failure does not hot-loop (the blocklist TTL bounds the blast radius). // +// It is the single place the SHARED part of a scope is derived — category +// (DenyAll vs capacity-scoped), the capacity tier, and the accelerator — so every +// adapter's ClassifyProvisionError delegates here and then only decorates with +// what is provider-specific (AWS adds its region). No scope is assembled anywhere +// else: the vnode handler resolves the accelerator off the Pod and passes it in, +// rather than mutating the returned scope. +// // capacityType is the tier the failing request used; it is stamped onto // accelerator-scoped scopes so the block is precise (a Spot failure does not -// block OnDemand). Adapters call this from ClassifyProvisionError after wrapping -// their own errors, passing the only tier they serve when they are single-tier. -func ClassifyError(err error, capacityType nebulav1alpha1.CapacityType) BlockScope { +// block OnDemand). accelerator is the canonical accelerator the request asked for +// ("" for a CPU-only Pod): on a capacity-scoped block a non-empty accelerator +// becomes an exact-match pointer (narrowing the block to the type that failed), +// while "" leaves AcceleratorType nil ("not applicable") so the block does not +// widen across every accelerator. Auth/quota (DenyAll) ignores both, since bad +// credentials fail for every accelerator and tier. +func ClassifyError(err error, capacityType nebulav1alpha1.CapacityType, accelerator string) BlockScope { if err == nil { return BlockScope{} } + // capacityScope builds an accelerator/tier-scoped block, promoting a non-empty + // accelerator to an exact-match pointer and leaving it nil otherwise. + capacityScope := func() BlockScope { + s := BlockScope{CapacityType: capacityType} + if accelerator != "" { + s.AcceleratorType = &accelerator + } + return s + } + switch { case errors.Is(err, ErrAuth), errors.Is(err, ErrQuota): return BlockScope{DenyAll: true} case errors.Is(err, ErrNoCapacity), errors.Is(err, ErrUnsupportedAccelerator): - return BlockScope{CapacityType: capacityType} + return capacityScope() } // Fall back to string heuristics for errors not wrapped with a sentinel. @@ -76,7 +97,7 @@ func ClassifyError(err error, capacityType nebulav1alpha1.CapacityType) BlockSco case containsAny(msg, "quota", "limit exceeded", "rate limit"): return BlockScope{DenyAll: true} case containsAny(msg, "no capacity", "capacity", "unavailable", "out of", "no gpu"): - return BlockScope{CapacityType: capacityType} + return capacityScope() default: return BlockScope{DenyAll: true} } diff --git a/pkg/provider/errors_test.go b/pkg/provider/errors_test.go index 3c91ceb..6e80b6c 100644 --- a/pkg/provider/errors_test.go +++ b/pkg/provider/errors_test.go @@ -18,6 +18,7 @@ package provider import ( "fmt" + "reflect" "testing" nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" @@ -25,6 +26,10 @@ import ( func TestClassifyError(t *testing.T) { const tier = nebulav1alpha1.CapacityOnDemand + const accel = "H100" + // capacityScope is the expected accelerator-scoped block: the accelerator is + // promoted to an exact-match pointer. + capacityScope := BlockScope{CapacityType: tier, AcceleratorType: ptrStr(accel)} tests := []struct { name string err error @@ -33,17 +38,17 @@ func TestClassifyError(t *testing.T) { {"nil", nil, BlockScope{}}, {"auth sentinel", ErrAuth, BlockScope{DenyAll: true}}, {"quota sentinel", ErrQuota, BlockScope{DenyAll: true}}, - {"no-capacity sentinel", ErrNoCapacity, BlockScope{CapacityType: tier}}, - {"unsupported sentinel", ErrUnsupportedAccelerator, BlockScope{CapacityType: tier}}, - {"wrapped sentinel", fmt.Errorf("provision failed: %w", ErrNoCapacity), BlockScope{CapacityType: tier}}, + {"no-capacity sentinel", ErrNoCapacity, capacityScope}, + {"unsupported sentinel", ErrUnsupportedAccelerator, capacityScope}, + {"wrapped sentinel", fmt.Errorf("provision failed: %w", ErrNoCapacity), capacityScope}, {"string unauthorized", fmt.Errorf("HTTP 401 unauthorized"), BlockScope{DenyAll: true}}, {"string quota", fmt.Errorf("account limit exceeded"), BlockScope{DenyAll: true}}, - {"string capacity", fmt.Errorf("no capacity available"), BlockScope{CapacityType: tier}}, + {"string capacity", fmt.Errorf("no capacity available"), capacityScope}, {"unknown conservative", fmt.Errorf("weird transient blip"), BlockScope{DenyAll: true}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if got := ClassifyError(tt.err, tier); got != tt.want { + if got := ClassifyError(tt.err, tier, accel); !reflect.DeepEqual(got, tt.want) { t.Fatalf("ClassifyError(%v) = %+v, want %+v", tt.err, got, tt.want) } }) @@ -53,8 +58,22 @@ func TestClassifyError(t *testing.T) { // Spot tier must be stamped onto accelerator-scoped blocks so a Spot failure // does not block OnDemand on the same provider. func TestClassifyError_TierStamped(t *testing.T) { - got := ClassifyError(ErrNoCapacity, nebulav1alpha1.CapacitySpot) + got := ClassifyError(ErrNoCapacity, nebulav1alpha1.CapacitySpot, "H100") if got.CapacityType != nebulav1alpha1.CapacitySpot || got.DenyAll { t.Fatalf("expected Spot accelerator-scoped block, got %+v", got) } } + +// An empty accelerator (CPU-only Pod) must leave AcceleratorType nil ("not +// applicable"), never a wildcard that would widen the block across every GPU. +func TestClassifyError_EmptyAcceleratorStaysNil(t *testing.T) { + got := ClassifyError(ErrNoCapacity, nebulav1alpha1.CapacityOnDemand, "") + if got.AcceleratorType != nil { + t.Fatalf("expected nil AcceleratorType for a CPU-only request, got %+v", got.AcceleratorType) + } + if got.DenyAll || got.CapacityType != nebulav1alpha1.CapacityOnDemand { + t.Fatalf("expected an OnDemand capacity scope, got %+v", got) + } +} + +func ptrStr(s string) *string { return &s } diff --git a/pkg/provider/fake/fake.go b/pkg/provider/fake/fake.go index baa03ec..a436606 100644 --- a/pkg/provider/fake/fake.go +++ b/pkg/provider/fake/fake.go @@ -148,8 +148,8 @@ func (p *Provider) List(_ context.Context) ([]provider.Instance, error) { // ClassifyProvisionError never really fails to provision, so any error it is // asked to classify is treated as a whole-provider block on OnDemand — the same // shared derivation the real adapters use. -func (p *Provider) ClassifyProvisionError(err error) provider.BlockScope { - return provider.ClassifyError(err, nebulav1alpha1.CapacityOnDemand) +func (p *Provider) ClassifyProvisionError(err error, accelerator string) provider.BlockScope { + return provider.ClassifyError(err, nebulav1alpha1.CapacityOnDemand, accelerator) } // fixedCatalog is a trivial catalog.Lookup: a static set of OnDemand offerings diff --git a/pkg/provider/modal/modal.go b/pkg/provider/modal/modal.go index 25e289a..d59961d 100644 --- a/pkg/provider/modal/modal.go +++ b/pkg/provider/modal/modal.go @@ -230,8 +230,8 @@ func (p *Provider) List(ctx context.Context) ([]provider.Instance, error) { // matching shared sentinel (e.g. fmt.Errorf("...: %w", provider.ErrNoCapacity)); // ClassifyError honours those first and falls back to string heuristics for raw // API messages, so no Modal-specific matching is duplicated here. -func (p *Provider) ClassifyProvisionError(err error) provider.BlockScope { - return provider.ClassifyError(err, nebulav1alpha1.CapacityOnDemand) +func (p *Provider) ClassifyProvisionError(err error, accelerator string) provider.BlockScope { + return provider.ClassifyError(err, nebulav1alpha1.CapacityOnDemand, accelerator) } // findByClaim returns the sandbox tagged with claimName, or nil if none. diff --git a/pkg/provider/modal/modal_test.go b/pkg/provider/modal/modal_test.go index 18c2610..bd23c41 100644 --- a/pkg/provider/modal/modal_test.go +++ b/pkg/provider/modal/modal_test.go @@ -278,7 +278,7 @@ func TestClassifyProvisionError(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := p.ClassifyProvisionError(tt.err) + got := p.ClassifyProvisionError(tt.err, "") if got != tt.want { t.Fatalf("got %+v, want %+v", got, tt.want) } diff --git a/pkg/provider/provider.go b/pkg/provider/provider.go index b8f75d5..b59fa9b 100644 --- a/pkg/provider/provider.go +++ b/pkg/provider/provider.go @@ -96,7 +96,14 @@ type Provider interface { // the failing placement should be blocklisted. This keeps failover precise: // a "no H100 capacity" error blocks only {provider, H100, capacityType}, // while an auth/quota error blocks the whole provider. See BlockScope. - ClassifyProvisionError(err error) BlockScope + // + // accelerator is the canonical accelerator the failing request asked for ("" + // for a CPU-only Pod). The provider needs it because the error alone does not + // carry it, and the provider is the single owner of the complete scope: it + // stamps the accelerator on a capacity block AND adds any provider-specific + // axis it knows (AWS confines the block to its region). No caller assembles a + // scope piecemeal after this returns. + ClassifyProvisionError(err error, accelerator string) BlockScope } // ProvisionRequest carries only the decisions the placement controller made @@ -114,6 +121,13 @@ type ProvisionRequest struct { // This is the one workload-independent decision that cannot be expressed on // the Pod, so it must be passed explicitly. CapacityType nebulav1alpha1.CapacityType + // Region is the provider region the optimizer chose to provision in, in the + // provider's own vocabulary (e.g. AWS "us-east-1"). Like CapacityType it is a + // workload-independent decision absent from the Pod. Empty means "use the + // provider's configured default region" — region-simple NeoClouds (Modal, + // RunPod) ignore it, and a region-aware adapter falls back to the region its + // client was built with (see the AWS adapter's NewSDKClient). + Region string } // Capabilities declares provider quirks as data, so the control plane filters @@ -140,6 +154,17 @@ type Capabilities struct { // this — they are pushed synchronously by CreatePod/DeletePod — so this only // bounds detection latency for events no provider pushes. PollInterval time.Duration + // ProvisionTimeout bounds a single Provision call end to end. It exists for + // region-aware providers that fail over across capacity pools (e.g. AWS tries + // each availability zone in turn on a capacity error): the deadline caps how + // long that inner failover loop may spend before giving up so the outer + // region-level failover can proceed, rather than a slow provider stalling a + // Pod indefinitely. It bounds the PROVISIONING attempt (the launch API calls + // and any per-zone retries), NOT "the workload became healthy" — that remains + // the poll loop's job. Zero means "no adapter-imposed deadline" (the caller's + // context still applies); providers with a single capacity pool (Modal) leave + // it zero. + ProvisionTimeout time.Duration } // Instance is the provider-agnostic view of one external instance, as observed. @@ -151,6 +176,9 @@ type Instance struct { Endpoint string // CapacityType reflects how the instance was provisioned, when known. CapacityType nebulav1alpha1.CapacityType + // Region is where the instance actually lives, in the provider's own + // vocabulary. Empty for region-simple providers that do not report one. + Region string } // InstanceState is the provider-agnostic lifecycle state, normalized from each @@ -177,19 +205,75 @@ type Offering struct { CapacityType nebulav1alpha1.CapacityType PricePerHour float64 Available bool + // Region is the provider region this row prices, in the provider's own + // vocabulary (e.g. AWS "us-east-1"). Empty for region-simple providers whose + // catalog is not region-partitioned (Modal, RunPod); a region-aware provider + // emits one row per {accelerator, capacityType, region}. + Region string + // AcceleratorID is this provider's own name for what serves the canonical + // AcceleratorType — the per-provider half of the mapping the canonical GPU + // vocabulary is translated through (e.g. AWS "p5.48xlarge" for H100). It is + // the lookup data a diverging provider's MapAccelerator override returns. Left + // unset by providers whose mapping is identity (Modal names its GPUs exactly + // like the canonical types), which reuse catalog.Base's default MapAccelerator + // and so need no per-name id. Mirrors the providerAcceleratorID that + // MapAccelerator returns. + AcceleratorID string + // GPUCount is how many accelerators the AcceleratorID provides. It matters for + // providers where the accelerator count is not a free knob but is baked into + // the offering: on AWS you do not ask for "N T4s", you pick an instance type + // whose GPU count is fixed (T4x1 = g4dn.xlarge, T4x8 = g4dn.metal), so the + // mapping key is (AcceleratorType, GPUCount) and the same accelerator appears + // once per count. Left 0 by providers that attach an arbitrary count to a + // single offering (Modal takes the count as a parameter), for whom it is not a + // lookup dimension. + GPUCount int32 } // BlockScope is the granularity at which a failed placement is excluded, matched -// to what actually failed (SkyPilot's blocklist-granularity rule). Empty fields -// act as wildcards in the in-memory blocklist match: a failed H100 request must -// not disqualify A100 requests on the same provider. +// to what actually failed (SkyPilot's blocklist-granularity rule). It is a partial +// MATCH PATTERN, not a value, and AcceleratorType/Region are THREE-STATE pointers +// so a pattern can distinguish "any value" from "no value": +// +// - nil => the axis is NOT APPLICABLE to this block: it matches only a +// candidate whose corresponding field is also empty. A CPU-only Pod (no +// accelerator) and a region-simple provider (no region) both leave the axis +// nil, so the block never spuriously widens across it. +// - &"" => WILDCARD: matches any value on that axis (a Spot-capacity block that +// spans every accelerator, say). +// - &"v" => EXACT: matches only candidates whose field equals "v" (so a failed +// H100 request must not disqualify A100 on the same provider). +// +// This is why the scope is assembled in two places: the adapter classifies from the +// error alone (which knows the region — AWS is bound to one — but never the +// accelerator, a property of the request), and the vnode handler fills the +// accelerator in from the failing Pod. Neither can produce the full scope, so the +// unset fields must be distinguishable from a deliberate wildcard — hence pointers. +// DenyAll is the broadest scope and ignores the pattern fields entirely. +// +// This is the opposite sense to a value field like NodePoolSpec's +// ProviderSpec.Regions, where empty means "the one default region" — because a +// pattern matches, while a request takes a single default. The blocklist never +// mixes the two: a NodePool candidate is resolved to fully-concrete values before +// it is matched against these patterns. type BlockScope struct { - // AcceleratorType empty => blocks all accelerator types on this provider. - AcceleratorType string + // AcceleratorType: nil => not applicable (CPU-only Pod, or a DenyAll scope); + // &"" => wildcard, matches every accelerator; &"H100" => exactly that one. + // A capacity error carries no accelerator (it is classified from the error + // alone), so the adapter leaves this nil and the vnode handler fills it in + // from the failing Pod's request — narrowing the block to the accelerator that + // actually failed. + AcceleratorType *string // CapacityType empty => blocks all capacity types. CapacityType nebulav1alpha1.CapacityType + // Region: nil => the provider has no region axis (a region-simple provider like + // Modal/RunPod, whose candidate carries an empty region too, so nil matches it); + // &"us-east-1" => confines the block to that one region, so a "no capacity in + // us-east-1" failure does not disqualify the same request in us-west-2. A + // region-aware provider (AWS) sets it; a region-simple one leaves it nil. + Region *string // DenyAll true => block everything on this provider (e.g. auth/quota errors), - // ignoring AcceleratorType/CapacityType. The scope is still this one provider; - // it never spans providers. + // ignoring AcceleratorType/CapacityType/Region. The scope is still this one + // provider; it never spans providers. DenyAll bool } diff --git a/pkg/provider/registry.go b/pkg/provider/registry.go index 776720d..d966197 100644 --- a/pkg/provider/registry.go +++ b/pkg/provider/registry.go @@ -12,14 +12,16 @@ import ( // ProviderLabel value on each provider's virtual node. Keeping them as consts in // one place stops those three call sites from drifting apart on a typo. // -// Hyperscalers (aws/gcp/azure) are a planned near-term expansion and will be -// added here when their adapters land; nothing else about the registry changes. +// AWS is the first hyperscaler (region-aware) backend; GCP/Azure are a planned +// near-term expansion and will be added here when their adapters land, with +// nothing else about the registry changing. const ( ProviderRunPod = "runpod" ProviderModal = "modal" ProviderCoreWeave = "coreweave" ProviderLambda = "lambda" ProviderKubernetes = "kubernetes" + ProviderAWS = "aws" ) // registry is the process-wide map of registered provider backends, keyed by diff --git a/pkg/vnode/handler.go b/pkg/vnode/handler.go index 4b0f330..e42630b 100644 --- a/pkg/vnode/handler.go +++ b/pkg/vnode/handler.go @@ -29,6 +29,7 @@ import ( "github.com/virtual-kubelet/virtual-kubelet/node/api/statsv1alpha1" corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + logf "sigs.k8s.io/controller-runtime/pkg/log" nebulav1alpha1 "github.com/InftyAI/Nebula/api/v1alpha1" "github.com/InftyAI/Nebula/pkg/provider" @@ -44,6 +45,32 @@ import ( // reclaims sooner; an OnDemand-only one can poll lazily). const defaultPollInterval = 15 * time.Second +// defaultBlocklistTTL bounds how long a failed placement is excluded when the Pod +// carries no BlocklistTTLAnnotation (an out-of-band Pod, or one placed before the +// annotation existed). It mirrors FailoverPolicy.BlocklistTTL's own default so the +// behaviour is identical whether or not the pool policy reached the handler. +const defaultBlocklistTTL = 10 * time.Minute + +// Blocklister records a failed placement so the placement controller can fail over +// to the next candidate instead of hot-looping against a provider that just +// rejected the request. It is the write half of pkg/failover.Blocklist; the +// handler depends on this narrow interface (not the concrete type) so it stays +// testable and a nil blocklist is a no-op. +type Blocklister interface { + Record(prov string, scope provider.BlockScope, ttl time.Duration) +} + +// defaultProvisionTimeout bounds a single Provision call when the provider does +// not override it (Capabilities.ProvisionTimeout). Provisioning is a network +// call to the backend (Modal's gRPC CreateSandbox, AWS's RunInstances and its +// per-zone capacity failover), and none of them are guaranteed to return +// promptly — a hung call would otherwise pin the pod controller's worker +// indefinitely. Enforcing it here, at the single call site, keeps every provider +// bounded by default; a provider that legitimately needs longer (AWS sweeping +// several zones) raises it via Capabilities. It is deliberately generous: the +// happy path returns in well under this, so it is a backstop, not a tuning knob. +const defaultProvisionTimeout = 90 * time.Second + // Handler bridges one provider into the virtual kubelet: it implements the VK // PodLifecycleHandler (+ PodNotifier) so the pod controller's CreatePod // provisions an external instance through the provider seam and DeletePod @@ -58,6 +85,12 @@ const defaultPollInterval = 15 * time.Second type Handler struct { prov provider.Provider + // blocklist records failed placements so the placement controller can fail over + // (zone → region → tier). Shared across every provider's handler and the + // placement controller. May be nil (a no-op), keeping tests and any + // blocklist-less wiring simple. + blocklist Blocklister + mu sync.Mutex tracked map[string]*trackedPod // key: namespace/name notify func(*corev1.Pod) @@ -79,14 +112,16 @@ type trackedPod struct { // NewHandler builds a Handler for the given provider backend. The poll cadence // comes from the provider's Capabilities.PollInterval, falling back to -// defaultPollInterval when the provider does not set one. -func NewHandler(prov provider.Provider) *Handler { +// defaultPollInterval when the provider does not set one. blocklist is the shared +// failover blocklist a Provision failure is recorded on; it may be nil (no-op). +func NewHandler(prov provider.Provider, blocklist Blocklister) *Handler { poll := prov.Capabilities().PollInterval if poll <= 0 { poll = defaultPollInterval } return &Handler{ prov: prov, + blocklist: blocklist, tracked: make(map[string]*trackedPod), nowFn: metav1.Now, pollEvery: poll, @@ -120,10 +155,28 @@ func (h *Handler) CreatePod(ctx context.Context, pod *corev1.Pod) error { req := provider.ProvisionRequest{ ClaimName: claim, CapacityType: nebulav1alpha1.CapacityType(pod.Annotations[nebulav1alpha1.CapacityTypeAnnotation]), + Region: pod.Annotations[nebulav1alpha1.RegionAnnotation], } + // Bound the provision call so a wedged backend cannot pin this worker forever. + // The provider may raise the deadline via Capabilities.ProvisionTimeout (AWS + // does, to leave room for cross-zone failover); zero means "use the default". + timeout := h.prov.Capabilities().ProvisionTimeout + if timeout <= 0 { + timeout = defaultProvisionTimeout + } + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + id, err := h.prov.Provision(ctx, pod, req) if err != nil { + // Record the failure on the shared blocklist so placement fails over to the + // next candidate (zone → region → tier) instead of hot-looping here. The + // provider classifies its own error into the precise BlockScope (a Spot + // no-capacity in one region blocks only that, not OnDemand or other regions; + // auth/quota blocks the whole provider); the TTL rides on the Pod from the + // pool's FailoverPolicy. + h.recordBlock(pod, err) // Surface the failure on the Pod so the placement controller can fail // over, and return the error so the pod controller retries with backoff. h.markStatus(pod, corev1.PodFailed, reasonProvisionFailed, err.Error()) @@ -230,11 +283,14 @@ func (h *Handler) NotifyPods(ctx context.Context, cb func(*corev1.Pod)) { // pollLoop periodically reconciles tracked pods against the provider's live // instance list and pushes any status change through the notify callback. func (h *Handler) pollLoop(ctx context.Context) { + log := logf.FromContext(ctx).WithName("vnode-poll").WithValues("provider", h.prov.Name()) + log.Info("poll loop started", "interval", h.pollEvery.String()) t := time.NewTicker(h.pollEvery) defer t.Stop() for { select { case <-ctx.Done(): + log.Info("poll loop stopped") return case <-t.C: h.reconcileOnce(ctx) @@ -246,9 +302,16 @@ func (h *Handler) pollLoop(ctx context.Context) { // from the matching instance (matched by claim name, since a pod maps 1:1 to a // claim). A tracked pod whose instance is absent from the list is treated as // terminated (preempted / externally torn down), per the provider contract. +// +// A List error is logged and skipped (retry next tick) rather than swallowed: +// a provider whose List fails every tick would otherwise strand every tracked +// pod in its last status forever with no signal, which is exactly the failure +// mode this logging exists to surface. func (h *Handler) reconcileOnce(ctx context.Context) { + log := logf.FromContext(ctx).WithName("vnode-poll").WithValues("provider", h.prov.Name()) instances, err := h.prov.List(ctx) if err != nil { + log.Error(err, "provider List failed; tracked pod statuses will not advance this tick") return // transient; retry on the next tick } byClaim := make(map[string]provider.Instance, len(instances)) @@ -258,12 +321,15 @@ func (h *Handler) reconcileOnce(ctx context.Context) { h.mu.Lock() changed := make([]*corev1.Pod, 0) + tracked := len(h.tracked) + matched := 0 for _, tp := range h.tracked { inst, present := byClaim[tp.claimName] before := tp.pod.Status.Phase if !present { applyState(tp.pod, provider.InstanceTerminated, "", h.nowFn()) } else { + matched++ applyState(tp.pod, inst.State, inst.Endpoint, h.nowFn()) } if tp.pod.Status.Phase != before { @@ -273,6 +339,12 @@ func (h *Handler) reconcileOnce(ctx context.Context) { notify := h.notify h.mu.Unlock() + // At V(1) so a healthy steady state stays quiet, but the counts are there when + // pods appear stuck: tracked>0 with matched==0 means List returns instances the + // claim names don't line up with (the classic "provisioned but never Running"). + log.V(1).Info("poll tick", + "listed", len(instances), "tracked", tracked, "matched", matched, "changed", len(changed)) + if notify != nil { for _, p := range changed { notify(p) @@ -307,6 +379,57 @@ func (h *Handler) markStatus(pod *corev1.Pod, phase corev1.PodPhase, reason, msg setPhase(pod, phase, reason, msg, h.nowFn()) } +// recordBlock classifies a Provision failure into a BlockScope (via the provider) +// and records it on the shared blocklist for the pool's BlocklistTTL, so placement +// fails over instead of retrying the same failing candidate. It is a no-op when no +// blocklist is wired or the error yields an empty scope (nothing to block). +// +// The provider owns the whole scope: the handler resolves the requested accelerator +// off the Pod (the error does not carry it) and passes it in, but does NOT assemble +// or mutate the scope itself — narrowing to the failing accelerator and any +// region axis is the provider's job (see ClassifyError / the adapters). This keeps +// scope derivation in one place per provider rather than spread across the handler. +// +// Because the provider stamps the accelerator, the block errs NARROW: a truly +// account/region-wide error (e.g. a per-region Spot limit that spans accelerators) +// is narrowed to just this accelerator too, since the error text does not +// distinguish it from an instance-type shortage; the cost is that a sibling +// accelerator gets one wasted re-probe before it re-records. That is the right +// trade — over-narrow costs a fast retry, over-broad wrongly excludes serviceable +// accelerators. DenyAll (auth/quota) ignores the accelerator: it fails for all. +func (h *Handler) recordBlock(pod *corev1.Pod, err error) { + if h.blocklist == nil { + return + } + // The accelerator is a property of the REQUEST, not the error; the provider + // cannot see it, so we resolve it off the Pod and hand it over ("" for a + // CPU-only Pod, which the provider treats as "not applicable"). + accel, _, _ := util.AcceleratorRequest(pod) + scope := h.prov.ClassifyProvisionError(err, accel) + if scope == (provider.BlockScope{}) { + // An empty scope means the error is not one we know how to blocklist; do not + // install a wildcard block that would exclude everything on the provider. + return + } + h.blocklist.Record(h.prov.Name(), scope, blocklistTTL(pod)) +} + +// blocklistTTL reads the pool's FailoverPolicy.BlocklistTTL off the Pod annotation +// the placement controller stamped, falling back to defaultBlocklistTTL when it is +// absent or unparseable (an out-of-band Pod, or a non-positive value that could +// otherwise install a permanent block). +func blocklistTTL(pod *corev1.Pod) time.Duration { + raw := pod.Annotations[nebulav1alpha1.BlocklistTTLAnnotation] + if raw == "" { + return defaultBlocklistTTL + } + d, err := time.ParseDuration(raw) + if err != nil || d <= 0 { + return defaultBlocklistTTL + } + return d +} + func ptrNow(t metav1.Time) *metav1.Time { return &t } // --- Unused nodeutil.Provider surface -------------------------------------- diff --git a/pkg/vnode/handler_test.go b/pkg/vnode/handler_test.go index fe78a50..e8da9f4 100644 --- a/pkg/vnode/handler_test.go +++ b/pkg/vnode/handler_test.go @@ -45,6 +45,12 @@ type fakeProvider struct { list []provider.Instance listErr error capabilities provider.Capabilities + // classifyScope is what ClassifyProvisionError returns for a failure; the zero + // value (empty scope) means "not blocklistable". classifyAccel records the + // accelerator the handler passed in, so a test can assert it resolved it off the + // Pod (the provider now owns the whole scope; the handler only supplies this). + classifyScope provider.BlockScope + classifyAccel string } func (f *fakeProvider) Name() string { return "fake" } @@ -72,8 +78,25 @@ func (f *fakeProvider) List(context.Context) ([]provider.Instance, error) { } func (f *fakeProvider) Offerings(context.Context) ([]provider.Offering, error) { return nil, nil } func (f *fakeProvider) MapAccelerator(c string) (string, bool) { return c, true } -func (f *fakeProvider) ClassifyProvisionError(error) provider.BlockScope { - return provider.BlockScope{} +func (f *fakeProvider) ClassifyProvisionError(_ error, accel string) provider.BlockScope { + f.classifyAccel = accel + return f.classifyScope +} + +// recordingBlocklist captures Record calls so a test can assert what the handler +// blocklisted on a Provision failure. +type recordingBlocklist struct { + prov string + scope provider.BlockScope + ttl time.Duration + calls int +} + +func (b *recordingBlocklist) Record(prov string, scope provider.BlockScope, ttl time.Duration) { + b.prov = prov + b.scope = scope + b.ttl = ttl + b.calls++ } func testPod(ns, name string) *corev1.Pod { @@ -85,7 +108,7 @@ func testPod(ns, name string) *corev1.Pod { func TestCreatePod_ProvisionsAndTracks(t *testing.T) { fp := &fakeProvider{provisionID: "inst-1"} - h := NewHandler(fp) + h := NewHandler(fp, nil) pod := testPod("default", "p1") pod.Annotations = map[string]string{nebulav1alpha1.CapacityTypeAnnotation: string(nebulav1alpha1.CapacitySpot)} @@ -113,7 +136,7 @@ func TestCreatePod_ProvisionsAndTracks(t *testing.T) { func TestCreatePod_ProvisionErrorSurfaces(t *testing.T) { fp := &fakeProvider{provisionErr: errors.New("no capacity")} - h := NewHandler(fp) + h := NewHandler(fp, nil) pod := testPod("default", "p1") if err := h.CreatePod(context.Background(), pod); err == nil { @@ -129,9 +152,76 @@ func TestCreatePod_ProvisionErrorSurfaces(t *testing.T) { } } +func TestCreatePod_ProvisionFailureRecordsBlock(t *testing.T) { + accel, region := "H100", "us-east-1" + scope := provider.BlockScope{ + AcceleratorType: &accel, + CapacityType: nebulav1alpha1.CapacitySpot, + Region: ®ion, + } + fp := &fakeProvider{provisionErr: errors.New("no capacity"), classifyScope: scope} + bl := &recordingBlocklist{} + h := NewHandler(fp, bl) + + pod := testPod("default", "p1") + pod.Annotations = map[string]string{nebulav1alpha1.BlocklistTTLAnnotation: "3m"} + pod.Labels = map[string]string{nebulav1alpha1.AcceleratorTypeLabel: "H100"} + + if err := h.CreatePod(context.Background(), pod); err == nil { + t.Fatal("expected CreatePod to return the provision error") + } + if bl.calls != 1 { + t.Fatalf("expected exactly one Record call, got %d", bl.calls) + } + // The handler must resolve the accelerator off the Pod and hand it to the + // provider (which owns scope assembly), not mutate the returned scope itself. + if fp.classifyAccel != "H100" { + t.Fatalf("handler passed accelerator %q to ClassifyProvisionError, want H100", fp.classifyAccel) + } + if bl.prov != "fake" { + t.Fatalf("recorded provider = %q, want fake", bl.prov) + } + if bl.scope != scope { + t.Fatalf("recorded scope = %+v, want %+v", bl.scope, scope) + } + // TTL comes from the Pod annotation the placement controller stamped. + if bl.ttl != 3*time.Minute { + t.Fatalf("recorded ttl = %v, want 3m from the annotation", bl.ttl) + } +} + +func TestCreatePod_EmptyScopeDoesNotBlock(t *testing.T) { + // A classifier that yields the zero scope must NOT install a wildcard block + // (which would exclude everything on the provider). + fp := &fakeProvider{provisionErr: errors.New("weird"), classifyScope: provider.BlockScope{}} + bl := &recordingBlocklist{} + h := NewHandler(fp, bl) + + if err := h.CreatePod(context.Background(), testPod("default", "p1")); err == nil { + t.Fatal("expected CreatePod to return the provision error") + } + if bl.calls != 0 { + t.Fatalf("expected no Record call for an empty scope, got %d", bl.calls) + } +} + +func TestCreatePod_MissingTTLAnnotationUsesDefault(t *testing.T) { + fp := &fakeProvider{provisionErr: errors.New("no capacity"), classifyScope: provider.BlockScope{DenyAll: true}} + bl := &recordingBlocklist{} + h := NewHandler(fp, bl) + + // No BlocklistTTLAnnotation on the Pod => the handler's built-in default. + if err := h.CreatePod(context.Background(), testPod("default", "p1")); err == nil { + t.Fatal("expected CreatePod to return the provision error") + } + if bl.ttl != defaultBlocklistTTL { + t.Fatalf("recorded ttl = %v, want default %v", bl.ttl, defaultBlocklistTTL) + } +} + func TestDeletePod_TerminatesAndUntracks(t *testing.T) { fp := &fakeProvider{provisionID: "inst-1"} - h := NewHandler(fp) + h := NewHandler(fp, nil) pod := testPod("default", "p1") if err := h.CreatePod(context.Background(), pod); err != nil { @@ -153,7 +243,7 @@ func TestDeletePod_TerminatesAndUntracks(t *testing.T) { func TestDeletePod_Idempotent(t *testing.T) { fp := &fakeProvider{provisionID: "inst-1"} - h := NewHandler(fp) + h := NewHandler(fp, nil) pod := testPod("default", "p1") _ = h.CreatePod(context.Background(), pod) @@ -167,7 +257,7 @@ func TestDeletePod_Idempotent(t *testing.T) { func TestReconcileOnce_ReportsRunning(t *testing.T) { fp := &fakeProvider{provisionID: "inst-1"} - h := NewHandler(fp) + h := NewHandler(fp, nil) pod := testPod("default", "p1") _ = h.CreatePod(context.Background(), pod) @@ -206,7 +296,7 @@ func TestReconcileOnce_ReportsRunning(t *testing.T) { func TestReconcileOnce_AbsentInstanceIsTerminated(t *testing.T) { fp := &fakeProvider{provisionID: "inst-1"} - h := NewHandler(fp) + h := NewHandler(fp, nil) pod := testPod("default", "p1") _ = h.CreatePod(context.Background(), pod) @@ -226,18 +316,18 @@ func TestReconcileOnce_AbsentInstanceIsTerminated(t *testing.T) { func TestNewHandler_PollIntervalFromCapabilities(t *testing.T) { // A provider that declares a cadence overrides the default. custom := &fakeProvider{capabilities: provider.Capabilities{PollInterval: 5 * time.Second}} - if got := NewHandler(custom).pollEvery; got != 5*time.Second { + if got := NewHandler(custom, nil).pollEvery; got != 5*time.Second { t.Fatalf("expected the provider's PollInterval, got %v", got) } // A provider that leaves it zero falls back to the vnode default. - if got := NewHandler(&fakeProvider{}).pollEvery; got != defaultPollInterval { + if got := NewHandler(&fakeProvider{}, nil).pollEvery; got != defaultPollInterval { t.Fatalf("expected the default cadence, got %v", got) } } func TestGetPods_ReturnsTracked(t *testing.T) { fp := &fakeProvider{provisionID: "inst-1"} - h := NewHandler(fp) + h := NewHandler(fp, nil) _ = h.CreatePod(context.Background(), testPod("default", "p1")) _ = h.CreatePod(context.Background(), testPod("default", "p2")) diff --git a/pkg/vnode/node.go b/pkg/vnode/node.go index d8133e7..f6f109c 100644 --- a/pkg/vnode/node.go +++ b/pkg/vnode/node.go @@ -72,17 +72,21 @@ func NodeName(providerName string) string { // 0.33 line this project pins. Nebula does not serve the kubelet API (logs/exec // are unsupported — see handler.go), so none of that surface is needed. type Runner struct { - prov provider.Provider - client kubernetes.Interface - nodeName string + prov provider.Provider + client kubernetes.Interface + blocklist Blocklister + nodeName string } -// NewRunner builds the virtual-node runner for one provider. -func NewRunner(prov provider.Provider, client kubernetes.Interface) *Runner { +// NewRunner builds the virtual-node runner for one provider. blocklist is the +// shared failover blocklist the handler records Provision failures on (may be +// nil). +func NewRunner(prov provider.Provider, client kubernetes.Interface, blocklist Blocklister) *Runner { return &Runner{ - prov: prov, - client: client, - nodeName: NodeName(prov.Name()), + prov: prov, + client: client, + blocklist: blocklist, + nodeName: NodeName(prov.Name()), } } @@ -93,7 +97,7 @@ var _ manager.Runnable = (*Runner)(nil) func (r *Runner) Start(ctx context.Context) error { log := logf.FromContext(ctx).WithValues("virtualNode", r.nodeName, "provider", r.prov.Name()) - handler := NewHandler(r.prov) + handler := NewHandler(r.prov, r.blocklist) nodeSpec := nodeSpec(r.nodeName, r.prov.Name()) // Pod informer scoped to this node only (spec.nodeName == nodeName), so the diff --git a/pkg/vnode/status.go b/pkg/vnode/status.go index ff5d5cd..1572564 100644 --- a/pkg/vnode/status.go +++ b/pkg/vnode/status.go @@ -51,6 +51,12 @@ const ( // Pending -> PodPending, Ready=False (starting) // Failed -> PodFailed // Terminated -> PodFailed (instance gone: torn down or reclaimed out-of-band) +// +// Running already means "reachable": a provider only reports InstanceRunning once +// the instance has passed its readiness bar (for AWS, the 2/2 EC2 status checks — +// see toState), so reaching Running is the point at which the Pod is both Running +// and Ready. Ready is the condition Kubernetes counts toward a Deployment's ready +// replicas. func applyState(pod *corev1.Pod, state provider.InstanceState, endpoint string, now metav1.Time) { switch state { case provider.InstanceRunning: