Skip to content

feat(accelerator): add an AWS Neuron (Trainium/Inferentia) plugin - #14580

Open
hoyajigi wants to merge 11 commits into
mainfrom
feature/aws-neuron-accelerator
Open

hoyajigi wants to merge 11 commits into
mainfrom
feature/aws-neuron-accelerator

Conversation

@hoyajigi

@hoyajigi hoyajigi commented Sep 12, 2026 •

Copy link
Copy Markdown
Member

Hardware verification status: the hardware captures below were taken on a
trn1.2xlarge probe instance. No Neuron workload has been run at any point,
and the machine the later revisions were written on has no Neuron device or
aws-neuronx-dkms driver.
The migrations were re-run at 4a5593641d
against a PostgreSQL 16.3 scratch database (the halfstack image); see
"The migrations were executed" below. What remains unverified is listed at the
bottom — reviewers with a trn1.32xlarge or inf2.24xlarge can close the
hardware gap quickly.

⚠️ One irreversible decision needs maintainer sign-off: the slot is neuron.core, not neuron.device

The slot name becomes a PK in resource_slot_types and lands in agent_resources / resource_allocations, so changing it later needs a data migration. I chose the NeuronCore on three independent pieces of evidence, all from a real trn1.2xlarge:

  1. AWS's own allocation primitive is the core. NEURON_RT_VISIBLE_CORES is the knob; NEURON_RT_NUM_CORES is explicitly deprecated in its favour (per nccom-test shipped on the Neuron DLAMI).
  2. Every runtime-varying metric is keyed per core. In the captured sysfs tree the device level exposes only static capacity plus driver-owned host memory; memory_usage, status and other_info counters all live under neuron_core{N}/.
  3. sysfs models cores as first-class objects — real nested directories neuron0/neuron_core{0,1}/, each with its own full stat tree and arch_type (NCv2 vs the device's NDv2).

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 — driver aws-neuronx-dkms 2.26.5.0, tools aws-neuronx-tools 2.28.23.0, PCI 1d0f:7164. Captured neuron-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_device is an int but numa_node is a string ("-1", clamped to 0 as tenstorrent does); connected_to is null on single-device instances; memory_size is bare bytes (34359738368 = exactly 32 GiB, though the human table prints "32 GB").

Notable implementation choices

  • CLI-absence handling follows rebellions, not tenstorrent/rngd. Those two catch only ImportError. neuron-ls is a CLI, so a missing tool raises FileNotFoundError — and since an agent loads all installed accelerator plugins unless allow-compute-plugins is set, that would abort agent startup on every host without Neuron tooling. This plugin resolves the path, logs, sets enabled = False and returns.

  • Absolute path /opt/aws/neuron/bin/neuron-ls. That directory is injected into PATH only by DLAMI profile scripts; under sudo, systemd or a container entrypoint the bare name is command not found (exit 127). Measured, not assumed.

  • gather_node_measures is implemented from sysfs, not stubbed. Per-core device_mem/host_mem are world-readable and populated with no workload attached, so no vendor daemon is needed. neuron-monitor is deliberately not run: it is a 5-second-cadence stream with no one-shot flag, no existing plugin runs a persistent vendor daemon, and gather_node_measures(ctx) is a per-tick pull. Per-core utilization needs neuron-monitor plus an attached runtime process, so gather_process_measures returns [] for now.

  • NEURON_RT_VISIBLE_CORES is injected via "Env", which I verified is honoured: agent/utils.py:91-111 update_nested_dict deep-merges and extends lists, docker/agent.py:1206 merges plugin args into container_config with no whitelist, and the agent sets its own Env at :1133 before that merge, so plugin entries are appended rather than clobbered. tpu/plugin.py:166-171, ipu:390, cuda_open:327 and mock:600 already rely on this.

    • Caveat: the Kubernetes backend drops it. kubernetes/agent.py:488-496 has generate_docker_args commented out entirely with # TODO: add support for accelerator allocation, so every key is dropped there — not just Env, and not specific to this plugin.
  • resource_slot_types seeding is mandatory, not cosmetic. agent_resources.slot_name has a hard FK to resource_slot_types.slot_name, and the manager upserts one agent_resources row 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 calls legacy_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.

    • Hosts without Neuron hardware need the row too. resource_slot_to_quantities preserves zero-valued slots, so a host with no Neuron hardware — where the plugin has already disabled itself — still reports neuron.core: 0 and still hits the same foreign key. The no-hardware case is the common one, not the exception.
    • fixtures/manager/example-resource-slot-types.json carries the same rows with the same uuids, so a fresh install and an upgraded deployment end up with the same slot identity.
    • enabled keeps 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.
    • Both migrations' downgrade() are a deliberate no-op: five tables carry an FK onto resource_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.md principle 1 limits backport migrations to schema fixes, so the missing tt-n300.device row gets its own revision ahead of the Neuron feature revision:

    a91c4e7d0b35  (main head: backfill_missing_model_store_projects)
          ↓
    d0a201e9be45  seed tt-n300.device   ← schema fix, backportable on its own
          ↓
    a1c7e4b93f20  seed neuron.core      ← feature
    

    I checked a live production manager DB — all 14 existing rows share one created_at and tt-n300.device is absent, so a Tenstorrent agent takes an FK violation on heartbeat today. Whoever backports cherry-picks d0a201e9be45 alone; its INSERT ... ON CONFLICT (slot_name) DO UPDATE is idempotent.

Validation

Everything below was run locally and passes:

check result
pants test tests/unit/accelerator/neuron:: ✅ pass
pants test tests/unit/plugin/test_accelerator_entrypoints.py ✅ pass — the neuron entry point resolves
python3 scripts/check-multiple-alembic-heads.py (branch merged with main at 35e8009bab) ✅ single head a1c7e4b93f20
pants fmt lint check on both migration files ✅ pass
pants check src/ai/backend/accelerator/neuron:: tests/unit/accelerator:: (mypy) ✅ pass — caught and fixed 2 real type errors
ruff check / ruff format --check (repo config) ✅ clean
pants tailor --check, target graph, accelerator/wheel tags ✅ correct
both alembic migrations against real PostgreSQL 16.3 ✅ see below

Getting pants test green needed a fix worth calling out: tools/pants-plugins/accelerator_wheels strips every src/ai/backend/* dependency from targets tagged accelerator so each wheel builds standalone, and that stripping also keeps ai.backend.agent out 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 oneshot as CI does, then stamped back to a91c4e7d0b35 so that neither row existed. A test agent row was inserted to exercise the FK the heartbeat hits.

step result
at a91c4e7d0b35: INSERT INTO agent_resources with neuron.core / tt-n300.device, capacity 0 ❌ both rejected by fk_agent_resources_slot_name_resource_slot_types
upgrade a91c4e7d0b35 -> d0a201e9be45 ✅ only tt-n300.device inserted; its insert is now accepted, neuron.core is still rejected
upgrade d0a201e9be45 -> a1c7e4b93f20 ✅ neuron.core inserted; all fields and uuids match the fixture, required = false, enabled = true; the capacity-0 insert is accepted
downgrade a1c7e4b93f20 -> d0a201e9be45 -> a91c4e7d0b35 ✅ both no-ops; rows and dependent agent_resources rows left in place
re-upgrade to head after overwriting both rows' display_name and rank ✅ no conflict error; both rows restored to the seeded values, uuids unchanged

Without these rows, an agent reporting neuron.core or tt-n300.device — including a zero-capacity report from a host without the hardware — cannot complete a heartbeat.

What is not verified — read before merging

  • Multi-device renumbering. The probe had exactly one device. The two-device fixture in the tests is synthesized, not captured, so it pins my renumbering logic rather than AWS's real trn1.32xlarge output. 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.
  • Live per-core utilization and memory. No workload was ever run; device_mem/* read 0 throughout, so the non-zero path is exercised only against synthetic values.
  • NEURON_RT_VISIBLE_CORES actually 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).
  • Pre-existing, not introduced here: agent/agent.py:2752-2754 skips mount_krunner when restarting, so apply_accelerator_allocation never runs and a plugin's whole generate_docker_args output appears to be dropped on kernel restart. Affects all accelerator plugins equally.

@hoyajigi
hoyajigi marked this pull request as ready for review September 12, 2026 12:09
@hoyajigi
hoyajigi requested a review from a team as a code owner September 12, 2026 12:09
Copilot AI balanced review requested due to automatic review settings September 12, 2026 12:09

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
Comment thread src/ai/backend/accelerator/neuron/plugin.py Outdated
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
hoyajigi force-pushed the feature/aws-neuron-accelerator branch from ce41883 to 08db97f Compare September 12, 2026 12:35
@github-actions github-actions Bot added size:XL 500~ LoC comp:manager Related to Manager component require:db-migration Automatically set when alembic migrations are added or updated labels Sep 12, 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
""")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)},
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@hoyajigi

Copy link
Copy Markdown
Member Author

리뷰 감사합니다. 지적하신 5건을 3200f55 에서 처리했습니다. 항목별로 정리합니다.


1. 마이그레이션이 슬롯 타입을 전역으로 삽입하는 문제 — 현행 유지, 근거 설명

말씀하신 "전역 설치가 아닌데 왜 모든 사이트에 넣느냐"는 지적은, resource_slot_types 가 설치된 하드웨어의 기록이 아니라 클러스터가 이름을 알고 렌더링할 수 있는 슬롯의 카탈로그라는 점에서 갈립니다.

  • 선례: ccf8ae5c90fe (add_device_metadata_to_resource_slot_types) 는 ipu.device, atom.device, atom-plus.device, atom-max.device, gaudi2.device, warboy.device, rngd.device, hyperaccel-lpu.device 8종을 해당 플러그인 설치 여부와 무관하게 모든 배포에 심습니다. neuron.core 만 다르게 취급할 근거가 없습니다.
  • 기능적 필요: agent_resources.slot_name 에 fk_agent_resources_slot_name_resource_slot_types 하드 FK 가 있고, 매니저는 모든 agent heartbeat 마다 보고된 슬롯별로 agent_resources 행을 upsert 합니다. 런타임에 resource_slot_types 로 insert 하는 경로는 없습니다. 즉 시드되지 않은 슬롯을 보고하는 에이전트는 heartbeat 자체가 FK 위반으로 실패합니다.
  • tt-n300.device 를 같이 넣은 것도 같은 이유입니다. Tenstorrent n300 플러그인은 추가된 시점부터 그 슬롯을 보고해 왔는데 시드된 적이 없어서, 지금도 Tenstorrent 에이전트는 이 FK 에 걸립니다.

행 하나가 추가되는 비용은 UI 카탈로그에 이름이 하나 더 있는 것뿐이고(enabled 는 서버 기본값을 그대로 씁니다), 빠졌을 때의 비용은 에이전트 기동 실패입니다. 이 비대칭 때문에 현행을 유지했습니다. 그래도 "설치된 플러그인만 시드" 정책을 원하시면, 그건 이 PR 하나가 아니라 위 8종 + tt-n300 을 포함한 별도 정리 PR 이 맞다고 보고, 그 경우 런타임 등록 경로부터 만들어야 합니다. 방향을 정해주시면 그쪽으로 따로 올리겠습니다.

2. 삭제하는 downgrade 가 위험하다 — 동의, no-op 으로 변경

동의합니다. 가드(NOT EXISTS)를 붙여도 결국 그 순간 무엇이 참조 중이었는지에 따라 결과가 달라지는, 재현되지 않는 downgrade 였습니다. downgrade() 를 명시적 no-op 으로 바꾸고 사유를 주석과 모듈 docstring 에 남겼습니다. 구버전 매니저는 이 두 행을 조회하지 않으므로 남겨도 무해하고, 실제 제거는 스키마 결정이 아니라 운영자 결정입니다. 이에 따라 _added_slot_names / _referencing_tables 상수와 sqlalchemy import 도 제거했습니다.

3. import 예외 처리 불필요 — 동의, 제거

확인 결과 get_resource_spec_from_container 는 트리 전체에서 ai.backend.agent.docker.resources 와 ai.backend.agent.kubernetes.resources 에만 정의돼 있고 ai.backend.agent.resources 에는 없습니다. 즉 try 절은 항상 실패하는 죽은 코드였고, 거기 붙은 # type: ignore 도 불필요했습니다. ai.backend.accelerator.ipu.plugin 이 이미 쓰는 방식대로 직접 import 로 바꿨습니다.

(같은 맥락에서 로거의 __spec__.name # type: ignore 도 __name__ 으로 바꿔 이 모듈의 마지막 # type: ignore 를 없앴습니다.)

4. 코어 격리 강제 — 강제 불가함을 인정하고, 완화 + 명시

지적이 맞습니다. /dev/neuronN 은 그 디바이스의 모든 코어를 담고 있고 Docker 가 넘길 수 있는 최소 단위가 노드이므로, 코어 단위 격리는 NEURON_RT_VISIBLE_CORES 라는 워크로드가 덮어쓸 수 있는 프로세스 환경변수에 의존합니다. 강제가 아니라 협조입니다. 숨기지 않고 다음과 같이 처리했습니다.

  • 완화: alloc map 의 할당 전략을 기본 EVENLY 에서 AllocationStrategy.FILL 로 바꿨습니다. 기본값은 코어를 디바이스에 분산시켜 노드 공유를 최대화하는 방향이었습니다. FILL 이면 한 디바이스를 채운 뒤 다음으로 넘어가므로, 완전히 빈 디바이스가 없을 때만 분할됩니다. 부수적으로 한 세션의 코어가 한 디바이스에 모이므로 collectives 가 NeuronLink 를 건너지 않습니다.
  • 명시: README.md 에 ## Isolation 절을 추가해 (a) 격리가 강제가 아니라는 점, (b) gather_container_measures 의 디바이스 노드 단위 귀속 때문에 한 디바이스를 나눠 쓰는 두 세션이 각각 디바이스 전체 사용량으로 보고된다는 점(Tenstorrent·Rebellions 플러그인과 동일)을 적었습니다.

참고로 이 신뢰 모델 자체는 이 플러그인이 처음 도입하는 것이 아닙니다. cuda.shares 도 /dev/nvidia{N} 를 여러 세션이 공유하며 소프트 강제에 의존합니다.

강제 격리를 이번 릴리스에 넣지 않은 이유: 방법이 두 가지인데 둘 다 이 PR 범위를 넘습니다.

  1. 슬롯을 neuron.device (디바이스 단위) 로 바꾸기 — 슬롯 이름이 resource_slot_types 의 PK 이고 agent_resources/resource_allocations 가 참조하므로 나중에 세분화가 불가능해집니다(README 의 설계 근거 절 참고). Trainium1 기준 최소 할당 단위가 2코어가 됩니다.
  2. 디바이스를 통째로 할당하되 요청한 코어 수만 과금 — DiscretePropertyAllocMap.allocate() 의 "요청량만큼만 할당한다" 계약을 깨므로 매니저 스케줄러 쪽 지원이 필요합니다.

1번을 원하시면 지금 바꾸는 편이 낫습니다(머지 후에는 데이터 마이그레이션 비용이 붙습니다). 어느 쪽을 택할지 결정해 주시면 반영하겠습니다. 결정 전까지는 core 단위 + 문서화된 한계로 두는 것을 제안합니다.

5. 디바이스 할당 부분 실패 처리 — 동의, 실패하도록 변경

동의합니다. 그리고 이 플러그인에서는 다른 NPU 플러그인보다 더 나쁩니다. 노드를 건너뛰어도 container_index_of 는 이미 그 디바이스를 번호에 포함시킨 뒤이므로, 남은 코어의 NEURON_RT_VISIBLE_CORES 컨테이너-로컬 인덱스가 런타임이 실제로 보는 번호와 어긋납니다. 즉 "가속기 없이 시작" 에 더해 "남은 가속기도 잘못된 번호로 지정" 이 됩니다.

할당된 /dev/neuron{N} 이 없으면 ai.backend.agent.errors.resources.ResourceError 를 던져 컨테이너 생성을 실패시키도록 바꿨습니다(RuntimeError 대신 BackendAIError 계열을 쓰라는 AGENTS.md 규칙에 맞춤). 기존 테스트 test_missing_device_node_is_skipped_not_fatal 은 test_missing_device_node_fails_container_creation 으로 교체했습니다.


Copilot 리뷰 중 반영하지 않은 항목

  • src/ 아래 BUILD 파일 금지: BUILDING.md 의 해당 규칙은 "새 파이썬 모듈/패키지"를 대상으로 하고, "최상위 컴포넌트 BUILD 파일만 존재한다" 고 명시합니다. src/ai/backend/accelerator/neuron/BUILD 는 cuda_open, rocm, tenstorrent, furiosa 등과 같은 레벨의 배포 단위 최상위 BUILD 입니다. 이게 없으면 python_distribution 과 backendai_accelerator_v21 엔트리포인트가 없어서 휠 자체가 만들어지지 않습니다. 유지했습니다.
  • 컨테이너 메트릭 중복 집계: 실재하는 한계가 맞지만 Tenstorrent·Rebellions 플러그인과 동일한 동작이고, Docker inspect 가 노드 단위 정보만 주는 데서 오는 제약입니다. 컨테이너별 neuron.core 할당을 읽어오려면 stat 수집 경로에 resource spec 조회를 추가해야 해서, 세 플러그인을 함께 고치는 별도 PR 이 맞다고 봅니다. 우선 README 에 한계로 명시했습니다.

검증 범위 (중요)

이 박스에는 Neuron 하드웨어도 aws-neuronx-dkms 드라이버도 없습니다. 실행한 것은 다음뿐입니다.

  • pants fmt lint — 통과
  • pants check (mypy, CPython 3.13.7) — Success: no issues found in 6 source files
  • pants test tests/unit/accelerator/neuron:: — 통과 (모두 neuron-ls JSON 과 sysfs 를 스텁으로 대체한 테스트입니다)

실기 검증은 하지 못했습니다. trn1.2xlarge 이상에서 다음 두 가지를 확인해 주시면 좋겠습니다.

  1. 한 디바이스의 코어 일부만 할당한 세션에서 echo $NEURON_RT_VISIBLE_CORES 와 컨테이너 내 /dev/neuron* 목록이 일치하는지 (FILL 전환으로 할당 순서가 바뀌었습니다).
  2. neuron.core 를 보고하는 에이전트의 heartbeat 가 마이그레이션 적용 후 FK 위반 없이 통과하는지.

알렘빅 데이터 마이그레이션 검증도 로컬 DB 에서 수행하지 못했습니다. 다만 이번 변경은 downgrade() 를 no-op 으로 만든 것뿐이고 upgrade() 의 SQL 은 손대지 않았습니다.

@hoyajigi
hoyajigi requested a review from fregataa September 15, 2026 00:16
hoyajigi and others added 2 commits September 15, 2026 00:17
`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>
@hoyajigi

Copy link
Copy Markdown
Member Author

추가로, check-alembic-migrations 가 실패해서 한 건 더 처리했습니다 (a441b63).

main 에 머지된 f4a1c9d20b73 (sync the seed roles) 가 이 PR 의 a1c7e4b93f20 과 같은 down_revision(c58b0d3a9e14) 을 쓰고 있어서, PR 머지 커밋 기준으로 alembic head 가 두 개가 되었습니다. models/alembic/AGENTS.md 의 diverged-heads 규칙대로 — main 자체는 head 가 하나이므로 merge migration 없이 — 이 브랜치의 미머지 리비전 쪽 down_revision 을 f4a1c9d20b73 으로 옮겼습니다. origin/main 은 rebase 가 아니라 merge 로 가져왔습니다(리뷰 앵커 보존).

python scripts/check-multiple-alembic-heads.py 로컬 실행 결과: Detected head revisions: a1c7e4b93f20 (단일 head).

@fregataa

Copy link
Copy Markdown
Member

Thanks for the detailed writeup — items 2 through 5 all look good to me. Switching the alloc map to FILL and adding an explicit ## Isolation section is exactly the right way to handle a limitation you can't actually remove, and I appreciate that you didn't paper over it.

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 resource_slot_types"

There is one, and it's wired across every layer:

  • src/ai/backend/manager/services/resource_slot/actions/create.py — CreateResourceSlotTypeAction
  • src/ai/backend/manager/api/gql/resource_slot/resolver.py (GraphQL)
  • src/ai/backend/manager/api/rest/v2/resource_slot/handler.py (REST v2)
  • src/ai/backend/client/cli/v2/resource_slot/slot_type.py:99 — ./bai resource-slot slot-type create

The narrower statement is the true one: the heartbeat path never creates the slot type for you.

That distinction matters a lot here, because "have the operator register the slot type when they install the plugin" is precisely the alternative a reviewer would reach for. Saying the path doesn't exist reads as arguing around the objection rather than answering it — which I don't think was the intent, but it's how it lands.

The accurate version is actually stronger for your case: the API exists, nothing calls it automatically, and until someone remembers to call it the agent's heartbeat fails on the FK with nothing pointing at the cause.

The argument I think you actually want

Take a look at the same function that performs the agent_resources upsert:

# src/ai/backend/manager/repositories/agent/repository.py:141-144
if upsert_result.need_resource_slot_update:
    await self._config_provider.legacy_etcd_config_loader.update_resource_slots(
        upsert_data.resource_info.slot_key_and_units
    )

and what that call does:

# src/ai/backend/manager/config/loader/legacy_etcd_loader.py:93-99
for k, v in slot_key_and_units.items():
    if k not in known_slots or v != known_slots[k]:
        updates[f"config/resource_slots/{k}"] = v.value

So a single heartbeat drives two slot registries side by side: the legacy etcd one self-registers any slot name it hasn't seen before, while the normalized table does not — it only enforces a hard FK.

That reframes the whole question. The seed isn't a catalogue policy you chose; it's the only mechanism that currently exists, because the slot-normalization work replaced a dynamic registry with a static one and didn't carry the auto-registration across. That's a much easier position to defend, and it moves the design question to where it belongs (restore auto-registration, or keep the catalogue curated) instead of onto this PR.

One more fact in your favour that you left out

src/ai/backend/manager/repositories/resource_slot/types.py:17 — resource_slot_to_quantities explicitly "Preserves zero-valued slots".

So a host with no Neuron hardware, where the plugin has already disabled itself, still reports neuron.core: 0, and that zero-capacity slot goes straight into the same FK. "Sites without Neuron don't need this row" isn't just weak, it's false — and the no-hardware case is the common one. Worth stating outright; it answers the objection more directly than the precedent argument does.

Small correction on the cost estimate

the cost of one extra row is just one more name in the UI catalogue (enabled keeps its server default)

src/ai/backend/manager/models/resource_slot/row.py:80-86 — enabled's server default is sa.true(). Keeping the default means the slot ships enabled on every deployment rather than hidden. If we want the mitigation that parenthetical implies, we'd have to seed enabled = false and let operators switch it on. I'm fine either way, but let's not describe it as a no-op.

On the "seed only installed plugins" offer

I don't think that's what's being asked, and I don't think a migration can do it — the database has no way to know what a future agent will have installed. The two questions worth separating are:

  1. whether this belongs in the seed fixtures / bootstrap rather than a schema migration (you already touch fixtures/manager/example-resource-slot-types.json, so it lives in both places today), and
  2. whether registration should happen at runtime, the way the etcd registry already does.

A cleanup PR for the eight existing rows answers neither. My preference: keep this PR's neuron.core seed as-is, and open a separate issue for (2).

Request

Could you split tt-n300.device into its own fix PR? It's a genuine bug fix that should be backportable, and src/ai/backend/manager/models/alembic/README.md principle 1 is "Backport migrations must contain schema fixes only ... Feature-level schema changes are never backported" — bundling it with the Neuron feature revision blocks that.


Separately, thank you for the amount of evidence behind the neuron.core decision. The three-way argument (the runtime's own allocation primitive, every runtime-varying metric being per-core, and sysfs modelling cores as first-class objects) is convincing, and I'd approve that choice as written.

@fregataa

Copy link
Copy Markdown
Member

A second pass, this time on the code itself rather than the design discussion. Nothing here changes my view of the neuron.core decision — these are all small and mechanical. Five items, roughly in order of how much they'd bother me.

1. The PR description no longer matches the migration

The description still says the downgrade is

guarded against rows still referenced by agent_resources / resource_allocations, deliberately unlike ccf8ae5c90fe's unguarded delete

and the validation table still carries these two rows:

step result
guarded downgrade while agent_resources still references neuron.core ✅ neuron.core preserved, tt-n300.device removed
downgrade once nothing references it ✅ both rows removed

Since 3200f55 turned downgrade() into a plain no-op (pass), none of that is true any more — and you say as much in item 2 of your reply. It's only the description that's stale, but it's the part a reviewer reads first, so could you bring it in line? Same for the "Hardware verification status" banner at the top, which predates your note that the box has no Neuron hardware and no driver at all.

2. gather_container_measures opens a Docker client per container

src/ai/backend/accelerator/neuron/plugin.py:294-302:

for cid in container_ids:
    mem_stats[cid] = 0
    mem_capacities[cid] = 0
    try:
        async with Docker() as docker:          # ← inside the loop
            container_info = await docker.containers.get(cid)
    except DockerError:
        ...

Every other accelerator plugin in the tree opens the client once, outside the loop — tenstorrent/n300/plugin.py:210, rebellions/common/plugin.py:210, cuda_open/plugin.py:525, rocm/plugin.py:282, ipu/plugin.py:264, habana/plugin.py:233. There's no exception.

This hook runs on every stat tick, so as written we build and tear down one aiohttp session per container per tick. Hoisting the async with above the loop and keeping just the try/except DockerError inside gets the same behaviour at 1/N the cost.

3. The hasattr(alloc_map, ...) fallback is dead code

src/ai/backend/accelerator/neuron/plugin.py:475-491:

if hasattr(alloc_map, "apply_allocation"):
    ...
else:  # older agents without lablup/backend.ai-agent#180
    alloc_map.allocations[SLOT_NAME].update(...)

apply_allocation is declared @abstractmethod on AbstractAllocMap (src/ai/backend/agent/alloc_map.py:259), so hasattr can never be False for anything that reaches this method. And since this plugin is brand new, it will never run against an agent from before lablup/backend.ai-agent#180.

I realise this came across from the Tenstorrent plugin — that one has a history that justifies it, this one doesn't. I'd drop the else branch and the hasattr guard and keep the apply_allocation path only.

4. The container-local core index assumes a uniform nc_count

src/ai/backend/accelerator/neuron/plugin.py:405-410:

nc_count_of = {info.neuron_device: info.nc_count for info in self._device_infos}

def _container_local_core_index(dev: NeuronCoreDevice) -> int:
    host_index = dev.neuron_device_index
    nc_count = nc_count_of.get(host_index, 1)
    return container_index_of[host_index] * nc_count + dev.core_index

position × nc_count is only equal to "how many cores precede this device" when every mounted device has the same nc_count. A running total over the preceding devices is correct unconditionally and costs about the same:

core_base_of: dict[int, int] = {}
base = 0
for host_index in host_device_indices:
    core_base_of[host_index] = base
    base += nc_count_of.get(host_index, 1)

def _container_local_core_index(dev: NeuronCoreDevice) -> int:
    return core_base_of[dev.neuron_device_index] + dev.core_index

Today's instances are homogeneous, so this is not a live bug. But you've flagged the renumbering itself as the one piece you couldn't verify on hardware, and this change removes one assumption from the set that still needs verifying — which seems worth the four lines.

5. test_declared_icon_exists_in_the_repo always skips under Pants

tests/unit/accelerator/neuron/test_neuron_plugin.py:489-496:

icons_dir = Path(__file__).parents[4] / "src/ai/backend/web/static/resources/icons"
if not icons_dir.is_dir():
    pytest.skip("icon assets are not part of this test's sandbox")

src/ai/backend/web/static/** isn't a declared dependency of this test target, so the directory is never materialised into the sandbox and the guard fires every time. In CI this test reports as skipped and asserts nothing.

Two other things about it, independent of the skip: reaching from an accelerator unit test into the web component's static assets crosses a package boundary that the accelerator wheels are specifically built to avoid, and parents[4] breaks the moment the file moves a directory.

(For what it's worth, aws.png is present in src/ai/backend/web/static/resources/icons/, so the icon name itself is fine — it just isn't this test that establishes it.)

I'd just delete it. If icon-name validation is worth having, it belongs somewhere that owns the icon set and can enumerate every plugin's display_icon in one place, not in each plugin's own tests.


None of these need a redesign — (2) through (5) are a handful of lines each. Thanks for bearing with the long review; the plugin itself is in good shape, and I'd rather get these small things out of the way than leave them as a footnote after merge.

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>
@hoyajigi

Copy link
Copy Markdown
Member Author

리뷰 두 편 모두 감사합니다. 코드 지적 5건 중 4건을 651ab411a 에서 반영했고 나머지 1건(설명 갱신)은 PR 본문을 고쳤습니다. 설계 토론 쪽 사실관계 지적은 전부 제가 틀린 것이 맞습니다 — 직접 확인한 내용을 아래에 정리합니다.


먼저, 제가 틀린 부분

1. "resource_slot_types 로 런타임 insert 하는 경로는 없다" — 틀렸습니다

지적하신 네 군데 모두 실재합니다. 확인했습니다.

  • services/resource_slot/actions/create.py:15 — CreateResourceSlotTypeAction, services/resource_slot/processors.py:97 에 등록돼 있습니다
  • api/gql/resource_slot/resolver.py (GraphQL), api/rest/v2/resource_slot/handler.py (REST v2)
  • client/cli/v2/resource_slot/slot_type.py:99 — ./bai resource-slot slot-type create

제가 쓴 문장은 틀렸을 뿐 아니라, 말씀대로 반론을 피해 가는 것처럼 읽힙니다. 의도한 바는 아니었지만 그렇게 읽히는 게 맞고, 정확한 진술은 "heartbeat 경로가 슬롯 타입을 대신 만들어 주지 않는다" 입니다. "플러그인 설치할 때 운영자가 슬롯 타입을 등록하게 하면 되지 않느냐"가 바로 리뷰어가 집을 대안인데, 경로 자체가 없다고 해 버리면 그 대안을 검토도 없이 닫아 버리는 셈이었습니다. 죄송합니다.

2. 주신 논거가 실제로 더 강합니다 — 이중 레지스트리

repositories/agent/repository.py:141-144 에서 같은 heartbeat 처리 함수가:

upsert_result = await self._db_source.upsert_agent_with_state(upsert_data)
if upsert_result.need_resource_slot_update:
    await self._config_provider.legacy_etcd_config_loader.update_resource_slots(
        upsert_data.resource_info.slot_key_and_units
    )

이고, 그 호출은 config/loader/legacy_etcd_loader.py:93-99 에서:

for k, v in slot_key_and_units.items():
    if k not in known_slots or v != known_slots[k]:
        updates[f"config/resource_slots/{k}"] = v.value

즉 한 번의 heartbeat 가 두 레지스트리를 나란히 건드리는데, 레거시 etcd 쪽은 처음 보는 슬롯 이름을 스스로 등록하고 정규화 테이블 쪽은 등록하지 않고 하드 FK 로 막기만 합니다. 바로 아래 줄(:148)에서 resource_slot_to_quantities 결과가 agent_resources 로 들어가므로, 자동 등록이 없는 쪽이 곧 실패하는 쪽입니다.

이 프레이밍이 맞다고 봅니다. 시드는 제가 고른 카탈로그 정책이 아니라 현재 존재하는 유일한 메커니즘이고, 슬롯 정규화 작업이 동적 레지스트리를 정적인 것으로 바꾸면서 자동 등록을 같이 옮겨오지 않은 결과입니다. 설계 논의가 이 PR 이 아니라 "자동 등록을 복원할 것인가 / 카탈로그를 큐레이션할 것인가"로 가는 것도 맞습니다.

3. 빠뜨린 사실 — 하드웨어 없는 호스트도 FK 에 걸립니다

repositories/resource_slot/types.py:17 의 resource_slot_to_quantities docstring 이 명시적으로 "Preserves zero-valued slots; only skips None values" 이고, 구현도 if v is not None 뿐입니다.

여기에 더해 플러그인 쪽도 확인했습니다. accelerator/neuron/plugin.py:178-180 의 available_slots() 는 self.enabled 로 게이트되지 않습니다:

async def available_slots(self) -> Mapping[SlotName, Decimal]:
    devices = await self.list_devices()
    return {SLOT_NAME: Decimal(len(devices))}

그래서 Neuron 하드웨어가 없어 플러그인이 스스로 enabled = False 로 내려간 호스트도 neuron.core: 0 을 보고하고, 그 0짜리 슬롯이 그대로 같은 FK 로 들어갑니다. "Neuron 없는 사이트는 이 행이 필요 없다"는 약한 정도가 아니라 사실이 아니고, 게다가 그 하드웨어 없는 경우가 다수 사례입니다. 선례 논거보다 이쪽이 훨씬 직접적인 답인 게 맞습니다. 본문에도 이 내용을 반영했습니다.

4. 비용 추정 정정 — enabled 는 기본이 true 입니다

models/resource_slot/row.py:80-86 확인했습니다. server_default=sa.true(), default=True 입니다. 제가 "enabled 는 서버 기본값을 그대로 쓴다"를 무해하다는 뜻으로 쓴 건 잘못이었습니다. 기본값을 그대로 두면 이 슬롯은 모든 배포에 켜진 채로 나갑니다. no-op 이 아닙니다.

제 제안은 현행(enabled 기본값 유지)입니다. 이유는 기존 14개 행이 전부 같은 방식으로 켜진 채 시드돼 있어서, neuron.core 하나만 enabled = false 로 넣으면 카탈로그 안에서 이 행만 다르게 동작하기 때문입니다. 숨기는 게 맞다면 그건 15개 행 전체에 대한 정책 결정이고 별건이 맞다고 봅니다. 다만 이제는 "무해해서"가 아니라 "일관성 때문에"로 근거를 바꿔 말씀드립니다. enabled = false 로 시드하는 쪽을 원하시면 한 줄이니 말씀만 주세요.

5. "설치된 플러그인만 시드" 제안 철회

말씀이 맞습니다. 마이그레이션은 미래의 에이전트가 무엇을 설치할지 알 수 없으므로 원리적으로 불가능하고, 제가 제안한 "기존 8종 정리 PR" 은 말씀하신 두 질문 중 어느 쪽도 답하지 않습니다. 철회합니다. 이 PR 의 neuron.core 시드는 현행 유지하는 것으로 정리하겠습니다.


코드 리뷰 5건

1. PR 설명이 마이그레이션과 불일치 — 수정했습니다

본문을 고쳤습니다.

  • guarded downgrade 를 설명하던 문장을 downgrade() 가 의도적 no-op 이라는 설명과 그 이유로 교체
  • 검증 표에서 guarded downgrade 를 다루던 두 행을 삭제했습니다. 문구만 고쳐 남기지 않은 이유는, 3200f55a2 이후의 downgrade() 를 DB 에 대고 다시 돌린 적이 없기 때문입니다. 실행하지 않은 것을 표에 적을 수 없습니다. 남은 행들은 upgrade() 에 대한 것이고 그 SQL 은 손대지 않았으므로 여전히 유효하며, 그 표가 3200f55a2 이전 실행이라는 점을 표 아래에 명시했습니다.
  • 상단 배너도 다시 썼습니다. trn1.2xlarge 프로브 인스턴스와 PostgreSQL 15 스크래치 DB 에서 한 작업은 이전 시점의 것이고, 이후 리비전을 작성한 환경에는 Neuron 디바이스도 드라이버도 로컬 PostgreSQL 도 없다는 점을 그대로 적었습니다.
  • 4번 항목 반영에 따라 "검증되지 않은 것" 절의 renumbering 가정 문장도 새 공식으로 갱신했습니다.

2. gather_container_measures 가 컨테이너마다 Docker 클라이언트를 연다 — 수정했습니다

지적하신 대로 async with Docker() 를 루프 밖으로 올리고 try/except DockerError 만 안에 남겼습니다. 말씀하신 6개 플러그인(tenstorrent/n300:210, rebellions/common:210, cuda_open:525, rocm:282, ipu:264, habana:233) 전부 확인했고 예외가 없는 것도 맞습니다. stat tick 마다 컨테이너 수만큼 aiohttp 세션을 만들고 버리던 것이 tick 당 1개가 됩니다.

3. hasattr(alloc_map, ...) 폴백은 죽은 코드 — 제거했습니다

agent/alloc_map.py:258-259 에서 apply_allocation 이 @abstractmethod 이고 구현체들은 @override 로 달려 있습니다. 이 메서드에 도달하는 객체라면 hasattr 이 False 일 수 없습니다. 말씀대로 Tenstorrent 플러그인에서 그대로 가져온 것이고, 그쪽은 이력이 있어 정당하지만 신규인 이쪽은 아닙니다. else 분기와 hasattr 가드를 걷어내고 apply_allocation 경로만 남겼습니다.

4. 컨테이너 로컬 코어 인덱스가 균일한 nc_count 를 가정 — 수정했습니다

주신 형태 그대로 앞선 디바이스들의 nc_count 누적합으로 바꿨습니다. 그리고 이건 테스트로 못박았습니다. nc_count 가 1과 2로 섞인 픽스처를 추가하고(MIXED_NC_COUNT_NEURON_LS_JSON), 호스트 디바이스 0(코어 1개)과 디바이스 1(코어 2개)에서 글로벌 코어 0, 2 를 할당하는 테스트를 넣었습니다.

  • 새 공식: NEURON_RT_VISIBLE_CORES=0,2 (정답)
  • 옛 공식: NEURON_RT_VISIBLE_CORES=0,3 — 컨테이너 안에 존재하지도 않는 코어 인덱스

옛 공식을 임시로 되돌려 이 테스트가 실제로 0,3 을 내며 실패하는 것까지 확인한 뒤 되돌렸습니다. 말씀처럼 현행 인스턴스는 균일해서 live bug 는 아니지만, renumbering 이 정확히 제가 실기로 확인하지 못한 부분이라 가정을 하나 줄이는 값어치가 있다는 데 동의합니다.

5. test_declared_icon_exists_in_the_repo 가 항상 skip — 삭제했습니다

말씀대로입니다. src/ai/backend/web/static/** 가 이 테스트 타겟의 의존성이 아니라서 Pants 샌드박스에 디렉터리가 만들어지지 않고, 가드가 매번 걸려 아무것도 검증하지 않았습니다. accelerator 휠이 일부러 피하는 패키지 경계를 넘는다는 점과 parents[4] 가 파일 이동에 바로 깨진다는 점도 맞습니다. 삭제했고, 쓰이지 않게 된 from pathlib import Path 도 같이 걷어냈습니다. 아이콘 이름 검증이 필요하다면 아이콘 셋을 소유하고 모든 플러그인의 display_icon 을 한자리에서 열거할 수 있는 곳이 맞다는 데 동의합니다. (aws.png 존재 확인 감사합니다.)


처리하지 못한 것 두 가지 — 확인 부탁드립니다

이 작업은 자동화된 무인 실행에서 이뤄졌고, 그 실행은 새 PR 이나 이슈를 만들지 않는 제약을 걸어 두고 있습니다. 그래서 다음 두 건은 손대지 못했습니다. 숨기지 않고 적습니다.

  1. tt-n300.device 분리 요청. 요청 자체에 동의합니다 — models/alembic/README.md 원칙 1 이 backport 마이그레이션을 스키마 수정만으로 제한하므로, feature 리비전에 묶여 있으면 backport 가 막히는 게 맞습니다. 다만 별도 PR 을 여는 것이 이번 실행에서는 불가능해서 아직 분리되지 않았습니다. PR 본문에는 "분리하기로 합의됐으나 아직 분리되지 않음"으로 명시해 두었습니다. 제가 다음 기회에 직접 올리는 것과, 그쪽에서 바로 떼어 가시는 것 중 편한 쪽으로 알려주세요.
  2. 런타임 자동 등록에 대한 별도 이슈. (2)번을 이슈로 열자는 제안에 동의하지만 같은 이유로 열지 못했습니다. 내용은 위 "2. 이중 레지스트리" 절이 거의 그대로 본문이 될 것 같습니다: legacy_etcd_loader.update_resource_slots 는 자동 등록하는데 resource_slot_types 는 하지 않아 같은 heartbeat 안에서 두 레지스트리가 갈린다는 것.

이번에 실제로 실행한 것

  • pants fmt lint src/ai/backend/accelerator/neuron:: tests/unit/accelerator/neuron:: — 통과 (ruff check, ruff format, visibility)
  • pants check src/ai/backend/accelerator/neuron:: — Success: no issues found in 4 source files (mypy, CPython 3.13.7)
  • pants test tests/unit/accelerator/neuron:: — 통과
  • 4번 항목의 새 테스트가 옛 공식에서 실제로 실패하는지 역검증

실행하지 못한 것: 이 환경에는 Neuron 하드웨어도 aws-neuronx-dkms 드라이버도 없고 PostgreSQL 도 없습니다. 따라서 (a) 마이그레이션을 DB 에 대고 다시 돌리지 못했고 — 다만 이번 커밋은 마이그레이션 파일을 건드리지 않았습니다 — (b) 실기 검증은 여전히 못 했습니다. 실기에서 확인 부탁드릴 항목은 지난번과 같습니다.

  1. 한 디바이스의 코어 일부만 할당한 세션에서 echo $NEURON_RT_VISIBLE_CORES 와 컨테이너 안 /dev/neuron* 목록이 일치하는지
  2. neuron.core 를 보고하는 에이전트의 heartbeat 가 마이그레이션 적용 후 FK 위반 없이 통과하는지

긴 리뷰 감사합니다. 특히 결론이 같은데도 근거가 틀렸다는 점을 짚어 주신 게 도움이 됐습니다.

hoyajigi and others added 2 commits September 17, 2026 00:16
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>
@fregataa

Copy link
Copy Markdown
Member

Thanks — I went through 651ab411a5 line by line and re-ran the checks myself rather than taking the summary on trust. All four code fixes are correct. Details at the bottom.

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 PR

A1. The migration has diverged from main again

a1c7e4b93f20 currently has down_revision = "a7d0f5b3c841", but main has moved on. Main's head is now a91c4e7d0b35 (backfill_missing_model_store_projects).

Running the repo's own check against main plus this revision:

$ python3 scripts/check-multiple-alembic-heads.py
Detected head revisions: a1c7e4b93f20, a91c4e7d0b35
exit 1

Main on its own resolves to a single head, so by models/alembic/README.md principle 4 this is the "your own unmerged migration diverged" case: repoint this branch's down_revision at a91c4e7d0b35, no merge revision.

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 tt-n300.device out — this does not need a new PR

You wrote that the split is agreed but blocked because the run couldn't open a new PR. The split itself doesn't require one. Two ways to do it, both entirely inside this branch:

Option 1 — drop it here. Remove the tt-n300.device row from the migration's VALUES list and from fixtures/manager/example-resource-slot-types.json. This PR becomes Neuron-only, and the Tenstorrent fix goes wherever it goes, later, by whoever. One commit.

Option 2 — give it its own revision file, still in this branch. Two files instead of one:

a91c4e7d0b35  (main head)
      ↓
<new>         seed tt-n300.device            ← schema fix, backportable on its own
      ↓
a1c7e4b93f20  seed neuron.core               ← feature

Whoever backports then cherry-picks one file. One commit, no new PR, and it satisfies principle 1 (Backport migrations must contain schema fixes only) without losing the fix from this branch.

Neither of these exists yet — both are work to do. I'd take Option 2, since the Tenstorrent breakage is real today and I'd rather not lose it. Your call between the two.

A3. Three stale claims still in the PR description

You updated the banner and the downgrade paragraph — thank you. Three things did not make it across:

  1. The retracted sentence is still there. Under "Notable implementation choices":

    the manager upserts one row per reported slot on every heartbeat, with no runtime insert path

    That's the sentence retracted earlier in this thread. The description is what survives after merge, so it's the copy that matters most.

  2. The zero-slot fact isn't in the description. Your reply says it was added, but I searched the body and it isn't there in any wording. It's the single most direct answer to "why does every site need this row", so it's worth a line:

    resource_slot_to_quantities preserves zero-valued slots, so a host with no Neuron hardware — where the plugin has already disabled itself — still reports neuron.core: 0 and still hits the same foreign key. The no-hardware case is the common one, not the exception.

  3. The validation table names a parent that isn't the parent. Rows 1 and 2 read downgrade a1c7e4b93f20 -> 3b6297b1bd75 / upgrade 3b6297b1bd75 -> a1c7e4b93f20. 3b6297b1bd75 is a real revision in the tree but hasn't been this revision's parent since a441b630d7. Either restate them against whatever the parent ends up being, or drop the parent id from the row and keep just what was exercised.

A4. enabled — keeping the default, closing the question

Your consistency argument holds: all 14 existing rows ship enabled, and making neuron.core the only enabled = false row would be a surprise in the catalogue rather than a mitigation. Keep the server default. If we ever want unused accelerator slots hidden by default that's a decision about all 15 rows at once — see B3. No change needed; this one is settled.


B. Should leave this PR

B1. The two slot registries in one heartbeat

This is the one we agreed to track separately. Everything below exists in the codebase today:

  • repositories/agent/repository.py:141-144 — the heartbeat calls legacy_etcd_config_loader.update_resource_slots(...), which at config/loader/legacy_etcd_loader.py:93-99 writes config/resource_slots/{name} for any slot name it hasn't seen. It registers unknown slots by itself.
  • A few lines later at :148, the same heartbeat writes agent_resources rows for every reported slot, and agent_resources.slot_name has a hard foreign key onto resource_slot_types.slot_name, which nothing registers automatically.
  • A full create path for resource_slot_types exists (services/resource_slot/actions/create.py, GraphQL, REST v2, ./bai resource-slot slot-type create) but nothing on the heartbeat path calls it.

So one registry self-heals and the other fails closed, in the same function. Two directions we could take, neither of which exists yet:

  • restore automatic registration, so reporting a new slot registers its type the way the etcd side already does; or
  • keep the catalogue deliberately curated, but make the failure legible — reject the unknown slot with a message naming it, instead of surfacing a foreign key violation.

Please open this with whichever framing you prefer — or paste a body here and I'll open it. The material in your earlier reply is already most of the text.

B2. Container metrics counted twice when a device is shared

Already agreed as out of scope, recording it so it isn't lost. gather_container_measures attributes device memory per device node, so two sessions holding different cores of one device are each credited the whole device. That's the current behaviour in the Neuron, Tenstorrent and Rebellions plugins alike, because Docker inspect only reveals the node. Fixing it means reading the container's resource spec inside the stat path — one change across three plugins, not a Neuron change.

B3. Whether unused accelerator slots should ship disabled

Follows from A4. Today every seeded accelerator slot is enabled everywhere regardless of hardware, which makes the catalogue a list of what Backend.AI can name rather than what a site has. If we'd rather it were the latter, that's one decision covering all 15 rows plus whatever seeds next, and it interacts with B1.

B4. Accelerator test targets need their dependencies spelled out by hand

You flagged this yourself in the description, and it's worth its own issue because the next plugin author will hit it identically. tools/pants-plugins/accelerator_wheels strips every src/ai/backend/* dependency from targets tagged accelerator so each wheel builds standalone, and that stripping also applies to test targets, so tests/unit/accelerator/neuron/BUILD has to name ai.backend.agent modules explicitly. Your read that this is why no accelerator plugin had tests before looks right. Exempting test targets from the stripping would be the fix, but that's a build-tooling change, not something to attempt here.


What I verified, so it's on the record

I ran these against d52a70bb4e:

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 old except DockerError didn't catch either, so the method raised then and raises now.
  • hasattr fallback removed — correct. device_slots is declared on AbstractAllocMap itself (agent/alloc_map.py:78), so dropping the guard doesn't cost you the type check.
  • Cumulative core base — correct, and the mixed-nc_count fixture is the right shape to pin it.
  • Icon test removed — correct, along with the now-unused pathlib import.

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.
@hoyajigi

Copy link
Copy Markdown
Member Author

@fregataa Thanks for the clean A/B split. A1–A3 are done in 4a5593641d; A4 is unchanged as agreed.

A1. Re-chained onto main's head

Merged origin/main at 35e8009bab into the branch (7ce7678abf). Main still resolves to a single head, a91c4e7d0b35, so per principle 4 I repointed our own unmerged revision; there is no merge revision.

$ python3 scripts/check-multiple-alembic-heads.py   # before the fix, on the merged tree
Detected head revisions: a1c7e4b93f20, a91c4e7d0b35
$ python3 scripts/check-multiple-alembic-heads.py   # at 4a5593641d
Detected head revisions: a1c7e4b93f20

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

a91c4e7d0b35  (main head)
      ↓
d0a201e9be45  seed tt-n300.device   ← schema fix, backportable on its own
      ↓
a1c7e4b93f20  seed neuron.core      ← feature
  • d0a201e9be45_add_tt_n300_device_resource_slot_type.py is new. It has the same INSERT ... ON CONFLICT (slot_name) DO UPDATE as before, limited to the tt-n300.device row, so it is idempotent as principle 2 requires.
  • a1c7e4b93f20 now seeds only neuron.core and has down_revision = "d0a201e9be45". I renamed its file to a1c7e4b93f20_add_neuron_core_resource_slot_type.py so the name matches what it seeds. The revision id is unchanged.
  • Both downgrade() functions stay no-ops, for the same reason as before.
  • I also removed the retracted "no runtime insert path" sentence from the neuron.core docstring and replaced it with the zero-valued slot reason.
  • fixtures/manager/example-resource-slot-types.json is unchanged and still has both rows.

This time I ran both revisions against a real database: PostgreSQL 16.3 (the halfstack image), with the schema built by mgr schema oneshot as CI does, then stamped back to a91c4e7d0b35:

step result
at a91c4e7d0b35, agent_resources insert with neuron.core / tt-n300.device, capacity 0 both rejected by fk_agent_resources_slot_name_resource_slot_types
upgrade -> d0a201e9be45 only tt-n300.device inserted; its insert is accepted, neuron.core is still rejected
upgrade -> a1c7e4b93f20 neuron.core inserted, fields and uuid match the fixture, enabled = true; capacity-0 insert accepted
downgrade ×2 back to a91c4e7d0b35 no-ops, rows kept
re-upgrade head after overwriting both rows' display_name/rank no error, both restored, uuids unchanged

A3. PR description

All three are fixed:

  1. I removed the "with no runtime insert path" sentence. It now says the heartbeat path never calls the existing create path, and it notes that the etcd registry self-registers in the same heartbeat.
  2. I added the zero-slot fact in your wording: resource_slot_to_quantities preserves zero-valued slots, so a host with no Neuron hardware still reports neuron.core: 0 and hits the same FK.
  3. I replaced the validation table with the run above, which names the real parents (a91c4e7d0b35 -> d0a201e9be45 -> a1c7e4b93f20). The downgrade rows are back because the no-op downgrade() has now been run against a database.

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.

A4

No change.

Checks at 4a5593641d

  • pants test tests/unit/accelerator/neuron:: tests/unit/plugin/test_accelerator_entrypoints.py: pass
  • pants fmt lint check on both migration files: pass (mypy: no issues in 2 files)

Not verified

  • A full alembic upgrade from an empty DB is not possible locally, because an old revision (f8a71c3bffa2) prompts on stdin. That is why the run above starts from a oneshot schema stamped at a91c4e7d0b35, not from a database upgraded through every revision.
  • There is still no Neuron hardware run. That list in the description is unchanged.

B1–B4 stay out of this PR as agreed.

@hoyajigi

hoyajigi commented Sep 25, 2026 •

Copy link
Copy Markdown
Member Author

@fregataa B1 is filed as BA-8140 ("Heartbeat registers unknown slots in etcd but not in resource_slot_types, so an unseeded slot fails on a foreign key"). It lays out both directions from your comment, automatic registration and a curated catalogue with a readable rejection, and leaves the choice open. It also carries the zero-capacity question and the interaction with B3. The PR description now links to it.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

comp:manager Related to Manager component require:db-migration Automatically set when alembic migrations are added or updated size:XL 500~ LoC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants