diff --git a/apps/api/cmd/gateway/main.go b/apps/api/cmd/gateway/main.go index a17228f..d55e092 100644 --- a/apps/api/cmd/gateway/main.go +++ b/apps/api/cmd/gateway/main.go @@ -111,17 +111,12 @@ func main() { var handler http.Handler if cfg.RunsCapacityController() { - kubernetes, kubernetesErr := capacitycontroller.NewKubernetesClient(capacitycontroller.KubernetesConfig{ - Namespace: cfg.CapacityControllerNamespace, - APIServer: cfg.CapacityControllerAPIServer, - TokenFile: cfg.CapacityControllerTokenFile, - CAFile: cfg.CapacityControllerCAFile, - }) - if kubernetesErr != nil { - logger.Error("initialize capacity controller Kubernetes client failed", "error", kubernetesErr) + orchestrator, orchestratorErr := capacitycontroller.NewConfiguredAdapter(cfg) + if orchestratorErr != nil { + logger.Error("initialize capacity orchestrator adapter failed", "error", orchestratorErr) os.Exit(1) } - controller := capacitycontroller.New(cfg, coordinationDB, kubernetes, logger) + controller := capacitycontroller.New(cfg, coordinationDB, orchestrator, logger) go controller.Run(ctx) handler = controller.Handler() } else { diff --git a/apps/api/docs/swagger.json b/apps/api/docs/swagger.json index d9712ac..eda3d99 100644 --- a/apps/api/docs/swagger.json +++ b/apps/api/docs/swagger.json @@ -2082,6 +2082,36 @@ } } }, + "/api/admin/runtime/execution-pools": { + "get": { + "security": [ + { + "BearerAuth": [] + } + ], + "produces": [ + "application/json" + ], + "tags": [ + "Admin Runtime" + ], + "summary": "查询逻辑执行池、Worker 容量和上游路由健康状态", + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/httpapi.AdminExecutionPoolResponse" + } + }, + "500": { + "description": "Internal Server Error", + "schema": { + "$ref": "#/definitions/httpapi.ErrorEnvelope" + } + } + } + } + }, "/api/admin/runtime/model-rate-limits": { "get": { "security": [ @@ -10699,6 +10729,177 @@ } } }, + "httpapi.AdminExecutionPool": { + "type": "object", + "properties": { + "capabilities": { + "type": "object", + "additionalProperties": {} + }, + "capacity": { + "$ref": "#/definitions/httpapi.AdminExecutionPoolCapacity" + }, + "labels": { + "type": "object", + "additionalProperties": { + "type": "string" + } + }, + "poolId": { + "type": "string" + }, + "state": { + "type": "string" + }, + "workers": { + "type": "array", + "items": { + "$ref": "#/definitions/httpapi.AdminExecutionPoolWorker" + } + } + } + }, + "httpapi.AdminExecutionPoolCapacity": { + "type": "object", + "properties": { + "activeTasks": { + "type": "integer" + }, + "allocated": { + "type": "integer" + }, + "critical": { + "type": "boolean" + }, + "estimatedQueueWaitMs": { + "type": "integer" + }, + "heavyCapacity": { + "type": "integer" + }, + "heavyTasks": { + "type": "integer" + }, + "queueDepth": { + "type": "integer" + }, + "resourceHeadroom": { + "type": "number" + }, + "safeCapacity": { + "type": "integer" + }, + "sampledAt": { + "type": "string" + }, + "stability": { + "type": "number" + }, + "workerCount": { + "type": "integer" + } + } + }, + "httpapi.AdminExecutionPoolResponse": { + "type": "object", + "properties": { + "capturedAt": { + "type": "string" + }, + "pools": { + "type": "array", + "items": { + "$ref": "#/definitions/httpapi.AdminExecutionPool" + } + }, + "routeHealth": { + "type": "array", + "items": { + "$ref": "#/definitions/httpapi.AdminRouteHealth" + } + } + } + }, + "httpapi.AdminExecutionPoolWorker": { + "type": "object", + "properties": { + "activeTasks": { + "type": "integer" + }, + "allocated": { + "type": "integer" + }, + "capabilities": { + "type": "object", + "additionalProperties": {} + }, + "heartbeatAt": { + "type": "string" + }, + "instanceId": { + "type": "string" + }, + "pressureState": { + "type": "string" + }, + "protocolVersion": { + "type": "string" + }, + "revision": { + "type": "string" + }, + "safeCapacity": { + "type": "integer" + }, + "workerId": { + "type": "string" + } + } + }, + "httpapi.AdminRouteHealth": { + "type": "object", + "properties": { + "connectTlsP95Ms": { + "type": "integer" + }, + "consecutiveFailures": { + "type": "integer" + }, + "consecutiveSuccesses": { + "type": "integer" + }, + "expiresAt": { + "type": "string" + }, + "firstByteP95Ms": { + "type": "integer" + }, + "jitterP95Ms": { + "type": "integer" + }, + "poolId": { + "type": "string" + }, + "routeProfileKey": { + "type": "string" + }, + "sampleCount": { + "type": "integer" + }, + "sampledAt": { + "type": "string" + }, + "state": { + "type": "string" + }, + "successRate": { + "type": "number" + }, + "uploadBytesPerSecond": { + "type": "number" + } + } + }, "httpapi.AdminTaskListResponse": { "type": "object", "properties": { @@ -17072,6 +17273,9 @@ "podUid": { "type": "string" }, + "poolId": { + "type": "string" + }, "preparingTasks": { "type": "integer" }, diff --git a/apps/api/docs/swagger.yaml b/apps/api/docs/swagger.yaml index 93c914b..f250c1b 100644 --- a/apps/api/docs/swagger.yaml +++ b/apps/api/docs/swagger.yaml @@ -72,6 +72,119 @@ definitions: $ref: '#/definitions/store.AccessRule' type: array type: object + httpapi.AdminExecutionPool: + properties: + capabilities: + additionalProperties: {} + type: object + capacity: + $ref: '#/definitions/httpapi.AdminExecutionPoolCapacity' + labels: + additionalProperties: + type: string + type: object + poolId: + type: string + state: + type: string + workers: + items: + $ref: '#/definitions/httpapi.AdminExecutionPoolWorker' + type: array + type: object + httpapi.AdminExecutionPoolCapacity: + properties: + activeTasks: + type: integer + allocated: + type: integer + critical: + type: boolean + estimatedQueueWaitMs: + type: integer + heavyCapacity: + type: integer + heavyTasks: + type: integer + queueDepth: + type: integer + resourceHeadroom: + type: number + safeCapacity: + type: integer + sampledAt: + type: string + stability: + type: number + workerCount: + type: integer + type: object + httpapi.AdminExecutionPoolResponse: + properties: + capturedAt: + type: string + pools: + items: + $ref: '#/definitions/httpapi.AdminExecutionPool' + type: array + routeHealth: + items: + $ref: '#/definitions/httpapi.AdminRouteHealth' + type: array + type: object + httpapi.AdminExecutionPoolWorker: + properties: + activeTasks: + type: integer + allocated: + type: integer + capabilities: + additionalProperties: {} + type: object + heartbeatAt: + type: string + instanceId: + type: string + pressureState: + type: string + protocolVersion: + type: string + revision: + type: string + safeCapacity: + type: integer + workerId: + type: string + type: object + httpapi.AdminRouteHealth: + properties: + connectTlsP95Ms: + type: integer + consecutiveFailures: + type: integer + consecutiveSuccesses: + type: integer + expiresAt: + type: string + firstByteP95Ms: + type: integer + jitterP95Ms: + type: integer + poolId: + type: string + routeProfileKey: + type: string + sampleCount: + type: integer + sampledAt: + type: string + state: + type: string + successRate: + type: number + uploadBytesPerSecond: + type: number + type: object httpapi.AdminTaskListResponse: properties: items: @@ -4449,6 +4562,8 @@ definitions: type: string podUid: type: string + poolId: + type: string preparingTasks: type: integer pressureReason: @@ -5816,6 +5931,24 @@ paths: summary: 重试计费结算 tags: - billing + /api/admin/runtime/execution-pools: + get: + produces: + - application/json + responses: + "200": + description: OK + schema: + $ref: '#/definitions/httpapi.AdminExecutionPoolResponse' + "500": + description: Internal Server Error + schema: + $ref: '#/definitions/httpapi.ErrorEnvelope' + security: + - BearerAuth: [] + summary: 查询逻辑执行池、Worker 容量和上游路由健康状态 + tags: + - Admin Runtime /api/admin/runtime/model-rate-limits: get: description: 管理端查看平台模型维度的限流和冷却状态。 diff --git a/apps/api/internal/capacitycontroller/adapter_factory.go b/apps/api/internal/capacitycontroller/adapter_factory.go new file mode 100644 index 0000000..63af733 --- /dev/null +++ b/apps/api/internal/capacitycontroller/adapter_factory.go @@ -0,0 +1,40 @@ +package capacitycontroller + +import ( + "errors" + "os" + "strings" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/config" +) + +func NewConfiguredAdapter(cfg config.Config) (OrchestratorAdapter, error) { + switch strings.ToLower(strings.TrimSpace(cfg.CapacityOrchestratorAdapter)) { + case "static": + pools, err := cfg.CapacityPools() + if err != nil { + return nil, err + } + initial := make(map[string]int, len(pools)) + for _, pool := range pools { + initial[pool.ID] = pool.BootstrapReplicas + } + return NewStaticAdapter(initial), nil + case "", "kubernetes": + return NewKubernetesClient(KubernetesConfig{ + Namespace: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_NAMESPACE", adapterEnv("POD_NAMESPACE", "easyai")), + APIServer: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_API_SERVER", "https://kubernetes.default.svc"), + TokenFile: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_TOKEN_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/token"), + CAFile: adapterEnv("AI_GATEWAY_CAPACITY_CONTROLLER_CA_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"), + }) + default: + return nil, errors.New("unsupported capacity orchestrator adapter") + } +} + +func adapterEnv(name, fallback string) string { + if value := strings.TrimSpace(os.Getenv(name)); value != "" { + return value + } + return fallback +} diff --git a/apps/api/internal/capacitycontroller/controller.go b/apps/api/internal/capacitycontroller/controller.go index 7d9877e..f33ed40 100644 --- a/apps/api/internal/capacitycontroller/controller.go +++ b/apps/api/internal/capacitycontroller/controller.go @@ -12,6 +12,7 @@ import ( "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/config" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" "github.com/jackc/pgx/v5" ) @@ -25,16 +26,18 @@ const ( type capacityStore interface { TryAcquireCapacityControllerLeadership(context.Context) (store.Leadership, bool, error) WorkerQueueRuntime(context.Context) (store.WorkerQueueRuntime, error) + ListPoolQueueRuntime(context.Context) ([]store.PoolQueueRuntime, error) ListWorkerInstanceRuntime(context.Context) ([]store.WorkerInstanceRuntime, error) CapacityDatabaseHealth(context.Context) (store.CapacityDatabaseHealth, error) MarkWorkerDraining(context.Context, string) error ReactivateWorkerInstance(context.Context, string) error + PublishDesiredCapacity(context.Context, executionpool.DesiredCapacity) error } -type capacityKubernetes interface { - SiteState(context.Context, string) (KubernetesSiteState, error) - ScaleWorkerDeployment(context.Context, string, int) error - SetPodDeletionCost(context.Context, string, int) error +type OrchestratorAdapter interface { + PoolState(context.Context, string, string) (PoolInfrastructureState, error) + ScalePool(context.Context, string, string, int) error + SetInstanceTerminationPriority(context.Context, string, int) error } type Status struct { @@ -49,7 +52,8 @@ type Status struct { type Controller struct { cfg config.Config store capacityStore - kubernetes capacityKubernetes + orchestrator OrchestratorAdapter + pools []config.ExecutionPoolCapacityConfig logger *slog.Logger expectedRevision string now func() time.Time @@ -62,16 +66,22 @@ type Controller struct { func New( cfg config.Config, db capacityStore, - kubernetes capacityKubernetes, + orchestrator OrchestratorAdapter, logger *slog.Logger, ) *Controller { return &Controller{ - cfg: cfg, store: db, kubernetes: kubernetes, logger: logger, + cfg: cfg, store: db, orchestrator: orchestrator, logger: logger, expectedRevision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), now: time.Now, + pools: mustCapacityPools(cfg), } } +func mustCapacityPools(cfg config.Config) []config.ExecutionPoolCapacityConfig { + pools, _ := cfg.CapacityPools() + return pools +} + func (controller *Controller) Run(ctx context.Context) { for ctx.Err() == nil { leadership, acquired, err := controller.store.TryAcquireCapacityControllerLeadership(ctx) @@ -128,6 +138,14 @@ func (controller *Controller) reconcile(ctx context.Context) error { if err != nil { return err } + poolQueues, err := controller.store.ListPoolQueueRuntime(ctx) + if err != nil { + return err + } + demandByPool := make(map[string]int, len(poolQueues)) + for _, poolQueue := range poolQueues { + demandByPool[poolQueue.PoolID] = poolQueue.Queued + poolQueue.Running + } instances, err := controller.store.ListWorkerInstanceRuntime(ctx) if err != nil { return err @@ -137,21 +155,21 @@ func (controller *Controller) reconcile(ctx context.Context) error { return err } now := controller.now() - sites := make([]SiteResources, 0, 2) - kubernetesStates := make(map[string]KubernetesSiteState, 2) - for _, site := range []string{"ningbo", "hongkong"} { - minReplicas, maxReplicas := controller.siteReplicaBounds(site) - if maxReplicas == 0 { + poolResources := make([]PoolResources, 0, len(controller.pools)) + infrastructureStates := make(map[string]PoolInfrastructureState, len(controller.pools)) + for _, pool := range controller.pools { + if pool.MaxReplicas == 0 { continue } - state, stateErr := controller.kubernetes.SiteState(ctx, site) + state, stateErr := controller.orchestrator.PoolState(ctx, pool.ID, pool.AdapterRef) if stateErr != nil { return stateErr } - kubernetesStates[site] = state - sites = append(sites, SiteResources{ - Site: site, CurrentReplicas: state.CurrentReplicas, - MinReplicas: minReplicas, MaxReplicas: maxReplicas, + infrastructureStates[pool.ID] = state + poolResources = append(poolResources, PoolResources{ + PoolID: pool.ID, CurrentReplicas: state.CurrentReplicas, + MinReplicas: pool.MinReplicas, MaxReplicas: pool.MaxReplicas, + Demand: demandByPool[pool.ID], AllocatableMemoryBytes: state.AllocatableMemoryBytes, UsedMemoryBytes: state.UsedMemoryBytes, WorkerRequestMemoryBytes: state.WorkerRequestMemoryBytes, @@ -163,8 +181,8 @@ func (controller *Controller) reconcile(ctx context.Context) error { }) } currentTotal := 0 - for _, site := range sites { - currentTotal += site.CurrentReplicas + for _, pool := range poolResources { + currentTotal += pool.CurrentReplicas } target := controller.cfg.WorkerTargetOutstandingPerReplica if target < 1 { @@ -185,7 +203,7 @@ func (controller *Controller) reconcile(ctx context.Context) error { controller.highSince = time.Time{} controller.lowSince = time.Time{} } - revisionHealthy := controller.revisionsMatch(instances, kubernetesStates) + revisionHealthy := controller.revisionsMatch(instances, infrastructureStates) scaleUpEligible := !controller.highSince.IsZero() && now.Sub(controller.highSince) >= time.Duration(controller.cfg.WorkerScaleUpWindowSeconds)*time.Second && revisionHealthy @@ -206,21 +224,33 @@ func (controller *Controller) reconcile(ctx context.Context) error { SynchronousDatabasePeers: database.SynchronousPeers, ScaleUpEligible: scaleUpEligible, ScaleDownEligible: scaleDownEligible, - Sites: sites, + Pools: poolResources, }) if !revisionHealthy && plan.FrozenReason == "" { plan.FrozenReason = "release_revision_mismatch" } + for _, poolPlan := range plan.Pools { + reason := "capacity_plan" + if plan.FrozenReason != "" { + reason = "frozen:" + plan.FrozenReason + } + if err := controller.store.PublishDesiredCapacity(ctx, executionpool.DesiredCapacity{ + PoolID: poolPlan.PoolID, Desired: poolPlan.DesiredReplicas, + Reason: reason, ValidUntil: now.Add(2 * controllerReconcileInterval), + }); err != nil { + return err + } + } if controller.cfg.WorkerAutoscalingEnabled { - for _, sitePlan := range plan.Sites { + for _, poolPlan := range plan.Pools { switch { - case sitePlan.DesiredReplicas > sitePlan.CurrentReplicas: - if err := controller.kubernetes.ScaleWorkerDeployment(ctx, sitePlan.Site, sitePlan.DesiredReplicas); err != nil { + case poolPlan.DesiredReplicas > poolPlan.CurrentReplicas: + if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.DesiredReplicas); err != nil { return err } - controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.DesiredReplicas, "scale_up") - case sitePlan.DesiredReplicas < sitePlan.CurrentReplicas: - if err := controller.reconcileScaleDown(ctx, now, sitePlan, instances); err != nil { + controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.DesiredReplicas, "scale_up") + case poolPlan.DesiredReplicas < poolPlan.CurrentReplicas: + if err := controller.reconcileScaleDown(ctx, now, poolPlan, instances); err != nil { return err } } @@ -256,21 +286,24 @@ func nodeResources(states []KubernetesNodeState) []NodeResources { func (controller *Controller) reconcileScaleDown( ctx context.Context, now time.Time, - sitePlan SitePlan, + poolPlan PoolPlan, instances []store.WorkerInstanceRuntime, ) error { for _, instance := range instances { - if instance.Site != sitePlan.Site || instance.Status != "draining" { + if workerPoolID(instance) != poolPlan.PoolID || instance.Status != "draining" { continue } if instance.RunningTasks == 0 && instance.ActiveLeases == 0 { - if err := controller.kubernetes.SetPodDeletionCost(ctx, instance.PodName, controllerDeletionCost); err != nil { + if strings.TrimSpace(instance.OrchestratorInstanceRef) == "" { + return errors.New("worker orchestrator instance reference is required for scale down") + } + if err := controller.orchestrator.SetInstanceTerminationPriority(ctx, instance.OrchestratorInstanceRef, controllerDeletionCost); err != nil { return err } - if err := controller.kubernetes.ScaleWorkerDeployment(ctx, sitePlan.Site, sitePlan.CurrentReplicas-1); err != nil { + if err := controller.orchestrator.ScalePool(ctx, poolPlan.PoolID, controller.adapterRef(poolPlan.PoolID), poolPlan.CurrentReplicas-1); err != nil { return err } - controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas-1, "drained_scale_down") + controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas-1, "drained_scale_down") return nil } if instance.DrainingAt != nil && @@ -278,14 +311,14 @@ func (controller *Controller) reconcileScaleDown( if err := controller.store.ReactivateWorkerInstance(ctx, instance.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) { return err } - controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas, "drain_timeout") + controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_timeout") } return nil } var candidate *store.WorkerInstanceRuntime for index := range instances { instance := &instances[index] - if instance.Site != sitePlan.Site || instance.Status != "active" { + if workerPoolID(*instance) != poolPlan.PoolID || instance.Status != "active" { continue } if candidate == nil || @@ -299,24 +332,29 @@ func (controller *Controller) reconcileScaleDown( if err := controller.store.MarkWorkerDraining(ctx, candidate.InstanceID); err != nil && !errors.Is(err, pgx.ErrNoRows) { return err } - controller.observeScale(sitePlan.Site, sitePlan.CurrentReplicas, sitePlan.CurrentReplicas, "drain_started") + controller.observeScale(poolPlan.PoolID, poolPlan.CurrentReplicas, poolPlan.CurrentReplicas, "drain_started") return nil } -func (controller *Controller) siteReplicaBounds(site string) (int, int) { - switch site { - case "ningbo": - return controller.cfg.WorkerMinReplicasNingbo, controller.cfg.WorkerMaxReplicasNingbo - case "hongkong": - return controller.cfg.WorkerMinReplicasHongkong, controller.cfg.WorkerMaxReplicasHongkong - default: - return 0, 0 +func workerPoolID(instance store.WorkerInstanceRuntime) string { + if strings.TrimSpace(instance.PoolID) != "" { + return instance.PoolID } + return instance.Site +} + +func (controller *Controller) adapterRef(poolID string) string { + for _, pool := range controller.pools { + if pool.ID == poolID { + return pool.AdapterRef + } + } + return poolID } func (controller *Controller) revisionsMatch( instances []store.WorkerInstanceRuntime, - states map[string]KubernetesSiteState, + states map[string]PoolInfrastructureState, ) bool { if controller.expectedRevision == "" { return true diff --git a/apps/api/internal/capacitycontroller/kubernetes.go b/apps/api/internal/capacitycontroller/kubernetes.go index caa001a..291e479 100644 --- a/apps/api/internal/capacitycontroller/kubernetes.go +++ b/apps/api/internal/capacitycontroller/kubernetes.go @@ -32,8 +32,8 @@ type KubernetesClient struct { client *http.Client } -type KubernetesSiteState struct { - Site string +type PoolInfrastructureState struct { + PoolID string NodeName string CurrentReplicas int Revision string @@ -45,10 +45,10 @@ type KubernetesSiteState struct { UsedMilliCPU int64 WorkerRequestMilliCPU int64 MemoryPressure bool - Nodes []KubernetesNodeState + Nodes []PoolNodeState } -type KubernetesNodeState struct { +type PoolNodeState struct { NodeName string CurrentReplicas int AllocatableMemoryBytes int64 @@ -60,6 +60,25 @@ type KubernetesNodeState struct { MemoryPressure bool } +type KubernetesSiteState = PoolInfrastructureState +type KubernetesNodeState = PoolNodeState + +var _ OrchestratorAdapter = (*KubernetesClient)(nil) + +func (client *KubernetesClient) PoolState(ctx context.Context, poolID, adapterRef string) (PoolInfrastructureState, error) { + state, err := client.SiteState(ctx, adapterRef) + state.PoolID = poolID + return state, err +} + +func (client *KubernetesClient) ScalePool(ctx context.Context, _, adapterRef string, replicas int) error { + return client.ScaleWorkerDeployment(ctx, adapterRef, replicas) +} + +func (client *KubernetesClient) SetInstanceTerminationPriority(ctx context.Context, instanceRef string, priority int) error { + return client.SetPodDeletionCost(ctx, instanceRef, priority) +} + func NewKubernetesClient(config KubernetesConfig) (*KubernetesClient, error) { if strings.TrimSpace(config.Namespace) == "" { return nil, errors.New("capacity controller Kubernetes namespace is required") @@ -144,7 +163,7 @@ func (client *KubernetesClient) SiteState(ctx context.Context, site string) (Kub if err := client.getJSON(ctx, deploymentPath, &deployment); err != nil { return KubernetesSiteState{}, err } - state := KubernetesSiteState{Site: site, CurrentReplicas: deployment.Spec.Replicas} + state := KubernetesSiteState{PoolID: site, CurrentReplicas: deployment.Spec.Replicas} for _, container := range deployment.Spec.Template.Spec.Containers { if container.Name != "worker" { continue diff --git a/apps/api/internal/capacitycontroller/kubernetes_test.go b/apps/api/internal/capacitycontroller/kubernetes_test.go index 6db6a5f..8a11f09 100644 --- a/apps/api/internal/capacitycontroller/kubernetes_test.go +++ b/apps/api/internal/capacitycontroller/kubernetes_test.go @@ -129,7 +129,7 @@ func TestKubernetesClientAllowsDisabledSiteWithoutWorkerNodes(t *testing.T) { if err != nil { t.Fatal(err) } - if state.Site != "ningbo" || state.CurrentReplicas != 0 || len(state.Nodes) != 0 { + if state.PoolID != "ningbo" || state.CurrentReplicas != 0 || len(state.Nodes) != 0 { t.Fatalf("site state=%+v", state) } } diff --git a/apps/api/internal/capacitycontroller/planner.go b/apps/api/internal/capacitycontroller/planner.go index 4ba103c..7636cd1 100644 --- a/apps/api/internal/capacitycontroller/planner.go +++ b/apps/api/internal/capacitycontroller/planner.go @@ -5,8 +5,9 @@ import ( "sort" ) -type SiteResources struct { - Site string +type PoolResources struct { + PoolID string + Demand int CurrentReplicas int MinReplicas int MaxReplicas int @@ -47,11 +48,12 @@ type PlanInput struct { SynchronousDatabasePeers int ScaleUpEligible bool ScaleDownEligible bool - Sites []SiteResources + Pools []PoolResources } -type SitePlan struct { - Site string `json:"site"` +type PoolPlan struct { + PoolID string `json:"poolId"` + Demand int `json:"demand"` CurrentReplicas int `json:"currentReplicas"` DesiredReplicas int `json:"desiredReplicas"` ResourceMax int `json:"resourceMax"` @@ -64,7 +66,7 @@ type Plan struct { DesiredTotal int `json:"desiredTotal"` CurrentTotal int `json:"currentTotal"` FrozenReason string `json:"frozenReason,omitempty"` - Sites []SitePlan `json:"sites"` + Pools []PoolPlan `json:"pools"` } func CalculatePlan(input PlanInput) Plan { @@ -88,33 +90,34 @@ func CalculatePlan(input PlanInput) Plan { plan := Plan{RawDesired: rawDesired} minTotal := 0 resourceMaxTotal := 0 - for _, site := range input.Sites { - for _, node := range site.Nodes { + for _, pool := range input.Pools { + for _, node := range pool.Nodes { if node.MemoryPressure { - site.MemoryPressure = true + pool.MemoryPressure = true } } - current := max(site.CurrentReplicas, 0) - minReplicas := max(site.MinReplicas, 0) - configMax := max(site.MaxReplicas, minReplicas) - resourceMax, memoryPercent, cpuPercent := siteResourceMaximum( - site, + current := max(pool.CurrentReplicas, 0) + minReplicas := max(pool.MinReplicas, 0) + configMax := max(pool.MaxReplicas, minReplicas) + resourceMax, memoryPercent, cpuPercent := poolResourceMaximum( + pool, input.MemoryTargetPercent, input.CPUTargetPercent, ) resourceMax = min(resourceMax, configMax) if resourceMax < minReplicas && plan.FrozenReason == "" { - plan.FrozenReason = "site_resource_budget" + plan.FrozenReason = "pool_resource_budget" } resourceMax = max(resourceMax, minReplicas) - plan.Sites = append(plan.Sites, SitePlan{ - Site: site.Site, CurrentReplicas: current, DesiredReplicas: minReplicas, + plan.Pools = append(plan.Pools, PoolPlan{ + PoolID: pool.PoolID, CurrentReplicas: current, DesiredReplicas: minReplicas, + Demand: pool.Demand, ResourceMax: resourceMax, MemoryPercent: memoryPercent, CPUPercent: cpuPercent, }) plan.CurrentTotal += current minTotal += minReplicas resourceMaxTotal += resourceMax - if site.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) { + if pool.MemoryPressure || memoryPercent >= float64(input.MemoryHardPercent) { plan.FrozenReason = "node_memory_pressure" } } @@ -123,10 +126,10 @@ func CalculatePlan(input PlanInput) Plan { (input.DatabaseConnectionBudget-input.NonWorkerConnectionBudget)/input.WorkerDatabasePoolMax, 0, ) - capSiteResourceMaxima(plan.Sites, workerReplicaBudget) + capPoolResourceMaxima(plan.Pools, workerReplicaBudget) resourceMaxTotal = 0 - for _, site := range plan.Sites { - resourceMaxTotal += site.ResourceMax + for _, pool := range plan.Pools { + resourceMaxTotal += pool.ResourceMax } if workerReplicaBudget < minTotal && plan.FrozenReason == "" { plan.FrozenReason = "database_connection_budget" @@ -167,47 +170,47 @@ func CalculatePlan(input PlanInput) Plan { } desired = max(desired, minTotal) plan.DesiredTotal = desired - distributeDesiredReplicas(plan.Sites, desired) + distributeDesiredReplicas(plan.Pools, desired) return plan } -func capSiteResourceMaxima(sites []SitePlan, total int) { +func capPoolResourceMaxima(pools []PoolPlan, total int) { minimum := 0 - original := make(map[string]int, len(sites)) - for index := range sites { - original[sites[index].Site] = sites[index].ResourceMax - sites[index].ResourceMax = sites[index].DesiredReplicas - minimum += sites[index].ResourceMax + original := make(map[string]int, len(pools)) + for index := range pools { + original[pools[index].PoolID] = pools[index].ResourceMax + pools[index].ResourceMax = pools[index].DesiredReplicas + minimum += pools[index].ResourceMax } target := max(total, minimum) assigned := minimum for assigned < target { - sort.SliceStable(sites, func(left, right int) bool { - leftRoom := original[sites[left].Site] - sites[left].ResourceMax - rightRoom := original[sites[right].Site] - sites[right].ResourceMax + sort.SliceStable(pools, func(left, right int) bool { + leftRoom := original[pools[left].PoolID] - pools[left].ResourceMax + rightRoom := original[pools[right].PoolID] - pools[right].ResourceMax if leftRoom != rightRoom { return leftRoom > rightRoom } - return sites[left].Site < sites[right].Site + return pools[left].PoolID < pools[right].PoolID }) - if original[sites[0].Site] <= sites[0].ResourceMax { + if original[pools[0].PoolID] <= pools[0].ResourceMax { break } - sites[0].ResourceMax++ + pools[0].ResourceMax++ assigned++ } } -func siteResourceMaximum(site SiteResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) { - if len(site.Nodes) > 0 { - if site.WorkerRequestMemoryBytes <= 0 || site.WorkerRequestMilliCPU <= 0 { +func poolResourceMaximum(pool PoolResources, memoryTargetPercent int, cpuTargetPercent int) (int, float64, float64) { + if len(pool.Nodes) > 0 { + if pool.WorkerRequestMemoryBytes <= 0 || pool.WorkerRequestMilliCPU <= 0 { return 0, 0, 0 } memoryMax := 0 cpuMax := 0 memoryPercent := float64(0) cpuPercent := float64(0) - for _, node := range site.Nodes { + for _, node := range pool.Nodes { nodeMemoryPercent := usagePercent(node.UsedMemoryBytes, node.AllocatableMemoryBytes) nodeCPUPercent := usagePercent(node.UsedMilliCPU, node.AllocatableMilliCPU) memoryPercent = max(memoryPercent, nodeMemoryPercent) @@ -217,60 +220,65 @@ func siteResourceMaximum(site SiteResources, memoryTargetPercent int, cpuTargetP 0, ) memoryBudget := node.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorkerMemory - memoryMax += int(max(memoryBudget, 0) / site.WorkerRequestMemoryBytes) + memoryMax += int(max(memoryBudget, 0) / pool.WorkerRequestMemoryBytes) nonWorkerCPU := max( node.UsedMilliCPU-node.WorkerUsedMilliCPU, 0, ) cpuBudget := node.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorkerCPU - cpuMax += int(max(cpuBudget, 0) / site.WorkerRequestMilliCPU) + cpuMax += int(max(cpuBudget, 0) / pool.WorkerRequestMilliCPU) } - return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent + return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent } - memoryPercent := usagePercent(site.UsedMemoryBytes, site.AllocatableMemoryBytes) - cpuPercent := usagePercent(site.UsedMilliCPU, site.AllocatableMilliCPU) - memoryMax := site.MaxReplicas - if site.AllocatableMemoryBytes > 0 && site.WorkerRequestMemoryBytes > 0 { - nonWorker := max(site.UsedMemoryBytes-int64(site.CurrentReplicas)*site.WorkerRequestMemoryBytes, 0) - budget := site.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker - memoryMax = int(max(budget, 0) / site.WorkerRequestMemoryBytes) + memoryPercent := usagePercent(pool.UsedMemoryBytes, pool.AllocatableMemoryBytes) + cpuPercent := usagePercent(pool.UsedMilliCPU, pool.AllocatableMilliCPU) + memoryMax := pool.MaxReplicas + if pool.AllocatableMemoryBytes > 0 && pool.WorkerRequestMemoryBytes > 0 { + nonWorker := max(pool.UsedMemoryBytes-int64(pool.CurrentReplicas)*pool.WorkerRequestMemoryBytes, 0) + budget := pool.AllocatableMemoryBytes*int64(memoryTargetPercent)/100 - nonWorker + memoryMax = int(max(budget, 0) / pool.WorkerRequestMemoryBytes) } - cpuMax := site.MaxReplicas - if site.AllocatableMilliCPU > 0 && site.WorkerRequestMilliCPU > 0 { - nonWorker := max(site.UsedMilliCPU-int64(site.CurrentReplicas)*site.WorkerRequestMilliCPU, 0) - budget := site.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker - cpuMax = int(max(budget, 0) / site.WorkerRequestMilliCPU) + cpuMax := pool.MaxReplicas + if pool.AllocatableMilliCPU > 0 && pool.WorkerRequestMilliCPU > 0 { + nonWorker := max(pool.UsedMilliCPU-int64(pool.CurrentReplicas)*pool.WorkerRequestMilliCPU, 0) + budget := pool.AllocatableMilliCPU*int64(cpuTargetPercent)/100 - nonWorker + cpuMax = int(max(budget, 0) / pool.WorkerRequestMilliCPU) } - return min(site.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent + return min(pool.MaxReplicas, min(memoryMax, cpuMax)), memoryPercent, cpuPercent } -func distributeDesiredReplicas(sites []SitePlan, desired int) { - if len(sites) == 0 { +func distributeDesiredReplicas(pools []PoolPlan, desired int) { + if len(pools) == 0 { return } assigned := 0 - for index := range sites { - assigned += sites[index].DesiredReplicas + for index := range pools { + assigned += pools[index].DesiredReplicas } for assigned < desired { - sort.SliceStable(sites, func(left, right int) bool { - leftRoom := sites[left].ResourceMax - sites[left].DesiredReplicas - rightRoom := sites[right].ResourceMax - sites[right].DesiredReplicas + sort.SliceStable(pools, func(left, right int) bool { + leftDemandPerReplica := float64(max(pools[left].Demand, 0)) / float64(max(pools[left].DesiredReplicas, 1)) + rightDemandPerReplica := float64(max(pools[right].Demand, 0)) / float64(max(pools[right].DesiredReplicas, 1)) + if leftDemandPerReplica != rightDemandPerReplica { + return leftDemandPerReplica > rightDemandPerReplica + } + leftRoom := pools[left].ResourceMax - pools[left].DesiredReplicas + rightRoom := pools[right].ResourceMax - pools[right].DesiredReplicas if leftRoom != rightRoom { return leftRoom > rightRoom } - if sites[left].DesiredReplicas != sites[right].DesiredReplicas { - return sites[left].DesiredReplicas < sites[right].DesiredReplicas + if pools[left].DesiredReplicas != pools[right].DesiredReplicas { + return pools[left].DesiredReplicas < pools[right].DesiredReplicas } - return sites[left].Site < sites[right].Site + return pools[left].PoolID < pools[right].PoolID }) - if sites[0].DesiredReplicas >= sites[0].ResourceMax { + if pools[0].DesiredReplicas >= pools[0].ResourceMax { break } - sites[0].DesiredReplicas++ + pools[0].DesiredReplicas++ assigned++ } - sort.Slice(sites, func(left, right int) bool { return sites[left].Site < sites[right].Site }) + sort.Slice(pools, func(left, right int) bool { return pools[left].PoolID < pools[right].PoolID }) } func usagePercent(used int64, allocatable int64) float64 { diff --git a/apps/api/internal/capacitycontroller/planner_test.go b/apps/api/internal/capacitycontroller/planner_test.go index 64b0c17..00b3ad4 100644 --- a/apps/api/internal/capacitycontroller/planner_test.go +++ b/apps/api/internal/capacitycontroller/planner_test.go @@ -12,15 +12,15 @@ func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) { MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70, DatabaseConnections: 80, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1, ScaleUpEligible: true, - Sites: []SiteResources{ + Pools: []PoolResources{ { - Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, + PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 6 << 30, WorkerRequestMemoryBytes: 1 << 30, AllocatableMilliCPU: 4000, UsedMilliCPU: 2000, WorkerRequestMilliCPU: 500, }, { - Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, + PoolID: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30, WorkerRequestMemoryBytes: 1 << 30, AllocatableMilliCPU: 4000, UsedMilliCPU: 500, WorkerRequestMilliCPU: 500, @@ -34,8 +34,8 @@ func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) { if plan.DesiredTotal != 4 { t.Fatalf("desired total=%d, want step-limited 4: %+v", plan.DesiredTotal, plan) } - if plan.Sites[0].Site != "hongkong" || plan.Sites[0].DesiredReplicas != 3 { - t.Fatalf("site allocation=%+v, want hongkong=3 ningbo=1", plan.Sites) + if plan.Pools[0].PoolID != "hongkong" || plan.Pools[0].DesiredReplicas != 3 { + t.Fatalf("pool allocation=%+v, want hongkong=3 ningbo=1", plan.Pools) } } @@ -44,8 +44,8 @@ func TestCalculatePlanFreezesScaleUpOnHardMemoryPressure(t *testing.T) { Queued: 200, InstanceSlots: 24, DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1, ScaleUpEligible: true, - Sites: []SiteResources{{ - Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, + Pools: []PoolResources{{ + PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 7 << 30, WorkerRequestMemoryBytes: 1 << 30, }}, @@ -60,14 +60,14 @@ func TestCalculatePlanReportsConfiguredMinimumOutsideResourceBudget(t *testing.T Queued: 200, InstanceSlots: 24, DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1, ScaleUpEligible: true, - Sites: []SiteResources{{ - Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, + Pools: []PoolResources{{ + PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4, AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 13 << 29, WorkerRequestMemoryBytes: 2 << 30, AllocatableMilliCPU: 4000, UsedMilliCPU: 2800, WorkerRequestMilliCPU: 500, }}, }) - if plan.DesiredTotal != 1 || plan.FrozenReason != "site_resource_budget" { + if plan.DesiredTotal != 1 || plan.FrozenReason != "pool_resource_budget" { t.Fatalf("plan=%+v, want the existing minimum preserved and expansion frozen", plan) } } @@ -77,9 +77,9 @@ func TestCalculatePlanWaitsForSafeScaleDownWindow(t *testing.T) { InstanceSlots: 24, DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1, ScaleDownEligible: false, - Sites: []SiteResources{ - {Site: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4}, - {Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4}, + Pools: []PoolResources{ + {PoolID: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4}, + {PoolID: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4}, }, } plan := CalculatePlan(input) @@ -99,8 +99,8 @@ func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) { MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70, DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 20, SynchronousDatabasePeers: 1, ScaleUpEligible: true, - Sites: []SiteResources{{ - Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8, + Pools: []PoolResources{{ + PoolID: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8, WorkerRequestMemoryBytes: 1 << 30, WorkerRequestMilliCPU: 500, Nodes: []NodeResources{ { @@ -116,8 +116,8 @@ func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) { }, }}, }) - if plan.Sites[0].ResourceMax != 7 { - t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Sites[0].ResourceMax, plan) + if plan.Pools[0].ResourceMax != 7 { + t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Pools[0].ResourceMax, plan) } if plan.DesiredTotal != 4 { t.Fatalf("desired=%d, want two-wave step from 2 to 4", plan.DesiredTotal) @@ -130,17 +130,17 @@ func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) { DatabaseConnections: 20, DatabaseConnectionBudget: 150, NonWorkerConnectionBudget: 72, WorkerDatabasePoolMax: 32, SynchronousDatabasePeers: 1, ScaleUpEligible: true, - Sites: []SiteResources{ - {Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4}, - {Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4}, + Pools: []PoolResources{ + {PoolID: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4}, + {PoolID: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4}, }, }) if plan.DesiredTotal != 2 { t.Fatalf("desired=%d, want the current 1+1 database-safe topology: %+v", plan.DesiredTotal, plan) } resourceMax := 0 - for _, site := range plan.Sites { - resourceMax += site.ResourceMax + for _, pool := range plan.Pools { + resourceMax += pool.ResourceMax } if resourceMax != 2 { t.Fatalf("resource maximum=%d, want floor((150-72)/32)=2: %+v", resourceMax, plan) @@ -150,13 +150,13 @@ func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) { func TestPlanJSONUsesAcceptanceContractFieldNames(t *testing.T) { payload, err := json.Marshal(Plan{ RawDesired: 3, - Sites: []SitePlan{{Site: "hongkong", ResourceMax: 2}}, + Pools: []PoolPlan{{PoolID: "hongkong", ResourceMax: 2}}, }) if err != nil { t.Fatal(err) } text := string(payload) - for _, field := range []string{`"rawDesired":3`, `"sites"`, `"site":"hongkong"`, `"resourceMax":2`} { + for _, field := range []string{`"rawDesired":3`, `"pools"`, `"poolId":"hongkong"`, `"resourceMax":2`} { if !strings.Contains(text, field) { t.Fatalf("plan JSON %s does not contain %s", text, field) } diff --git a/apps/api/internal/capacitycontroller/static.go b/apps/api/internal/capacitycontroller/static.go new file mode 100644 index 0000000..dc8c2f0 --- /dev/null +++ b/apps/api/internal/capacitycontroller/static.go @@ -0,0 +1,42 @@ +package capacitycontroller + +import ( + "context" + "sync" +) + +// StaticAdapter proves that capacity coordination does not require Kubernetes. +// It reports configured fixed replicas and intentionally ignores scale writes. +type StaticAdapter struct { + mu sync.RWMutex + replicas map[string]int +} + +var _ OrchestratorAdapter = (*StaticAdapter)(nil) + +func NewStaticAdapter(initial map[string]int) *StaticAdapter { + copyOfInitial := make(map[string]int, len(initial)) + for poolID, replicas := range initial { + copyOfInitial[poolID] = replicas + } + return &StaticAdapter{replicas: copyOfInitial} +} + +func (adapter *StaticAdapter) PoolState(_ context.Context, poolID, _ string) (PoolInfrastructureState, error) { + adapter.mu.RLock() + replicas := adapter.replicas[poolID] + adapter.mu.RUnlock() + return PoolInfrastructureState{ + PoolID: poolID, CurrentReplicas: replicas, + AllocatableMemoryBytes: 1 << 60, WorkerRequestMemoryBytes: 1, + AllocatableMilliCPU: 1 << 50, WorkerRequestMilliCPU: 1, + }, nil +} + +func (adapter *StaticAdapter) ScalePool(context.Context, string, string, int) error { + return nil +} + +func (adapter *StaticAdapter) SetInstanceTerminationPriority(context.Context, string, int) error { + return nil +} diff --git a/apps/api/internal/capacitycontroller/static_test.go b/apps/api/internal/capacitycontroller/static_test.go new file mode 100644 index 0000000..61b0f88 --- /dev/null +++ b/apps/api/internal/capacitycontroller/static_test.go @@ -0,0 +1,24 @@ +package capacitycontroller + +import ( + "context" + "testing" +) + +func TestStaticAdapterDoesNotRequireKubernetes(t *testing.T) { + adapter := NewStaticAdapter(map[string]int{"pool-a": 2}) + state, err := adapter.PoolState(context.Background(), "pool-a", "ignored") + if err != nil { + t.Fatal(err) + } + if state.PoolID != "pool-a" || state.CurrentReplicas != 2 { + t.Fatalf("state=%+v", state) + } + if err := adapter.ScalePool(context.Background(), "pool-a", "ignored", 4); err != nil { + t.Fatal(err) + } + state, _ = adapter.PoolState(context.Background(), "pool-a", "ignored") + if state.CurrentReplicas != 2 { + t.Fatalf("static adapter scaled to %d", state.CurrentReplicas) + } +} diff --git a/apps/api/internal/config/config.go b/apps/api/internal/config/config.go index f72ea81..5079554 100644 --- a/apps/api/internal/config/config.go +++ b/apps/api/internal/config/config.go @@ -1,6 +1,7 @@ package config import ( + "encoding/json" "errors" "log/slog" "net/url" @@ -84,13 +85,25 @@ type Config struct { AsyncAdmissionMicrobatchSize int AsyncAdmissionDispatcherEnabled bool AsyncAdmissionDispatcherConfigured bool + RoutingMode string + ExecutionPoolID string + ExecutionPoolLabels string + WorkerID string + WorkerAdvertiseEndpoint string + WorkerOrchestratorInstanceRef string + WorkerEndpointAllowedSuffixes string + WorkerEndpointAllowPrivate bool + WorkerExecutionSecret string + WorkerExecutionCAFile string + WorkerExecutionCertFile string + WorkerExecutionKeyFile string + RouteProbeEnabled bool + RouteProbeHotIntervalSeconds int + RouteProbeColdIntervalSeconds int + RouteProbeTimeoutMS int WorkerAutoscalingEnabled bool - WorkerReplicasNingbo int - WorkerReplicasHongkong int - WorkerMinReplicasNingbo int - WorkerMinReplicasHongkong int - WorkerMaxReplicasNingbo int - WorkerMaxReplicasHongkong int + CapacityOrchestratorAdapter string + CapacityPoolsJSON string WorkerTargetOutstandingPerReplica int WorkerScaleUpWindowSeconds int WorkerScaleDownStabilizationSeconds int @@ -101,10 +114,6 @@ type Config struct { PostgresConnectionBudget int PostgresNonWorkerConnectionBudget int WorkerDatabaseMaxConns int - CapacityControllerNamespace string - CapacityControllerAPIServer string - CapacityControllerTokenFile string - CapacityControllerCAFile string } func Load() Config { @@ -198,13 +207,25 @@ func Load() Config { AsyncAdmissionDispatcherConfigured: envValue( "AI_GATEWAY_ASYNC_ADMISSION_DISPATCHER_ENABLED", ) != "", + RoutingMode: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_ROUTING_MODE", "legacy"))), + ExecutionPoolID: strings.TrimSpace(env("AI_GATEWAY_EXECUTION_POOL_ID", "legacy-default")), + ExecutionPoolLabels: strings.TrimSpace(env("AI_GATEWAY_EXECUTION_POOL_LABELS", "{}")), + WorkerID: strings.TrimSpace(env("AI_GATEWAY_WORKER_ID", "")), + WorkerAdvertiseEndpoint: workerAdvertiseEndpoint(), + WorkerOrchestratorInstanceRef: strings.TrimSpace(env("AI_GATEWAY_ORCHESTRATOR_INSTANCE_REF", "")), + WorkerEndpointAllowedSuffixes: strings.TrimSpace(env("AI_GATEWAY_WORKER_ENDPOINT_ALLOWED_SUFFIXES", "svc,cluster.local")), + WorkerEndpointAllowPrivate: env("AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE", "true") == "true", + WorkerExecutionSecret: env("AI_GATEWAY_WORKER_EXECUTION_SECRET", env("SERVER_MAIN_INTERNAL_SECRET", env("SERVER_MAIN_INTERNAL_TOKEN", ""))), + WorkerExecutionCAFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_CA_FILE", "")), + WorkerExecutionCertFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_CERT_FILE", "")), + WorkerExecutionKeyFile: strings.TrimSpace(env("AI_GATEWAY_WORKER_EXECUTION_KEY_FILE", "")), + RouteProbeEnabled: env("AI_GATEWAY_ROUTE_PROBE_ENABLED", "true") == "true", + RouteProbeHotIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS", 15), + RouteProbeColdIntervalSeconds: envIntValidated("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS", 60), + RouteProbeTimeoutMS: envIntValidated("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS", 3000), WorkerAutoscalingEnabled: env("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "false") == "true", - WorkerReplicasNingbo: envOptionalIntValidated("AI_GATEWAY_WORKER_REPLICAS_NINGBO", 1), - WorkerReplicasHongkong: envOptionalIntValidated("AI_GATEWAY_WORKER_REPLICAS_HONGKONG", 0), - WorkerMinReplicasNingbo: envOptionalIntValidated("AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO", 1), - WorkerMinReplicasHongkong: envOptionalIntValidated("AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG", 0), - WorkerMaxReplicasNingbo: envOptionalIntValidated("AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO", 1), - WorkerMaxReplicasHongkong: envOptionalIntValidated("AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG", 0), + CapacityOrchestratorAdapter: strings.ToLower(strings.TrimSpace(env("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER", "kubernetes"))), + CapacityPoolsJSON: strings.TrimSpace(env("AI_GATEWAY_CAPACITY_POOLS", "")), WorkerTargetOutstandingPerReplica: envOptionalIntValidated("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA", 0), WorkerScaleUpWindowSeconds: envIntValidated("AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS", 20), WorkerScaleDownStabilizationSeconds: envIntValidated( @@ -220,11 +241,7 @@ func Load() Config { "AI_GATEWAY_POSTGRES_NON_WORKER_CONNECTION_BUDGET", 72, ), - WorkerDatabaseMaxConns: envIntValidated("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS", 32), - CapacityControllerNamespace: env("POD_NAMESPACE", env("AI_GATEWAY_CAPACITY_CONTROLLER_NAMESPACE", "easyai")), - CapacityControllerAPIServer: env("AI_GATEWAY_CAPACITY_CONTROLLER_API_SERVER", "https://kubernetes.default.svc"), - CapacityControllerTokenFile: env("AI_GATEWAY_CAPACITY_CONTROLLER_TOKEN_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/token"), - CapacityControllerCAFile: env("AI_GATEWAY_CAPACITY_CONTROLLER_CA_FILE", "/var/run/secrets/kubernetes.io/serviceaccount/ca.crt"), + WorkerDatabaseMaxConns: envIntValidated("AI_GATEWAY_WORKER_DATABASE_MAX_CONNS", 32), } } @@ -308,20 +325,46 @@ func (c Config) Validate() error { default: return errors.New("AI_GATEWAY_WORKER_LOAD_MODE must be adaptive or legacy") } + switch strings.ToLower(strings.TrimSpace(c.RoutingMode)) { + case "", "legacy", "shadow", "enforced": + default: + return errors.New("AI_GATEWAY_ROUTING_MODE must be legacy, shadow, or enforced") + } + if strings.TrimSpace(c.ExecutionPoolID) != "" && !validExecutionPoolID(c.ExecutionPoolID) { + return errors.New("AI_GATEWAY_EXECUTION_POOL_ID is invalid") + } + if strings.TrimSpace(c.ExecutionPoolLabels) != "" && !jsonObject(c.ExecutionPoolLabels) { + return errors.New("AI_GATEWAY_EXECUTION_POOL_LABELS must be a JSON object") + } + if c.RouteProbeHotIntervalSeconds != 0 && (c.RouteProbeHotIntervalSeconds < 5 || c.RouteProbeHotIntervalSeconds > 300) { + return errors.New("AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS must be between 5 and 300") + } + if c.RouteProbeColdIntervalSeconds != 0 && (c.RouteProbeColdIntervalSeconds < c.RouteProbeHotIntervalSeconds || c.RouteProbeColdIntervalSeconds > 3600) { + return errors.New("AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS must be between the hot interval and 3600") + } + if c.RouteProbeTimeoutMS != 0 && (c.RouteProbeTimeoutMS < 250 || c.RouteProbeTimeoutMS > 10000) { + return errors.New("AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS must be between 250 and 10000") + } + if strings.EqualFold(c.RoutingMode, "enforced") && len(c.WorkerExecutionSecret) < 32 { + return errors.New("AI_GATEWAY_WORKER_EXECUTION_SECRET must be at least 32 bytes in enforced routing mode") + } + if (c.WorkerExecutionCertFile == "") != (c.WorkerExecutionKeyFile == "") { + return errors.New("AI_GATEWAY_WORKER_EXECUTION_CERT_FILE and AI_GATEWAY_WORKER_EXECUTION_KEY_FILE must be configured together") + } if c.AsyncAdmissionMicrobatchSize < 1 || c.AsyncAdmissionMicrobatchSize > 32 { return errors.New("AI_GATEWAY_ASYNC_ADMISSION_MICROBATCH_SIZE must be between 1 and 32") } - if c.WorkerReplicasNingbo < 0 || c.WorkerReplicasHongkong < 0 || - c.WorkerMinReplicasNingbo < 0 || c.WorkerMinReplicasHongkong < 0 || - c.WorkerMaxReplicasNingbo < c.WorkerMinReplicasNingbo || - c.WorkerMaxReplicasHongkong < c.WorkerMinReplicasHongkong || - c.WorkerMaxReplicasNingbo > 64 || c.WorkerMaxReplicasHongkong > 64 { - return errors.New("Worker replica configuration must be between 0 and 64 with min not greater than max") + pools, err := c.CapacityPools() + if err != nil { + return err } - if c.WorkerAutoscalingEnabled && - (c.WorkerReplicasNingbo < c.WorkerMinReplicasNingbo || c.WorkerReplicasNingbo > c.WorkerMaxReplicasNingbo || - c.WorkerReplicasHongkong < c.WorkerMinReplicasHongkong || c.WorkerReplicasHongkong > c.WorkerMaxReplicasHongkong) { - return errors.New("Worker bootstrap replicas must be within the autoscaling min/max range") + if c.RunsCapacityController() && len(pools) == 0 { + return errors.New("AI_GATEWAY_CAPACITY_POOLS must contain at least one pool for the capacity-controller role") + } + switch strings.ToLower(strings.TrimSpace(c.CapacityOrchestratorAdapter)) { + case "", "kubernetes", "static": + default: + return errors.New("AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER must be kubernetes or static") } if c.WorkerTargetOutstandingPerReplica < 0 || c.WorkerTargetOutstandingPerReplica > 100000 { return errors.New("AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA must be between 0 and 100000") @@ -503,11 +546,79 @@ func (c Config) RunsCapacityController() bool { return c.EffectiveProcessRole() == "capacity-controller" } +func (c Config) RunsRouteProber() bool { + return c.RunsAsyncExecutionWorker() && c.RouteProbeEnabled && !strings.EqualFold(c.RoutingMode, "legacy") +} + +func workerAdvertiseEndpoint() string { + return strings.TrimRight(strings.TrimSpace(os.Getenv("AI_GATEWAY_WORKER_ADVERTISE_ENDPOINT")), "/") +} + +func validExecutionPoolID(value string) bool { + value = strings.TrimSpace(value) + if value == "" || len(value) > 128 { + return false + } + for index, character := range value { + valid := (character >= 'a' && character <= 'z') || + (character >= 'A' && character <= 'Z') || + (character >= '0' && character <= '9') || + character == '.' || character == '_' || character == ':' || character == '/' || character == '-' + if !valid || (index == 0 && !((character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') || (character >= '0' && character <= '9'))) { + return false + } + } + return true +} + +func jsonObject(value string) bool { + var object map[string]any + return json.Unmarshal([]byte(value), &object) == nil && object != nil +} + type GlobalHTTPProxyStatus struct { HTTPProxy string Source string } +type ExecutionPoolCapacityConfig struct { + ID string `json:"id"` + AdapterRef string `json:"adapterRef"` + BootstrapReplicas int `json:"bootstrapReplicas"` + MinReplicas int `json:"minReplicas"` + MaxReplicas int `json:"maxReplicas"` +} + +func (c Config) CapacityPools() ([]ExecutionPoolCapacityConfig, error) { + raw := strings.TrimSpace(c.CapacityPoolsJSON) + if raw == "" { + return nil, nil + } + var pools []ExecutionPoolCapacityConfig + if err := json.Unmarshal([]byte(raw), &pools); err != nil { + return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS must be a JSON array") + } + if len(pools) == 0 { + return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS must contain at least one pool") + } + seen := make(map[string]struct{}, len(pools)) + for index := range pools { + pools[index].ID = strings.TrimSpace(pools[index].ID) + pools[index].AdapterRef = strings.TrimSpace(pools[index].AdapterRef) + if pools[index].AdapterRef == "" { + pools[index].AdapterRef = pools[index].ID + } + if !validExecutionPoolID(pools[index].ID) || pools[index].MinReplicas < 0 || pools[index].MaxReplicas < pools[index].MinReplicas || pools[index].MaxReplicas > 64 || pools[index].BootstrapReplicas < pools[index].MinReplicas || pools[index].BootstrapReplicas > pools[index].MaxReplicas { + return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS contains an invalid pool definition") + } + if _, duplicate := seen[pools[index].ID]; duplicate { + return nil, errors.New("AI_GATEWAY_CAPACITY_POOLS contains duplicate pool IDs") + } + seen[pools[index].ID] = struct{}{} + } + return pools, nil +} + func LoadGlobalHTTPProxyStatus() GlobalHTTPProxyStatus { for _, key := range []string{ "AI_GATEWAY_GLOBAL_HTTP_PROXY", diff --git a/apps/api/internal/config/config_test.go b/apps/api/internal/config/config_test.go index b613e37..c0fc3a7 100644 --- a/apps/api/internal/config/config_test.go +++ b/apps/api/internal/config/config_test.go @@ -198,42 +198,27 @@ func TestProcessRolePrecedenceAndCompatibility(t *testing.T) { func TestWorkerAutoscalingConfiguration(t *testing.T) { t.Setenv("AI_GATEWAY_WORKER_AUTOSCALING_ENABLED", "true") - t.Setenv("AI_GATEWAY_WORKER_REPLICAS_NINGBO", "1") - t.Setenv("AI_GATEWAY_WORKER_REPLICAS_HONGKONG", "2") - t.Setenv("AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO", "1") - t.Setenv("AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG", "1") - t.Setenv("AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO", "2") - t.Setenv("AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG", "3") + t.Setenv("AI_GATEWAY_CAPACITY_POOLS", `[{"id":"pool-a","bootstrapReplicas":2,"minReplicas":1,"maxReplicas":3}]`) cfg := Load() if err := cfg.Validate(); err != nil { t.Fatalf("valid autoscaling configuration was rejected: %v", err) } - if !cfg.WorkerAutoscalingEnabled || cfg.WorkerReplicasHongkong != 2 || cfg.WorkerMaxReplicasHongkong != 3 { + pools, err := cfg.CapacityPools() + if err != nil || !cfg.WorkerAutoscalingEnabled || len(pools) != 1 || pools[0].BootstrapReplicas != 2 || pools[0].MaxReplicas != 3 { t.Fatalf("autoscaling configuration=%+v", cfg) } - cfg.WorkerReplicasHongkong = 4 - if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "bootstrap") { + cfg.CapacityPoolsJSON = `[{"id":"pool-a","bootstrapReplicas":4,"minReplicas":1,"maxReplicas":3}]` + if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "invalid pool") { t.Fatalf("Validate() error=%v, want bootstrap bound failure", err) } } -func TestWorkerTopologyDefaultsToNingboSingleNode(t *testing.T) { - for _, name := range []string{ - "AI_GATEWAY_WORKER_REPLICAS_NINGBO", - "AI_GATEWAY_WORKER_REPLICAS_HONGKONG", - "AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO", - "AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG", - "AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO", - "AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG", - } { - t.Setenv(name, "") - } +func TestWorkerTopologyHasNoBuiltInSites(t *testing.T) { + t.Setenv("AI_GATEWAY_CAPACITY_POOLS", "") cfg := Load() - if cfg.WorkerReplicasNingbo != 1 || cfg.WorkerMinReplicasNingbo != 1 || cfg.WorkerMaxReplicasNingbo != 1 { - t.Fatalf("Ningbo Worker defaults=%d/%d/%d, want 1/1/1", cfg.WorkerReplicasNingbo, cfg.WorkerMinReplicasNingbo, cfg.WorkerMaxReplicasNingbo) - } - if cfg.WorkerReplicasHongkong != 0 || cfg.WorkerMinReplicasHongkong != 0 || cfg.WorkerMaxReplicasHongkong != 0 { - t.Fatalf("Hong Kong Worker defaults=%d/%d/%d, want disabled", cfg.WorkerReplicasHongkong, cfg.WorkerMinReplicasHongkong, cfg.WorkerMaxReplicasHongkong) + pools, err := cfg.CapacityPools() + if err != nil || len(pools) != 0 { + t.Fatalf("default pools=%+v error=%v, want no built-in sites", pools, err) } } @@ -405,6 +390,7 @@ func TestLoadAndValidateDirectMediaOSS(t *testing.T) { func TestCapacityControllerDoesNotRequireMediaOSSCredentials(t *testing.T) { cfg := Load() cfg.ProcessRole = "capacity-controller" + cfg.CapacityPoolsJSON = `[{"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}]` cfg.MediaOSSDirectEnabled = true cfg.MediaOSSAccessKeyID = "" cfg.MediaOSSAccessKeySecret = "" diff --git a/apps/api/internal/config/execution_pools_test.go b/apps/api/internal/config/execution_pools_test.go new file mode 100644 index 0000000..a9b38da --- /dev/null +++ b/apps/api/internal/config/execution_pools_test.go @@ -0,0 +1,27 @@ +package config + +import "testing" + +func TestCapacityPoolsAreDataDriven(t *testing.T) { + cfg := Config{CapacityPoolsJSON: `[ + {"id":"pool-a","adapterRef":"deployment-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":3}, + {"id":"pool-b","adapterRef":"deployment-b","bootstrapReplicas":0,"minReplicas":0,"maxReplicas":2} + ]`} + pools, err := cfg.CapacityPools() + if err != nil { + t.Fatal(err) + } + if len(pools) != 2 || pools[0].ID != "pool-a" || pools[1].AdapterRef != "deployment-b" { + t.Fatalf("pools=%+v", pools) + } +} + +func TestCapacityPoolsRejectDuplicates(t *testing.T) { + cfg := Config{CapacityPoolsJSON: `[ + {"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}, + {"id":"pool-a","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1} + ]`} + if _, err := cfg.CapacityPools(); err == nil { + t.Fatal("expected duplicate pool rejection") + } +} diff --git a/apps/api/internal/executionpool/identity.go b/apps/api/internal/executionpool/identity.go new file mode 100644 index 0000000..476e6c1 --- /dev/null +++ b/apps/api/internal/executionpool/identity.go @@ -0,0 +1,137 @@ +package executionpool + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "net" + "net/url" + "strings" + "time" +) + +func QueueName(poolID string) string { + poolID = strings.TrimSpace(poolID) + if poolID == "" { + return "gateway_tasks" + } + digest := sha256.Sum256([]byte(poolID)) + return "gateway_pool_" + fmt.Sprintf("%x", digest[:8]) +} + +func RouteProfileKey(provider, protocol, endpointHost, proxyMode, configRevision string) string { + value := strings.Join([]string{ + strings.ToLower(strings.TrimSpace(provider)), + strings.ToLower(strings.TrimSpace(protocol)), + strings.ToLower(strings.TrimSpace(endpointHost)), + strings.ToLower(strings.TrimSpace(proxyMode)), + strings.TrimSpace(configRevision), + }, "\x00") + digest := sha256.Sum256([]byte(value)) + return fmt.Sprintf("route_%x", digest[:16]) +} + +func EndpointHost(rawURL string) string { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil { + return "" + } + return strings.ToLower(parsed.Hostname()) +} + +func ValidateAdvertisedEndpoint(rawURL string, allowedSuffixes []string, allowPrivate bool) error { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil || parsed.Host == "" || parsed.User != nil { + return errors.New("worker endpoint must be an absolute URL without user info") + } + if parsed.Scheme != "https" && parsed.Scheme != "http" { + return errors.New("worker endpoint scheme must be http or https") + } + host := strings.ToLower(parsed.Hostname()) + if host == "" { + return errors.New("worker endpoint host is required") + } + if ip := net.ParseIP(host); ip != nil { + if allowPrivate && (ip.IsPrivate() || ip.IsLoopback()) { + return nil + } + return errors.New("worker endpoint IP is outside the configured trust boundary") + } + for _, suffix := range allowedSuffixes { + suffix = strings.ToLower(strings.TrimSpace(suffix)) + if suffix != "" && (host == strings.TrimPrefix(suffix, ".") || strings.HasSuffix(host, "."+strings.TrimPrefix(suffix, "."))) { + return nil + } + } + return errors.New("worker endpoint host is outside the configured trust boundary") +} + +type ExecutionClaims struct { + Audience string `json:"aud"` + TaskID string `json:"task_id"` + PoolID string `json:"pool_id"` + WorkerID string `json:"worker_id"` + Nonce string `json:"nonce"` + ExpiresAt int64 `json:"exp"` +} + +type TokenSigner struct { + Secret []byte + Now func() time.Time +} + +func (s TokenSigner) Sign(claims ExecutionClaims) (string, error) { + if len(s.Secret) < 32 { + return "", errors.New("execution token secret must be at least 32 bytes") + } + if claims.Audience == "" || claims.TaskID == "" || claims.PoolID == "" || claims.WorkerID == "" || claims.Nonce == "" || claims.ExpiresAt <= 0 { + return "", errors.New("execution token claims are incomplete") + } + payload, err := json.Marshal(claims) + if err != nil { + return "", err + } + encoded := base64.RawURLEncoding.EncodeToString(payload) + mac := hmac.New(sha256.New, s.Secret) + _, _ = mac.Write([]byte(encoded)) + signature := base64.RawURLEncoding.EncodeToString(mac.Sum(nil)) + return encoded + "." + signature, nil +} + +func (s TokenSigner) Verify(token, audience string) (ExecutionClaims, error) { + var claims ExecutionClaims + if len(s.Secret) < 32 { + return claims, errors.New("execution token secret must be at least 32 bytes") + } + parts := strings.Split(token, ".") + if len(parts) != 2 { + return claims, errors.New("invalid execution token") + } + signature, err := base64.RawURLEncoding.DecodeString(parts[1]) + if err != nil { + return claims, errors.New("invalid execution token signature") + } + mac := hmac.New(sha256.New, s.Secret) + _, _ = mac.Write([]byte(parts[0])) + if !hmac.Equal(signature, mac.Sum(nil)) { + return claims, errors.New("invalid execution token signature") + } + payload, err := base64.RawURLEncoding.DecodeString(parts[0]) + if err != nil || json.Unmarshal(payload, &claims) != nil { + return ExecutionClaims{}, errors.New("invalid execution token payload") + } + now := time.Now() + if s.Now != nil { + now = s.Now() + } + if claims.ExpiresAt <= now.Unix() { + return ExecutionClaims{}, errors.New("execution token expired") + } + if claims.Audience != audience { + return ExecutionClaims{}, errors.New("execution token audience mismatch") + } + return claims, nil +} diff --git a/apps/api/internal/executionpool/identity_test.go b/apps/api/internal/executionpool/identity_test.go new file mode 100644 index 0000000..456978b --- /dev/null +++ b/apps/api/internal/executionpool/identity_test.go @@ -0,0 +1,44 @@ +package executionpool + +import ( + "strings" + "testing" + "time" +) + +func TestQueueNameDoesNotExposePoolID(t *testing.T) { + name := QueueName("region-sensitive-name") + if strings.Contains(name, "sensitive") || name == QueueName("other") { + t.Fatalf("unexpected queue name %q", name) + } +} + +func TestTokenSigner(t *testing.T) { + now := time.Unix(1000, 0) + signer := TokenSigner{Secret: []byte("01234567890123456789012345678901"), Now: func() time.Time { return now }} + token, err := signer.Sign(ExecutionClaims{ + Audience: "worker", TaskID: "task", PoolID: "pool", WorkerID: "worker", Nonce: "nonce", ExpiresAt: now.Add(30 * time.Second).Unix(), + }) + if err != nil { + t.Fatal(err) + } + claims, err := signer.Verify(token, "worker") + if err != nil || claims.TaskID != "task" { + t.Fatalf("claims=%+v err=%v", claims, err) + } + if _, err := signer.Verify(token, "other"); err == nil { + t.Fatal("expected audience mismatch") + } +} + +func TestValidateAdvertisedEndpoint(t *testing.T) { + if err := ValidateAdvertisedEndpoint("http://10.0.0.2:8088", nil, true); err != nil { + t.Fatal(err) + } + if err := ValidateAdvertisedEndpoint("https://worker.internal.example", []string{"internal.example"}, false); err != nil { + t.Fatal(err) + } + if err := ValidateAdvertisedEndpoint("https://public.example", []string{"internal.example"}, false); err == nil { + t.Fatal("expected untrusted endpoint rejection") + } +} diff --git a/apps/api/internal/executionpool/preference.go b/apps/api/internal/executionpool/preference.go new file mode 100644 index 0000000..ef83db7 --- /dev/null +++ b/apps/api/internal/executionpool/preference.go @@ -0,0 +1,63 @@ +package executionpool + +import ( + "strings" + "time" +) + +type RoutePreference struct { + RouteProfileKey string + CurrentPoolID string + CurrentSince time.Time + ChallengerPoolID string + ChallengerWins int + UpdatedAt time.Time +} + +func AdvanceRoutePreference( + previous RoutePreference, + proposedPoolID string, + scores map[string]float64, + now time.Time, +) RoutePreference { + if now.IsZero() { + now = time.Now() + } + proposedPoolID = strings.TrimSpace(proposedPoolID) + next := previous + next.UpdatedAt = now + if proposedPoolID == "" { + return next + } + if _, eligible := scores[next.CurrentPoolID]; next.CurrentPoolID == "" || !eligible { + next.CurrentPoolID = proposedPoolID + next.CurrentSince = now + next.ChallengerPoolID = "" + next.ChallengerWins = 0 + return next + } + if proposedPoolID == next.CurrentPoolID { + next.ChallengerPoolID = "" + next.ChallengerWins = 0 + return next + } + if relativeImprovement(scores[proposedPoolID], scores[next.CurrentPoolID]) < defaultSwitchImprovement { + next.ChallengerPoolID = "" + next.ChallengerWins = 0 + return next + } + if next.ChallengerPoolID == proposedPoolID { + next.ChallengerWins++ + } else { + next.ChallengerPoolID = proposedPoolID + next.ChallengerWins = 1 + } + if next.ChallengerWins >= defaultBetterWindowCount && + !next.CurrentSince.IsZero() && now.Sub(next.CurrentSince) >= defaultMinimumDwell { + next.CurrentPoolID = proposedPoolID + next.CurrentSince = now + next.ChallengerPoolID = "" + next.ChallengerWins = 0 + } + return next +} diff --git a/apps/api/internal/executionpool/preference_test.go b/apps/api/internal/executionpool/preference_test.go new file mode 100644 index 0000000..48e18b8 --- /dev/null +++ b/apps/api/internal/executionpool/preference_test.go @@ -0,0 +1,32 @@ +package executionpool + +import ( + "testing" + "time" +) + +func TestAdvanceRoutePreferenceRequiresThreeWinsAndMinimumDwell(t *testing.T) { + now := time.Date(2026, 8, 5, 12, 0, 0, 0, time.UTC) + state := RoutePreference{CurrentPoolID: "slow", CurrentSince: now} + scores := map[string]float64{"slow": 0.5, "fast": 0.7} + for window := 1; window <= 3; window++ { + state = AdvanceRoutePreference(state, "fast", scores, now.Add(time.Duration(window)*10*time.Second)) + if state.CurrentPoolID != "slow" { + t.Fatalf("switched before minimum dwell in window %d", window) + } + } + state = AdvanceRoutePreference(state, "fast", scores, now.Add(2*time.Minute)) + if state.CurrentPoolID != "fast" { + t.Fatalf("current pool=%q, want fast", state.CurrentPoolID) + } +} + +func TestAdvanceRoutePreferenceImmediatelyLeavesIneligibleCurrentPool(t *testing.T) { + now := time.Now() + state := AdvanceRoutePreference(RoutePreference{ + CurrentPoolID: "offline", CurrentSince: now.Add(-time.Minute), + }, "healthy", map[string]float64{"healthy": 0.6}, now) + if state.CurrentPoolID != "healthy" || state.ChallengerWins != 0 { + t.Fatalf("unexpected preference: %#v", state) + } +} diff --git a/apps/api/internal/executionpool/probe.go b/apps/api/internal/executionpool/probe.go new file mode 100644 index 0000000..f6c244e --- /dev/null +++ b/apps/api/internal/executionpool/probe.go @@ -0,0 +1,101 @@ +package executionpool + +import ( + "context" + "crypto/tls" + "net/http" + "net/http/httptrace" + "time" +) + +type ProbeTarget struct { + RouteProfile RouteProfile + URL string + Method string + Timeout time.Duration +} + +type ProbeResult struct { + Reachable bool + StatusCode int + DNSDuration time.Duration + TCPDuration time.Duration + TLSDuration time.Duration + FirstByte time.Duration + SampledAt time.Time + ErrorClass string +} + +type NetworkProbe interface { + Probe(context.Context, ProbeTarget) ProbeResult +} + +type HTTPProbe struct { + Client *http.Client + Now func() time.Time +} + +func (p HTTPProbe) Probe(ctx context.Context, target ProbeTarget) ProbeResult { + result := ProbeResult{} + now := time.Now + if p.Now != nil { + now = p.Now + } + result.SampledAt = now() + method := target.Method + if method == "" { + method = http.MethodHead + } + timeout := target.Timeout + if timeout <= 0 { + timeout = 3 * time.Second + } + probeCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + var dnsStart, tcpStart, tlsStart time.Time + requestStart := now() + trace := &httptrace.ClientTrace{ + DNSStart: func(httptrace.DNSStartInfo) { dnsStart = now() }, + DNSDone: func(httptrace.DNSDoneInfo) { + if !dnsStart.IsZero() { + result.DNSDuration = now().Sub(dnsStart) + } + }, + ConnectStart: func(_, _ string) { tcpStart = now() }, + ConnectDone: func(_, _ string, _ error) { + if !tcpStart.IsZero() { + result.TCPDuration = now().Sub(tcpStart) + } + }, + TLSHandshakeStart: func() { tlsStart = now() }, + TLSHandshakeDone: func(tls.ConnectionState, error) { + if !tlsStart.IsZero() { + result.TLSDuration = now().Sub(tlsStart) + } + }, + GotFirstResponseByte: func() { result.FirstByte = now().Sub(requestStart) }, + } + request, err := http.NewRequestWithContext(httptrace.WithClientTrace(probeCtx, trace), method, target.URL, nil) + if err != nil { + result.ErrorClass = "invalid_target" + return result + } + client := p.Client + if client == nil { + client = http.DefaultClient + } + response, err := client.Do(request) + if err != nil { + if probeCtx.Err() != nil { + result.ErrorClass = "timeout" + } else { + result.ErrorClass = "network" + } + return result + } + defer response.Body.Close() + result.StatusCode = response.StatusCode + // Any HTTP response, including 401/403/404/405, proves transport reachability. + result.Reachable = true + return result +} diff --git a/apps/api/internal/executionpool/selector.go b/apps/api/internal/executionpool/selector.go new file mode 100644 index 0000000..3031f30 --- /dev/null +++ b/apps/api/internal/executionpool/selector.go @@ -0,0 +1,202 @@ +package executionpool + +import ( + "errors" + "math" + "sort" + "strings" + "time" +) + +const ( + defaultHealthMaxAge = 45 * time.Second + defaultMinimumDwell = 2 * time.Minute + defaultSwitchImprovement = 0.20 + defaultBetterWindowCount = 3 +) + +var ErrNoEligiblePool = errors.New("no eligible execution pool") + +type SelectionCandidate struct { + Pool ExecutionPool + Health RouteHealth + Capacity CapacitySnapshot + CapabilityMatched bool + BetterWindows int +} + +type SelectionRequest struct { + Candidates []SelectionCandidate + CurrentPoolID string + CurrentPoolSince time.Time + Now time.Time +} + +type Decision struct { + PoolID string + Score float64 + Reason string + Components map[string]float64 + Rejected map[string]string +} + +type Selector struct { + HealthMaxAge time.Duration + MinimumDwell time.Duration + SwitchImprovement float64 + BetterWindows int +} + +func NewSelector() Selector { + return Selector{ + HealthMaxAge: defaultHealthMaxAge, MinimumDwell: defaultMinimumDwell, + SwitchImprovement: defaultSwitchImprovement, BetterWindows: defaultBetterWindowCount, + } +} + +type scoredCandidate struct { + candidate SelectionCandidate + score float64 + components map[string]float64 +} + +func (s Selector) Select(request SelectionRequest) (Decision, error) { + now := request.Now + if now.IsZero() { + now = time.Now() + } + if s.HealthMaxAge <= 0 { + s.HealthMaxAge = defaultHealthMaxAge + } + if s.MinimumDwell <= 0 { + s.MinimumDwell = defaultMinimumDwell + } + if s.SwitchImprovement <= 0 { + s.SwitchImprovement = defaultSwitchImprovement + } + if s.BetterWindows <= 0 { + s.BetterWindows = defaultBetterWindowCount + } + rejected := make(map[string]string) + scored := make([]scoredCandidate, 0, len(request.Candidates)) + for _, candidate := range request.Candidates { + if reason := rejectionReason(candidate, now, s.HealthMaxAge); reason != "" { + rejected[candidate.Pool.ID] = reason + continue + } + network := routeQuality(candidate.Health) + queue := queueQuality(candidate.Capacity) + resources := clamp01(candidate.Capacity.ResourceHeadroom) + stability := clamp01(candidate.Capacity.Stability) + components := map[string]float64{ + "network": network, "queue": queue, "resources": resources, "stability": stability, + } + scored = append(scored, scoredCandidate{ + candidate: candidate, + score: 0.55*network + 0.25*queue + 0.15*resources + 0.05*stability, + components: components, + }) + } + if len(scored) == 0 { + return Decision{Rejected: rejected}, ErrNoEligiblePool + } + sort.SliceStable(scored, func(i, j int) bool { + if math.Abs(scored[i].score-scored[j].score) > 0.000001 { + return scored[i].score > scored[j].score + } + return scored[i].candidate.Pool.ID < scored[j].candidate.Pool.ID + }) + selected := scored[0] + if currentIndex := indexOfPool(scored, strings.TrimSpace(request.CurrentPoolID)); currentIndex >= 0 && currentIndex != 0 { + current := scored[currentIndex] + withinDwell := !request.CurrentPoolSince.IsZero() && now.Sub(request.CurrentPoolSince) < s.MinimumDwell + improvement := relativeImprovement(selected.score, current.score) + if withinDwell || improvement < s.SwitchImprovement || selected.candidate.BetterWindows < s.BetterWindows { + selected = current + } + } + return Decision{ + PoolID: selected.candidate.Pool.ID, + Score: selected.score, + Reason: "network_quality_then_safe_capacity", + Components: selected.components, + Rejected: rejected, + }, nil +} + +func rejectionReason(candidate SelectionCandidate, now time.Time, maxAge time.Duration) string { + if strings.TrimSpace(candidate.Pool.ID) == "" { + return "invalid_pool" + } + if candidate.Pool.State != PoolActive { + return "pool_not_active" + } + if !candidate.CapabilityMatched { + return "capability_mismatch" + } + if candidate.Health.State == RouteUnreachable { + return "route_unreachable" + } + if candidate.Health.State == RouteUnknown || candidate.Health.State == "" { + return "route_unknown" + } + if candidate.Health.SampledAt.IsZero() || now.Sub(candidate.Health.SampledAt) > maxAge || (!candidate.Health.ExpiresAt.IsZero() && !now.Before(candidate.Health.ExpiresAt)) { + return "route_stale" + } + if candidate.Capacity.Critical { + return "resource_critical" + } + if candidate.Capacity.WorkerCount < 1 { + return "no_worker" + } + if candidate.Capacity.SafeCapacity-candidate.Capacity.ActiveTasks < 1 { + return "capacity_exhausted" + } + return "" +} + +func routeQuality(health RouteHealth) float64 { + success := clamp01(health.SuccessRate) + connect := latencyQuality(health.ConnectTLSP95, 100*time.Millisecond) + firstByte := latencyQuality(health.FirstByteP95, 250*time.Millisecond) + throughput := clamp01(health.UploadBytesPerSecond / (10 * 1024 * 1024)) + jitter := latencyQuality(health.JitterP95, 100*time.Millisecond) + return 0.35*success + 0.25*connect + 0.25*firstByte + 0.10*throughput + 0.05*jitter +} + +func queueQuality(capacity CapacitySnapshot) float64 { + available := capacity.SafeCapacity - capacity.ActiveTasks + if available <= 0 { + return 0 + } + headroom := float64(available) / float64(max(capacity.SafeCapacity, 1)) + waitPenalty := 1 / (1 + capacity.EstimatedQueueWait.Seconds()) + return clamp01(0.6*headroom + 0.4*waitPenalty) +} + +func latencyQuality(value, target time.Duration) float64 { + if value <= 0 { + return 0 + } + return clamp01(1 / (1 + float64(value)/float64(target))) +} + +func relativeImprovement(candidate, current float64) float64 { + if current <= 0 { + return 1 + } + return (candidate - current) / current +} + +func indexOfPool(candidates []scoredCandidate, poolID string) int { + for index := range candidates { + if candidates[index].candidate.Pool.ID == poolID { + return index + } + } + return -1 +} + +func clamp01(value float64) float64 { + return math.Max(0, math.Min(1, value)) +} diff --git a/apps/api/internal/executionpool/selector_test.go b/apps/api/internal/executionpool/selector_test.go new file mode 100644 index 0000000..f50cbf9 --- /dev/null +++ b/apps/api/internal/executionpool/selector_test.go @@ -0,0 +1,86 @@ +package executionpool + +import ( + "errors" + "testing" + "time" +) + +func TestSelectorRejectsUnreachableAndPrefersNetwork(t *testing.T) { + now := time.Now() + selector := NewSelector() + decision, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{ + candidate("slow", RouteHealthy, now, 120*time.Millisecond, 8, 1), + candidate("fast", RouteHealthy, now, 20*time.Millisecond, 8, 1), + candidate("down", RouteUnreachable, now, time.Millisecond, 32, 0), + }}) + if err != nil { + t.Fatal(err) + } + if decision.PoolID != "fast" { + t.Fatalf("pool=%s, want fast", decision.PoolID) + } + if decision.Rejected["down"] != "route_unreachable" { + t.Fatalf("down rejection=%q", decision.Rejected["down"]) + } +} + +func TestSelectorUsesCapacityAsHardGate(t *testing.T) { + now := time.Now() + selector := NewSelector() + decision, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{ + candidate("fast-full", RouteHealthy, now, 10*time.Millisecond, 2, 2), + candidate("slower-free", RouteHealthy, now, 80*time.Millisecond, 8, 1), + }}) + if err != nil { + t.Fatal(err) + } + if decision.PoolID != "slower-free" { + t.Fatalf("pool=%s, want slower-free", decision.PoolID) + } +} + +func TestSelectorRejectsUnknownAndStale(t *testing.T) { + now := time.Now() + selector := NewSelector() + _, err := selector.Select(SelectionRequest{Now: now, Candidates: []SelectionCandidate{ + candidate("unknown", RouteUnknown, now, time.Millisecond, 8, 0), + candidate("stale", RouteHealthy, now.Add(-time.Minute), time.Millisecond, 8, 0), + }}) + if !errors.Is(err, ErrNoEligiblePool) { + t.Fatalf("err=%v, want ErrNoEligiblePool", err) + } +} + +func TestSelectorHonorsHysteresis(t *testing.T) { + now := time.Now() + selector := NewSelector() + fast := candidate("fast", RouteHealthy, now, 20*time.Millisecond, 8, 1) + fast.BetterWindows = 3 + current := candidate("current", RouteHealthy, now, 80*time.Millisecond, 8, 1) + decision, err := selector.Select(SelectionRequest{ + Now: now, Candidates: []SelectionCandidate{fast, current}, + CurrentPoolID: "current", CurrentPoolSince: now.Add(-time.Minute), + }) + if err != nil { + t.Fatal(err) + } + if decision.PoolID != "current" { + t.Fatalf("pool=%s, want current during dwell", decision.PoolID) + } +} + +func candidate(id string, state RouteState, sampled time.Time, latency time.Duration, capacity, active int) SelectionCandidate { + return SelectionCandidate{ + Pool: ExecutionPool{ID: id, State: PoolActive}, CapabilityMatched: true, + Health: RouteHealth{ + PoolID: id, State: state, SuccessRate: 1, ConnectTLSP95: latency, + FirstByteP95: latency, UploadBytesPerSecond: 10 * 1024 * 1024, + JitterP95: latency / 10, SampledAt: sampled, ExpiresAt: sampled.Add(45 * time.Second), + }, + Capacity: CapacitySnapshot{ + PoolID: id, WorkerCount: 1, SafeCapacity: capacity, ActiveTasks: active, + ResourceHeadroom: 0.8, Stability: 1, + }, + } +} diff --git a/apps/api/internal/executionpool/types.go b/apps/api/internal/executionpool/types.go new file mode 100644 index 0000000..0662125 --- /dev/null +++ b/apps/api/internal/executionpool/types.go @@ -0,0 +1,169 @@ +package executionpool + +import ( + "context" + "io" + "time" +) + +const ProtocolVersion = "v1" + +type PoolState string + +const ( + PoolActive PoolState = "active" + PoolDraining PoolState = "draining" + PoolDisabled PoolState = "disabled" +) + +type RouteState string + +const ( + RouteHealthy RouteState = "healthy" + RouteDegraded RouteState = "degraded" + RouteUnreachable RouteState = "unreachable" + RouteUnknown RouteState = "unknown" +) + +type ExecutionPool struct { + ID string + Labels map[string]string + Capabilities map[string]any + State PoolState + UpdatedAt time.Time +} + +type WorkerDescriptor struct { + WorkerID string + InstanceID string + PoolID string + Endpoint string + ProtocolVersion string + Revision string + Capabilities map[string]any + Allocated int + SafeCapacity int + HeavyCapacity int + ActiveTasks int + PressureState string + HeartbeatAt time.Time + LoadSampledAt time.Time +} + +func (w WorkerDescriptor) AvailableCapacity() int { + limit := w.Allocated + if w.SafeCapacity < limit { + limit = w.SafeCapacity + } + if limit < 0 { + return 0 + } + available := limit - w.ActiveTasks + if available < 0 { + return 0 + } + return available +} + +type RouteProfile struct { + Key string + Provider string + Protocol string + EndpointHost string + ProxyMode string + ConfigRevision string +} + +type RouteHealth struct { + PoolID string + RouteProfileKey string + State RouteState + SuccessRate float64 + ConnectTLSP95 time.Duration + FirstByteP95 time.Duration + UploadBytesPerSecond float64 + JitterP95 time.Duration + ConsecutiveFailures int + ConsecutiveSuccesses int + SampleCount int + SampledAt time.Time + ExpiresAt time.Time +} + +// RouteObservation contains only transport-phase facts. Model generation time +// is deliberately excluded from this structure and therefore from routing +// quality updates. +type RouteObservation struct { + PoolID string + RouteProfileKey string + SampleCount int + SuccessCount int + ConnectTLSP95 time.Duration + UploadBytesPerSecond float64 +} + +type CapacitySnapshot struct { + PoolID string + WorkerCount int + Allocated int + SafeCapacity int + ActiveTasks int + HeavyCapacity int + HeavyTasks int + QueueDepth int + EstimatedQueueWait time.Duration + ResourceHeadroom float64 + Stability float64 + Critical bool + SampledAt time.Time +} + +type DesiredCapacity struct { + PoolID string + Desired int + Reason string + ValidUntil time.Time +} + +// WorkerDirectory is the platform-neutral discovery boundary. Implementations +// may use PostgreSQL, Consul, etcd, or another registry. +type WorkerDirectory interface { + ListWorkers(context.Context, time.Time) ([]WorkerDescriptor, error) +} + +// ExecutionBroker hides the durable queue implementation from routing code. +type ExecutionBroker interface { + Publish(context.Context, string, string, time.Time) (int64, error) +} + +// ExecutionTransport hides HTTP, gRPC, or service-mesh transport details. +type ExecutionTransport interface { + Execute(context.Context, WorkerDescriptor, ExecutionRequest) (ExecutionResponse, error) +} + +type ExecutionRequest struct { + TaskID string + PoolID string + WorkerID string + LeaseID string + AuthorizationToken string + Deadline time.Time + Stream bool +} + +type ExecutionResponse struct { + StatusCode int + Headers map[string][]string + Body io.ReadCloser +} + +type RouteHealthRepository interface { + ListRouteHealth(context.Context, string, time.Time) ([]RouteHealth, error) + RecordRouteHealth(context.Context, RouteHealth) error + RecordRouteObservation(context.Context, RouteObservation) error +} + +type CapacityProvider interface { + ListCapacity(context.Context, time.Time) ([]CapacitySnapshot, error) + PublishDesiredCapacity(context.Context, DesiredCapacity) error +} diff --git a/apps/api/internal/httpapi/admin_execution_pools.go b/apps/api/internal/httpapi/admin_execution_pools.go new file mode 100644 index 0000000..54c2930 --- /dev/null +++ b/apps/api/internal/httpapi/admin_execution_pools.go @@ -0,0 +1,143 @@ +package httpapi + +import ( + "net/http" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" +) + +type AdminExecutionPoolWorker struct { + WorkerID string `json:"workerId"` + InstanceID string `json:"instanceId"` + Revision string `json:"revision,omitempty"` + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]any `json:"capabilities"` + Allocated int `json:"allocated"` + SafeCapacity int `json:"safeCapacity"` + ActiveTasks int `json:"activeTasks"` + PressureState string `json:"pressureState"` + HeartbeatAt time.Time `json:"heartbeatAt"` +} + +type AdminExecutionPool struct { + PoolID string `json:"poolId"` + Labels map[string]string `json:"labels"` + Capabilities map[string]any `json:"capabilities"` + State string `json:"state"` + Capacity *AdminExecutionPoolCapacity `json:"capacity,omitempty"` + Workers []AdminExecutionPoolWorker `json:"workers"` +} + +type AdminExecutionPoolCapacity struct { + WorkerCount int `json:"workerCount"` + Allocated int `json:"allocated"` + SafeCapacity int `json:"safeCapacity"` + ActiveTasks int `json:"activeTasks"` + HeavyCapacity int `json:"heavyCapacity"` + HeavyTasks int `json:"heavyTasks"` + QueueDepth int `json:"queueDepth"` + EstimatedQueueWaitMilli int64 `json:"estimatedQueueWaitMs"` + ResourceHeadroom float64 `json:"resourceHeadroom"` + Stability float64 `json:"stability"` + Critical bool `json:"critical"` + SampledAt time.Time `json:"sampledAt"` +} + +type AdminRouteHealth struct { + PoolID string `json:"poolId"` + RouteProfileKey string `json:"routeProfileKey"` + State string `json:"state"` + SuccessRate float64 `json:"successRate"` + ConnectTLSP95Milli int64 `json:"connectTlsP95Ms"` + FirstByteP95Milli int64 `json:"firstByteP95Ms"` + UploadBytesPerSecond float64 `json:"uploadBytesPerSecond"` + JitterP95Milli int64 `json:"jitterP95Ms"` + ConsecutiveFailures int `json:"consecutiveFailures"` + ConsecutiveSuccesses int `json:"consecutiveSuccesses"` + SampleCount int `json:"sampleCount"` + SampledAt time.Time `json:"sampledAt"` + ExpiresAt time.Time `json:"expiresAt"` +} + +type AdminExecutionPoolResponse struct { + Pools []AdminExecutionPool `json:"pools"` + RouteHealth []AdminRouteHealth `json:"routeHealth"` + CapturedAt time.Time `json:"capturedAt"` +} + +// listExecutionPools godoc +// @Summary 查询逻辑执行池、Worker 容量和上游路由健康状态 +// @Tags Admin Runtime +// @Produce json +// @Success 200 {object} AdminExecutionPoolResponse +// @Failure 500 {object} ErrorEnvelope +// @Security BearerAuth +// @Router /api/admin/runtime/execution-pools [get] +func (s *Server) listExecutionPools(w http.ResponseWriter, r *http.Request) { + now := time.Now() + pools, err := s.coordinationStore.ListExecutionPools(r.Context()) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + workers, err := s.coordinationStore.ListWorkers(r.Context(), now) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + capacities, err := s.coordinationStore.ListCapacity(r.Context(), now) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + health, err := s.coordinationStore.ListAllRouteHealth(r.Context()) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + capacityByPool := make(map[string]executionpool.CapacitySnapshot, len(capacities)) + for _, capacity := range capacities { + capacityByPool[capacity.PoolID] = capacity + } + workersByPool := make(map[string][]AdminExecutionPoolWorker) + for _, worker := range workers { + workersByPool[worker.PoolID] = append(workersByPool[worker.PoolID], AdminExecutionPoolWorker{ + WorkerID: worker.WorkerID, InstanceID: worker.InstanceID, Revision: worker.Revision, + ProtocolVersion: worker.ProtocolVersion, Capabilities: worker.Capabilities, + Allocated: worker.Allocated, SafeCapacity: worker.SafeCapacity, + ActiveTasks: worker.ActiveTasks, PressureState: worker.PressureState, + HeartbeatAt: worker.HeartbeatAt, + }) + } + response := AdminExecutionPoolResponse{ + Pools: make([]AdminExecutionPool, 0, len(pools)), RouteHealth: make([]AdminRouteHealth, 0, len(health)), CapturedAt: now, + } + for _, item := range health { + response.RouteHealth = append(response.RouteHealth, AdminRouteHealth{ + PoolID: item.PoolID, RouteProfileKey: item.RouteProfileKey, State: string(item.State), + SuccessRate: item.SuccessRate, ConnectTLSP95Milli: item.ConnectTLSP95.Milliseconds(), + FirstByteP95Milli: item.FirstByteP95.Milliseconds(), UploadBytesPerSecond: item.UploadBytesPerSecond, + JitterP95Milli: item.JitterP95.Milliseconds(), ConsecutiveFailures: item.ConsecutiveFailures, + ConsecutiveSuccesses: item.ConsecutiveSuccesses, SampleCount: item.SampleCount, + SampledAt: item.SampledAt, ExpiresAt: item.ExpiresAt, + }) + } + for _, pool := range pools { + item := AdminExecutionPool{ + PoolID: pool.ID, Labels: pool.Labels, Capabilities: pool.Capabilities, + State: string(pool.State), Workers: workersByPool[pool.ID], + } + if capacity, ok := capacityByPool[pool.ID]; ok { + item.Capacity = &AdminExecutionPoolCapacity{ + WorkerCount: capacity.WorkerCount, Allocated: capacity.Allocated, SafeCapacity: capacity.SafeCapacity, + ActiveTasks: capacity.ActiveTasks, HeavyCapacity: capacity.HeavyCapacity, HeavyTasks: capacity.HeavyTasks, + QueueDepth: capacity.QueueDepth, EstimatedQueueWaitMilli: capacity.EstimatedQueueWait.Milliseconds(), + ResourceHeadroom: capacity.ResourceHeadroom, Stability: capacity.Stability, + Critical: capacity.Critical, SampledAt: capacity.SampledAt, + } + } + response.Pools = append(response.Pools, item) + } + writeJSON(w, http.StatusOK, response) +} diff --git a/apps/api/internal/httpapi/compat_protocol_test.go b/apps/api/internal/httpapi/compat_protocol_test.go index d9d04f4..a53f60e 100644 --- a/apps/api/internal/httpapi/compat_protocol_test.go +++ b/apps/api/internal/httpapi/compat_protocol_test.go @@ -125,10 +125,10 @@ func TestProtocolErrorsUseCompatibleShapesWithStandardPublicErrors(t *testing.T) }, }, { - name: "volces", protocol: clients.ProtocolVolcesContents, status: http.StatusBadGateway, + name: "volces", protocol: clients.ProtocolVolcesContents, status: http.StatusGatewayTimeout, assertBody: func(t *testing.T, body map[string]any) { errorBody := requireObject(t, body["error"]) - if errorBody["code"] != "upstream_submission_unknown" || errorBody["httpStatus"] != float64(http.StatusBadGateway) || errorBody["retryable"] != true { + if errorBody["code"] != "upstream_timeout" || errorBody["httpStatus"] != float64(http.StatusGatewayTimeout) || errorBody["retryable"] != true { t.Fatalf("unexpected Volces error: %+v", body) } assertNoKeys(t, errorBody, "status", "taskId", "gateway_status") @@ -151,7 +151,7 @@ func TestProtocolErrorsUseCompatibleShapesWithStandardPublicErrors(t *testing.T) if test.name == "gemini" { code = "rate_limit" } else if test.name == "volces" { - code = "upstream_submission_unknown" + code = "upstream_timeout" } else if test.name == "keling" { code = "gateway_rate_limited" } diff --git a/apps/api/internal/httpapi/handlers.go b/apps/api/internal/httpapi/handlers.go index ead87b1..e91e1a6 100644 --- a/apps/api/internal/httpapi/handlers.go +++ b/apps/api/internal/httpapi/handlers.go @@ -1892,6 +1892,11 @@ func applyRunErrorHeaders(w http.ResponseWriter, err error) { if limitRetryAfter := store.RateLimitRetryAfter(err); limitRetryAfter > 0 { retryAfter = limitRetryAfter } + if retryAfter <= 0 { + if seconds := retryAfterSecondsFromDetails(clients.ErrorDetails(err)); seconds > 0 { + retryAfter = time.Duration(seconds) * time.Second + } + } if retryAfter > 0 { seconds := int((retryAfter + time.Second - 1) / time.Second) if seconds < 1 { @@ -1901,6 +1906,24 @@ func applyRunErrorHeaders(w http.ResponseWriter, err error) { } } +func retryAfterSecondsFromDetails(details map[string]any) int { + if details == nil { + return 0 + } + switch value := details["retryAfterSeconds"].(type) { + case int: + return value + case int32: + return int(value) + case int64: + return int(value) + case float64: + return int(value) + default: + return 0 + } +} + func rateLimitErrorSummary(err error) string { var limitErr *store.RateLimitExceededError if !errors.As(err, &limitErr) { diff --git a/apps/api/internal/httpapi/internal_execution.go b/apps/api/internal/httpapi/internal_execution.go new file mode 100644 index 0000000..3f8ec29 --- /dev/null +++ b/apps/api/internal/httpapi/internal_execution.go @@ -0,0 +1,50 @@ +package httpapi + +import ( + "encoding/json" + "io" + "net/http" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/runner" +) + +const maxInternalExecutionRequestBytes = 64 * 1024 + +func (s *Server) internalExecution(w http.ResponseWriter, r *http.Request) { + body := http.MaxBytesReader(w, r.Body, maxInternalExecutionRequestBytes) + defer body.Close() + var input runner.InternalExecutionRequest + decoder := json.NewDecoder(body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&input); err != nil { + http.Error(w, "invalid internal execution request", http.StatusBadRequest) + return + } + var trailing any + if err := decoder.Decode(&trailing); err != io.EOF { + http.Error(w, "invalid internal execution request", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/x-ndjson") + w.Header().Set("Cache-Control", "no-store") + encoder := json.NewEncoder(w) + flusher, _ := w.(http.Flusher) + writeFrame := func(frame runner.InternalExecutionFrame) error { + if err := encoder.Encode(frame); err != nil { + return err + } + if flusher != nil { + flusher.Flush() + } + return nil + } + result, err := s.runner.ExecuteInternal(r.Context(), r.Header.Get("Authorization"), input, func(delta clients.StreamDeltaEvent) error { + return writeFrame(runner.InternalExecutionFrame{Type: "delta", Delta: &delta}) + }) + if err != nil { + _ = writeFrame(runner.InternalExecutionFrameForError(err)) + return + } + _ = writeFrame(runner.InternalExecutionFrame{Type: "result", Output: result.Output, Wire: result.Wire}) +} diff --git a/apps/api/internal/httpapi/route_unavailable_error_test.go b/apps/api/internal/httpapi/route_unavailable_error_test.go new file mode 100644 index 0000000..dad7947 --- /dev/null +++ b/apps/api/internal/httpapi/route_unavailable_error_test.go @@ -0,0 +1,18 @@ +package httpapi + +import ( + "net/http/httptest" + "testing" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" +) + +func TestApplyRunErrorHeadersUsesRouteRetryAfter(t *testing.T) { + recorder := httptest.NewRecorder() + applyRunErrorHeaders(recorder, &clients.ClientError{ + Code: "upstream_route_unavailable", Details: map[string]any{"retryAfterSeconds": 10}, + }) + if got := recorder.Header().Get("Retry-After"); got != "10" { + t.Fatalf("Retry-After=%q, want 10", got) + } +} diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index 0cf3259..9f2617c 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -138,6 +138,7 @@ func NewServerWithStores( logger.Info("asynchronous queue worker disabled for this process") } } + server.runner.StartRouteHealthProber(ctx) if cfg.RunsAsyncAdmissionDispatcher() { server.runner.StartAsyncAdmissionDispatcher(ctx) } @@ -156,6 +157,9 @@ func NewServerWithStores( mux.Handle("GET /metrics", securityEventMetrics.DynamicHandler(postgresMetricsProvider{ Store: db, critical: coordinationDB, river: riverDB, })) + if cfg.RunsAsyncExecutionWorker() { + mux.HandleFunc("POST /internal/v1/executions", server.internalExecution) + } if !cfg.RunsPublicHTTP() { return server.recover(mux) } @@ -247,6 +251,7 @@ func NewServerWithStores( mux.Handle("PATCH /api/admin/runtime/policy-sets/{policySetID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateRuntimePolicySet))) mux.Handle("DELETE /api/admin/runtime/policy-sets/{policySetID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.deleteRuntimePolicySet))) mux.Handle("GET /api/admin/runtime/runner-policy", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.getRunnerPolicy))) + mux.Handle("GET /api/admin/runtime/execution-pools", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listExecutionPools))) mux.Handle("PATCH /api/admin/runtime/runner-policy", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.updateRunnerPolicy))) mux.Handle("GET /api/admin/runtime/billing-settlements", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listBillingSettlements))) mux.Handle("POST /api/admin/runtime/billing-settlements/{settlementId}/retry", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.retryBillingSettlement))) diff --git a/apps/api/internal/publicerror/public_error.go b/apps/api/internal/publicerror/public_error.go index 1db9dd4..ac599c7 100644 --- a/apps/api/internal/publicerror/public_error.go +++ b/apps/api/internal/publicerror/public_error.go @@ -79,10 +79,12 @@ func mappedError(code string, message string, status int, retryable bool) (Error return newError(code, firstNonEmptyMessage(message, "New production tasks are paused while validation is running."), "gateway", http.StatusServiceUnavailable, true, "retry_after"), true case "traffic_gate_unavailable": return newError(code, "The gateway traffic admission service is temporarily unavailable.", "gateway", http.StatusServiceUnavailable, true, "retry"), true + case "upstream_route_unavailable": + return newError(code, "No execution pool can currently reach the upstream service.", "gateway", http.StatusServiceUnavailable, true, "retry_after"), true case "response_read_error", "stream_read_error", "network", "connection_reset", "upstream_connection_interrupted", "unexpected_eof", "http2_stream_closed", "request_asset_fetch_failed": return newError("upstream_connection_interrupted", "The upstream connection was interrupted before a complete response was received.", "upstream", http.StatusBadGateway, true, "retry"), true - case "upstream_timeout", "timeout", "context_deadline_exceeded": - return newError("upstream_timeout", "The upstream service did not respond in time.", "upstream", http.StatusGatewayTimeout, true, "retry"), true + case "upstream_timeout", "upstream_submission_unknown", "timeout", "context_deadline_exceeded": + return newError("upstream_timeout", "The upstream service did not respond in time. The gateway does not resubmit after submission begins.", "upstream", http.StatusGatewayTimeout, true, "retry_status"), true case "gateway_rate_limited": return newError("gateway_rate_limited", "The gateway rate limit was reached.", "rate_limit", http.StatusTooManyRequests, true, "retry_after"), true case "rate_limit", "upstream_rate_limited", "too_many_requests": @@ -131,10 +133,10 @@ func mappedError(code string, message string, status int, retryable bool) (Error if strings.HasPrefix(code, "http_") || (code == "provider_failed" && status > 0 && status < 500) { return upstreamHTTPError(status, originalUpstreamErrorDetail(code, message, status)), true } - if status >= 500 && code != "upstream_submission_unknown" && (strings.Contains(code, "upstream") || strings.Contains(code, "provider")) { + if status >= 500 && (strings.Contains(code, "upstream") || strings.Contains(code, "provider")) { return newError("upstream_unavailable", "The upstream service is temporarily unavailable.", "upstream", http.StatusServiceUnavailable, true, "retry"), true } - if status >= 500 && code != "upstream_submission_unknown" { + if status >= 500 { return newError("gateway_error", defaultMessageForStatus(status), "gateway", status, retryable, retryAction(retryable)), true } return Error{}, false diff --git a/apps/api/internal/publicerror/submission_unknown_compat_test.go b/apps/api/internal/publicerror/submission_unknown_compat_test.go new file mode 100644 index 0000000..c8ba0b7 --- /dev/null +++ b/apps/api/internal/publicerror/submission_unknown_compat_test.go @@ -0,0 +1,10 @@ +package publicerror + +import "testing" + +func TestLegacySubmissionUnknownNeverLeavesPublicBoundary(t *testing.T) { + value := FromFields("upstream_submission_unknown", "legacy value", 500, true) + if value.Code != "upstream_timeout" { + t.Fatalf("public code=%q, want upstream_timeout", value.Code) + } +} diff --git a/apps/api/internal/runner/admission.go b/apps/api/internal/runner/admission.go index 4cf0e93..33f04c7 100644 --- a/apps/api/internal/runner/admission.go +++ b/apps/api/internal/runner/admission.go @@ -685,6 +685,10 @@ func (s *Service) SubmitAsyncTask(ctx context.Context, task store.GatewayTask) e _, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error()) return err } + if err := s.routeAsyncTask(ctx, &task, user, plan.Candidate); err != nil { + _, _ = s.store.FailQueuedTask(context.WithoutCancel(ctx), task.ID, clients.ErrorCode(err), err.Error()) + return err + } if !plan.Eligible { if err := s.EnqueueAsyncTask(ctx, task); err != nil { if s.cancelAsyncSubmissionIfDisconnected(ctx, task.ID) { @@ -830,7 +834,10 @@ func (s *Service) dispatchWaitingAsyncTasks(ctx context.Context, admissions []st if !ok { return fmt.Errorf("async admission batch task %s was not prepared", input.TaskID) } - return s.enqueueAsyncTaskTx(ctx, tx, task.ID, asyncTaskInsertOpts(task)) + _, publishErr := (riverExecutionBroker{service: s}).publishTx( + ctx, tx, task.AssignedPoolID, task.ID, time.Time{}, 2, + ) + return publishErr } outcomes, batchErr := s.store.TryTaskAdmissionAtomicBatchWithAdmittedHook( ctx, diff --git a/apps/api/internal/runner/http_execution_transport.go b/apps/api/internal/runner/http_execution_transport.go new file mode 100644 index 0000000..b48ee26 --- /dev/null +++ b/apps/api/internal/runner/http_execution_transport.go @@ -0,0 +1,60 @@ +package runner + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "net/http" + "strings" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" +) + +// httpExecutionTransport is the first platform-neutral Worker transport +// adapter. The routing core sees only WorkerDescriptor and executionpool's +// versioned request/response types; Pod DNS and other orchestrator concepts do +// not cross this boundary. +type httpExecutionTransport struct { + client *http.Client +} + +func (transport httpExecutionTransport) Execute( + ctx context.Context, + worker executionpool.WorkerDescriptor, + input executionpool.ExecutionRequest, +) (executionpool.ExecutionResponse, error) { + if transport.client == nil { + return executionpool.ExecutionResponse{}, errors.New("worker HTTP client is required") + } + payload, err := json.Marshal(InternalExecutionRequest{ + TaskID: input.TaskID, + LeaseID: input.LeaseID, + Stream: input.Stream, + }) + if err != nil { + return executionpool.ExecutionResponse{}, err + } + request, err := http.NewRequestWithContext( + ctx, + http.MethodPost, + strings.TrimRight(worker.Endpoint, "/")+workerExecutionPath, + bytes.NewReader(payload), + ) + if err != nil { + return executionpool.ExecutionResponse{}, err + } + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Worker "+input.AuthorizationToken) + response, err := transport.client.Do(request) + if err != nil { + return executionpool.ExecutionResponse{}, err + } + return executionpool.ExecutionResponse{ + StatusCode: response.StatusCode, + Headers: response.Header, + Body: response.Body, + }, nil +} + +var _ executionpool.ExecutionTransport = httpExecutionTransport{} diff --git a/apps/api/internal/runner/http_execution_transport_test.go b/apps/api/internal/runner/http_execution_transport_test.go new file mode 100644 index 0000000..3f186f3 --- /dev/null +++ b/apps/api/internal/runner/http_execution_transport_test.go @@ -0,0 +1,53 @@ +package runner + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" +) + +func TestHTTPExecutionTransportUsesAdvertisedEndpointAndProtocol(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) { + if request.URL.Path != workerExecutionPath { + t.Fatalf("path=%q", request.URL.Path) + } + if got := request.Header.Get("Authorization"); got != "Worker signed-token" { + t.Fatalf("authorization=%q", got) + } + var payload InternalExecutionRequest + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Fatal(err) + } + if payload.TaskID != "task-1" || payload.LeaseID != "lease-1" || !payload.Stream { + t.Fatalf("payload=%+v", payload) + } + response.Header().Set("X-Worker", "accepted") + response.WriteHeader(http.StatusAccepted) + _, _ = response.Write([]byte("ok")) + })) + defer server.Close() + + transport := httpExecutionTransport{client: server.Client()} + result, err := transport.Execute(context.Background(), executionpool.WorkerDescriptor{ + WorkerID: "worker-1", PoolID: "pool-1", Endpoint: server.URL, + }, executionpool.ExecutionRequest{ + TaskID: "task-1", PoolID: "pool-1", WorkerID: "worker-1", LeaseID: "lease-1", + AuthorizationToken: "signed-token", Stream: true, + }) + if err != nil { + t.Fatal(err) + } + defer result.Body.Close() + body, err := io.ReadAll(result.Body) + if err != nil { + t.Fatal(err) + } + if result.StatusCode != http.StatusAccepted || len(result.Headers["X-Worker"]) != 1 || result.Headers["X-Worker"][0] != "accepted" || string(body) != "ok" { + t.Fatalf("result=%+v body=%q", result, body) + } +} diff --git a/apps/api/internal/runner/passive_route_observer.go b/apps/api/internal/runner/passive_route_observer.go new file mode 100644 index 0000000..d62ec34 --- /dev/null +++ b/apps/api/internal/runner/passive_route_observer.go @@ -0,0 +1,137 @@ +package runner + +import ( + "context" + "crypto/tls" + "net/http" + "net/http/httptrace" + "sort" + "sync" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +type passiveRouteObserver struct { + mu sync.Mutex + samples int + successes int + connectTLS []time.Duration + uploadBytesRate []float64 +} + +type passiveRouteRoundTripper struct { + base http.RoundTripper + observer *passiveRouteObserver +} + +func (observer *passiveRouteObserver) wrap(client *http.Client) *http.Client { + if client == nil { + client = http.DefaultClient + } + cloned := *client + base := cloned.Transport + if base == nil { + base = http.DefaultTransport + } + cloned.Transport = passiveRouteRoundTripper{base: base, observer: observer} + return &cloned +} + +func (transport passiveRouteRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + startedAt := time.Now() + var connectStartedAt time.Time + var connectDuration time.Duration + var tlsStartedAt time.Time + var tlsDuration time.Duration + var wroteRequestAt time.Time + trace := &httptrace.ClientTrace{ + ConnectStart: func(_, _ string) { connectStartedAt = time.Now() }, + ConnectDone: func(_, _ string, _ error) { + if !connectStartedAt.IsZero() { + connectDuration = time.Since(connectStartedAt) + } + }, + TLSHandshakeStart: func() { tlsStartedAt = time.Now() }, + TLSHandshakeDone: func(tls.ConnectionState, error) { + if !tlsStartedAt.IsZero() { + tlsDuration = time.Since(tlsStartedAt) + } + }, + WroteRequest: func(httptrace.WroteRequestInfo) { wroteRequestAt = time.Now() }, + } + traced := request.Clone(httptrace.WithClientTrace(request.Context(), trace)) + response, err := transport.base.RoundTrip(traced) + transport.observer.record( + err == nil, + connectDuration+tlsDuration, + request.ContentLength, + startedAt, + wroteRequestAt, + ) + return response, err +} + +func (observer *passiveRouteObserver) record(success bool, connectTLS time.Duration, contentLength int64, startedAt, wroteRequestAt time.Time) { + observer.mu.Lock() + defer observer.mu.Unlock() + observer.samples++ + if success { + observer.successes++ + } + if connectTLS > 0 { + observer.connectTLS = append(observer.connectTLS, connectTLS) + } + if contentLength > 0 && !wroteRequestAt.IsZero() && wroteRequestAt.After(startedAt) { + observer.uploadBytesRate = append(observer.uploadBytesRate, float64(contentLength)/wroteRequestAt.Sub(startedAt).Seconds()) + } +} + +func (observer *passiveRouteObserver) snapshot(poolID, routeProfileKey string) executionpool.RouteObservation { + observer.mu.Lock() + defer observer.mu.Unlock() + return executionpool.RouteObservation{ + PoolID: poolID, RouteProfileKey: routeProfileKey, + SampleCount: observer.samples, SuccessCount: observer.successes, + ConnectTLSP95: durationP95(observer.connectTLS), + UploadBytesPerSecond: floatP95(observer.uploadBytesRate), + } +} + +func durationP95(values []time.Duration) time.Duration { + if len(values) == 0 { + return 0 + } + copyOfValues := append([]time.Duration(nil), values...) + sort.Slice(copyOfValues, func(i, j int) bool { return copyOfValues[i] < copyOfValues[j] }) + return copyOfValues[(len(copyOfValues)*95+99)/100-1] +} + +func floatP95(values []float64) float64 { + if len(values) == 0 { + return 0 + } + copyOfValues := append([]float64(nil), values...) + sort.Float64s(copyOfValues) + return copyOfValues[(len(copyOfValues)*95+99)/100-1] +} + +func (s *Service) recordPassiveRouteObservation(task store.GatewayTask, observer *passiveRouteObserver) { + if observer == nil || !s.routingEnabled() || task.RouteProfileKey == "" { + return + } + poolID := task.AssignedPoolID + if poolID == "" && s.cfg.RunsAsyncExecutionWorker() { + poolID = s.cfg.ExecutionPoolID + } + observation := observer.snapshot(poolID, task.RouteProfileKey) + if observation.PoolID == "" || observation.SampleCount == 0 { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond) + defer cancel() + if err := s.coordinationStore.RecordRouteObservation(ctx, observation); err != nil && s.logger != nil { + s.logger.Warn("record passive route observation failed", "error_category", "route_observation_failed") + } +} diff --git a/apps/api/internal/runner/passive_route_observer_test.go b/apps/api/internal/runner/passive_route_observer_test.go new file mode 100644 index 0000000..9847f9e --- /dev/null +++ b/apps/api/internal/runner/passive_route_observer_test.go @@ -0,0 +1,25 @@ +package runner + +import ( + "testing" + "time" +) + +func TestPassiveRouteObserverRecordsOnlyTransportPhaseFacts(t *testing.T) { + observer := &passiveRouteObserver{} + startedAt := time.Now() + observer.record(true, 10*time.Millisecond, 1024, startedAt, startedAt.Add(time.Millisecond)) + observer.record(false, 20*time.Millisecond, 2048, startedAt, startedAt.Add(2*time.Millisecond)) + observer.record(true, 30*time.Millisecond, 0, startedAt, time.Time{}) + + observation := observer.snapshot("pool-a", "route-a") + if observation.SampleCount != 3 || observation.SuccessCount != 2 { + t.Fatalf("observation=%+v", observation) + } + if observation.ConnectTLSP95 != 30*time.Millisecond { + t.Fatalf("connect/TLS p95=%s", observation.ConnectTLSP95) + } + if observation.UploadBytesPerSecond < 1_024_000-1 || observation.UploadBytesPerSecond > 1_024_000+1 { + t.Fatalf("upload rate=%f", observation.UploadBytesPerSecond) + } +} diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index 05b048b..7f3e6e6 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -2,6 +2,7 @@ package runner import ( "context" + "encoding/json" "errors" "fmt" "os" @@ -9,6 +10,7 @@ import ( "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" "github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload" "github.com/google/uuid" @@ -17,7 +19,6 @@ import ( "github.com/riverqueue/river" "github.com/riverqueue/river/riverdriver/riverpgxv5" "github.com/riverqueue/river/rivermigrate" - "github.com/riverqueue/river/rivertype" ) const ( @@ -59,6 +60,11 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs if task.Status == "succeeded" || task.Status == "failed" || task.Status == "cancelled" { return nil } + if assignedPool := strings.TrimSpace(task.AssignedPoolID); assignedPool != "" && assignedPool != strings.TrimSpace(w.service.cfg.ExecutionPoolID) { + w.service.logger.Warn("river task claimed by a worker outside its assigned pool", + "taskID", task.ID, "assignedPool", assignedPool, "workerPool", w.service.cfg.ExecutionPoolID) + return river.JobSnooze(time.Second) + } loadLease, admitted := w.service.tryStartWorkerTask() if !admitted { return river.JobSnooze(workerLoadRetryDelay(task.ID)) @@ -87,6 +93,10 @@ func (w *asyncTaskWorker) Work(ctx context.Context, job *river.Job[asyncTaskArgs task.ExecutionToken = executionToken queued, queueErr := w.service.requeueInterruptedAsyncTask(context.WithoutCancel(ctx), task) if queueErr != nil { + if errors.Is(queueErr, store.ErrTaskExecutionManualReview) { + w.service.logger.Warn("interrupted task held after upstream submission began", "taskID", task.ID, "error_category", "upstream_timeout") + return nil + } return queueErr } w.service.logger.Debug("river async task interrupted and requeued", "taskID", task.ID, "status", queued.Status, "riverJobID", job.ID) @@ -309,14 +319,19 @@ func (s *Service) newRiverAsyncExecutionClient(capacity int) (*river.Client[pgx. if err := river.AddWorkerSafely(workers, &asyncTaskWorker{service: s}); err != nil { return nil, err } + queues := map[string]river.QueueConfig{ + asyncTaskQueueName: {MaxWorkers: capacity}, + } + poolQueue := executionpool.QueueName(s.cfg.ExecutionPoolID) + if poolQueue != asyncTaskQueueName { + queues[poolQueue] = river.QueueConfig{MaxWorkers: capacity} + } return river.NewClient(riverpgxv5.New(s.riverStore.Pool()), &river.Config{ ID: fmt.Sprintf("%s-exec-%d-%d", s.workerInstanceID, capacity, time.Now().UnixNano()), JobTimeout: -1, Logger: s.logger, CompletedJobRetentionPeriod: 24 * time.Hour, - Queues: map[string]river.QueueConfig{ - asyncTaskQueueName: {MaxWorkers: capacity}, - }, + Queues: queues, // Provider-backed media jobs commonly poll for 10-20 minutes. River may // execute a still-running job again once this window elapses, so keep the // rescue horizon above the longest configured provider poll timeout. @@ -346,25 +361,40 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke return store.AsyncWorkerCapacitySnapshot{}, err } loadSnapshot := s.sampleWorkerLoad() + labels := map[string]string{} + if raw := strings.TrimSpace(s.cfg.ExecutionPoolLabels); raw != "" { + if err := json.Unmarshal([]byte(raw), &labels); err != nil { + return store.AsyncWorkerCapacitySnapshot{}, fmt.Errorf("decode execution pool labels: %w", err) + } + } + if endpoint := strings.TrimSpace(s.cfg.WorkerAdvertiseEndpoint); endpoint != "" { + if err := executionpool.ValidateAdvertisedEndpoint(endpoint, splitTrimmed(s.cfg.WorkerEndpointAllowedSuffixes), s.cfg.WorkerEndpointAllowPrivate); err != nil { + return store.AsyncWorkerCapacitySnapshot{}, fmt.Errorf("validate worker advertised endpoint: %w", err) + } + } allocation, err := s.coordinationStore.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{ - InstanceID: s.workerInstanceID, - PodUID: strings.TrimSpace(os.Getenv("POD_UID")), - PodName: strings.TrimSpace(os.Getenv("POD_NAME")), - Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")), - Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), - DesiredCapacity: snapshot.Capacity, - CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit, - LoadMode: loadSnapshot.Mode, - SafeCapacity: loadSnapshot.SafeCapacity, - HeavyCapacity: loadSnapshot.HeavyLimit, - ActiveTasks: loadSnapshot.ActiveTasks, - PreparingTasks: loadSnapshot.PreparingTasks, - WaitingUpstreamTasks: loadSnapshot.WaitingUpstreamTasks, - FinalizingTasks: loadSnapshot.FinalizingTasks, - PressureState: string(loadSnapshot.PressureState), - PressureReason: loadSnapshot.PressureReason, - LoadSampledAt: loadSnapshot.SampledAt, - HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second, + InstanceID: s.workerInstanceID, + WorkerID: firstNonEmptyString(s.cfg.WorkerID, s.workerInstanceID), + PoolID: s.cfg.ExecutionPoolID, + Endpoint: s.cfg.WorkerAdvertiseEndpoint, + Labels: labels, + Capabilities: defaultWorkerCapabilities(), + ProtocolVersion: executionpool.ProtocolVersion, + OrchestratorInstanceRef: s.cfg.WorkerOrchestratorInstanceRef, + Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), + DesiredCapacity: snapshot.Capacity, + CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit, + LoadMode: loadSnapshot.Mode, + SafeCapacity: loadSnapshot.SafeCapacity, + HeavyCapacity: loadSnapshot.HeavyLimit, + ActiveTasks: loadSnapshot.ActiveTasks, + PreparingTasks: loadSnapshot.PreparingTasks, + WaitingUpstreamTasks: loadSnapshot.WaitingUpstreamTasks, + FinalizingTasks: loadSnapshot.FinalizingTasks, + PressureState: string(loadSnapshot.PressureState), + PressureReason: loadSnapshot.PressureReason, + LoadSampledAt: loadSnapshot.SampledAt, + HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second, }) if err != nil { return store.AsyncWorkerCapacitySnapshot{}, err @@ -599,7 +629,8 @@ func (s *Service) observeAsyncWorkerResize(outcome string) { } func (s *Service) EnqueueAsyncTask(ctx context.Context, task store.GatewayTask) error { - return s.enqueueAsyncTaskWithOptions(ctx, task.ID, asyncTaskInsertOpts(task)) + _, err := (riverExecutionBroker{service: s}).Publish(ctx, task.AssignedPoolID, task.ID, time.Time{}) + return err } func (s *Service) enqueueAsyncTaskWithOptions(ctx context.Context, taskID string, opts *river.InsertOpts) error { @@ -761,22 +792,25 @@ func asyncTaskInsertOpts(task store.GatewayTask) *river.InsertOpts { if task.ID == "" { priority = 3 } - return &river.InsertOpts{ - MaxAttempts: 1000, - Priority: priority, - Queue: asyncTaskQueueName, - Tags: []string{"gateway-task"}, - UniqueOpts: river.UniqueOpts{ - ByArgs: true, - ByQueue: true, - ByState: []rivertype.JobState{ - rivertype.JobStateAvailable, - rivertype.JobStatePending, - rivertype.JobStateRetryable, - rivertype.JobStateRunning, - rivertype.JobStateScheduled, - }, - }, + return riverExecutionInsertOptions(task.AssignedPoolID, time.Time{}, priority) +} + +func splitTrimmed(value string) []string { + items := make([]string, 0) + for item := range strings.SplitSeq(value, ",") { + if item = strings.TrimSpace(item); item != "" { + items = append(items, item) + } + } + return items +} + +func defaultWorkerCapabilities() map[string]any { + return map[string]any{ + "protocolVersion": executionpool.ProtocolVersion, + "streaming": true, + "media": true, + "taskKinds": []any{"*"}, } } diff --git a/apps/api/internal/runner/remote_execution.go b/apps/api/internal/runner/remote_execution.go new file mode 100644 index 0000000..f087767 --- /dev/null +++ b/apps/api/internal/runner/remote_execution.go @@ -0,0 +1,264 @@ +package runner + +import ( + "bufio" + "context" + "crypto/tls" + "crypto/x509" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "strings" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" + "github.com/google/uuid" +) + +const ( + workerExecutionAudience = "easyai-worker-execution" + workerExecutionPath = "/internal/v1/executions" +) + +type InternalExecutionRequest struct { + TaskID string `json:"task_id"` + LeaseID string `json:"lease_id"` + Stream bool `json:"stream"` +} + +type InternalExecutionFrame struct { + Type string `json:"type"` + Delta *clients.StreamDeltaEvent `json:"delta,omitempty"` + Output map[string]any `json:"output,omitempty"` + Wire *clients.WireResponse `json:"wire,omitempty"` + Code string `json:"code,omitempty"` + Message string `json:"message,omitempty"` + Status int `json:"status,omitempty"` + Retryable bool `json:"retryable,omitempty"` + Details map[string]any `json:"details,omitempty"` +} + +func (s *Service) executeRouted(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) { + if task.AsyncMode || !s.cfg.RunsPublicHTTP() || !s.routingEnabled() || task.RunMode == "simulation" { + return s.executeLocal(ctx, task, user, onDelta) + } + candidate, err := s.routingCandidateForTask(ctx, task, user) + if err != nil { + if s.routingEnforced() { + s.observeExecutionPoolRouting("unavailable") + return Result{}, upstreamRouteUnavailable(err) + } + return s.executeLocal(ctx, task, user, onDelta) + } + decision, profile, err := s.selectExecutionPool(ctx, task, candidate) + if err != nil { + _ = s.coordinationStore.RequestRouteProbe(context.WithoutCancel(ctx), profile.Key) + if s.routingEnforced() { + s.observeExecutionPoolRouting("unavailable") + return Result{}, upstreamRouteUnavailable(err) + } + _ = s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{ + TaskID: task.ID, RouteProfileKey: profile.Key, RoutingVersion: routingVersion, + PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID, + Reason: "shadow_no_eligible_pool", Snapshot: map[string]any{"mode": "shadow", "error": safeRoutingError(err)}, + }) + s.observeExecutionPoolRouting("shadow") + return s.executeLocal(ctx, task, user, onDelta) + } + snapshot := routingDecisionSnapshot(decision) + if !s.routingEnforced() { + snapshot["suggestedPoolId"] = decision.PoolID + _ = s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{ + TaskID: task.ID, RouteProfileKey: profile.Key, RoutingVersion: routingVersion, + PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID, + Reason: decision.Reason, Snapshot: snapshot, + }) + s.observeExecutionPoolRouting("shadow") + return s.executeLocal(ctx, task, user, onDelta) + } + if err := s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{ + TaskID: task.ID, PoolID: decision.PoolID, RouteProfileKey: profile.Key, + PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID, + RoutingVersion: routingVersion, Reason: decision.Reason, Snapshot: snapshot, + }); err != nil { + return Result{}, err + } + s.observeExecutionPoolRouting("selected") + return s.executeThroughWorker(ctx, task, decision.PoolID, onDelta) +} + +func (s *Service) executeLocal(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) { + return s.executeWithToken(ctx, task, user, onDelta, uuid.NewString()) +} + +func (s *Service) executeThroughWorker(ctx context.Context, task store.GatewayTask, poolID string, onDelta clients.StreamDelta) (Result, error) { + nonce := uuid.NewString() + lease, err := s.coordinationStore.ReserveWorkerExecution(ctx, task.ID, poolID, nonce, 30*time.Second) + if err != nil { + s.observeExecutionPoolRouting("capacity_rejected") + return Result{}, upstreamRouteUnavailable(err) + } + if err := executionpool.ValidateAdvertisedEndpoint( + lease.Endpoint, splitTrimmed(s.cfg.WorkerEndpointAllowedSuffixes), s.cfg.WorkerEndpointAllowPrivate, + ); err != nil { + _ = s.coordinationStore.ReleaseWorkerExecutionLease(context.WithoutCancel(ctx), lease.LeaseID) + return Result{}, &clients.ClientError{Code: "worker_endpoint_untrusted", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: false} + } + signer := executionpool.TokenSigner{Secret: []byte(s.cfg.WorkerExecutionSecret)} + token, err := signer.Sign(executionpool.ExecutionClaims{ + Audience: workerExecutionAudience, TaskID: task.ID, PoolID: poolID, + WorkerID: lease.WorkerID, Nonce: nonce, ExpiresAt: time.Now().Add(30 * time.Second).Unix(), + }) + if err != nil { + _ = s.coordinationStore.ReleaseWorkerExecutionLease(context.WithoutCancel(ctx), lease.LeaseID) + return Result{}, err + } + client, err := s.workerExecutionHTTPClient() + if err != nil { + return Result{}, err + } + transport := httpExecutionTransport{client: client} + response, err := transport.Execute(ctx, executionpool.WorkerDescriptor{ + WorkerID: lease.WorkerID, InstanceID: lease.InstanceID, PoolID: poolID, Endpoint: lease.Endpoint, + }, executionpool.ExecutionRequest{ + TaskID: task.ID, PoolID: poolID, WorkerID: lease.WorkerID, LeaseID: lease.LeaseID, + AuthorizationToken: token, Stream: onDelta != nil, + }) + if err != nil { + latest, readErr := s.store.GetTask(context.WithoutCancel(ctx), task.ID) + if readErr == nil && latest.SubmissionState != "not_started" { + return Result{Task: latest}, &clients.ClientError{ + Code: "upstream_timeout", Message: "worker transport ended after upstream submission began", + StatusCode: http.StatusGatewayTimeout, Retryable: true, + } + } + return Result{}, &clients.ClientError{Code: "worker_transport_error", Message: err.Error(), StatusCode: http.StatusServiceUnavailable, Retryable: true} + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + payload, _ := io.ReadAll(io.LimitReader(response.Body, 64*1024)) + return Result{}, &clients.ClientError{Code: "worker_transport_error", Message: strings.TrimSpace(string(payload)), StatusCode: response.StatusCode, Retryable: response.StatusCode >= 500} + } + var output map[string]any + var wire *clients.WireResponse + scanner := bufio.NewScanner(response.Body) + scanner.Buffer(make([]byte, 4096), 1024*1024) + for scanner.Scan() { + var frame InternalExecutionFrame + if err := json.Unmarshal(scanner.Bytes(), &frame); err != nil { + return Result{}, &clients.ClientError{Code: "worker_protocol_error", Message: "invalid worker execution frame", StatusCode: http.StatusBadGateway, Retryable: false} + } + switch frame.Type { + case "delta": + if onDelta != nil && frame.Delta != nil { + if err := onDelta(*frame.Delta); err != nil { + return Result{}, err + } + } + case "result": + output, wire = frame.Output, frame.Wire + case "error": + return Result{}, &clients.ClientError{ + Code: frame.Code, Message: frame.Message, StatusCode: frame.Status, + Retryable: frame.Retryable, Details: frame.Details, + } + } + } + if err := scanner.Err(); err != nil { + return Result{}, &clients.ClientError{Code: "worker_transport_error", Message: err.Error(), StatusCode: http.StatusBadGateway, Retryable: true} + } + finished, err := s.store.GetTask(ctx, task.ID) + if err != nil { + return Result{}, err + } + return Result{Task: finished, Output: output, Wire: wire}, nil +} + +func (s *Service) ExecuteInternal(ctx context.Context, authorization string, input InternalExecutionRequest, onDelta clients.StreamDelta) (Result, error) { + if s.cfg.RunsPublicHTTP() || !s.cfg.RunsAsyncExecutionWorker() { + return Result{}, errors.New("internal execution is only available on worker processes") + } + token := strings.TrimSpace(strings.TrimPrefix(authorization, "Worker ")) + if token == authorization || token == "" { + return Result{}, errors.New("worker execution authorization is required") + } + signer := executionpool.TokenSigner{Secret: []byte(s.cfg.WorkerExecutionSecret)} + claims, err := signer.Verify(token, workerExecutionAudience) + if err != nil { + return Result{}, err + } + localWorkerID := firstNonEmptyString(s.cfg.WorkerID, s.workerInstanceID) + if claims.TaskID != input.TaskID || claims.PoolID != s.cfg.ExecutionPoolID || claims.WorkerID != localWorkerID { + return Result{}, errors.New("worker execution claims do not match this worker") + } + lease, err := s.coordinationStore.ConsumeWorkerExecutionLease(ctx, input.LeaseID, claims.Nonce) + if err != nil { + return Result{}, err + } + defer func() { + _ = s.coordinationStore.ReleaseWorkerExecutionLease(context.WithoutCancel(ctx), lease.LeaseID) + }() + if lease.TaskID != input.TaskID || lease.PoolID != s.cfg.ExecutionPoolID || lease.WorkerID != localWorkerID || lease.InstanceID != s.workerInstanceID { + return Result{}, errors.New("worker execution lease does not match this worker") + } + loadLease, admitted := s.tryStartWorkerTask() + if !admitted { + return Result{}, &clients.ClientError{Code: "worker_capacity_unavailable", Message: "worker has no safe execution capacity", StatusCode: http.StatusServiceUnavailable, Retryable: true} + } + defer loadLease.Release() + ctx = context.WithValue(ctx, workerLoadLeaseContextKey{}, loadLease) + task, err := s.store.GetTask(ctx, input.TaskID) + if err != nil { + return Result{}, err + } + return s.executeLocal(ctx, task, authUserFromTask(task), onDelta) +} + +func (s *Service) workerExecutionHTTPClient() (*http.Client, error) { + transport := http.DefaultTransport.(*http.Transport).Clone() + transport.Proxy = nil + if strings.TrimSpace(s.cfg.WorkerExecutionCAFile) != "" { + caBytes, err := os.ReadFile(s.cfg.WorkerExecutionCAFile) + if err != nil { + return nil, err + } + roots := x509.NewCertPool() + if !roots.AppendCertsFromPEM(caBytes) { + return nil, errors.New("worker execution CA file does not contain a certificate") + } + transport.TLSClientConfig = &tls.Config{RootCAs: roots, MinVersion: tls.VersionTLS12} + if strings.TrimSpace(s.cfg.WorkerExecutionCertFile) != "" || strings.TrimSpace(s.cfg.WorkerExecutionKeyFile) != "" { + certificate, err := tls.LoadX509KeyPair(s.cfg.WorkerExecutionCertFile, s.cfg.WorkerExecutionKeyFile) + if err != nil { + return nil, err + } + transport.TLSClientConfig.Certificates = []tls.Certificate{certificate} + } + } + return &http.Client{Transport: transport}, nil +} + +func internalExecutionStatus(err error) int { + var clientErr *clients.ClientError + if errors.As(err, &clientErr) && clientErr.StatusCode > 0 { + return clientErr.StatusCode + } + return http.StatusInternalServerError +} + +func InternalExecutionFrameForError(err error) InternalExecutionFrame { + var clientErr *clients.ClientError + if errors.As(err, &clientErr) { + return InternalExecutionFrame{ + Type: "error", Code: clientErr.Code, Message: clientErr.Message, + Status: internalExecutionStatus(err), Retryable: clientErr.Retryable, Details: clientErr.Details, + } + } + return InternalExecutionFrame{Type: "error", Code: "worker_execution_failed", Message: fmt.Sprint(err), Status: http.StatusInternalServerError} +} diff --git a/apps/api/internal/runner/river_execution_broker.go b/apps/api/internal/runner/river_execution_broker.go new file mode 100644 index 0000000..b48796c --- /dev/null +++ b/apps/api/internal/runner/river_execution_broker.go @@ -0,0 +1,99 @@ +package runner + +import ( + "context" + "errors" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/jackc/pgx/v5" + "github.com/riverqueue/river" + "github.com/riverqueue/river/rivertype" +) + +// riverExecutionBroker is the River adapter for the platform-neutral broker +// port. Queue names and River uniqueness semantics do not escape this adapter. +type riverExecutionBroker struct { + service *Service +} + +func (broker riverExecutionBroker) Publish( + ctx context.Context, + poolID string, + taskID string, + notBefore time.Time, +) (int64, error) { + tx, err := broker.service.store.Pool().Begin(ctx) + if err != nil { + return 0, err + } + defer func() { + rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = tx.Rollback(rollbackCtx) + }() + jobID, err := broker.publishTx(ctx, tx, poolID, taskID, notBefore, 2) + if err != nil { + return 0, err + } + if err := tx.Commit(ctx); err != nil { + return 0, err + } + return jobID, nil +} + +func (broker riverExecutionBroker) publishTx( + ctx context.Context, + tx pgx.Tx, + poolID string, + taskID string, + notBefore time.Time, + priority int, +) (int64, error) { + riverClient := broker.service.asyncControlClient() + if riverClient == nil { + return 0, errors.New("River execution broker is not started") + } + opts := riverExecutionInsertOptions(poolID, notBefore, priority) + result, err := riverClient.InsertTx(ctx, tx, asyncTaskArgs{TaskID: taskID}, opts) + if err != nil { + return 0, err + } + if result.Job == nil { + return 0, nil + } + if _, err := tx.Exec(ctx, ` +UPDATE gateway_tasks +SET river_job_id = $2, + updated_at = now() +WHERE id = $1::uuid`, taskID, result.Job.ID); err != nil { + return 0, err + } + return result.Job.ID, nil +} + +func riverExecutionInsertOptions(poolID string, notBefore time.Time, priority int) *river.InsertOpts { + if priority < 1 { + priority = 2 + } + opts := &river.InsertOpts{ + MaxAttempts: 1000, + Priority: priority, + Queue: executionpool.QueueName(poolID), + Tags: []string{"gateway-task"}, + UniqueOpts: river.UniqueOpts{ + ByArgs: true, ByQueue: true, + ByState: []rivertype.JobState{ + rivertype.JobStateAvailable, rivertype.JobStatePending, + rivertype.JobStateRetryable, rivertype.JobStateRunning, + rivertype.JobStateScheduled, + }, + }, + } + if !notBefore.IsZero() { + opts.ScheduledAt = notBefore + } + return opts +} + +var _ executionpool.ExecutionBroker = riverExecutionBroker{} diff --git a/apps/api/internal/runner/route_prober.go b/apps/api/internal/runner/route_prober.go new file mode 100644 index 0000000..416c3ac --- /dev/null +++ b/apps/api/internal/runner/route_prober.go @@ -0,0 +1,168 @@ +package runner + +import ( + "context" + "errors" + "log/slog" + "math" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +type routeProbeState struct { + next time.Time + lastRequestObserved time.Time +} + +func (s *Service) StartRouteHealthProber(ctx context.Context) { + if !s.cfg.RunsRouteProber() { + return + } + go s.runRouteHealthProber(ctx) +} + +func (s *Service) runRouteHealthProber(ctx context.Context) { + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + states := make(map[string]routeProbeState) + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + targets, err := s.coordinationStore.ListRouteProbeTargets(ctx) + if err != nil { + s.logRouteProbeFailure("list_targets", err) + continue + } + requested, requestErr := s.coordinationStore.ListRequestedRouteProbes(ctx) + if requestErr != nil { + s.logRouteProbeFailure("list_requests", requestErr) + continue + } + now := time.Now() + for _, target := range targets { + profile, profileErr := routeProfileForCandidate(target.Candidate) + if profileErr != nil { + continue + } + state := states[profile.Key] + requestedAt := requested[profile.Key] + requestDue := !requestedAt.IsZero() && requestedAt.After(state.lastRequestObserved) + if !requestDue && now.Before(state.next) { + continue + } + interval := time.Duration(s.cfg.RouteProbeColdIntervalSeconds) * time.Second + if target.Hot { + interval = time.Duration(s.cfg.RouteProbeHotIntervalSeconds) * time.Second + } + states[profile.Key] = routeProbeState{next: now.Add(interval), lastRequestObserved: requestedAt} + go s.probeRouteTarget(ctx, target, profile) + } + } + } +} + +func (s *Service) probeRouteTarget(ctx context.Context, target store.RouteProbeTarget, profile executionpool.RouteProfile) { + leadership, locked, err := s.coordinationStore.TryAcquireRouteProbeLeadership(ctx, s.cfg.ExecutionPoolID, profile.Key) + if err != nil || !locked { + return + } + defer leadership.Release() + client, err := s.httpClientForCandidate(target.Candidate, false, "") + if err != nil { + s.logRouteProbeFailure("prepare_client", err) + return + } + probe := executionpool.HTTPProbe{Client: client} + result := probe.Probe(ctx, executionpool.ProbeTarget{ + RouteProfile: profile, + URL: routeEndpointForCandidate(target.Candidate), + Timeout: time.Duration(s.cfg.RouteProbeTimeoutMS) * time.Millisecond, + }) + previous := executionpool.RouteHealth{} + if history, historyErr := s.coordinationStore.ListRouteHealth(ctx, profile.Key, time.Now()); historyErr == nil { + for _, item := range history { + if item.PoolID == s.cfg.ExecutionPoolID { + previous = item + break + } + } + } + health := nextRouteHealth(previous, s.cfg.ExecutionPoolID, profile.Key, result) + s.observeRouteProbe(string(health.State)) + if err := s.coordinationStore.RecordRouteHealth(context.WithoutCancel(ctx), health); err != nil { + s.logRouteProbeFailure("record", err) + } +} + +func (s *Service) observeRouteProbe(state string) { + observer, ok := s.billingMetrics.(interface{ ObserveRouteProbe(string) }) + if ok { + observer.ObserveRouteProbe(state) + } +} + +func nextRouteHealth(previous executionpool.RouteHealth, poolID, profileKey string, result executionpool.ProbeResult) executionpool.RouteHealth { + health := previous + health.PoolID = poolID + health.RouteProfileKey = profileKey + health.SampleCount++ + health.SampledAt = result.SampledAt + health.ExpiresAt = result.SampledAt.Add(45 * time.Second) + observation := 0.0 + if result.Reachable { + observation = 1 + health.ConsecutiveSuccesses++ + health.ConsecutiveFailures = 0 + if health.ConsecutiveSuccesses >= 2 || previous.State == executionpool.RouteHealthy { + health.State = executionpool.RouteHealthy + } else { + health.State = executionpool.RouteDegraded + } + } else { + health.ConsecutiveFailures++ + health.ConsecutiveSuccesses = 0 + switch { + case health.ConsecutiveFailures >= 3: + health.State = executionpool.RouteUnreachable + case previous.State == executionpool.RouteHealthy || previous.State == executionpool.RouteDegraded: + health.State = executionpool.RouteDegraded + default: + health.State = executionpool.RouteUnknown + } + } + if previous.SampleCount == 0 { + health.SuccessRate = observation + } else { + health.SuccessRate = 0.8*previous.SuccessRate + 0.2*observation + } + connectTLS := result.TCPDuration + result.TLSDuration + if connectTLS > 0 { + health.ConnectTLSP95 = ewmaDuration(previous.ConnectTLSP95, connectTLS) + } + if result.FirstByte > 0 { + if previous.FirstByteP95 > 0 { + delta := math.Abs(float64(result.FirstByte - previous.FirstByteP95)) + health.JitterP95 = ewmaDuration(previous.JitterP95, time.Duration(delta)) + } + health.FirstByteP95 = ewmaDuration(previous.FirstByteP95, result.FirstByte) + } + return health +} + +func ewmaDuration(previous, current time.Duration) time.Duration { + if previous <= 0 { + return current + } + return time.Duration(0.8*float64(previous) + 0.2*float64(current)) +} + +func (s *Service) logRouteProbeFailure(stage string, err error) { + if err == nil || errors.Is(err, context.Canceled) || s.logger == nil { + return + } + s.logger.Log(context.Background(), slog.LevelWarn, "route health probe failed", "stage", stage, "error_category", "route_probe_failed") +} diff --git a/apps/api/internal/runner/route_prober_test.go b/apps/api/internal/runner/route_prober_test.go new file mode 100644 index 0000000..024b637 --- /dev/null +++ b/apps/api/internal/runner/route_prober_test.go @@ -0,0 +1,49 @@ +package runner + +import ( + "testing" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +func TestNextRouteHealthRequiresFailuresAndRecoverySamples(t *testing.T) { + now := time.Now() + health := executionpool.RouteHealth{} + for index := 0; index < 2; index++ { + health = nextRouteHealth(health, "pool", "route", executionpool.ProbeResult{SampledAt: now.Add(time.Duration(index) * time.Second), ErrorClass: "network"}) + if health.State == executionpool.RouteUnreachable { + t.Fatalf("became unreachable after %d failures", index+1) + } + } + health = nextRouteHealth(health, "pool", "route", executionpool.ProbeResult{SampledAt: now.Add(2 * time.Second), ErrorClass: "network"}) + if health.State != executionpool.RouteUnreachable { + t.Fatalf("state=%s, want unreachable", health.State) + } + health = nextRouteHealth(health, "pool", "route", executionpool.ProbeResult{SampledAt: now.Add(3 * time.Second), Reachable: true, TCPDuration: 5 * time.Millisecond, FirstByte: 20 * time.Millisecond}) + if health.State != executionpool.RouteDegraded { + t.Fatalf("state=%s, want degraded during recovery", health.State) + } + health = nextRouteHealth(health, "pool", "route", executionpool.ProbeResult{SampledAt: now.Add(4 * time.Second), Reachable: true, TCPDuration: 5 * time.Millisecond, FirstByte: 20 * time.Millisecond}) + if health.State != executionpool.RouteHealthy { + t.Fatalf("state=%s, want healthy", health.State) + } + if health.ExpiresAt.Sub(health.SampledAt) != 45*time.Second { + t.Fatalf("health TTL=%s", health.ExpiresAt.Sub(health.SampledAt)) + } +} + +func TestPinCandidatesToRoutingAssignment(t *testing.T) { + candidates := []store.RuntimeModelCandidate{ + {PlatformID: "platform-a", PlatformModelID: "model-a"}, + {PlatformID: "platform-b", PlatformModelID: "model-b"}, + } + pinned := pinCandidatesToRoutingAssignment(candidates, "platform-b", "model-b") + if len(pinned) != 1 || pinned[0].PlatformModelID != "model-b" { + t.Fatalf("pinned=%+v", pinned) + } + if got := pinCandidatesToRoutingAssignment(candidates, "platform-a", "missing"); len(got) != 0 { + t.Fatalf("unexpected fallback: %+v", got) + } +} diff --git a/apps/api/internal/runner/routing.go b/apps/api/internal/runner/routing.go new file mode 100644 index 0000000..4b6a7bf --- /dev/null +++ b/apps/api/internal/runner/routing.go @@ -0,0 +1,332 @@ +package runner + +import ( + "context" + "errors" + "net/http" + "sort" + "strings" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/clients" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/netproxy" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +const ( + routingVersion = "execution-pool-v1" + asyncRouteWaitMaximum = 10 * time.Second +) + +func (s *Service) routingEnabled() bool { + mode := strings.ToLower(strings.TrimSpace(s.cfg.RoutingMode)) + return mode == "shadow" || mode == "enforced" +} + +func (s *Service) routingEnforced() bool { + return strings.EqualFold(strings.TrimSpace(s.cfg.RoutingMode), "enforced") +} + +func (s *Service) routeAsyncTask(ctx context.Context, task *store.GatewayTask, user *auth.User, preferred store.RuntimeModelCandidate) error { + if task == nil || !task.AsyncMode || !s.routingEnabled() || task.RunMode == "simulation" { + return nil + } + startedAt := time.Now() + delays := []time.Duration{0, 2 * time.Second, 3 * time.Second, 4 * time.Second} + var lastErr error + for _, delay := range delays { + if delay > 0 { + timer := time.NewTimer(delay) + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return ctx.Err() + case <-timer.C: + } + } + candidate := preferred + if strings.TrimSpace(candidate.PlatformID) == "" { + resolved, err := s.routingCandidateForTask(ctx, *task, user) + if err != nil { + lastErr = err + continue + } + candidate = resolved + } + decision, profile, err := s.selectExecutionPool(ctx, *task, candidate) + if err == nil { + poolID := decision.PoolID + snapshot := routingDecisionSnapshot(decision) + if !s.routingEnforced() { + s.observeExecutionPoolRouting("shadow") + snapshot["suggestedPoolId"] = poolID + poolID = "" + } + if s.routingEnforced() { + s.observeExecutionPoolRouting("selected") + } + if err := s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{ + TaskID: task.ID, PoolID: poolID, RouteProfileKey: profile.Key, + PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID, + RoutingVersion: routingVersion, Reason: decision.Reason, Snapshot: snapshot, + }); err != nil { + return err + } + task.AssignedPoolID = poolID + task.RouteProfileKey = profile.Key + task.RoutingVersion = routingVersion + task.RoutingReason = decision.Reason + task.RoutingSnapshot = snapshot + return nil + } + _ = s.coordinationStore.RequestRouteProbe(context.WithoutCancel(ctx), profile.Key) + lastErr = err + if !s.routingEnforced() { + s.observeExecutionPoolRouting("shadow") + _ = s.store.AssignTaskRouting(ctx, store.TaskRoutingDecision{ + TaskID: task.ID, RouteProfileKey: profile.Key, RoutingVersion: routingVersion, + PlatformID: candidate.PlatformID, PlatformModelID: candidate.PlatformModelID, + Reason: "shadow_no_eligible_pool", Snapshot: map[string]any{"mode": "shadow", "error": safeRoutingError(err)}, + }) + return nil + } + } + s.observeExecutionPoolRouting("unavailable") + if remaining := asyncRouteWaitMaximum - time.Since(startedAt); remaining > 0 { + timer := time.NewTimer(remaining) + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return ctx.Err() + case <-timer.C: + } + } + return upstreamRouteUnavailable(lastErr) +} + +func (s *Service) observeExecutionPoolRouting(outcome string) { + observer, ok := s.billingMetrics.(interface{ ObserveExecutionPoolRouting(string) }) + if ok { + observer.ObserveExecutionPoolRouting(outcome) + } +} + +func (s *Service) routingCandidateForTask(ctx context.Context, task store.GatewayTask, user *auth.User) (store.RuntimeModelCandidate, error) { + body := normalizeRequest(task.Kind, task.Request) + modelType := modelTypeFromKind(task.Kind, body) + candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user) + if err != nil { + return store.RuntimeModelCandidate{}, err + } + candidates, err = filterCandidatesByRequestedPlatform(candidates, body) + if err == nil { + candidates, _, err = filterRuntimeCandidatesByRequest(task.Kind, task.Model, modelType, body, candidates) + } + if err == nil { + candidates, _, err = filterRuntimeCandidatesByOutputTokens(task.Kind, task.Model, modelType, body, candidates) + } + if err != nil { + return store.RuntimeModelCandidate{}, err + } + if len(candidates) == 0 { + return store.RuntimeModelCandidate{}, store.ErrNoModelCandidate + } + return candidates[0], nil +} + +func (s *Service) selectExecutionPool(ctx context.Context, task store.GatewayTask, candidate store.RuntimeModelCandidate) (executionpool.Decision, executionpool.RouteProfile, error) { + profile, err := routeProfileForCandidate(candidate) + if err != nil { + return executionpool.Decision{}, profile, err + } + now := time.Now() + pools, err := s.coordinationStore.ListExecutionPools(ctx) + if err != nil { + return executionpool.Decision{}, profile, err + } + workers, err := s.coordinationStore.ListWorkers(ctx, now) + if err != nil { + return executionpool.Decision{}, profile, err + } + health, err := s.coordinationStore.ListRouteHealth(ctx, profile.Key, now) + if err != nil { + return executionpool.Decision{}, profile, err + } + capacity, err := s.coordinationStore.ListCapacity(ctx, now) + if err != nil { + return executionpool.Decision{}, profile, err + } + workersByPool := make(map[string][]executionpool.WorkerDescriptor) + for _, worker := range workers { + workersByPool[worker.PoolID] = append(workersByPool[worker.PoolID], worker) + } + healthByPool := make(map[string]executionpool.RouteHealth) + for _, item := range health { + healthByPool[item.PoolID] = item + } + capacityByPool := make(map[string]executionpool.CapacitySnapshot) + for _, item := range capacity { + capacityByPool[item.PoolID] = item + } + selection := make([]executionpool.SelectionCandidate, 0, len(pools)) + for _, pool := range pools { + poolWorkers := workersByPool[pool.ID] + selection = append(selection, executionpool.SelectionCandidate{ + Pool: pool, Health: healthByPool[pool.ID], Capacity: capacityByPool[pool.ID], + CapabilityMatched: workerCapabilityMatched(poolWorkers, task), + }) + } + selector := executionpool.NewSelector() + returnDecision, selectErr := selector.Select(executionpool.SelectionRequest{ + Candidates: selection, CurrentPoolID: task.AssignedPoolID, Now: now, + }) + if selectErr != nil { + return returnDecision, profile, selectErr + } + poolDecisions := make(map[string]executionpool.Decision, len(selection)) + scores := make(map[string]float64, len(selection)) + for _, candidate := range selection { + decision, candidateErr := selector.Select(executionpool.SelectionRequest{Candidates: []executionpool.SelectionCandidate{candidate}, Now: now}) + if candidateErr == nil { + poolDecisions[decision.PoolID] = decision + scores[decision.PoolID] = decision.Score + } + } + preference, err := s.coordinationStore.ResolveRoutePreference(ctx, profile.Key, returnDecision.PoolID, scores, now) + if err != nil { + return executionpool.Decision{}, profile, err + } + if preference.CurrentPoolID != returnDecision.PoolID { + preferred, ok := poolDecisions[preference.CurrentPoolID] + if !ok { + return executionpool.Decision{}, profile, executionpool.ErrNoEligiblePool + } + preferred.Reason = "route_preference_hysteresis" + preferred.Rejected = returnDecision.Rejected + returnDecision = preferred + } + return returnDecision, profile, nil +} + +func routeProfileForCandidate(candidate store.RuntimeModelCandidate) (executionpool.RouteProfile, error) { + baseURL := routeEndpointForCandidate(candidate) + host := executionpool.EndpointHost(baseURL) + if host == "" { + return executionpool.RouteProfile{}, errors.New("candidate does not expose a probeable upstream endpoint") + } + proxyMode := "none" + if proxyConfig, err := netproxy.Normalize(netproxy.FromPlatformConfig(candidate.PlatformConfig)); err == nil { + proxyMode = string(proxyConfig.Mode) + } + protocol := firstNonEmptyString(candidate.ResponseProtocol, candidate.SpecType, candidate.Provider) + profile := executionpool.RouteProfile{ + Provider: candidate.Provider, Protocol: protocol, EndpointHost: host, + ProxyMode: proxyMode, ConfigRevision: candidate.PlatformID + ":" + candidate.PlatformModelID, + } + profile.Key = executionpool.RouteProfileKey(profile.Provider, profile.Protocol, profile.EndpointHost, profile.ProxyMode, profile.ConfigRevision) + return profile, nil +} + +func routeEndpointForCandidate(candidate store.RuntimeModelCandidate) string { + baseURL := strings.TrimSpace(candidate.BaseURL) + if baseURL == "" { + baseURL = firstString(candidate.PlatformConfig, "endpoint", "baseURL", "base_url") + } + return baseURL +} + +func workerCapabilityMatched(workers []executionpool.WorkerDescriptor, task store.GatewayTask) bool { + for _, worker := range workers { + if worker.ProtocolVersion != executionpool.ProtocolVersion || worker.AvailableCapacity() < 1 { + continue + } + if requestStreamEnabled(task.Request) && !boolValue(worker.Capabilities["streaming"]) { + continue + } + if !workerSupportsTaskKind(worker.Capabilities["taskKinds"], task.Kind) { + continue + } + return true + } + return false +} + +func workerSupportsTaskKind(raw any, taskKind string) bool { + taskKind = strings.ToLower(strings.TrimSpace(taskKind)) + if taskKind == "" { + return false + } + values, ok := raw.([]any) + if !ok { + if stringsList, stringsOK := raw.([]string); stringsOK { + values = make([]any, 0, len(stringsList)) + for _, value := range stringsList { + values = append(values, value) + } + } + } + for _, value := range values { + capability, _ := value.(string) + capability = strings.ToLower(strings.TrimSpace(capability)) + if capability == "*" || capability == taskKind || + (strings.HasSuffix(capability, ".*") && strings.HasPrefix(taskKind, strings.TrimSuffix(capability, "*"))) { + return true + } + } + return false +} + +func requestStreamEnabled(body map[string]any) bool { + value, _ := body["stream"].(bool) + return value +} + +func routingDecisionSnapshot(decision executionpool.Decision) map[string]any { + rejected := make([]string, 0, len(decision.Rejected)) + for poolID, reason := range decision.Rejected { + rejected = append(rejected, poolID+":"+reason) + } + sort.Strings(rejected) + return map[string]any{ + "mode": "selected", "score": decision.Score, "components": decision.Components, + "rejected": rejected, + } +} + +func safeRoutingError(err error) string { + if errors.Is(err, executionpool.ErrNoEligiblePool) { + return "no_eligible_pool" + } + return "routing_unavailable" +} + +func upstreamRouteUnavailable(cause error) error { + details := map[string]any{"retryAfterSeconds": 10} + if cause != nil { + details["reason"] = safeRoutingError(cause) + } + return &clients.ClientError{ + Code: "upstream_route_unavailable", Message: "no execution pool can currently reach the upstream route", + Retryable: true, StatusCode: http.StatusServiceUnavailable, Details: details, + } +} + +func firstString(values map[string]any, keys ...string) string { + for _, key := range keys { + if value, ok := values[key].(string); ok && strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} + +func boolValue(value any) bool { + result, _ := value.(bool) + return result +} diff --git a/apps/api/internal/runner/routing_candidate.go b/apps/api/internal/runner/routing_candidate.go new file mode 100644 index 0000000..efc7efd --- /dev/null +++ b/apps/api/internal/runner/routing_candidate.go @@ -0,0 +1,12 @@ +package runner + +import "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" + +func pinCandidatesToRoutingAssignment(candidates []store.RuntimeModelCandidate, platformID, platformModelID string) []store.RuntimeModelCandidate { + for _, candidate := range candidates { + if candidate.PlatformModelID == platformModelID && (platformID == "" || candidate.PlatformID == platformID) { + return []store.RuntimeModelCandidate{candidate} + } + } + return nil +} diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 7216fba..7e42211 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -21,7 +21,6 @@ import ( scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script" "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" "github.com/easyai/easyai-ai-gateway/apps/api/internal/workerload" - "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/riverqueue/river" ) @@ -77,24 +76,24 @@ type TaskQueuedError struct { Delay time.Duration } -type upstreamSubmissionUnknownError struct { +type submissionConfirmationPendingError struct { AttemptID string Cause error } -func shouldClassifyUpstreamSubmissionUnknown(simulated bool, submissionStatus string, err error) bool { +func shouldEnterSubmissionConfirmationPending(simulated bool, submissionStatus string, err error) bool { return !simulated && submissionStatus == "submitting" && clients.ErrorCode(err) != "timeout" } -func (e *upstreamSubmissionUnknownError) Error() string { - return "upstream submission result is unknown" +func (e *submissionConfirmationPendingError) Error() string { + return "upstream submission confirmation timed out" } -func (e *upstreamSubmissionUnknownError) ErrorCode() string { - return "upstream_submission_unknown" +func (e *submissionConfirmationPendingError) ErrorCode() string { + return "upstream_timeout" } -func (e *upstreamSubmissionUnknownError) Unwrap() error { +func (e *submissionConfirmationPendingError) Unwrap() error { return e.Cause } @@ -131,6 +130,18 @@ func NewWithStores( if cfg.AsyncWorkerHardLimit == 0 { cfg.AsyncWorkerHardLimit = 2048 } + if strings.TrimSpace(cfg.ExecutionPoolID) == "" { + cfg.ExecutionPoolID = "legacy-default" + } + if cfg.RouteProbeHotIntervalSeconds == 0 { + cfg.RouteProbeHotIntervalSeconds = 15 + } + if cfg.RouteProbeColdIntervalSeconds == 0 { + cfg.RouteProbeColdIntervalSeconds = 60 + } + if cfg.RouteProbeTimeoutMS == 0 { + cfg.RouteProbeTimeoutMS = 3000 + } if cfg.AsyncWorkerInstanceHardLimit == 0 { cfg.AsyncWorkerInstanceHardLimit = 32 } @@ -259,15 +270,15 @@ func (s *Service) observeResultURLSigning(outcome string) { } func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *auth.User) (Result, error) { - return s.execute(ctx, task, user, nil) + return s.executeRouted(ctx, task, user, nil) } func (s *Service) ExecuteStream(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) { - return s.execute(ctx, task, user, onDelta) + return s.executeRouted(ctx, task, user, onDelta) } func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) { - return s.executeWithToken(ctx, task, user, onDelta, uuid.NewString()) + return s.executeRouted(ctx, task, user, onDelta) } func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta, executionToken string) (Result, error) { @@ -511,6 +522,17 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, return Result{Task: failed, Output: failed.Result}, err } } + if task.RoutingPlatformModelID != "" { + candidates = pinCandidatesToRoutingAssignment(candidates, task.RoutingPlatformID, task.RoutingPlatformModelID) + if len(candidates) == 0 { + err = &clients.ClientError{Code: "routing_candidate_unavailable", Message: "the routed upstream candidate is no longer available", Retryable: true} + failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(err), err.Error(), task.RunMode == "simulation", err) + if finishErr != nil { + return Result{}, finishErr + } + return Result{Task: failed, Output: failed.Result}, err + } + } for _, candidate := range candidates { if candidate.LoadAvoided { s.observeCandidateRouting("full_avoided") @@ -1096,11 +1118,12 @@ candidatesLoop: } return Result{Task: finished, Output: output, Wire: response.Wire}, nil } - var submissionUnknown *upstreamSubmissionUnknownError - if errors.As(err, &submissionUnknown) { + var confirmationPending *submissionConfirmationPendingError + if errors.As(err, &confirmationPending) { + _ = s.store.SetTaskSubmissionState(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, "submission_confirmation_pending") review, reviewErr := s.store.FinishTaskManualReview(context.WithoutCancel(ctx), store.FinishTaskManualReviewInput{ - TaskID: task.ID, ExecutionToken: task.ExecutionToken, AttemptID: submissionUnknown.AttemptID, TaskStatus: "failed", - Code: "upstream_submission_unknown", Message: submissionUnknown.Error(), + TaskID: task.ID, ExecutionToken: task.ExecutionToken, AttemptID: confirmationPending.AttemptID, TaskStatus: "failed", + Code: "upstream_timeout", Message: confirmationPending.Error(), PricingSnapshot: candidatePricing.Snapshot, RequestFingerprint: pricingRequestFingerprint(task.Kind, task.Model, candidateBody), }) @@ -1108,8 +1131,8 @@ candidatesLoop: return Result{}, reviewErr } walletReservationFinalized = true - s.logger.Warn("upstream submission requires manual review", "taskID", task.ID, "attemptID", submissionUnknown.AttemptID, "error_category", "upstream_submission_unknown") - return Result{Task: review, Output: review.Result}, submissionUnknown + s.logger.Warn("upstream submission confirmation pending", "taskID", task.ID, "attemptID", confirmationPending.AttemptID, "error_category", "upstream_timeout") + return Result{Task: review, Output: review.Result}, confirmationPending } if isLocalRateLimitError(err) { lastErr = err @@ -1592,6 +1615,11 @@ func (s *Service) runCandidate( return nil } var submissionWire *clients.WireResponse + var routeObserver *passiveRouteObserver + if task.RouteProfileKey != "" { + routeObserver = &passiveRouteObserver{} + requestHTTPClient = routeObserver.wrap(requestHTTPClient) + } runCtx, stopLeaseRenewal := s.startConcurrencyLeaseRenewal(ctx, task.ID, limitResult.Leases) response, err := client.Run(runCtx, clients.Request{ Kind: task.Kind, @@ -1620,6 +1648,9 @@ func (s *Service) runCandidate( if err := setSubmissionStatus("response_received"); err != nil { return err } + if err := s.store.SetTaskSubmissionState(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, "submitted"); err != nil { + return err + } return enterWorkerWaiting(ctx) }, OnRemoteTaskPolled: func(remoteTaskID string, payload map[string]any) error { @@ -1641,11 +1672,17 @@ func (s *Service) runCandidate( if err := setSubmissionStatus("submitting"); err != nil { return err } + if err := s.store.SetTaskSubmissionState(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, "submitting"); err != nil { + return err + } markUpstreamSubmissionStarted(ctx) return enterWorkerWaiting(ctx) }, OnUpstreamResponseReceived: func() error { submissionStatus = "response_received" + if err := s.store.SetTaskSubmissionState(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, "submitted"); err != nil { + return err + } return enterWorkerFinalizing(ctx) }, OnUpstreamWireResponse: func(wire *clients.WireResponse) error { @@ -1660,6 +1697,7 @@ func (s *Service) runCandidate( UpstreamPreviousResponseID: responseExecution.UpstreamPreviousResponseID, PreviousResponseTurns: responseExecution.PreviousTurns, }) + s.recordPassiveRouteObservation(task, routeObserver) if phaseErr := enterWorkerFinalizing(runCtx); err == nil && phaseErr != nil { err = phaseErr } @@ -1733,8 +1771,8 @@ func (s *Service) runCandidate( ErrorMessage: err.Error(), }) _ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable, "requestId": requestID, "statusCode": clients.ErrorResponseMetadata(err).StatusCode, "metrics": metrics}, simulated) - if shouldClassifyUpstreamSubmissionUnknown(simulated, submissionStatus, err) { - return clients.Response{}, &upstreamSubmissionUnknownError{AttemptID: attemptID, Cause: err} + if shouldEnterSubmissionConfirmationPending(simulated, submissionStatus, err) { + return clients.Response{}, &submissionConfirmationPendingError{AttemptID: attemptID, Cause: err} } return clients.Response{}, err } @@ -2352,7 +2390,7 @@ func (s *Service) observeConcurrencyLeaseRenewal(outcome string) { } func (s *Service) requeueInterruptedAsyncTask(ctx context.Context, task store.GatewayTask) (store.GatewayTask, error) { - queued, err := s.store.RequeueTask(ctx, task.ID, task.ExecutionToken, 0, "") + queued, err := s.store.ResolveInterruptedTaskExecution(ctx, task.ID, task.ExecutionToken) if err != nil { return store.GatewayTask{}, err } diff --git a/apps/api/internal/runner/submission_timeout_test.go b/apps/api/internal/runner/submission_timeout_test.go index 5a025e3..5dd98d5 100644 --- a/apps/api/internal/runner/submission_timeout_test.go +++ b/apps/api/internal/runner/submission_timeout_test.go @@ -8,11 +8,11 @@ import ( func TestTimeoutDoesNotBecomeUpstreamSubmissionUnknown(t *testing.T) { timeoutErr := &clients.ClientError{Code: "timeout", Message: "upstream request timed out", Retryable: false} - if shouldClassifyUpstreamSubmissionUnknown(false, "submitting", timeoutErr) { + if shouldEnterSubmissionConfirmationPending(false, "submitting", timeoutErr) { t.Fatal("a definitive provider timeout must remain timeout instead of manual-review unknown") } networkErr := &clients.ClientError{Code: "network", Message: "connection reset", Retryable: true} - if !shouldClassifyUpstreamSubmissionUnknown(false, "submitting", networkErr) { + if !shouldEnterSubmissionConfirmationPending(false, "submitting", networkErr) { t.Fatal("an ambiguous non-timeout disconnect must retain manual-review protection") } } diff --git a/apps/api/internal/runner/worker_capability_test.go b/apps/api/internal/runner/worker_capability_test.go new file mode 100644 index 0000000..9187afa --- /dev/null +++ b/apps/api/internal/runner/worker_capability_test.go @@ -0,0 +1,23 @@ +package runner + +import ( + "testing" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" +) + +func TestWorkerCapabilityNegotiationRejectsLegacyDescriptor(t *testing.T) { + worker := executionpool.WorkerDescriptor{ + ProtocolVersion: executionpool.ProtocolVersion, Allocated: 1, SafeCapacity: 1, + Capabilities: map[string]any{}, HeartbeatAt: time.Now(), + } + if workerCapabilityMatched([]executionpool.WorkerDescriptor{worker}, store.GatewayTask{Kind: "images.generations"}) { + t.Fatal("legacy worker without taskKinds capability was accepted") + } + worker.Capabilities = map[string]any{"taskKinds": []any{"images.*"}} + if !workerCapabilityMatched([]executionpool.WorkerDescriptor{worker}, store.GatewayTask{Kind: "images.generations"}) { + t.Fatal("compatible worker was rejected") + } +} diff --git a/apps/api/internal/securityevents/metrics.go b/apps/api/internal/securityevents/metrics.go index 8b5c781..0f36b7d 100644 --- a/apps/api/internal/securityevents/metrics.go +++ b/apps/api/internal/securityevents/metrics.go @@ -96,6 +96,14 @@ type Metrics struct { candidateDisabledSkipped atomic.Uint64 candidateAllFullQueued atomic.Uint64 candidateRoutingOther atomic.Uint64 + executionPoolSelected atomic.Uint64 + executionPoolShadow atomic.Uint64 + executionPoolUnavailable atomic.Uint64 + executionPoolCapacityReject atomic.Uint64 + routeProbeHealthy atomic.Uint64 + routeProbeDegraded atomic.Uint64 + routeProbeUnreachable atomic.Uint64 + routeProbeUnknown atomic.Uint64 storageAliyunWrites atomic.Uint64 storageAliyunFailures atomic.Uint64 storageAliyunWriteNanos atomic.Uint64 @@ -324,6 +332,32 @@ func (m *Metrics) ObserveCandidateRouting(reason string) { } } +func (m *Metrics) ObserveExecutionPoolRouting(outcome string) { + switch outcome { + case "selected": + m.executionPoolSelected.Add(1) + case "shadow": + m.executionPoolShadow.Add(1) + case "capacity_rejected": + m.executionPoolCapacityReject.Add(1) + default: + m.executionPoolUnavailable.Add(1) + } +} + +func (m *Metrics) ObserveRouteProbe(state string) { + switch state { + case "healthy": + m.routeProbeHealthy.Add(1) + case "degraded": + m.routeProbeDegraded.Add(1) + case "unreachable": + m.routeProbeUnreachable.Add(1) + default: + m.routeProbeUnknown.Add(1) + } +} + func (m *Metrics) ObserveObjectStorage(event string, provider string, bytes int64, duration time.Duration) { provider = strings.ToLower(strings.TrimSpace(provider)) if duration < 0 { @@ -612,6 +646,18 @@ func (m *Metrics) Handler(provider MetricsSnapshotProvider, issuer, audience str {"all_full_queued", m.candidateAllFullQueued.Load()}, {"other", m.candidateRoutingOther.Load()}, }) + outcomeCounters(w, "easyai_gateway_execution_pool_routing_total", "Execution-pool routing decisions by bounded outcome.", []outcomeValue{ + {"selected", m.executionPoolSelected.Load()}, + {"shadow", m.executionPoolShadow.Load()}, + {"unavailable", m.executionPoolUnavailable.Load()}, + {"capacity_rejected", m.executionPoolCapacityReject.Load()}, + }) + outcomeCounters(w, "easyai_gateway_route_probe_total", "Non-billable route probes by resulting state.", []outcomeValue{ + {"healthy", m.routeProbeHealthy.Load()}, + {"degraded", m.routeProbeDegraded.Load()}, + {"unreachable", m.routeProbeUnreachable.Load()}, + {"unknown", m.routeProbeUnknown.Load()}, + }) storageProtocolCounters(w, "easyai_gateway_storage_write_attempts_total", "Object storage write attempts by bounded protocol.", m.storageAliyunWrites.Load(), m.storageS3Writes.Load(), m.storageServerMainWrites.Load()) storageProtocolCounters(w, "easyai_gateway_storage_write_failures_total", "Object storage write failures by bounded protocol.", diff --git a/apps/api/internal/store/billing_v2_integration_test.go b/apps/api/internal/store/billing_v2_integration_test.go index 3651dfb..bec6c1e 100644 --- a/apps/api/internal/store/billing_v2_integration_test.go +++ b/apps/api/internal/store/billing_v2_integration_test.go @@ -151,7 +151,7 @@ func TestExpiredExecutionLeaseDoesNotReplayAmbiguousUpstreamSubmission(t *testin if err != nil { t.Fatal(err) } - if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" { + if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_timeout" { t.Fatalf("manual review task=%+v", review) } var attemptStatus string @@ -162,7 +162,7 @@ FROM gateway_task_attempts WHERE id=$1::uuid`, attemptID).Scan(&attemptStatus, &attemptErrorCode); err != nil { t.Fatal(err) } - if attemptStatus != "failed" || attemptErrorCode != "upstream_submission_unknown" { + if attemptStatus != "failed" || attemptErrorCode != "upstream_timeout" { t.Fatalf("manual review attempt status=%s error=%s", attemptStatus, attemptErrorCode) } var outboxStatus string @@ -174,11 +174,84 @@ FROM settlement_outbox WHERE task_id=$1::uuid AND event_type='task.billing.review'`, created.ID).Scan(&outboxStatus, &outboxAction, &reviewReason); err != nil { t.Fatal(err) } - if outboxStatus != "manual_review" || outboxAction != "release" || reviewReason != "upstream_submission_unknown" { + if outboxStatus != "manual_review" || outboxAction != "release" || reviewReason != "upstream_timeout" { t.Fatalf("review outbox status=%s action=%s reason=%s", outboxStatus, outboxAction, reviewReason) } } +func TestResolveInterruptedTaskExecutionPreservesSubmissionFence(t *testing.T) { + db := billingV2IntegrationStore(t) + ctx := context.Background() + user := &auth.User{ID: "interrupted-worker-" + uuid.NewString(), Source: "gateway"} + + type testCase struct { + name string + submissionStatus string + remoteTaskID string + wantManualReview bool + } + for _, tc := range []testCase{ + {name: "not submitted requeues", submissionStatus: "not_submitted"}, + {name: "ambiguous submission stops", submissionStatus: "submitting", wantManualReview: true}, + {name: "known remote task requeues for polling", submissionStatus: "response_received", remoteTaskID: "remote-" + uuid.NewString()}, + } { + t.Run(tc.name, func(t *testing.T) { + created, err := db.CreateTask(ctx, CreateTaskInput{ + Kind: "images.generations", Model: "interrupted-worker-model", RunMode: "production", Async: true, + Request: map[string]any{"model": "interrupted-worker-model"}, + }, user) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + _, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE id=$1::uuid`, created.ID) + }) + + executionToken := uuid.NewString() + if _, err := db.ClaimTaskExecution(ctx, created.ID, executionToken, 5*time.Minute); err != nil { + t.Fatal(err) + } + attemptID, err := db.CreateTaskAttempt(ctx, CreateTaskAttemptInput{ + TaskID: created.ID, ExecutionToken: executionToken, AttemptNo: 1, Status: "running", + }) + if err != nil { + t.Fatal(err) + } + if err := db.SetAttemptUpstreamSubmissionStatusForExecution( + ctx, attemptID, created.ID, executionToken, tc.submissionStatus, + ); err != nil { + t.Fatal(err) + } + if tc.remoteTaskID != "" { + if err := db.SetTaskRemoteTask(ctx, created.ID, executionToken, attemptID, tc.remoteTaskID, nil); err != nil { + t.Fatal(err) + } + } + + resolved, err := db.ResolveInterruptedTaskExecution(ctx, created.ID, executionToken) + if tc.wantManualReview { + if !errors.Is(err, ErrTaskExecutionManualReview) { + t.Fatalf("resolve error=%v, want ErrTaskExecutionManualReview", err) + } + failed, getErr := db.GetTask(ctx, created.ID) + if getErr != nil { + t.Fatal(getErr) + } + if failed.Status != "failed" || failed.ErrorCode != "upstream_timeout" || failed.SubmissionState != "submission_confirmation_pending" { + t.Fatalf("ambiguous task=%+v", failed) + } + return + } + if err != nil { + t.Fatal(err) + } + if resolved.Status != "queued" || resolved.ExecutionToken != "" || resolved.RemoteTaskID != tc.remoteTaskID { + t.Fatalf("resolved task=%+v", resolved) + } + }) + } +} + func TestQueuedPreparationDoesNotReplayInterruptedUpstreamSubmission(t *testing.T) { db := billingV2IntegrationStore(t) ctx := context.Background() @@ -244,7 +317,7 @@ WHERE id=$1::uuid`, created.ID); err != nil { if err != nil { t.Fatal(err) } - if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_submission_unknown" { + if review.Status != "failed" || review.BillingStatus != "manual_review" || review.ErrorCode != "upstream_timeout" { t.Fatalf("manual review task=%+v", review) } }) @@ -365,7 +438,7 @@ func TestFinishTaskManualReviewCreatesVisibleBillingRecord(t *testing.T) { } if _, err := db.FinishTaskManualReview(ctx, FinishTaskManualReviewInput{ TaskID: created.ID, ExecutionToken: token, AttemptID: attemptID, TaskStatus: "failed", - Code: "upstream_submission_unknown", Message: "upstream submission result is unknown", + Code: "upstream_timeout", Message: "upstream submission confirmation timed out", }); err != nil { t.Fatal(err) } @@ -374,7 +447,7 @@ func TestFinishTaskManualReviewCreatesVisibleBillingRecord(t *testing.T) { SELECT EXISTS ( SELECT 1 FROM settlement_outbox WHERE task_id=$1::uuid AND status='manual_review' - AND action='release' AND manual_review_reason='upstream_submission_unknown' + AND action='release' AND manual_review_reason='upstream_timeout' )`, created.ID).Scan(&visible); err != nil { t.Fatal(err) } diff --git a/apps/api/internal/store/execution_pools.go b/apps/api/internal/store/execution_pools.go new file mode 100644 index 0000000..52f79af --- /dev/null +++ b/apps/api/internal/store/execution_pools.go @@ -0,0 +1,537 @@ +package store + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "strings" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool" + "github.com/google/uuid" + "github.com/jackc/pgx/v5" +) + +type TaskRoutingDecision struct { + TaskID string + PoolID string + WorkerID string + RouteProfileKey string + PlatformID string + PlatformModelID string + RoutingVersion string + Reason string + Snapshot map[string]any +} + +type WorkerExecutionLease struct { + LeaseID string + TaskID string + PoolID string + WorkerID string + InstanceID string + Endpoint string + Nonce string + ExpiresAt time.Time +} + +func (s *Store) UpsertExecutionPool(ctx context.Context, pool executionpool.ExecutionPool) error { + pool.ID = strings.TrimSpace(pool.ID) + if pool.ID == "" { + return errors.New("execution pool ID is required") + } + if pool.Labels == nil { + pool.Labels = map[string]string{} + } + if pool.Capabilities == nil { + pool.Capabilities = map[string]any{} + } + labels, err := json.Marshal(pool.Labels) + if err != nil { + return err + } + capabilities, err := json.Marshal(pool.Capabilities) + if err != nil { + return err + } + state := pool.State + if state == "" { + state = executionpool.PoolActive + } + _, err = s.pool.Exec(ctx, ` +INSERT INTO gateway_execution_pools (pool_id, labels, capabilities, state, updated_at) +VALUES ($1, $2::jsonb, $3::jsonb, $4, now()) +ON CONFLICT (pool_id) DO UPDATE +SET labels = EXCLUDED.labels, + capabilities = EXCLUDED.capabilities, + state = CASE + WHEN gateway_execution_pools.state = 'disabled' THEN 'disabled' + ELSE EXCLUDED.state + END, + updated_at = now()`, pool.ID, labels, capabilities, string(state)) + return err +} + +func (s *Store) ListExecutionPools(ctx context.Context) ([]executionpool.ExecutionPool, error) { + rows, err := s.pool.Query(ctx, ` +SELECT pool_id, labels, capabilities, state, updated_at +FROM gateway_execution_pools +ORDER BY pool_id`) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]executionpool.ExecutionPool, 0) + for rows.Next() { + var item executionpool.ExecutionPool + var labels, capabilities []byte + if err := rows.Scan(&item.ID, &labels, &capabilities, &item.State, &item.UpdatedAt); err != nil { + return nil, err + } + item.Labels = decodeStringMap(labels) + item.Capabilities = decodeObject(capabilities) + items = append(items, item) + } + return items, rows.Err() +} + +func (s *Store) ListWorkers(ctx context.Context, now time.Time) ([]executionpool.WorkerDescriptor, error) { + if now.IsZero() { + now = time.Now() + } + rows, err := s.pool.Query(ctx, ` +SELECT worker_id, instance_id, pool_id, endpoint, protocol_version, revision, + capabilities, allocated_capacity, safe_capacity, heavy_capacity, + active_tasks, pressure_state, heartbeat_at, load_sampled_at +FROM gateway_worker_instances +WHERE status = 'active' + AND heartbeat_at > $1 - $2::interval +ORDER BY pool_id, instance_id`, now, workerHeartbeatStaleAfter.String()) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]executionpool.WorkerDescriptor, 0) + for rows.Next() { + var item executionpool.WorkerDescriptor + var capabilities []byte + var loadSampledAt *time.Time + if err := rows.Scan( + &item.WorkerID, &item.InstanceID, &item.PoolID, &item.Endpoint, + &item.ProtocolVersion, &item.Revision, &capabilities, &item.Allocated, + &item.SafeCapacity, &item.HeavyCapacity, &item.ActiveTasks, + &item.PressureState, &item.HeartbeatAt, &loadSampledAt, + ); err != nil { + return nil, err + } + item.Capabilities = decodeObject(capabilities) + if loadSampledAt != nil { + item.LoadSampledAt = *loadSampledAt + } + items = append(items, item) + } + return items, rows.Err() +} + +func (s *Store) RecordRouteHealth(ctx context.Context, health executionpool.RouteHealth) error { + _, err := s.pool.Exec(ctx, ` +INSERT INTO gateway_route_health ( + pool_id, route_profile_key, state, success_rate, + connect_tls_p95_ms, first_byte_p95_ms, upload_bytes_per_second, jitter_p95_ms, + consecutive_failures, consecutive_successes, sample_count, sampled_at, expires_at, updated_at +) +VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, now()) +ON CONFLICT (pool_id, route_profile_key) DO UPDATE +SET state = EXCLUDED.state, + success_rate = EXCLUDED.success_rate, + connect_tls_p95_ms = EXCLUDED.connect_tls_p95_ms, + first_byte_p95_ms = EXCLUDED.first_byte_p95_ms, + upload_bytes_per_second = EXCLUDED.upload_bytes_per_second, + jitter_p95_ms = EXCLUDED.jitter_p95_ms, + consecutive_failures = EXCLUDED.consecutive_failures, + consecutive_successes = EXCLUDED.consecutive_successes, + sample_count = EXCLUDED.sample_count, + sampled_at = EXCLUDED.sampled_at, + expires_at = EXCLUDED.expires_at, + updated_at = now() +WHERE gateway_route_health.sampled_at <= EXCLUDED.sampled_at`, + health.PoolID, health.RouteProfileKey, string(health.State), health.SuccessRate, + health.ConnectTLSP95.Milliseconds(), health.FirstByteP95.Milliseconds(), health.UploadBytesPerSecond, + health.JitterP95.Milliseconds(), health.ConsecutiveFailures, health.ConsecutiveSuccesses, + health.SampleCount, health.SampledAt, health.ExpiresAt) + return err +} + +func (s *Store) RecordRouteObservation(ctx context.Context, observation executionpool.RouteObservation) error { + if observation.SampleCount < 1 || observation.SuccessCount < 0 || observation.SuccessCount > observation.SampleCount { + return errors.New("invalid passive route observation") + } + successRate := float64(observation.SuccessCount) / float64(observation.SampleCount) + _, err := s.pool.Exec(ctx, ` +UPDATE gateway_route_health +SET success_rate = 0.8 * success_rate + 0.2 * $3, + connect_tls_p95_ms = CASE + WHEN $4::bigint > 0 THEN (0.8 * connect_tls_p95_ms + 0.2 * $4::bigint)::bigint + ELSE connect_tls_p95_ms + END, + upload_bytes_per_second = CASE + WHEN $5::double precision > 0 THEN 0.8 * upload_bytes_per_second + 0.2 * $5::double precision + ELSE upload_bytes_per_second + END, + sample_count = sample_count + $6, + updated_at = now() +WHERE pool_id = $1 + AND route_profile_key = $2`, + strings.TrimSpace(observation.PoolID), strings.TrimSpace(observation.RouteProfileKey), + successRate, observation.ConnectTLSP95.Milliseconds(), observation.UploadBytesPerSecond, + observation.SampleCount, + ) + return err +} + +func (s *Store) RequestRouteProbe(ctx context.Context, routeProfileKey string) error { + _, err := s.pool.Exec(ctx, ` +INSERT INTO gateway_route_probe_requests (route_profile_key, requested_at, expires_at) +VALUES ($1, now(), now() + interval '30 seconds') +ON CONFLICT (route_profile_key) DO UPDATE +SET requested_at = now(), expires_at = now() + interval '30 seconds'`, strings.TrimSpace(routeProfileKey)) + return err +} + +func (s *Store) ListRequestedRouteProbes(ctx context.Context) (map[string]time.Time, error) { + rows, err := s.pool.Query(ctx, ` +SELECT route_profile_key, requested_at +FROM gateway_route_probe_requests +WHERE expires_at > now()`) + if err != nil { + return nil, err + } + defer rows.Close() + items := make(map[string]time.Time) + for rows.Next() { + var key string + var requestedAt time.Time + if err := rows.Scan(&key, &requestedAt); err != nil { + return nil, err + } + items[key] = requestedAt + } + return items, rows.Err() +} + +func (s *Store) ListRouteHealth(ctx context.Context, routeProfileKey string, _ time.Time) ([]executionpool.RouteHealth, error) { + return s.listRouteHealth(ctx, "WHERE route_profile_key = $1", strings.TrimSpace(routeProfileKey)) +} + +func (s *Store) ListAllRouteHealth(ctx context.Context) ([]executionpool.RouteHealth, error) { + return s.listRouteHealth(ctx, "", nil) +} + +func (s *Store) ResolveRoutePreference( + ctx context.Context, + routeProfileKey string, + proposedPoolID string, + scores map[string]float64, + now time.Time, +) (executionpool.RoutePreference, error) { + if now.IsZero() { + now = time.Now() + } + tx, err := s.pool.Begin(ctx) + if err != nil { + return executionpool.RoutePreference{}, err + } + defer rollbackTransaction(tx) + if _, err := tx.Exec(ctx, ` +INSERT INTO gateway_route_preferences (route_profile_key, updated_at) +VALUES ($1, $2) +ON CONFLICT (route_profile_key) DO NOTHING`, strings.TrimSpace(routeProfileKey), now); err != nil { + return executionpool.RoutePreference{}, err + } + var previous executionpool.RoutePreference + var currentSince *time.Time + if err := tx.QueryRow(ctx, ` +SELECT route_profile_key, current_pool_id, current_since, + challenger_pool_id, challenger_wins, updated_at +FROM gateway_route_preferences +WHERE route_profile_key = $1 +FOR UPDATE`, strings.TrimSpace(routeProfileKey)).Scan( + &previous.RouteProfileKey, &previous.CurrentPoolID, ¤tSince, + &previous.ChallengerPoolID, &previous.ChallengerWins, &previous.UpdatedAt, + ); err != nil { + return executionpool.RoutePreference{}, err + } + if currentSince != nil { + previous.CurrentSince = *currentSince + } + next := executionpool.AdvanceRoutePreference(previous, proposedPoolID, scores, now) + if _, err := tx.Exec(ctx, ` +UPDATE gateway_route_preferences +SET current_pool_id = $2, + current_since = $3, + challenger_pool_id = $4, + challenger_wins = $5, + updated_at = $6 +WHERE route_profile_key = $1`, next.RouteProfileKey, next.CurrentPoolID, nullablePreferenceTime(next.CurrentSince), + next.ChallengerPoolID, next.ChallengerWins, next.UpdatedAt); err != nil { + return executionpool.RoutePreference{}, err + } + if err := tx.Commit(ctx); err != nil { + return executionpool.RoutePreference{}, err + } + return next, nil +} + +func nullablePreferenceTime(value time.Time) any { + if value.IsZero() { + return nil + } + return value +} + +func (s *Store) listRouteHealth(ctx context.Context, where string, argument any) ([]executionpool.RouteHealth, error) { + query := ` +SELECT pool_id, route_profile_key, state, success_rate, + connect_tls_p95_ms, first_byte_p95_ms, upload_bytes_per_second, jitter_p95_ms, + consecutive_failures, consecutive_successes, sample_count, sampled_at, expires_at +FROM gateway_route_health +` + where + ` +ORDER BY route_profile_key, pool_id` + var rows pgx.Rows + var err error + if where == "" { + rows, err = s.pool.Query(ctx, query) + } else { + rows, err = s.pool.Query(ctx, query, argument) + } + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]executionpool.RouteHealth, 0) + for rows.Next() { + var item executionpool.RouteHealth + var connectMS, firstByteMS, jitterMS int64 + if err := rows.Scan( + &item.PoolID, &item.RouteProfileKey, &item.State, &item.SuccessRate, + &connectMS, &firstByteMS, &item.UploadBytesPerSecond, &jitterMS, + &item.ConsecutiveFailures, &item.ConsecutiveSuccesses, &item.SampleCount, + &item.SampledAt, &item.ExpiresAt, + ); err != nil { + return nil, err + } + item.ConnectTLSP95 = time.Duration(connectMS) * time.Millisecond + item.FirstByteP95 = time.Duration(firstByteMS) * time.Millisecond + item.JitterP95 = time.Duration(jitterMS) * time.Millisecond + items = append(items, item) + } + return items, rows.Err() +} + +func (s *Store) PublishDesiredCapacity(ctx context.Context, desired executionpool.DesiredCapacity) error { + _, err := s.pool.Exec(ctx, ` +INSERT INTO gateway_pool_capacity_desires (pool_id, desired, reason, valid_until, updated_at) +VALUES ($1, $2, $3, $4, now()) +ON CONFLICT (pool_id) DO UPDATE +SET desired = EXCLUDED.desired, + reason = EXCLUDED.reason, + valid_until = EXCLUDED.valid_until, + updated_at = now()`, desired.PoolID, desired.Desired, desired.Reason, desired.ValidUntil) + return err +} + +func (s *Store) ListCapacity(ctx context.Context, now time.Time) ([]executionpool.CapacitySnapshot, error) { + if now.IsZero() { + now = time.Now() + } + rows, err := s.pool.Query(ctx, ` +SELECT pool.pool_id, + COUNT(worker.instance_id)::int, + COALESCE(SUM(worker.allocated_capacity), 0)::int, + COALESCE(SUM(worker.safe_capacity), 0)::int, + COALESCE(SUM(worker.active_tasks), 0)::int, + COALESCE(SUM(worker.heavy_capacity), 0)::int, + COALESCE(SUM(worker.preparing_tasks + worker.finalizing_tasks), 0)::int, + COALESCE(BOOL_OR(worker.pressure_state = 'critical'), false), + COALESCE(MAX(worker.load_sampled_at), $1) +FROM gateway_execution_pools pool +LEFT JOIN gateway_worker_instances worker + ON worker.pool_id = pool.pool_id + AND worker.status = 'active' + AND worker.heartbeat_at > $1 - $2::interval +WHERE pool.state = 'active' +GROUP BY pool.pool_id +ORDER BY pool.pool_id`, now, workerHeartbeatStaleAfter.String()) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]executionpool.CapacitySnapshot, 0) + for rows.Next() { + var item executionpool.CapacitySnapshot + if err := rows.Scan( + &item.PoolID, &item.WorkerCount, &item.Allocated, &item.SafeCapacity, + &item.ActiveTasks, &item.HeavyCapacity, &item.HeavyTasks, &item.Critical, + &item.SampledAt, + ); err != nil { + return nil, err + } + if item.SafeCapacity > 0 { + item.ResourceHeadroom = float64(max(item.SafeCapacity-item.ActiveTasks, 0)) / float64(item.SafeCapacity) + } + item.Stability = 1 + items = append(items, item) + } + return items, rows.Err() +} + +func (s *Store) AssignTaskRouting(ctx context.Context, decision TaskRoutingDecision) error { + snapshot, err := json.Marshal(decision.Snapshot) + if err != nil { + return err + } + _, err = s.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET assigned_pool_id = NULLIF($2, ''), + assigned_worker_id = NULLIF($3, ''), + route_profile_key = NULLIF($4, ''), + routing_platform_id = NULLIF($5, '')::uuid, + routing_platform_model_id = NULLIF($6, '')::uuid, + routing_version = NULLIF($7, ''), + routing_reason = NULLIF($8, ''), + routing_snapshot = $9::jsonb, + updated_at = now() +WHERE id = $1::uuid`, decision.TaskID, decision.PoolID, decision.WorkerID, + decision.RouteProfileKey, decision.PlatformID, decision.PlatformModelID, + decision.RoutingVersion, decision.Reason, snapshot) + return err +} + +func (s *Store) SetTaskSubmissionState(ctx context.Context, taskID, executionToken, state string) error { + command, err := s.pool.Exec(ctx, ` +UPDATE gateway_tasks +SET submission_state = $3, updated_at = now() +WHERE id = $1::uuid + AND execution_token = $2::uuid + AND status = 'running'`, taskID, executionToken, state) + if err != nil { + return err + } + if command.RowsAffected() == 0 { + return ErrTaskExecutionLeaseLost + } + return nil +} + +func (s *Store) ReserveWorkerExecution(ctx context.Context, taskID, poolID, nonce string, ttl time.Duration) (WorkerExecutionLease, error) { + if ttl <= 0 || ttl > time.Minute { + ttl = 30 * time.Second + } + nonceDigest := sha256.Sum256([]byte(nonce)) + nonceHash := hex.EncodeToString(nonceDigest[:]) + tx, err := s.pool.Begin(ctx) + if err != nil { + return WorkerExecutionLease{}, err + } + defer rollbackTransaction(tx) + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1, 0))`, "gateway-worker-execution:"+strings.TrimSpace(poolID)); err != nil { + return WorkerExecutionLease{}, err + } + var lease WorkerExecutionLease + lease.LeaseID = uuid.NewString() + lease.TaskID = taskID + lease.PoolID = poolID + lease.Nonce = nonce + lease.ExpiresAt = time.Now().Add(ttl) + err = tx.QueryRow(ctx, ` +WITH active_leases AS ( + SELECT instance_id, COUNT(*)::int AS count + FROM gateway_worker_execution_leases + WHERE released_at IS NULL AND expires_at > now() + GROUP BY instance_id +), candidate AS ( + SELECT worker.worker_id, worker.instance_id, worker.endpoint + FROM gateway_worker_instances worker + LEFT JOIN active_leases lease ON lease.instance_id = worker.instance_id + WHERE worker.pool_id = $1 + AND worker.status = 'active' + AND worker.heartbeat_at > now() - $2::interval + AND worker.protocol_version = $3 + AND NULLIF(worker.endpoint, '') IS NOT NULL + AND worker.pressure_state <> 'critical' + AND worker.active_tasks + COALESCE(lease.count, 0) < LEAST(worker.allocated_capacity, worker.safe_capacity) + ORDER BY worker.active_tasks + COALESCE(lease.count, 0), worker.instance_id + LIMIT 1 + FOR UPDATE OF worker SKIP LOCKED +) +SELECT worker_id, instance_id, endpoint FROM candidate`, poolID, workerHeartbeatStaleAfter.String(), executionpool.ProtocolVersion).Scan( + &lease.WorkerID, &lease.InstanceID, &lease.Endpoint, + ) + if err != nil { + return WorkerExecutionLease{}, err + } + if _, err := tx.Exec(ctx, ` +INSERT INTO gateway_worker_execution_leases ( + lease_id, task_id, pool_id, worker_id, instance_id, nonce_hash, expires_at +) +VALUES ($1::uuid, $2::uuid, $3, $4, $5, $6, $7)`, lease.LeaseID, taskID, poolID, + lease.WorkerID, lease.InstanceID, nonceHash, lease.ExpiresAt); err != nil { + return WorkerExecutionLease{}, err + } + if _, err := tx.Exec(ctx, ` +UPDATE gateway_tasks +SET assigned_pool_id = $2, + assigned_worker_id = $3, + updated_at = now() +WHERE id = $1::uuid`, taskID, poolID, lease.WorkerID); err != nil { + return WorkerExecutionLease{}, err + } + if err := tx.Commit(ctx); err != nil { + return WorkerExecutionLease{}, err + } + return lease, nil +} + +func (s *Store) ConsumeWorkerExecutionLease(ctx context.Context, leaseID, nonce string) (WorkerExecutionLease, error) { + digest := sha256.Sum256([]byte(nonce)) + nonceHash := hex.EncodeToString(digest[:]) + var lease WorkerExecutionLease + err := s.pool.QueryRow(ctx, ` +UPDATE gateway_worker_execution_leases +SET state = 'running', updated_at = now() +WHERE lease_id = $1::uuid + AND nonce_hash = $2 + AND state = 'reserved' + AND released_at IS NULL + AND expires_at > now() +RETURNING lease_id::text, task_id::text, pool_id, worker_id, instance_id, expires_at`, leaseID, nonceHash).Scan( + &lease.LeaseID, &lease.TaskID, &lease.PoolID, &lease.WorkerID, &lease.InstanceID, &lease.ExpiresAt, + ) + return lease, err +} + +func (s *Store) ReleaseWorkerExecutionLease(ctx context.Context, leaseID string) error { + command, err := s.pool.Exec(ctx, ` +UPDATE gateway_worker_execution_leases +SET state = 'released', released_at = now(), updated_at = now() +WHERE lease_id = $1::uuid AND released_at IS NULL`, leaseID) + if err != nil { + return err + } + if command.RowsAffected() == 0 { + return pgx.ErrNoRows + } + return nil +} + +func decodeStringMap(payload []byte) map[string]string { + result := map[string]string{} + _ = json.Unmarshal(payload, &result) + return result +} + +var _ executionpool.WorkerDirectory = (*Store)(nil) +var _ executionpool.RouteHealthRepository = (*Store)(nil) +var _ executionpool.CapacityProvider = (*Store)(nil) diff --git a/apps/api/internal/store/leadership.go b/apps/api/internal/store/leadership.go index 92e8ec3..a7a8778 100644 --- a/apps/api/internal/store/leadership.go +++ b/apps/api/internal/store/leadership.go @@ -42,6 +42,10 @@ func (s *Store) TryAcquireCapacityControllerLeadership(ctx context.Context) (Lea return s.tryAcquireLeadership(ctx, capacityControllerLeadershipKey) } +func (s *Store) TryAcquireRouteProbeLeadership(ctx context.Context, poolID, routeProfileKey string) (Leadership, bool, error) { + return s.tryAcquireLeadership(ctx, "easyai-gateway-route-probe:"+poolID+":"+routeProfileKey) +} + func (s *Store) tryAcquireLeadership(ctx context.Context, key string) (Leadership, bool, error) { conn, err := s.pool.Acquire(ctx) if err != nil { diff --git a/apps/api/internal/store/postgres.go b/apps/api/internal/store/postgres.go index 706cbbd..fd7d065 100644 --- a/apps/api/internal/store/postgres.go +++ b/apps/api/internal/store/postgres.go @@ -634,6 +634,15 @@ type GatewayTask struct { CompatibilitySubmitHTTPStatus int `json:"compatibilitySubmitHttpStatus,omitempty"` CompatibilitySubmitHeaders map[string]any `json:"compatibilitySubmitHeaders,omitempty"` CompatibilitySubmitBody map[string]any `json:"compatibilitySubmitBody,omitempty"` + AssignedPoolID string `json:"-"` + AssignedWorkerID string `json:"-"` + RouteProfileKey string `json:"-"` + RoutingPlatformID string `json:"-"` + RoutingPlatformModelID string `json:"-"` + RoutingVersion string `json:"-"` + RoutingReason string `json:"-"` + RoutingSnapshot map[string]any `json:"-"` + SubmissionState string `json:"-"` Attempts []TaskAttempt `json:"attempts,omitempty"` CreatedAt time.Time `json:"createdAt"` UpdatedAt time.Time `json:"updatedAt"` @@ -662,6 +671,10 @@ COALESCE(error_code, ''), COALESCE(error_message, ''), COALESCE(public_error, '{ COALESCE(compatibility_protocol, ''), COALESCE(compatibility_public_id, ''), COALESCE(compatibility_source_protocol, ''), COALESCE(compatibility_submit_http_status, 0), COALESCE(compatibility_submit_headers, '{}'::jsonb), COALESCE(compatibility_submit_body, '{}'::jsonb), +COALESCE(assigned_pool_id, ''), COALESCE(assigned_worker_id, ''), COALESCE(route_profile_key, ''), +COALESCE(routing_platform_id::text, ''), COALESCE(routing_platform_model_id::text, ''), +COALESCE(routing_version, ''), COALESCE(routing_reason, ''), COALESCE(routing_snapshot, '{}'::jsonb), +COALESCE(submission_state, 'not_started'), created_at, updated_at, COALESCE(finished_at::text, '')` type TaskEvent struct { @@ -2243,6 +2256,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) { var remoteTaskPayloadBytes []byte var compatibilitySubmitHeadersBytes []byte var compatibilitySubmitBodyBytes []byte + var routingSnapshotBytes []byte var publicErrorBytes []byte if err := scanner.Scan( &task.ID, @@ -2306,6 +2320,15 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) { &task.CompatibilitySubmitHTTPStatus, &compatibilitySubmitHeadersBytes, &compatibilitySubmitBodyBytes, + &task.AssignedPoolID, + &task.AssignedWorkerID, + &task.RouteProfileKey, + &task.RoutingPlatformID, + &task.RoutingPlatformModelID, + &task.RoutingVersion, + &task.RoutingReason, + &routingSnapshotBytes, + &task.SubmissionState, &task.CreatedAt, &task.UpdatedAt, &task.FinishedAt, @@ -2322,6 +2345,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) { task.PricingSnapshot = decodeObject(pricingSnapshotBytes) task.CompatibilitySubmitHeaders = decodeObject(compatibilitySubmitHeadersBytes) task.CompatibilitySubmitBody = decodeObject(compatibilitySubmitBodyBytes) + task.RoutingSnapshot = decodeObject(routingSnapshotBytes) if len(publicErrorBytes) > 0 { var snapshot publicerror.Error if json.Unmarshal(publicErrorBytes, &snapshot) == nil && snapshot.Code != "" { diff --git a/apps/api/internal/store/rate_limits.go b/apps/api/internal/store/rate_limits.go index 3cefc37..3094605 100644 --- a/apps/api/internal/store/rate_limits.go +++ b/apps/api/internal/store/rate_limits.go @@ -691,13 +691,13 @@ ON CONFLICT (task_id, event_type) DO NOTHING`, releaseBillingTaskIDs); err != ni tag, err := tx.Exec(ctx, ` UPDATE gateway_task_attempts attempt SET status = 'failed', - retryable = task.error_code IS DISTINCT FROM 'upstream_submission_unknown', + retryable = task.submission_state IS DISTINCT FROM 'submission_confirmation_pending', error_code = CASE - WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream_submission_unknown' + WHEN task.submission_state = 'submission_confirmation_pending' THEN 'upstream_timeout' ELSE 'execution_lease_expired' END, error_message = CASE - WHEN task.error_code = 'upstream_submission_unknown' THEN 'upstream submission result is unknown' + WHEN task.submission_state = 'submission_confirmation_pending' THEN 'upstream submission confirmation timed out' ELSE 'attempt execution lease expired after worker heartbeat became stale' END, finished_at = now() diff --git a/apps/api/internal/store/route_probe_targets.go b/apps/api/internal/store/route_probe_targets.go new file mode 100644 index 0000000..b4e1742 --- /dev/null +++ b/apps/api/internal/store/route_probe_targets.go @@ -0,0 +1,68 @@ +package store + +import ( + "context" +) + +type RouteProbeTarget struct { + Candidate RuntimeModelCandidate + Hot bool +} + +func (s *Store) ListRouteProbeTargets(ctx context.Context) ([]RouteProbeTarget, error) { + rows, err := s.pool.Query(ctx, ` +SELECT platform.id::text, + platform.provider, + COALESCE(NULLIF(platform.config->>'specType', ''), NULLIF(catalog.provider_type, ''), NULLIF(platform.config->>'sourceSpecType', ''), platform.provider), + COALESCE( + NULLIF(platform.base_url, ''), + NULLIF(platform.config->>'endpoint', ''), + NULLIF(platform.config->>'baseURL', ''), + NULLIF(platform.config->>'base_url', ''), + NULLIF(catalog.default_base_url, ''), + '' + ), + platform.config, + model.id::text, + COALESCE(model.model_type->>0, ''), + EXISTS ( + SELECT 1 + FROM gateway_task_attempts attempt + WHERE attempt.platform_model_id = model.id + AND attempt.started_at > now() - interval '10 minutes' + ) +FROM integration_platforms platform +JOIN platform_models model ON model.platform_id = platform.id +LEFT JOIN model_catalog_providers catalog + ON catalog.provider_key = platform.provider OR catalog.provider_code = platform.provider +WHERE platform.status = 'enabled' + AND platform.deleted_at IS NULL + AND model.enabled = true + AND COALESCE( + NULLIF(platform.base_url, ''), + NULLIF(platform.config->>'endpoint', ''), + NULLIF(platform.config->>'baseURL', ''), + NULLIF(platform.config->>'base_url', ''), + NULLIF(catalog.default_base_url, '') + ) IS NOT NULL +ORDER BY platform.id, model.id`) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]RouteProbeTarget, 0) + for rows.Next() { + var item RouteProbeTarget + var config []byte + if err := rows.Scan( + &item.Candidate.PlatformID, &item.Candidate.Provider, &item.Candidate.SpecType, + &item.Candidate.BaseURL, &config, &item.Candidate.PlatformModelID, + &item.Candidate.ModelType, &item.Hot, + ); err != nil { + return nil, err + } + item.Candidate.PlatformConfig = decodeObject(config) + items = append(items, item) + } + return items, rows.Err() +} diff --git a/apps/api/internal/store/tasks_runtime.go b/apps/api/internal/store/tasks_runtime.go index 7e55ffc..ce179af 100644 --- a/apps/api/internal/store/tasks_runtime.go +++ b/apps/api/internal/store/tasks_runtime.go @@ -392,8 +392,8 @@ func markTaskExecutionManualReviewTx( UPDATE gateway_task_attempts SET status = 'failed', retryable = false, - error_code = 'upstream_submission_unknown', - error_message = 'upstream submission result is unknown', + error_code = 'upstream_timeout', + error_message = 'upstream submission confirmation timed out', finished_at = COALESCE(finished_at, now()), upstream_submission_updated_at = now() WHERE task_id = $1::uuid @@ -407,14 +407,15 @@ SET status = 'failed', billing_status = CASE WHEN $2 THEN 'manual_review' ELSE 'not_required' END, billing_updated_at = now(), error = NULL, - error_code = 'upstream_submission_unknown', - error_message = 'upstream submission result is unknown', + error_code = 'upstream_timeout', + error_message = 'upstream submission confirmation timed out', locked_by = NULL, locked_at = NULL, heartbeat_at = NULL, execution_token = NULL, execution_lease_expires_at = NULL, remote_task_payload = '{}'::jsonb, + submission_state = 'submission_confirmation_pending', finished_at = now(), updated_at = now() WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil { @@ -424,7 +425,7 @@ WHERE id = $1::uuid`, taskID, hasGatewayUser); err != nil { return nil } payloadJSON, _ := json.Marshal(map[string]any{ - "taskId": taskID, "classification": "upstream_submission_unknown", + "taskId": taskID, "classification": "upstream_timeout", }) if _, err := tx.Exec(ctx, ` INSERT INTO settlement_outbox ( @@ -432,7 +433,7 @@ INSERT INTO settlement_outbox ( status, next_attempt_at, manual_review_reason ) SELECT id, 'task.billing.review', 'release', reservation_amount, billing_currency, - pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_submission_unknown' + pricing_snapshot, $2::jsonb, 'manual_review', now(), 'upstream_timeout' FROM gateway_tasks WHERE id = $1::uuid ON CONFLICT (task_id, event_type) DO NOTHING`, taskID, string(payloadJSON)); err != nil { @@ -866,6 +867,75 @@ WHERE task_id = $1::uuid return task, changed, nil } +// ResolveInterruptedTaskExecution atomically preserves the no-duplicate-submit +// invariant when a Worker disappears or its execution context is cancelled. +// A known remote task is safe to requeue for polling takeover. An ambiguous +// submission without a remote ID is moved to manual review instead of replayed. +func (s *Store) ResolveInterruptedTaskExecution( + ctx context.Context, + taskID string, + executionToken string, +) (GatewayTask, error) { + tx, err := s.pool.Begin(ctx) + if err != nil { + return GatewayTask{}, err + } + defer rollbackTransaction(tx) + var remoteTaskID string + var hasGatewayUser bool + if err := tx.QueryRow(ctx, ` +SELECT COALESCE(remote_task_id, ''), gateway_user_id IS NOT NULL +FROM gateway_tasks +WHERE id = $1::uuid + AND status = 'running' + AND execution_token = $2::uuid +FOR UPDATE`, taskID, executionToken).Scan(&remoteTaskID, &hasGatewayUser); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return GatewayTask{}, ErrTaskExecutionLeaseLost + } + return GatewayTask{}, err + } + ambiguous, err := taskExecutionRequiresManualReviewTx(ctx, tx, taskID) + if err != nil { + return GatewayTask{}, err + } + if remoteTaskID == "" && ambiguous { + if err := markTaskExecutionManualReviewTx(ctx, tx, taskID, hasGatewayUser); err != nil { + return GatewayTask{}, err + } + if err := tx.Commit(ctx); err != nil { + return GatewayTask{}, err + } + return GatewayTask{}, ErrTaskExecutionManualReview + } + nextRunAt := time.Now().Add(time.Second) + queued, err := scanGatewayTask(tx.QueryRow(ctx, ` +UPDATE gateway_tasks +SET status = 'queued', + locked_by = NULL, + locked_at = NULL, + heartbeat_at = NULL, + execution_token = NULL, + execution_lease_expires_at = NULL, + next_run_at = $3, + error = NULL, + error_code = NULL, + error_message = NULL, + public_error = NULL, + updated_at = now() +WHERE id = $1::uuid + AND status = 'running' + AND execution_token = $2::uuid +RETURNING `+gatewayTaskColumns, taskID, executionToken, nextRunAt)) + if err != nil { + return GatewayTask{}, err + } + if err := tx.Commit(ctx); err != nil { + return GatewayTask{}, err + } + return queued, nil +} + func (s *Store) SetTaskRiverJobID(ctx context.Context, taskID string, riverJobID int64) error { if riverJobID <= 0 { return nil diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index 0de79ef..207ce2c 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -2,6 +2,7 @@ package store import ( "context" + "encoding/json" "errors" "strings" "time" @@ -9,31 +10,37 @@ import ( "github.com/jackc/pgx/v5" ) -const workerHeartbeatStaleAfter = 30 * time.Second +const workerHeartbeatStaleAfter = 15 * time.Second type WorkerRegistrationInput struct { - InstanceID string - PodUID string - PodName string - Site string - Revision string - DesiredCapacity int - CapacityLimit int - LoadMode string - SafeCapacity int - HeavyCapacity int - ActiveTasks int - PreparingTasks int - WaitingUpstreamTasks int - FinalizingTasks int - PressureState string - PressureReason string - LoadSampledAt time.Time - HeartbeatStaleAfter time.Duration + InstanceID string + WorkerID string + PoolID string + Endpoint string + Labels map[string]string + Capabilities map[string]any + ProtocolVersion string + OrchestratorInstanceRef string + Revision string + DesiredCapacity int + CapacityLimit int + LoadMode string + SafeCapacity int + HeavyCapacity int + ActiveTasks int + PreparingTasks int + WaitingUpstreamTasks int + FinalizingTasks int + PressureState string + PressureReason string + LoadSampledAt time.Time + HeartbeatStaleAfter time.Duration } type WorkerAllocation struct { InstanceID string + WorkerID string + PoolID string DesiredCapacity int Allocated int GlobalAllocated int @@ -63,6 +70,31 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra if input.InstanceID == "" { return WorkerAllocation{}, errors.New("worker instance ID is required") } + input.WorkerID = strings.TrimSpace(input.WorkerID) + if input.WorkerID == "" { + input.WorkerID = input.InstanceID + } + input.PoolID = strings.TrimSpace(input.PoolID) + if input.PoolID == "" { + input.PoolID = "legacy-default" + } + if strings.TrimSpace(input.ProtocolVersion) == "" { + input.ProtocolVersion = "v1" + } + if input.Labels == nil { + input.Labels = map[string]string{} + } + if input.Capabilities == nil { + input.Capabilities = map[string]any{} + } + labels, err := json.Marshal(input.Labels) + if err != nil { + return WorkerAllocation{}, err + } + capabilities, err := json.Marshal(input.Capabilities) + if err != nil { + return WorkerAllocation{}, err + } if input.DesiredCapacity < 0 { return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative") } @@ -110,22 +142,38 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra return WorkerAllocation{}, err } if _, err := tx.Exec(ctx, ` +INSERT INTO gateway_execution_pools (pool_id, labels, capabilities, state, updated_at) +VALUES ($1, $2::jsonb, $3::jsonb, 'active', now()) +ON CONFLICT (pool_id) DO UPDATE +SET labels = EXCLUDED.labels, + capabilities = EXCLUDED.capabilities, + updated_at = now() +WHERE gateway_execution_pools.state <> 'disabled'`, input.PoolID, labels, capabilities); err != nil { + return WorkerAllocation{}, err + } + if _, err := tx.Exec(ctx, ` INSERT INTO gateway_worker_instances ( - instance_id, pod_uid, pod_name, site, revision, status, + instance_id, worker_id, pool_id, endpoint, labels, capabilities, protocol_version, + orchestrator_instance_ref, revision, status, desired_capacity, capacity_limit, hard_capacity_limit, safe_capacity, heavy_capacity, active_tasks, preparing_tasks, waiting_upstream_tasks, finalizing_tasks, pressure_state, pressure_reason, load_sampled_at, allocated_capacity, started_at, heartbeat_at, updated_at ) VALUES ( - $1, $2, $3, $4, $5, 'active', $6, $7, $8, $9, $10, - $11, $12, $13, $14, $15, $16, $17, + $1, $2, $3, $4, $5::jsonb, $6::jsonb, $7, + $8, $9, 'active', $10, $11, $12, $13, $14, + $15, $16, $17, $18, $19, $20, $21, 0, now(), now(), now() ) ON CONFLICT (instance_id) DO UPDATE -SET pod_uid = EXCLUDED.pod_uid, - pod_name = EXCLUDED.pod_name, - site = EXCLUDED.site, +SET worker_id = EXCLUDED.worker_id, + pool_id = EXCLUDED.pool_id, + endpoint = EXCLUDED.endpoint, + labels = EXCLUDED.labels, + capabilities = EXCLUDED.capabilities, + protocol_version = EXCLUDED.protocol_version, + orchestrator_instance_ref = EXCLUDED.orchestrator_instance_ref, revision = EXCLUDED.revision, status = CASE WHEN gateway_worker_instances.status = 'draining' THEN 'draining' @@ -146,9 +194,13 @@ SET pod_uid = EXCLUDED.pod_uid, heartbeat_at = now(), updated_at = now()`, input.InstanceID, - strings.TrimSpace(input.PodUID), - strings.TrimSpace(input.PodName), - strings.TrimSpace(input.Site), + input.WorkerID, + input.PoolID, + strings.TrimRight(strings.TrimSpace(input.Endpoint), "/"), + labels, + capabilities, + strings.TrimSpace(input.ProtocolVersion), + strings.TrimSpace(input.OrchestratorInstanceRef), strings.TrimSpace(input.Revision), input.DesiredCapacity, input.CapacityLimit, @@ -235,6 +287,8 @@ WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil { } return WorkerAllocation{ InstanceID: input.InstanceID, + WorkerID: input.WorkerID, + PoolID: input.PoolID, DesiredCapacity: input.DesiredCapacity, Allocated: allocated, GlobalAllocated: globalAllocated, @@ -310,28 +364,30 @@ WHERE instance_id = $1 } type WorkerInstanceRuntime struct { - InstanceID string `json:"instanceId"` - PodUID string `json:"podUid,omitempty"` - PodName string `json:"podName,omitempty"` - Site string `json:"site,omitempty"` - Revision string `json:"revision,omitempty"` - Status string `json:"status"` - Allocated int `json:"allocatedCapacity"` - CapacityLimit int `json:"capacityLimit"` - HardCapacityLimit int `json:"hardCapacityLimit"` - SafeCapacity int `json:"safeCapacity"` - HeavyCapacity int `json:"heavyCapacity"` - ReportedActiveTasks int `json:"reportedActiveTasks"` - PreparingTasks int `json:"preparingTasks"` - WaitingUpstreamTasks int `json:"waitingUpstreamTasks"` - FinalizingTasks int `json:"finalizingTasks"` - PressureState string `json:"pressureState"` - PressureReason string `json:"pressureReason,omitempty"` - LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"` - RunningTasks int `json:"runningTasks"` - ActiveLeases int `json:"activeLeases"` - HeartbeatAt time.Time `json:"heartbeatAt"` - DrainingAt *time.Time `json:"drainingAt,omitempty"` + InstanceID string `json:"instanceId"` + PodUID string `json:"podUid,omitempty"` + PodName string `json:"podName,omitempty"` + OrchestratorInstanceRef string `json:"-"` + Site string `json:"site,omitempty"` + PoolID string `json:"poolId,omitempty"` + Revision string `json:"revision,omitempty"` + Status string `json:"status"` + Allocated int `json:"allocatedCapacity"` + CapacityLimit int `json:"capacityLimit"` + HardCapacityLimit int `json:"hardCapacityLimit"` + SafeCapacity int `json:"safeCapacity"` + HeavyCapacity int `json:"heavyCapacity"` + ReportedActiveTasks int `json:"reportedActiveTasks"` + PreparingTasks int `json:"preparingTasks"` + WaitingUpstreamTasks int `json:"waitingUpstreamTasks"` + FinalizingTasks int `json:"finalizingTasks"` + PressureState string `json:"pressureState"` + PressureReason string `json:"pressureReason,omitempty"` + LoadSampledAt *time.Time `json:"loadSampledAt,omitempty"` + RunningTasks int `json:"runningTasks"` + ActiveLeases int `json:"activeLeases"` + HeartbeatAt time.Time `json:"heartbeatAt"` + DrainingAt *time.Time `json:"drainingAt,omitempty"` } type WorkerQueueRuntime struct { @@ -340,6 +396,12 @@ type WorkerQueueRuntime struct { OldestWaitSeconds float64 `json:"oldestWaitSeconds"` } +type PoolQueueRuntime struct { + PoolID string `json:"poolId"` + Queued int `json:"queued"` + Running int `json:"running"` +} + type WorkerClusterRuntime struct { Workers []WorkerInstanceRuntime `json:"workers"` Queue WorkerQueueRuntime `json:"queue"` @@ -380,12 +442,40 @@ WHERE status IN ('queued', 'running') return snapshot, err } +func (s *Store) ListPoolQueueRuntime(ctx context.Context) ([]PoolQueueRuntime, error) { + rows, err := s.pool.Query(ctx, ` +SELECT assigned_pool_id, + count(*) FILTER (WHERE status = 'queued')::int, + count(*) FILTER (WHERE status = 'running')::int +FROM gateway_tasks +WHERE assigned_pool_id IS NOT NULL + AND status IN ('queued', 'running') + AND run_mode IN ('production', 'acceptance', 'acceptance_canary') +GROUP BY assigned_pool_id +ORDER BY assigned_pool_id`) + if err != nil { + return nil, err + } + defer rows.Close() + items := make([]PoolQueueRuntime, 0) + for rows.Next() { + var item PoolQueueRuntime + if err := rows.Scan(&item.PoolID, &item.Queued, &item.Running); err != nil { + return nil, err + } + items = append(items, item) + } + return items, rows.Err() +} + func (s *Store) ListWorkerInstanceRuntime(ctx context.Context) ([]WorkerInstanceRuntime, error) { rows, err := s.pool.Query(ctx, ` SELECT worker.instance_id, worker.pod_uid, worker.pod_name, + worker.orchestrator_instance_ref, worker.site, + worker.pool_id, worker.revision, worker.status, worker.allocated_capacity, @@ -414,7 +504,7 @@ LEFT JOIN gateway_concurrency_leases lease WHERE worker.status IN ('active', 'draining') AND worker.heartbeat_at > now() - $1::interval GROUP BY worker.instance_id -ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`, +ORDER BY worker.pool_id ASC, worker.status DESC, worker.instance_id ASC`, runtimeWorkerStaleAfter.String(), ) if err != nil { @@ -428,7 +518,9 @@ ORDER BY worker.site ASC, worker.status DESC, worker.instance_id ASC`, &instance.InstanceID, &instance.PodUID, &instance.PodName, + &instance.OrchestratorInstanceRef, &instance.Site, + &instance.PoolID, &instance.Revision, &instance.Status, &instance.Allocated, diff --git a/apps/api/migrations/0105_execution_pool_routing.sql b/apps/api/migrations/0105_execution_pool_routing.sql new file mode 100644 index 0000000..bf4b69e --- /dev/null +++ b/apps/api/migrations/0105_execution_pool_routing.sql @@ -0,0 +1,192 @@ +CREATE TABLE IF NOT EXISTS gateway_execution_pools ( + pool_id text PRIMARY KEY, + labels jsonb NOT NULL DEFAULT '{}'::jsonb, + capabilities jsonb NOT NULL DEFAULT '{}'::jsonb, + state text NOT NULL DEFAULT 'active', + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now(), + CONSTRAINT gateway_execution_pools_id_check + CHECK (pool_id ~ '^[A-Za-z0-9][A-Za-z0-9._:/-]{0,127}$'), + CONSTRAINT gateway_execution_pools_state_check + CHECK (state IN ('active', 'draining', 'disabled')), + CONSTRAINT gateway_execution_pools_labels_object_check + CHECK (jsonb_typeof(labels) = 'object'), + CONSTRAINT gateway_execution_pools_capabilities_object_check + CHECK (jsonb_typeof(capabilities) = 'object') +); + +INSERT INTO gateway_execution_pools (pool_id, labels) +VALUES ('legacy-default', '{"migration":"legacy"}'::jsonb) +ON CONFLICT (pool_id) DO NOTHING; + +INSERT INTO gateway_execution_pools (pool_id, labels) +SELECT DISTINCT site, jsonb_build_object('legacy_site', site) +FROM gateway_worker_instances +WHERE NULLIF(site, '') IS NOT NULL +ON CONFLICT (pool_id) DO NOTHING; + +ALTER TABLE gateway_worker_instances + ADD COLUMN pool_id text DEFAULT 'legacy-default', + ADD COLUMN worker_id text DEFAULT '', + ADD COLUMN endpoint text DEFAULT '', + ADD COLUMN orchestrator_instance_ref text DEFAULT '', + ADD COLUMN labels jsonb DEFAULT '{}'::jsonb, + ADD COLUMN capabilities jsonb DEFAULT '{}'::jsonb, + ADD COLUMN protocol_version text DEFAULT 'v1'; + +UPDATE gateway_worker_instances +SET pool_id = site +WHERE NULLIF(site, '') IS NOT NULL + AND pool_id = 'legacy-default'; + +UPDATE gateway_worker_instances +SET worker_id = instance_id +WHERE worker_id = ''; + +UPDATE gateway_worker_instances +SET orchestrator_instance_ref = pod_name +WHERE orchestrator_instance_ref = '' + AND NULLIF(pod_name, '') IS NOT NULL; + +ALTER TABLE gateway_worker_instances + ADD CONSTRAINT gateway_worker_instances_pool_fk + FOREIGN KEY (pool_id) REFERENCES gateway_execution_pools(pool_id) ON DELETE RESTRICT, + ADD CONSTRAINT gateway_worker_instances_routing_fields_not_null_check + CHECK (num_nonnulls(pool_id, worker_id, endpoint, orchestrator_instance_ref, labels, capabilities, protocol_version) = 7) NOT VALID, + ADD CONSTRAINT gateway_worker_instances_labels_object_check + CHECK (jsonb_typeof(labels) = 'object') NOT VALID, + ADD CONSTRAINT gateway_worker_instances_capabilities_object_check + CHECK (jsonb_typeof(capabilities) = 'object') NOT VALID; + +ALTER TABLE gateway_worker_instances + VALIDATE CONSTRAINT gateway_worker_instances_routing_fields_not_null_check; + +ALTER TABLE gateway_worker_instances + VALIDATE CONSTRAINT gateway_worker_instances_labels_object_check; + +ALTER TABLE gateway_worker_instances + VALIDATE CONSTRAINT gateway_worker_instances_capabilities_object_check; + +CREATE INDEX IF NOT EXISTS idx_worker_instances_pool_heartbeat + ON gateway_worker_instances(pool_id, status, heartbeat_at, instance_id); + +CREATE TABLE IF NOT EXISTS gateway_route_health ( + pool_id text NOT NULL REFERENCES gateway_execution_pools(pool_id) ON DELETE CASCADE, + route_profile_key text NOT NULL, + state text NOT NULL DEFAULT 'unknown', + success_rate double precision NOT NULL DEFAULT 0, + connect_tls_p95_ms bigint NOT NULL DEFAULT 0, + first_byte_p95_ms bigint NOT NULL DEFAULT 0, + upload_bytes_per_second double precision NOT NULL DEFAULT 0, + jitter_p95_ms bigint NOT NULL DEFAULT 0, + consecutive_failures integer NOT NULL DEFAULT 0, + consecutive_successes integer NOT NULL DEFAULT 0, + sample_count integer NOT NULL DEFAULT 0, + sampled_at timestamptz NOT NULL, + expires_at timestamptz NOT NULL, + updated_at timestamptz NOT NULL DEFAULT now(), + PRIMARY KEY (pool_id, route_profile_key), + CONSTRAINT gateway_route_health_state_check + CHECK (state IN ('healthy', 'degraded', 'unreachable', 'unknown')), + CONSTRAINT gateway_route_health_values_check + CHECK ( + success_rate >= 0 AND success_rate <= 1 + AND connect_tls_p95_ms >= 0 + AND first_byte_p95_ms >= 0 + AND upload_bytes_per_second >= 0 + AND jitter_p95_ms >= 0 + AND consecutive_failures >= 0 + AND consecutive_successes >= 0 + AND sample_count >= 0 + AND expires_at >= sampled_at + ) +); + +CREATE INDEX IF NOT EXISTS idx_gateway_route_health_fresh + ON gateway_route_health(route_profile_key, expires_at, state, pool_id); + +CREATE TABLE IF NOT EXISTS gateway_route_preferences ( + route_profile_key text PRIMARY KEY, + current_pool_id text NOT NULL DEFAULT '', + current_since timestamptz, + challenger_pool_id text NOT NULL DEFAULT '', + challenger_wins integer NOT NULL DEFAULT 0, + updated_at timestamptz NOT NULL DEFAULT now(), + CONSTRAINT gateway_route_preferences_wins_check CHECK (challenger_wins >= 0) +); + +CREATE TABLE IF NOT EXISTS gateway_route_probe_requests ( + route_profile_key text PRIMARY KEY, + requested_at timestamptz NOT NULL DEFAULT now(), + expires_at timestamptz NOT NULL DEFAULT now() + interval '30 seconds' +); + +CREATE INDEX IF NOT EXISTS idx_gateway_route_probe_requests_due + ON gateway_route_probe_requests(expires_at, requested_at); + +CREATE TABLE IF NOT EXISTS gateway_pool_capacity_desires ( + pool_id text PRIMARY KEY REFERENCES gateway_execution_pools(pool_id) ON DELETE CASCADE, + desired integer NOT NULL, + reason text NOT NULL DEFAULT '', + valid_until timestamptz NOT NULL, + updated_at timestamptz NOT NULL DEFAULT now(), + CONSTRAINT gateway_pool_capacity_desires_value_check CHECK (desired >= 0) +); + +CREATE TABLE IF NOT EXISTS gateway_worker_execution_leases ( + lease_id uuid PRIMARY KEY, + task_id uuid NOT NULL REFERENCES gateway_tasks(id) ON DELETE CASCADE, + pool_id text NOT NULL REFERENCES gateway_execution_pools(pool_id) ON DELETE RESTRICT, + worker_id text NOT NULL, + instance_id text NOT NULL REFERENCES gateway_worker_instances(instance_id) ON DELETE CASCADE, + nonce_hash text NOT NULL UNIQUE, + state text NOT NULL DEFAULT 'reserved', + expires_at timestamptz NOT NULL, + released_at timestamptz, + created_at timestamptz NOT NULL DEFAULT now(), + updated_at timestamptz NOT NULL DEFAULT now(), + CONSTRAINT gateway_worker_execution_leases_state_check + CHECK (state IN ('reserved', 'running', 'released', 'expired')) +); + +CREATE INDEX IF NOT EXISTS idx_worker_execution_leases_active + ON gateway_worker_execution_leases(pool_id, instance_id, expires_at) + WHERE released_at IS NULL; + +ALTER TABLE gateway_tasks + ADD COLUMN assigned_pool_id text, + ADD COLUMN assigned_worker_id text, + ADD COLUMN route_profile_key text, + ADD COLUMN routing_platform_id uuid, + ADD COLUMN routing_platform_model_id uuid, + ADD COLUMN routing_version text, + ADD COLUMN routing_reason text, + ADD COLUMN routing_snapshot jsonb, + ADD COLUMN submission_state text DEFAULT 'not_started'; + +ALTER TABLE gateway_tasks + ADD CONSTRAINT gateway_tasks_assigned_pool_fk + FOREIGN KEY (assigned_pool_id) REFERENCES gateway_execution_pools(pool_id) ON DELETE SET NULL, + ADD CONSTRAINT gateway_tasks_routing_platform_fk + FOREIGN KEY (routing_platform_id) REFERENCES integration_platforms(id) ON DELETE SET NULL, + ADD CONSTRAINT gateway_tasks_routing_platform_model_fk + FOREIGN KEY (routing_platform_model_id) REFERENCES platform_models(id) ON DELETE SET NULL, + ADD CONSTRAINT gateway_tasks_routing_snapshot_object_check + CHECK (routing_snapshot IS NULL OR jsonb_typeof(routing_snapshot) = 'object') NOT VALID, + ADD CONSTRAINT gateway_tasks_submission_state_not_null_check + CHECK (num_nonnulls(submission_state) = 1) NOT VALID, + ADD CONSTRAINT gateway_tasks_submission_state_check + CHECK (submission_state IN ('not_started', 'submitting', 'submitted', 'submission_confirmation_pending', 'completed')) NOT VALID; + +ALTER TABLE gateway_tasks + VALIDATE CONSTRAINT gateway_tasks_submission_state_not_null_check; + +ALTER TABLE gateway_tasks + VALIDATE CONSTRAINT gateway_tasks_routing_snapshot_object_check; + +ALTER TABLE gateway_tasks + VALIDATE CONSTRAINT gateway_tasks_submission_state_check; + +CREATE INDEX IF NOT EXISTS idx_gateway_tasks_assigned_pool_queue + ON gateway_tasks(assigned_pool_id, status, next_run_at, priority, created_at) + WHERE async_mode = true AND status = 'queued'; diff --git a/deploy/kubernetes/local-acceptance/local-config.yaml b/deploy/kubernetes/local-acceptance/local-config.yaml index 567bd46..d454458 100644 --- a/deploy/kubernetes/local-acceptance/local-config.yaml +++ b/deploy/kubernetes/local-acceptance/local-config.yaml @@ -44,6 +44,7 @@ data: AI_GATEWAY_WORKER_MIN_REPLICAS_HONGKONG: "1" AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO: "2" AI_GATEWAY_WORKER_MAX_REPLICAS_HONGKONG: "2" + AI_GATEWAY_CAPACITY_POOLS: '[{"id":"ningbo","adapterRef":"ningbo","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":2},{"id":"hongkong","adapterRef":"hongkong","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":2}]' AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA: "48" AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS: "20" AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS: "600" diff --git a/deploy/kubernetes/production/application-config.yaml b/deploy/kubernetes/production/application-config.yaml index 7134bea..bd96a76 100644 --- a/deploy/kubernetes/production/application-config.yaml +++ b/deploy/kubernetes/production/application-config.yaml @@ -47,10 +47,16 @@ data: AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "48" AI_GATEWAY_ASYNC_WORKER_GLOBAL_HARD_LIMIT: "48" AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS: "5" + AI_GATEWAY_ROUTING_MODE: shadow + AI_GATEWAY_ROUTE_PROBE_ENABLED: "true" + AI_GATEWAY_ROUTE_PROBE_HOT_INTERVAL_SECONDS: "15" + AI_GATEWAY_ROUTE_PROBE_COLD_INTERVAL_SECONDS: "60" + AI_GATEWAY_ROUTE_PROBE_TIMEOUT_MS: "3000" + AI_GATEWAY_WORKER_ENDPOINT_ALLOW_PRIVATE: "true" + AI_GATEWAY_CAPACITY_ORCHESTRATOR_ADAPTER: kubernetes + AI_GATEWAY_CAPACITY_POOLS: >- + [{"id":"pool-cn-east-1","adapterRef":"ningbo","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1},{"id":"pool-hk-1","adapterRef":"hongkong","bootstrapReplicas":1,"minReplicas":1,"maxReplicas":1}] AI_GATEWAY_WORKER_AUTOSCALING_ENABLED: "false" - AI_GATEWAY_WORKER_REPLICAS_NINGBO: "1" - AI_GATEWAY_WORKER_MIN_REPLICAS_NINGBO: "1" - AI_GATEWAY_WORKER_MAX_REPLICAS_NINGBO: "1" AI_GATEWAY_WORKER_TARGET_OUTSTANDING_PER_REPLICA: "48" AI_GATEWAY_WORKER_SCALE_UP_WINDOW_SECONDS: "20" AI_GATEWAY_WORKER_SCALE_DOWN_STABILIZATION_SECONDS: "600" diff --git a/deploy/kubernetes/production/application.yaml b/deploy/kubernetes/production/application.yaml index f8b657c..9268d2e 100644 --- a/deploy/kubernetes/production/application.yaml +++ b/deploy/kubernetes/production/application.yaml @@ -94,6 +94,10 @@ spec: valueFrom: fieldRef: fieldPath: metadata.uid + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP - name: EASYAI_SITE value: ningbo ports: @@ -363,6 +367,10 @@ spec: valueFrom: fieldRef: fieldPath: metadata.uid + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP - name: EASYAI_SITE value: hongkong ports: @@ -540,8 +548,22 @@ spec: valueFrom: fieldRef: fieldPath: metadata.uid - - name: EASYAI_SITE - value: ningbo + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: AI_GATEWAY_WORKER_ID + value: pool-cn-east-1-worker + - name: AI_GATEWAY_WORKER_ADVERTISE_ENDPOINT + value: http://$(POD_IP):8088 + - name: AI_GATEWAY_ORCHESTRATOR_INSTANCE_REF + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: AI_GATEWAY_EXECUTION_POOL_ID + value: pool-cn-east-1 + - name: AI_GATEWAY_EXECUTION_POOL_LABELS + value: '{"region":"cn-east","zone":"ningbo","provider":"self-hosted","egress_profile":"direct"}' ports: - name: health containerPort: 8088 @@ -717,8 +739,22 @@ spec: valueFrom: fieldRef: fieldPath: metadata.uid - - name: EASYAI_SITE - value: hongkong + - name: POD_IP + valueFrom: + fieldRef: + fieldPath: status.podIP + - name: AI_GATEWAY_WORKER_ID + value: pool-hk-1-worker + - name: AI_GATEWAY_WORKER_ADVERTISE_ENDPOINT + value: http://$(POD_IP):8088 + - name: AI_GATEWAY_ORCHESTRATOR_INSTANCE_REF + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: AI_GATEWAY_EXECUTION_POOL_ID + value: pool-hk-1 + - name: AI_GATEWAY_EXECUTION_POOL_LABELS + value: '{"region":"ap-east","zone":"hongkong","provider":"self-hosted","egress_profile":"direct"}' - name: AI_GATEWAY_PLATFORM_PROXY_BYPASS_IDS value: 99372d7c-f2a4-472a-987f-30cb76c7962c ports: