Work with ML Training Jobs with Viam's ML Training API

The ML training API allows you to get information about and cancel ML training jobs taking place on the Viam app.

Establish a connection

To use the Viam ML training client API, you first need to instantiate a ViamClient and then instantiate an MLTrainingClient. See the following example for reference.

Select Include secret on the Code sample page of the CONNECT tab of the Viam app to obtain your API key and API key ID values.

import asyncio

from viam.rpc.dial import DialOptions, Credentials
from viam.app.viam_client import ViamClient


async def connect() -> ViamClient:
    dial_options = DialOptions(
      credentials=Credentials(
        type="api-key",
        # Replace "<API-KEY>" (including brackets) with your machine's API key
        payload='<API-KEY>',
      ),
      # Replace "<API-KEY-ID>" (including brackets) with your machine's
      # API key ID
      auth_entity='<API-KEY-ID>'
    )
    return await ViamClient.create_from_dial_options(dial_options)


async def main():

    # Make a ViamClient
    viam_client = await connect()
    # Instantiate an MLTrainingClient to run ML training client API methods on
    ml_training_client = viam_client.ml_training_client

    viam_client.close()

if __name__ == '__main__':
    asyncio.run(main())

Once you have instantiated an MLTrainingClient, you can run the following API methods against the MLTrainingClient object (named ml_training_client in the examples).

API

The ML training client API supports the following methods:

Method NameDescription
GetTrainingJobGet training job metadata by ID.
ListTrainingJobsGet training job metadata for all jobs within an organization.
SubmitTrainingJobSubmit a training job.
CancelTrainingJobCancel the specified training job.

GetTrainingJob

Get training job metadata.

Parameters:

Returns:

job_metadata = await ml_training_client.get_training_job(
    id="INSERT YOUR JOB ID")

For more information, see the Python SDK Docs.

ListTrainingJobs

Get training job metadata for all jobs within an organization.

Parameters:

  • org_id (str): The ID of your organization to request training job metadata from. Retrieve this value with the fleet management API’s ListOrganizations().
  • training_status (Optional[TrainingStatus.ValueType]): The status of training jobs you want to filter the list by. If you leave this unspecified, all training jobs for your organization are returned.

Returns:

jobs_metadata = await ml_training_client.list_training_jobs(
    org_id="INSERT YOUR ORG ID")

first_job_id = jobs_metadata[1].id

For more information, see the Python SDK Docs.

SubmitTrainingJob

Submit a training job.

Parameters:

  • org_id (str) (required): The ID of the organization to submit the training job to. To retrieve this, expand your organization’s dropdown in the top right corner of the Viam app, select Settings, and copy Organization ID.
  • dataset_id (str) (required): The ID of the dataset to train the ML model on. To retrieve this, navigate to your dataset’s page in the Viam app, click in the left-hand menu, and click Copy dataset ID.
  • model_name (str) (required): The name of the ML model you wish to train.
  • model_version (str) (required): The version of the ML model you’re training. This string must be unique from any previous versions you’ve set.
  • model_type (ModelType.ValueType) (required): The type of the ML model. Options: ModelType.MODEL_TYPE_SINGLE_LABEL_CLASSIFICATION, ModelType.MODEL_TYPE_MULTI_LABEL_CLASSIFICATION, ModelType.MODEL_TYPE_OBJECT_DETECTION.
  • tags (List(str)) (required): The labels to train on.

Returns:

  • (str): The ID of the training job created.
job_id = await ml_training_client.submit_training_job(
  "3988fd7d-b1c0-4f36-9842-0c666f374d3c",
  "66491a0666a6c8aa4a33e43e",
  "your-model-name",
  "1",
  "ModelType.MODEL_TYPE_SINGLE_LABEL_CLASSIFICATION",
  tags)

For more information, see the Python SDK Docs.

Parameters:

  • datasetId (String) (required): The ID of the dataset to train the ML model on. To retrieve this, navigate to your dataset’s page in the Viam app, click in the left-hand menu, and click Copy dataset ID.
  • modelName (String) (required): The name of the ML model to train.
  • modelType (ModelType) (required): The type of the ML model. Options: MODEL_TYPE_SINGLE_LABEL_CLASSIFICATION, MODEL_TYPE_MULTI_LABEL_CLASSIFICATION, or MODEL_TYPE_OBJECT_DETECTION.
  • modelVersion (String) (required): The version of the ML model you’re training. This string must be unique from any previous versions you’ve set.
  • organizationId (String) (required): The ID of the organization to submit the training job to. To retrieve this, expand your organization’s dropdown in the top right corner of the Viam app, select Settings, and copy Organization ID.
  • tags (List)<String> (required): The labels to train on.

Returns:

  • (String) (required): The ID of the training job created.
final tags = <String>['A', 'B', 'C']; // Creates list of tags.

String jobID = ml_training.submitTrainingJob("66491a0666a6c8aa4a33e43e", "your-model-name", "MODEL_TYPE_SINGLE_LABEL_CLASSIFICATION", "1", "3988fd7d-b1c0-4f36-9842-0c666f374d3c", tags);

For more information, see the Flutter SDK Docs.

CancelTrainingJob

Cancel the specified training job.

Parameters:

Returns:

  • None

Raises:

  • GRPCError: If no training job exists with the given ID or cancellation was otherwise unsuccessful.
await ml_training_client.cancel_training_job(
    id="INSERT YOUR JOB ID")

For more information, see the Python SDK Docs.

Have questions, or want to meet other people working on robots? Join our Community Discord.

If you notice any issues with the documentation, feel free to file an issue or edit this file.