From ab15dce3c887b8934ac4717a5ca0a3a98c63412f Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Mon, 3 Aug 2026 13:42:10 -0700 Subject: [PATCH 1/4] move tokenizer to a separate thread, so that it does not block the asyncio event loop --- .../endpoints/chat_completions.py | 53 ++++++++++++++----- .../text_generation_server.py | 11 ++++ tools/run_dynamic_text_generation_server.py | 18 +++++++ 3 files changed, 68 insertions(+), 14 deletions(-) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index e81b5273e84..a039c6cf3de 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -7,6 +7,7 @@ import traceback import uuid import warnings +from functools import partial from megatron.core.inference.inference_request import unwrap_serialized_tensors from megatron.core.inference.sampling_params import SamplingParams @@ -355,6 +356,25 @@ def _replace_prefix_tokens( return previous_turn_token_ids + current_turn_additional_token_ids +def _apply_chat_template_sync( + tokenizer, messages, tools, chat_template_kwargs, add_generation_prompt=True +): + """Apply the chat template and coerce to `list[int]`, for use in a worker thread. + + The coercion runs here too: it walks every token, so leaving it on the event loop + would keep part of the stall this offload exists to remove. + """ + return _coerce_to_token_id_list( + tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=add_generation_prompt, + tools=tools, + **chat_template_kwargs, + ) + ) + + def _coerce_to_token_id_list(result): """Convert the return value of `tokenizer.apply_chat_template` to `list[int]`. @@ -460,14 +480,15 @@ async def chat_completions(): hasattr(tokenizer, 'apply_chat_template') and getattr(tokenizer, "chat_template", None) is not None ): - prompt_tokens = _coerce_to_token_id_list( - tokenizer.apply_chat_template( + prompt_tokens = await asyncio.get_running_loop().run_in_executor( + current_app.config['tokenize_executor'], + partial( + _apply_chat_template_sync, + current_app.config['tokenize_tokenizer'], template_messages, - tokenize=True, - add_generation_prompt=True, - tools=template_tools, - **chat_template_kwargs, - ) + template_tools, + chat_template_kwargs, + ), ) if req.get("prevent_retokenization", True): @@ -508,13 +529,17 @@ async def chat_completions(): ] # Get the templated tokenization of just the previous generation - retokenized_previous_turn_token_ids = _coerce_to_token_id_list( - tokenizer.apply_chat_template( - messages_to_last_assistant_message, - tokenize=True, - add_generation_prompt=False, - tools=template_tools, - **chat_template_kwargs, + retokenized_previous_turn_token_ids = ( + await asyncio.get_running_loop().run_in_executor( + current_app.config['tokenize_executor'], + partial( + _apply_chat_template_sync, + current_app.config['tokenize_tokenizer'], + messages_to_last_assistant_message, + template_tools, + chat_template_kwargs, + add_generation_prompt=False, + ), ) ) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py index af12bc5bbc1..9e7a8da716e 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py @@ -1,9 +1,11 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. import asyncio +import copy import logging import multiprocessing as mp import socket +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from typing import List, Optional @@ -81,6 +83,15 @@ async def _run_text_gen_server( app.config['parsers'] = parsers app.config['verbose'] = verbose + # Applying the chat template is synchronous and O(prompt); on the event loop it + # stalls every other request this replica owns, including delivery of responses + # that already finished. One worker is enough - the point is the yield, not + # throughput. The copy is required: HF tokenizers are not thread-safe. + app.config['tokenize_executor'] = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="tokenize" + ) + app.config['tokenize_tokenizer'] = copy.deepcopy(tokenizer) + # Register all blueprints from the 'endpoints' package for endpoint in endpoints.__all__: app.register_blueprint(endpoint) diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index a33b066e8ea..a07ff444d5c 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -2,9 +2,11 @@ import argparse import asyncio +import logging import torch +from megatron.core import parallel_state from megatron.core.inference.engines import DynamicInferenceEngine from megatron.core.inference.text_generation_server.dynamic_text_gen_server import ( start_text_gen_server, @@ -30,6 +32,12 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser): parser.add_argument( "--parsers", type=str, nargs="+", default=[], help="Parsers to use for parsing the response" ) + parser.add_argument( + "--frontend-replicas", type=int, default=-1, + help="Number of HTTP frontend processes spawned on the primary rank. " + "-1 (default) uses max(data parallel size, 4), so frontend capacity " + "scales with the number of engines.", + ) return parser @@ -54,6 +62,15 @@ async def run_text_generation_server( hostname=hostname, ) + num_replicas = args.frontend_replicas + if num_replicas < 0: + # Each replica is a single event loop, so frontend capacity has to scale with + # the number of engines it feeds. The floor of 4 preserves the previous default + # for small deployments. + num_replicas = max(parallel_state.get_data_parallel_world_size(), 4) + if rank == 0: + logging.info("Starting %d HTTP frontend replica(s).", num_replicas) + try: if rank == 0: start_text_gen_server( @@ -63,6 +80,7 @@ async def run_text_generation_server( rank=rank, server_port=server_port, verbose=args.inference_text_gen_server_logging, + num_replicas=num_replicas, hostname=hostname, ) From 77c62ff6f23d61d1d8d3a3c9fc4b05d27dccfa5b Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Mon, 3 Aug 2026 14:24:09 -0700 Subject: [PATCH 2/4] move detokenization from the coordinator to the HTTP server --- .../coordinator.py | 3 +++ megatron/core/inference/sampling_params.py | 5 +++++ .../endpoints/chat_completions.py | 12 +++++++++++- 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py index 67eb1bb5829..ac590bdb8ca 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator/coordinator.py @@ -464,6 +464,9 @@ def detokenize(self, finished_request): finished_request (dict): The serialized merged request containing the generated tokens to be detokenized. It is modified in place. """ + if not finished_request["sampling_params"]["detokenize_generations"]: + return + detokenize_stop_sequence = (finished_request.get("sampling_params", {}) or {}).get( "detokenize_stop_sequence", False ) diff --git a/megatron/core/inference/sampling_params.py b/megatron/core/inference/sampling_params.py index 2f7d5bb4551..48a008aa694 100644 --- a/megatron/core/inference/sampling_params.py +++ b/megatron/core/inference/sampling_params.py @@ -38,6 +38,11 @@ class SamplingParams: # drops prompt_tokens before serializing the finished request, saving the ZMQ # transmission cost for long prompts. Opt in when the client needs them. return_prompt_tokens: bool = False + # When False, the DP coordinator skips detokenizing this request's output. The + # coordinator is a single process serving every DP rank, so per-request work there + # is a throughput ceiling; callers that can detokenize themselves (e.g. the HTTP + # frontend, which is replicated) should set this to False. + detokenize_generations: bool = True streaming: bool = False # Emit incremental ENGINE_REPLY_PARTIAL frames. streaming_interval: int = 1 # Minimum unsent tokens per ENGINE_REPLY_PARTIAL. diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py index a039c6cf3de..35ae357d00d 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py @@ -11,6 +11,9 @@ from megatron.core.inference.inference_request import unwrap_serialized_tensors from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) from megatron.core.tokenizers.text.parsers import PARSER_MAPPING from ..incremental_detokenizer import HuggingFaceFastIncrementalDetokenizer @@ -632,6 +635,9 @@ async def chat_completions(): termination_id=-1 if ignore_eos else None, return_prompt_tokens=return_prompt_tokens, streaming_interval=int(_get_non_none(req, "streaming_interval", 1)), + # This frontend detokenizes its own output below. Keeping it off the + # coordinator matters because that is one process shared by all DP ranks. + detokenize_generations=False, ) except ValueError as e: return Response(f"Invalid sampling parameter: {e}", status=400) @@ -730,7 +736,11 @@ async def chat_completions(): for result_item in batch_results: result = unwrap_serialized_tensors(result_item) - text_output = result["generated_text"] + text_output = TextGenerationController.detokenize( + tokenizer, + result["generated_tokens"], + remove_EOD=not sampling_params.detokenize_stop_sequence, + ) # The engine always reports prompt_length (for usage), but drops the # prompt_tokens tensor unless return_prompt_tokens was set. prompt_tokens_count = result.get("prompt_length") From 52094e5e29673deb2bbb3a3f83735f88a4e1a2d1 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Thu, 6 Aug 2026 13:27:15 -0700 Subject: [PATCH 3/4] take DP size from engine --- tools/run_dynamic_text_generation_server.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index a07ff444d5c..bd169818ef3 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -6,13 +6,12 @@ import torch -from megatron.core import parallel_state from megatron.core.inference.engines import DynamicInferenceEngine from megatron.core.inference.text_generation_server.dynamic_text_gen_server import ( start_text_gen_server, stop_text_gen_server, ) -from megatron.core.utils import configure_nvtx_profiling, trace_async_exceptions +from megatron.core.utils import configure_nvtx_profiling, get_pg_size, trace_async_exceptions from megatron.inference.utils import add_inference_args, get_dynamic_inference_engine from megatron.post_training.arguments import add_modelopt_args from megatron.training import get_args @@ -67,7 +66,7 @@ async def run_text_generation_server( # Each replica is a single event loop, so frontend capacity has to scale with # the number of engines it feeds. The floor of 4 preserves the previous default # for small deployments. - num_replicas = max(parallel_state.get_data_parallel_world_size(), 4) + num_replicas = max(get_pg_size(engine.pg_collection.dp), 4) if rank == 0: logging.info("Starting %d HTTP frontend replica(s).", num_replicas) From 54522c4cd7cd47d1bfb75bb62bf99ce024ac08fe Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Thu, 13 Aug 2026 23:23:54 -0700 Subject: [PATCH 4/4] Give each frontend replica its own SO_REUSEPORT socket The replicas shared one listening socket: the parent bound it and passed the same fd to every forked worker, so all of them accepted from a single queue. That does not balance. Whichever worker is already running tends to win the wakeup, and it keeps winning, because an event loop with work in flight polls more often than one blocked in accept. SO_REUSEPORT was set on that socket but was inert -- the kernel only load-balances when several sockets are bound to the port and it can hash a connection's 4-tuple to choose between them. Measured with 32 replicas and a fresh connection per request, ~90% of traffic landed on 5 of them, 20 replicas served exactly one request each, and throughput was 3.7x lower than the same server under a pooled client that opened its connections up front. Load made it worse rather than averaging it out: at 2048 requests the busiest replica took 604x the quietest. Each replica now binds its own socket on the shared port, so every one gets its own accept queue. Spread became max/min 1.5x with all replicas serving, and throughput 2.5x on the fresh-connection path. A pooled client is unaffected in steady state, which is the point: how well the frontend spreads no longer depends on connection behaviour the server cannot observe. start_text_gen_server now returns the base URL it is serving on. Callers that start a frontend on more than one rank need the addresses to spread requests over; previously they had to reconstruct them. The signature is otherwise unchanged, including sock, which still fixes the port -- it is closed rather than shared, since replicas bind their own. tools/run_dynamic_text_generation_server.py gains --frontend-on-all-ranks, which hosts a frontend on every rank and gathers the URLs. Frontend work is CPU-bound and otherwise confined to one rank's CPU allocation while the rest of the job's cores go unused. Signed-off-by: Siddharth Singh (cherry picked from commit 1001ca96e8e11c1faf79c53ed5a1b63aef6ecc25) --- .../text_generation_server.py | 152 ++++++++++++------ tools/run_dynamic_text_generation_server.py | 66 ++++++-- 2 files changed, 151 insertions(+), 67 deletions(-) diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py index 9e7a8da716e..2ba3cdbbb7e 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/text_generation_server.py @@ -26,7 +26,6 @@ # Global reference to manage the background server processes _SERVER_PROCESSES: List[mp.Process] = [] -_SHARED_SOCKET = None @contextmanager @@ -49,7 +48,6 @@ async def _run_text_gen_server( server_port: int, parsers: Optional[List[str]] = None, verbose: bool = False, - fd: Optional[int] = None, hostname: Optional[str] = None, ): """ @@ -65,6 +63,10 @@ async def _run_text_gen_server( logger.info(f"Rank {rank}: InferenceClient connected.") try: + # Bind what the caller asked for -- None means every interface, which is + # not the single address gethostname() resolves to. The resolved name is + # for the log line only. + bind_host = hostname if hostname is None: try: hostname = socket.gethostname() @@ -103,18 +105,20 @@ async def _run_text_gen_server( 2**14 ) # Allow many concurrent streams for HTTP/2 clients. - if fd is not None: - config.bind = [f"fd://{fd}"] - else: - config.bind = [f"{hostname}:{server_port}"] + # Held for this worker's lifetime; closing it would drop the listener. + own_socket = _bind_reuseport_socket(server_port, bind_host) + config.bind = [f"fd://{own_socket.fileno()}"] with temp_log_level(logging.INFO, logger): logger.info(f"Starting text generation server on http://{hostname}:{server_port}") logger.info(f"Using tokenizer: {type(tokenizer)}") logger.info(f"Using parsers: {parsers}") - # Quart is natively ASGI, so we can serve the app directly - await serve(app, config) + try: + # Quart is natively ASGI, so we can serve the app directly + await serve(app, config) + finally: + own_socket.close() finally: # Gracefully shut down the client when the server stops @@ -129,7 +133,6 @@ def _server_process_worker( server_port: int, parsers: Optional[List[str]] = None, verbose: bool = False, - fd: Optional[int] = None, hostname: Optional[str] = None, ): """Synchronous worker function that sets up a new event loop for the separate process.""" @@ -138,7 +141,7 @@ def _server_process_worker( try: loop.run_until_complete( _run_text_gen_server( - coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname + coordinator_addr, tokenizer, rank, server_port, parsers, verbose, hostname ) ) except KeyboardInterrupt: @@ -152,6 +155,38 @@ def _server_process_worker( loop.close() +def _bind_reuseport_socket(server_port: int, hostname: Optional[str]) -> socket.socket: + """Bind this worker's own socket on the shared port, with SO_REUSEPORT. + + Unlike inheriting one fd, this gives every worker its own accept queue and + lets the kernel hash each connection's 4-tuple across them. It is a large + improvement on sharing (measured 604x -> 2x spread at 32 replicas) but is + still hashing, not balancing: it cannot see that a replica is already busy, + so the spread is statistical rather than exact. + """ + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + # Required on every socket sharing the port; without it the second bind fails. + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) + sock.bind((hostname if hostname is not None else "0.0.0.0", server_port)) + sock.setblocking(False) + return sock + + +def _reserve_port(hostname: Optional[str]) -> int: + """Pick a free port for the replicas to bind individually. + + Replicas each bind the port themselves, so the parent cannot hold the socket + and hand out its fd; it binds only long enough to learn a free port. The gap + before the replicas bind is a small race with unrelated processes, which is + why an explicit port is preferred when one is available. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + probe.bind((hostname if hostname is not None else "0.0.0.0", 0)) + return probe.getsockname()[1] + + def start_text_gen_server( coordinator_addr: str, tokenizer, @@ -162,77 +197,88 @@ def start_text_gen_server( num_replicas: int = 4, hostname: Optional[str] = None, sock: Optional[socket.socket] = None, -): - """Start the text generation server.""" +) -> Optional[str]: + """Start the text generation server. + + Every replica binds its own socket on ``server_port`` with SO_REUSEPORT, so + each gets its own accept queue and the kernel spreads connections across + them. Sharing one inherited socket does not balance -- replicas race to + accept from a single queue and whichever is already running keeps winning, + which concentrates most traffic on a handful of them as replica count grows. + + Call this on every rank that should host a frontend. Frontend work (chat + template, detokenize, parsers, JSON) is CPU-bound, so hosting on a single + rank confines it to that rank's CPU allocation and leaves the rest of the + job's cores unused. Each caller gets its own URL back; collecting them and + spreading requests over the result is the caller's business. + + Args: + server_port: Port to listen on. Overridden by ``sock`` when given; 0 + asks the OS to choose a free one. + sock: A socket the caller already bound, used only to fix the port. + Replicas bind that port themselves, so it is closed here rather than + shared with them. + + Returns: + The base URL this rank serves on, or None if the server was already + running. + """ global _SERVER_PROCESSES - global _SHARED_SOCKET if _SERVER_PROCESSES: logger.warning("Text gen server processes are already running.") - return + return None - # The caller may pass in a socket it has already bound ahead of time. if sock is not None: - bound_port = sock.getsockname()[1] - if bound_port == 0: + # Take the port and release the socket: replicas each bind their own with + # SO_REUSEPORT, which one shared socket cannot provide. + server_port = sock.getsockname()[1] + if server_port == 0: raise ValueError( "socket must be bound to a real port before being passed to start_text_gen_server" ) - _SHARED_SOCKET = sock - server_port = bound_port - _SHARED_SOCKET.setblocking(False) - else: - _SHARED_SOCKET = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - if hasattr(socket, 'SO_REUSEPORT'): - try: - _SHARED_SOCKET.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1) - except OSError: - pass - - bind_address = hostname if hostname is not None else "0.0.0.0" - _SHARED_SOCKET.bind((bind_address, server_port)) - _SHARED_SOCKET.setblocking(False) - - _SHARED_SOCKET.set_inheritable(True) - fd = _SHARED_SOCKET.fileno() + sock.close() + elif server_port == 0: + server_port = _reserve_port(hostname) for i in range(num_replicas): p = mp.Process( target=_server_process_worker, - args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, fd, hostname), + args=(coordinator_addr, tokenizer, rank, server_port, parsers, verbose, hostname), daemon=True, ) p.start() _SERVER_PROCESSES.append(p) - logger.info(f"Started text gen frontend replica {i+1}/{num_replicas} (PID: {p.pid})") + logger.info( + f"Started text gen frontend replica {i+1}/{num_replicas} " + f"on port {server_port} (PID: {p.pid})" + ) + return f"http://{hostname or socket.gethostname()}:{server_port}" -def stop_text_gen_server(): - """Stop the text generation server.""" - global _SERVER_PROCESSES - global _SHARED_SOCKET - if not _SERVER_PROCESSES: +def _terminate(processes: List[mp.Process], what: str): + """Terminate a group of worker processes, escalating to kill if needed.""" + if not processes: return - - logger.info(f"Terminating {len(_SERVER_PROCESSES)} Text Gen frontend processes...") - - for p in _SERVER_PROCESSES: + logger.info(f"Terminating {len(processes)} {what} processes...") + for p in processes: if p.is_alive(): p.terminate() - - for p in _SERVER_PROCESSES: + for p in processes: p.join(timeout=3) if p.is_alive(): p.kill() p.join() - # Clean up the master socket - if _SHARED_SOCKET is not None: - _SHARED_SOCKET.close() - _SHARED_SOCKET = None +def stop_text_gen_server(): + """Stop this rank's frontend replica processes.""" + global _SERVER_PROCESSES + + if not _SERVER_PROCESSES: + return + + _terminate(_SERVER_PROCESSES, "Text Gen frontend") _SERVER_PROCESSES = [] logger.info("All text gen frontend processes terminated.") diff --git a/tools/run_dynamic_text_generation_server.py b/tools/run_dynamic_text_generation_server.py index bd169818ef3..8aca26ba537 100644 --- a/tools/run_dynamic_text_generation_server.py +++ b/tools/run_dynamic_text_generation_server.py @@ -33,9 +33,20 @@ def add_text_generation_server_args(parser: argparse.ArgumentParser): ) parser.add_argument( "--frontend-replicas", type=int, default=-1, - help="Number of HTTP frontend processes spawned on the primary rank. " - "-1 (default) uses max(data parallel size, 4), so frontend capacity " - "scales with the number of engines.", + help="Number of HTTP frontend processes spawned per hosting rank. " + "-1 (default) uses max(data parallel size, 4), or a flat 4 with " + "--frontend-on-all-ranks, where capacity already scales with the " + "number of ranks hosting a frontend.", + ) + parser.add_argument( + "--frontend-on-all-ranks", action="store_true", + help="Run HTTP frontends on every rank instead of only rank 0, and " + "return every rank's URL for the caller to spread requests over. " + "Frontend work (chat template, detokenize, parsers, JSON) is " + "CPU-bound and otherwise confined to the hosting rank's CPU " + "allocation, which leaves the rest of the job's cores unused. " + "Ranks still share one DP coordinator; only the HTTP tier is " + "replicated.", ) return parser @@ -63,33 +74,60 @@ async def run_text_generation_server( num_replicas = args.frontend_replicas if num_replicas < 0: - # Each replica is a single event loop, so frontend capacity has to scale with - # the number of engines it feeds. The floor of 4 preserves the previous default - # for small deployments. - num_replicas = max(get_pg_size(engine.pg_collection.dp), 4) + if args.frontend_on_all_ranks: + # Capacity now scales with the number of ranks, so the per-rank + # replica count stays flat rather than tracking DP size on top of it. + num_replicas = 4 + else: + # Each replica is a single event loop, so frontend capacity has to scale with + # the number of engines it feeds. The floor of 4 preserves the previous default + # for small deployments. + num_replicas = max(get_pg_size(engine.pg_collection.dp), 4) if rank == 0: - logging.info("Starting %d HTTP frontend replica(s).", num_replicas) + logging.info("Starting %d HTTP frontend replica(s) per hosting rank.", num_replicas) + + if args.frontend_on_all_ranks: + # Only the DP coordinator rank learns the coordinator's address: the + # engine broadcasts it over the DP group, which is a singleton when data + # parallel size is 1. Every rank needs it here, since every rank's + # frontend opens its own client. + address = [coordinator_addr] + torch.distributed.broadcast_object_list(address, src=0) + coordinator_addr = address[0] + assert coordinator_addr is not None, "no rank published a DP coordinator address" try: - if rank == 0: - start_text_gen_server( + url = None + if args.frontend_on_all_ranks or rank == 0: + url = start_text_gen_server( coordinator_addr=coordinator_addr, tokenizer=engine.controller.tokenizer, parsers=args.parsers, rank=rank, - server_port=server_port, + server_port=0 if args.frontend_on_all_ranks else server_port, verbose=args.inference_text_gen_server_logging, num_replicas=num_replicas, hostname=hostname, ) + if args.frontend_on_all_ranks: + # Unlike callers that already collect a URL per worker, this entry + # point has to gather them itself before it can report the set. + urls = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(urls, url) + if rank == 0: + for entry in [u for u in urls if u]: + logging.info("Frontend: %s", entry) + elif rank == 0: + logging.info("Frontend: %s", url) + # Await the engine loop directly since the server is running in a separate process await engine.engine_loop_task finally: - # Guarantee that the separate process is terminated when the engine loop stops or is interrupted - if rank == 0: - stop_text_gen_server() + # Guarantee that the separate processes are terminated when the engine loop + # stops or is interrupted. Every rank may now own frontend processes. + stop_text_gen_server() if __name__ == "__main__":