[DEV] Support nested wrapper #70

mark14wu · 2025-05-30T00:59:52Z

tests/test_trace_add_clients.py
- new test validating client deduplication and dynamic addition via stacked @trace decorators.
triton_viz/clients/sanitizer/sanitizer.py
- Replaced function-level factory with class Sanitizer(ABC); concrete implementation selected in new.
- Added NullSanitizer for backend == "off".
- Registered SanitizerBruteForce, SanitizerSymbolicExecution, and NullSanitizer with Sanitizer.register(...).
- Removed unused imports and switched to cfg for backend lookup.
triton_viz/core/client.py
- ClientManager.__init__ defaults clients to an empty list.
- Added add_clients() to append new clients with type-based deduplication.
triton_viz/core/config.py
- Stored _sanitizer_backend lazily; raises at first access if unset instead of on module import.
- Environment variable fetch simplified; removed immediate ValueError.
triton_viz/core/patch.py
- Replaced direct imports of report_grid_execution_progress and sanitizer_backend with cfg.*.
triton_viz/core/trace.py
- Introduced _normalize_client() and add_clients() for runtime client management and deduplication.
- trace() decorator:
  - Supports stacking:
    - if target is already a Trace, it appends clients
    - if a JITFunction, it wraps once.
- Trace.__init__ now initializes an empty ClientManager then adds clients via add_clients().
Header/Import cleanup across modified files to consistently reference configuration via cfg and drop unused modules.

Copilot

Pull Request Overview

This PR enables support for nested trace wrappers while improving dynamic client addition and deduplication for instrumentation. Key changes include refactoring the Trace class to use an internal add_clients method, updating client instantiation to use a factory pattern via the Sanitizer abstract class, and cleaning up configuration and patching code.

Updated Trace class to support stacking and client deduplication
Refactored sanitizer instantiation to use a class-based factory with NullSanitizer
Header and import cleanup to consistently reference configuration via cfg

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
triton_viz/core/trace.py	Refactored Trace: added _normalize_client and add_clients methods
triton_viz/core/patch.py	Updated progress reporting and sanitizer backend lookups
triton_viz/core/config.py	Made sanitizer_backend lazy, deferring its check until first access
triton_viz/core/client.py	Modified ClientManager: added add_clients; mutable default parameter
triton_viz/clients/sanitizer/sanitizer.py	Switched to class-based sanitizer factory and added NullSanitizer
tests/test_trace_add_clients.py	Added tests for client deduplication across stacked trace decorators

triton_viz/core/client.py

Jokeren · 2025-05-30T01:05:38Z

tests/test_trace_add_clients.py

+    """
+    @triton_viz.trace(("sanitizer", "profiler"))
+    @triton_viz.trace("tracer")
+    @triton_viz.trace(("sanitizer",))    # Duplicate Sanitizer (should be ignored)


It seems too complicated isn't it...Why not we just standardize the api to support only one of the possible use cases? It can be either an array of modes or nested annotations, but not both

It's not hard for me to support both...

I want to support nested wrapper because I want to deal with the legacy code that use @triton_viz.sanitizer wrapper. With the new command-line tool triton-sanitizer, there may be nested wrappers.

It's not hard for me to support both...

It will be error-prone and complicate users. Let's just modify it to be either the first or the second convention.

I prefer the second convention because

I want to deal with the legacy code that use @triton_viz.sanitizer wrapper.

It's OK. Please update your code accordingly

Jokeren · 2025-05-30T22:23:36Z

triton_viz/core/trace.py

        raise ValueError("At least one client must be specified!")

+    if not isinstance(client, (str, Client)):
+        raise TypeError(f"Expected str or Client, got {type(client)}")
+
    def decorator(kernel) -> Trace:
        # When sanitizer is off, skip tracing and return the original kernel unchanged
        if cfg.sanitizer_backend == "off":
            return kernel

        # First-time wrapping
        if isinstance(kernel, JITFunction):


Have you tested this with @triton.autotune?

No. Let me write a unittest with autotune.

Maybe we can consider running all tests using CI.

I just ran test/test_autotune_add.py and it worked well.

Let's fix the bugs before merging

…; add unittest for adding clients

Co-authored-by: Copilot <[email protected]>

Jokeren · 2025-06-04T19:56:13Z

.github/workflows/python-app.yml

@@ -5,9 +5,9 @@ name: Python application

 on:
  push:
-    branches: [ "main" ]
+    branches: []


Those changes are suspicious

This change has been fixed in #71.
We use

branches-ignore: - '**'

instead of

branches: []

to ignore all ci tests for now.

Jokeren · 2025-06-04T19:56:19Z

.github/workflows/python-app.yml

-        cd triton_viz
-        pip install pre-commit
-        pre-commit run --all-files
+    # - name: Lint with pre-commit


And why merge without making all tests passed?

I want to focus on indirect load first and fix CI later.

As for CI problem, I figured out that once you set TRITON_INTERPRET=1, triton will skip driver checking.
See https://github.com/triton-lang/triton/blob/d141ab8b1bfa8e8dc703459412cd827b09f80b21/python/triton/_internal_testing.py#L30
However, triton-viz cannot work with TRITON_INTERPRET=1 and needs to be fixed.

I will do this after I finish indirect load.

Sure, it has to be compatible with TRITON_INTERPRET=1

mark14wu requested a review from Copilot May 30, 2025 00:59

Copilot AI reviewed May 30, 2025

View reviewed changes

triton_viz/core/client.py Outdated Show resolved Hide resolved

Jokeren reviewed May 30, 2025

View reviewed changes

mark14wu changed the title ~~Support nested wrapper~~ [DEV] Support nested wrapper May 30, 2025

Jokeren reviewed May 30, 2025

View reviewed changes

mark14wu and others added 7 commits May 30, 2025 16:15

delay backend check until first use

ff65acf

support add a client in client manager and trace.py

147495d

refactored sanitizer to a factory class; fix config usage in patch.py…

8253301

…; add unittest for adding clients

Use a safer default argument in __init__ of triton_viz/core/client.py

53f496e

Co-authored-by: Copilot <[email protected]>

stop supporting array of modes. use nested instead

3052148

conform to legacy naming

3b70f43

add ci test

06741bf

mark14wu force-pushed the support_nested_wrapper branch from 91e1409 to 06741bf Compare May 30, 2025 23:15

mark14wu added 7 commits May 30, 2025 19:22

skip pre-commit linting for now

4cc1002

restrict triton==3.1.0 in ci

13cedb0

specify backend in sanitizer unittest

ddd4410

remove cuda requirements in unittests

1a0bf24

temporarily skip symbolic execution tests

1bab829

specify backend in test_wrapper.py

0eaac5a

ci: remove GitHub Actions unit-test workflow

cf26a92

mark14wu merged commit 12dc9fc into keren/v2.0 Jun 4, 2025
0 of 2 checks passed

mark14wu deleted the support_nested_wrapper branch June 4, 2025 17:31

Jokeren reviewed Jun 4, 2025

View reviewed changes

[DEV] Support nested wrapper #70

[DEV] Support nested wrapper #70

Uh oh!

Conversation

mark14wu commented May 30, 2025

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mark14wu Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mark14wu Jun 4, 2025 •

edited

Loading