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

.gitignore

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

Dockerfile.runtime+cuda

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

skynet.ini.example

@@ -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

skynet/brain.py

@@ -1,35 +1,24 @@
 #!/usr/bin/python
 
-import time
-import json
-import uuid
-import zlib
 import logging
-import traceback
 
-from uuid import UUID
-from pathlib import Path
-from functools import partial
 from contextlib import asynccontextmanager as acm
 from collections import OrderedDict
 
 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 .protobuf import *
 
 
+class SkynetRPCBadRequest(BaseException):
+    ...
+
 class SkynetDGPUOffline(BaseException):
     ...
 
@@ -44,39 +33,71 @@ class SkynetShutdownRequested(BaseException):
 
 
 @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()
-    wip_reqs = {}
-    fin_reqs = {}
     heartbeats = {}
     next_worker: Optional[int] = None
-    security = len(tls_whitelist) > 0
 
-    def connect_node(uid):
+    def connect_node(req: SkynetRPCRequest):
         nonlocal next_worker
-        nodes[uid] = {
-            'task': None
-        }
-        logging.info(f'dgpu online: {uid}')
 
-        if not next_worker:
+        node_params = MessageToDict(req.params)
-            next_worker = 0
+        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
+            }
+            node.update(node_params)
+
+            nodes[req.uid] = node
+            logging.info(f'DGPU node online: {req.uid}')
+
+            if not next_worker:
+                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):
         nonlocal next_worker
         if uid not in nodes:
+            logging.warning(f'Attempt to disconnect unknown node {uid}')
             return
+
         i = list(nodes.keys()).index(uid)
+        nodes[uid]['session'].disconnect()
         del nodes[uid]
 
         if i < next_worker:
             next_worker -= 1
 
+        logging.warning(f'DGPU node offline: {uid}')
+
         if len(nodes) == 0:
-            logging.info('nw: None')
+            logging.info('All nodes disconnected.')
             next_worker = None
 
-        logging.warning(f'dgpu offline: {uid}')
 
     def is_worker_busy(nid: str):
         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():
         nonlocal next_worker
-        logging.info('get next_worker called')
-        logging.info(f'pre next_worker: {next_worker}')
 
         if next_worker == None:
             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):
             next_worker = 0
 
-        logging.info(f'post next_worker: {next_worker}')
-
         return nid
 
-    async def dgpu_heartbeat_service():
+    async def rpc_handler(req: SkynetRPCRequest, ctx: Context):
-        nonlocal heartbeats
+        result = {'ok': {}}
-        while True:
-            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():
-        nonlocal wip_reqs, fin_reqs, heartbeats
-        while True:
-            raw_msg = await dgpu_bus.arecv()
-            logging.info(f'streamer got {len(raw_msg)} bytes.')
-            msg = DGPUBusMessage()
-            msg.ParseFromString(raw_msg)
-
-            if security:
-                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()
-        idx = list(nodes.keys()).index(nid)
-        logging.info(f'dgpu_stream_one_img {idx}/{len(nodes)} {nid}')
-        rid = uuid.uuid4().hex
-        ack_event = trio.Event()
-        img_event = trio.Event()
-        wip_reqs[rid] = ack_event
-
-        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)
-            raise SkynetDGPUOffline('dgpu failed to acknowledge request')
-
-        ack_msg = fin_reqs[rid]
-        if 'ack' not in ack_msg.params:
-            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:
-            async with db_pool.acquire() as conn:
-                user = await get_or_create_user(conn, req.uid)
-
-                result = {}
-
-                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 = {
-                                'id': rid,
-                                'img': img.hex(),
-                                'meta': meta
-                            }
-                            await update_user_stats(conn, user)
-                            logging.info('updated user stats.')
-
-                        else:
-                            result = {
-                                'error': 'skynet_no_last_prompt',
-                                'message': 'No prompt to redo, do txt2img first'
-                            }
-
-                    case 'config':
-                        logging.info('config')
-                        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 = {}
 
+        try:
             match req.method:
-                case 'skynet_shutdown':
-                    raise SkynetShutdownRequested
-
                 case 'dgpu_online':
