rafiki.client

class rafiki.client.Client(admin_host: str = '127.0.0.1', admin_port: int = '3000')

Initializes the Client to connect to a running Rafiki Admin instance that the Client connects to.

Parameters

• admin_host – Host of Rafiki Admin

• admin_port – Port of Rafiki Admin

login(email: str, password: str) → Dict[str, Any]

Creates a login session as a Rafiki user. You will have to be logged in to perform any actions.

App developers can create, list and stop train and inference jobs, as well as list models. Model developers can create and list models.

The login session (the session token) expires in 1 hour.

Parameters

• email – User’s email

• password – User’s password

Returns

Logged-in user as dictionary

get_current_user() → Dict[str, Any]

Gets currently logged in user’s data.

Returns

Current user as dictionary, or None if client is not logged in

logout()

Clears the current login session.

create_user(email: str, password: str, user_type: rafiki.constants.UserType) → Dict[str, Any]

Creates a Rafiki user.

Only admins can create users (except for admins). Only superadmins can create admins.

Parameters

• email – The new user’s email

• password – The new user’s password

• user_type – The new user’s type

Returns

Created user as dictionary

get_users() → List[Dict[str, Any]]

Lists all Rafiki users.

Only admins can list all users.

Returns

List of users as list of dictionaries

ban_user(email: str) → Dict[str, Any]

Bans a Rafiki user, disallowing logins.

This action is irrevisible. Only admins can ban users (except for admins). Only superadmins can ban admins.

Parameters

email – The user’s email

Returns

Banned user as dictionary

create_dataset(name: str, task: str, dataset_path: str = None, dataset_url: str = None) → Dict[str, Any]

Creates a dataset on Rafiki, either by uploading the dataset file from your filesystem or specifying a URL where the dataset file can be downloaded. The dataset should be in a format specified by the task Either dataset_url or dataset_path should be specified.

Only admins, model developers and app developers can manage their own datasets.

Parameters

• name – Name for the dataset, does not need to be unique

• task – Task associated to the dataset

• dataset_path – Path to the dataset file to upload from the local filesystem

• dataset_url – Publicly accessible URL where the dataset file can be downloaded

Returns

Created dataset as dictionary

get_datasets(task: str = None) → List[Dict[str, Any]]

Lists all datasets owned by the current user, optionally filtering by task.

Parameters

task – Task name

Returns

List of datasets as list of dictionaries

create_model(name: str, task: str, model_file_path: str, model_class: str, dependencies: Dict[str, str] = None, access_right: rafiki.constants.ModelAccessRight = 'PRIVATE', docker_image: str = None) → Dict[str, Any]

Creates a model on Rafiki.

Only admins & model developers can manage models.

Parameters

• name – Name of the model, which must be unique across all models added by the current user

• task – Task associated with the model, where the model must adhere to the specification of the task

• model_file_path – Path to a single Python file that contains the definition for the model class

• model_class – The name of the model class inside the Python file. This class should implement rafiki.model.BaseModel

• dependencies – List of Python dependencies & their versions

• access_right – Model access right

• docker_image – A custom Docker image that extends rafikiai/rafiki_worker, publicly available on Docker Hub.

Returns

Created model as dictionary

Refer to Model Development Guide for more details on how to write & test models for Rafiki.

model_file_path should point to a single file that contains all necessary Python code for the model’s implementation. If the Python file imports any external Python modules, you should list it in dependencies or create a custom docker_image.

If a model’s access_right is set to PUBLIC, this model will be publicly available to all other users on Rafiki for training and inference. By default, a model’s access is PRIVATE.

dependencies should be a dictionary of { <dependency_name>: <dependency_version> }, where <dependency_name> corresponds to the name of the Python Package Index (PyPI) package (e.g. tensorflow) and <dependency_version> corresponds to the version of the PyPI package (e.g. 1.12.0). Refer to Configuring the Model’s Environment to understand more about this option.

get_model(model_id: str) → Dict[str, Any]

Retrieves details of a single model.

Model developers can only view their own models.

Parameters

model_id – ID of model

Returns

Model as dictionary

download_model_file(model_id: str, out_model_file_path: str) → Dict[str, any]

Downloads the Python model class file for the Rafiki model.

Model developers can only download their own models.

Parameters

• model_id – ID of model

• out_model_file_path – Absolute/relative path to save model class file to

Returns

Model as dictionary

get_available_models(task: str = None) → List[Dict[str, Any]]

Lists all Rafiki models available to the current user, optionally filtering by task.

Parameters

task – Task name

Returns

Available models as list of dictionaries

delete_model(model_id: str) → Dict[str, Any]

Deletes a single model. Models that have been used in train jobs cannot be deleted.

Model developers can only delete their own models.

Parameters

model_id (str) – ID of model

Returns

Deleted model as dictionary

create_train_job(app: str, task: str, train_dataset_id: str, val_dataset_id: str, budget: Dict[rafiki.constants.BudgetOption, Any], models: List[str] = None, train_args: Dict[str, any] = None) → Dict[str, Any]

Creates and starts a train job on Rafiki.

A train job is uniquely identified by user, its associated app, and the app version (returned in output).

Only admins, model developers & app developers can manage train jobs. Model developers & app developers can only manage their own train jobs.

Parameters

• app – Name of the app associated with the train job

• task – Task associated with the train job, the train job will train models associated with the task

• train_dataset_id – ID of the train dataset, previously created on Rafiki

• val_dataset_id – ID of the validation dataset, previously created on Rafiki

• budget – Budget for train job

• models – List of IDs of model to use for train job. Defaults to all available models

• train_args – Additional arguments to pass to models during training, if any. Refer to the task’s specification for appropriate arguments

