mirror of https://github.com/skygpu/skynet.git
130 lines
3.4 KiB
Python
130 lines
3.4 KiB
Python
#!/usr/bin/python
|
|
|
|
import os
|
|
import json
|
|
import time
|
|
import random
|
|
import string
|
|
import logging
|
|
|
|
from functools import partial
|
|
from pathlib import Path
|
|
|
|
import trio
|
|
import pytest
|
|
import psycopg2
|
|
import trio_asyncio
|
|
|
|
from docker.types import Mount, DeviceRequest
|
|
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
|
|
|
|
from skynet.constants import *
|
|
from skynet.brain import run_skynet
|
|
|
|
|
|
@pytest.fixture(scope='session')
|
|
def postgres_db(dockerctl):
|
|
rpassword = ''.join(
|
|
random.choice(string.ascii_lowercase)
|
|
for i in range(12))
|
|
password = ''.join(
|
|
random.choice(string.ascii_lowercase)
|
|
for i in range(12))
|
|
|
|
with dockerctl.run(
|
|
'postgres',
|
|
name='skynet-test-postgres',
|
|
ports={'5432/tcp': None},
|
|
environment={
|
|
'POSTGRES_PASSWORD': rpassword
|
|
}
|
|
) as containers:
|
|
container = containers[0]
|
|
# ip = container.attrs['NetworkSettings']['IPAddress']
|
|
port = container.ports['5432/tcp'][0]['HostPort']
|
|
host = f'localhost:{port}'
|
|
|
|
for log in container.logs(stream=True):
|
|
log = log.decode().rstrip()
|
|
logging.info(log)
|
|
if ('database system is ready to accept connections' in log or
|
|
'database system is shut down' in log):
|
|
break
|
|
|
|
# why print the system is ready to accept connections when its not
|
|
# postgres? wtf
|
|
time.sleep(1)
|
|
logging.info('creating skynet db...')
|
|
|
|
conn = psycopg2.connect(
|
|
user='postgres',
|
|
password=rpassword,
|
|
host='localhost',
|
|
port=port
|
|
)
|
|
logging.info('connected...')
|
|
conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
f'CREATE USER {DB_USER} WITH PASSWORD \'{password}\'')
|
|
cursor.execute(
|
|
f'CREATE DATABASE {DB_NAME}')
|
|
cursor.execute(
|
|
f'GRANT ALL PRIVILEGES ON DATABASE {DB_NAME} TO {DB_USER}')
|
|
|
|
conn.close()
|
|
|
|
logging.info('done.')
|
|
yield container, password, host
|
|
|
|
|
|
@pytest.fixture
|
|
async def skynet_running(postgres_db):
|
|
db_container, db_pass, db_host = postgres_db
|
|
|
|
async with run_skynet(
|
|
db_pass=db_pass,
|
|
db_host=db_host
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def dgpu_workers(request, dockerctl, skynet_running):
|
|
devices = [DeviceRequest(capabilities=[['gpu']])]
|
|
mounts = [Mount(
|
|
'/skynet', str(Path().resolve()), type='bind')]
|
|
|
|
num_containers, initial_algos = request.param
|
|
|
|
cmds = []
|
|
for i in range(num_containers):
|
|
cmd = f'''
|
|
pip install -e . && \
|
|
skynet run dgpu \
|
|
--algos=\'{json.dumps(initial_algos)}\' \
|
|
--uid=dgpu-{i}
|
|
'''
|
|
cmds.append(['bash', '-c', cmd])
|
|
|
|
logging.info(f'launching: \n{cmd}')
|
|
|
|
with dockerctl.run(
|
|
DOCKER_RUNTIME_CUDA,
|
|
name='skynet-test-runtime-cuda',
|
|
commands=cmds,
|
|
environment={
|
|
'HF_TOKEN': os.environ['HF_TOKEN'],
|
|
'HF_HOME': '/skynet/hf_home'
|
|
},
|
|
network='host',
|
|
mounts=mounts,
|
|
device_requests=devices,
|
|
num=num_containers
|
|
) as containers:
|
|
yield containers
|
|
|
|
#for i, container in enumerate(containers):
|
|
# logging.info(f'container {i} logs:')
|
|
# logging.info(container.logs().decode())
|