Skip to content

Commit

Permalink
feat(ml): bringing ml screens up to date w/ the modern ui layout & im…
Browse files Browse the repository at this point in the history 
…proving ml lineage (datahub-project#4651)

* backend ml changes

* updating ml model UI

* more work on the UI

* ml primary key joining the party

* more progress on UI

* making progress on lineage

* finalizing UI experience

* remove irrelevant test

* fixing lint

* fixups

* add tests and fix what the issues they discovered

* internal > core
  • Loading branch information
@gabe-lyons @maggiehays
gabe-lyons authored and maggiehays committed Aug 1, 2022
1 parent 1095c22 commit 3029c23
Show file tree
Hide file tree
Showing 57 changed files with 1,729 additions and 681 deletions.
51 changes: 51 additions & 0 deletions datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1175,6 +1175,13 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
new LoadableTypeResolver<>(dataPlatformType,
(env) -> ((MLFeatureTable) env.getSource()).getPlatform().getUrn()))
)
.dataFetcher("domain",
new LoadableTypeResolver<>(
domainType,
(env) -> {
final MLFeatureTable entity = env.getSource();
return entity.getDomain() != null ? entity.getDomain().getUrn() : null;
}))
)
.type("MLFeatureTableProperties", typeWiring -> typeWiring
.dataFetcher("mlFeatures", new AuthenticatedResolver<>(
Expand Down Expand Up @@ -1217,6 +1224,13 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
new LoadableTypeResolver<>(dataPlatformType,
(env) -> ((MLModel) env.getSource()).getPlatform().getUrn()))
)
.dataFetcher("domain",
new LoadableTypeResolver<>(
domainType,
(env) -> {
final MLModel mlModel = env.getSource();
return mlModel.getDomain() != null ? mlModel.getDomain().getUrn() : null;
}))
)
.type("MLModelProperties", typeWiring -> typeWiring
.dataFetcher("groups", new AuthenticatedResolver<>(
Expand All @@ -1243,6 +1257,43 @@ private void configureMLFeatureTableResolvers(final RuntimeWiring.Builder builde
new LoadableTypeResolver<>(dataPlatformType,
(env) -> ((MLModelGroup) env.getSource()).getPlatform().getUrn()))
)
.dataFetcher("domain",
new LoadableTypeResolver<>(
domainType,
(env) -> {
final MLModelGroup entity = env.getSource();
return entity.getDomain() != null ? entity.getDomain().getUrn() : null;
}))
)
.type("MLFeature", typeWiring -> typeWiring
.dataFetcher("relationships", new AuthenticatedResolver<>(
new EntityRelationshipsResultResolver(graphClient)
))
.dataFetcher("lineage", new AuthenticatedResolver<>(
new EntityLineageResultResolver(graphClient)
))
.dataFetcher("domain",
new LoadableTypeResolver<>(
domainType,
(env) -> {
final MLFeature entity = env.getSource();
return entity.getDomain() != null ? entity.getDomain().getUrn() : null;
}))
)
.type("MLPrimaryKey", typeWiring -> typeWiring
.dataFetcher("relationships", new AuthenticatedResolver<>(
new EntityRelationshipsResultResolver(graphClient)
))
.dataFetcher("lineage", new AuthenticatedResolver<>(
new EntityLineageResultResolver(graphClient)
))
.dataFetcher("domain",
new LoadableTypeResolver<>(
domainType,
(env) -> {
final MLPrimaryKey entity = env.getSource();
return entity.getDomain() != null ? entity.getDomain().getUrn() : null;
}))
);
}

Expand Down
59 changes: 59 additions & 0 deletions ...ql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/DescriptionUtils.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,11 @@
import com.linkedin.metadata.Constants;
import com.linkedin.metadata.authorization.PoliciesConfig;
import com.linkedin.metadata.entity.EntityService;
import com.linkedin.ml.metadata.EditableMLFeatureProperties;
import com.linkedin.ml.metadata.EditableMLFeatureTableProperties;
import com.linkedin.ml.metadata.EditableMLModelGroupProperties;
import com.linkedin.ml.metadata.EditableMLModelProperties;
import com.linkedin.ml.metadata.EditableMLPrimaryKeyProperties;
import com.linkedin.notebook.EditableNotebookProperties;
import com.linkedin.schema.EditableSchemaFieldInfo;
import com.linkedin.schema.EditableSchemaMetadata;
Expand Down Expand Up @@ -261,4 +266,58 @@ public static boolean isAuthorizedToUpdateDescription(@Nonnull QueryContext cont
targetUrn.toString(),
orPrivilegeGroups);
}

