yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-03 14:16:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ml): bringing ml screens up to date w/ the modern ui layout & improving ml lineage (#4651)

Browse Source 
* 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
This commit is contained in:
Gabe Lyons 2022-04-12 22:42:12 -07:00 committed by GitHub
parent 08c34bfe15
commit c92990d32b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
57 changed files with 1729 additions and 681 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
View File

@ -1175,6 +1175,13 @@ public class GmsGraphQLEngine {
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<>(
@ -1217,6 +1224,13 @@ public class GmsGraphQLEngine {
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<>(
@ -1243,6 +1257,43 @@ public class GmsGraphQLEngine {
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;
}))
);
}

59
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/DescriptionUtils.java
View File

@ -15,6 +15,11 @@ import com.linkedin.identity.CorpGroupEditableInfo;
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;
@ -261,4 +266,58 @@ public class DescriptionUtils {
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
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/UpdateDescriptionResolver.java
View File

@ -40,6 +40,16 @@ public class UpdateDescriptionResolver implements DataFetcher<CompletableFuture<
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));
@ -219,4 +229,129 @@ public class UpdateDescriptionResolver implements DataFetcher<CompletableFuture<
}
});
}
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
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureMapper.java
View File

@ -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;
@ -54,6 +64,12 @@ public class MLFeatureMapper implements ModelMapper<EntityResponse, MLFeature> {
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();
}
@ -66,9 +82,35 @@ public class MLFeatureMapper implements ModelMapper<EntityResponse, MLFeature> {
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);
}
}

41
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLFeatureTableMapper.java
View File

@ -1,23 +1,32 @@
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.DataPlatform;
import com.linkedin.datahub.graphql.generated.Domain;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.MLFeatureTable;
import com.linkedin.datahub.graphql.generated.MLFeatureTableEditableProperties;
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.MLFeatureTableKey;
import com.linkedin.ml.metadata.EditableMLFeatureTableProperties;
import com.linkedin.ml.metadata.MLFeatureTableProperties;
import javax.annotation.Nonnull;
@ -54,6 +63,12 @@ public class MLFeatureTableMapper implements ModelMapper<EntityResponse, MLFeatu
mappingHelper.mapToResult(DEPRECATION_ASPECT_NAME, (mlFeatureTable, dataMap) ->
mlFeatureTable.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_TABLE_EDITABLE_PROPERTIES_ASPECT_NAME, this::mapEditableProperties);
return mappingHelper.getResult();
}
@ -68,6 +83,32 @@ public class MLFeatureTableMapper implements ModelMapper<EntityResponse, MLFeatu
private void mapMLFeatureTableProperties(@Nonnull MLFeatureTable mlFeatureTable, @Nonnull DataMap dataMap) {
MLFeatureTableProperties featureTableProperties = new MLFeatureTableProperties(dataMap);
mlFeatureTable.setFeatureTableProperties(MLFeatureTablePropertiesMapper.map(featureTableProperties));
mlFeatureTable.setProperties(MLFeatureTablePropertiesMapper.map(featureTableProperties));
mlFeatureTable.setDescription(featureTableProperties.getDescription());
}
private void mapGlobalTags(MLFeatureTable 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 MLFeatureTable 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(MLFeatureTable entity, DataMap dataMap) {
EditableMLFeatureTableProperties input = new EditableMLFeatureTableProperties(dataMap);
MLFeatureTableEditableProperties editableProperties = new MLFeatureTableEditableProperties();
if (input.hasDescription()) {
editableProperties.setDescription(input.getDescription());
}
entity.setEditableProperties(editableProperties);
}
}

39
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupMapper.java
View File

@ -1,22 +1,30 @@
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.Ownership;
import com.linkedin.common.Status;
import com.linkedin.data.DataMap;
import com.linkedin.data.template.RecordTemplate;
import com.linkedin.datahub.graphql.generated.DataPlatform;
import com.linkedin.datahub.graphql.generated.Domain;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.FabricType;
import com.linkedin.datahub.graphql.generated.MLModelGroup;
import com.linkedin.datahub.graphql.generated.MLModelGroupEditableProperties;
import com.linkedin.datahub.graphql.types.common.mappers.DeprecationMapper;
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.MLModelGroupKey;
import com.linkedin.ml.metadata.EditableMLModelGroupProperties;
import com.linkedin.ml.metadata.MLModelGroupProperties;
import javax.annotation.Nonnull;
@ -51,6 +59,12 @@ public class MLModelGroupMapper implements ModelMapper<EntityResponse, MLModelGr
mappingHelper.mapToResult(DEPRECATION_ASPECT_NAME, (mlModelGroup, dataMap) ->
mlModelGroup.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_MODEL_GROUP_EDITABLE_PROPERTIES_ASPECT_NAME, this::mapEditableProperties);
return mappingHelper.getResult();
}
@ -70,4 +84,29 @@ public class MLModelGroupMapper implements ModelMapper<EntityResponse, MLModelGr
mlModelGroup.setDescription(modelGroupProperties.getDescription());
}
}
private void mapGlobalTags(MLModelGroup 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 MLModelGroup 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(MLModelGroup entity, DataMap dataMap) {
EditableMLModelGroupProperties input = new EditableMLModelGroupProperties(dataMap);
MLModelGroupEditableProperties editableProperties = new MLModelGroupEditableProperties();
if (input.hasDescription()) {
editableProperties.setDescription(input.getDescription());
}
entity.setEditableProperties(editableProperties);
}
}

30
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelMapper.java
View File

@ -3,27 +3,33 @@ package com.linkedin.datahub.graphql.types.mlmodel.mappers;
import com.linkedin.common.Cost;
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.DataPlatform;
import com.linkedin.datahub.graphql.generated.Domain;
import com.linkedin.datahub.graphql.generated.EntityType;
import com.linkedin.datahub.graphql.generated.FabricType;
import com.linkedin.datahub.graphql.generated.MLModel;
import com.linkedin.datahub.graphql.generated.MLModelEditableProperties;
import com.linkedin.datahub.graphql.types.common.mappers.CostMapper;
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.MLModelKey;
import com.linkedin.ml.metadata.CaveatsAndRecommendations;
import com.linkedin.ml.metadata.EditableMLModelProperties;
import com.linkedin.ml.metadata.EthicalConsiderations;
import com.linkedin.ml.metadata.EvaluationData;
import com.linkedin.ml.metadata.IntendedUse;
@ -92,6 +98,10 @@ public class MLModelMapper implements ModelMapper<EntityResponse, MLModel> {
mlModel.setCost(CostMapper.map(new Cost(dataMap))));
mappingHelper.mapToResult(DEPRECATION_ASPECT_NAME, (mlModel, dataMap) ->
mlModel.setDeprecation(DeprecationMapper.map(new Deprecation(dataMap))));
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_MODEL_EDITABLE_PROPERTIES_ASPECT_NAME, this::mapEditableProperties);
return mappingHelper.getResult();
}
@ -128,4 +138,24 @@ public class MLModelMapper implements ModelMapper<EntityResponse, MLModel> {
.map(SourceCodeUrlMapper::map).collect(Collectors.toList()));
mlModel.setSourceCode(graphQlSourceCode);
}
private void mapDomains(@Nonnull MLModel 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(MLModel entity, DataMap dataMap) {
EditableMLModelProperties input = new EditableMLModelProperties(dataMap);
MLModelEditableProperties editableProperties = new MLModelEditableProperties();
if (input.hasDescription()) {
editableProperties.setDescription(input.getDescription());
}
entity.setEditableProperties(editableProperties);
}
}

39
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLPrimaryKeyMapper.java
View File

@ -1,23 +1,31 @@
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.MLFeatureDataType;
import com.linkedin.datahub.graphql.generated.MLPrimaryKey;
import com.linkedin.datahub.graphql.generated.MLPrimaryKeyEditableProperties;
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.MLPrimaryKeyKey;
import com.linkedin.ml.metadata.EditableMLPrimaryKeyProperties;
import com.linkedin.ml.metadata.MLPrimaryKeyProperties;
import javax.annotation.Nonnull;
@ -54,6 +62,11 @@ public class MLPrimaryKeyMapper implements ModelMapper<EntityResponse, MLPrimary
mappingHelper.mapToResult(DEPRECATION_ASPECT_NAME, (mlPrimaryKey, dataMap) ->
mlPrimaryKey.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_PRIMARY_KEY_EDITABLE_PROPERTIES_ASPECT_NAME, this::mapEditableProperties);
return mappingHelper.getResult();
}
@ -66,9 +79,35 @@ public class MLPrimaryKeyMapper implements ModelMapper<EntityResponse, MLPrimary
private void mapMLPrimaryKeyProperties(MLPrimaryKey mlPrimaryKey, DataMap dataMap) {
MLPrimaryKeyProperties primaryKeyProperties = new MLPrimaryKeyProperties(dataMap);
mlPrimaryKey.setPrimaryKeyProperties(MLPrimaryKeyPropertiesMapper.map(primaryKeyProperties));
mlPrimaryKey.setProperties(MLPrimaryKeyPropertiesMapper.map(primaryKeyProperties));
mlPrimaryKey.setDescription(primaryKeyProperties.getDescription());
if (primaryKeyProperties.getDataType() != null) {
mlPrimaryKey.setDataType(MLFeatureDataType.valueOf(primaryKeyProperties.getDataType().toString()));
}
}
private void mapGlobalTags(MLPrimaryKey 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 MLPrimaryKey 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(MLPrimaryKey entity, DataMap dataMap) {
EditableMLPrimaryKeyProperties input = new EditableMLPrimaryKeyProperties(dataMap);
MLPrimaryKeyEditableProperties editableProperties = new MLPrimaryKeyEditableProperties();
if (input.hasDescription()) {
editableProperties.setDescription(input.getDescription());
}
entity.setEditableProperties(editableProperties);
}
}

153
datahub-graphql-core/src/main/resources/entity.graphql
View File

