mirror of
https://github.com/azaion/annotations.git
synced 2026-04-22 10:16:30 +00:00
add azaion loader
This commit is contained in:
Alex Bezdieniezhnykh
@@ -1,20 +0,0 @@
|
||||
from user cimport User
|
||||
from credentials cimport Credentials
|
||||
from cdn_manager cimport CDNManager
|
||||
|
||||
|
||||
cdef class ApiClient:
|
||||
cdef Credentials credentials
|
||||
cdef CDNManager cdn_manager
|
||||
cdef str token, folder, api_url
|
||||
cdef User user
|
||||
|
||||
cdef set_credentials(self, Credentials credentials)
|
||||
cdef login(self)
|
||||
cdef set_token(self, str token)
|
||||
cdef get_user(self)
|
||||
|
||||
cdef load_bytes(self, str filename, str folder)
|
||||
cdef upload_file(self, str filename, bytes resource, str folder)
|
||||
cdef load_big_small_resource(self, str resource_name, str folder, str key)
|
||||
cdef upload_big_small_resource(self, bytes resource, str resource_name, str folder, str key)
|
||||
@@ -1,177 +0,0 @@
|
||||
import json
|
||||
import os
|
||||
from http import HTTPStatus
|
||||
from os import path
|
||||
from uuid import UUID
|
||||
import jwt
|
||||
import requests
|
||||
cimport constants
|
||||
import yaml
|
||||
from requests import HTTPError
|
||||
|
||||
from cdn_manager cimport CDNManager, CDNCredentials
|
||||
from hardware_service cimport HardwareService
|
||||
from security cimport Security
|
||||
from user cimport User, RoleEnum
|
||||
|
||||
cdef class ApiClient:
|
||||
"""Handles API authentication and downloading of the AI model."""
|
||||
def __init__(self, str api_url):
|
||||
self.credentials = None
|
||||
self.user = None
|
||||
self.token = None
|
||||
self.cdn_manager = None
|
||||
self.api_url = api_url
|
||||
|
||||
cdef set_credentials(self, Credentials credentials):
|
||||
self.credentials = credentials
|
||||
if self.cdn_manager is not None:
|
||||
return
|
||||
|
||||
yaml_bytes = self.load_bytes(constants.CDN_CONFIG, <str>'')
|
||||
yaml_config = yaml.safe_load(yaml_bytes)
|
||||
creds = CDNCredentials(yaml_config["host"],
|
||||
yaml_config["downloader_access_key"],
|
||||
yaml_config["downloader_access_secret"],
|
||||
yaml_config["uploader_access_key"],
|
||||
yaml_config["uploader_access_secret"])
|
||||
|
||||
self.cdn_manager = CDNManager(creds)
|
||||
|
||||
cdef login(self):
|
||||
response = None
|
||||
try:
|
||||
response = requests.post(f"{self.api_url}/login",
|
||||
json={"email": self.credentials.email, "password": self.credentials.password})
|
||||
response.raise_for_status()
|
||||
token = response.json()["token"]
|
||||
self.set_token(token)
|
||||
except HTTPError as e:
|
||||
print(response.json())
|
||||
if response.status_code == HTTPStatus.CONFLICT:
|
||||
res = response.json()
|
||||
raise Exception(res['Message'])
|
||||
|
||||
|
||||
cdef set_token(self, str token):
|
||||
self.token = token
|
||||
claims = jwt.decode(token, options={"verify_signature": False})
|
||||
|
||||
try:
|
||||
id = str(UUID(claims.get("nameid", "")))
|
||||
except ValueError:
|
||||
raise ValueError("Invalid GUID format in claims")
|
||||
|
||||
email = claims.get("unique_name", "")
|
||||
|
||||
role_str = claims.get("role", "")
|
||||
if role_str == "ApiAdmin":
|
||||
role = RoleEnum.ApiAdmin
|
||||
elif role_str == "Admin":
|
||||
role = RoleEnum.Admin
|
||||
elif role_str == "ResourceUploader":
|
||||
role = RoleEnum.ResourceUploader
|
||||
elif role_str == "Validator":
|
||||
role = RoleEnum.Validator
|
||||
elif role_str == "Operator":
|
||||
role = RoleEnum.Operator
|
||||
else:
|
||||
role = RoleEnum.NONE
|
||||
self.user = User(id, email, role)
|
||||
|
||||
cdef get_user(self):
|
||||
if self.user is None:
|
||||
self.login()
|
||||
return self.user
|
||||
|
||||
cdef upload_file(self, str filename, bytes resource, str folder):
|
||||
if self.token is None:
|
||||
self.login()
|
||||
url = f"{self.api_url}/resources/{folder}"
|
||||
headers = { "Authorization": f"Bearer {self.token}" }
|
||||
files = {'data': (filename, resource)}
|
||||
try:
|
||||
r = requests.post(url, headers=headers, files=files, allow_redirects=True)
|
||||
r.raise_for_status()
|
||||
constants.log(f"Uploaded {filename} to {self.api_url}/{folder} successfully: {r.status_code}.")
|
||||
except Exception as e:
|
||||
constants.log(f"Upload fail: {e}")
|
||||
|
||||
cdef load_bytes(self, str filename, str folder):
|
||||
hardware_service = HardwareService()
|
||||
cdef str hardware = hardware_service.get_hardware_info()
|
||||
if self.token is None:
|
||||
self.login()
|
||||
url = f"{self.api_url}/resources/get/{folder}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.token}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = json.dumps(
|
||||
{
|
||||
"password": self.credentials.password,
|
||||
"hardware": hardware,
|
||||
"fileName": filename
|
||||
}, indent=4)
|
||||
response = requests.post(url, data=payload, headers=headers, stream=True)
|
||||
if response.status_code == HTTPStatus.UNAUTHORIZED or response.status_code == HTTPStatus.FORBIDDEN:
|
||||
self.login()
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.token}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
response = requests.post(url, data=payload, headers=headers, stream=True)
|
||||
|
||||
if response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR:
|
||||
print('500!')
|
||||
|
||||
hw_hash = Security.get_hw_hash(hardware)
|
||||
key = Security.get_api_encryption_key(self.credentials, hw_hash)
|
||||
|
||||
resp_bytes = response.raw.read()
|
||||
data = Security.decrypt_to(resp_bytes, key)
|
||||
constants.log(<str>f'Downloaded file: {filename}, {len(data)} bytes')
|
||||
return data
|
||||
|
||||
cdef load_big_file_cdn(self, str folder, str big_part):
|
||||
print(f'downloading file {folder}\\{big_part} from cdn...')
|
||||
if self.cdn_manager.download(folder, big_part):
|
||||
with open(path.join(<str> folder, big_part), 'rb') as binary_file:
|
||||
encrypted_bytes_big = binary_file.read()
|
||||
else:
|
||||
return None
|
||||
|
||||
cdef load_big_small_resource(self, str resource_name, str folder, str key):
|
||||
cdef str big_part = f'{resource_name}.big'
|
||||
cdef str small_part = f'{resource_name}.small'
|
||||
|
||||
encrypted_bytes_small = self.load_bytes(small_part, folder)
|
||||
|
||||
print(f'checking on existence for {folder}\\{big_part}')
|
||||
if os.path.exists(os.path.join(<str> folder, big_part)):
|
||||
with open(path.join(<str> folder, big_part), 'rb') as binary_file:
|
||||
local_bytes_big = binary_file.read()
|
||||
print(f'local file {folder}\\{big_part} is found!')
|
||||
try:
|
||||
return Security.decrypt_to(encrypted_bytes_small + local_bytes_big, key)
|
||||
except Exception as ex:
|
||||
print('Local file doesnt match with api file, old version')
|
||||
|
||||
remote_bytes_big = self.load_big_file_cdn(folder, big_part)
|
||||
return Security.decrypt_to(encrypted_bytes_small + remote_bytes_big)
|
||||
|
||||
cdef upload_big_small_resource(self, bytes resource, str resource_name, str folder, str key):
|
||||
cdef str big_part_name = f'{resource_name}.big'
|
||||
cdef str small_part_name = f'{resource_name}.small'
|
||||
|
||||
resource_encrypted = Security.encrypt_to(<bytes>resource, key)
|
||||
part_small_size = min(constants.SMALL_SIZE_KB * 1024, int(0.3 * len(resource_encrypted)))
|
||||
part_small = resource_encrypted[:part_small_size] # slice bytes for part1
|
||||
|
||||
part_big = resource_encrypted[part_small_size:]
|
||||
|
||||
self.cdn_manager.upload(folder, <str>big_part_name, part_big)
|
||||
with open(path.join(<str>folder, <str>big_part_name), 'wb') as f:
|
||||
f.write(part_big)
|
||||
self.upload_file(small_part_name, part_small, folder)
|
||||
@@ -26,8 +26,6 @@ tmp_ret = collect_all('pycuda')
|
||||
datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2]
|
||||
tmp_ret = collect_all('pynvml')
|
||||
datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2]
|
||||
tmp_ret = collect_all('boto3')
|
||||
datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2]
|
||||
tmp_ret = collect_all('jwt')
|
||||
datas += tmp_ret[0]; binaries += tmp_ret[1]; hiddenimports += tmp_ret[2]
|
||||
|
||||
|
||||
@@ -33,7 +33,6 @@ venv\Scripts\pyinstaller --name=azaion-inference ^
|
||||
--collect-all tensorrt ^
|
||||
--collect-all pycuda ^
|
||||
--collect-all pynvml ^
|
||||
--collect-all boto3 ^
|
||||
--collect-all jwt ^
|
||||
--hidden-import constants ^
|
||||
--hidden-import annotation ^
|
||||
|
||||
@@ -1,14 +0,0 @@
|
||||
cdef class CDNCredentials:
|
||||
cdef str host
|
||||
cdef str downloader_access_key
|
||||
cdef str downloader_access_secret
|
||||
cdef str uploader_access_key
|
||||
cdef str uploader_access_secret
|
||||
|
||||
cdef class CDNManager:
|
||||
cdef CDNCredentials creds
|
||||
cdef object download_client
|
||||
cdef object upload_client
|
||||
|
||||
cdef upload(self, str bucket, str filename, bytes file_bytes)
|
||||
cdef download(self, str bucket, str filename)
|
||||
@@ -1,41 +0,0 @@
|
||||
import io
|
||||
import boto3
|
||||
|
||||
|
||||
cdef 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
|
||||
|
||||
|
||||
cdef class CDNManager:
|
||||
def __init__(self, CDNCredentials credentials):
|
||||
|
||||
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)
|
||||
|
||||
cdef upload(self, str bucket, str filename, bytes file_bytes):
|
||||
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
|
||||
|
||||
cdef download(self, str folder, str filename):
|
||||
try:
|
||||
self.download_client.download_file(folder, filename, f'{folder}\\{filename}')
|
||||
print(f'downloaded {filename} from the {folder} to current folder')
|
||||
return True
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return False
|
||||
@@ -2,12 +2,7 @@ import time
|
||||
|
||||
cdef str CONFIG_FILE = "config.yaml" # Port for the zmq
|
||||
|
||||
cdef int QUEUE_MAXSIZE = 1000 # Maximum size of the command queue
|
||||
cdef str COMMANDS_QUEUE = "azaion-commands"
|
||||
cdef str ANNOTATIONS_QUEUE = "azaion-annotations"
|
||||
|
||||
cdef str QUEUE_CONFIG_FILENAME = "secured-config.json"
|
||||
|
||||
cdef str AI_ONNX_MODEL_FILE = "azaion.onnx"
|
||||
|
||||
cdef str CDN_CONFIG = "cdn.yaml"
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
cdef class Credentials:
|
||||
cdef public str email
|
||||
cdef public str password
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data)
|
||||
@@ -1,15 +0,0 @@
|
||||
from msgpack import unpackb
|
||||
|
||||
cdef class Credentials:
|
||||
|
||||
def __init__(self, str email, str password):
|
||||
self.email = email
|
||||
self.password = password
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data):
|
||||
unpacked = unpackb(data, strict_map_key=False)
|
||||
return Credentials(
|
||||
unpacked.get("Email"),
|
||||
unpacked.get("Password"))
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
cdef class FileData:
|
||||
cdef public str folder
|
||||
cdef public str filename
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data)
|
||||
@@ -1,14 +0,0 @@
|
||||
from msgpack import unpackb
|
||||
|
||||
cdef class FileData:
|
||||
|
||||
def __init__(self, str folder, str filename):
|
||||
self.folder = folder
|
||||
self.filename = filename
|
||||
|
||||
@staticmethod
|
||||
cdef from_msgpack(bytes data):
|
||||
unpacked = unpackb(data, strict_map_key=False)
|
||||
return FileData(
|
||||
unpacked.get("Folder"),
|
||||
unpacked.get("Filename"))
|
||||
@@ -1,5 +0,0 @@
|
||||
cdef class HardwareService:
|
||||
|
||||
@staticmethod
|
||||
cdef has_nvidia_gpu()
|
||||
cdef str get_hardware_info(self)
|
||||
@@ -1,61 +0,0 @@
|
||||
import os
|
||||
import subprocess
|
||||
import pynvml
|
||||
|
||||
|
||||
cdef class HardwareService:
|
||||
|
||||
@staticmethod
|
||||
cdef has_nvidia_gpu():
|
||||
try:
|
||||
pynvml.nvmlInit()
|
||||
device_count = pynvml.nvmlDeviceGetCount()
|
||||
|
||||
if device_count > 0:
|
||||
print(f"Found NVIDIA GPU(s).")
|
||||
return True
|
||||
else:
|
||||
print("No NVIDIA GPUs found by NVML.")
|
||||
return False
|
||||
|
||||
except pynvml.NVMLError as error:
|
||||
print(f"Failed to find NVIDIA GPU")
|
||||
return False
|
||||
finally:
|
||||
try:
|
||||
pynvml.nvmlShutdown()
|
||||
except:
|
||||
print('Failed to shutdown pynvml cause probably no NVidia GPU')
|
||||
pass
|
||||
|
||||
cdef str get_hardware_info(self):
|
||||
if os.name == 'nt': # windows
|
||||
os_command = (
|
||||
"powershell -Command \""
|
||||
"Get-CimInstance -ClassName Win32_Processor | Select-Object -ExpandProperty Name | Write-Output; "
|
||||
"Get-CimInstance -ClassName Win32_VideoController | Select-Object -ExpandProperty Name | Write-Output; "
|
||||
"Get-CimInstance -ClassName Win32_OperatingSystem | Select-Object -ExpandProperty TotalVisibleMemorySize | Write-Output; "
|
||||
"(Get-Disk | Where-Object {$_.IsSystem -eq $true}).SerialNumber"
|
||||
"\""
|
||||
)
|
||||
else:
|
||||
os_command = (
|
||||
"lscpu | grep 'Model name:' | cut -d':' -f2 && "
|
||||
"lspci | grep VGA | cut -d':' -f3 && "
|
||||
"free -k | awk '/^Mem:/ {print $2}' && "
|
||||
"cat /sys/block/sda/device/vpd_pg80 2>/dev/null || cat /sys/block/sda/device/serial 2>/dev/null"
|
||||
)
|
||||
|
||||
result = subprocess.check_output(os_command, shell=True).decode('utf-8', errors='ignore')
|
||||
lines = [line.replace(" ", " ").replace("Name=", "").strip('\x00\x14 \t\n\r\v\f') for line in result.splitlines() if line.strip()]
|
||||
|
||||
cdef str cpu = lines[0]
|
||||
cdef str gpu = lines[1]
|
||||
# could be multiple gpus
|
||||
|
||||
len_lines = len(lines)
|
||||
cdef str memory = lines[len_lines-2].replace("TotalVisibleMemorySize=", "").replace(" ", " ")
|
||||
cdef str drive_serial = lines[len_lines-1]
|
||||
|
||||
cdef str res = f'CPU: {cpu}. GPU: {gpu}. Memory: {memory}. DriveSerial: {drive_serial}'
|
||||
return res
|
||||
@@ -6,27 +6,24 @@ from threading import Thread
|
||||
|
||||
import yaml
|
||||
|
||||
from api_client cimport ApiClient
|
||||
from annotation cimport Annotation
|
||||
from inference cimport Inference
|
||||
from remote_command cimport RemoteCommand, CommandType
|
||||
from remote_command_handler cimport RemoteCommandHandler
|
||||
from credentials cimport Credentials
|
||||
from file_data cimport FileData
|
||||
|
||||
cdef class CommandProcessor:
|
||||
cdef ApiClient api_client
|
||||
|
||||
cdef RemoteCommandHandler remote_handler
|
||||
cdef object inference_queue
|
||||
cdef bint running
|
||||
cdef Inference inference
|
||||
|
||||
def __init__(self, int zmq_port, str api_url):
|
||||
self.api_client = ApiClient(api_url)
|
||||
self.remote_handler = RemoteCommandHandler(zmq_port, self.on_command)
|
||||
self.inference_queue = Queue(maxsize=constants.QUEUE_MAXSIZE)
|
||||
self.remote_handler.start()
|
||||
self.running = True
|
||||
#TODO: replace api_client to azaion_loader.exe call
|
||||
self.inference = Inference(self.api_client, self.on_annotation)
|
||||
|
||||
def start(self):
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
from credentials cimport Credentials
|
||||
|
||||
cdef class Security:
|
||||
@staticmethod
|
||||
cdef encrypt_to(input_stream, key)
|
||||
|
||||
@staticmethod
|
||||
cdef decrypt_to(input_bytes, key)
|
||||
|
||||
@staticmethod
|
||||
cdef get_hw_hash(str hardware)
|
||||
|
||||
@staticmethod
|
||||
cdef get_api_encryption_key(Credentials credentials, str hardware_hash)
|
||||
|
||||
@staticmethod
|
||||
cdef get_model_encryption_key()
|
||||
|
||||
@staticmethod
|
||||
cdef calc_hash(str key)
|
||||
@@ -1,68 +0,0 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import os
|
||||
from hashlib import sha384
|
||||
from credentials cimport Credentials
|
||||
from cryptography.hazmat.backends import default_backend
|
||||
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
|
||||
from cryptography.hazmat.primitives import padding
|
||||
|
||||
BUFFER_SIZE = 64 * 1024 # 64 KB
|
||||
|
||||
cdef class Security:
|
||||
@staticmethod
|
||||
cdef encrypt_to(input_bytes, key):
|
||||
cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
|
||||
iv = os.urandom(16)
|
||||
|
||||
cipher = Cipher(algorithms.AES(<bytes> aes_key), modes.CBC(iv), backend=default_backend())
|
||||
encryptor = cipher.encryptor()
|
||||
padder = padding.PKCS7(128).padder()
|
||||
|
||||
padded_plaintext = padder.update(input_bytes) + padder.finalize()
|
||||
ciphertext = encryptor.update(padded_plaintext) + encryptor.finalize()
|
||||
|
||||
return iv + ciphertext
|
||||
|
||||
@staticmethod
|
||||
cdef decrypt_to(ciphertext_with_iv_bytes, key):
|
||||
cdef bytes aes_key = hashlib.sha256(key.encode('utf-8')).digest()
|
||||
iv = ciphertext_with_iv_bytes[:16]
|
||||
ciphertext_bytes = ciphertext_with_iv_bytes[16:]
|
||||
|
||||
cipher = Cipher(algorithms.AES(<bytes>aes_key), modes.CBC(<bytes>iv), backend=default_backend())
|
||||
decryptor = cipher.decryptor()
|
||||
|
||||
decrypted_padded_bytes = decryptor.update(ciphertext_bytes) + decryptor.finalize()
|
||||
|
||||
# Manual PKCS7 unpadding check and removal
|
||||
padding_value = decrypted_padded_bytes[-1] # Get the last byte, which indicates padding length
|
||||
if 1 <= padding_value <= 16: # Valid PKCS7 padding value range for AES-128
|
||||
padding_length = padding_value
|
||||
plaintext_bytes = decrypted_padded_bytes[:-padding_length] # Remove padding bytes
|
||||
else:
|
||||
plaintext_bytes = decrypted_padded_bytes
|
||||
|
||||
return bytes(plaintext_bytes)
|
||||
|
||||
@staticmethod
|
||||
cdef get_hw_hash(str hardware):
|
||||
cdef str key = f'Azaion_{hardware}_%$$$)0_'
|
||||
return Security.calc_hash(key)
|
||||
|
||||
@staticmethod
|
||||
cdef get_api_encryption_key(Credentials creds, str hardware_hash):
|
||||
cdef str key = f'{creds.email}-{creds.password}-{hardware_hash}-#%@AzaionKey@%#---'
|
||||
return Security.calc_hash(key)
|
||||
|
||||
@staticmethod
|
||||
cdef get_model_encryption_key():
|
||||
cdef str key = '-#%@AzaionKey@%#---234sdfklgvhjbnn'
|
||||
return Security.calc_hash(key)
|
||||
|
||||
@staticmethod
|
||||
cdef calc_hash(str key):
|
||||
str_bytes = key.encode('utf-8')
|
||||
hash_bytes = sha384(str_bytes).digest()
|
||||
cdef str h = base64.b64encode(hash_bytes).decode('utf-8')
|
||||
return h
|
||||
@@ -1,15 +0,0 @@
|
||||
cdef enum RoleEnum:
|
||||
NONE = 0
|
||||
Operator = 10
|
||||
Validator = 20
|
||||
CompanionPC = 30
|
||||
Admin = 40
|
||||
ResourceUploader = 50
|
||||
ApiAdmin = 1000
|
||||
|
||||
cdef class User:
|
||||
cdef public str id
|
||||
cdef public str email
|
||||
cdef public RoleEnum role
|
||||
|
||||
cdef bytes serialize(self)
|
||||
@@ -1,15 +0,0 @@
|
||||
import msgpack
|
||||
|
||||
cdef class User:
|
||||
|
||||
def __init__(self, str id, str email, RoleEnum role):
|
||||
self.id = id
|
||||
self.email = email
|
||||
self.role = role
|
||||
|
||||
cdef bytes serialize(self):
|
||||
return msgpack.packb({
|
||||
"i": self.id,
|
||||
"e": self.email,
|
||||
"r": self.role
|
||||
})
|
||||
Reference in New Issue
Block a user