Custom model serving using Databricks Asset Bundles

MLOperator
New Contributor II

I am using MLFlow to register custom model (python model) in Unity Catalog, and Databricks Asset Bundle to create a serving endpoint for that custom model. 
I was able to create the serving endpoint using DABs, but I want to deploy the model by using custom alias instead of the entity_version.

Is it possible to create a serving endpoint of a registered model by alias using DABs? 
I want something like this to define the entity_name: 'myCatalog.mySchema.my-ads-model@myAlias'

Bundle sample

bundle:
  name: model_bundle

resources:
  model_serving_endpoints:
    uc_model_serving_endpoint:
      name: 'uc-model-endpoint'
      config:
        served_entities:
          - entity_name: 'myCatalog.mySchema.my-ads-model'
            entity_version: '10'
            workload_size: 'Small'
            scale_to_zero_enabled: 'true'
        traffic_config:
          routes:
            - served_model_name: 'my-ads-model-10'
              traffic_percentage: '100'
      tags:
        - key: 'team'
          value: 'data science'

 

koji_kawamura
Databricks Employee
Databricks Employee

Hi @MLOperator 

Since model_serving_endpoints only accepts a version number of a served entity, I think that is not possible. However, the get-by-alias version API can be used to retrieve a version number from a model alias name.  Then the model name and its version can be passed as variables.

As an example, I tested the following configuration:

 
variables:
  model_name:
  model_champion_version:

resources:
  model_serving_endpoints:
    uc_model_serving_endpoint:
      name: 'labuser9602087_1742260663_test'
      config:
        served_entities:
          - entity_name: '${var.model_name}'
            entity_version: '${var.model_champion_version}'
            workload_size: 'Small'
            scale_to_zero_enabled: 'true'
      tags:
        - key: 'team'
          value: 'data science'

When deploying it, the variables can be assigned like below. I'm using "champion" as the alias name:

export BUNDLE_VAR_model_name="catalog_name.schema_name.model_name"
export BUNDLE_VAR_model_champion_version=`databricks model-versions get-by-alias $BUNDLE_VAR_model_name champion |grep '"version":' |sed -r 's/.+:(.+)/\1/'`
databricks bundle deploy

Not ideal, but it worked. I hope this helps!

 

View solution in original post