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
« 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."""
17from __future__ import annotations
19import logging
20import os
21from pathlib import Path
22import subprocess
23import sys
24import threading
25from typing import Optional, Sequence
27logger = logging.getLogger(__name__)
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)
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
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 )
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.
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.
77 Args:
78 config_path: Training YAML used by the native Trainer.
79 overrides: Trainer CLI overrides forwarded unchanged to generation.
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
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
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
124__all__ = ["start_codegen_sidecar"]