Train AI and ML models

Preview

This product is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the Service Specific Terms, and the Additional Terms for Generative AI Preview Products. Pre-GA products are available "as is" and might have limited support. For more information, see the launch stage descriptions.

Learn how to train AI and ML models in the Google Cloud Data Agent Kit extension for Visual Studio Code.

In this quickstart, you use a session template and a sample Jupyter notebook to predict New York City taxi tip amounts. Using a remote Jupyter kernel with PySpark, you try out various models like linear regression, random forest, and XGBoost. This process lets you perform distributed training and inference. It demonstrates scalability across multiple machines using Spark ML and the XGBoost library.

Although not covered in this quickstart, there are multiple ways to train AI and ML models using Google Cloud Data Agent Kit extension for Visual Studio Code:

• If your training dataset is large or you want the distributed training capabilities that Apache Spark offers, you can use Spark notebooks with remote kernels.
• If your dataset is in BigQuery and BigQuery ML supports your use case, you can use a BigQuery DataFrames notebook.
• If your dataset is small and you want to train your model locally, you can use a Python notebook.

Create a Spark runtime template

Serverless Spark Runtime templates let you start an Apache Spark session with a given set of configurations. To create a new Serverless Runtime template, complete the following steps:

1. In the IDE activity bar, click the Google Cloud Data Agent Kit icon.
2. In the Google Cloud Data Agent Kit menu, expand Apache Spark.
3. Expand Serverless and then click + Create serverless runtimes. A Serverless Runtime creation form appears.
4. In the Display Name field, enter ai-ml-tutorial.
5. Go to the Auto Scaling section.
6. Set spark.dynamicAllocation.enabled to false in the drop-down list. This setting is needed for XGBoost to work with Apache Spark.
7. Leave all other fields set to the default.
8. Click Submit.

Create a new notebook

Note: In this quickstart, you use a prebuilt sample notebook. You can also use the agent in your IDE to generate a Spark notebook with a prompt that describes your specific ML use case.

Next, create a new Spark notebook:

1. Under Apache Spark in the Google Cloud Data Agent Kit tab, click + New Spark Notebook.
2. Choose Remote Kernel for the kernel type.
3. Click Start with a sample notebook.
4. In the list of samples, select Data Science with PySpark and Distributed XGBoost. An untitled Jupyter notebook appears.

Train your model

1. In the notebook tab, click Run All. The kernel picker asks you to select a kernel to execute the notebook with.
2. Click Select Another Kernel.
3. Click Remote Spark Kernels.
4. Select the ai-ml-tutorial on Serverless Spark, the Runtime template that you created earlier.

Note: On Windows, Long Path support must be enabled to work with Spark Notebooks on Remote Kernels.

You see the following notification while the system creates your Serverless Spark session: Connecting to kernel: ai-ml-tutorial on Serverless Spark. When the notebook connects to the remote PySpark kernel, execution starts at the first cell. This process takes approximately two to three minutes.

Key Point: Managed Service for Apache Spark executes your training code entirely. The system doesn't download any part of your training dataset to your local machine.

Inspect your Spark session

1. In the Google Cloud Data Agent Kit tab, under Apache Spark, expand the ai-ml-tutorial Runtime template. The IDE displays the list of interactive sessions that you have created with this runtime template.
2. Locate the session that the system created by executing the notebook at the top of the list. Click the session to see its details. You can review the session configuration and the resources that the system consumed to execute your notebook.

Clean up

After successfully executing the notebook, perform the following cleanup steps.

1. In the Google Cloud Data Agent Kit tab, under Apache Spark, right-click Serverless and select List Serverless Runtimes. The list of Serverless Runtimes appears.
2. Click the Action menu for ai-ml-tutorial to list all the Interactive Sessions that the system created from your template.
3. Under Actions, click Delete.
4. Go back to the Serverless Runtimes window.
5. Under Actions for ai-ml-tutorial, click Delete.
6. Click Confirm to delete the template that you created for this tutorial.

What's next

• Managed Service for Apache Spark overview
• Transform data