From b2754b871faba9ecb493ec32381e87abd932a0f1 Mon Sep 17 00:00:00 2001 From: toki Date: Sun, 5 Jul 2026 19:26:06 +0900 Subject: [PATCH] feat: long context admission support - refactor edge server, add input estimator --- .../code_review_local_G08_0.log | 152 ++++++++++++++++++ .../02+01_long_capacity/complete.log | 46 ++++++ .../02+01_long_capacity/plan_local_G08_0.log} | 0 .../CODE_REVIEW-local-G08.md | 103 ------------ apps/edge/internal/configrefresh/classify.go | 20 +++ .../internal/configrefresh/classify_test.go | 94 +++++++++++ apps/edge/internal/service/model_queue.go | 118 ++++++++++---- .../edge/internal/service/model_queue_test.go | 139 +++++++++++++++- apps/edge/internal/service/run_dispatch.go | 28 ++-- .../internal/service/service_internal_test.go | 63 ++++++++ configs/edge.yaml | 9 ++ packages/go/config/config.go | 52 ++++++ packages/go/config/config_test.go | 128 +++++++++++++++ 13 files changed, 803 insertions(+), 149 deletions(-) create mode 100644 agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/code_review_local_G08_0.log create mode 100644 agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/complete.log rename agent-task/{m-model-group-long-context-admission/02+01_long_capacity/PLAN-local-G08.md => archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/plan_local_G08_0.log} (100%) delete mode 100644 agent-task/m-model-group-long-context-admission/02+01_long_capacity/CODE_REVIEW-local-G08.md diff --git a/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/code_review_local_G08_0.log b/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/code_review_local_G08_0.log new file mode 100644 index 0000000..8b13cbb --- /dev/null +++ b/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/code_review_local_G08_0.log @@ -0,0 +1,152 @@ + + +# Code Review Reference - LONG_CAPACITY + +> **[IMPLEMENTING AGENT — READ FIRST] Filling in this file is the mandatory final step of implementation.** +> Complete the checklist, fill implementation-owned sections, then stop with active files in place and report ready for review. Do not run review finalization. Do not ask the user directly during implementation. + +## 개요 + +date=2026-07-05 +task=m-model-group-long-context-admission/02+01_long_capacity, plan=0, tag=LONG_CAPACITY + +## Roadmap Targets + +- Milestone: `agent-roadmap/phase/operational-observability-provider-management/milestones/model-group-long-context-admission.md` +- Task ids: + - `provider-long-capacity`: provider별 total context budget, long capacity, long_in_flight 상태와 동시 점유 + - `provider-exclusion`: long slot이 찬 provider를 long 후보에서 제외하고 normal 후보는 유지 +- Completion mode: check-on-pass + +## 이 파일을 읽는 리뷰 에이전트에게 + +> **[REVIEW AGENT ONLY]** 종결 절차는 코드리뷰 에이전트 전용이다. PASS이면 완료 이벤트만 보고하고 roadmap을 직접 수정하지 않는다. + +## 구현 항목별 완료 여부 + +| 항목 | 완료 여부 | +|------|---------| +| [LONG_CAPACITY-1] Provider long capacity config와 validation | [x] | +| [LONG_CAPACITY-2] Long slot reserve/release와 provider exclusion | [x] | + +## 구현 체크리스트 + +- [x] provider config에 `total_context_tokens`, `long_context_capacity`를 추가하고 `total_context_tokens >= context_window_tokens * long_context_capacity`를 config load/check에서 검증한다. +- [x] config refresh classifier가 provider long capacity 관련 field를 live-apply로 분류하게 한다. +- [x] long 요청 dispatch 시 일반 capacity와 long capacity를 함께 reserve/release하고 `long_in_flight` 상태를 유지한다. 검증: long dispatch 중 두 counter가 증가하고 완료/오류/취소 후 회복된다. +- [x] long slot이 full인 provider는 long 후보에서 제외하되, normal 요청 후보에서는 일반 capacity만 기준으로 유지한다. 검증: long은 다른 provider 또는 queue 대기, normal은 남은 일반 slot으로 dispatch된다. +- [x] config/service/configrefresh regression tests와 config check를 실행한다. +- [x] CODE_REVIEW-*-G??.md의 구현 에이전트 소유 섹션을 실제 구현 내용과 검증 출력으로 채운다. 이 항목이 완료되기 전에는 구현이 완료된 것이 아니다. + +## 구현 요약 + +- `packages/go/config/config.go`: `NodeProviderConf`에 `TotalContextTokens`, `LongContextCapacity` 추가. provider-level `Validate()`에 non-negative 검증과 `long_context_capacity > 0 → total_context_tokens > 0` 불변식 추가. `LoadEdge`에서 `providerByID` 맵을 만들고 `validateProviderLongContextBudget`로 `context_window_tokens > 0`인 model group이 참조하는 provider마다 `total_context_tokens >= context_window_tokens * long_context_capacity`를 검증. +- `configs/edge.yaml`: 두 field 주석 문서화, `mac-mlx-vllm` provider에 `total_context_tokens: 524288`, `long_context_capacity: 2` 예시 추가(262144*2 만족). +- `apps/edge/internal/configrefresh/classify.go`: `providerKey`에 두 field 추가, `buildProviderIndex`에서 채우고, capacity/priority와 동일하게 `StatusApplied`(live-apply)로 diff. +- `apps/edge/internal/service/model_queue.go`: `candidateNode.longContextCapacity`, `modelQueueGroup.longInflight`, `queueItem.long`, `inflightRec.long` 추가. `admit`/`trackInflight`에 variadic long flag(파일 기존 `releaseSlot` variadic 관례 준수), `findAvailableNodeLocked`/`releaseSlotLocked`에 명시적 long 파라미터. long 요청은 일반 `inflight`와 `longInflight`를 함께 예약/해제하고, long slot full provider는 long 후보에서만 제외. `longInFlightForProvider` 조회 helper 추가. 미사용 `releaseSlotByCandidate` 제거. +- `apps/edge/internal/service/run_dispatch.go`: `contextClassLong` 상수, provider-pool candidate에 `longContextCapacity` 전달, `submitRunQueued`에서 `ContextClass == "long"`를 admit/track/release 경로에 연결. + +## 코드리뷰 전용 체크리스트 + +- [x] 판정을 append한다. +- [x] active review/plan을 `code_review_local_G08_0.log`, `plan_local_G08_0.log`로 아카이브한다. +- [x] PASS이면 `complete.log` 작성 후 active task 디렉터리를 archive로 이동한다. +- [x] PASS split 작업이면 빈 active parent 처리 여부를 확인한다. active parent에는 sibling subtask가 남아 있어 유지했다. +- [ ] WARN/FAIL이면 다음 active plan/review 또는 `USER_REVIEW.md`를 작성한다. + +## 계획 대비 변경 사항 + +- 계획의 gate 식(`req.ContextClass == "long" && cand.longContextCapacity > 0 && cand.longInFlight >= cand.longContextCapacity → continue`)은 그대로 `findAvailableNodeLocked`에 적용했다. 다만 상태를 후보 struct 필드가 아니라 `modelQueueGroup.longInflight` 카운터로 보관해 기존 `inflight` 관리 방식과 대칭을 맞췄다. +- `admit`/`trackInflight`의 long flag는 별도 파라미터를 추가하되 기존 테스트/호출부 20여 곳을 깨지 않도록 파일에 이미 존재하는 `releaseSlot(... providerID ...string)` variadic 관례를 따라 variadic으로 넣었다. 시그니처가 내부 전용인 `findAvailableNodeLocked`/`releaseSlotLocked`는 명시적 파라미터로 바꾸고 해당 테스트 호출부만 갱신했다. +- 계획에 없던 정리: 호출부가 없는 `releaseSlotByCandidate`를 제거했다(grep 결과 주석 1개 외 참조 없음). 범위 내 비동작성 정리로 판단. + +## 주요 설계 결정 + +- `long_context_capacity == 0` provider는 long 요청 후보에서 제외하지 않고 long slot도 예약하지 않는다(계획 gate의 `> 0` 조건 그대로). 즉 long slot 카운터는 항상 capacity 이내에서만 의미를 가진다. long을 서빙할 provider는 SDD 표준선대로 `long_context_capacity > 0`을 설정한다. +- long in-flight 조회 helper(`longInFlightForProvider`)는 `group.longInflight` 예약 카운터가 아니라 기존 `getStatsForProviderLocked`와 동일하게 `inflightByRun`의 `long` flag를 세어 도출한다. proto `ProviderSnapshot` 노출은 04 범위이므로 service-local 상태/조회에 한정했다. +- budget 검증은 `context_window_tokens > 0`인 model group이 참조하는 provider에만 적용한다. `long_context_capacity == 0`이면 필요 budget이 0이라 자연히 통과하므로 무조건 비교로 단순화했다. +- config refresh 분류는 계획대로 두 field 모두 `StatusApplied`(live-apply)로 두어 capacity/priority와 일관되게 했다. + +## 사용자 리뷰 요청 + +_기본값은 `없음`이다. 구현 중 새 결정이 필요해 보여도 직접 질문하거나 선택지를 제시하거나 `request_user_input`을 호출하지 않는다. 이 섹션은 선택된 Milestone `구현 잠금 > 결정 필요` 항목이 실구현을 차단할 때만 채운다. 외부 환경/secret/서비스 준비, 검증 증거 공백, 반복 실패, 일반 범위 조정은 사용자 리뷰 요청이 아니며 `검증 결과`, `계획 대비 변경 사항`, 또는 code-review의 일반 follow-up plan으로 처리한다._ + +- 상태: 없음 +- 사유 유형: 없음 +- 연결 대상: 없음 +- 결정 필요: 없음 +- 차단 근거: 없음 +- 실행한 검증/명령: 없음 +- 자동 후속 불가 이유: 없음 +- 재개 조건: 없음 + +## 리뷰어를 위한 체크포인트 + +- `total_context_tokens >= context_window_tokens * long_context_capacity` 검증이 model provider mapping 기준으로 동작하는지 확인한다. +- long request가 일반 capacity와 long capacity를 함께 점유하고 terminal path에서 둘 다 회복하는지 확인한다. +- long slot full provider가 normal 후보에서 제외되지 않는지 확인한다. + +## 검증 결과 + +### LONG_CAPACITY-1 중간 검증 + +```bash +$ go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh +ok iop/packages/go/config 0.034s +ok iop/apps/edge/internal/configrefresh 0.015s +``` + +새 테스트: `TestLoadEdgeProviderLongContextCapacityValidation`(valid + negative total/long + long-without-total), `TestLoadEdgeRejectsProviderLongContextBudgetBelowModelWindow`, `TestClassifyProviderLongContextCapacityApplied`. + +### LONG_CAPACITY-2 중간 검증 + +```bash +$ go test -count=1 ./apps/edge/internal/service +ok iop/apps/edge/internal/service 0.463s +``` + +새 테스트: `TestModelQueueLongDispatchReservesBothCounters`(complete/error/cancelled 3 subtest), `TestModelQueueNormalRequestDoesNotReserveLongSlot`, `TestModelQueueLongSlotFullExcludesLongKeepsNormal`, `TestRefreshProviderLongCapacityAffectsCandidate`. + +### 최종 검증 + +```bash +$ go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh ./apps/edge/internal/edgevalidate ./apps/edge/internal/service +ok iop/packages/go/config 0.036s +ok iop/apps/edge/internal/configrefresh 0.022s +ok iop/apps/edge/internal/edgevalidate 0.006s +ok iop/apps/edge/internal/service 0.463s +$ go run ./apps/edge/cmd/edge config check --config configs/edge.yaml +OK configs/edge.yaml +``` + +보조 확인: `go build ./...` 통과, `gofmt -l`(변경 파일) 출력 없음, `git diff --check` whitespace 오류 없음. + +--- + +> **[IMPLEMENTING AGENT — BEFORE SAVING] Have you filled in every implementation-owned section: completion table, implementation checklist, changes from plan, design decisions, user review request state, and verification output?** + +## 코드리뷰 결과 + +- 종합 판정: PASS +- 차원별 평가: + - correctness: Pass + - completeness: Pass + - test coverage: Pass + - API contract: Pass + - code quality: Pass + - implementation deviation: Pass + - verification trust: Pass + - spec conformance: Pass +- 발견된 문제: 없음 +- 다음 단계: PASS 종결 절차를 수행한다. active review/plan을 로그로 아카이브하고 `complete.log`를 작성한 뒤 task directory를 archive로 이동한다. + +검토 메모: +- `total_context_tokens`/`long_context_capacity` schema, validation, refresh classification, provider-pool candidate propagation이 plan과 `iop.edge-config-runtime-refresh` 계약에 맞게 반영됐다. +- long 요청은 일반 `inflight`와 `longInflight`를 함께 reserve/release하고, terminal `complete`/`error`/`cancelled` 경로에서 둘 다 회복된다. +- long slot full provider는 long 후보에서 제외되지만 normal 요청 후보에서는 유지된다. status/proto 노출은 `04+01,02_status_logs` 범위로 분리되어 있어 이번 task의 service-local evidence 범위와 충돌하지 않는다. +- 재검증: + - `go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh ./apps/edge/internal/edgevalidate ./apps/edge/internal/service` PASS + - `go run ./apps/edge/cmd/edge config check --config configs/edge.yaml` PASS, `OK configs/edge.yaml` + - `gofmt -l` on changed Go files PASS, no output + - `git diff --check` PASS, no output + - `go build ./...` PASS diff --git a/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/complete.log b/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/complete.log new file mode 100644 index 0000000..5259968 --- /dev/null +++ b/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/complete.log @@ -0,0 +1,46 @@ +# Complete - m-model-group-long-context-admission/02+01_long_capacity + +## 완료 일시 + +2026-07-05 + +## 요약 + +Model Group Long-Context Admission의 provider long capacity/config admission 작업을 1회 루프로 완료했다. 최종 판정: PASS. + +## 루프 이력 + +| Plan | Review | Verdict | 메모 | +|------|--------|---------|------| +| `plan_local_G08_0.log` | `code_review_local_G08_0.log` | PASS | provider long capacity schema/validation, refresh classification, long slot reserve/release, provider exclusion 검증 완료 | + +## 구현/정리 내용 + +- `NodeProviderConf`에 `total_context_tokens`, `long_context_capacity`를 추가하고 provider/model catalog 기준 budget validation을 추가했다. +- config refresh classifier가 provider long capacity 관련 field를 live-apply 변경으로 분류하게 했다. +- provider-pool long 요청이 일반 `inflight`와 `longInflight`를 함께 예약/해제하고, long slot이 full인 provider를 long 후보에서 제외하되 normal 후보에는 유지하게 했다. +- 관련 config/service/configrefresh regression tests와 `configs/edge.yaml` 예시를 갱신했다. + +## 최종 검증 + +- `go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh ./apps/edge/internal/edgevalidate ./apps/edge/internal/service` - PASS; `packages/go/config`, `configrefresh`, `edgevalidate`, `service` 모두 `ok` +- `go run ./apps/edge/cmd/edge config check --config configs/edge.yaml` - PASS; `OK configs/edge.yaml` +- `gofmt -l packages/go/config/config.go packages/go/config/config_test.go apps/edge/internal/configrefresh/classify.go apps/edge/internal/configrefresh/classify_test.go apps/edge/internal/service/run_dispatch.go apps/edge/internal/service/model_queue.go apps/edge/internal/service/model_queue_test.go apps/edge/internal/service/service_internal_test.go` - PASS; output 없음 +- `git diff --check` - PASS; output 없음 +- `go build ./...` - PASS; output 없음 + +## Roadmap Completion + +- Milestone: `agent-roadmap/phase/operational-observability-provider-management/milestones/model-group-long-context-admission.md` +- Completed task ids: + - `provider-long-capacity`: PASS; evidence=`agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/plan_local_G08_0.log`, `agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/code_review_local_G08_0.log`; verification=`go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh ./apps/edge/internal/edgevalidate ./apps/edge/internal/service`, `go run ./apps/edge/cmd/edge config check --config configs/edge.yaml` + - `provider-exclusion`: PASS; evidence=`agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/plan_local_G08_0.log`, `agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/code_review_local_G08_0.log`; verification=`go test -count=1 ./apps/edge/internal/service` +- Not completed task ids: 없음 + +## 잔여 Nit + +- 없음 + +## 후속 작업 + +- 없음 diff --git a/agent-task/m-model-group-long-context-admission/02+01_long_capacity/PLAN-local-G08.md b/agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/plan_local_G08_0.log similarity index 100% rename from agent-task/m-model-group-long-context-admission/02+01_long_capacity/PLAN-local-G08.md rename to agent-task/archive/2026/07/m-model-group-long-context-admission/02+01_long_capacity/plan_local_G08_0.log diff --git a/agent-task/m-model-group-long-context-admission/02+01_long_capacity/CODE_REVIEW-local-G08.md b/agent-task/m-model-group-long-context-admission/02+01_long_capacity/CODE_REVIEW-local-G08.md deleted file mode 100644 index b24df22..0000000 --- a/agent-task/m-model-group-long-context-admission/02+01_long_capacity/CODE_REVIEW-local-G08.md +++ /dev/null @@ -1,103 +0,0 @@ - - -# Code Review Reference - LONG_CAPACITY - -> **[IMPLEMENTING AGENT — READ FIRST] Filling in this file is the mandatory final step of implementation.** -> Complete the checklist, fill implementation-owned sections, then stop with active files in place and report ready for review. Do not run review finalization. Do not ask the user directly during implementation. - -## 개요 - -date=2026-07-05 -task=m-model-group-long-context-admission/02+01_long_capacity, plan=0, tag=LONG_CAPACITY - -## Roadmap Targets - -- Milestone: `agent-roadmap/phase/operational-observability-provider-management/milestones/model-group-long-context-admission.md` -- Task ids: - - `provider-long-capacity`: provider별 total context budget, long capacity, long_in_flight 상태와 동시 점유 - - `provider-exclusion`: long slot이 찬 provider를 long 후보에서 제외하고 normal 후보는 유지 -- Completion mode: check-on-pass - -## 이 파일을 읽는 리뷰 에이전트에게 - -> **[REVIEW AGENT ONLY]** 종결 절차는 코드리뷰 에이전트 전용이다. PASS이면 완료 이벤트만 보고하고 roadmap을 직접 수정하지 않는다. - -## 구현 항목별 완료 여부 - -| 항목 | 완료 여부 | -|------|---------| -| [LONG_CAPACITY-1] Provider long capacity config와 validation | [ ] | -| [LONG_CAPACITY-2] Long slot reserve/release와 provider exclusion | [ ] | - -## 구현 체크리스트 - -- [ ] provider config에 `total_context_tokens`, `long_context_capacity`를 추가하고 `total_context_tokens >= context_window_tokens * long_context_capacity`를 config load/check에서 검증한다. -- [ ] config refresh classifier가 provider long capacity 관련 field를 live-apply로 분류하게 한다. -- [ ] long 요청 dispatch 시 일반 capacity와 long capacity를 함께 reserve/release하고 `long_in_flight` 상태를 유지한다. 검증: long dispatch 중 두 counter가 증가하고 완료/오류/취소 후 회복된다. -- [ ] long slot이 full인 provider는 long 후보에서 제외하되, normal 요청 후보에서는 일반 capacity만 기준으로 유지한다. 검증: long은 다른 provider 또는 queue 대기, normal은 남은 일반 slot으로 dispatch된다. -- [ ] config/service/configrefresh regression tests와 config check를 실행한다. -- [ ] CODE_REVIEW-*-G??.md의 구현 에이전트 소유 섹션을 실제 구현 내용과 검증 출력으로 채운다. 이 항목이 완료되기 전에는 구현이 완료된 것이 아니다. - -## 코드리뷰 전용 체크리스트 - -- [ ] 판정을 append한다. -- [ ] active review/plan을 `code_review_local_G08_0.log`, `plan_local_G08_0.log`로 아카이브한다. -- [ ] PASS이면 `complete.log` 작성 후 active task 디렉터리를 archive로 이동한다. -- [ ] PASS split 작업이면 빈 active parent 처리 여부를 확인한다. -- [ ] WARN/FAIL이면 다음 active plan/review 또는 `USER_REVIEW.md`를 작성한다. - -## 계획 대비 변경 사항 - -_구현 에이전트가 계획과 다르게 구현한 부분을 이유와 함께 기록한다._ - -## 주요 설계 결정 - -_구현 에이전트가 주요 설계 결정 사항을 기록한다._ - -## 사용자 리뷰 요청 - -_기본값은 `없음`이다. 구현 중 새 결정이 필요해 보여도 직접 질문하거나 선택지를 제시하거나 `request_user_input`을 호출하지 않는다. 이 섹션은 선택된 Milestone `구현 잠금 > 결정 필요` 항목이 실구현을 차단할 때만 채운다. 외부 환경/secret/서비스 준비, 검증 증거 공백, 반복 실패, 일반 범위 조정은 사용자 리뷰 요청이 아니며 `검증 결과`, `계획 대비 변경 사항`, 또는 code-review의 일반 follow-up plan으로 처리한다._ - -- 상태: 없음 -- 사유 유형: 없음 -- 연결 대상: 없음 -- 결정 필요: 없음 -- 차단 근거: 없음 -- 실행한 검증/명령: 없음 -- 자동 후속 불가 이유: 없음 -- 재개 조건: 없음 - -## 리뷰어를 위한 체크포인트 - -- `total_context_tokens >= context_window_tokens * long_context_capacity` 검증이 model provider mapping 기준으로 동작하는지 확인한다. -- long request가 일반 capacity와 long capacity를 함께 점유하고 terminal path에서 둘 다 회복하는지 확인한다. -- long slot full provider가 normal 후보에서 제외되지 않는지 확인한다. - -## 검증 결과 - -### LONG_CAPACITY-1 중간 검증 - -```bash -$ go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh -(output) -``` - -### LONG_CAPACITY-2 중간 검증 - -```bash -$ go test -count=1 ./apps/edge/internal/service -(output) -``` - -### 최종 검증 - -```bash -$ go test -count=1 ./packages/go/config ./apps/edge/internal/configrefresh ./apps/edge/internal/edgevalidate ./apps/edge/internal/service -(output) -$ go run ./apps/edge/cmd/edge config check --config configs/edge.yaml -(output) -``` - ---- - -> **[IMPLEMENTING AGENT — BEFORE SAVING] Have you filled in every implementation-owned section: completion table, implementation checklist, changes from plan, design decisions, user review request state, and verification output?** diff --git a/apps/edge/internal/configrefresh/classify.go b/apps/edge/internal/configrefresh/classify.go index 6004dc3..f796e6d 100644 --- a/apps/edge/internal/configrefresh/classify.go +++ b/apps/edge/internal/configrefresh/classify.go @@ -87,6 +87,8 @@ type providerKey struct { Models []string Health string Capacity int + TotalContextTokens int + LongContextCapacity int Priority int MaxQueue int QueueTimeoutMS int @@ -124,6 +126,8 @@ func buildProviderIndex(cfg *config.EdgeConfig) map[string]providerKey { Models: append([]string(nil), p.Models...), Health: p.Health, Capacity: p.Capacity, + TotalContextTokens: p.TotalContextTokens, + LongContextCapacity: p.LongContextCapacity, Priority: p.Priority, MaxQueue: p.MaxQueue, QueueTimeoutMS: p.QueueTimeoutMS, @@ -324,6 +328,22 @@ func Classify(current, candidate *config.EdgeConfig) Result { Next: fmt.Sprintf("%d", cp.Capacity), }) } + if cur.TotalContextTokens != cp.TotalContextTokens { + changes = append(changes, Change{ + Path: fmt.Sprintf("nodes[].providers[%q].total_context_tokens", provID), + Class: StatusApplied, + Previous: fmt.Sprintf("%d", cur.TotalContextTokens), + Next: fmt.Sprintf("%d", cp.TotalContextTokens), + }) + } + if cur.LongContextCapacity != cp.LongContextCapacity { + changes = append(changes, Change{ + Path: fmt.Sprintf("nodes[].providers[%q].long_context_capacity", provID), + Class: StatusApplied, + Previous: fmt.Sprintf("%d", cur.LongContextCapacity), + Next: fmt.Sprintf("%d", cp.LongContextCapacity), + }) + } if cur.Priority != cp.Priority { changes = append(changes, Change{ Path: fmt.Sprintf("nodes[].providers[%q].priority", provID), diff --git a/apps/edge/internal/configrefresh/classify_test.go b/apps/edge/internal/configrefresh/classify_test.go index 247191b..30c3ef5 100644 --- a/apps/edge/internal/configrefresh/classify_test.go +++ b/apps/edge/internal/configrefresh/classify_test.go @@ -1218,3 +1218,97 @@ nodes: t.Errorf("providers[].priority change path not found in: %+v", result.Changes) } } + +// TestClassifyProviderLongContextCapacityApplied verifies that provider +// total_context_tokens and long_context_capacity changes are classified as +// live-apply (StatusApplied), consistent with capacity/priority. +func TestClassifyProviderLongContextCapacityApplied(t *testing.T) { + dir := t.TempDir() + currentYAML := ` +server: + listen: "0.0.0.0:9090" +nodes: + - id: "node-1" + alias: "n1" + token: "tok-1" + adapters: + cli: + enabled: true + providers: + - id: "prov-a" + type: "ollama" + category: "local_inference" + adapter: "cli" + models: ["llama3.1"] + capacity: 2 + total_context_tokens: 262144 + long_context_capacity: 1 +` + currentPath := writeYAML(t, dir, "current.yaml", currentYAML) + // Candidate only changes the two long-context fields. + candidateYAML := ` +server: + listen: "0.0.0.0:9090" +nodes: + - id: "node-1" + alias: "n1" + token: "tok-1" + adapters: + cli: + enabled: true + providers: + - id: "prov-a" + type: "ollama" + category: "local_inference" + adapter: "cli" + models: ["llama3.1"] + capacity: 2 + total_context_tokens: 524288 + long_context_capacity: 2 +` + candidatePath := writeYAML(t, dir, "candidate.yaml", candidateYAML) + + current := buildNormalizedCurrent(t, currentPath) + result, _, err := configrefresh.Evaluate(context.Background(), current, configrefresh.Request{ + Mode: configrefresh.ModeDryRun, + ConfigPath: candidatePath, + RequestID: "test-long-capacity-applied", + }) + if err != nil { + t.Fatalf("Evaluate: %v", err) + } + if result.Status != configrefresh.StatusApplied { + t.Fatalf("expected status=applied for long-context-only change, got %q (summary: %s)", result.Status, result.Summary) + } + foundTotal := false + foundLong := false + for _, c := range result.Changes { + switch c.Path { + case `nodes[].providers["prov-a"].total_context_tokens`: + foundTotal = true + if c.Class != configrefresh.StatusApplied { + t.Errorf("total_context_tokens class: got %q, want applied", c.Class) + } + if c.Previous != "262144" || c.Next != "524288" { + t.Errorf("total_context_tokens change: prev=%q next=%q", c.Previous, c.Next) + } + case `nodes[].providers["prov-a"].long_context_capacity`: + foundLong = true + if c.Class != configrefresh.StatusApplied { + t.Errorf("long_context_capacity class: got %q, want applied", c.Class) + } + if c.Previous != "1" || c.Next != "2" { + t.Errorf("long_context_capacity change: prev=%q next=%q", c.Previous, c.Next) + } + } + if c.Class == configrefresh.StatusRestartRequired { + t.Errorf("unexpected restart_required change for long-context-only diff: %s", c.Path) + } + } + if !foundTotal { + t.Errorf("providers[].total_context_tokens change path not found in: %+v", result.Changes) + } + if !foundLong { + t.Errorf("providers[].long_context_capacity change path not found in: %+v", result.Changes) + } +} diff --git a/apps/edge/internal/service/model_queue.go b/apps/edge/internal/service/model_queue.go index cba08ad..ae2a366 100644 --- a/apps/edge/internal/service/model_queue.go +++ b/apps/edge/internal/service/model_queue.go @@ -38,6 +38,10 @@ type candidateNode struct { providerID string // non-empty for provider-pool candidates adapter string // non-empty for provider-pool candidates; dispatch adapter key servedTarget string // concrete served model name; used for target rewrite + // longContextCapacity is the provider's concurrent long-context slot limit. + // Zero means the provider declares no dedicated long-slot limit and long + // requests are not gated on it. + longContextCapacity int } // slotKey returns a unique slot key for inflight accounting. @@ -60,6 +64,7 @@ type inflightRec struct { groupKey string nodeID string providerID string // non-empty for provider-pool dispatches + long bool // true when a long-context slot was reserved for this run } type admitResult struct { @@ -71,6 +76,7 @@ type queueItem struct { candidates []candidateNode waitCh chan admitResult deadline time.Time + long bool // true when the queued request is a long-context request } type modelQueueGroup struct { @@ -78,6 +84,7 @@ type modelQueueGroup struct { policy groupPolicy queue []*queueItem inflight map[string]int // slotKey → current in-flight count + longInflight map[string]int // slotKey → current long-context in-flight count lastSelectedSlot string adapter string target string @@ -156,9 +163,10 @@ func (m *modelQueueManager) getOrCreateGroupLocked(key string, policy groupPolic policy.queueTimeout = defaultQueueTimeout } g = &modelQueueGroup{ - key: key, - policy: policy, - inflight: make(map[string]int), + key: key, + policy: policy, + inflight: make(map[string]int), + longInflight: make(map[string]int), } m.groups[key] = g } @@ -188,7 +196,7 @@ func (m *modelQueueManager) updateGroupPolicyLocked(g *modelQueueGroup, policy g // in_flight and priority are equal, all tied candidates remain in best[] so // that the rotation logic in the caller picks the next slot after the last // selected one. -func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode) *candidateNode { +func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode, long bool) *candidateNode { var best []candidateNode bestInflight := 0 bestPriority := 0 @@ -202,6 +210,13 @@ func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, cand if inflight >= c.capacity { continue } + // Long-context gate: a long request is excluded from a provider whose + // long slots are full. Providers with long_context_capacity == 0 declare + // no dedicated long-slot limit and are not gated here. Normal requests + // (long == false) never consult long slot state. + if long && c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity { + continue + } switch { case len(best) == 0 || inflight < bestInflight: best = []candidateNode{c} @@ -250,7 +265,10 @@ func candidateLess(a, b *candidateNode) bool { // the request until a slot opens. Blocks until a candidate is assigned, the // queue timeout expires, or ctx is cancelled. Returns the selected candidateNode // so callers can rewrite the dispatch target for provider-pool requests. -func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy) (*candidateNode, error) { +// The optional long flag marks the request as long-context so it is gated on +// provider long slots and reserves a long slot on dispatch. +func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, error) { + long := len(longFlag) > 0 && longFlag[0] m.mu.Lock() group := m.getOrCreateGroupLocked(groupKey, policy) @@ -265,10 +283,13 @@ func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target group.target = target } - candidate := m.findAvailableNodeLocked(group, candidates) + candidate := m.findAvailableNodeLocked(group, candidates, long) if candidate != nil { slot := candidate.slotKey() group.inflight[slot]++ + if long && candidate.longContextCapacity > 0 { + group.longInflight[slot]++ + } group.lastSelectedSlot = slot m.mu.Unlock() return candidate, nil @@ -283,6 +304,7 @@ func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target candidates: candidates, waitCh: make(chan admitResult, 1), deadline: time.Now().Add(group.policy.queueTimeout), + long: long, } group.queue = append(group.queue, item) timeout := group.policy.queueTimeout @@ -314,7 +336,8 @@ func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target } else { nodeID = slot } - m.releaseSlot(groupKey, nodeID, providerID) + longSlot := item.long && res.candidate.longContextCapacity > 0 + m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot) } default: } @@ -337,7 +360,8 @@ func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target } else { nodeID = slot } - m.releaseSlot(groupKey, nodeID, providerID) + longSlot := item.long && res.candidate.longContextCapacity > 0 + m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot) } default: } @@ -347,10 +371,13 @@ func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target // trackInflight records a dispatched run so release events can find it. // providerID is non-empty for provider-pool dispatches and is used to -// attribute in-flight counts to the correct provider snapshot. -func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID, providerID string) { +// attribute in-flight counts to the correct provider snapshot. The optional +// longSlot flag records that a long-context slot was reserved so the terminal +// release also frees the long slot. +func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID, providerID string, longSlot ...bool) { + long := len(longSlot) > 0 && longSlot[0] m.mu.Lock() - m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID, providerID: providerID} + m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID, providerID: providerID, long: long} m.mu.Unlock() } @@ -364,7 +391,7 @@ func (m *modelQueueManager) releaseRun(runID, reason string) { return } delete(m.inflightByRun, runID) - m.releaseSlotLocked(rec.groupKey, rec.nodeID, rec.providerID) + m.releaseSlotLocked(rec.groupKey, rec.nodeID, rec.providerID, rec.long) m.mu.Unlock() } @@ -396,18 +423,22 @@ func (m *modelQueueManager) releaseNode(nodeID, reason string) { // Clear all slot keys that start with this nodeID. // For provider-pool: "nodeID:providerID". For legacy: "nodeID". + // Long-context in-flight counters are keyed identically, so clear them + // alongside the normal counters. for slot := range group.inflight { if colonIdx := findLastColon(slot); colonIdx > 0 { candidateNodeID := slot[:colonIdx] // Only clear if the node part matches. if candidateNodeID == nodeID { delete(group.inflight, slot) + delete(group.longInflight, slot) continue } } else { // Legacy slot: exact nodeID match. if slot == nodeID { delete(group.inflight, slot) + delete(group.longInflight, slot) } } } @@ -427,16 +458,23 @@ func (m *modelQueueManager) releaseNode(nodeID, reason string) { // For provider-pool dispatches (nodeID:providerID slot), pass providerID. // For legacy dispatches, pass empty providerID so nodeID is used directly. func (m *modelQueueManager) releaseSlot(groupKey, nodeID string, providerID ...string) { - m.mu.Lock() var pid string if len(providerID) > 0 { pid = providerID[0] } - m.releaseSlotLocked(groupKey, nodeID, pid) + m.releaseSlotWithLong(groupKey, nodeID, pid, false) +} + +// releaseSlotWithLong releases a reserved slot, additionally freeing a long-context +// slot when longSlot is true. Used by admit-path failure/race handling where the +// caller knows whether a long slot was reserved. +func (m *modelQueueManager) releaseSlotWithLong(groupKey, nodeID, providerID string, longSlot bool) { + m.mu.Lock() + m.releaseSlotLocked(groupKey, nodeID, providerID, longSlot) m.mu.Unlock() } -func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerID string) { +func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerID string, longSlot bool) { group, ok := m.groups[groupKey] if !ok { return @@ -449,25 +487,10 @@ func (m *modelQueueManager) releaseSlotLocked(groupKey, nodeID string, providerI if group.inflight[slot] > 0 { group.inflight[slot]-- } - m.tryDispatchLocked(group) -} - -// releaseSlotByCandidate releases the slot identified by the candidate directly. -// IMPORTANT: This function MUST be called with m.mu held because it calls -// releaseSlotLocked which modifies shared state (group.inflight, tryDispatchLocked). -// Use releaseSlot instead when releasing from an unlocked context. -func (m *modelQueueManager) releaseSlotByCandidate(groupKey string, candidate *candidateNode) { - slot := candidate.slotKey() - // For provider-pool candidates, slot is "nodeID:providerID". - // We need to split it back for the locked release. - var nodeID, providerID string - if colonIdx := findLastColon(slot); colonIdx > 0 { - nodeID = slot[:colonIdx] - providerID = slot[colonIdx+1:] - } else { - nodeID = slot + if longSlot && group.longInflight[slot] > 0 { + group.longInflight[slot]-- } - m.releaseSlotLocked(groupKey, nodeID, providerID) + m.tryDispatchLocked(group) } // findLastColon returns the index of the last ':' in s, or -1 if not found. @@ -495,14 +518,18 @@ func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) { continue } - candidate := m.findAvailableNodeLocked(group, head.candidates) + candidate := m.findAvailableNodeLocked(group, head.candidates, head.long) if candidate == nil { break } group.queue = group.queue[1:] slot := candidate.slotKey() + reserveLong := head.long && candidate.longContextCapacity > 0 group.inflight[slot]++ + if reserveLong { + group.longInflight[slot]++ + } select { case head.waitCh <- admitResult{candidate: candidate}: @@ -511,6 +538,9 @@ func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) { default: // Caller already timed out or cancelled; release the reserved slot and try next. group.inflight[slot]-- + if reserveLong { + group.longInflight[slot]-- + } } } } @@ -707,6 +737,26 @@ func (m *modelQueueManager) getStatsForProviderLocked(nodeID, providerID string) return } +// getLongStatsForProviderLocked returns the long-context in-flight count for a +// provider-pool provider identified by (nodeID, providerID), derived from tracked +// runs. Must be called with m.mu held. +func (m *modelQueueManager) getLongStatsForProviderLocked(nodeID, providerID string) (longInFlight int) { + for _, rec := range m.inflightByRun { + if rec.nodeID == nodeID && rec.providerID == providerID && rec.long { + longInFlight++ + } + } + return +} + +// longInFlightForProvider returns the long-context in-flight count for a provider, +// acquiring the manager lock. Exposed for status/snapshot reporting and tests. +func (m *modelQueueManager) longInFlightForProvider(nodeID, providerID string) int { + m.mu.Lock() + defer m.mu.Unlock() + return m.getLongStatsForProviderLocked(nodeID, providerID) +} + func (m *modelQueueManager) getStatsForAdapterLocked(nodeID string, rec *edgenode.NodeRecord, adapterName string) (inFlight, queued int) { for _, group := range m.groups { canonical, ok := resolveSnapshotAdapterName(rec, group.adapter, group.target) diff --git a/apps/edge/internal/service/model_queue_test.go b/apps/edge/internal/service/model_queue_test.go index 49632c8..afbe653 100644 --- a/apps/edge/internal/service/model_queue_test.go +++ b/apps/edge/internal/service/model_queue_test.go @@ -77,7 +77,7 @@ func TestModelQueueFIFOOrdering(t *testing.T) { // Release one slot — item1 (head) must be dispatched first. m.mu.Lock() - m.releaseSlotLocked("g-fifo", "node-q1", "") + m.releaseSlotLocked("g-fifo", "node-q1", "", false) m.mu.Unlock() select { @@ -101,7 +101,7 @@ func TestModelQueueFIFOOrdering(t *testing.T) { // Release the slot item1 holds so item2 gets dispatched. m.mu.Lock() - m.releaseSlotLocked("g-fifo", "node-q1", "") + m.releaseSlotLocked("g-fifo", "node-q1", "", false) m.mu.Unlock() select { @@ -592,7 +592,7 @@ func TestModelQueueProviderLevelingPriorityAdmissionOrder(t *testing.T) { m.mu.Lock() g := m.getOrCreateGroupLocked("qwen3.6:35b", groupPolicy{}) - if candidate := m.findAvailableNodeLocked(g, cands); candidate != nil { + if candidate := m.findAvailableNodeLocked(g, cands, false); candidate != nil { t.Fatalf("expected no available candidate after filling capacity, got %s", candidate.providerID) } m.mu.Unlock() @@ -1708,3 +1708,136 @@ func TestGetSnapshotForNodeLegacyAdapterFallback(t *testing.T) { t.Error("expected ollama-local adapter snapshot in legacy fallback") } } + +// longSlotCount reads the group's long in-flight counter for a provider slot. +func longSlotCount(m *modelQueueManager, groupKey, nodeID, providerID string) int { + m.mu.Lock() + defer m.mu.Unlock() + g, ok := m.groups[groupKey] + if !ok { + return -1 + } + return g.longInflight[nodeID+":"+providerID] +} + +func inflightSlotCount(m *modelQueueManager, groupKey, nodeID, providerID string) int { + m.mu.Lock() + defer m.mu.Unlock() + g, ok := m.groups[groupKey] + if !ok { + return -1 + } + return g.inflight[nodeID+":"+providerID] +} + +// TestModelQueueLongDispatchReservesBothCounters verifies S04: a long request +// occupies both the normal in-flight counter and the long in-flight counter, and +// every terminal reason (complete/error/cancelled) recovers both. +func TestModelQueueLongDispatchReservesBothCounters(t *testing.T) { + for _, reason := range []string{"complete", "error", "cancelled"} { + t.Run(reason, func(t *testing.T) { + entry := &edgenode.NodeEntry{NodeID: "node-long"} + cands := []candidateNode{{entry: entry, capacity: 3, providerID: "prov-long", longContextCapacity: 2}} + m := newModelQueueManager(nil) + + sel, err := m.admit(context.Background(), "g-long", "", "", cands, groupPolicy{}, true) + if err != nil || sel == nil { + t.Fatalf("long admit: %v", err) + } + if got := inflightSlotCount(m, "g-long", "node-long", "prov-long"); got != 1 { + t.Errorf("normal in-flight after admit = %d, want 1", got) + } + if got := longSlotCount(m, "g-long", "node-long", "prov-long"); got != 1 { + t.Errorf("long in-flight after admit = %d, want 1", got) + } + + longReserved := sel.longContextCapacity > 0 + m.trackInflight("g-long", "run-long-1", sel.entry.NodeID, sel.providerID, longReserved) + if got := m.longInFlightForProvider("node-long", "prov-long"); got != 1 { + t.Errorf("longInFlightForProvider after track = %d, want 1", got) + } + + m.releaseRun("run-long-1", reason) + if got := inflightSlotCount(m, "g-long", "node-long", "prov-long"); got != 0 { + t.Errorf("normal in-flight after %s = %d, want 0", reason, got) + } + if got := longSlotCount(m, "g-long", "node-long", "prov-long"); got != 0 { + t.Errorf("long in-flight after %s = %d, want 0", reason, got) + } + if got := m.longInFlightForProvider("node-long", "prov-long"); got != 0 { + t.Errorf("longInFlightForProvider after %s = %d, want 0", reason, got) + } + }) + } +} + +// TestModelQueueNormalRequestDoesNotReserveLongSlot verifies a normal request +// occupies only the normal in-flight counter, never the long counter, even on a +// provider that declares long capacity. +func TestModelQueueNormalRequestDoesNotReserveLongSlot(t *testing.T) { + entry := &edgenode.NodeEntry{NodeID: "node-n"} + cands := []candidateNode{{entry: entry, capacity: 3, providerID: "prov-n", longContextCapacity: 2}} + m := newModelQueueManager(nil) + + sel, err := m.admit(context.Background(), "g-n", "", "", cands, groupPolicy{}) // long omitted → normal + if err != nil || sel == nil { + t.Fatalf("normal admit: %v", err) + } + if got := inflightSlotCount(m, "g-n", "node-n", "prov-n"); got != 1 { + t.Errorf("normal in-flight = %d, want 1", got) + } + if got := longSlotCount(m, "g-n", "node-n", "prov-n"); got != 0 { + t.Errorf("long in-flight for normal request = %d, want 0", got) + } +} + +// TestModelQueueLongSlotFullExcludesLongKeepsNormal verifies S05 and D06: a +// provider whose long slots are full is excluded from long candidates (the long +// request moves to another provider or queues), but the same provider remains a +// valid candidate for normal requests while its normal capacity has room. +func TestModelQueueLongSlotFullExcludesLongKeepsNormal(t *testing.T) { + entryA := &edgenode.NodeEntry{NodeID: "node-a"} + entryB := &edgenode.NodeEntry{NodeID: "node-b"} + cands := []candidateNode{ + {entry: entryA, capacity: 3, priority: 0, providerID: "prov-a", longContextCapacity: 1}, + {entry: entryB, capacity: 3, priority: 1, providerID: "prov-b", longContextCapacity: 1}, + } + m := newModelQueueManager(nil) + + // Long #1 → prov-a (priority 0 preferred), filling prov-a's single long slot. + sel1, err := m.admit(context.Background(), "g-x", "", "", cands, groupPolicy{}, true) + if err != nil || sel1 == nil { + t.Fatalf("long1 admit: %v", err) + } + if sel1.providerID != "prov-a" { + t.Fatalf("long1 provider = %q, want prov-a", sel1.providerID) + } + + // Long #2 must skip the long-full prov-a and land on prov-b. + sel2, err := m.admit(context.Background(), "g-x", "", "", cands, groupPolicy{}, true) + if err != nil || sel2 == nil { + t.Fatalf("long2 admit: %v", err) + } + if sel2.providerID != "prov-b" { + t.Fatalf("long2 provider = %q, want prov-b (prov-a long slot full)", sel2.providerID) + } + + // A normal request restricted to the long-full prov-a is still admitted, + // because normal admission ignores long slot state (D06). + normalCands := []candidateNode{cands[0]} + selN, err := m.admit(context.Background(), "g-x", "", "", normalCands, groupPolicy{}) + if err != nil || selN == nil { + t.Fatalf("normal admit on long-full provider: %v", err) + } + if selN.providerID != "prov-a" { + t.Fatalf("normal provider = %q, want prov-a", selN.providerID) + } + + // A further long request restricted to the long-full prov-a cannot be + // admitted immediately; it queues and times out. + longOnlyA := []candidateNode{cands[0]} + _, err = m.admit(context.Background(), "g-x", "", "", longOnlyA, groupPolicy{maxQueue: 1, queueTimeout: 20 * time.Millisecond}, true) + if !errors.Is(err, errQueueTimeout) { + t.Fatalf("long request on long-full provider: expected queue timeout, got %v", err) + } +} diff --git a/apps/edge/internal/service/run_dispatch.go b/apps/edge/internal/service/run_dispatch.go index c302ee6..edea34a 100644 --- a/apps/edge/internal/service/run_dispatch.go +++ b/apps/edge/internal/service/run_dispatch.go @@ -15,6 +15,10 @@ import ( iop "iop/proto/gen/iop" ) +// contextClassLong is the ContextClass value that marks a request for +// long-context admission (long slot gating and reservation). +const contextClassLong = "long" + type SubmitRunRequest struct { NodeRef string RunID string @@ -132,11 +136,16 @@ func (s *Service) submitRunQueued(ctx context.Context, req SubmitRunRequest) (Ru return nil, err } - selected, err := s.queue.admit(ctx, req.ModelGroupKey, req.Adapter, req.Target, candidates, policy) + long := req.ContextClass == contextClassLong + selected, err := s.queue.admit(ctx, req.ModelGroupKey, req.Adapter, req.Target, candidates, policy, long) if err != nil { return nil, err } + // A long slot is reserved only for long requests on providers that declare a + // long-context limit; the release/track path must match that reservation. + longReserved := long && selected.longContextCapacity > 0 + // Rewrite adapter and target for provider-pool dispatch: the winning candidate // carries the concrete adapter and served model name determined at selection time. if selected.adapter != "" { @@ -148,13 +157,13 @@ func (s *Service) submitRunQueued(ctx context.Context, req SubmitRunRequest) (Ru runReq, runID, err := BuildRunRequest(req) if err != nil { - s.queue.releaseSlot(req.ModelGroupKey, selected.entry.NodeID, selected.providerID) + s.queue.releaseSlotWithLong(req.ModelGroupKey, selected.entry.NodeID, selected.providerID, longReserved) return nil, err } // Track inflight before send so the event watcher can release the slot // even if a terminal event arrives before the Send call completes. - s.queue.trackInflight(req.ModelGroupKey, runID, selected.entry.NodeID, selected.providerID) + s.queue.trackInflight(req.ModelGroupKey, runID, selected.entry.NodeID, selected.providerID, longReserved) var runEvents <-chan *iop.RunEvent var unregisterRun func() @@ -615,12 +624,13 @@ func (s *Service) resolveProviderPoolCandidates(req SubmitRunRequest, store *edg } cap := prov.Capacity candidates = append(candidates, candidateNode{ - entry: entry, - capacity: cap, - priority: prov.Priority, - providerID: prov.ID, - adapter: adapterKey, - servedTarget: servedModel, + entry: entry, + capacity: cap, + priority: prov.Priority, + providerID: prov.ID, + adapter: adapterKey, + servedTarget: servedModel, + longContextCapacity: prov.LongContextCapacity, }) if !policySet && (prov.MaxQueue > 0 || prov.QueueTimeoutMS > 0) { policy = groupPolicy{ diff --git a/apps/edge/internal/service/service_internal_test.go b/apps/edge/internal/service/service_internal_test.go index bdf7dc6..18713f0 100644 --- a/apps/edge/internal/service/service_internal_test.go +++ b/apps/edge/internal/service/service_internal_test.go @@ -70,6 +70,69 @@ func TestRefreshProviderCapacityAffectsNextDispatch(t *testing.T) { } } +// TestRefreshProviderLongCapacityAffectsCandidate verifies that a provider +// long_context_capacity change applied via SetRuntimeConfig is reflected in the +// candidate longContextCapacity used for the next long-context admission. +func TestRefreshProviderLongCapacityAffectsCandidate(t *testing.T) { + reg := edgenode.NewRegistry() + reg.Register(&edgenode.NodeEntry{NodeID: "node-cap"}) + svc := New(reg, nil) + + newStore := func(longCapacity int) *edgenode.NodeStore { + store := edgenode.NewNodeStore() + store.Add(&edgenode.NodeRecord{ + ID: "node-cap", + Adapters: config.AdaptersConf{ + VllmInstances: []config.VllmInstanceConf{ + {Name: "vllm-gpu", Enabled: true, Endpoint: "http://127.0.0.1:8000/v1"}, + }, + }, + Providers: []config.NodeProviderConf{ + { + ID: "prov-a", + Type: "vllm", + Category: config.CategoryAPI, + Adapter: "vllm-gpu", + Models: []string{"served-a"}, + Health: "available", + Capacity: 4, + TotalContextTokens: 524288, + LongContextCapacity: longCapacity, + }, + }, + }) + return store + } + catalog := []config.ModelCatalogEntry{{ + ID: "qwen3.6:35b", + ContextWindowTokens: 262144, + Providers: map[string]string{"prov-a": "served-a"}, + }} + + req := SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true, ContextClass: "long"} + + svc.SetRuntimeConfig(newStore(1), catalog) + store, cat := svc.runtimeConfigSnapshot() + cands, _, err := svc.resolveProviderPoolCandidates(req, store, cat) + if err != nil { + t.Fatalf("resolve before refresh: %v", err) + } + if len(cands) != 1 || cands[0].longContextCapacity != 1 { + t.Fatalf("before refresh: got candidates=%+v, want one with longContextCapacity=1", cands) + } + + // Refresh raises long capacity to 2; the next candidate build must observe it. + svc.SetRuntimeConfig(newStore(2), catalog) + store, cat = svc.runtimeConfigSnapshot() + cands, _, err = svc.resolveProviderPoolCandidates(req, store, cat) + if err != nil { + t.Fatalf("resolve after refresh: %v", err) + } + if len(cands) != 1 || cands[0].longContextCapacity != 2 { + t.Fatalf("after refresh: got candidates=%+v, want one with longContextCapacity=2", cands) + } +} + func TestRuntimeConfigSnapshotConcurrentReplace(t *testing.T) { reg := edgenode.NewRegistry() reg.Register(&edgenode.NodeEntry{NodeID: "node-refresh"}) diff --git a/configs/edge.yaml b/configs/edge.yaml index 57c0253..3bfc267 100644 --- a/configs/edge.yaml +++ b/configs/edge.yaml @@ -99,6 +99,12 @@ openai: # nodes[].providers[].models — served model names this provider can serve. # nodes[].providers[].health — observed health state string. # nodes[].providers[].capacity — provider-pool max concurrent execution slots; 0 is not dispatchable. +# nodes[].providers[].total_context_tokens — provider runtime total KV/context budget used for +# long-context admission. Must satisfy total_context_tokens >= context_window_tokens * +# long_context_capacity for every model group referencing this provider. Live-apply on refresh. +# nodes[].providers[].long_context_capacity — concurrent long-context slots (context_window-sized +# requests). Long requests occupy a normal slot and a long slot; 0 means no dedicated long limit. +# Live-apply on refresh. # nodes[].providers[].max_queue — max queue depth (per-provider). # nodes[].providers[].queue_timeout_ms — queue timeout in milliseconds. # nodes[].providers[].lifecycle_capabilities — coarse lifecycle capabilities list. @@ -184,6 +190,9 @@ nodes: - "mlx-community/Qwen3.6-35B-A3B-4bit" health: "healthy" capacity: 2 + # Long-context admission budget: total_context_tokens >= 262144 * 2. + total_context_tokens: 524288 + long_context_capacity: 2 # Legacy adapters configuration (compat override example): # adapters: # cli: diff --git a/packages/go/config/config.go b/packages/go/config/config.go index 509e395..36535d0 100644 --- a/packages/go/config/config.go +++ b/packages/go/config/config.go @@ -150,6 +150,17 @@ type NodeProviderConf struct { // Provider-pool dispatch treats 0 as not dispatchable rather than falling // back to node runtime concurrency metadata. Capacity int `mapstructure:"capacity" yaml:"capacity,omitempty"` + // TotalContextTokens is the provider runtime's total KV/context budget used + // as the long-context admission budget. Runtime-specific options (e.g. + // ctx_size) map onto this value. Must satisfy + // total_context_tokens >= context_window_tokens * long_context_capacity for + // every model group that references this provider. + TotalContextTokens int `mapstructure:"total_context_tokens" yaml:"total_context_tokens,omitempty"` + // LongContextCapacity is the number of concurrent long-context slots this + // provider allows (context_window_tokens-sized requests). Long requests + // occupy both a normal capacity slot and a long slot. Zero means the + // provider declares no dedicated long-slot limit. + LongContextCapacity int `mapstructure:"long_context_capacity" yaml:"long_context_capacity,omitempty"` // MaxQueue is the maximum queue depth. MaxQueue int `mapstructure:"max_queue" yaml:"max_queue,omitempty"` // QueueTimeoutMS is the queue timeout in milliseconds. @@ -214,6 +225,15 @@ func (p NodeProviderConf) Validate() error { if p.Priority < 0 { return fmt.Errorf("nodes[].providers[%q].priority must be non-negative", id) } + if p.TotalContextTokens < 0 { + return fmt.Errorf("nodes[].providers[%q].total_context_tokens must be non-negative", id) + } + if p.LongContextCapacity < 0 { + return fmt.Errorf("nodes[].providers[%q].long_context_capacity must be non-negative", id) + } + if p.LongContextCapacity > 0 && p.TotalContextTokens <= 0 { + return fmt.Errorf("nodes[].providers[%q].total_context_tokens must be positive when long_context_capacity > 0", id) + } return nil } @@ -588,6 +608,7 @@ func LoadEdge(cfgFile string) (*EdgeConfig, error) { // Collect all provider IDs from nodes[].providers[] for cross-referencing // and validate uniqueness within each node. providerIDs := make(map[string]struct{}) + providerByID := make(map[string]NodeProviderConf) for i := range cfg.Nodes { kind, err := NormalizeAgentKind(cfg.Nodes[i].AgentKind) if err != nil { @@ -617,6 +638,7 @@ func LoadEdge(cfgFile string) (*EdgeConfig, error) { return nil, fmt.Errorf("nodes[%d].providers[%d]: duplicate provider id %q across nodes (global uniqueness violation)", i, j, p.ID) } providerIDs[p.ID] = struct{}{} + providerByID[p.ID] = p } } @@ -643,11 +665,41 @@ func LoadEdge(cfgFile string) (*EdgeConfig, error) { if err := m.Validate(providerIDs, serveModels); err != nil { return nil, fmt.Errorf("models[%d]: %w", i, err) } + if err := validateProviderLongContextBudget(m, providerByID); err != nil { + return nil, fmt.Errorf("models[%d]: %w", i, err) + } } return &cfg, nil } +// validateProviderLongContextBudget enforces the long-context admission budget: +// every provider referenced by a model group with a positive context window must +// have total_context_tokens >= context_window_tokens * long_context_capacity. +// Providers with long_context_capacity == 0 declare no long-slot budget and are +// therefore unconstrained by this check. +func validateProviderLongContextBudget(m ModelCatalogEntry, providerByID map[string]NodeProviderConf) error { + if m.ContextWindowTokens <= 0 { + return nil + } + for pid := range m.Providers { + p := strings.TrimSpace(pid) + prov, ok := providerByID[p] + if !ok { + continue + } + if prov.LongContextCapacity <= 0 { + continue + } + required := m.ContextWindowTokens * prov.LongContextCapacity + if prov.TotalContextTokens < required { + return fmt.Errorf("providers[%q].total_context_tokens (%d) must be >= context_window_tokens (%d) * long_context_capacity (%d) = %d", + p, prov.TotalContextTokens, m.ContextWindowTokens, prov.LongContextCapacity, required) + } + } + return nil +} + // NormalizeAdapters promotes legacy single-instance Ollama/Vllm fields into the // typed instance slices and validates that all instance names are unique. func NormalizeAdapters(a *AdaptersConf) error { diff --git a/packages/go/config/config_test.go b/packages/go/config/config_test.go index f9006d8..827f987 100644 --- a/packages/go/config/config_test.go +++ b/packages/go/config/config_test.go @@ -3138,3 +3138,131 @@ nodes: t.Fatalf("expected error mentioning priority non-negative, got %v", err) } } + +// TestLoadEdgeProviderLongContextCapacityValidation covers the provider-level +// long-context fields: a valid config loads the values, and the field-level +// invariants (non-negative, and total>0 when long_context_capacity>0) are enforced. +func TestLoadEdgeProviderLongContextCapacityValidation(t *testing.T) { + validYAML := ` +server: + listen: "0.0.0.0:9090" +models: + - id: "qwen3.6:35b" + context_window_tokens: 262144 + providers: + vllm-gpu: "nvidia/Qwen3.6-35B" +nodes: + - id: "node-gpu-01" + alias: "gpu-node" + providers: + - id: "vllm-gpu" + type: "vllm" + category: "api" + models: + - "nvidia/Qwen3.6-35B" + capacity: 4 + total_context_tokens: 524288 + long_context_capacity: 2 +` + t.Run("valid", func(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "edge.yaml") + if err := os.WriteFile(f, []byte(validYAML), 0o600); err != nil { + t.Fatalf("write yaml: %v", err) + } + cfg, err := config.LoadEdge(f) + if err != nil { + t.Fatalf("load: %v", err) + } + prov := cfg.Nodes[0].Providers[0] + if prov.TotalContextTokens != 524288 { + t.Errorf("total_context_tokens = %d, want 524288", prov.TotalContextTokens) + } + if prov.LongContextCapacity != 2 { + t.Errorf("long_context_capacity = %d, want 2", prov.LongContextCapacity) + } + }) + + cases := []struct { + name string + total string + long string + wantMsg string + }{ + {"negative total", "total_context_tokens: -1", "long_context_capacity: 0", "total_context_tokens must be non-negative"}, + {"negative long", "total_context_tokens: 524288", "long_context_capacity: -1", "long_context_capacity must be non-negative"}, + {"long without total", "total_context_tokens: 0", "long_context_capacity: 2", "total_context_tokens must be positive when long_context_capacity"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "edge.yaml") + yaml := fmt.Sprintf(` +server: + listen: "0.0.0.0:9090" +nodes: + - id: "node-gpu-01" + alias: "gpu-node" + providers: + - id: "vllm-gpu" + type: "vllm" + category: "api" + models: + - "nvidia/Qwen3.6-35B" + capacity: 4 + %s + %s +`, tc.total, tc.long) + if err := os.WriteFile(f, []byte(yaml), 0o600); err != nil { + t.Fatalf("write yaml: %v", err) + } + _, err := config.LoadEdge(f) + if err == nil { + t.Fatalf("expected error for %s", tc.name) + } + if !strings.Contains(err.Error(), tc.wantMsg) { + t.Fatalf("expected error mentioning %q, got %v", tc.wantMsg, err) + } + }) + } +} + +// TestLoadEdgeRejectsProviderLongContextBudgetBelowModelWindow verifies the +// cross-entry budget rule: a provider referenced by a model group must have +// total_context_tokens >= context_window_tokens * long_context_capacity. +func TestLoadEdgeRejectsProviderLongContextBudgetBelowModelWindow(t *testing.T) { + dir := t.TempDir() + f := filepath.Join(dir, "edge.yaml") + // 262144 * 2 = 524288 required, but total is only 262144. + yaml := ` +server: + listen: "0.0.0.0:9090" +models: + - id: "qwen3.6:35b" + context_window_tokens: 262144 + providers: + vllm-gpu: "nvidia/Qwen3.6-35B" +nodes: + - id: "node-gpu-01" + alias: "gpu-node" + providers: + - id: "vllm-gpu" + type: "vllm" + category: "api" + models: + - "nvidia/Qwen3.6-35B" + capacity: 4 + total_context_tokens: 262144 + long_context_capacity: 2 +` + if err := os.WriteFile(f, []byte(yaml), 0o600); err != nil { + t.Fatalf("write yaml: %v", err) + } + _, err := config.LoadEdge(f) + if err == nil { + t.Fatal("expected error for insufficient total context budget") + } + if !strings.Contains(err.Error(), "total_context_tokens") || !strings.Contains(err.Error(), "long_context_capacity") { + t.Fatalf("expected budget error mentioning total_context_tokens and long_context_capacity, got %v", err) + } +}