public static void updateMlModelDescription(
String newDescription,
Urn resourceUrn,
Urn actor,
EntityService entityService) {
EditableMLModelProperties editableProperties = (EditableMLModelProperties) getAspectFromEntity(
resourceUrn.toString(), Constants.ML_MODEL_EDITABLE_PROPERTIES_ASPECT_NAME, entityService, new EditableMLModelProperties());
editableProperties.setDescription(newDescription);
persistAspect(resourceUrn, Constants.ML_MODEL_EDITABLE_PROPERTIES_ASPECT_NAME, editableProperties, actor, entityService);
}

public static void updateMlModelGroupDescription(
String newDescription,
Urn resourceUrn,
Urn actor,
EntityService entityService) {
EditableMLModelGroupProperties editableProperties = (EditableMLModelGroupProperties) getAspectFromEntity(
resourceUrn.toString(), Constants.ML_MODEL_GROUP_EDITABLE_PROPERTIES_ASPECT_NAME, entityService, new EditableMLModelGroupProperties());
editableProperties.setDescription(newDescription);
persistAspect(resourceUrn, Constants.ML_MODEL_GROUP_EDITABLE_PROPERTIES_ASPECT_NAME, editableProperties, actor, entityService);
}
public static void updateMlFeatureDescription(
String newDescription,
Urn resourceUrn,
Urn actor,
EntityService entityService) {
EditableMLFeatureProperties editableProperties = (EditableMLFeatureProperties) getAspectFromEntity(
resourceUrn.toString(), Constants.ML_FEATURE_EDITABLE_PROPERTIES_ASPECT_NAME, entityService, new EditableMLFeatureProperties());
editableProperties.setDescription(newDescription);
persistAspect(resourceUrn, Constants.ML_FEATURE_EDITABLE_PROPERTIES_ASPECT_NAME, editableProperties, actor, entityService);
}

public static void updateMlFeatureTableDescription(
String newDescription,
Urn resourceUrn,
Urn actor,
EntityService entityService) {
EditableMLFeatureTableProperties editableProperties = (EditableMLFeatureTableProperties) getAspectFromEntity(
resourceUrn.toString(), Constants.ML_FEATURE_TABLE_EDITABLE_PROPERTIES_ASPECT_NAME, entityService, new EditableMLFeatureTableProperties());
editableProperties.setDescription(newDescription);
persistAspect(resourceUrn, Constants.ML_FEATURE_TABLE_EDITABLE_PROPERTIES_ASPECT_NAME, editableProperties, actor, entityService);
}

public static void updateMlPrimaryKeyDescription(
String newDescription,
Urn resourceUrn,
Urn actor,
EntityService entityService) {
EditableMLPrimaryKeyProperties editableProperties = (EditableMLPrimaryKeyProperties) getAspectFromEntity(
resourceUrn.toString(), Constants.ML_PRIMARY_KEY_EDITABLE_PROPERTIES_ASPECT_NAME, entityService, new EditableMLPrimaryKeyProperties());
editableProperties.setDescription(newDescription);
persistAspect(resourceUrn, Constants.ML_PRIMARY_KEY_EDITABLE_PROPERTIES_ASPECT_NAME, editableProperties, actor, entityService);
}
}
135 changes: 135 additions & 0 deletions ...rc/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateDescriptionResolver.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,16 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
return updateCorpGroupDescription(targetUrn, input, environment.getContext());
case Constants.NOTEBOOK_ENTITY_NAME:
return updateNotebookDescription(targetUrn, input, environment.getContext());
case Constants.ML_MODEL_ENTITY_NAME:
return updateMlModelDescription(targetUrn, input, environment.getContext());
case Constants.ML_MODEL_GROUP_ENTITY_NAME:
return updateMlModelGroupDescription(targetUrn, input, environment.getContext());
case Constants.ML_FEATURE_TABLE_ENTITY_NAME:
return updateMlFeatureTableDescription(targetUrn, input, environment.getContext());
case Constants.ML_FEATURE_ENTITY_NAME:
return updateMlFeatureDescription(targetUrn, input, environment.getContext());
case Constants.ML_PRIMARY_KEY_ENTITY_NAME:
return updateMlPrimaryKeyDescription(targetUrn, input, environment.getContext());
default:
throw new RuntimeException(
String.format("Failed to update description. Unsupported resource type %s provided.", targetUrn));
Expand Down Expand Up @@ -219,4 +229,129 @@ private CompletableFuture<Boolean> updateNotebookDescription(Urn targetUrn, Desc
}
});
}

