Skip to content

Commit 4475762

Browse files
shreyas-goenkashreyas-goenka
authored
Add support for models in direct deployment (#3625)
## Changes This PR adds support for direct deployment for MLflow models, matching the existing Terraform behaviour today. ## Why To make direct deployments possible. ## Tests Test asserting that TF and direct behaviour match for all relevent fields.
1 parent 739b1f7 commit 4475762

File tree

13 files changed

+481
-15
lines changed

13 files changed

+481
-15
lines changed

‎acceptance/bundle/deploy/models/basic/out.test.toml‎

Lines changed: 5 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎acceptance/bundle/deploy/models/basic/output.txt‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
2+
>>> export MODEL_NAME=original-name-[UNIQUE_NAME]
3+
4+
>>> export MODEL_DESCRIPTION=original-description-[UNIQUE_NAME]
5+
6+
=== create a model
7+
>>> [CLI] bundle plan
8+
create models.my_model
9+
10+
Plan: 1 to add, 0 to change, 0 to delete, 0 unchanged
11+
12+
>>> [CLI] bundle deploy
13+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
14+
Deploying resources...
15+
Updating deployment state...
16+
Deployment complete!
17+
18+
>>> [CLI] model-registry get-model original-name-[UNIQUE_NAME]
19+
{
20+
"name": "original-name-[UNIQUE_NAME]",
21+
"description": "original-description-[UNIQUE_NAME]",
22+
"tags": [
23+
{
24+
"key": "key1",
25+
"value": "value1"
26+
}
27+
]
28+
}
29+
30+
>>> export MODEL_DESCRIPTION=new-description-[UNIQUE_NAME]
31+
32+
=== update the description, this should update the description remotely as well
33+
>>> [CLI] bundle plan
34+
update models.my_model
35+
36+
Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged
37+
38+
>>> [CLI] bundle deploy
39+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
40+
Deploying resources...
41+
Updating deployment state...
42+
Deployment complete!
43+
44+
>>> [CLI] model-registry get-model original-name-[UNIQUE_NAME]
45+
{
46+
"name": "original-name-[UNIQUE_NAME]",
47+
"description": "new-description-[UNIQUE_NAME]",
48+
"tags": [
49+
{
50+
"key": "key1",
51+
"value": "value1"
52+
}
53+
]
54+
}
55+
56+
>>> export MODEL_NAME=new-name-[UNIQUE_NAME]
57+
58+
=== update the name, this should recreate the model with the new name
59+
>>> [CLI] bundle plan
60+
recreate models.my_model
61+
62+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
63+
64+
>>> [CLI] bundle deploy
65+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
66+
Deploying resources...
67+
Updating deployment state...
68+
Deployment complete!
69+
70+
>>> [CLI] model-registry get-model new-name-[UNIQUE_NAME]
71+
{
72+
"name": "new-name-[UNIQUE_NAME]",
73+
"description": "new-description-[UNIQUE_NAME]",
74+
"tags": [
75+
{
76+
"key": "key1",
77+
"value": "value1"
78+
}
79+
]
80+
}
81+
82+
=== add a new tag, this should be a no-op
83+
>>> [CLI] bundle plan
84+
update models.my_model
85+
86+
Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged
87+
88+
>>> [CLI] bundle deploy
89+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default/files...
90+
Deploying resources...
91+
Updating deployment state...
92+
Deployment complete!
93+
94+
>>> [CLI] model-registry get-model new-name-[UNIQUE_NAME]
95+
{
96+
"name": "new-name-[UNIQUE_NAME]",
97+
"description": "new-description-[UNIQUE_NAME]",
98+
"tags": [
99+
{
100+
"key": "key1",
101+
"value": "value1"
102+
}
103+
]
104+
}
105+
106+
>>> [CLI] bundle destroy --auto-approve
107+
The following resources will be deleted:
108+
delete model my_model
109+
110+
All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/deploy-models-basic-[UNIQUE_NAME]/default
111+
112+
Deleting files...
113+
Destroy complete!

‎acceptance/bundle/deploy/models/basic/script‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
cleanup() {
2+
trace $CLI bundle destroy --auto-approve
3+
}
4+
trap cleanup EXIT
5+
6+
trace export MODEL_NAME=original-name-$UNIQUE_NAME
7+
trace export MODEL_DESCRIPTION=original-description-$UNIQUE_NAME
8+
envsubst < templates/one_tag.tmpl > databricks.yml
9+
title "create a model"
10+
trace $CLI bundle plan
11+
trace $CLI bundle deploy
12+
trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'
13+
14+
trace export MODEL_DESCRIPTION=new-description-$UNIQUE_NAME
15+
envsubst < templates/one_tag.tmpl > databricks.yml
16+
title "update the description, this should update the description remotely as well"
17+
trace $CLI bundle plan
18+
trace $CLI bundle deploy
19+
trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'
20+
21+
trace export MODEL_NAME=new-name-$UNIQUE_NAME
22+
envsubst < templates/one_tag.tmpl > databricks.yml
23+
title "update the name, this should recreate the model with the new name"
24+
trace $CLI bundle plan
25+
trace $CLI bundle deploy
26+
trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'
27+
28+
title "add a new tag, this should be a no-op"
29+
envsubst < templates/two_tag.tmpl > databricks.yml
30+
trace $CLI bundle plan
31+
trace $CLI bundle deploy
32+
trace $CLI model-registry get-model $MODEL_NAME | jq '.registered_model_databricks | {name, description, tags}'

‎acceptance/bundle/deploy/models/basic/templates/one_tag.tmpl‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
bundle:
2+
name: deploy-models-basic-$UNIQUE_NAME
3+
4+
resources:
5+
models:
6+
my_model:
7+
name: $MODEL_NAME
8+
description: $MODEL_DESCRIPTION
9+
tags:
10+
- key: "key1"
11+
value: "value1"

‎acceptance/bundle/deploy/models/basic/templates/two_tag.tmpl‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
bundle:
2+
name: deploy-models-basic-$UNIQUE_NAME
3+
4+
resources:
5+
models:
6+
my_model:
7+
name: $MODEL_NAME
8+
description: $MODEL_DESCRIPTION
9+
tags:
10+
- key: "key1"
11+
value: "value1"
12+
- key: "key2"
13+
value: "value2"

‎acceptance/bundle/deploy/models/basic/test.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Cloud = true
2+
Local = true

‎acceptance/bundle/refschema/out.fields.txt‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1919,6 +1919,45 @@ resources.jobs.*.webhook_notifications.on_streaming_backlog_exceeded[*].id strin
19191919
resources.jobs.*.webhook_notifications.on_success []jobs.Webhook INPUT STATE
19201920
resources.jobs.*.webhook_notifications.on_success[*] jobs.Webhook INPUT STATE
19211921
resources.jobs.*.webhook_notifications.on_success[*].id string INPUT STATE
1922+
resources.models.*.creation_timestamp int64 REMOTE
1923+
resources.models.*.description string ALL
1924+
resources.models.*.id string INPUT REMOTE
1925+
resources.models.*.last_updated_timestamp int64 REMOTE
1926+
resources.models.*.latest_versions []ml.ModelVersion REMOTE
1927+
resources.models.*.latest_versions[*] ml.ModelVersion REMOTE
1928+
resources.models.*.latest_versions[*].creation_timestamp int64 REMOTE
1929+
resources.models.*.latest_versions[*].current_stage string REMOTE
1930+
resources.models.*.latest_versions[*].description string REMOTE
1931+
resources.models.*.latest_versions[*].last_updated_timestamp int64 REMOTE
1932+
resources.models.*.latest_versions[*].name string REMOTE
1933+
resources.models.*.latest_versions[*].run_id string REMOTE
1934+
resources.models.*.latest_versions[*].run_link string REMOTE
1935+
resources.models.*.latest_versions[*].source string REMOTE
1936+
resources.models.*.latest_versions[*].status ml.ModelVersionStatus REMOTE
1937+
resources.models.*.latest_versions[*].status_message string REMOTE
1938+
resources.models.*.latest_versions[*].tags []ml.ModelVersionTag REMOTE
1939+
resources.models.*.latest_versions[*].tags[*] ml.ModelVersionTag REMOTE
1940+
resources.models.*.latest_versions[*].tags[*].key string REMOTE
1941+
resources.models.*.latest_versions[*].tags[*].value string REMOTE
1942+
resources.models.*.latest_versions[*].user_id string REMOTE
1943+
resources.models.*.latest_versions[*].version string REMOTE
1944+
resources.models.*.lifecycle resources.Lifecycle INPUT
1945+
resources.models.*.lifecycle.prevent_destroy bool INPUT
1946+
resources.models.*.modified_status string INPUT
1947+
resources.models.*.name string ALL
1948+
resources.models.*.permission_level ml.PermissionLevel REMOTE
1949+
resources.models.*.permissions []resources.MlflowModelPermission INPUT
1950+
resources.models.*.permissions[*] resources.MlflowModelPermission INPUT
1951+
resources.models.*.permissions[*].group_name string INPUT
1952+
resources.models.*.permissions[*].level resources.MlflowModelPermissionLevel INPUT
1953+
resources.models.*.permissions[*].service_principal_name string INPUT
1954+
resources.models.*.permissions[*].user_name string INPUT
1955+
resources.models.*.tags []ml.ModelTag ALL
1956+
resources.models.*.tags[*] ml.ModelTag ALL
1957+
resources.models.*.tags[*].key string ALL
1958+
resources.models.*.tags[*].value string ALL
1959+
resources.models.*.url string INPUT
1960+
resources.models.*.user_id string REMOTE
19221961
resources.pipelines.*.allow_duplicate_names bool INPUT STATE
19231962
resources.pipelines.*.budget_policy_id string INPUT STATE
19241963
resources.pipelines.*.catalog string INPUT STATE

‎bundle/direct/dresources/all.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ var SupportedResources = map[string]any{
1111
"pipelines": (*ResourcePipeline)(nil),
1212
"schemas": (*ResourceSchema)(nil),
1313
"volumes": (*ResourceVolume)(nil),
14+
"models": (*ResourceMlflowModel)(nil),
1415
"apps": (*ResourceApp)(nil),
1516
"sql_warehouses": (*ResourceSqlWarehouse)(nil),
1617
"database_instances": (*ResourceDatabaseInstance)(nil),

‎bundle/direct/dresources/all_test.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import (
1616
"github.com/databricks/databricks-sdk-go/service/apps"
1717
"github.com/databricks/databricks-sdk-go/service/catalog"
1818
"github.com/databricks/databricks-sdk-go/service/database"
19+
"github.com/databricks/databricks-sdk-go/service/ml"
1920
"github.com/stretchr/testify/assert"
2021
"github.com/stretchr/testify/require"
2122
)
@@ -60,6 +61,18 @@ var testConfig map[string]any = map[string]any{
6061
Name: "main.myschema.my_synced_table",
6162
},
6263
},
64+
"models": &resources.MlflowModel{
65+
CreateModelRequest: ml.CreateModelRequest{
66+
Name: "my_mlflow_model",
67+
Description: "my_mlflow_model_description",
68+
Tags: []ml.ModelTag{
69+
{
70+
Key: "k1",
71+
Value: "v1",
72+
},
73+
},
74+
},
75+
},
6376
}
6477

