diff --git a/.nextchanges/bundles/ai-gateway-model-service.md b/.nextchanges/bundles/ai-gateway-model-service.md new file mode 100644 index 00000000000..85baad946ad --- /dev/null +++ b/.nextchanges/bundles/ai-gateway-model-service.md @@ -0,0 +1 @@ +* Add bundle support for the AI Gateway `model_service` resource (direct engine). ([#6525](https://github.com/databricks/cli/pull/6525)) diff --git a/acceptance/bundle/deployment/bind/model_service/databricks.yml b/acceptance/bundle/deployment/bind/model_service/databricks.yml new file mode 100644 index 00000000000..8d56e3b54d4 --- /dev/null +++ b/acceptance/bundle/deployment/bind/model_service/databricks.yml @@ -0,0 +1,9 @@ +bundle: + name: test-bundle + +resources: + model_services: + ms1: + parent: schemas/main.myschema + model_service_id: mysvc + comment: bound service diff --git a/acceptance/bundle/deployment/bind/model_service/out.test.toml b/acceptance/bundle/deployment/bind/model_service/out.test.toml new file mode 100644 index 00000000000..27ec2a7fcd6 --- /dev/null +++ b/acceptance/bundle/deployment/bind/model_service/out.test.toml @@ -0,0 +1,3 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] +EnvMatrix.DMS = [""] diff --git a/acceptance/bundle/deployment/bind/model_service/output.txt b/acceptance/bundle/deployment/bind/model_service/output.txt new file mode 100644 index 00000000000..bd402a482c2 --- /dev/null +++ b/acceptance/bundle/deployment/bind/model_service/output.txt @@ -0,0 +1,30 @@ + +>>> [CLI] bundle deployment bind ms1 main.myschema.mysvc --auto-approve +Successfully bound model_service with an id 'main.myschema.mysvc' +Run 'bundle deploy' to deploy changes to your workspace + +>>> [CLI] bundle summary +Name: test-bundle +Target: default +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/test-bundle/default +Resources: + Model services: + ms1: + Name: mysvc + URL: [DATABRICKS_URL]/explore/data/model-services/main/myschema/mysvc?w=[NUMID] + +>>> [CLI] bundle deployment unbind ms1 + +>>> [CLI] bundle summary +Name: test-bundle +Target: default +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/test-bundle/default +Resources: + Model services: + ms1: + Name: mysvc + URL: (not deployed) diff --git a/acceptance/bundle/deployment/bind/model_service/script b/acceptance/bundle/deployment/bind/model_service/script new file mode 100644 index 00000000000..e886f5dcd85 --- /dev/null +++ b/acceptance/bundle/deployment/bind/model_service/script @@ -0,0 +1,5 @@ +trace $CLI bundle deployment bind ms1 main.myschema.mysvc --auto-approve +trace $CLI bundle summary + +trace $CLI bundle deployment unbind ms1 +trace $CLI bundle summary diff --git a/acceptance/bundle/deployment/bind/model_service/test.toml b/acceptance/bundle/deployment/bind/model_service/test.toml new file mode 100644 index 00000000000..dc03ed31a1b --- /dev/null +++ b/acceptance/bundle/deployment/bind/model_service/test.toml @@ -0,0 +1,18 @@ +# AI Gateway securables are direct-engine only. +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +Cloud = false + +Ignore = [ + ".databricks", +] + +# The bind flow issues a GET to confirm the remote resource exists before binding. +[[Server]] +Pattern = "GET /api/2.1/unity-catalog/model-services/{name}" +Response.Body = ''' +{ + "name": "model-services/main.myschema.mysvc", + "comment": "bound service" +} +''' diff --git a/acceptance/bundle/invariant/configs/model_service.yml.tmpl b/acceptance/bundle/invariant/configs/model_service.yml.tmpl new file mode 100644 index 00000000000..e4ea6c39dc8 --- /dev/null +++ b/acceptance/bundle/invariant/configs/model_service.yml.tmpl @@ -0,0 +1,9 @@ +bundle: + name: test-bundle-$UNIQUE_NAME + +resources: + model_services: + foo: + parent: schemas/main.default + model_service_id: test-model-service-$UNIQUE_NAME + comment: test model service diff --git a/acceptance/bundle/invariant/continue_293/out.test.toml b/acceptance/bundle/invariant/continue_293/out.test.toml index 20fe46d85de..e23caeac873 100644 --- a/acceptance/bundle/invariant/continue_293/out.test.toml +++ b/acceptance/bundle/invariant/continue_293/out.test.toml @@ -45,5 +45,6 @@ EnvMatrix.INPUT_CONFIG = [ "synced_database_table.yml.tmpl", "volume.yml.tmpl", "volume_external.yml.tmpl", - "volume_uppercase_name.yml.tmpl" + "volume_uppercase_name.yml.tmpl", + "model_service.yml.tmpl" ] diff --git a/acceptance/bundle/invariant/delete_idempotent/out.test.toml b/acceptance/bundle/invariant/delete_idempotent/out.test.toml index 59018deb1ce..2554ae2ce3f 100644 --- a/acceptance/bundle/invariant/delete_idempotent/out.test.toml +++ b/acceptance/bundle/invariant/delete_idempotent/out.test.toml @@ -60,6 +60,7 @@ EnvMatrix.INPUT_CONFIG = [ "volume.yml.tmpl", "volume_external.yml.tmpl", "volume_path_job_ref.yml.tmpl", - "volume_uppercase_name.yml.tmpl" + "volume_uppercase_name.yml.tmpl", + "model_service.yml.tmpl" ] EnvMatrix.READPLAN = ["", "1"] diff --git a/acceptance/bundle/invariant/destroy_idempotent/out.test.toml b/acceptance/bundle/invariant/destroy_idempotent/out.test.toml index 59018deb1ce..2554ae2ce3f 100644 --- a/acceptance/bundle/invariant/destroy_idempotent/out.test.toml +++ b/acceptance/bundle/invariant/destroy_idempotent/out.test.toml @@ -60,6 +60,7 @@ EnvMatrix.INPUT_CONFIG = [ "volume.yml.tmpl", "volume_external.yml.tmpl", "volume_path_job_ref.yml.tmpl", - "volume_uppercase_name.yml.tmpl" + "volume_uppercase_name.yml.tmpl", + "model_service.yml.tmpl" ] EnvMatrix.READPLAN = ["", "1"] diff --git a/acceptance/bundle/invariant/migrate/test.toml b/acceptance/bundle/invariant/migrate/test.toml index 2924a3dec1c..0c94059c56d 100644 --- a/acceptance/bundle/invariant/migrate/test.toml +++ b/acceptance/bundle/invariant/migrate/test.toml @@ -21,6 +21,8 @@ EnvMatrixExclude.no_genie_space = ["INPUT_CONFIG=genie_space.yml.tmpl"] EnvMatrixExclude.no_instance_pool = ["INPUT_CONFIG=instance_pool.yml.tmpl"] # Cluster policies are direct-only; the terraform deploy that seeds the migration fails for them. EnvMatrixExclude.no_cluster_policy = ["INPUT_CONFIG=cluster_policy.yml.tmpl"] +# AI Gateway securables are direct-only; the terraform deploy that seeds the migration fails for them. +EnvMatrixExclude.no_model_service = ["INPUT_CONFIG=model_service.yml.tmpl"] # Cross-resource permission references (e.g. ${resources.jobs.job_b.permissions[0].level}) # don't work in terraform mode: the terraform interpolator converts the path to diff --git a/acceptance/bundle/invariant/no_drift/out.test.toml b/acceptance/bundle/invariant/no_drift/out.test.toml index 8439713320e..5ea98dcf5a2 100644 --- a/acceptance/bundle/invariant/no_drift/out.test.toml +++ b/acceptance/bundle/invariant/no_drift/out.test.toml @@ -60,6 +60,7 @@ EnvMatrix.INPUT_CONFIG = [ "volume.yml.tmpl", "volume_external.yml.tmpl", "volume_path_job_ref.yml.tmpl", - "volume_uppercase_name.yml.tmpl" + "volume_uppercase_name.yml.tmpl", + "model_service.yml.tmpl" ] EnvMatrix.READPLAN = ["", "1"] diff --git a/acceptance/bundle/invariant/test.toml b/acceptance/bundle/invariant/test.toml index 37bbb2c5103..780f79774b0 100644 --- a/acceptance/bundle/invariant/test.toml +++ b/acceptance/bundle/invariant/test.toml @@ -78,6 +78,7 @@ EnvMatrix.INPUT_CONFIG = [ "volume_external.yml.tmpl", "volume_path_job_ref.yml.tmpl", "volume_uppercase_name.yml.tmpl", + "model_service.yml.tmpl", ] EnvMatrixExclude.no_alert_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=alert.yml.tmpl"] @@ -108,6 +109,10 @@ EnvMatrixExclude.no_external_location_on_cloud = ["CONFIG_Cloud=true", "INPUT_CO EnvMatrixExclude.no_external_volume_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=volume_external.yml.tmpl"] # Telemetry requires a real model, which cloud invariant tests do not provision. EnvMatrixExclude.no_model_serving_endpoint_telemetry_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=model_serving_endpoint_telemetry.yml.tmpl"] +# AI Gateway securables (direct-only) need a real routing target / provider config / +# UC connection to create on cloud, which the invariant tests don't provision; the +# local mock server accepts them, so these run locally only. +EnvMatrixExclude.no_model_service_on_cloud = ["CONFIG_Cloud=true", "INPUT_CONFIG=model_service.yml.tmpl"] # Fake SQL endpoint for local tests [[Server]] diff --git a/acceptance/bundle/python/model_services-support/databricks.yml b/acceptance/bundle/python/model_services-support/databricks.yml new file mode 100644 index 00000000000..37acd94d33a --- /dev/null +++ b/acceptance/bundle/python/model_services-support/databricks.yml @@ -0,0 +1,17 @@ +bundle: + name: my_project + +sync: {paths: []} # don't need to copy files + +python: + resources: + - "resources:load_resources" + mutators: + - "mutators:update_model_service" + +resources: + model_services: + my_model_service_1: + parent: "schemas/main.default" + model_service_id: "my_model_service_1" + comment: "My model service" diff --git a/acceptance/bundle/python/model_services-support/mutators.py b/acceptance/bundle/python/model_services-support/mutators.py new file mode 100644 index 00000000000..fdd96009e1b --- /dev/null +++ b/acceptance/bundle/python/model_services-support/mutators.py @@ -0,0 +1,11 @@ +from dataclasses import replace + +from databricks.bundles.model_services import ModelService +from databricks.bundles.core import model_service_mutator + + +@model_service_mutator +def update_model_service(model_service: ModelService) -> ModelService: + assert isinstance(model_service.comment, str) + + return replace(model_service, comment=f"{model_service.comment} (updated)") diff --git a/acceptance/bundle/python/model_services-support/out.test.toml b/acceptance/bundle/python/model_services-support/out.test.toml new file mode 100644 index 00000000000..02da5baab89 --- /dev/null +++ b/acceptance/bundle/python/model_services-support/out.test.toml @@ -0,0 +1,4 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] +EnvMatrix.DMS = ["", "true"] +EnvMatrix.PYDAB_VERSION = ["current"] diff --git a/acceptance/bundle/python/model_services-support/output.txt b/acceptance/bundle/python/model_services-support/output.txt new file mode 100644 index 00000000000..52935d3c124 --- /dev/null +++ b/acceptance/bundle/python/model_services-support/output.txt @@ -0,0 +1,28 @@ + +>>> uv run [UV_ARGS] -q [CLI] bundle validate --output json +{ + "experimental": { + "python": { + "mutators": [ + "mutators:update_model_service" + ], + "resources": [ + "resources:load_resources" + ] + } + }, + "resources": { + "model_services": { + "my_model_service_1": { + "comment": "My model service (updated)", + "model_service_id": "my_model_service_1", + "parent": "schemas/main.default" + }, + "my_model_service_2": { + "comment": "My model service (2) (updated)", + "model_service_id": "my_model_service_2", + "parent": "schemas/main.default" + } + } + } +} diff --git a/acceptance/bundle/python/model_services-support/resources.py b/acceptance/bundle/python/model_services-support/resources.py new file mode 100644 index 00000000000..58a92215744 --- /dev/null +++ b/acceptance/bundle/python/model_services-support/resources.py @@ -0,0 +1,16 @@ +from databricks.bundles.core import Resources + + +def load_resources() -> Resources: + resources = Resources() + + resources.add_model_service( + "my_model_service_2", + { + "parent": "schemas/main.default", + "model_service_id": "my_model_service_2", + "comment": "My model service (2)", + }, + ) + + return resources diff --git a/acceptance/bundle/python/model_services-support/script b/acceptance/bundle/python/model_services-support/script new file mode 100644 index 00000000000..e273fb45a53 --- /dev/null +++ b/acceptance/bundle/python/model_services-support/script @@ -0,0 +1,5 @@ + +trace uv run $UV_ARGS -q $CLI bundle validate --output json | \ + jq "pick(.experimental.python, .resources)" + +rm -fr .databricks __pycache__ diff --git a/acceptance/bundle/python/model_services-support/test.toml b/acceptance/bundle/python/model_services-support/test.toml new file mode 100644 index 00000000000..092080306b1 --- /dev/null +++ b/acceptance/bundle/python/model_services-support/test.toml @@ -0,0 +1,4 @@ +Cloud = false # tests don't interact with APIs + +# model_services are only supported in the current version of the wheel +EnvMatrix.PYDAB_VERSION = ["current"] diff --git a/acceptance/bundle/refschema/out.fields.txt b/acceptance/bundle/refschema/out.fields.txt index 158fcc5d1c0..1b26570ab5a 100644 --- a/acceptance/bundle/refschema/out.fields.txt +++ b/acceptance/bundle/refschema/out.fields.txt @@ -2077,6 +2077,68 @@ resources.jobs.*.permissions[*].group_name string ALL resources.jobs.*.permissions[*].level iam.PermissionLevel ALL resources.jobs.*.permissions[*].service_principal_name string ALL resources.jobs.*.permissions[*].user_name string ALL +resources.model_services.*.comment string ALL +resources.model_services.*.config *catalog.ModelServiceConfig ALL +resources.model_services.*.config.inference_table *catalog.InferenceTableConfig ALL +resources.model_services.*.config.inference_table.disabled bool ALL +resources.model_services.*.config.inference_table.is_deleted bool ALL +resources.model_services.*.config.inference_table.parent string ALL +resources.model_services.*.config.inference_table.table string ALL +resources.model_services.*.config.inference_table.table_name_prefix string ALL +resources.model_services.*.config.rate_limits []catalog.RateLimit ALL +resources.model_services.*.config.rate_limits[*] catalog.RateLimit ALL +resources.model_services.*.config.rate_limits[*].key catalog.RateLimitRateLimitKey ALL +resources.model_services.*.config.rate_limits[*].principal string ALL +resources.model_services.*.config.rate_limits[*].renewal_period catalog.RateLimitRateLimitRenewalPeriod ALL +resources.model_services.*.config.rate_limits[*].request_tag_key string ALL +resources.model_services.*.config.rate_limits[*].request_tag_value string ALL +resources.model_services.*.config.rate_limits[*].requests int64 ALL +resources.model_services.*.config.rate_limits[*].tokens int64 ALL +resources.model_services.*.config.routing *catalog.ModelServiceConfigRoutingConfig ALL +resources.model_services.*.config.routing.destinations []catalog.ModelServiceConfigDestinationConfig ALL +resources.model_services.*.config.routing.destinations[*] catalog.ModelServiceConfigDestinationConfig ALL +resources.model_services.*.config.routing.destinations[*].destination_type catalog.ModelServiceConfigDestinationConfigDestinationType ALL +resources.model_services.*.config.routing.destinations[*].external_model_config *catalog.ModelServiceConfigExternalModelConfig ALL +resources.model_services.*.config.routing.destinations[*].external_model_config.model_provider_service string ALL +resources.model_services.*.config.routing.destinations[*].external_model_config.target catalog.ModelProviderServiceConfigModelTargetConfig ALL +resources.model_services.*.config.routing.destinations[*].external_model_config.target.model string ALL +resources.model_services.*.config.routing.destinations[*].external_model_config.target.native_api_types []string ALL +resources.model_services.*.config.routing.destinations[*].external_model_config.target.native_api_types[*] string ALL +resources.model_services.*.config.routing.destinations[*].is_deleted bool ALL +resources.model_services.*.config.routing.destinations[*].name string ALL +resources.model_services.*.config.routing.destinations[*].pay_per_token_config *catalog.ModelServiceConfigPayPerTokenConfig ALL +resources.model_services.*.config.routing.destinations[*].pay_per_token_config.model string ALL +resources.model_services.*.config.routing.destinations[*].provisioned_throughput_config *catalog.ModelServiceConfigProvisionedThroughputConfig ALL +resources.model_services.*.config.routing.destinations[*].provisioned_throughput_config.model string ALL +resources.model_services.*.config.routing.destinations[*].provisioned_throughput_config.model_serving_endpoint string ALL +resources.model_services.*.config.routing.destinations[*].traffic_percentage int ALL +resources.model_services.*.config.routing.fallback *catalog.ModelServiceConfigFallbackConfig ALL +resources.model_services.*.config.routing.fallback.destinations []catalog.ModelServiceConfigDestinationConfig ALL +resources.model_services.*.config.routing.fallback.destinations[*] catalog.ModelServiceConfigDestinationConfig ALL +resources.model_services.*.config.routing.fallback.destinations[*].destination_type catalog.ModelServiceConfigDestinationConfigDestinationType ALL +resources.model_services.*.config.routing.fallback.destinations[*].external_model_config *catalog.ModelServiceConfigExternalModelConfig ALL +resources.model_services.*.config.routing.fallback.destinations[*].external_model_config.model_provider_service string ALL +resources.model_services.*.config.routing.fallback.destinations[*].external_model_config.target catalog.ModelProviderServiceConfigModelTargetConfig ALL +resources.model_services.*.config.routing.fallback.destinations[*].external_model_config.target.model string ALL +resources.model_services.*.config.routing.fallback.destinations[*].external_model_config.target.native_api_types []string ALL +resources.model_services.*.config.routing.fallback.destinations[*].external_model_config.target.native_api_types[*] string ALL +resources.model_services.*.config.routing.fallback.destinations[*].is_deleted bool ALL +resources.model_services.*.config.routing.fallback.destinations[*].name string ALL +resources.model_services.*.config.routing.fallback.destinations[*].pay_per_token_config *catalog.ModelServiceConfigPayPerTokenConfig ALL +resources.model_services.*.config.routing.fallback.destinations[*].pay_per_token_config.model string ALL +resources.model_services.*.config.routing.fallback.destinations[*].provisioned_throughput_config *catalog.ModelServiceConfigProvisionedThroughputConfig ALL +resources.model_services.*.config.routing.fallback.destinations[*].provisioned_throughput_config.model string ALL +resources.model_services.*.config.routing.fallback.destinations[*].provisioned_throughput_config.model_serving_endpoint string ALL +resources.model_services.*.config.routing.fallback.destinations[*].traffic_percentage int ALL +resources.model_services.*.config.routing.first_token_timeout *duration.Duration ALL +resources.model_services.*.config.routing.traffic_splitting *catalog.ModelServiceConfigRoutingConfigTrafficSplitting ALL +resources.model_services.*.id string INPUT +resources.model_services.*.lifecycle resources.Lifecycle INPUT +resources.model_services.*.lifecycle.prevent_destroy bool INPUT +resources.model_services.*.model_service_id string ALL +resources.model_services.*.modified_status string INPUT +resources.model_services.*.parent string ALL +resources.model_services.*.url string INPUT resources.model_serving_endpoints.*.ai_gateway *serving.AiGatewayConfig ALL resources.model_serving_endpoints.*.ai_gateway.fallback_config *serving.FallbackConfig ALL resources.model_serving_endpoints.*.ai_gateway.fallback_config.enabled bool ALL diff --git a/acceptance/bundle/resources/model_services/basic/databricks.yml b/acceptance/bundle/resources/model_services/basic/databricks.yml new file mode 100644 index 00000000000..3f82564f271 --- /dev/null +++ b/acceptance/bundle/resources/model_services/basic/databricks.yml @@ -0,0 +1,9 @@ +bundle: + name: test-bundle + +resources: + model_services: + ms1: + parent: schemas/main.myschema + model_service_id: myservice + comment: COMMENT1 diff --git a/acceptance/bundle/resources/model_services/basic/out.test.toml b/acceptance/bundle/resources/model_services/basic/out.test.toml new file mode 100644 index 00000000000..59b56a2037c --- /dev/null +++ b/acceptance/bundle/resources/model_services/basic/out.test.toml @@ -0,0 +1,3 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] +EnvMatrix.DMS = ["", "true"] diff --git a/acceptance/bundle/resources/model_services/basic/output.txt b/acceptance/bundle/resources/model_services/basic/output.txt new file mode 100644 index 00000000000..af2231f70aa --- /dev/null +++ b/acceptance/bundle/resources/model_services/basic/output.txt @@ -0,0 +1,92 @@ + +=== Initial summary before deploy +>>> [CLI] bundle summary -o json +{ + "comment": "COMMENT1", + "model_service_id": "myservice", + "modified_status": "created", + "parent": "schemas/main.myschema" +} + +=== Verify it does not exist yet +>>> musterr [CLI] ai-gateway get-model-service model-services/main.myschema.myservice +Error: Resource catalog.ModelService not found: main.myschema.myservice + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Created model_services.ms1 +Files: 5 uploaded, 0 deleted +Resources: 1 created, 0 changed, 0 deleted, 0 unchanged + +>>> print_requests.py //unity +{ + "method": "POST", + "path": "/api/2.1/unity-catalog/model-services", + "q": { + "model_service_id": "myservice", + "parent": "schemas/main.myschema" + }, + "body": { + "comment": "COMMENT1" + } +} + +=== Summary should show the id and the Catalog Explorer url +>>> [CLI] bundle summary -o json +{ + "id": "main.myschema.myservice", + "url": "[DATABRICKS_URL]/explore/data/model-services/main/myschema/myservice?w=[NUMID]" +} + +=== Verify deployment +>>> [CLI] ai-gateway get-model-service model-services/main.myschema.myservice +{ + "name": "model-services/main.myschema.myservice", + "comment": "COMMENT1" +} + +=== Update comment (should update in place, not recreate) +>>> update_file.py databricks.yml COMMENT1 COMMENT2 + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Updated model_services.ms1 +Files: 3 uploaded, 0 deleted +Resources: 0 created, 1 changed, 0 deleted, 0 unchanged + +>>> print_requests.py //unity +{ + "method": "PATCH", + "path": "/api/2.1/unity-catalog/model-services/main.myschema.myservice", + "q": { + "update_mask": "*" + }, + "body": { + "comment": "COMMENT2" + } +} + +>>> [CLI] ai-gateway get-model-service model-services/main.myschema.myservice +"COMMENT2" + +=== Change an immutable field (should plan a recreate) +>>> update_file.py databricks.yml myservice myservice-renamed + +>>> [CLI] bundle plan +recreate model_services.ms1 + +Plan: 1 to add, 0 to change, 1 to delete, 0 unchanged + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.model_services.ms1 + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/test-bundle/default + +Destroy: 1 deleted + +>>> print_requests.py //unity +{ + "method": "DELETE", + "path": "/api/2.1/unity-catalog/model-services/main.myschema.myservice" +} diff --git a/acceptance/bundle/resources/model_services/basic/script b/acceptance/bundle/resources/model_services/basic/script new file mode 100644 index 00000000000..f52a481086a --- /dev/null +++ b/acceptance/bundle/resources/model_services/basic/script @@ -0,0 +1,29 @@ +title "Initial summary before deploy" +trace $CLI bundle summary -o json | jq .resources.model_services.ms1 + +title "Verify it does not exist yet" +trace musterr $CLI ai-gateway get-model-service model-services/main.myschema.myservice + +trace $CLI bundle deploy +trace print_requests.py //unity + +title "Summary should show the id and the Catalog Explorer url" +trace $CLI bundle summary -o json | jq ".resources.model_services.ms1 | {id, url}" + +title "Verify deployment" +trace $CLI ai-gateway get-model-service model-services/main.myschema.myservice | jq '{name, comment}' + +title "Update comment (should update in place, not recreate)" +trace update_file.py databricks.yml COMMENT1 COMMENT2 +trace $CLI bundle deploy +trace print_requests.py //unity +trace $CLI ai-gateway get-model-service model-services/main.myschema.myservice | jq .comment + +title "Change an immutable field (should plan a recreate)" +trace update_file.py databricks.yml myservice myservice-renamed +trace $CLI bundle plan + +trace $CLI bundle destroy --auto-approve +trace print_requests.py //unity + +rm -f out.requests.txt diff --git a/acceptance/bundle/resources/model_services/basic/test.toml b/acceptance/bundle/resources/model_services/basic/test.toml new file mode 100644 index 00000000000..116f2013783 --- /dev/null +++ b/acceptance/bundle/resources/model_services/basic/test.toml @@ -0,0 +1,7 @@ +# Local only: this test inspects the recorded request stream (print_requests), +# which has no equivalent against a real workspace. +Cloud = false + +Ignore = [ + ".databricks", +] diff --git a/acceptance/bundle/resources/model_services/lifecycle/databricks.yml.tmpl b/acceptance/bundle/resources/model_services/lifecycle/databricks.yml.tmpl new file mode 100644 index 00000000000..ad0a522fee8 --- /dev/null +++ b/acceptance/bundle/resources/model_services/lifecycle/databricks.yml.tmpl @@ -0,0 +1,17 @@ +bundle: + name: deploy-model-service-test-$UNIQUE_NAME + +resources: + model_services: + ms: + parent: schemas/main.default + model_service_id: test_ms_$UNIQUE_NAME + comment: "Routes chat traffic to a foundation model" + config: + routing: + destinations: + - name: primary + destination_type: DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL + pay_per_token_config: + model: models/system.ai.databricks-gpt-5 + traffic_percentage: 100 diff --git a/acceptance/bundle/resources/model_services/lifecycle/out.test.toml b/acceptance/bundle/resources/model_services/lifecycle/out.test.toml new file mode 100644 index 00000000000..ae5c7bd798f --- /dev/null +++ b/acceptance/bundle/resources/model_services/lifecycle/out.test.toml @@ -0,0 +1,3 @@ +Cloud = true +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] +EnvMatrix.DMS = ["", "true"] diff --git a/acceptance/bundle/resources/model_services/lifecycle/output.txt b/acceptance/bundle/resources/model_services/lifecycle/output.txt new file mode 100644 index 00000000000..7538b2d32b6 --- /dev/null +++ b/acceptance/bundle/resources/model_services/lifecycle/output.txt @@ -0,0 +1,23 @@ + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/deploy-model-service-test-[UNIQUE_NAME]/default/files... +Created model_services.ms +Files: 5 uploaded, 0 deleted +Resources: 1 created, 0 changed, 0 deleted, 0 unchanged + +>>> [CLI] ai-gateway get-model-service model-services/main.default.test_ms_[UNIQUE_NAME] +{ + "name": "model-services/main.default.test_ms_[UNIQUE_NAME]", + "comment": "Routes chat traffic to a foundation model" +} + +>>> [CLI] bundle plan +Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.model_services.ms + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/deploy-model-service-test-[UNIQUE_NAME]/default + +Destroy: 1 deleted diff --git a/acceptance/bundle/resources/model_services/lifecycle/script b/acceptance/bundle/resources/model_services/lifecycle/script new file mode 100644 index 00000000000..ed9c256d01c --- /dev/null +++ b/acceptance/bundle/resources/model_services/lifecycle/script @@ -0,0 +1,18 @@ +envsubst < databricks.yml.tmpl > databricks.yml + +cleanup() { + trace $CLI bundle destroy --auto-approve +} +trap cleanup EXIT + +trace $CLI bundle deploy + +MODEL_SERVICE_ID=$($CLI bundle summary --output json | jq -r '.resources.model_services.ms.id') + +trace $CLI ai-gateway get-model-service "model-services/$MODEL_SERVICE_ID" | jq '{name, comment}' + +# Verify there is no drift right after deploy. Text plan (not -o json) keeps the +# output identical between the mock and a real workspace: the server populates +# output-only fields (owner, timestamps, resolved config) but the resource +# action stays "unchanged", so the rendered plan is the same. +trace $CLI bundle plan diff --git a/acceptance/bundle/resources/model_services/lifecycle/test.toml b/acceptance/bundle/resources/model_services/lifecycle/test.toml new file mode 100644 index 00000000000..402ecc18c7f --- /dev/null +++ b/acceptance/bundle/resources/model_services/lifecycle/test.toml @@ -0,0 +1,9 @@ +# Model service lifecycle also runs against a real workspace. It uses a minimal +# pay-per-token destination to a system.ai foundation model, the smallest config +# the backend accepts (a model service cannot be created with only a comment). +RecordRequests = false + +Ignore = [ + "databricks.yml", + ".databricks", +] diff --git a/acceptance/bundle/resources/model_services/remote-delete/databricks.yml b/acceptance/bundle/resources/model_services/remote-delete/databricks.yml new file mode 100644 index 00000000000..3f82564f271 --- /dev/null +++ b/acceptance/bundle/resources/model_services/remote-delete/databricks.yml @@ -0,0 +1,9 @@ +bundle: + name: test-bundle + +resources: + model_services: + ms1: + parent: schemas/main.myschema + model_service_id: myservice + comment: COMMENT1 diff --git a/acceptance/bundle/resources/model_services/remote-delete/out.test.toml b/acceptance/bundle/resources/model_services/remote-delete/out.test.toml new file mode 100644 index 00000000000..59b56a2037c --- /dev/null +++ b/acceptance/bundle/resources/model_services/remote-delete/out.test.toml @@ -0,0 +1,3 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] +EnvMatrix.DMS = ["", "true"] diff --git a/acceptance/bundle/resources/model_services/remote-delete/output.txt b/acceptance/bundle/resources/model_services/remote-delete/output.txt new file mode 100644 index 00000000000..c7f74e3b69f --- /dev/null +++ b/acceptance/bundle/resources/model_services/remote-delete/output.txt @@ -0,0 +1,20 @@ + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/test-bundle/default/files... +Created model_services.ms1 +Files: 4 uploaded, 0 deleted +Resources: 1 created, 0 changed, 0 deleted, 0 unchanged + +=== Delete the resource out of band +>>> [CLI] ai-gateway delete-model-service model-services/main.myschema.myservice + +=== Plan should detect the resource is gone and re-create it +>>> [CLI] bundle plan +create model_services.ms1 + +Plan: 1 to add, 0 to change, 0 to delete, 0 unchanged + +>>> [CLI] bundle destroy --auto-approve +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/test-bundle/default + +Destroy: 0 deleted diff --git a/acceptance/bundle/resources/model_services/remote-delete/script b/acceptance/bundle/resources/model_services/remote-delete/script new file mode 100644 index 00000000000..ad9170ef687 --- /dev/null +++ b/acceptance/bundle/resources/model_services/remote-delete/script @@ -0,0 +1,9 @@ +trace $CLI bundle deploy + +title "Delete the resource out of band" +trace $CLI ai-gateway delete-model-service model-services/main.myschema.myservice + +title "Plan should detect the resource is gone and re-create it" +trace $CLI bundle plan + +trace $CLI bundle destroy --auto-approve diff --git a/acceptance/bundle/resources/model_services/remote-delete/test.toml b/acceptance/bundle/resources/model_services/remote-delete/test.toml new file mode 100644 index 00000000000..074377ab341 --- /dev/null +++ b/acceptance/bundle/resources/model_services/remote-delete/test.toml @@ -0,0 +1,9 @@ +# Local only: simulates an out-of-band delete with a fixed resource name and no +# cloud-valid config, so it can't run against a real workspace. +Cloud = false + +RecordRequests = false + +Ignore = [ + ".databricks", +] diff --git a/acceptance/bundle/resources/model_services/test.toml b/acceptance/bundle/resources/model_services/test.toml new file mode 100644 index 00000000000..90109467fba --- /dev/null +++ b/acceptance/bundle/resources/model_services/test.toml @@ -0,0 +1,7 @@ +# AI Gateway securables are only deployable via the direct deployment engine +# (there is no Terraform provider path for them in bundles). +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +# Lifecycle tests also run against a real workspace. Tests that inspect the +# recorded request stream (e.g. basic) can't run on cloud and override this. +Cloud = true diff --git a/acceptance/experimental/open/output.txt b/acceptance/experimental/open/output.txt index 75591ed7423..881d659b71d 100644 --- a/acceptance/experimental/open/output.txt +++ b/acceptance/experimental/open/output.txt @@ -9,7 +9,7 @@ === unknown resource type >>> [CLI] experimental open --url unknown 123 -Error: unknown resource type "unknown", must be one of: alerts, apps, catalogs, cluster_policies, clusters, dashboards, database_catalogs, database_instances, experiments, genie_spaces, instance_pools, jobs, model_serving_endpoints, models, notebooks, pipelines, postgres_catalogs, postgres_synced_tables, quality_monitors, queries, registered_models, schemas, secrets, synced_database_tables, vector_search_endpoints, vector_search_indexes, volumes, warehouses +Error: unknown resource type "unknown", must be one of: alerts, apps, catalogs, cluster_policies, clusters, dashboards, database_catalogs, database_instances, experiments, genie_spaces, instance_pools, jobs, model_services, model_serving_endpoints, models, notebooks, pipelines, postgres_catalogs, postgres_synced_tables, quality_monitors, queries, registered_models, schemas, secrets, synced_database_tables, vector_search_endpoints, vector_search_indexes, volumes, warehouses === test auto-completion handler >>> [CLI] __complete experimental open , @@ -25,6 +25,7 @@ experiments genie_spaces instance_pools jobs +model_services model_serving_endpoints models notebooks diff --git a/bundle/config/mutator/resourcemutator/apply_bundle_permissions_test.go b/bundle/config/mutator/resourcemutator/apply_bundle_permissions_test.go index 55125ac916c..7aba169bd5e 100644 --- a/bundle/config/mutator/resourcemutator/apply_bundle_permissions_test.go +++ b/bundle/config/mutator/resourcemutator/apply_bundle_permissions_test.go @@ -25,6 +25,7 @@ var unsupportedResources = []string{ "schemas", "quality_monitors", "registered_models", + "model_services", "database_catalogs", "synced_database_tables", "postgres_branches", diff --git a/bundle/config/mutator/resourcemutator/apply_target_mode_test.go b/bundle/config/mutator/resourcemutator/apply_target_mode_test.go index 43483a50c77..d3645660eea 100644 --- a/bundle/config/mutator/resourcemutator/apply_target_mode_test.go +++ b/bundle/config/mutator/resourcemutator/apply_target_mode_test.go @@ -108,6 +108,9 @@ func mockBundle(mode config.Mode) *bundle.Bundle { ModelServingEndpoints: map[string]*resources.ModelServingEndpoint{ "servingendpoint1": {CreateServingEndpoint: serving.CreateServingEndpoint{Name: "servingendpoint1"}}, }, + ModelServices: map[string]*resources.ModelService{ + "modelservice1": {ModelServiceConfig: resources.ModelServiceConfig{Parent: "schemas/main.default", ModelServiceId: "modelservice1"}}, + }, RegisteredModels: map[string]*resources.RegisteredModel{ "registeredmodel1": {CreateRegisteredModelRequest: catalog.CreateRegisteredModelRequest{Name: "registeredmodel1"}}, }, diff --git a/bundle/config/mutator/resourcemutator/capture_uc_dependencies.go b/bundle/config/mutator/resourcemutator/capture_uc_dependencies.go index c4b387b76c2..1373978bfb9 100644 --- a/bundle/config/mutator/resourcemutator/capture_uc_dependencies.go +++ b/bundle/config/mutator/resourcemutator/capture_uc_dependencies.go @@ -88,6 +88,23 @@ func resolveCatalog(b *bundle.Bundle, catalogName string) string { return catalogName } +// resolveParent rewrites a `schemas/{catalog}.{schema}` parent reference so that +// a catalog or schema defined in the same bundle becomes an explicit deploy-time +// dependency. AI Gateway securables address their parent schema with this +// compound field rather than separate catalog/schema fields. +func resolveParent(b *bundle.Bundle, parent string) string { + rest, ok := strings.CutPrefix(parent, "schemas/") + if !ok { + return parent + } + parts := strings.SplitN(rest, ".", 2) + if len(parts) != 2 { + return parent + } + catalogName, schemaName := parts[0], parts[1] + return "schemas/" + resolveCatalog(b, catalogName) + "." + resolveSchema(b, catalogName, schemaName) +} + func (m *captureUCDependencies) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { // Resolve resources that depend on schemas before resolving schemas themselves. // The schema resolution below modifies schema.CatalogName, and findSchema @@ -154,6 +171,12 @@ func (m *captureUCDependencies) Apply(ctx context.Context, b *bundle.Bundle) dia acc.CatalogName = resolveCatalog(b, acc.CatalogName) } } + for _, ms := range b.Config.Resources.ModelServices { + if ms == nil { + continue + } + ms.Parent = resolveParent(b, ms.Parent) + } // Schemas are resolved last because the schema catalog resolution modifies // schema.CatalogName, and findSchema (used by resolveSchema above) matches diff --git a/bundle/config/mutator/resourcemutator/capture_uc_dependencies_test.go b/bundle/config/mutator/resourcemutator/capture_uc_dependencies_test.go index 6f5463f399d..a99db361848 100644 --- a/bundle/config/mutator/resourcemutator/capture_uc_dependencies_test.go +++ b/bundle/config/mutator/resourcemutator/capture_uc_dependencies_test.go @@ -131,6 +131,11 @@ func TestCaptureUCDependencies(t *testing.T) { }, }}, }, + ModelServices: map[string]*resources.ModelService{ + "my_model_service": {ModelServiceConfig: resources.ModelServiceConfig{ + Parent: "schemas/mycatalog.myschema", ModelServiceId: "myservice", + }}, + }, }, }, } @@ -163,6 +168,9 @@ func TestCaptureUCDependencies(t *testing.T) { itc := b.Config.Resources.ModelServingEndpoints["my_endpoint"].AiGateway.InferenceTableConfig assert.Equal(t, schemaRef, itc.SchemaName) assert.Equal(t, catalogRef, itc.CatalogName) + + // Model service (compound "schemas/{catalog}.{schema}" parent field). + assert.Equal(t, "schemas/"+catalogRef+"."+schemaRef, b.Config.Resources.ModelServices["my_model_service"].Parent) } // Pipeline schema and target are mutually exclusive; only the populated field diff --git a/bundle/config/mutator/resourcemutator/run_as_test.go b/bundle/config/mutator/resourcemutator/run_as_test.go index 356d1c6ef17..13a3491a9f3 100644 --- a/bundle/config/mutator/resourcemutator/run_as_test.go +++ b/bundle/config/mutator/resourcemutator/run_as_test.go @@ -48,6 +48,7 @@ func allResourceTypes(t *testing.T) []string { "internal_immutable_snapshots", "job_runs", "jobs", + "model_services", "model_serving_endpoints", "models", "pipelines", @@ -187,6 +188,7 @@ var allowList = []string{ "jobs", "pipelines", "models", + "model_services", "postgres_branches", "postgres_catalogs", "postgres_databases", diff --git a/bundle/config/resources.go b/bundle/config/resources.go index d62b89d9384..8f6a9de9a46 100644 --- a/bundle/config/resources.go +++ b/bundle/config/resources.go @@ -18,6 +18,7 @@ type Resources struct { Models map[string]*resources.MlflowModel `json:"models,omitempty"` Experiments map[string]*resources.MlflowExperiment `json:"experiments,omitempty"` ModelServingEndpoints map[string]*resources.ModelServingEndpoint `json:"model_serving_endpoints,omitempty"` + ModelServices map[string]*resources.ModelService `json:"model_services,omitempty"` RegisteredModels map[string]*resources.RegisteredModel `json:"registered_models,omitempty"` QualityMonitors map[string]*resources.QualityMonitor `json:"quality_monitors,omitempty"` Catalogs map[string]*resources.Catalog `json:"catalogs,omitempty"` @@ -111,6 +112,7 @@ func (r *Resources) AllResources() []ResourceGroup { collectResourceMap(descriptions["models"], r.Models), collectResourceMap(descriptions["experiments"], r.Experiments), collectResourceMap(descriptions["model_serving_endpoints"], r.ModelServingEndpoints), + collectResourceMap(descriptions["model_services"], r.ModelServices), collectResourceMap(descriptions["registered_models"], r.RegisteredModels), collectResourceMap(descriptions["quality_monitors"], r.QualityMonitors), collectResourceMap(descriptions["catalogs"], r.Catalogs), @@ -183,6 +185,7 @@ func SupportedResources() map[string]resources.ResourceDescription { "experiments": (&resources.MlflowExperiment{}).ResourceDescription(), "instance_pools": (&resources.InstancePool{}).ResourceDescription(), "model_serving_endpoints": (&resources.ModelServingEndpoint{}).ResourceDescription(), + "model_services": (&resources.ModelService{}).ResourceDescription(), "registered_models": (&resources.RegisteredModel{}).ResourceDescription(), "quality_monitors": (&resources.QualityMonitor{}).ResourceDescription(), "catalogs": (&resources.Catalog{}).ResourceDescription(), diff --git a/bundle/config/resources/model_service.go b/bundle/config/resources/model_service.go new file mode 100644 index 00000000000..6f98ffbefb9 --- /dev/null +++ b/bundle/config/resources/model_service.go @@ -0,0 +1,103 @@ +package resources + +import ( + "context" + "net/url" + + "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/workspaceurls" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/marshal" + "github.com/databricks/databricks-sdk-go/service/catalog" +) + +// ModelServiceConfig is the bundle-authored state for an AI Gateway model +// service. +// +// The SDK models the create inputs `parent` and `model_service_id` as URL +// parameters (`json:"-"`) that sit outside the ModelService body, and derives +// the resource `name` (`model-services/{catalog}.{schema}.{model_service}`) +// server-side. We therefore cannot embed catalog.CreateModelServiceRequest the +// way volume embeds catalog.CreateVolumeRequestContent: its identity fields +// would be invisible to the bundle schema. Instead we expose a flat struct with +// the immutable identity (parent + model_service_id) plus the mutable body. +// +// Owner is intentionally not exposed yet (mirrors volume, which does not manage +// owner): the API returns effective_owner rather than owner on read, so round +// tripping it needs extra care. See the direct engine resource for the CRUD. +type ModelServiceConfig struct { + // Parent schema, format `schemas/{catalog}.{schema}`. Immutable: the server + // derives `name` from parent + model_service_id, so changing it recreates + // the resource. + Parent string `json:"parent"` + // Leaf id of the model service, e.g. "my_model_service". Immutable. + ModelServiceId string `json:"model_service_id"` + // User-provided description. + Comment string `json:"comment,omitempty"` + // Operational configuration: destinations, routing, rate limits, inference + // table. + Config *catalog.ModelServiceConfig `json:"config,omitempty"` + + ForceSendFields []string `json:"-" url:"-"` +} + +func (c *ModelServiceConfig) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, c) +} + +func (c ModelServiceConfig) MarshalJSON() ([]byte, error) { + return marshal.Marshal(c) +} + +type ModelService struct { + BaseResource + ModelServiceConfig +} + +// UnmarshalJSON / MarshalJSON are defined on the wrapper so it does not inherit +// ModelServiceConfig's promoted marshaler, which would silently drop the +// BaseResource fields (id, url, lifecycle, modified_status). +func (m *ModelService) UnmarshalJSON(b []byte) error { + return marshal.Unmarshal(b, m) +} + +func (m ModelService) MarshalJSON() ([]byte, error) { + return marshal.Marshal(m) +} + +func (m *ModelService) Exists(ctx context.Context, w *databricks.WorkspaceClient, id string) (bool, error) { + // The engine tracks the id as the bare {catalog}.{schema}.{model_service}; + // the API addresses the resource by its full name. + _, err := w.AiGateway.GetModelService(ctx, catalog.GetModelServiceRequest{Name: "model-services/" + id}) + if err != nil { + log.Debugf(ctx, "model service %s does not exist", id) + if apierr.IsMissing(err) { + return false, nil + } + return false, err + } + return true, nil +} + +func (*ModelService) ResourceDescription() ResourceDescription { + return ResourceDescription{ + SingularName: "model_service", + PluralName: "model_services", + SingularTitle: "Model service", + PluralTitle: "Model services", + } +} + +func (m *ModelService) InitializeURL(baseURL url.URL) { + if m.ID == "" { + return + } + // The id is the bare {catalog}.{schema}.{model_service}; ResourceURL splits + // it into the Catalog Explorer path explore/data/model-services/... + m.URL = workspaceurls.ResourceURL(baseURL, "model_services", m.ID) +} + +func (m *ModelService) GetName() string { + return m.ModelServiceId +} diff --git a/bundle/config/resources_test.go b/bundle/config/resources_test.go index 153f89b3700..cc9d972374e 100644 --- a/bundle/config/resources_test.go +++ b/bundle/config/resources_test.go @@ -243,6 +243,11 @@ func TestResourcesBindSupport(t *testing.T) { CreateServingEndpoint: serving.CreateServingEndpoint{}, }, }, + ModelServices: map[string]*resources.ModelService{ + "my_model_service": { + ModelServiceConfig: resources.ModelServiceConfig{}, + }, + }, SecretScopes: map[string]*resources.SecretScope{ "my_secret_scope": { Name: "0", @@ -390,6 +395,7 @@ func TestResourcesBindSupport(t *testing.T) { m.GetMockAlertsV2API().EXPECT().GetAlertById(mock.Anything, mock.Anything).Return(nil, nil) m.GetMockQualityMonitorsAPI().EXPECT().Get(mock.Anything, mock.Anything).Return(nil, nil) m.GetMockServingEndpointsAPI().EXPECT().Get(mock.Anything, mock.Anything).Return(nil, nil) + m.GetMockAiGatewayAPI().EXPECT().GetModelService(mock.Anything, mock.Anything).Return(nil, nil) m.GetMockSecretsAPI().EXPECT().ListScopesAll(mock.Anything).Return([]workspace.SecretScope{ {Name: "0"}, }, nil) diff --git a/bundle/deploy/terraform/lifecycle_test.go b/bundle/deploy/terraform/lifecycle_test.go index 502ee1a0d8e..1a1fd92ba5e 100644 --- a/bundle/deploy/terraform/lifecycle_test.go +++ b/bundle/deploy/terraform/lifecycle_test.go @@ -25,6 +25,8 @@ func TestConvertLifecycleForAllResources(t *testing.T) { // so it is deployed through the direct engine only (see the resource's // acceptance test.toml). "postgres_snapshot_schedules", + // AI Gateway model service is deployed through the direct engine only. + "model_services", "secrets", "vector_search_endpoints", "vector_search_indexes", diff --git a/bundle/direct/dresources/all.go b/bundle/direct/dresources/all.go index 9f9e1badd43..bb4c0314640 100644 --- a/bundle/direct/dresources/all.go +++ b/bundle/direct/dresources/all.go @@ -36,6 +36,7 @@ var SupportedResources = map[string]any{ "genie_spaces": (*ResourceGenieSpace)(nil), "secret_scopes": (*ResourceSecretScope)(nil), "model_serving_endpoints": (*ResourceModelServingEndpoint)(nil), + "model_services": (*ResourceModelService)(nil), "quality_monitors": (*ResourceQualityMonitor)(nil), "vector_search_endpoints": (*ResourceVectorSearchEndpoint)(nil), "vector_search_indexes": (*ResourceVectorSearchIndex)(nil), diff --git a/bundle/direct/dresources/all_test.go b/bundle/direct/dresources/all_test.go index b48974eb7ed..e3771faff66 100644 --- a/bundle/direct/dresources/all_test.go +++ b/bundle/direct/dresources/all_test.go @@ -98,6 +98,14 @@ var testConfig map[string]any = map[string]any{ }, }, + "model_services": &resources.ModelService{ + ModelServiceConfig: resources.ModelServiceConfig{ + Parent: "schemas/main.default", + ModelServiceId: "my_model_service", + Comment: "Test model service", + }, + }, + "registered_models": &resources.RegisteredModel{ CreateRegisteredModelRequest: catalog.CreateRegisteredModelRequest{ Name: "my_registered_model", diff --git a/bundle/direct/dresources/apitypes.generated.yml b/bundle/direct/dresources/apitypes.generated.yml index ebb50c43c49..73d2f523c07 100644 --- a/bundle/direct/dresources/apitypes.generated.yml +++ b/bundle/direct/dresources/apitypes.generated.yml @@ -28,6 +28,8 @@ job_runs: jobs.RunNow jobs: jobs.JobSettings +model_services: catalog.CreateModelServiceRequest + model_serving_endpoints: serving.CreateServingEndpoint models: ml.CreateModelRequest diff --git a/bundle/direct/dresources/model_service.go b/bundle/direct/dresources/model_service.go new file mode 100644 index 00000000000..bef56d5e0b2 --- /dev/null +++ b/bundle/direct/dresources/model_service.go @@ -0,0 +1,144 @@ +package dresources + +import ( + "context" + "fmt" + "strings" + + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/common/types/fieldmask" + "github.com/databricks/databricks-sdk-go/service/catalog" +) + +// AI Gateway model service. +// API: https://docs.databricks.com/api/workspace/aigateway +// Terraform: databricks_ai_gateway_model_service +// +// The remote type returned by DoRead is the same bundle-local +// resources.ModelServiceConfig used for state, so RemapState is not needed. +// DoRead reconstructs the create-time identity (parent + model_service_id) from +// the server-derived resource name so those fields participate in normal drift +// detection rather than being suppressed as missing-in-remote. +// modelServiceNamePrefix is the fixed prefix of the resource name +// (model-services/{catalog}.{schema}.{model_service}). The engine tracks the id +// as the bare {catalog}.{schema}.{model_service} portion and this prefix is +// re-added when addressing the resource through the SDK. +const modelServiceNamePrefix = "model-services/" + +type ResourceModelService struct { + client *databricks.WorkspaceClient +} + +func (*ResourceModelService) New(client *databricks.WorkspaceClient) *ResourceModelService { + return &ResourceModelService{client: client} +} + +func (*ResourceModelService) PrepareState(input *resources.ModelService) *resources.ModelServiceConfig { + return &input.ModelServiceConfig +} + +// modelServiceIdentityFromName reconstructs the create-time parent and leaf id +// from the server-derived resource name +// `model-services/{catalog}.{schema}.{model_service}`. +func modelServiceIdentityFromName(name string) (parent, modelServiceId string, err error) { + rest, ok := strings.CutPrefix(name, "model-services/") + if !ok { + return "", "", fmt.Errorf("unexpected model service name %q (want model-services/{catalog}.{schema}.{model_service})", name) + } + parts := strings.Split(rest, ".") + if len(parts) != 3 { + return "", "", fmt.Errorf("unexpected model service name %q (want three dot-separated components)", name) + } + return "schemas/" + parts[0] + "." + parts[1], parts[2], nil +} + +func responseToModelServiceConfig(ms *catalog.ModelService) (*resources.ModelServiceConfig, error) { + parent, id, err := modelServiceIdentityFromName(ms.Name) + if err != nil { + return nil, err + } + return &resources.ModelServiceConfig{ + Parent: parent, + ModelServiceId: id, + Comment: ms.Comment, + Config: ms.Config, + ForceSendFields: nil, + }, nil +} + +// modelServiceBody builds the ModelService write payload from the bundle config. +// Only comment and config are client-settable; every other field is OUTPUT_ONLY +// (server-derived) and sent as its zero value. +func modelServiceBody(config *resources.ModelServiceConfig) catalog.ModelService { + return catalog.ModelService{ + Comment: config.Comment, + Config: config.Config, + CreateTime: nil, + CreatedBy: "", + EffectiveOwner: "", + Etag: "", + MetastoreId: "", + Name: "", + Owner: "", + SupportedApiTypes: nil, + UpdateTime: nil, + UpdatedBy: "", + ForceSendFields: nil, + } +} + +func (r *ResourceModelService) DoRead(ctx context.Context, id string) (*resources.ModelServiceConfig, error) { + ms, err := r.client.AiGateway.GetModelService(ctx, catalog.GetModelServiceRequest{Name: modelServiceNamePrefix + id}) + if err != nil { + return nil, err + } + return responseToModelServiceConfig(ms) +} + +func (r *ResourceModelService) DoCreate(ctx context.Context, config *resources.ModelServiceConfig) (string, *resources.ModelServiceConfig, error) { + resp, err := r.client.AiGateway.CreateModelService(ctx, catalog.CreateModelServiceRequest{ + Parent: config.Parent, + ModelServiceId: config.ModelServiceId, + ModelService: modelServiceBody(config), + }) + if err != nil { + return "", nil, err + } + state, err := responseToModelServiceConfig(resp) + if err != nil { + return "", nil, err + } + return strings.TrimPrefix(resp.Name, modelServiceNamePrefix), state, nil +} + +// DoUpdate sends update_mask "*" on every update. name, parent and +// model_service_id are immutable (recreate_on_changes in resources.yml), so the +// wildcard replaces every client-settable field (comment + a full config +// replace), matching the mask the Terraform provider generates. +// +// Etag is intentionally left empty here and in DoDelete: an empty etag means no +// If-Match precondition (last-write-wins), matching the Terraform provider, +// which also does not send etag. We deliberately do not do optimistic +// concurrency on these resources. +func (r *ResourceModelService) DoUpdate(ctx context.Context, id string, config *resources.ModelServiceConfig, _ *PlanEntry) (*resources.ModelServiceConfig, error) { + resp, err := r.client.AiGateway.UpdateModelService(ctx, catalog.UpdateModelServiceRequest{ + Etag: "", + ModelService: modelServiceBody(config), + Name: modelServiceNamePrefix + id, + UpdateMask: fieldmask.FieldMask{Paths: []string{"*"}}, + ForceSendFields: nil, + }) + if err != nil { + return nil, err + } + return responseToModelServiceConfig(resp) +} + +func (r *ResourceModelService) DoDelete(ctx context.Context, id string, _ *resources.ModelServiceConfig) error { + return r.client.AiGateway.DeleteModelService(ctx, catalog.DeleteModelServiceRequest{ + Etag: "", + Name: modelServiceNamePrefix + id, + ForceSendFields: nil, + }) +} diff --git a/bundle/direct/dresources/resources.generated.yml b/bundle/direct/dresources/resources.generated.yml index 85a0466d499..c297a90cba9 100644 --- a/bundle/direct/dresources/resources.generated.yml +++ b/bundle/direct/dresources/resources.generated.yml @@ -209,6 +209,8 @@ resources: # jobs: no api field behaviors + # model_services: no api field behaviors + model_serving_endpoints: ignore_remote_changes: diff --git a/bundle/direct/dresources/resources.yml b/bundle/direct/dresources/resources.yml index 69bf2d93aef..084a95e3f72 100644 --- a/bundle/direct/dresources/resources.yml +++ b/bundle/direct/dresources/resources.yml @@ -371,6 +371,41 @@ resources: # and the remote diverges, normal drift detection still applies. - field: config.traffic_config + model_services: + provided_id_fields: + # parent + model_service_id compose the server-derived resource name + # (model-services/{catalog}.{schema}.{model_service}), which is the ID the + # resource is fetched by. Both are immutable; a local change recreates. + # DoRead reconstructs them from the returned name, so a remote-only + # difference can only be normalization and is skipped. + - field: parent + reason: id_field + - field: model_service_id + reason: id_field + ignore_remote_changes: + # Server-resolved / tombstone fields returned on read, not user-authored. + # The generator maps model_services to the create-request type (config + # nested under model_service.*) so it can't emit these config.* paths; + # classify them by hand to avoid phantom drift. + - field: config.inference_table.table + reason: output_only + - field: config.inference_table.is_deleted + reason: output_only + # A routing target (model / MPS) that gets soft-deleted stays visible on + # read with is_deleted=true so the broken dependency is identifiable; the + # field is server-computed, not user-authored, so suppress it like above. + - field: config.routing.destinations[*].is_deleted + reason: output_only + - field: config.routing.fallback.destinations[*].is_deleted + reason: output_only + # A provisioned-throughput destination resolves the backing UC model FQN + # from model_serving_endpoint server-side; it is returned on read but + # never user-authored, so suppress it like the tombstone fields above. + - field: config.routing.destinations[*].provisioned_throughput_config.model + reason: output_only + - field: config.routing.fallback.destinations[*].provisioned_throughput_config.model + reason: output_only + registered_models: ignore_remote_changes: # Output-only timestamp/user fields populated by the backend on read. diff --git a/bundle/internal/schema/annotations.yml b/bundle/internal/schema/annotations.yml index f50d3f5bc98..a27b15632bf 100644 --- a/bundle/internal/schema/annotations.yml +++ b/bundle/internal/schema/annotations.yml @@ -1191,6 +1191,25 @@ resources: "table_update": "description": |- PLACEHOLDER + "model_services": + "description": |- + PLACEHOLDER + "$fields": + "comment": + "description": |- + PLACEHOLDER + "config": + "description": |- + PLACEHOLDER + "lifecycle": + "description": |- + PLACEHOLDER + "model_service_id": + "description": |- + PLACEHOLDER + "parent": + "description": |- + PLACEHOLDER "model_serving_endpoints": "description": |- The model serving endpoint definitions for the bundle, where each key is the name of the model serving endpoint. diff --git a/bundle/internal/validation/generated/enum_fields.go b/bundle/internal/validation/generated/enum_fields.go index 2eba9886740..2ac14ca4ebe 100644 --- a/bundle/internal/validation/generated/enum_fields.go +++ b/bundle/internal/validation/generated/enum_fields.go @@ -158,6 +158,11 @@ var EnumFields = map[string][]string{ "resources.jobs.*.triggers[*].sql_condition.trigger_mode": {"QUERY_RETURNS_ROWS", "RESULT_VALUE_CHANGES"}, "resources.jobs.*.triggers[*].table_update.condition": {"ALL_UPDATED", "ANY_UPDATED"}, + "resources.model_services.*.config.rate_limits[*].key": {"RATE_LIMIT_KEY_REQUEST_TAG", "RATE_LIMIT_KEY_SERVICE", "RATE_LIMIT_KEY_SERVICE_PRINCIPAL", "RATE_LIMIT_KEY_USER", "RATE_LIMIT_KEY_USER_DEFAULT", "RATE_LIMIT_KEY_USER_GROUP"}, + "resources.model_services.*.config.rate_limits[*].renewal_period": {"RATE_LIMIT_RENEWAL_PERIOD_HOUR", "RATE_LIMIT_RENEWAL_PERIOD_MINUTE"}, + "resources.model_services.*.config.routing.destinations[*].destination_type": {"DESTINATION_TYPE_EXTERNAL_FOUNDATION_MODEL", "DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL", "DESTINATION_TYPE_PROVISIONED_THROUGHPUT_FOUNDATION_MODEL"}, + "resources.model_services.*.config.routing.fallback.destinations[*].destination_type": {"DESTINATION_TYPE_EXTERNAL_FOUNDATION_MODEL", "DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL", "DESTINATION_TYPE_PROVISIONED_THROUGHPUT_FOUNDATION_MODEL"}, + "resources.model_serving_endpoints.*.ai_gateway.guardrails.input.pii.behavior": {"BLOCK", "MASK", "NONE"}, "resources.model_serving_endpoints.*.ai_gateway.guardrails.output.pii.behavior": {"BLOCK", "MASK", "NONE"}, "resources.model_serving_endpoints.*.ai_gateway.rate_limits[*].key": {"endpoint", "service_principal", "user", "user_group"}, diff --git a/bundle/internal/validation/generated/required_fields.go b/bundle/internal/validation/generated/required_fields.go index 7a8a027a10d..d48791e645a 100644 --- a/bundle/internal/validation/generated/required_fields.go +++ b/bundle/internal/validation/generated/required_fields.go @@ -203,6 +203,20 @@ var RequiredFields = map[string][]string{ "resources.jobs.*.webhook_notifications.on_streaming_backlog_exceeded[*]": {"id"}, "resources.jobs.*.webhook_notifications.on_success[*]": {"id"}, + "resources.model_services.*": {"parent", "model_service_id"}, + "resources.model_services.*.config.inference_table": {"parent"}, + "resources.model_services.*.config.rate_limits[*]": {"key", "renewal_period"}, + "resources.model_services.*.config.routing.destinations[*]": {"destination_type", "name"}, + "resources.model_services.*.config.routing.destinations[*].external_model_config": {"model_provider_service", "target"}, + "resources.model_services.*.config.routing.destinations[*].external_model_config.target": {"model"}, + "resources.model_services.*.config.routing.destinations[*].pay_per_token_config": {"model"}, + "resources.model_services.*.config.routing.destinations[*].provisioned_throughput_config": {"model_serving_endpoint"}, + "resources.model_services.*.config.routing.fallback.destinations[*]": {"destination_type", "name"}, + "resources.model_services.*.config.routing.fallback.destinations[*].external_model_config": {"model_provider_service", "target"}, + "resources.model_services.*.config.routing.fallback.destinations[*].external_model_config.target": {"model"}, + "resources.model_services.*.config.routing.fallback.destinations[*].pay_per_token_config": {"model"}, + "resources.model_services.*.config.routing.fallback.destinations[*].provisioned_throughput_config": {"model_serving_endpoint"}, + "resources.model_serving_endpoints.*": {"name"}, "resources.model_serving_endpoints.*.ai_gateway.fallback_config": {"enabled"}, "resources.model_serving_endpoints.*.ai_gateway.rate_limits[*]": {"renewal_period"}, diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index 6db8c5cc7e2..b291da0f146 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -1741,6 +1741,39 @@ } ] }, + "resources.ModelService": { + "oneOf": [ + { + "type": "object", + "properties": { + "comment": { + "$ref": "#/$defs/string" + }, + "config": { + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfig" + }, + "lifecycle": { + "$ref": "#/$defs/github.com/databricks/cli/bundle/config/resources.Lifecycle" + }, + "model_service_id": { + "$ref": "#/$defs/string" + }, + "parent": { + "$ref": "#/$defs/string" + } + }, + "additionalProperties": false, + "required": [ + "parent", + "model_service_id" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "resources.ModelServingEndpoint": { "oneOf": [ { @@ -3922,6 +3955,9 @@ "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.Job", "markdownDescription": "The job definitions for the bundle, where each key is the name of the job. See [jobs](https://docs.databricks.com/dev-tools/bundles/resources.html#jobs)." }, + "model_services": { + "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.ModelService" + }, "model_serving_endpoints": { "description": "The model serving endpoint definitions for the bundle, where each key is the name of the model serving endpoint.", "$ref": "#/$defs/map/github.com/databricks/cli/bundle/config/resources.ModelServingEndpoint", @@ -5335,6 +5371,311 @@ } ] }, + "catalog.InferenceTableConfig": { + "oneOf": [ + { + "type": "object", + "description": "Inference table configuration for payload logging on a model service.\n\n`parent` is always REQUIRED when the sub-message is set; the destination\nUC schema is needed to construct or rebind the payload TABLE regardless of\nwhether payload logging is currently active. Payload logging is active by\ndefault; set `disabled = true` to pause runtime logging without dropping the\ntable or the binding.", + "properties": { + "disabled": { + "description": "[Beta] Indicates whether payload logging is disabled (opt-out). Unset means that\npayload logging is active (the on-by-default state coincides with the proto\nzero-value, so the server never fills this field for a client that leaves it\nunset). Set `disabled = true` to pause runtime logging while keeping the\nsub-message attached (preserving `parent` and `table_name_prefix` for a\nlater flip back to active). `parent` remains required either way.", + "$ref": "#/$defs/bool", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "parent": { + "description": "[Beta] Parent UC schema where the inference table is created.\nFormat: `schemas/{catalog}.{schema}`. Set at create time and immutable\nthereafter; changing it on an existing service is rejected.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "table_name_prefix": { + "description": "[Beta] Prefix for the inference-table's UC-registered name. The actual leaf name UC\nstores is `\u003ctable_name_prefix\u003e_payload`; the `_payload` suffix is appended\nautomatically. To find the actual UC table after Create, read the `table`\nfield on the response. Defaults to `\u003cmodel_service_name\u003e_payload` when unset.\nSet at create time and immutable thereafter; changing it on an existing\nservice is rejected.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "parent" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelProviderServiceConfigModelTargetConfig": { + "oneOf": [ + { + "type": "object", + "description": "Model target configuration for an external model destination.", + "properties": { + "model": { + "description": "[Beta] Provider-side model identifier (e.g. \"gpt-5\", \"claude-opus-4-7\"). This is\na string on the LLM provider's side, not a UC entity. The UC governance\nhook for external destinations is the ModelProviderService referenced by\n`ExternalModelConfig.model_provider_service`, not the model itself.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "native_api_types": { + "description": "[Beta] Provider-native API types the model supports (e.g.\n\"openai/v1/chat/completions\"). Used by the platform for request/response\ntranslation from the unified API type. At most 64 entries of at most 256\ncharacters each; the list is persisted into the destination binding's\nbounded storage envelope.", + "$ref": "#/$defs/slice/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "model" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfig": { + "oneOf": [ + { + "type": "object", + "description": "Operational configuration wrapped around the ModelService resource.", + "properties": { + "inference_table": { + "description": "[Beta] Inference table config for payload logging.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.InferenceTableConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "rate_limits": { + "description": "[Beta] Rate limits applied to requests routed through this model service.", + "$ref": "#/$defs/slice/github.com/databricks/databricks-sdk-go/service/catalog.RateLimit", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "routing": { + "description": "[Beta] Routing configuration: destinations, routing strategy, and fallback.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigRoutingConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigDestinationConfig": { + "oneOf": [ + { + "type": "object", + "description": "A destination the model service can route traffic to. Exactly one of the\nper-type configs inside `type_config` must be set, and it must match\n`destination_type`.", + "properties": { + "destination_type": { + "description": "[Beta] Backing-model category. Determines which oneof variant is populated.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigDestinationConfigDestinationType", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "external_model_config": { + "description": "[Beta] Configuration for an external-foundation-model destination. Provider auth\nand provider-specific cloud configuration are owned by a separate, governed\nModelProviderService entity referenced via `model_provider_service`; the\nplatform resolves the provider at invocation time.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigExternalModelConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "name": { + "description": "[Beta] User-facing label for this destination, used in routing references.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "pay_per_token_config": { + "description": "[Beta] Configuration for a pay-per-token foundation-model destination. Identifies\nthe foundation model by its UC resource name; the platform resolves it to a\nModel Serving endpoint at request time.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigPayPerTokenConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "provisioned_throughput_config": { + "description": "[Beta] Configuration for a provisioned-throughput foundation-model destination.\nReferences a pre-existing Model Serving endpoint that serves the model;\nsizing (provisioned throughput, burst scaling, model version) is owned by\nthe Model Serving endpoint itself, not by this message.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigProvisionedThroughputConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "traffic_percentage": { + "description": "[Beta] Share of traffic sent to this destination, 0-100. Optional on fallback\ndestinations; see FallbackConfig.", + "$ref": "#/$defs/int", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "destination_type", + "name" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigDestinationConfigDestinationType": { + "oneOf": [ + { + "type": "string", + "description": "Backing-model category for a model service destination.", + "enum": [ + "DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL", + "DESTINATION_TYPE_PROVISIONED_THROUGHPUT_FOUNDATION_MODEL", + "DESTINATION_TYPE_EXTERNAL_FOUNDATION_MODEL" + ], + "enumDescriptions": [ + "[Beta]", + "[Beta]", + "[Beta]" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigExternalModelConfig": { + "oneOf": [ + { + "type": "object", + "description": "Configuration for an external-foundation-model destination. Provider auth\nand provider-specific cloud configuration are owned by a separate, governed\nModelProviderService entity referenced via `model_provider_service`; the\nplatform resolves the provider at invocation time.", + "properties": { + "model_provider_service": { + "description": "[Beta] Resource name of the governed ModelProviderService that owns provider\nauth and provider-specific configuration. The referenced\nModelProviderService also carries the provider type, so this message\ndoes not surface it directly.\nFormat: `model-provider-services/{catalog}.{schema}.{model_provider_service}`.\nEach `{...}` component is capped at 255 characters individually.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "target": { + "description": "[Beta] Routing target for the destination: the provider-side model selected from\nthe referenced ModelProviderService's `targets` catalog, plus the unified\nAPI types the platform should translate to/from at request time.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelProviderServiceConfigModelTargetConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "model_provider_service", + "target" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigFallbackConfig": { + "oneOf": [ + { + "type": "object", + "description": "Fallback routing, applied after the primary destination returns a retryable\nerror. Traversal is in list order; the attempt count is the length of the\nlist.", + "properties": { + "destinations": { + "description": "[Beta] Ordered list of fallback destinations. Traversal is in list order; the\nattempt count is the length of the list. At most 5 are allowed.", + "$ref": "#/$defs/slice/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigDestinationConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigPayPerTokenConfig": { + "oneOf": [ + { + "type": "object", + "description": "Configuration for a pay-per-token foundation-model destination. Identifies\nthe foundation model by its UC resource name; the platform resolves it to a\nModel Serving endpoint at request time.", + "properties": { + "model": { + "description": "[Beta] Resource name of the UC model.\nFormat: `models/{catalog}.{schema}.{model}`.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "model" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigProvisionedThroughputConfig": { + "oneOf": [ + { + "type": "object", + "description": "Configuration for a provisioned-throughput foundation-model destination.\nReferences a pre-existing Model Serving endpoint that serves the model;\nsizing (provisioned throughput, burst scaling, model version) is owned by\nthe Model Serving endpoint itself, not by this message.", + "properties": { + "model_serving_endpoint": { + "description": "[Beta] Name of the backing Model Serving endpoint serving the provisioned-\nthroughput foundation model, as the AIP-122 typed resource name\n`serving-endpoints/{name}`. The same UC model can be served on multiple\nModel Serving endpoints (different throughput / region / config); the\ncaller picks which one this destination routes to. The endpoint must\nexist at create time.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "model_serving_endpoint" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigRoutingConfig": { + "oneOf": [ + { + "type": "object", + "description": "Routing configuration for a model service, nesting destinations, routing\nstrategy, and fallback under a single sub-message.", + "properties": { + "destinations": { + "description": "[Beta] Primary routing destinations. At most 10 are allowed. At least one is\nrequired on CreateModelService; on UpdateModelService it is required only\nwhen `config.routing` (or a `config.routing.*` subpath) appears in\n`update_mask`.", + "$ref": "#/$defs/slice/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigDestinationConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "fallback": { + "description": "[Beta] Fallback routing config, applied after primary destinations fail.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigFallbackConfig", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "first_token_timeout": { + "description": "[Beta] Timeout for the first token of a streaming response. If a destination does\nnot return its first token within this duration, AI Gateway aborts the\nattempt and fails over to the next destination. Applies to streaming\nrequests only. Leave unset for no first-token timeout.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/common/types/duration.Duration", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "traffic_splitting": { + "description": "[Beta] Marker message selecting request-based traffic splitting. Traffic is\ndistributed according to each destination's traffic_percentage value;\nno configuration lives on this message itself.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigRoutingConfigTrafficSplitting", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.ModelServiceConfigRoutingConfigTrafficSplitting": { + "oneOf": [ + { + "type": "object", + "description": "Marker message selecting request-based traffic splitting across primary\ndestinations. Split weights are read from each DestinationConfig.traffic_percentage.", + "additionalProperties": false + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "catalog.MonitorCronSchedule": { "oneOf": [ { @@ -5715,6 +6056,108 @@ } ] }, + "catalog.RateLimit": { + "oneOf": [ + { + "type": "object", + "description": "A rate limit applied to service requests. Leave `requests` or `tokens`\nunset to impose no limit on that dimension; set a value to cap that dimension\nwithin the renewal period.", + "properties": { + "key": { + "description": "[Beta] Scope key. Determines whether `principal` is required.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.RateLimitRateLimitKey", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "principal": { + "description": "[Beta] Principal this limit applies to: user email, group name, or service\nprincipal application ID. Required unless `key` is\n`RATE_LIMIT_KEY_SERVICE`, `RATE_LIMIT_KEY_USER_DEFAULT`, or\n`RATE_LIMIT_KEY_REQUEST_TAG` (which must not set a principal).", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "renewal_period": { + "description": "[Beta] Renewal period.", + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.RateLimitRateLimitRenewalPeriod", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "request_tag_key": { + "description": "[Beta] Request tag key this limit applies to. Required when `key` is\n`RATE_LIMIT_KEY_REQUEST_TAG`, forbidden otherwise.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "request_tag_value": { + "description": "[Beta] Request tag value this limit applies to. Only valid when `key` is\n`RATE_LIMIT_KEY_REQUEST_TAG`. Leave unset to apply the limit to every\nvalue of `request_tag_key` (an any-value default); a set value is a\nspecific override for that value.", + "$ref": "#/$defs/string", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "requests": { + "description": "[Beta] Max requests allowed within a renewal period. Leave unset for no request limit.", + "$ref": "#/$defs/int64", + "x-databricks-launch-stage": "PUBLIC_BETA" + }, + "tokens": { + "description": "[Beta] Max tokens allowed within a renewal period. Leave unset for no token limit.", + "$ref": "#/$defs/int64", + "x-databricks-launch-stage": "PUBLIC_BETA" + } + }, + "additionalProperties": false, + "required": [ + "key", + "renewal_period" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.RateLimitRateLimitKey": { + "oneOf": [ + { + "type": "string", + "description": "Scope key for a rate limit.", + "enum": [ + "RATE_LIMIT_KEY_USER", + "RATE_LIMIT_KEY_USER_GROUP", + "RATE_LIMIT_KEY_SERVICE_PRINCIPAL", + "RATE_LIMIT_KEY_SERVICE", + "RATE_LIMIT_KEY_USER_DEFAULT", + "RATE_LIMIT_KEY_REQUEST_TAG" + ], + "enumDescriptions": [ + "[Beta]", + "[Beta]", + "[Beta]", + "[Beta]", + "[Beta]", + "[Beta]" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, + "catalog.RateLimitRateLimitRenewalPeriod": { + "oneOf": [ + { + "type": "string", + "description": "Renewal period for a rate limit.", + "enum": [ + "RATE_LIMIT_RENEWAL_PERIOD_MINUTE", + "RATE_LIMIT_RENEWAL_PERIOD_HOUR" + ], + "enumDescriptions": [ + "[Beta]", + "[Beta]" + ] + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "catalog.RegisteredModelAlias": { "oneOf": [ { @@ -16743,6 +17186,20 @@ } ] }, + "resources.ModelService": { + "oneOf": [ + { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/github.com/databricks/cli/bundle/config/resources.ModelService" + } + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "resources.ModelServingEndpoint": { "oneOf": [ { @@ -17404,6 +17861,20 @@ } ] }, + "catalog.ModelServiceConfigDestinationConfig": { + "oneOf": [ + { + "type": "array", + "items": { + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.ModelServiceConfigDestinationConfig" + } + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "catalog.MonitorMetric": { "oneOf": [ { @@ -17446,6 +17917,20 @@ } ] }, + "catalog.RateLimit": { + "oneOf": [ + { + "type": "array", + "items": { + "$ref": "#/$defs/github.com/databricks/databricks-sdk-go/service/catalog.RateLimit" + } + }, + { + "type": "string", + "pattern": "\\$\\{(var(\\._*\\p{L}+([-_]*[\\p{L}\\p{N}]+)*(\\[[0-9]+\\])*)+)\\}" + } + ] + }, "catalog.RegisteredModelAlias": { "oneOf": [ { diff --git a/bundle/statemgmt/state_load_test.go b/bundle/statemgmt/state_load_test.go index 32027bf0914..c4730fd9ae9 100644 --- a/bundle/statemgmt/state_load_test.go +++ b/bundle/statemgmt/state_load_test.go @@ -32,6 +32,7 @@ func TestStateToBundleEmptyLocalResources(t *testing.T) { "resources.models.test_mlflow_model": {ID: "1"}, "resources.experiments.test_mlflow_experiment": {ID: "1"}, "resources.model_serving_endpoints.test_model_serving": {ID: "1"}, + "resources.model_services.test_model_service": {ID: "main.default.test_model_service"}, "resources.registered_models.test_registered_model": {ID: "1"}, "resources.quality_monitors.test_monitor": {ID: "1"}, "resources.catalogs.test_catalog": {ID: "1"}, @@ -165,6 +166,9 @@ func TestStateToBundleEmptyLocalResources(t *testing.T) { assert.Equal(t, "main.default.test_secret", config.Resources.Secrets["test_secret"].ID) assert.Equal(t, resources.ModifiedStatusDeleted, config.Resources.Secrets["test_secret"].ModifiedStatus) + assert.Equal(t, "main.default.test_model_service", config.Resources.ModelServices["test_model_service"].ID) + assert.Equal(t, resources.ModifiedStatusDeleted, config.Resources.ModelServices["test_model_service"].ModifiedStatus) + AssertFullResourceCoverage(t, &config) } @@ -424,6 +428,14 @@ func TestStateToBundleEmptyRemoteResources(t *testing.T) { }, }, }, + ModelServices: map[string]*resources.ModelService{ + "test_model_service": { + ModelServiceConfig: resources.ModelServiceConfig{ + Parent: "schemas/main.default", + ModelServiceId: "test_model_service", + }, + }, + }, }, } @@ -535,6 +547,9 @@ func TestStateToBundleEmptyRemoteResources(t *testing.T) { assert.Empty(t, config.Resources.ClusterPolicies["test_cluster_policy"].ID) assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.ClusterPolicies["test_cluster_policy"].ModifiedStatus) + assert.Empty(t, config.Resources.ModelServices["test_model_service"].ID) + assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.ModelServices["test_model_service"].ModifiedStatus) + AssertFullResourceCoverage(t, &config) } @@ -961,6 +976,20 @@ func TestStateToBundleModifiedResources(t *testing.T) { }, }, }, + ModelServices: map[string]*resources.ModelService{ + "test_model_service": { + ModelServiceConfig: resources.ModelServiceConfig{ + Parent: "schemas/main.default", + ModelServiceId: "test_model_service", + }, + }, + "test_model_service_new": { + ModelServiceConfig: resources.ModelServiceConfig{ + Parent: "schemas/main.default", + ModelServiceId: "test_model_service_new", + }, + }, + }, }, } state := ExportedResourcesMap{ @@ -1022,6 +1051,8 @@ func TestStateToBundleModifiedResources(t *testing.T) { "resources.instance_pools.test_instance_pool_old": {ID: "2"}, "resources.cluster_policies.test_cluster_policy": {ID: "cp-1"}, "resources.cluster_policies.test_cluster_policy_old": {ID: "cp-2"}, + "resources.model_services.test_model_service": {ID: "main.default.test_model_service"}, + "resources.model_services.test_model_service_old": {ID: "main.default.test_model_service_old"}, "resources.secrets.test_secret": {ID: "main.default.test_secret"}, "resources.secrets.test_secret_old": {ID: "main.default.test_secret_old"}, } @@ -1233,6 +1264,13 @@ func TestStateToBundleModifiedResources(t *testing.T) { assert.Empty(t, config.Resources.ClusterPolicies["test_cluster_policy_new"].ID) assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.ClusterPolicies["test_cluster_policy_new"].ModifiedStatus) + assert.Equal(t, "main.default.test_model_service", config.Resources.ModelServices["test_model_service"].ID) + assert.Empty(t, config.Resources.ModelServices["test_model_service"].ModifiedStatus) + assert.Equal(t, "main.default.test_model_service_old", config.Resources.ModelServices["test_model_service_old"].ID) + assert.Equal(t, resources.ModifiedStatusDeleted, config.Resources.ModelServices["test_model_service_old"].ModifiedStatus) + assert.Empty(t, config.Resources.ModelServices["test_model_service_new"].ID) + assert.Equal(t, resources.ModifiedStatusCreated, config.Resources.ModelServices["test_model_service_new"].ModifiedStatus) + assert.Equal(t, "main.default.test_secret", config.Resources.Secrets["test_secret"].ID) assert.Empty(t, config.Resources.Secrets["test_secret"].ModifiedStatus) assert.Equal(t, "main.default.test_secret_old", config.Resources.Secrets["test_secret_old"].ID) diff --git a/cmd/experimental/workspace_open_test.go b/cmd/experimental/workspace_open_test.go index 95502694472..1504fed6a03 100644 --- a/cmd/experimental/workspace_open_test.go +++ b/cmd/experimental/workspace_open_test.go @@ -67,7 +67,7 @@ func TestBuildWorkspaceURLFragmentBasedResources(t *testing.T) { func TestBuildWorkspaceURLUnknownResourceType(t *testing.T) { _, err := workspaceurls.BuildResourceURL("https://myworkspace.databricks.com", "unknown", "123", "") assert.ErrorContains(t, err, "unknown resource type \"unknown\"") - assert.ErrorContains(t, err, "alerts, apps, catalogs, cluster_policies, clusters, dashboards, database_catalogs, database_instances, experiments, genie_spaces, instance_pools, jobs, model_serving_endpoints, models, notebooks, pipelines, postgres_catalogs, postgres_synced_tables, quality_monitors, queries, registered_models, schemas, secrets, synced_database_tables, vector_search_endpoints, vector_search_indexes, volumes, warehouses") + assert.ErrorContains(t, err, "alerts, apps, catalogs, cluster_policies, clusters, dashboards, database_catalogs, database_instances, experiments, genie_spaces, instance_pools, jobs, model_services, model_serving_endpoints, models, notebooks, pipelines, postgres_catalogs, postgres_synced_tables, quality_monitors, queries, registered_models, schemas, secrets, synced_database_tables, vector_search_endpoints, vector_search_indexes, volumes, warehouses") } func TestBuildWorkspaceURLHostWithTrailingSlash(t *testing.T) { @@ -119,6 +119,7 @@ func TestWorkspaceOpenCommandCompletion(t *testing.T) { "genie_spaces", "instance_pools", "jobs", + "model_services", "model_serving_endpoints", "models", "notebooks", @@ -149,7 +150,7 @@ func TestWorkspaceOpenCommandCompletionSecondArg(t *testing.T) { func TestWorkspaceOpenCommandHelpText(t *testing.T) { cmd := newWorkspaceOpenCommand() - assert.Contains(t, cmd.Long, "Supported resource types: alerts, apps, catalogs, cluster_policies, clusters, dashboards, database_catalogs, database_instances, experiments, genie_spaces, instance_pools, jobs, model_serving_endpoints, models, notebooks, pipelines, postgres_catalogs, postgres_synced_tables, quality_monitors, queries, registered_models, schemas, secrets, synced_database_tables, vector_search_endpoints, vector_search_indexes, volumes, warehouses.") + assert.Contains(t, cmd.Long, "Supported resource types: alerts, apps, catalogs, cluster_policies, clusters, dashboards, database_catalogs, database_instances, experiments, genie_spaces, instance_pools, jobs, model_services, model_serving_endpoints, models, notebooks, pipelines, postgres_catalogs, postgres_synced_tables, quality_monitors, queries, registered_models, schemas, secrets, synced_database_tables, vector_search_endpoints, vector_search_indexes, volumes, warehouses.") assert.Contains(t, cmd.Long, "databricks experimental open jobs 123456789") assert.Contains(t, cmd.Long, "databricks experimental open notebooks /Users/user@example.com/my-notebook") assert.Contains(t, cmd.Long, "databricks experimental open registered_models catalog.schema.my_model") diff --git a/libs/testserver/fake_workspace.go b/libs/testserver/fake_workspace.go index d5bfbe796b1..36ed8db3da2 100644 --- a/libs/testserver/fake_workspace.go +++ b/libs/testserver/fake_workspace.go @@ -204,6 +204,7 @@ type FakeWorkspace struct { Catalogs map[string]catalog.CatalogInfo ExternalLocations map[string]catalog.ExternalLocationInfo RegisteredModels map[string]catalog.RegisteredModelInfo + ModelServices map[string]catalog.ModelService ServingEndpoints map[string]serving.ServingEndpointDetailed VectorSearchEndpoints map[string]vectorsearch.EndpointInfo VectorSearchIndexes map[string]fakeVectorSearchIndex @@ -479,6 +480,7 @@ func NewFakeWorkspace(url, token string) *FakeWorkspace { ExternalLocations: map[string]catalog.ExternalLocationInfo{}, Schemas: map[string]catalog.SchemaInfo{}, RegisteredModels: map[string]catalog.RegisteredModelInfo{}, + ModelServices: map[string]catalog.ModelService{}, Volumes: map[string]catalog.VolumeInfo{}, Dashboards: NewEventualMap[string, *fakeDashboard](strings.HasPrefix(token, EventualConsistencyTokenPrefix)), PublishedDashboards: map[string]dashboards.PublishedDashboard{}, diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 500b1f6f453..ce2e6411102 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -630,6 +630,24 @@ func AddDefaultHandlers(server *Server) { return MapDelete(req.Workspace, req.Workspace.RegisteredModels, req.Vars["full_name"]) }) + // Model Services (AI Gateway): + + server.Handle("POST", "/api/2.1/unity-catalog/model-services", func(req Request) any { + return req.Workspace.ModelServicesCreate(req) + }) + + server.Handle("GET", "/api/2.1/unity-catalog/model-services/{name}", func(req Request) any { + return MapGet(req.Workspace, req.Workspace.ModelServices, req.Vars["name"]) + }) + + server.Handle("PATCH", "/api/2.1/unity-catalog/model-services/{name}", func(req Request) any { + return req.Workspace.ModelServicesUpdate(req, req.Vars["name"]) + }) + + server.Handle("DELETE", "/api/2.1/unity-catalog/model-services/{name}", func(req Request) any { + return MapDelete(req.Workspace, req.Workspace.ModelServices, req.Vars["name"]) + }) + // Volumes: server.Handle("GET", "/api/2.1/unity-catalog/volumes/{full_name}", func(req Request) any { diff --git a/libs/testserver/model_services.go b/libs/testserver/model_services.go new file mode 100644 index 00000000000..ce653543ac5 --- /dev/null +++ b/libs/testserver/model_services.go @@ -0,0 +1,77 @@ +package testserver + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/databricks/databricks-sdk-go/service/catalog" +) + +// ModelServicesCreate fakes POST /api/2.1/unity-catalog/model-services. +// +// `parent` (schemas/{catalog}.{schema}) and `model_service_id` arrive as query +// parameters; the ModelService body is sent directly (not wrapped in +// "model_service"). The server derives the resource name +// model-services/{catalog}.{schema}.{model_service}. +// The map is keyed by the {catalog}.{schema}.{model_service} portion, which is +// the path segment used on subsequent get/update/delete. +func (s *FakeWorkspace) ModelServicesCreate(req Request) Response { + defer s.LockUnlock()() + + // The SDK sends the ModelService body directly (not wrapped in + // "model_service"); parent and model_service_id are query parameters. + var ms catalog.ModelService + if err := json.Unmarshal(req.Body, &ms); err != nil { + return Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: http.StatusInternalServerError, + } + } + + schema := strings.TrimPrefix(req.URL.Query().Get("parent"), "schemas/") + key := schema + "." + req.URL.Query().Get("model_service_id") + + ms.Name = "model-services/" + key + ms.CreatedBy = s.CurrentUser().UserName + ms.UpdatedBy = s.CurrentUser().UserName + ms.EffectiveOwner = s.CurrentUser().UserName + ms.MetastoreId = nextUUID() + + s.ModelServices[key] = ms + return Response{ + Body: ms, + } +} + +func (s *FakeWorkspace) ModelServicesUpdate(req Request, name string) Response { + defer s.LockUnlock()() + + existing, ok := s.ModelServices[name] + if !ok { + return Response{ + StatusCode: http.StatusNotFound, + Body: fmt.Sprintf("model service %s not found", name), + } + } + + // The SDK sends the ModelService body directly; update_mask is a query param. + var incoming catalog.ModelService + if err := json.Unmarshal(req.Body, &incoming); err != nil { + return Response{ + Body: fmt.Sprintf("internal error: %s", err), + StatusCode: http.StatusInternalServerError, + } + } + + // Apply the mutable fields carried in the update mask (comment, config). + existing.Comment = incoming.Comment + existing.Config = incoming.Config + existing.UpdatedBy = s.CurrentUser().UserName + + s.ModelServices[name] = existing + return Response{ + Body: existing, + } +} diff --git a/libs/workspaceurls/urls.go b/libs/workspaceurls/urls.go index 0a1680a0753..1eb8dcd034f 100644 --- a/libs/workspaceurls/urls.go +++ b/libs/workspaceurls/urls.go @@ -22,6 +22,7 @@ var resourceURLPatterns = map[string]string{ "jobs": "jobs/%s", "models": "ml/models/%s", "model_serving_endpoints": "ml/endpoints/%s", + "model_services": "explore/data/model-services/%s", "notebooks": "#notebook/%s", "pipelines": "pipelines/%s", "postgres_catalogs": "explore/data/%s", @@ -53,6 +54,7 @@ var resourceAliases = map[string]string{ // requires slash-separated segments. var dotSeparatedResources = map[string]bool{ "catalogs": true, + "model_services": true, "postgres_synced_tables": true, "quality_monitors": true, "registered_models": true, diff --git a/python/databricks/bundles/core/__init__.py b/python/databricks/bundles/core/__init__.py index bf0fd205dc5..d3991c493de 100644 --- a/python/databricks/bundles/core/__init__.py +++ b/python/databricks/bundles/core/__init__.py @@ -30,6 +30,7 @@ "load_resources_from_package_module", "mlflow_experiment_mutator", "mlflow_model_mutator", + "model_service_mutator", "model_serving_endpoint_mutator", "pipeline_mutator", "postgres_snapshot_schedule_mutator", @@ -65,6 +66,7 @@ job_run_mutator, mlflow_experiment_mutator, mlflow_model_mutator, + model_service_mutator, model_serving_endpoint_mutator, pipeline_mutator, postgres_snapshot_schedule_mutator, diff --git a/python/databricks/bundles/core/_generated/__init__.py b/python/databricks/bundles/core/_generated/__init__.py index a28d59b6c7d..f785073c1ce 100644 --- a/python/databricks/bundles/core/_generated/__init__.py +++ b/python/databricks/bundles/core/_generated/__init__.py @@ -37,6 +37,10 @@ job_run_mutator, ) from databricks.bundles.core._generated.jobs import _JobResources, job_mutator +from databricks.bundles.core._generated.model_services import ( + _ModelServiceResources, + model_service_mutator, +) from databricks.bundles.core._generated.model_serving_endpoints import ( _ModelServingEndpointResources, model_serving_endpoint_mutator, @@ -103,6 +107,7 @@ "job_run_mutator", "mlflow_experiment_mutator", "mlflow_model_mutator", + "model_service_mutator", "model_serving_endpoint_mutator", "pipeline_mutator", "postgres_snapshot_schedule_mutator", @@ -131,6 +136,7 @@ class _GeneratedResources( _InstancePoolResources, _JobRunResources, _JobResources, + _ModelServiceResources, _ModelServingEndpointResources, _MlflowModelResources, _PipelineResources, @@ -162,6 +168,7 @@ def _all_resource_types() -> "tuple[_ResourceType, ...]": instance_pools, job_runs, jobs, + model_services, model_serving_endpoints, models, pipelines, @@ -190,6 +197,7 @@ def _all_resource_types() -> "tuple[_ResourceType, ...]": instance_pools._resource_type(), job_runs._resource_type(), jobs._resource_type(), + model_services._resource_type(), model_serving_endpoints._resource_type(), models._resource_type(), pipelines._resource_type(), diff --git a/python/databricks/bundles/core/_generated/model_services.py b/python/databricks/bundles/core/_generated/model_services.py new file mode 100644 index 00000000000..2b37ebe7d6b --- /dev/null +++ b/python/databricks/bundles/core/_generated/model_services.py @@ -0,0 +1,118 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from collections.abc import Callable +from typing import TYPE_CHECKING, Optional, overload + +from databricks.bundles.core._bundle import Bundle +from databricks.bundles.core._location import Location +from databricks.bundles.core._resource_mutator import ResourceMutator +from databricks.bundles.core._transform import _transform + +if TYPE_CHECKING: + from databricks.bundles.core._resource_type import _ResourceType + from databricks.bundles.model_services._models.model_service import ( + ModelService, + ModelServiceParam, + ) + + +def _resource_type() -> "_ResourceType": + from databricks.bundles.core._resource_type import _ResourceType + from databricks.bundles.model_services._models.model_service import ModelService + + return _ResourceType( + resource_type=ModelService, + singular_name="model_service", + plural_name="model_services", + ) + + +class _ModelServiceResources: + """ + Generated model_service accessors, mixed into Resources. + """ + + # Provided by the Resources subclass; declared here so the generated methods + # below type-check. + _resources: dict[str, dict] + + if TYPE_CHECKING: + + def add_location(self, path: tuple[str, ...], location: Location) -> None: ... + + def add_diagnostic_error( + self, + msg: str, + *, + detail: Optional[str] = None, + path: Optional[tuple[str, ...]] = None, + location: Optional[Location] = None, + ) -> None: ... + + @property + def model_services(self) -> dict[str, "ModelService"]: + return self._resources["model_services"] + + def add_model_service( + self, + resource_name: str, + model_service: "ModelServiceParam", + *, + location: Optional[Location] = None, + ) -> None: + """ + Adds the resource model_service to the collection of resources. Resource name must be unique across all model_services. + + :param resource_name: unique identifier for the model_service + :param model_service: the model_service to add, can be ModelService or dict + :param location: optional location of the model_service in the source code + """ + from databricks.bundles.model_services._models.model_service import ModelService + + model_service = _transform(ModelService, model_service) + path = ("resources", "model_services", resource_name) + location = location or Location.from_stack_frame(depth=1) + + if self._resources["model_services"].get(resource_name): + self.add_diagnostic_error( + msg=f"Duplicate resource name '{resource_name}' for resource 'model_service'. Resource names must be unique.", + location=location, + path=path, + ) + else: + if location: + self.add_location(path, location) + + self._resources["model_services"][resource_name] = model_service + + +@overload +def model_service_mutator( + function: Callable[[Bundle, "ModelService"], "ModelService"], +) -> ResourceMutator["ModelService"]: ... + + +@overload +def model_service_mutator( + function: Callable[["ModelService"], "ModelService"], +) -> ResourceMutator["ModelService"]: ... + + +def model_service_mutator(function: Callable) -> ResourceMutator["ModelService"]: + """ + Decorator for defining mutator for model_services. Function should return a new instance of the model_service + with the desired changes, instead of mutating the input model_service. + + Example: + + .. code-block:: python + + @model_service_mutator + def my_model_service_mutator(bundle: Bundle, model_service: ModelService) -> ModelService: + return replace(model_service, ...) + + :param function: Function that mutates model_services. + """ + from databricks.bundles.model_services._models.model_service import ModelService + + return ResourceMutator(resource_type=ModelService, function=function) diff --git a/python/databricks/bundles/model_services/__init__.py b/python/databricks/bundles/model_services/__init__.py new file mode 100644 index 00000000000..0a3be24da20 --- /dev/null +++ b/python/databricks/bundles/model_services/__init__.py @@ -0,0 +1,128 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +__all__ = [ + "InferenceTableConfig", + "InferenceTableConfigDict", + "InferenceTableConfigParam", + "Lifecycle", + "LifecycleDict", + "LifecycleParam", + "ModelProviderServiceConfigModelTargetConfig", + "ModelProviderServiceConfigModelTargetConfigDict", + "ModelProviderServiceConfigModelTargetConfigParam", + "ModelService", + "ModelServiceConfig", + "ModelServiceConfigDestinationConfig", + "ModelServiceConfigDestinationConfigDestinationType", + "ModelServiceConfigDestinationConfigDestinationTypeParam", + "ModelServiceConfigDestinationConfigDict", + "ModelServiceConfigDestinationConfigParam", + "ModelServiceConfigDict", + "ModelServiceConfigExternalModelConfig", + "ModelServiceConfigExternalModelConfigDict", + "ModelServiceConfigExternalModelConfigParam", + "ModelServiceConfigFallbackConfig", + "ModelServiceConfigFallbackConfigDict", + "ModelServiceConfigFallbackConfigParam", + "ModelServiceConfigParam", + "ModelServiceConfigPayPerTokenConfig", + "ModelServiceConfigPayPerTokenConfigDict", + "ModelServiceConfigPayPerTokenConfigParam", + "ModelServiceConfigProvisionedThroughputConfig", + "ModelServiceConfigProvisionedThroughputConfigDict", + "ModelServiceConfigProvisionedThroughputConfigParam", + "ModelServiceConfigRoutingConfig", + "ModelServiceConfigRoutingConfigDict", + "ModelServiceConfigRoutingConfigParam", + "ModelServiceConfigRoutingConfigTrafficSplitting", + "ModelServiceConfigRoutingConfigTrafficSplittingDict", + "ModelServiceConfigRoutingConfigTrafficSplittingParam", + "ModelServiceDict", + "ModelServiceParam", + "RateLimit", + "RateLimitDict", + "RateLimitParam", + "RateLimitRateLimitKey", + "RateLimitRateLimitKeyParam", + "RateLimitRateLimitRenewalPeriod", + "RateLimitRateLimitRenewalPeriodParam", +] + + +from databricks.bundles.model_services._models.inference_table_config import ( + InferenceTableConfig, + InferenceTableConfigDict, + InferenceTableConfigParam, +) +from databricks.bundles.model_services._models.lifecycle import ( + Lifecycle, + LifecycleDict, + LifecycleParam, +) +from databricks.bundles.model_services._models.model_provider_service_config_model_target_config import ( + ModelProviderServiceConfigModelTargetConfig, + ModelProviderServiceConfigModelTargetConfigDict, + ModelProviderServiceConfigModelTargetConfigParam, +) +from databricks.bundles.model_services._models.model_service import ( + ModelService, + ModelServiceDict, + ModelServiceParam, +) +from databricks.bundles.model_services._models.model_service_config import ( + ModelServiceConfig, + ModelServiceConfigDict, + ModelServiceConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_destination_config import ( + ModelServiceConfigDestinationConfig, + ModelServiceConfigDestinationConfigDict, + ModelServiceConfigDestinationConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_destination_config_destination_type import ( + ModelServiceConfigDestinationConfigDestinationType, + ModelServiceConfigDestinationConfigDestinationTypeParam, +) +from databricks.bundles.model_services._models.model_service_config_external_model_config import ( + ModelServiceConfigExternalModelConfig, + ModelServiceConfigExternalModelConfigDict, + ModelServiceConfigExternalModelConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_fallback_config import ( + ModelServiceConfigFallbackConfig, + ModelServiceConfigFallbackConfigDict, + ModelServiceConfigFallbackConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_pay_per_token_config import ( + ModelServiceConfigPayPerTokenConfig, + ModelServiceConfigPayPerTokenConfigDict, + ModelServiceConfigPayPerTokenConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_provisioned_throughput_config import ( + ModelServiceConfigProvisionedThroughputConfig, + ModelServiceConfigProvisionedThroughputConfigDict, + ModelServiceConfigProvisionedThroughputConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_routing_config import ( + ModelServiceConfigRoutingConfig, + ModelServiceConfigRoutingConfigDict, + ModelServiceConfigRoutingConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_routing_config_traffic_splitting import ( + ModelServiceConfigRoutingConfigTrafficSplitting, + ModelServiceConfigRoutingConfigTrafficSplittingDict, + ModelServiceConfigRoutingConfigTrafficSplittingParam, +) +from databricks.bundles.model_services._models.rate_limit import ( + RateLimit, + RateLimitDict, + RateLimitParam, +) +from databricks.bundles.model_services._models.rate_limit_rate_limit_key import ( + RateLimitRateLimitKey, + RateLimitRateLimitKeyParam, +) +from databricks.bundles.model_services._models.rate_limit_rate_limit_renewal_period import ( + RateLimitRateLimitRenewalPeriod, + RateLimitRateLimitRenewalPeriodParam, +) diff --git a/python/databricks/bundles/model_services/_models/inference_table_config.py b/python/databricks/bundles/model_services/_models/inference_table_config.py new file mode 100644 index 00000000000..77e4ff25560 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/inference_table_config.py @@ -0,0 +1,104 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr, VariableOrOptional + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class InferenceTableConfig: + """ + Inference table configuration for payload logging on a model service. + + `parent` is always REQUIRED when the sub-message is set; the destination + UC schema is needed to construct or rebind the payload TABLE regardless of + whether payload logging is currently active. Payload logging is active by + default; set `disabled = true` to pause runtime logging without dropping the + table or the binding. + """ + + parent: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Parent UC schema where the inference table is created. + Format: `schemas/{catalog}.{schema}`. Set at create time and immutable + thereafter; changing it on an existing service is rejected. + """ + + disabled: VariableOrOptional[bool] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Indicates whether payload logging is disabled (opt-out). Unset means that + payload logging is active (the on-by-default state coincides with the proto + zero-value, so the server never fills this field for a client that leaves it + unset). Set `disabled = true` to pause runtime logging while keeping the + sub-message attached (preserving `parent` and `table_name_prefix` for a + later flip back to active). `parent` remains required either way. + """ + + table_name_prefix: VariableOrOptional[str] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Prefix for the inference-table's UC-registered name. The actual leaf name UC + stores is `_payload`; the `_payload` suffix is appended + automatically. To find the actual UC table after Create, read the `table` + field on the response. Defaults to `_payload` when unset. + Set at create time and immutable thereafter; changing it on an existing + service is rejected. + """ + + @classmethod + def from_dict(cls, value: "InferenceTableConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "InferenceTableConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class InferenceTableConfigDict(TypedDict, total=False): + """""" + + parent: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Parent UC schema where the inference table is created. + Format: `schemas/{catalog}.{schema}`. Set at create time and immutable + thereafter; changing it on an existing service is rejected. + """ + + disabled: VariableOrOptional[bool] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Indicates whether payload logging is disabled (opt-out). Unset means that + payload logging is active (the on-by-default state coincides with the proto + zero-value, so the server never fills this field for a client that leaves it + unset). Set `disabled = true` to pause runtime logging while keeping the + sub-message attached (preserving `parent` and `table_name_prefix` for a + later flip back to active). `parent` remains required either way. + """ + + table_name_prefix: VariableOrOptional[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Prefix for the inference-table's UC-registered name. The actual leaf name UC + stores is `_payload`; the `_payload` suffix is appended + automatically. To find the actual UC table after Create, read the `table` + field on the response. Defaults to `_payload` when unset. + Set at create time and immutable thereafter; changing it on an existing + service is rejected. + """ + + +InferenceTableConfigParam = InferenceTableConfigDict | InferenceTableConfig diff --git a/python/databricks/bundles/model_services/_models/lifecycle.py b/python/databricks/bundles/model_services/_models/lifecycle.py new file mode 100644 index 00000000000..697776a198e --- /dev/null +++ b/python/databricks/bundles/model_services/_models/lifecycle.py @@ -0,0 +1,40 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOrOptional + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class Lifecycle: + """""" + + prevent_destroy: VariableOrOptional[bool] = None + """ + Lifecycle setting to prevent the resource from being destroyed. + """ + + @classmethod + def from_dict(cls, value: "LifecycleDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "LifecycleDict": + return _transform_to_json_value(self) # type:ignore + + +class LifecycleDict(TypedDict, total=False): + """""" + + prevent_destroy: VariableOrOptional[bool] + """ + Lifecycle setting to prevent the resource from being destroyed. + """ + + +LifecycleParam = LifecycleDict | Lifecycle diff --git a/python/databricks/bundles/model_services/_models/model_provider_service_config_model_target_config.py b/python/databricks/bundles/model_services/_models/model_provider_service_config_model_target_config.py new file mode 100644 index 00000000000..d2538a33196 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_provider_service_config_model_target_config.py @@ -0,0 +1,79 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr, VariableOrList + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelProviderServiceConfigModelTargetConfig: + """ + Model target configuration for an external model destination. + """ + + model: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Provider-side model identifier (e.g. "gpt-5", "claude-opus-4-7"). This is + a string on the LLM provider's side, not a UC entity. The UC governance + hook for external destinations is the ModelProviderService referenced by + `ExternalModelConfig.model_provider_service`, not the model itself. + """ + + native_api_types: VariableOrList[str] = field(default_factory=list) + """ + :meta private: [EXPERIMENTAL] + + [Beta] Provider-native API types the model supports (e.g. + "openai/v1/chat/completions"). Used by the platform for request/response + translation from the unified API type. At most 64 entries of at most 256 + characters each; the list is persisted into the destination binding's + bounded storage envelope. + """ + + @classmethod + def from_dict( + cls, value: "ModelProviderServiceConfigModelTargetConfigDict" + ) -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelProviderServiceConfigModelTargetConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelProviderServiceConfigModelTargetConfigDict(TypedDict, total=False): + """""" + + model: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Provider-side model identifier (e.g. "gpt-5", "claude-opus-4-7"). This is + a string on the LLM provider's side, not a UC entity. The UC governance + hook for external destinations is the ModelProviderService referenced by + `ExternalModelConfig.model_provider_service`, not the model itself. + """ + + native_api_types: VariableOrList[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Provider-native API types the model supports (e.g. + "openai/v1/chat/completions"). Used by the platform for request/response + translation from the unified API type. At most 64 entries of at most 256 + characters each; the list is persisted into the destination binding's + bounded storage envelope. + """ + + +ModelProviderServiceConfigModelTargetConfigParam = ( + ModelProviderServiceConfigModelTargetConfigDict + | ModelProviderServiceConfigModelTargetConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service.py b/python/databricks/bundles/model_services/_models/model_service.py new file mode 100644 index 00000000000..d5eab192811 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service.py @@ -0,0 +1,59 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._resource import Resource +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr, VariableOrOptional +from databricks.bundles.model_services._models.lifecycle import ( + Lifecycle, + LifecycleParam, +) +from databricks.bundles.model_services._models.model_service_config import ( + ModelServiceConfig, + ModelServiceConfigParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelService(Resource): + """""" + + model_service_id: VariableOr[str] + + parent: VariableOr[str] + + comment: VariableOrOptional[str] = None + + config: VariableOrOptional[ModelServiceConfig] = None + + lifecycle: VariableOrOptional[Lifecycle] = None + + @classmethod + def from_dict(cls, value: "ModelServiceDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceDict(TypedDict, total=False): + """""" + + model_service_id: VariableOr[str] + + parent: VariableOr[str] + + comment: VariableOrOptional[str] + + config: VariableOrOptional[ModelServiceConfigParam] + + lifecycle: VariableOrOptional[LifecycleParam] + + +ModelServiceParam = ModelServiceDict | ModelService diff --git a/python/databricks/bundles/model_services/_models/model_service_config.py b/python/databricks/bundles/model_services/_models/model_service_config.py new file mode 100644 index 00000000000..5726a1509de --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config.py @@ -0,0 +1,86 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOrList, VariableOrOptional +from databricks.bundles.model_services._models.inference_table_config import ( + InferenceTableConfig, + InferenceTableConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_routing_config import ( + ModelServiceConfigRoutingConfig, + ModelServiceConfigRoutingConfigParam, +) +from databricks.bundles.model_services._models.rate_limit import ( + RateLimit, + RateLimitParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfig: + """ + Operational configuration wrapped around the ModelService resource. + """ + + inference_table: VariableOrOptional[InferenceTableConfig] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Inference table config for payload logging. + """ + + rate_limits: VariableOrList[RateLimit] = field(default_factory=list) + """ + :meta private: [EXPERIMENTAL] + + [Beta] Rate limits applied to requests routed through this model service. + """ + + routing: VariableOrOptional[ModelServiceConfigRoutingConfig] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Routing configuration: destinations, routing strategy, and fallback. + """ + + @classmethod + def from_dict(cls, value: "ModelServiceConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigDict(TypedDict, total=False): + """""" + + inference_table: VariableOrOptional[InferenceTableConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Inference table config for payload logging. + """ + + rate_limits: VariableOrList[RateLimitParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Rate limits applied to requests routed through this model service. + """ + + routing: VariableOrOptional[ModelServiceConfigRoutingConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Routing configuration: destinations, routing strategy, and fallback. + """ + + +ModelServiceConfigParam = ModelServiceConfigDict | ModelServiceConfig diff --git a/python/databricks/bundles/model_services/_models/model_service_config_destination_config.py b/python/databricks/bundles/model_services/_models/model_service_config_destination_config.py new file mode 100644 index 00000000000..d0b0138f2e5 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_destination_config.py @@ -0,0 +1,164 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr, VariableOrOptional +from databricks.bundles.model_services._models.model_service_config_destination_config_destination_type import ( + ModelServiceConfigDestinationConfigDestinationType, + ModelServiceConfigDestinationConfigDestinationTypeParam, +) +from databricks.bundles.model_services._models.model_service_config_external_model_config import ( + ModelServiceConfigExternalModelConfig, + ModelServiceConfigExternalModelConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_pay_per_token_config import ( + ModelServiceConfigPayPerTokenConfig, + ModelServiceConfigPayPerTokenConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_provisioned_throughput_config import ( + ModelServiceConfigProvisionedThroughputConfig, + ModelServiceConfigProvisionedThroughputConfigParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigDestinationConfig: + """ + A destination the model service can route traffic to. Exactly one of the + per-type configs inside `type_config` must be set, and it must match + `destination_type`. + """ + + destination_type: VariableOr[ModelServiceConfigDestinationConfigDestinationType] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Backing-model category. Determines which oneof variant is populated. + """ + + name: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] User-facing label for this destination, used in routing references. + """ + + external_model_config: VariableOrOptional[ModelServiceConfigExternalModelConfig] = ( + None + ) + """ + :meta private: [EXPERIMENTAL] + + [Beta] Configuration for an external-foundation-model destination. Provider auth + and provider-specific cloud configuration are owned by a separate, governed + ModelProviderService entity referenced via `model_provider_service`; the + platform resolves the provider at invocation time. + """ + + pay_per_token_config: VariableOrOptional[ModelServiceConfigPayPerTokenConfig] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Configuration for a pay-per-token foundation-model destination. Identifies + the foundation model by its UC resource name; the platform resolves it to a + Model Serving endpoint at request time. + """ + + provisioned_throughput_config: VariableOrOptional[ + ModelServiceConfigProvisionedThroughputConfig + ] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Configuration for a provisioned-throughput foundation-model destination. + References a pre-existing Model Serving endpoint that serves the model; + sizing (provisioned throughput, burst scaling, model version) is owned by + the Model Serving endpoint itself, not by this message. + """ + + traffic_percentage: VariableOrOptional[int] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Share of traffic sent to this destination, 0-100. Optional on fallback + destinations; see FallbackConfig. + """ + + @classmethod + def from_dict(cls, value: "ModelServiceConfigDestinationConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigDestinationConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigDestinationConfigDict(TypedDict, total=False): + """""" + + destination_type: VariableOr[ + ModelServiceConfigDestinationConfigDestinationTypeParam + ] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Backing-model category. Determines which oneof variant is populated. + """ + + name: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] User-facing label for this destination, used in routing references. + """ + + external_model_config: VariableOrOptional[ + ModelServiceConfigExternalModelConfigParam + ] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Configuration for an external-foundation-model destination. Provider auth + and provider-specific cloud configuration are owned by a separate, governed + ModelProviderService entity referenced via `model_provider_service`; the + platform resolves the provider at invocation time. + """ + + pay_per_token_config: VariableOrOptional[ModelServiceConfigPayPerTokenConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Configuration for a pay-per-token foundation-model destination. Identifies + the foundation model by its UC resource name; the platform resolves it to a + Model Serving endpoint at request time. + """ + + provisioned_throughput_config: VariableOrOptional[ + ModelServiceConfigProvisionedThroughputConfigParam + ] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Configuration for a provisioned-throughput foundation-model destination. + References a pre-existing Model Serving endpoint that serves the model; + sizing (provisioned throughput, burst scaling, model version) is owned by + the Model Serving endpoint itself, not by this message. + """ + + traffic_percentage: VariableOrOptional[int] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Share of traffic sent to this destination, 0-100. Optional on fallback + destinations; see FallbackConfig. + """ + + +ModelServiceConfigDestinationConfigParam = ( + ModelServiceConfigDestinationConfigDict | ModelServiceConfigDestinationConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_destination_config_destination_type.py b/python/databricks/bundles/model_services/_models/model_service_config_destination_config_destination_type.py new file mode 100644 index 00000000000..a98c2ca8161 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_destination_config_destination_type.py @@ -0,0 +1,30 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from enum import Enum +from typing import Literal + + +class ModelServiceConfigDestinationConfigDestinationType(Enum): + """ + Backing-model category for a model service destination. + """ + + DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL = ( + "DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL" + ) + DESTINATION_TYPE_PROVISIONED_THROUGHPUT_FOUNDATION_MODEL = ( + "DESTINATION_TYPE_PROVISIONED_THROUGHPUT_FOUNDATION_MODEL" + ) + DESTINATION_TYPE_EXTERNAL_FOUNDATION_MODEL = ( + "DESTINATION_TYPE_EXTERNAL_FOUNDATION_MODEL" + ) + + +ModelServiceConfigDestinationConfigDestinationTypeParam = ( + Literal[ + "DESTINATION_TYPE_PAY_PER_TOKEN_FOUNDATION_MODEL", + "DESTINATION_TYPE_PROVISIONED_THROUGHPUT_FOUNDATION_MODEL", + "DESTINATION_TYPE_EXTERNAL_FOUNDATION_MODEL", + ] + | ModelServiceConfigDestinationConfigDestinationType +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_external_model_config.py b/python/databricks/bundles/model_services/_models/model_service_config_external_model_config.py new file mode 100644 index 00000000000..f277f2829cf --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_external_model_config.py @@ -0,0 +1,83 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr +from databricks.bundles.model_services._models.model_provider_service_config_model_target_config import ( + ModelProviderServiceConfigModelTargetConfig, + ModelProviderServiceConfigModelTargetConfigParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigExternalModelConfig: + """ + Configuration for an external-foundation-model destination. Provider auth + and provider-specific cloud configuration are owned by a separate, governed + ModelProviderService entity referenced via `model_provider_service`; the + platform resolves the provider at invocation time. + """ + + model_provider_service: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Resource name of the governed ModelProviderService that owns provider + auth and provider-specific configuration. The referenced + ModelProviderService also carries the provider type, so this message + does not surface it directly. + Format: `model-provider-services/{catalog}.{schema}.{model_provider_service}`. + Each `{...}` component is capped at 255 characters individually. + """ + + target: VariableOr[ModelProviderServiceConfigModelTargetConfig] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Routing target for the destination: the provider-side model selected from + the referenced ModelProviderService's `targets` catalog, plus the unified + API types the platform should translate to/from at request time. + """ + + @classmethod + def from_dict(cls, value: "ModelServiceConfigExternalModelConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigExternalModelConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigExternalModelConfigDict(TypedDict, total=False): + """""" + + model_provider_service: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Resource name of the governed ModelProviderService that owns provider + auth and provider-specific configuration. The referenced + ModelProviderService also carries the provider type, so this message + does not surface it directly. + Format: `model-provider-services/{catalog}.{schema}.{model_provider_service}`. + Each `{...}` component is capped at 255 characters individually. + """ + + target: VariableOr[ModelProviderServiceConfigModelTargetConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Routing target for the destination: the provider-side model selected from + the referenced ModelProviderService's `targets` catalog, plus the unified + API types the platform should translate to/from at request time. + """ + + +ModelServiceConfigExternalModelConfigParam = ( + ModelServiceConfigExternalModelConfigDict | ModelServiceConfigExternalModelConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_fallback_config.py b/python/databricks/bundles/model_services/_models/model_service_config_fallback_config.py new file mode 100644 index 00000000000..b5a7922f04d --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_fallback_config.py @@ -0,0 +1,58 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOrList +from databricks.bundles.model_services._models.model_service_config_destination_config import ( + ModelServiceConfigDestinationConfig, + ModelServiceConfigDestinationConfigParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigFallbackConfig: + """ + Fallback routing, applied after the primary destination returns a retryable + error. Traversal is in list order; the attempt count is the length of the + list. + """ + + destinations: VariableOrList[ModelServiceConfigDestinationConfig] = field( + default_factory=list + ) + """ + :meta private: [EXPERIMENTAL] + + [Beta] Ordered list of fallback destinations. Traversal is in list order; the + attempt count is the length of the list. At most 5 are allowed. + """ + + @classmethod + def from_dict(cls, value: "ModelServiceConfigFallbackConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigFallbackConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigFallbackConfigDict(TypedDict, total=False): + """""" + + destinations: VariableOrList[ModelServiceConfigDestinationConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Ordered list of fallback destinations. Traversal is in list order; the + attempt count is the length of the list. At most 5 are allowed. + """ + + +ModelServiceConfigFallbackConfigParam = ( + ModelServiceConfigFallbackConfigDict | ModelServiceConfigFallbackConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_pay_per_token_config.py b/python/databricks/bundles/model_services/_models/model_service_config_pay_per_token_config.py new file mode 100644 index 00000000000..988201bbe82 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_pay_per_token_config.py @@ -0,0 +1,52 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigPayPerTokenConfig: + """ + Configuration for a pay-per-token foundation-model destination. Identifies + the foundation model by its UC resource name; the platform resolves it to a + Model Serving endpoint at request time. + """ + + model: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Resource name of the UC model. + Format: `models/{catalog}.{schema}.{model}`. + """ + + @classmethod + def from_dict(cls, value: "ModelServiceConfigPayPerTokenConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigPayPerTokenConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigPayPerTokenConfigDict(TypedDict, total=False): + """""" + + model: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Resource name of the UC model. + Format: `models/{catalog}.{schema}.{model}`. + """ + + +ModelServiceConfigPayPerTokenConfigParam = ( + ModelServiceConfigPayPerTokenConfigDict | ModelServiceConfigPayPerTokenConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_provisioned_throughput_config.py b/python/databricks/bundles/model_services/_models/model_service_config_provisioned_throughput_config.py new file mode 100644 index 00000000000..34c9f4ce321 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_provisioned_throughput_config.py @@ -0,0 +1,64 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigProvisionedThroughputConfig: + """ + Configuration for a provisioned-throughput foundation-model destination. + References a pre-existing Model Serving endpoint that serves the model; + sizing (provisioned throughput, burst scaling, model version) is owned by + the Model Serving endpoint itself, not by this message. + """ + + model_serving_endpoint: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Name of the backing Model Serving endpoint serving the provisioned- + throughput foundation model, as the AIP-122 typed resource name + `serving-endpoints/{name}`. The same UC model can be served on multiple + Model Serving endpoints (different throughput / region / config); the + caller picks which one this destination routes to. The endpoint must + exist at create time. + """ + + @classmethod + def from_dict( + cls, value: "ModelServiceConfigProvisionedThroughputConfigDict" + ) -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigProvisionedThroughputConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigProvisionedThroughputConfigDict(TypedDict, total=False): + """""" + + model_serving_endpoint: VariableOr[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Name of the backing Model Serving endpoint serving the provisioned- + throughput foundation model, as the AIP-122 typed resource name + `serving-endpoints/{name}`. The same UC model can be served on multiple + Model Serving endpoints (different throughput / region / config); the + caller picks which one this destination routes to. The endpoint must + exist at create time. + """ + + +ModelServiceConfigProvisionedThroughputConfigParam = ( + ModelServiceConfigProvisionedThroughputConfigDict + | ModelServiceConfigProvisionedThroughputConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_routing_config.py b/python/databricks/bundles/model_services/_models/model_service_config_routing_config.py new file mode 100644 index 00000000000..893b9d894fd --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_routing_config.py @@ -0,0 +1,125 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOrList, VariableOrOptional +from databricks.bundles.model_services._models.model_service_config_destination_config import ( + ModelServiceConfigDestinationConfig, + ModelServiceConfigDestinationConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_fallback_config import ( + ModelServiceConfigFallbackConfig, + ModelServiceConfigFallbackConfigParam, +) +from databricks.bundles.model_services._models.model_service_config_routing_config_traffic_splitting import ( + ModelServiceConfigRoutingConfigTrafficSplitting, + ModelServiceConfigRoutingConfigTrafficSplittingParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigRoutingConfig: + """ + Routing configuration for a model service, nesting destinations, routing + strategy, and fallback under a single sub-message. + """ + + destinations: VariableOrList[ModelServiceConfigDestinationConfig] = field( + default_factory=list + ) + """ + :meta private: [EXPERIMENTAL] + + [Beta] Primary routing destinations. At most 10 are allowed. At least one is + required on CreateModelService; on UpdateModelService it is required only + when `config.routing` (or a `config.routing.*` subpath) appears in + `update_mask`. + """ + + fallback: VariableOrOptional[ModelServiceConfigFallbackConfig] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Fallback routing config, applied after primary destinations fail. + """ + + first_token_timeout: VariableOrOptional[str] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Timeout for the first token of a streaming response. If a destination does + not return its first token within this duration, AI Gateway aborts the + attempt and fails over to the next destination. Applies to streaming + requests only. Leave unset for no first-token timeout. + """ + + traffic_splitting: VariableOrOptional[ + ModelServiceConfigRoutingConfigTrafficSplitting + ] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Marker message selecting request-based traffic splitting. Traffic is + distributed according to each destination's traffic_percentage value; + no configuration lives on this message itself. + """ + + @classmethod + def from_dict(cls, value: "ModelServiceConfigRoutingConfigDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigRoutingConfigDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigRoutingConfigDict(TypedDict, total=False): + """""" + + destinations: VariableOrList[ModelServiceConfigDestinationConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Primary routing destinations. At most 10 are allowed. At least one is + required on CreateModelService; on UpdateModelService it is required only + when `config.routing` (or a `config.routing.*` subpath) appears in + `update_mask`. + """ + + fallback: VariableOrOptional[ModelServiceConfigFallbackConfigParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Fallback routing config, applied after primary destinations fail. + """ + + first_token_timeout: VariableOrOptional[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Timeout for the first token of a streaming response. If a destination does + not return its first token within this duration, AI Gateway aborts the + attempt and fails over to the next destination. Applies to streaming + requests only. Leave unset for no first-token timeout. + """ + + traffic_splitting: VariableOrOptional[ + ModelServiceConfigRoutingConfigTrafficSplittingParam + ] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Marker message selecting request-based traffic splitting. Traffic is + distributed according to each destination's traffic_percentage value; + no configuration lives on this message itself. + """ + + +ModelServiceConfigRoutingConfigParam = ( + ModelServiceConfigRoutingConfigDict | ModelServiceConfigRoutingConfig +) diff --git a/python/databricks/bundles/model_services/_models/model_service_config_routing_config_traffic_splitting.py b/python/databricks/bundles/model_services/_models/model_service_config_routing_config_traffic_splitting.py new file mode 100644 index 00000000000..3a728660b2d --- /dev/null +++ b/python/databricks/bundles/model_services/_models/model_service_config_routing_config_traffic_splitting.py @@ -0,0 +1,37 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class ModelServiceConfigRoutingConfigTrafficSplitting: + """ + Marker message selecting request-based traffic splitting across primary + destinations. Split weights are read from each DestinationConfig.traffic_percentage. + """ + + @classmethod + def from_dict( + cls, value: "ModelServiceConfigRoutingConfigTrafficSplittingDict" + ) -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "ModelServiceConfigRoutingConfigTrafficSplittingDict": + return _transform_to_json_value(self) # type:ignore + + +class ModelServiceConfigRoutingConfigTrafficSplittingDict(TypedDict, total=False): + """""" + + +ModelServiceConfigRoutingConfigTrafficSplittingParam = ( + ModelServiceConfigRoutingConfigTrafficSplittingDict + | ModelServiceConfigRoutingConfigTrafficSplitting +) diff --git a/python/databricks/bundles/model_services/_models/rate_limit.py b/python/databricks/bundles/model_services/_models/rate_limit.py new file mode 100644 index 00000000000..b1d9715bea0 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/rate_limit.py @@ -0,0 +1,154 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from dataclasses import dataclass +from typing import TYPE_CHECKING, TypedDict + +from databricks.bundles.core._transform import _transform +from databricks.bundles.core._transform_to_json import _transform_to_json_value +from databricks.bundles.core._variable import VariableOr, VariableOrOptional +from databricks.bundles.model_services._models.rate_limit_rate_limit_key import ( + RateLimitRateLimitKey, + RateLimitRateLimitKeyParam, +) +from databricks.bundles.model_services._models.rate_limit_rate_limit_renewal_period import ( + RateLimitRateLimitRenewalPeriod, + RateLimitRateLimitRenewalPeriodParam, +) + +if TYPE_CHECKING: + from typing_extensions import Self + + +@dataclass(kw_only=True) +class RateLimit: + """ + A rate limit applied to service requests. Leave `requests` or `tokens` + unset to impose no limit on that dimension; set a value to cap that dimension + within the renewal period. + """ + + key: VariableOr[RateLimitRateLimitKey] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Scope key. Determines whether `principal` is required. + """ + + renewal_period: VariableOr[RateLimitRateLimitRenewalPeriod] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Renewal period. + """ + + principal: VariableOrOptional[str] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Principal this limit applies to: user email, group name, or service + principal application ID. Required unless `key` is + `RATE_LIMIT_KEY_SERVICE`, `RATE_LIMIT_KEY_USER_DEFAULT`, or + `RATE_LIMIT_KEY_REQUEST_TAG` (which must not set a principal). + """ + + request_tag_key: VariableOrOptional[str] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Request tag key this limit applies to. Required when `key` is + `RATE_LIMIT_KEY_REQUEST_TAG`, forbidden otherwise. + """ + + request_tag_value: VariableOrOptional[str] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Request tag value this limit applies to. Only valid when `key` is + `RATE_LIMIT_KEY_REQUEST_TAG`. Leave unset to apply the limit to every + value of `request_tag_key` (an any-value default); a set value is a + specific override for that value. + """ + + requests: VariableOrOptional[int] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Max requests allowed within a renewal period. Leave unset for no request limit. + """ + + tokens: VariableOrOptional[int] = None + """ + :meta private: [EXPERIMENTAL] + + [Beta] Max tokens allowed within a renewal period. Leave unset for no token limit. + """ + + @classmethod + def from_dict(cls, value: "RateLimitDict") -> "Self": + return _transform(cls, value) + + def as_dict(self) -> "RateLimitDict": + return _transform_to_json_value(self) # type:ignore + + +class RateLimitDict(TypedDict, total=False): + """""" + + key: VariableOr[RateLimitRateLimitKeyParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Scope key. Determines whether `principal` is required. + """ + + renewal_period: VariableOr[RateLimitRateLimitRenewalPeriodParam] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Renewal period. + """ + + principal: VariableOrOptional[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Principal this limit applies to: user email, group name, or service + principal application ID. Required unless `key` is + `RATE_LIMIT_KEY_SERVICE`, `RATE_LIMIT_KEY_USER_DEFAULT`, or + `RATE_LIMIT_KEY_REQUEST_TAG` (which must not set a principal). + """ + + request_tag_key: VariableOrOptional[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Request tag key this limit applies to. Required when `key` is + `RATE_LIMIT_KEY_REQUEST_TAG`, forbidden otherwise. + """ + + request_tag_value: VariableOrOptional[str] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Request tag value this limit applies to. Only valid when `key` is + `RATE_LIMIT_KEY_REQUEST_TAG`. Leave unset to apply the limit to every + value of `request_tag_key` (an any-value default); a set value is a + specific override for that value. + """ + + requests: VariableOrOptional[int] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Max requests allowed within a renewal period. Leave unset for no request limit. + """ + + tokens: VariableOrOptional[int] + """ + :meta private: [EXPERIMENTAL] + + [Beta] Max tokens allowed within a renewal period. Leave unset for no token limit. + """ + + +RateLimitParam = RateLimitDict | RateLimit diff --git a/python/databricks/bundles/model_services/_models/rate_limit_rate_limit_key.py b/python/databricks/bundles/model_services/_models/rate_limit_rate_limit_key.py new file mode 100644 index 00000000000..0e4c11862a8 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/rate_limit_rate_limit_key.py @@ -0,0 +1,30 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from enum import Enum +from typing import Literal + + +class RateLimitRateLimitKey(Enum): + """ + Scope key for a rate limit. + """ + + RATE_LIMIT_KEY_USER = "RATE_LIMIT_KEY_USER" + RATE_LIMIT_KEY_USER_GROUP = "RATE_LIMIT_KEY_USER_GROUP" + RATE_LIMIT_KEY_SERVICE_PRINCIPAL = "RATE_LIMIT_KEY_SERVICE_PRINCIPAL" + RATE_LIMIT_KEY_SERVICE = "RATE_LIMIT_KEY_SERVICE" + RATE_LIMIT_KEY_USER_DEFAULT = "RATE_LIMIT_KEY_USER_DEFAULT" + RATE_LIMIT_KEY_REQUEST_TAG = "RATE_LIMIT_KEY_REQUEST_TAG" + + +RateLimitRateLimitKeyParam = ( + Literal[ + "RATE_LIMIT_KEY_USER", + "RATE_LIMIT_KEY_USER_GROUP", + "RATE_LIMIT_KEY_SERVICE_PRINCIPAL", + "RATE_LIMIT_KEY_SERVICE", + "RATE_LIMIT_KEY_USER_DEFAULT", + "RATE_LIMIT_KEY_REQUEST_TAG", + ] + | RateLimitRateLimitKey +) diff --git a/python/databricks/bundles/model_services/_models/rate_limit_rate_limit_renewal_period.py b/python/databricks/bundles/model_services/_models/rate_limit_rate_limit_renewal_period.py new file mode 100644 index 00000000000..58636c784c0 --- /dev/null +++ b/python/databricks/bundles/model_services/_models/rate_limit_rate_limit_renewal_period.py @@ -0,0 +1,19 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from enum import Enum +from typing import Literal + + +class RateLimitRateLimitRenewalPeriod(Enum): + """ + Renewal period for a rate limit. + """ + + RATE_LIMIT_RENEWAL_PERIOD_MINUTE = "RATE_LIMIT_RENEWAL_PERIOD_MINUTE" + RATE_LIMIT_RENEWAL_PERIOD_HOUR = "RATE_LIMIT_RENEWAL_PERIOD_HOUR" + + +RateLimitRateLimitRenewalPeriodParam = ( + Literal["RATE_LIMIT_RENEWAL_PERIOD_MINUTE", "RATE_LIMIT_RENEWAL_PERIOD_HOUR"] + | RateLimitRateLimitRenewalPeriod +) diff --git a/python/databricks_tests/core/_generated/__init__.py b/python/databricks_tests/core/_generated/__init__.py index fca10520617..94627b5b0a0 100644 --- a/python/databricks_tests/core/_generated/__init__.py +++ b/python/databricks_tests/core/_generated/__init__.py @@ -12,6 +12,7 @@ instance_pools, job_runs, jobs, + model_services, model_serving_endpoints, models, pipelines, @@ -42,6 +43,7 @@ instance_pools._test_case(), job_runs._test_case(), jobs._test_case(), + model_services._test_case(), model_serving_endpoints._test_case(), models._test_case(), pipelines._test_case(), diff --git a/python/databricks_tests/core/_generated/model_services.py b/python/databricks_tests/core/_generated/model_services.py new file mode 100644 index 00000000000..42ae3099700 --- /dev/null +++ b/python/databricks_tests/core/_generated/model_services.py @@ -0,0 +1,32 @@ +# Code generated by pydabs-codegen. DO NOT EDIT. + +from databricks.bundles.core import Resources, model_service_mutator +from databricks.bundles.core._generated.model_services import _resource_type +from databricks.bundles.model_services._models.lifecycle import Lifecycle +from databricks.bundles.model_services._models.model_service import ModelService +from databricks.bundles.model_services._models.model_service_config import ( + ModelServiceConfig, +) +from databricks_tests.core._resource_test_case import ResourceTestCase + + +def _test_case(): + return ( + ResourceTestCase( + add_resource=Resources.add_model_service, + dict_example={ + "config": {}, + "lifecycle": {}, + "model_service_id": "model_service_id", + "parent": "parent", + }, + dataclass_example=ModelService( + config=ModelServiceConfig(), + lifecycle=Lifecycle(), + model_service_id="model_service_id", + parent="parent", + ), + mutator=model_service_mutator, + ), + _resource_type(), + ) diff --git a/python/databricks_tests/core/public_api.txt b/python/databricks_tests/core/public_api.txt index 05b42854dfa..52a32c56ec8 100644 --- a/python/databricks_tests/core/public_api.txt +++ b/python/databricks_tests/core/public_api.txt @@ -29,6 +29,7 @@ __all__ = [ load_resources_from_package_module, mlflow_experiment_mutator, mlflow_model_mutator, + model_service_mutator, model_serving_endpoint_mutator, pipeline_mutator, postgres_snapshot_schedule_mutator, @@ -100,6 +101,7 @@ class Resources: def add_location(self, path: tuple[str, ...], location: Location) -> None def add_mlflow_experiment(self, resource_name: str, mlflow_experiment: MlflowExperimentParam, *, location: Union[Location, None] = None) -> None def add_mlflow_model(self, resource_name: str, mlflow_model: MlflowModelParam, *, location: Union[Location, None] = None) -> None + def add_model_service(self, resource_name: str, model_service: ModelServiceParam, *, location: Union[Location, None] = None) -> None def add_model_serving_endpoint(self, resource_name: str, model_serving_endpoint: ModelServingEndpointParam, *, location: Union[Location, None] = None) -> None def add_pipeline(self, resource_name: str, pipeline: PipelineParam, *, location: Union[Location, None] = None) -> None def add_postgres_snapshot_schedule(self, resource_name: str, postgres_snapshot_schedule: PostgresSnapshotScheduleParam, *, location: Union[Location, None] = None) -> None @@ -127,6 +129,7 @@ class Resources: @property instance_pools -> dict[str, InstancePool] @property job_runs -> dict[str, JobRun] @property jobs -> dict[str, Job] + @property model_services -> dict[str, ModelService] @property model_serving_endpoints -> dict[str, ModelServingEndpoint] @property models -> dict[str, MlflowModel] @property pipelines -> dict[str, Pipeline] @@ -215,6 +218,10 @@ def mlflow_experiment_mutator(function: Callable) -> ResourceMutator[MlflowExper @overload def mlflow_model_mutator(function: Callable[[MlflowModel], MlflowModel]) -> ResourceMutator[MlflowModel] def mlflow_model_mutator(function: Callable) -> ResourceMutator[MlflowModel] +@overload def model_service_mutator(function: Callable[[Bundle, ModelService], ModelService]) -> ResourceMutator[ModelService] +@overload def model_service_mutator(function: Callable[[ModelService], ModelService]) -> ResourceMutator[ModelService] +def model_service_mutator(function: Callable) -> ResourceMutator[ModelService] + @overload def model_serving_endpoint_mutator(function: Callable[[Bundle, ModelServingEndpoint], ModelServingEndpoint]) -> ResourceMutator[ModelServingEndpoint] @overload def model_serving_endpoint_mutator(function: Callable[[ModelServingEndpoint], ModelServingEndpoint]) -> ResourceMutator[ModelServingEndpoint] def model_serving_endpoint_mutator(function: Callable) -> ResourceMutator[ModelServingEndpoint] @@ -282,6 +289,7 @@ singular_name=job plural_name=jobs resource_type=Job singular_name=job_run plural_name=job_runs resource_type=JobRun singular_name=mlflow_experiment plural_name=experiments resource_type=MlflowExperiment singular_name=mlflow_model plural_name=models resource_type=MlflowModel +singular_name=model_service plural_name=model_services resource_type=ModelService singular_name=model_serving_endpoint plural_name=model_serving_endpoints resource_type=ModelServingEndpoint singular_name=pipeline plural_name=pipelines resource_type=Pipeline singular_name=postgres_snapshot_schedule plural_name=postgres_snapshot_schedules resource_type=PostgresSnapshotSchedule