def test_execute_code_raises_err_with_abnormal_status(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.send.return_value = mock.MagicMock(status=400)
mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj)
kernel_id = token_hex(12)
run_id = token_hex(8)
with pytest.raises(BackendAPIError):
k = Kernel(kernel_id)
k.execute(run_id, 'hello')
async def test_execute_code_url(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = asynctest.MagicMock(status=200)
kernel_id = token_hex(12)
run_id = token_hex(8)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).execute(run_id, 'hello')
mock_req_cls.assert_called_once_with(
session, 'POST', '/kernel/{}'.format(kernel_id),
{'mode': 'query', 'runId': run_id, 'code': 'hello'})
async def test_execute_code_url(self, mocker):
return_value = {'result': 'hi'}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=200,
json=mock_json_coro)
kernel_id = token_hex(12)
run_id = token_hex(8)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).execute(run_id, 'hello')
mock_req_cls.assert_called_once_with(
session, 'POST', '/kernel/{}'.format(kernel_id), params={})
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 exec_loop(kernel, mode, code, opts=None, user_inputs=None):
# The server may return continuation if kernel preparation
# takes a little longer time (a few seconds).
console = []
num_queries = 0
run_id = token_hex(8)
if user_inputs is None:
user_inputs = []
while True:
result = kernel.execute(
run_id,
code=code if num_queries == 0 or mode == 'input' else '',
mode=mode,
opts=opts)
num_queries += 1
console.extend(result['console'])
if result['status'] == 'finished':
break
elif result['status'] == 'waiting-input':
mode = 'input'
code = user_inputs.pop(0)
opts = None
else:
mode = 'continue'
code = ''
opts = None
return aggregate_console(console), num_queries
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 exec_loop(kernel, mode, code, opts=None, user_inputs=None):
# The server may return continuation if kernel preparation
# takes a little longer time (a few seconds).
console = []
num_queries = 0
run_id = token_hex(8)
if user_inputs is None:
user_inputs = []
while True:
result = kernel.execute(
run_id,
code=code if num_queries == 0 or mode == 'input' else '',
mode=mode,
opts=opts)
num_queries += 1
console.extend(result['console'])
if result['status'] == 'finished':
break
elif result['status'] == 'waiting-input':
mode = 'input'
code = user_inputs.pop(0)
opts = None
else:
mode = 'continue'
code = ''
opts = None
return aggregate_console(console), num_queries
async def test_execute_code_raises_err_with_abnormal_status():
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.asend.return_value = asynctest.MagicMock(status=400)
run_id = token_hex(8)
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
with pytest.raises(BackendAPIError):
await Kernel('mykernel').execute(run_id, 'hello')
def test_restart_kernel_raises_err_with_abnormal_status(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.send.return_value = mock.MagicMock(status=400)
mocker.patch('ai.backend.client.kernel.Request', return_value=mock_req_obj)
kernel_id = token_hex(12)
with pytest.raises(BackendAPIError):
restart_kernel(kernel_id)
def test_execute_code_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.send.return_value = mock.MagicMock(status=200)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
kernel_id = token_hex(12)
k = Kernel(kernel_id)
run_id = token_hex(8)
k.execute(run_id, 'hello')
mock_req.assert_called_once_with('POST', '/kernel/{}'.format(kernel_id), {
'mode': 'query',
'runId': run_id,
'code': 'hello'
})
mock_req_obj.send.assert_called_once_with()
mock_req_obj.send.return_value.json.assert_called_once_with()
async def test_restart_kernel_url():
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.asend.return_value = asynctest.MagicMock(status=204)
kernel_id = token_hex(12)
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await Kernel(kernel_id).restart()
mock_req_cls.assert_called_once_with('PATCH',
'/kernel/{}'.format(kernel_id))
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()
async def test_destroy_kernel_url(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.afetch.return_value = asynctest.MagicMock(status=204)
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).destroy()
mock_req_cls.assert_called_once_with(
session, 'DELETE', '/kernel/{}'.format(kernel_id))
def test_kernel_execution_query_mode_user_input(self, py3_kernel):
name = token_hex(8)
code = 'name = input("your name? "); print(f"hello, {name}!")'
console, n = exec_loop(py3_kernel,
'query',
code,
None,
user_inputs=[name])
assert 'your name?' in console['stdout']
assert 'hello, {}!'.format(name) in console['stdout']
def test_restart_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.send.return_value = mock.MagicMock(status=204)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
kernel_id = token_hex(12)
restart_kernel(kernel_id)
mock_req.assert_called_once_with('PATCH', '/kernel/{}'.format(kernel_id))
mock_req_obj.send.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()
async def test_execute_code_url():
mock_req_obj = asynctest.MagicMock(spec=Request)
mock_req_obj.asend.return_value = asynctest.MagicMock(status=200)
kernel_id = token_hex(12)
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await execute_code(kernel_id, 'hello')
mock_req_cls.assert_called_once_with('POST',
'/kernel/{}'.format(kernel_id), {
'mode': 'query',
'code': 'hello'
})
def test_execute_code_url(self, mocker):
return_value = {'result': 'hi'}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(
status=200, json=mock_json_coro)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
kernel_id = token_hex(12)
k = session.Kernel(kernel_id)
run_id = token_hex(8)
k.execute(run_id, 'hello')
mock_req.assert_called_once_with(
session, 'POST', '/kernel/{}'.format(kernel_id),
params={}
)
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_get_kernel_info_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.send.return_value = mock.MagicMock(status=200)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
kernel_id = token_hex(12)
get_kernel_info(kernel_id)
mock_req.assert_called_once_with('GET', '/kernel/{}'.format(kernel_id))
mock_req_obj.send.assert_called_once_with()
mock_req_obj.send.return_value.json.assert_called_once_with()
def test_kernel_get_or_create_reuse(intgr_config):
with Session(config=intgr_config) as sess:
try:
# Sessions with same token and same language must be reused.
t = token_hex(6)
kernel1 = sess.Kernel.get_or_create('python:latest',
client_token=t)
kernel2 = sess.Kernel.get_or_create('python:latest',
client_token=t)
assert kernel1.kernel_id == kernel2.kernel_id
finally:
kernel1.destroy()
def test_kernel_execution_batch_mode_user_input(self, py3_kernel):
name = token_hex(8)
with tempfile.NamedTemporaryFile('w', suffix='.py', dir=Path.cwd()) as f:
f.write('name = input("your name? "); print(f"hello, {name}!")')
f.flush()
f.seek(0)
py3_kernel.upload([f.name])
console, _ = exec_loop(py3_kernel, 'batch', '', {
'build': '',
'exec': 'python {}'.format(Path(f.name).name),
}, user_inputs=[name])
assert 'your name?' in console['stdout']
assert 'hello, {}!'.format(name) in console['stdout']
async def test_get_kernel_info_url(self, mocker):
return_value = {}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=200,
json=mock_json_coro)
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).get_info()
mock_req_cls.assert_called_once_with(
session, 'GET', '/kernel/{}'.format(kernel_id), params={})
async def test_restart_kernel_url(self, mocker):
mock_req_obj = asynctest.MagicMock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=204)
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
await session.Kernel(kernel_id).restart()
mock_req_cls.assert_called_once_with(
session,
'PATCH',
'/kernel/{}'.format(kernel_id),
params={})
def exec_loop(kernel, code):
# The server may return continuation if kernel preparation
# takes a little longer time (a few seconds).
console = []
num_queries = 0
run_id = token_hex(8)
while True:
result = kernel.execute(run_id, code if num_queries == 0 else '')
num_queries += 1
console.extend(result['console'])
if result['status'] == 'finished':
break
return aggregate_console(console), num_queries
async def test_stream_pty(mocker):
mock_req_obj = asynctest.MagicMock(spec=Request)
ws = object()
mock_req_obj.connect_websocket.return_value = None, ws
kernel_id = token_hex(12)
async with AsyncSession() as session:
with asynctest.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj) as mock_req_cls:
stream = await session.Kernel(kernel_id).stream_pty()
mock_req_cls.assert_called_once_with(
session, 'GET', '/stream/kernel/{}/pty'.format(kernel_id))
mock_req_obj.connect_websocket.assert_called_once_with()
assert isinstance(stream, StreamPty)
assert stream.ws is ws
def test_destroy_kernel_url(mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = mock.MagicMock(status=204)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
kernel_id = token_hex(12)
k = session.Kernel(kernel_id)
k.destroy()
mock_req.assert_called_once_with(session, 'DELETE',
'/kernel/{}'.format(kernel_id))
mock_req_obj.fetch.assert_called_once_with()
def test_restart_kernel_url(self, mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=204)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
kernel_id = token_hex(12)
k = session.Kernel(kernel_id)
k.restart()
mock_req.assert_called_once_with(session,
'PATCH', '/kernel/{}'.format(kernel_id),
params={})
mock_req_obj.fetch.assert_called_once_with()
def run_execute_code(kid):
if kid is not None:
begin = time.monotonic()
console = []
run_id = token_hex(8)
while True:
result = Kernel(kid).execute(run_id, sample_code)
console.extend(result['console'])
if result['status'] == 'finished':
break
stdout = ''.join(rec[1] for rec in console if rec[0] == 'stdout')
end = time.monotonic()
print(stdout)
return end - begin
return None
def test_kernel_execution_batch_mode_user_input(py3_kernel):
name = token_hex(8)
with tempfile.NamedTemporaryFile('w', suffix='.py', dir=Path.cwd()) as f:
f.write('name = input("your name? "); print(f"hello, {name}!")')
f.flush()
f.seek(0)
py3_kernel.upload([f.name])
console, _ = exec_loop(py3_kernel,
'batch',
'', {
'build': '',
'exec': 'python {}'.format(Path(f.name).name),
},
user_inputs=[name])
assert 'your name?' in console['stdout']
assert 'hello, {}!'.format(name) in console['stdout']
def test_restart_kernel_url(self, mocker):
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=204)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
kernel_id = token_hex(12)
k = session.Kernel(kernel_id)
k.restart()
mock_req.assert_called_once_with(session,
'PATCH',
'/kernel/{}'.format(kernel_id),
params={})
mock_req_obj.fetch.assert_called_once_with()
def test_get_kernel_info_url(self, mocker):
return_value = {}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(
status=200, json=mock_json_coro)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
kernel_id = token_hex(12)
k = session.Kernel(kernel_id)
k.get_info()
mock_req.assert_called_once_with(session,
'GET', '/kernel/{}'.format(kernel_id),
params={})
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()
def test_get_kernel_info_url(self, mocker):
return_value = {}
mock_json_coro = asynctest.CoroutineMock(return_value=return_value)
mock_req_obj = mock.Mock()
mock_req_obj.fetch.return_value = ContextMagicMock(status=200,
json=mock_json_coro)
mock_req = mocker.patch('ai.backend.client.kernel.Request',
return_value=mock_req_obj)
with Session() as session:
kernel_id = token_hex(12)
k = session.Kernel(kernel_id)
k.get_info()
mock_req.assert_called_once_with(session,
'GET',
'/kernel/{}'.format(kernel_id),
params={})
mock_req_obj.fetch.assert_called_once_with()
mock_req_obj.fetch.return_value.json.assert_called_once_with()
def test_kernel_execution_query_mode_user_input(self, py3_kernel):
name = token_hex(8)
code = 'name = input("your name? "); print(f"hello, {name}!")'
console, n = exec_loop(py3_kernel, 'query', code, None, user_inputs=[name])
assert 'your name?' in console['stdout']
assert 'hello, {}!'.format(name) in console['stdout']
