async def test_fetch(dummy_endpoint):
with aioresponses() as m, Session() as session:
body = b'hello world'
m.post(
dummy_endpoint + 'function', status=200, body=body,
headers={'Content-Type': 'text/plain; charset=utf-8',
'Content-Length': str(len(body))},
)
rqst = Request('POST', 'function')
async with rqst.fetch() as resp:
assert isinstance(resp, Response)
assert resp.status == 200
assert resp.content_type == 'text/plain'
assert await resp.text() == body.decode()
assert resp.content_length == len(body)
with aioresponses() as m, Session() as session:
body = b'{"a": 1234, "b": null}'
m.post(
dummy_endpoint + 'function', status=200, body=body,
headers={'Content-Type': 'application/json; charset=utf-8',
'Content-Length': str(len(body))},
)
rqst = Request('POST', 'function')
async with rqst.fetch() as resp:
assert isinstance(resp, Response)
assert resp.status == 200
assert resp.content_type == 'application/json'
assert await resp.text() == body.decode()
assert await resp.json() == {'a': 1234, 'b': None}
assert resp.content_length == len(body)
def test_api_function_metaclass():
# Here, we repeat intentionally the same stuffs
# to check if our metaclass works across multiple
# re-definition and re-instantiation scenarios.
with Session() as session:
Dummy = type('DummyFunction', (BaseFunction, ), {
**DummyFunction.__dict__,
'session': session,
})
assert Dummy.session is session
assert Dummy().session is session
assert Dummy.get_or_create() == 'created'
assert Dummy().calculate() == 'done'
assert Dummy.get_or_create() == 'created'
assert Dummy().calculate() == 'done'
with Session() as session:
Dummy = type('DummyFunction', (BaseFunction, ), {
**DummyFunction.__dict__,
'session': session,
})
assert Dummy.session is session
assert Dummy().session is session
assert Dummy.get_or_create() == 'created'
assert Dummy().calculate() == 'done'
assert Dummy.get_or_create() == 'created'
assert Dummy().calculate() == 'done'
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.log.info(
'Backend.AI kernel starting with client session ID: {0}'.format(
self.ident))
self.backend_session = Session()
self.kernel = self.backend_session.Kernel.get_or_create(
self.backend_lang, self.ident)
def test_vfolder_list_files():
with Session() as session:
with aioresponses() as m:
vfolder_name = 'fake-vfolder-name'
payload = {
"files": [{
"mode": "-rw-r--r--",
"size": 4751244,
"ctime": 1528277299.2744732,
"mtime": 1528277299.2744732,
"atime": 1528277300.7658687,
"filename": "bigtxt.txt",
}, {
"mode": "-rw-r--r--",
"size": 200000,
"ctime": 1528333257.6576185,
"mtime": 1528288069.625786,
"atime": 1528332829.692922,
"filename": "200000",
}],
"folder_path":
"/mnt/local/1f6bd27fde1248cabfb50306ea83fc0a",
}
m.get(build_url(session.config,
'/folders/{}/files'.format(vfolder_name)),
status=200,
payload=payload)
resp = session.VFolder(vfolder_name).list_files('.')
assert resp == payload
def test_create_with_config(mocker, api_version):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=201,
json=AsyncMock())
mock_req = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
myconfig = APIConfig(
endpoint='https://localhost:9999',
access_key='1234',
secret_key='asdf',
user_agent='BAIClientTest',
version=f'v{api_version[0]}.{api_version[1]}',
)
with Session(config=myconfig) as session:
prefix = get_naming(session.api_version, 'path')
if api_version[0] == 4:
assert prefix == 'kernel'
else:
assert prefix == 'session'
assert session.config is myconfig
cs = session.ComputeSession.get_or_create('python')
mock_req.assert_called_once_with(session, 'POST', f'/{prefix}')
assert str(cs.session.config.endpoint) == 'https://localhost:9999'
assert cs.session.config.user_agent == 'BAIClientTest'
assert cs.session.config.access_key == '1234'
assert cs.session.config.secret_key == 'asdf'
def test_delete_vfolder():
with Session() as session, aioresponses() as m:
vfolder_name = 'fake-vfolder-name'
m.delete(build_url(session.config, '/folders/{}'.format(vfolder_name)),
status=204)
resp = session.VFolder(vfolder_name).delete()
assert resp == {}
def test_kernel_execution_with_vfolder_mounts(self):
with Session() as sess:
vfname = 'vftest-' + token_hex(4)
sess.VFolder.create(vfname)
vfolder = sess.VFolder(vfname)
try:
with tempfile.NamedTemporaryFile('w',
suffix='.py',
dir=Path.cwd()) as f:
f.write('print("hello world")\nx = 1 / 0\n')
f.flush()
f.seek(0)
vfolder.upload([f.name])
kernel = sess.Kernel.get_or_create('python:3.6-ubuntu18.04',
mounts=[vfname])
try:
console, n = exec_loop(
kernel, 'batch', '', {
'build':
'',
'exec':
'python {}/{}'.format(vfname,
Path(f.name).name),
})
assert 'hello world' in console['stdout']
assert 'ZeroDivisionError' in console['stderr']
assert len(console['media']) == 0
finally:
kernel.destroy()
finally:
vfolder.delete()
def test_auth_missing_body():
with Session() as sess:
request = Request('GET', '/auth')
with pytest.raises(BackendAPIError) as e:
with request.fetch():
pass
assert e.value.status == 400
async def test_user_cannot_mutate_alias_dealias(self, userconfig):
with Session() as sess:
test_alias = 'testalias-b9f1ce136f584ca892d5fef3e78dd11d'
with pytest.raises(BackendAPIError):
sess.Image.aliasImage(test_alias, 'lua:5.1-alpine3.8')
with pytest.raises(BackendAPIError):
sess.Image.dealiasImage(test_alias)
def test_unfreeze(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=204)
mocker.patch('ai.backend.client.func.manager.Request', return_value=mock_req_obj)
with Session() as session:
session.Manager.unfreeze()
mock_req_obj.fetch.assert_called_once_with()
def test_create_kernel_raises_err_with_abnormal_status(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = mock.MagicMock(status=400)
mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj)
with Session() as session:
with pytest.raises(BackendAPIError):
session.Kernel.get_or_create('python')
async def test_list_images_by_user(self, userconfig):
with Session() as sess:
images = sess.Image.list()
image = images[0]
assert len(images) > 0
assert 'name' in image
assert 'tag' in image
assert 'hash' in image
def test_vfolder_delete_invitation():
with Session() as session, aioresponses() as m:
payload = {'msg': 'Vfolder invitation is deleted: fake-inv-id.'}
m.delete(build_url(session.config, '/folders/invitations/delete'),
status=200,
payload=payload)
resp = session.VFolder.delete_invitation('inv-id')
assert resp == payload
async def test_manipulate_resource_policy(self):
access_key = get_config().access_key
rpname = 'testrp-' + uuid.uuid4().hex
with Session() as sess:
try:
rp = sess.ResourcePolicy(access_key)
assert rp.info(rpname) is None
rps = sess.ResourcePolicy.list()
original_count = len(rps)
# Create resource policy
sess.ResourcePolicy.create(name=rpname,
default_for_unspecified='LIMITED',
total_resource_slots='{}',
max_concurrent_sessions=1,
max_containers_per_session=1,
max_vfolder_count=1,
max_vfolder_size=1,
idle_timeout=1,
allowed_vfolder_hosts=['local'])
rps = sess.ResourcePolicy.list()
assert len(rps) == original_count + 1
info = rp.info(rpname)
assert info['name'] == rpname
assert info['total_resource_slots'] == '{}'
assert info['max_concurrent_sessions'] == 1
assert info['max_vfolder_count'] == 1
assert info['max_vfolder_size'] == 1
assert info['idle_timeout'] == 1
# Update resource policy
sess.ResourcePolicy.update(
name=rpname,
default_for_unspecified='LIMITED',
total_resource_slots='{"cpu": "count"}',
max_concurrent_sessions=2,
max_containers_per_session=2,
max_vfolder_count=2,
max_vfolder_size=2,
idle_timeout=2,
allowed_vfolder_hosts=['local'])
rps = sess.ResourcePolicy.list()
assert len(rps) == original_count + 1
info = rp.info(rpname)
assert info['name'] == rpname
assert info['total_resource_slots'] == '{"cpu": "count"}'
assert info['max_concurrent_sessions'] == 2
assert info['max_vfolder_count'] == 2
assert info['max_vfolder_size'] == 2
assert info['idle_timeout'] == 2
# Delete ResourcePolicy
sess.ResourcePolicy.delete(rpname)
rps = sess.ResourcePolicy.list()
assert len(rps) == original_count
except Exception:
sess.ResourcePolicy.delete(rpname)
raise
async def test_user_cannot_create_keypair(self, userconfig):
email = 'testion' + uuid.uuid4().hex + '@test.mars'
with Session() as sess:
with pytest.raises(BackendAPIError):
sess.KeyPair.create(user_id=email,
is_active=True,
is_admin=False,
resource_policy='default',
rate_limit=1)
def test_freeze_opt_force_kill(self, mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=204)
mocker.patch('ai.backend.client.manager.Request',
return_value=mock_req_obj)
with Session() as session:
session.Manager.freeze(force_kill=True)
mock_req_obj.fetch.assert_called_once_with()
def test_destroy_kernel_raises_err_with_abnormal_status(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = mock.MagicMock(status=400)
mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj)
kernel_id = token_hex(12)
with Session() as session:
with pytest.raises(BackendAPIError):
k = session.Kernel(kernel_id)
k.destroy()
def test_vfolder_delete_files():
with Session() as session, aioresponses() as m:
vfolder_name = 'fake-vfolder-name'
files = ['fake-file1', 'fake-file2']
m.delete(build_url(session.config,
'/folders/{}/delete_files'.format(vfolder_name)),
status=200,
payload={})
resp = session.VFolder(vfolder_name).delete_files(files)
assert resp == '{}'
def test_create_kernel_return_id_only(mocker):
mock_json_func = lambda: {'kernelId': 'mock_kernel_id'}
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = mock.MagicMock(status=201,
json=mock_json_func)
mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj)
with Session() as session:
k = session.Kernel.get_or_create('python')
assert k.kernel_id == mock_json_func()['kernelId']
def test_vfolder_upload(tmp_path: Path):
with Session() as session, aioresponses() as m:
mock_file = tmp_path / 'example.bin'
mock_file.write_bytes(secrets.token_bytes(32))
vfolder_name = 'fake-vfolder-name'
m.post(build_url(session.config,
'/folders/{}/upload'.format(vfolder_name)),
status=201)
resp = session.VFolder(vfolder_name).upload([mock_file],
basedir=tmp_path)
assert resp == ''
def test_vfolder_invite():
with Session() as session, aioresponses() as m:
vfolder_name = 'fake-vfolder-name'
user_ids = ['*****@*****.**', '*****@*****.**']
payload = {'invited_ids': user_ids}
m.post(build_url(session.config,
'/folders/{}/invite'.format(vfolder_name)),
status=201,
payload=payload)
resp = session.VFolder(vfolder_name).invite('rw', user_ids)
assert resp == payload
def test_auth_malformed():
with Session() as sess:
request = Request('GET', '/auth')
request.set_content(
b'<this is not json>',
content_type='application/json',
)
with pytest.raises(BackendAPIError) as e:
with request.fetch():
pass
assert e.value.status == 400
def test_vfolder_accept_invitation():
with Session() as session, aioresponses() as m:
payload = {
'msg': ('User [email protected] now can access'
' vfolder fake-vfolder-id'),
}
m.post(build_url(session.config, '/folders/invitations/accept'),
status=200,
payload=payload)
resp = session.VFolder.accept_invitation('inv-id')
assert resp == payload
def test_vfolder_upload(tmpdir):
with Session() as session:
with aioresponses() as m:
mockfile = tmpdir.join('example.jpg')
mockfile.write('mock file')
vfolder_name = 'fake-vfolder-name'
m.post(build_url(session.config,
'/folders/{}/upload'.format(vfolder_name)),
status=201)
resp = session.VFolder(vfolder_name).upload([mockfile.strpath],
basedir=tmpdir.strpath)
assert resp.status == 201
def test_create_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = AsyncContextMock(status=201,
json=AsyncMock())
mock_req = mocker.patch('ai.backend.client.func.session.Request',
return_value=mock_req_obj)
with Session() as session:
prefix = get_naming(session.api_version, 'path')
session.ComputeSession.get_or_create('python:3.6-ubuntu18.04')
mock_req.assert_called_once_with(session, 'POST', f'/{prefix}')
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()
def test_kernel_get_or_create_reuse(self):
with Session() as sess:
try:
# Sessions with same token and same language must be reused.
t = token_hex(6)
kernel1 = sess.Kernel.get_or_create('python:3.6-ubuntu18.04',
client_token=t)
kernel2 = sess.Kernel.get_or_create('python:3.6-ubuntu18.04',
client_token=t)
assert kernel1.kernel_id == kernel2.kernel_id
finally:
kernel1.destroy()
def test_auth():
random_msg = uuid.uuid4().hex
with Session() as sess:
request = Request('GET', '/auth')
request.set_json({
'echo': random_msg,
})
with request.fetch() as resp:
assert resp.status == 200
data = resp.json()
assert data['authorized'] == 'yes'
assert data['echo'] == random_msg
def test_create_kernel_url(self, mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(
status=201, json=asynctest.CoroutineMock())
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
session.Kernel.get_or_create('python:3.6-ubuntu18.04')
mock_req.assert_called_once_with(session, 'POST', '/kernel/create')
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()
def test_not_found(self):
with Session() as sess:
request = Request(sess, 'GET', '/invalid-url-wow')
with pytest.raises(BackendAPIError) as e:
with request.fetch():
pass
assert e.value.status == 404
request = Request(sess, 'GET', '/auth/uh-oh')
with pytest.raises(BackendAPIError) as e:
with request.fetch():
pass
assert e.value.status == 404
def test_create_kernel_return_id_only(self, mocker):
return_value = {'kernelId': 'mock_kernel_id'}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=201,
json=mock_json_coro)
mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
k = session.Kernel.get_or_create('python:3.6-ubuntu18.04')
assert k.kernel_id == return_value['kernelId']
