Torch Run 的使用方式

1. 背景

最近使用 Bernini(https://github.com/bytedance/Bernini)进行视频编辑,在其官方文档中,通过 torchrun 使用多 GPU 进行推理:

# Multi-GPU video editing
torchrun --nproc-per-node 8 infer_multi_gpu.py \
    --config pretrained_models/Bernini-R-Diffusers --ulysses 8 \
    --case assets/testcases/v2v/v2v_case1.json --guidance_mode v2v_apg

上面的使用方式的问题是:如果有许多任务,那么每次都要 Fork-Execute-Destroy,也就是开启进程 → 加载权重 → 一次推理 → 释放显存 → 销毁进程。理想的情况是:开启进程 → 加载权重 → 推理 → … → 推理 → 释放显存 → 销毁进程。下文将讨论如何实现该理想情况。


2. 第一步

2.1. 封装 torchrun 的底层 API

封装 torchrun 底层使用的 API elastic_launch,而非直接使用 torchrun 命令行。将 elastic_launch 封装进单独的进程中,该封装进程是:

如果主进程意外退出,并且 torchrun 子进程未退出,那么它们将脱离“掌控”,导致无法释放显存。这种情况通过如下方式解决:

对于 torchrun 而言,本机 Elastic Agent 盯着每个 Worker 的 PID。有一个非零退出,Agent 将给其余 Worker 发 SIGTERM,超时再 SIGKILL。所以可以通过封装进程的状态,判断 torchrun 子进程的状态。

2.2. torch_run_wrapper.py

from typing_extensions import Self
from typing import Any, Callable
import os
import signal
import multiprocessing
import logging

from torch.distributed.run import elastic_launch, LaunchConfig
import prctl
import psutil

LOGGER = logging.getLogger(__name__)


def procs_in_pgrp(pgid: int) -> list[psutil.Process]:
    """
    获取指定进程组内的所有进程。
    """
    procs = []
    for proc in psutil.process_iter(["pid", "name"]):
        try:
            if os.getpgid(proc.pid) == pgid:
                procs.append(proc)
        except (psutil.NoSuchProcess, ProcessLookupError):
            continue
    return procs


class TorchRunWrapper:
    def __init__(
        self: Self,
        launch_config: LaunchConfig,
        worker_function: Callable[..., Any],
        worker_function_args: tuple[Any, ...] = tuple(),
        worker_function_kwargs: dict[str, Any] = {},
        worker_env: dict[str, str] = os.environ.copy(),
    ) -> None:
        """
        Args:
            launch_config: 用于启动 Torch Run 的配置。
            worker_function: 在 Torch Run 进程中执行的函数。
            worker_function_args: 传递给 worker function 的位置参数。
            worker_function_kwargs: 传递给 worker function 的关键字参数。
            worker_env: 在 Torch Run 进程中设置的环境变量。
        """
        self._launch_config = launch_config
        self._worker_function = worker_function
        self._worker_function_args = worker_function_args
        self._worker_function_kwargs = worker_function_kwargs
        self._worker_env = worker_env

        # 包装进程 - 主进程的子进程;torch run 进程的父进程
        self._process: multiprocessing.Process | None = None

    def _worker_function_wrapper(
        self: Self,
    ) -> Any:
        """
        包装 worker_function 函数。
        """
        # 设置环境变量
        for k, v in self._worker_env.items():
            os.environ[k] = v
        # 设置日志
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - [%(filename)s:%(lineno)d] - %(levelname)s - %(message)s",  #noqa: E501
            datefmt="%Y-%m-%d %H:%M:%S"
        )
        # 执行 worker function
        return self._worker_function(
            *self._worker_function_args,
            **self._worker_function_kwargs,
        )

    def _elastic_launch_wrapper(
        self: Self,
    ) -> Any:
        """
        包装 elastic launch 函数,运行在包装进程中。
        """
        # 确保父进程结束时,包装进程也结束。
        # 包装进程结束时,torch run 进程将自动退出。
        # 从而实现级联退出。
        os.setpgrp()
        prctl.set_pdeathsig(signal.SIGKILL)
        LOGGER.info(
            "Torch run wrapper process running with PID %d, parent PID is %d",
            os.getpid(),
            os.getppid(),
        )
        # torch run 进程结束时,包装进程将退出。
        return elastic_launch(
            self._launch_config,
            self._worker_function_wrapper,
        )()

    def running(self: Self) -> bool:
        """
        检查 torch run 进程是否在运行。
        """
        # 检查包装进程即可确定 torch run 进程是否在运行。
        return self._process is not None and self._process.is_alive()

    def start(self, timeout: float | None = None) -> None:
        """
        启动 torch run 进程。
        """
        if self._process is not None:
            raise RuntimeError("Already started")
        # 在包装进程中启动 torch run
        self._process = multiprocessing.Process(
            target=self._elastic_launch_wrapper,
        )
        self._process.start()

    def join(self: Self, timeout: float | None = None) -> None:
        """
        等待 torch run 进程结束。
        """
        if not self.running():
            raise RuntimeError("Not running")
        # torch run 进程结束时,包装进程将退出,
        # 因此等待包装进程结束即可。
        self._process.join(timeout=timeout)  # type: ignore[union-attr]

    def kill(self: Self, timeout: float = 4.) -> None:
        """
        杀死 torch run 进程。
        """
        if not self.running():
            return
        # 杀死包装进程,将导致 torch run 进程退出,
        # 因此杀死包装进程即可。
        LOGGER.info("Killing Torch Run Wrapper %d", self._process.pid)  # type: ignore # noqa:E501
        procs = procs_in_pgrp(os.getpgid(self._process.pid))  # type: ignore
        self._process.kill()  # type: ignore[union-attr]
        _, alive = psutil.wait_procs(procs, timeout=timeout)
        for proc in alive:
            LOGGER.info(
                "Killing process %s(%d) forcely", proc.name(), proc.pid,
            )
            proc.kill()