private CompletableFuture<Boolean> updateMlModelDescription(Urn targetUrn, DescriptionUpdateInput input,
QueryContext context) {
return CompletableFuture.supplyAsync(() -> {

if (!DescriptionUtils.isAuthorizedToUpdateDescription(context, targetUrn)) {
throw new AuthorizationException(
"Unauthorized to perform this action. Please contact your DataHub administrator.");
}
DescriptionUtils.validateLabelInput(targetUrn, _entityService);

try {
Urn actor = CorpuserUrn.createFromString(context.getActorUrn());
DescriptionUtils.updateMlModelDescription(
input.getDescription(),
targetUrn,
actor,
_entityService);
return true;
} catch (Exception e) {
log.error("Failed to perform update against input {}, {}", input.toString(), e.getMessage());
throw new RuntimeException(String.format("Failed to perform update against input %s", input.toString()), e);
}
});
}

private CompletableFuture<Boolean> updateMlModelGroupDescription(Urn targetUrn, DescriptionUpdateInput input,
QueryContext context) {
return CompletableFuture.supplyAsync(() -> {

if (!DescriptionUtils.isAuthorizedToUpdateDescription(context, targetUrn)) {
throw new AuthorizationException(
"Unauthorized to perform this action. Please contact your DataHub administrator.");
}
DescriptionUtils.validateLabelInput(targetUrn, _entityService);

try {
Urn actor = CorpuserUrn.createFromString(context.getActorUrn());
DescriptionUtils.updateMlModelGroupDescription(
input.getDescription(),
targetUrn,
actor,
_entityService);
return true;
} catch (Exception e) {
log.error("Failed to perform update against input {}, {}", input.toString(), e.getMessage());
throw new RuntimeException(String.format("Failed to perform update against input %s", input.toString()), e);
}
});
}

private CompletableFuture<Boolean> updateMlFeatureDescription(Urn targetUrn, DescriptionUpdateInput input,
QueryContext context) {
return CompletableFuture.supplyAsync(() -> {

if (!DescriptionUtils.isAuthorizedToUpdateDescription(context, targetUrn)) {
throw new AuthorizationException(
"Unauthorized to perform this action. Please contact your DataHub administrator.");
}
DescriptionUtils.validateLabelInput(targetUrn, _entityService);

try {
Urn actor = CorpuserUrn.createFromString(context.getActorUrn());
DescriptionUtils.updateMlFeatureDescription(
input.getDescription(),
targetUrn,
actor,
_entityService);
return true;
} catch (Exception e) {
log.error("Failed to perform update against input {}, {}", input.toString(), e.getMessage());
throw new RuntimeException(String.format("Failed to perform update against input %s", input.toString()), e);
}
});
}

private CompletableFuture<Boolean> updateMlPrimaryKeyDescription(Urn targetUrn, DescriptionUpdateInput input,
QueryContext context) {
return CompletableFuture.supplyAsync(() -> {

if (!DescriptionUtils.isAuthorizedToUpdateDescription(context, targetUrn)) {
throw new AuthorizationException(
"Unauthorized to perform this action. Please contact your DataHub administrator.");
}
DescriptionUtils.validateLabelInput(targetUrn, _entityService);

try {
Urn actor = CorpuserUrn.createFromString(context.getActorUrn());
DescriptionUtils.updateMlPrimaryKeyDescription(
input.getDescription(),
targetUrn,
actor,
_entityService);
return true;
} catch (Exception e) {
log.error("Failed to perform update against input {}, {}", input.toString(), e.getMessage());
throw new RuntimeException(String.format("Failed to perform update against input %s", input.toString()), e);
}
});
}

private CompletableFuture<Boolean> updateMlFeatureTableDescription(Urn targetUrn, DescriptionUpdateInput input,
QueryContext context) {
return CompletableFuture.supplyAsync(() -> {

if (!DescriptionUtils.isAuthorizedToUpdateDescription(context, targetUrn)) {
throw new AuthorizationException(
"Unauthorized to perform this action. Please contact your DataHub administrator.");
}
DescriptionUtils.validateLabelInput(targetUrn, _entityService);

try {
Urn actor = CorpuserUrn.createFromString(context.getActorUrn());
DescriptionUtils.updateMlFeatureTableDescription(
input.getDescription(),
targetUrn,
actor,
_entityService);
return true;
} catch (Exception e) {
log.error("Failed to perform update against input {}, {}", input.toString(), e.getMessage());
throw new RuntimeException(String.format("Failed to perform update against input %s", input.toString()), e);
}
});
}
}
42 changes: 42 additions & 0 deletions ...ore/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java
View file
Original file line number Diff line number Diff line change
@@ -1,23 +1,33 @@
package com.linkedin.datahub.graphql.types.mlmodel.mappers;

