def make_model_info():
task = gen()
name = gen()
model_file_path = MODEL_FILE_PATH
model_class = MODEL_CLASS
dependencies = {}
return (task, name, model_file_path, model_class, dependencies)
def make_inference_job_info(client: Client, task=None, model_id=None):
task = task or gen()
app = gen()
train_job_id = make_train_job(client,
task=task,
app=app,
model_id=model_id)
wait_for_train_job_status(client, app, TrainJobStatus.STOPPED)
return (task, app, train_job_id)
def make_train_job_info(client: Client, task=None):
task = task or gen()
app = gen()
train_dataset_id = make_dataset(client, task=task)
val_dataset_id = make_dataset(client, task=task)
model_id = make_model(task=task)
# 1 trial & no GPU
budget = {BudgetOption.MODEL_TRIAL_COUNT: 1, BudgetOption.GPU_COUNT: 0}
return (task, app, model_id, train_dataset_id, val_dataset_id, budget)
def test_app_dev_cant_use_running_train_job(self):
app_dev = make_app_dev()
app = gen()
train_job_id = make_train_job(app_dev, app=app)
# Will throw error as train job is still running
with pytest.raises(Exception, match='InvalidTrainJobError'):
app_dev.create_inference_job(app)
def test_app_dev_cant_manage_users(self):
ban_email = f'{gen()}@rafiki'
email = gen_email()
password = gen()
model_dev = make_model_dev(email=ban_email)
app_dev = make_app_dev()
# Can't create user
with pytest.raises(Exception):
app_dev.create_user(email, password, UserType.ADMIN)
# Can't list users
with pytest.raises(Exception):
app_dev.get_users()
# Can't ban user
with pytest.raises(Exception):
app_dev.ban_user(ban_email)
def test_admin_ban_users(self):
admin = make_admin()
model_dev_email = gen_email()
app_dev_email = gen_email()
password = gen()
model_dev = make_model_dev(email=model_dev_email, password=password)
app_dev = make_app_dev(email=app_dev_email, password=password)
# Ban both users
admin.ban_user(model_dev_email)
admin.ban_user(app_dev_email)
# Both users cannot login again
with pytest.raises(Exception):
model_dev.login(model_dev_email, password)
with pytest.raises(Exception):
app_dev.login(app_dev_email, password)
def test_admin_create_users(self):
admin = make_admin()
model_dev_email = gen_email()
app_dev_email = gen_email()
password = gen()
# Create model dev
model_dev = admin.create_user(model_dev_email, password, UserType.MODEL_DEVELOPER)
assert 'id' in model_dev
model_dev_id = model_dev['id']
# Create app dev
app_dev = admin.create_user(app_dev_email, password, UserType.APP_DEVELOPER)
assert 'id' in app_dev
app_dev_id = model_dev['id']
# Created users are in list of users
users = admin.get_users()
assert any([x['id'] == model_dev_id for x in users])
assert any([x['id'] == app_dev_id for x in users])
def test_app_dev_informed_model_error_in_inference(self):
app_dev = make_app_dev()
task = gen()
model_id = make_model(task=task,
model_file_path=MODEL_INVALID_LOAD_PARAMS_PATH)
(task, app, train_job_id) = make_inference_job_info(app_dev,
task=task,
model_id=model_id)
# Create inference job
inference_job = app_dev.create_inference_job(app)
# Will eventuallly throw error as inference job is no longer running
with pytest.raises(Exception, match='InvalidRunningInferenceJobError'):
wait_for_inference_job_status(app_dev, app,
InferenceJobStatus.ERRORED)
# Inference job errored
inference_jobs = app_dev.get_inference_jobs_of_app(app)
inference_job = next(x for x in inference_jobs
if x['id'] == inference_job['id'])
assert inference_job['status'] == InferenceJobStatus.ERRORED
def make_dataset_info():
name = gen()
task = gen()
file_path = DATASET_FILE_PATH
return (name, task, file_path)