如何在同一Python脚本中调用环境冲突的Docker容器内的DL模型？

针对你的需求，这里有几个适合实验场景的简单方案，无需修改现有模型代码或环境配置：

方案1：将模型封装为API服务（最接近你的示例脚本体验） #

把两个模型分别做成轻量的HTTP API服务，在各自容器中运行，主脚本通过网络请求调用模型获取结果。这种方式能模拟你想要的“直接调用”体验，且无需处理文件读写的同步问题。

步骤： #

1. 给每个容器添加API服务
   以FastAPI为例，在Container1中创建api_a.py：

   from fastapi import FastAPI
   from envA import modelA
   import torch
   
   app = FastAPI()
   
   @app.post("/predict")
   def predict(data: list):
       # 将传入的列表转为模型需要的张量
       data_tensor = torch.tensor(data)
       output = modelA(data_tensor)
       # 将张量转为可序列化的列表返回
       return {"output": output.detach().cpu().numpy().tolist()}
   

   同理在Container2中创建api_b.py，替换为modelB的逻辑。

2. 启动容器并映射端口
   给两个容器创建自定义网络（避免端口冲突，且容器间可直接访问）：

   docker network create model-network
   

   启动Container1的API服务：

   docker run -p 8000:8000 --network model-network container1-image uvicorn api_a:app --host 0.0.0.0
   

   启动Container2的API服务：

   docker run -p 8001:8001 --network model-network container2-image uvicorn api_b:app --host 0.0.0.0
   

3. 主脚本调用API
   主脚本中用requests库发起请求，获取结果后转为张量：

   import requests
   import torch
   
   def call_model_a(data):
       response = requests.post("http://localhost:8000/predict", json={"data": data.tolist()})
       return torch.tensor(response.json()["output"])
   
   def call_model_b(data):
       response = requests.post("http://localhost:8001/predict", json={"data": data.tolist()})
       return torch.tensor(response.json()["output"])
   
   # 通用代码
   dataA = torch.randn(1, 3, 224, 224)
   dataB = torch.randn(1, 10)
   
   outputA = call_model_a(dataA)
   outputB = call_model_b(dataB)
   
   # 后续处理逻辑
   print("ModelA输出形状:", outputA.shape)
   print("ModelB输出形状:", outputB.shape)
   

方案2：用Docker Compose+共享卷自动化文件传递 #

基于你提到的临时pkl方案，用Docker Compose实现容器间的输入输出自动化，无需手动运行容器和处理文件。

步骤： #

1. 创建docker-compose.yml

   version: '3.8'
   services:
     model_a:
       image: container1-image
       volumes:
         - ./data:/data
         - ./outputs:/outputs
       command: python run_model_a.py --input /data/dataA.pt --output /outputs/outputA.pt
   
     model_b:
       image: container2-image
       volumes:
         - ./data:/data
         - ./outputs:/outputs
       command: python run_model_b.py --input /data/dataB.pt --output /outputs/outputB.pt
   

   其中run_model_a.py是你预先写好的脚本，负责读取输入文件、运行ModelA、保存输出文件。

2. 主脚本控制流程

   import subprocess
   import torch
   import os
   
   # 准备输入数据目录
   os.makedirs("./data", exist_ok=True)
   os.makedirs("./outputs", exist_ok=True)
   
   # 生成并保存输入数据
   dataA = torch.randn(1, 3, 224, 224)
   dataB = torch.randn(1, 10)
   torch.save(dataA, "./data/dataA.pt")
   torch.save(dataB, "./data/dataB.pt")
   
   # 启动两个模型容器，等待运行完成
   subprocess.run(["docker-compose", "up", "--abort-on-container-exit"], check=True)
   
   # 读取输出结果
   outputA = torch.load("./outputs/outputA.pt")
   outputB = torch.load("./outputs/outputB.pt")
   
   # 清理容器（可选）
   subprocess.run(["docker-compose", "down"])
   
   # 后续处理逻辑
   print("ModelA输出形状:", outputA.shape)
   print("ModelB输出形状:", outputB.shape)
   

方案3：用Docker Python SDK直接控制容器 #

通过Docker官方的Python SDK，在主脚本中直接启动、管理两个容器，灵活控制输入输出的传递。

步骤： #

1. 安装Docker SDK

   pip install docker
   

2. 主脚本代码

   import docker
   import torch
   import os
   
   client = docker.from_env()
   
   # 准备临时目录存储输入输出
   temp_dir = "./temp_workspace"
   input_dir = f"{temp_dir}/input"
   output_dir = f"{temp_dir}/output"
   os.makedirs(input_dir, exist_ok=True)
   os.makedirs(output_dir, exist_ok=True)
   
   # 生成输入数据
   dataA = torch.randn(1, 3, 224, 224)
   dataB = torch.randn(1, 10)
   torch.save(dataA, f"{input_dir}/dataA.pt")
   torch.save(dataB, f"{input_dir}/dataB.pt")
   
   # 运行ModelA容器
   container_a = client.containers.run(
       "container1-image",
       command=f"python run_model_a.py --input /input/dataA.pt --output /output/outputA.pt",
       volumes={
           os.path.abspath(input_dir): {"bind": "/input", "mode": "ro"},
           os.path.abspath(output_dir): {"bind": "/output", "mode": "rw"}
       },
       detach=True
   )
   # 等待容器运行完成
   container_a.wait()
   
   # 运行ModelB容器
   container_b = client.containers.run(
       "container2-image",
       command=f"python run_model_b.py --input /input/dataB.pt --output /output/outputB.pt",
       volumes={
           os.path.abspath(input_dir): {"bind": "/input", "mode": "ro"},
           os.path.abspath(output_dir): {"bind": "/output", "mode": "rw"}
       },
       detach=True
   )
   container_b.wait()
   
   # 读取输出
   outputA = torch.load(f"{output_dir}/outputA.pt")
   outputB = torch.load(f"{output_dir}/outputB.pt")
   
   # 清理容器
   container_a.remove()
   container_b.remove()
   
   # 后续处理逻辑
   print("ModelA输出形状:", outputA.shape)
   print("ModelB输出形状:", outputB.shape)
   

方案选择建议 #

• 优先选方案1：最贴合你想要的“在同一个脚本中调用模型”的需求，支持异步调用加速实验，且逻辑清晰。
• 若你更习惯文件传递的方式，选方案2：通过Docker Compose简化容器管理，适合批量处理数据的实验。

内容的提问来源于stack exchange，提问作者MufasaChan