3. 第二步

3.1. 实现轻量级广播队列

在拉任务时,所有 torchrun 进程都要拉到同一任务,因此需要支持“广播”的队列。处理完成后,所有进程都可以上报处理结果,但通常只使用 RANK 0 的结果。

为避免引入额外组件,在单独的进程中封装实现“广播”队列语义的 HTTP 服务,该封装进程:

3.2. model.py

from typing import Any, Literal

from pydantic import BaseModel


class Command(BaseModel):
    command_id: str
    opname: str
    args: dict[str, Any] = {}


class Response(BaseModel):
    command_id: str
    status: Literal["success", "error"] = "success"
    result: dict[str, Any] = {}
    # 是否为响应流中的最后一个响应
    is_last: bool = True

3.3. internal_web.py

from typing_extensions import Self
from typing import Any
import multiprocessing
import time
import asyncio

import requests
from requests import Response as RequestsResponse
from fastapi import Depends, FastAPI, Query, Body
import uvicorn

from model import Command, Response


DEFAULT_HOST: str = "0.0.0.0"
DEFAULT_PORT: int = 2322

HEALTH_ENDPOINT: str = "/health"
REGISTER_WORKER_ENDPOINT: str = "/register/worker"
WORKERS_ENDPOINT: str = "/workers"
COMMAND_ENDPOINT: str = "/command"
RESPONSE_ENDPOINT: str = "/response"
CLEAR_STATE_ENDPOINT: str = "/clear/state"


class SharedStateManager:
    def __init__(self: Self) -> None:
        self._cond = asyncio.Lock()
        # 每个 Worker 用其 Rank 唯一标识
        self._workers: dict[int, dict[str, Any]] = {}
        # rank -> commands
        self._commands: dict[int, list[Command]] = {}
        # command id -> rank -> responses
        self._responses: dict[str, dict[int, list[Response]]] = {}

    async def register_worker(self: Self, rank: int) -> None:
        async with self._cond:
            self._workers[rank] = {
                "timestamp": time.time(),
            }

    async def workers(self: Self) -> list[int]:
        async with self._cond:
            return list(self._workers.keys())

    async def send_command(self: Self, command: Command) -> None:
        async with self._cond:
            # 将命令分发到所有 Worker
            for rank in self._workers:
                queue: list[Command] = self._commands.setdefault(rank, [])
                queue.append(command)

    async def get_command(
        self: Self, rank: int,
    ) -> Command | None:
        async with self._cond:
            if not self._commands.get(rank, []):
                return None
            return self._commands[rank].pop(0)

    async def send_response(self: Self, rank: int, response: Response) -> None:
        async with self._cond:
            command_id: str = response.command_id
            if command_id not in self._responses:
                self._responses[command_id] = {}
            if rank not in self._responses[command_id]:
                self._responses[command_id][rank] = []
            self._responses[command_id][rank].append(response)

    async def get_response(
        self: Self, command_id: str,
    ) -> Response | None:
        async with self._cond:
            # 只获取 Rank 0 的响应
            responses = self._responses.get(command_id, {}).get(0, [])
            if not responses:
                return None
            response = responses.pop(0)
            if response.is_last:
                self._responses.pop(command_id)
            return response

    async def clear_state(self: Self) -> None:
        async with self._cond:
            self._workers.clear()
            self._commands.clear()
            self._responses.clear()


