""" RPyC Service wrapper around the engine supporting tensor parallelism. If there are 4 GPUs, initialize 4 services. Rank 0 is the main, the rest are helpers. Main runs the full engine and broadcasts tensors to helpers. Start each EngineService in a separate process. Connect and init each. Use RPyC to communicate with main worker to add jobs and poll generated tokens. """ import rpyc from rpyc.utils.server import ThreadedServer import os import torch import torch.multiprocessing as mp import torch.distributed as dist import time import pickle from tqdm import tqdm import datetime from suno_utils.gpt.engine import Engine, Request def pickle_result(func): """ Decorator to pickle the result of a function. """ def wrapped(*args, **kwargs): return pickle.dumps(func(*args, **kwargs)) return wrapped def terminate_processes(): import os import signal import psutil def kill_child_processes(parent_pid, sig=signal.SIGTERM): try: parent = psutil.Process(parent_pid) except psutil.NoSuchProcess: return children = parent.children(recursive=True) for process in children: process.send_signal(sig) # Get the current process ID pid = os.getpid() # Kill all child processes kill_child_processes(pid) class EngineService(rpyc.Service): def on_connect(self, conn): from torch._inductor import config config.triton.cudagraph_trees = False def exposed_init( self, ckpt_path: str, tokenizer_path: str, max_sequences: int, compile: bool, min_batch_size: int, batch_increment: int, rank: int, world_size: int, max_length_s: float = None, ): self.rank = rank self.world_size = world_size os.environ["LOCAL_WORLD_SIZE"] = str(world_size) os.environ["LOCAL_RANK"] = str(rank) os.environ["ENABLE_INTRA_NODE_COMM"] = "1" print(f"Rank {rank} world size {world_size}, initing process group") torch.cuda.set_device(rank) # leaving backend blank inits both gloo and nccl backends. gloo is necessary for cpu data dist.init_process_group( init_method="tcp://localhost:28765", rank=rank, world_size=world_size, timeout=datetime.timedelta(hours=10000), # dont timeout while waiting for requests ) print( f"Initialized process group for rank {dist.get_rank()} and world size {dist.get_world_size()}" ) dist.barrier() # sanity check process group initialization self.exposed_engine = Engine( ckpt_path, tokenizer_path, max_sequences=max_sequences, compile=compile, min_batch_size=min_batch_size, batch_increment=batch_increment, use_tp=world_size > 1, max_length_s=max_length_s, ) print(f"Rank {rank} world size {world_size}, starting engine") self.exposed_engine.start() self.exposed_engine.wait_for_warmup() print(f"Rank {rank} world size {world_size}, engine started") def get_engine(self): return self.exposed_engine @pickle_result def get_generated_codes(self, job_id: str, start=0): """ Retrieve generated codes for a given job starting from a specified index. Args: job_id (str): The unique identifier for the job. start (int, optional): The starting index from which to retrieve the generated codes. Defaults to 0. Returns: list or None: A list of generated codes if available, an empty list if the job is not completed and no codes are available, or None if the job is completed and no codes are available. """ job = self.exposed_engine.jobs[job_id] codes = job.generated_tokens[start:] if len(codes) > 0: return codes if job.completed: return None # signal completed with None else: return [] @pickle_result def get_job(self, job_id: str): return self.exposed_engine.jobs[job_id] @pickle_result def get_model_cfg(self): return self.exposed_engine.model.config @pickle_result def get_tokenizer(self): return self.exposed_engine.tokenizer def get_free_caches(self): return len(self.exposed_engine.free_caches()) def get_used_caches(self): return len(self.exposed_engine.used_caches()) def add_request(self, pickled_request: bytes): request: Request = pickle.loads(pickled_request) return self.exposed_engine.add_request_blocking(request) def run_request(self, picked_requests: bytes, **kwargs): requests = pickle.loads(picked_requests) return self.exposed_engine.run_request(requests, **kwargs) def init_env(port: int): t = ThreadedServer( EngineService, port=port, protocol_config={"allow_all_attrs": True, "allow_pickle": True}, ) t.start() class EngineRpcClient: """ A client for interacting with the EngineService via RPC. This client provides methods to add and run requests on the EngineService. It handles the serialization and deserialization of requests using pickle. """ def __init__(self, engine_service: EngineService): self.engine_service = engine_service def add_request(self, request: Request): return self.engine_service.add_request(pickle.dumps(request)) def run_request(self, request: Request, **kwargs): return self.engine_service.run_request(pickle.dumps(request), **kwargs) def get_generated_codes(self, job_id: str, start=0): return pickle.loads(self.engine_service.get_generated_codes(job_id, start)) def get_job(self, job_id: str): return pickle.loads(self.engine_service.get_job(job_id)) def get_model_cfg(self): return pickle.loads(self.engine_service.get_model_cfg()) def get_tokenizer(self): return pickle.loads(self.engine_service.get_tokenizer()) def start_service_processes( ckpt_path: str, tokenizer_path: str, max_sequences: int, min_batch_size: int = 4, batch_increment: int = 8, max_length_s: float = None, compile=True, port: int = 18861, world_size: int = 1, ): """ Initialize and start multiple service processes for the GPT engine. This function sets up and starts a specified number of service processes, each running on a different port. It also waits for all services to start and initializes/warms up the engines. Returns: EngineRpcClient: A client to interact with the initialized engine services. """ processes = [] for i in range(world_size): p = mp.get_context("forkserver").Process(target=init_env, args=(port + i,)) p.start() processes.append(p) # wait for all the services to start for i in tqdm(range(world_size), desc="Waiting for services to start"): service_started = False while not service_started: try: c = rpyc.connect( "localhost", port + i, config={"allow_all_attrs": True, "allow_pickle": True}, ) service_started = True except BaseException: print(f"Service {i} not started yet, waiting...") time.sleep(1) # connect to all the services and initialize the engines connections = [] inits = [] for i in tqdm(range(world_size), desc="Initializing engines"): c = rpyc.connect( "localhost", port + i, config={"allow_all_attrs": True, "allow_pickle": True}, ) async_init = rpyc.async_(c.root.init) init_state = async_init( ckpt_path, tokenizer_path, max_sequences, compile=compile, min_batch_size=min_batch_size, batch_increment=batch_increment, rank=i, world_size=world_size, max_length_s=max_length_s, ) init_state.set_expiry(2 * 60) connections.append(c) inits.append(init_state) for i in range(world_size): inits[i].wait() if inits[i].error or inits[i].expired: terminate_processes() raise Exception(f"Rank {i} init failed: {inits[i].error}, {inits[i].value}") print(f"Rank {i} init done.") connections[0]._config["sync_request_timeout"] = None # No timeout time.sleep(5) return EngineRpcClient(connections[0].root)