def list_models(request: Request, project_id: int) -> JSONResponse:
"""Get models list.
Args:
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
resp = requests.get(f'{url}/api/2.0/preview/mlflow/registered-models/list')
registered_models = []
for model in resp.json().get('registered_models_detailed', []):
registered_models.append({
'id': model.get('registered_model', {}).get('name'),
'project_id': project_id,
'creation_timestamp': model.get('creation_timestamp'),
'last_updated_timestamp': model.get('last_updated_timestamp')
})
return JSONResponse(registered_models)
def get_model(request: Request, model_id: Text, project_id: int) -> JSONResponse:
"""Get model.
Args:
model_id {Text}: model id (name)
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
model_resp = requests.post(
url=f'{url}/api/2.0/preview/mlflow/registered-models/get-details',
json={'registered_model': {'name': model_id}}
)
if model_resp.status_code != HTTPStatus.OK:
return error_response(
http_response_code=model_resp.status_code,
message=model_resp.json().get('message')
)
registered_model_detailed = model_resp.json().get('registered_model_detailed', {})
model = {
'id': registered_model_detailed.get('registered_model', {}).get('name'),
'project_id': project_id,
'creation_timestamp': registered_model_detailed.get('creation_timestamp'),
'last_updated_timestamp': registered_model_detailed.get('last_updated_timestamp')
}
return JSONResponse(model)
def delete_model(request: Request, model_id: Text, project_id: int) -> JSONResponse:
"""Delete model.
Args:
model_id {Text}: model id (name)
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request, {
'project_id': project_id,
'model_id': model_id
})
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
model_resp = requests.delete(
url=f'{url}/api/2.0/preview/mlflow/registered-models/delete',
json={'registered_model': {'name': model_id}}
)
if model_resp.status_code != HTTPStatus.OK:
return error_response(
http_response_code=model_resp.status_code,
message=model_resp.json().get('message')
)
return JSONResponse({'model_id': model_id})
def list_runs(request: Request, project_id: int,
experiment_id: Text) -> JSONResponse:
"""Get runs list.
Args:
project_id {int}: project id
experiment_id {Text}: experiment id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
resp = requests.post(url=f'{url}/api/2.0/preview/mlflow/runs/search',
json={'experiment_ids': [experiment_id]})
if resp.status_code != HTTPStatus.OK:
return error_response(http_response_code=resp.status_code,
message=resp.json().get('message'))
runs = resp.json().get('runs', [])
for run in runs:
run['id'] = run.get('info', {}).get('run_id')
return JSONResponse(runs)
def get_model_version_uri(project_id: int, model_id: Text,
version: Text) -> Text:
"""Get model version URI.
Args:
project_id {int}: project id
model_id {Text}: model id (name)
version {Text}: model version
Returns:
Text: path to model package for model version
"""
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
model_resp = requests.post(
url=f'{url}/api/2.0/preview/mlflow/registered-models/get-details',
json={'registered_model': {
'name': model_id
}})
if model_resp.status_code != HTTPStatus.OK:
if model_resp.status_code == HTTPStatus.NOT_FOUND:
raise RegisteredModelNotFoundError(f'Model {model_id} not found')
raise Exception(model_resp.text)
this_model_versions = filter_model_versions(get_model_versions(project_id),
model_id)
for ver in this_model_versions:
if ver.get('model_version', {}).get('version') == version:
return ver.get('source')
raise RegisteredModelNotFoundError(
f'Version {version} of {model_id} not found')
def delete_experiment(request: Request, experiment_id: Text,
project_id: int) -> JSONResponse:
"""Delete experiment.
Args:
experiment_id {Text}: experiment id
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request, {
'project_id': project_id,
'experiment_id': experiment_id
})
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
experiment_resp = requests.post(
url=f'{url}/api/2.0/preview/mlflow/experiments/delete',
json={'experiment_id': experiment_id})
if experiment_resp.status_code != HTTPStatus.OK:
return error_response(http_response_code=experiment_resp.status_code,
message=experiment_resp.json().get('message'))
return JSONResponse({'experiment_id': experiment_id})
def list_experiments(request: Request, project_id: int) -> JSONResponse:
"""Get experiments list.
Args:
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
resp = requests.get(url=f'{url}/api/2.0/preview/mlflow/experiments/list')
experiments = []
for exp in resp.json().get('experiments'):
experiment_id = exp.get('experiment_id')
runs_resp = requests.get(
f'{url}/api/2.0/preview/mlflow/runs/search?experiment_ids=[{experiment_id}]'
)
runs = runs_resp.json().get('runs', [])
creation_time = ''
last_update_time = ''
"""
if corresponding tags are empty then fields:
* creation_time = start_time of the first run;
* last_update_time = end_time of the last run.
"""
if len(runs) > 0:
creation_time = runs[len(runs) - 1].get('info',
{}).get('start_time')
last_update_time = runs[0].get('info', {}).get('end_time')
tags = {tag['key']: tag['value'] for tag in exp.get('tags', [])}
experiments.append({
'id':
experiment_id,
'user_id':
tags.get('user_id', ''),
'name':
exp.get('name'),
'artifact_location':
exp.get('artifact_location'),
'lifecycle_stage':
exp.get('lifecycle_stage'),
'last_update_time':
tags.get('last_update_time', last_update_time),
'creation_time':
tags.get('creation_time', creation_time),
'description':
tags.get('mlflow.note.content', ''),
'project_id':
tags.get('project_id', project_id)
})
return JSONResponse(experiments)
def register_model(
request: Request,
project_id: int,
name: Text = Form(...),
source: Text = Form(...),
run_id: Text = Form(...)
) -> JSONResponse:
"""Register model.
Args:
project_id {int}: project id
name {Text}: model name
source {Text}: path to model package
run_id {Text}: run id
Returns:
starlette.responses.JSONResponse
"""
log_request(request, {
'project_id': project_id,
'name': name,
'source': source,
'run_id': run_id
})
project_manager = ProjectManager()
if not is_model(source):
raise ModelDoesNotExistError(f'Model {source} does not exist or is not MLflow model')
url = project_manager.get_internal_tracking_uri(project_id)
requests.post(
url=f'{url}/api/2.0/preview/mlflow/registered-models/create',
json={'name': name}
)
requests.post(
url=f'{url}/api/2.0/preview/mlflow/model-versions/create',
json={
'name': name,
'source': source,
'run_id': run_id
}
)
model_resp = requests.post(
url=f'{url}/api/2.0/preview/mlflow/registered-models/get-details',
json={'registered_model': {'name': name}}
)
registered_model_detailed = model_resp.json().get('registered_model_detailed', {})
model = {
'id': name,
'project_id': project_id,
'creation_timestamp': registered_model_detailed.get('creation_timestamp'),
'last_updated_timestamp': registered_model_detailed.get('last_updated_timestamp')
}
return JSONResponse(model, HTTPStatus.CREATED)
def get_experiment(request: Request, experiment_id: Text,
project_id: int) -> JSONResponse:
"""Get experiment.
Args:
experiment_id {Text}: experiment id
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
experiment_resp = requests.get(
url=
f'{url}/api/2.0/preview/mlflow/experiments/get?experiment_id={experiment_id}'
)
if experiment_resp.status_code != HTTPStatus.OK:
return error_response(http_response_code=experiment_resp.status_code,
message=experiment_resp.json().get('message'))
experiment = experiment_resp.json().get('experiment')
experiment_id = experiment.get('experiment_id')
runs_resp = requests.get(
f'{url}/api/2.0/preview/mlflow/runs/search?experiment_ids=[{experiment_id}]'
)
runs = runs_resp.json().get('runs', [])
creation_time = ''
last_update_time = ''
"""
if corresponding tags are empty then fields:
* creation_time = start_time of the first run;
* last_update_time = end_time of the last run.
"""
if len(runs) > 0:
creation_time = runs[len(runs) - 1].get('info', {}).get('start_time')
last_update_time = runs[0].get('info', {}).get('end_time')
experiment['id'] = experiment.pop('experiment_id')
tags = {tag['key']: tag['value'] for tag in experiment.pop('tags', [])}
experiment['description'] = tags.get('mlflow.note.content', '')
experiment['user_id'] = tags.get('user_id', '')
experiment['project_id'] = tags.get('project_id', project_id)
experiment['creation_time'] = tags.get('creation_time', creation_time)
experiment['last_update_time'] = tags.get('last_update_time',
last_update_time)
return JSONResponse(experiment)
def get_model_versions(project_id: int) -> List[Dict]:
"""Get all model versions by tracking server uri
Args:
project_id {int}: project id
Returns:
List[Dict]: list of model versions
"""
project_manager = ProjectManager()
tracking_uri = project_manager.get_internal_tracking_uri(project_id)
model_versions_resp = requests.get(
url=f'{tracking_uri}/api/2.0/preview/mlflow/model-versions/search')
model_versions = model_versions_resp.json().get('model_versions_detailed')
return model_versions
def list_artifacts(request: Request, project_id: int,
run_id: Text) -> JSONResponse:
"""Get artifacts list.
Args:
project_id {int}: project id
run_id {int}: run_id
Returns:
starlette.responses.JSONResponse
"""
log_request(request, {'project_id': project_id, 'run_id': run_id})
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
runs_resp = requests.get(
url=f'{url}/api/2.0/preview/mlflow/artifacts/list?run_id={run_id}')
if runs_resp.status_code != HTTPStatus.OK:
return error_response(http_response_code=runs_resp.status_code,
message=runs_resp.json().get('message'))
runs = runs_resp.json()
root_uri = runs.get('root_uri')
files = runs.get('files', [])
runs_list = []
for i, file in enumerate(files):
runs_resp = requests.get(
url=f'{url}/api/2.0/preview/mlflow/runs/get?run_id={run_id}', )
run = runs_resp.json()
run_info = run.get('run', {}).get('info', {})
experiment_id = run_info.get('experiment_id')
runs_list.append({
'id': f'{project_id}{experiment_id}{run_id}{i}',
'project_id': project_id,
'experiment_id': experiment_id,
'run_id': run_id,
'type': str(get_artifact_type(root_uri, file)),
'creation_timestamp': run_info.get('start_time'),
'root_uri': root_uri,
'path': file.get('path')
})
return JSONResponse(runs_list)
def ping(request: Request, project_id: int) -> JSONResponse:
"""Ping project's tracking server.
Args:
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
project = project_manager.get_project(project_id)
try:
requests.get(url)
return JSONResponse(project, HTTPStatus.OK)
except requests.exceptions.ConnectionError:
return JSONResponse(project, HTTPStatus.BAD_REQUEST)
def get_run(request: Request, run_id: Text, project_id: int) -> JSONResponse:
"""Get run.
Args:
run_id {Text}: run id
project_id {int}: project id
Returns:
starlette.responses.JSONResponse
"""
log_request(request)
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
resp = requests.get(
url=f'{url}/api/2.0/preview/mlflow/runs/get?run_id={run_id}', )
if resp.status_code != HTTPStatus.OK:
return error_response(http_response_code=resp.status_code,
message=resp.json().get('message'))
run = resp.json().get('run')
run['id'] = run.get('info', {}).get('run_id')
return JSONResponse(run)
def check_if_project_and_model_exist(project_id: int, model_id: Text) -> None:
"""Check if project and model are exist.
Args:
project_id {int}: project id
model_id {Text}: model id (name)
Raises:
ProjectNotFoundError: if project does not exists;
RegisteredModelNotFoundError: if project exists, but model does not exist
"""
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
if url is None:
raise ProjectNotFoundError(f'Project with ID {project_id} not found')
model_resp = requests.post(
url=f'{url}/api/2.0/preview/mlflow/registered-models/get-details',
json={'registered_model': {
'name': model_id
}})
if model_resp.status_code != HTTPStatus.OK:
raise RegisteredModelNotFoundError(f'Model {model_id} not found')
def create_experiment(
request: Request,
project_id: int,
user_id: Text = Form(''),
name: Text = Form(''),
description: Text = Form('')
) -> JSONResponse:
"""Create experiment
Args:
project_id {int}: project id
user_id {Text}: user id (name)
name {Text}: experiment name
description {Text}: experiment description
Returns:
starlette.responses.JSONResponse
"""
log_request(
request, {
'project_id': project_id,
'user_id': user_id,
'name': name,
'description': description
})
project_manager = ProjectManager()
url = project_manager.get_internal_tracking_uri(project_id)
creation_resp = requests.post(
url=f'{url}/api/2.0/preview/mlflow/experiments/create',
json={'name': name})
creation_json = creation_resp.json()
experiment_id = creation_json.get('experiment_id')
if creation_resp.status_code != HTTPStatus.OK:
return error_response(http_response_code=creation_resp.status_code,
message=creation_json.get('message'))
utc_timestamp = get_utc_timestamp()
tags = {
'mlflow.note.content': description,
'user_id': user_id,
'project_id': str(project_id),
'creation_time': utc_timestamp,
'last_update_time': utc_timestamp
}
for key, value in tags.items():
requests.post(
url=f'{url}/api/2.0/preview/mlflow/experiments/set-experiment-tag',
json={
'experiment_id': experiment_id,
'key': key,
'value': value
})
experiment_request = requests.get(
url=
f'{url}/api/2.0/preview/mlflow/experiments/get?experiment_id={experiment_id}'
)
experiment = experiment_request.json().get('experiment')
tags = {tag['key']: tag['value'] for tag in experiment.pop('tags', [])}
experiment['description'] = tags.get('mlflow.note.content', '')
experiment['user_id'] = tags.get('user_id', '')
experiment['project_id'] = tags.get('project_id', '')
experiment['creation_time'] = tags.get('creation_time', '')
experiment['last_update_time'] = tags.get('last_update_time', '')
return JSONResponse(experiment, HTTPStatus.CREATED)