diff --git a/deepinfer/engine/profile_server.py b/deepinfer/engine/profile_server.py index 3fd5d94..79f6416 100644 --- a/deepinfer/engine/profile_server.py +++ b/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 +传入,浏览器里的 /
就无法伪造请求触发采集。默认只绑回环;需要从别的节点 +采集时把 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 参数:否则 + # 仍能把参数从 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: diff --git a/deepinfer/environ.py b/deepinfer/environ.py index e3a55bb..a5af40e 100644 --- a/deepinfer/environ.py +++ b/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 -- diff --git a/requirements-dev.txt b/requirements-dev.txt index 4941093..a762124 100644 --- a/requirements-dev.txt +++ b/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 diff --git a/tests/engine/test_profile_server.py b/tests/engine/test_profile_server.py new file mode 100644 index 0000000..deca4be --- /dev/null +++ b/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: + # 是 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