-                    connect_node(req.uid)
+                    connect_node(req)
+
+                case 'dgpu_call':
+                    nid = get_next_worker()
+                    idx = list(nodes.keys()).index(nid)
+                    node = nodes[nid]
+                    logging.info(f'dgpu_call {idx}/{len(nodes)} {nid} @ {node["dgpu_addr"]}')
+                    dgpu_time = await node['session'].rpc('dgpu_time')
+                    if 'ok' not in dgpu_time.result:
+                        status = MessageToDict(dgpu_time.result)
+                        logging.warning(json.dumps(status, indent=4))
+                        disconnect_node(nid)
+                        raise SkynetDGPUComputeError(status['error'])
+
+                    dgpu_time = dgpu_time.result['ok']
+                    logging.info(f'ping to {nid}: {time_ms() - dgpu_time} ms')
+
+                    try:
+                        dgpu_result = await node['session'].rpc(
+                            timeout=45,  # give this 45 sec to run cause its compute
+                            binext=req.bin,
+                            **req.params
+                        )
+                        result = MessageToDict(dgpu_result.result)
+
+                        if dgpu_result.bin:
+                            resp.bin = dgpu_result.bin
+
+                    except trio.TooSlowError:
+                        result = {'error': 'timeout while processing request'}
 
                 case 'dgpu_offline':
                     disconnect_node(req.uid)
 
                 case 'dgpu_workers':
-                    result = len(nodes)
+                    result = {'ok': len(nodes)}
 
                 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 _:
-                    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
-        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.')

skynet/cli.py

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

skynet/config.py

@@ -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

skynet/constants.py

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

skynet/db/__init__.py

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

skynet/db/functions.py

@@ -1,18 +1,21 @@
 #!/usr/bin/python
 
+import time
+import random
+import string
 import logging
 
 from typing import Optional
 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 psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
 
