Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / __init__.py: 92%

1# Copyright 2025-2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Public interfaces for SAPP-ND memory estimation.""" 

16 

17import importlib 

18from typing import TYPE_CHECKING, Any 

19 

20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.hook_base import MemEvalHook, hook_runner 

21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory, Unit 

22from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType 

23 

24if TYPE_CHECKING: 

25 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2, estimate_memory 

26 

27# The estimator loads on first use. It reaches the cost model's config, whose 

28# variables import this package's size module, so loading it here closed an 

29# import cycle for whichever module of that ring a process imported first. 

30_ESTIMATOR = "hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2" 

31_ESTIMATOR_NAMES = ("EvaluatorV2", "estimate_memory") 

32 

33 

34# PEP 562 names this hook: a module-level dunder, which the naming rule for functions does not allow. 

35def __getattr__(name: str) -> Any: # pylint: disable=invalid-name 

36 """Load the estimator's public names when they are first asked for (PEP 562). 

37 

38 Args: 

39 name: The attribute looked up on this package. 

40 

41 Returns: 

42 The estimator's class or function of that name. 

43 

44 Raises: 

45 AttributeError: For a name this package does not provide. 

46 """ 

47 if name in _ESTIMATOR_NAMES: 

48 return getattr(importlib.import_module(_ESTIMATOR), name) 

49 raise AttributeError(f"module {__name__!r} has no attribute {name!r}") 

50 

51 

52__all__ = [ 

53 "EvaluatorV2", 

54 "LayerType", 

55 "MemEvalHook", 

56 "Memory", 

57 "Unit", 

58 "estimate_memory", 

59 "hook_runner", 

60]