shared_state = SharedStateManager()


def get_shared_state() -> SharedStateManager:
    return shared_state


app = FastAPI()


@app.get(HEALTH_ENDPOINT)
async def health():
    return {"message": "OK"}


@app.put(REGISTER_WORKER_ENDPOINT)
async def register_worker(
    rank: int = Query(..., description="Rank of the worker"),
    state: SharedStateManager = Depends(get_shared_state)
):
    await state.register_worker(rank)
    return {"message": f"Worker {rank} registered"}


@app.get(WORKERS_ENDPOINT)
async def get_workers(
    state: SharedStateManager = Depends(get_shared_state)
):
    return await state.workers()


@app.put(COMMAND_ENDPOINT)
async def send_command(
    command: Command,
    state: SharedStateManager = Depends(get_shared_state)
):
    await state.send_command(command)
    return {"message": f"Command {command.command_id} received"}


@app.get(COMMAND_ENDPOINT)
async def get_command(
    rank: int = Query(..., description="Rank of the worker"),
    state: SharedStateManager = Depends(get_shared_state)
):
    return await state.get_command(rank)


@app.put(RESPONSE_ENDPOINT)
async def send_response(
    rank: int = Query(..., description="Rank of the worker"),
    response: Response = Body(..., description="Response to send"),
    state: SharedStateManager = Depends(get_shared_state)
):
    await state.send_response(rank, response)
    return {"message": f"Response {response.command_id} from rank {rank} received"}  # noqa: E501


@app.get(RESPONSE_ENDPOINT)
async def get_response(
    command_id: str = Query(..., description="Command ID"),
    state: SharedStateManager = Depends(get_shared_state)
):
    return await state.get_response(command_id)


@app.delete(CLEAR_STATE_ENDPOINT)
async def clear_state(
    state: SharedStateManager = Depends(get_shared_state)
):
    await state.clear_state()
    return {"message": "State cleared"}


class InternalWebServer:
    def __init__(
        self: Self,
        host: str = DEFAULT_HOST,
        port: int = DEFAULT_PORT,
    ) -> None:
        self._host: str = host
        self._port: int = port

        self._app: FastAPI = app
        self._process: multiprocessing.Process | None = None

    def start(self: Self) -> None:
        if self._process is not None:
            raise RuntimeError("Already started")
        self._process = multiprocessing.Process(
            target=uvicorn.run,
            kwargs={
                "app": self._app,
                "host": self._host,
                "port": self._port,
                "access_log": False,
            },
        )
        self._process.start()

    def running(self: Self) -> bool:
        """
        检查 Web 服务是否正在运行。
        """
        return self._process is not None and self._process.is_alive()

    def terminate(
        self: Self,
        force: bool = False,
        timeout: float | None = None,
    ) -> None:
        """
        终止 Web 服务。

        Args:
            force: 是否强制终止。
            timeout: 等待终止的时间,仅在非强制终止时有效。
        """
        if not self.running():
            return
        if force:
            self._process.kill()  # type: ignore[union-attr]
        else:
            self._process.terminate()  # type: ignore[union-attr]
            self._process.join(timeout=timeout)  # type: ignore[union-attr]


