alice0321/nyx
0
1#!/bin/sh2 3# 添加更详细的日志4log_info() {5 echo "[$(date '+%Y-%m-%d %H:%M:%S')] INFO: $1"6}7 8log_error() {9 echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $1"10}11 12# 检查环境变量13if [ -z "$HF_TOKEN" ] || [ -z "$DATASET_ID" ]; then14 log_error "未启用备份功能 - 缺少HF_TOKEN或DATASET_ID环境变量"15 log_info "HF_TOKEN=${HF_TOKEN:0:3}... DATASET_ID=${DATASET_ID}"16 exit 017fi18 19# 创建临时目录20TEMP_DIR="/tmp/sillytavern_backup"21DATA_DIR="/home/node/app/data"22 23# 确保目录存在并有正确权限24mkdir -p $TEMP_DIR25chmod -R 777 $TEMP_DIR26mkdir -p $DATA_DIR27chmod -R 777 $DATA_DIR28 29log_info "临时目录: $TEMP_DIR"30log_info "数据目录: $DATA_DIR"31log_info "HF_TOKEN: ${HF_TOKEN:0:5}..."32log_info "DATASET_ID: $DATASET_ID"33 34# 安装python和huggingface_hub35if ! command -v python3 > /dev/null 2>&1; then36 log_info "正在安装Python..."37 apk add --no-cache python3 py3-pip38else39 log_info "Python3已安装: $(python3 --version)"40fi41 42# 确保pip已安装43if ! command -v pip3 > /dev/null 2>&1; then44 log_info "正在安装pip..."45 apk add --no-cache py3-pip46else47 log_info "Pip3已安装: $(pip3 --version)"48fi49 50# 安装或更新huggingface_hub51log_info "正在安装/更新huggingface_hub..."52pip3 install --no-cache-dir --upgrade huggingface_hub53log_info "huggingface_hub安装完成"54 55# 测试huggingface_hub是否正常工作56if ! python3 -c "import huggingface_hub; print('huggingface_hub版本:', huggingface_hub.__version__)" > /dev/null 2>&1; then57 log_error "huggingface_hub导入失败，正在重试安装..."58 pip3 install --no-cache-dir huggingface_hub59fi60 61# 测试权限是否正常62touch "${TEMP_DIR}/test_file" && rm "${TEMP_DIR}/test_file"63if [ $? -ne 0 ]; then64 log_error "临时目录权限测试失败，正在修复权限..."65 chmod -R 777 $TEMP_DIR66fi67 68# 测试与HuggingFace API的连接69log_info "正在测试与HuggingFace API的连接..."70python3 -c "71from huggingface_hub import HfApi72try:73 api = HfApi(token='$HF_TOKEN')74 user_info = api.whoami()75 print(f'成功连接到HuggingFace API，用户: {user_info}')76except Exception as e:77 print(f'连接HuggingFace API失败: {str(e)}')78 exit(1)79"80if [ $? -ne 0 ]; then81 log_error "HuggingFace API连接测试失败，请检查令牌是否有效"82else83 log_info "HuggingFace API连接测试成功"84fi85 86# 生成唯一的测试文件名87TEST_FILE_NAME="test_file_$(date +%s)"88 89# 测试数据集权限90log_info "正在测试Dataset权限..."91python3 -c "92from huggingface_hub import HfApi93try:94 api = HfApi(token='$HF_TOKEN')95 96 # 创建本地测试文件97 with open('$TEMP_DIR/test_file', 'w') as f:98 f.write('test')99 100 # 上传测试文件101 test_file_name = '$TEST_FILE_NAME'102 print(f'正在上传测试文件: {test_file_name}')103 104 api.upload_file(105 path_or_fileobj='$TEMP_DIR/test_file',106 path_in_repo=test_file_name,107 repo_id='$DATASET_ID',108 repo_type='dataset'109 )110 print('成功上传测试文件到Dataset')111 112 # 删除已上传的测试文件113 print('正在删除测试文件...')114 api.delete_file(115 path_in_repo=test_file_name,116 repo_id='$DATASET_ID',117 repo_type='dataset'118 )119 print('已成功删除测试文件')120 121except Exception as e:122 print(f'Dataset权限测试失败: {str(e)}')123 exit(1)124"125if [ $? -ne 0 ]; then126 log_error "Dataset权限测试失败，请检查DATASET_ID是否正确且有写入权限"127else128 log_info "Dataset权限测试成功，测试文件已清理"129fi130 131# 确保本地测试文件被删除132rm -f "$TEMP_DIR/test_file"133 134# 上传备份135upload_backup() {136 file_path="$1"137 file_name="$2"138 139 if [ ! -f "$file_path" ]; then140 log_error "备份文件不存在: $file_path"141 return 1142 fi143 144 log_info "开始上传备份: $file_name ($(du -h $file_path | cut -f1))"145 146 python3 -c "147from huggingface_hub import HfApi148import sys149import os150import time151 152def manage_backups(api, repo_id, max_files=10):153 try:154 files = api.list_repo_files(repo_id=repo_id, repo_type='dataset')155 backup_files = [f for f in files if f.startswith('sillytavern_backup_') and f.endswith('.tar.gz')]156 backup_files.sort()157 158 if len(backup_files) >= max_files:159 files_to_delete = backup_files[:(len(backup_files) - max_files + 1)]160 for file_to_delete in files_to_delete:161 try:162 api.delete_file(path_in_repo=file_to_delete, repo_id=repo_id, repo_type='dataset')163 print(f'已删除旧备份: {file_to_delete}')164 except Exception as e:165 print(f'删除 {file_to_delete} 时出错: {str(e)}')166 except Exception as e:167 print(f'管理备份文件时出错: {str(e)}')168 169token='$HF_TOKEN'170repo_id='$DATASET_ID'171 172try:173 api = HfApi(token=token)174 175 # 检查文件大小176 file_size = os.path.getsize('$file_path')177 print(f'备份文件大小: {file_size / (1024*1024):.2f} MB')178 179 # 确认Dataset存在180 try:181 dataset_info = api.dataset_info(repo_id=repo_id)182 print(f'Dataset信息: {dataset_info.id}')183 except Exception as e:184 print(f'获取Dataset信息失败: {str(e)}')185 186 start_time = time.time()187 print(f'开始上传: {start_time}')188 189 # 上传文件190 api.upload_file(191 path_or_fileobj='$file_path',192 path_in_repo='$file_name',193 repo_id=repo_id,194 repo_type='dataset'195 )196 197 end_time = time.time()198 print(f'上传完成，耗时: {end_time - start_time:.2f} 秒')199 print(f'成功上传 $file_name')200 201 # 管理备份202 manage_backups(api, repo_id)203except Exception as e:204 print(f'上传文件时出错: {str(e)}')205 sys.exit(1)206"207 if [ $? -ne 0 ]; then208 log_error "备份上传失败"209 return 1210 else211 log_info "备份上传成功"212 return 0213 fi214}215 216# 下载最新备份217download_latest_backup() {218 log_info "开始下载最新备份..."219 220 python3 -c "221from huggingface_hub import HfApi222import sys223import os224import tarfile225import tempfile226import time227 228try:229 api = HfApi(token='$HF_TOKEN')230 print('已创建API实例')231 232 # 列出仓库文件233 try:234 files = api.list_repo_files(repo_id='$DATASET_ID', repo_type='dataset')235 print(f'仓库文件数量: {len(files)}')236 except Exception as e:237 print(f'列出仓库文件失败: {str(e)}')238 sys.exit(1)239 240 backup_files = [f for f in files if f.startswith('sillytavern_backup_') and f.endswith('.tar.gz')]241 print(f'找到备份文件数量: {len(backup_files)}')242 243 if not backup_files:244 print('未找到备份文件')245 sys.exit(0)246 247 # 按名称排序（实际上是按时间戳排序）248 latest_backup = sorted(backup_files)[-1]249 print(f'最新备份文件: {latest_backup}')250 251 with tempfile.TemporaryDirectory() as temp_dir:252 print(f'创建临时目录: {temp_dir}')253 254 start_time = time.time()255 print(f'开始下载: {start_time}')256 257 # 下载文件258 try:259 filepath = api.hf_hub_download(260 repo_id='$DATASET_ID',261 filename=latest_backup,262 repo_type='dataset',263 local_dir=temp_dir264 )265 print(f'文件下载到: {filepath}')266 except Exception as e:267 print(f'下载文件失败: {str(e)}')268 sys.exit(1)269 270 end_time = time.time()271 print(f'下载完成，耗时: {end_time - start_time:.2f} 秒')272 273 if filepath and os.path.exists(filepath):274 # 确保目标目录存在275 os.makedirs('$DATA_DIR', exist_ok=True)276 277 # 检查文件权限278 print(f'文件权限: {oct(os.stat(filepath).st_mode)[-3:]}')279 280 # 解压文件281 try:282 with tarfile.open(filepath, 'r:gz') as tar:283 print('开始解压文件...')284 tar.extractall('$DATA_DIR')285 print('文件解压完成')286 except Exception as e:287 print(f'解压文件失败: {str(e)}')288 sys.exit(1)289 290 print(f'成功从 {latest_backup} 恢复备份')291 else:292 print('下载的文件路径无效')293 sys.exit(1)294except Exception as e:295 print(f'下载备份过程中出错: {str(e)}')296 sys.exit(1)297"298 if [ $? -ne 0 ]; then299 log_error "备份下载失败"300 return 1301 else302 log_info "备份下载成功"303 return 0304 fi305}306 307# 首次启动时下载最新备份308log_info "正在从HuggingFace下载最新备份..."309download_latest_backup310 311# 同步函数312sync_data() {313 log_info "数据同步服务已启动"314 315 while true; do316 log_info "开始同步进程，时间: $(date)"317 318 if [ -d "$DATA_DIR" ]; then319 timestamp=$(date +%Y%m%d_%H%M%S)320 backup_file="sillytavern_backup_${timestamp}.tar.gz"321 backup_path="${TEMP_DIR}/${backup_file}"322 323 log_info "创建备份文件: $backup_path"324 325 # 检查数据目录内容326 file_count=$(find "$DATA_DIR" -type f | wc -l)327 log_info "数据目录文件数量: $file_count"328 329 if [ "$file_count" -eq 0 ]; then330 log_info "数据目录为空，跳过备份"331 else332 # 压缩数据目录333 tar -czf "$backup_path" -C "$DATA_DIR" .334 if [ $? -ne 0 ]; then335 log_error "创建压缩文件失败"336 else337 log_info "压缩文件创建成功: $(du -h $backup_path | cut -f1)"338 339 # 上传备份340 log_info "正在上传备份到HuggingFace..."341 upload_backup "$backup_path" "$backup_file"342 343 # 删除临时备份文件344 rm -f "$backup_path"345 log_info "已删除临时备份文件"346 fi347 fi348 else349 log_error "数据目录不存在: $DATA_DIR"350 mkdir -p "$DATA_DIR"351 chmod -R 777 "$DATA_DIR"352 fi353 354 # 设置同步间隔355 SYNC_INTERVAL=${SYNC_INTERVAL:-3600}356 log_info "下次同步将在 ${SYNC_INTERVAL} 秒后进行..."357 sleep $SYNC_INTERVAL358 done359}360 361# 启动同步进程362sync_data