add an example to create serving endpoint

databricks · Sep 28, 2023 · fc5a7ea · fc5a7ea
1 parent 3492167
commit fc5a7ea
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 0 deletions.
diff --git a/examples/serving/create_serving_endpoint.py b/examples/serving/create_serving_endpoint.py
@@ -0,0 +1,25 @@
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.serving import (EndpointCoreConfigInput, Route, ServedModelInput, TrafficConfig)
+
+w = WorkspaceClient()
+
+served_model_1 = ServedModelInput(model_name="model_name",
+                                  model_version="model_version",
+                                  workload_size="Small",
+                                  scale_to_zero_enabled=True)
+served_model_2 = ServedModelInput(model_name="model_name",
+                                  model_version="model_version",
+                                  workload_size="Small",
+                                  scale_to_zero_enabled=True)
+
+route_1 = Route(served_model_name=f"{served_model_1.model_name}-{served_model_1.model_version}",
+                traffic_percentage=50)
+route_2 = Route(served_model_name=f"{served_model_2.model_name}-{served_model_2.model_version}",
+                traffic_percentage=50)
+
+traffic_config = TrafficConfig(routes=[route_1, route_2])
+
+endpoint_config = EndpointCoreConfigInput(served_models=[served_model_1, served_model_2],
+                                          traffic_config=traffic_config)
+
+w.serving_endpoints.create_and_wait(name="endpoint-name", config=endpoint_config)
diff --git a/examples/serving/create_serving_endpoint_with_config_dict.py b/examples/serving/create_serving_endpoint_with_config_dict.py
@@ -0,0 +1,30 @@
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.serving import EndpointCoreConfigInput
+
+w = WorkspaceClient()
+
+endpoint_config_dict = {
+    "served_models": [
+        {
+            "model_name": "<model_name>",
+            "model_version": "<model_version>",
+            "scale_to_zero_enabled": True,
+            "workload_size": "Small",
+        },
+        {
+            "model_name": "<model_name>",
+            "model_version": "<model_version>",
+            "scale_to_zero_enabled": True,
+            "workload_size": "Small",
+        },
+    ],
+    "traffic_config": {
+        "routes": [
+            {"served_model_name": "<model_name>-<model_version>", "traffic_percentage": 50},
+            {"served_model_name": "<model_name>-<model_version>", "traffic_percentage": 50},
+        ]
+    },
+}
+
+endpoint_config = EndpointCoreConfigInput.from_dict(endpoint_config_dict)
+w.serving_endpoints.create_and_wait(name="endpoint-name", config=endpoint_config)