-from .constants import *
+from ..constants import *
 
 
 DB_INIT_SQL = '''
@@ -75,29 +78,67 @@ def try_decode_uid(uid: str):
         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):

skynet/db/proxy.py

@@ -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

skynet/dgpu.py

@@ -2,29 +2,17 @@
 
 import gc
 import io
-import trio
 import json
-import uuid
-import time
-import zlib
 import random
 import logging
-import traceback
 
 from PIL import Image
 from typing import List, Optional
-from pathlib import Path
-from contextlib import ExitStack
 
-import pynng
+import trio
 import torch
 
-from pynng import TLSConfig
+from pynng import Context
-from OpenSSL.crypto import (
-    load_privatekey,
-    load_certificate,
-    FILETYPE_PEM
-)
 from diffusers import (
     StableDiffusionPipeline,
     StableDiffusionImg2ImgPipeline,
@@ -34,12 +22,9 @@ from realesrgan import RealESRGANer
 from basicsr.archs.rrdbnet_arch import RRDBNet
 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 .frontend import open_skynet_rpc
 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(
     cert_name: str,
     unique_id: str,
     key_name: Optional[str],
     rpc_address: str = DEFAULT_RPC_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!')
-
-    name = uuid.uuid4()
-
     logging.info(f'loading models...')
 
     upscaler = init_upscaler()
@@ -141,241 +77,140 @@ async def open_dgpu_node(
     logging.info('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:
-        if algo not in models:
+            case 'diffuse':
-            least_used = list(models.keys())[0]
+                image = None
-            for model in models:
+                algo = params['algo']
-                if models[least_used]['generated'] > models[model]['generated']:
+                if binext:
-                    least_used = model
+                    algo += 'img'
+                    image = Image.open(io.BytesIO(binext))
+                    w, h = image.size
+                    logging.info(f'user sent img of size {image.size}')
 
-            del models[least_used]
+                    if w > 512 or h > 512:
-            gc.collect()
+                        image.thumbnail((512, 512))
+                        logging.info(f'resized it to {image.size}')
 
-            models[algo] = {
+                if algo not in models:
-                'pipe': pipeline_for(ireq.algo, image=True if image else False),
+                    logging.info(f'{algo} not in loaded models, swapping...')
-                'generated': 0
+                    least_used = list(models.keys())[0]
-            }
+                    for model in models:
+                        if models[least_used]['generated'] > models[model]['generated']:
+                            least_used = model
 
-        _params = {}
+                    del models[least_used]
-        if ireq.image:
+                    gc.collect()
-            _params['image'] = image
-            _params['strength'] = ireq.strength
 
-        else:
+                    models[algo] = {
-            _params['width'] = int(ireq.width)
+                        'pipe': pipeline_for(params['algo'], image=True if binext else False),
-            _params['height'] = int(ireq.height)
+                        'generated': 0
+                    }
+                    logging.info(f'swapping done.')
 
-        try:
+                _params = {}
-            image = models[algo]['pipe'](
+                logging.info(method)
-                ireq.prompt,
+                logging.info(json.dumps(params, indent=4))
-                **_params,
+                logging.info(f'binext: {len(binext) if binext else 0} bytes')
-                guidance_scale=ireq.guidance,
+                if binext:
-                num_inference_steps=int(ireq.step),
+                    _params['image'] = image
-                generator=torch.Generator("cuda").manual_seed(ireq.seed)
+                    _params['strength'] = params['strength']
-            ).images[0]
 
-            if ireq.upscaler == 'x4':
+                else:
-                logging.info(f'size: {len(image.tobytes())}')
+                    _params['width'] = int(params['width'])
-                logging.info('performing upscale...')
+                    _params['height'] = int(params['height'])
-                input_img = image.convert('RGB')
-                up_img, _ = upscaler.enhance(
-                    convert_from_image_to_cv2(input_img), outscale=4)
 
-                image = convert_from_cv2_to_image(up_img)
+                try:
-                logging.info('done')
+                    image = models[algo]['pipe'](
+                        params['prompt'],
+                        **_params,
+                        guidance_scale=params['guidance'],
+                        num_inference_steps=int(params['step']),
+                        generator=torch.Generator("cuda").manual_seed(
+                            int(params['seed']) if params['seed'] else random.randint(0, 2 ** 64)
+                        )
+                    ).images[0]
 
-            img_byte_arr = io.BytesIO()
+                    if params['upscaler'] == 'x4':
-            image.save(img_byte_arr, format='PNG')
+                        logging.info(f'size: {len(image.tobytes())}')
-            raw_img = img_byte_arr.getvalue()
+                        logging.info('performing upscale...')
-            logging.info(f'final img size {len(raw_img)} bytes.')
+                        input_img = image.convert('RGB')
+                        up_img, _ = upscaler.enhance(
+                            convert_from_image_to_cv2(input_img), outscale=4)
 
-            return raw_img
+                        image = convert_from_cv2_to_image(up_img)
+                        logging.info('done')
 
-        except BaseException as e:
+                    img_byte_arr = io.BytesIO()
-            logging.error(e)
+                    image.save(img_byte_arr, format='PNG')
-            raise DGPUComputeError(str(e))
+                    raw_img = img_byte_arr.getvalue()
+                    logging.info(f'final img size {len(raw_img)} bytes.')
 
-        finally:
+                    return raw_img
-            torch.cuda.empty_cache()
+
+                except BaseException as e:
+                    logging.error(e)
+                    raise DGPUComputeError(str(e))
+
+                finally:
+                    torch.cuda.empty_cache()
+
+            case _:
+                raise DGPUComputeError('Unsupported compute method')
+
+    async def rpc_handler(req: SkynetRPCRequest, ctx: Context):
+        result = {}
+        resp = SkynetRPCResponse()
+
+        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,
+        key_name=key_name
+    )
+    skynet_rpc = SessionClient(
+        rpc_address,
+        unique_id,
+        cert_name=cert_name,
+        key_name=key_name
+    )
+    skynet_rpc.connect()
 
 
-    async with (
+    async with rpc_server.open() as rpc_server:
-        open_skynet_rpc(
+        res = await skynet_rpc.rpc(
-            unique_id,
+            'dgpu_online', {
-            rpc_address=rpc_address,
+                'dgpu_addr': rpc_server.addr,
-            security=security,
+                'cert': cert_name
-            cert_name=cert_name,
+            })
-            key_name=key_name
-        ) as rpc_call,
-        trio.open_nursery() as n
-    ):
 
-        tls_config = None
-        if security:
-            # load tls certs
-            if not key_name:
-                key_name = cert_name
-
-            certs_dir = Path(DEFAULT_CERTS_DIR).resolve()
-
-            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
 
         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:
             logging.info('interrupt caught, stopping...')
-            n.cancel_scope.cancel()
-            dgpu_bus.close()
 
         finally:
-            res = await rpc_call('dgpu_offline')
+            res = await skynet_rpc.rpc('dgpu_offline')
             assert 'ok' in res.result

skynet/frontend/__init__.py

@@ -4,7 +4,7 @@ import json
 
 from typing import Union, Optional
 from pathlib import Path
-from contextlib import asynccontextmanager as acm
+from contextlib import contextmanager as cm
 
 import pynng
 
@@ -17,6 +17,7 @@ from OpenSSL.crypto import (
 
 from google.protobuf.struct_pb2 import Struct
 
+from ..network import SessionClient
 from ..constants 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,
     rpc_address: str = DEFAULT_RPC_ADDR,
-    security: bool = False,
     cert_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):
     params = req.split(' ')

skynet/frontend/telegram.py

@@ -6,8 +6,6 @@ import logging
 
 from datetime import datetime
 
-import pynng
-
 from PIL import Image
 from trio_asyncio import aio_as_trio
 
@@ -16,6 +14,7 @@ from telebot.types import (
 )
 from telebot.async_telebot import AsyncTeleBot
 
+from ..db import open_database_connection
 from ..constants import *
 
 from . import *
@@ -56,228 +55,274 @@ def prepare_metainfo_caption(tguser, meta: dict) -> str:
 
 
 async def run_skynet_telegram(
+    name: 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)
     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
-        rpc_address=rpc_address,
+    ) as db_call:
-        security=True,
+        with open_skynet_rpc(
-        cert_name=cert_name,
+            f'skynet-telegram-{name}',
-        key_name=key_name
+            rpc_address=rpc_address,
-    ) as rpc_call:
+            cert_name=cert_name,
+            key_name=key_name
+        ) as session:
 
-        async def _rpc_call(
+            @bot.message_handler(commands=['help'])
-            uid: int,
+            async def send_help(message):
-            method: str,
+                splt_msg = message.text.split(' ')
-            params: dict = {}
-        ):
-            return await rpc_call(
-                method, params, uid=f'{PREFIX}+{uid}')
 
-        @bot.message_handler(commands=['help'])
+                if len(splt_msg) == 1:
-        async def send_help(message):
+                    await bot.reply_to(message, HELP_TEXT)
-            splt_msg = message.text.split(' ')
-
-            if len(splt_msg) == 1:
-                await bot.reply_to(message, HELP_TEXT)
-
-            else:
-                param = splt_msg[1]
-                if param in HELP_TOPICS:
-                    await bot.reply_to(message, HELP_TOPICS[param])
 
                 else:
-                    await bot.reply_to(message, HELP_UNKWNOWN_PARAM)
+                    param = splt_msg[1]
+                    if param in HELP_TOPICS:
+                        await bot.reply_to(message, HELP_TOPICS[param])
 
-        @bot.message_handler(commands=['cool'])
+                    else:
-        async def send_cool_words(message):
+                        await bot.reply_to(message, HELP_UNKWNOWN_PARAM)
-            await bot.reply_to(message, '\n'.join(COOL_WORDS))
 
-        @bot.message_handler(commands=['txt2img'])
+            @bot.message_handler(commands=['cool'])
-        async def send_txt2img(message):
+            async def send_cool_words(message):
-            chat = message.chat
+                await bot.reply_to(message, '\n'.join(COOL_WORDS))
 
-            prompt = ' '.join(message.text.split(' ')[1:])
+            @bot.message_handler(commands=['txt2img'])
+            async def send_txt2img(message):
+                chat = message.chat
+                reply_id = None
+                if chat.type == 'group' and chat.id == GROUP_ID:
+                    reply_id = message.message_id
 
-            if len(prompt) == 0:
+                user_id = f'tg+{message.from_user.id}'
-                await bot.reply_to(message, 'Empty text prompt ignored.')
-                return
 
-            logging.info(f'mid: {message.id}')
+                prompt = ' '.join(message.text.split(' ')[1:])
-            resp = await _rpc_call(
-                message.from_user.id,
-                'txt2img',
-                {'prompt': prompt}
-            )
-            logging.info(f'resp to {message.id} arrived')
 
-            resp_txt = ''
+                if len(prompt) == 0:
-            result = MessageToDict(resp.result)
+                    await bot.reply_to(message, 'Empty text prompt ignored.')
-            if 'error' in resp.result:
+                    return
-                resp_txt = resp.result['message']
 
-            else:
+                logging.info(f'mid: {message.id}')
-                logging.info(result['id'])
+                user = await db_call('get_or_create_user', user_id)
-                img_raw = zlib.decompress(bytes.fromhex(result['img']))
+                user_config = {**(await db_call('get_user_config', user))}
-                logging.info(f'got image of size: {len(img_raw)}')
+                del user_config['id']
-                img = Image.open(io.BytesIO(img_raw))
 
-                await bot.send_photo(
+                resp = await session.rpc(
-                    GROUP_ID,
+                    'dgpu_call', {
-                    caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
+                        'method': 'diffuse',
-                    photo=img,
+                        'params': {
-                    reply_markup=build_redo_menu()
+                            'prompt': prompt,
+                            **user_config
+                        }
+                    },
+                    timeout=60
                 )
-                return
+                logging.info(f'resp to {message.id} arrived')
 
-            await bot.reply_to(message, resp_txt)
+                resp_txt = ''
+                result = MessageToDict(resp.result)
+                if 'error' in resp.result:
+                    resp_txt = resp.result['message']
+                    await bot.reply_to(message, resp_txt)
 
-        @bot.message_handler(func=lambda message: True, content_types=['photo'])
+                else:
-        async def send_img2img(message):
+                    logging.info(result['id'])
-            chat = message.chat
+                    img_raw = resp.bin
+                    logging.info(f'got image of size: {len(img_raw)}')
+                    img = Image.open(io.BytesIO(img_raw))
 
-            if not message.caption.startswith('/img2img'):
+                    await bot.send_photo(
-                return
+                        GROUP_ID,
+                        caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
+                        photo=img,
+                        reply_to_message_id=reply_id,
+                        reply_markup=build_redo_menu()
+                    )
+                    return
 
-            prompt = ' '.join(message.caption.split(' ')[1:])
 
-            if len(prompt) == 0:
+            @bot.message_handler(func=lambda message: True, content_types=['photo'])
-                await bot.reply_to(message, 'Empty text prompt ignored.')
+            async def send_img2img(message):
-                return
+                chat = message.chat
+                reply_id = None
+                if chat.type == 'group' and chat.id == GROUP_ID:
+                    reply_id = message.message_id
 
-            file_id = message.photo[-1].file_id
+                user_id = f'tg+{message.from_user.id}'
-            file_path = (await bot.get_file(file_id)).file_path
-            file_raw = await bot.download_file(file_path)
-            img = zlib.compress(file_raw)
 
-            logging.info(f'mid: {message.id}')
+                if not message.caption.startswith('/img2img'):
-            resp = await _rpc_call(
+                    await bot.reply_to(
-                message.from_user.id,
+                        message,
-                'img2img',
+                        'For image to image you need to add /img2img to the beggining of your caption'
-                {'prompt': prompt, 'img': img.hex()}
+                    )
-            )
+                    return
-            logging.info(f'resp to {message.id} arrived')
 
-            resp_txt = ''
+                prompt = ' '.join(message.caption.split(' ')[1:])
-            result = MessageToDict(resp.result)
-            if 'error' in resp.result:
-                resp_txt = resp.result['message']
 
-            else:
+                if len(prompt) == 0:
-                logging.info(result['id'])
+                    await bot.reply_to(message, 'Empty text prompt ignored.')
-                img_raw = zlib.decompress(bytes.fromhex(result['img']))
+                    return
-                logging.info(f'got image of size: {len(img_raw)}')
-                img = Image.open(io.BytesIO(img_raw))
 
-                await bot.send_media_group(
+                file_id = message.photo[-1].file_id
-                    GROUP_ID,
+                file_path = (await bot.get_file(file_id)).file_path
-                    media=[
+                file_raw = await bot.download_file(file_path)
-                        InputMediaPhoto(file_id),
+
-                        InputMediaPhoto(
+                logging.info(f'mid: {message.id}')
-                            img,
+
-                            caption=prepare_metainfo_caption(message.from_user, result['meta']['meta'])
+                user = await db_call('get_or_create_user', user_id)
-                        )
+                user_config = {**(await db_call('get_user_config', user))}
-                    ]
+                del user_config['id']
+
+                resp = await session.rpc(
+                    'dgpu_call', {
+                        'method': 'diffuse',
+                        'params': {
+                            'prompt': prompt,
+                            **user_config
+                        }
+                    },
+                    binext=file_raw,
+                    timeout=60
                 )
-                return
+                logging.info(f'resp to {message.id} arrived')
 
-            await bot.reply_to(message, resp_txt)
+                resp_txt = ''
+                result = MessageToDict(resp.result)
+                if 'error' in resp.result:
+                    resp_txt = resp.result['message']
+                    await bot.reply_to(message, resp_txt)
 
-        @bot.message_handler(commands=['img2img'])
+                else:
-        async def redo_txt2img(message):
+                    logging.info(result['id'])
-            await bot.reply_to(
+                    img_raw = resp.bin
-                message,
+                    logging.info(f'got image of size: {len(img_raw)}')
-                'seems you tried to do an img2img command without sending image'
+                    img = Image.open(io.BytesIO(img_raw))
-            )
 
-        async def _redo(message):
+                    await bot.send_media_group(
-            resp = await _rpc_call(message.from_user.id, 'redo')
+                        GROUP_ID,
+                        media=[
+                            InputMediaPhoto(file_id),
+                            InputMediaPhoto(
+                                img,
+                                caption=prepare_metainfo_caption(message.from_user, result['meta']['meta'])
+                            )
+                        ],
+                        reply_to_message_id=reply_id
+                    )
+                    return
 
-            resp_txt = ''
-            result = MessageToDict(resp.result)
-            if 'error' in resp.result:
-                resp_txt = resp.result['message']
 
-            else:
+            @bot.message_handler(commands=['img2img'])
-                logging.info(result['id'])
+            async def img2img_missing_image(message):
-                img_raw = zlib.decompress(bytes.fromhex(result['img']))
+                await bot.reply_to(
-                logging.info(f'got image of size: {len(img_raw)}')
+                    message,
-                img = Image.open(io.BytesIO(img_raw))
+                    'seems you tried to do an img2img command without sending image'
-
-                await bot.send_photo(
-                    GROUP_ID,
-                    caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
-                    photo=img,
-                    reply_markup=build_redo_menu()
                 )
-                return
 
-            await bot.reply_to(message, resp_txt)
+            @bot.message_handler(commands=['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
 
-        @bot.message_handler(commands=['redo'])
+                user_config = {**(await db_call('get_user_config', user))}
-        async def redo_txt2img(message):
+                del user_config['id']
-            await _redo(message)
+                prompt = await db_call('get_last_prompt_of', user)
 
-        @bot.message_handler(commands=['config'])
+                resp = await session.rpc(
-        async def set_config(message):
+                    'dgpu_call', {
-            rpc_params = {}
+                        'method': 'diffuse',
-            try:
+                        'params': {
-                attr, val, reply_txt = validate_user_config_request(
+                            'prompt': prompt,
-                    message.text)
+                            **user_config
+                        }
+                    },
+                    timeout=60
+                )
+                logging.info(f'resp to {message.id} arrived')
 
-                resp = await _rpc_call(
+                resp_txt = ''
-                    message.from_user.id,
+                result = MessageToDict(resp.result)
-                    'config', {'attr': attr, 'val': val})
+                if 'error' in resp.result:
+                    resp_txt = resp.result['message']
+                    await bot.reply_to(message, resp_txt)
 
-            except BaseException as e:
+                else:
-                reply_txt = str(e)
+                    logging.info(result['id'])
+                    img_raw = resp.bin
+                    logging.info(f'got image of size: {len(img_raw)}')
+                    img = Image.open(io.BytesIO(img_raw))
 
-            finally:
+                    await bot.send_photo(
-                await bot.reply_to(message, reply_txt)
+                        GROUP_ID,
+                        caption=prepare_metainfo_caption(message.from_user, result['meta']['meta']),
+                        photo=img,
+                        reply_to_message_id=reply_id
+                    )
+                    return
 
-        @bot.message_handler(commands=['stats'])
+            @bot.message_handler(commands=['config'])
-        async def user_stats(message):
+            async def set_config(message):
-            resp = await _rpc_call(
+                rpc_params = {}
-                message.from_user.id,
+                try:
-                'stats',
+                    attr, val, reply_txt = validate_user_config_request(
-                {}
+                        message.text)
-            )
-            stats = resp.result
 
-            stats_str = f'generated: {stats["generated"]}\n'
+                    logging.info(f'user config update: {attr} to {val}')
-            stats_str += f'joined: {stats["joined"]}\n'
+                    await db_call('update_user_config',
-            stats_str += f'role: {stats["role"]}\n'
+                        user, req.params['attr'], req.params['val'])
+                    logging.info('done')
 
-            await bot.reply_to(
+                except BaseException as e:
-                message, stats_str)
+                    reply_txt = str(e)
 
-        @bot.message_handler(commands=['donate'])
+                finally:
-        async def donation_info(message):
+                    await bot.reply_to(message, reply_txt)
-            await bot.reply_to(
-                message, DONATION_INFO)
 
-        @bot.message_handler(commands=['say'])
+            @bot.message_handler(commands=['stats'])
-        async def say(message):
+            async def user_stats(message):
-            chat = message.chat
-            user = message.from_user
 
-            if (chat.type == 'group') or (user.id != 383385940):
+                generated, joined, role = await db_call('get_user_stats', user)
-                return
 
-            await bot.send_message(GROUP_ID, message.text[4:])
+                stats_str = f'generated: {generated}\n'
+                stats_str += f'joined: {joined}\n'
+                stats_str += f'role: {role}\n'
+
+                await bot.reply_to(
+                    message, stats_str)
+
+            @bot.message_handler(commands=['donate'])
+            async def donation_info(message):
+                await bot.reply_to(
+                    message, DONATION_INFO)
+
+            @bot.message_handler(commands=['say'])
+            async def say(message):
+                chat = message.chat
+                user = message.from_user
+
+                if (chat.type == 'group') or (user.id != 383385940):
+                    return
+
+                await bot.send_message(GROUP_ID, message.text[4:])
 
 
-        @bot.message_handler(func=lambda message: True)
+            @bot.message_handler(func=lambda message: True)
-        async def echo_message(message):
+            async def echo_message(message):
-            if message.text[0] == '/':
+                if message.text[0] == '/':
-                await bot.reply_to(message, UNKNOWN_CMD_TEXT)
+                    await bot.reply_to(message, UNKNOWN_CMD_TEXT)
 
         @bot.callback_query_handler(func=lambda call: True)
         async def callback_query(call):
@@ -289,4 +334,4 @@ async def run_skynet_telegram(
                     await _redo(call)
 
 
-        await aio_as_trio(bot.infinity_polling())
+            await aio_as_trio(bot.infinity_polling)()

skynet/network.py

@@ -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

skynet/protobuf/__init__.py

@@ -1,29 +1,4 @@
 #!/usr/bin/python
 
-from typing import Optional
-from dataclasses import dataclass, asdict
-
-from google.protobuf.json_format import MessageToDict
-
 from .auth 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]

skynet/protobuf/auth.py

@@ -7,7 +7,8 @@ from hashlib import sha256
 from collections import OrderedDict
 
 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 *
 
@@ -46,20 +47,23 @@ def serialize_msg_deterministic(msg):
         if field_descriptor.message_type.name == 'Struct':
             hash_dict(MessageToDict(getattr(msg, field_name)))
 
-    deterministic_msg = shasum.hexdigest()
+    deterministic_msg = shasum.digest()
 
     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),
         serialize_msg_deterministic(msg),
-        'sha256'
+        padding.PKCS1v15(),
+        hashes.SHA256()
     )

skynet/protobuf/skynet.proto

@@ -13,18 +13,12 @@ message SkynetRPCRequest {
     string uid = 1;
     string method = 2;
     google.protobuf.Struct params = 3;
-    optional Auth auth = 4;
+    optional bytes bin = 4;
+    optional Auth auth = 5;
 }
 
 message SkynetRPCResponse {
     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;
 }

skynet/protobuf/skynet_pb2.py

@@ -14,7 +14,7 @@ _sym_db = _symbol_database.Default()
 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.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'skynet_pb2', globals())
@@ -24,9 +24,7 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _AUTH._serialized_start=54
   _AUTH._serialized_end=87
   _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)

skynet/utils.py

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

tests/conftest.py

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

tests/test_dgpu.py

@@ -12,29 +12,26 @@ from functools import partial
 
 import trio
 import pytest
-import trio_asyncio
 
 from PIL import Image
 from google.protobuf.json_format import MessageToDict
 
 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.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
-    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()
@@ -48,34 +45,33 @@ async def check_request_img(
 ):
     global _images
 
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         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',
-                'prompt': 'red old tractor in a sunny wheat field',
+                'params': {
-                'step': 28,
+                    'prompt': 'red old tractor in a sunny wheat field',
-                'width': width, 'height': height,
+                    'step': 28,
-                'guidance': 7.5,
+                    'width': width, 'height': height,
-                'seed': None,
+                    'guidance': 7.5,
-                'algo': list(ALGOS.keys())[i],
+                    'seed': None,
-                'upscaler': upscaler
+                    'algo': list(ALGOS.keys())[i],
-            })
+                    'upscaler': upscaler
+                }
+            },
+            timeout=60
+        )
 
         if 'error' in 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 = Image.frombytes(
+        img = Image.open(io.BytesIO(img_raw))
-            'RGB', (width, height), img_raw)
 
         if expect_unique and img_sha in _images:
             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
     '''
 
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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:
             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)
 
 
