[Code Health] DFX collectors are worker-level objects with run-level thread lifetime

[ChaoWao]

Related: #995

Category

Technical Debt (cleanup, refactor)

Component

Host Runtime

Description

The five DFX collectors are already members of the worker's device context, but their
threads are started and joined per run. That mismatch is the root of three separate
costs.

The objects:

Collector	Owner
ChipSwimlaneCollector / ArgsDumpCollector / PmuCollector / ScopeStatsCollector	DeviceRunnerBase
DepGenCollector	per-arch DeviceRunner

The lifetime: start_shared_collectors_for_run() runs inside launch_run's arming block;
teardown_shared_collectors_after_run() runs in reap_run, joining the threads and doing
the actual artifact export. apply_call_config() binds both the five enables and
output_prefix_ from that run's CallConfig.

1. Enabling any DFX channel disables pipelining.

const bool allow_prepared_successor =
    concurrent_native_prepare_supported_impl() != 0 && !config->diagnostics_any();

A single per-run collector set cannot serve two overlapping runs, so diagnostics_any()
forces depth back to 1. This is a cost paid today, on every diagnostic run.

2. Thread create/join on every run. Four to five threads spawned in the arming block
and joined at drain, per run.

3. The device-side producer has a non-optional dependency on a host consumer, and its
acquire loop does not time out when that consumer is absent.

dfx_backpressure_device.h states the contract: "Block-on-contention is the only
behavior — no opt-out gate." The free-queue depths are small
(PLATFORM_DEP_GEN_SLOT_COUNT = 4, PLATFORM_DUMP_SLOT_COUNT = 4,
PLATFORM_SCOPE_STATS_SLOT_COUNT = 8), so a producer starves quickly without a consumer.

The acquire loop in profiler_device_engine.h:

do {
    if (head != tail) { ...; return true; }
    mark_fq_contended(header, &contended_signalled);
    if (!pop_freeze_barrier(header, Module::kBackpressureWaitCycles)) break;
} while (true);

pop_freeze_barrier spins only while fq_freeze_active != 0, and only the host sets
that flag. With no host consumer the barrier returns immediately every iteration, so the
30 s PLATFORM_DFX_BACKPRESSURE_TIMEOUT_CYCLES never arms — the header's own comment says
it is armed "only then and is the sole give-up (host dead/hung mid-freeze)". The result is
an unbounded spin on the AICPU, reaped only by the OS op-execute timeout. args_dump takes
a different path (wait_for_release, whose predicate is already true via fq_contended) and
does time out at 30 s — same subsystem, two failure shapes.

This third one is latent on main: every path that arms the device side also starts the
host side. It becomes reachable as soon as device work can be re-submitted without
re-entering the host bracket that owns the collector threads.

Location

• src/common/platform/onboard/host/device_runner_base.h:1244,1250-1252 — collectors are worker-level members
• src/a2a3/platform/onboard/host/device_runner.h:285, src/a5/platform/onboard/host/device_runner.h:180 — DepGenCollector
• src/common/platform/onboard/host/device_runner_base.cpp:1852 — start_shared_collectors_for_run()
• src/common/platform/onboard/host/device_runner_base.cpp:1887 — teardown_shared_collectors_after_run(), which also performs export
• src/a2a3/platform/onboard/host/device_runner.cpp:606,705 — the per-run start / teardown bracket
• src/common/platform/onboard/host/c_api_shared.cpp:747-748 — diagnostics_any() disabling allow_prepared_successor
• src/common/platform/include/common/dfx_backpressure_device.h — "no opt-out gate"
• src/common/platform/include/aicpu/profiler_device_engine.h:88-130 — wait_for_free_queue_entry acquire loop
• src/common/platform/include/host/profiler_base.h:1198-1240 — drain loop; its idle timeout only logs and keeps the consumer alive
• src/a2a3/platform/include/common/platform_config.h — free-queue depths and PLATFORM_DFX_BACKPRESSURE_TIMEOUT_CYCLES

Observed on 3e08ede76478a6e672c8dda44587ed6b3e4e7e01.

Proposed Fix

Move the thread bracket and the configuration binding from the run to the worker; keep a
separate, explicit boundary for record attribution.

Service = worker lifetime. Threads, device-side headers and queues live from
simpler_init to finalize_device. A consumer is then always present, which removes the
liveness hazard for every caller shape, removes the per-run thread churn, and removes the
reason diagnostics_any() has to force depth 1.

