model-adapt-strategy/main.py

"""
Model Adaptation Strategy for ModelHub XC.

Runs inside a ModelHub container. Receives model info via environment variables.
Performs model download, validation, and produces adaptation results.
"""

import json
import os
import sys
import subprocess
import traceback
from datetime import datetime


def log(msg: str) -> None:
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {msg}", flush=True)


def run_cmd(cmd: str, cwd: str = "/workspace") -> tuple[int, str]:
    log(f"Running: {cmd}")
    result = subprocess.run(
        cmd, shell=True, cwd=cwd,
        capture_output=True, text=True, timeout=3600,
    )
    if result.stdout:
        print(result.stdout)
    if result.stderr:
        print(result.stderr, file=sys.stderr)
    return result.returncode, result.stdout


def download_model(model_address: str, output_dir: str) -> bool:
    log(f"Downloading model from: {model_address}")

    if "huggingface.co" in model_address or "hf.co" in model_address:
        from huggingface_hub import snapshot_download
        model_id = model_address.split("huggingface.co/")[-1]
        snapshot_download(repo_id=model_id, local_dir=output_dir)
    elif "modelscope.cn" in model_address:
        from modelscope import snapshot_download
        model_id = model_address.split("modelscope.cn/models/")[-1]
        snapshot_download(model_id, cache_dir=output_dir)
    else:
        log(f"Unknown source, trying git clone: {model_address}")
        run_cmd(f"git clone {model_address} {output_dir}")

    if os.path.isdir(output_dir) and os.listdir(output_dir):
        log(f"Model downloaded to {output_dir}")
        return True
    log("Model download failed - empty output directory")
    return False


def main():
    log("=" * 60)
    log("ModelHub Adaptation Strategy Starting")
    log("=" * 60)

    model_address = os.environ.get("MODEL_ADDRESS", "")
    config_params = os.environ.get("CONFIG_PARAMS", "{}")

    try:
        config = json.loads(config_params)
    except json.JSONDecodeError:
        config = {}

    log(f"Model: {model_address}")
    log(f"Config: {json.dumps(config, indent=2, ensure_ascii=False)}")

    if not model_address:
        log("ERROR: MODEL_ADDRESS not set")
        sys.exit(1)

    os.makedirs("/workspace/output", exist_ok=True)
    os.makedirs("/workspace/model", exist_ok=True)

    result = {
        "status": "success",
        "model_address": model_address,
        "started_at": datetime.now().isoformat(),
        "steps": [],
    }

    try:
        ok = download_model(model_address, "/workspace/model")
        result["steps"].append({"step": "download", "status": "success" if ok else "failed"})
        if not ok:
            raise RuntimeError("Model download failed")

        log("Validating model files...")
        model_files = []
        for root, dirs, files in os.walk("/workspace/model"):
            for f in files:
                model_files.append(os.path.join(root, f))
        size_mb = sum(os.path.getsize(f) for f in model_files) / (1024 * 1024)
        log(f"Files: {len(model_files)}, Size: {size_mb:.1f} MB")

        result["model_files"] = len(model_files)
        result["model_size_mb"] = round(size_mb, 2)
        result["steps"].append({"step": "validate", "status": "success"})

    except Exception as e:
        log(f"ERROR: {e}")
        traceback.print_exc()
        result["status"] = "failed"
        result["error"] = str(e)

    result["finished_at"] = datetime.now().isoformat()
    with open("/workspace/output/result.json", "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2, ensure_ascii=False)

    log(f"Result: {json.dumps(result, indent=2, ensure_ascii=False)}")
    return 0 if result["status"] == "success" else 1


if __name__ == "__main__":
    sys.exit(main())