Skip to content

Commit 877d53b

Browse files
shreyas-goenkashreyas-goenka
authored
Add support for registered models for direct deployment (#3670)
## Changes This PR adds support for registered models, matching the same behaviour as TF. This PR also adds test server handles for registered models and UC catalogs, as was necessary for the integration test. ## Why To complete coverage for direct deployment. ## Tests New integration test.
1 parent 51d3d4c commit 877d53b

File tree

13 files changed

+561
-6
lines changed

13 files changed

+561
-6
lines changed

‎acceptance/bundle/deploy/registered_models/basic/databricks.yml.tmpl‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
bundle:
2+
name: deploy-registered-models-basic-$UNIQUE_NAME
3+
4+
resources:
5+
registered_models:
6+
my_registered_model:
7+
name: $NAME
8+
comment: $COMMENT
9+
catalog_name: $CATALOG_NAME
10+
schema_name: $SCHEMA_NAME

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

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

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

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
2+
>>> export NAME=my-registered-model-[UNIQUE_NAME]
3+
4+
>>> export COMMENT=original comment
5+
6+
>>> export CATALOG_NAME=main
7+
8+
>>> export SCHEMA_NAME=default
9+
10+
=== create catalog and schema to test diff functionality
11+
>>> [CLI] catalogs create mycatalog-[UNIQUE_NAME]
12+
{
13+
"full_name": "mycatalog-[UNIQUE_NAME]"
14+
}
15+
16+
>>> [CLI] schemas create myschema-[UNIQUE_NAME] mycatalog-[UNIQUE_NAME]
17+
{
18+
"full_name": "mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME]"
19+
}
20+
21+
=== create the registered model
22+
>>> [CLI] bundle plan
23+
create registered_models.my_registered_model
24+
25+
Plan: 1 to add, 0 to change, 0 to delete, 0 unchanged
26+
27+
>>> [CLI] bundle deploy
28+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
29+
Deploying resources...
30+
Updating deployment state...
31+
Deployment complete!
32+
33+
>>> [CLI] registered-models get main.default.my-registered-model-[UNIQUE_NAME]
34+
{
35+
"name": "my-registered-model-[UNIQUE_NAME]",
36+
"comment": "original comment",
37+
"catalog_name": "main",
38+
"schema_name": "default"
39+
}
40+
41+
=== update the comment, this should not recreate
42+
>>> [CLI] bundle plan
43+
update registered_models.my_registered_model
44+
45+
Plan: 0 to add, 1 to change, 0 to delete, 0 unchanged
46+
47+
>>> [CLI] bundle deploy
48+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
49+
Deploying resources...
50+
Updating deployment state...
51+
Deployment complete!
52+
53+
>>> [CLI] registered-models get main.default.my-registered-model-[UNIQUE_NAME]
54+
{
55+
"name": "my-registered-model-[UNIQUE_NAME]",
56+
"comment": "updated comment",
57+
"catalog_name": "main",
58+
"schema_name": "default"
59+
}
60+
61+
=== update the name, this should recreate
62+
>>> [CLI] bundle plan
63+
recreate registered_models.my_registered_model
64+
65+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
66+
67+
>>> [CLI] bundle deploy
68+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
69+
Deploying resources...
70+
Updating deployment state...
71+
Deployment complete!
72+
73+
>>> [CLI] registered-models get main.default.my-registered-model-updated-[UNIQUE_NAME]
74+
{
75+
"name": "my-registered-model-updated-[UNIQUE_NAME]",
76+
"comment": "updated comment",
77+
"catalog_name": "main",
78+
"schema_name": "default"
79+
}
80+
81+
=== update the catalog name, this should recreate
82+
>>> [CLI] bundle plan
83+
recreate registered_models.my_registered_model
84+
85+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
86+
87+
>>> [CLI] bundle deploy
88+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
89+
Deploying resources...
90+
Updating deployment state...
91+
Deployment complete!
92+
93+
>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].default.my-registered-model-updated-[UNIQUE_NAME]
94+
{
95+
"name": "my-registered-model-updated-[UNIQUE_NAME]",
96+
"comment": "updated comment",
97+
"catalog_name": "mycatalog-[UNIQUE_NAME]",
98+
"schema_name": "default"
99+
}
100+
101+
=== update the schema name, this should recreate
102+
>>> [CLI] bundle plan
103+
recreate registered_models.my_registered_model
104+
105+
Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged
106+
107+
>>> [CLI] bundle deploy
108+
Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default/files...
109+
Deploying resources...
110+
Updating deployment state...
111+
Deployment complete!
112+
113+
>>> [CLI] registered-models get mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME].my-registered-model-updated-[UNIQUE_NAME]
114+
{
115+
"name": "my-registered-model-updated-[UNIQUE_NAME]",
116+
"comment": "updated comment",
117+
"catalog_name": "mycatalog-[UNIQUE_NAME]",
118+
"schema_name": "myschema-[UNIQUE_NAME]"
119+
}
120+
121+
>>> [CLI] bundle destroy --auto-approve
122+
The following resources will be deleted:
123+
delete registered_model my_registered_model
124+
125+
All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/deploy-registered-models-basic-[UNIQUE_NAME]/default
126+
127+
Deleting files...
128+
Destroy complete!
129+
130+
>>> [CLI] schemas delete mycatalog-[UNIQUE_NAME].myschema-[UNIQUE_NAME] --force
131+
132+
>>> [CLI] catalogs delete mycatalog-[UNIQUE_NAME] --force

