def wrapped(*args, **kwargs):
auth_header = request.headers.get('authorization', None)
token = extract_token_from_header(auth_header)
auth = decode_token(token)
if auth.get('user_type') not in user_types:
raise UnauthorizedError()
metastore = MetaStore()
with metastore:
user = metastore.get_user_by_id(user_id=auth.get("user_id"))
if user is None:
raise UnauthorizedError("user is not exist")
return f(auth, *args, **kwargs)
def ban_user(auth, params):
admin = g.admin
with admin:
user = admin.get_user_by_email(params['email'])
if user is not None:
# Only superadmins can ban admins
if auth['user_type'] != UserType.SUPERADMIN and \
user['user_type'] in [UserType.ADMIN, UserType.SUPERADMIN]:
raise UnauthorizedError()
# Cannot ban yourself
if auth['user_id'] == user['id']:
raise UnauthorizedError()
return jsonify(admin.ban_user(email=params['email']))
def delete_model(auth, model_id):
admin = g.admin
with admin:
# Non-admins cannot delete others' models
if auth['user_type'] in [UserType.MODEL_DEVELOPER]:
model = admin.get_model(model_id)
if auth['user_id'] != model['user_id']:
raise UnauthorizedError()
return jsonify(admin.delete_model(model_id))
def get_inference_jobs_by_user(auth, params):
admin = g.admin
assert 'user_id' in params
# Non-admins can only get their own jobs
if auth['user_type'] in [UserType.APP_DEVELOPER, UserType.MODEL_DEVELOPER] \
and auth['user_id'] != params['user_id']:
raise UnauthorizedError()
with admin:
return jsonify(admin.get_inference_jobs_by_user(**params))
def get_train_jobs_by_user(auth, params):
admin = g.admin
if 'user_id' in params:
# Non-admins can only get their own jobs
if auth['user_type'] in [UserType.APP_DEVELOPER, UserType.MODEL_DEVELOPER] \
and auth['user_id'] != params['user_id']:
raise UnauthorizedError()
user_id = auth['user_id']
with admin:
return jsonify(admin.get_train_jobs_by_user(user_id=user_id))
def download_model_file(auth, model_id):
admin = g.admin
with admin:
# Non-admins cannot access others' models
if auth['user_type'] in [UserType.MODEL_DEVELOPER]:
model = admin.get_model(model_id)
if auth['user_id'] != model['user_id']:
raise UnauthorizedError()
model_file = admin.get_model_file(model_id)
res = make_response(model_file)
res.headers.set('Content-Type', 'application/octet-stream')
return res
def create_user(auth, params):
admin = g.admin
# Only superadmins can create admins
if auth['user_type'] != UserType.SUPERADMIN and \
params['user_type'] in [UserType.ADMIN, UserType.SUPERADMIN]:
raise UnauthorizedError()
with admin:
return jsonify(
admin.create_user(
email=params['email'],
password=params['password'],
user_type=params['user_type'],
))
def generate_user_token(params):
admin = g.admin
# Error will be thrown here if credentials are invalid
with admin:
user = admin.authenticate_user(**params)
# User cannot be banned
if user.get('banned_date'
) is not None and datetime.now() > user.get('banned_date'):
raise UnauthorizedError('User is banned')
token = generate_token(user)
return jsonify({
'user_id': user['id'],
'user_type': user['user_type'],
'token': token
})
def create_train_job(auth, params):
admin = g.admin
budget = params['budget'] if 'budget' in params else {}
budget = {
BudgetOption.TIME_HOURS: 0.1,
BudgetOption.GPU_COUNT: 0,
**budget
}
feed_params = {}
feed_params['user_id'] = auth['user_id']
feed_params['app'] = params['app']
feed_params['task'] = params['task']
feed_params['train_dataset_id'] = params['train_dataset_id']
feed_params['val_dataset_id'] = params['val_dataset_id']
feed_params['budget'] = budget
if "annotation_dataset_id" in params:
feed_params['annotation_dataset_id'] = params['annotation_dataset_id']
if 'model_ids' in params:
feed_params['model_ids'] = params['model_ids']
if 'train_args' in params:
feed_params['train_args'] = params['train_args']
with admin:
admin._services_manager.service_app_name = params['app']
# Ensure that datasets are owned by current user
dataset_attrs = ['train_dataset_id', 'val_dataset_id']
for attr in dataset_attrs:
if attr in params:
dataset_id = params[attr]
dataset = admin.get_dataset(dataset_id)
if auth['user_id'] != dataset['owner_id']:
raise UnauthorizedError(
'You have no access to dataset of ID "{}"'.format(
dataset_id))
return jsonify(admin.create_train_job(**feed_params))