import logging import os import boto3 import botocore import boto3.s3.transfer as s3transfer import sys from urllib.parse import urlparse import requests import json import gradio as gr sys.path.append(os.getcwd()) import tarfile import shutil from pathlib import Path import psutil LOGGING_LEVEL = os.environ.get('LOG_LEVEL') or logging.ERROR logger = logging.getLogger(__name__) logger.setLevel(LOGGING_LEVEL) s3_client = boto3.client('s3') class ModelsRef: def __init__(self): self.models_ref = {} def get_models_ref_dict(self): return self.models_ref def add_models_ref(self, model_name): if model_name in self.models_ref: self.models_ref[model_name] += 1 else: self.models_ref[model_name] = 0 def remove_model_ref(self,model_name): if self.models_ref.get(model_name): del self.models_ref[model_name] def get_models_ref(self, model_name): return self.models_ref.get(model_name) def get_least_ref_model(self): sorted_models = sorted(self.models_ref.items(), key=lambda item: item[1]) if sorted_models: least_ref_model, least_counter = sorted_models[0] return least_ref_model,least_counter else: return None,None def pop_least_ref_model(self): sorted_models = sorted(self.models_ref.items(), key=lambda item: item[1]) if sorted_models: least_ref_model, least_counter = sorted_models[0] del self.models_ref[least_ref_model] return least_ref_model,least_counter else: return None,None def get_sorted_models(self, key_list=None): print('!!!!!!!!!!!', key_list) if key_list is None: return sorted(self.models_ref.items(), key=lambda item: item[1]) else: models_ref_tmp = {} for key_value in key_list: if key_value not in self.models_ref.keys(): models_ref_tmp[key_value] = -1 else: models_ref_tmp[key_value] = self.models_ref[key_value] models_sorted_info = sorted(models_ref_tmp.items(), key=lambda item: item[1]) models_sorted = [] for model_info in models_sorted_info: models_sorted.append(model_info[0]) return models_sorted # sd_models_Ref = ModelsRef() # cn_models_Ref = ModelsRef() # lora_models_Ref = ModelsRef() # hyper_models_Ref = ModelsRef() # embedding_Ref = ModelsRef() def upload_folder_to_s3(local_folder_path, bucket_name, s3_folder_path): for root, dirs, files in os.walk(local_folder_path): for file in files: local_file_path = os.path.join(root, file) s3_file_path = os.path.join(s3_folder_path, local_file_path) s3_client.upload_file(local_file_path, bucket_name, s3_file_path) def upload_folder_to_s3_by_tar(local_folder_path, bucket_name, s3_folder_path): tar_name = f"{os.path.basename(local_folder_path)}.tar" # os.system(f'tar cvf {tar_name} {local_folder_path}') tar(mode='c', archive=tar_name, sfiles=local_folder_path, verbose=True) # tar = tarfile.open(tar_path, "w:gz") # for root, dirs, files in os.walk(local_folder_path): # for file in files: # local_file_path = os.path.join(root, file) # tar.add(local_file_path) # tar.close() s3_client.upload_file(tar_name, bucket_name, os.path.join(s3_folder_path, tar_name)) # os.system(f"rm {tar_name}") rm(tar_name, recursive=True) def upload_file_to_s3(file_name, bucket, directory=None, object_name=None): # If S3 object_name was not specified, use file_name if object_name is None: object_name = file_name # Add the directory to the object_name if directory: object_name = f"{directory}/{object_name}" # Upload the file try: s3_client.upload_file(file_name, bucket, object_name) print(f"File {file_name} uploaded to {bucket}/{object_name}") except Exception as e: print(f"Error occurred while uploading {file_name} to {bucket}/{object_name}: {e}") return False return True def upload_file_to_s3_by_presign_url(local_path, s3_presign_url): response = requests.put(s3_presign_url, open(local_path, "rb")) response.raise_for_status() def upload_multipart_files_to_s3_by_signed_url(local_path, signed_urls, part_size): integral_uploaded = False with open(local_path, "rb") as f: parts = [] try: for i, signed_url in enumerate(signed_urls): file_data = f.read(part_size) response = requests.put(signed_url, data=file_data) response.raise_for_status() etag = response.headers['ETag'] parts.append({ 'ETag': etag, 'PartNumber': i + 1 }) print(f'model upload part {i+1}: {response}') integral_uploaded = True return parts except Exception as e: print(e) gr.Error(f'Upload file{local_path} failed, please try again. If still not work, contact your admin') finally: if not integral_uploaded: gr.Error(f'Upload file {local_path} not complete, please try again or create new one.') raise Exception('failed at multipart') def download_folder_from_s3(bucket_name, s3_folder_path, local_folder_path): s3_resource = boto3.resource('s3') bucket = s3_resource.Bucket(bucket_name) for obj in bucket.objects.filter(Prefix=s3_folder_path): obj_dirname = os.sep.join(os.path.dirname(obj.key).split("/")[1:]) obj_basename = os.path.basename(obj.key) local_sub_folder_path = os.path.join(local_folder_path, obj_dirname) if not os.path.exists(local_sub_folder_path): os.makedirs(local_sub_folder_path) bucket.download_file(obj.key, os.path.join(local_sub_folder_path, obj_basename)) # save to same path def download_folder_from_s3_by_tar(bucket_name, s3_tar_path, local_tar_path, target_dir="."): s3_client.download_file(bucket_name, s3_tar_path, local_tar_path) # tar_name = os.path.basename(s3_tar_path) # os.system(f"tar xvf {local_tar_path} -C {target_dir}") tar(mode='x', archive=local_tar_path, verbose=True, change_dir=target_dir) # tar = tarfile.open(local_tar_path, "r") # tar.extractall() # tar.close() # os.system(f"rm {local_tar_path}") rm(local_tar_path, recursive=True) def download_file_from_s3(bucket_name, s3_file_path, local_file_path): s3_client.download_file(bucket_name, s3_file_path, local_file_path) def get_bucket_name_from_s3_url(s3_path) -> str: o = urlparse(s3_path, allow_fragments=False) return o.netloc def get_bucket_name_from_s3_path(s3_path) -> str: s3_path = s3_path.replace("s3://", "") return s3_path.split("/")[0] def get_path_from_s3_path(s3_path) -> str: s3_path = s3_path.replace("s3://", "") return "/".join(s3_path.split("/")[1:]) def split_s3_path(s3_path): path_parts = s3_path.replace("s3://", "").split("/") bucket = path_parts.pop(0) key = "/".join(path_parts) return bucket, key def read_from_s3(s3_path): bucket, key = split_s3_path(s3_path) s3_resp = s3_client.get_object( Bucket=bucket, Key=key, ) if s3_resp['ContentLength'] > 0: return s3_resp['Body'].read() raise Exception(f'no content for file {s3_path}') def fast_upload(session, bucketname, s3dir, filelist, progress_func=None, workers=10): # timer = Timer() botocore_config = botocore.config.Config(max_pool_connections=workers) s3client = session.client('s3', config=botocore_config) transfer_config = s3transfer.TransferConfig( use_threads=True, max_concurrency=workers, ) s3t = s3transfer.create_transfer_manager(s3client, transfer_config) # timer.record("init") for src in filelist: dst = os.path.join(s3dir, os.path.basename(src)) s3t.upload( src, bucketname, dst, subscribers=[ s3transfer.ProgressCallbackInvoker(progress_func), ] if progress_func else None, ) s3t.shutdown() # wait for all the upload tasks to finish # timer.record("upload") # print(timer.summary()) def save_variable_to_json(variable_name, variable_value, filename='sagemaker_ui.json'): data = {} if os.path.exists(filename): with open(filename, 'r') as json_file: data = json.load(json_file) data[variable_name] = variable_value with open(filename, 'w') as json_file: json.dump(data, json_file) def get_variable_from_json(variable_name, filename='sagemaker_ui.json'): if not os.path.exists(filename): initial_data = { "api_gateway_url": "", "api_token": "", "username": "" } with open(filename, 'w') as json_file: json.dump(initial_data, json_file) with open(filename, 'r') as json_file: data = json.load(json_file) variable_value = data.get(variable_name) return variable_value def host_url(): return get_variable_from_json('api_gateway_url') def api_key(): return get_variable_from_json('api_token') def is_gcr(): api_url = get_variable_from_json('api_gateway_url') return api_url and '.execute-api.cn-' in api_url def has_config(): return host_url() and api_key() """ Description: Below functions are used to replace existing shell command implementation with os.system method, which is not os agonostic and not recommended. """ def tar(mode, archive, sfiles=None, verbose=False, change_dir=None): """ Description: Create or extract a tar archive. Args: mode: 'c' for create or 'x' for extract archive: the archive file name files: a list of files to add to the archive (when creating) or extract (when extracting); None to extract all files verbose: whether to print the names of the files as they are being processed change_dir: the directory to change to before performing any other operations; None to use the current directory Usage: # Create a new archive tar(mode='c', archive='archive.tar', sfiles=['file1.txt', 'file2.txt']) # Extract files from an archive tar(mode='x', archive='archive.tar') # Create a new archive with verbose mode and input directory tar(mode='c', archive='archive.tar', sfiles='./some_directory', verbose=True) # Extract files from an archive with verbose mode and change directory tar(mode='x', archive='archive.tar', verbose=True, change_dir='./some_directory') """ if mode == 'c': # os.chdir(change_dir) with tarfile.open(archive, mode='w') as tar: # check if input option file is a list or string if isinstance(sfiles, list): for file in sfiles: if verbose: print(f"Adding {file} to {archive}") tar.add(file) # take it as a folder name string else: for folder_path, subfolders, files in os.walk(sfiles): for file in files: if verbose: print(f"Adding {os.path.join(folder_path, file)} to {archive}") tar.add(os.path.join(folder_path, file)) elif mode == 'x': with tarfile.open(archive, mode='r') as tar: # sfiles is set to all files in the archive if not specified if not sfiles: sfiles = tar.getnames() for file in sfiles: if verbose: print(f"Extracting {file} from {archive}") # extra to specified directory if change_dir: tar.extract(file, path=change_dir) else: tar.extract(file) def rm(path, force=False, recursive=False): """ Description: Remove a file or directory. Args: path (str): Path of the file or directory to remove. force (bool): If True, ignore non-existent files and errors. Default is False. recursive (bool): If True, remove directories and their contents recursively. Default is False. Usage: # Remove the file rm(dst) # Remove a directory recursively rm("directory/path", recursive=True) """ path_obj = Path(path) try: if path_obj.is_file() or (path_obj.is_symlink() and not path_obj.is_dir()): path_obj.unlink() elif path_obj.is_dir() and recursive: shutil.rmtree(path) elif path_obj.is_dir(): raise ValueError("Cannot remove directory without recursive=True") else: raise ValueError("File or directory does not exist") except Exception as e: if not force: raise e def cp(src, dst, recursive=False, dereference=False, preserve=True): """ Description: Copy a file or directory from source path to destination path. Args: src (str): Source file or directory path. dst (str): Destination file or directory path. recursive (bool): If True, copy directory and its contents recursively. Default is False. dereference (bool): If True, always dereference symbolic links. Default is False. preserve (bool): If True, preserve file metadata. Default is True. Usage: src = "source/file/path.txt" dst = "destination/file/path.txt" # Copy the file cp(src, dst) # Copy a directory recursively and dereference symlinks cp("source/directory", "destination/directory", recursive=True, dereference=True) """ src_path = Path(src) dst_path = Path(dst) try: if dereference: src_path = src_path.resolve() if src_path.is_dir() and recursive: if preserve: shutil.copytree(src_path, dst_path, copy_function=shutil.copy2, symlinks=not dereference) else: shutil.copytree(src_path, dst_path, symlinks=not dereference) elif src_path.is_file(): if preserve: shutil.copy2(src_path, dst_path) else: shutil.copy(src_path, dst_path) else: raise ValueError("Source must be a file or a directory with recursive=True") except shutil.SameFileError: print("Source and destination represents the same file.") def mv(src, dest, force=False): """ Description: Move or rename files and directories. Args: src (str): Source file or directory path. dest (str): Destination file or directory path. force (bool): If True, overwrite the destination if it exists. Default is False. Usage: # Rename a file mv('old_name.txt', 'new_name.txt') # Move a file to a new directory mv('file.txt', 'new_directory/file.txt') # Move a directory to another directory mv('source_directory', 'destination_directory') # Force move (overwrite) a file or directory mv('source_file.txt', 'existing_destination_file.txt', force=True) """ src_path = Path(src) dest_path = Path(dest) if src_path.exists(): if dest_path.exists() and not force: raise FileExistsError(f"Destination path '{dest}' already exists and 'force' is not set") else: if dest_path.is_file(): dest_path.unlink() elif dest_path.is_dir(): shutil.rmtree(dest_path) if src_path.is_file() or src_path.is_dir(): shutil.move(src, dest) else: raise FileNotFoundError(f"Source path '{src}' does not exist") def format_size(size, human_readable): if human_readable: for unit in ['B', 'K', 'M', 'G', 'T', 'P']: if size < 1024: return f"{size:.1f}{unit}" size /= 1024 else: return str(size) def df(show_all=False, human_readable=False): """ Description: Get disk usage statistics. Args: show_all (bool): If True, include all filesystems. Default is False. human_readable (bool): If True, format sizes in human readable format. Default is False. Usage: filesystems = df(show_all=True, human_readable=True) for filesystem in filesystems: print(f"Filesystem: {filesystem['filesystem']}") print(f"Total: {filesystem['total']}") print(f"Used: {filesystem['used']}") print(f"Free: {filesystem['free']}") print(f"Percent: {filesystem['percent']}%") print(f"Mountpoint: {filesystem['mountpoint']}") """ partitions = psutil.disk_partitions(all=show_all) result = [] for partition in partitions: usage = psutil.disk_usage(partition.mountpoint) partition_info = { 'filesystem': partition.device, 'total': format_size(usage.total, human_readable), 'used': format_size(usage.used, human_readable), 'free': format_size(usage.free, human_readable), 'percent': usage.percent, 'mountpoint': partition.mountpoint, } result.append(partition_info) return result if __name__ == '__main__': import sys # upload_file_to_s3(sys.argv[1], 'aws-gcr-csdc-atl-exp-us-west-2', sys.argv[2]) # fast_upload(boto3.Session(), 'aws-gcr-csdc-atl-exp-us-west-2', sys.argv[2], [sys.argv[1]]) download_folder_from_s3_by_tar('aws-gcr-csdc-atl-exp-us-west-2', 'aigc-webui-test-samples/samples.tar', 'samples.tar')