Pi0模型部署指南：Docker容器化最佳实践

Pi0模型部署指南：Docker容器化最佳实践

最近在机器人圈子里，Pi0这个模型挺火的。它是个视觉-语言-动作模型，简单说就是能让机器人看懂图像、理解指令，然后执行动作。听起来很酷对吧？但真要把这个模型用起来，部署是个大问题。

我最近正好在折腾Pi0的部署，发现官方文档虽然详细，但真要把它放到生产环境里稳定运行，还是有不少坑要填。特别是用Docker容器化部署这块，网上能找到的完整方案不多。

今天我就把自己折腾出来的这套Docker容器化方案分享出来，从镜像构建到服务编排，一步步带你走通整个流程。这套方案已经在我们的测试环境里跑了小半个月，稳定性还不错。

1. 为什么选择Docker容器化？

先说说为什么非要搞容器化。Pi0这个模型对运行环境要求不低，特别是GPU这块。我们团队里有几个人都在用，每个人的开发环境配置都不一样，有的用Ubuntu 22.04，有的用20.04，Python版本、CUDA版本更是五花八门。

结果就是，在我机器上跑得好好的代码，到同事那儿就各种报错。光是解决环境依赖问题，就浪费了不少时间。后来我们决定统一用Docker，问题一下子就少了很多。

用Docker有几个明显的好处：

环境一致性：不管在谁的机器上，容器里的环境都是一模一样的，彻底告别“在我这儿能跑”的问题。

资源隔离：可以精确控制每个容器能用多少GPU内存、多少CPU核心，避免一个模型把整个机器的资源都吃光。

快速部署：镜像构建好之后，在任何支持Docker的机器上都是一条命令就能跑起来，特别适合团队协作。

版本管理：不同版本的模型可以打包成不同的镜像，随时切换回滚，管理起来很方便。

2. 环境准备与基础镜像选择

开始之前，你得先确保机器上装好了Docker和NVIDIA Container Toolkit。后者是让Docker容器能用上GPU的关键。

# 安装NVIDIA Container Toolkit distribution=$(. /etc/os-release;echo $ID$VERSION_ID) curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit sudo systemctl restart docker

装好之后，跑个测试命令看看GPU能不能在容器里用：

docker run --gpus all nvidia/cuda:12.1.1-base-ubuntu22.04 nvidia-smi

如果能看到GPU信息，说明环境配置对了。

接下来要选基础镜像。Pi0官方推荐用Ubuntu 22.04，Python版本建议3.9以上。我试了几个组合，最后选了nvidia/cuda:12.1.1-runtime-ubuntu22.04这个镜像。它已经装好了CUDA 12.1，大小也比较合适，大概3GB左右。

为什么不选更小的镜像？因为Pi0依赖的库不少，特别是PyTorch、JAX这些，如果从零开始装，反而更麻烦。这个镜像已经包含了大部分基础依赖，我们只需要补充Pi0特定的部分就行。

3. 构建Pi0专用Docker镜像

镜像的Dockerfile我放在项目根目录下，结构是这样的：

# Dockerfile FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04 # 设置环境变量 ENV DEBIAN_FRONTEND=noninteractive ENV PYTHONUNBUFFERED=1 ENV PYTHONPATH=/app # 安装系统依赖 RUN apt-get update && apt-get install -y \ git \ wget \ curl \ build-essential \ python3.10 \ python3.10-dev \ python3-pip \ python3.10-venv \ && rm -rf /var/lib/apt/lists/* # 设置Python 3.10为默认 RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 # 安装uv（Python包管理工具） RUN curl -LsSf https://astral.sh/uv/install.sh | sh ENV PATH="/root/.cargo/bin:${PATH}" # 创建工作目录 WORKDIR /app # 复制项目文件 COPY . /app/ # 设置环境变量，跳过Git LFS下载（加速依赖安装） ENV GIT_LFS_SKIP_SMUDGE=1 # 安装Python依赖 RUN uv sync # 安装openpi包 RUN uv pip install -e . # 创建模型缓存目录 RUN mkdir -p /root/.cache/openpi # 设置默认命令 CMD ["bash"]

