From b0732018e75d22f2e26fa7284bcd5d21b98cd362 Mon Sep 17 00:00:00 2001 From: LeoBorcherding Date: Mon, 11 May 2026 05:59:13 -0500 Subject: [PATCH] fix: give mod stubs __path__ and pre-stub _tensor to fix 'not a package' import error _make_mod_stub now sets __path__=[] so Python treats stub modules as packages. Without it, any import of a submodule raises "is not a package". Also pre-stub torch.distributed._tensor and its submodules so that _tensor/__init__.py (which re-exports from torch.distributed.tensor) never runs and torchao's `from torch.distributed._tensor import DTensor` gets a harmless stub instead of crashing. --- studio/backend/core/training/worker.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index a654b51f9b..c2009b1230 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1100,8 +1100,14 @@ def run_training_process( } # Helper: build a ModuleType stub whose __getattr__ auto-creates child stubs. + # __path__ is set to [] so Python treats the stub as a package — without it, + # any attempt to import a submodule (e.g. "torch.distributed.tensor._foo") + # raises "is not a package" because Python checks __path__ before looking + # in sys.modules for the child. def _make_mod_stub(mod_name): m = _types.ModuleType(mod_name) + m.__path__ = [] # marks this as a package to the import system + m.__package__ = mod_name def _ga(attr, _m=m, _n=mod_name): if attr.startswith("__"): raise AttributeError(attr) @@ -1194,6 +1200,15 @@ def run_training_process( "torch.distributed.tensor._ops._conv_ops", "torch.distributed.tensor._dtensor_spec", "torch.distributed.tensor.placement_types", + # torch.distributed._tensor is the canonical private package; + # its __init__.py tries to re-export submodules from + # torch.distributed.tensor (which we stubbed above), causing + # "is not a package" errors. Stubbing _tensor directly + # short-circuits that __init__ so torchao's + # `from torch.distributed._tensor import DTensor` gets a stub. + "torch.distributed._tensor", + "torch.distributed._tensor.placement_types", + "torch.distributed._tensor.api", ): if _dist_name not in sys.modules: sys.modules[_dist_name] = _make_mod_stub(_dist_name)