class InternalWebClient:
    def __init__(
        self: Self,
        host: str = DEFAULT_HOST,
        port: int = DEFAULT_PORT,
        rank: int | None = None,
        timeout: float = 10.0,
    ) -> None:
        self._host: str = host
        self._port: int = port
        self._rank: int | None = rank
        self._timeout: float = timeout

    def _build_url(self: Self, endpoint: str) -> str:
        return f"http://{self._host}:{self._port}/{endpoint.lstrip('/')}"

    def health(
        self: Self, wait_for: float,
    ) -> bool:
        start_time: float = time.time()
        while time.time() - start_time < wait_for:
            try:
                response = requests.get(
                    self._build_url(HEALTH_ENDPOINT),
                    timeout=self._timeout,
                )
                response.raise_for_status()
                return True
            except (
                requests.exceptions.ConnectionError,
                requests.exceptions.Timeout,
            ):
                continue
            except requests.exceptions.RequestException:
                return False
        return False

    def register_worker(self: Self, rank: int | None = None) -> None:
        if rank is None:
            rank = self._rank
        if rank is None:
            raise ValueError("Rank is required")
        response = requests.put(
            self._build_url(REGISTER_WORKER_ENDPOINT),
            params={"rank": rank},
            timeout=self._timeout,
        )
        response.raise_for_status()

    def get_workers(self: Self) -> list[int]:
        response = requests.get(
            self._build_url(WORKERS_ENDPOINT),
            timeout=self._timeout,
        )
        response.raise_for_status()
        return response.json()

    def send_command(self: Self, command: Command) -> None:
        response = requests.put(
            self._build_url(COMMAND_ENDPOINT),
            json=command.model_dump(),
            timeout=self._timeout,
        )
        response.raise_for_status()

    def get_command(
        self: Self, rank: int | None = None,
        wait_for: float | None = None,
        poll_interval: float = 0.005,
    ) -> Command | None:
        if rank is None:
            rank = self._rank
        if rank is None:
            raise ValueError("Rank is required")
        need_wait: bool = wait_for and wait_for > 0  # type: ignore  # noqa: E501
        start_time: float = time.time()
        remaining: float = 0
        while True:
            response = requests.get(
                self._build_url(COMMAND_ENDPOINT),
                params={"rank": rank},
                timeout=self._timeout,
            )
            response.raise_for_status()
            data = response.json()
            if need_wait:
                remaining = wait_for - (time.time() - start_time)  # type: ignore  # noqa: E501
            if data is None:
                if remaining > 0:
                    time.sleep(poll_interval)
                    continue
                return None
            return Command.model_validate(data)

    def send_response(
        self: Self, response: Response, rank: int | None = None,
    ) -> None:
        if rank is None:
            rank = self._rank
        if rank is None:
            raise ValueError("Rank is required")
        requests_response: RequestsResponse = requests.put(
            self._build_url(RESPONSE_ENDPOINT),
            params={"rank": rank},
            json=response.model_dump(),
            timeout=self._timeout,
        )
        requests_response.raise_for_status()

    def get_response(
        self: Self, command_id: str,
        wait_for: float | None = None,
        poll_interval: float = 0.005,
    ) -> Response | None:
        need_wait: bool = wait_for and wait_for > 0  # type: ignore  # noqa: E501
        start_time: float = time.time()
        remaining: float = 0
        while True:
            response = requests.get(
                self._build_url(RESPONSE_ENDPOINT),
                params={"command_id": command_id},
                timeout=self._timeout,
            )
            response.raise_for_status()
            data = response.json()
            if need_wait:
                remaining = wait_for - (time.time() - start_time)  # type: ignore  # noqa: E501
            if data is None:
                if remaining > 0:
                    time.sleep(poll_interval)
                    continue
                return None
            return Response.model_validate(data)

    def clear_state(self: Self) -> None:
        response = requests.delete(
            self._build_url(CLEAR_STATE_ENDPOINT),
            timeout=self._timeout,
        )
        response.raise_for_status()

4. 组合

4.1. 各进程之间的关系

4.2. 安装依赖

pip install torch python-prctl psutil fastapi uvicorn requests pydantic typing_extensions

4.3. main.py

"""串联 InternalWeb 和 TorchRunWrapper。
"""
import argparse
import logging
import os
import time
import uuid

from torch.distributed.run import LaunchConfig

from internal_web import DEFAULT_PORT, InternalWebClient, InternalWebServer
from model import Command, Response
from torch_run_wrapper import TorchRunWrapper

LOGGER = logging.getLogger(__name__)

OP_PING = "ping"
OP_SHUTDOWN = "shutdown"