‎acceptance/bundle/deploy/registered_models/basic/script‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
trace export NAME="my-registered-model-$UNIQUE_NAME"
2+
trace export COMMENT="original comment"
3+
trace export CATALOG_NAME="main"
4+
trace export SCHEMA_NAME="default"
5+
envsubst < databricks.yml.tmpl > databricks.yml
6+
7+
title "create catalog and schema to test diff functionality"
8+
catalog_name="mycatalog-${UNIQUE_NAME}"
9+
schema_name="myschema-${UNIQUE_NAME}"
10+
trace $CLI catalogs create ${catalog_name} | jq '{full_name}'
11+
trace $CLI schemas create ${schema_name} ${catalog_name} | jq '{full_name}'
12+
13+
cleanup() {
14+
trace $CLI bundle destroy --auto-approve
15+
trace $CLI schemas delete ${catalog_name}.${schema_name} --force
16+
trace $CLI catalogs delete ${catalog_name} --force
17+
}
18+
trap cleanup EXIT
19+
20+
deploy_registered_model() {
21+
trace $CLI bundle plan
22+
trace $CLI bundle deploy
23+
registered_model_id=$($CLI bundle summary --output json | jq -r '.resources.registered_models.my_registered_model.id')
24+
trace $CLI registered-models get "${registered_model_id}" | jq '{name, comment, catalog_name, schema_name}'
25+
}
26+
27+
title "create the registered model"
28+
deploy_registered_model
29+
30+
export COMMENT="updated comment"
31+
envsubst < databricks.yml.tmpl > databricks.yml
32+
33+
title "update the comment, this should not recreate"
34+
deploy_registered_model
35+
36+
export NAME="my-registered-model-updated-$UNIQUE_NAME"
37+
envsubst < databricks.yml.tmpl > databricks.yml
38+
39+
title "update the name, this should recreate"
40+
deploy_registered_model
41+
42+
title "update the catalog name, this should recreate"
43+
export CATALOG_NAME="${catalog_name}"
44+
envsubst < databricks.yml.tmpl > databricks.yml
45+
deploy_registered_model
46+
47+
title "update the schema name, this should recreate"
48+
export SCHEMA_NAME="${schema_name}"
49+
envsubst < databricks.yml.tmpl > databricks.yml
50+
deploy_registered_model

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

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

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2571,6 +2571,33 @@ resources.pipelines.*.trigger.cron.quartz_cron_schedule string INPUT STATE
25712571
resources.pipelines.*.trigger.cron.timezone_id string INPUT STATE
25722572
resources.pipelines.*.trigger.manual *pipelines.ManualTrigger INPUT STATE
25732573
resources.pipelines.*.url string INPUT
2574+
resources.registered_models.*.aliases []catalog.RegisteredModelAlias REMOTE
2575+
resources.registered_models.*.aliases[*] catalog.RegisteredModelAlias REMOTE
2576+
resources.registered_models.*.aliases[*].alias_name string REMOTE
2577+
resources.registered_models.*.aliases[*].version_num int REMOTE
2578+
resources.registered_models.*.browse_only bool REMOTE
2579+
resources.registered_models.*.catalog_name string ALL
2580+
resources.registered_models.*.comment string ALL
2581+
resources.registered_models.*.created_at int64 REMOTE
2582+
resources.registered_models.*.created_by string REMOTE
2583+
resources.registered_models.*.full_name string REMOTE
2584+
resources.registered_models.*.grants []resources.Grant INPUT
2585+
resources.registered_models.*.grants[*] resources.Grant INPUT
2586+
resources.registered_models.*.grants[*].principal string INPUT
2587+
resources.registered_models.*.grants[*].privileges []string INPUT
2588+
resources.registered_models.*.grants[*].privileges[*] string INPUT
2589+
resources.registered_models.*.id string INPUT
2590+
resources.registered_models.*.lifecycle resources.Lifecycle INPUT
2591+
resources.registered_models.*.lifecycle.prevent_destroy bool INPUT
2592+
resources.registered_models.*.metastore_id string REMOTE
2593+
resources.registered_models.*.modified_status string INPUT
2594+
resources.registered_models.*.name string ALL
2595+
resources.registered_models.*.owner string REMOTE
2596+
resources.registered_models.*.schema_name string ALL
2597+
resources.registered_models.*.storage_location string ALL
2598+
resources.registered_models.*.updated_at int64 REMOTE
2599+
resources.registered_models.*.updated_by string REMOTE
2600+
resources.registered_models.*.url string INPUT
25742601
resources.schemas.*.browse_only bool REMOTE
25752602
resources.schemas.*.catalog_name string ALL
25762603
resources.schemas.*.catalog_type catalog.CatalogType REMOTE

