Skip to content

test(neuron): integration coverage + fixture for multi-runtime per-co… #749

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
spanaik wants to merge 2 commits into
base: main
Choose a base branch
from
+611 −8
Open

test(neuron): integration coverage + fixture for multi-runtime per-co… #749

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
17c44f9
test(neuron): integration coverage + fixture for multi-runtime per-co…
spanaik Sep 22, 2026
9484c62
Merge branch 'main' into neuron-multi-runtime-tests
spanaik Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 148 additions & 7 deletions terraform/eks/daemon/otel-neuron/main.tf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,103 @@ resource "null_resource" "neuron_burn_core" {
}
}

# --- neuron-burn-peer Deployment: the SECOND runtime on the workload node ---
#
# Gives the node two Neuron runtimes burning different cores, the only shape in
# which the per-core data-loss defect appears (test/otel/neuron/multi_runtime_test.go).
# inf2.xlarge has 1 device x 2 cores; burn-core takes one, this takes the other.
#
# Three constraints, each of which breaks an existing test if changed:
# - NOT named neuron-burn-core-*: TestNeuronBurnWorkloadLabels and
# TestNeuronBurnCorePodColor match HasPrefix(pod, "neuron-burn-core").
# - no `neuron-test: "true"` label: burn-core's podAntiAffinity targets it, which
# would make the two mutually exclusive on a host and push burn-core onto the
# idle node, breaking the idle-node tests.
# - podAffinity, not just nodeSelector: both node groups use var.instance_type, so
# a nodeSelector alone could land this on the idle node.
#
# Sits Pending on first apply until burn-core is scheduled; the scheduler retries.

resource "null_resource" "neuron_burn_peer" {
depends_on = [
helm_release.neuron_device_plugin,
null_resource.kubectl,
null_resource.neuron_burn_core,
]
provisioner "local-exec" {
command = <<-EOT
cat <<'EOF' | kubectl apply -f -
apiVersion: apps/v1
kind: Deployment
metadata:
name: neuron-burn-peer
namespace: default
spec:
replicas: 1
revisionHistoryLimit: 2
progressDeadlineSeconds: 300
strategy:
type: RollingUpdate
rollingUpdate:
maxSurge: 0
maxUnavailable: 1
selector:
matchLabels:
app: neuron-burn-peer
template:
metadata:
labels:
app: neuron-burn-peer
ci-test.example.com/pod-color: green
spec:
tolerations:
- key: aws.amazon.com/neuron
operator: Exists
effect: NoSchedule
affinity:
podAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: app
operator: In
values: ["neuron-burn-core"]
topologyKey: kubernetes.io/hostname
nodeSelector:
node.kubernetes.io/instance-type: ${var.instance_type}
containers:
- name: neuron-burn
image: public.ecr.aws/neuron/pytorch-inference-neuronx:2.1.2-neuronx-py310-sdk2.20.2-ubuntu20.04
command: ["python3", "-c"]
args:
- |
import torch
import torch_neuronx
import time
print("Compiling neuron trace (this takes a minute)...")
x = torch.randn(256, 256)
model = torch.nn.Linear(256, 256, bias=False)
traced = torch_neuronx.trace(model, x)
print("Trace compiled. Starting burn loop...")
iteration = 0
while True:
start = time.time()
for _ in range(1000):
_ = traced(x)
elapsed = time.time() - start
iteration += 1
print(f"Iteration {iteration}: 1000 inferences in {elapsed:.2f}s")
resources:
limits:
aws.amazon.com/neuroncore: "1"
requests:
cpu: "1"
memory: 4Gi
EOF
EOT
}
}

# --- neuron-burn-multi-device Deployment (inf2.24xlarge, uses multiple devices) ---

