120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
"""
|
|
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())
|