这个Dockerfile有几个关键点需要注意：

Python版本：Pi0官方说支持Python 3.9+，但我实测下来3.10最稳定。有些依赖包在3.11上会有兼容性问题。

uv工具：这是Pi0项目推荐的包管理工具，比pip快不少，而且能更好地处理依赖冲突。

Git LFS跳过：设置GIT_LFS_SKIP_SMUDGE=1能大幅加速依赖安装过程，因为Pi0依赖的LeRobot子模块里有大文件。

缓存目录：提前创建好缓存目录，避免运行时权限问题。

构建镜像的命令很简单：

docker build -t pi0-docker:latest .

第一次构建可能会花点时间，主要是下载基础镜像和安装依赖。如果网络不好，可以考虑配置Docker镜像加速。

4. 运行Pi0模型的基础容器

镜像构建好之后，先跑个简单的测试容器，看看模型能不能正常加载：

docker run --gpus all -it \ --name pi0-test \ -v $(pwd)/models:/root/.cache/openpi \ pi0-docker:latest \ python -c " from openpi.training import config as _config from openpi.policies import policy_config from openpi.shared import download print('正在加载配置...') config = _config.get_config('pi05_droid') print('配置加载成功') print('正在下载模型...') checkpoint_dir = download.maybe_download('gs://openpi-assets/checkpoints/pi05_droid') print(f'模型下载到: {checkpoint_dir}') print('正在创建策略...') policy = policy_config.create_trained_policy(config, checkpoint_dir) print('策略创建成功，模型加载完成！') "

这个命令做了几件事：

1. --gpus all：把宿主机的所有GPU都分配给容器
2. -v $(pwd)/models:/root/.cache/openpi：把本地的models目录挂载到容器的缓存目录，这样下载的模型文件会保存在本地，下次启动就不用重新下载了
3. 最后那段Python代码是测试模型能不能正常加载

如果一切顺利，你会看到类似这样的输出：

正在加载配置... 配置加载成功 正在下载模型... 模型下载到: /root/.cache/openpi/checkpoints/pi05_droid 正在创建策略... 策略创建成功，模型加载完成！

第一次运行会下载模型文件，Pi0的模型大概有几十GB，需要耐心等一会儿。下载完成后，模型文件就保存在本地的models目录里了。

5. 配置资源限制与优化

在生产环境里，不能任由容器随便用资源。特别是GPU内存，Pi0模型推理大概需要8-10GB，如果同时跑多个任务，很容易把GPU内存吃光。

下面是个更生产化的运行脚本：

#!/bin/bash # run_pi0.sh CONTAINER_NAME="pi0-service" MODEL_CACHE_DIR="./model_cache" LOG_DIR="./logs" # 创建必要的目录 mkdir -p ${MODEL_CACHE_DIR} ${LOG_DIR} # 停止并删除已有的容器 docker stop ${CONTAINER_NAME} 2>/dev/null || true docker rm ${CONTAINER_NAME} 2>/dev/null || true # 运行新容器 docker run -d \ --name ${CONTAINER_NAME} \ --gpus '"device=0"' \ # 只使用第一块GPU --cpus 4 \ # 限制使用4个CPU核心 --memory 16g \ # 限制16GB内存 --memory-swap 20g \ # 交换空间限制 --shm-size 2g \ # 共享内存大小 -p 8000:8000 \ # 暴露API端口 -v ${MODEL_CACHE_DIR}:/root/.cache/openpi \ -v ${LOG_DIR}:/app/logs \ -e CUDA_VISIBLE_DEVICES=0 \ -e XLA_PYTHON_CLIENT_MEM_FRACTION=0.8 \ # 限制JAX内存使用 --restart unless-stopped \ pi0-docker:latest \ uv run scripts/serve_policy.py policy:checkpoint \ --policy.config=pi05_droid \ --policy.dir=/root/.cache/openpi/checkpoints/pi05_droid \ --port=8000 \ --host=0.0.0.0