Attribution = explicit session. Today the run boundary is the collector boundary —
that is how records reach a run's output_prefix_. With a resident service that identity is
gone, so it needs to be stated: a session_begin(config, output_prefix) /
session_end() pair, where session_end reuses the existing open-drain-release machinery
(ProfilerAlgorithms::update_backpressure_freeze) to quiesce the pipeline and then export.
Draining rather than tagging avoids adding an epoch field to the record wire format and
avoids the race where a record produced in one session is drained in the next. The
precondition — the caller has quiesced the device — is natural for a diagnostic boundary and
is already satisfied at the point reap_run calls teardown today.

Wrapping the existing prepare_run / reap_run points in a session should leave current
behavior byte-identical; that equivalence is the regression bar.

Two behavior changes to decide explicitly:

• Per-run config becomes per-session. Today apply_call_config re-binds the five
  enables and output_prefix_ on every run, so run 1 can have PMU on and run 2 off. That
  capability has to move onto the session or be dropped.
• finish_clock_correlation_session() is per-run today and needs the same treatment.

[ChaoWao]

Implementation plan

Two corrections to the Proposed Fix above, then the concrete edit list.

Correction 1 — the session_begin / session_end pair is not needed

The Proposed Fix introduces a session because moving the service to worker
lifetime leaves record attribution without a carrier. That reasoning is right,
but the carrier it needs already exists.

Today the run boundary carries two jobs: (a) switching the collection
service on and off, and (b) deciding which output_prefix_ a batch of
records belongs to. Residency only takes (a) away. Whether (b) needs a new
boundary depends on a decision that has not been made yet:

	attribution	pipelining	cost
A. session ≡ run	exact	still disabled	none — and no new concept: the "new boundary" coincides with the run, so apply_call_config and teardown_shared_collectors_after_run keep their current names and roles
B. session spans several runs	overlapping runs merge into one export	enabled	deps.json / swimlane.json stop being per-run artifacts
C. epoch field on the record	exact	enabled	changes the host↔device wire format, which the Proposed Fix explicitly wants to avoid

Recommendation: A. It already removes two of the three costs, and it is the
only option compatible with the regression bar this issue sets for itself
("wrapping the existing prepare_run / reap_run points ... byte-identical").
Under A there is no new API surface at all — the change is entirely in the
implementation.

Naming note if B or C is ever chosen: "session" collides with the remote-L3
session already in chip_worker.cpp / worker.py / remote_l3_session.py.
The only existing DFX use of the word is ClockCorrelationSession
(clock_correlation.h:59), which is per-run and serves the swimlane alone.

Correction 2 — cost #1 does not fall out of residency

The issue says residency "removes the reason diagnostics_any() has to force
depth 1". It does not, on its own.

Attribution under A and B is by drain, which requires the device to be
quiesced. Pipelining means run N is still executing while run N+1 is being
prepared — so there is no quiescent point between them at which to export.
Residency removes the thread churn and the liveness hazard; the depth-1 gate at
c_api_shared.cpp:747-748 survives it and needs the A/B/C decision above.

Cost #2 (thread churn) and cost #3 (unbounded AICPU spin) do fall out of it.

Also: the lifetime mismatch is wider than the threads

The issue frames this as thread lifetime. The device shared memory is per-run
too — init_chip_swimlane / init_args_dump / init_pmu / init_dep_gen /
init_scope_stats all run inside prepare_execution, and
finalize_collectors() releases them from cleanup_execution
(a2a3 device_runner.cpp:474, a5:661).

That pairing is load-bearing today: ChipSwimlaneCollector::initialize returns
an error when shm_host_ != nullptr (chip_swimlane_collector.cpp:92), so
without the per-run release the next run cannot re-init. a5's
finalize_collectors says so directly in its comment.

What changes

Nothing above DeviceRunnerBase. Python, the nanobind layer, ChipWorker, and
all six exported C API entry points keep their signatures. Inside the C API only
c_api_shared.cpp:808 changes, and only to push a few more parameters.

One new primitive — ProfilerBase::quiesce().

stop() (profiler_base.h:981) currently discharges two responsibilities at
once, and its drain guarantee is paid for by the threads exiting — its own
comment explains the ordering: stop+join mgmt first so its final-drain pass
lands in L2 before poll is told to exit. With threads resident that guarantee
needs an epoch handshake instead:

