feat(routing): 引入多执行池智能调度
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
|
||||
@@ -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: 管理端查看平台模型维度的限流和冷却状态。
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
@@ -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 = ""
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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{}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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{"*"},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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}
|
||||
}
|
||||
@@ -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{}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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.",
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 != "" {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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';
|
||||
Reference in New Issue
Block a user