这个脚本里的参数都很重要：

GPU限制：--gpus '"device=0"'只让容器用第一块GPU。如果你有多块GPU，可以改成"device=0,1"。

CPU和内存：根据你的机器配置调整。Pi0推理对CPU要求不高，4核足够了。内存给16GB比较保险。

共享内存：--shm-size 2g设置共享内存大小，有些Python库会用到。

端口映射：-p 8000:8000把容器的8000端口映射到宿主机，这样就能从外面访问Pi0的服务了。

环境变量：XLA_PYTHON_CLIENT_MEM_FRACTION=0.8限制JAX最多用80%的GPU内存，留点余量给系统和其他程序。

重启策略：--restart unless-stopped让容器自动重启，除非你手动停止它。

跑起来之后，可以用这个命令检查服务状态：

curl http://localhost:8000/health

如果返回{"status":"healthy"}，说明服务启动成功了。

6. 编写Docker Compose编排文件

如果只是跑一个容器，上面的脚本就够了。但实际项目中，我们可能需要同时跑多个服务，比如Pi0推理服务、数据预处理服务、结果存储服务等等。这时候用Docker Compose会更方便。

创建一个docker-compose.yml文件：

version: '3.8' services: pi0-inference: build: . image: pi0-docker:latest container_name: pi0-inference runtime: nvidia # 使用NVIDIA容器运行时 deploy: resources: reservations: devices: - driver: nvidia count: 1 capabilities: [gpu] ports: - "8000:8000" volumes: - ./model_cache:/root/.cache/openpi - ./logs:/app/logs - ./config:/app/config environment: - CUDA_VISIBLE_DEVICES=0 - XLA_PYTHON_CLIENT_MEM_FRACTION=0.8 - LOG_LEVEL=INFO command: > uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_droid --policy.dir=/root/.cache/openpi/checkpoints/pi05_droid --port=8000 --host=0.0.0.0 healthcheck: test: ["CMD", "curl", "-f", "http://localhost:8000/health"] interval: 30s timeout: 10s retries: 3 start_period: 40s restart: unless-stopped networks: - pi0-network pi0-monitor: image: grafana/grafana:latest container_name: pi0-monitor ports: - "3000:3000" volumes: - ./monitor/grafana:/var/lib/grafana environment: - GF_SECURITY_ADMIN_PASSWORD=admin restart: unless-stopped networks: - pi0-network pi0-redis: image: redis:alpine container_name: pi0-redis ports: - "6379:6379" volumes: - ./data/redis:/data command: redis-server --appendonly yes restart: unless-stopped networks: - pi0-network networks: pi0-network: driver: bridge

这个编排文件定义了三个服务：

1. pi0-inference：Pi0模型推理服务，这是核心
2. pi0-monitor：Grafana监控界面，用来查看服务运行状态
3. pi0-redis：Redis缓存，可以用来缓存推理结果或者做任务队列

用Docker Compose启动所有服务：

# 启动所有服务 docker-compose up -d # 查看服务状态 docker-compose ps # 查看日志 docker-compose logs -f pi0-inference # 停止所有服务 docker-compose down

7. 实现高可用部署

单点部署有个问题，万一机器挂了或者容器崩溃了，服务就中断了。对于生产环境，我们需要高可用方案。

我设计了一个简单的高可用架构，用Nginx做负载均衡，多个Pi0容器同时运行：

