mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 21:46:30 +00:00
Alex Bezdieniezhnykh
e798af470b
automate tensorrt model conversion in case of no existing one for user's gpu
87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
import io
|
|
import sys
|
|
|
|
import boto3
|
|
import yaml
|
|
import os
|
|
|
|
|
|
class CDNCredentials:
|
|
def __init__(self, host, downloader_access_key, downloader_access_secret, uploader_access_key, uploader_access_secret):
|
|
self.host = host
|
|
self.downloader_access_key = downloader_access_key
|
|
self.downloader_access_secret = downloader_access_secret
|
|
self.uploader_access_key = uploader_access_key
|
|
self.uploader_access_secret = uploader_access_secret
|
|
|
|
|
|
class CDNManager:
|
|
def __init__(self, credentials: CDNCredentials):
|
|
self.creds = credentials
|
|
self.download_client = boto3.client('s3', endpoint_url=self.creds.host,
|
|
aws_access_key_id=self.creds.downloader_access_key,
|
|
aws_secret_access_key=self.creds.downloader_access_secret)
|
|
self.upload_client = boto3.client('s3', endpoint_url=self.creds.host,
|
|
aws_access_key_id=self.creds.uploader_access_key,
|
|
aws_secret_access_key=self.creds.uploader_access_secret)
|
|
|
|
def upload(self, bucket: str, filename: str, file_bytes: bytearray):
|
|
try:
|
|
self.upload_client.upload_fileobj(io.BytesIO(file_bytes), bucket, filename)
|
|
print(f'uploaded {filename} ({len(file_bytes)} bytes) to the {bucket}')
|
|
return True
|
|
except Exception as e:
|
|
print(e)
|
|
return False
|
|
|
|
def download(self, bucket: str, filename: str):
|
|
try:
|
|
if filename is not None:
|
|
self.download_client.download_file(bucket, filename, os.path.join(bucket, filename))
|
|
print(f'downloaded {filename} from the {bucket} to current folder')
|
|
else:
|
|
response = self.download_client.list_objects_v2(Bucket=bucket)
|
|
if 'Contents' in response:
|
|
for obj in response['Contents']:
|
|
object_key = obj['Key']
|
|
local_filepath = os.path.join(bucket, object_key)
|
|
local_dir = os.path.dirname(local_filepath)
|
|
if local_dir:
|
|
os.makedirs(local_dir, exist_ok=True)
|
|
|
|
if not object_key.endswith('/'):
|
|
try:
|
|
self.download_client.download_file(bucket, object_key, local_filepath)
|
|
except Exception as e_file:
|
|
all_successful = False # Mark as failed if any file fails
|
|
return True
|
|
except Exception as e:
|
|
print(e)
|
|
return False
|
|
|
|
|
|
config_file = os.path.join('downloader_config.yaml')
|
|
with open(config_file, 'r') as f:
|
|
config = yaml.safe_load(f)
|
|
|
|
cdn_manager = CDNManager(CDNCredentials(
|
|
config['host'],
|
|
config['downloader_access_key'],
|
|
config['downloader_access_secret'],
|
|
config['uploader_access_key'],
|
|
config['uploader_access_secret']
|
|
))
|
|
|
|
input_action = sys.argv[1]
|
|
input_bucket = sys.argv[2]
|
|
input_filename = sys.argv[3] if len(sys.argv) > 3 else None
|
|
if len(sys.argv) > 4: # 0 is this script's path, hence 5 args is max
|
|
input_path = sys.argv[4]
|
|
|
|
if input_action == 'upload':
|
|
with open(input_path, 'rb') as f:
|
|
cdn_manager.upload(input_bucket, input_filename, bytearray(f.read()))
|
|
elif input_action == 'download':
|
|
cdn_manager.download(input_bucket, input_filename)
|
|
|