276 lines
9.4 KiB
Python
276 lines
9.4 KiB
Python
import os
|
||
import sys
|
||
import shutil
|
||
from modelscope.hub.snapshot_download import snapshot_download
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
# 使用持久化的缓存目录（挂载到宿主机）
|
||
CACHE_DIR = os.getenv('MODEL_CACHE_DIR', '/app/model_cache')
|
||
os.makedirs(CACHE_DIR, exist_ok=True)
|
||
|
||
|
||
def discover_models():
|
||
"""从环境变量中发现所有模型"""
|
||
models = {}
|
||
for key, value in os.environ.items():
|
||
if key.endswith('_MODEL_ID'):
|
||
prefix = key[:-9]
|
||
service_name = os.getenv(f'{prefix}_SERVICE_NAME', f'{prefix.lower()}_classifier')
|
||
service_folder = service_name.replace('_classifier', '')
|
||
|
||
models[prefix.lower()] = {
|
||
'prefix': prefix,
|
||
'model_id': value,
|
||
'port': os.getenv(f'{prefix}_PORT', '5002'),
|
||
'service_name': service_name,
|
||
'model_dir': os.getenv(f'{prefix}_MODEL_DIR', f'/app/services/bert_{service_folder}/model'),
|
||
'service_folder': service_folder,
|
||
'cache_dir': os.path.join(CACHE_DIR, prefix.lower())
|
||
}
|
||
return models
|
||
|
||
|
||
def check_model_complete(model_dir):
|
||
"""检查模型是否完整"""
|
||
if not os.path.exists(model_dir):
|
||
return False
|
||
|
||
try:
|
||
files = os.listdir(model_dir)
|
||
if not files:
|
||
return False
|
||
|
||
has_weights = any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files)
|
||
has_config = any(f == 'config.json' for f in files)
|
||
|
||
if has_weights and has_config:
|
||
return True
|
||
return False
|
||
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def find_model_files(cache_dir):
|
||
"""在下载的缓存目录中查找实际的模型文件"""
|
||
# 检查是否直接有模型文件
|
||
direct_files = [f for f in os.listdir(cache_dir) if
|
||
not f.startswith('.') and os.path.isfile(os.path.join(cache_dir, f))]
|
||
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in direct_files):
|
||
return cache_dir
|
||
|
||
# 检查 models/ 子目录（ModelScope 的标准结构）
|
||
models_dir = os.path.join(cache_dir, 'models')
|
||
if os.path.exists(models_dir):
|
||
for root, dirs, files in os.walk(models_dir):
|
||
if 'snapshots' in root:
|
||
for f in files:
|
||
if f.endswith(('.safetensors', '.bin', '.pt', '.pth')):
|
||
return root
|
||
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files):
|
||
return root
|
||
|
||
# 检查是否存在单层子目录
|
||
subdirs = [d for d in os.listdir(cache_dir) if os.path.isdir(os.path.join(cache_dir, d)) and not d.startswith('.')]
|
||
for subdir in subdirs:
|
||
subpath = os.path.join(cache_dir, subdir)
|
||
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in os.listdir(subpath)):
|
||
return subpath
|
||
|
||
# 递归查找任何包含模型文件的目录
|
||
for root, dirs, files in os.walk(cache_dir):
|
||
if any(f.endswith(('.safetensors', '.bin', '.pt', '.pth')) for f in files):
|
||
return root
|
||
|
||
return None
|
||
|
||
|
||
def copy_model_from_cache(cache_dir, target_dir):
|
||
"""从缓存目录复制模型到目标目录"""
|
||
src_dir = find_model_files(cache_dir)
|
||
|
||
if src_dir is None:
|
||
print(f" ✗ Could not find model files in cache: {cache_dir}")
|
||
return False
|
||
|
||
print(f" ✓ Found cached model files in: {src_dir}")
|
||
|
||
os.makedirs(target_dir, exist_ok=True)
|
||
|
||
copied_count = 0
|
||
for file in os.listdir(src_dir):
|
||
src_file = os.path.join(src_dir, file)
|
||
if os.path.isfile(src_file):
|
||
dest_file = os.path.join(target_dir, file)
|
||
shutil.copy2(src_file, dest_file)
|
||
copied_count += 1
|
||
|
||
print(f" ✓ Copied {copied_count} files to {target_dir}")
|
||
return copied_count > 0
|
||
|
||
|
||
def download_model_with_fallback(name, config):
|
||
"""下载单个模型，带备选方案"""
|
||
print("=" * 50)
|
||
print(f"Processing {name} model...")
|
||
print(f"Model ID: {config['model_id']}")
|
||
print(f"Target: {config['model_dir']}")
|
||
print(f"Cache: {config['cache_dir']}")
|
||
print("=" * 50)
|
||
|
||
try:
|
||
os.makedirs(config['model_dir'], exist_ok=True)
|
||
os.makedirs(config['cache_dir'], exist_ok=True)
|
||
|
||
# 1. 检查目标目录是否已有完整模型
|
||
if check_model_complete(config['model_dir']):
|
||
print(f"✅ Model already exists in target, skipping")
|
||
return True
|
||
|
||
# 2. 检查缓存目录是否有模型
|
||
if check_model_complete(config['cache_dir']):
|
||
print(f"✅ Found complete model in local cache")
|
||
print(f" Copying from cache to target...")
|
||
|
||
if copy_model_from_cache(config['cache_dir'], config['model_dir']):
|
||
if check_model_complete(config['model_dir']):
|
||
print(f"✅ Model copied successfully from cache")
|
||
return True
|
||
else:
|
||
print(f"⚠️ Model copied but appears incomplete, will re-download")
|
||
|
||
# 3. 缓存中没有，需要下载
|
||
print(f"📥 Downloading model from ModelScope...")
|
||
|
||
download_success = False
|
||
|
||
# 尝试方式1：直接下载
|
||
try:
|
||
snapshot_download(
|
||
model_id=config['model_id'],
|
||
cache_dir=config['cache_dir'],
|
||
revision="master",
|
||
ignore_file_pattern=[".git", ".gitattributes"]
|
||
)
|
||
print(f"✅ Download completed")
|
||
download_success = True
|
||
except Exception as e:
|
||
print(f"⚠️ Primary download failed: {e}")
|
||
|
||
# 尝试方式2：使用 HubApi（如果有 token）
|
||
token = os.getenv('MODELSCOPE_TOKEN', '')
|
||
if token:
|
||
print(" Trying alternative download with HubApi...")
|
||
try:
|
||
from modelscope.hub.api import HubApi
|
||
api = HubApi()
|
||
api.login(token=token)
|
||
|
||
snapshot_download(
|
||
model_id=config['model_id'],
|
||
cache_dir=config['cache_dir'],
|
||
revision="master",
|
||
ignore_file_pattern=[".git", ".gitattributes"]
|
||
)
|
||
print(f"✅ Alternative download completed")
|
||
download_success = True
|
||
except Exception as alt_error:
|
||
print(f"✗ Alternative download also failed: {alt_error}")
|
||
raise
|
||
else:
|
||
print("✗ No MODELSCOPE_TOKEN set, cannot use alternative method")
|
||
raise
|
||
|
||
if not download_success:
|
||
return False
|
||
|
||
# 验证并复制
|
||
if not check_model_complete(config['cache_dir']):
|
||
print(f"⚠️ Downloaded model appears incomplete in cache")
|
||
print(f" Attempting to copy anyway...")
|
||
|
||
print(f"Copying from cache to target...")
|
||
if copy_model_from_cache(config['cache_dir'], config['model_dir']):
|
||
if check_model_complete(config['model_dir']):
|
||
print(f"✅ Model deployed successfully!")
|
||
return True
|
||
else:
|
||
print(f"⚠️ Model deployed but appears incomplete")
|
||
return True
|
||
else:
|
||
print(f"✗ Failed to copy model to target")
|
||
return False
|
||
|
||
except Exception as e:
|
||
print(f"✗ Failed to process {name} model: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return False
|
||
|
||
|
||
def main():
|
||
print("=" * 50)
|
||
print("Model Download Script")
|
||
print("=" * 50)
|
||
print(f"Cache directory: {CACHE_DIR}")
|
||
|
||
force_download = os.getenv('FORCE_MODEL_DOWNLOAD', 'false').lower() == 'true'
|
||
if force_download:
|
||
print("⚠️ FORCE_MODEL_DOWNLOAD=true - will re-download all models")
|
||
|
||
print("\nDiscovering models from .env...")
|
||
models = discover_models()
|
||
|
||
if not models:
|
||
print("⚠️ No models found in .env file!")
|
||
sys.exit(1)
|
||
|
||
print(f"\nFound {len(models)} model(s):")
|
||
all_exist = True
|
||
for name, config in models.items():
|
||
target_exists = check_model_complete(config['model_dir'])
|
||
cache_exists = check_model_complete(config['cache_dir'])
|
||
|
||
print(f" {name}: {config['model_id']}")
|
||
print(f" Target: {'✅' if target_exists else '❌'} {config['model_dir']}")
|
||
print(f" Cache: {'✅' if cache_exists else '❌'} {config['cache_dir']}")
|
||
|
||
if not target_exists and not cache_exists:
|
||
all_exist = False
|
||
print("")
|
||
|
||
if all_exist and not force_download:
|
||
print("✅ All models exist!")
|
||
print(" Skipping downloads.")
|
||
sys.exit(0)
|
||
|
||
if force_download:
|
||
print("⚠️ Force download enabled, clearing caches...")
|
||
for name, config in models.items():
|
||
if os.path.exists(config['cache_dir']):
|
||
print(f" Clearing cache for {name}")
|
||
shutil.rmtree(config['cache_dir'])
|
||
os.makedirs(config['cache_dir'], exist_ok=True)
|
||
|
||
print("\n" + "=" * 50)
|
||
print("Processing models...")
|
||
print("=" * 50)
|
||
|
||
success_count = 0
|
||
for name, config in models.items():
|
||
if download_model_with_fallback(name, config):
|
||
success_count += 1
|
||
print("")
|
||
|
||
print("=" * 50)
|
||
print(f"Processing completed: {success_count}/{len(models)} models")
|
||
print("=" * 50)
|
||
|
||
if success_count < len(models):
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |