skygpu
/
skynet
mirror of https://github.com/skygpu/skynet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Network rework in progress, also swap db to be on frontend layer 

Simplify protobuf proto
Add manual telegram test
Add new .ini config system add example and gitignore entry
Expose ports on dgpu Dockerfile
pull/4/head
Guillermo Rodriguez 2023-01-22 12:12:33 -03:00
parent 71c6270a03
commit 4feed81662
No known key found for this signature in database
GPG Key ID: EC3AB66D5D83B392
23 changed files with 1333 additions and 1284 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -1,3 +1,4 @@
skynet.ini
.python-version .python-version
hf_home hf_home
outputs outputs

1
Dockerfile.runtime+cuda
View File

@ -32,3 +32,4 @@ env HF_HOME /hf_home
copy scripts scripts copy scripts scripts
copy tests tests copy tests tests
expose 40000-45000

12
skynet.ini.example 100644
View File

@ -0,0 +1,12 @@
[skynet]
certs_dir = certs
[skynet.dgpu]
hf_home = hf_home
hf_token = hf_XxXxXxXxXxXxXxXxXxXxXxXxXxXxXxXxXx
[skynet.telegram]
token = XXXXXXXXXX:xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
[skynet.telegram-test]
token = XXXXXXXXXX:xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

484
skynet/brain.py
View File

@ -1,35 +1,24 @@
#!/usr/bin/python #!/usr/bin/python
import time
import json
import uuid
import zlib
import logging import logging
import traceback
from uuid import UUID
from pathlib import Path
from functools import partial
from contextlib import asynccontextmanager as acm from contextlib import asynccontextmanager as acm
from collections import OrderedDict from collections import OrderedDict
import trio import trio
import pynng
import trio_asyncio
from pynng import TLSConfig from pynng import Context
from OpenSSL.crypto import (
load_privatekey,
load_certificate,
FILETYPE_PEM
)
from .db import * from .utils import time_ms
from .network import *
from .protobuf import *
from .constants import * from .constants import *
from .protobuf import *
class SkynetRPCBadRequest(BaseException):
...
class SkynetDGPUOffline(BaseException): class SkynetDGPUOffline(BaseException):
... ...
@ -44,39 +33,71 @@ class SkynetShutdownRequested(BaseException):
@acm @acm
async def open_rpc_service(sock, dgpu_bus, db_pool, tls_whitelist, tls_key): async def run_skynet(
rpc_address: str = DEFAULT_RPC_ADDR
):
logging.basicConfig(level=logging.INFO)
logging.info('skynet is starting')
nodes = OrderedDict() nodes = OrderedDict()
wip_reqs = {}
fin_reqs = {}
heartbeats = {} heartbeats = {}
next_worker: Optional[int] = None next_worker: Optional[int] = None
security = len(tls_whitelist) > 0
def connect_node(uid): def connect_node(req: SkynetRPCRequest):
nonlocal next_worker nonlocal next_worker
nodes[uid] = {
'task': None node_params = MessageToDict(req.params)
logging.info(f'got node params {node_params}')
if 'dgpu_addr' not in node_params:
raise SkynetRPCBadRequest(
f'DGPU connection params don\'t include dgpu addr')
session = SessionClient(
node_params['dgpu_addr'],
'skynet',
cert_name='brain.cert',
key_name='brain.key',
ca_name=node_params['cert']
)
try:
session.connect()
node = {
'task': None,
'session': session
} }
logging.info(f'dgpu online: {uid}') node.update(node_params)
nodes[req.uid] = node
logging.info(f'DGPU node online: {req.uid}')
if not next_worker: if not next_worker:
next_worker = 0 next_worker = 0
except pynng.exceptions.ConnectionRefused:
logging.warning(f'error while dialing dgpu node... dropping...')
raise SkynetDGPUOffline('Connection to dgpu node addr failed.')
def disconnect_node(uid): def disconnect_node(uid):
nonlocal next_worker nonlocal next_worker
if uid not in nodes: if uid not in nodes:
logging.warning(f'Attempt to disconnect unknown node {uid}')
return return
i = list(nodes.keys()).index(uid) i = list(nodes.keys()).index(uid)
nodes[uid]['session'].disconnect()
del nodes[uid] del nodes[uid]
if i < next_worker: if i < next_worker:
next_worker -= 1 next_worker -= 1
logging.warning(f'DGPU node offline: {uid}')
if len(nodes) == 0: if len(nodes) == 0:
logging.info('nw: None') logging.info('All nodes disconnected.')
next_worker = None next_worker = None
logging.warning(f'dgpu offline: {uid}')
def is_worker_busy(nid: str): def is_worker_busy(nid: str):
return nodes[nid]['task'] != None return nodes[nid]['task'] != None
@ -90,8 +111,6 @@ async def open_rpc_service(sock, dgpu_bus, db_pool, tls_whitelist, tls_key):
def get_next_worker(): def get_next_worker():
nonlocal next_worker nonlocal next_worker
logging.info('get next_worker called')
logging.info(f'pre next_worker: {next_worker}')
if next_worker == None: if next_worker == None:
raise SkynetDGPUOffline('No workers connected, try again later') raise SkynetDGPUOffline('No workers connected, try again later')
@ -113,392 +132,79 @@ async def open_rpc_service(sock, dgpu_bus, db_pool, tls_whitelist, tls_key):
if next_worker >= len(nodes): if next_worker >= len(nodes):
next_worker = 0 next_worker = 0
logging.info(f'post next_worker: {next_worker}')
return nid return nid
async def dgpu_heartbeat_service(): async def rpc_handler(req: SkynetRPCRequest, ctx: Context):
nonlocal heartbeats result = {'ok': {}}
while True: resp = SkynetRPCResponse()
await trio.sleep(60)
rid = uuid.uuid4().hex
beat_msg = DGPUBusMessage(
rid=rid,
nid='',
method='heartbeat'
)
heartbeats.clear()
heartbeats[rid] = int(time.time() * 1000)
await dgpu_bus.asend(beat_msg.SerializeToString())
logging.info('sent heartbeat')
async def dgpu_bus_streamer(): try:
nonlocal wip_reqs, fin_reqs, heartbeats match req.method:
while True: case 'dgpu_online':
raw_msg = await dgpu_bus.arecv() connect_node(req)
logging.info(f'streamer got {len(raw_msg)} bytes.')
msg = DGPUBusMessage()
msg.ParseFromString(raw_msg)
if security: case 'dgpu_call':
verify_protobuf_msg(msg, tls_whitelist[msg.auth.cert])
rid = msg.rid
if msg.method == 'heartbeat':
sent_time = heartbeats[rid]
delta = msg.params['time'] - sent_time
logging.info(f'got heartbeat reply from {msg.nid}, ping: {delta}')
continue
if rid not in wip_reqs:
continue
if msg.method == 'binary-reply':
logging.info('bin reply, recv extra data')
raw_img = await dgpu_bus.arecv()
msg = (msg, raw_img)
fin_reqs[rid] = msg
event = wip_reqs[rid]
event.set()
del wip_reqs[rid]
async def dgpu_stream_one_img(req: DiffusionParameters, img_buf=None):
nonlocal wip_reqs, fin_reqs, next_worker
nid = get_next_worker() nid = get_next_worker()
idx = list(nodes.keys()).index(nid) idx = list(nodes.keys()).index(nid)
logging.info(f'dgpu_stream_one_img {idx}/{len(nodes)} {nid}') node = nodes[nid]
rid = uuid.uuid4().hex logging.info(f'dgpu_call {idx}/{len(nodes)} {nid} @ {node["dgpu_addr"]}')
ack_event = trio.Event() dgpu_time = await node['session'].rpc('dgpu_time')
img_event = trio.Event() if 'ok' not in dgpu_time.result:
wip_reqs[rid] = ack_event status = MessageToDict(dgpu_time.result)
logging.warning(json.dumps(status, indent=4))
nodes[nid]['task'] = rid
dgpu_req = DGPUBusMessage(
rid=rid,
nid=nid,
method='diffuse')
dgpu_req.params.update(req.to_dict())
if security:
dgpu_req.auth.cert = 'skynet'
dgpu_req.auth.sig = sign_protobuf_msg(dgpu_req, tls_key)
msg = dgpu_req.SerializeToString()
if img_buf:
logging.info(f'sending img of size {len(img_buf)} as attachment')
logging.info(img_buf[:10])
msg = f'BINEXT%$%$'.encode() + msg + b'%$%$' + img_buf
await dgpu_bus.asend(msg)
with trio.move_on_after(4):
await ack_event.wait()
logging.info(f'ack event: {ack_event.is_set()}')
if not ack_event.is_set():
disconnect_node(nid) disconnect_node(nid)
raise SkynetDGPUOffline('dgpu failed to acknowledge request') raise SkynetDGPUComputeError(status['error'])
ack_msg = fin_reqs[rid] dgpu_time = dgpu_time.result['ok']
if 'ack' not in ack_msg.params: logging.info(f'ping to {nid}: {time_ms() - dgpu_time} ms')
disconnect_node(nid)
raise SkynetDGPUOffline('dgpu failed to acknowledge request')
wip_reqs[rid] = img_event
with trio.move_on_after(30):
await img_event.wait()
logging.info(f'img event: {ack_event.is_set()}')
if not img_event.is_set():
disconnect_node(nid)
raise SkynetDGPUComputeError('30 seconds timeout while processing request')
nodes[nid]['task'] = None
resp = fin_reqs[rid]
del fin_reqs[rid]
if isinstance(resp, tuple):
meta, img = resp
return rid, img, meta.params
raise SkynetDGPUComputeError(MessageToDict(resp.params))
async def handle_user_request(rpc_ctx, req):
try: try:
async with db_pool.acquire() as conn: dgpu_result = await node['session'].rpc(
user = await get_or_create_user(conn, req.uid) timeout=45, # give this 45 sec to run cause its compute
binext=req.bin,
result = {} **req.params
match req.method:
case 'txt2img':
logging.info('txt2img')
user_config = {**(await get_user_config(conn, user))}
del user_config['id']
user_config.update(MessageToDict(req.params))
req = DiffusionParameters(**user_config, image=False)
rid, img, meta = await dgpu_stream_one_img(req)
logging.info(f'done streaming {rid}')
result = {
'id': rid,
'img': img.hex(),
'meta': meta
}
await update_user_stats(conn, user, last_prompt=user_config['prompt'])
logging.info('updated user stats.')
case 'img2img':
logging.info('img2img')
user_config = {**(await get_user_config(conn, user))}
del user_config['id']
params = MessageToDict(req.params)
img_buf = bytes.fromhex(params['img'])
del params['img']
user_config.update(params)
req = DiffusionParameters(**user_config, image=True)
if not req.image:
raise AssertionError('Didn\'t enable image flag for img2img?')
rid, img, meta = await dgpu_stream_one_img(req, img_buf=img_buf)
logging.info(f'done streaming {rid}')
result = {
'id': rid,
'img': img.hex(),
'meta': meta
}
await update_user_stats(conn, user, last_prompt=user_config['prompt'])
logging.info('updated user stats.')
case 'redo':
logging.info('redo')
user_config = {**(await get_user_config(conn, user))}
del user_config['id']
prompt = await get_last_prompt_of(conn, user)
if prompt:
req = DiffusionParameters(
prompt=prompt,
**user_config,
image=False
) )
rid, img, meta = await dgpu_stream_one_img(req) result = MessageToDict(dgpu_result.result)
result = {
'id': rid,
'img': img.hex(),
'meta': meta
}
await update_user_stats(conn, user)
logging.info('updated user stats.')
else: if dgpu_result.bin:
result = { resp.bin = dgpu_result.bin
'error': 'skynet_no_last_prompt',
'message': 'No prompt to redo, do txt2img first'
}
case 'config': except trio.TooSlowError:
logging.info('config') result = {'error': 'timeout while processing request'}
if req.params['attr'] in CONFIG_ATTRS:
logging.info(f'update: {req.params}')
await update_user_config(
conn, user, req.params['attr'], req.params['val'])
logging.info('done')
else:
logging.warning(f'{req.params["attr"]} not in {CONFIG_ATTRS}')
case 'stats':
logging.info('stats')
generated, joined, role = await get_user_stats(conn, user)
result = {
'generated': generated,
'joined': joined.strftime(DATE_FORMAT),
'role': role
}
case _:
logging.warn('unknown method')
except SkynetDGPUOffline as e:
result = {
'error': 'skynet_dgpu_offline',
'message': str(e)
}
except SkynetDGPUOverloaded as e:
result = {
'error': 'skynet_dgpu_overloaded',
'message': str(e),
'nodes': len(nodes)
}
except SkynetDGPUComputeError as e:
result = {
'error': 'skynet_dgpu_compute_error',
'message': str(e)
}
except BaseException as e:
traceback.print_exception(type(e), e, e.__traceback__)
result = {
'error': 'skynet_internal_error',
'message': str(e)
}
resp = SkynetRPCResponse()
resp.result.update(result)
if security:
resp.auth.cert = 'skynet'
resp.auth.sig = sign_protobuf_msg(resp, tls_key)
logging.info('sending response')
await rpc_ctx.asend(resp.SerializeToString())
rpc_ctx.close()
logging.info('done')
async def request_service(n):
nonlocal next_worker
while True:
ctx = sock.new_context()
req = SkynetRPCRequest()
req.ParseFromString(await ctx.arecv())
if security:
if req.auth.cert not in tls_whitelist:
logging.warning(
f'{req.cert} not in tls whitelist and security=True')
continue
try:
verify_protobuf_msg(req, tls_whitelist[req.auth.cert])
except ValueError:
logging.warning(
f'{req.cert} sent an unauthenticated msg with security=True')
continue
result = {}
match req.method:
case 'skynet_shutdown':
raise SkynetShutdownRequested
case 'dgpu_online':
connect_node(req.uid)
case 'dgpu_offline': case 'dgpu_offline':
disconnect_node(req.uid) disconnect_node(req.uid)
case 'dgpu_workers': case 'dgpu_workers':
result = len(nodes) result = {'ok': len(nodes)}
case 'dgpu_next': case 'dgpu_next':
result = next_worker result = {'ok': next_worker}
case 'heartbeat': case 'skynet_shutdown':
logging.info('beat') raise SkynetShutdownRequested
result = {'time': time.time()}
case _: case _:
n.start_soon( logging.warning(f'Unknown method {req.method}')
handle_user_request, ctx, req) result = {'error': 'unknown method'}
continue
resp = SkynetRPCResponse() except BaseException as e:
resp.result.update({'ok': result}) result = {'error': str(e)}
if security: resp.result.update(result)
resp.auth.cert = 'skynet'
resp.auth.sig = sign_protobuf_msg(resp, tls_key)
await ctx.asend(resp.SerializeToString()) return resp
ctx.close() rpc_server = SessionServer(
rpc_address,
rpc_handler,
cert_name='brain.cert',
key_name='brain.key'
)
async with rpc_server.open():
async with trio.open_nursery() as n: logging.info('rpc server is up')
n.start_soon(dgpu_bus_streamer)
n.start_soon(dgpu_heartbeat_service)
n.start_soon(request_service, n)
logging.info('starting rpc service')
yield yield
logging.info('stopping rpc service') logging.info('skynet is shuting down...')
n.cancel_scope.cancel()
logging.info('skynet down.')
@acm
async def run_skynet(
db_user: str = DB_USER,
db_pass: str = DB_PASS,
db_host: str = DB_HOST,
rpc_address: str = DEFAULT_RPC_ADDR,
dgpu_address: str = DEFAULT_DGPU_ADDR,
security: bool = True
):
logging.basicConfig(level=logging.INFO)
logging.info('skynet is starting')
tls_config = None
if security:
# load tls certs
certs_dir = Path(DEFAULT_CERTS_DIR).resolve()
tls_key_data = (certs_dir / DEFAULT_CERT_SKYNET_PRIV).read_text()
tls_key = load_privatekey(FILETYPE_PEM, tls_key_data)
tls_cert_data = (certs_dir / DEFAULT_CERT_SKYNET_PUB).read_text()
tls_cert = load_certificate(FILETYPE_PEM, tls_cert_data)
tls_whitelist = {}
for cert_path in (certs_dir / 'whitelist').glob('*.cert'):
tls_whitelist[cert_path.stem] = load_certificate(
FILETYPE_PEM, cert_path.read_text())
cert_start = tls_cert_data.index('\n') + 1
logging.info(f'tls_cert: {tls_cert_data[cert_start:cert_start+64]}...')
logging.info(f'tls_whitelist len: {len(tls_whitelist)}')
rpc_address = 'tls+' + rpc_address
dgpu_address = 'tls+' + dgpu_address
tls_config = TLSConfig(
TLSConfig.MODE_SERVER,
own_key_string=tls_key_data,
own_cert_string=tls_cert_data)
with (
pynng.Rep0(recv_max_size=0) as rpc_sock,
pynng.Bus0(recv_max_size=0) as dgpu_bus
):
async with open_database_connection(
db_user, db_pass, db_host) as db_pool:
logging.info('connected to db.')
if security:
rpc_sock.tls_config = tls_config
dgpu_bus.tls_config = tls_config
rpc_sock.listen(rpc_address)
dgpu_bus.listen(dgpu_address)
try:
async with open_rpc_service(
rpc_sock, dgpu_bus, db_pool, tls_whitelist, tls_key):
yield
except SkynetShutdownRequested:
...
logging.info('disconnected from db.')

36
skynet/cli.py
View File

@ -17,8 +17,8 @@ if torch_enabled:
from .dgpu import open_dgpu_node from .dgpu import open_dgpu_node
from .brain import run_skynet from .brain import run_skynet
from .config import *
from .constants import ALGOS, DEFAULT_RPC_ADDR, DEFAULT_DGPU_ADDR from .constants import ALGOS, DEFAULT_RPC_ADDR, DEFAULT_DGPU_ADDR
from .frontend.telegram import run_skynet_telegram from .frontend.telegram import run_skynet_telegram
@ -38,8 +38,8 @@ def skynet(*args, **kwargs):
@click.option('--steps', '-s', default=26) @click.option('--steps', '-s', default=26)
@click.option('--seed', '-S', default=None) @click.option('--seed', '-S', default=None)
def txt2img(*args, **kwargs): def txt2img(*args, **kwargs):
assert 'HF_TOKEN' in os.environ _, hf_token, _, cfg = init_env_from_config()
utils.txt2img(os.environ['HF_TOKEN'], **kwargs) utils.txt2img(hf_token, **kwargs)
@click.command() @click.command()
@click.option('--model', '-m', default='midj') @click.option('--model', '-m', default='midj')
@ -52,9 +52,9 @@ def txt2img(*args, **kwargs):
@click.option('--steps', '-s', default=26) @click.option('--steps', '-s', default=26)
@click.option('--seed', '-S', default=None) @click.option('--seed', '-S', default=None)
def img2img(model, prompt, input, output, strength, guidance, steps, seed): def img2img(model, prompt, input, output, strength, guidance, steps, seed):
assert 'HF_TOKEN' in os.environ _, hf_token, _, cfg = init_env_from_config()
utils.img2img( utils.img2img(
os.environ['HF_TOKEN'], hf_token,
model=model, model=model,
prompt=prompt, prompt=prompt,
img_path=input, img_path=input,
@ -85,29 +85,17 @@ def run(*args, **kwargs):
@click.option('--loglevel', '-l', default='warning', help='Logging level') @click.option('--loglevel', '-l', default='warning', help='Logging level')
@click.option( @click.option(
'--host', '-H', default=DEFAULT_RPC_ADDR) '--host', '-H', default=DEFAULT_RPC_ADDR)
@click.option(
'--host-dgpu', '-D', default=DEFAULT_DGPU_ADDR)
@click.option(
'--db-host', '-h', default='localhost:5432')
@click.option(
'--db-pass', '-p', default='password')
def brain( def brain(
loglevel: str, loglevel: str,
host: str, host: str
host_dgpu: str,
db_host: str,
db_pass: str
): ):
async def _run_skynet(): async def _run_skynet():
async with run_skynet( async with run_skynet(
db_host=db_host, rpc_address=host
db_pass=db_pass,
rpc_address=host,
dgpu_address=host_dgpu
): ):
await trio.sleep_forever() await trio.sleep_forever()
trio_asyncio.run(_run_skynet) trio.run(_run_skynet)
@run.command() @run.command()
@ -115,9 +103,9 @@ def brain(
@click.option( @click.option(
'--uid', '-u', required=True) '--uid', '-u', required=True)
@click.option( @click.option(
'--key', '-k', default='dgpu') '--key', '-k', default='dgpu.key')
@click.option( @click.option(
'--cert', '-c', default='whitelist/dgpu') '--cert', '-c', default='whitelist/dgpu.cert')
@click.option( @click.option(
'--algos', '-a', default=json.dumps(['midj'])) '--algos', '-a', default=json.dumps(['midj']))
@click.option( @click.option(
@ -159,11 +147,11 @@ def telegram(
cert: str, cert: str,
rpc: str rpc: str
): ):
assert 'TG_TOKEN' in os.environ _, _, tg_token, cfg = init_env_from_config()
trio_asyncio.run( trio_asyncio.run(
partial( partial(
run_skynet_telegram, run_skynet_telegram,
os.environ['TG_TOKEN'], tg_token,
key_name=key, key_name=key,
cert_name=cert, cert_name=cert,
rpc_address=rpc rpc_address=rpc

36
skynet/config.py 100644
View File

@ -0,0 +1,36 @@
#!/usr/bin/python
from pathlib import Path
from configparser import ConfigParser
from .constants import DEFAULT_CONFIG_PATH
def load_skynet_ini(
file_path=DEFAULT_CONFIG_PATH
):
config = ConfigParser()
config.read(file_path)
return config
def init_env_from_config(
file_path=DEFAULT_CONFIG_PATH
):
config = load_skynet_ini()
if 'HF_TOKEN' in os.environ:
hf_token = os.environ['HF_TOKEN']
else:
hf_token = config['skynet']['dgpu']['hf_token']
if 'HF_HOME' in os.environ:
hf_home = os.environ['HF_HOME']
else:
hf_home = config['skynet']['dgpu']['hf_home']
if 'TG_TOKEN' in os.environ:
tg_token = os.environ['TG_TOKEN']
else:
tg_token = config['skynet']['telegram']['token']
return hf_home, hf_token, tg_token, config

8
skynet/constants.py
View File

@ -1,14 +1,9 @@
#!/usr/bin/python #!/usr/bin/python
VERSION = '0.1a8' VERSION = '0.1a9'
DOCKER_RUNTIME_CUDA = 'skynet:runtime-cuda' DOCKER_RUNTIME_CUDA = 'skynet:runtime-cuda'
DB_HOST = 'localhost:5432'
DB_USER = 'skynet'
DB_PASS = 'password'
DB_NAME = 'skynet'
ALGOS = { ALGOS = {
'midj': 'prompthero/openjourney', 'midj': 'prompthero/openjourney',
'stable': 'runwayml/stable-diffusion-v1-5', 'stable': 'runwayml/stable-diffusion-v1-5',
@ -118,6 +113,7 @@ DEFAULT_ALGO = 'midj'
DEFAULT_ROLE = 'pleb' DEFAULT_ROLE = 'pleb'
DEFAULT_UPSCALER = None DEFAULT_UPSCALER = None
DEFAULT_CONFIG_PATH = 'skynet.ini'
DEFAULT_CERTS_DIR = 'certs' DEFAULT_CERTS_DIR = 'certs'
DEFAULT_CERT_WHITELIST_DIR = 'whitelist' DEFAULT_CERT_WHITELIST_DIR = 'whitelist'
DEFAULT_CERT_SKYNET_PUB = 'brain.cert' DEFAULT_CERT_SKYNET_PUB = 'brain.cert'

5
skynet/db/__init__.py 100644
View File

@ -0,0 +1,5 @@
#!/usr/bin/python
from .proxy import open_database_connection
from .functions import open_new_database

95
skynet/db.py → skynet/db/functions.py
View File

@ -1,18 +1,21 @@
#!/usr/bin/python #!/usr/bin/python
import time
import random
import string
import logging import logging
from typing import Optional from typing import Optional
from datetime import datetime from datetime import datetime
from contextlib import asynccontextmanager as acm from contextlib import contextmanager as cm
import trio import docker
import triopg import psycopg2
import trio_asyncio
from asyncpg.exceptions import UndefinedColumnError from asyncpg.exceptions import UndefinedColumnError
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
from .constants import * from ..constants import *
DB_INIT_SQL = ''' DB_INIT_SQL = '''
@ -75,29 +78,67 @@ def try_decode_uid(uid: str):
return None, None return None, None
@acm @cm
async def open_database_connection( def open_new_database():
db_user: str = DB_USER, rpassword = ''.join(
db_pass: str = DB_PASS, random.choice(string.ascii_lowercase)
db_host: str = DB_HOST, for i in range(12))
db_name: str = DB_NAME password = ''.join(
): random.choice(string.ascii_lowercase)
async with trio_asyncio.open_loop() as loop: for i in range(12))
async with triopg.create_pool(
dsn=f'postgres://{db_user}:{db_pass}@{db_host}/{db_name}'
) as pool_conn:
async with pool_conn.acquire() as conn:
res = await conn.execute(f'''
select distinct table_schema
from information_schema.tables
where table_schema = \'{db_name}\'
''')
if '1' in res:
logging.info('schema already in db, skipping init')
else:
await conn.execute(DB_INIT_SQL)
yield pool_conn dclient = docker.from_env()
container = dclient.containers.run(
'postgres',
name='skynet-test-postgres',
ports={'5432/tcp': None},
environment={
'POSTGRES_PASSWORD': rpassword
},
detach=True,
remove=True
)
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
# ip = container.attrs['NetworkSettings']['IPAddress']
container.reload()
port = container.ports['5432/tcp'][0]['HostPort']
host = f'localhost:{port}'
# 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 skynet WITH PASSWORD \'{password}\'')
cursor.execute(
f'CREATE DATABASE skynet')
cursor.execute(
f'GRANT ALL PRIVILEGES ON DATABASE skynet TO skynet')
conn.close()
logging.info('done.')
yield container, password, host
container.stop()
async def get_user(conn, uid: str): async def get_user(conn, uid: str):

123
skynet/db/proxy.py 100644
View File

@ -0,0 +1,123 @@
#!/usr/bin/python
import importlib
from contextlib import asynccontextmanager as acm
import trio
import tractor
import asyncpg
import asyncio
import trio_asyncio
_spawn_kwargs = {
'infect_asyncio': True,
}
async def aio_db_proxy(
to_trio: trio.MemorySendChannel,
from_trio: asyncio.Queue,
db_user: str = 'skynet',
db_pass: str = 'password',
db_host: str = 'localhost:5432',
db_name: str = 'skynet'
) -> None:
db = importlib.import_module('skynet.db.functions')
pool = await asyncpg.create_pool(
dsn=f'postgres://{db_user}:{db_pass}@{db_host}/{db_name}')
async with pool_conn.acquire() as conn:
res = await conn.execute(f'''
select distinct table_schema
from information_schema.tables
where table_schema = \'{db_name}\'
''')
if '1' in res:
logging.info('schema already in db, skipping init')
else:
await conn.execute(DB_INIT_SQL)
# a first message must be sent **from** this ``asyncio``
# task or the ``trio`` side will never unblock from
# ``tractor.to_asyncio.open_channel_from():``
to_trio.send_nowait('start')
# XXX: this uses an ``from_trio: asyncio.Queue`` currently but we
# should probably offer something better.
while True:
msg = await from_trio.get()
method = getattr(db, msg.get('method'))
args = getattr(db, msg.get('args', []))
kwargs = getattr(db, msg.get('kwargs', {}))
async with pool_conn.acquire() as conn:
result = await method(conn, *args, **kwargs)
to_trio.send_nowait(result)
@tractor.context
async def trio_to_aio_db_proxy(
ctx: tractor.Context,
db_user: str = 'skynet',
db_pass: str = 'password',
db_host: str = 'localhost:5432',
db_name: str = 'skynet'
):
# this will block until the ``asyncio`` task sends a "first"
# message.
async with tractor.to_asyncio.open_channel_from(
aio_db_proxy,
db_user=db_user,
db_pass=db_pass,
db_host=db_host,
db_name=db_name
) as (first, chan):
assert first == 'start'
await ctx.started(first)
async with ctx.open_stream() as stream:
async for msg in stream:
await chan.send(msg)
out = await chan.receive()
# echo back to parent actor-task
await stream.send(out)
@acm
async def open_database_connection(
db_user: str = 'skynet',
db_pass: str = 'password',
db_host: str = 'localhost:5432',
db_name: str = 'skynet'
):
async with tractor.open_nursery() as n:
p = await n.start_actor(
'aio_db_proxy',
enable_modules=[__name__],
infect_asyncio=True,
)
async with p.open_context(
trio_to_aio_db_proxy,
db_user=db_user,
db_pass=db_pass,
db_host=db_host,
db_name=db_name
) as (ctx, first):
async with ctx.open_stream() as stream:
async def _db_pc(method: str, *args, **kwargs):
await stream.send({
'method': method,
'args': args,
'kwargs': kwargs
})
return await stream.receive()
yield _db_pc

329
skynet/dgpu.py
View File

@ -2,29 +2,17 @@
import gc import gc
import io import io
import trio
import json import json
import uuid
import time
import zlib
import random import random
import logging import logging
import traceback
from PIL import Image from PIL import Image
from typing import List, Optional from typing import List, Optional
from pathlib import Path
from contextlib import ExitStack
import pynng import trio
import torch import torch
from pynng import TLSConfig from pynng import Context
from OpenSSL.crypto import (
load_privatekey,
load_certificate,
FILETYPE_PEM
)
from diffusers import ( from diffusers import (
StableDiffusionPipeline, StableDiffusionPipeline,
StableDiffusionImg2ImgPipeline, StableDiffusionImg2ImgPipeline,
@ -34,12 +22,9 @@ from realesrgan import RealESRGANer
from basicsr.archs.rrdbnet_arch import RRDBNet from basicsr.archs.rrdbnet_arch import RRDBNet
from diffusers.models import UNet2DConditionModel from diffusers.models import UNet2DConditionModel
from .utils import ( from .utils import *
pipeline_for, from .network import *
convert_from_cv2_to_image, convert_from_image_to_cv2
)
from .protobuf import * from .protobuf import *
from .frontend import open_skynet_rpc
from .constants import * from .constants import *
@ -64,65 +49,16 @@ class DGPUComputeError(BaseException):
... ...
class ReconnectingBus:
def __init__(self, address: str, tls_config: Optional[TLSConfig]):
self.address = address
self.tls_config = tls_config
self._stack = ExitStack()
self._sock = None
self._closed = True
def connect(self):
self._sock = self._stack.enter_context(
pynng.Bus0(recv_max_size=0))
self._sock.tls_config = self.tls_config
self._sock.dial(self.address)
self._closed = False
async def arecv(self):
while True:
try:
return await self._sock.arecv()
except pynng.exceptions.Closed:
if self._closed:
raise
async def asend(self, msg):
while True:
try:
return await self._sock.asend(msg)
except pynng.exceptions.Closed:
if self._closed:
raise
def close(self):
self._stack.close()
self._stack = ExitStack()
self._closed = True
def reconnect(self):
self.close()
self.connect()
async def open_dgpu_node( async def open_dgpu_node(
cert_name: str, cert_name: str,
unique_id: str, unique_id: str,
key_name: Optional[str], key_name: Optional[str],
rpc_address: str = DEFAULT_RPC_ADDR, rpc_address: str = DEFAULT_RPC_ADDR,
dgpu_address: str = DEFAULT_DGPU_ADDR, dgpu_address: str = DEFAULT_DGPU_ADDR,
initial_algos: Optional[List[str]] = None, initial_algos: Optional[List[str]] = None
security: bool = True
): ):
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.DEBUG)
logging.info(f'starting dgpu node!') logging.info(f'starting dgpu node!')
name = uuid.uuid4()
logging.info(f'loading models...') logging.info(f'loading models...')
upscaler = init_upscaler() upscaler = init_upscaler()
@ -141,9 +77,23 @@ async def open_dgpu_node(
logging.info('memory summary:') logging.info('memory summary:')
logging.info('\n' + torch.cuda.memory_summary()) logging.info('\n' + torch.cuda.memory_summary())
async def gpu_compute_one(ireq: DiffusionParameters, image=None): async def gpu_compute_one(method: str, params: dict, binext: Optional[bytes] = None):
algo = ireq.algo + 'img' if image else ireq.algo match method:
case 'diffuse':
image = None
algo = params['algo']
if binext:
algo += 'img'
image = Image.open(io.BytesIO(binext))
w, h = image.size
logging.info(f'user sent img of size {image.size}')
if w > 512 or h > 512:
image.thumbnail((512, 512))
logging.info(f'resized it to {image.size}')
if algo not in models: if algo not in models:
logging.info(f'{algo} not in loaded models, swapping...')
least_used = list(models.keys())[0] least_used = list(models.keys())[0]
for model in models: for model in models:
if models[least_used]['generated'] > models[model]['generated']: if models[least_used]['generated'] > models[model]['generated']:
@ -153,29 +103,35 @@ async def open_dgpu_node(
gc.collect() gc.collect()
models[algo] = { models[algo] = {
'pipe': pipeline_for(ireq.algo, image=True if image else False), 'pipe': pipeline_for(params['algo'], image=True if binext else False),
'generated': 0 'generated': 0
} }
logging.info(f'swapping done.')
_params = {} _params = {}
if ireq.image: logging.info(method)
logging.info(json.dumps(params, indent=4))
logging.info(f'binext: {len(binext) if binext else 0} bytes')
if binext:
_params['image'] = image _params['image'] = image
_params['strength'] = ireq.strength _params['strength'] = params['strength']
else: else:
_params['width'] = int(ireq.width) _params['width'] = int(params['width'])
_params['height'] = int(ireq.height) _params['height'] = int(params['height'])
try: try:
image = models[algo]['pipe']( image = models[algo]['pipe'](
ireq.prompt, params['prompt'],
**_params, **_params,
guidance_scale=ireq.guidance, guidance_scale=params['guidance'],
num_inference_steps=int(ireq.step), num_inference_steps=int(params['step']),
generator=torch.Generator("cuda").manual_seed(ireq.seed) generator=torch.Generator("cuda").manual_seed(
int(params['seed']) if params['seed'] else random.randint(0, 2 ** 64)
)
).images[0] ).images[0]
if ireq.upscaler == 'x4': if params['upscaler'] == 'x4':
logging.info(f'size: {len(image.tobytes())}') logging.info(f'size: {len(image.tobytes())}')
logging.info('performing upscale...') logging.info('performing upscale...')
input_img = image.convert('RGB') input_img = image.convert('RGB')
@ -199,183 +155,62 @@ async def open_dgpu_node(
finally: finally:
torch.cuda.empty_cache() torch.cuda.empty_cache()
case _:
raise DGPUComputeError('Unsupported compute method')
async with ( async def rpc_handler(req: SkynetRPCRequest, ctx: Context):
open_skynet_rpc( result = {}
unique_id, resp = SkynetRPCResponse()
rpc_address=rpc_address,
security=security, match req.method:
case 'dgpu_time':
result = {'ok': time_ms()}
case _:
logging.debug(f'dgpu got one request: {req.method}')
try:
resp.bin = await gpu_compute_one(
req.method, MessageToDict(req.params),
binext=req.bin if req.bin else None
)
logging.debug(f'dgpu processed one request')
except DGPUComputeError as e:
result = {'error': str(e)}
resp.result.update(result)
return resp
rpc_server = SessionServer(
dgpu_address,
rpc_handler,
cert_name=cert_name, cert_name=cert_name,
key_name=key_name key_name=key_name
) as rpc_call, )
trio.open_nursery() as n skynet_rpc = SessionClient(
): rpc_address,
unique_id,
cert_name=cert_name,
key_name=key_name
)
skynet_rpc.connect()
tls_config = None
if security:
# load tls certs
if not key_name:
key_name = cert_name
certs_dir = Path(DEFAULT_CERTS_DIR).resolve() async with rpc_server.open() as rpc_server:
res = await skynet_rpc.rpc(
'dgpu_online', {
'dgpu_addr': rpc_server.addr,
'cert': cert_name
})
skynet_cert_path = certs_dir / 'brain.cert'
tls_cert_path = certs_dir / f'{cert_name}.cert'
tls_key_path = certs_dir / f'{key_name}.key'
cert_name = tls_cert_path.stem
skynet_cert_data = skynet_cert_path.read_text()
skynet_cert = load_certificate(FILETYPE_PEM, skynet_cert_data)
tls_cert_data = tls_cert_path.read_text()
tls_key_data = tls_key_path.read_text()
tls_key = load_privatekey(FILETYPE_PEM, tls_key_data)
logging.info(f'skynet cert: {skynet_cert_path}')
logging.info(f'dgpu cert: {tls_cert_path}')
logging.info(f'dgpu key: {tls_key_path}')
dgpu_address = 'tls+' + dgpu_address
tls_config = TLSConfig(
TLSConfig.MODE_CLIENT,
own_key_string=tls_key_data,
own_cert_string=tls_cert_data,
ca_string=skynet_cert_data)
logging.info(f'connecting to {dgpu_address}')
dgpu_bus = ReconnectingBus(dgpu_address, tls_config)
dgpu_bus.connect()
last_msg = time.time()
async def connection_refresher(refresh_time: int = 120):
nonlocal last_msg
while True:
now = time.time()
last_msg_time_delta = now - last_msg
logging.info(f'time since last msg: {last_msg_time_delta}')
if last_msg_time_delta > refresh_time:
dgpu_bus.reconnect()
logging.info('reconnected!')
last_msg = now
await trio.sleep(refresh_time)
n.start_soon(connection_refresher)
res = await rpc_call('dgpu_online')
assert 'ok' in res.result assert 'ok' in res.result
try: try:
while True: await trio.sleep_forever()
msg = await dgpu_bus.arecv()
img = None
if b'BINEXT' in msg:
header, msg, img_raw = msg.split(b'%$%$')
logging.info(f'got img attachment of size {len(img_raw)}')
logging.info(img_raw[:10])
raw_img = zlib.decompress(img_raw)
logging.info(raw_img[:10])
img = Image.open(io.BytesIO(raw_img))
w, h = img.size
logging.info(f'user sent img of size {img.size}')
if w > 512 or h > 512:
img.thumbnail((512, 512))
logging.info(f'resized it to {img.size}')
req = DGPUBusMessage()
req.ParseFromString(msg)
last_msg = time.time()
if req.method == 'heartbeat':
rep = DGPUBusMessage(
rid=req.rid,
nid=unique_id,
method=req.method
)
rep.params.update({'time': int(time.time() * 1000)})
if security:
rep.auth.cert = cert_name
rep.auth.sig = sign_protobuf_msg(rep, tls_key)
await dgpu_bus.asend(rep.SerializeToString())
logging.info('heartbeat reply')
continue
if req.nid != unique_id:
logging.info(
f'witnessed msg {req.rid}, node involved: {req.nid}')
continue
if security:
verify_protobuf_msg(req, skynet_cert)
ack_resp = DGPUBusMessage(
rid=req.rid,
nid=req.nid
)
ack_resp.params.update({'ack': {}})
if security:
ack_resp.auth.cert = cert_name
ack_resp.auth.sig = sign_protobuf_msg(ack_resp, tls_key)
# send ack
await dgpu_bus.asend(ack_resp.SerializeToString())
logging.info(f'sent ack, processing {req.rid}...')
try:
img_req = DiffusionParameters(**req.params)
if not img_req.seed:
img_req.seed = random.randint(0, 2 ** 64)
img = await gpu_compute_one(img_req, image=img)
img_resp = DGPUBusMessage(
rid=req.rid,
nid=req.nid,
method='binary-reply'
)
img_resp.params.update({
'len': len(img),
'meta': img_req.to_dict()
})
except DGPUComputeError as e:
traceback.print_exception(type(e), e, e.__traceback__)
img_resp = DGPUBusMessage(
rid=req.rid,
nid=req.nid
)
img_resp.params.update({'error': str(e)})
if security:
img_resp.auth.cert = cert_name
img_resp.auth.sig = sign_protobuf_msg(img_resp, tls_key)
# send final image
logging.info('sending img back...')
raw_msg = img_resp.SerializeToString()
await dgpu_bus.asend(raw_msg)
logging.info(f'sent {len(raw_msg)} bytes.')
if img_resp.method == 'binary-reply':
await dgpu_bus.asend(zlib.compress(img))
logging.info(f'sent {len(img)} bytes.')
except KeyboardInterrupt: except KeyboardInterrupt:
logging.info('interrupt caught, stopping...') logging.info('interrupt caught, stopping...')
n.cancel_scope.cancel()
dgpu_bus.close()
finally: finally:
res = await rpc_call('dgpu_offline') res = await skynet_rpc.rpc('dgpu_offline')
assert 'ok' in res.result assert 'ok' in res.result

79
skynet/frontend/__init__.py
View File

@ -4,7 +4,7 @@ import json
from typing import Union, Optional from typing import Union, Optional
from pathlib import Path from pathlib import Path
from contextlib import asynccontextmanager as acm from contextlib import contextmanager as cm
import pynng import pynng
@ -17,6 +17,7 @@ from OpenSSL.crypto import (
from google.protobuf.struct_pb2 import Struct from google.protobuf.struct_pb2 import Struct
from ..network import SessionClient
from ..constants import * from ..constants import *
from ..protobuf.auth import * from ..protobuf.auth import *
@ -39,75 +40,23 @@ class ConfigSizeDivisionByEight(BaseException):
... ...
@acm @cm
async def open_skynet_rpc( def open_skynet_rpc(
unique_id: str, unique_id: str,
rpc_address: str = DEFAULT_RPC_ADDR, rpc_address: str = DEFAULT_RPC_ADDR,
security: bool = False,
cert_name: Optional[str] = None, cert_name: Optional[str] = None,
key_name: Optional[str] = None key_name: Optional[str] = None
): ):
tls_config = None sesh = SessionClient(
rpc_address,
if security: unique_id,
# load tls certs cert_name=cert_name,
if not key_name: key_name=key_name
key_name = cert_name )
logging.debug(f'opening skynet rpc...')
certs_dir = Path(DEFAULT_CERTS_DIR).resolve() sesh.connect()
yield sesh
skynet_cert_data = (certs_dir / 'brain.cert').read_text() sesh.disconnect()
skynet_cert = load_certificate(FILETYPE_PEM, skynet_cert_data)
tls_cert_path = certs_dir / f'{cert_name}.cert'
tls_cert_data = tls_cert_path.read_text()
tls_cert = load_certificate(FILETYPE_PEM, tls_cert_data)
cert_name = tls_cert_path.stem
tls_key_data = (certs_dir / f'{key_name}.key').read_text()
tls_key = load_privatekey(FILETYPE_PEM, tls_key_data)
rpc_address = 'tls+' + rpc_address
tls_config = TLSConfig(
TLSConfig.MODE_CLIENT,
own_key_string=tls_key_data,
own_cert_string=tls_cert_data,
ca_string=skynet_cert_data)
with pynng.Req0(recv_max_size=0) as sock:
if security:
sock.tls_config = tls_config
sock.dial(rpc_address)
async def _rpc_call(
method: str,
params: dict = {},
uid: Optional[str] = None
):
req = SkynetRPCRequest()
req.uid = uid if uid else unique_id
req.method = method
req.params.update(params)
if security:
req.auth.cert = cert_name
req.auth.sig = sign_protobuf_msg(req, tls_key)
ctx = sock.new_context()
await ctx.asend(req.SerializeToString())
resp = SkynetRPCResponse()
resp.ParseFromString(await ctx.arecv())
ctx.close()
if security:
verify_protobuf_msg(resp, skynet_cert)
return resp
yield _rpc_call
def validate_user_config_request(req: str): def validate_user_config_request(req: str):
params = req.split(' ') params = req.split(' ')

155
skynet/frontend/telegram.py
View File

@ -6,8 +6,6 @@ import logging
from datetime import datetime from datetime import datetime
import pynng
from PIL import Image from PIL import Image
from trio_asyncio import aio_as_trio from trio_asyncio import aio_as_trio
@ -16,6 +14,7 @@ from telebot.types import (
) )
from telebot.async_telebot import AsyncTeleBot from telebot.async_telebot import AsyncTeleBot
from ..db import open_database_connection
from ..constants import * from ..constants import *
from . import * from . import *
@ -56,30 +55,29 @@ def prepare_metainfo_caption(tguser, meta: dict) -> str:
async def run_skynet_telegram( async def run_skynet_telegram(
name: str,
tg_token: str, tg_token: str,
key_name: str = 'telegram-frontend', key_name: str = 'telegram-frontend.key',
cert_name: str = 'whitelist/telegram-frontend', cert_name: str = 'whitelist/telegram-frontend.cert',
rpc_address: str = DEFAULT_RPC_ADDR rpc_address: str = DEFAULT_RPC_ADDR,
db_host: str = 'localhost:5432',
db_user: str = 'skynet',
db_pass: str = 'password'
): ):
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
bot = AsyncTeleBot(tg_token) bot = AsyncTeleBot(tg_token)
logging.info(f'tg_token: {tg_token}')
async with open_skynet_rpc( async with open_database_connection(
'skynet-telegram-0', db_user, db_pass, db_host
) as db_call:
with open_skynet_rpc(
f'skynet-telegram-{name}',
rpc_address=rpc_address, rpc_address=rpc_address,
security=True,
cert_name=cert_name, cert_name=cert_name,
key_name=key_name key_name=key_name
) as rpc_call: ) as session:
async def _rpc_call(
uid: int,
method: str,
params: dict = {}
):
return await rpc_call(
method, params, uid=f'{PREFIX}+{uid}')
@bot.message_handler(commands=['help']) @bot.message_handler(commands=['help'])
async def send_help(message): async def send_help(message):
@ -103,6 +101,11 @@ async def run_skynet_telegram(
@bot.message_handler(commands=['txt2img']) @bot.message_handler(commands=['txt2img'])
async def send_txt2img(message): async def send_txt2img(message):
chat = message.chat chat = message.chat
reply_id = None
if chat.type == 'group' and chat.id == GROUP_ID:
reply_id = message.message_id
user_id = f'tg+{message.from_user.id}'
prompt = ' '.join(message.text.split(' ')[1:]) prompt = ' '.join(message.text.split(' ')[1:])
@ -111,10 +114,19 @@ async def run_skynet_telegram(
return return
logging.info(f'mid: {message.id}') logging.info(f'mid: {message.id}')
resp = await _rpc_call( user = await db_call('get_or_create_user', user_id)
message.from_user.id, user_config = {**(await db_call('get_user_config', user))}
'txt2img', del user_config['id']
{'prompt': prompt}
resp = await session.rpc(
'dgpu_call', {
'method': 'diffuse',
'params': {
'prompt': prompt,
**user_config
}
},
timeout=60
) )
logging.info(f'resp to {message.id} arrived') logging.info(f'resp to {message.id} arrived')
@ -122,10 +134,11 @@ async def run_skynet_telegram(
result = MessageToDict(resp.result) result = MessageToDict(resp.result)
if 'error' in resp.result: if 'error' in resp.result:
resp_txt = resp.result['message'] resp_txt = resp.result['message']
await bot.reply_to(message, resp_txt)
else: else:
logging.info(result['id']) logging.info(result['id'])
img_raw = zlib.decompress(bytes.fromhex(result['img'])) img_raw = resp.bin
logging.info(f'got image of size: {len(img_raw)}') logging.info(f'got image of size: {len(img_raw)}')
img = Image.open(io.BytesIO(img_raw)) img = Image.open(io.BytesIO(img_raw))
@ -133,17 +146,26 @@ async def run_skynet_telegram(
GROUP_ID, GROUP_ID,
caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']), caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
photo=img, photo=img,
reply_to_message_id=reply_id,
reply_markup=build_redo_menu() reply_markup=build_redo_menu()
) )
return return
await bot.reply_to(message, resp_txt)
@bot.message_handler(func=lambda message: True, content_types=['photo']) @bot.message_handler(func=lambda message: True, content_types=['photo'])
async def send_img2img(message): async def send_img2img(message):
chat = message.chat chat = message.chat
reply_id = None
if chat.type == 'group' and chat.id == GROUP_ID:
reply_id = message.message_id
user_id = f'tg+{message.from_user.id}'
if not message.caption.startswith('/img2img'): if not message.caption.startswith('/img2img'):
await bot.reply_to(
message,
'For image to image you need to add /img2img to the beggining of your caption'
)
return return
prompt = ' '.join(message.caption.split(' ')[1:]) prompt = ' '.join(message.caption.split(' ')[1:])
@ -155,13 +177,23 @@ async def run_skynet_telegram(
file_id = message.photo[-1].file_id file_id = message.photo[-1].file_id
file_path = (await bot.get_file(file_id)).file_path file_path = (await bot.get_file(file_id)).file_path
file_raw = await bot.download_file(file_path) file_raw = await bot.download_file(file_path)
img = zlib.compress(file_raw)
logging.info(f'mid: {message.id}') logging.info(f'mid: {message.id}')
resp = await _rpc_call(
message.from_user.id, user = await db_call('get_or_create_user', user_id)
'img2img', user_config = {**(await db_call('get_user_config', user))}
{'prompt': prompt, 'img': img.hex()} del user_config['id']
resp = await session.rpc(
'dgpu_call', {
'method': 'diffuse',
'params': {
'prompt': prompt,
**user_config
}
},
binext=file_raw,
timeout=60
) )
logging.info(f'resp to {message.id} arrived') logging.info(f'resp to {message.id} arrived')
@ -169,10 +201,11 @@ async def run_skynet_telegram(
result = MessageToDict(resp.result) result = MessageToDict(resp.result)
if 'error' in resp.result: if 'error' in resp.result:
resp_txt = resp.result['message'] resp_txt = resp.result['message']
await bot.reply_to(message, resp_txt)
else: else:
logging.info(result['id']) logging.info(result['id'])
img_raw = zlib.decompress(bytes.fromhex(result['img'])) img_raw = resp.bin
logging.info(f'got image of size: {len(img_raw)}') logging.info(f'got image of size: {len(img_raw)}')
img = Image.open(io.BytesIO(img_raw)) img = Image.open(io.BytesIO(img_raw))
@ -184,30 +217,51 @@ async def run_skynet_telegram(
img, img,
caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']) caption=prepare_metainfo_caption(message.from_user, result['meta']['meta'])
) )
] ],
reply_to_message_id=reply_id
) )
return return
await bot.reply_to(message, resp_txt)
@bot.message_handler(commands=['img2img']) @bot.message_handler(commands=['img2img'])
async def redo_txt2img(message): async def img2img_missing_image(message):
await bot.reply_to( await bot.reply_to(
message, message,
'seems you tried to do an img2img command without sending image' 'seems you tried to do an img2img command without sending image'
) )
async def _redo(message): @bot.message_handler(commands=['redo'])
resp = await _rpc_call(message.from_user.id, 'redo') async def redo(message):
chat = message.chat
reply_id = None
if chat.type == 'group' and chat.id == GROUP_ID:
reply_id = message.message_id
user_config = {**(await db_call('get_user_config', user))}
del user_config['id']
prompt = await db_call('get_last_prompt_of', user)
resp = await session.rpc(
'dgpu_call', {
'method': 'diffuse',
'params': {
'prompt': prompt,
**user_config
}
},
timeout=60
)
logging.info(f'resp to {message.id} arrived')
resp_txt = '' resp_txt = ''
result = MessageToDict(resp.result) result = MessageToDict(resp.result)
if 'error' in resp.result: if 'error' in resp.result:
resp_txt = resp.result['message'] resp_txt = resp.result['message']
await bot.reply_to(message, resp_txt)
else: else:
logging.info(result['id']) logging.info(result['id'])
img_raw = zlib.decompress(bytes.fromhex(result['img'])) img_raw = resp.bin
logging.info(f'got image of size: {len(img_raw)}') logging.info(f'got image of size: {len(img_raw)}')
img = Image.open(io.BytesIO(img_raw)) img = Image.open(io.BytesIO(img_raw))
@ -215,16 +269,10 @@ async def run_skynet_telegram(
GROUP_ID, GROUP_ID,
caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']), caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
photo=img, photo=img,
reply_markup=build_redo_menu() reply_to_message_id=reply_id
) )
return return
await bot.reply_to(message, resp_txt)
@bot.message_handler(commands=['redo'])
async def redo_txt2img(message):
await _redo(message)
@bot.message_handler(commands=['config']) @bot.message_handler(commands=['config'])
async def set_config(message): async def set_config(message):
rpc_params = {} rpc_params = {}
@ -232,9 +280,10 @@ async def run_skynet_telegram(
attr, val, reply_txt = validate_user_config_request( attr, val, reply_txt = validate_user_config_request(
message.text) message.text)
resp = await _rpc_call( logging.info(f'user config update: {attr} to {val}')
message.from_user.id, await db_call('update_user_config',
'config', {'attr': attr, 'val': val}) user, req.params['attr'], req.params['val'])
logging.info('done')
except BaseException as e: except BaseException as e:
reply_txt = str(e) reply_txt = str(e)
@ -244,16 +293,12 @@ async def run_skynet_telegram(
@bot.message_handler(commands=['stats']) @bot.message_handler(commands=['stats'])
async def user_stats(message): async def user_stats(message):
resp = await _rpc_call(
message.from_user.id,
'stats',
{}
)
stats = resp.result
stats_str = f'generated: {stats["generated"]}\n' generated, joined, role = await db_call('get_user_stats', user)
stats_str += f'joined: {stats["joined"]}\n'
stats_str += f'role: {stats["role"]}\n' stats_str = f'generated: {generated}\n'
stats_str += f'joined: {joined}\n'
stats_str += f'role: {role}\n'
await bot.reply_to( await bot.reply_to(
message, stats_str) message, stats_str)
@ -289,4 +334,4 @@ async def run_skynet_telegram(
await _redo(call) await _redo(call)
await aio_as_trio(bot.infinity_polling()) await aio_as_trio(bot.infinity_polling)()

341
skynet/network.py 100644
View File

@ -0,0 +1,341 @@
#!/usr/bin/python
import zlib
import socket
from typing import Callable, Awaitable, Optional
from pathlib import Path
from contextlib import asynccontextmanager as acm
from cryptography import x509
from cryptography.hazmat.primitives import serialization
import trio
import pynng
from pynng import TLSConfig, Context
from .protobuf import *
from .constants import *
def get_random_port():
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
s.bind(('', 0))
return s.getsockname()[1]
def load_certs(
certs_dir: str,
cert_name: str,
key_name: str
):
certs_dir = Path(certs_dir).resolve()
tls_key_data = (certs_dir / key_name).read_bytes()
tls_key = serialization.load_pem_private_key(
tls_key_data,
password=None
)
tls_cert_data = (certs_dir / cert_name).read_bytes()
tls_cert = x509.load_pem_x509_certificate(
tls_cert_data
)
tls_whitelist = {}
for cert_path in (*(certs_dir / 'whitelist').glob('*.cert'), certs_dir / 'brain.cert'):
tls_whitelist[cert_path.stem] = x509.load_pem_x509_certificate(
cert_path.read_bytes()
)
return (
SessionTLSConfig(
TLSConfig.MODE_SERVER,
own_key_string=tls_key_data,
own_cert_string=tls_cert_data
),
tls_whitelist
)
def load_certs_client(
certs_dir: str,
cert_name: str,
key_name: str,
ca_name: Optional[str] = None
):
certs_dir = Path(certs_dir).resolve()
if not ca_name:
ca_name = 'brain.cert'
ca_cert_data = (certs_dir / ca_name).read_bytes()
tls_key_data = (certs_dir / key_name).read_bytes()
tls_cert_data = (certs_dir / cert_name).read_bytes()
tls_whitelist = {}
for cert_path in (*(certs_dir / 'whitelist').glob('*.cert'), certs_dir / 'brain.cert'):
tls_whitelist[cert_path.stem] = x509.load_pem_x509_certificate(
cert_path.read_bytes()
)
return (
SessionTLSConfig(
TLSConfig.MODE_CLIENT,
own_key_string=tls_key_data,
own_cert_string=tls_cert_data,
ca_string=ca_cert_data
),
tls_whitelist
)
class SessionError(BaseException):
...
class SessionTLSConfig(TLSConfig):
def __init__(
self,
mode,
server_name=None,
ca_string=None,
own_key_string=None,
own_cert_string=None,
auth_mode=None,
ca_files=None,
cert_key_file=None,
passwd=None
):
super().__init__(
mode,
server_name=server_name,
ca_string=ca_string,
own_key_string=own_key_string,
own_cert_string=own_cert_string,
auth_mode=auth_mode,
ca_files=ca_files,
cert_key_file=cert_key_file,
passwd=passwd
)
if ca_string:
self.ca_cert = x509.load_pem_x509_certificate(ca_string)
self.cert = x509.load_pem_x509_certificate(own_cert_string)
self.key = serialization.load_pem_private_key(
own_key_string,
password=passwd
)
class SessionServer:
def __init__(
self,
addr: str,
msg_handler: Callable[
[SkynetRPCRequest, Context], Awaitable[SkynetRPCResponse]
],
cert_name: Optional[str] = None,
key_name: Optional[str] = None,
cert_dir: str = DEFAULT_CERTS_DIR,
recv_max_size = 0
):
self.addr = addr
self.msg_handler = msg_handler
self.cert_name = cert_name
self.tls_config = None
self.tls_whitelist = None
if cert_name and key_name:
self.cert_name = cert_name
self.tls_config, self.tls_whitelist = load_certs(
cert_dir, cert_name, key_name)
self.addr = 'tls+' + self.addr
self.recv_max_size = recv_max_size
async def _handle_msg(self, req: SkynetRPCRequest, ctx: Context):
resp = await self.msg_handler(req, ctx)
if self.tls_config:
resp.auth.cert = 'skynet'
resp.auth.sig = sign_protobuf_msg(
resp, self.tls_config.key)
raw_msg = zlib.compress(resp.SerializeToString())
await ctx.asend(raw_msg)
ctx.close()
async def _listener (self, sock):
async with trio.open_nursery() as n:
while True:
ctx = sock.new_context()
raw_msg = await ctx.arecv()
raw_size = len(raw_msg)
logging.debug(f'rpc server new msg {raw_size} bytes')
try:
msg = zlib.decompress(raw_msg)
msg_size = len(msg)
except zlib.error:
logging.warning(f'Zlib decompress error, dropping msg of size {len(raw_msg)}')
continue
logging.debug(f'msg after decompress {msg_size} bytes, +{msg_size - raw_size} bytes')
req = SkynetRPCRequest()
try:
req.ParseFromString(msg)
except google.protobuf.message.DecodeError:
logging.warning(f'Dropping malfomed msg of size {len(msg)}')
continue
logging.debug(f'msg method: {req.method}')
if self.tls_config:
if req.auth.cert not in self.tls_whitelist:
logging.warning(
f'{req.auth.cert} not in tls whitelist')
continue
try:
verify_protobuf_msg(req, self.tls_whitelist[req.auth.cert])
except ValueError:
logging.warning(
f'{req.cert} sent an unauthenticated msg')
continue
n.start_soon(self._handle_msg, req, ctx)
@acm
async def open(self):
with pynng.Rep0(
recv_max_size=self.recv_max_size
) as sock:
if self.tls_config:
sock.tls_config = self.tls_config
sock.listen(self.addr)
logging.debug(f'server socket listening at {self.addr}')
async with trio.open_nursery() as n:
n.start_soon(self._listener, sock)
try:
yield self
finally:
n.cancel_scope.cancel()
logging.debug('server socket is off.')
class SessionClient:
def __init__(
self,
connect_addr: str,
uid: str,
cert_name: Optional[str] = None,
key_name: Optional[str] = None,
ca_name: Optional[str] = None,
cert_dir: str = DEFAULT_CERTS_DIR,
recv_max_size = 0
):
self.uid = uid
self.connect_addr = connect_addr
self.cert_name = None
self.tls_config = None
self.tls_whitelist = None
self.tls_cert = None
self.tls_key = None
if cert_name and key_name:
self.cert_name = Path(cert_name).stem
self.tls_config, self.tls_whitelist = load_certs_client(
cert_dir, cert_name, key_name, ca_name=ca_name)
if not self.connect_addr.startswith('tls'):
self.connect_addr = 'tls+' + self.connect_addr
self.recv_max_size = recv_max_size
self._connected = False
self._sock = None
def connect(self):
self._sock = pynng.Req0(
recv_max_size=0,
name=self.uid
)
if self.tls_config:
self._sock.tls_config = self.tls_config
logging.debug(f'client is dialing {self.connect_addr}...')
self._sock.dial(self.connect_addr, block=True)
self._connected = True
logging.debug(f'client is connected to {self.connect_addr}')
def disconnect(self):
self._sock.close()
self._connected = False
logging.debug(f'client disconnected.')
async def rpc(
self,
method: str,
params: dict = {},
binext: Optional[bytes] = None,
timeout: float = 2.
):
if not self._connected:
raise SessionError('tried to use rpc without connecting')
req = SkynetRPCRequest()
req.uid = self.uid
req.method = method
req.params.update(params)
if binext:
logging.debug('added binary extension')
req.bin = binext
if self.tls_config:
req.auth.cert = self.cert_name
req.auth.sig = sign_protobuf_msg(req, self.tls_config.key)
with trio.fail_after(timeout):
ctx = self._sock.new_context()
raw_req = zlib.compress(req.SerializeToString())
logging.debug(f'rpc client sending new msg {method} of size {len(raw_req)}')
await ctx.asend(raw_req)
logging.debug('sent, awaiting response...')
raw_resp = await ctx.arecv()
logging.debug(f'rpc client got response of size {len(raw_resp)}')
raw_resp = zlib.decompress(raw_resp)
resp = SkynetRPCResponse()
resp.ParseFromString(raw_resp)
ctx.close()
if self.tls_config:
verify_protobuf_msg(resp, self.tls_config.ca_cert)
return resp

25
skynet/protobuf/__init__.py
View File

@ -1,29 +1,4 @@
#!/usr/bin/python #!/usr/bin/python
from typing import Optional
from dataclasses import dataclass, asdict
from google.protobuf.json_format import MessageToDict
from .auth import * from .auth import *
from .skynet_pb2 import * from .skynet_pb2 import *
class Struct:
def to_dict(self):
return asdict(self)
@dataclass
class DiffusionParameters(Struct):
algo: str
prompt: str
step: int
width: int
height: int
guidance: float
strength: float
seed: Optional[int]
image: bool # if true indicates a bytestream is next msg
upscaler: Optional[str]

22
skynet/protobuf/auth.py
View File

@ -7,7 +7,8 @@ from hashlib import sha256
from collections import OrderedDict from collections import OrderedDict
from google.protobuf.json_format import MessageToDict from google.protobuf.json_format import MessageToDict
from OpenSSL.crypto import PKey, X509, verify, sign from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from .skynet_pb2 import * from .skynet_pb2 import *
@ -46,20 +47,23 @@ def serialize_msg_deterministic(msg):
if field_descriptor.message_type.name == 'Struct': if field_descriptor.message_type.name == 'Struct':
hash_dict(MessageToDict(getattr(msg, field_name))) hash_dict(MessageToDict(getattr(msg, field_name)))
deterministic_msg = shasum.hexdigest() deterministic_msg = shasum.digest()
return deterministic_msg return deterministic_msg
def sign_protobuf_msg(msg, key: PKey): def sign_protobuf_msg(msg, key):
return sign( return key.sign(
key, serialize_msg_deterministic(msg), 'sha256').hex() serialize_msg_deterministic(msg),
padding.PKCS1v15(),
hashes.SHA256()
).hex()
def verify_protobuf_msg(msg, cert: X509): def verify_protobuf_msg(msg, cert):
return verify( return cert.public_key().verify(
cert,
bytes.fromhex(msg.auth.sig), bytes.fromhex(msg.auth.sig),
serialize_msg_deterministic(msg), serialize_msg_deterministic(msg),
'sha256' padding.PKCS1v15(),
hashes.SHA256()
) )

14
skynet/protobuf/skynet.proto
View File

@ -13,18 +13,12 @@ message SkynetRPCRequest {
string uid = 1; string uid = 1;
string method = 2; string method = 2;
google.protobuf.Struct params = 3; google.protobuf.Struct params = 3;
optional Auth auth = 4; optional bytes bin = 4;
optional Auth auth = 5;
} }
message SkynetRPCResponse { message SkynetRPCResponse {
google.protobuf.Struct result = 1; google.protobuf.Struct result = 1;
optional Auth auth = 2; optional bytes bin = 2;
} optional Auth auth = 3;
message DGPUBusMessage {
string rid = 1;
string nid = 2;
string method = 3;
google.protobuf.Struct params = 4;
optional Auth auth = 5;
} }

10
skynet/protobuf/skynet_pb2.py
View File

@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0cskynet.proto\x12\x06skynet\x1a\x1cgoogle/protobuf/struct.proto\"!\n\x04\x41uth\x12\x0c\n\x04\x63\x65rt\x18\x01 \x01(\t\x12\x0b\n\x03sig\x18\x02 \x01(\t\"\x82\x01\n\x10SkynetRPCRequest\x12\x0b\n\x03uid\x18\x01 \x01(\t\x12\x0e\n\x06method\x18\x02 \x01(\t\x12\'\n\x06params\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x1f\n\x04\x61uth\x18\x04 \x01(\x0b\x32\x0c.skynet.AuthH\x00\x88\x01\x01\x42\x07\n\x05_auth\"f\n\x11SkynetRPCResponse\x12\'\n\x06result\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x1f\n\x04\x61uth\x18\x02 \x01(\x0b\x32\x0c.skynet.AuthH\x00\x88\x01\x01\x42\x07\n\x05_auth\"\x8d\x01\n\x0e\x44GPUBusMessage\x12\x0b\n\x03rid\x18\x01 \x01(\t\x12\x0b\n\x03nid\x18\x02 \x01(\t\x12\x0e\n\x06method\x18\x03 \x01(\t\x12\'\n\x06params\x18\x04 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x1f\n\x04\x61uth\x18\x05 \x01(\x0b\x32\x0c.skynet.AuthH\x00\x88\x01\x01\x42\x07\n\x05_authb\x06proto3') DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x0cskynet.proto\x12\x06skynet\x1a\x1cgoogle/protobuf/struct.proto\"!\n\x04\x41uth\x12\x0c\n\x04\x63\x65rt\x18\x01 \x01(\t\x12\x0b\n\x03sig\x18\x02 \x01(\t\"\x9c\x01\n\x10SkynetRPCRequest\x12\x0b\n\x03uid\x18\x01 \x01(\t\x12\x0e\n\x06method\x18\x02 \x01(\t\x12\'\n\x06params\x18\x03 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x10\n\x03\x62in\x18\x04 \x01(\x0cH\x00\x88\x01\x01\x12\x1f\n\x04\x61uth\x18\x05 \x01(\x0b\x32\x0c.skynet.AuthH\x01\x88\x01\x01\x42\x06\n\x04_binB\x07\n\x05_auth\"\x80\x01\n\x11SkynetRPCResponse\x12\'\n\x06result\x18\x01 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x10\n\x03\x62in\x18\x02 \x01(\x0cH\x00\x88\x01\x01\x12\x1f\n\x04\x61uth\x18\x03 \x01(\x0b\x32\x0c.skynet.AuthH\x01\x88\x01\x01\x42\x06\n\x04_binB\x07\n\x05_authb\x06proto3')
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'skynet_pb2', globals()) _builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'skynet_pb2', globals())
@ -24,9 +24,7 @@ if _descriptor._USE_C_DESCRIPTORS == False:
_AUTH._serialized_start=54 _AUTH._serialized_start=54
_AUTH._serialized_end=87 _AUTH._serialized_end=87
_SKYNETRPCREQUEST._serialized_start=90 _SKYNETRPCREQUEST._serialized_start=90
_SKYNETRPCREQUEST._serialized_end=220 _SKYNETRPCREQUEST._serialized_end=246
_SKYNETRPCRESPONSE._serialized_start=222 _SKYNETRPCRESPONSE._serialized_start=249
_SKYNETRPCRESPONSE._serialized_end=324 _SKYNETRPCRESPONSE._serialized_end=377
_DGPUBUSMESSAGE._serialized_start=327
_DGPUBUSMESSAGE._serialized_end=468
# @@protoc_insertion_point(module_scope) # @@protoc_insertion_point(module_scope)

5
skynet/utils.py
View File

@ -1,5 +1,6 @@
#!/usr/bin/python #!/usr/bin/python
import time
import random import random
from typing import Optional from typing import Optional
@ -21,6 +22,10 @@ from huggingface_hub import login
from .constants import ALGOS from .constants import ALGOS
def time_ms():
return int(time.time() * 1000)
def convert_from_cv2_to_image(img: np.ndarray) -> Image: def convert_from_cv2_to_image(img: np.ndarray) -> Image:
# return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) # return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
return Image.fromarray(img) return Image.fromarray(img)

81
tests/conftest.py
View File

@ -3,89 +3,30 @@
import os import os
import json import json
import time import time
import random
import string
import logging import logging
from functools import partial
from pathlib import Path from pathlib import Path
from functools import partial
import trio
import pytest import pytest
import psycopg2
import trio_asyncio
from docker.types import Mount, DeviceRequest from docker.types import Mount, DeviceRequest
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
from skynet.constants import * from skynet.db import open_new_database
from skynet.brain import run_skynet from skynet.brain import run_skynet
from skynet.network import get_random_port
from skynet.constants import *
@pytest.fixture(scope='session') @pytest.fixture(scope='session')
def postgres_db(dockerctl): def postgres_db(dockerctl):
rpassword = ''.join( with open_new_database() as db_params:
random.choice(string.ascii_lowercase) yield db_params
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 @pytest.fixture
async def skynet_running(postgres_db): async def skynet_running():
db_container, db_pass, db_host = postgres_db async with run_skynet():
async with run_skynet(
db_pass=db_pass,
db_host=db_host
):
yield yield
@ -99,11 +40,13 @@ def dgpu_workers(request, dockerctl, skynet_running):
cmds = [] cmds = []
for i in range(num_containers): for i in range(num_containers):
dgpu_addr = f'tcp://127.0.0.1:{get_random_port()}'
cmd = f''' cmd = f'''
pip install -e . && \ pip install -e . && \
skynet run dgpu \ skynet run dgpu \
--algos=\'{json.dumps(initial_algos)}\' \ --algos=\'{json.dumps(initial_algos)}\' \
--uid=dgpu-{i} --uid=dgpu-{i} \
--dgpu={dgpu_addr}
''' '''
cmds.append(['bash', '-c', cmd]) cmds.append(['bash', '-c', cmd])
@ -120,7 +63,7 @@ def dgpu_workers(request, dockerctl, skynet_running):
network='host', network='host',
mounts=mounts, mounts=mounts,
device_requests=devices, device_requests=devices,
num=num_containers num=num_containers,
) as containers: ) as containers:
yield containers yield containers

272
tests/test_dgpu.py
View File

@ -12,29 +12,26 @@ from functools import partial
import trio import trio
import pytest import pytest
import trio_asyncio
from PIL import Image from PIL import Image
from google.protobuf.json_format import MessageToDict from google.protobuf.json_format import MessageToDict
from skynet.brain import SkynetDGPUComputeError from skynet.brain import SkynetDGPUComputeError
from skynet.constants import * from skynet.network import get_random_port, SessionServer
from skynet.protobuf import SkynetRPCResponse
from skynet.frontend import open_skynet_rpc from skynet.frontend import open_skynet_rpc
from skynet.constants import *
async def wait_for_dgpus(rpc, amount: int, timeout: float = 30.0): async def wait_for_dgpus(session, amount: int, timeout: float = 30.0):
gpu_ready = False gpu_ready = False
start_time = time.time() with trio.fail_after(timeout):
current_time = time.time() while not gpu_ready:
while not gpu_ready and (current_time - start_time) < timeout: res = await session.rpc('dgpu_workers')
res = await rpc('dgpu_workers')
if res.result['ok'] >= amount: if res.result['ok'] >= amount:
break break
await trio.sleep(1) await trio.sleep(1)
current_time = time.time()
assert (current_time - start_time) < timeout
_images = set() _images = set()
@ -48,14 +45,15 @@ async def check_request_img(
): ):
global _images global _images
async with open_skynet_rpc( with open_skynet_rpc(
uid, uid,
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as rpc_call: res = await session.rpc(
res = await rpc_call( 'dgpu_call', {
'txt2img', { 'method': 'diffuse',
'params': {
'prompt': 'red old tractor in a sunny wheat field', 'prompt': 'red old tractor in a sunny wheat field',
'step': 28, 'step': 28,
'width': width, 'height': height, 'width': width, 'height': height,
@ -63,19 +61,17 @@ async def check_request_img(
'seed': None, 'seed': None,
'algo': list(ALGOS.keys())[i], 'algo': list(ALGOS.keys())[i],
'upscaler': upscaler 'upscaler': upscaler
}) }
},
timeout=60
)
if 'error' in res.result: if 'error' in res.result:
raise SkynetDGPUComputeError(MessageToDict(res.result)) raise SkynetDGPUComputeError(MessageToDict(res.result))
if upscaler == 'x4': img_raw = res.bin
width *= 4
height *= 4
img_raw = zlib.decompress(bytes.fromhex(res.result['img']))
img_sha = sha256(img_raw).hexdigest() img_sha = sha256(img_raw).hexdigest()
img = Image.frombytes( img = Image.open(io.BytesIO(img_raw))
'RGB', (width, height), img_raw)
if expect_unique and img_sha in _images: if expect_unique and img_sha in _images:
raise ValueError('Duplicated image sha: {img_sha}') raise ValueError('Duplicated image sha: {img_sha}')
@ -96,13 +92,12 @@ async def test_dgpu_worker_compute_error(dgpu_workers):
then generate a smaller image to show gpu worker recovery then generate a smaller image to show gpu worker recovery
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 1)
await wait_for_dgpus(test_rpc, 1)
with pytest.raises(SkynetDGPUComputeError) as e: with pytest.raises(SkynetDGPUComputeError) as e:
await check_request_img(0, width=4096, height=4096) await check_request_img(0, width=4096, height=4096)
@ -112,20 +107,35 @@ async def test_dgpu_worker_compute_error(dgpu_workers):
await check_request_img(0) await check_request_img(0)
@pytest.mark.parametrize(
'dgpu_workers', [(1, ['midj'])], indirect=True)
async def test_dgpu_worker(dgpu_workers):
'''Generate one image in a single dgpu worker
'''
with open_skynet_rpc(
'test-ctx',
cert_name='whitelist/testing.cert',
key_name='testing.key'
) as session:
await wait_for_dgpus(session, 1)
await check_request_img(0)
@pytest.mark.parametrize( @pytest.mark.parametrize(
'dgpu_workers', [(1, ['midj', 'stable'])], indirect=True) 'dgpu_workers', [(1, ['midj', 'stable'])], indirect=True)
async def test_dgpu_workers(dgpu_workers): async def test_dgpu_worker_two_models(dgpu_workers):
'''Generate two images in a single dgpu worker using '''Generate two images in a single dgpu worker using
two different models. two different models.
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 1)
await wait_for_dgpus(test_rpc, 1)
await check_request_img(0) await check_request_img(0)
await check_request_img(1) await check_request_img(1)
@ -138,14 +148,12 @@ async def test_dgpu_worker_upscale(dgpu_workers):
two different models. two different models.
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 1)
await wait_for_dgpus(test_rpc, 1)
logging.error('UPSCALE')
img = await check_request_img(0, upscaler='x4') img = await check_request_img(0, upscaler='x4')
@ -157,13 +165,12 @@ async def test_dgpu_worker_upscale(dgpu_workers):
async def test_dgpu_workers_two(dgpu_workers): async def test_dgpu_workers_two(dgpu_workers):
'''Generate two images in two separate dgpu workers '''Generate two images in two separate dgpu workers
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 2, timeout=60)
await wait_for_dgpus(test_rpc, 2)
async with trio.open_nursery() as n: async with trio.open_nursery() as n:
n.start_soon(check_request_img, 0) n.start_soon(check_request_img, 0)
@ -175,13 +182,12 @@ async def test_dgpu_workers_two(dgpu_workers):
async def test_dgpu_worker_algo_swap(dgpu_workers): async def test_dgpu_worker_algo_swap(dgpu_workers):
'''Generate an image using a non default model '''Generate an image using a non default model
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 1)
await wait_for_dgpus(test_rpc, 1)
await check_request_img(5) await check_request_img(5)
@ -191,33 +197,32 @@ async def test_dgpu_rotation_next_worker(dgpu_workers):
'''Connect three dgpu workers, disconnect and check next_worker '''Connect three dgpu workers, disconnect and check next_worker
rotation happens correctly rotation happens correctly
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 3)
await wait_for_dgpus(test_rpc, 3)
res = await test_rpc('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 0 assert res.result['ok'] == 0
await check_request_img(0) await check_request_img(0)
res = await test_rpc('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 1 assert res.result['ok'] == 1
await check_request_img(0) await check_request_img(0)
res = await test_rpc('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 2 assert res.result['ok'] == 2
await check_request_img(0) await check_request_img(0)
res = await test_rpc('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 0 assert res.result['ok'] == 0
@ -228,13 +233,12 @@ async def test_dgpu_rotation_next_worker_disconnect(dgpu_workers):
'''Connect three dgpu workers, disconnect the first one and check '''Connect three dgpu workers, disconnect the first one and check
next_worker rotation happens correctly next_worker rotation happens correctly
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 3)
await wait_for_dgpus(test_rpc, 3)
await trio.sleep(3) await trio.sleep(3)
@ -245,7 +249,7 @@ async def test_dgpu_rotation_next_worker_disconnect(dgpu_workers):
dgpu_workers[0].wait() dgpu_workers[0].wait()
res = await test_rpc('dgpu_workers') res = await session.rpc('dgpu_workers')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 2 assert res.result['ok'] == 2
@ -258,24 +262,41 @@ async def test_dgpu_no_ack_node_disconnect(skynet_running):
'''Mock a node that connects, gets a request but fails to '''Mock a node that connects, gets a request but fails to
acknowledge it, then check skynet correctly drops the node acknowledge it, then check skynet correctly drops the node
''' '''
async with open_skynet_rpc(
'test-ctx',
security=True,
cert_name='whitelist/testing',
key_name='testing'
) as rpc_call:
res = await rpc_call('dgpu_online') async def mock_rpc(req, ctx):
resp = SkynetRPCResponse()
resp.result.update({'error': 'can\'t do it mate'})
return resp
dgpu_addr = f'tcp://127.0.0.1:{get_random_port()}'
mock_server = SessionServer(
dgpu_addr,
mock_rpc,
cert_name='whitelist/testing.cert',
key_name='testing.key'
)
async with mock_server.open():
with open_skynet_rpc(
'test-ctx',
cert_name='whitelist/testing.cert',
key_name='testing.key'
) as session:
res = await session.rpc('dgpu_online', {
'dgpu_addr': dgpu_addr,
'cert': 'whitelist/testing.cert'
})
assert 'ok' in res.result assert 'ok' in res.result
await wait_for_dgpus(rpc_call, 1) await wait_for_dgpus(session, 1)
with pytest.raises(SkynetDGPUComputeError) as e: with pytest.raises(SkynetDGPUComputeError) as e:
await check_request_img(0) await check_request_img(0)
assert 'dgpu failed to acknowledge request' in str(e) assert 'can\'t do it mate' in str(e.value)
res = await rpc_call('dgpu_workers') res = await session.rpc('dgpu_workers')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 0 assert res.result['ok'] == 0
@ -286,13 +307,12 @@ async def test_dgpu_timeout_while_processing(dgpu_workers):
'''Stop node while processing request to cause timeout and '''Stop node while processing request to cause timeout and
then check skynet correctly drops the node. then check skynet correctly drops the node.
''' '''
async with open_skynet_rpc( with open_skynet_rpc(
'test-ctx', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as test_rpc: await wait_for_dgpus(session, 1)
await wait_for_dgpus(test_rpc, 1)
async def check_request_img_raises(): async def check_request_img_raises():
with pytest.raises(SkynetDGPUComputeError) as e: with pytest.raises(SkynetDGPUComputeError) as e:
@ -308,36 +328,23 @@ async def test_dgpu_timeout_while_processing(dgpu_workers):
assert ec == 0 assert ec == 0
@pytest.mark.parametrize(
'dgpu_workers', [(1, ['midj'])], indirect=True)
async def test_dgpu_heartbeat(dgpu_workers):
'''
'''
async with open_skynet_rpc(
'test-ctx',
security=True,
cert_name='whitelist/testing',
key_name='testing'
) as test_rpc:
await wait_for_dgpus(test_rpc, 1)
await trio.sleep(120)
@pytest.mark.parametrize( @pytest.mark.parametrize(
'dgpu_workers', [(1, ['midj'])], indirect=True) 'dgpu_workers', [(1, ['midj'])], indirect=True)
async def test_dgpu_img2img(dgpu_workers): async def test_dgpu_img2img(dgpu_workers):
async with open_skynet_rpc( with open_skynet_rpc(
'1', 'test-ctx',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as rpc_call: await wait_for_dgpus(session, 1)
await wait_for_dgpus(rpc_call, 1)
await trio.sleep(2)
res = await rpc_call( res = await session.rpc(
'txt2img', { 'dgpu_call', {
'method': 'diffuse',
'params': {
'prompt': 'red old tractor in a sunny wheat field', 'prompt': 'red old tractor in a sunny wheat field',
'step': 28, 'step': 28,
'width': 512, 'height': 512, 'width': 512, 'height': 512,
@ -345,35 +352,38 @@ async def test_dgpu_img2img(dgpu_workers):
'seed': None, 'seed': None,
'algo': list(ALGOS.keys())[0], 'algo': list(ALGOS.keys())[0],
'upscaler': None 'upscaler': None
}) }
},
timeout=60
)
if 'error' in res.result: if 'error' in res.result:
raise SkynetDGPUComputeError(MessageToDict(res.result)) raise SkynetDGPUComputeError(MessageToDict(res.result))
img_raw = res.result['img'] img_raw = res.bin
img = zlib.decompress(bytes.fromhex(img_raw)) img = Image.open(io.BytesIO(img_raw))
logging.info(img[:10])
img = Image.open(io.BytesIO(img))
img.save('txt2img.png') img.save('txt2img.png')
res = await rpc_call( res = await session.rpc(
'img2img', { 'dgpu_call', {
'prompt': 'red sports car in a sunny wheat field', 'method': 'diffuse',
'params': {
'prompt': 'red ferrari in a sunny wheat field',
'step': 28, 'step': 28,
'img': img_raw, 'guidance': 8,
'guidance': 12, 'strength': 0.7,
'seed': None, 'seed': None,
'algo': list(ALGOS.keys())[0], 'algo': list(ALGOS.keys())[0],
'upscaler': 'x4' 'upscaler': 'x4'
}) }
},
binext=img_raw,
timeout=60
)
if 'error' in res.result: if 'error' in res.result:
raise SkynetDGPUComputeError(MessageToDict(res.result)) raise SkynetDGPUComputeError(MessageToDict(res.result))
img_raw = res.result['img'] img_raw = res.bin
img = zlib.decompress(bytes.fromhex(img_raw)) img = Image.open(io.BytesIO(img_raw))
logging.info(img[:10])
img = Image.open(io.BytesIO(img))
img.save('img2img.png') img.save('img2img.png')

49
tests/test_skynet.py
View File

@ -9,6 +9,7 @@ import trio_asyncio
from skynet.brain import run_skynet from skynet.brain import run_skynet
from skynet.structs import * from skynet.structs import *
from skynet.network import SessionServer
from skynet.frontend import open_skynet_rpc from skynet.frontend import open_skynet_rpc
@ -17,54 +18,66 @@ async def test_skynet(skynet_running):
async def test_skynet_attempt_insecure(skynet_running): async def test_skynet_attempt_insecure(skynet_running):
with pytest.raises(pynng.exceptions.NNGException) as e: with pytest.raises(trio.TooSlowError) as e:
async with open_skynet_rpc('bad-actor'): with open_skynet_rpc('bad-actor') as session:
... with trio.fail_after(5):
await session.rpc('skynet_shutdown')
assert str(e.value) == 'Connection shutdown'
async def test_skynet_dgpu_connection_simple(skynet_running): async def test_skynet_dgpu_connection_simple(skynet_running):
async with open_skynet_rpc(
async def rpc_handler(req, ctx):
...
fake_dgpu_addr = 'tcp://127.0.0.1:41001'
rpc_server = SessionServer(
fake_dgpu_addr,
rpc_handler,
cert_name='whitelist/testing.cert',
key_name='testing.key'
)
with open_skynet_rpc(
'dgpu-0', 'dgpu-0',
security=True, cert_name='whitelist/testing.cert',
cert_name='whitelist/testing', key_name='testing.key'
key_name='testing' ) as session:
) as rpc_call:
# check 0 nodes are connected # check 0 nodes are connected
res = await rpc_call('dgpu_workers') res = await session.rpc('dgpu_workers')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 0 assert res.result['ok'] == 0
# check next worker is None # check next worker is None
res = await rpc_call('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == None assert res.result['ok'] == None
async with rpc_server.open() as rpc_server:
# connect 1 dgpu # connect 1 dgpu
res = await rpc_call('dgpu_online') res = await session.rpc(
'dgpu_online', {'dgpu_addr': fake_dgpu_addr})
assert 'ok' in res.result assert 'ok' in res.result
# check 1 node is connected # check 1 node is connected
res = await rpc_call('dgpu_workers') res = await session.rpc('dgpu_workers')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 1 assert res.result['ok'] == 1
# check next worker is 0 # check next worker is 0
res = await rpc_call('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 0 assert res.result['ok'] == 0
# disconnect 1 dgpu # disconnect 1 dgpu
res = await rpc_call('dgpu_offline') res = await session.rpc('dgpu_offline')
assert 'ok' in res.result assert 'ok' in res.result
# check 0 nodes are connected # check 0 nodes are connected
res = await rpc_call('dgpu_workers') res = await session.rpc('dgpu_workers')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == 0 assert res.result['ok'] == 0
# check next worker is None # check next worker is None
res = await rpc_call('dgpu_next') res = await session.rpc('dgpu_next')
assert 'ok' in res.result assert 'ok' in res.result
assert res.result['ok'] == None assert res.result['ok'] == None

28
tests/test_telegram.py 100644
View File

@ -0,0 +1,28 @@
#!/usr/bin/python
import trio
from functools import partial
from skynet.db import open_new_database
from skynet.brain import run_skynet
from skynet.config import load_skynet_ini
from skynet.frontend.telegram import run_skynet_telegram
if __name__ == '__main__':
'''You will need a telegram bot token configured on skynet.ini for this
'''
with open_new_database() as db_params:
db_container, db_pass, db_host = db_params
config = load_skynet_ini()
async def main():
await run_skynet_telegram(
'telegram-test',
config['skynet.telegram-test']['token'],
db_host=db_host,
db_pass=db_pass
)
trio.run(main)