add some debugging to vllm docker (#454)

* add some debugging to vllm docker

* update

* check processes using GPU

* lint

Author: yunfeng-scale

model-engine/model_engine_server/inference/vllm/Dockerfile

@@ -1,5 +1,12 @@
 FROM nvcr.io/nvidia/pytorch:23.09-py3
 
+RUN apt-get update \
+    && apt-get install -y \
+        gdb \
+        psmisc \
+    && apt-get autoremove -y \
+    && rm -rf /var/lib/apt/lists/*
+
 RUN pip uninstall torch -y
 COPY requirements.txt /workspace/requirements.txt
 RUN pip install -r requirements.txt

model-engine/model_engine_server/inference/vllm/vllm_server.py

@@ -1,5 +1,9 @@
 import argparse
+import code
 import json
+import signal
+import subprocess
+import traceback
 from typing import AsyncGenerator
 
 import uvicorn
@@ -46,9 +50,9 @@ async def stream_results() -> AsyncGenerator[str, None]:
                 "text": request_output.outputs[-1].text[len(last_output_text) :],
                 "count_prompt_tokens": len(request_output.prompt_token_ids),
                 "count_output_tokens": len(request_output.outputs[0].token_ids),
-                "log_probs": request_output.outputs[0].logprobs[-1]
-                if sampling_params.logprobs
-                else None,
+                "log_probs": (
+                    request_output.outputs[0].logprobs[-1] if sampling_params.logprobs else None
+                ),
                 "finished": request_output.finished,
             }
             last_output_text = request_output.outputs[-1].text
@@ -88,7 +92,47 @@ async def abort_request() -> None:
     return Response(content=json.dumps(ret))
 
 
+def get_gpu_free_memory():
+    """Get GPU free memory using nvidia-smi."""
+    try:
+        output = subprocess.check_output(
+            ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"]
+        ).decode("utf-8")
+        gpu_memory = [int(x) for x in output.strip().split("\n")]
+        return gpu_memory
+    except subprocess.CalledProcessError:
+        return None
+
+
+def check_unknown_startup_memory_usage():
+    """Check for unknown memory usage at startup."""
+    gpu_free_memory = get_gpu_free_memory()
+    if gpu_free_memory is not None:
+        min_mem = min(gpu_free_memory)
+        max_mem = max(gpu_free_memory)
+        if max_mem - min_mem > 10:
+            print(
+                f"WARNING: Unbalanced GPU memory usage at start up. This may cause OOM. Memory usage per GPU in MB: {gpu_free_memory}."
+            )
+            output = subprocess.check_output(["fuser -v /dev/nvidia*"], shell=True).decode("utf-8")
+            print(f"Processes using GPU: {output}")
+
+
+def debug(sig, frame):
+    """Interrupt running process, and provide a python prompt for
+    interactive debugging."""
+    d = {"_frame": frame}  # Allow access to frame object.
+    d.update(frame.f_globals)  # Unless shadowed by global
+    d.update(frame.f_locals)
+
+    i = code.InteractiveConsole(d)
+    message = "Signal received : entering python shell.\nTraceback:\n"
+    message += "".join(traceback.format_stack(frame))
+    i.interact(message)
+
+
 if __name__ == "__main__":
+    check_unknown_startup_memory_usage()
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default=None)  # None == IPv4 / IPv6 dualstack
     parser.add_argument("--port", type=int, default=5005)
@@ -98,6 +142,8 @@ async def abort_request() -> None:
     engine_args = AsyncEngineArgs.from_cli_args(args)
     engine = AsyncLLMEngine.from_engine_args(engine_args)
 
+    signal.signal(signal.SIGUSR1, debug)
+
     uvicorn.run(
         app,
         host=args.host,