Initial model adaptation strategy
This commit is contained in:
16
Dockerfile
Normal file
16
Dockerfile
Normal 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
119
main.py
Normal 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
4
requirements.txt
Normal file
@@ -0,0 +1,4 @@
|
|||||||
|
torch>=2.0.0
|
||||||
|
transformers>=4.30.0
|
||||||
|
huggingface_hub>=0.20.0
|
||||||
|
modelscope>=1.10.0
|
||||||
Reference in New Issue
Block a user