Returns

Created train job as dictionary

If models is unspecified, all models accessible to the user for the specified task will be used.

budget should be a dictionary of { <budget_type>: <budget_amount> }, where <budget_type> is one of rafiki.constants.BudgetOption and <budget_amount> specifies the amount for the associated budget option.

The following describes the budget options available:

Budget Option	Description
TIME_HOURS	Max no. of hours to train (soft target). Defaults to 0.1.
GPU_COUNT	No. of GPUs to allocate for training, across all models. Defaults to 0.
MODEL_TRIAL_COUNT	Max no. of trials to conduct for each model (soft target). -1 for unlimited. Defaults to -1.

get_train_jobs_by_user(user_id: str) → List[Dict[str, Any]]

Lists all of user’s train jobs on Rafiki.

Parameters

user_id – ID of the user

Returns

Train jobs as list of dictionaries

get_train_jobs_of_app(app: str) → List[Dict[str, Any]]

Lists all of current user’s train jobs associated to the app name on Rafiki.

Parameters

app – Name of the app

Returns

Train jobs as list of dictionaries

get_train_job(app: str, app_version: int = -1) → Dict[str, Any]

Retrieves details of the current user’s train job identified by an app and an app version, including workers’ details.

Parameters

• app – Name of the app

• app_version – Version of the app (-1 for latest version)

Returns

Train job as dictionary

stop_train_job(app: str, app_version: int = -1) → Dict[str, Any]

Prematurely stops the current user’s train job identified by an app and an app version. Otherwise, the train job should stop by itself when its budget is reached.

Parameters

• app – Name of the app

• app_version – Version of the app (-1 for latest version)

Returns

Stopped train job as dictionary

get_trial(trial_id: str) → Dict[str, Any]

Gets a specific trial.

Parameters

trial_id – ID of trial

Returns

Trial as dictionary

get_best_trials_of_train_job(app: str, app_version: int = -1, max_count: int = 2) → List[Dict[str, Any]]

Lists the best scoring trials of the current user’s train job identified by an app and an app version, ordered by descending score.

Parameters

• app – Name of the app

• app_version – Version of the app (-1 for latest version)

• max_count – Maximum number of trials to return

Returns

Trials as list of dictionaries

get_trials_of_train_job(app: str, app_version: int = -1) → List[Dict[str, Any]]

Lists all trials of the current user’s train job identified by an app and an app version, ordered by when the trial started.

Parameters

• app – Name of the app

• app_version – Version of the app (-1 for latest version)

Returns

Trials as list of dictionaries

get_trial_logs(trial_id: str) → Dict[str, Any]

Gets the logs for a specific trial.

Parameters

trial_id – ID of trial

Returns

Logs of trial as dictionary

get_trial_parameters(trial_id: str) → Dict[str, Union[str, int, float, numpy.ndarray]]

Gets parameters of the model associated with the trial. The trial’s model parameters must have been saved.

Parameters

trial_id – ID of trial

Returns

Parameters of the trained model associated with the trial

load_trial_model(trial_id: str, ModelClass: Type[rafiki.model.model.BaseModel]) → rafiki.model.model.BaseModel

Loads an instance of a trial’s model with the trial’s knobs & parameters.

Before this, you must have the trial’s model class file already in your local filesystem, the dependencies of the model must have been installed separately, and the model class must have been imported and passed into this method.

Wraps get_trial_parameters() and get_trial().

Parameters

• trial_id – ID of trial

• ModelClass – model class that conincides with the trial’s model class

Returns

A trained model instance of ModelClass, loaded with the trial’s knobs and parameters

create_inference_job(app: str, app_version: int = -1, budget: Dict[rafiki.constants.InferenceBudgetOption, Any] = None) → Dict[str, Any]

Creates and starts a inference job on Rafiki with the best-scoring trials of the associated train job. The train job must have the status of STOPPED.The inference job would be tagged with the train job’s app and app version. Throws an error if an inference job of the same train job is already running.

In this method’s response, predictor_host is this inference job’s predictor’s host.

Only admins, model developers & app developers can manage inference jobs. Model developers & app developers can only manage their own inference jobs.

Parameters

• app – Name of the app identifying the train job to use

• app_version – Version of the app identifying the train job to use

• budget – Budget for inference job

Returns

Created inference job as dictionary

budget should be a dictionary of { <budget_type>: <budget_amount> }, where <budget_type> is one of rafiki.constants.InferenceBudgetOption and <budget_amount> specifies the amount for the associated budget option.

The following describes the budget options available:

Budget Option	Description
GPU_COUNT	No. of GPUs to allocate for inference, across all trials. Defaults to 0.

get_inference_jobs_by_user(user_id: str) → List[Dict[str, Any]]

Lists all of user’s inference jobs on Rafiki.

Parameters

user_id – ID of the user

Returns

Inference jobs as list of dictionaries

get_inference_jobs_of_app(app: str) → List[Dict[str, Any]]

Lists all inference jobs associated to an app on Rafiki.

Parameters

app – Name of the app

Returns

Inference jobs as list of dictionaries

get_running_inference_job(app: str, app_version: int = -1) → Dict[str, Any]

Retrieves details of the running inference job identified by an app and an app version, including workers’ details.

Parameters

• app – Name of the app

• app_version – Version of the app (-1 for latest version)

Returns

Inference job as dictionary

stop_inference_job(app: str, app_version: int = -1) → Dict[str, Any]

Stops the inference job identified by an app and an app version.

Parameters

• app – Name of the app

• app_version – Version of the app (-1 for latest version)

Returns

Stopped inference job as dictionary

stop_all_jobs()

Stops all train and inference jobs on Rafiki.

Only the superadmin can call this.