wangbo
7786692d32
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
265 lines
11 KiB
Go
265 lines
11 KiB
Go
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}
|
|
}
|