+@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(
     '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
     two different models.
     '''
 
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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(1)
@@ -138,14 +148,12 @@ async def test_dgpu_worker_upscale(dgpu_workers):
     two different models.
     '''
 
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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')
 
@@ -157,13 +165,12 @@ async def test_dgpu_worker_upscale(dgpu_workers):
 async def test_dgpu_workers_two(dgpu_workers):
     '''Generate two images in two separate dgpu workers
     '''
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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:
             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):
     '''Generate an image using a non default model
     '''
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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)
 
 
@@ -191,33 +197,32 @@ async def test_dgpu_rotation_next_worker(dgpu_workers):
     '''Connect three dgpu workers, disconnect and check next_worker
     rotation happens correctly
     '''
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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 res.result['ok'] == 0
 
         await check_request_img(0)
 
-        res = await test_rpc('dgpu_next')
+        res = await session.rpc('dgpu_next')
         assert 'ok' in res.result
         assert res.result['ok'] == 1
 
         await check_request_img(0)
 
-        res = await test_rpc('dgpu_next')
+        res = await session.rpc('dgpu_next')
         assert 'ok' in res.result
         assert res.result['ok'] == 2
 
         await check_request_img(0)
 
-        res = await test_rpc('dgpu_next')
+        res = await session.rpc('dgpu_next')
         assert 'ok' in res.result
         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
     next_worker rotation happens correctly
     '''
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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)
 
@@ -245,7 +249,7 @@ async def test_dgpu_rotation_next_worker_disconnect(dgpu_workers):
 
         dgpu_workers[0].wait()
 
-        res = await test_rpc('dgpu_workers')
+        res = await session.rpc('dgpu_workers')
         assert 'ok' in res.result
         assert res.result['ok'] == 2
 
@@ -258,26 +262,43 @@ async def test_dgpu_no_ack_node_disconnect(skynet_running):
     '''Mock a node that connects, gets a request but fails to
     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):
