import os import requests 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()) # from modules.timer import Timer import tarfile 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): s3_client = boto3.client('s3') 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 = 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 = boto3.client('s3') s3_client.upload_file(tar_name, bucket_name, os.path.join(s3_folder_path, tar_name)) os.system(f"rm {tar_name}") 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 = boto3.client('s3') 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 = "/".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 = boto3.client('s3') 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 = tarfile.open(local_tar_path, "r") # tar.extractall() # tar.close() os.system(f"rm {local_tar_path}") def download_file_from_s3(bucket_name, s3_file_path, local_file_path): s3_client = boto3.client('s3') 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 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": "" } 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 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]]) upload_folder_to_s3_by_tar('models/dreambooth/sagemaker_test/samples', 'aws-gcr-csdc-atl-exp-us-west-2', 'aigc-webui-test-samples') download_folder_from_s3_by_tar('aws-gcr-csdc-atl-exp-us-west-2', 'aigc-webui-test-samples/samples.tar', 'samples.tar')