[REFACTOR] Rename Triton-Viz to TileLens - #477
Conversation
Performance Benchmark
Iterations: 1 warmup + 20 measured |
| @@ -271,7 +271,7 @@ def test_trace_nested_jit_calls(): | |||
| This test ensures compatibility with the command line triton-sanitizer wrapper. | |||
There was a problem hiding this comment.
triton-sanitizer->tile-sanitizer
| @@ -101,11 +101,11 @@ def test_triton_sanitizer_injects_trace_outermost(tmp_path: Path, monkeypatch): | |||
| def test_triton_profiler_injects_trace_outermost(tmp_path: Path, monkeypatch): | |||
| """ | |||
| Black-box verification for triton-profiler: | |||
There was a problem hiding this comment.
triton-profiler -> tile-profiler
| @@ -177,11 +177,11 @@ def test_triton_profiler_injects_trace_outermost(tmp_path: Path, monkeypatch): | |||
| def test_triton_race_detector_injects_trace_outermost(tmp_path: Path, monkeypatch): | |||
| """ | |||
| Black-box verification for triton-race-detector: | |||
There was a problem hiding this comment.
triton-race-detector->tile-race
|
Also note that https://github.com/gpu-mode/Triton-Puzzles depends on it. So we probably want to keep the old package link or use other approaches for backward compatibility |
|
Should we preserve backward compatibility with the old names so dependent projects continue to work? I can add that if needed. I’ll also update the CLI names. |
Just make sure triton-puzzles is working after this fix. |
|
@pd-hung Please let me know once the triton-puzzle is working. |
| from tilelens.utils.dtypes import STORAGE_DTYPES | ||
| except ModuleNotFoundError: | ||
| pytest.skip( | ||
| "NeuronX dependencies are missing. Install triton-viz[nki] to run these tests.", |
There was a problem hiding this comment.
@pd-hung still triton-viz here. That's why I suggested you invoke gpt-6 to thoroughly go through the codebase
There was a problem hiding this comment.
My understanding is that we’re keeping the PyPI package name as triton-viz, so triton-viz[nki] is still the install command. The triton-puzle should also work.
There was a problem hiding this comment.
we’re keeping the PyPI package name as triton-viz
No this is not true. I meant we can have a legacy pypi package with triton-viz but going forward we can use tilelens as the only reference name
|
Sorry for the confusion, I have made the change. Existing triton-viz import from triton-puzzle should also work. |
|
@mark14wu take another pass please |
mark14wu
left a comment
mark14wu
left a comment
There was a problem hiding this comment.
I took another pass with the focus on Triton-Puzzles. Short version: the current gpu-mode/Triton-Puzzles notebook keeps working, but the backward-compat story only covers the 8 top-level names, and the distribution rename introduces an in-place-upgrade trap. Details and repro below; inline comments point at the exact lines.
What I verified works
- The notebook (
pip install -U git+.../triton-viz,import triton_viz,triton_viz.clear(),triton_viz.trace("tracer")as call and as decorator,triton_viz.launch(share=True/False)) was replayed cell by cell ona56864dvs the merge-base203b0c7: 4 demo cells + 11 solved puzzles (2-D/3-D grids, masks, loops,tl.dot), with and withoutTRITON_INTERPRET=1. Results identical.triton_viz.trace is tilelens.traceandtriton_viz.config is tilelens.confighold. - Fresh venv + real
pipinstall from the git URL: onlytilelens 3.0+ the shim get installed, re-running the install cell is idempotent, all 8 console scripts (tile-*andtriton-*) exist and thetriton-sanitizer/profiler/race-detectoraliases run end to end. Flask/,/api/data,/api/sbuf, static assets all 200. - Legacy
.tvzloading with a real archive written by203b0c7(OOB records, TracebackInfo, nested dataclasses) loads and renders; the legacy parametrization intest_trace_io.pydoes hit the remap path. - Rename sweep: the only remaining
triton_vizstrings are the intentional ones (shim, trace_io legacy handling, test_trace_io, pyproject include, README transition text, paper citation)..tvzextension, citation titles and the Triton-Puzzles link were not touched. Recompilingfrontend/with tsc and diffing againsttilelens/staticis clean.
What breaks
1. triton_viz.* submodules are gone (major). The shim only re-exports 8 names; every submodule import raises ModuleNotFoundError:
$ python -c "from triton_viz.clients import Sanitizer"
ModuleNotFoundError: No module named 'triton_viz.clients'
$ python -c "from triton_viz.core.trace import launches"
ModuleNotFoundError: No module named 'triton_viz.core'
On main there are 48 importable triton_viz.* modules; from triton_viz.clients import Sanitizer was the documented way to use the sanitizer, and GitHub code search finds course notebooks, puzzle forks and experiment repos doing exactly that. So the README statement "continues to expose the same public API" is only true for the top-level names. A lazy sys.meta_path finder that aliases triton_viz.X to the same tilelens.X module object fixes this with no double execution (prototype in the inline comment on the shim; verified against the puzzles flow and the old import surface).
2. In-place upgrade leaves two overlapping distributions (major). Because the dist is now tilelens, pip does not treat it as an upgrade of triton-viz. On any environment that already has triton-viz installed, the notebook's pip install -U git+.../triton-viz yields:
$ pip list | grep -i -E "triton-viz|tilelens"
tilelens 3.0
triton-viz 3.0
$ ls site-packages/triton_viz/
__init__.py clients/ core/ static/ templates/ utils/ visualizer/ wrapper.py ...
Both dists own triton_viz/__init__.py and bin/triton-*. The stale triton_viz/clients etc. are importable but built on the old Client base class:
>>> from triton_viz.clients import Sanitizer # silently the stale copy
>>> triton_viz.trace(client=Sanitizer(abort_on_error=True))(k)[(1,)](x, B=8)
TypeError: Expected str or Client, got <class 'triton_viz.clients.sanitizer.sanitizer.SymbolicSanitizer'>
And pip uninstall triton-viz afterwards deletes the shim, so import triton_viz fails; re-running the install cell is a no-op (same version 3.0), only --force-reinstall recovers. Fresh Colab is unaffected. Minimal fix: tell users to pip uninstall -y triton-viz before installing in the transition section. Cleaner long-term fix: a transitional triton-viz release on PyPI that depends on tilelens and owns the shim, so tilelens itself stops shipping triton_viz/.
3. TRITON_VIZ_VERBOSE / TRITON_VIZ_NUM_SMS / TRITON_VIZ_PORT are silently ignored (minor). No fallback, and the transition section doesn't mention them. Everything else (imports, .tvz, CLI argv[0]) got a compat path; this is the one thing that didn't.
4. Nothing tests the compat shim (minor). grep -rn "import triton_viz" tests/ .github/ is empty, and the pip-install CI job now imports tilelens, so the promise this PR makes about Triton-Puzzles is unguarded.
5. Description / docs vs code (minor). The PR body still says "Keep the original PyPI package name" but a56864d renames the dist to tilelens; pip install tilelens is a 404 on PyPI today and github.com/Deep-Learning-Profiling-Tools/tilelens does not exist yet, while README/CONTRIBUTING/docs/index.html (GitHub Pages, live on merge) already point there. "Breaking Changes" is empty; it should list: dist renamed, triton_viz.* submodules removed, env vars renamed, .tvz written by TileLens not readable by older triton-viz, CLI primary names now tile-*, frontend localStorage key triton-viz-theme -> tilelens-theme.
Suggested minimal set before merge
- Submodule aliasing in
triton_viz/__init__.py(or, if we decide not to support it, fix the README wording and list it as breaking). pip uninstall -y triton-viznote in the transition section.- Fallback to the
TRITON_VIZ_*env vars inconfig.py/interface.py. - A
tests/unit/test_legacy_import.py(importtriton_viz, assert identity withtilelensobjects, runtrace("tracer")+clear()) andimport triton_vizin the pip-install CI job. - Update the PR description and agree on the order of repo rename / PyPI registration / release.
| return _launch(*args, **kwargs) | ||
| """Backward compatibility with triton_viz""" | ||
|
|
||
| from tilelens import ( |
There was a problem hiding this comment.
This only aliases the 8 top-level names, so triton_viz.clients, triton_viz.core, triton_viz.visualizer, triton_viz.wrapper all raise ModuleNotFoundError (48 modules on main, 0 here). The current Triton-Puzzles notebook doesn't need them, but from triton_viz.clients import Sanitizer is what our own README and most downstream code used.
A lazy meta-path finder keeps one module object per name (same launches list, same config, same Client classes), and because it sits first in sys.meta_path it also shadows a stale site-packages/triton_viz/<subpkg> left behind by an old triton-viz install. Verified against the puzzles flow and the pre-PR import surface:
prototype
"""Backward compatibility with triton_viz (prototype shim with submodule aliasing)."""
from __future__ import annotations
import importlib
import importlib.abc
import importlib.util
import sys
from tilelens import (
__all__ as __all__,
__version__ as __version__,
clear as clear,
config as config,
git_version as git_version,
launch as launch,
load as load,
save as save,
trace as trace,
)
_OLD, _NEW = __name__, "tilelens"
class _AliasLoader(importlib.abc.Loader):
def __init__(self, target: str) -> None:
self._target = target
def create_module(self, spec):
# Return the *existing* tilelens module object so both names share
# one module (one `launches` list, one `config`, identical classes).
return importlib.import_module(self._target)
def exec_module(self, module):
pass # already executed under its canonical tilelens name
class _AliasFinder(importlib.abc.MetaPathFinder):
def find_spec(self, fullname, path=None, target=None):
if not fullname.startswith(_OLD + "."):
return None
new_name = _NEW + fullname[len(_OLD):]
try:
importlib.import_module(new_name)
except ModuleNotFoundError as exc:
if exc.name in (new_name, None) or (exc.name and new_name.startswith(exc.name)):
return None # -> ModuleNotFoundError: No module named 'triton_viz.x'
raise
spec = importlib.util.spec_from_loader(fullname, _AliasLoader(new_name))
spec.submodule_search_locations = getattr(
sys.modules[new_name], "__path__", None
)
return spec
if not any(isinstance(f, _AliasFinder) for f in sys.meta_path):
# insert(0): also shadows a stale site-packages/triton_viz/<subpkg> left
# behind by an older 'triton-viz' distribution.
sys.meta_path.insert(0, _AliasFinder())Please avoid __path__ = tilelens.__path__ as the shortcut: that executes every module twice under two names, so isinstance(client, Client) fails across the boundary.
|
|
||
| [project] | ||
| name = "triton-viz" | ||
| name = "tilelens" |
There was a problem hiding this comment.
Renaming the distribution has two consequences worth handling explicitly:
- pip does not treat
tilelensas an upgrade oftriton-viz. On an env that already has triton-viz, the notebook'spip install -U git+.../triton-vizleaves both dists installed, both owningtriton_viz/__init__.pyandbin/triton-*, with staletriton_viz/clients,core, ... still importable (repro in the review summary).pip uninstall triton-vizthen removes the shim. At minimum the README transition section should say to uninstall triton-viz first; longer term a transitionaltriton-vizrelease that depends ontilelensand owns the shim would let this wheel droptriton_viz*frompackages.find. tilelensis not on PyPI yet (404) and the PR description still says the PyPI name is kept. Please update the description and list this under Breaking Changes.
| """Reload configuration from environment variables and apply defaults.""" | ||
| self.verbose: bool = _is_one("TRITON_VIZ_VERBOSE") | ||
| self.num_sms: int = _get_int_env("TRITON_VIZ_NUM_SMS", 1, minimum=1) | ||
| self.verbose: bool = _is_one("TILELENS_VERBOSE") |
There was a problem hiding this comment.
TRITON_VIZ_VERBOSE / TRITON_VIZ_NUM_SMS are now silently ignored (there is no TRITON_VIZ_ string left anywhere in the tree):
$ TRITON_VIZ_VERBOSE=1 TRITON_VIZ_NUM_SMS=4 python -c "import triton_viz; print(triton_viz.config.verbose, triton_viz.config.num_sms)"
False 1 # 203b0c7 prints: True 4
Suggest reading TILELENS_* first and falling back to TRITON_VIZ_* (optionally with a DeprecationWarning), and documenting the alias in the README env-var table.
| """ | ||
| default_port = 8000 if share else 5001 | ||
| actual_port = port or int(os.getenv("TRITON_VIZ_PORT", default_port)) | ||
| actual_port = port or int(os.getenv("TILELENS_PORT", default_port)) |
There was a problem hiding this comment.
Same as config.py: TRITON_VIZ_PORT=9999 now binds 5001. Please fall back to the old name here too.
|
|
||
| - name: Verify installation | ||
| run: python -c "import triton_viz; print(triton_viz.__version__)" | ||
| run: python -c "import tilelens; print(tilelens.__version__)" |
There was a problem hiding this comment.
This was the only place that exercised import triton_viz from an installed wheel, and it now imports tilelens instead, so the compatibility this PR promises for Triton-Puzzles is untested. Suggest python -c "import triton_viz, tilelens; assert triton_viz.trace is tilelens.trace" here (ideally against the built wheel rather than pip install -e), plus a tests/unit/test_legacy_import.py that runs triton_viz.trace("tracer") + triton_viz.clear() on a small kernel.
|
|
||
| The GitHub repository and PyPI distribution are named `tilelens`. Use | ||
| `import tilelens` in new code. The legacy `triton_viz` import continues to expose | ||
| the same public API for existing code, including Triton-Puzzles. |
There was a problem hiding this comment.
Only true for the 8 top-level names; every triton_viz.* submodule import fails (see the comment on triton_viz/__init__.py). Either make it true via submodule aliasing or narrow the wording to "the top-level API (trace, clear, save, load, launch, config)".
Also worth adding here: users upgrading an existing install must pip uninstall -y triton-viz before pip install tilelens, otherwise both dists end up installed and later uninstalling triton-viz deletes the triton_viz shim. And TRITON_VIZ_VERBOSE/NUM_SMS/PORT were renamed to TILELENS_*; the env-var section should mention the old names.
|
@mark14wu Please let me know if you have additional comments on this change. |
_LegacyFinder.find_spec called importlib.util.find_spec(), which imports parent packages. CPython calls meta path finders while holding the global import lock, so that import could wait on a module lock held by a thread in _LegacyLoader.exec_module that was itself waiting for the global lock, hanging both threads permanently. find_spec now imports nothing: loaded modules are read from sys.modules, and the other meta path finders are asked for a module only once its canonical parent is loaded, as importlib itself guarantees. While another thread is still importing the parent, a deferred spec is returned and exec_module resolves it, which also avoids a KeyError for namespace packages such as tilelens.clients.tracer. exec_module imports the canonical ancestors top-down so it never holds a child lock while waiting for a parent. A thread that raced a first import could also bind names from the empty placeholder module, because CPython keeps the module object it found in sys.modules before waiting for the import lock, so "from triton_viz.clients import Sanitizer" raised ImportError. The placeholder now forwards attribute lookups to the canonical module. The finder is installed before the tilelens re-export. This narrows, but cannot close, the window in which a concurrent first import of a triton_viz submodule fails with ModuleNotFoundError.
Use tile-race as the race detector command, as requested in review. triton-race-detector stays installed as the compatibility alias. The wrapper previously derived the legacy name by replacing "tile-" with "triton-", which would map tile-race to a nonexistent triton-race command. It now looks the legacy name up in an explicit table.
--force-reinstall also reinstalls every dependency and can upgrade Triton. Adding --no-deps restores only the files shared with triton-viz.
The v3.0 tag and triton-viz 3.0 on PyPI already exist, so the first release under the tilelens name is 3.1.
3 participants
Summary
triton_viz→tilelens), the PyPI distribution (triton-viz→tilelens), CLI commands, environment variables, tests, examples, and documentation.triton_vizcompatibility package.import triton_vizand everytriton_viz.*submodule import resolve to the sametilelensmodule objects, including when several threads import them at the same time..tvztraces still load, and the oldTRITON_VIZ_*environment variables andtriton-*CLI commands still work.Test Plan
-eand checks, from outside the checkout, thattriton_viz.trace,triton_viz.config,triton_viz.clients.Sanitizer, andtriton_viz.core.trace.launchesare the same objects as theirtilelenscounterparts.import triton_viz,triton_viz.trace("tracer"),clear(),launch()) against this branch.Related Issues
Breaking Changes
tilelens.triton-viz3.0 is the last release under the old name, andimportlib.metadata.version("triton-viz")no longer finds this package.triton-vizandtilelensship the sametriton_viz/__init__.pyandtriton-*scripts. Uninstalltriton-vizbefore installingtilelens. Installingtriton-vizaftertilelensoverwrites the compatibility package, and uninstalling either one while both are installed removes the shared files.tile-sanitizer,tile-profiler,tile-race, andtile-visualizer.TILELENS_VERBOSE,TILELENS_NUM_SMS, andTILELENS_PORT."format": "tilelens_trace") cannot be loaded by older Triton-Viz releases.tilelens-themeinstead oftriton-viz-theme, so a saved theme resets once. The API base override is nowwindow.__TILELENS_API__instead ofwindow.__TRITON_VIZ_API__.triton_vizis now a thin compatibility package.pkgutil.iter_modules(triton_viz.__path__)returns nothing,importlib.resources.files("triton_viz")no longer contains the templates and static assets, and after a bareimport triton_viz, submodules such astriton_viz.coremust be imported explicitly. In multithreaded programs, importtriton_vizin the main thread before other threads import its submodules.github.com/Deep-Learning-Profiling-Tools/tilelensanddeep-learning-profiling-tools.github.io/tilelens/.Still working:
import triton_vizandtriton_viz.*submodule imports; theTRITON_VIZ_*variables (theTILELENS_*value wins when both are set); thetriton-sanitizer,triton-profiler,triton-race-detector, andtriton-visualizercommands; loading existing.tvztraces.Rollout
The README and the docs site already use the new names, so these steps should all happen on the day this PR is merged. Do not publish a release or run
publish.ymlbefore the merge.v3.1and publish a GitHub release. This PR already bumpsversionto 3.1, because thev3.0tag andtriton-viz3.0 already exist.publish.ymlthen uploadstilelens3.1 through the pending trusted publisher on PyPI (repositorytriton-viz, workflowpublish.yml).tilelens.Checklist
npm run build:frontendif the PR modified any TypeScript code.