优化docker sdk

utils/odm_monitor.py

@@ -1,6 +1,6 @@
 import os
 import logging
-import subprocess
+import docker
 from typing import Tuple
 
 
@@ -11,6 +11,7 @@ class ODMProcessMonitor:
         self.max_retries = max_retries
         self.logger = logging.getLogger('UAV_Preprocess.ODMMonitor')
         self.mode = mode
+        self.client = docker.from_env()
 
     def _check_success(self, grid_dir: str) -> bool:
         """检查ODM是否执行成功"""
@@ -20,55 +21,55 @@ class ODMProcessMonitor:
         return all(os.path.exists(os.path.join(grid_dir, 'project', marker)) for marker in success_markers)
 
     def run_odm_with_monitor(self, grid_dir: str, grid_idx: int, fast_mode: bool = True) -> Tuple[bool, str]:
-        """运行ODM命令"""
+        """运行ODM容器"""
         try:
             self.logger.info(f"开始处理网格 {grid_idx + 1}")
-                
+
-            # 构建基础命令
+            # 准备容器配置
-            cmd = [
+            volumes = {
-                "docker", "run", "--rm",
+                grid_dir: {'bind': '/datasets', 'mode': 'rw'}
-                "-v", f"{grid_dir}:/datasets",
+            }
-            "opendronemap/odm",
+
+            # 准备命令参数
+            command = [
                 "--project-path", "/datasets", "project",
                 "--max-concurrency", "10",
                 "--force-gps",
                 "--rerun-all"
             ]
 
-                if fast_mode:
+            if fast_mode:
-                    command.extend([
+                command.extend([
-                        "--feature-quality", "lowest",
+                    "--feature-quality", "lowest",
-                        "--orthophoto-resolution", "8",
+                    "--orthophoto-resolution", "8",
-                        "--fast-orthophoto",
+                    "--fast-orthophoto",
-                        "--skip-3dmodel"
+                    "--skip-3dmodel"
-                    ])
+                ])
 
-                # 运行容器
+            # 运行容器并等待完成
-                container = self.client.containers.run(
+            container = self.client.containers.run(
-                    "opendronemap/odm",
+                "opendronemap/odm",
-                    command=command,
+                command=command,
-                    volumes=volumes,
+                volumes=volumes,
-                    detach=True,
+                detach=True,
-                    remove=True
+                remove=True,
-                )
+                environment={"PYTHONUNBUFFERED": "1"},
+                mem_limit="0",  # 不限制内存
+                cpu_count=0,    # 使用所有CPU
+                network_mode="host"  # 使用主机网络模式
+            )
 
-                # 等待容器完成
+            # 等待容器完成并获取状态码
-                result = container.wait()
+            result = container.wait()
 
-                # 只在失败时获取日志
+            # 检查是否成功完成
-                if result['StatusCode'] != 0:
+            if result['StatusCode'] == 0 and self._check_success(grid_dir):
-                    logs = container.logs().decode('utf-8')
+                self.logger.info(f"网格 {grid_idx + 1} 处理成功")
-                    self.logger.error("容器执行失败，最后10行日志：")
+                return True, ""
-                    self.logger.error(''.join(logs.split('\n')[-10:]))
 
-                # 检查是否成功完成
+            return False, f"网格 {grid_idx + 1} 处理失败"
-                if result['StatusCode'] == 0 and self._check_success(grid_dir):
-                    self.logger.info(f"网格 {grid_idx + 1} ODM处理成功")
-                    return True, ""
-
-                self.logger.warning(f"网格 {grid_idx + 1} 第 {attempt + 1} 次尝试失败")
 
         except Exception as e:
             error_msg = f"执行异常: {str(e)}"
             self.logger.error(error_msg)
-            return False, error_msg
+            return False, error_msg