feat: update mlflow UI

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchUtils.java

@@ -75,7 +75,8 @@ private SearchUtils() {}
           EntityType.DATA_PRODUCT,
           EntityType.NOTEBOOK,
           EntityType.BUSINESS_ATTRIBUTE,
-          EntityType.SCHEMA_FIELD);
+          EntityType.SCHEMA_FIELD,
+          EntityType.DATA_PROCESS_INSTANCE);
 
   /** Entities that are part of autocomplete by default in Auto Complete Across Entities */
   public static final List<EntityType> AUTO_COMPLETE_ENTITY_TYPES =

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/DataPlatformInstanceAspectMapper.java

@@ -1,6 +1,7 @@
 package com.linkedin.datahub.graphql.types.common.mappers;
 
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.DataPlatform;
 import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
 import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
@@ -28,6 +29,11 @@ public DataPlatformInstance apply(
       result.setType(EntityType.DATA_PLATFORM_INSTANCE);
       result.setUrn(input.getInstance().toString());
     }
+    result.setPlatform(
+      DataPlatform.builder()
+          .setUrn(input.getPlatform().toString())
+          .setType(EntityType.DATA_PLATFORM)
+          .build());
     return result;
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/TimeStampMapper.java

@@ -0,0 +1,28 @@
+package com.linkedin.datahub.graphql.types.common.mappers;
+
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.TimeStamp;
+import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+public class TimeStampMapper implements ModelMapper<com.linkedin.common.TimeStamp, TimeStamp> {
+
+  public static final TimeStampMapper INSTANCE = new TimeStampMapper();
+
+  public static TimeStamp map(
+      @Nullable QueryContext context, @Nonnull final com.linkedin.common.TimeStamp timestamp) {
+    return INSTANCE.apply(context, timestamp);
+  }
+
+  @Override
+  public TimeStamp apply(
+      @Nullable QueryContext context, @Nonnull final com.linkedin.common.TimeStamp timestamp) {
+    final TimeStamp result = new TimeStamp();
+    result.setTime(timestamp.getTime());
+    if (timestamp.hasActor()) {
+      result.setActor(timestamp.getActor().toString());
+    }
+    return result;
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/UrnToEntityMapper.java

@@ -18,6 +18,7 @@
 import com.linkedin.datahub.graphql.generated.DataJob;
 import com.linkedin.datahub.graphql.generated.DataPlatform;
 import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
+import com.linkedin.datahub.graphql.generated.DataProcessInstance;
 import com.linkedin.datahub.graphql.generated.DataProduct;
 import com.linkedin.datahub.graphql.generated.Dataset;
 import com.linkedin.datahub.graphql.generated.Domain;
@@ -225,6 +226,11 @@ public Entity apply(@Nullable QueryContext context, Urn input) {
       ((BusinessAttribute) partialEntity).setUrn(input.toString());
       ((BusinessAttribute) partialEntity).setType(EntityType.BUSINESS_ATTRIBUTE);
     }
+    if (input.getEntityType().equals(DATA_PROCESS_INSTANCE_ENTITY_NAME)) {
+      partialEntity = new DataProcessInstance();
+      ((DataProcessInstance) partialEntity).setUrn(input.toString());
+      ((DataProcessInstance) partialEntity).setType(EntityType.DATA_PROCESS_INSTANCE);
+    }
     return partialEntity;
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinst/mappers/DataProcessInstanceMapper.java

@@ -2,25 +2,33 @@
 
 import static com.linkedin.metadata.Constants.*;
 
+import com.linkedin.common.DataPlatformInstance;
+import com.linkedin.common.SubTypes;
+import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.DataPlatform;
 import com.linkedin.datahub.graphql.generated.DataProcessInstance;
 import com.linkedin.datahub.graphql.generated.EntityType;
-import com.linkedin.datahub.graphql.types.common.mappers.AuditStampMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.DataPlatformInstanceAspectMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.SubTypesMapper;
 import com.linkedin.datahub.graphql.types.common.mappers.util.MappingHelper;
+import com.linkedin.datahub.graphql.types.common.mappers.AuditStampMapper;
 import com.linkedin.datahub.graphql.types.mappers.ModelMapper;
+import com.linkedin.datahub.graphql.types.mlmodel.mappers.MLHyperParamMapper;
+import com.linkedin.datahub.graphql.types.mlmodel.mappers.MLMetricMapper;
 import com.linkedin.dataprocess.DataProcessInstanceProperties;
+import com.linkedin.ml.metadata.MLTrainingRunProperties;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.EnvelopedAspectMap;
+import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
+import lombok.extern.slf4j.Slf4j;
 
-/**
- * Maps Pegasus {@link RecordTemplate} objects to objects conforming to the GQL schema.
- *
- * <p>To be replaced by auto-generated mappers implementations
- */
+@Slf4j
 public class DataProcessInstanceMapper implements ModelMapper<EntityResponse, DataProcessInstance> {
 
   public static final DataProcessInstanceMapper INSTANCE = new DataProcessInstanceMapper();
@@ -30,31 +38,115 @@ public static DataProcessInstance map(
     return INSTANCE.apply(context, entityResponse);
   }
 
+  private void mapContainers(
+      @Nullable final QueryContext context,
+      @Nonnull DataProcessInstance dataProcessInstance,
+      @Nonnull DataMap dataMap) {
+    final com.linkedin.container.Container gmsContainer =
+        new com.linkedin.container.Container(dataMap);
+    dataProcessInstance.setContainer(
+        com.linkedin.datahub.graphql.generated.Container.builder()
+            .setType(EntityType.CONTAINER)
+            .setUrn(gmsContainer.getContainer().toString())
+            .build());
+  }
+
   @Override
   public DataProcessInstance apply(
       @Nullable QueryContext context, @Nonnull final EntityResponse entityResponse) {
     final DataProcessInstance result = new DataProcessInstance();
     result.setUrn(entityResponse.getUrn().toString());
     result.setType(EntityType.DATA_PROCESS_INSTANCE);
 
+    Urn entityUrn = entityResponse.getUrn();
     EnvelopedAspectMap aspectMap = entityResponse.getAspects();
     MappingHelper<DataProcessInstance> mappingHelper = new MappingHelper<>(aspectMap, result);
     mappingHelper.mapToResult(
-        context, DATA_PROCESS_INSTANCE_PROPERTIES_ASPECT_NAME, this::mapDataProcessProperties);
-
+        DATA_PROCESS_INSTANCE_PROPERTIES_ASPECT_NAME,
+        (dataProcessInstance, dataMap) -> {
+          try {
+            mapTrainingRunProperties(context, dataProcessInstance, dataMap, entityUrn);
+          } catch (Exception e) {
+            mapDataProcessProperties(context, dataProcessInstance, dataMap, entityUrn);
+          }
+        });
+    mappingHelper.mapToResult(
+        DATA_PLATFORM_INSTANCE_ASPECT_NAME,
+        (dataProcessInstance, dataMap) -> {
+          DataPlatformInstance dataPlatformInstance = new DataPlatformInstance(dataMap);
+          dataProcessInstance.setDataPlatformInstance(
+              DataPlatformInstanceAspectMapper.map(context, dataPlatformInstance));
+          DataPlatform dataPlatform = new DataPlatform();
+          dataPlatform.setUrn(dataPlatformInstance.getPlatform().toString());
+          dataPlatform.setType(EntityType.DATA_PLATFORM);
+          dataProcessInstance.setPlatform(dataPlatform);
+        });
+    mappingHelper.mapToResult(
+        SUB_TYPES_ASPECT_NAME,
+        (dataProcessInstance, dataMap) ->
+            dataProcessInstance.setSubTypes(SubTypesMapper.map(context, new SubTypes(dataMap))));   
+    mappingHelper.mapToResult(
+        CONTAINER_ASPECT_NAME,
+        (dataProcessInstance, dataMap) -> mapContainers(context, dataProcessInstance, dataMap));
     return mappingHelper.getResult();
   }
 
+  private void mapTrainingRunProperties(
+      @Nonnull QueryContext context,
+      @Nonnull DataProcessInstance dpi,
+      @Nonnull DataMap dataMap,
+      @Nonnull Urn entityUrn) {
+    MLTrainingRunProperties trainingProperties = new MLTrainingRunProperties(dataMap);
+
+    com.linkedin.datahub.graphql.generated.MLTrainingRunProperties properties =
+        new com.linkedin.datahub.graphql.generated.MLTrainingRunProperties();
+    if (trainingProperties.hasOutputUrls()) {
+        properties.setOutputUrls(
+            trainingProperties.getOutputUrls()
+                .stream()
+                .map(url -> url.toString()) 
+                .collect(Collectors.toList())
+        );
+    }
+    if (trainingProperties.getHyperParams() != null) {
+      properties.setHyperParams(
+          trainingProperties.getHyperParams().stream()
+              .map(param -> MLHyperParamMapper.map(context, param))
+              .collect(Collectors.toList()));
+    }
+    if (trainingProperties.getTrainingMetrics() != null) {
+      properties.setTrainingMetrics(
+          trainingProperties.getTrainingMetrics().stream()
+              .map(metric -> MLMetricMapper.map(context, metric))
+              .collect(Collectors.toList()));
+    }
+    if (trainingProperties.hasId()) {
+      properties.setId(trainingProperties.getId());
+    }
+  }
+
   private void mapDataProcessProperties(
-      @Nonnull QueryContext context, @Nonnull DataProcessInstance dpi, @Nonnull DataMap dataMap) {
+      @Nonnull QueryContext context,
+      @Nonnull DataProcessInstance dpi,
+      @Nonnull DataMap dataMap,
+      @Nonnull Urn entityUrn) {
     DataProcessInstanceProperties dataProcessInstanceProperties =
         new DataProcessInstanceProperties(dataMap);
     dpi.setName(dataProcessInstanceProperties.getName());
-    if (dataProcessInstanceProperties.hasCreated()) {
-      dpi.setCreated(AuditStampMapper.map(context, dataProcessInstanceProperties.getCreated()));
-    }
+    com.linkedin.datahub.graphql.generated.DataProcessInstanceProperties properties =
+        new com.linkedin.datahub.graphql.generated.DataProcessInstanceProperties();
     if (dataProcessInstanceProperties.hasExternalUrl()) {
       dpi.setExternalUrl(dataProcessInstanceProperties.getExternalUrl().toString());
+      properties.setExternalUrl(dataProcessInstanceProperties.getExternalUrl().toString());
+    }
+    if (dataProcessInstanceProperties.hasCustomProperties()) {
+      properties.setCustomProperties(
+          CustomPropertiesMapper.map(
+              dataProcessInstanceProperties.getCustomProperties(), entityUrn));
+    }
+    if (dataProcessInstanceProperties.hasCreated()) {
+      dpi.setCreated(AuditStampMapper.map(context, dataProcessInstanceProperties.getCreated()));
     }
+    dpi.setProperties(properties);
   }
-}
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataprocessinstance/DataProcessInstanceType.java

@@ -0,0 +1,96 @@
+package com.linkedin.datahub.graphql.types.dataprocessinstance;
+
+import static com.linkedin.metadata.Constants.*;
+
+import com.google.common.collect.ImmutableSet;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
+import com.linkedin.datahub.graphql.generated.DataProcessInstance;
+import com.linkedin.datahub.graphql.generated.Entity;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.types.dataprocessinst.mappers.DataProcessInstanceMapper;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.client.EntityClient;
+import graphql.execution.DataFetcherResult;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import javax.annotation.Nonnull;
+import lombok.RequiredArgsConstructor;
+
+@RequiredArgsConstructor
+public class DataProcessInstanceType
+    implements com.linkedin.datahub.graphql.types.EntityType<DataProcessInstance, String> {
+
+  public static final Set<String> ASPECTS_TO_FETCH =
+      ImmutableSet.of(
+          DATA_PLATFORM_INSTANCE_ASPECT_NAME,
+          DATA_PROCESS_INSTANCE_PROPERTIES_ASPECT_NAME,
+          DATA_PROCESS_INSTANCE_RELATIONSHIPS_ASPECT_NAME,
+          SUB_TYPES_ASPECT_NAME,
+          CONTAINER_ASPECT_NAME);
+
+  private final EntityClient _entityClient;
+  private final FeatureFlags _featureFlags;
+
+  @Override
+  public EntityType type() {
+    return EntityType.DATA_PROCESS_INSTANCE;
+  }
+
+  @Override
+  public Function<Entity, String> getKeyProvider() {
+    return Entity::getUrn;
+  }
+
+  @Override
+  public Class<DataProcessInstance> objectClass() {
+    return DataProcessInstance.class;
+  }
+
+  @Override
+  public List<DataFetcherResult<DataProcessInstance>> batchLoad(
+      @Nonnull List<String> urns, @Nonnull QueryContext context) throws Exception {
+    final List<Urn> dataProcessInstanceUrns =
+        urns.stream().map(UrnUtils::getUrn).collect(Collectors.toList());
+
+    try {
+      Map<Urn, EntityResponse> entities = new HashMap<>();
+      if (_featureFlags.isDataProcessInstanceEntityEnabled()) {
+        entities =
+            _entityClient.batchGetV2(
+                context.getOperationContext(),
+                DATA_PROCESS_INSTANCE_ENTITY_NAME,
+                new HashSet<>(dataProcessInstanceUrns),
+                ASPECTS_TO_FETCH);
+      }
+
+      final List<EntityResponse> gmsResults = new ArrayList<>();
+      for (Urn urn : dataProcessInstanceUrns) {
+        if (_featureFlags.isDataProcessInstanceEntityEnabled()) {
+          gmsResults.add(entities.getOrDefault(urn, null));
+        }
+      }
+
+      return gmsResults.stream()
+          .map(
+              gmsResult ->
+                  gmsResult == null
+                      ? null
+                      : DataFetcherResult.<DataProcessInstance>newResult()
+                          .data(DataProcessInstanceMapper.map(context, gmsResult))
+                          .build())
+          .collect(Collectors.toList());
+
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to load schemaField entity", e);
+    }
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java

@@ -77,6 +77,9 @@ public class EntityTypeUrnMapper {
           .put(
               Constants.BUSINESS_ATTRIBUTE_ENTITY_NAME,
               "urn:li:entityType:datahub.businessAttribute")
+          .put(
+            Constants.DATA_PROCESS_INSTANCE_ENTITY_NAME,
+            "urn:li:entityType:datahub.dataProcessInstance")
           .build();
 
   private static final Map<String, String> ENTITY_TYPE_URN_TO_NAME =

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupPropertiesMapper.java

@@ -5,6 +5,7 @@
 import com.linkedin.datahub.graphql.generated.MLModelGroupProperties;
 import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
 import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
+import com.linkedin.datahub.graphql.types.common.mappers.TimeStampMapper;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -27,13 +28,19 @@ public MLModelGroupProperties apply(
       @Nonnull final com.linkedin.ml.metadata.MLModelGroupProperties mlModelGroupProperties,
       @Nonnull Urn entityUrn) {
     final MLModelGroupProperties result = new MLModelGroupProperties();
-
+
+    if (mlModelGroupProperties.getName() != null) {
+      result.setName(mlModelGroupProperties.getName());
+    } else {
+      // backfill name from URN for backwards compatibility
+      result.setName(entityUrn.getEntityKey().get(1)); // indexed access is safe here
+    }
     result.setDescription(mlModelGroupProperties.getDescription());
     if (mlModelGroupProperties.getVersion() != null) {
       result.setVersion(VersionTagMapper.map(context, mlModelGroupProperties.getVersion()));
     }
-    result.setCreatedAt(mlModelGroupProperties.getCreatedAt());
-
+    result.setCreated(TimeStampMapper.map(context, mlModelGroupProperties.getCreated()));
+    result.setLastModified(TimeStampMapper.map(context, mlModelGroupProperties.getLastModified()));
     result.setCustomProperties(
         CustomPropertiesMapper.map(mlModelGroupProperties.getCustomProperties(), entityUrn));