@ -6442,6 +6442,21 @@ type MLModel implements EntityWithRelationships & Entity {
Edges extending from this entity grouped by direction in the lineage graph
"""
lineage(input: LineageInput!): EntityLineageResult
"""
The structured glossary terms associated with the entity
"""
glossaryTerms: GlossaryTerms
"""
The Domain associated with the entity
"""
domain: Domain
"""
An additional set of of read write properties
"""
editableProperties: MLModelEditableProperties
}
"""
@ -6508,11 +6523,31 @@ type MLModelGroup implements EntityWithRelationships & Entity {
Edges extending from this entity grouped by direction in the lineage graph
"""
lineage(input: LineageInput!): EntityLineageResult
"""
Tags applied to entity
"""
tags: GlobalTags
"""
The structured glossary terms associated with the entity
"""
glossaryTerms: GlossaryTerms
"""
The Domain associated with the entity
"""
domain: Domain
"""
An additional set of of read write properties
"""
editableProperties: MLModelGroupEditableProperties
}
type MLModelGroupProperties {
description: String
description: String
createdAt: Long
@ -6522,7 +6557,7 @@ type MLModelGroupProperties {
"""
An ML Feature Metadata Entity Note that this entity is incubating
"""
type MLFeature implements Entity {
type MLFeature implements EntityWithRelationships & Entity {
"""
The primary key of the ML Feature
"""
@ -6561,7 +6596,12 @@ type MLFeature implements Entity {
"""
ModelProperties metadata of the MLFeature
"""
featureProperties: MLFeatureProperties
featureProperties: MLFeatureProperties @deprecated
"""
ModelProperties metadata of the MLFeature
"""
properties: MLFeatureProperties
"""
References to internal resources related to the MLFeature
@ -6587,12 +6627,32 @@ type MLFeature implements Entity {
Edges extending from this entity grouped by direction in the lineage graph
"""
lineage(input: LineageInput!): EntityLineageResult
"""
Tags applied to entity
"""
tags: GlobalTags
"""
The structured glossary terms associated with the entity
"""
glossaryTerms: GlossaryTerms
"""
The Domain associated with the entity
"""
domain: Domain
"""
An additional set of of read write properties
"""
editableProperties: MLFeatureEditableProperties
}
type MLHyperParam {
name: String
name: String
description: String
description: String
value: String
@ -6648,7 +6708,7 @@ type MLFeatureProperties {
"""
An ML Primary Key Entity Note that this entity is incubating
"""
type MLPrimaryKey implements Entity {
type MLPrimaryKey implements EntityWithRelationships & Entity {
"""
The primary key of the ML Primary Key
"""
@ -6719,11 +6779,31 @@ type MLPrimaryKey implements Entity {
Edges extending from this entity grouped by direction in the lineage graph
"""
lineage(input: LineageInput!): EntityLineageResult
"""
Tags applied to entity
"""
tags: GlobalTags
"""
The structured glossary terms associated with the entity
"""
glossaryTerms: GlossaryTerms
"""
The Domain associated with the entity
"""
domain: Domain
"""
An additional set of of read write properties
"""
editableProperties: MLPrimaryKeyEditableProperties
}
type MLPrimaryKeyProperties {
description: String
description: String
dataType: MLFeatureDataType
@ -6735,7 +6815,7 @@ type MLPrimaryKeyProperties {
"""
An ML Feature Table Entity Note that this entity is incubating
"""
type MLFeatureTable implements Entity {
type MLFeatureTable implements EntityWithRelationships & Entity {
"""
The primary key of the ML Feature Table
"""
@ -6801,11 +6881,66 @@ type MLFeatureTable implements Entity {
Edges extending from this entity grouped by direction in the lineage graph
"""
lineage(input: LineageInput!): EntityLineageResult
"""
Tags applied to entity
"""
tags: GlobalTags
"""
The structured glossary terms associated with the entity
"""
glossaryTerms: GlossaryTerms
"""
The Domain associated with the entity
"""
domain: Domain
"""
An additional set of of read write properties
"""
editableProperties: MLFeatureTableEditableProperties
}
type MLFeatureTableEditableProperties {
"""
The edited description
"""
description: String
}
type MLFeatureEditableProperties {
"""
The edited description
"""
description: String
}
type MLPrimaryKeyEditableProperties {
"""
The edited description
"""
description: String
}
type MLModelEditableProperties {
"""
The edited description
"""
description: String
}
type MLModelGroupEditableProperties {
"""
The edited description
"""
description: String
}
type MLFeatureTableProperties {
description: String
description: String
mlFeatures: [MLFeature]

5
datahub-web-react/src/app/entity/EntityPage.tsx
View File

@ -60,6 +60,11 @@ export const EntityPage = ({ entityType }: Props) => {
entityType === EntityType.Chart ||
entityType === EntityType.DataFlow ||
entityType === EntityType.DataJob ||
entityType === EntityType.Mlmodel ||
entityType === EntityType.Mlfeature ||
entityType === EntityType.MlprimaryKey ||
entityType === EntityType.MlfeatureTable ||
entityType === EntityType.MlmodelGroup ||
entityType === EntityType.GlossaryTerm;
return (

1
datahub-web-react/src/app/entity/dataset/DatasetEntity.tsx
View File

@ -112,7 +112,6 @@ export class DatasetEntity implements Entity<Dataset> {
display: {
visible: (_, _1) => true,
enabled: (_, dataset: GetDatasetQuery) => {
console.log(dataset?.dataset?.upstream, dataset?.dataset?.downstream);
return (
(dataset?.dataset?.upstream?.total || 0) > 0 ||
(dataset?.dataset?.downstream?.total || 0) > 0

69
datahub-web-react/src/app/entity/mlFeature/MLFeatureEntity.tsx
View File

@ -1,10 +1,18 @@
import * as React from 'react';
import { DotChartOutlined } from '@ant-design/icons';
import { MlFeature, EntityType, SearchResult } from '../../../types.generated';
import { MlFeature, EntityType, SearchResult, OwnershipType } from '../../../types.generated';
import { Preview } from './preview/Preview';
import { MLFeatureProfile } from './profile/MLFeatureProfile';
import { Entity, IconStyleType, PreviewType } from '../Entity';
import { getDataForEntityType } from '../shared/containers/profile/utils';
import { EntityProfile } from '../shared/containers/profile/EntityProfile';
import { GenericEntityProperties } from '../shared/types';
import { useGetMlFeatureQuery } from '../../../graphql/mlFeature.generated';
import { SidebarAboutSection } from '../shared/containers/profile/sidebar/SidebarAboutSection';
import { SidebarTagsSection } from '../shared/containers/profile/sidebar/SidebarTagsSection';
import { SidebarOwnerSection } from '../shared/containers/profile/sidebar/Ownership/SidebarOwnerSection';
import { SidebarDomainSection } from '../shared/containers/profile/sidebar/Domain/SidebarDomainSection';
import { DocumentationTab } from '../shared/tabs/Documentation/DocumentationTab';
import { FeatureTableTab } from '../shared/tabs/ML/MlFeatureFeatureTableTab';
/**
* Definition of the DataHub MLFeature entity.
@ -45,7 +53,50 @@ export class MLFeatureEntity implements Entity<MlFeature> {
getCollectionName = () => 'Features';
renderProfile = (urn: string) => <MLFeatureProfile urn={urn} />;
getOverridePropertiesFromEntity = (_?: MlFeature | null): GenericEntityProperties => {
return {};
};
renderProfile = (urn: string) => (
<EntityProfile
urn={urn}
key={urn}
entityType={EntityType.Mlfeature}
useEntityQuery={useGetMlFeatureQuery}
getOverrideProperties={this.getOverridePropertiesFromEntity}
tabs={[
{
name: 'Feature Tables',
component: FeatureTableTab,
},
{
name: 'Documentation',
component: DocumentationTab,
},
]}
sidebarSections={[
{
component: SidebarAboutSection,
},
{
component: SidebarTagsSection,
properties: {
hasTags: true,
hasTerms: true,
},
},
{
component: SidebarOwnerSection,
properties: {
defaultOwnerType: OwnershipType.TechnicalOwner,
},
},
{
component: SidebarDomainSection,
},
]}
/>
);
renderPreview = (_: PreviewType, data: MlFeature) => {
return (
@ -79,4 +130,16 @@ export class MLFeatureEntity implements Entity<MlFeature> {
getGenericEntityProperties = (mlFeature: MlFeature) => {
return getDataForEntityType({ data: mlFeature, entityType: this.type, getOverrideProperties: (data) => data });
};
getLineageVizConfig = (entity: MlFeature) => {
return {
urn: entity.urn,
name: entity.name,
type: EntityType.Mlfeature,
// eslint-disable-next-line
icon: entity?.['featureTables']?.relationships?.[0]?.entity?.platform?.properties?.logoUrl || undefined,
// eslint-disable-next-line
platform: entity?.['featureTables']?.relationships?.[0]?.entity?.platform?.name,
};
};
}

34
datahub-web-react/src/app/entity/mlFeature/profile/MLFeatureProfile.tsx
View File

@ -1,34 +0,0 @@
import React from 'react';
import { LegacyEntityProfile } from '../../../shared/LegacyEntityProfile';
import { EntityType } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import analytics, { EventType } from '../../../analytics';
export enum TabType {
Features = 'Features',
Sources = 'Sources',
Ownership = 'Ownership',
}
/**
* Responsible for display the MLFeature Page
*/
export const MLFeatureProfile = ({ urn }: { urn: string }): JSX.Element => {
const entityRegistry = useEntityRegistry();
return (
<LegacyEntityProfile
titleLink={`/${entityRegistry.getPathName(EntityType.Mlfeature)}/${urn}`}
title={urn}
header={<></>}
onTabChange={(tab: string) => {
analytics.event({
type: EventType.EntitySectionViewEvent,
entityType: EntityType.Mlfeature,
entityUrn: urn,
section: tab,
});
}}
/>
);
};

64
datahub-web-react/src/app/entity/mlFeatureTable/MLFeatureTableEntity.tsx
View File

@ -1,10 +1,19 @@
import * as React from 'react';
import { DotChartOutlined } from '@ant-design/icons';
import { MlFeatureTable, EntityType, SearchResult } from '../../../types.generated';
import { MlFeatureTable, EntityType, SearchResult, OwnershipType } from '../../../types.generated';
import { Preview } from './preview/Preview';
import { MLFeatureTableProfile } from './profile/MLFeatureTableProfile';
import { Entity, IconStyleType, PreviewType } from '../Entity';
import { getDataForEntityType } from '../shared/containers/profile/utils';
import { GenericEntityProperties } from '../shared/types';
import { useGetMlFeatureTableQuery } from '../../../graphql/mlFeatureTable.generated';
import { EntityProfile } from '../shared/containers/profile/EntityProfile';
import { SidebarDomainSection } from '../shared/containers/profile/sidebar/Domain/SidebarDomainSection';
import { SidebarOwnerSection } from '../shared/containers/profile/sidebar/Ownership/SidebarOwnerSection';
import { SidebarAboutSection } from '../shared/containers/profile/sidebar/SidebarAboutSection';
import { SidebarTagsSection } from '../shared/containers/profile/sidebar/SidebarTagsSection';
import MlFeatureTableFeatures from './profile/features/MlFeatureTableFeatures';
import Sources from './profile/Sources';
import { DocumentationTab } from '../shared/tabs/Documentation/DocumentationTab';
/**
* Definition of the DataHub MLFeatureTable entity.
@ -45,7 +54,54 @@ export class MLFeatureTableEntity implements Entity<MlFeatureTable> {
getCollectionName = () => 'Feature Tables';
renderProfile = (urn: string) => <MLFeatureTableProfile urn={urn} />;
getOverridePropertiesFromEntity = (_?: MlFeatureTable | null): GenericEntityProperties => {
return {};
};
renderProfile = (urn: string) => (
<EntityProfile
urn={urn}
key={urn}
entityType={EntityType.MlfeatureTable}
useEntityQuery={useGetMlFeatureTableQuery}
getOverrideProperties={this.getOverridePropertiesFromEntity}
tabs={[
{
name: 'Features',
component: MlFeatureTableFeatures,
},
{
name: 'Sources',
component: Sources,
},
{
name: 'Documentation',
component: DocumentationTab,
},
]}
sidebarSections={[
{
component: SidebarAboutSection,
},
{
component: SidebarTagsSection,
properties: {
hasTags: true,
hasTerms: true,
},
},
{
component: SidebarOwnerSection,
properties: {
defaultOwnerType: OwnershipType.TechnicalOwner,
},
},
{
component: SidebarDomainSection,
},
]}
/>
);
renderPreview = (_: PreviewType, data: MlFeatureTable) => {
return (
@ -79,8 +135,6 @@ export class MLFeatureTableEntity implements Entity<MlFeatureTable> {
urn: entity.urn,
name: entity.name,
type: EntityType.MlfeatureTable,
upstreamChildren: [],
downstreamChildren: [],
icon: entity.platform.properties?.logoUrl || undefined,
platform: entity.platform.name,
};

99
datahub-web-react/src/app/entity/mlFeatureTable/profile/MLFeatureTableProfile.tsx
View File

@ -1,99 +0,0 @@
import React from 'react';
import { Alert } from 'antd';
import { useGetMlFeatureTableQuery } from '../../../../graphql/mlFeatureTable.generated';
import { LegacyEntityProfile } from '../../../shared/LegacyEntityProfile';
import { MlFeatureTable, MlFeature, MlPrimaryKey, EntityType } from '../../../../types.generated';
import MLFeatureTableHeader from './MLFeatureTableHeader';
import { Message } from '../../../shared/Message';
import { Ownership as OwnershipView } from '../../shared/components/legacy/Ownership';
import { useEntityRegistry } from '../../../useEntityRegistry';
import analytics, { EventType } from '../../../analytics';
import { notEmpty } from '../../shared/utils';
import MlFeatureTableFeatures from './features/MlFeatureTableFeatures';
import SourcesView from './Sources';
export enum TabType {
Features = 'Features',
Sources = 'Sources',
Ownership = 'Ownership',
}
/**
* Responsible for display the MLFeatureTable Page
*/
export const MLFeatureTableProfile = ({ urn }: { urn: string }): JSX.Element => {
const entityRegistry = useEntityRegistry();
const { loading, error, data } = useGetMlFeatureTableQuery({ variables: { urn } });
if (error || (!loading && !error && !data)) {
return <Alert type="error" message={error?.message || 'Entity failed to load'} />;
}
const getHeader = (mlFeatureTable: MlFeatureTable) => <MLFeatureTableHeader mlFeatureTable={mlFeatureTable} />;
const getTabs = ({ ownership, featureTableProperties }: MlFeatureTable) => {
const features: Array<MlFeature | MlPrimaryKey> =
featureTableProperties && (featureTableProperties?.mlFeatures || featureTableProperties?.mlPrimaryKeys)
? [
...(featureTableProperties?.mlPrimaryKeys || []),
...(featureTableProperties?.mlFeatures || []),
].filter(notEmpty)
: [];
return [
{
name: TabType.Features,
path: TabType.Features.toLowerCase(),
content: <MlFeatureTableFeatures features={features} />,
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
{
name: TabType.Sources,
path: TabType.Sources.toLowerCase(),
content: <SourcesView features={features} />,
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
{
name: TabType.Ownership,
path: TabType.Ownership.toLowerCase(),
content: (
<OwnershipView
owners={(ownership && ownership.owners) || []}
lastModifiedAt={(ownership && ownership.lastModified?.time) || 0}
/>
),
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
];
};
return (
<>
{loading && <Message type="loading" content="Loading..." style={{ marginTop: '10%' }} />}
{data && data.mlFeatureTable && (
<LegacyEntityProfile
titleLink={`/${entityRegistry.getPathName(EntityType.MlfeatureTable)}/${urn}`}
title={data.mlFeatureTable?.name || ''}
tabs={getTabs(data.mlFeatureTable as MlFeatureTable)}
header={getHeader(data.mlFeatureTable as MlFeatureTable)}
onTabChange={(tab: string) => {
analytics.event({
type: EventType.EntitySectionViewEvent,
entityType: EntityType.MlfeatureTable,
entityUrn: urn,
section: tab,
});
}}
/>
)}
</>
);
};

35
datahub-web-react/src/app/entity/mlFeatureTable/profile/Sources.tsx
View File

@ -1,35 +1,48 @@
import { List, Typography } from 'antd';
import React, { useMemo } from 'react';
import styled from 'styled-components';
import { MlFeature, MlPrimaryKey, Dataset, EntityType } from '../../../../types.generated';
import { GetMlFeatureTableQuery } from '../../../../graphql/mlFeatureTable.generated';
import { Dataset, EntityType } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import { PreviewType } from '../../Entity';
export type Props = {
features?: Array<MlFeature | MlPrimaryKey>;
};
import { useBaseEntity } from '../../shared/EntityContext';
import { notEmpty } from '../../shared/utils';
const ViewRawButtonContainer = styled.div`
display: flex;
justify-content: flex-end;
`;
export default function SourcesView({ features }: Props) {
export default function SourcesView() {
const entityRegistry = useEntityRegistry();
const baseEntity = useBaseEntity<GetMlFeatureTableQuery>();
const featureTable = baseEntity?.mlFeatureTable;
const features = useMemo(
() =>
featureTable?.properties &&
(featureTable?.properties?.mlFeatures || featureTable?.properties?.mlPrimaryKeys)
? [
...(featureTable?.properties?.mlPrimaryKeys || []),
...(featureTable?.properties?.mlFeatures || []),
].filter(notEmpty)
: [],
[featureTable?.properties],
);
const sources = useMemo(
() =>
features?.reduce((accumulator: Array<Dataset>, feature: MlFeature | MlPrimaryKey) => {
if (feature.__typename === 'MLFeature' && feature.featureProperties?.sources) {
features?.reduce((accumulator: Array<Dataset>, feature) => {
if (feature.__typename === 'MLFeature' && feature.properties?.sources) {
// eslint-disable-next-line array-callback-return
feature.featureProperties?.sources.map((source: Dataset | null) => {
feature.properties?.sources.map((source: Dataset | null) => {
if (source && accumulator.findIndex((dataset) => dataset.urn === source?.urn) === -1) {
accumulator.push(source);
}
});
} else if (feature.__typename === 'MLPrimaryKey' && feature.primaryKeyProperties?.sources) {
} else if (feature.__typename === 'MLPrimaryKey' && feature.properties?.sources) {
// eslint-disable-next-line array-callback-return
feature.primaryKeyProperties?.sources.map((source: Dataset | null) => {
feature.properties?.sources.map((source: Dataset | null) => {
if (source && accumulator.findIndex((dataset) => dataset.urn === source?.urn) === -1) {
accumulator.push(source);
}

130
datahub-web-react/src/app/entity/mlFeatureTable/profile/features/MlFeatureTableFeatures.tsx
View File

@ -1,20 +1,24 @@
import React from 'react';
import React, { useState } from 'react';
import { Table, Typography } from 'antd';
import { CheckSquareOutlined } from '@ant-design/icons';
import { AlignType } from 'rc-table/lib/interface';
import styled from 'styled-components';
import { Link } from 'react-router-dom';
import MlFeatureDataTypeIcon from './MlFeatureDataTypeIcon';
import { MlFeatureDataType, MlPrimaryKey, MlFeature } from '../../../../../types.generated';
import MarkdownViewer from '../../../shared/components/legacy/MarkdownViewer';
import { GetMlFeatureTableQuery } from '../../../../../graphql/mlFeatureTable.generated';
import { useBaseEntity, useRefetch } from '../../../shared/EntityContext';
import { notEmpty } from '../../../shared/utils';
import TagTermGroup from '../../../../shared/tags/TagTermGroup';
import SchemaDescriptionField from '../../../dataset/profile/schema/components/SchemaDescriptionField';
import { useUpdateDescriptionMutation } from '../../../../../graphql/mutations.generated';
import { useEntityRegistry } from '../../../../useEntityRegistry';
const FeaturesContainer = styled.div`
margin-bottom: 100px;
`;
export type Props = {
features: Array<MlFeature | MlPrimaryKey>;
};
const defaultColumns = [
{
title: 'Type',
@ -26,36 +30,128 @@ const defaultColumns = [
return <MlFeatureDataTypeIcon dataType={dataType} />;
},
},
{
];
export default function MlFeatureTableFeatures() {
const baseEntity = useBaseEntity<GetMlFeatureTableQuery>();
const refetch = useRefetch();
const featureTable = baseEntity?.mlFeatureTable;
const [updateDescription] = useUpdateDescriptionMutation();
const entityRegistry = useEntityRegistry();
const [tagHoveredIndex, setTagHoveredIndex] = useState<string | undefined>(undefined);
const features =
featureTable?.properties && (featureTable?.properties?.mlFeatures || featureTable?.properties?.mlPrimaryKeys)
? [
...(featureTable?.properties?.mlPrimaryKeys || []),
...(featureTable?.properties?.mlFeatures || []),
].filter(notEmpty)
: [];
const onTagTermCell = (record: any, rowIndex: number | undefined) => ({
onMouseEnter: () => {
setTagHoveredIndex(`${record.urn}-${rowIndex}`);
},
onMouseLeave: () => {
setTagHoveredIndex(undefined);
},
});
const nameColumn = {
title: 'Name',
dataIndex: 'name',
key: 'name',
width: 100,
render: (name: string) => <Typography.Text strong>{name}</Typography.Text>,
},
{
render: (name: string, feature: MlFeature | MlPrimaryKey) => (
<Link to={entityRegistry.getEntityUrl(feature.type, feature.urn)}>
<Typography.Text strong>{name}</Typography.Text>
</Link>
),
};
const descriptionColumn = {
title: 'Description',
dataIndex: 'description',
key: 'description',
render: (description: string) => <MarkdownViewer source={description} />,
render: (_, feature: MlFeature | MlPrimaryKey) => (
<SchemaDescriptionField
description={feature?.editableProperties?.description || feature?.properties?.description || ''}
original={feature?.properties?.description}
isEdited={!!feature?.editableProperties?.description}
onUpdate={(updatedDescription) =>
updateDescription({
variables: {
input: {
description: updatedDescription,
resourceUrn: feature.urn,
},
},
}).then(refetch)
}
/>
),
width: 300,
},
{
};
const tagColumn = {
width: 125,
title: 'Tags',
dataIndex: 'tags',
key: 'tags',
render: (_, feature: MlFeature | MlPrimaryKey, rowIndex: number) => (
<TagTermGroup
editableTags={feature.tags}
canRemove
buttonProps={{ size: 'small' }}
canAddTag={tagHoveredIndex === `${feature.urn}-${rowIndex}`}
onOpenModal={() => setTagHoveredIndex(undefined)}
entityUrn={feature.urn}
entityType={feature.type}
refetch={refetch}
/>
),
onCell: onTagTermCell,
};
const termColumn = {
width: 125,
title: 'Terms',
dataIndex: 'glossaryTerms',
key: 'glossarTerms',
render: (_, feature: MlFeature | MlPrimaryKey, rowIndex: number) => (
<TagTermGroup
editableGlossaryTerms={feature.glossaryTerms}
canRemove
buttonProps={{ size: 'small' }}
canAddTerm={tagHoveredIndex === `${feature.urn}-${rowIndex}`}
onOpenModal={() => setTagHoveredIndex(undefined)}
entityUrn={feature.urn}
entityType={feature.type}
refetch={refetch}
/>
),
onCell: onTagTermCell,
};
const primaryKeyColumn = {
title: 'Primary Key',
dataIndex: 'primaryKey',
key: 'primaryKey',
render: (_: any, record: MlFeature | MlPrimaryKey) =>
record.__typename === 'MLPrimaryKey' ? <CheckSquareOutlined /> : null,
width: 50,
},
];
};
const allColumns = [...defaultColumns, nameColumn, descriptionColumn, tagColumn, termColumn, primaryKeyColumn];
export default function MlFeatureTableFeatures({ features }: Props) {
return (
<FeaturesContainer>
{features && features.length > 0 && (
<Table
columns={defaultColumns}
columns={allColumns}
// eslint-disable-next-line @typescript-eslint/ban-ts-comment
// @ts-ignore
dataSource={features}
rowKey={(record) => `${record.dataType}-${record.name}`}
expandable={{ defaultExpandAllRows: true, expandRowByClick: true }}

72
datahub-web-react/src/app/entity/mlModel/MLModelEntity.tsx
View File

@ -1,10 +1,21 @@
import * as React from 'react';
import { CodeSandboxOutlined } from '@ant-design/icons';
import { MlModel, EntityType, SearchResult } from '../../../types.generated';
import { MlModel, EntityType, SearchResult, OwnershipType } from '../../../types.generated';
import { Preview } from './preview/Preview';
import { MLModelProfile } from './profile/MLModelProfile';
import { Entity, IconStyleType, PreviewType } from '../Entity';
import { getDataForEntityType } from '../shared/containers/profile/utils';
import { EntityProfile } from '../shared/containers/profile/EntityProfile';
import { useGetMlModelQuery } from '../../../graphql/mlModel.generated';
import { GenericEntityProperties } from '../shared/types';
import MLModelSummary from './profile/MLModelSummary';
import MLModelGroupsTab from './profile/MLModelGroupsTab';
import { SidebarTagsSection } from '../shared/containers/profile/sidebar/SidebarTagsSection';
import { SidebarAboutSection } from '../shared/containers/profile/sidebar/SidebarAboutSection';
import { SidebarOwnerSection } from '../shared/containers/profile/sidebar/Ownership/SidebarOwnerSection';
import { SidebarDomainSection } from '../shared/containers/profile/sidebar/Domain/SidebarDomainSection';
import { PropertiesTab } from '../shared/tabs/Properties/PropertiesTab';
import { DocumentationTab } from '../shared/tabs/Documentation/DocumentationTab';
import MlModelFeaturesTab from './profile/MlModelFeaturesTab';
/**
* Definition of the DataHub MlModel entity.
@ -45,7 +56,62 @@ export class MLModelEntity implements Entity<MlModel> {
getCollectionName = () => 'ML Models';
renderProfile = (urn: string) => <MLModelProfile urn={urn} />;
getOverridePropertiesFromEntity = (_?: MlModel | null): GenericEntityProperties => {
return {};
};
renderProfile = (urn: string) => (
<EntityProfile
urn={urn}
key={urn}
entityType={EntityType.Mlmodel}
useEntityQuery={useGetMlModelQuery}
getOverrideProperties={this.getOverridePropertiesFromEntity}
tabs={[
{
name: 'Summary',
component: MLModelSummary,
},
{
name: 'Documentation',
component: DocumentationTab,
},
{
name: 'Group',
component: MLModelGroupsTab,
},
{
name: 'Features',
component: MlModelFeaturesTab,
},
{
name: 'Properties',
component: PropertiesTab,
},
]}
sidebarSections={[
{
component: SidebarAboutSection,
},
{
component: SidebarTagsSection,
properties: {
hasTags: true,
hasTerms: true,
},
},
{
component: SidebarOwnerSection,
properties: {
defaultOwnerType: OwnershipType.TechnicalOwner,
},
},
{
component: SidebarDomainSection,
},
]}
/>
);
renderPreview = (_: PreviewType, data: MlModel) => {
return <Preview model={data} />;

24
datahub-web-react/src/app/entity/mlModel/__tests__/MlModelProfile.test.tsx
View File

@ -1,24 +0,0 @@
import React from 'react';
import { render, waitFor } from '@testing-library/react';
import { MockedProvider } from '@apollo/client/testing';
import { MLModelProfile } from '../profile/MLModelProfile';
import TestPageContainer from '../../../../utils/test-utils/TestPageContainer';
import { mocks } from '../../../../Mocks';
describe('MlModelProfile', () => {
it('renders', async () => {
const { getByText, queryAllByText } = render(
<MockedProvider mocks={mocks} addTypename={false}>
<TestPageContainer
initialEntries={['/mlModel/urn:li:mlModel:(urn:li:dataPlatform:sagemaker,trustmodel,PROD)']}
>
<MLModelProfile urn="urn:li:mlModel:(urn:li:dataPlatform:sagemaker,trustmodel,PROD)" />
</TestPageContainer>
</MockedProvider>,
);
await waitFor(() => expect(queryAllByText('trust model').length).toBeGreaterThanOrEqual(1));
expect(getByText('a ml trust model')).toBeInTheDocument();
});
});

35
datahub-web-react/src/app/entity/mlModel/profile/MLModelGroupsTab.tsx
View File

@ -2,16 +2,21 @@ import React from 'react';
import { Space, Table, Typography } from 'antd';
import Link from 'antd/lib/typography/Link';
import { ColumnsType } from 'antd/es/table';
import styled from 'styled-components';
import { EntityType, MlModel, MlModelGroup } from '../../../../types.generated';
import { EntityType, MlModelGroup } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import { useBaseEntity } from '../../shared/EntityContext';
import { GetMlModelQuery } from '../../../../graphql/mlModel.generated';
export type Props = {
model?: MlModel;
};
const TabContent = styled.div`
padding: 16px;
`;
export default function MLModelGroupsTab() {
const baseEntity = useBaseEntity<GetMlModelQuery>();
const model = baseEntity?.mlModel;
export default function MLModelGroupsTab({ model }: Props) {
console.log(model?.properties);
const entityRegistry = useEntityRegistry();
const propertyTableColumns: ColumnsType<MlModelGroup> = [
@ -29,13 +34,15 @@ export default function MLModelGroupsTab({ model }: Props) {
];
return (
<Space direction="vertical" style={{ width: '100%' }} size="large">
<Typography.Title level={3}>Groups</Typography.Title>
<Table
pagination={false}
columns={propertyTableColumns}
dataSource={model?.properties?.groups as MlModelGroup[]}
/>
</Space>
<TabContent>
<Space direction="vertical" style={{ width: '100%' }} size="large">
<Typography.Title level={3}>Groups</Typography.Title>
<Table
pagination={false}
columns={propertyTableColumns}
dataSource={model?.properties?.groups as MlModelGroup[]}
/>
</Space>
</TabContent>
);
}

65
datahub-web-react/src/app/entity/mlModel/profile/MLModelHeader.tsx
View File

@ -1,65 +0,0 @@
import { Image, Row, Space, Typography } from 'antd';
import React from 'react';
import styled from 'styled-components';
import { MlModel } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import CompactContext from '../../../shared/CompactContext';
import { AvatarsGroup } from '../../../shared/avatar';
import MarkdownViewer from '../../shared/components/legacy/MarkdownViewer';
const HeaderInfoItem = styled.div`
display: inline-block;
text-align: left;
width: 125px;
vertical-align: top;
`;
const PlatformName = styled(Typography.Text)`
font-size: 16px;
`;
const PreviewImage = styled(Image)`
max-height: 20px;
padding-top: 3px;
width: auto;
object-fit: contain;
`;
export type Props = {
mlModel: MlModel;
};
export default function MLModelHeader({ mlModel: { ownership, platform, properties } }: Props) {
const entityRegistry = useEntityRegistry();
const isCompact = React.useContext(CompactContext);
return (
<>
<Space direction="vertical" size="middle">
<Row justify="space-between">
{platform ? (
<HeaderInfoItem>
<div>
<Typography.Text strong type="secondary" style={{ fontSize: 11 }}>
Platform
</Typography.Text>
</div>
<Space direction="horizontal">
{platform.properties?.logoUrl ? (
<PreviewImage
preview={false}
src={platform.properties?.logoUrl}
placeholder
alt={platform.name}
/>
) : null}
<PlatformName>{platform.name}</PlatformName>
</Space>
</HeaderInfoItem>
) : null}
</Row>
<MarkdownViewer isCompact={isCompact} source={properties?.description || ''} />
<AvatarsGroup owners={ownership?.owners} entityRegistry={entityRegistry} />
</Space>
</>
);
}

110
datahub-web-react/src/app/entity/mlModel/profile/MLModelProfile.tsx
View File

@ -1,110 +0,0 @@
import React from 'react';
import { Alert } from 'antd';
import { useGetMlModelQuery } from '../../../../graphql/mlModel.generated';
import { LegacyEntityProfile } from '../../../shared/LegacyEntityProfile';
import { MlModel, EntityType } from '../../../../types.generated';
import MLModelHeader from './MLModelHeader';
import { Message } from '../../../shared/Message';
import { Ownership as OwnershipView } from '../../shared/components/legacy/Ownership';
import { useEntityRegistry } from '../../../useEntityRegistry';
import analytics, { EventType } from '../../../analytics';
import MLModelSummary from './MLModelSummary';
import MLModelGroupsTab from './MLModelGroupsTab';
import { Properties } from '../../shared/components/legacy/Properties';
const EMPTY_ARR: never[] = [];
export enum TabType {
Summary = 'Summary',
Groups = 'Groups',
Deployments = 'Deployments',
Features = 'Features',
Ownership = 'Ownership',
CustomProperties = 'Properties',
}
/**
* Responsible for display the MLModel Page
*/
export const MLModelProfile = ({ urn }: { urn: string }): JSX.Element => {
const entityRegistry = useEntityRegistry();
const { loading, error, data } = useGetMlModelQuery({ variables: { urn } });
if (error || (!loading && !error && !data)) {
return <Alert type="error" message={error?.message || 'Entity failed to load'} />;
}
const getHeader = (mlModel: MlModel) => <MLModelHeader mlModel={mlModel} />;
const getTabs = (model: MlModel) => {
return [
{
name: TabType.Summary,
path: TabType.Summary.toLowerCase(),
content: <MLModelSummary model={model} />,
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
{
name: TabType.Groups,
path: TabType.Groups.toLowerCase(),
content: <MLModelGroupsTab model={model} />,
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
{
name: TabType.CustomProperties,
path: TabType.CustomProperties.toLowerCase(),
content: <Properties properties={model?.properties?.customProperties || EMPTY_ARR} />,
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
// {
// name: TabType.Deployments,
// path: TabType.Deployments.toLowerCase(),
// content: <SourcesView features={features} />,
// },
{
name: TabType.Ownership,
path: TabType.Ownership.toLowerCase(),
content: (
<OwnershipView
owners={model?.ownership?.owners || []}
lastModifiedAt={model?.ownership?.lastModified?.time || 0}
/>
),
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
];
};
return (
<>
{loading && <Message type="loading" content="Loading..." style={{ marginTop: '10%' }} />}
{data && data.mlModel && (
<LegacyEntityProfile
titleLink={`/${entityRegistry.getPathName(EntityType.Mlmodel)}/${urn}`}
title={data.mlModel?.name || ''}
tabs={getTabs(data.mlModel as MlModel)}
header={getHeader(data.mlModel as MlModel)}
onTabChange={(tab: string) => {
analytics.event({
type: EventType.EntitySectionViewEvent,
entityType: EntityType.Mlmodel,
entityUrn: urn,
section: tab,
});
}}
/>
)}
</>
);
};

47
datahub-web-react/src/app/entity/mlModel/profile/MLModelSummary.tsx
View File

@ -1,12 +1,19 @@
import React from 'react';
import styled from 'styled-components';
import { Space, Table, Typography } from 'antd';
import { MlHyperParam, MlMetric, MlModel } from '../../../../types.generated';
export type Props = {
model?: MlModel;
};
import { MlHyperParam, MlMetric } from '../../../../types.generated';
import { useBaseEntity } from '../../shared/EntityContext';
import { GetMlModelQuery } from '../../../../graphql/mlModel.generated';
const TabContent = styled.div`
padding: 16px;
`;
export default function MLModelSummary() {
const baseEntity = useBaseEntity<GetMlModelQuery>();
const model = baseEntity?.mlModel;
export default function MLModelSummary({ model }: Props) {
const propertyTableColumns = [
{
title: 'Name',
@ -20,19 +27,21 @@ export default function MLModelSummary({ model }: Props) {
];
return (
<Space direction="vertical" style={{ width: '100%' }} size="large">
<Typography.Title level={3}>Training Metrics</Typography.Title>
<Table
pagination={false}
columns={propertyTableColumns}
dataSource={model?.properties?.trainingMetrics as MlMetric[]}
/>
<Typography.Title level={3}>Hyper Parameters</Typography.Title>
<Table
pagination={false}
columns={propertyTableColumns}
dataSource={model?.properties?.hyperParams as MlHyperParam[]}
/>
</Space>
<TabContent>
<Space direction="vertical" style={{ width: '100%' }} size="large">
<Typography.Title level={3}>Training Metrics</Typography.Title>
<Table
pagination={false}
columns={propertyTableColumns}
dataSource={model?.properties?.trainingMetrics as MlMetric[]}
/>
<Typography.Title level={3}>Hyper Parameters</Typography.Title>
<Table
pagination={false}
columns={propertyTableColumns}
dataSource={model?.properties?.hyperParams as MlHyperParam[]}
/>
</Space>
</TabContent>
);
}

18
datahub-web-react/src/app/entity/mlModel/profile/MlModelFeaturesTab.tsx Normal file
View File

@ -0,0 +1,18 @@
import React from 'react';
import { EntityType } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import { useBaseEntity } from '../../shared/EntityContext';
import { GetMlModelQuery } from '../../../../graphql/mlModel.generated';
import { EntityList } from '../../shared/tabs/Entity/components/EntityList';
export default function MlModelFeaturesTab() {
const entity = useBaseEntity() as GetMlModelQuery;
const entityRegistry = useEntityRegistry();
const model = entity && entity.mlModel;
const featureTables = model?.features?.relationships.map((relationship) => relationship.entity);
const title = `Part of ${entityRegistry.getEntityName(EntityType.MlfeatureTable)}`;
return <EntityList title={title} type={EntityType.MlfeatureTable} entities={featureTables || []} />;
}

57
datahub-web-react/src/app/entity/mlModelGroup/MLModelGroupEntity.tsx
View File

@ -1,10 +1,18 @@
import * as React from 'react';
import { CodeSandboxOutlined } from '@ant-design/icons';
import { MlModelGroup, EntityType, SearchResult } from '../../../types.generated';
import { MlModelGroup, EntityType, SearchResult, OwnershipType } from '../../../types.generated';
import { Preview } from './preview/Preview';
import { Entity, IconStyleType, PreviewType } from '../Entity';
import { MLModelGroupProfile } from './profile/MLModelGroupProfile';
import { getDataForEntityType } from '../shared/containers/profile/utils';
import { GenericEntityProperties } from '../shared/types';
import { EntityProfile } from '../shared/containers/profile/EntityProfile';
import { SidebarDomainSection } from '../shared/containers/profile/sidebar/Domain/SidebarDomainSection';
import { SidebarOwnerSection } from '../shared/containers/profile/sidebar/Ownership/SidebarOwnerSection';
import { SidebarAboutSection } from '../shared/containers/profile/sidebar/SidebarAboutSection';
import { SidebarTagsSection } from '../shared/containers/profile/sidebar/SidebarTagsSection';
import { useGetMlModelGroupQuery } from '../../../graphql/mlModelGroup.generated';
import ModelGroupModels from './profile/ModelGroupModels';
import { DocumentationTab } from '../shared/tabs/Documentation/DocumentationTab';
/**
* Definition of the DataHub MlModelGroup entity.
@ -45,7 +53,50 @@ export class MLModelGroupEntity implements Entity<MlModelGroup> {
getCollectionName = () => 'ML Groups';
renderProfile = (urn: string) => <MLModelGroupProfile urn={urn} />;
getOverridePropertiesFromEntity = (_?: MlModelGroup | null): GenericEntityProperties => {
return {};
};
renderProfile = (urn: string) => (
<EntityProfile
urn={urn}
key={urn}
entityType={EntityType.MlmodelGroup}
useEntityQuery={useGetMlModelGroupQuery}
getOverrideProperties={this.getOverridePropertiesFromEntity}
tabs={[
{
name: 'Models',
component: ModelGroupModels,
},
{
name: 'Documentation',
component: DocumentationTab,
},
]}
sidebarSections={[
{
component: SidebarAboutSection,
},
{
component: SidebarTagsSection,
properties: {
hasTags: true,
hasTerms: true,
},
},
{
component: SidebarOwnerSection,
properties: {
defaultOwnerType: OwnershipType.TechnicalOwner,
},
},
{
component: SidebarDomainSection,
},
]}
/>
);
renderPreview = (_: PreviewType, data: MlModelGroup) => {
return <Preview group={data} />;

26
datahub-web-react/src/app/entity/mlModelGroup/__tests__/MlModelGroupProfile.test.tsx
View File

@ -1,26 +0,0 @@
import React from 'react';
import { render, waitFor } from '@testing-library/react';
import { MockedProvider } from '@apollo/client/testing';
import { MLModelGroupProfile } from '../profile/MLModelGroupProfile';
import TestPageContainer from '../../../../utils/test-utils/TestPageContainer';
import { mocks } from '../../../../Mocks';
describe('MlModelGroupProfile', () => {
it('renders', async () => {
const { getByText, queryAllByText } = render(
<MockedProvider mocks={mocks} addTypename={false}>
<TestPageContainer
initialEntries={[
'/mlModelGroup/urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,another-group,PROD)',
]}
>
<MLModelGroupProfile urn="urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,another-group,PROD)" />
</TestPageContainer>
</MockedProvider>,
);
await waitFor(() => expect(queryAllByText('trust model group').length).toBeGreaterThanOrEqual(1));
expect(getByText('a ml trust model group')).toBeInTheDocument();
});
});

84
datahub-web-react/src/app/entity/mlModelGroup/profile/MLModelGroupProfile.tsx
View File

@ -1,84 +0,0 @@
import React from 'react';
import { Alert } from 'antd';
import { LegacyEntityProfile } from '../../../shared/LegacyEntityProfile';
import { EntityType, MlModelGroup } from '../../../../types.generated';
import { Message } from '../../../shared/Message';
import { Ownership as OwnershipView } from '../../shared/components/legacy/Ownership';
import { useEntityRegistry } from '../../../useEntityRegistry';
import analytics, { EventType } from '../../../analytics';
import { useGetMlModelGroupQuery } from '../../../../graphql/mlModelGroup.generated';
import ModelGroupHeader from './ModelGroupHeader';
import MLGroupModels from './ModelGroupModels';
export enum TabType {
Models = 'Models',
Ownership = 'Ownership',
}
/**
* Responsible for display the MLModel Page
*/
export const MLModelGroupProfile = ({ urn }: { urn: string }): JSX.Element => {
const entityRegistry = useEntityRegistry();
const { loading, error, data } = useGetMlModelGroupQuery({ variables: { urn } });
if (error || (!loading && !error && !data)) {
return <Alert type="error" message={error?.message || 'Entity failed to load'} />;
}
const getHeader = (group: MlModelGroup) => <ModelGroupHeader mlModelGroup={group} />;
const getTabs = (group: MlModelGroup) => {
return [
{
name: TabType.Models,
path: TabType.Models.toLowerCase(),
content: (
<MLGroupModels
// eslint-disable-next-line @typescript-eslint/dot-notation
models={group?.['incoming']?.relationships?.map((relationship) => relationship.entity) || []}
/>
),
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
{
name: TabType.Ownership,
path: TabType.Ownership.toLowerCase(),
content: (
<OwnershipView
owners={group?.ownership?.owners || []}
lastModifiedAt={group?.ownership?.lastModified?.time || 0}
/>
),
display: {
visible: (_, _1) => true,
enabled: (_, _1) => true,
},
},
];
};
return (
<>
{loading && <Message type="loading" content="Loading..." style={{ marginTop: '10%' }} />}
{data && data.mlModelGroup && (
<LegacyEntityProfile
titleLink={`/${entityRegistry.getPathName(EntityType.MlmodelGroup)}/${urn}`}
title={data.mlModelGroup?.name || ''}
tabs={getTabs(data.mlModelGroup as MlModelGroup)}
header={getHeader(data.mlModelGroup as MlModelGroup)}
onTabChange={(tab: string) => {
analytics.event({
type: EventType.EntitySectionViewEvent,
entityType: EntityType.Mlmodel,
entityUrn: urn,
section: tab,
});
}}
/>
)}
</>
);
};

65
datahub-web-react/src/app/entity/mlModelGroup/profile/ModelGroupHeader.tsx
View File

@ -1,65 +0,0 @@
import { Image, Row, Space, Typography } from 'antd';
import React from 'react';
import styled from 'styled-components';
import { MlModelGroup } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import CompactContext from '../../../shared/CompactContext';
import { AvatarsGroup } from '../../../shared/avatar';
import MarkdownViewer from '../../shared/components/legacy/MarkdownViewer';
const HeaderInfoItem = styled.div`
display: inline-block;
text-align: left;
width: 125px;
vertical-align: top;
`;
const PlatformName = styled(Typography.Text)`
font-size: 16px;
`;
const PreviewImage = styled(Image)`
max-height: 20px;
padding-top: 3px;
width: auto;
object-fit: contain;
`;
export type Props = {
mlModelGroup: MlModelGroup;
};
export default function ModelGroupHeader({ mlModelGroup: { description, ownership, platform } }: Props) {
const entityRegistry = useEntityRegistry();
const isCompact = React.useContext(CompactContext);
return (
<>
<Space direction="vertical" size="middle">
<Row justify="space-between">
{platform ? (
<HeaderInfoItem>
<div>
<Typography.Text strong type="secondary" style={{ fontSize: 11 }}>
Platform
</Typography.Text>
</div>
<Space direction="horizontal">
{platform.properties?.logoUrl ? (
<PreviewImage
preview={false}
src={platform.properties?.logoUrl}
placeholder
alt={platform.name}
/>
) : null}
<PlatformName>{platform.name}</PlatformName>
</Space>
</HeaderInfoItem>
) : null}
</Row>
<MarkdownViewer isCompact={isCompact} source={description || ''} />
<AvatarsGroup owners={ownership?.owners} entityRegistry={entityRegistry} />
</Space>
</>
);
}

34
datahub-web-react/src/app/entity/mlModelGroup/profile/ModelGroupModels.tsx
View File

@ -1,44 +1,28 @@
import { Button, List, Space, Typography } from 'antd';
import { List, Space, Typography } from 'antd';
import React from 'react';
import { useHistory, useLocation } from 'react-router-dom';
import styled from 'styled-components';
import { Entity, EntityType } from '../../../../types.generated';
import { navigateToLineageUrl } from '../../../lineage/utils/navigateToLineageUrl';
import { GetMlModelGroupQuery } from '../../../../graphql/mlModelGroup.generated';
import { EntityType } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import { PreviewType } from '../../Entity';
import { useBaseEntity } from '../../shared/EntityContext';
export type Props = {
models?: Entity[];
};
export default function MLGroupModels() {
const baseEntity = useBaseEntity<GetMlModelGroupQuery>();
const models = baseEntity?.mlModelGroup?.incoming?.relationships?.map((relationship) => relationship.entity) || [];
const ViewRawButtonContainer = styled.div`
display: flex;
justify-content: flex-end;
`;
export default function MLGroupModels({ models }: Props) {
const entityRegistry = useEntityRegistry();
const history = useHistory();
const location = useLocation();
return (
<>
<div>
<ViewRawButtonContainer>
<Button onClick={() => navigateToLineageUrl({ location, history, isLineageMode: true })}>
View Graph
</Button>
</ViewRawButtonContainer>
</div>
<Space direction="vertical" style={{ width: '100%' }} size="large">
<List
style={{ marginTop: '24px', padding: '16px 32px' }}
style={{ padding: '16px 16px' }}
bordered
dataSource={models}
header={<Typography.Title level={3}>Models</Typography.Title>}
renderItem={(item) => (
<List.Item style={{ paddingTop: '20px' }}>
{entityRegistry.renderPreview(item?.type || EntityType.Mlmodel, PreviewType.PREVIEW, item)}
{entityRegistry.renderPreview(EntityType.Mlmodel, PreviewType.PREVIEW, item)}
</List.Item>
)}
/>

69
datahub-web-react/src/app/entity/mlPrimaryKey/MLPrimaryKeyEntity.tsx
View File

@ -1,10 +1,18 @@
import * as React from 'react';
import { DotChartOutlined } from '@ant-design/icons';
import { MlPrimaryKey, EntityType, SearchResult } from '../../../types.generated';
import { MlPrimaryKey, EntityType, SearchResult, OwnershipType } from '../../../types.generated';
import { Preview } from './preview/Preview';
import { MLPrimaryKeyProfile } from './profile/MLPrimaryKeyProfile';
import { Entity, IconStyleType, PreviewType } from '../Entity';
import { getDataForEntityType } from '../shared/containers/profile/utils';
import { GenericEntityProperties } from '../shared/types';
import { useGetMlPrimaryKeyQuery } from '../../../graphql/mlPrimaryKey.generated';
import { EntityProfile } from '../shared/containers/profile/EntityProfile';
import { FeatureTableTab } from '../shared/tabs/ML/MlPrimaryKeyFeatureTableTab';
import { DocumentationTab } from '../shared/tabs/Documentation/DocumentationTab';
import { SidebarTagsSection } from '../shared/containers/profile/sidebar/SidebarTagsSection';
import { SidebarAboutSection } from '../shared/containers/profile/sidebar/SidebarAboutSection';
import { SidebarDomainSection } from '../shared/containers/profile/sidebar/Domain/SidebarDomainSection';
import { SidebarOwnerSection } from '../shared/containers/profile/sidebar/Ownership/SidebarOwnerSection';
/**
* Definition of the DataHub MLPrimaryKey entity.
@ -45,7 +53,50 @@ export class MLPrimaryKeyEntity implements Entity<MlPrimaryKey> {
getCollectionName = () => 'ML Primary Keys';
renderProfile = (urn: string) => <MLPrimaryKeyProfile urn={urn} />;
getOverridePropertiesFromEntity = (_?: MlPrimaryKey | null): GenericEntityProperties => {
return {};
};
renderProfile = (urn: string) => (
<EntityProfile
urn={urn}
key={urn}
entityType={EntityType.MlprimaryKey}
useEntityQuery={useGetMlPrimaryKeyQuery}
getOverrideProperties={this.getOverridePropertiesFromEntity}
tabs={[
{
name: 'Feature Tables',
component: FeatureTableTab,
},
{
name: 'Documentation',
component: DocumentationTab,
},
]}
sidebarSections={[
{
component: SidebarAboutSection,
},
{
component: SidebarTagsSection,
properties: {
hasTags: true,
hasTerms: true,
},
},
{
component: SidebarOwnerSection,
properties: {
defaultOwnerType: OwnershipType.TechnicalOwner,
},
},
{
component: SidebarDomainSection,
},
]}
/>
);
renderPreview = (_: PreviewType, data: MlPrimaryKey) => {
return (
@ -83,4 +134,16 @@ export class MLPrimaryKeyEntity implements Entity<MlPrimaryKey> {
getOverrideProperties: (data) => data,
});
};
getLineageVizConfig = (entity: MlPrimaryKey) => {
return {
urn: entity.urn,
name: entity.name,
type: EntityType.MlprimaryKey,
// eslint-disable-next-line
icon: entity?.['featureTables']?.relationships?.[0]?.entity?.platform?.properties?.logoUrl || undefined,
// eslint-disable-next-line
platform: entity?.['featureTables']?.relationships?.[0]?.entity?.platform?.name,
};
};
}

34
datahub-web-react/src/app/entity/mlPrimaryKey/profile/MLPrimaryKeyProfile.tsx
View File

@ -1,34 +0,0 @@
import React from 'react';
import { LegacyEntityProfile } from '../../../shared/LegacyEntityProfile';
import { EntityType } from '../../../../types.generated';
import { useEntityRegistry } from '../../../useEntityRegistry';
import analytics, { EventType } from '../../../analytics';
export enum TabType {
Features = 'Features',
Sources = 'Sources',
Ownership = 'Ownership',
}
/**
* Responsible for display the MLPrimaryKey Page
*/
export const MLPrimaryKeyProfile = ({ urn }: { urn: string }): JSX.Element => {
const entityRegistry = useEntityRegistry();
return (
<LegacyEntityProfile
titleLink={`/${entityRegistry.getPathName(EntityType.MlprimaryKey)}/${urn}`}
title={urn}
header={<></>}
onTabChange={(tab: string) => {
analytics.event({
type: EventType.EntitySectionViewEvent,
entityType: EntityType.MlprimaryKey,
entityUrn: urn,
section: tab,
});
}}
/>
);
};

42
datahub-web-react/src/app/entity/shared/tabs/Dataset/Schema/utils/useTagsAndTermsRendererFeatureTable.tsx Normal file
View File

@ -0,0 +1,42 @@
import React from 'react';
import { EditableSchemaMetadata, EntityType, GlobalTags, SchemaField } from '../../../../../../../types.generated';
import TagTermGroup from '../../../../../../shared/tags/TagTermGroup';
import { pathMatchesNewPath } from '../../../../../dataset/profile/schema/utils/utils';
import { useEntityData, useRefetch } from '../../../../EntityContext';
export default function useTagsAndTermsRendererFeatureTable(
editableSchemaMetadata: EditableSchemaMetadata | null | undefined,
tagHoveredIndex: string | undefined,
setTagHoveredIndex: (index: string | undefined) => void,
options: { showTags: boolean; showTerms: boolean },
) {
const { urn } = useEntityData();
const refetch = useRefetch();
const tagAndTermRender = (tags: GlobalTags, record: SchemaField, rowIndex: number | undefined) => {
const relevantEditableFieldInfo = editableSchemaMetadata?.editableSchemaFieldInfo.find(
(candidateEditableFieldInfo) => pathMatchesNewPath(candidateEditableFieldInfo.fieldPath, record.fieldPath),
);
return (
<div data-testid={`schema-field-${record.fieldPath}-${options.showTags ? 'tags' : 'terms'}`}>
<TagTermGroup
uneditableTags={options.showTags ? tags : null}
editableTags={options.showTags ? relevantEditableFieldInfo?.globalTags : null}
uneditableGlossaryTerms={options.showTerms ? record.glossaryTerms : null}
editableGlossaryTerms={options.showTerms ? relevantEditableFieldInfo?.glossaryTerms : null}
canRemove
buttonProps={{ size: 'small' }}
canAddTag={tagHoveredIndex === `${record.fieldPath}-${rowIndex}` && options.showTags}
canAddTerm={tagHoveredIndex === `${record.fieldPath}-${rowIndex}` && options.showTerms}
onOpenModal={() => setTagHoveredIndex(undefined)}
entityUrn={urn}
entityType={EntityType.Dataset}
entitySubresource={record.fieldPath}
refetch={refetch}
/>
</div>
);
};
return tagAndTermRender;
}

17
datahub-web-react/src/app/entity/shared/tabs/ML/MlFeatureFeatureTableTab.tsx Normal file
View File

@ -0,0 +1,17 @@
import React from 'react';
import { useBaseEntity } from '../../EntityContext';
import { EntityType } from '../../../../../types.generated';
import { EntityList } from '../Entity/components/EntityList';
import { useEntityRegistry } from '../../../../useEntityRegistry';
import { GetMlFeatureQuery } from '../../../../../graphql/mlFeature.generated';
export const FeatureTableTab = () => {
const entity = useBaseEntity() as GetMlFeatureQuery;
const entityRegistry = useEntityRegistry();
const feature = entity && entity.mlFeature;
const featureTables = feature?.featureTables?.relationships.map((relationship) => relationship.entity);
const title = `Part of ${entityRegistry.getEntityName(EntityType.MlfeatureTable)}`;
return <EntityList title={title} type={EntityType.MlfeatureTable} entities={featureTables || []} />;
};

17
datahub-web-react/src/app/entity/shared/tabs/ML/MlPrimaryKeyFeatureTableTab.tsx Normal file
View File

@ -0,0 +1,17 @@
import React from 'react';
import { useBaseEntity } from '../../EntityContext';
import { EntityType } from '../../../../../types.generated';
import { EntityList } from '../Entity/components/EntityList';
import { useEntityRegistry } from '../../../../useEntityRegistry';
import { GetMlPrimaryKeyQuery } from '../../../../../graphql/mlPrimaryKey.generated';
export const FeatureTableTab = () => {
const entity = useBaseEntity() as GetMlPrimaryKeyQuery;
const entityRegistry = useEntityRegistry();
const feature = entity && entity.mlPrimaryKey;
const featureTables = feature?.featureTables?.relationships.map((relationship) => relationship.entity);
const title = `Part of ${entityRegistry.getEntityName(EntityType.MlfeatureTable)}`;
return <EntityList title={title} type={EntityType.MlfeatureTable} entities={featureTables || []} />;
};

8
datahub-web-react/src/app/lineage/utils/useGetEntityQuery.ts
View File

@ -157,6 +157,10 @@ export default function useGetEntityQuery(urn: string, entityType?: EntityType)
allResults[EntityType.Mlmodel],
// eslint-disable-next-line react-hooks/exhaustive-deps
allResults[EntityType.MlmodelGroup],
// eslint-disable-next-line react-hooks/exhaustive-deps
allResults[EntityType.Mlfeature],
// eslint-disable-next-line react-hooks/exhaustive-deps
allResults[EntityType.MlprimaryKey],
]);
const returnObject = useMemo(() => {
@ -192,6 +196,10 @@ export default function useGetEntityQuery(urn: string, entityType?: EntityType)
allResults[EntityType.Mlmodel],
// eslint-disable-next-line react-hooks/exhaustive-deps
allResults[EntityType.MlmodelGroup],
// eslint-disable-next-line react-hooks/exhaustive-deps
allResults[EntityType.Mlfeature],
// eslint-disable-next-line react-hooks/exhaustive-deps
allResults[EntityType.MlprimaryKey],
]);
return returnObject;

4
datahub-web-react/src/app/lineage/utils/useLazyGetEntityQuery.ts
View File

@ -45,11 +45,11 @@ export default function useLazyGetEntityQuery() {
setFetchedEntityType(type);
getAsyncMLFeatureTable({ variables: { urn } });
}
if (type === EntityType.MlfeatureTable) {
if (type === EntityType.Mlfeature) {
setFetchedEntityType(type);
getAsyncMLFeature({ variables: { urn } });
}
if (type === EntityType.MlfeatureTable) {
if (type === EntityType.MlprimaryKey) {
setFetchedEntityType(type);
getAsyncMLPrimaryKey({ variables: { urn } });
}

69
datahub-web-react/src/graphql/fragments.graphql
View File

@ -361,7 +361,7 @@ fragment nonRecursiveMLFeature on MLFeature {
featureNamespace
description
dataType
featureProperties {
properties {
description
dataType
version {
@ -389,6 +389,18 @@ fragment nonRecursiveMLFeature on MLFeature {
status {
removed
}
glossaryTerms {
...glossaryTerms
}
domain {
...entityDomain
}
tags {
...globalTagsFields
}
editableProperties {
description
}
deprecation {
...deprecationFields
}
@ -401,7 +413,7 @@ fragment nonRecursiveMLPrimaryKey on MLPrimaryKey {
featureNamespace
description
dataType
primaryKeyProperties {
properties {
description
dataType
version {
@ -429,6 +441,18 @@ fragment nonRecursiveMLPrimaryKey on MLPrimaryKey {
status {
removed
}
glossaryTerms {
...glossaryTerms
}
domain {
...entityDomain
}
tags {
...globalTagsFields
}
editableProperties {
description
}
deprecation {
...deprecationFields
}
@ -442,7 +466,7 @@ fragment nonRecursiveMLFeatureTable on MLFeatureTable {
...platformFields
}
description
featureTableProperties {
properties {
description
mlFeatures {
...nonRecursiveMLFeature
@ -460,6 +484,18 @@ fragment nonRecursiveMLFeatureTable on MLFeatureTable {
status {
removed
}
glossaryTerms {
...glossaryTerms
}
domain {
...entityDomain
}
tags {
...globalTagsFields
}
editableProperties {
description
}
deprecation {
...deprecationFields
}
@ -578,6 +614,18 @@ fragment nonRecursiveMLModel on MLModel {
status {
removed
}
glossaryTerms {
...glossaryTerms
}
domain {
...entityDomain
}
tags {
...globalTagsFields
}
editableProperties {
description
}
deprecation {
...deprecationFields
}
@ -598,9 +646,24 @@ fragment nonRecursiveMLModelGroupFields on MLModelGroup {
status {
removed
}
glossaryTerms {
...glossaryTerms
}
domain {
...entityDomain
}
tags {
...globalTagsFields
}
editableProperties {
description
}
deprecation {
...deprecationFields
}
properties {
description
}
}
fragment platformFields on DataPlatform {

29
datahub-web-react/src/graphql/lineage.graphql
View File

@ -130,6 +130,35 @@ fragment relationshipFields on EntityWithRelationships {
removed
}
}
... on MLFeatureTable {
...nonRecursiveMLFeatureTable
}
... on MLFeature {
...nonRecursiveMLFeature
featureTables: relationships(input: { types: ["Contains"], direction: INCOMING, start: 0, count: 100 }) {
relationships {
type
entity {
... on MLFeatureTable {
...nonRecursiveMLFeatureTable
}
}
}
}
}
... on MLPrimaryKey {
...nonRecursiveMLPrimaryKey
featureTables: relationships(input: { types: ["KeyedBy"], direction: INCOMING, start: 0, count: 100 }) {
relationships {
type
entity {
... on MLFeatureTable {
...nonRecursiveMLFeatureTable
}
}
}
}
}
upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
...leafLineageResults
}

9
datahub-web-react/src/graphql/mlFeature.graphql
View File

@ -1,5 +1,14 @@
query getMLFeature($urn: String!) {
mlFeature(urn: $urn) {
...nonRecursiveMLFeature
upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
featureTables: relationships(input: { types: ["Contains"], direction: INCOMING, start: 0, count: 100 }) {
...fullRelationshipResults
}
}
}

6
datahub-web-react/src/graphql/mlFeatureTable.graphql
View File

@ -1,5 +1,11 @@
query getMLFeatureTable($urn: String!) {
mlFeatureTable(urn: $urn) {
...nonRecursiveMLFeatureTable
upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
}
}

23
datahub-web-react/src/graphql/mlModel.graphql
View File

@ -1,24 +1,13 @@
query getMLModel($urn: String!) {
mlModel(urn: $urn) {
...nonRecursiveMLModel
incoming: relationships(
input: {
types: ["DownstreamOf", "Consumes", "Produces", "TrainedBy", "MemberOf"]
direction: INCOMING
start: 0
count: 100
}
) {
...fullRelationshipResults
upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
outgoing: relationships(
input: {
types: ["DownstreamOf", "Consumes", "Produces", "TrainedBy", "MemberOf"]
direction: OUTGOING
start: 0
count: 100
}
) {
downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
features: relationships(input: { types: ["Consumes"], direction: OUTGOING, start: 0, count: 100 }) {
...fullRelationshipResults
}
}

6
datahub-web-react/src/graphql/mlModelGroup.graphql
View File

@ -21,5 +21,11 @@ query getMLModelGroup($urn: String!) {
) {
...fullRelationshipResults
}
upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
}
}

9
datahub-web-react/src/graphql/mlPrimaryKey.graphql
View File

@ -1,5 +1,14 @@
query getMLPrimaryKey($urn: String!) {
mlPrimaryKey(urn: $urn) {
...nonRecursiveMLPrimaryKey
upstream: lineage(input: { direction: UPSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
downstream: lineage(input: { direction: DOWNSTREAM, start: 0, count: 100 }) {
...fullLineageResults
}
featureTables: relationships(input: { types: ["KeyedBy"], direction: INCOMING, start: 0, count: 100 }) {
...fullRelationshipResults
}
}
}

15
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/EditableMLFeatureProperties.pdl Normal file
View File

@ -0,0 +1,15 @@
namespace com.linkedin.ml.metadata
/**
* Properties associated with a MLFeature editable from the UI
*/
@Aspect = {
"name": "editableMlFeatureProperties"
}
record EditableMLFeatureProperties {
/**
* Documentation of the MLFeature
*/
description: optional string
}

15
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/EditableMLFeatureTableProperties.pdl Normal file
View File

@ -0,0 +1,15 @@
namespace com.linkedin.ml.metadata
/**
* Properties associated with a MLFeatureTable editable from the ui
*/
@Aspect = {
"name": "editableMlFeatureTableProperties"
}
record EditableMLFeatureTableProperties {
/**
* Documentation of the MLFeatureTable
*/
description: optional string
}

18
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/EditableMLModelGroupProperties.pdl Normal file
View File

@ -0,0 +1,18 @@
namespace com.linkedin.ml.metadata
/**
* Properties associated with an ML Model Group editable from the UI
*/
@Aspect = {
"name": "editableMlModelGroupProperties"
}
record EditableMLModelGroupProperties {
/**
* Documentation of the ml model group
*/
@Searchable = {
"fieldType": "TEXT",
"fieldName": "editedDescription",
}
description: optional string
}

18
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/EditableMLModelProperties.pdl Normal file
View File

@ -0,0 +1,18 @@
namespace com.linkedin.ml.metadata
/**
* Properties associated with a ML Model editable from the UI
*/
@Aspect = {
"name": "editableMlModelProperties"
}
record EditableMLModelProperties {
/**
* Documentation of the ml model
*/
@Searchable = {
"fieldType": "TEXT",
"fieldName": "editedDescription",
}
description: optional string
}

15
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/EditableMLPrimaryKeyProperties.pdl Normal file
View File

@ -0,0 +1,15 @@
namespace com.linkedin.ml.metadata
/**
* Properties associated with a MLPrimaryKey editable from the UI
*/
@Aspect = {
"name": "editableMlPrimaryKeyProperties"
}
record EditableMLPrimaryKeyProperties {
/**
* Documentation of the MLPrimaryKey
*/
description: optional string
}

3
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/MLFeatureProperties.pdl
View File

@ -33,7 +33,8 @@ record MLFeatureProperties {
@Relationship = {
"/*": {
"name": "DerivedFrom",
"entityTypes": [ "dataset" ]
"entityTypes": [ "dataset" ],
"isLineage": true
}
}
sources: optional array[Urn]

5
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/MLFeatureTableProperties.pdl
View File

@ -22,8 +22,7 @@ record MLFeatureTableProperties includes CustomProperties {
@Relationship = {
"/*": {
"name": "Contains",
"entityTypes": [ "mlFeature" ],
"isLineage": true
"entityTypes": [ "mlFeature" ]
}
}
@Searchable = {
@ -50,4 +49,4 @@ record MLFeatureTableProperties includes CustomProperties {
}
}
mlPrimaryKeys: optional array[Urn]
}
}

7
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/MLModelProperties.pdl
View File

@ -67,6 +67,13 @@ record MLModelProperties includes CustomProperties, ExternalReference {
/**
* List of features used for MLModel training
*/
@Relationship = {
"/*": {
"name": "Consumes",
"entityTypes": [ "mlFeature" ],
"isLineage": true
}
}
mlFeatures: optional array[MLFeatureUrn]
/**

3
metadata-models/src/main/pegasus/com/linkedin/ml/metadata/MLPrimaryKeyProperties.pdl
View File

@ -33,7 +33,8 @@ record MLPrimaryKeyProperties {
@Relationship = {
"/*": {
"name": "DerivedFrom",
"entityTypes": [ "dataset" ]
"entityTypes": [ "dataset" ],
"isLineage": true
}
}
sources: array[Urn]

35
metadata-models/src/main/resources/entity-registry.yml
View File

@ -174,4 +174,39 @@ entities:
keyAspect: dataPlatformInstanceKey
aspects:
- status
- name: mlModel
category: core
keyAspect: mlModelKey
aspects:
- glossaryTerms
- editableMlModelProperties
- domains
- name: mlModelGroup
category: core
keyAspect: mlModelGroupKey
aspects:
- glossaryTerms
- editableMlModelGroupProperties
- domains
- name: mlFeatureTable
category: core
keyAspect: mlFeatureTableKey
aspects:
- glossaryTerms
- editableMlFeatureTableProperties
- domains
- name: mlFeature
category: core
keyAspect: mlFeatureKey
aspects:
- glossaryTerms
- editableMlFeatureProperties
- domains
- name: mlPrimaryKey
category: core
keyAspect: mlPrimaryKeyKey
aspects:
- glossaryTerms
- editableMlPrimaryKeyProperties
- domains
events:

5
metadata-utils/src/main/java/com/linkedin/metadata/Constants.java
View File

@ -117,14 +117,17 @@ public class Constants {
// ML Feature
public static final String ML_FEATURE_KEY_ASPECT_NAME = "mlFeatureKey";
public static final String ML_FEATURE_PROPERTIES_ASPECT_NAME = "mlFeatureProperties";
public static final String ML_FEATURE_EDITABLE_PROPERTIES_ASPECT_NAME = "editableMlFeatureProperties";
// ML Feature Table
public static final String ML_FEATURE_TABLE_KEY_ASPECT_NAME = "mlFeatureTableKey";
public static final String ML_FEATURE_TABLE_PROPERTIES_ASPECT_NAME = "mlFeatureTableProperties";
public static final String ML_FEATURE_TABLE_EDITABLE_PROPERTIES_ASPECT_NAME = "editableMlFeatureTableProperties";
//ML Model
public static final String ML_MODEL_KEY_ASPECT_NAME = "mlModelKey";
public static final String ML_MODEL_PROPERTIES_ASPECT_NAME = "mlModelProperties";
public static final String ML_MODEL_EDITABLE_PROPERTIES_ASPECT_NAME = "editableMlModelProperties";
public static final String INTENDED_USE_ASPECT_NAME = "intendedUse";
public static final String ML_MODEL_FACTOR_PROMPTS_ASPECT_NAME = "mlModelFactorPrompts";
public static final String METRICS_ASPECT_NAME = "metrics";
@ -139,10 +142,12 @@ public class Constants {
// ML Model Group
public static final String ML_MODEL_GROUP_KEY_ASPECT_NAME = "mlModelGroupKey";
public static final String ML_MODEL_GROUP_PROPERTIES_ASPECT_NAME = "mlModelGroupProperties";
public static final String ML_MODEL_GROUP_EDITABLE_PROPERTIES_ASPECT_NAME = "editableMlModelGroupProperties";
// ML Primary Key
public static final String ML_PRIMARY_KEY_KEY_ASPECT_NAME = "mlPrimaryKeyKey";
public static final String ML_PRIMARY_KEY_PROPERTIES_ASPECT_NAME = "mlPrimaryKeyProperties";
public static final String ML_PRIMARY_KEY_EDITABLE_PROPERTIES_ASPECT_NAME = "editableMlPrimaryKeyProperties";
// Policy
public static final String DATAHUB_POLICY_INFO_ASPECT_NAME = "dataHubPolicyInfo";

46
smoke-test/tests/cypress/cypress/integration/ml/feature_table.js Normal file
View File

@ -0,0 +1,46 @@
describe('features', () => {
it('can visit feature tables and see features', () => {
cy.visit('/')
cy.login();
cy.visit('/featureTables/urn:li:mlFeatureTable:(urn:li:dataPlatform:sagemaker,cypress-feature-table)/Features?is_lineage_mode=false');
// the feature table descriptions should be there
cy.contains('Yet another test feature group');
cy.contains('this is a description from source system');
// additional properties are visible
cy.contains('CypressPrimaryKeyTag');
cy.contains('CypressFeatureTag');
// navigate to sources
cy.contains('Sources').click();
// feature & primary key sources are visible
cy.contains('SampleCypressHdfsDataset');
cy.contains('SampleCypressKafkaDataset');
});
it('can visit feature page', () => {
cy.visit('/')
cy.login();
cy.visit('/features/urn:li:mlFeature:(cypress-test-2,some-cypress-feature-1)/Feature%20Tables?is_lineage_mode=false');
// Shows the parent table
cy.contains('cypress-feature-table');
// Has upstream & downstream lineage
cy.contains('1 upstream, 1 downstream');
});
it('can visit primary key page', () => {
cy.visit('/')
cy.login();
cy.visit('/mlPrimaryKeys/urn:li:mlPrimaryKey:(cypress-test-2,some-cypress-feature-2)/Feature%20Tables?is_lineage_mode=false');
// Shows the parent table
cy.contains('cypress-feature-table');
// Has upstream from its sources
cy.contains('1 upstream, 0 downstream');
});
})

30
smoke-test/tests/cypress/cypress/integration/ml/model.js Normal file
View File

@ -0,0 +1,30 @@
describe('models', () => {
it('can visit models and groups', () => {
cy.visit('/')
cy.login();
cy.visit('/mlModels/urn:li:mlModel:(urn:li:dataPlatform:sagemaker,cypress-model,PROD)/Summary?is_lineage_mode=false');
cy.contains('ml model description');
// the model has metrics & hyper params
cy.contains('another-metric');
cy.contains('parameter-1');
// the model has features
cy.contains('Features').click();
cy.contains('some-cypress-feature-1');
// the model has a group
cy.visit('/mlModels/urn:li:mlModel:(urn:li:dataPlatform:sagemaker,cypress-model,PROD)/Group?is_lineage_mode=false');
cy.contains('cypress-model-package-group');
});
it('can visit models and groups', () => {
cy.visit('/')
cy.login();
cy.visit('/mlModelGroup/urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,cypress-model-package-group,PROD)');
// the model group has its model
cy.contains('cypress-model');
cy.contains('Just a model package group.');
});
})

190
smoke-test/tests/cypress/data.json
View File

@ -1630,5 +1630,195 @@
"contentType": "application/json"
},
"systemMetadata": null
},
{
"auditHeader": null,
"proposedSnapshot": {
"com.linkedin.pegasus2avro.metadata.snapshot.MLFeatureSnapshot": {
"urn": "urn:li:mlFeature:(cypress-test-2,some-cypress-feature-1)",
"aspects": [
{
"com.linkedin.pegasus2avro.ml.metadata.MLFeatureProperties": {
"description": null,
"dataType": "TEXT",
"version": null,
"sources": ["urn:li:dataset:(urn:li:dataPlatform:kafka,SampleCypressKafkaDataset,PROD)"]
}
},
{
"com.linkedin.pegasus2avro.common.GlobalTags": {
"tags": [{ "tag": "urn:li:tag:CypressFeatureTag" }]
}
}
]
}
},
"proposedDelta": null,
"systemMetadata": null
},
{
"auditHeader": null,
"proposedSnapshot": {
"com.linkedin.pegasus2avro.metadata.snapshot.MLPrimaryKeySnapshot": {
"urn": "urn:li:mlPrimaryKey:(cypress-test-2,some-cypress-feature-2)",
"aspects": [
{
"com.linkedin.pegasus2avro.ml.metadata.MLPrimaryKeyProperties": {
"description": "this is a description from source system",
"dataType": "ORDINAL",
"version": null,
"sources": ["urn:li:dataset:(urn:li:dataPlatform:hdfs,SampleCypressHdfsDataset,PROD)"]
}
},
{
"com.linkedin.pegasus2avro.common.GlobalTags": {
"tags": [{ "tag": "urn:li:tag:CypressPrimaryKeyTag" }]
}
}
]
}
},
"proposedDelta": null,
"systemMetadata": null
},
{
"auditHeader": null,
"proposedSnapshot": {
"com.linkedin.pegasus2avro.metadata.snapshot.MLFeatureTableSnapshot": {
"urn": "urn:li:mlFeatureTable:(urn:li:dataPlatform:sagemaker,cypress-feature-table)",
"aspects": [
{
"com.linkedin.pegasus2avro.common.BrowsePaths": {
"paths": [
"/sagemaker/cypress-feature-table"
]
}
},
{
"com.linkedin.pegasus2avro.ml.metadata.MLFeatureTableProperties": {
"customProperties": {
"arn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-2",
"creation_time": "2021-06-24 09:48:37.035000",
"status": "Created"
},
"description": "Yet another test feature group",
"mlFeatures": [
"urn:li:mlFeature:(cypress-test-2,some-cypress-feature-1)"
],
"mlPrimaryKeys": [
"urn:li:mlPrimaryKey:(cypress-test-2,some-cypress-feature-2)"
]
}
}
]
}
},
"proposedDelta": null,
"systemMetadata": null
},
{
"auditHeader": null,
"proposedSnapshot": {
"com.linkedin.pegasus2avro.metadata.snapshot.MLModelSnapshot": {
"urn": "urn:li:mlModel:(urn:li:dataPlatform:sagemaker,cypress-model,PROD)",
"aspects": [
{
"com.linkedin.pegasus2avro.ml.metadata.MLModelProperties": {
"customProperties": {
"EnableNetworkIsolation": "True"
},
"externalUrl": "https://us-west-2.console.aws.amazon.com/sagemaker/home?region=us-west-2#/models/the-first-model",
"description": "ml model description",
"date": 1420070400000,
"version": null,
"type": null,
"hyperParameters": null,
"hyperParams": [
{
"name": "parameter-1",
"description": null,
"value": "some-value",
"createdAt": null
}
],
"trainingMetrics": [
{
"name": "another-metric",
"description": null,
"value": "1.0",
"createdAt": null
}
],
"onlineMetrics": null,
"mlFeatures": ["urn:li:mlFeature:(cypress-test-2,some-cypress-feature-1)"],
"tags": [],
"deployments": [],
"trainingJobs": [],
"downstreamJobs": [],
"groups": [
"urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,cypress-model-package-group,PROD)"
]
}
},
{
"com.linkedin.pegasus2avro.common.BrowsePaths": {
"paths": [
"/sagemaker/cypress-model-package-group/cypress-model"
]
}
}
]
}
},
"proposedDelta": null,
"systemMetadata": null
},
{
"auditHeader": null,
"proposedSnapshot": {
"com.linkedin.pegasus2avro.metadata.snapshot.MLModelGroupSnapshot": {
"urn": "urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,cypress-model-package-group,PROD)",
"aspects": [
{
"com.linkedin.pegasus2avro.ml.metadata.MLModelGroupProperties": {
"customProperties": {
"ModelPackageGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:model-package-group/a-model-package-group",
"ModelPackageGroupDescription": "Just a model package group.",
"CreatedBy": "{'UserProfileArn': 'arn:aws:sagemaker:us-west-2:123412341234:user-profile/some-domain/some-user', 'UserProfileName': 'some-user', 'DomainId': 'some-domain'}",
"ModelPackageGroupStatus": "Completed"
},
"description": "Just a model package group.",
"createdAt": 1420070400000,
"version": null
}
},
{
"com.linkedin.pegasus2avro.common.Ownership": {
"owners": [
{
"owner": "urn:li:corpuser:some-user",
"type": "DATAOWNER",
"source": null
}
],
"lastModified": {
"time": 0,
"actor": "urn:li:corpuser:unknown",
"impersonator": null
}
}
},
{
"com.linkedin.pegasus2avro.common.BrowsePaths": {
"paths": [
"/sagemaker/cypress-model-package-group"
]
}
}
]
}
},
"proposedDelta": null,
"systemMetadata": null
}
]