How to edit endpoint settings on Unified AI Platform using Python?

I have successfully created an endpoint on Unified Cloud AI Platform and have deployed two Models to it – Model A and Model B with 20% and 80% of the traffic respectively. Now, on the Cloud Console (the UI) I get an option to Edit Settings and change the traffic split to 30% and 70% respectively and the Models are deployed. However I am unable to figure out how to do this using the Python Client API.

The documentation provided over here is inadequate to understand how we could do this. Any help will be appreciated.

Answer

The documentation for AI Platform Unified does not yet have an example on how to edit the traffic using python. Here is the code for that:

NOTE: Don’t forget to update the values for end_point (endpoint ID),project (project ID), model_id_1 and model_id_2 before running the code.

from google.cloud import aiplatform
from google.cloud import aiplatform_v1


def update_endpoint_traffic(
    end_point: str, 
    project: str, 
    location: str = "us-central1",
    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
    timeout: int = 7200,
):
    # The AI Platform services require regional API endpoints.
    client_options = {"api_endpoint": api_endpoint}
    # Initialize client that will be used to create and send requests.
    # This client only needs to be created once, and can be reused for multiple requests.
    client = aiplatform.gapic.EndpointServiceClient(client_options=client_options)
    client_model = aiplatform_v1.services.model_service.ModelServiceClient(client_options=client_options)

    deployed_model_id_list = []
    model_id_1 = 'xxxxxxxx' # place your model id here
    model_id_2 = 'xxxxxxxx' # place your model id here
    model_list = [f'projects/{project}/locations/{location}/models/{model_id_1}',f'projects/{project}/locations/{location}/models/{model_id_2}']

    for model in model_list:
        model_request = aiplatform_v1.types.GetModelRequest(name=model)
        model_info = client_model.get_model(request=model_request)
        deployed_models_info = model_info.deployed_models
        deployed_model_id=model_info.deployed_models[0].deployed_model_id
        deployed_model_id_list.append(deployed_model_id)

    traffic_split = {deployed_model_id_list[0]: 60, deployed_model_id_list[1]:40} #update values of 60 and 40 to desired traffic split ex.(30 70)

    name=f'projects/{project}/locations/{location}/endpoints/{end_point}'

    endpoint = aiplatform_v1.types.Endpoint(name=name,traffic_split=traffic_split)
    update_endpoint = aiplatform_v1.types.UpdateEndpointRequest(endpoint=endpoint)
    client.update_endpoint(request=update_endpoint)

update_endpoint_traffic(end_point='your-endpoint-id',project='your-project-id')