-        assert 'ok' in res.result
+        resp = SkynetRPCResponse()
+        resp.result.update({'error': 'can\'t do it mate'})
+        return resp
 
-        await wait_for_dgpus(rpc_call, 1)
+    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'
+    )
 
-        with pytest.raises(SkynetDGPUComputeError) as e:
+    async with mock_server.open():
-            await check_request_img(0)
+        with open_skynet_rpc(
+            'test-ctx',
+            cert_name='whitelist/testing.cert',
+            key_name='testing.key'
+        ) as session:
 
-        assert 'dgpu failed to acknowledge request' in str(e)
+            res = await session.rpc('dgpu_online', {
+                'dgpu_addr': dgpu_addr,
+                'cert': 'whitelist/testing.cert'
+            })
+            assert 'ok' in res.result
 
-        res = await rpc_call('dgpu_workers')
+            await wait_for_dgpus(session, 1)
-        assert 'ok' in res.result
+
-        assert res.result['ok'] == 0
+            with pytest.raises(SkynetDGPUComputeError) as e:
+                await check_request_img(0)
+
+            assert 'can\'t do it mate' in str(e.value)
+
+            res = await session.rpc('dgpu_workers')
+            assert 'ok' in res.result
+            assert res.result['ok'] == 0
 
 
 @pytest.mark.parametrize(
@@ -286,13 +307,12 @@ async def test_dgpu_timeout_while_processing(dgpu_workers):
     '''Stop node while processing request to cause timeout and
     then check skynet correctly drops the node.
     '''
-    async with open_skynet_rpc(
+    with open_skynet_rpc(
         '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():
             with pytest.raises(SkynetDGPUComputeError) as e:
@@ -308,72 +328,62 @@ async def test_dgpu_timeout_while_processing(dgpu_workers):
             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(
     'dgpu_workers', [(1, ['midj'])], indirect=True)
 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', {
-                'prompt': 'red old tractor in a sunny wheat field',
+                'method': 'diffuse',
-                'step': 28,
+                'params': {
-                'width': 512, 'height': 512,
+                    'prompt': 'red old tractor in a sunny wheat field',
-                'guidance': 7.5,
+                    'step': 28,
-                'seed': None,
+                    'width': 512, 'height': 512,
-                'algo': list(ALGOS.keys())[0],
+                    'guidance': 7.5,
-                'upscaler': None
+                    'seed': None,
-            })
+                    'algo': list(ALGOS.keys())[0],
+                    'upscaler': None
+                }
+            },
+            timeout=60
+        )
 
         if 'error' in 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')
 
