Skip to content

Commit 9487527

Browse files
authored
[python][backends] Setup sys.modules with any loaded backend (triton-lang#6468)
From [importlib's docs](https://docs.python.org/3.13/library/importlib.html#importing-a-source-file-directly) you're supposed to setup `sys.modules` with the imported module after you import it, and the name should be the fully qualified import path. Besides following the docs, this enables any external user to access/reference a backend module (otherwise there's no easy way-- meaning without traversing the `gc`-- to access those modules).
1 parent 191ece3 commit 9487527

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

python/triton/backends/__init__.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
1-
import os
21
import importlib.util
32
import inspect
3+
import os
4+
import sys
45
from dataclasses import dataclass
56
from .driver import DriverBase
67
from .compiler import BaseBackend
@@ -9,6 +10,7 @@
910
def _load_module(name, path):
1011
spec = importlib.util.spec_from_file_location(name, path)
1112
module = importlib.util.module_from_spec(spec)
13+
sys.modules[name] = module
1214
spec.loader.exec_module(module)
1315
return module
1416

@@ -40,8 +42,8 @@ def _discover_backends():
4042
continue
4143
if name.startswith('__'):
4244
continue
43-
compiler = _load_module(name, os.path.join(root, name, 'compiler.py'))
44-
driver = _load_module(name, os.path.join(root, name, 'driver.py'))
45+
compiler = _load_module(f"triton.backends.{name}.compiler", os.path.join(root, name, 'compiler.py'))
46+
driver = _load_module(f"triton.backends.{name}.driver", os.path.join(root, name, 'driver.py'))
4547
backends[name] = Backend(_find_concrete_subclasses(compiler, BaseBackend),
4648
_find_concrete_subclasses(driver, DriverBase))
4749
return backends

0 commit comments

Comments
 (0)