File size: 10,907 Bytes
13ad5f7
 
7d26d69
13ad5f7
 
 
 
85c4208
 
 
 
 
 
 
 
 
 
7d26d69
13ad5f7
 
 
 
7d26d69
 
13ad5f7
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13ad5f7
 
7d26d69
13ad5f7
 
7d26d69
 
 
 
 
 
 
 
 
85c4208
 
 
 
 
 
 
 
 
7d26d69
 
 
 
13ad5f7
85c4208
7d26d69
 
13ad5f7
 
7d26d69
85c4208
13ad5f7
85c4208
 
7d26d69
85c4208
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13ad5f7
85c4208
13ad5f7
7d26d69
 
13ad5f7
 
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13ad5f7
 
7d26d69
 
 
 
 
 
 
 
b10884e
7d26d69
 
 
 
 
 
 
 
b10884e
13ad5f7
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85c4208
7d26d69
 
 
 
 
 
dce165e
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85c4208
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85c4208
13ad5f7
 
 
85c4208
 
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13ad5f7
7d26d69
dce165e
7d26d69
85c4208
7d26d69
 
85c4208
7d26d69
 
 
85c4208
 
7d26d69
 
 
 
 
 
 
 
 
 
 
 
13ad5f7
 
 
7d26d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
#!/bin/bash

# 此脚本仅创建同步脚本,不在构建阶段运行同步进程

# 创建数据目录
mkdir -p ./data

# 定义哈希计算函数
get_file_hash() {
    local file_path="$1"
    if [ -f "$file_path" ]; then
        md5sum "$file_path" | awk '{print $1}'
    else
        echo "文件不存在"
    fi
}

# 创建 Hugging Face 同步脚本,优化内存使用
cat > /tmp/hf_sync.py << 'EOL'
from huggingface_hub import HfApi
import sys
import os
import gc
import time

def manage_backups(api, repo_id, max_files=10):
    """管理备份文件,保留最新的max_files个文件,内存优化版"""
    try:
        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
        backup_files = [f for f in files if f.startswith('webui_backup_') and f.endswith('.db')]
        
        # 按日期分组文件(从文件名中提取日期)
        backup_by_date = {}
        for file in backup_files:
            try:
                date_part = file.split('_')[2].split('.')[0]
                backup_by_date[date_part] = file
            except:
                continue
        
        # 保留最新的max_files个文件
        sorted_dates = sorted(backup_by_date.keys(), reverse=True)
        if len(sorted_dates) > max_files:
            files_to_delete = [backup_by_date[date] for date in sorted_dates[max_files:]]
            
            # 分批删除文件以减少内存使用
            batch_size = 3
            for i in range(0, len(files_to_delete), batch_size):
                batch = files_to_delete[i:i+batch_size]
                for file in batch:
                    api.delete_file(path_in_repo=file, repo_id=repo_id, repo_type="dataset")
                    print(f"已删除旧备份: {file}")
                # 强制垃圾回收
                gc.collect()
                # 批次间暂停
                if i + batch_size < len(files_to_delete):
                    time.sleep(2)
    except Exception as e:
        print(f"管理备份错误: {str(e)}")
    finally:
        # 确保垃圾回收
        gc.collect()

def upload_backup(file_path, file_name, token, repo_id):
    """上传备份文件到Hugging Face,内存优化版"""
    api = HfApi(token=token)
    try:
        # 检查文件是否存在
        try:
            files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
            if file_name in files:
                api.delete_file(path_in_repo=file_name, repo_id=repo_id, repo_type="dataset")
                print(f"已删除同名文件: {file_name}")
        except Exception as e:
            print(f"检查文件存在错误: {str(e)}")
        gc.collect()
        
        # 上传新文件
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=file_name,
            repo_id=repo_id,
            repo_type="dataset"
        )
        print(f"成功上传: {file_name}")
        
        # 管理备份,可选执行
        if os.environ.get("MANAGE_BACKUPS", "true").lower() == "true":
            manage_backups(api, repo_id)
    except Exception as e:
        print(f"上传失败: {str(e)}")
    finally:
        gc.collect()

