Initial model adaptation strategy

This commit is contained in:
7355608
2026-07-03 21:24:00 +08:00
commit 4059cd2604
3 changed files with 139 additions and 0 deletions

16
Dockerfile Normal file
View File

@@ -0,0 +1,16 @@
FROM python:3.10-slim
WORKDIR /workspace
RUN apt-get update && apt-get install -y --no-install-recommends \
git git-lfs \
&& rm -rf /var/lib/apt/lists/*
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY main.py .
ENV PYTHONUNBUFFERED=1
ENTRYPOINT ["python", "main.py"]

119
main.py Normal file
View File

@@ -0,0 +1,119 @@
"""
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())

4
requirements.txt Normal file
View File

@@ -0,0 +1,4 @@
torch>=2.0.0
transformers>=4.30.0
huggingface_hub>=0.20.0
modelscope>=1.10.0