// Drain to a quiescent point: L1 ring and every host ready-queue shard empty,
// on_buffer_collected called for every entry that was in either. Threads stay.
void quiesce();

void stop();   // = quiesce() + join, used by the teardown path

ProfilerAlgorithms::update_backpressure_freeze (profiler_base.h:541) is
useful here as optional hardening — it parks device producers — but it is not
the mechanism: teardown's precondition is that the caller already quiesced the
device, which is exactly what holds where reap_run calls it today.

Three edits in the runners:

1. initialize() gains idempotence. The five "already initialized" error
   returns become silent skips. This is what unloads the per-run
   finalize_collectors.
2. initialize()'s per-run parameters move out. output_prefix,
   chip_swimlane_level_, dump_args_level_, PMU's csv_path and
   event_type are burned into a one-shot resource call today
   (chip_swimlane_collector.cpp:123). They move to apply_call_config, which
   already binds output_prefix_ (device_runner_base.cpp:1255) — so this
   removes a double binding rather than adding one. Sizing switches to
   PLATFORM_MAX_CORES / PLATFORM_MAX_AICPU_THREADS; header->num_cores and
   the shard counts still reflect the run, via the existing set_core_types
   path. This is the only step with real work in it.
3. Two call sites are deleted, one is swapped. Delete
   finalize_collectors() from cleanup_execution and
   start_shared_collectors_for_run() from the arming block; swap stop() for
   quiesce() inside teardown_shared_collectors_after_run. The
   finalize_collectors() already present in finalize_device
   (a2a3:1071, a5:939) is promoted from "no-run-since-init backstop" to the
   only path.

Lazily, not at simpler_init. prewarm_config is optional
(chip_worker.cpp:324: "Null => no prewarm") and is documented as
fork-constant, whereas the five DFX enables are re-bound per run — so
simpler_init cannot know whether this worker will ever want DFX. Standing the
collectors up there would charge every non-DFX worker. Build on the first run
with diagnostics_any() true, then never tear down until finalize_device.

Sizing: allocating at the platform maximum is affordable

Measured by compiling the headers directly (PLATFORM_* constants ×
sizeof of each buffer type):

collector	scales with	a2a3 at max	a5 at max
swimlane	num_aicore	26.78 MiB (72)	40.21 MiB (108)
pmu	num_aicore	4.51 MiB	6.76 MiB
dep_gen	aicpu_thread_num	74.00 MiB	92.50 MiB
dump	aicpu_thread_num	1.00 MiB	1.25 MiB
scope_stats	aicpu_thread_num	0.63 MiB	0.78 MiB

The delta is what matters, and it is small:

• The three that scale with aicpu_thread_num cost nothing extra —
  PLATFORM_MAX_AICPU_THREADS is 4 (a2a3) / 5 (a5), which real runs already
  use. dep_gen's 74/92.5 MiB is what it allocates today.
• Only swimlane and pmu have a real delta: a typical 24-core run growing to the
  platform max is about +21 MiB (a2a3) / +36 MiB (a5).
• With lazy construction, a worker that never enables DFX pays zero.

num_aicore is block_dim * cores_per_blockdim_, derived from
runtime.get_worker_count(), so it genuinely varies per callable — allocating
at the maximum (rather than rebuilding when it grows) is what removes the
rebuild logic entirely.

Sequencing

#	step	acceptance
0	a5 dep_gen host-graph gap	done — #2082
1	quiesce() split	cpput: after quiesce the threads are alive and a second round of records is not lost
2	initialize() parameters + idempotence	artifacts byte-identical
3	delete the two call sites, swap stop() → quiesce()	artifacts byte-identical; a multi-run DFX case shows no re-allocation
4	(optional) fold the four runners' duplicated sequences together	artifacts byte-identical

Step 0 was an independent defect found while surveying these collectors, not a
prerequisite in the dependency sense — but it touched the same a5 files and was
a behaviour change, so landing it first is what lets steps 2–4 hold a
byte-identical bar.

Step 4 is worth doing on its own merits: the teardown sequence exists in four
near-identical copies (DeviceRunnerBase + SimDeviceRunnerBase, which are
separate classes rather than one hierarchy), and they have already drifted —
the sim runners skip the export entirely on a failed run where onboard exports
it, and the two orders artifacts are written in differ. A shared component
(composition, not a base-class method — the two bases share no ancestor) would
fix both by construction.

Cost #1 is deliberately not in this list; it needs the A/B/C decision first.