def download_latest_backup(token, repo_id):
    """从Hugging Face下载最新备份,内存优化版"""
    api = HfApi(token=token)
    try:
        files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
        backup_files = [f for f in files if f.startswith('webui_backup_') and f.endswith('.db')]
        
        if not backup_files:
            return False
        
        # 找到最新的文件(按日期排序)
        latest_file = max(backup_files, key=lambda x: x.split('_')[2].split('.')[0])
        file_path = api.hf_hub_download(
            repo_id=repo_id,
            filename=latest_file,
            repo_type="dataset"
        )
        
        if file_path and os.path.exists(file_path):
            os.makedirs('./data', exist_ok=True)
            os.system(f'cp "{file_path}" ./data/webui.db')
            print(f"成功从Hugging Face恢复: {latest_file}")
            return True
        else:
            return False
    except Exception as e:
        print(f"下载失败: {str(e)}")
        return False
    finally:
        gc.collect()

if __name__ == "__main__":
    try:
        action = sys.argv[1]
        token = sys.argv[2]
        repo_id = sys.argv[3]
        
        if action == "upload":
            file_path = sys.argv[4]
            file_name = sys.argv[5]
            upload_backup(file_path, file_name, token, repo_id)
        elif action == "download":
            download_latest_backup(token, repo_id)
    except Exception as e:
        print(f"脚本执行错误: {str(e)}")
    finally:
        # 最终垃圾回收
        gc.collect()
EOL

# 创建同步服务启动脚本(不在构建时执行)
cat > /tmp/start_sync.sh << 'EOL'
#!/bin/bash

# 检查必要的环境变量
if [ -z "$WEBDAV_URL" ] || [ -z "$WEBDAV_USERNAME" ] || [ -z "$WEBDAV_PASSWORD" ]; then
    echo "缺少必要的环境变量: WEBDAV_URL、WEBDAV_USERNAME 或 WEBDAV_PASSWORD"
    export WEBDAV_ENABLED="false"
else
    export WEBDAV_ENABLED="true"
fi

if [ -z "$HF_TOKEN" ] || [ -z "$DATASET_ID" ]; then
    echo "缺少必要的环境变量: HF_TOKEN 或 DATASET_ID"
    export HF_ENABLED="false"
else
    export HF_ENABLED="true"
fi

# 初始化数据恢复策略
echo "初始化数据恢复..."

# 尝试恢复数据
restore_data() {
    # 首先尝试从 WebDAV 恢复
    if [ "$WEBDAV_ENABLED" = "true" ]; then
        echo "尝试从 WebDAV 获取文件列表..."
        webdav_files=$(curl -s -X PROPFIND --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" -H "Depth: 1" "$WEBDAV_URL/openwebui/" | grep '<d:href>' | grep 'webui_[0-9]\{8\}.db' | sed 's|</?d:href>||g')

        if [ -n "$webdav_files" ]; then
            latest_file=$(echo "$webdav_files" | sort -r | head -n 1)
            download_url="$WEBDAV_URL/openwebui/$latest_file"
            curl -L -o "./data/webui.db" --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" "$download_url" && {
                echo "成功从 WebDAV 下载最新数据库: $latest_file"
                return 0
            }
        fi
        echo "WebDAV 恢复失败"
    fi

    # 如果 WebDAV 失败,尝试从 Hugging Face 恢复
    if [ "$HF_ENABLED" = "true" ]; then
        echo "尝试从 Hugging Face 恢复..."
        python /tmp/hf_sync.py download "$HF_TOKEN" "$DATASET_ID" && {
            echo "成功从 Hugging Face 恢复"
            return 0
        }
    fi

    # 所有恢复方法都失败
    echo "所有恢复失败,创建空数据库..."
    touch ./data/webui.db
    return 1
}

