File size: 8,106 Bytes
74da989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | # SPDX-License-Identifier: Apache-2.0
import atexit
import contextlib
import multiprocessing as mp
import os
import signal
import socket
import time
from collections.abc import Callable
from multiprocessing.process import BaseProcess
from typing import Any
import trainer.envs as envs
from trainer.trainer_args import TrainerArgs
from trainer.logger import init_logger
from trainer.pipelines.pipeline_batch_info import ForwardBatch
from trainer.worker.executor import Executor
from trainer.worker.gpu_worker import run_worker_process
logger = init_logger(__name__)
class MultiprocExecutor(Executor):
def _init_executor(self) -> None:
self.world_size = self.trainer_args.num_gpus
self.shutting_down = False
# this will force the use of the `spawn` multiprocessing start if cuda
# is initialized
mp.set_start_method("spawn", force=True)
self.workers: list[BaseProcess] = []
self.worker_pipes = []
# Check if master_port is provided in trainer_args
if hasattr(
self.trainer_args,
'master_port') and self.trainer_args.master_port is not None:
self.master_port = self.trainer_args.master_port
logger.info("Using provided master port: %s", self.master_port)
else:
# Auto-find available port
import random
for port in range(29503 + random.randint(0, 10000), 65535):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
if s.connect_ex(('localhost', port)) != 0:
self.master_port = port
break
else:
raise ValueError("No unused port found to use as master port")
logger.info("Auto-selected master port: %s", self.master_port)
# Create pipes and start workers
for rank in range(self.world_size):
executor_pipe, worker_pipe = mp.Pipe(duplex=True)
self.worker_pipes.append(executor_pipe)
worker = mp.Process(target=run_worker_process,
name=f"FVWorkerProc-{rank}",
kwargs=dict(trainer_args=self.trainer_args,
local_rank=rank,
rank=rank,
pipe=worker_pipe,
master_port=self.master_port))
worker.start()
self.workers.append(worker)
# Wait for all workers to be ready
for idx, pipe in enumerate(self.worker_pipes):
data = pipe.recv()
if data["status"] != "ready" or data["local_rank"] != idx:
raise RuntimeError(f"Worker {idx} failed to start")
logger.info("%d workers ready", self.world_size)
# Register shutdown on exit
atexit.register(self.shutdown)
def execute_forward(self, forward_batch: ForwardBatch,
trainer_args: TrainerArgs) -> ForwardBatch:
responses = self.collective_rpc("execute_forward",
kwargs={
"forward_batch": forward_batch,
"trainer_args": trainer_args
})
output = responses[0]["output_batch"]
logging_info = None
if envs.TRAINER_STAGE_LOGGING:
logging_info = responses[0]["logging_info"]
else:
logging_info = None
result_batch = ForwardBatch(data_type=forward_batch.data_type,
output=output,
logging_info=logging_info)
return result_batch
def set_lora_adapter(self,
lora_nickname: str,
lora_path: str | None = None) -> None:
responses = self.collective_rpc("set_lora_adapter",
kwargs={
"lora_nickname": lora_nickname,
"lora_path": lora_path
})
for i, response in enumerate(responses):
if response["status"] != "lora_adapter_set":
raise RuntimeError(
f"Worker {i} failed to set LoRA adapter to {lora_path}")
def collective_rpc(self,
method: str | Callable,
timeout: float | None = None,
args: tuple = (),
kwargs: dict | None = None) -> list[Any]:
kwargs = kwargs or {}
try:
for pipe in self.worker_pipes:
pipe.send({"method": method, "args": args, "kwargs": kwargs})
responses = []
for pipe in self.worker_pipes:
response = pipe.recv()
responses.append(response)
return responses
except TimeoutError as e:
raise TimeoutError(f"RPC call to {method} timed out.") from e
except KeyboardInterrupt as e:
# if we catch a KeyboardInterrupt, user wants to stop the execution.
# we need to send a signal to all workers to stop.
logger.info(
"Received KeyboardInterrupt, sending SIGINT to all workers")
for worker in self.workers:
if worker.pid is not None:
os.kill(worker.pid, signal.SIGINT)
raise e
except Exception as e:
raise e
def shutdown(self) -> None:
"""Properly shut down the executor and its workers"""
if hasattr(self, 'shutting_down') and self.shutting_down:
return # Prevent multiple shutdown calls
logger.info("Shutting down MultiprocExecutor...")
self.shutting_down = True
# First try gentle termination
try:
# Send termination message to all workers
for pipe in self.worker_pipes:
with contextlib.suppress(Exception):
pipe.send({"method": "shutdown", "args": (), "kwargs": {}})
# Give workers some time to exit gracefully
start_time = time.perf_counter()
while time.perf_counter() - start_time < 5.0: # 5 seconds timeout
if all(not worker.is_alive() for worker in self.workers):
break
time.sleep(0.1)
# Force terminate any remaining workers
for worker in self.workers:
if worker.is_alive():
worker.terminate()
# Final timeout for terminate
start_time = time.perf_counter()
while time.perf_counter() - start_time < 2.0: # 2 seconds timeout
if all(not worker.is_alive() for worker in self.workers):
break
time.sleep(0.1)
# Kill if still alive
for worker in self.workers:
if worker.is_alive():
worker.kill()
worker.join(timeout=1.0)
except Exception as e:
logger.error("Error during shutdown: %s", e)
# Last resort, try to kill all workers
for worker in self.workers:
with contextlib.suppress(Exception):
if worker.is_alive():
worker.kill()
# Clean up pipes
for pipe in self.worker_pipes:
with contextlib.suppress(Exception):
pipe.close()
self.workers = []
self.worker_pipes = []
logger.info("MultiprocExecutor shutdown complete")
def __del__(self):
"""Ensure cleanup on garbage collection"""
self.shutdown()
def __enter__(self):
"""Support for context manager protocol"""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Ensure cleanup when exiting context"""
self.shutdown()
|