From ba154e4ec2fd22305556c8ac7e810e2639219330 Mon Sep 17 00:00:00 2001 From: Mikhail Date: Mon, 24 Aug 2026 19:34:35 +0300 Subject: [PATCH] fix: keep local readiness responsive --- cmd/cognitive-worker/fleet_status.go | 71 ++++++++++++++--------- cmd/cognitive-worker/fleet_status_test.go | 24 ++++++++ cmd/cognitive-worker/main.go | 8 ++- 3 files changed, 72 insertions(+), 31 deletions(-) diff --git a/cmd/cognitive-worker/fleet_status.go b/cmd/cognitive-worker/fleet_status.go index bd007cf..51f4073 100644 --- a/cmd/cognitive-worker/fleet_status.go +++ b/cmd/cognitive-worker/fleet_status.go @@ -6,6 +6,7 @@ import ( "fmt" "net/http" "strings" + "sync" "time" "git.cin.su/m/Assistant/contracts/capability" @@ -31,35 +32,16 @@ type fleetStatusResponse struct { func fleetReadinessHandler(client *http.Client, preferred capability.FleetTier, endpoints []fleetEndpoint) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { members := make([]capability.FleetMember, 0, len(endpoints)) - for _, endpoint := range endpoints { - member := capability.FleetMember{ProviderInstanceID: capability.ProviderInstanceID(endpoint.ProviderID), ModelDescriptorID: capability.ModelDescriptorID(endpoint.ModelID), Tier: endpoint.Tier, Reason: "NOT_CONFIGURED"} - if strings.TrimSpace(endpoint.URL) == "" { - members = append(members, member) - continue - } - request, err := http.NewRequestWithContext(r.Context(), http.MethodGet, endpoint.URL, nil) - if err == nil { - for name, value := range endpoint.Headers { - request.Header.Set(name, value) - } - probeClient := endpoint.Client - if probeClient == nil { - probeClient = client - } - response, requestErr := probeClient.Do(request) - if requestErr == nil { - response.Body.Close() - if response.StatusCode >= 200 && response.StatusCode < 300 { - member.Ready, member.Reason = true, "AVAILABLE" - } else { - member.Reason = fmt.Sprintf("status=%s", response.Status) - } - } else { - member.Reason = requestErr.Error() - } - } - members = append(members, member) + members = make([]capability.FleetMember, len(endpoints)) + var waitGroup sync.WaitGroup + for index, endpoint := range endpoints { + waitGroup.Add(1) + go func(index int, endpoint fleetEndpoint) { + defer waitGroup.Done() + members[index] = probeFleetEndpoint(r.Context(), client, endpoint) + }(index, endpoint) } + waitGroup.Wait() readiness, err := fleetdomain.BuildFleetReadiness(time.Now().UnixMilli(), members) response := fleetStatusResponse{Readiness: readiness, Preferred: preferred} if err == nil { @@ -76,4 +58,37 @@ func fleetReadinessHandler(client *http.Client, preferred capability.FleetTier, }) } +func probeFleetEndpoint(parent context.Context, fallbackClient *http.Client, endpoint fleetEndpoint) capability.FleetMember { + member := capability.FleetMember{ProviderInstanceID: capability.ProviderInstanceID(endpoint.ProviderID), ModelDescriptorID: capability.ModelDescriptorID(endpoint.ModelID), Tier: endpoint.Tier, Reason: "NOT_CONFIGURED"} + if strings.TrimSpace(endpoint.URL) == "" { + return member + } + ctx, cancel := context.WithTimeout(parent, 2*time.Second) + defer cancel() + request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.URL, nil) + if err != nil { + member.Reason = err.Error() + return member + } + for name, value := range endpoint.Headers { + request.Header.Set(name, value) + } + probeClient := endpoint.Client + if probeClient == nil { + probeClient = fallbackClient + } + response, requestErr := probeClient.Do(request) + if requestErr != nil { + member.Reason = requestErr.Error() + return member + } + defer response.Body.Close() + if response.StatusCode >= 200 && response.StatusCode < 300 { + member.Ready, member.Reason = true, "AVAILABLE" + } else { + member.Reason = fmt.Sprintf("status=%s", response.Status) + } + return member +} + func fleetContext(ctx context.Context) context.Context { return ctx } diff --git a/cmd/cognitive-worker/fleet_status_test.go b/cmd/cognitive-worker/fleet_status_test.go index 45f1296..0f223e8 100644 --- a/cmd/cognitive-worker/fleet_status_test.go +++ b/cmd/cognitive-worker/fleet_status_test.go @@ -5,6 +5,7 @@ import ( "net/http/httptest" "strings" "testing" + "time" "git.cin.su/m/Assistant/contracts/capability" ) @@ -26,6 +27,29 @@ func TestFleetReadinessSelectsBestAvailablePrimaryTierBeforeLifeboat(t *testing. } } +func TestFleetReadinessDoesNotWaitForSlowMembersSerially(t *testing.T) { + slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(150 * time.Millisecond) + w.WriteHeader(http.StatusServiceUnavailable) + })) + defer slow.Close() + fast := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })) + defer fast.Close() + handler := fleetReadinessHandler(fast.Client(), capability.FleetTierFast, []fleetEndpoint{ + {ProviderID: "slow", ModelID: "slow-model", Tier: capability.FleetTierMedium, URL: slow.URL}, + {ProviderID: "lifeboat", ModelID: "lifeboat-model", Tier: capability.FleetTierLifeboat, URL: fast.URL}, + }) + started := time.Now() + recorder := httptest.NewRecorder() + handler.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/v1/fleet/readiness", nil)) + if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), `"selected_tier":"LIFEBOAT"`) { + t.Fatalf("status=%d body=%s", recorder.Code, recorder.Body.String()) + } + if elapsed := time.Since(started); elapsed >= 300*time.Millisecond { + t.Fatalf("fleet readiness was serialized: elapsed=%s", elapsed) + } +} + func TestFleetReadinessSendsConfiguredProviderHeaders(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Header.Get("Authorization") != "Bearer fleet-key" { diff --git a/cmd/cognitive-worker/main.go b/cmd/cognitive-worker/main.go index ae53a85..9b679b5 100644 --- a/cmd/cognitive-worker/main.go +++ b/cmd/cognitive-worker/main.go @@ -126,7 +126,9 @@ func main() { } cpuEndpoint := strings.TrimRight(envOr("LOCAL_CPU_PROVIDER_URL", "http://docker.cin.su:18082/v1"), "/") + "/models" cpuHealthProbe := health.OptionalHTTPProbeWithHeaders(cpuReadinessClient, cpuEndpoint, cpuFleetHeaders) - cpuHealthProbe.Timeout = 10 * time.Second + // Opportunistic providers must not hide a ready local lifeboat. Keep their + // readiness budget short; their detailed fleet state is reported separately. + cpuHealthProbe.Timeout = 2 * time.Second mediumEndpoint := "" mediumHeaders := map[string]string{} mediumHealthProbe := health.Probe{} @@ -136,7 +138,7 @@ func main() { mediumHeaders["Authorization"] = "Bearer " + mediumAPIKey } mediumHealthProbe = health.OptionalHTTPProbeWithHeaders(mediumReadinessClient, mediumEndpoint, mediumHeaders) - mediumHealthProbe.Timeout = 10 * time.Second + mediumHealthProbe.Timeout = 2 * time.Second } readinessProbes := []health.Probe{ health.HTTPProbe(readinessClient, strings.TrimRight(envOr("CORE_RUNTIME_URL", "http://core-runtime:8080"), "/")+"/healthz"), @@ -149,7 +151,7 @@ func main() { if gpuConfigured { gpuHeaders := map[string]string{"Authorization": "Bearer " + gpuKey} gpuProbe := health.OptionalHTTPProbeWithHeaders(readinessClient, strings.TrimRight(gpuURL, "/")+"/models", gpuHeaders) - gpuProbe.Timeout = 10 * time.Second + gpuProbe.Timeout = 2 * time.Second readinessProbes = append(readinessProbes, gpuProbe) } -- 2.54.0