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