-        res = await rpc_call(
+        res = await session.rpc(
-            'img2img', {
+            'dgpu_call', {
-                'prompt': 'red sports car in a sunny wheat field',
+                'method': 'diffuse',
-                'step': 28,
+                'params': {
-                'img': img_raw,
+                    'prompt': 'red ferrari in a sunny wheat field',
-                'guidance': 12,
+                    'step': 28,
-                'seed': None,
+                    'guidance': 8,
-                'algo': list(ALGOS.keys())[0],
+                    'strength': 0.7,
-                'upscaler': 'x4'
+                    'seed': None,
-            })
+                    'algo': list(ALGOS.keys())[0],
+                    'upscaler': 'x4'
+                }
+            },
+            binext=img_raw,
+            timeout=60
+        )
 
         if 'error' in 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')

tests/test_skynet.py

@@ -9,6 +9,7 @@ import trio_asyncio
 
 from skynet.brain import run_skynet
 from skynet.structs import *
+from skynet.network import SessionServer
 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):
-    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 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',
-        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
-        res = await rpc_call('dgpu_workers')
+        res = await session.rpc('dgpu_workers')
         assert 'ok' in res.result
         assert res.result['ok'] == 0
 
         # check next worker is None
-        res = await rpc_call('dgpu_next')
+        res = await session.rpc('dgpu_next')
         assert 'ok' in res.result
         assert res.result['ok'] == None
 
-        # connect 1 dgpu
+        async with rpc_server.open() as rpc_server:
-        res = await rpc_call('dgpu_online')
+            # connect 1 dgpu
-        assert 'ok' in res.result
+            res = await session.rpc(
+                'dgpu_online', {'dgpu_addr': fake_dgpu_addr})
+            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
-        res = await rpc_call('dgpu_workers')
+        res = await session.rpc('dgpu_workers')
         assert 'ok' in res.result
         assert res.result['ok'] == 0
 
         # check next worker is None
-        res = await rpc_call('dgpu_next')
+        res = await session.rpc('dgpu_next')
         assert 'ok' in res.result
         assert res.result['ok'] == None

tests/test_telegram.py

@@ -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)