已合并
fix lazy_init and api export #36707
bellatan创建于 5月26日
fix lazy_init and api export #36707
已合并
共 4 个文件变更+123-19
| @@ -1,7 +1,5 @@ | |||
| 1 | # Owner(s): ["module: unknown"] | 1 | # Owner(s): ["module: unknown"] |
| 2 | -import json | ||
| 3 | import os | 2 | import os |
| 4 | -import statistics | ||
| 5 | import subprocess | 3 | import subprocess |
| 6 | import sys | 4 | import sys |
| 7 | import textwrap | 5 | import textwrap |
| @@ -28,13 +26,17 @@ EXPECTED_LOADED_MODULES = [ | |||
| 28 | "torch_npu.profiler", | 26 | "torch_npu.profiler", |
| 29 | "torch_npu.distributed", | 27 | "torch_npu.distributed", |
| 30 | "torch_npu.distributed.rpc", | 28 | "torch_npu.distributed.rpc", |
| 29 | + "torch_npu.distributed.nn", | ||
| 30 | + "torch_npu.distributed.nn.functional", | ||
| 31 | "torch_npu.op_plugin", | 31 | "torch_npu.op_plugin", |
| 32 | "torch_npu.op_plugin.meta", | 32 | "torch_npu.op_plugin.meta", |
| 33 | "torch_npu.op_plugin.meta._meta_registrations", | 33 | "torch_npu.op_plugin.meta._meta_registrations", |
| 34 | + "torch_npu.asd.checksum", | ||
| 34 | "torch_npu.utils._dynamo", | 35 | "torch_npu.utils._dynamo", |
| 35 | "torch_npu.utils._inductor", | 36 | "torch_npu.utils._inductor", |
| 36 | "torch_npu.utils.custom_ops", | 37 | "torch_npu.utils.custom_ops", |
| 37 | "torch_npu.utils.patch_getenv", | 38 | "torch_npu.utils.patch_getenv", |
| 39 | + "torch_npu.utils.syncbatchnorm", | ||
| 38 | ] | 40 | ] |
| 39 | 41 | ||
| 40 | EXPECTED_NOT_LOADED_MODULES = [ | 42 | EXPECTED_NOT_LOADED_MODULES = [ |
| @@ -469,6 +471,83 @@ class TestTorchNpuBootstrap(TestCase): | |||
| 469 | """ | 471 | """ |
| 470 | ) | 472 | ) |
| 471 | 473 | ||
| 474 | + def test_09_legacy_submodule_attribute_compatibility(self): | ||
| 475 | + self._run_python( | ||
| 476 | + """ | ||
| 477 | + import sys | ||
| 478 | + import torch_npu | ||
| 479 | + | ||
| 480 | + # Old behavior: importing torch_npu also exposed these child modules | ||
| 481 | + # as attributes on their parent packages. | ||
| 482 | + | ||
| 483 | + assert "torch_npu.asd.checksum" in sys.modules | ||
| 484 | + assert hasattr(torch_npu, "asd") | ||
| 485 | + assert hasattr(torch_npu.asd, "checksum") | ||
| 486 | + assert torch_npu.asd.checksum is sys.modules["torch_npu.asd.checksum"] | ||
| 487 | + assert hasattr(torch_npu.asd.checksum, "_matmul_checksum") | ||
| 488 | + | ||
| 489 | + assert "torch_npu.utils.syncbatchnorm" in sys.modules | ||
| 490 | + assert hasattr(torch_npu.utils, "syncbatchnorm") | ||
| 491 | + assert torch_npu.utils.syncbatchnorm is ( | ||
| 492 | + sys.modules["torch_npu.utils.syncbatchnorm"] | ||
| 493 | + ) | ||
| 494 | + """ | ||
| 495 | + ) | ||
| 496 | + | ||
| 497 | + def test_10_legacy_top_level_distributed_api_compatibility(self): | ||
| 498 | + self._run_python( | ||
| 499 | + """ | ||
| 500 | + import torch_npu | ||
| 501 | + from torch.distributed.fsdp import sharded_grad_scaler | ||
| 502 | + from torch_npu._C._distributed_c10d import ParallelStore | ||
| 503 | + from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler | ||
| 504 | + | ||
| 505 | + # Old behavior: these names were visible on torch_npu top-level | ||
| 506 | + # due to module-scope imports in the old monolithic __init__.py. | ||
| 507 | + assert hasattr(torch_npu, "ParallelStore") | ||
| 508 | + assert torch_npu.ParallelStore is ParallelStore | ||
| 509 | + assert "ParallelStore" not in torch_npu.__all__ | ||
| 510 | + | ||
| 511 | + assert hasattr(torch_npu, "_ShardedGradScaler") | ||
| 512 | + assert torch_npu._ShardedGradScaler is _ShardedGradScaler | ||
| 513 | + assert "_ShardedGradScaler" not in torch_npu.__all__ | ||
| 514 | + | ||
| 515 | + # The FSDP patch behavior should still be preserved. | ||
| 516 | + assert sharded_grad_scaler.ShardedGradScaler is _ShardedGradScaler | ||
| 517 | + """ | ||
| 518 | + ) | ||
| 519 | + | ||
| 520 | + def test_11_import_does_not_trigger_device_count(self): | ||
| 521 | + cases = [ | ||
| 522 | + "import torch; import torch_npu", | ||
| 523 | + "import torch_npu", | ||
| 524 | + ] | ||
| 525 | + | ||
| 526 | + for import_code in cases: | ||
| 527 | + self._run_python( | ||
| 528 | + f""" | ||
| 529 | + import os | ||
| 530 | + | ||
| 531 | + {import_code} | ||
| 532 | + | ||
| 533 | + assert torch_npu.npu.is_initialized() is False, ( | ||
| 534 | + "import torch_npu unexpectedly triggered NPU lazy init" | ||
| 535 | + ) | ||
| 536 | + | ||
| 537 | + # Regression test for import-time low-level device probing. | ||
| 538 | + # If import torch_npu has already called low-level NPU device count, | ||
| 539 | + # changing ASCEND_RT_VISIBLE_DEVICES here will no longer take effect. | ||
| 540 | + os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "32" | ||
| 541 | + | ||
| 542 | + raw_count = torch_npu._C._npu_getDeviceCount() | ||
| 543 | + assert raw_count == 0, ( | ||
| 544 | + f"import torch_npu unexpectedly triggered low-level NPU " | ||
| 545 | + f"device probing, raw_count={{raw_count}}" | ||
| 546 | + ) | ||
| 547 | + | ||
| 548 | + assert torch_npu.npu.is_initialized() is False | ||
| 549 | + """ | ||
| 550 | + ) | ||
| 472 | 551 | ||
| 473 | if __name__ == "__main__": | 552 | if __name__ == "__main__": |
| 474 | run_tests() | 553 | run_tests() |
| @@ -28,15 +28,6 @@ os.environ["TORCH_DEVICE_BACKEND_AUTOLOAD"] = "0" | |||
| 28 | 28 | ||
| 29 | import torch | 29 | import torch |
| 30 | 30 | ||
| 31 | -# Import-time env access logging patch. Keep early to capture initialization-time getenv. | ||
| 32 | -import torch_npu.utils.patch_getenv | ||
| 33 | -from torch_npu._init.core.module_loader import _load_core_modules | ||
| 34 | -from torch_npu._init.core.optional_features import _enable_optional_features | ||
| 35 | -from torch_npu._init.core.runtime_lifecycle import _initialize_runtime_lifecycle | ||
| 36 | -from torch_npu._init.patches.patch_manager import _apply_all_patches | ||
| 37 | -from torch_npu._init.registry.registry_manager import _register_components | ||
| 38 | -from torch_npu.version import __version__ as __version__ | ||
| 39 | - | ||
| 40 | 31 | ||
| 41 | def _check_device_conflict(): | 32 | def _check_device_conflict(): |
| 42 | acc = torch._C._get_accelerator() | 33 | acc = torch._C._get_accelerator() |
| @@ -59,23 +50,33 @@ def _check_device_conflict(): | |||
| 59 | ) | 50 | ) |
| 60 | 51 | ||
| 61 | 52 | ||
| 62 | -def _initialize(): | 53 | +_check_device_conflict() |
| 63 | - # 1. pre-init checks | ||
| 64 | - _check_device_conflict() | ||
| 65 | 54 | ||
| 66 | - # 2. core modules, registration side effects and public API export | 55 | + |
| 56 | +# Import-time env access logging patch. Keep early to capture initialization-time getenv. | ||
| 57 | +import torch_npu.utils.patch_getenv | ||
| 58 | +from torch_npu._init.core.module_loader import _load_core_modules | ||
| 59 | +from torch_npu._init.core.optional_features import _enable_optional_features | ||
| 60 | +from torch_npu._init.core.runtime_lifecycle import _initialize_runtime_lifecycle | ||
| 61 | +from torch_npu._init.patches.patch_manager import _apply_all_patches | ||
| 62 | +from torch_npu._init.registry.registry_manager import _register_components | ||
| 63 | +from torch_npu.version import __version__ as __version__ | ||
| 64 | + | ||
| 65 | + | ||
| 66 | +def _initialize(): | ||
| 67 | + # 1. core modules, registration side effects and public API export | ||
| 67 | _load_core_modules() | 68 | _load_core_modules() |
| 68 | 69 | ||
| 69 | - # 3. backend and framework integration registration | 70 | + # 2. backend and framework integration registration |
| 70 | _register_components() | 71 | _register_components() |
| 71 | 72 | ||
| 72 | - # 4. apply patches | 73 | + # 3. apply patches |
| 73 | _apply_all_patches() | 74 | _apply_all_patches() |
| 74 | 75 | ||
| 75 | - # 5. final extension barrier and shutdown hook | 76 | + # 4. final extension barrier and shutdown hook |
| 76 | _initialize_runtime_lifecycle() | 77 | _initialize_runtime_lifecycle() |
| 77 | 78 | ||
| 78 | - # 6. optional runtime features | 79 | + # 5. optional runtime features |
| 79 | _enable_optional_features() | 80 | _enable_optional_features() |
| 80 | 81 | ||
| 81 | 82 | ||
| @@ -95,6 +95,20 @@ def _export_lazy_python_apis(globals_dict, all_list): | |||
| 95 | _append_unique(all_list, _LAZY_PYTHON_SYMBOLS.keys()) | 95 | _append_unique(all_list, _LAZY_PYTHON_SYMBOLS.keys()) |
| 96 | 96 | ||
| 97 | 97 | ||
| 98 | +def _export_legacy_distributed_apis(globals_dict): | ||
| 99 | + """ | ||
| 100 | + Export legacy distributed-related top-level APIs. | ||
| 101 | + Rule: | ||
| 102 | + - torch_npu._C._distributed_c10d.ParallelStore -> torch_npu.ParallelStore | ||
| 103 | + - torch_npu.npu.amp.sharded_grad_scaler._ShardedGradScaler -> torch_npu._ShardedGradScaler | ||
| 104 | + """ | ||
| 105 | + from torch_npu._C._distributed_c10d import ParallelStore | ||
| 106 | + from torch_npu.npu.amp.sharded_grad_scaler import _ShardedGradScaler | ||
| 107 | + | ||
| 108 | + globals_dict["ParallelStore"] = ParallelStore | ||
| 109 | + globals_dict["_ShardedGradScaler"] = _ShardedGradScaler | ||
| 110 | + | ||
| 111 | + | ||
| 98 | def _export_public_apis(): | 112 | def _export_public_apis(): |
| 99 | """ | 113 | """ |
| 100 | Export torch_npu public APIs. | 114 | Export torch_npu public APIs. |
| @@ -109,6 +123,10 @@ def _export_public_apis(): | |||
| 109 | 123 | ||
| 110 | 3. DType symbols: | 124 | 3. DType symbols: |
| 111 | - torch_npu._C._cd.DType.<dtype_name> -> torch_npu.<dtype_name> | 125 | - torch_npu._C._cd.DType.<dtype_name> -> torch_npu.<dtype_name> |
| 126 | + | ||
| 127 | + 4. legacy distributed APIs: | ||
| 128 | + - torch_npu._C._distributed_c10d.ParallelStore -> torch_npu.ParallelStore | ||
| 129 | + - torch_npu.npu.amp.sharded_grad_scaler._ShardedGradScaler -> torch_npu._ShardedGradScaler | ||
| 112 | """ | 130 | """ |
| 113 | 131 | ||
| 114 | _export_dtype_symbols() | 132 | _export_dtype_symbols() |
| @@ -120,3 +138,4 @@ def _export_public_apis(): | |||
| 120 | _export_lazy_python_apis(globals_dict, all_list) | 138 | _export_lazy_python_apis(globals_dict, all_list) |
| 121 | _export_npu_ops(globals_dict, all_list) | 139 | _export_npu_ops(globals_dict, all_list) |
| 122 | _export_dtype_symbols() | 140 | _export_dtype_symbols() |
| 141 | + _export_legacy_distributed_apis(globals_dict) | ||
| @@ -197,10 +197,15 @@ def _load_registration_modules(): | |||
| 197 | import torch_npu.npu.aclnn # noqa: F401 | 197 | import torch_npu.npu.aclnn # noqa: F401 |
| 198 | import torch_npu.op_plugin | 198 | import torch_npu.op_plugin |
| 199 | import torch_npu.optim # noqa: F401 | 199 | import torch_npu.optim # noqa: F401 |
| 200 | + | ||
| 200 | from torch_npu.op_plugin.meta import _meta_registrations # noqa: F401 | 201 | from torch_npu.op_plugin.meta import _meta_registrations # noqa: F401 |
| 201 | from torch_npu.utils import custom_ops # noqa: F401 | 202 | from torch_npu.utils import custom_ops # noqa: F401 |
| 202 | from torch_npu.utils._afd_ops import initialize_afd_bindings | 203 | from torch_npu.utils._afd_ops import initialize_afd_bindings |
| 203 | 204 | ||
| 205 | + # export for submodule | ||
| 206 | + import torch_npu.asd.checksum # noqa: F401 | ||
| 207 | + import torch_npu.utils.syncbatchnorm # noqa: F401 | ||
| 208 | + | ||
OO 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。   【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。 | |||
| 204 | importlib.import_module("torch_npu._op_plugin_docs") | 209 | importlib.import_module("torch_npu._op_plugin_docs") |
| 205 | if hasattr(torch_npu, "_op_plugin_docs"): | 210 | if hasattr(torch_npu, "_op_plugin_docs"): |
| 206 | delattr(torch_npu, "_op_plugin_docs") | 211 | delattr(torch_npu, "_op_plugin_docs") |
此条代码评论区间+200至+208
【openlibing.ci】识别到代码检查告警抑制注释，匹配工具：flake8，请Committer检视其合理性。