# docker-compose-ha.yml version: '3.8' services: nginx: image: nginx:alpine container_name: pi0-nginx ports: - "8080:80" volumes: - ./nginx/nginx.conf:/etc/nginx/nginx.conf - ./nginx/conf.d:/etc/nginx/conf.d depends_on: - pi0-1 - pi0-2 networks: - pi0-ha-network pi0-1: build: . image: pi0-docker:latest container_name: pi0-1 runtime: nvidia deploy: resources: reservations: devices: - driver: nvidia count: 1 capabilities: [gpu] volumes: - ./model_cache:/root/.cache/openpi - ./logs/pi0-1:/app/logs environment: - CUDA_VISIBLE_DEVICES=0 - XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 # 两个容器共享GPU，各用40% - INSTANCE_ID=pi0-1 command: > uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_droid --policy.dir=/root/.cache/openpi/checkpoints/pi05_droid --port=8000 --host=0.0.0.0 healthcheck: test: ["CMD", "curl", "-f", "http://localhost:8000/health"] interval: 10s timeout: 5s retries: 3 networks: - pi0-ha-network pi0-2: build: . image: pi0-docker:latest container_name: pi0-2 runtime: nvidia deploy: resources: reservations: devices: - driver: nvidia count: 1 capabilities: [gpu] volumes: - ./model_cache:/root/.cache/openpi - ./logs/pi0-2:/app/logs environment: - CUDA_VISIBLE_DEVICES=0 - XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 - INSTANCE_ID=pi0-2 command: > uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_droid --policy.dir=/root/.cache/openpi/checkpoints/pi05_droid --port=8000 --host=0.0.0.0 healthcheck: test: ["CMD", "curl", "-f", "http://localhost:8000/health"] interval: 10s timeout: 5s retries: 3 networks: - pi0-ha-network haproxy-exporter: image: prom/haproxy-exporter:latest container_name: haproxy-exporter ports: - "9101:9101" command: - "--haproxy.scrape-uri=http://nginx:80/stats;csv" networks: - pi0-ha-network networks: pi0-ha-network: driver: bridge

对应的Nginx配置：

# nginx/nginx.conf events { worker_connections 1024; } http { upstream pi0_backend { least_conn; # 使用最少连接负载均衡 server pi0-1:8000 max_fails=3 fail_timeout=30s; server pi0-2:8000 max_fails=3 fail_timeout=30s; # 健康检查 check interval=5000 rise=2 fall=3 timeout=3000; } server { listen 80; location / { proxy_pass http://pi0_backend; proxy_set_header Host $host; proxy_set_header X-Real-IP $remote_addr; proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; # 超时设置 proxy_connect_timeout 60s; proxy_send_timeout 60s; proxy_read_timeout 60s; } # 状态页面 location /nginx_status { stub_status; allow 127.0.0.1; deny all; } } }

这个高可用方案的特点：

负载均衡：Nginx把请求分发给两个Pi0容器，如果一个容器挂了，流量会自动切到另一个。

共享GPU：两个容器共享同一块GPU，通过XLA_PYTHON_CLIENT_MEM_FRACTION控制各自的内存使用量。

健康检查：Nginx会定期检查容器健康状态，不健康的容器会被移出负载均衡。

监控：HAProxy Exporter把负载均衡状态暴露给Prometheus，方便监控。

启动高可用集群：

docker-compose -f docker-compose-ha.yml up -d

现在访问http://localhost:8080，请求会被自动分配到可用的Pi0容器上。

8. 监控与日志管理

服务跑起来之后，监控和日志很重要。这里我用了几个简单的工具：

日志收集：每个容器的日志都挂载到宿主机，用ELK或者Loki+Granafa可以集中查看。

性能监控：用NVTOP监控GPU使用情况：

# 安装nvtop sudo apt install nvtop # 查看GPU状态 nvtop

容器监控：用cAdvisor监控容器资源使用：

docker run \ --volume=/:/rootfs:ro \ --volume=/var/run:/var/run:ro \ --volume=/sys:/sys:ro \ --volume=/var/lib/docker/:/var/lib/docker:ro \ --volume=/dev/disk/:/dev/disk:ro \ --publish=8081:8080 \ --detach=true \ --name=cadvisor \ --privileged \ --device=/dev/kmsg \ gcr.io/cadvisor/cadvisor:latest

