a1warming/haloo

1#!/bin/bash2 3# 检查环境变量4if [[ -z "$HF_TOKEN" ]] || [[ -z "$DATASET_ID" ]]; then5    echo "Starting without backup functionality - missing HF_TOKEN or DATASET_ID"6    exec java ${JVM_OPTS} -jar /opt/halo/halo.jar7    exit 08fi9 10# 激活虚拟环境11source /opt/venv/bin/activate12 13# Python 函数: 上传备份14upload_backup() {15  file_path="$1"16  file_name="$2"17  token="$HF_TOKEN"18  repo_id="$DATASET_ID"19 20  python3 -c "21from huggingface_hub import HfApi22import sys23import os24import tarfile25import tempfile26api = HfApi(token='$token')27try:28    api.upload_file(29        path_or_fileobj='$file_path',30        path_in_repo='$file_name',31        repo_id='$repo_id',32        repo_type='dataset'33    )34    print(f'Successfully uploaded $file_name')35except Exception as e:36    print(f'Error uploading file: {str(e)}')37"38}39 40# Python 函数: 下载最新备份41download_latest_backup() {42  token="$HF_TOKEN"43  repo_id="$DATASET_ID"44 45  python3 -c "46from huggingface_hub import HfApi47import sys48import os49import tarfile50import tempfile51api = HfApi(token='$token')52try:53    files = api.list_repo_files(repo_id='$repo_id', repo_type='dataset')54    backup_files = [f for f in files if f.startswith('halo_backup_') and f.endswith('.tar.gz')]55    56    if not backup_files:57        print('No backup files found')58        sys.exit()59        60    latest_backup = sorted(backup_files)[-1]61    62    with tempfile.TemporaryDirectory() as temp_dir:63        filepath = api.hf_hub_download(64            repo_id='$repo_id',65            filename=latest_backup,66            repo_type='dataset',67            local_dir=temp_dir68        )69        70        if filepath and os.path.exists(filepath):71            with tarfile.open(filepath, 'r:gz') as tar:72                tar.extractall(os.path.expanduser('~/.halo2'))73            print(f'Successfully restored backup from {latest_backup}')74        75except Exception as e:76    print(f'Error downloading backup: {str(e)}')77"78}79 80# 首次启动时下载最新备份81echo "Downloading latest backup from HuggingFace..."82download_latest_backup83 84# 同步函数85sync_data() {86    while true; do87        echo "Starting sync process at $(date)"88        89        if [ -d ~/.halo2 ]; then90            timestamp=$(date +%Y%m%d_%H%M%S)91            backup_file="halo_backup_${timestamp}.tar.gz"92            93            # 压缩数据目录94            tar -czf "/tmp/${backup_file}" -C ~/.halo2 .95            96            echo "Uploading backup to HuggingFace..."97            upload_backup "/tmp/${backup_file}" "${backup_file}"98            99            rm -f "/tmp/${backup_file}"100        else101            echo "Data directory does not exist yet, waiting for next sync..."102        fi103        104        SYNC_INTERVAL=${SYNC_INTERVAL:-7200}105        echo "Next sync in ${SYNC_INTERVAL} seconds..."106        sleep $SYNC_INTERVAL107    done108}109 110# 后台启动同步进程111sync_data &112 113# 启动 Halo114exec java ${JVM_OPTS} -jar /opt/halo/halo.jar