def worker(web_port: int = DEFAULT_PORT) -> None:
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    client = InternalWebClient(port=web_port, rank=rank)
    client.register_worker()
    LOGGER.info("rank %d/%d registered, pid=%d", rank, world_size, os.getpid())

    while True:
        command = client.get_command(wait_for=1.0)
        if command is None:
            continue
        LOGGER.info("rank %d got op=%s id=%s", rank, command.opname, command.command_id)
        if command.opname == OP_SHUTDOWN:
            client.send_response(
                Response(
                    command_id=command.command_id,
                    result={"rank": rank, "pid": os.getpid()},
                )
            )
            break
        if command.opname == OP_PING:
            client.send_response(
                Response(
                    command_id=command.command_id,
                    result={
                        "rank": rank,
                        "world_size": world_size,
                        "pid": os.getpid(),
                    },
                )
            )
            continue
        client.send_response(
            Response(
                command_id=command.command_id,
                status="error",
                result={"message": f"unknown op {command.opname}"},
            )
        )


def build_launch_config(nproc_per_node: int, rdzv_port: int) -> LaunchConfig:
    return LaunchConfig(
        min_nodes=1,
        max_nodes=1,
        nproc_per_node=nproc_per_node,
        role="worker",
        rdzv_backend="c10d",
        rdzv_endpoint=f"127.0.0.1:{rdzv_port}",
        max_restarts=0,
        monitor_interval=1,
        start_method="spawn",
    )


def wait_workers(client: InternalWebClient, expected: int, timeout: float) -> list[int]:
    deadline = time.time() + timeout
    while time.time() < deadline:
        workers = client.get_workers()
        if len(workers) >= expected:
            return workers
        time.sleep(0.1)
    raise TimeoutError(f"only {client.get_workers()} registered, expected {expected}")


def run(nproc: int, port: int, rdzv_port: int, npings: int, sleep: float) -> None:
    server = InternalWebServer(port=port)
    client = InternalWebClient(port=port)
    wrapper = TorchRunWrapper(
        launch_config=build_launch_config(nproc, rdzv_port),
        worker_function=worker,
        worker_function_kwargs={"web_port": port},
    )

    server.start()
    try:
        if not client.health(wait_for=10.0):
            raise RuntimeError(f"internal web did not come up on port {port}")
        LOGGER.info("internal web ready on port %d", port)

        wrapper.start()
        workers = wait_workers(client, expected=nproc, timeout=30.0)
        LOGGER.info("workers registered: %s", workers)

        for i in range(npings):
            ping = Command(command_id=str(uuid.uuid4()), opname=OP_PING)
            client.send_command(ping)
            ping_resp = client.get_response(ping.command_id, wait_for=10.0)
            LOGGER.info("ping %d/%d response (rank 0): %s", i + 1, npings, ping_resp)
            time.sleep(sleep)

        shutdown = Command(command_id=str(uuid.uuid4()), opname=OP_SHUTDOWN)
        client.send_command(shutdown)
        shutdown_resp = client.get_response(shutdown.command_id, wait_for=10.0)
        LOGGER.info("shutdown response (rank 0): %s", shutdown_resp)

        wrapper.join(timeout=30.0)
        LOGGER.info("torchrun wrapper joined")
    finally:
        if wrapper.running():
            wrapper.kill()
        server.terminate(timeout=5.0)


def main() -> None:
    parser = argparse.ArgumentParser(description="Integrate InternalWeb + TorchRunWrapper")
    parser.add_argument("--nproc", type=int, default=2, help="workers on this node")
    parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="internal web port")
    parser.add_argument("--npings", type=int, default=3, help="how many ping commands to send")
    parser.add_argument("--sleep", type=float, default=1.0, help="seconds to sleep after each ping")
    parser.add_argument("--rdzv-port", type=int, default=29400, help="c10d rendezvous port")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - [%(filename)s:%(lineno)d] - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    run(
        nproc=args.nproc,
        port=args.port,
        rdzv_port=args.rdzv_port,
        npings=args.npings,
        sleep=args.sleep,
    )


if __name__ == "__main__":
    main()

4.4. 运行示例