resource "null_resource" "neuron_burn_multi_device" {
Expand Down Expand Up @@ -501,10 +598,39 @@ resource "null_resource" "restart_pods" {
# --- Wait for neuron-monitor pods to be ready ---

resource "null_resource" "wait_neuron_monitor" {
depends_on = [null_resource.restart_pods, null_resource.neuron_burn_core, null_resource.neuron_burn_multi_device]
triggers = { timestamp = timestamp() }
depends_on = [
null_resource.restart_pods,
null_resource.neuron_burn_core,
null_resource.neuron_burn_peer,
null_resource.neuron_burn_multi_device,
]
triggers = { timestamp = timestamp() }
provisioner "local-exec" {
command = <<-EOT
# A runtime_tag only exists once the runtime is up, and the trace compile takes
# ~a minute, so without this the tests race the fixture and see one runtime.
echo "Waiting for burn deployments to become Available..."
for d in neuron-burn-core neuron-burn-peer; do
kubectl -n default rollout status deployment/$d --timeout=600s || {
echo "ERROR: deployment/$d did not become Available"
kubectl -n default describe deployment/$d || true
kubectl -n default get pods -l app=$d -o wide || true
exit 1
}
done

# A silent split across nodes would make every multi-runtime test vacuous.
NODES=$(kubectl -n default get pods -l 'app in (neuron-burn-core,neuron-burn-peer)' \
-o jsonpath='{range .items[*]}{.spec.nodeName}{"\n"}{end}' | sort -u | grep -c . || true)
if [ "$NODES" != "1" ]; then
echo "ERROR: neuron-burn-core and neuron-burn-peer are on $NODES nodes, expected 1."
echo "The multi-runtime per-core regression cannot be exercised unless both"
echo "runtimes share a node. Check the podAffinity on neuron-burn-peer."
kubectl -n default get pods -l 'app in (neuron-burn-core,neuron-burn-peer)' -o wide || true
exit 1
fi
echo "Both burn runtimes are co-located on one node."

echo "Waiting for neuron-monitor pods to be ready..."
READY=0
for i in $(seq 1 30); do
Expand Down Expand Up @@ -538,42 +664,57 @@ resource "null_resource" "validator" {
echo "This can take up to 15 minutes on a cold start due to PyTorch-Neuron image pull + trace compilation."
echo "Readiness signal: 'Iteration N' log lines from the active burn loop."
CORE_READY=0
PEER_READY=0
MULTI_READY=0
# 90 iterations × 10s = 15 minutes max.
for i in $(seq 1 90); do
# Match "Iteration" in the tail — the burn loop emits one per second once
# the trace is compiled. This is a more robust readiness signal than
# grepping for "Trace compiled" which scrolls off the tail quickly.
CORE_READY=$(kubectl logs -n default -l app=neuron-burn-core --tail=5 2>/dev/null | grep -c "^Iteration " || true)
# neuron-burn-peer is the second runtime on the core node. wait_neuron_monitor
# only gets it to Available (container Running); the trace compile that starts
# it emitting runtime_tag takes ~1 more minute, so gate on it here too --
# otherwise the multi-runtime tests can run against a single-runtime surface.
PEER_READY=$(kubectl logs -n default -l app=neuron-burn-peer --tail=5 2>/dev/null | grep -c "^Iteration " || true)
MULTI_READY=$(kubectl logs -n default -l app=neuron-burn-multi --tail=5 2>/dev/null | grep -c "^Iteration " || true)
if [ "$CORE_READY" -gt 0 ] && [ "$MULTI_READY" -gt 0 ]; then
echo "Neuron runtime active on both burn workloads (after $((i*10))s)"
if [ "$CORE_READY" -gt 0 ] && [ "$PEER_READY" -gt 0 ] && [ "$MULTI_READY" -gt 0 ]; then
echo "Neuron runtime active on all burn workloads (after $((i*10))s)"
break
fi
# Every 60s dump pod status so we can see image-pull vs crash vs runtime-init.
if [ $((i % 6)) -eq 0 ]; then
echo "--- Attempt $i ($((i*10))s elapsed): core_iter_lines=$CORE_READY multi_iter_lines=$MULTI_READY ---"
echo "--- Attempt $i ($((i*10))s elapsed): core_iter_lines=$CORE_READY peer_iter_lines=$PEER_READY multi_iter_lines=$MULTI_READY ---"
kubectl get pods -n default -l neuron-test=true -o wide 2>&1 | head -10 || true
# neuron-burn-peer carries no neuron-test label (the anti-affinity on
# neuron-burn-core targets it), so the dump above omits it.
kubectl get pods -n default -l app=neuron-burn-peer -o wide 2>&1 | head -5 || true
fi
sleep 10
done
if [ "$CORE_READY" -eq 0 ] || [ "$MULTI_READY" -eq 0 ]; then
echo "ERROR: Neuron burn loop not active after 15 minutes (core_iter=$CORE_READY multi_iter=$MULTI_READY)"
if [ "$CORE_READY" -eq 0 ] || [ "$PEER_READY" -eq 0 ] || [ "$MULTI_READY" -eq 0 ]; then
echo "ERROR: Neuron burn loop not active after 15 minutes (core_iter=$CORE_READY peer_iter=$PEER_READY multi_iter=$MULTI_READY)"
echo "=== Final pod status ==="
kubectl get pods -n default -l neuron-test=true -o wide 2>&1 || true
kubectl get pods -n default -l app=neuron-burn-peer -o wide 2>&1 || true
echo "=== burn-core events ==="
kubectl describe pod -n default -l app=neuron-burn-core 2>&1 | tail -40 || true
echo "=== burn-peer events ==="
kubectl describe pod -n default -l app=neuron-burn-peer 2>&1 | tail -40 || true
echo "=== burn-multi events ==="
kubectl describe pod -n default -l app=neuron-burn-multi 2>&1 | tail -40 || true
echo "=== burn-core logs ==="
kubectl logs -n default -l app=neuron-burn-core --tail=40 2>&1 || true
echo "=== burn-peer logs ==="
kubectl logs -n default -l app=neuron-burn-peer --tail=40 2>&1 || true
echo "=== burn-multi logs ==="
kubectl logs -n default -l app=neuron-burn-multi --tail=40 2>&1 || true
exit 1
fi

echo "=== Burn workload pod status ==="
kubectl get pods -n default -l neuron-test=true -o wide
kubectl get pods -n default -l app=neuron-burn-peer -o wide

echo "Waiting 6 minutes for metrics to propagate (covers Zeus 5-min staleness window)..."
sleep 360
Expand Down
Loading
Loading