Fine-tuning/backend/app/services/model_test_service.py

from pathlib import Path
from typing import Any

from app.config import get_settings
from app.core.logging import logger

settings = get_settings()


async def test_model(model_id: str, prompt: str, max_new_tokens: int = 128, temperature: float = 0.8, top_p: float = 0.95) -> dict[str, Any]:
    """加载已缓存模型并生成测试响应。"""
    if settings.use_remote_compute:
        return _test_model_remote(model_id, prompt, max_new_tokens, temperature, top_p)
    return _test_model_local(model_id, prompt, max_new_tokens, temperature, top_p)


def _test_model_remote(model_id: str, prompt: str, max_new_tokens: int, temperature: float, top_p: float) -> dict[str, Any]:
    """通过 SSH 在算力节点执行模型测试。

    流程：scp 到远端宿主机 → docker cp 传入容器 → docker exec 执行 → 清理
    """
    import json
    import os
    import tempfile
    from app.core.remote_executor import scp_to_remote, ssh_exec

    container = settings.compute_node_docker_container
    python = settings.compute_node_python
    workdir = settings.compute_node_workdir

    # 独立的模型测试脚本内容（零 app/db 依赖）
    python_script = """\
import json, sys
from pathlib import Path
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel

def find_model_path(model_id):
    candidates = [
        '/root/.cache/huggingface/hub',
        '/root/.cache/modelscope/hub',
        '/root/models',
    ]
    for base in candidates:
        bp = Path(base)
        if not bp.is_dir():
            continue
        for child in bp.rglob('config.json'):
            parent = child.parent
            if parent.is_dir():
                return str(parent)
    return None

model_id = sys.argv[1]
prompt = sys.argv[2]
max_new_tokens = int(sys.argv[3])
temperature = float(sys.argv[4])
top_p = float(sys.argv[5])

model_path = find_model_path(model_id)
if model_path is None:
    print(json.dumps({'error': 'Model not found in cache'}))
    sys.exit(1)

t = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
t.pad_token = t.pad_token or t.eos_token

m = None
for cls, kw in [
    (AutoModelForCausalLM, {'trust_remote_code': True}),
    (AutoModel, {'trust_remote_code': True}),
]:
    try:
        m = cls.from_pretrained(model_path, torch_dtype=torch.float16, device_map='auto', **kw)
        break
    except Exception:
        pass

if m is None:
    print(json.dumps({'error': 'Unable to load model'}))
    sys.exit(1)

m.eval()
inp = t(prompt, return_tensors='pt').to(m.device)
out = m.generate(**inp, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, do_sample=%s, pad_token_id=t.eos_token_id)
gen = t.decode(out[0][inp['input_ids'].shape[1]:], skip_special_tokens=True)
print(json.dumps({'generated_text': gen}))
""" % str(temperature > 0).lower()

    remote_script = "/tmp/_model_test.py"
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as tmp:
        tmp.write(python_script)
        tmp.flush()
        tmp_path = tmp.name

    try:
        # Step 1: SCP 到远端宿主机
        host_tmp = "/tmp/_model_test_host.py"
        code, out, err = scp_to_remote(tmp_path, host_tmp)
        if code != 0:
            logger.error(f"SCP failed: {err}")
            return {"error": f"Failed to upload script: {err.strip()}"}

        # Step 2: docker cp 把文件从宿主机传入容器
        cp_cmd = f"docker cp {host_tmp} {container}:/tmp/_model_test.py"
        code, out, err = ssh_exec(cp_cmd, timeout=10)
        if code != 0:
            logger.error(f"docker cp failed: {err}")
            return {"error": f"Failed to copy script to container: {err.strip()}"}

        # Step 3: docker exec 执行容器内的脚本
        safe_prompt = prompt.replace("'", "\\'")
        run_cmd = f"docker exec -w {workdir} {container} {python} /tmp/_model_test.py '{model_id}' '{safe_prompt}' {max_new_tokens} {temperature} {top_p}"
        code, stdout, stderr = ssh_exec(run_cmd, timeout=600)

        if code != 0:
            logger.error(f"Remote model test failed: {stderr}")
            return {"error": stderr.strip() or "Remote test failed"}
    finally:
        os.unlink(tmp_path)
        ssh_exec(f"rm -f /tmp/_model_test_host.py", timeout=10)
        ssh_exec(f"docker exec {container} rm -f /tmp/_model_test.py", timeout=10)

    logger.info(f"Remote test result: code={code}, stdout_len={len(stdout)}, stderr_len={len(stderr)}")
    if stdout:
        logger.info(f"stdout (first 500): {stdout[:500]}")
    if stderr:
        logger.info(f"stderr (first 500): {stderr[:500]}")

    if code != 0:
        logger.error(f"Remote model test failed: {stderr}")
        return {"error": stderr.strip() or "Remote test failed"}

    # 提取最后一行 JSON
    for line in reversed(stdout.strip().split("\n")):
        line = line.strip()
        if line.startswith("{"):
            try:
                result = json.loads(line)
                result["model_id"] = model_id
                result["prompt"] = prompt
                return result
            except json.JSONDecodeError:
                continue

    return {"error": f"Invalid response: {stdout[:500]}"}


async def _test_model_local(model_id: str, prompt: str, max_new_tokens: int, temperature: float, top_p: float) -> dict[str, Any]:
    """本地执行模型测试（仅用于开发环境）。"""
    import torch
    from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, AutoConfig

    from app.services.model_service import resolve_model_path
    model_path = await resolve_model_path(model_id)
    if not model_path:
        return {"error": f"Model not found in cache: {model_id}"}

    model_dir = Path(model_path)
    if not (model_dir / "config.json").exists():
        return {"error": f"Model directory not found: {model_dir}"}

    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # 通用加载策略：尝试多种加载方式，自动兼容各种新架构
    model = None
    for loader_cls, kwargs in [
        (AutoModelForCausalLM, {"trust_remote_code": True}),
        (AutoModel, {"trust_remote_code": True}),
    ]:
        try:
            model = loader_cls.from_pretrained(
                model_dir,
                torch_dtype=torch.float16,
                device_map="auto",
                **kwargs,
            )
            break
        except Exception:
            continue

    if model is None:
        return {"error": f"Unable to load model with any available loader. Model type may not be supported yet."}
    model.eval()

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_text = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)

    return {
        "model_id": model_id,
        "prompt": prompt,
        "generated_text": generated_text,
    }