访问http://localhost:8081就能看到容器的CPU、内存、网络使用情况。

自定义健康检查：除了容器自带的健康检查，我们还可以写个脚本检查模型服务是否正常：

# health_check.py import requests import time import sys def check_pi0_health(url="http://localhost:8000/health", timeout=10): try: response = requests.get(url, timeout=timeout) if response.status_code == 200: data = response.json() if data.get("status") == "healthy": return True, "服务正常" else: return False, f"服务状态异常: {data}" else: return False, f"HTTP状态码异常: {response.status_code}" except requests.exceptions.RequestException as e: return False, f"请求失败: {str(e)}" if __name__ == "__main__": healthy, message = check_pi0_health() if healthy: print(f"✓ {message}") sys.exit(0) else: print(f"✗ {message}") sys.exit(1)

把这个脚本放到容器里，定时执行，就能及时发现服务异常。

9. 实际部署中的问题与解决

在实际部署中，我遇到了一些问题，这里分享下解决方案：

问题1：模型下载太慢

Pi0的模型文件存在Google Cloud Storage上，国内下载很慢。解决方案是提前下载好，或者用代理。

# 提前下载模型到本地 gsutil -m cp -r gs://openpi-assets/checkpoints/pi05_base ./model_cache/ # 或者用wget（如果有直链） wget -P ./model_cache https://storage.googleapis.com/openpi-assets/checkpoints/pi05_base/params

问题2：GPU内存不足

如果GPU内存不够，可以尝试这些方法：

1. 使用更小的模型，比如pi0-small
2. 调整XLA_PYTHON_CLIENT_MEM_FRACTION，给模型少分点内存
3. 启用梯度检查点，减少内存占用：

# 在训练配置里加这个 config.gradient_checkpointing = True

问题3：容器启动慢

Pi0容器第一次启动要加载模型，可能需要几十秒。解决方案是预热：

# 预热脚本 warmup.py import requests import time def warmup_model(): # 发送一个简单的推理请求 test_data = { "observation/exterior_image_1_left": "...", # 这里放测试图像数据 "prompt": "pick up the object" } start_time = time.time() try: response = requests.post( "http://localhost:8000/infer", json=test_data, timeout=120 ) elapsed = time.time() - start_time print(f"预热完成，耗时: {elapsed:.2f}秒") return True except Exception as e: print(f"预热失败: {e}") return False if __name__ == "__main__": warmup_model()

在容器启动后自动运行这个脚本，让模型提前加载到GPU内存里。

问题4：版本升级

Pi0还在快速迭代，经常有新版本。为了平滑升级，我用了这样的策略：

# 使用特定版本标签 FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04 # 固定Pi0版本 RUN git clone https://github.com/Physical-Intelligence/openpi.git && \ cd openpi && \ git checkout v1.0.0 # 指定版本

每次升级前，先用新版本创建测试容器，跑通后再更新生产环境。

10. 总结

折腾了这么久，总算把Pi0的Docker容器化部署跑通了。整体感觉下来，用Docker确实省心不少，特别是环境一致性和资源隔离这两点，对团队协作帮助很大。

这套方案从最简单的单容器部署，到多容器编排，再到高可用集群，基本上覆盖了从开发到生产的各种场景。你可以根据自己的需求选择合适的方案。

如果只是自己学习研究，用单容器方案就够了，简单直接。如果是团队开发，建议用Docker Compose编排，方便管理多个服务。如果是生产环境，高可用方案更稳妥，虽然配置复杂点，但能保证服务稳定性。

实际用下来，Pi0模型的效果确实不错，但部署和运维的工作量也不小。特别是GPU资源管理、模型版本控制这些，需要仔细规划。希望这篇文章能帮你少走些弯路。

最后提醒一点，这套方案是基于目前Pi0的版本设计的，如果以后官方有大的更新，可能要做些调整。建议部署前先看看官方文档有没有变化，特别是依赖库的版本要求。

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。