import com.linkedin.common.Deprecation;


import com.linkedin.common.GlobalTags;
import com.linkedin.common.GlossaryTerms;
import com.linkedin.common.InstitutionalMemory;
import com.linkedin.common.Ownership;
import com.linkedin.common.Status;
import com.linkedin.data.DataMap;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.datahub.graphql.generated.Domain;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.MLFeature;
import com.linkedin.datahub.graphql.generated.MLFeatureDataType;
import com.linkedin.datahub.graphql.generated.MLFeatureEditableProperties;
import com.linkedin.datahub.graphql.types.common.mappers.DeprecationMapper;
import com.linkedin.datahub.graphql.types.common.mappers.InstitutionalMemoryMapper;
import com.linkedin.datahub.graphql.types.common.mappers.OwnershipMapper;
import com.linkedin.datahub.graphql.types.common.mappers.StatusMapper;
import com.linkedin.datahub.graphql.types.common.mappers.util.MappingHelper;
import com.linkedin.datahub.graphql.types.glossary.mappers.GlossaryTermsMapper;
import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
import com.linkedin.datahub.graphql.types.tag.mappers.GlobalTagsMapper;
import com.linkedin.domain.Domains;
import com.linkedin.entity.EntityResponse;
import com.linkedin.entity.EnvelopedAspectMap;
import com.linkedin.metadata.key.MLFeatureKey;
import com.linkedin.ml.metadata.EditableMLFeatureProperties;
import com.linkedin.ml.metadata.MLFeatureProperties;
import javax.annotation.Nonnull;

Expand Down Expand Up @@ -54,6 +64,12 @@ public MLFeature apply(@Nonnull final EntityResponse entityResponse) {
mappingHelper.mapToResult(DEPRECATION_ASPECT_NAME, (mlFeature, dataMap) ->
mlFeature.setDeprecation(DeprecationMapper.map(new Deprecation(dataMap))));

mappingHelper.mapToResult(GLOBAL_TAGS_ASPECT_NAME, this::mapGlobalTags);
mappingHelper.mapToResult(GLOSSARY_TERMS_ASPECT_NAME, (entity, dataMap) ->
entity.setGlossaryTerms(GlossaryTermsMapper.map(new GlossaryTerms(dataMap))));
mappingHelper.mapToResult(DOMAINS_ASPECT_NAME, this::mapDomains);
mappingHelper.mapToResult(ML_FEATURE_EDITABLE_PROPERTIES_ASPECT_NAME, this::mapEditableProperties);

return mappingHelper.getResult();
}

Expand All @@ -66,9 +82,35 @@ private void mapMLFeatureKey(@Nonnull MLFeature mlFeature, @Nonnull DataMap data
private void mapMLFeatureProperties(@Nonnull MLFeature mlFeature, @Nonnull DataMap dataMap) {
MLFeatureProperties featureProperties = new MLFeatureProperties(dataMap);
mlFeature.setFeatureProperties(MLFeaturePropertiesMapper.map(featureProperties));
mlFeature.setProperties(MLFeaturePropertiesMapper.map(featureProperties));
mlFeature.setDescription(featureProperties.getDescription());
if (featureProperties.getDataType() != null) {
mlFeature.setDataType(MLFeatureDataType.valueOf(featureProperties.getDataType().toString()));
}
}

private void mapGlobalTags(MLFeature entity, DataMap dataMap) {
GlobalTags globalTags = new GlobalTags(dataMap);
com.linkedin.datahub.graphql.generated.GlobalTags graphQlGlobalTags = GlobalTagsMapper.map(globalTags);
entity.setTags(graphQlGlobalTags);
}

private void mapDomains(@Nonnull MLFeature entity, @Nonnull DataMap dataMap) {
final Domains domains = new Domains(dataMap);
// Currently we only take the first domain if it exists.
if (domains.getDomains().size() > 0) {
entity.setDomain(Domain.builder()
.setType(EntityType.DOMAIN)
.setUrn(domains.getDomains().get(0).toString()).build());
}
}

private void mapEditableProperties(MLFeature entity, DataMap dataMap) {
EditableMLFeatureProperties input = new EditableMLFeatureProperties(dataMap);
MLFeatureEditableProperties editableProperties = new MLFeatureEditableProperties();
if (input.hasDescription()) {
editableProperties.setDescription(input.getDescription());
}
entity.setEditableProperties(editableProperties);
}
}
Loading

0 comments on commit 3029c23

Please sign in to comment.