6578
type prepareWorkspace func(client *databricks.WorkspaceClient) error

‎bundle/direct/dresources/model.go‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
package dresources
2+
3+
import (
4+
"context"
5+
6+
"github.com/databricks/cli/bundle/config/resources"
7+
"github.com/databricks/cli/bundle/deployplan"
8+
"github.com/databricks/databricks-sdk-go"
9+
"github.com/databricks/databricks-sdk-go/service/ml"
10+
)
11+
12+
type ResourceMlflowModel struct {
13+
client *databricks.WorkspaceClient
14+
}
15+
16+
func (*ResourceMlflowModel) New(client *databricks.WorkspaceClient) *ResourceMlflowModel {
17+
return &ResourceMlflowModel{client: client}
18+
}
19+
20+
func (*ResourceMlflowModel) PrepareState(input *resources.MlflowModel) *ml.CreateModelRequest {
21+
return &input.CreateModelRequest
22+
}
23+
24+
func (*ResourceMlflowModel) RemapState(model *ml.ModelDatabricks) *ml.CreateModelRequest {
25+
return &ml.CreateModelRequest{
26+
Name: model.Name,
27+
Tags: model.Tags,
28+
Description: model.Description,
29+
ForceSendFields: filterFields[ml.CreateModelRequest](model.ForceSendFields),
30+
}
31+
}
32+
33+
func (r *ResourceMlflowModel) DoRefresh(ctx context.Context, id string) (*ml.ModelDatabricks, error) {
34+
response, err := r.client.ModelRegistry.GetModel(ctx, ml.GetModelRequest{
35+
Name: id,
36+
})
37+
if err != nil {
38+
return nil, err
39+
}
40+
return response.RegisteredModelDatabricks, nil
41+
}
42+
43+
func (r *ResourceMlflowModel) DoCreate(ctx context.Context, config *ml.CreateModelRequest) (string, *ml.ModelDatabricks, error) {
44+
response, err := r.client.ModelRegistry.CreateModel(ctx, *config)
45+
if err != nil {
46+
return "", nil, err
47+
}
48+
// Create API call returns [ml.Model] while DoRefresh returns [ml.ModelDatabricks].
49+
// Thus we need to convert the response to the expected type.
50+
modelDatabricks := &ml.ModelDatabricks{
51+
Name: response.RegisteredModel.Name,
52+
Description: response.RegisteredModel.Description,
53+
Tags: response.RegisteredModel.Tags,
54+
ForceSendFields: filterFields[ml.ModelDatabricks](response.RegisteredModel.ForceSendFields, "CreationTimestamp", "Id", "LastUpdatedTimestamp", "LatestVersions", "PermissionLevel", "UserId"),
55+
56+
// Coping the fields only to satisfy the linter. These fields are not
57+
// part of the configuration tree so they don't need to be copied.
58+
// The linter works as a safeguard to ensure we add new fields to the bundle config tree
59+
// to the mapping logic here as well.
60+
CreationTimestamp: 0,
61+
Id: "",
62+
LastUpdatedTimestamp: 0,
63+
LatestVersions: nil,
64+
PermissionLevel: "",
65+
UserId: "",
66+
}
67+
return response.RegisteredModel.Name, modelDatabricks, nil
68+
}
69+
70+
func (r *ResourceMlflowModel) DoUpdate(ctx context.Context, id string, config *ml.CreateModelRequest) (*ml.ModelDatabricks, error) {
71+
updateRequest := ml.UpdateModelRequest{
72+
Name: id,
73+
Description: config.Description,
74+
ForceSendFields: filterFields[ml.UpdateModelRequest](config.ForceSendFields),
75+
}
76+
77+
response, err := r.client.ModelRegistry.UpdateModel(ctx, updateRequest)
78+
if err != nil {
79+
return nil, err
80+
}
81+
82+
// Update API call returns [ml.Model] while DoRefresh returns [ml.ModelDatabricks].
83+
// Thus we need to convert the response to the expected type.
84+
modelDatabricks := &ml.ModelDatabricks{
85+
Name: response.RegisteredModel.Name,
86+
Description: response.RegisteredModel.Description,
87+
Tags: response.RegisteredModel.Tags,
88+
ForceSendFields: filterFields[ml.ModelDatabricks](response.RegisteredModel.ForceSendFields, "CreationTimestamp", "Id", "LastUpdatedTimestamp", "LatestVersions", "PermissionLevel", "UserId"),
89+
90+
// Coping the fields only to satisfy the linter. These fields are not
91+
// part of the configuration tree so they don't need to be copied.
92+
// The linter works as a safeguard to ensure we add new fields to the bundle config tree
93+
// to the mapping logic here as well.
94+
CreationTimestamp: 0,
95+
Id: "",
96+
LastUpdatedTimestamp: 0,
97+
LatestVersions: nil,
98+
PermissionLevel: "",
99+
UserId: "",
100+
}
101+
return modelDatabricks, nil
102+
}
103+
104+
func (r *ResourceMlflowModel) DoDelete(ctx context.Context, id string) error {
105+
return r.client.ModelRegistry.DeleteModel(ctx, ml.DeleteModelRequest{
106+
Name: id,
107+
})
108+
}
109+
110+
func (*ResourceMlflowModel) FieldTriggers() map[string]deployplan.ActionType {
111+
return map[string]deployplan.ActionType{
112+
// Recreate matches current behavior of Terraform. It is possible to rename without recreate
113+
// but that would require dynamic select of the method during update since
114+
// the [ml.RenameModel] needs to be called instead of [ml.UpdateModel].
115+
//
116+
// We might reasonably choose to never fix this because this is a legacy resource.
117+
"name": deployplan.ActionTypeRecreate,
118+
119+
// Allowing updates for tags requires dynamic selection of the method since
120+
// tags can only be updated by calling [ml.SetModelTag] or [ml.DeleteModelTag] methods.
121+
//
122+
// Skip annotation matches the current behavior of Terraform where tags changes are showed
123+
// in plan but are just ignored / not applied. Since this is a legacy resource we might
124+
// reasonably choose to not fix it here as well.
125+
"tags": deployplan.ActionTypeSkip,
126+
}
127+
}

0 commit comments

Comments
 (0)