Files
bttoxin-pipeline/backend/app/core/docker_client.py

73 lines
2.1 KiB
Python

"""Docker 客户端管理"""
import docker
from typing import Dict, Any
from pathlib import Path
import logging
logger = logging.getLogger(__name__)
class DockerManager:
"""Docker 容器管理器"""
def __init__(self, image: str = None):
from ..config import settings
self.client = docker.from_env()
self.image = image or settings.DOCKER_IMAGE
def ensure_image(self) -> bool:
"""确保镜像存在"""
try:
self.client.images.get(self.image)
return True
except docker.errors.ImageNotFound:
logger.info(f"Pulling image {self.image}...")
self.client.images.pull(self.image)
return True
def run_bttoxin_digger(
self,
input_dir: Path,
output_dir: Path,
sequence_type: str = "nucl",
scaf_suffix: str = ".fna",
threads: int = 4
) -> Dict[str, Any]:
"""运行 BtToxin_Digger"""
self.ensure_image()
volumes = {
str(input_dir.absolute()): {'bind': '/data', 'mode': 'ro'},
str(output_dir.absolute()): {'bind': '/results', 'mode': 'rw'}
}
command = [
"/usr/local/env-execute", "BtToxin_Digger",
"--SeqPath", "/data",
"--SequenceType", sequence_type,
"--Scaf_suffix", scaf_suffix,
"--threads", str(threads)
]
try:
container = self.client.containers.run(
self.image,
command=command,
volumes=volumes,
platform="linux/amd64",
detach=True,
remove=False
)
result = container.wait()
logs = container.logs().decode('utf-8')
container.remove()
return {
'success': result['StatusCode'] == 0,
'logs': logs,
'exit_code': result['StatusCode']
}
except Exception as e:
logger.error(f"Error: {e}")
return {'success': False, 'error': str(e)}