Horizontal Scaling for AI Model Serving

Question

Pulumi · Accepted Answer

Horizontal scaling, in the context of AI model serving, involves spreading the workloads across multiple instances of the model that are running in parallel. This approach allows the system to handle more requests simultaneously and provides high availability and redundancy. There are multiple cloud services that offer the ability to scale AI model serving horizontally, depending on the specific requirements of the application, such as expected load, response time, and deployment preferences.

In this guide, we are going to write a Pulumi Python program that sets up AI model serving with horizontal scaling capabilities. For demonstration purposes, we'll focus on setting up a scalable AI model endpoint using Azure's Machine Learning compute resources, as Azure is a common platform for ML workloads, and it also appears in the Pulumi Registry Results you've provided.

Firstly, an Azure Machine Learning workspace must be set up; this provides a foundational component for managing machine learning artifacts. Using the `Workspace` class from the `azure-native.machinelearningservices` module, we create a workspace resource.

Secondly, we need a compute cluster where the model will be deployed, which can scale out (increase the number of compute nodes) or scale in (decrease the number of compute nodes) automatically. The `AmlCompute` resource in the same module can be used here.

Thirdly, the model is packaged into a "container image" using the `Model`, `Environment`, and `InferenceConfig` classes (not specifically mentioned in the Registry Results but are necessary components of the Azure ML SDK).

Finally, we deploy this packaged model onto the compute cluster through an `InferenceEndpoint`, which allows clients to consume the model over a REST API. This endpoint is capable of handling HTTP requests to serve predictions from the deployed model.

Let's write a Pulumi program below that sets up a horizontally scalable AI model serving endpoint on Azure:

```python
import pulumi
import pulumi_azure_native.machinelearningservices as azure_ml
from pulumi_azure_native.resources import ResourceGroup
from pulumi_azure_native.authorization import RoleAssignment, get_builtin_role_definition
from pulumi import Output

# Set up a resource group for our resources
resource_group = ResourceGroup("rg")

# Create an Azure ML Workspace
ml_workspace = azure_ml.Workspace(
    "mlw",
    resource_group_name=resource_group.name,
    location=resource_group.location,
    sku=azure_ml.SkuArgs(name="Enterprise")
)

# Create a scalable AML Compute cluster as the compute target where the model will be deployed
aml_compute = azure_ml.Compute(
    "compute",
    location=resource_group.location,
    resource_group_name=resource_group.name,
    # The name of the created Workspace
    workspace_name=ml_workspace.name,
    properties=azure_ml.ComputePropertiesArgs(
        compute_type="AmlCompute",
        properties=azure_ml.AmlComputeArgs(
            vm_size="STANDARD_D2_V2", # You can choose a different VM size as per your need
            scale_settings=azure_ml.ScaleSettingsArgs(
                min_node_count=0, # Start with no compute nodes - it will scale up as needed
                max_node_count=4, # Define max number of nodes to which cluster can scale out
            )
        ),
    ),
)

# Assume we have already registered the model and the scoring environment.
# Container image is built using registered model and environment.
# Deployment is done through creating an OnlineDeployment that references the container image.

# Assuming the rest of the setup (like model registration, environments, and container image creation)
# has already been completed, you can now create an endpoint for serving the model.

# Note: The following components (not explicitly listed in the provided Registry Results) are
# generally required to complete the setup. In practice, you would need to define the Model,
# Environment, and InferenceConfig resources. Since these are complex operations involving
# specific configurations, we'll focus on the Compute and InferenceEndpoint resources primarily.

# The InferenceEndpoint is created with autoscaling enabled via the container_resource_requirements
# which would allow horizontal scaling based on the number of incoming requests.
inference_endpoint = azure_ml.InferenceEndpoint(
    "infer-endpoint",
    kind="Deployment",
    location=resource_group.location,
    resource_group_name=resource_group.name,
    compute_name=aml_compute.name,
    deployment_name="model-deployment",
    properties=Output.all(resource_group.location, ml_workspace, aml_compute).apply(lambda args: {
        "Description": "Scalable model deployment",
        "Model": "modelName",  # This and the next params should match the model you have registered
        "Environment": "scoringEnvironment",
        "ContainerResourceRequirements": {
            "Cpu": 1.0,  # You define CPU and memory requests per instance here
            "MemoryInGB": 2.0,
            "ScoringTimeoutMs": 60000,  # Scoring timeout in milliseconds
        },
    })
)

# Export the Azure ML Model Endpoint URL
pulumi.export("endpoint_url", inference_endpoint.properties.apply(lambda props: props.get("scoringUri", "")))
```

In the above program, we have defined and created several Azure resources needed for serving an AI model. The `aml_compute` is the resource that provides the compute power for our model and is set up to scale between 0 to 4 nodes, depending on demand. The `inference_endpoint` is where the model will be accessed and served to clients. It's important to note that additional setup is required to register the model and set up the necessary environments and container images, which have not been explicitly covered in this code due to complexity.

We also have an export statement at the end which will provide us with the endpoint's URL when the deployment is successful. This URL can then be used to access the AI model for predictions.

In a real-world scenario, you would need to set up authentication for secure access, manage deployment versions, monitor the endpoints for usage and performance, and possibly use more advanced scaling options provided by the cloud provider. Azure Machine Learning service has many features to support the complete lifecycle of building, training, and deploying machine learning models at scale.