feat: AWS Neuron (Trainium/Inferentia) device backend

[hoyajigi]

Hardware verification status: all four CI gates pass on macOS and on
Linux (where the module actually compiles — it is entirely #[cfg(target_os = "linux")]), plus --features mock, which neither the default build nor CI's
clippy covers and which caught a non-exhaustive match on PlatformType.
What remains unverified is hardware-bound: the probe instance had exactly
one Neuron device and no workload was ever run, so multi-device numbering and
the populated neuron-monitor path are pinned by synthesised fixtures.
Details at the bottom. A trn1.32xlarge or inf2.24xlarge closes most of
that gap in one short session.

feat: AWS Neuron (Trainium / Inferentia) device backend

Adds a Neuron backend to all-smi: device reader, Prometheus exporter,
doctor checks, and a mock platform. Branch: feature/aws-neuron-reader.

Hardware this was built against

Everything below was measured on a real instance, not inferred from docs:

Instance	trn1.2xlarge
Devices	1 Trainium device, 2 NeuronCores (NDv2 / NCv2, device_name = Trainium1)
Host CPU	Intel Xeon Platinum 8375C @ 2.90 GHz, 8 vCPU
OS	Ubuntu 22.04.5, kernel 6.8.0-1046-aws, Neuron DLAMI
Driver	aws-neuronx-dkms 2.26.5.0 (neuron module, /sys/module/neuron/version)
Tools	aws-neuronx-tools 2.28.23.0 (neuron-ls, neuron-monitor 2.28.23.0)

The full capture (63 artifacts: command output, exit codes, every sysfs
leaf and its value, non-root and device-absent behaviour) lives outside
this repo and is the source for every claim in this description.

What is implemented

One GpuInfo row per NeuronCore, device_type = "NPU" (which is what
makes CommonNpuExporter emit the shared all_smi_npu_* family for free).

GpuInfo field	Source
uuid	neuron-<serial>-nc<flat core>; falls back to the PCI BDF when sysfs has no serial
name	sysfs info/architecture/device_name → "AWS Trainium1"; falls back to neuron-monitor's neuron_device_type, then "AWS Neuron Device"
total_memory	neuron-ls memory_size ÷ nc_count (34359738368 B ÷ 2 = 16 GiB)
used_memory	sysfs neuron_core<N>/stats/memory_usage/device_mem/present
utilization	neuron-monitor neuroncore_counters.neuroncores_in_use[<flat idx>].neuroncore_utilization — best effort
numa_node_id	neuron-ls numa_node string; "-1" canonicalised to None
temperature	absent (0) — no source exists
frequency	absent (0) — no source exists
power_consumption	absent (-1.0) — see below
detail	BDF, CPU affinity, NUMA node, device index, global flat core index, local core index, core count, device HBM bytes, serial, arch_type, device_name, instance_type, core arch_type, connected_to (when non-null), driver version, raw power line, utilization_source, memory_usage_source, and neuron-monitor hardware info
13 NVIDIA-only fields	None / empty, as in furiosa.rs and tenstorrent.rs

Exporter emits all_smi_neuron_device_info, _serial_info, _arch_info,
_instance_info, _driver_info, _device_index, _core_index,
_core_count, _device_memory_bytes. Every series is conditional on its
detail key being present, so an unsourced value is absent from the
exposition rather than exported as 0.

doctor adds neuron.dev_node, neuron.driver, neuron.sysfs,
neuron.tools. The mock server gains --platform neuron (aliases
trainium, inferentia).

Design decisions worth reviewing

neuron-ls + sysfs, not neuron-monitor, for everything static.
neuron-ls --json-output is a natural one-shot (exits 0) and per-core
memory is in sysfs, world-readable and populated with no workload
attached. Only per-core utilization needs neuron-monitor.

neuron-monitor is shelled out one-shot, not backgrounded. It has no
--once / --count / --period flag, but it emits a complete NDJSON
record immediately at t≈0 (then every 5.000 s — measured: 02:41:18.673,
:23.678, :28.678). The reader spawns it, reads the first line with a 2 s
deadline on a helper thread, then kills the child. A single failure sets
a process-global flag so a host without the tool never pays the timeout
again. This deliberately avoids a background-process subsystem; the
existing sanctioned command helpers could not be used here because
run_command_with_timeout discards stdout when it kills a child, and
neuron-monitor never exits on its own.

Core keys are a global flat index. Per AWS's own shipped
neuron-monitor-device-view.py: nc_idx = nd_idx * neuroncore_per_device_count + nc_idx_counter, then
neuroncores_in_use[str(nc_idx)]. On a 16-device trn1.32xlarge, device 3
core 1 is key "7". The reader computes the same expression rather than
assuming device-local keys.

Binaries by absolute path. /opt/aws/neuron/bin is added to PATH
only by the DLAMI profile scripts. Measured: an unqualified neuron-ls
under sudo -u exits 127. Both the reader and the doctor check use
/opt/aws/neuron/bin/neuron-ls and .../neuron-monitor.

Detection. /dev/neuron0 (mode 0666, world rw — no group gate;
every tool works fully as nobody, verified) else lspci -nn containing
1d0f:7164 (Amazon.com, Inc. NeuronDevice (Trainium)). No CLI probe in
the detector, because the PATH trap above would make it a false negative.

Platform gate. Whole-module #[cfg(target_os = "linux")] at the mod
declaration and every use site, like tenstorrent. No target_arch gate
(no reader has one) and no cargo feature (this repo's rule is that a
feature exists only for an optional external crate; this links nothing).

Absence-contract choices (issue #325)

• Temperature: absent. trn1 exposes no temperature field anywhere —
  not in neuron-ls, not in neuron-monitor, not in sysfs (checked every
  leaf; lm-sensors is not even installed). Reported as 0, the
  unavailable marker for temperature.
• Frequency: absent. No clock probe exists for NeuronCores.
• Power: absent (-1.0), deliberately. A field does exist:
  /sys/devices/virtual/neuron_device/neuron0/stats/power/utilization =
  POWER_STATUS_VALID,1789180860,0.00,0.00,0.00. Three unlabelled floats,
  all 0.00 on real idle hardware, and a Trainium device plainly does not
  draw 0 W. Rather than guess which float is watts, the raw line is kept
  in detail["power_utilization_raw"] and the reading is marked absent.
  This is the most reviewable call in the PR — if someone can confirm
  the field order against a loaded device, promoting it to a real reading
  is a small change.
• Utilization: absent unless a runtime is attached. With
  neuron_runtime_data: [] the map is empty and the row reports absent,
  never 0 %.
• No device: no rows. Measured: device absent → exit 1, stdout empty
  and clean, one logfmt line on stderr. "Empty stdout + non-zero rc" is
  therefore the absence test. A bad flag prints its message on stdout
  with exit 1, and the non-zero rc catches that too. Missing tools: 127.
• Memory has no sentinel in GpuInfo (u64), so an unreadable sysfs
  leaf yields 0 like every other reader. detail["memory_usage_source"]
  distinguishes sysfs from unavailable so the 0 is attributable.

Judgement calls a reviewer might overrule

1. One row per NeuronCore, not per device. Chosen because memory and
   utilization are both per-core in the driver and in AWS's own tooling.
   The cost is item 2.
2. Device HBM split evenly across cores (32 GiB → 16 GiB per core).
   neuron-ls reports HBM per device; every consumer of total_memory
   sums rows, so charging the full device size to each core would double
   the host's reported memory. The even split is a derivation, not a
   reading. The undivided value is preserved in
   detail["device_memory_bytes"] and exported as
   all_smi_neuron_device_memory_bytes.
3. neuron-monitor is spawned on every refresh tick when the binary
   exists (~2 process spawns per tick total, in line with the Furiosa
   reader's 3). An alternative is to skip it unless neuron-ls reports a
   non-empty neuron_processes, which would be cheaper but relies on an
   array whose element shape is unverified.
4. new_command is used directly for neuron-monitor instead of
   execute_command_default, for the stdout-on-kill reason above.
   neuron-ls goes through execute_command_default as normal. No
   std::process::Command::new anywhere.
5. No privileges.dev_neuron doctor check, unlike
   privileges.dev_tenstorrent. /dev/neuron0 is world-rw with no group
   gate, so there is no privilege question to report, and
   neuron.dev_node already covers presence.
6. No TUI status-message notification (get_*_status_message /
   *_notification_shown), unlike Tenstorrent and TPU. Neuron has no
   "initializing"/"failed" state to surface; add one later if a real
   failure mode turns up.
7. detail keys are snake_case (pci_bdf, serial_number),
   following furiosa.rs. Some other readers use Title Case keys; the
   repo is not consistent and this picked the NPU-family convention.
8. DeviceType::AwsNeuronNpu was added to the public client.rs
   enum. Additive, but it is public API.
9. Mock emits no temperature/power/frequency series. Consistent with
   the hardware, but it makes the Neuron mock the only one whose metric
   set is a strict subset of the others.

What could NOT be verified

• Multi-device numbering. The instance had exactly one device.
  flat_core_index implements AWS's own formula and is unit-tested
  (device 3 core 1 → 7), but no multi-device output was ever observed.
  Whether neuroncore_ids is global or device-local on a 16-device box is
  unknown; the reader does not depend on it except as a fallback for a
  missing nc_count.
• connected_to topology. null on a single-device instance. Its
  shape on trn1.32xlarge is unknown, so it is carried through as opaque
  JSON (stringified into detail) rather than modelled. It is not fed
  into the topology tab.
• Live per-core utilization. No workload was run, so
  neuron_runtime_data was always []. The populated-record test literal
  is synthesised from AWS's own shipped frontends
  (neuron-monitor-prometheus.py, neuron-monitor-device-view.py), which
  are authoritative for the field names but are not captured output. It is
  labelled as such in the test. The empty-record literal is verbatim
  capture.
• neuron_processes element shape. Always [] with no workload
  attached, so get_process_info parses pid / command / name
  leniently and drops any entry without a PID rather than emitting PID 0.
  It may simply return nothing on a busy host; that degrades to "no
  process rows", never to wrong rows.
• Reboot stability of identifiers. serial_number
  (9ff5434815c8bd80) and the PCI BDF both look stable and the UUID is
  composed from them, but the instance was never rebooted. No comment in
  the code claims stability.
• Inferentia. Only Trainium (NDv2) hardware was available. Inf1/Inf2
  use the same driver, tools, and sysfs layout, and the code paths are
  generic, but no Inferentia device was tested. The 1d0f:7164 PCI ID in
  the fallback detector is Trainium's; an Inferentia-only host with no
  /dev/neuron0 would not be detected.
• logical_neuroncore_config > 1 (LNC / "logical NeuronCore" mode on
  trn2). Observed as 1. It is recorded in detail but does not affect
  row expansion.

Validation

Run on macOS (aarch64-apple-darwin) and on a Linux box, because the whole
Neuron module is behind #[cfg(target_os = "linux")] and therefore is not
compiled at all on macOS.

Command	macOS	Linux (x86_64, rustc 1.97.1)
cargo fmt --all -- --check	exit 0	exit 0
cargo clippy --all-targets -- -D warnings	exit 0	exit 0
cargo test	exit 0	exit 0 (lib 1724 passed / 1 ignored, bin 1903 passed / 2 ignored)
cargo build --no-default-features --lib	exit 0	exit 0

Additionally on Linux, because src/mock sits behind the non-default
mock feature and neither plain cargo test nor CI's
cargo clippy --all-targets builds it:

• cargo clippy --all-targets --features mock -- -D warnings — exit 0
  (this is what caught a non-exhaustive match on PlatformType in
  src/mock/generator.rs)
• cargo test --features mock --bin all-smi-mock-server — 65 passed

Caveat on the Linux run: cargo test there was invoked with
--skip test_ensure_sudo_permissions. Those two tests call
ensure_sudo_permissions(), which spawns sudo -v and exit(1)s when it
fails; they abort the harness on any host without an interactive TTY.
Verified pre-existing by running them from a pristine main on the same
box: sudo: A terminal is required to authenticate, same failure. They do
not run on macOS (one is cfg(not(target_os = "macos")), the other is
ignored there), which is why baseline main is green locally.

Tests added (14 in the lib, 2 in the mock bin)

device::readers::neuron: neuron-ls array parses; memory_size read as
bare bytes (34359738368 → 32 GiB, not multiplied); numa_node: "-1"
string parses and canonicalises to absent; connected_to: null tolerated;
empty stdout yields an empty list (not a row of zeros); per-core rows split
device memory and mark every unsourced metric absent; monitor record with
no runtime reports no utilization; monitor record with a runtime keys cores
by global flat index; flat index arithmetic; core-count fallback;
device-name fallback chain; processes without a PID dropped.

api::metrics::npu: every vendor's hardcoded pool index pinned (Gaudi,
Rebellions, Furiosa, TPU, Tenstorrent, Neuron, plus pool length); Neuron
rows route to the Neuron exporter.

doctor::checks::neuron: ids namespaced and every check total on a host
with no Neuron hardware.

mock::templates::neuron: cores numbered globally and grouped by device;
temperature/power/frequency series absent.

Real captured output is embedded as string literals (house pattern, cf.
src/device/hlsmi/parser.rs), with the one synthesised literal explicitly
marked.