Project 2: A Vector-Aware Allocation Decider

OpenSearch's allocation layer decides which node each shard copy lives on. It balances by shard count and by disk, and it consults a chain of AllocationDeciders that can each veto or throttle a move. None of them know that a knn_vector shard using the faiss engine carries a large native-memory footprint that lives entirely outside the JVM heap — loaded by the plugin on first query or via the warmup API, capped by a native-memory circuit breaker. So the balancer will happily pack three heavy vector shards onto one node and starve another, because from its point of view they are "just shards."

This project builds a vector-aware AllocationDecider (or an allocation-influencing awareness attribute) that accounts for the native-memory footprint of k-NN shards when placing and balancing them. It sits exactly on the seam between two subsystems most contributors treat as separate worlds: cluster coordination / allocation and k-NN native memory. Bridging them is the whole skill.

Note: This is the project to take if you came through shard allocation and want vectors, or vice versa. You already met an allocation decider in Lab 4: Fix-It Allocation Decider — this is that muscle, scaled up to a real cross-plugin design.

Problem & motivation

The allocation deciders you know enforce real constraints: SameShardAllocationDecider (don't put primary and replica together), DiskThresholdDecider (watermark-aware), AwarenessAllocationDecider (zone spreading), ThrottlingAllocationDecider (rate-limit recoveries), FilterAllocationDecider (include/exclude). They reason about shard count, on-disk bytes, and topology. None reason about native memory.

For a vector workload that is a real, operational gap:

• A faiss HNSW index loads its graph into native memory, sized by the number of vectors, the dimension, and M. That footprint can dwarf the JVM heap and is capped by knn.memory.circuit_breaker.limit. When too many heavy vector shards land on one node, that node's circuit breaker trips on first query — queries to those shards fail — while a sibling node sits half-empty. The balancer caused it and cannot see it.
• "Just add more shard-count balancing" does not help: ten tiny lexical shards and three huge vector shards have the same count weight but radically different native-memory cost.
• The footprint is known — the plugin exposes it through GET /_plugins/_knn/stats and accounts it in NativeMemoryCacheManager. The information exists; allocation just never asked for it.

The motivating constraint that makes this hard (and worth a design note): allocation runs in core, the footprint lives in a plugin, and a decider must not make a remote call mid-allocation round. You need the signal to arrive in cluster state (or node stats already gossiped), so the decider can read it cheaply and deterministically.

Real-world grounding

There is no single merged "vector-aware allocation decider." That is why this is a build-the-capability-then-RFC project. Ground it in two real bodies of work:

• The native-memory circuit-breaker rearchitecture — k-NN #1582: https://github.com/opensearch-project/k-NN/issues/1582 — the discussion of how native memory should be accounted and capped, which is precisely the signal your decider consumes.
• A concrete native-memory circuit-breaker config bug — k-NN #585: https://github.com/opensearch-project/k-NN/issues/585 — evidence that getting this accounting right is genuinely subtle.

For the allocation side, ground in the existing deciders themselves (read the source) and search live for prior art before you design:

# in opensearch-project/OpenSearch
is:issue allocation decider memory OR "native memory" OR knn
is:issue rebalance vector OR knn
# in opensearch-project/k-NN
is:issue allocation OR rebalance OR "shard placement"

Citation discipline: cite #1582 and #585 (real). For the allocation-side issue, link the current one you find — do not invent a number. If you find none, that is itself a finding for your design note: this gap is unaddressed, which is the case for an RFC.

Subsystems you'll touch

Cross-references: shard allocation deep dive · cluster state · k-NN native JNI and memory · circuit breakers and memory · the allocation-decider lab.

Phased plan

Phase 0 — Read the chain, run the failure (1 day)

Read the decider chain and the balancer weight function. Then reproduce the imbalance.

# In an OpenSearch clone:
grep -rn "class .*AllocationDecider" server/src/main/java/org/opensearch/cluster/routing/allocation/decider
grep -rn "canAllocate\|canRemain\|canRebalance" server/src/main/java/org/opensearch/cluster/routing/allocation/decider/DiskThresholdDecider.java
grep -rn "class BalancedShardsAllocator\|WeightFunction\|weight(" server/src/main/java/org/opensearch/cluster/routing/allocation/allocator

# Reproduce: bring up a 3-node cluster with k-NN, create several heavy vector indices,
# warm them, and watch the native-memory stats skew across nodes.
curl -s localhost:9200/_cat/shards?v
curl -s "localhost:9200/_plugins/_knn/stats?pretty"      # per-node graph memory usage
curl -s "localhost:9200/_nodes/stats/breaker?pretty"     # where the circuit breaker shows

Write capstone-work/allocation-trace.md: how a reroute round flows through AllocationService → AllocationDeciders → BalancedShardsAllocator, with the file:line where the count and disk weights are applied. This is your execution-path artifact.

Phase 1 — A read-only decider that observes (the scoped slice)

Build a custom AllocationDecider (registered via the plugin's ClusterPlugin getAllocationDeciders extension point — grep getAllocationDeciders) that, for now, only logs/decides on a synthetic per-node footprint you put in cluster state by hand. It returns Decision.YES always but records what it would have decided. This proves the wiring without risking allocation correctness.

public class VectorMemoryAllocationDecider extends AllocationDecider {
    static final String NAME = "vector_memory";
    private volatile long perNodeLimitBytes;   // from a dynamic cluster setting

    @Override
    public Decision canAllocate(ShardRouting shardRouting, RoutingNode node, RoutingAllocation allocation) {
        long projected = projectedNativeMemory(node, shardRouting, allocation);
        if (projected > perNodeLimitBytes) {
            return allocation.decision(Decision.NO, NAME,
                "node [%s] projected native vector memory [%d] would exceed limit [%d]",
                node.nodeId(), projected, perNodeLimitBytes);
        }
        return allocation.decision(Decision.YES, NAME, "within native vector memory budget");
    }
}

Register it:

// in your plugin (a ClusterPlugin):
@Override
public Collection<AllocationDecider> createAllocationDeciders(Settings settings, ClusterSettings cs) {
    return List.of(new VectorMemoryAllocationDecider(settings, cs));
}

Test with the allocation test harness used by the built-in deciders:

// extends OpenSearchAllocationTestCase
public void testVetoesNodeOverNativeMemoryBudget() {
    AllocationService service = createAllocationService(/* with the decider */);
    ClusterState state = /* 2 nodes, set footprints so node1 is full */;
    state = service.reroute(state, "test");
    assertThat(shardsOn(state, "node1"), /* shard did not land on the full node */);
}

./gradlew :server:test --tests "*VectorMemoryAllocationDecider*"
./gradlew :server:check -x internalClusterTest

Phase 2 — Get the real footprint into allocation

Replace the synthetic number with the actual k-NN footprint. The clean path: have the k-NN plugin publish per-node native-memory usage into a place the decider can read without an RPC — either via ClusterInfo (the same mechanism DiskThresholdDecider uses for disk usage) or a small cluster-state custom updated from node stats.

grep -rn "ClusterInfo\b\|ClusterInfoService\|getNodeLeastAvailableDiskUsages" server/src/main/java/org/opensearch/cluster
grep -rn "getShardSize\|shardSizes\|NodeStats" server/src/main/java/org/opensearch/cluster/ClusterInfo.java

Decide and document the mechanism in your design note — this is the architecturally load-bearing choice and the thing an RFC would argue about (latency of the signal, staleness, who owns the field, BWC of any new cluster-state custom).

Phase 3 — Influence balancing, not just placement

A canAllocate veto stops bad placements but does not rebalance an already-skewed cluster. For that, the native-memory cost must enter the BalancedShardsAllocator weight function, so the balancer actively moves heavy shards off hot nodes.

grep -rn "WeightFunction\|float weight\|theta\|indexBalance\|shardBalance" \
  server/src/main/java/org/opensearch/cluster/routing/allocation/allocator/BalancedShardsAllocator.java

This is the deepest and most contentious change — touching the weight function affects every workload, not just vectors. Scope it behind a setting, off by default, and prove with a test that a skewed cluster converges to balanced native memory without destabilizing count/disk balance.

Phase 4 — Prove it converges and does no harm

Build an OpenSearchAllocationTestCase / InternalTestCluster scenario: a skewed start state, run reroute to convergence, assert native-memory variance across nodes drops below a threshold and shard-count/disk balance does not regress. The "does no harm" half is as important as the "it balances" half — a decider that fixes vectors by wrecking lexical balance won't land.

Deliverables

• capstone-work/allocation-trace.md — reroute → deciders → balancer, with file:line citations
• capstone-work/design.md — the gap, the signal-delivery mechanism chosen (ClusterInfo vs custom), alternatives rejected, BWC analysis
• Phase 1: a registered read-only/veto decider with an OpenSearchAllocationTestCase test
• Phase 2: real per-node native-memory footprint delivered into allocation without an RPC
• (Phase 3) the native-memory weight in the balancer, behind a default-off setting
• (Phase 4) a convergence test that also proves count/disk balance does not regress
• capstone-work/validation.md — ./gradlew :server:check output, test commands, seeds
• An upstreaming decision: a written RFC draft (this almost certainly needs one) + the issue comment you would open
• A 500–1000 word write-up: the cross-subsystem seam and how you bridged it

Difficulty & time

Warning: Touching BalancedShardsAllocator's weight function changes behaviour for every index in every cluster. That is exactly why Phases 1–2 (a vetoing decider behind a setting) are the realistic upstream slice and Phase 3 is RFC-gated.

Stretch goals

• Make the decider respect the circuit-breaker limit per node, so it leaves headroom rather than packing to 100% — directly addressing the spirit of k-NN #585.
• Add an awareness-attribute mode: spread vector shards across attributes the way AwarenessAllocationDecider spreads across zones, but weighted by footprint.
• Expose a _cat/allocation-style column showing per-node native vector memory.
• Account for the warmup state: an unwarmed shard's footprint is projected, a warmed one is real — model both and let the decider use the real number when available.

Evaluation

Self-grade against the 100-point rubric:

How to turn this into a real contribution

1. This is an RFC-first project. Allocation behaviour is core; you do not change the balancer by surprise. Write the design note as an RFC and open it on the OpenSearch repo.
2. Phase 1–2 is the mergeable slice: a registered, setting-gated, default-off vetoing decider plus a clean footprint signal. That can land as an opt-in feature.
3. The cross-plugin signal is the crux. Argue the mechanism (ClusterInfo vs cluster-state custom) in the RFC; both have BWC and staleness consequences. Maintainers will care more about how the number arrives than about the decider logic.
4. Bring the reproduction. A _cat/shards + _plugins/_knn/stats before/after that shows the skew and then the fix is worth more than any prose.
5. DCO applies (GitHub-native, git commit -s, CHANGELOG.md, backport bot) for whatever slice lands in core.

Even if nothing merges, a clean RFC + a working setting-gated decider + a convergence proof is a strong portfolio artifact in the hardest part of the engine.