diff --git a/requirements.test.txt b/requirements.test.txt index ce51dec..f39926b 100644 --- a/requirements.test.txt +++ b/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 diff --git a/requirements.txt b/requirements.txt index 7afc143..c773225 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,5 @@ protobuf pyOpenSSL trio_asyncio pyTelegramBotAPI + +git+https://github.com/goodboy/tractor.git@master#egg=tractor diff --git a/skynet/cli.py b/skynet/cli.py index 856835f..021e1e6 100644 --- a/skynet/cli.py +++ b/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 diff --git a/skynet/config.py b/skynet/config.py index 65158c6..91d6101 100644 --- a/skynet/config.py +++ b/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 diff --git a/skynet/utils.py b/skynet/utils.py index f84c0ef..637078b 100644 --- a/skynet/utils.py +++ b/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) + diff --git a/tests/conftest.py b/tests/conftest.py index ac2f4be..0b4c335 100644 --- a/tests/conftest.py +++ b/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()) diff --git a/tests/test_skynet.py b/tests/test_skynet.py index ad1c488..1587d5d 100644 --- a/tests/test_skynet.py +++ b/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