# WebDAV 同步函数(使用 cron 风格的调度)
webdav_sync() {
    if [ "$WEBDAV_ENABLED" != "true" ]; then
        echo "WebDAV 同步已禁用"
        return
    fi

    echo "执行 WebDAV 同步: $(date)"
    
    if [ ! -f "./data/webui.db" ]; then
        echo "未找到 webui.db,跳过同步"
        return
    fi

    # 生成文件名(包含年月日)
    current_date=$(date +'%Y%m%d')
    file_name="webui_${current_date}.db"
    upload_url="$WEBDAV_URL/openwebui/${file_name}"
    
    # 计算本地文件哈希
    local_hash=$(get_file_hash "./data/webui.db")
    
    # 获取远程文件哈希(通过临时下载)
    remote_temp="/tmp/webui_remote.db"
    curl -s -o "$remote_temp" --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" "$upload_url" > /dev/null 2>&1
    remote_hash=$(get_file_hash "$remote_temp")
    rm -f "$remote_temp"
    
    if [ "$local_hash" = "$remote_hash" ]; then
        echo "文件未变化,跳过 WebDAV 上传"
        return
    fi

    echo "检测到文件变化,开始上传到 WebDAV..."
    curl -L -T "./data/webui.db" --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" "$upload_url" && {
        echo "WebDAV 上传成功: $file_name"
        
        # 更新主文件(覆盖 webui.db)
        main_url="$WEBDAV_URL/openwebui/webui.db"
        curl -L -T "./data/webui.db" --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" "$main_url" && {
            echo "主文件更新成功"
        } || {
            echo "主文件更新失败"
        }
    } || {
        echo "WebDAV 上传失败"
    }

    # 清理过期 WebDAV 文件(保留最近 7 天)
    echo "清理过期 WebDAV 文件..."
    webdav_files=$(curl -s -X PROPFIND --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" -H "Depth: 1" "$WEBDAV_URL/openwebui/" | grep '<d:href>' | grep 'webui_[0-9]\{8\}.db' | sed 's|</?d:href>||g')
    cleanup_days=7
    cutoff_date=$(date -d "-${cleanup_days} days" +%Y%m%d)
    
    for file in $webdav_files; do
        file_date=$(echo "$file" | grep -oE '[0-9]{8}')
        if [ -n "$file_date" ] && [ "$file_date" -lt "$cutoff_date" ]; then
            delete_url="$WEBDAV_URL/openwebui/$file"
            curl -X DELETE --user "$WEBDAV_USERNAME:$WEBDAV_PASSWORD" "$delete_url" && echo "删除过期文件: $file"
        fi
    done
}

# Hugging Face 同步函数
hf_sync() {
    if [ "$HF_ENABLED" != "true" ]; then
        echo "Hugging Face 同步已禁用"
        return
    fi

    echo "执行 Hugging Face 同步: $(date)"
    
    if [ ! -f "./data/webui.db" ]; then
        echo "未找到数据库文件,跳过 Hugging Face 同步"
        return
    fi

    current_date=$(date +'%Y%m%d')
    backup_file="webui_backup_${current_date}.db"
    temp_path="/tmp/${backup_file}"
    cp "./data/webui.db" "$temp_path"
    
    echo "正在上传到 Hugging Face..."
    python /tmp/hf_sync.py upload "$HF_TOKEN" "$DATASET_ID" "$temp_path" "$backup_file"
    rm -f "$temp_path"
}

# 主函数
main() {
    # 恢复数据
    restore_data
    
    # 设置同步间隔(默认2小时)
    SYNC_INTERVAL=${SYNC_INTERVAL:-7200}
    echo "同步间隔设置为: ${SYNC_INTERVAL} 秒"
    
    # 循环执行同步,但使用更高效的方式
    while true; do
        # 每次同步前先休眠,避免启动时立即同步
        sleep $SYNC_INTERVAL
        
        # 执行WebDAV同步
        if [ "$WEBDAV_ENABLED" = "true" ]; then
            webdav_sync
        fi
        
        # 清理内存
        sync
        echo 3 > /proc/sys/vm/drop_caches 2>/dev/null || true
        
        # 执行Hugging Face同步
        if [ "$HF_ENABLED" = "true" ]; then
            hf_sync
        fi
        
        # 清理内存
        sync
        echo 3 > /proc/sys/vm/drop_caches 2>/dev/null || true
    done
}

# 以非阻塞方式启动主函数
main &
EOL

# 确保脚本可执行
chmod +x /tmp/start_sync.sh

# 修改启动脚本以包含同步功能,但在容器启动时启动而不是构建时
cat > /tmp/sync_starter.sh << 'EOL'
# 在容器启动后启动同步服务
/bin/bash /tmp/start_sync.sh &
EOL

# 注意:此处只是准备脚本,不执行它们
echo "同步脚本已准备就绪,将在容器启动时执行"