From 3e99d290856fc0af521a71ad6bf2be04c316ee47 Mon Sep 17 00:00:00 2001 From: Mikezhang001 <2826398872@qq.com> Date: Fri, 5 Jun 2026 18:20:58 +0800 Subject: [PATCH] fix: use importlib to load scripts/generate_pyi.py instead of package import The top-level 'from scripts.generate_pyi import generate_pyi_file' in setup.py requires scripts/ to be a Python package (with __init__.py), but scripts/ is a utility directory and was never intended to be one. This was introduced in PR #231 and causes ModuleNotFoundError when running install.sh or 'python setup.py bdist_wheel'. Replace the package-style import with importlib.util.spec_from_file_location to load the module by file path, which does not require __init__.py. --- setup.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index c4d74ae929..2c57a8a974 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,6 @@ from pathlib import Path from torch.utils.cpp_extension import CUDAExtension, CUDA_HOME from wheel.bdist_wheel import bdist_wheel as _bdist_wheel -from scripts.generate_pyi import generate_pyi_file DG_SKIP_CUDA_BUILD = int(os.getenv('DG_SKIP_CUDA_BUILD', '0')) == 1 @@ -126,7 +125,13 @@ def run(self): build_py.run(self) def generate_pyi_file(self): - generate_pyi_file(name='_C', root='./csrc', output_dir='./stubs') + import importlib.util + spec = importlib.util.spec_from_file_location( + 'generate_pyi', os.path.join(current_dir, 'scripts', 'generate_pyi.py') + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + mod.generate_pyi_file(name='_C', root='./csrc', output_dir='./stubs') pyi_source = os.path.join(current_dir, 'stubs', '_C.pyi') pyi_target = os.path.join(self.build_lib, 'deep_gemm', '_C.pyi')