‎bundle/direct/dresources/all.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ var SupportedResources = map[string]any{
1919
"database_catalogs": (*ResourceDatabaseCatalog)(nil),
2020
"synced_database_tables": (*ResourceSyncedDatabaseTable)(nil),
2121
"alerts": (*ResourceAlert)(nil),
22+
"registered_models": (*ResourceRegisteredModel)(nil),
2223
}
2324

2425
func InitAll(client *databricks.WorkspaceClient) (map[string]*Adapter, error) {

‎bundle/direct/dresources/all_test.go‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,17 @@ var testConfig map[string]any = map[string]any{
6262
Name: "main.myschema.my_synced_table",
6363
},
6464
},
65+
66+
"registered_models": &resources.RegisteredModel{
67+
CreateRegisteredModelRequest: catalog.CreateRegisteredModelRequest{
68+
Name: "my_registered_model",
69+
Comment: "Test registered model",
70+
CatalogName: "main",
71+
SchemaName: "default",
72+
StorageLocation: "s3://my-bucket/my-path",
73+
},
74+
},
75+
6576
"experiments": &resources.MlflowExperiment{
6677
CreateExperiment: ml.CreateExperiment{
6778
Name: "my-experiment",
@@ -74,6 +85,7 @@ var testConfig map[string]any = map[string]any{
7485
ArtifactLocation: "s3://my-bucket/my-experiment",
7586
},
7687
},
88+
7789
"models": &resources.MlflowModel{
7890
CreateModelRequest: ml.CreateModelRequest{
7991
Name: "my_mlflow_model",
@@ -163,22 +175,26 @@ func testCRUD(t *testing.T, group string, adapter *Adapter, client *databricks.W
163175
require.Equal(t, remote, remoteStateFromWaitCreate)
164176
}
165177

178+
remappedState, err := adapter.RemapState(remote)
179+
require.NoError(t, err)
180+
require.NotNil(t, remappedState)
181+
166182
remoteStateFromUpdate, err := adapter.DoUpdate(ctx, createdID, newState)
167183
require.NoError(t, err, "DoUpdate failed")
168184
if remoteStateFromUpdate != nil {
169-
require.Equal(t, remote, remoteStateFromUpdate)
185+
remappedStateFromUpdate, err := adapter.RemapState(remoteStateFromUpdate)
186+
require.NoError(t, err)
187+
require.Equal(t, remappedState, remappedStateFromUpdate)
170188
}
171189

172190
remoteStateFromWaitUpdate, err := adapter.WaitAfterUpdate(ctx, newState)
173191
require.NoError(t, err)
174192
if remoteStateFromWaitUpdate != nil {
175-
require.Equal(t, remote, remoteStateFromWaitUpdate)
193+
remappedStateFromWaitUpdate, err := adapter.RemapState(remoteStateFromWaitUpdate)
194+
require.NoError(t, err)
195+
require.Equal(t, remappedState, remappedStateFromWaitUpdate)
176196
}
177197

178-
remappedState, err := adapter.RemapState(remote)
179-
require.NoError(t, err)
180-
require.NotNil(t, remappedState)
181-
182198
require.NoError(t, structwalk.Walk(newState, func(path *structpath.PathNode, val any, field *reflect.StructField) {
183199
remoteValue, err := structaccess.Get(remappedState, path)
184200
if err != nil {

0 commit comments

Comments
 (0)