From 7a28f576efc8762d7de07eb1d8d157671a1c728d Mon Sep 17 00:00:00 2001
From: Alex Ott <alexey.ott@databricks.com>
Date: Fri, 2 Aug 2024 14:44:59 +0200
Subject: [PATCH] [Exporter] Don't export model serving endpoints with
 foundational models (#3845)

## Changes
<!-- Summary of your changes that are easy to understand -->

We need to skip model serving endpoints that serve foundational models
because they are provided by Databricks and we don't have ability to
create serving endpoint for foundational models.

## Tests
<!--
How is this tested? Please see the checklist below and also describe any
other relevant tests
-->

- [x] `make test` run locally
- [ ] relevant change in `docs/` folder
- [ ] covered with integration tests in `internal/acceptance`
- [ ] relevant acceptance tests are passing
- [ ] using Go SDK
---
 exporter/exporter_test.go |  5 +++--
 exporter/importables.go   | 18 ++++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/exporter/exporter_test.go b/exporter/exporter_test.go
index 8cc973c4e1..249c65fa50 100644
--- a/exporter/exporter_test.go
+++ b/exporter/exporter_test.go
@@ -2303,8 +2303,9 @@ func TestImportingModelServing(t *testing.T) {
 				},
 			},
 			{
-				Method:   "GET",
-				Resource: "/api/2.0/serving-endpoints/abc?",
+				Method:       "GET",
+				Resource:     "/api/2.0/serving-endpoints/abc?",
+				ReuseRequest: true,
 				Response: serving.ServingEndpointDetailed{
 					Name: "abc",
 					Id:   "1234",
diff --git a/exporter/importables.go b/exporter/importables.go
index b154541bf4..f6f0f217f7 100644
--- a/exporter/importables.go
+++ b/exporter/importables.go
@@ -2156,6 +2156,12 @@ var resourcesMap map[string]importable = map[string]importable{
 				return err
 			}
 			for offset, endpoint := range endpointsList {
+				if endpoint.Config != nil && endpoint.Config.ServedEntities != nil && len(endpoint.Config.ServedEntities) > 0 {
+					if endpoint.Config.ServedEntities[0].FoundationModel != nil {
+						log.Printf("[INFO] skipping endpoint %s that is foundation model", endpoint.Name)
+						continue
+					}
+				}
 				ic.EmitIfUpdatedAfterMillis(&resource{
 					Resource: "databricks_model_serving",
 					ID:       endpoint.Name,
@@ -2228,6 +2234,18 @@ var resourcesMap map[string]importable = map[string]importable{
 			}
 			return nil
 		},
+		Ignore: func(ic *importContext, r *resource) bool {
+			// We need to ignore endpoints that are foundation models
+			endpoint, err := ic.workspaceClient.ServingEndpoints.GetByName(ic.Context, r.ID)
+			if err != nil {
+				log.Printf("[WARN] Can't get endpoint by name %s: %s", r.ID, err)
+				return false
+			}
+			res := (endpoint.Config != nil && endpoint.Config.ServedEntities != nil &&
+				len(endpoint.Config.ServedEntities) > 0 && endpoint.Config.ServedEntities[0].FoundationModel != nil)
+			log.Printf("[DEBUG] Ignore serving endpoint %s? %t", r.ID, res)
+			return res
+		},
 		ShouldOmitField: func(ic *importContext, pathString string, as *schema.Schema, d *schema.ResourceData) bool {
 			if pathString == "config.0.traffic_config" || pathString == "config.0.auto_capture_config.0.enabled" ||
 				(pathString == "config.0.auto_capture_config.0.table_name_prefix" && d.Get(pathString).(string) != "") {