async def agent(request, tmpdir, event_loop):
config = argparse.Namespace()
config.namespace = os.environ.get('BACKEND_NAMESPACE', 'testing')
config.agent_host = '127.0.0.1'
config.agent_port = 6001 # default 6001
config.stat_port = 6002
config.kernel_host_override = '127.0.0.1'
etcd_addr = os.environ.get('BACKEND_ETCD_ADDR', '127.0.0.1:2379')
redis_addr = os.environ.get('BACKEND_REDIS_ADDR', '127.0.0.1:6379')
config.etcd_addr = host_port_pair(etcd_addr)
config.redis_addr = host_port_pair(redis_addr)
config.event_addr = '127.0.0.1:5000' # dummy value
config.idle_timeout = 600
config.debug = True
config.debug_kernel = None
config.kernel_aliases = None
config.scratch_root = Path(tmpdir)
agent = None
config.instance_id = await identity.get_instance_id()
config.inst_type = await identity.get_instance_type()
config.region = await identity.get_instance_region()
print(f'serving test agent: {config.instance_id} ({config.inst_type}),'
f' ip: {config.agent_host}')
agent = AgentRPCServer(config, loop=event_loop)
await agent.init(skip_detect_manager=True)
await asyncio.sleep(0)
yield agent
print('shutting down test agent...')
if agent:
await agent.shutdown()
await asyncio.sleep(3)
async def app(event_loop, test_ns, test_db, unused_tcp_port):
""" For tests that do not require actual server running.
"""
app = web.Application(middlewares=[
exception_middleware,
api_middleware,
])
app['config'] = load_config(argv=[], extra_args_funcs=(gw_args, ))
app['config'].debug = True
# Override basic settings.
# Change these configs if local servers have different port numbers.
app['config'].redis_addr = host_port_pair(os.environ['BACKEND_REDIS_ADDR'])
app['config'].db_addr = host_port_pair(os.environ['BACKEND_DB_ADDR'])
app['config'].db_name = test_db
app['config'].docker_registry = 'lablup'
# Override extra settings
app['config'].namespace = test_ns
app['config'].heartbeat_timeout = 10.0
app['config'].service_ip = '127.0.0.1'
app['config'].service_port = unused_tcp_port
app['config'].verbose = False
# import ssl
# app['config'].ssl_cert = here / 'sample-ssl-cert' / 'sample.crt'
# app['config'].ssl_key = here / 'sample-ssl-cert' / 'sample.key'
# app['sslctx'] = ssl.SSLContext(ssl.PROTOCOL_SSLv23)
# app['sslctx'].load_cert_chain(str(app['config'].ssl_cert),
# str(app['config'].ssl_key))
# num_workers = 1
app['pidx'] = 0
return app
def test_args_parse_by_load_config():
# basic args
agent_port = 6003
redis_addr = '127.0.0.1:6381'
db_addr = '127.0.0.1:5434'
db_name = 'backendai-test'
db_user = '******'
db_password = '******'
# extra args
namespace = 'local-test'
etcd_addr = '127.0.0.1:2381'
events_port = 5002
argv = [
'--agent-port',
str(agent_port),
'--redis-addr',
redis_addr,
'--db-addr',
db_addr,
'--db-name',
db_name,
'--db-user',
db_user,
'--db-password',
db_password,
'--namespace',
namespace,
'--etcd-addr',
etcd_addr,
'--events-port',
str(events_port),
]
args = load_config(argv, extra_args_funcs=(gw_args, ))
assert args.agent_port == agent_port
assert args.redis_addr == host_port_pair(redis_addr)
assert args.db_addr == host_port_pair(db_addr)
assert args.db_name == db_name
assert args.db_user == db_user
assert args.db_password == db_password
assert args.namespace == namespace
assert args.etcd_addr == host_port_pair(etcd_addr)
assert args.events_port == port_no(events_port)
assert args.docker_registry is None
assert args.heartbeat_timeout == 5.0
assert args.service_ip == ip_address('0.0.0.0')
assert args.service_port == 0
def dirty_etcd(test_ns):
etcd_addr = host_port_pair(os.environ['BACKEND_ETCD_ADDR'])
args = Namespace(key='manager/status',
value='frozen',
etcd_addr=etcd_addr,
namespace=test_ns)
put(args)
async def maker(modules=None, ev_router=False, spawn_agent=False):
nonlocal client, runner, extra_proc
if modules is None:
modules = []
scheduler_opts = {
'close_timeout': 10,
}
cors_opts = {
'*': aiohttp_cors.ResourceOptions(
allow_credentials=False,
expose_headers="*", allow_headers="*"),
}
aiojobs.aiohttp.setup(app, **scheduler_opts)
await gw_init(app, cors_opts)
for mod in modules:
target_module = import_module(f'.{mod}', 'ai.backend.gateway')
subapp, mw = getattr(target_module, 'create_app', None)(cors_opts)
assert isinstance(subapp, web.Application)
for key in PUBLIC_INTERFACES:
subapp[key] = app[key]
prefix = subapp.get('prefix', mod.replace('_', '-'))
aiojobs.aiohttp.setup(subapp, **scheduler_opts)
app.add_subapp('/' + prefix, subapp)
app.middlewares.extend(mw)
# TODO: refactor to avoid duplicates with gateway.server
# Add legacy version-prefixed routes to the root app with some hacks
for r in subapp.router.routes():
for version in subapp['api_versions']:
subpath = r.resource.canonical
if subpath == f'/{prefix}':
subpath += '/'
legacy_path = f'/v{version}{subpath}'
handler = _get_legacy_handler(r.handler, subapp, version)
app.router.add_route(r.method, legacy_path, handler)
server_params = {}
client_params = {}
runner = web.AppRunner(app)
await runner.setup()
site = web.TCPSite(
runner,
app['config']['manager']['service-addr'].host,
app['config']['manager']['service-addr'].port,
ssl_context=app.get('sslctx'),
**server_params,
)
await site.start()
if ev_router:
# Run event_router proc. Is it enough? No way to get return values
# (app, client, etc) by using aiotools.start_server.
args = (app['config'],)
extra_proc = mp.Process(target=event_router,
args=('', 0, args,),
daemon=True)
extra_proc.start()
# Launch an agent daemon
if spawn_agent:
etcd_addr = host_port_pair(os.environ['BACKEND_ETCD_ADDR'])
os.makedirs(f'/tmp/backend.ai/scratches-{test_id}', exist_ok=True)
agent_proc = subprocess.Popen([
'python', '-m', 'ai.backend.agent.server',
'--etcd-addr', str(etcd_addr),
'--namespace', test_ns,
'--scratch-root', f'/tmp/backend.ai/scratches-{test_id}',
'--idle-timeout', '30',
])
def finalize_agent():
agent_proc.terminate()
try:
agent_proc.wait(timeout=5.0)
except subprocess.TimeoutExpired:
agent_proc.kill()
shutil.rmtree(f'/tmp/backend.ai/scratches-{test_id}')
request.addfinalizer(finalize_agent)
async def wait_for_agent():
while True:
all_ids = [inst_id async for inst_id in
app['registry'].enumerate_instances()]
if len(all_ids) > 0:
break
await asyncio.sleep(0.2)
task = event_loop.create_task(wait_for_agent())
with timeout(10.0):
await task
port = app['config']['manager']['service-addr'].port
if app.get('sslctx'):
url = f'https://localhost:{port}'
client_params['connector'] = aiohttp.TCPConnector(verify_ssl=False)
else:
url = f'http://localhost:{port}'
http_session = aiohttp.ClientSession(
loop=event_loop, **client_params)
client = Client(http_session, url)
return app, client
def prepare_and_cleanup_databases(request, test_ns, test_db, folder_mount,
folder_host):
os.environ['BACKEND_NAMESPACE'] = test_ns
os.environ['BACKEND_DB_NAME'] = test_db
# Clear and reset etcd namespace using CLI functions.
etcd_addr = host_port_pair(os.environ['BACKEND_ETCD_ADDR'])
samples = here / '..' / 'sample-configs'
args = Namespace(file=samples / 'image-metadata.yml',
etcd_addr=etcd_addr,
namespace=test_ns)
update_images(args)
args = Namespace(file=samples / 'image-aliases.yml',
etcd_addr=etcd_addr,
namespace=test_ns)
update_aliases(args)
args = Namespace(key='volumes/_mount',
value=str(folder_mount),
etcd_addr=etcd_addr,
namespace=test_ns)
put(args)
args = Namespace(key='volumes/_default_host',
value=str(folder_host),
etcd_addr=etcd_addr,
namespace=test_ns)
put(args)
args = Namespace(key='nodes/docker_registry',
value='lablup',
etcd_addr=etcd_addr,
namespace=test_ns)
put(args)
def finalize_etcd():
args = Namespace(key='',
prefix=True,
etcd_addr=etcd_addr,
namespace=test_ns)
delete(args)
request.addfinalizer(finalize_etcd)
# Create database using low-level psycopg2 API.
db_addr = host_port_pair(os.environ['BACKEND_DB_ADDR'])
db_user = os.environ['BACKEND_DB_USER']
db_pass = os.environ['BACKEND_DB_PASSWORD']
if db_pass:
# TODO: escape/urlquote db_pass
db_url = f'postgresql://{db_user}:{db_pass}@{db_addr}'
else:
db_url = f'postgresql://{db_user}@{db_addr}'
conn = pg.connect(db_url)
conn.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
cur = conn.cursor()
cur.execute(f'CREATE DATABASE "{test_db}";')
cur.close()
conn.close()
def finalize_db():
conn = pg.connect(db_url)
conn.set_isolation_level(pg.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
cur = conn.cursor()
cur.execute(f'REVOKE CONNECT ON DATABASE "{test_db}" FROM public;')
cur.execute('SELECT pg_terminate_backend(pid) FROM pg_stat_activity '
'WHERE pid <> pg_backend_pid();')
cur.execute(f'DROP DATABASE "{test_db}";')
cur.close()
conn.close()
request.addfinalizer(finalize_db)
# Load the database schema using CLI function.
alembic_url = db_url + '/' + test_db
with tempfile.NamedTemporaryFile(mode='w', encoding='utf8') as alembic_cfg:
alembic_sample_cfg = here / '..' / 'alembic.ini.sample'
alembic_cfg_data = alembic_sample_cfg.read_text()
alembic_cfg_data = re.sub(r'^sqlalchemy.url = .*$',
f'sqlalchemy.url = {alembic_url}',
alembic_cfg_data,
flags=re.M)
alembic_cfg.write(alembic_cfg_data)
alembic_cfg.flush()
args = Namespace(config=Path(alembic_cfg.name), schema_version='head')
oneshot(args)
# Populate example_keypair fixture
fixture = getattr(fixtures, 'example_keypair')
engine = sa.create_engine(alembic_url)
conn = engine.connect()
for rowset in fixture:
table = getattr(models, rowset[0])
conn.execute(table.insert(), rowset[1])
conn.close()
engine.dispose()