Add download command

Fix ini config loading
Update test requirements
Fix attempt_insecure and connection simple dgpu tests

requirements.test.txt

@@ -3,4 +3,4 @@ pytest
 pytest-trio
 psycopg2-binary
 
-git+https://github.com/guilledk/pytest-dockerctl.git@host_network#egg=pytest-dockerctl
+git+https://github.com/guilledk/pytest-dockerctl.git@multi_names#egg=pytest-dockerctl

requirements.txt

@@ -9,3 +9,5 @@ protobuf
 pyOpenSSL
 trio_asyncio
 pyTelegramBotAPI
+
+git+https://github.com/goodboy/tractor.git@master#egg=tractor

skynet/cli.py

@@ -76,6 +76,12 @@ def upscale(input, output, model):
         model_path=model)
 
 
+@skynet.command()
+def download():
+    _, hf_token, _, cfg = init_env_from_config()
+    utils.download_all_models(hf_token)
+
+
 @skynet.group()
 def run(*args, **kwargs):
     pass

skynet/config.py

@@ -1,5 +1,7 @@
 #!/usr/bin/python
 
+import os
+
 from pathlib import Path
 from configparser import ConfigParser
 
@@ -18,19 +20,20 @@ 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']
+        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']
+        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']
+        tg_token = config['skynet.telegram']['token']
 
     return hf_home, hf_token, tg_token, config

skynet/utils.py

@@ -169,3 +169,13 @@ def upscale(
 
 
     image.save(output)
+
+
+def download_all_models(hf_token: str):
+    assert torch.cuda.is_available()
+
+    login(token=hf_token)
+    for model in ALGOS:
+        print(f'DOWNLOADING {model.upper()}')
+        pipeline_for(model)
+

tests/conftest.py

@@ -57,7 +57,6 @@ def dgpu_workers(request, dockerctl, skynet_running):
         name='skynet-test-runtime-cuda',
         commands=cmds,
         environment={
-            'HF_TOKEN': os.environ['HF_TOKEN'],
             'HF_HOME': '/skynet/hf_home'
         },
         network='host',
@@ -67,6 +66,6 @@ def dgpu_workers(request, dockerctl, skynet_running):
     ) as containers:
         yield containers
 
-        #for i, container in enumerate(containers):
-        #    logging.info(f'container {i} logs:')
-        #    logging.info(container.logs().decode())
+        for i, container in enumerate(containers):
+            logging.info(f'container {i} logs:')
+            logging.info(container.logs().decode())

tests/test_skynet.py

@@ -18,7 +18,7 @@ async def test_skynet(skynet_running):
 
 
 async def test_skynet_attempt_insecure(skynet_running):
-    with pytest.raises(trio.TooSlowError) as e:
+    with pytest.raises(pynng.exceptions.NNGException) as e:
         with open_skynet_rpc('bad-actor') as session:
             with trio.fail_after(5):
                 await session.rpc('skynet_shutdown')
@@ -44,40 +44,43 @@ async def test_skynet_dgpu_connection_simple(skynet_running):
     ) as session:
         # check 0 nodes are connected
         res = await session.rpc('dgpu_workers')
-        assert 'ok' in res.result
+        assert 'ok' in res.result.keys()
         assert res.result['ok'] == 0
 
         # check next worker is None
         res = await session.rpc('dgpu_next')
-        assert 'ok' in res.result
+        assert 'ok' in res.result.keys()
         assert res.result['ok'] == None
 
         async with rpc_server.open() as rpc_server:
             # connect 1 dgpu
             res = await session.rpc(
-                'dgpu_online', {'dgpu_addr': fake_dgpu_addr})
-            assert 'ok' in res.result
+                'dgpu_online', {
+                'dgpu_addr': fake_dgpu_addr,
+                'cert': 'whitelist/testing.cert'
+            })
+            assert 'ok' in res.result.keys()
 
             # check 1 node is connected
             res = await session.rpc('dgpu_workers')
-            assert 'ok' in res.result
+            assert 'ok' in res.result.keys()
             assert res.result['ok'] == 1
 
             # check next worker is 0
             res = await session.rpc('dgpu_next')
-            assert 'ok' in res.result
+            assert 'ok' in res.result.keys()
             assert res.result['ok'] == 0
 
             # disconnect 1 dgpu
             res = await session.rpc('dgpu_offline')
-            assert 'ok' in res.result
+            assert 'ok' in res.result.keys()
 
         # check 0 nodes are connected
         res = await session.rpc('dgpu_workers')
-        assert 'ok' in res.result
+        assert 'ok' in res.result.keys()
         assert res.result['ok'] == 0
 
         # check next worker is None
         res = await session.rpc('dgpu_next')
-        assert 'ok' in res.result
+        assert 'ok' in res.result.keys()
         assert res.result['ok'] == None