yujunjun/OpenMetadata
1
0
Fork 0
mirror of https://github.com/open-metadata/OpenMetadata.git synced 2025-08-04 23:28:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

[issue-1284] - Add MlStore and Server to MlModel (#1439)

Browse Source 
* Add mlStore and server to MlModel entity

* Add mlStore and server to Create MlModel entity

* Add MlStore and Server to MlModel properties
This commit is contained in:
Pere Miquel Brull 2021-11-29 19:08:00 +01:00 committed by GitHub
parent a1e9f986fb
commit c2efcb1107
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 105 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
catalog-rest-service/src/main/java/org/openmetadata/catalog/jdbi3/MlModelRepository.java
View File

@ -43,7 +43,7 @@ import java.util.UUID;
public class MlModelRepository extends EntityRepository<MlModel> {
private static final Logger LOG = LoggerFactory.getLogger(MlModelRepository.class);
private static final Fields MODEL_UPDATE_FIELDS = new Fields(MlModelResource.FIELD_LIST,
"owner,algorithm,dashboard,mlHyperParameters,mlFeatures,tags");
"owner,algorithm,dashboard,mlHyperParameters,mlFeatures,mlStore,server,tags");
private static final Fields MODEL_PATCH_FIELDS = new Fields(MlModelResource.FIELD_LIST,
"owner,algorithm,dashboard,mlHyperParameters,mlFeatures,tags");
private final CollectionDAO dao;
@ -77,6 +77,8 @@ public class MlModelRepository extends EntityRepository<MlModel> {
mlModel.setDashboard(fields.contains("dashboard") ? getDashboard(mlModel) : null);
mlModel.setMlFeatures(fields.contains("mlFeatures") ? mlModel.getMlFeatures(): null);
mlModel.setMlHyperParameters(fields.contains("mlHyperParameters") ? mlModel.getMlHyperParameters(): null);
mlModel.setMlStore(fields.contains("mlStore") ? mlModel.getMlStore(): null);
mlModel.setServer(fields.contains("server") ? mlModel.getServer(): null);
mlModel.setFollowers(fields.contains("followers") ? getFollowers(mlModel) : null);
mlModel.setTags(fields.contains("tags") ? getTags(mlModel.getFullyQualifiedName()) : null);
mlModel.setUsageSummary(fields.contains("usageSummary") ? EntityUtil.getLatestUsage(dao.usageDAO(),
@ -338,6 +340,8 @@ public class MlModelRepository extends EntityRepository<MlModel> {
updateDashboard(origMlModel, updatedMlModel);
updateMlFeatures(origMlModel, updatedMlModel);
updateMlHyperParameters(origMlModel, updatedMlModel);
updateMlStore(origMlModel, updatedMlModel);
updateServer(origMlModel, updatedMlModel);
}
private void updateAlgorithm(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
@ -352,6 +356,14 @@ public class MlModelRepository extends EntityRepository<MlModel> {
recordChange("mlHyperParameters", origModel.getMlHyperParameters(), updatedModel.getMlHyperParameters());
}
private void updateMlStore(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
recordChange("mlStore", origModel.getMlStore(), updatedModel.getMlStore(), true);
}
private void updateServer(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
recordChange("server", origModel.getServer(), updatedModel.getServer());
}
private void updateDashboard(MlModel origModel, MlModel updatedModel) throws JsonProcessingException {
String modelId = updatedModel.getId().toString();

5
catalog-rest-service/src/main/java/org/openmetadata/catalog/resources/mlmodels/MlModelResource.java
View File

@ -110,7 +110,8 @@ public class MlModelResource {
}
}
static final String FIELDS = "owner,dashboard,algorithm,mlFeatures,mlHyperParameters,followers,tags,usageSummary";
static final String FIELDS = "owner,dashboard,algorithm,mlFeatures,mlHyperParameters,mlStore,server," +
"followers,tags,usageSummary";
public static final List<String> FIELD_LIST = Arrays.asList(FIELDS.replaceAll(" ", "")
.split(","));
@ -357,6 +358,8 @@ public class MlModelResource {
.withAlgorithm(create.getAlgorithm())
.withMlFeatures(create.getMlFeatures())
.withMlHyperParameters(create.getMlHyperParameters())
.withMlStore(create.getMlStore())
.withServer(create.getServer())
.withTags(create.getTags())
.withOwner(create.getOwner())
.withUpdatedBy(securityContext.getUserPrincipal().getName())

8
catalog-rest-service/src/main/resources/json/schema/api/data/createMlModel.json
View File

@ -43,6 +43,14 @@
"description": "Performance Dashboard URL to track metric evolution",
"$ref" : "../../type/entityReference.json"
},
"mlStore": {
"description": "Location containing the ML Model. It can be a storage layer and/or a container repository.",
"$ref": "../../entity/data/mlmodel.json#/definitions/mlStore"
},
"server": {
"description": "Endpoint that makes the ML Model available, e.g,. a REST API serving the data or computing predictions.",
"$ref": "../../type/basic.json#/definitions/href"
},
"tags": {
"description": "Tags for this ML Model",
"type": "array",

24
catalog-rest-service/src/main/resources/json/schema/entity/data/mlmodel.json
View File

@ -155,6 +155,22 @@
"type": "string"
}
}
},
"mlStore": {
"type": "object",
"javaType": "org.openmetadata.catalog.type.MlStore",
"description": "Location containing the ML Model. It can be a storage layer and/or a container repository.",
"additionalProperties": false,
"properties": {
"storage": {
"description": "Storage Layer containing the ML Model data.",
"$ref": "../../type/basic.json#/definitions/href"
},
"imageRepository": {
"description": "Container Repository with the ML Model image.",
"$ref": "../../type/basic.json#/definitions/href"
}
}
}
},
"properties" : {
@ -206,6 +222,14 @@
"description": "Performance Dashboard URL to track metric evolution.",
"$ref" : "../../type/entityReference.json"
},
"mlStore": {
"description": "Location containing the ML Model. It can be a storage layer and/or a container repository.",
"$ref": "#/definitions/mlStore"
},
"server": {
"description": "Endpoint that makes the ML Model available, e.g,. a REST API serving the data or computing predictions.",
"$ref": "../../type/basic.json#/definitions/href"
},
"href": {
"description": "Link to the resource corresponding to this entity.",
"$ref": "../../type/basic.json#/definitions/href"

61
catalog-rest-service/src/test/java/org/openmetadata/catalog/resources/mlmodels/MlModelResourceTest.java
View File

@ -36,6 +36,7 @@ import org.openmetadata.catalog.type.FeatureSourceDataType;
import org.openmetadata.catalog.type.FieldChange;
import org.openmetadata.catalog.type.MlFeature;
import org.openmetadata.catalog.type.MlFeatureDataType;
import org.openmetadata.catalog.type.MlStore;
import org.openmetadata.catalog.type.MlFeatureSource;
import org.openmetadata.catalog.type.MlHyperParameter;
import org.openmetadata.catalog.entity.services.DashboardService;
@ -52,6 +53,7 @@ import org.openmetadata.catalog.util.JsonUtils;
import javax.ws.rs.client.WebTarget;
import javax.ws.rs.core.Response.Status;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.Arrays;
import java.util.Collections;
@ -83,6 +85,12 @@ public class MlModelResourceTest extends EntityResourceTest<MlModel> {
public static String ALGORITHM = "regression";
public static Dashboard DASHBOARD;
public static EntityReference DASHBOARD_REFERENCE;
public static URI SERVER = URI.create("http://localhost.com/mlModel");
public static MlStore ML_STORE = new MlStore()
.withStorage(URI.create("s3://my-bucket.com/mlModel"))
.withImageRepository(URI.create("https://12345.dkr.ecr.region.amazonaws.com"));
public static List<MlFeature> ML_FEATURES = Arrays.asList(
new MlFeature()
.withName("age")
@ -196,6 +204,18 @@ public class MlModelResourceTest extends EntityResourceTest<MlModel> {
createAndCheckEntity(create, adminAuthHeaders());
}
@Test
public void post_MlModelWitMlStore_200_ok(TestInfo test) throws IOException {
CreateMlModel create = create(test).withMlStore(ML_STORE);
createAndCheckEntity(create, adminAuthHeaders());
}
@Test
public void post_MlModelWitServer_200_ok(TestInfo test) throws IOException {
CreateMlModel create = create(test).withServer(SERVER);
createAndCheckEntity(create, adminAuthHeaders());
}
@Test
public void post_MlModel_as_non_admin_401(TestInfo test) {
CreateMlModel create = create(test);
@ -257,6 +277,30 @@ public class MlModelResourceTest extends EntityResourceTest<MlModel> {
);
}
@Test
public void put_MlModelAddServer_200(TestInfo test) throws IOException {
CreateMlModel request = create(test);
MlModel model = createAndCheckEntity(request, adminAuthHeaders());
ChangeDescription change = getChangeDescription(model.getVersion());
change.getFieldsAdded().add(new FieldChange().withName("server").withNewValue(SERVER));
updateAndCheckEntity(
request.withServer(SERVER), Status.OK, adminAuthHeaders(), MINOR_UPDATE, change
);
}
@Test
public void put_MlModelAddMlStore_200(TestInfo test) throws IOException {
CreateMlModel request = create(test);
MlModel model = createAndCheckEntity(request, adminAuthHeaders());
ChangeDescription change = getChangeDescription(model.getVersion());
change.getFieldsAdded().add(new FieldChange().withName("mlStore").withNewValue(ML_STORE));
updateAndCheckEntity(
request.withMlStore(ML_STORE), Status.OK, adminAuthHeaders(), MINOR_UPDATE, change
);
}
@Test
public void get_nonExistentMlModel_404_notFound() {
HttpResponseException exception = assertThrows(HttpResponseException.class, () ->
@ -469,23 +513,30 @@ public class MlModelResourceTest extends EntityResourceTest<MlModel> {
if (expected == actual) {
return;
}
if (fieldName.contains("mlFeatures") && !fieldName.endsWith("tags") && !fieldName.endsWith("description")) {
if (fieldName.contains("mlFeatures")) {
List<MlFeature> expectedFeatures = (List<MlFeature>) expected;
List<MlFeature> actualFeatures = JsonUtils.readObjects(actual.toString(), MlFeature.class);
assertEquals(expectedFeatures, actualFeatures);
} else if (fieldName.contains("mlHyperParameters") && !fieldName.endsWith("tags")
&& !fieldName.endsWith("description")) {
} else if (fieldName.contains("mlHyperParameters")) {
List<MlHyperParameter> expectedConstraints = (List<MlHyperParameter>) expected;
List<MlHyperParameter> actualConstraints = JsonUtils.readObjects(actual.toString(), MlHyperParameter.class);
assertEquals(expectedConstraints, actualConstraints);
} else if (fieldName.endsWith("algorithm")) {
} else if (fieldName.contains("algorithm")) {
String expectedAlgorithm = (String) expected;
String actualAlgorithm = actual.toString();
assertEquals(expectedAlgorithm, actualAlgorithm);
} else if (fieldName.endsWith("dashboard")) {
} else if (fieldName.contains("dashboard")) {
EntityReference expectedDashboard = (EntityReference) expected;
EntityReference actualDashboard = JsonUtils.readValue(actual.toString(), EntityReference.class);
assertEquals(expectedDashboard, actualDashboard);
} else if (fieldName.contains("server")) {
URI expectedServer = (URI) expected;
URI actualServer = URI.create(actual.toString());
assertEquals(expectedServer, actualServer);
} else if (fieldName.contains("mlStore")) {
MlStore expectedMlStore = (MlStore) expected;
MlStore actualMlStore = JsonUtils.readValue(actual.toString(), MlStore.class);
assertEquals(expectedMlStore, actualMlStore);
} else {
assertCommonFieldChange(fieldName, expected, actual);
}