From e84771d499dc12620b6f96e5af1e52ec4197b8af Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Mon, 22 Jun 2026 22:52:46 -0700 Subject: [PATCH] add profile endpoints Signed-off-by: Siddharth Singh --- .../data_parallel_inference_coordinator.py | 10 +++++ .../core/inference/engines/dynamic_engine.py | 6 +++ megatron/core/inference/headers.py | 2 + megatron/core/inference/inference_client.py | 13 ++++++ .../endpoints/__init__.py | 3 +- .../endpoints/profile.py | 41 +++++++++++++++++++ 6 files changed, 74 insertions(+), 1 deletion(-) create mode 100644 megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/profile.py diff --git a/megatron/core/inference/data_parallel_inference_coordinator.py b/megatron/core/inference/data_parallel_inference_coordinator.py index 50f586cc598..6c090bc62df 100644 --- a/megatron/core/inference/data_parallel_inference_coordinator.py +++ b/megatron/core/inference/data_parallel_inference_coordinator.py @@ -533,6 +533,16 @@ def start(self): if header == Headers.STOP: self.state = self.CoordinatorState.RUNNING + elif header in (Headers.START_CUDA_PROFILER, Headers.STOP_CUDA_PROFILER): + # Profiler control: broadcast to every connected DP engine. Not a + # state transition, so no CoordinatorState checks — just forward. + if sender_identity not in known_clients: + logging.warning("Coordinator: ignoring profiler signal from unknown client.") + continue + broadcast_payload = msgpack.packb(deserialized_payload, use_bin_type=True) + for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks): + self._send_to_engine(data_parallel_rank_id, broadcast_payload) + elif header == Headers.ENGINE_REPLY: # This is the output of a single engine step on some data parallel rank. assert sender_identity in self.identities_of_data_parallel_ranks diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index b8ac617bfd9..6a5a7210ab0 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -2349,6 +2349,12 @@ def schedule_requests(self) -> int: nvtx_range_pop("add_request") elif header == Headers.SET_GENERATION_EPOCH: new_generation_epoch = data[1] + elif header == Headers.START_CUDA_PROFILER: + # Side-effect, not a state transition: apply immediately on every + # rank so an outer nsys --capture-range=cudaProfilerApi starts here. + torch.cuda.cudart().cudaProfilerStart() + elif header == Headers.STOP_CUDA_PROFILER: + torch.cuda.cudart().cudaProfilerStop() else: # Control signal: queue for second pass. self._pending_signals.append(message) diff --git a/megatron/core/inference/headers.py b/megatron/core/inference/headers.py index 8ad1913e6b1..107a200818c 100644 --- a/megatron/core/inference/headers.py +++ b/megatron/core/inference/headers.py @@ -21,6 +21,8 @@ class Headers(Enum): DISCONNECT = auto() SHUTDOWN = auto() TP_BROADCAST = auto() + START_CUDA_PROFILER = auto() + STOP_CUDA_PROFILER = auto() class UnknownHeaderError(Exception): diff --git a/megatron/core/inference/inference_client.py b/megatron/core/inference/inference_client.py index f5a78d1b16c..c3017737146 100644 --- a/megatron/core/inference/inference_client.py +++ b/megatron/core/inference/inference_client.py @@ -201,6 +201,19 @@ def unpause_engines(self) -> None: """Sends UNPAUSE to all engines. No synchronization needed.""" self._send_signal_to_engines(Headers.UNPAUSE) + def start_cuda_profiler(self) -> None: + """Sends START_CUDA_PROFILER to all engines via coordinator. + + Each engine calls ``torch.cuda.profiler.start()`` (cudaProfilerStart) on + its next loop iteration, so an outer ``nsys profile --capture-range= + cudaProfilerApi`` begins recording. No synchronization needed. + """ + self._send_signal_to_engines(Headers.START_CUDA_PROFILER) + + def stop_cuda_profiler(self) -> None: + """Sends STOP_CUDA_PROFILER to all engines (cudaProfilerStop).""" + self._send_signal_to_engines(Headers.STOP_CUDA_PROFILER) + def set_generation_epoch(self, generation_epoch: int): """Sends a signal to stamp all in-flight requests with the given generation epoch. diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py index f2b0661dace..251039fa60a 100644 --- a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/__init__.py @@ -5,7 +5,8 @@ from .chat_completions import bp as ChatCompletions from .completions import bp as Completions from .health import bp as Health + from .profile import bp as Profile - __all__ = [Completions, ChatCompletions, Health] + __all__ = [Completions, ChatCompletions, Health, Profile] except ImportError: __all__ = [] diff --git a/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/profile.py b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/profile.py new file mode 100644 index 00000000000..71478aace70 --- /dev/null +++ b/megatron/core/inference/text_generation_server/dynamic_text_gen_server/endpoints/profile.py @@ -0,0 +1,41 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +"""CUDA profiler control endpoints. + +POST /start_profile and /stop_profile relay a control signal through the +InferenceClient -> data-parallel coordinator -> every connected EP/DP engine, +which calls cudaProfilerStart()/cudaProfilerStop(). Pair with an outer +`nsys profile --capture-range=cudaProfilerApi` to bracket a capture window. +""" + +import logging + +logger = logging.getLogger(__name__) + +try: + from quart import Blueprint, current_app, jsonify + + bp = Blueprint('profile_api', __name__) + + @bp.route('/start_profile', methods=['POST']) + @bp.route('/v1/start_profile', methods=['POST']) + async def start_profile(): + """Broadcast cudaProfilerStart to all engines.""" + client = current_app.config.get('client') + if client is None: + return jsonify({"status": "error", "details": "client not initialized"}), 503 + client.start_cuda_profiler() + return jsonify({"status": "ok", "action": "start_profile"}), 200 + + @bp.route('/stop_profile', methods=['POST']) + @bp.route('/v1/stop_profile', methods=['POST']) + async def stop_profile(): + """Broadcast cudaProfilerStop to all engines.""" + client = current_app.config.get('client') + if client is None: + return jsonify({"status": "error", "details": "client not initialized"}), 503 + client.stop_cuda_profiler() + return jsonify({"status": "ok", "action": "stop_profile"}), 200 + +except ImportError as e: + logger.warning(f"Could not import quart: {e}")