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

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 if _get_rank() == 0: 

35 kwargs.setdefault("stacklevel", 2) 

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

37 

38 

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

40 if _get_rank() == 0: 

41 kwargs.setdefault("stacklevel", 2) 

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

43 

44 

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

46 if _get_rank() == 0: 

47 kwargs.setdefault("stacklevel", 2) 

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

49 

50 

51def get_device_count() -> int: 

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

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

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

55 return npu.device_count() 

56 if torch.cuda.is_available(): 

57 return torch.cuda.device_count() 

58 return 1 

59 

60 

61def get_current_device() -> torch.device: 

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

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

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

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

66 if torch.cuda.is_available(): 

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

68 return torch.device("cpu") 

69 

70 

71def empty_accelerator_cache() -> None: 

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

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

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

75 npu.empty_cache() 

76 elif torch.cuda.is_available(): 

77 torch.cuda.empty_cache() 

78 

79 

80_INSTALLED = False 

81 

82 

83def _install_logger_methods() -> None: 

84 global _INSTALLED 

85 if _INSTALLED: 

86 return 

87 logging.Logger.info_rank0 = info_rank0 

88 logging.Logger.warning_rank0 = warning_rank0 

89 logging.Logger.debug_rank0 = debug_rank0 

90 _INSTALLED = True 

91 

92 

93_install_logger_methods()