import json import shutil import subprocess import sys import time from pathlib import Path class BatchPreflight: def __init__(self, base_dir, output_dir, temp_dir, ffmpeg_path): self.base_dir = Path(base_dir) self.output_dir = Path(output_dir) self.temp_dir = Path(temp_dir) self.ffmpeg_path = Path(ffmpeg_path) self.python_exe = self._resolve_python_executable() def _resolve_python_executable(self): env_python = self.base_dir / "env" / "Scripts" / "python.exe" if env_python.exists(): return env_python return Path(sys.executable) def _gpu_env(self): import os env = os.environ.copy() env["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" env["PYTHONUTF8"] = "1" nvidia_root = self.base_dir / "env" / "Lib" / "site-packages" / "nvidia" extra_paths = [] if nvidia_root.exists(): for child in nvidia_root.iterdir(): for sub in ("bin", "lib"): p = child / sub if p.exists(): extra_paths.append(str(p)) if extra_paths: # FIX: dùng os.pathsep (Linux/Docker là ":", Windows là ";") env["PATH"] = os.pathsep.join(extra_paths + [env.get("PATH", "")]) prev_py = env.get("PYTHONPATH", "") env["PYTHONPATH"] = str(self.base_dir) + (os.pathsep + prev_py if prev_py else "") return env def run(self, asr_engine, trans_engine, tts_engine, strict_gpu, log_fn=None, required_checks=None): required_checks = set(required_checks or []) checks = [] checks.append(self._check_writable("output_writable", self.output_dir, True)) checks.append(self._check_writable("temp_writable", self.temp_dir, True)) checks.append(self._check_disk_space("disk_space", self.output_dir, min_free_gb=5, required=True)) checks.append(self._check_nvenc(required="ffmpeg_h264_nvenc" in required_checks)) checks.append(self._check_python_cuda("torch_cuda", required="torch_cuda" in required_checks)) checks.append(self._check_asr_cuda(required="faster_whisper_ctranslate2_cuda" in required_checks)) if "PaddleOCR" in asr_engine or "Quét chữ" in asr_engine: checks.append(self._check_paddle_cuda(required="paddleocr_cuda" in required_checks)) checks.append(self._check_onnx_cuda(required="onnxruntime_cuda_provider" in required_checks)) checks.append(self._check_provider_config(trans_engine, required=True)) report = { "created_at": time.strftime("%Y-%m-%d %H:%M:%S"), "strict_gpu": strict_gpu, "asr_engine": asr_engine, "trans_engine": trans_engine, "tts_engine": tts_engine, "checks": checks, "status": "PASS" if all(c["ok"] or not c["required"] for c in checks) else "FAILED", } self._write_reports(report) if log_fn: for c in checks: state = "PASS" if c["ok"] else ("FAIL" if c["required"] else "WARN") log_fn(f"[PREFLIGHT] {state}: {c['name']} - {c['message']}") log_fn(f"[PREFLIGHT] report: {self.output_dir / 'preflight_report.json'}") return report def _run_cmd(self, cmd, timeout=20): startupinfo = None if sys.platform == "win32": startupinfo = subprocess.STARTUPINFO() startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW try: res = subprocess.run( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="utf-8", errors="ignore", timeout=timeout, startupinfo=startupinfo, env=self._gpu_env(), ) return res.returncode, (res.stdout or ""), (res.stderr or "") except Exception as e: return 99, "", str(e) def _check_writable(self, name, path, required): try: path.mkdir(parents=True, exist_ok=True) probe = path / ".preflight_write_test" probe.write_text("ok", encoding="utf-8") probe.unlink(missing_ok=True) return self._result(name, True, required, f"{path} writable") except Exception as e: return self._result(name, False, required, str(e)) def _check_disk_space(self, name, path, min_free_gb, required): try: usage = shutil.disk_usage(path) free_gb = usage.free / (1024 ** 3) return self._result(name, free_gb >= min_free_gb, required, f"{free_gb:.1f} GB free") except Exception as e: return self._result(name, False, required, str(e)) def _check_nvenc(self, required): cmd = [ str(self.ffmpeg_path), "-y", "-f", "lavfi", "-i", "color=c=black:s=256x256", "-t", "1", "-c:v", "h264_nvenc", "-f", "null", "-" ] code, out, err = self._run_cmd(cmd, timeout=15) msg = "h264_nvenc OK" if code == 0 else (err or out)[-300:] return self._result("ffmpeg_h264_nvenc", code == 0, required, msg) def _check_python_cuda(self, name, required): code = ( "import torch; " "assert torch.cuda.is_available(), 'torch cuda unavailable'; " "x=torch.randn(1, device='cuda'); " "print(torch.cuda.get_device_name(0))" ) rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=30) msg = out.strip() if rc == 0 else (err or out)[-300:] return self._result(name, rc == 0, required, msg) def _check_asr_cuda(self, required): code = ( "import os, tempfile, wave; " "import ctranslate2; " "n=ctranslate2.get_cuda_device_count(); " "assert n>0, 'ctranslate2 cuda device count is 0'; " "from faster_whisper import WhisperModel; " "fd,path=tempfile.mkstemp(suffix='.wav'); os.close(fd); " "wf=wave.open(path,'wb'); wf.setnchannels(1); wf.setsampwidth(2); wf.setframerate(16000); " "wf.writeframes(b'\\x00\\x00'*16000); wf.close(); " "model=WhisperModel('base', device='cuda', compute_type='float16'); " "segments,info=model.transcribe(path, vad_filter=True, beam_size=1); " "list(segments); " "os.remove(path); " "print('ctranslate2_cuda_transcribe_ok', ctranslate2.__version__, 'devices', n)" ) rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=90) msg = out.strip() if rc == 0 else (err or out)[-300:] return self._result("faster_whisper_ctranslate2_cuda", rc == 0, required, msg) def _check_paddle_cuda(self, required): code = ( "import os, json; " "os.environ.setdefault('PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION','python'); " "from app.core.ocr_worker_cli import register_ocr_gpu_dll_paths, _run_ocr_quality_probe; " "register_ocr_gpu_dll_paths(); " "import paddle; " "ok=paddle.device.is_compiled_with_cuda(); " "assert ok, 'paddle is not compiled with cuda'; " "paddle.set_device('gpu:0'); " "from paddleocr import PaddleOCR; " "import inspect; " "kw=dict(use_angle_cls=True, lang='ch', show_log=False); " "kw.update({'device': 'gpu:0'} if 'device' in inspect.signature(PaddleOCR).parameters else {'use_gpu': True}); " "ocr=PaddleOCR(**kw); " "quality_ok, details=_run_ocr_quality_probe(ocr); " "assert quality_ok, details; " "print(paddle.device.get_device(), json.dumps(details, ensure_ascii=False))" ) rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=60) msg = out.strip() if rc == 0 else (err or out)[-300:] return self._result("paddleocr_cuda", rc == 0, required, msg) def _check_onnx_cuda(self, required): code = ( "import onnxruntime as ort; " "providers=ort.get_available_providers(); " "assert 'CUDAExecutionProvider' in providers, providers; " "print(providers)" ) rc, out, err = self._run_cmd([str(self.python_exe), "-c", code], timeout=15) msg = out.strip() if rc == 0 else (err or out)[-300:] return self._result("onnxruntime_cuda_provider", rc == 0, required, msg) @staticmethod def _has_pool_key(cfg): """ApiPool đọc key từ environ + root .env (OPENROUTER/GROQ/NVIDIA/GEMINI/..._KEY_N), config.json (hf_key/groq_key/nvidia_key) và gate .env. Check đúng nguồn pipeline dùng.""" import os as _os prefixes = ("OPENROUTER_KEY", "GROQ_KEY", "NVIDIA_KEY", "GEMINI_KEY", "SILICONFLOW_KEY", "TOGETHERAI_KEY", "XKIRO_KEY", "SUPER_AI_GATE_KEY", "NIM_API_KEY") for k, v in _os.environ.items(): if v and str(v).strip() and k.startswith(prefixes): return True for name in ("hf_key", "groq_key", "nvidia_key", "nim_api_key", "gemini_key", "openrouter_key"): if cfg.get(name): return True for env_file in (Path(__file__).resolve().parents[2] / ".env", Path("D:/TOOL GOM API/.env")): try: if env_file.exists(): for line in env_file.read_text(encoding="utf-8", errors="ignore").splitlines(): line = line.strip() if not line or line.startswith("#") or "=" not in line: continue k, v = line.split("=", 1) if v.strip().strip('"').strip("'") and k.strip().startswith(prefixes): return True except Exception: pass return False def _check_provider_config(self, trans_engine, required): config_path = self.base_dir / "config.json" try: cfg = json.loads(config_path.read_text(encoding="utf-8")) if config_path.exists() else {} eng = str(trans_engine or "") # FIX: "Google" Reserve cho bản miễn phí thật (Google Free/Dịch web). # Gemini (dù tên có chữ Google) vẫn cần GEMINI_KEY. if "Gemini" in eng and "Web" not in eng: if self._has_pool_key(cfg): return self._result("translation_provider_config", True, required, "Gemini key present") return self._result("translation_provider_config", False, required, "missing GEMINI_KEY for Gemini engine") if "Google" in eng: return self._result("translation_provider_config", True, required, "Google (free) selected") if "Ollama" in eng: model = cfg.get("ollama_model") or cfg.get("translation", {}).get("ollama_model") return self._result("translation_provider_config", bool(model), required, "Ollama model configured" if model else "missing Ollama model") # FIX: Ultimate/Pool/9Router/SuperAI/Nvidia/Groq/... chạy qua ApiPool — # check đúng nguồn key của pool thay vì chỉ hf_key/groq_key. if any(k in eng for k in ("Pool", "Ultimate", "Tối thượng", "9Router", "Super AI", "Nvidia", "NIM", "Groq", "OpenRouter", "SiliconFlow", "TogetherAI", "Nemotron", "DeepSeek", "xKiro", "Custom Trans")): if self._has_pool_key(cfg): return self._result("translation_provider_config", True, required, "API Pool key present") return self._result("translation_provider_config", False, required, "missing API Pool key (.env *_KEY_N)") has_key = bool(cfg.get("hf_key") or cfg.get("groq_key")) return self._result("translation_provider_config", has_key, required, "API key present" if has_key else "missing provider API key") except Exception as e: return self._result("translation_provider_config", False, required, str(e)) def _write_reports(self, report): self.output_dir.mkdir(parents=True, exist_ok=True) (self.output_dir / "preflight_report.json").write_text( json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8", ) lines = [ "=== PREFLIGHT REPORT ===", f"Status: {report['status']}", f"Strict GPU: {report['strict_gpu']}", ] for c in report["checks"]: state = "PASS" if c["ok"] else ("FAIL" if c["required"] else "WARN") lines.append(f"- {state} {c['name']}: {c['message']}") (self.output_dir / "preflight_report.txt").write_text("\n".join(lines) + "\n", encoding="utf-8") @staticmethod def _result(name, ok, required, message): return { "name": name, "ok": bool(ok), "required": bool(required), "message": str(message).strip(), }