Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / train_sidecar.py: 94%

36 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-11-15 07:18 +0800

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"""Best-effort Codegen sidecar for the standard training entry point.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20import os 

21from pathlib import Path 

22import subprocess 

23import sys 

24import threading 

25from typing import Optional, Sequence 

26 

27logger = logging.getLogger(__name__) 

28 

29_DISTRIBUTED_ENV_NAMES = ( 

30 "RANK", 

31 "LOCAL_RANK", 

32 "WORLD_SIZE", 

33 "LOCAL_WORLD_SIZE", 

34 "GROUP_RANK", 

35 "ROLE_RANK", 

36 "ROLE_WORLD_SIZE", 

37 "MASTER_ADDR", 

38 "MASTER_PORT", 

39 "TORCHELASTIC_RUN_ID", 

40 "TORCHELASTIC_RESTART_COUNT", 

41 "TORCHELASTIC_MAX_RESTARTS", 

42 "TORCHELASTIC_ERROR_FILE", 

43) 

44 

45 

46def _codegen_environment() -> dict[str, str]: 

47 """Return a child environment without torchrun distributed identity.""" 

48 child_env = os.environ.copy() 

49 for name in _DISTRIBUTED_ENV_NAMES: 

50 child_env.pop(name, None) 

51 return child_env 

52 

53 

54def _monitor_sidecar(process: subprocess.Popen, config_path: str) -> None: 

55 """Reap the Codegen process and report failure without stopping training.""" 

56 return_code = process.wait() 

57 if return_code == 0: 

58 logger.info("Codegen sidecar completed for %s", config_path) 

59 return 

60 logger.warning( 

61 "Codegen sidecar failed for %s with exit code %s; native training continues", 

62 config_path, 

63 return_code, 

64 ) 

65 

66 

67def start_codegen_sidecar( 

68 config_path: str | Path, 

69 overrides: Sequence[str] = (), 

70) -> Optional[subprocess.Popen]: 

71 """Start device-free bundle generation on global rank 0. 

72 

73 The child receives a copy of the parent environment with torchrun identity 

74 removed. The parent environment and native training process are unchanged. 

75 Startup and generation failures are best effort and never block training. 

76 

77 Args: 

78 config_path: Training YAML used by the native Trainer. 

79 overrides: Trainer CLI overrides forwarded unchanged to generation. 

80 

81 Returns: 

82 The rank-0 Codegen child process, or ``None`` on nonzero ranks and when 

83 the child cannot be started. 

84 """ 

85 if int(os.environ.get("RANK", "0")) != 0: 

86 return None 

87 

88 resolved_config = str(Path(config_path).resolve()) 

89 command = [ 

90 sys.executable, 

91 "-m", 

92 "hyper_parallel.codegen", 

93 "--verbose", 

94 "generate", 

95 "--config", 

96 resolved_config, 

97 *overrides, 

98 ] 

99 try: 

100 process = subprocess.Popen( # pylint: disable=consider-using-with 

101 command, 

102 env=_codegen_environment(), 

103 close_fds=True, 

104 ) 

105 except OSError as exc: 

106 logger.warning( 

107 "Codegen sidecar could not start for %s: %s; native training continues", 

108 resolved_config, 

109 exc, 

110 ) 

111 return None 

112 

113 monitor = threading.Thread( 

114 target=_monitor_sidecar, 

115 args=(process, resolved_config), 

116 name="hyper-parallel-codegen-sidecar", 

117 daemon=True, 

118 ) 

119 monitor.start() 

120 logger.info("Started Codegen sidecar process %s for %s", process.pid, resolved_config) 

121 return process 

122 

123 

124__all__ = ["start_codegen_sidecar"]