Conversation
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Core isolation is not enforceable, container accounting is inaccurate for shared devices, and several lifecycle and build-policy issues remain.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds an AWS Neuron accelerator plugin with per-NeuronCore discovery, allocation, metrics, container integration, and manager slot registration.
Changes:
- Implements Neuron host discovery, sysfs metrics, and Docker allocation.
- Seeds Neuron and Tenstorrent resource slot metadata.
- Adds documentation, packaging metadata, and unit tests.
File summaries
| File | Description |
|---|---|
README.md |
Lists AWS Neuron support. |
changes/14580.feature.md |
Adds the feature changelog. |
fixtures/manager/example-resource-slot-types.json |
Seeds new slot metadata. |
src/ai/backend/install/fixtures/example-resource-slot-types.json |
Mirrors installer slot metadata. |
src/ai/backend/accelerator/neuron/__init__.py |
Exposes package version. |
src/ai/backend/accelerator/neuron/BUILD |
Defines package and wheel targets. |
src/ai/backend/accelerator/neuron/README.md |
Documents design and operation. |
src/ai/backend/accelerator/neuron/neuron_api.py |
Wraps Neuron CLI and sysfs access. |
src/ai/backend/accelerator/neuron/plugin.py |
Implements the accelerator plugin. |
src/ai/backend/accelerator/neuron/py.typed |
Marks the package as typed. |
src/ai/backend/accelerator/neuron/types.py |
Defines NeuronCore devices. |
src/ai/backend/manager/models/alembic/versions/a1c7e4b93f20_add_neuron_core_and_tt_n300_resource_slot_types.py |
Migrates slot definitions. |
tests/unit/accelerator/BUILD |
Adds accelerator test targets. |
tests/unit/accelerator/__init__.py |
Initializes the test package. |
tests/unit/accelerator/neuron/BUILD |
Configures Neuron unit tests. |
tests/unit/accelerator/neuron/__init__.py |
Initializes Neuron tests. |
tests/unit/accelerator/neuron/test_neuron_plugin.py |
Tests discovery, metrics, and container arguments. |
Review details
- Files reviewed: 15/18 changed files
- Comments generated: 7
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+375
to
+378
| # A device node carries *all* of its cores, so allocating a subset of a | ||
| # device's cores still mounts the whole device node -- the container can | ||
| # see sibling cores it was not allocated. NEURON_RT_VISIBLE_CORES below | ||
| # is what confines the runtime to the allocated ones. |
Member
There was a problem hiding this comment.
Please check this review about enforcing core isolation
Comment on lines
+389
to
+393
| if not _device_node_exists(host_path): | ||
| # Just skip mounting without raising an error, matching the | ||
| # other NPU plugins' behaviour for hot-removed devices. | ||
| log.warning("device node {} is missing; not mounting it", host_path) | ||
| continue |
Member
There was a problem hiding this comment.
Please check this review, about partial failure handling when assigning a device
| @@ -0,0 +1,40 @@ | |||
| python_sources( | |||
Comment on lines
+288
to
+294
| for dev in devices: | ||
| stats = core_stats[(dev.neuron_device_index, dev.core_index)] | ||
| node_path = dev.device_node_path | ||
| usage_by_node[node_path] = ( | ||
| usage_by_node.get(node_path, 0) + stats.device_mem_present | ||
| ) | ||
| capacity_by_node[node_path] = capacity_by_node.get(node_path, 0) + dev.memory_size |
Comment on lines
+544
to
+545
| async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: | ||
| pass |
Comment on lines
+17
to
+20
| try: | ||
| from ai.backend.agent.resources import get_resource_spec_from_container # type: ignore | ||
| except ImportError: | ||
| from ai.backend.agent.docker.resources import get_resource_spec_from_container |
Adds `backend.ai-accelerator-neuron`, registering the `neuron` entry point under `backendai_accelerator_v21`. The allocation unit is the NeuronCore (slot `neuron.core`, SlotTypes.COUNT), not the Neuron device. Three independent pieces of evidence point at the core: AWS's own allocation primitive is `NEURON_RT_VISIBLE_CORES` (with `NEURON_RT_NUM_CORES` deprecated in its favour); every runtime-varying metric the driver exposes is keyed per core while the device level carries only static capacity; and sysfs models cores as first-class nested objects. Device identity (`bdf`, `serial_number`, the device index) is kept as metadata on each core rather than as the allocation key, so cores aggregate upward to whole devices without a later data migration. Discovery shells out to `neuron-ls --json-output`, whose payload is a top-level array with an int `neuron_device`, a *string* `numa_node`, a null `connected_to` on single-device instances and a `memory_size` in bare bytes. The device capacity is split evenly across the device's cores. `neuron-ls` is resolved at /opt/aws/neuron/bin/neuron-ls before $PATH, because only the DLAMI profile scripts put that directory on PATH and the bare name is unresolvable under systemd or a container entrypoint. A missing tool or an unloaded driver logs the reason and leaves the plugin disabled instead of raising: agents load every installed accelerator plugin unless `allow-compute-plugins` is set, and the plugin loader calls `entrypoint.load()` with no exception handling, so an escaping FileNotFoundError would abort the agent's whole accelerator plugin load. `gather_node_measures` is implemented from the per-core sysfs memory counters, which are world-readable and populated with no workload attached, rather than stubbed. `neuron-monitor` is deliberately not used: it is a streaming collector with no one-shot mode, so consuming it would mean running a persistent vendor daemon behind a per-tick pull hook. Per-core utilization, which needs `neuron-monitor` plus an attached process, is therefore left out of v1. `generate_docker_args` mounts each allocated device node renumbered to /dev/neuron0..k-1, sets IPC_LOCK, host IPC and an unlimited memlock for the runtime's pinned-memory registration, and confines the runtime with `NEURON_RT_VISIBLE_CORES`. A device node carries all of its cores, so allocating a subset of a device's cores still mounts the whole node.
…types `agent_resources.slot_name` carries a hard FK to `resource_slot_types.slot_name`, and the manager upserts one `agent_resources` row per reported slot on every agent heartbeat. There is no runtime insert path into `resource_slot_types`, so a slot type that is not seeded makes the heartbeat of any agent reporting it fail on the FK. Register `neuron.core` in the example fixture and in a new migration. Also adds the missing `tt-n300.device` row as a drive-by fix: the Tenstorrent n300 plugin has reported that slot since it was added but it was never seeded, so a Tenstorrent agent takes the same FK violation today. This part is separable from the Neuron work if a reviewer would rather see it split out. `display_icon` uses asset names that actually exist under web/static/resources/icons: `aws` for Neuron, and `npu_generic` for tt-n300 (no `npu.svg` or `tenstorrent.svg` exists, despite the plugin declaring `npu` and web/static/resources/device_metadata.json declaring `tenstorrent`). The uuids are pinned to the ones the fixture assigns, following `8f21c46a0b73`, so an upgraded deployment ends up with the same slot identity a fresh install gets instead of a random one per database. `required` and `enabled` keep their server defaults, matching every other accelerator slot row. The downgrade only deletes rows that nothing references, so it does not abort on a deployment with a live agent, a historical allocation, or a model card / preset / deployment revision naming either slot. All five tables that carry an FK onto `resource_slot_types.slot_name` are guarded.
Creates `tests/unit/accelerator/`, which did not exist: no accelerator plugin in this repo had tests. `tests/unit/accelerator/neuron/` exercises discovery against the verbatim `neuron-ls --json-output` payload captured from a trn1.2xlarge, with the subprocess call monkeypatched so no hardware is needed. It pins the value shapes that are easy to get wrong (int `neuron_device` vs string `numa_node`, null `connected_to`, `memory_size` as bare bytes), that one device with `nc_count: 2` yields two core devices splitting the capacity, that a missing CLI and a driver-absent host each leave the plugin disabled without raising, and the container device renumbering and `NEURON_RT_VISIBLE_CORES` value. Making `pants test` green needed an explicit dependency in the test target: `tools/pants-plugins/accelerator_wheels` strips every `src/ai/backend/*` dependency from targets tagged `accelerator` so each wheel can build standalone. That stripping also keeps `ai.backend.agent` out of a test sandbox that only depends on the accelerator lib, so the plugin's base classes have to be named explicitly. This is why no accelerator plugin in the tree has had tests.
hoyajigi
force-pushed
the
feature/aws-neuron-accelerator
branch
from
September 12, 2026 12:35
ce41883 to
08db97f
Compare
github-actions
Bot
added
size:XL
fregataa
requested changes
Sep 14, 2026
Comment on lines
+52
to
+76
| def upgrade() -> None: | ||
| # Use exec_driver_sql to avoid sa.text() parsing the JSON colons as bind params. | ||
| # `required` and `enabled` keep their server defaults (false / true), matching | ||
| # every other accelerator slot row. | ||
| conn = op.get_bind() | ||
| conn.exec_driver_sql(""" | ||
| INSERT INTO resource_slot_types | ||
| (uuid, slot_name, slot_type, display_name, description, | ||
| display_unit, display_icon, number_format, rank) | ||
| VALUES | ||
| ('ef63fa11-609e-4b96-8f90-a08f97a5f04b'::uuid, | ||
| 'tt-n300.device','count','Tenstorrent n300 Device','Tenstorrent n300', | ||
| 'n300', 'npu_generic', '{"binary":false,"round_length":0}', 1500), | ||
| ('7b968b58-f7cc-472d-b191-d4b19f417efd'::uuid, | ||
| 'neuron.core','count','AWS Neuron Core','AWS Neuron NeuronCore', | ||
| 'Core', 'aws', '{"binary":false,"round_length":0}', 1600) | ||
| ON CONFLICT (slot_name) DO UPDATE SET | ||
| slot_type = EXCLUDED.slot_type, | ||
| display_name = EXCLUDED.display_name, | ||
| description = EXCLUDED.description, | ||
| display_unit = EXCLUDED.display_unit, | ||
| display_icon = EXCLUDED.display_icon, | ||
| number_format = EXCLUDED.number_format, | ||
| rank = EXCLUDED.rank | ||
| """) |
Member
There was a problem hiding this comment.
This migration seems to insert the new resource slot type for every site but I don't think it is installed globally
Comment on lines
+84
to
+97
| guards = "\n".join( | ||
| f" AND NOT EXISTS (" | ||
| f"SELECT 1 FROM {table} r WHERE r.slot_name = resource_slot_types.slot_name)" | ||
| for table in _referencing_tables | ||
| ) | ||
| conn = op.get_bind() | ||
| conn.execute( | ||
| sa.text(f""" | ||
| DELETE FROM resource_slot_types | ||
| WHERE slot_name = ANY(:names) | ||
| {guards} | ||
| """), | ||
| {"names": list(_added_slot_names)}, | ||
| ) |
Member
There was a problem hiding this comment.
Such deleting migration can be dangerous
Comment on lines
+17
to
+20
| try: | ||
| from ai.backend.agent.resources import get_resource_spec_from_container # type: ignore | ||
| except ImportError: | ||
| from ai.backend.agent.docker.resources import get_resource_spec_from_container |
Member
There was a problem hiding this comment.
Such import error handling is not required anymore
Comment on lines
+375
to
+378
| # A device node carries *all* of its cores, so allocating a subset of a | ||
| # device's cores still mounts the whole device node -- the container can | ||
| # see sibling cores it was not allocated. NEURON_RT_VISIBLE_CORES below | ||
| # is what confines the runtime to the allocated ones. |
Member
There was a problem hiding this comment.
Please check this review about enforcing core isolation
Comment on lines
+389
to
+393
| if not _device_node_exists(host_path): | ||
| # Just skip mounting without raising an error, matching the | ||
| # other NPU plugins' behaviour for hot-removed devices. | ||
| log.warning("device node {} is missing; not mounting it", host_path) | ||
| continue |
Member
There was a problem hiding this comment.
Please check this review, about partial failure handling when assigning a device
Five points from the review:
- The optional `ai.backend.agent.resources.get_resource_spec_from_container`
import no longer resolves anywhere in the tree, so the try/except was dead
and its `# type: ignore` unnecessary. Import from `agent.docker.resources`
directly, as the IPU plugin already does.
- A missing `/dev/neuron{N}` was skipped when building the container spec.
Unlike the other NPU plugins this is not safe here: the device has already
been counted when renumbering, so the NEURON_RT_VISIBLE_CORES indices assume
it is mounted, and the kernel would start with a reserved core absent and
the remaining ones addressed under wrong numbers. Raise ResourceError.
- Core isolation cannot be enforced through a device node that carries every
core of its device. Allocate with AllocationStrategy.FILL so a device is
split between sessions only when no free device is left, and state the
residual limitation in README.md rather than implying it is enforced.
- The downgrade deleted the seeded rows under NOT EXISTS guards. Five tables
carry an FK onto `resource_slot_types.slot_name`, which makes the delete
destructive rather than reversible; leave both rows in place instead.
- `update_plugin_config` was a no-op while config watching stayed on, so etcd
changes were acknowledged and dropped. Set `config_watch_enabled = False`,
matching the ROCm and CUDA plugins.
Also replaces `__spec__.name` with `__name__` for the logger, dropping the
last `# type: ignore` in the module.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Member
Author
|
리뷰 감사합니다. 지적하신 5건을 3200f55 에서 처리했습니다. 항목별로 정리합니다. 1. 마이그레이션이 슬롯 타입을 전역으로 삽입하는 문제 — 현행 유지, 근거 설명말씀하신 "전역 설치가 아닌데 왜 모든 사이트에 넣느냐"는 지적은,
행 하나가 추가되는 비용은 UI 카탈로그에 이름이 하나 더 있는 것뿐이고( 2. 삭제하는 downgrade 가 위험하다 — 동의, no-op 으로 변경동의합니다. 가드( 3. import 예외 처리 불필요 — 동의, 제거확인 결과 (같은 맥락에서 로거의 4. 코어 격리 강제 — 강제 불가함을 인정하고, 완화 + 명시지적이 맞습니다.
참고로 이 신뢰 모델 자체는 이 플러그인이 처음 도입하는 것이 아닙니다. 강제 격리를 이번 릴리스에 넣지 않은 이유: 방법이 두 가지인데 둘 다 이 PR 범위를 넘습니다.
1번을 원하시면 지금 바꾸는 편이 낫습니다(머지 후에는 데이터 마이그레이션 비용이 붙습니다). 어느 쪽을 택할지 결정해 주시면 반영하겠습니다. 결정 전까지는 core 단위 + 문서화된 한계로 두는 것을 제안합니다. 5. 디바이스 할당 부분 실패 처리 — 동의, 실패하도록 변경동의합니다. 그리고 이 플러그인에서는 다른 NPU 플러그인보다 더 나쁩니다. 노드를 건너뛰어도 할당된 Copilot 리뷰 중 반영하지 않은 항목
검증 범위 (중요)이 박스에는 Neuron 하드웨어도
실기 검증은 하지 못했습니다.
알렘빅 데이터 마이그레이션 검증도 로컬 DB 에서 수행하지 못했습니다. 다만 이번 변경은 |
`f4a1c9d20b73` (sync the seed roles) landed on main with the same `down_revision` as this PR's `a1c7e4b93f20`, leaving two alembic heads on the merge commit and failing `check-alembic-migrations`. Repoint this branch's own unmerged revision onto it, per the diverged-heads rule in `models/alembic/AGENTS.md` -- no merge migration, since main itself has a single head. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Member
Author
|
추가로, main 에 머지된
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
neuron.core, notneuron.deviceThe slot name becomes a PK in
resource_slot_typesand lands inagent_resources/resource_allocations, so changing it later needs a data migration. I chose the NeuronCore on three independent pieces of evidence, all from a realtrn1.2xlarge:NEURON_RT_VISIBLE_CORESis the knob;NEURON_RT_NUM_CORESis explicitly deprecated in its favour (pernccom-testshipped on the Neuron DLAMI).memory_usage,statusandother_infocounters all live underneuron_core{N}/.neuron0/neuron_core{0,1}/, each with its own full stat tree andarch_type(NCv2vs the device'sNDv2).Device identity (
bdf,serial_number, device index) is carried as metadata on each core device, not as the key. Core-granular slots aggregate upward to whole devices for free; device-granular could never be subdivided without a migration.If maintainers prefer
neuron.device, the change is mechanical but must happen before this merges.What this was built against
A real
trn1.2xlarge(1 Trainium device, 2 NeuronCores) on the Neuron DLAMI — driveraws-neuronx-dkms 2.26.5.0, toolsaws-neuronx-tools 2.28.23.0, PCI1d0f:7164. Capturedneuron-ls --json-output,neuron-monitor, and the complete sysfs tree, and used them as test fixtures.[{"neuron_device": 0, "bdf": "0000:00:1e.0", "cpu_affinity": "0-7", "numa_node": "-1", "connected_to": null, "nc_count": 2, "memory_size": 34359738368, "neuroncore_ids": [0, 1], "neuron_processes": []}]Shape traps handled:
neuron_deviceis an int butnuma_nodeis a string ("-1", clamped to 0 as tenstorrent does);connected_toisnullon single-device instances;memory_sizeis bare bytes (34359738368 = exactly 32 GiB, though the human table prints "32 GB").Notable implementation choices
rebellions, nottenstorrent/rngd. Those two catch onlyImportError.neuron-lsis a CLI, so a missing tool raisesFileNotFoundError— and since an agent loads all installed accelerator plugins unlessallow-compute-pluginsis set, that would abort agent startup on every host without Neuron tooling. This plugin resolves the path, logs, setsenabled = Falseand returns./opt/aws/neuron/bin/neuron-ls. That directory is injected intoPATHonly by DLAMI profile scripts; undersudo, systemd or a container entrypoint the bare name iscommand not found(exit 127). Measured, not assumed.gather_node_measuresis implemented from sysfs, not stubbed. Per-coredevice_mem/host_memare world-readable and populated with no workload attached, so no vendor daemon is needed.neuron-monitoris deliberately not run: it is a 5-second-cadence stream with no one-shot flag, no existing plugin runs a persistent vendor daemon, andgather_node_measures(ctx)is a per-tick pull. Per-core utilization needsneuron-monitorplus an attached runtime process, sogather_process_measuresreturns[]for now.NEURON_RT_VISIBLE_CORESis injected via"Env", which I verified is honoured:agent/utils.py:91-111update_nested_dictdeep-merges and extends lists,docker/agent.py:1206merges plugin args intocontainer_configwith no whitelist, and the agent sets its ownEnvat:1133before that merge, so plugin entries are appended rather than clobbered.tpu/plugin.py:166-171,ipu:390,cuda_open:327andmock:600already rely on this.kubernetes/agent.py:488-496hasgenerate_docker_argscommented out entirely with# TODO: add support for accelerator allocation, so every key is dropped there — not justEnv, and not specific to this plugin.resource_slot_typesseeding is mandatory, not cosmetic.agent_resources.slot_namehas a hard FK toresource_slot_types.slot_nameand the manager upserts one row per reported slot on every heartbeat, with no runtime insert path. Both fixtures are updated (kept byte-identical) plus a new alembic migration. Its downgrade is guarded against rows still referenced byagent_resources/resource_allocations, deliberately unlikeccf8ae5c90fe's unguarded delete.tt-n300.devicerow. I checked a live production manager DB — all 14 existing rows share onecreated_at, andtt-n300.deviceis absent, so a Tenstorrent agent takes an FK violation on heartbeat today. Happy to split this into its own PR if preferred.Validation
Everything below was run locally and passes:
pants test tests/unit/accelerator/neuron::pants check src/ai/backend/accelerator/neuron:: tests/unit/accelerator::(mypy)ruff check/ruff format --check(repo config)pants tailor --check, target graph,accelerator/wheeltagsGetting
pants testgreen needed a fix worth calling out:tools/pants-plugins/accelerator_wheelsstrips everysrc/ai/backend/*dependency from targets taggedacceleratorso each wheel builds standalone, and that stripping also keepsai.backend.agentout of a test sandbox that only depends on the accelerator lib. This is very likely why no accelerator plugin in the tree has had tests. The test target now names the agent modules explicitly.The migration was executed, not just written
Against PostgreSQL 15.19 with the full 73-table schema built from the models:
downgrade a1c7e4b93f20 -> 3b6297b1bd75upgrade 3b6297b1bd75 -> a1c7e4b93f20INSERT INTO agent_resources ... 'neuron.core'(what a heartbeat does)fk_agent_resources_slot_name_resource_slot_typesdowngradewhileagent_resourcesstill referencesneuron.coreneuron.corepreserved,tt-n300.deviceremoveddowngradeonce nothing references itThe fourth row is the point: without this migration an agent reporting
neuron.corecannot complete a heartbeat. The same is true oftt-n300.devicetoday, which is why it is seeded here as a drive-by.What is not verified — read before merging
trn1.32xlargeoutput. Specifically, I assume the runtime numbers container-local cores as(position of device among visible devices) * nc_count + core index. Untested on hardware.connected_to/ NeuronLink topology is parsed and stored but used for nothing, so it does not influence allocation.device_mem/*read 0 throughout, so the non-zero path is exercised only against synthetic values.NEURON_RT_VISIBLE_CORESactually confining the runtime. I verified the variable reaches the container by reading the merge path; I did not run a Neuron workload to confirm the runtime honours it when the whole device node is mounted (each/dev/neuron{N}carries all of its cores).agent/agent.py:2752-2754skipsmount_krunnerwhenrestarting, soapply_accelerator_allocationnever runs and a plugin's wholegenerate_docker_argsoutput appears to be dropped on kernel restart. Affects all accelerator plugins equally.