perf(worker): 按实例内存上限扩展异步容量

将策略推导出的集群目标与单 Worker 内存安全上限分离,避免失活实例的全部容量转移到单个 2 GiB Pod。

生产环境单实例上限和媒体物化并发设为 24;两实例正常时总容量由 16 提升至 48,后续可通过增加 Worker 节点继续扩展。新增 0093 迁移保存实例容量上限,并兼容滚动发布中的旧实例。

风险:更高媒体并发会增加 Worker 和节点内存压力,保留单实例 24 的硬边界并由持续监控观察 RSS、OOM、队列与节点余量。

验证:go test ./... -count=1;go vet ./...;真实 PostgreSQL 分配与故障转移测试;迁移安全检查;kubectl kustomize;gofmt -l;git diff --check。
This commit is contained in:
2026-07-30 15:08:54 +08:00
parent ea58d21d03
commit 3aa78d3e5c
12 changed files with 180 additions and 25 deletions
+1 -1
View File
@@ -162,7 +162,7 @@ AI_GATEWAY_DATABASE_URL=postgresql://easyai:easyai2025@localhost:5432/easyai_ai_
如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。 如果现有 `easyai-pgvector` 没有把 `5432` 映射到宿主机,就需要补端口映射,或者把 AI Gateway 后端容器化后接入同一个 `easyai` Docker network。
异步队列 worker 不使用固定并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules``AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`仅作为单进程资源安全边界;平台模型和用户组的 PostgreSQL concurrency lease 是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。 异步队列 worker 不使用固定业务并发。服务分别汇总启用平台模型和活跃用户组的有效 `concurrent` 策略,采用两者中更严格的集群容量,默认每 5 秒在线调整 River 执行容量。策略解析兼容历史 `platformLimits/modelLimits.max_concurrent_requests`,运行时统一转换为 `rules``AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT` 默认 `2048`限制策略推导出的集群目标;`AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT` 默认 `32`,按单 Worker 的内存安全容量限制实例分配,即使其他实例失活也不会突破。平台模型和用户组的 PostgreSQL concurrency lease 是业务并发真值。可通过 `AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS` 调整刷新周期。
## 迁移原则 ## 迁移原则
+5
View File
@@ -63,6 +63,7 @@ type Config struct {
MediaMaterializationConcurrency int MediaMaterializationConcurrency int
AsyncQueueWorkerEnabled bool AsyncQueueWorkerEnabled bool
AsyncWorkerHardLimit int AsyncWorkerHardLimit int
AsyncWorkerInstanceHardLimit int
AsyncWorkerRefreshIntervalSeconds int AsyncWorkerRefreshIntervalSeconds int
} }
@@ -123,6 +124,7 @@ func Load() Config {
MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8), MediaMaterializationConcurrency: envIntValidated("AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY", 8),
AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true", AsyncQueueWorkerEnabled: env("AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED", "true") == "true",
AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048), AsyncWorkerHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", 2048),
AsyncWorkerInstanceHardLimit: envIntValidated("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", 32),
AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5), AsyncWorkerRefreshIntervalSeconds: envIntValidated("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS", 5),
} }
} }
@@ -150,6 +152,9 @@ func (c Config) Validate() error {
if c.AsyncWorkerHardLimit < 1 || c.AsyncWorkerHardLimit > 10000 { if c.AsyncWorkerHardLimit < 1 || c.AsyncWorkerHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT must be between 1 and 10000") return errors.New("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT must be between 1 and 10000")
} }
if c.AsyncWorkerInstanceHardLimit < 1 || c.AsyncWorkerInstanceHardLimit > 10000 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT must be between 1 and 10000")
}
if c.AsyncWorkerRefreshIntervalSeconds < 1 { if c.AsyncWorkerRefreshIntervalSeconds < 1 {
return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive") return errors.New("AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS must be positive")
} }
+24 -2
View File
@@ -26,7 +26,12 @@ func TestLoadIdentitySecretStoreUsesNewEnvironmentNamesAndIgnoresLegacyBusinessV
} }
func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) { func TestValidateIdentityFileSecretStoreRequiresDirectory(t *testing.T) {
cfg := Config{IdentitySecretStore: "file", AsyncWorkerHardLimit: 2048, AsyncWorkerRefreshIntervalSeconds: 5} cfg := Config{
IdentitySecretStore: "file",
AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "IDENTITY_SECRET_DIR") {
t.Fatalf("Validate() error = %v, want missing identity secret directory", err) t.Fatalf("Validate() error = %v, want missing identity secret directory", err)
} }
@@ -41,6 +46,7 @@ func TestValidateIdentityKubernetesSecretStore(t *testing.T) {
IdentitySecretStore: "kubernetes", IdentitySecretStore: "kubernetes",
IdentityKubernetesSecretName: "easyai-gateway-identity", IdentityKubernetesSecretName: "easyai-gateway-identity",
AsyncWorkerHardLimit: 2048, AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5, AsyncWorkerRefreshIntervalSeconds: 5,
} }
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "namespace") {
@@ -58,6 +64,7 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
IdentitySecurityEventStaleAfterSeconds: 60, IdentitySecurityEventStaleAfterSeconds: 60,
IdentitySecurityEventClockSkewSeconds: 60, IdentitySecurityEventClockSkewSeconds: 60,
AsyncWorkerHardLimit: 2048, AsyncWorkerHardLimit: 2048,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5, AsyncWorkerRefreshIntervalSeconds: 5,
} }
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "heartbeat") {
@@ -66,11 +73,20 @@ func TestValidateIdentitySecurityEventTiming(t *testing.T) {
} }
func TestValidateAsyncWorkerSettings(t *testing.T) { func TestValidateAsyncWorkerSettings(t *testing.T) {
cfg := Config{AsyncWorkerHardLimit: 10001, AsyncWorkerRefreshIntervalSeconds: 5} cfg := Config{
AsyncWorkerHardLimit: 10001,
AsyncWorkerInstanceHardLimit: 32,
AsyncWorkerRefreshIntervalSeconds: 5,
}
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid hard limit", err) t.Fatalf("Validate() error = %v, want invalid hard limit", err)
} }
cfg.AsyncWorkerHardLimit = 2048 cfg.AsyncWorkerHardLimit = 2048
cfg.AsyncWorkerInstanceHardLimit = 10001
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid instance hard limit", err)
}
cfg.AsyncWorkerInstanceHardLimit = 32
cfg.AsyncWorkerRefreshIntervalSeconds = 0 cfg.AsyncWorkerRefreshIntervalSeconds = 0
if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") { if err := cfg.Validate(); err == nil || !strings.Contains(err.Error(), "REFRESH_INTERVAL") {
t.Fatalf("Validate() error = %v, want invalid refresh interval", err) t.Fatalf("Validate() error = %v, want invalid refresh interval", err)
@@ -80,6 +96,12 @@ func TestValidateAsyncWorkerSettings(t *testing.T) {
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") { if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err) t.Fatalf("Validate() error = %v, want invalid non-integer hard limit", err)
} }
t.Setenv("AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT", "2048")
t.Setenv("AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT", "not-an-integer")
loaded = Load()
if err := loaded.Validate(); err == nil || !strings.Contains(err.Error(), "INSTANCE_HARD_LIMIT") {
t.Fatalf("Validate() error = %v, want invalid non-integer instance hard limit", err)
}
} }
func TestLoadAsyncQueueWorkerEnabled(t *testing.T) { func TestLoadAsyncQueueWorkerEnabled(t *testing.T) {
+8 -1
View File
@@ -220,6 +220,10 @@ func (s *Service) startRiverQueue(ctx context.Context, workerEnabled bool) error
s.logger.Info("async worker capacity initialized", s.logger.Info("async worker capacity initialized",
"capacity", snapshot.Capacity, "capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired, "desiredCapacity", snapshot.Desired,
"activeInstances", snapshot.ActiveInstances,
"globalCapacity", snapshot.GlobalCapacity,
"instanceID", snapshot.InstanceID,
"instanceHardLimit", s.cfg.AsyncWorkerInstanceHardLimit,
"hardLimit", snapshot.HardLimit, "hardLimit", snapshot.HardLimit,
"enabledModels", snapshot.EnabledModels, "enabledModels", snapshot.EnabledModels,
"unlimitedModels", snapshot.UnlimitedModels, "unlimitedModels", snapshot.UnlimitedModels,
@@ -289,13 +293,14 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")), Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")),
Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
DesiredCapacity: snapshot.Capacity, DesiredCapacity: snapshot.Capacity,
CapacityLimit: s.cfg.AsyncWorkerInstanceHardLimit,
HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second, HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second,
}) })
if err != nil { if err != nil {
return store.AsyncWorkerCapacitySnapshot{}, err return store.AsyncWorkerCapacitySnapshot{}, err
} }
snapshot.Capacity = allocation.Allocated snapshot.Capacity = allocation.Allocated
snapshot.GlobalCapacity = allocation.DesiredCapacity snapshot.GlobalCapacity = allocation.GlobalAllocated
snapshot.ActiveInstances = allocation.ActiveInstances snapshot.ActiveInstances = allocation.ActiveInstances
snapshot.InstanceID = allocation.InstanceID snapshot.InstanceID = allocation.InstanceID
return snapshot, nil return snapshot, nil
@@ -372,7 +377,9 @@ func (s *Service) resizeAsyncWorkerCapacity(ctx context.Context) {
"capacity", snapshot.Capacity, "capacity", snapshot.Capacity,
"desiredCapacity", snapshot.Desired, "desiredCapacity", snapshot.Desired,
"activeInstances", snapshot.ActiveInstances, "activeInstances", snapshot.ActiveInstances,
"globalCapacity", snapshot.GlobalCapacity,
"instanceID", snapshot.InstanceID, "instanceID", snapshot.InstanceID,
"instanceHardLimit", s.cfg.AsyncWorkerInstanceHardLimit,
"hardLimit", snapshot.HardLimit, "hardLimit", snapshot.HardLimit,
"modelDesired", snapshot.ModelDesired, "modelDesired", snapshot.ModelDesired,
"groupDesired", snapshot.GroupDesired, "groupDesired", snapshot.GroupDesired,
+3
View File
@@ -93,6 +93,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
if cfg.AsyncWorkerHardLimit == 0 { if cfg.AsyncWorkerHardLimit == 0 {
cfg.AsyncWorkerHardLimit = 2048 cfg.AsyncWorkerHardLimit = 2048
} }
if cfg.AsyncWorkerInstanceHardLimit == 0 {
cfg.AsyncWorkerInstanceHardLimit = 32
}
if cfg.AsyncWorkerRefreshIntervalSeconds == 0 { if cfg.AsyncWorkerRefreshIntervalSeconds == 0 {
cfg.AsyncWorkerRefreshIntervalSeconds = 5 cfg.AsyncWorkerRefreshIntervalSeconds = 5
} }
@@ -625,4 +625,17 @@ WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).Str
if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 { if err != nil || first.Allocated != 5 || first.ActiveInstances != 1 {
t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err) t.Fatalf("single-worker failover allocation = %+v, err=%v", first, err)
} }
first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
InstanceID: firstID, DesiredCapacity: 100, CapacityLimit: 3,
})
if err != nil || first.Allocated != 3 || first.GlobalAllocated != 3 {
t.Fatalf("bounded single-worker allocation = %+v, err=%v", first, err)
}
second, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{
InstanceID: secondID, DesiredCapacity: 100, CapacityLimit: 2,
})
if err != nil || second.Allocated != 2 || second.GlobalAllocated != 5 || second.ActiveInstances != 2 {
t.Fatalf("bounded two-worker allocation = %+v, err=%v", second, err)
}
} }
+60 -18
View File
@@ -18,6 +18,7 @@ type WorkerRegistrationInput struct {
Site string Site string
Revision string Revision string
DesiredCapacity int DesiredCapacity int
CapacityLimit int
HeartbeatStaleAfter time.Duration HeartbeatStaleAfter time.Duration
} }
@@ -25,10 +26,16 @@ type WorkerAllocation struct {
InstanceID string InstanceID string
DesiredCapacity int DesiredCapacity int
Allocated int Allocated int
GlobalAllocated int
ActiveInstances int ActiveInstances int
HeartbeatAt time.Time HeartbeatAt time.Time
} }
type activeWorkerCapacity struct {
InstanceID string
CapacityLimit int
}
// ActiveWorkerCapacity returns the cluster-wide execution capacity currently // ActiveWorkerCapacity returns the cluster-wide execution capacity currently
// owned by live Worker instances. // owned by live Worker instances.
func (s *Store) ActiveWorkerCapacity(ctx context.Context) (int, error) { func (s *Store) ActiveWorkerCapacity(ctx context.Context) (int, error) {
@@ -49,6 +56,12 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
if input.DesiredCapacity < 0 { if input.DesiredCapacity < 0 {
return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative") return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
} }
if input.CapacityLimit < 0 {
return WorkerAllocation{}, errors.New("worker capacity limit cannot be negative")
}
if input.CapacityLimit == 0 {
input.CapacityLimit = input.DesiredCapacity
}
staleAfter := input.HeartbeatStaleAfter staleAfter := input.HeartbeatStaleAfter
if staleAfter < workerHeartbeatStaleAfter { if staleAfter < workerHeartbeatStaleAfter {
staleAfter = workerHeartbeatStaleAfter staleAfter = workerHeartbeatStaleAfter
@@ -70,9 +83,9 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
if _, err := tx.Exec(ctx, ` if _, err := tx.Exec(ctx, `
INSERT INTO gateway_worker_instances ( INSERT INTO gateway_worker_instances (
instance_id, pod_uid, pod_name, site, revision, status, instance_id, pod_uid, pod_name, site, revision, status,
desired_capacity, allocated_capacity, started_at, heartbeat_at, updated_at desired_capacity, capacity_limit, allocated_capacity, started_at, heartbeat_at, updated_at
) )
VALUES ($1, $2, $3, $4, $5, 'active', $6, 0, now(), now(), now()) VALUES ($1, $2, $3, $4, $5, 'active', $6, $7, 0, now(), now(), now())
ON CONFLICT (instance_id) DO UPDATE ON CONFLICT (instance_id) DO UPDATE
SET pod_uid = EXCLUDED.pod_uid, SET pod_uid = EXCLUDED.pod_uid,
pod_name = EXCLUDED.pod_name, pod_name = EXCLUDED.pod_name,
@@ -80,6 +93,7 @@ SET pod_uid = EXCLUDED.pod_uid,
revision = EXCLUDED.revision, revision = EXCLUDED.revision,
status = 'active', status = 'active',
desired_capacity = EXCLUDED.desired_capacity, desired_capacity = EXCLUDED.desired_capacity,
capacity_limit = EXCLUDED.capacity_limit,
heartbeat_at = now(), heartbeat_at = now(),
updated_at = now()`, updated_at = now()`,
input.InstanceID, input.InstanceID,
@@ -88,6 +102,7 @@ SET pod_uid = EXCLUDED.pod_uid,
strings.TrimSpace(input.Site), strings.TrimSpace(input.Site),
strings.TrimSpace(input.Revision), strings.TrimSpace(input.Revision),
input.DesiredCapacity, input.DesiredCapacity,
input.CapacityLimit,
); err != nil { ); err != nil {
return WorkerAllocation{}, err return WorkerAllocation{}, err
} }
@@ -101,7 +116,7 @@ WHERE status <> 'active'
} }
rows, err := tx.Query(ctx, ` rows, err := tx.Query(ctx, `
SELECT instance_id SELECT instance_id, capacity_limit
FROM gateway_worker_instances FROM gateway_worker_instances
WHERE status = 'active' WHERE status = 'active'
AND heartbeat_at > now() - $1::interval AND heartbeat_at > now() - $1::interval
@@ -110,41 +125,37 @@ FOR UPDATE`, staleAfter.String())
if err != nil { if err != nil {
return WorkerAllocation{}, err return WorkerAllocation{}, err
} }
activeIDs := make([]string, 0) activeWorkers := make([]activeWorkerCapacity, 0)
for rows.Next() { for rows.Next() {
var instanceID string var worker activeWorkerCapacity
if err := rows.Scan(&instanceID); err != nil { if err := rows.Scan(&worker.InstanceID, &worker.CapacityLimit); err != nil {
rows.Close() rows.Close()
return WorkerAllocation{}, err return WorkerAllocation{}, err
} }
activeIDs = append(activeIDs, instanceID) activeWorkers = append(activeWorkers, worker)
} }
if err := rows.Err(); err != nil { if err := rows.Err(); err != nil {
rows.Close() rows.Close()
return WorkerAllocation{}, err return WorkerAllocation{}, err
} }
rows.Close() rows.Close()
if len(activeIDs) == 0 { if len(activeWorkers) == 0 {
return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat") return WorkerAllocation{}, errors.New("worker registration was not active after heartbeat")
} }
base := input.DesiredCapacity / len(activeIDs) allocations, globalAllocated := allocateWorkerCapacities(activeWorkers, input.DesiredCapacity)
remainder := input.DesiredCapacity % len(activeIDs)
allocated := 0 allocated := 0
for index, instanceID := range activeIDs { for _, worker := range activeWorkers {
capacity := base capacity := allocations[worker.InstanceID]
if index < remainder {
capacity++
}
if _, err := tx.Exec(ctx, ` if _, err := tx.Exec(ctx, `
UPDATE gateway_worker_instances UPDATE gateway_worker_instances
SET desired_capacity = $2, SET desired_capacity = $2,
allocated_capacity = $3, allocated_capacity = $3,
updated_at = now() updated_at = now()
WHERE instance_id = $1`, instanceID, input.DesiredCapacity, capacity); err != nil { WHERE instance_id = $1`, worker.InstanceID, input.DesiredCapacity, capacity); err != nil {
return WorkerAllocation{}, err return WorkerAllocation{}, err
} }
if instanceID == input.InstanceID { if worker.InstanceID == input.InstanceID {
allocated = capacity allocated = capacity
} }
} }
@@ -162,11 +173,42 @@ WHERE instance_id = $1`, input.InstanceID).Scan(&heartbeatAt); err != nil {
InstanceID: input.InstanceID, InstanceID: input.InstanceID,
DesiredCapacity: input.DesiredCapacity, DesiredCapacity: input.DesiredCapacity,
Allocated: allocated, Allocated: allocated,
ActiveInstances: len(activeIDs), GlobalAllocated: globalAllocated,
ActiveInstances: len(activeWorkers),
HeartbeatAt: heartbeatAt, HeartbeatAt: heartbeatAt,
}, nil }, nil
} }
func allocateWorkerCapacities(workers []activeWorkerCapacity, desired int) (map[string]int, int) {
allocations := make(map[string]int, len(workers))
if desired <= 0 || len(workers) == 0 {
return allocations, 0
}
remaining := desired
for remaining > 0 {
progressed := false
for _, worker := range workers {
limit := worker.CapacityLimit
if limit <= 0 {
limit = desired
}
if allocations[worker.InstanceID] >= limit {
continue
}
allocations[worker.InstanceID]++
remaining--
progressed = true
if remaining == 0 {
break
}
}
if !progressed {
break
}
}
return allocations, desired - remaining
}
func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error { func (s *Store) MarkWorkerDraining(ctx context.Context, instanceID string) error {
result, err := s.pool.Exec(ctx, ` result, err := s.pool.Exec(ctx, `
UPDATE gateway_worker_instances UPDATE gateway_worker_instances
@@ -0,0 +1,44 @@
package store
import "testing"
func TestAllocateWorkerCapacitiesHonorsInstanceLimits(t *testing.T) {
workers := []activeWorkerCapacity{
{InstanceID: "worker-a", CapacityLimit: 24},
{InstanceID: "worker-b", CapacityLimit: 24},
}
allocations, global := allocateWorkerCapacities(workers, 310)
if global != 48 || allocations["worker-a"] != 24 || allocations["worker-b"] != 24 {
t.Fatalf("allocations=%v global=%d, want 24/24 and 48", allocations, global)
}
}
func TestAllocateWorkerCapacitiesBalancesBelowLimits(t *testing.T) {
workers := []activeWorkerCapacity{
{InstanceID: "worker-a", CapacityLimit: 24},
{InstanceID: "worker-b", CapacityLimit: 24},
}
allocations, global := allocateWorkerCapacities(workers, 5)
if global != 5 || allocations["worker-a"] != 3 || allocations["worker-b"] != 2 {
t.Fatalf("allocations=%v global=%d, want 3/2 and 5", allocations, global)
}
}
func TestAllocateWorkerCapacitiesKeepsSurvivorBounded(t *testing.T) {
workers := []activeWorkerCapacity{{InstanceID: "worker-a", CapacityLimit: 24}}
allocations, global := allocateWorkerCapacities(workers, 310)
if global != 24 || allocations["worker-a"] != 24 {
t.Fatalf("allocations=%v global=%d, want survivor capped at 24", allocations, global)
}
}
func TestAllocateWorkerCapacitiesSupportsUnequalLimits(t *testing.T) {
workers := []activeWorkerCapacity{
{InstanceID: "worker-a", CapacityLimit: 1},
{InstanceID: "worker-b", CapacityLimit: 4},
}
allocations, global := allocateWorkerCapacities(workers, 5)
if global != 5 || allocations["worker-a"] != 1 || allocations["worker-b"] != 4 {
t.Fatalf("allocations=%v global=%d, want 1/4 and 5", allocations, global)
}
}
@@ -0,0 +1,8 @@
ALTER TABLE gateway_worker_instances
ADD COLUMN capacity_limit integer NOT NULL DEFAULT 0,
ADD CONSTRAINT gateway_worker_instances_capacity_limit_check
CHECK (capacity_limit >= 0);
UPDATE gateway_worker_instances
SET capacity_limit = GREATEST(allocated_capacity, 1)
WHERE capacity_limit = 0;
@@ -35,5 +35,5 @@ data:
AI_GATEWAY_LOCAL_RESULT_MAX_BYTES: "268435456" AI_GATEWAY_LOCAL_RESULT_MAX_BYTES: "268435456"
AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES: "536870912" AI_GATEWAY_LOCAL_RESULT_MAX_TASK_BYTES: "536870912"
AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED: "true" AI_GATEWAY_ASYNC_QUEUE_WORKER_ENABLED: "true"
AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "16" AI_GATEWAY_ASYNC_WORKER_HARD_LIMIT: "2048"
AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS: "5" AI_GATEWAY_ASYNC_WORKER_REFRESH_INTERVAL_SECONDS: "5"
@@ -312,8 +312,10 @@ spec:
value: "true" value: "true"
- name: AI_GATEWAY_DATABASE_MAX_CONNS - name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "24" value: "24"
- name: AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT
value: "24"
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY - name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
value: "8" value: "24"
- name: POD_NAMESPACE - name: POD_NAMESPACE
valueFrom: valueFrom:
fieldRef: fieldRef:
@@ -439,8 +441,10 @@ spec:
value: "true" value: "true"
- name: AI_GATEWAY_DATABASE_MAX_CONNS - name: AI_GATEWAY_DATABASE_MAX_CONNS
value: "24" value: "24"
- name: AI_GATEWAY_ASYNC_WORKER_INSTANCE_HARD_LIMIT
value: "24"
- name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY - name: AI_GATEWAY_MEDIA_MATERIALIZATION_CONCURRENCY
value: "8" value: "24"
- name: POD_NAMESPACE - name: POD_NAMESPACE
valueFrom: valueFrom:
fieldRef: fieldRef:
+7
View File
@@ -8,6 +8,13 @@
"non-null column addition" "non-null column addition"
], ],
"reason": "The migration replaces two existing constraints and one index with expanded definitions. Production preflight confirmed six schema-v1 identity revisions, no OIDC sessions or revocation watermarks, and small affected relations; the non-null columns use constant defaults." "reason": "The migration replaces two existing constraints and one index with expanded definitions. Production preflight confirmed six schema-v1 identity revisions, no OIDC sessions or revocation watermarks, and small affected relations; the non-null columns use constant defaults."
},
"apps/api/migrations/0093_worker_instance_capacity_limit.sql": {
"sha256": "4b6a39ef90e7032555739ccf0568bf7ae0bce23ef98c3bee9c2c8e3bfa13291e",
"allowedViolations": [
"non-null column addition"
],
"reason": "The table stores ephemeral Worker heartbeats and currently has only two production rows. PostgreSQL adds the integer column with a constant default, then the migration backfills each live row from its existing allocation so a rolling deployment cannot assign unsafe failover capacity."
} }
} }
} }