Add download command

Fix ini config loading
Update test requirements
Fix attempt_insecure and connection simple dgpu tests
pull/4/head
Guillermo Rodriguez 2023-03-20 10:58:35 -03:00
parent 4feed81662
commit 5e8787b4fc
No known key found for this signature in database
GPG Key ID: EC3AB66D5D83B392
7 changed files with 41 additions and 18 deletions

View File

@ -3,4 +3,4 @@ pytest
pytest-trio pytest-trio
psycopg2-binary 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

View File

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

View File

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

View File

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

View File

@ -169,3 +169,13 @@ def upscale(
image.save(output) 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)

View File

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

View File

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