kvcache-ai/AgentENV

Fork 0

mirror of https://github.com/kvcache-ai/AgentENV.git synced 2026-11-05 01:17:36 +00:00

[PR #51] [CLOSED] feat(scheduler): add snapshot artifact locality preference #141

New issue

Closed

opened 2026-10-29 14:20:07 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 14:20:07 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/AgentENV/pull/51
Author: @shudorcl
Created: 7/28/2026
Status: ❌ Closed

Base: main ← Head: issue-15-snapshot-locality

📝 Commits (10+)

a762ada feat(gateway): add snapshot locality requirements
1c1af6c refactor(scheduler): expose node locality context
086e71d feat(scheduler): score artifact locality
af1a96e feat(scheduler): apply locality before node selection
1b82399 chore(api): format generated scheduler protobuf
211fbf4 docs(services): document snapshot locality scheduling
ecd391b Update services/scheduler/internal/locality.go
89f1dad Update services/scheduler/internal/service.go
460c590 Update services/scheduler/internal/locality.go
186df1d fix(services): harden snapshot locality scheduling

📊 Changes

18 files changed (+1123 additions, -229 deletions)

View changed files

📝 services/README.md (+29 -0)
📝 services/api/proto/scheduler.pb.go (+277 -207)
📝 services/api/proto/scheduler.proto (+9 -0)
📝 services/gateway/internal/schedule_hint.go (+18 -2)
📝 services/gateway/internal/schedule_hint_test.go (+49 -0)
📝 services/gateway/internal/server.go (+2 -1)
➕ services/gateway/internal/template_resolution.go (+47 -0)
➕ services/gateway/internal/template_resolution_test.go (+95 -0)
📝 services/go.mod (+1 -1)
➕ services/scheduler/internal/locality.go (+116 -0)
➕ services/scheduler/internal/locality_test.go (+184 -0)
📝 services/scheduler/internal/node_registry.go (+12 -11)
📝 services/scheduler/internal/node_registry_test.go (+44 -0)
📝 services/scheduler/internal/service.go (+31 -3)
📝 services/scheduler/internal/service_test.go (+133 -0)
📝 services/scheduler/internal/store.go (+43 -3)
📝 services/scheduler/internal/store_test.go (+22 -0)
📝 services/scheduler/internal/types.go (+11 -1)

📄 Description

What

This PR implements the first P2P-backed snapshot locality preference slice for #15.

For POST /sandboxes, when templateID is a canonical snapshot UUID in the standard hyphenated form, the Gateway derives two advisory locality requirements using the existing fixed snapshot artifact identities:

snapshot/v1/artifacts/{snapshot_id}/vm_state.bin
snapshot/v1/artifacts/{snapshot_id}/firecracker-manifest.json

After resource-limit filtering, the Scheduler prefers eligible nodes that advertise the highest coverage of those artifacts in the existing P2P artifact index. The configured scheduling strategy performs the final selection among tied nodes.

Template aliases remain supported, but intentionally do not produce locality hints because the Gateway has no authoritative pre-scheduling alias resolver. Alias requests therefore use the existing scheduling behavior.

Why

The scheduler currently ignores whether an eligible node already provides the snapshot artifacts needed to start a sandbox. It can therefore select a cold node even when another feasible node advertises the required artifacts through the existing P2P index.

This change introduces a bounded, advisory locality preference without duplicating load-aware selection or changing runtime fallback behavior. It establishes the request and inventory path needed for later locality work while keeping the scheduler's artifact store limited to artifact-key-to-node mappings.

Refs #15

This PR is one focused part of #15 and does not close the full issue.

Scope and non-goals

In scope:

derive fixed snapshot artifact requirements from canonical snapshot IDs;
carry advisory locality requirements through the internal scheduling protocol;
prefer the eligible nodes with the highest positive artifact coverage;
preserve existing selection behavior when locality data is absent or stale;
bound request inspection, requirement count, and artifact lookup work;
remove a node's artifact mappings when it is unregistered;
document and test the new scheduling flow.

Non-goals:

resolving template aliases before scheduling;
resolving OCI image references into OverlayBD layer requirements;
indexing non-P2P local caches;
adding locality requirements for POST /sandboxes-cold;
introducing a monolithic cache_aware strategy;
implementing or duplicating load-aware node selection;
changing hard resource filtering or runtime artifact fallback behavior.

Design and behavior changes

The scheduling flow becomes:

ready nodes
    -> hard resource-limit filtering
    -> snapshot artifact locality preference
    -> configured strategy selection
    -> selected node

The Gateway inspects only exact POST /sandboxes creation requests. It buffers at most 64 KiB while extracting templateID and restores the original body for upstream forwarding. Oversized, malformed, empty, alias-based, and non-canonical template IDs produce no locality requirements and continue through normal scheduling.

For a canonical snapshot UUID, the Gateway normalizes the ID and emits the two fixed artifact keys listed above. It does not query runtime nodes to resolve aliases. Such fan-out would be non-authoritative, could observe inconsistent alias mappings, would amplify unauthenticated request work, and could expand credential exposure.

The Scheduler validates the request-controlled locality requirement count before node discovery. After resource filtering, PreferLocalNodes:

trims and deduplicates requirement keys;
isolates lookups by cluster ID and P2P backend;
scores only already-eligible candidates;
retains all nodes tied for the highest positive coverage;
preserves the complete eligible set when no useful locality data exists.

Artifact lookup intersects provider records with the eligible node-ID set inside the store, so scheduling work scales with current candidates rather than all historical providers. Locality does not create reservations or other request-lifetime state.

Locality remains advisory. Stale or missing records do not fail sandbox creation, and the selected runtime continues through its normal local cache, P2P, snapshot repository, or registry fallback path.

Compatibility and operations

Public API or generated protocol: Adds optional ScheduleRequestHint.locality_requirements, LocalityRequirement.key, and NewSandboxHint.template_id fields to the internal scheduler protobuf. The additions are wire-compatible. Older senders omit them, and receivers without locality support ignore the unknown fields.
Configuration or defaults: No configuration fields or defaults change. Existing round_robin and random strategy behavior is unchanged when locality requirements or provider records are unavailable.
Snapshot manifest, artifact layout, or storage format: No format or layout changes. The implementation reuses the existing snapshot/v1/artifacts/{snapshot_id}/... identities.
Upgrade and rollback: No persistent migration is required. Gateway and Scheduler versions can be rolled back without changing snapshot data or the artifact index; rollback removes the placement preference and restores the previous selection behavior.
Host requirements, permissions, ports, or dependencies: No new host permissions, services, ports, or runtime packages. github.com/google/uuid is promoted from an indirect to a direct Go dependency at the same version.

Validation

make fmt
make clippy
make test-unit
Relevant Rust integration tests
make -C services test (required when services/ changes)
Generated clients/server regenerated with the documented make target
Documentation updated
Benchmarks or performance comparison completed

Commands and results:

make -C services proto PROTOC=<protoc 3.21.12>
gofmt -w services/api/proto/scheduler.pb.go services/api/proto/scheduler_grpc.pb.go
# PASS: protoc-gen-go 1.36.11 and protoc-gen-go-grpc 1.6.2; no tracked diff

make -C services build
make -C services test
make -C services fmt-check
make -C services vet
# PASS with Go 1.25.0

cd services
go test -race ./gateway/internal ./scheduler/internal
# PASS with Go 1.25.0

git diff --check
# PASS

The Services commands above were run locally. fmt-check was verified from an LF-normalized export of the same commit because this WSL checkout uses core.autocrlf=true.

The following remote Services CI and integration jobs also passed for commit b2f9346577b001bba28fb038d46ae0beb233c956:

Services CI: https://github.com/shudorcl/AgentENV/actions/runs/30394461927
Integration Tests: https://github.com/shudorcl/AgentENV/actions/runs/30394461784

Skipped checks and reasons:

Root Rust formatting, Clippy, and unit-test targets were not run because this PR changes only the Go control-plane module and its generated Go protobuf. The remote AgentENV integration, single-node E2E, and docker-compose E2E jobs passed.
No benchmark was run. The scheduling hot path is bounded to at most 32 locality requirements, and provider lookup is intersected with the current eligible candidate set. The PR adds focused boundary and stale-record tests, but does not claim a measured performance improvement.

Risks and reviewer notes

Locality metadata is eventually consistent and can be stale. The preference is deliberately advisory, and no positive match preserves the complete eligible candidate set.
Requests with aliases receive no locality optimization. Supporting aliases safely requires an authoritative control-plane resolver and is left for a follow-up.
Hint extraction occurs before upstream authentication. Request-body inspection is therefore bounded to 64 KiB and does not fan out to runtime nodes.
The locality requirement list is request-controlled. The Scheduler rejects more than 32 requirements before node snapshotting or enrichment.
Popular artifacts may have many historical providers. LookupEligible bounds work to the current eligible candidate IDs, and unregistering a node removes its artifact mappings.
Reviewers should start with services/gateway/internal/schedule_hint.go, services/gateway/internal/template_resolution.go, services/scheduler/internal/locality.go, and services/scheduler/internal/service.go.

Checklist

The PR contains one coherent change and no unrelated formatting or refactoring.
New behavior is covered by tests, or I explained why testing is impractical.
Logs and examples contain no credentials, tokens, or private registry information.
I did not manually edit generated code without updating its source and regenerating it.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/AgentENV/pull/51 **Author:** [@shudorcl](https://github.com/shudorcl) **Created:** 7/28/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `issue-15-snapshot-locality` --- ### 📝 Commits (10+) - [`a762ada`](https://github.com/kvcache-ai/AgentENV/commit/a762ada38fe19a872be906ee1eb566393dbd5666) feat(gateway): add snapshot locality requirements - [`1c1af6c`](https://github.com/kvcache-ai/AgentENV/commit/1c1af6c3c3d669d73b43fc4bdb3d34da84e75575) refactor(scheduler): expose node locality context - [`086e71d`](https://github.com/kvcache-ai/AgentENV/commit/086e71d8ebc0ec06ba3da0ff0eabec9226faa4eb) feat(scheduler): score artifact locality - [`af1a96e`](https://github.com/kvcache-ai/AgentENV/commit/af1a96e9059ef38585562e5ffa7fffe9cd0c9943) feat(scheduler): apply locality before node selection - [`1b82399`](https://github.com/kvcache-ai/AgentENV/commit/1b8239933410735a2ae43bb85c1d6cba46a1aa02) chore(api): format generated scheduler protobuf - [`211fbf4`](https://github.com/kvcache-ai/AgentENV/commit/211fbf41a616dd69372ccf6fbeb02a5ba9e84deb) docs(services): document snapshot locality scheduling - [`ecd391b`](https://github.com/kvcache-ai/AgentENV/commit/ecd391b1904879c0c7ea15e3bdf57c9ad9e78f58) Update services/scheduler/internal/locality.go - [`89f1dad`](https://github.com/kvcache-ai/AgentENV/commit/89f1dadbf19614ece9f6c2a21b075b41e3dd8e0a) Update services/scheduler/internal/service.go - [`460c590`](https://github.com/kvcache-ai/AgentENV/commit/460c59066e63962b5080ae37269dfd67ea94a13d) Update services/scheduler/internal/locality.go - [`186df1d`](https://github.com/kvcache-ai/AgentENV/commit/186df1d327199b6452b6578086636ee8e3eee465) fix(services): harden snapshot locality scheduling ### 📊 Changes **18 files changed** (+1123 additions, -229 deletions) <details> <summary>View changed files</summary> 📝 `services/README.md` (+29 -0) 📝 `services/api/proto/scheduler.pb.go` (+277 -207) 📝 `services/api/proto/scheduler.proto` (+9 -0) 📝 `services/gateway/internal/schedule_hint.go` (+18 -2) 📝 `services/gateway/internal/schedule_hint_test.go` (+49 -0) 📝 `services/gateway/internal/server.go` (+2 -1) ➕ `services/gateway/internal/template_resolution.go` (+47 -0) ➕ `services/gateway/internal/template_resolution_test.go` (+95 -0) 📝 `services/go.mod` (+1 -1) ➕ `services/scheduler/internal/locality.go` (+116 -0) ➕ `services/scheduler/internal/locality_test.go` (+184 -0) 📝 `services/scheduler/internal/node_registry.go` (+12 -11) 📝 `services/scheduler/internal/node_registry_test.go` (+44 -0) 📝 `services/scheduler/internal/service.go` (+31 -3) 📝 `services/scheduler/internal/service_test.go` (+133 -0) 📝 `services/scheduler/internal/store.go` (+43 -3) 📝 `services/scheduler/internal/store_test.go` (+22 -0) 📝 `services/scheduler/internal/types.go` (+11 -1) </details> ### 📄 Description ## What This PR implements the first P2P-backed snapshot locality preference slice for #15. For `POST /sandboxes`, when `templateID` is a canonical snapshot UUID in the standard hyphenated form, the Gateway derives two advisory locality requirements using the existing fixed snapshot artifact identities: - `snapshot/v1/artifacts/{snapshot_id}/vm_state.bin` - `snapshot/v1/artifacts/{snapshot_id}/firecracker-manifest.json` After resource-limit filtering, the Scheduler prefers eligible nodes that advertise the highest coverage of those artifacts in the existing P2P artifact index. The configured scheduling strategy performs the final selection among tied nodes. Template aliases remain supported, but intentionally do not produce locality hints because the Gateway has no authoritative pre-scheduling alias resolver. Alias requests therefore use the existing scheduling behavior. ## Why The scheduler currently ignores whether an eligible node already provides the snapshot artifacts needed to start a sandbox. It can therefore select a cold node even when another feasible node advertises the required artifacts through the existing P2P index. This change introduces a bounded, advisory locality preference without duplicating load-aware selection or changing runtime fallback behavior. It establishes the request and inventory path needed for later locality work while keeping the scheduler's artifact store limited to artifact-key-to-node mappings. ## Related issue Refs #15 This PR is one focused part of #15 and does not close the full issue. ## Scope and non-goals In scope: - derive fixed snapshot artifact requirements from canonical snapshot IDs; - carry advisory locality requirements through the internal scheduling protocol; - prefer the eligible nodes with the highest positive artifact coverage; - preserve existing selection behavior when locality data is absent or stale; - bound request inspection, requirement count, and artifact lookup work; - remove a node's artifact mappings when it is unregistered; - document and test the new scheduling flow. Non-goals: - resolving template aliases before scheduling; - resolving OCI image references into OverlayBD layer requirements; - indexing non-P2P local caches; - adding locality requirements for `POST /sandboxes-cold`; - introducing a monolithic `cache_aware` strategy; - implementing or duplicating load-aware node selection; - changing hard resource filtering or runtime artifact fallback behavior. ## Design and behavior changes The scheduling flow becomes: ```text ready nodes -> hard resource-limit filtering -> snapshot artifact locality preference -> configured strategy selection -> selected node ``` The Gateway inspects only exact `POST /sandboxes` creation requests. It buffers at most 64 KiB while extracting `templateID` and restores the original body for upstream forwarding. Oversized, malformed, empty, alias-based, and non-canonical template IDs produce no locality requirements and continue through normal scheduling. For a canonical snapshot UUID, the Gateway normalizes the ID and emits the two fixed artifact keys listed above. It does not query runtime nodes to resolve aliases. Such fan-out would be non-authoritative, could observe inconsistent alias mappings, would amplify unauthenticated request work, and could expand credential exposure. The Scheduler validates the request-controlled locality requirement count before node discovery. After resource filtering, `PreferLocalNodes`: - trims and deduplicates requirement keys; - isolates lookups by cluster ID and P2P backend; - scores only already-eligible candidates; - retains all nodes tied for the highest positive coverage; - preserves the complete eligible set when no useful locality data exists. Artifact lookup intersects provider records with the eligible node-ID set inside the store, so scheduling work scales with current candidates rather than all historical providers. Locality does not create reservations or other request-lifetime state. Locality remains advisory. Stale or missing records do not fail sandbox creation, and the selected runtime continues through its normal local cache, P2P, snapshot repository, or registry fallback path. ## Compatibility and operations - Public API or generated protocol: Adds optional `ScheduleRequestHint.locality_requirements`, `LocalityRequirement.key`, and `NewSandboxHint.template_id` fields to the internal scheduler protobuf. The additions are wire-compatible. Older senders omit them, and receivers without locality support ignore the unknown fields. - Configuration or defaults: No configuration fields or defaults change. Existing `round_robin` and `random` strategy behavior is unchanged when locality requirements or provider records are unavailable. - Snapshot manifest, artifact layout, or storage format: No format or layout changes. The implementation reuses the existing `snapshot/v1/artifacts/{snapshot_id}/...` identities. - Upgrade and rollback: No persistent migration is required. Gateway and Scheduler versions can be rolled back without changing snapshot data or the artifact index; rollback removes the placement preference and restores the previous selection behavior. - Host requirements, permissions, ports, or dependencies: No new host permissions, services, ports, or runtime packages. `github.com/google/uuid` is promoted from an indirect to a direct Go dependency at the same version. ## Validation - [ ] `make fmt` - [ ] `make clippy` - [ ] `make test-unit` - [x] Relevant Rust integration tests - [x] `make -C services test` (required when `services/` changes) - [x] Generated clients/server regenerated with the documented `make` target - [x] Documentation updated - [ ] Benchmarks or performance comparison completed Commands and results: ```text make -C services proto PROTOC=<protoc 3.21.12> gofmt -w services/api/proto/scheduler.pb.go services/api/proto/scheduler_grpc.pb.go # PASS: protoc-gen-go 1.36.11 and protoc-gen-go-grpc 1.6.2; no tracked diff make -C services build make -C services test make -C services fmt-check make -C services vet # PASS with Go 1.25.0 cd services go test -race ./gateway/internal ./scheduler/internal # PASS with Go 1.25.0 git diff --check # PASS ``` The Services commands above were run locally. `fmt-check` was verified from an LF-normalized export of the same commit because this WSL checkout uses `core.autocrlf=true`. The following remote Services CI and integration jobs also passed for commit `b2f9346577b001bba28fb038d46ae0beb233c956`: - Services CI: https://github.com/shudorcl/AgentENV/actions/runs/30394461927 - Integration Tests: https://github.com/shudorcl/AgentENV/actions/runs/30394461784 Skipped checks and reasons: - Root Rust formatting, Clippy, and unit-test targets were not run because this PR changes only the Go control-plane module and its generated Go protobuf. The remote AgentENV integration, single-node E2E, and docker-compose E2E jobs passed. - No benchmark was run. The scheduling hot path is bounded to at most 32 locality requirements, and provider lookup is intersected with the current eligible candidate set. The PR adds focused boundary and stale-record tests, but does not claim a measured performance improvement. ## Risks and reviewer notes - Locality metadata is eventually consistent and can be stale. The preference is deliberately advisory, and no positive match preserves the complete eligible candidate set. - Requests with aliases receive no locality optimization. Supporting aliases safely requires an authoritative control-plane resolver and is left for a follow-up. - Hint extraction occurs before upstream authentication. Request-body inspection is therefore bounded to 64 KiB and does not fan out to runtime nodes. - The locality requirement list is request-controlled. The Scheduler rejects more than 32 requirements before node snapshotting or enrichment. - Popular artifacts may have many historical providers. `LookupEligible` bounds work to the current eligible candidate IDs, and unregistering a node removes its artifact mappings. - Reviewers should start with `services/gateway/internal/schedule_hint.go`, `services/gateway/internal/template_resolution.go`, `services/scheduler/internal/locality.go`, and `services/scheduler/internal/service.go`. ## Checklist - [x] The PR contains one coherent change and no unrelated formatting or refactoring. - [x] New behavior is covered by tests, or I explained why testing is impractical. - [x] Logs and examples contain no credentials, tokens, or private registry information. - [x] I did not manually edit generated code without updating its source and regenerating it. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-29 14:20:07 +00:00