$ python main.py --npings 5 --sleep 1
INFO:     Started server process [3542450]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:2322 (Press CTRL+C to quit)
2026-08-31 23:58:05 - [main.py:101] - INFO - internal web ready on port 2322
2026-08-31 23:58:05 - [torch_run_wrapper.py:87] - INFO - Torch run wrapper process running with PID 3542461, parent PID is 3542350
W0831 23:58:05.294985 3542461 /mnt/nvme0n1/home/zhoujingjiang/miniconda3/envs/vidu-infer/lib/python3.10/site-packages/torch/distributed/launcher/api.py:190] config has no run_id, generated a random run_id: 14939821356600118451770181188213592777
2026-08-31 23:58:06 - [main.py:28] - INFO - rank 1/2 registered, pid=3542468
2026-08-31 23:58:06 - [main.py:28] - INFO - rank 0/2 registered, pid=3542467
2026-08-31 23:58:06 - [main.py:105] - INFO - workers registered: [1, 0]
2026-08-31 23:58:06 - [main.py:34] - INFO - rank 0 got op=ping id=de784512-2689-40ab-bac7-44c4ba45b8de
2026-08-31 23:58:06 - [main.py:34] - INFO - rank 1 got op=ping id=de784512-2689-40ab-bac7-44c4ba45b8de
2026-08-31 23:58:06 - [main.py:111] - INFO - ping 1/5 response (rank 0): command_id='de784512-2689-40ab-bac7-44c4ba45b8de' status='success' result={'rank': 0, 'world_size': 2, 'pid': 3542467} is_last=True
2026-08-31 23:58:07 - [main.py:34] - INFO - rank 1 got op=ping id=341d1355-6ae3-49fc-98b7-1edd13382994
2026-08-31 23:58:07 - [main.py:34] - INFO - rank 0 got op=ping id=341d1355-6ae3-49fc-98b7-1edd13382994
2026-08-31 23:58:07 - [main.py:111] - INFO - ping 2/5 response (rank 0): command_id='341d1355-6ae3-49fc-98b7-1edd13382994' status='success' result={'rank': 0, 'world_size': 2, 'pid': 3542467} is_last=True
2026-08-31 23:58:08 - [main.py:34] - INFO - rank 1 got op=ping id=73d0b3c2-484a-456c-a172-52454ce55c1b
2026-08-31 23:58:08 - [main.py:34] - INFO - rank 0 got op=ping id=73d0b3c2-484a-456c-a172-52454ce55c1b
2026-08-31 23:58:08 - [main.py:111] - INFO - ping 3/5 response (rank 0): command_id='73d0b3c2-484a-456c-a172-52454ce55c1b' status='success' result={'rank': 0, 'world_size': 2, 'pid': 3542467} is_last=True
2026-08-31 23:58:09 - [main.py:34] - INFO - rank 1 got op=ping id=306ce671-adaf-40c3-8394-d867332e3bf9
2026-08-31 23:58:09 - [main.py:34] - INFO - rank 0 got op=ping id=306ce671-adaf-40c3-8394-d867332e3bf9
2026-08-31 23:58:09 - [main.py:111] - INFO - ping 4/5 response (rank 0): command_id='306ce671-adaf-40c3-8394-d867332e3bf9' status='success' result={'rank': 0, 'world_size': 2, 'pid': 3542467} is_last=True
2026-08-31 23:58:10 - [main.py:34] - INFO - rank 0 got op=ping id=d0b6efac-e1ef-4148-a6e5-3d6e1c7e1b0a
2026-08-31 23:58:10 - [main.py:34] - INFO - rank 1 got op=ping id=d0b6efac-e1ef-4148-a6e5-3d6e1c7e1b0a
2026-08-31 23:58:10 - [main.py:111] - INFO - ping 5/5 response (rank 0): command_id='d0b6efac-e1ef-4148-a6e5-3d6e1c7e1b0a' status='success' result={'rank': 0, 'world_size': 2, 'pid': 3542467} is_last=True
2026-08-31 23:58:11 - [main.py:34] - INFO - rank 1 got op=shutdown id=7e13361e-c058-495c-ab12-8b84195ea4ad
2026-08-31 23:58:11 - [main.py:34] - INFO - rank 0 got op=shutdown id=7e13361e-c058-495c-ab12-8b84195ea4ad
2026-08-31 23:58:11 - [main.py:117] - INFO - shutdown response (rank 0): command_id='7e13361e-c058-495c-ab12-8b84195ea4ad' status='success' result={'rank': 0, 'pid': 3542467} is_last=True
2026-08-31 23:58:12 - [main.py:120] - INFO - torchrun wrapper joined
INFO:     Shutting down
INFO:     Waiting for application shutdown.
INFO:     Application shutdown complete.
INFO:     Finished server process [3542450]