Conversation
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.
| # 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. |
There was a problem hiding this comment.
Please check this review about enforcing core isolation
| 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 |
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( | |||
| 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 |
| async def update_plugin_config(self, new_plugin_config: Mapping[str, Any]) -> None: | ||
| pass |
| 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.
ce41883 to
08db97f
Compare
| 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 | ||
| """) |
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
| 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)}, | ||
| ) |
There was a problem hiding this comment.
Such deleting migration can be dangerous
| 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 |
There was a problem hiding this comment.
Such import error handling is not required anymore
| # 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. |
There was a problem hiding this comment.
Please check this review about enforcing core isolation
| 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 |
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>
|
리뷰 감사합니다. 지적하신 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>
|
추가로, main 에 머지된
|
|
Thanks for the detailed writeup — items 2 through 5 all look good to me. Switching the alloc map to I'd like to come back to item 1, though. I think your conclusion (keep the seed as it is) is correct, but a couple of the supporting facts don't hold up, and I'd rather we land this on the argument that survives scrutiny. "There is no runtime insert path into
|
|
A second pass, this time on the code itself rather than the design discussion. Nothing here changes my view of the 1. The PR description no longer matches the migrationThe description still says the downgrade is
and the validation table still carries these two rows:
Since 3200f55 turned 2.
|
Four mechanical fixes from the review on the plugin itself: - gather_container_measures opened a Docker client inside the per-container loop, building and tearing down one aiohttp session per container per stat tick. Every other accelerator plugin in the tree opens it once outside the loop; do the same and keep only the DockerError handling inside. - restore_from_container guarded apply_allocation with hasattr and kept a fallback for agents predating lablup/backend.ai-agent#180. apply_allocation is an abstractmethod on AbstractAllocMap, so the guard can never be False, and this plugin is new enough that it will never meet such an agent. The guard came across from the Tenstorrent plugin, whose history justifies it. - The container-local core index computed position * nc_count, which equals "cores preceding this device" only when every mounted device carries the same nc_count. Accumulate the preceding devices' nc_count instead, which is correct unconditionally. No instance type ships a mixed nc_count today, so this is not a live bug, but the renumbering is precisely the piece that could not be verified on hardware and this removes one assumption from it. Covered by a new test that fails with the old formula. - test_declared_icon_exists_in_the_repo always skipped: src/ai/backend/web static assets are not a dependency of this test target, so the directory is never materialised into the Pants sandbox and the guard fires every run. Reaching into the web component also crosses a boundary the accelerator wheels exist to avoid. Removed; icon-name validation belongs somewhere that owns the icon set and can check every plugin's display_icon at once. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
리뷰 두 편 모두 감사합니다. 코드 지적 5건 중 4건을 먼저, 제가 틀린 부분1. "
|
main advanced to a7d0f5b3c841 (drop_vfolder_invitations) while this branch was under review, so a1c7e4b93f20's down_revision of f4a1c9d20b73 was no longer the tip and check-alembic-migrations saw two heads. main itself has a single head, so per models/alembic/AGENTS.md this is the "your migration diverged from main" case: repoint this branch's own unmerged revision at the current head, no merge migration. origin/main was brought in with a merge rather than a rebase to keep the review anchors attached. check-multiple-alembic-heads.py now reports a1c7e4b93f20 as the sole head. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Thanks — I went through Since the thread has grown, here is a clean split of what is left: what has to happen inside this PR, and what should leave this PR and be tracked on its own. A. Must happen in this PRA1. The migration has diverged from main again
Running the repo's own check against main plus this revision: Main on its own resolves to a single head, so by This is the third time this has happened while the PR sits open, and it will happen a fourth time if the PR stays open much longer. Worth doing as the last step before merge rather than now, or worth merging promptly after fixing it. A2. Split
|
| check | result |
|---|---|
pytest tests/unit/accelerator/neuron/ |
33 passed |
pytest tests/unit/plugin/test_accelerator_entrypoints.py (landed on main after this branch was written) |
3 passed — the neuron entry point resolves |
reverted the core-index formula to position * nc_count and re-ran |
fails with CORES=0,3 vs expected 0,2 — the new test is a real regression pin, and 3 is an index that doesn't exist in a container holding 3 cores |
migration values vs fixtures/manager/example-resource-slot-types.json |
all 8 fields match for both rows; no duplicate rank, uuid or slot name |
git diff --stat origin/main... |
17 files, all additive, nothing unrelated touched |
On the individual fixes:
- Docker client hoisted — correct, and no behaviour change at the edges: a failure constructing the client raises
ValueError, which the oldexcept DockerErrordidn't catch either, so the method raised then and raises now. hasattrfallback removed — correct.device_slotsis declared onAbstractAllocMapitself (agent/alloc_map.py:78), so dropping the guard doesn't cost you the type check.- Cumulative core base — correct, and the mixed-
nc_countfixture is the right shape to pin it. - Icon test removed — correct, along with the now-unused
pathlibimport.
Nice work on these. The plugin itself I'm happy with; what's left is A1–A3 plus deciding between the two options in A2.
The tt-n300.device row is a schema fix for the Tenstorrent n300 plugin and should be backportable on its own, which models/alembic/README.md principle 1 rules out while it shares a revision with the Neuron feature seed. - d0a201e9be45 seeds only tt-n300.device and chains onto main's current head a91c4e7d0b35 (main moved again while this branch was open). - a1c7e4b93f20 seeds only neuron.core and now revises d0a201e9be45; its file is renamed to match what it adds. Both keep the previous upsert and the no-op downgrade. The neuron.core docstring drops the retracted "no runtime insert path" claim and states the zero-valued slot reason instead. The fixture keeps both rows unchanged. check-multiple-alembic-heads.py reports a1c7e4b93f20 as the sole head.
|
@fregataa Thanks for the clean A/B split. A1–A3 are done in A1. Re-chained onto main's headMerged Point taken on timing. If main moves again before merge, I'll redo the re-chain as the last step. A2. Option 2: two revisions in this branch
This time I ran both revisions against a real database: PostgreSQL 16.3 (the halfstack image), with the schema built by
A3. PR descriptionAll three are fixed:
The description also shows the two-revision chain. I fixed one more stale claim while I was there: it said "both fixtures", but only one fixture file is involved. A4No change. Checks at
|
|
@fregataa B1 is filed as BA-8140 ("Heartbeat registers unknown slots in etcd but not in |
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
CLI-absence handling follows
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.Absolute path
/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_name, and the manager upserts oneagent_resourcesrow per reported slot on every heartbeat. A create path for slot types exists (CreateResourceSlotTypeAction, GraphQL, REST v2,./bai resource-slot slot-type create), but the heartbeat path never calls it, so an unseeded slot fails the heartbeat on the FK with nothing pointing at the cause. The same heartbeat also callslegacy_etcd_config_loader.update_resource_slots(), which does register unknown slot names in etcd; the seed exists because the normalized table has no equivalent. Reconciling the two registries is out of scope here and is tracked in BA-8140.resource_slot_to_quantitiespreserves zero-valued slots, so a host with no Neuron hardware — where the plugin has already disabled itself — still reportsneuron.core: 0and still hits the same foreign key. The no-hardware case is the common one, not the exception.fixtures/manager/example-resource-slot-types.jsoncarries the same rows with the same uuids, so a fresh install and an upgraded deployment end up with the same slot identity.enabledkeeps its server default (true), like all 14 existing rows. Whether unused accelerator slots should ship disabled is a decision about the whole catalogue, not this row.downgrade()are a deliberate no-op: five tables carry an FK ontoresource_slot_types.slot_name, so any guard would make the result depend on what happened to reference the rows at that instant, which is not a reproducible downgrade. Leaving the rows behind is harmless — older managers never read them — and removing them is an operator decision, not a schema one.The migration is split into two revisions so the Tenstorrent fix can be backported on its own.
models/alembic/README.mdprinciple 1 limits backport migrations to schema fixes, so the missingtt-n300.devicerow gets its own revision ahead of the Neuron feature revision:I checked a live production manager DB — all 14 existing rows share one
created_atandtt-n300.deviceis absent, so a Tenstorrent agent takes an FK violation on heartbeat today. Whoever backports cherry-picksd0a201e9be45alone; itsINSERT ... ON CONFLICT (slot_name) DO UPDATEis idempotent.Validation
Everything below was run locally and passes:
pants test tests/unit/accelerator/neuron::pants test tests/unit/plugin/test_accelerator_entrypoints.pyneuronentry point resolvespython3 scripts/check-multiple-alembic-heads.py(branch merged with main at35e8009bab)a1c7e4b93f20pants fmt lint checkon both migration filespants 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 migrations were executed, not just written
At
4a5593641d, against PostgreSQL 16.3 (postgres:16.3-alpine, the halfstack image). The schema was built with./backend.ai mgr schema oneshotas CI does, then stamped back toa91c4e7d0b35so that neither row existed. A test agent row was inserted to exercise the FK the heartbeat hits.a91c4e7d0b35:INSERT INTO agent_resourceswithneuron.core/tt-n300.device, capacity0fk_agent_resources_slot_name_resource_slot_typesupgrade a91c4e7d0b35 -> d0a201e9be45tt-n300.deviceinserted; its insert is now accepted,neuron.coreis still rejectedupgrade d0a201e9be45 -> a1c7e4b93f20neuron.coreinserted; all fields and uuids match the fixture,required = false,enabled = true; the capacity-0insert is accepteddowngrade a1c7e4b93f20 -> d0a201e9be45 -> a91c4e7d0b35agent_resourcesrows left in placeupgradeto head after overwriting both rows'display_nameandrankWithout these rows, an agent reporting
neuron.coreortt-n300.device— including a zero-capacity report from a host without the hardware — cannot complete a heartbeat.What is not verified — read before merging
trn1.32xlargeoutput. Specifically, I assume the runtime numbers container-local cores in the order of the visible devices, so a core's container-local index is the number of cores carried by the preceding mounted devices plus its own index within its device. 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.