import os import sys import time import hashlib import asyncssh import asyncio from tqdm import tqdm from loguru import logger from file import get_files_to_upload # 配置日志输出到文件和控制台 logger.add("output.log") logger.add(sys.stderr, level="WARNING") # 上传文件的函数 async def upload_file(local_path, local_md5, remote_path, sftp): file_size = os.path.getsize(local_path) async with await sftp.open(remote_path, 'wb') as file: bytes_uploaded = 0 with open(local_path, 'rb') as local_file: with tqdm(total=file_size, unit='B', unit_scale=True, ncols=80, desc=os.path.basename(local_path)) as progress_bar: while True: chunk = local_file.read(65536) # 读取文件块 if not chunk: break await file.write(chunk) # 写入文件块 bytes_uploaded += len(chunk) progress_bar.update(len(chunk)) if bytes_uploaded == file_size: remote_md5 = await asyncio.wait_for(calculate_remote_md5(sftp, remote_path), timeout=calculate_timeout) # 设置超时时间为5秒 if remote_md5 == local_md5: await write_remote_meta_md5(remote_path, remote_md5, sftp) logger.info("文件上传成功") else: logger.info("本地与远程MD5不一致，请检查上传过程中是否发生错误！") else: logger.info("文件上传失败，请检查文件后再试验。") # 递归创建远程目录 async def create_remote_directory_recursive(directory, sftp): parent_dir = os.path.dirname(directory) try: await sftp.stat(parent_dir) except asyncssh.SFTPError as e: if "No such file" in str(e): await create_remote_directory_recursive(parent_dir, sftp) else: logger.error(f"检查远程目录失败: {parent_dir}. 错误信息: {e}") return try: await sftp.mkdir(directory) except asyncssh.SFTPError as e: logger.error(f"创建远程目录失败: {directory}. 错误信息: {e}") # 创建一个锁对象 lock = asyncio.Lock() # 创建目录并上传文件的函数 async def create_directory_and_upload(local_path, local_md5, remote_path, sftp): directory = os.path.dirname(remote_path) # 检查远程目录是否存在，如果不存在则创建 try: await lock.acquire() # 获取锁对象 attempts = 5 # 尝试次数 while attempts > 0: try: await sftp.stat(directory) break # 目录已存在，跳出循环 except asyncssh.SFTPError as e: if "No such file" in str(e): await create_remote_directory_recursive(directory, sftp) break # 目录创建成功，跳出循环 elif "Failure" in str(e): # 这里可以根据实际情况修改判断条件 await asyncio.sleep(5) # 等待5秒 attempts -= 1 else: logger.error(f"检查远程目录失败: {directory}. 错误信息: {e}") return finally: lock.release() # 上传文件 for attempt in range(3): # 最多尝试3次上传操作 try: await upload_file(local_path, local_md5, remote_path, sftp) break # 上传成功，退出循环 except Exception as e: logger.error(f"上传文件失败: {os.path.basename(local_path)}。错误信息: {e}") if attempt < 2: logger.info("重新尝试上传文件...") else: logger.info("多次上传文件失败.") return async def compare_md5(local_path, remote_path, sftp): logger.info(f"正在比对文件：{os.path.basename(local_path)}中") try: local_md5 = calculate_md5(local_path) # 获取远程meta文件中的md5 remote_meta_md5 = await get_remote_meta_md5(remote_path, sftp) if remote_meta_md5 and local_md5 == remote_meta_md5: logger.info("Meta文件校验成功 - 跳过远程md5计算。") return # 计算远程MD5码，并设置较长的超时时间 remote_md5 = await asyncio.wait_for(calculate_remote_md5(sftp, remote_path), timeout=calculate_timeout) # 设置超时时间为5秒 if remote_md5 is None: # 执行上传操作 logger.info(f'准备文件上传 - 文件名为：{os.path.basename(local_path)}') await create_directory_and_upload(local_path, local_md5, remote_path, sftp) elif local_md5 != remote_md5: logger.info(f"MD5校验失败 - 校验结果为：{local_md5} - {remote_md5}") # 删除远程文件 for attempt in range(3): # 最多尝试3次删除操作 try: await sftp.remove(remote_path) logger.info(f'删除远程文件成功：{remote_path}') break # 删除成功，退出循环 except Exception as e: logger.error(f"删除远程文件失败: {remote_path}. 错误信息: {e}") if attempt < 2: logger.info("重新尝试删除远程文件...") else: return # 执行上传操作 logger.info(f'准备文件上传 - 文件名为：{os.path.basename(local_path)}') await create_directory_and_upload(local_path, local_md5, remote_path, sftp) else: logger.info(f"MD5码校验成功 - 校验结果为：{local_md5} - {remote_md5}") await write_remote_meta_md5(remote_path, remote_md5, sftp) except asyncio.TimeoutError: logger.error(f"Timeout occurred while comparing MD5 for {local_path} and {remote_path}") return async def get_remote_meta_md5(remote_path, sftp): meta_path = f"{remote_path}.meta" try: async with await sftp.open(meta_path, 'r') as meta_file: meta_content = await meta_file.read() if 'MD5:' in meta_content: remote_md5 = meta_content.split('MD5:')[1].strip() return remote_md5 except asyncssh.SFTPNoSuchFile: logger.warning(f"Meta文件不存在: {os.path.basename(meta_path)}") pass return None async def write_remote_meta_md5(remote_path, remote_md5, sftp): try: meta_content = f"MD5: {remote_md5}" meta_path = f"{remote_path}.meta" async with await sftp.open(meta_path, 'w') as meta_file: await meta_file.write(meta_content) logger.info(f'Meta文件创建成功：{os.path.basename(meta_path)}') except Exception as e: logger.error(f"写入远程.meta文件失败：{remote_path}. 错误信息：{e}") async def calculate_remote_md5(sftp, remote_path): hasher = hashlib.md5() try: async with sftp.open(remote_path, 'rb') as file: while True: data = await file.read(65536) # Read the file in chunks of 64KB if not data: break hasher.update(data) except asyncssh.SFTPNoSuchFile: logger.warning(f"远程文件不存在: {os.path.basename(remote_path)}") return None return hasher.hexdigest() def calculate_md5(file_path): hasher = hashlib.md5() with open(file_path, 'rb') as file: while True: data = file.read(65536) # Read the file in chunks of 64KB if not data: break hasher.update(data) return hasher.hexdigest() # Example usage: hostname = '192.168.188.252' username = 'neozhang' password = '58920912a' local_directory = 'E:\\AI\\stable-diffusion-webui-master\\models\\Stable-diffusion' remote_directory = '/home/neozhang/work/project/stable-diffusion-webui/models/Stable-diffusion' file_extension = '*' depth = -1 # 包括多少级子目录下的文件（-1代表包括所有子目录） calculate_timeout = 1000 # 计算远程MD5的超时时间 coroutine_count = 5 files_to_upload = get_files_to_upload(local_directory, remote_directory, depth, file_extension) logger.info(files_to_upload) async def main(): logger.info('======================================================================') async with asyncssh.connect(hostname, username=username, password=password, known_hosts=None) as conn: sftp = await conn.start_sftp_client() sem = asyncio.Semaphore(coroutine_count) # 设置最大并发协程数量为5 async def limited_compare_md5(local_path, remote_path): async with sem: await compare_md5(local_path, remote_path, sftp) coroutines = [limited_compare_md5(file_info['local_path'], file_info['remote_path']) for file_info in files_to_upload] await asyncio.gather(*coroutines) # Run the main coroutine loop = asyncio.get_event_loop() loop.run_until_complete(main())