Training and Serving Custom ML Models

Aizen recommends auto-ML for model training and serving because it handles most scenarios well and is easy to execute. Use the configure training and configure prediction commands to build an auto-ML pipeline.

If you need specialized ML models, which are not supported by auto-ML, then you may create custom ML models. Aizen provides a configure runtime command that supports custom ML model training and serving. This command deploys a runtime service that provides a JupyterLab notebook interface for you to train a custom ML model. This is a mechanism very similar to Google’s Colab service.

Training a Custom ML Model

To train a custom ML model, follow these steps:

  1. Log in to the Aizen Jupyter console. See Using the Aizen Jupyter Console.

  2. Set the current working project.

    set project <project name>
  3. Configure a runtime using the configure runtime command. In the notebook, you will be guided through a template form with boxes and drop-down lists that you can complete to define the runtime. Specify the runtime name and select Model Training. Specify any pip install requirements file that needs to be run during initialization. Specify any datasets that need to be made available in the runtime. If required, specify a previously trained model name and the run that needs to be made available in the runtime.

    configure runtime
  4. Configure a resource for the runtime and start the runtime:

    configure resource
    start runtime <runtime name>
  5. Check the status of the runtime. This output will provide URLs for you to access a JupyterLab notebook connected to the runtime:

    status runtime <runtime name>
  6. Connect to the JupyterLab notebook URL, which is displayed in the status runtime output. This is a separate JupyterLab notebook from the Aizen Jupyter console. You may perform pip installs, data analysis, and any custom ML model training in the runtime JupyterLab notebook.

    Aizen provides a Python module with some helper functions. There are functions to list datasets, load datasets, save a trained custom ML model to the Aizen platform, and load a saved model from the Aizen platform. Use the import aizen command to access this module in the runtime JupyterLab notebook.

    import aizen
    aizen.help()
    aizen.help(“save_model”)
    aizen.help(“list_datasets”)
  7. After training your custom ML model, save the model using the Aizen Python module:

    import aizen
    aizen.save_model(…)
  8. If you plan to serve the trained model, create a serving Python module with pydantic functions to serve the model for inference requests. The serving module will be needed later when deploying the model. Each function is advertised as a FastAPI endpoint for REST requests. The functions must be adequately annotated. The serving Python module must load the model by calling the Aizen read_model function.

    import aizen
    aizen.read_model(…)

    Test the model serving module using the start and stop model serve functions:

    import aizen
    aizen.start_model_serve(…)
    aizen.stop_model_serve(…)

    After you have tested the serving Python module, download the serving Python module to your laptop and upload it to the Aizen Jupyter console notebook.

  9. Stop the runtime when it is no longer needed. This will terminate the runtime and any data associated with it. The runtime does not provide permanent storage. Any models not saved using the Aizen save_model will be lost. Any data downloaded into the runtime, other than via Aizen datasets, will not be retained. Any code developed in the runtime, such as serving Python modules, will not be retained. You must download serving Python modules to your laptop and upload them to the console notebook. The runtime JupyterLab notebook that trained the custom ML model will not be retained. It is recommended that you save the notebook in GitHub, cloud storage, or some other permanent storage.

    stop runtime <runtime name>

Serving a Custom ML Model

To serve a custom ML model, follow these steps:

  1. Log in to the Aizen Jupyter console. See Using the Aizen Jupyter Console.

  2. Set the current working project.

    set project <project name>
  3. Configure a runtime using the configure runtime command. In the notebook, you will be guided through a template form with boxes and drop-down lists that you can complete to define the runtime. Specify the runtime name and select Model Serving. Specify any pip install requirements file that needs to be run during initialization. Specify the trained model name and the run that needs to be made available in the runtime. Select the serving Python module and pydantic functions that serve the trained model. Each function is advertised as a FastAPI endpoint for REST requests. The functions must be adequately annotated. The serving Python module must load the model by calling the Aizen read_model function. During development, the serving module can be coded and tested in a runtime that was configured for training.

    configure runtime
  4. Configure a resource for the runtime and start the runtime:

    configure resource
    start runtime <runtime name>
  5. Check the status of the runtime. This output will provide URLs that provide REST endpoints for inference requests to the serving module.

    status runtime <runtime name>

Last updated