Skip to content

[FIX] Keep every client a trace is given - #487

Open
mark14wu wants to merge 1 commit into
ir-mode-loweringfrom
claude/duplicate-client-instance-loss-ir
Open

mark14wu wants to merge 1 commit into
ir-mode-loweringfrom
claude/duplicate-client-instance-loss-ir

Conversation

@mark14wu

@mark14wu mark14wu commented Oct 3, 2026

Copy link
Copy Markdown
Collaborator

Stacked on #482 (ir-mode-lowering); the reproduction needs Sanitizer(compile=True) from #480. The core change is in tilelens/core/client.py, which is identical from #479 to #482, so it can also be folded lower in the stack.

Problem

ClientManager.add_clients silently dropped a client when one of its class was already in the trace, and a client of another class with the same NAME silently replaced the first. So two compiled sanitizers for different targets kept only the first:

sm80 = Sanitizer(compile=True, target="cuda:80", abort_on_error=False)
sm90 = Sanitizer(compile=True, target="cuda:90", abort_on_error=False)
traced = tilelens.trace(sm90)(tilelens.trace(sm80)(kernel))
traced[(4,)](x, 100, BLOCK=32)
# before: only cuda:80 is compiled; sm90.last_status is None, no error
# after:  both targets are compiled and checked, each with its own verdict

The same loss hit eager clients, e.g. Sanitizer(abort_on_error=False) stacked on an existing Sanitizer was dropped without a word.

Change

A client passed to a trace is now in the trace afterwards, or add_clients raises; nothing is dropped silently:

  • Adding the same object again changes nothing.
  • IR clients are all kept, so a trace can hold one instance of a client class per target. The core already compiled once per distinct target and delivered each client only its own target's events; the dedup rule was what kept a second instance out.
  • A second interpreting client of one NAME raises ValueError (with the "trace the kernel twice" advice the launch-conflict error gives), since one interpreted run serves one client per name.
  • A client named by string (trace("sanitizer")) is a no-op when the trace already has a client of that name: no settings of the caller's are lost, and the documented stacked-decorator case in test_trace_decorator_add_clients keeps working. The string is resolved to its class before anything is constructed.

API notes:

  • ClientManager.clients is now a list in trace order (it was a NAME-keyed dict); get_client(name) returns the first client of that NAME.
  • A same-NAME IR client no longer replaces the existing one; it is added beside it, so a skip/run conflict between them is now raised.
  • README: the compiled-sanitizer target section shows stacking one sanitizer per target.

Tests

  • New: test_sanitizers_for_two_targets_share_one_trace (end to end, cuda:80 ok / cuda:90 violations in one launch), test_instances_of_one_ir_client_class_keep_their_own_targets, test_add_clients_keeps_every_ir_client_instance, test_add_clients_refuses_a_second_interpreting_client_of_one_name, plus a refusal check in test_trace_decorator_add_clients. All of these fail on the base branch for the reason above.
  • Updated tests that read clients as a dict, and test_launch_conflict_check_sees_every_ir_client (formerly asserted the replacement).

Commands (Triton 3.6, CPU only):

CUDA_VISIBLE_DEVICES="" python -m pytest tests/unit/test_ir_lifecycle.py tests/unit/sanitizer_compiled/test_client.py tests/unit/test_wrapper.py tests/end_to_end/test_core.py::test_trace_decorator_add_clients tests/end_to_end/test_compiled_sanitizer.py
CUDA_VISIBLE_DEVICES="" python -m pytest tests/ --ignore=tests/end_to_end/test_gluon.py

Full suite: 1607 passed, 14 failed. The 14 fail identically on the base branch and are environment-related here: Gluon gfx1250.cluster import (6), tile-* executables not on PATH (5), and the three known test_sanitizer.py failures. test_gluon.py was ignored because it fails to collect for the same Gluon import reason.

Not in this PR

  • Two compiled-sanitizer verdicts in Launch.records both carry client="compiled_sanitizer" and no target; they are told apart by trace order or by each instance's last_verdict.
  • Separate, pre-existing: when two different interpreting clients share a trace, patch_op re-wraps the original op for each client, so only the last client's op callbacks fire.

ClientManager.add_clients dropped a client silently when one of its
class was already in the trace, and a client of another class with the
same NAME replaced the first one. So two Sanitizer(compile=True)
instances for different targets left only the first: the second never
compiled, its last_status stayed None, and nothing said so.

A client passed to a trace is now in it afterwards, or add_clients
raises:

- adding the same object again changes nothing;
- IR clients are all kept, so a trace can hold one instance of a class
  per target, each compiled for its own target with its own verdict;
- a second interpreting client of one NAME raises ValueError, since one
  interpreted run serves one client per name;
- a client named by string ("sanitizer") is a no-op when the trace
  already has one of that name, as no settings of the caller's are lost.

ClientManager.clients is now a list in trace order, and get_client
returns the first client of a NAME.

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

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant