Merge 4e75db9c23bc87131c6f9b7a931f27df7b50dd7e into 553741557a57df44a46fb3ab72b6aafeaf945337

deepinfer/engine/profile_server.py

@@ -4,6 +4,11 @@ POST /profile 采 num_steps 个 step 并落盘，阻塞到文件写完才返回
 接口。监听 DEEPINFER_PROFILE_HTTP_PORT_BASE + driver_gpu_id，默认关，设
 DEEPINFER_PROFILE_HTTP_ENABLE=1 开启。
 
+接口只接受 POST + JSON body（不接受 query 参数，GET 一律 405）：参数只能通过 body
+传入，浏览器里的 <img>/<form> 就无法伪造请求触发采集。默认只绑回环；需要从别的节点
+采集时把 DEEPINFER_PROFILE_HTTP_BIND 设为 0.0.0.0（或具体网卡 ip）。trace 只允许
+写到 DEEPINFER_TORCH_PROFILER_DIR（或其子目录），客户端传的 output_dir 越界即 400。
+
 只采 driver worker 这一张卡：torch profiler 是进程内的，CP > 1 时 slave worker 在别的
 进程里，采不到。
 
@@ -146,7 +151,7 @@ class ProfileManager:
             return
 
         try:
-            sock = socket.create_server(("0.0.0.0", self._port))
+            sock = socket.create_server((envs.DEEPINFER_PROFILE_HTTP_BIND, self._port))
         except OSError as e:
             logger.error(f"profile http server failed to bind port {self._port}: {e}. Profiling API disabled.")
             return
@@ -155,8 +160,10 @@ class ProfileManager:
         config = uvicorn.Config(app, log_config=None, access_log=False, log_level="warning")
         server = uvicorn.Server(config)
         threading.Thread(target=server.run, kwargs={"sockets": [sock]}, name="profile-http-server", daemon=True).start()
-        # 打出本机 ip 而不只是 0.0.0.0，否则不知道该从别的节点往哪儿发请求
-        logger.info(f"profile http server up, try: curl -XPOST http://{self._ip}:{self._port}{_ENDPOINT}")
+        # 绑在回环上时机器 ip 连不上，提示里如实打监听地址；绑 0.0.0.0 时打本机 ip，
+        # 否则不知道该从别的节点往哪儿发请求
+        display_host = self._ip if envs.DEEPINFER_PROFILE_HTTP_BIND == "0.0.0.0" else envs.DEEPINFER_PROFILE_HTTP_BIND
+        logger.info(f"profile http server up, try: curl -XPOST http://{display_host}:{self._port}{_ENDPOINT}")
 
     def _build_app(self) -> Any:
         from fastapi import FastAPI, Request
@@ -190,22 +197,33 @@ class ProfileManager:
             detail = f"{exc.detail} ({request.method} {request.url.path})"
             return as_response(_Reply.error(HTTPStatus(exc.status_code), detail))
 
-        @app.api_route(_ENDPOINT, methods=["GET", "POST"])
+        @app.api_route(_ENDPOINT, methods=["POST"])
         async def profile(request: Request) -> JSONResponse:
-            # query string 和 JSON body 合并，body 优先
-            params: dict[str, Any] = dict(request.query_params)
+            # 只接受 JSON body，不读 query 参数：否则 <form method="post" action="...?output_dir=...">
+            # 仍能把参数从 URL 塞进来，POST-only 就失去意义了
             try:
-                if raw := (await request.body()).strip():
-                    body = json.loads(raw)
-                    if not isinstance(body, dict):
-                        raise ValueError("request body must be a JSON object")
-                    params |= body
-                p = Params.model_validate(params)
+                raw = (await request.body()).strip()
+                body = json.loads(raw) if raw else {}
+                if not isinstance(body, dict):
+                    raise ValueError("request body must be a JSON object")
+                p = Params.model_validate(body)
             except (ValueError, ValidationError) as e:  # json.JSONDecodeError 也是 ValueError
                 return as_response(_Reply.error(HTTPStatus.BAD_REQUEST, _fmt_error(e)))
 
             loop = asyncio.get_running_loop()
-            output_dir = p.output_dir or envs.DEEPINFER_TORCH_PROFILER_DIR
+            base_dir = os.path.realpath(envs.DEEPINFER_TORCH_PROFILER_DIR)
+            if p.output_dir is None:
+                output_dir = base_dir
+            else:
+                # trace 只允许落到 DEEPINFER_TORCH_PROFILER_DIR 子树内；realpath 先解析掉
+                # ../ 和符号链接，防止 /tmp/link -> /etc 这类重定向
+                output_dir = os.path.realpath(p.output_dir)
+                if output_dir != base_dir and not output_dir.startswith(base_dir + os.sep):
+                    msg = (
+                        f"output_dir {p.output_dir!r} is outside DEEPINFER_TORCH_PROFILER_DIR "
+                        f"({base_dir!r}); traces may only be written inside it"
+                    )
+                    return as_response(_Reply.error(HTTPStatus.BAD_REQUEST, msg))
             try:
                 await loop.run_in_executor(self._io_pool, lambda: os.makedirs(output_dir, exist_ok=True))
             except OSError as e:

deepinfer/environ.py

@@ -208,6 +208,8 @@ class envs:
     DEEPINFER_ENABLE_MEMORY_TRACING = EnvVar.bool()
     DEEPINFER_PROFILE_HTTP_ENABLE = EnvVar.bool()
     DEEPINFER_PROFILE_HTTP_PORT_BASE = EnvVar.int(default=6070)
+    "profiling HTTP 接口监听的地址。默认只绑回环；多节点远程采集需显式设为 0.0.0.0（或具体网卡 ip）"
+    DEEPINFER_PROFILE_HTTP_BIND = EnvVar.str(default="127.0.0.1")
     DEEPINFER_TORCH_PROFILER_DIR = EnvVar.str(default="/tmp")
 
     # -- Multi-Modal --

requirements-dev.txt

@@ -17,3 +17,5 @@ matplotlib==3.11.2
 # Required by the engine's runtime profiling HTTP server (deepinfer/engine/profile_server.py)
 fastapi==0.128.0
 uvicorn==0.40.0
+# Required by fastapi.testclient (tests/engine/test_profile_server.py)
+httpx==0.28.1

tests/engine/test_profile_server.py

@@ -0,0 +1,98 @@
+import json
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+import pytest
+
+from deepinfer.engine.profile_server import ProfileManager
+
+pytest.importorskip("fastapi", reason="profile http server tests need the fastapi/uvicorn dev dependencies")
+
+
+@pytest.fixture
+def manager(tmp_path: Any, monkeypatch: Any) -> ProfileManager:
+    """Bare ProfileManager whose request slot is already occupied.
+
+    Occupying the slot (via _req_ready) makes every valid request return 409
+    immediately, so tests never wait on the engine-side pickup timeout.
+    """
+    monkeypatch.setenv("DEEPINFER_TORCH_PROFILER_DIR", str(tmp_path / "traces"))
+    from deepinfer import environ
+
+    environ.rearm_all()
+
+    manager = object.__new__(ProfileManager)
+    manager._pending_req = None
+    manager._req_ready = threading.Event()
+    manager._req_ready.set()
+    manager._io_pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="profile-io-test")
+    return manager
+
+
+@pytest.fixture
+def client(manager: ProfileManager) -> Any:
+    from fastapi.testclient import TestClient
+
+    with TestClient(manager._build_app()) as c:
+        yield c
+
+
+def _post(client: Any, body: dict[str, Any] | None = None, **kwargs: Any) -> Any:
+    content = json.dumps(body) if body is not None else None
+    return client.post("/profile", content=content, **kwargs)
+
+
+@pytest.mark.ci
+def test_get_is_rejected(client: Any) -> None:
+    # <img src="http://engine:port/profile?..."> 是 GET：必须 405，浏览器无法借此触发采集
+    response = client.get("/profile", params={"num_steps": 5})
+    assert response.status_code == 405
+
+
+@pytest.mark.ci
+def test_query_params_are_not_honored(client: Any, tmp_path: Any) -> None:
+    # form POST 可以把参数塞进 URL，但 query 已不被读取：output_dir 越界参数不生效，
+    # 请求按默认值（合法）处理并落入 409
+    response = client.post("/profile", params={"output_dir": str(tmp_path / "outside")})
+    assert response.status_code == 409
+
+
+@pytest.mark.ci
+def test_output_dir_outside_profiler_dir_is_rejected(client: Any, tmp_path: Any) -> None:
+    response = _post(client, {"output_dir": str(tmp_path / "outside")})
+    assert response.status_code == 400
+    assert "DEEPINFER_TORCH_PROFILER_DIR" in response.json()["message"]
+
+
+@pytest.mark.ci
+def test_output_dir_traversal_is_rejected(client: Any, tmp_path: Any) -> None:
+    # ../ 越界在 realpath 之后同样落到子树外
+    base = tmp_path / "traces"
+    response = _post(client, {"output_dir": f"{base}/../evil"})
+    assert response.status_code == 400
+
+
+@pytest.mark.ci
+def test_output_dir_inside_profiler_dir_is_accepted(client: Any, tmp_path: Any) -> None:
+    response = _post(client, {"output_dir": str(tmp_path / "traces" / "run1")})
+    assert response.status_code == 409  # slot occupied; request itself parsed fine
+
+
+@pytest.mark.ci
+def test_default_output_dir_is_profiler_dir(client: Any, tmp_path: Any) -> None:
+    response = _post(client, {"num_steps": 2})
+    assert response.status_code == 409
+    assert (tmp_path / "traces").is_dir()  # makedirs ran on the configured base dir
+
+
+@pytest.mark.ci
+def test_non_object_body_is_rejected(client: Any) -> None:
+    response = client.post("/profile", content=json.dumps([1, 2]))
+    assert response.status_code == 400
+
+
+@pytest.mark.ci
+def test_unknown_fields_are_rejected(client: Any) -> None:
+    response = _post(client, {"num_steps": 2, "evil": True})
+    assert response.status_code == 400