139 lines
4.8 KiB
Python
139 lines
4.8 KiB
Python
from lib import config as config_module
|
|
from lib import logging as logging_lib
|
|
from lib import get_specs
|
|
from lib import utils
|
|
import threading
|
|
import socket
|
|
import aiohttp
|
|
import asyncio
|
|
import json
|
|
import time
|
|
import sys
|
|
import os
|
|
|
|
config = config_module.config
|
|
specs = get_specs.Specs()
|
|
log = logging_lib.log
|
|
|
|
def complete_loader(fail=False):
|
|
sys.stdout.write("\r" + " " * 10 + "\r") # Clear loader
|
|
sys.stdout.flush()
|
|
if fail:
|
|
print("\rGetting server specifications - [ \x1b[31mERROR\x1b[0m")
|
|
else:
|
|
print("\rGetting server specifications - [ \x1b[32mCOMPLETE \x1b[0m]")
|
|
|
|
async def show_loader(loader_event):
|
|
# Your loader display logic here
|
|
animation = "|/-\\"
|
|
idx = 0
|
|
while True:
|
|
if loader_event.is_set():
|
|
sys.stdout.write("\r" + "🛠️ Getting server specifications " + animation[idx % len(animation)])
|
|
sys.stdout.flush()
|
|
idx += 1
|
|
await asyncio.sleep(0.1)
|
|
else:
|
|
break
|
|
|
|
async def register_server(data):
|
|
# Define the URL
|
|
url = "https://api.clore.ai/register_server"
|
|
|
|
# Convert the dictionary to JSON
|
|
json_data = json.dumps(data)
|
|
|
|
# Define the headers with content type as JSON
|
|
headers = {
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
connector = aiohttp.TCPConnector(family=socket.AF_INET)
|
|
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
try:
|
|
async with session.post(url, data=json_data, headers=headers, timeout=15) as response:
|
|
if response.status == 200:
|
|
# Successful response
|
|
response_data = await response.json()
|
|
if "result" in response_data and response_data["result"]=="updated":
|
|
with open(config.auth_file, 'w') as f:
|
|
f.write(f"{data['key']}:{data['private_communication_token']}")
|
|
print("\x1b[32mServer successfully logged in\x1b[0m")
|
|
utils.run_command_v2("systemctl restart clore-hosting.service")
|
|
os._exit(0)
|
|
elif "result" in response_data and response_data["result"]=="already_updated":
|
|
print("\x1b[31mThis token is already used by some server\x1b[0m")
|
|
elif "error" in response_data and response_data["error"]=="invalid_key":
|
|
print("\x1b[31mInvalid token\x1b[0m")
|
|
else:
|
|
print(response_data)
|
|
print("\x1b[31mIssues at CLORE.AI, please try again later\x1b[0m")
|
|
os._exit(1)
|
|
else:
|
|
# Failed response
|
|
print("\x1b[31mIssues connecting to CLORE.AI api, check your internet connection or try it later\x1b[0m")
|
|
os._exit(1)
|
|
except asyncio.TimeoutError:
|
|
# Request timed out
|
|
print("\x1b[31mIssues connecting to CLORE.AI api, check your internet connection or try it later\x1b[0m")
|
|
os._exit(1)
|
|
|
|
async def work_init(loader_event, init_token):
|
|
loader_event.set()
|
|
kernel = get_specs.get_kernel()
|
|
nvidia_info = get_specs.get_nvidia_version()
|
|
|
|
if utils.validate_cuda_version(nvidia_info.cuda_version):
|
|
loader_event.clear()
|
|
log.error("Cuda must be version 11.7+")
|
|
os._exit(1)
|
|
|
|
machine_specs = await specs.get(benchmark_internet=True, benchmark_disk=True, require_same_gpus=(not config.allow_mixed_gpus))
|
|
|
|
loader_event.clear()
|
|
complete_loader()
|
|
|
|
max_str_l = 64
|
|
for ckey in machine_specs.keys():
|
|
if ckey == "net":
|
|
pass
|
|
elif ckey == "disk" and len(str(machine_specs[ckey])) > max_str_l:
|
|
p = str(machine_specs[ckey]).split(' ')
|
|
ft = ''
|
|
req_size = p[-1]
|
|
for x in range(len(p) - 1):
|
|
ft += f"{p[x]} "
|
|
ft = ft[:max_str_l - len(req_size) - 1]
|
|
if ft[-3] == ' ':
|
|
machine_specs[ckey] = f"{ft} {req_size}"
|
|
else:
|
|
machine_specs[ckey] = f"{ft} {req_size}"
|
|
elif len(str(machine_specs[ckey])) > max_str_l:
|
|
machine_specs[ckey] = str(machine_specs[ckey])[:max_str_l]
|
|
|
|
|
|
del machine_specs["gpus"]
|
|
del machine_specs["nvml_error"]
|
|
del machine_specs["swap"]
|
|
del machine_specs["data_root_location"]
|
|
|
|
print(machine_specs)
|
|
|
|
await register_server({
|
|
"key": init_token,
|
|
"private_communication_token": utils.generate_random_string(32),
|
|
"specs": machine_specs
|
|
})
|
|
|
|
async def init(init_token):
|
|
loader_event = asyncio.Event()
|
|
|
|
# Run the loader asynchronously
|
|
loader_task = asyncio.create_task(show_loader(loader_event))
|
|
|
|
# Run the initialization asynchronously
|
|
await work_init(loader_event, init_token)
|
|
|
|
# Cancel the loader task once the initialization is done
|
|
loader_task.cancel() |