Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / optimizer / utils.py: 76%

1# Copyright 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 

16"""Install ``info_rank0`` / ``warning_rank0`` on ``logging.Logger``.""" 

17 

18import logging 

19import os 

20 

21import torch 

22 

23 

24def _get_rank() -> int: 

25 try: 

26 if torch.distributed.is_initialized(): 

27 return torch.distributed.get_rank() 

28 except (ImportError, RuntimeError): 

29 pass 

30 return int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) 

31 

32 

33def info_rank0(self, msg, *args, **kwargs) -> None: 

34 """Log an informational message only on rank zero.""" 

35 if _get_rank() == 0: 

36 kwargs.setdefault("stacklevel", 2) 

37 self.info(msg, *args, **kwargs) 

38 

39 

40def warning_rank0(self, msg, *args, **kwargs) -> None: 

41 """Log a warning message only on rank zero.""" 

42 if _get_rank() == 0: 

43 kwargs.setdefault("stacklevel", 2) 

44 self.warning(msg, *args, **kwargs) 

45 

46 

47def debug_rank0(self, msg, *args, **kwargs) -> None: 

48 """Log a debug message only on rank zero.""" 

49 if _get_rank() == 0: 

50 kwargs.setdefault("stacklevel", 2) 

51 self.debug(msg, *args, **kwargs) 

52 

53 

54def get_device_count() -> int: 

55 """Return the active accelerator count, defaulting to 1.""" 

56 npu = getattr(torch, "npu", None) 

57 if npu is not None and npu.is_available(): 

58 return npu.device_count() 

59 if torch.cuda.is_available(): 

60 return torch.cuda.device_count() 

61 return 1 

62 

63 

64def get_current_device() -> torch.device: 

65 """Return the current accelerator device, or CPU when none is available.""" 

66 npu = getattr(torch, "npu", None) 

67 if npu is not None and npu.is_available(): 

68 return torch.device("npu", npu.current_device()) 

69 if torch.cuda.is_available(): 

70 return torch.device("cuda", torch.cuda.current_device()) 

71 return torch.device("cpu") 

72 

73 

74def empty_accelerator_cache() -> None: 

75 """Clear the active accelerator cache when supported.""" 

76 npu = getattr(torch, "npu", None) 

77 if npu is not None and npu.is_available(): 

78 npu.empty_cache() 

79 elif torch.cuda.is_available(): 

80 torch.cuda.empty_cache() 

81 

82 

83_INSTALLED = False 

84 

85 

86def _install_logger_methods() -> None: 

87 global _INSTALLED 

88 if _INSTALLED: 

89 return 

90 logging.Logger.info_rank0 = info_rank0 

91 logging.Logger.warning_rank0 = warning_rank0 

92 logging.Logger.debug_rank0 = debug_rank0 

93 _INSTALLED = True 

94 

95 

96_install_logger_methods()