Skip to content

Commit 029a274

Browse files
yoonhyejinasikowitzRyanHolstienshirshankaryota-cloud
authored and
eagle-25
committed
feat: update ml system UI (datahub-project#12334)
Co-authored-by: Andrew Sikowitz <[email protected]> Co-authored-by: RyanHolstien <[email protected]> Co-authored-by: Shirshanka Das <[email protected]> Co-authored-by: ryota-cloud <[email protected]>
1 parent 88bbc51 commit 029a274

File tree

24 files changed

+989
-90
lines changed

24 files changed

+989
-90
lines changed

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

+33
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33
import com.linkedin.common.urn.Urn;
44
import com.linkedin.datahub.graphql.QueryContext;
55
import com.linkedin.datahub.graphql.generated.MLModelGroupProperties;
6+
import com.linkedin.datahub.graphql.generated.MLModelLineageInfo;
67
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
8+
import com.linkedin.datahub.graphql.types.common.mappers.TimeStampToAuditStampMapper;
79
import com.linkedin.datahub.graphql.types.mappers.EmbeddedModelMapper;
10+
import java.util.stream.Collectors;
811
import javax.annotation.Nonnull;
912
import javax.annotation.Nullable;
1013

@@ -33,10 +36,40 @@ public MLModelGroupProperties apply(
3336
result.setVersion(VersionTagMapper.map(context, mlModelGroupProperties.getVersion()));
3437
}
3538
result.setCreatedAt(mlModelGroupProperties.getCreatedAt());
39+
if (mlModelGroupProperties.hasCreated()) {
40+
result.setCreated(
41+
TimeStampToAuditStampMapper.map(context, mlModelGroupProperties.getCreated()));
42+
}
43+
if (mlModelGroupProperties.getName() != null) {
44+
result.setName(mlModelGroupProperties.getName());
45+
} else {
46+
// backfill name from URN for backwards compatibility
47+
result.setName(entityUrn.getEntityKey().get(1)); // indexed access is safe here
48+
}
49+
50+
if (mlModelGroupProperties.hasLastModified()) {
51+
result.setLastModified(
52+
TimeStampToAuditStampMapper.map(context, mlModelGroupProperties.getLastModified()));
53+
}
3654

3755
result.setCustomProperties(
3856
CustomPropertiesMapper.map(mlModelGroupProperties.getCustomProperties(), entityUrn));
3957

58+
final MLModelLineageInfo lineageInfo = new MLModelLineageInfo();
59+
if (mlModelGroupProperties.hasTrainingJobs()) {
60+
lineageInfo.setTrainingJobs(
61+
mlModelGroupProperties.getTrainingJobs().stream()
62+
.map(urn -> urn.toString())
63+
.collect(Collectors.toList()));
64+
}
65+
if (mlModelGroupProperties.hasDownstreamJobs()) {
66+
lineageInfo.setDownstreamJobs(
67+
mlModelGroupProperties.getDownstreamJobs().stream()
68+
.map(urn -> urn.toString())
69+
.collect(Collectors.toList()));
70+
}
71+
result.setMlModelLineageInfo(lineageInfo);
72+
4073
return result;
4174
}
4275
}

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

+15
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.linkedin.common.urn.Urn;
66
import com.linkedin.datahub.graphql.QueryContext;
77
import com.linkedin.datahub.graphql.generated.MLModelGroup;
8+
import com.linkedin.datahub.graphql.generated.MLModelLineageInfo;
89
import com.linkedin.datahub.graphql.generated.MLModelProperties;
910
import com.linkedin.datahub.graphql.types.common.mappers.CustomPropertiesMapper;
1011
import com.linkedin.datahub.graphql.types.common.mappers.TimeStampToAuditStampMapper;
@@ -87,6 +88,20 @@ public MLModelProperties apply(
8788
.collect(Collectors.toList()));
8889
}
8990
result.setTags(mlModelProperties.getTags());
91+
final MLModelLineageInfo lineageInfo = new MLModelLineageInfo();
92+
if (mlModelProperties.hasTrainingJobs()) {
93+
lineageInfo.setTrainingJobs(
94+
mlModelProperties.getTrainingJobs().stream()
95+
.map(urn -> urn.toString())
96+
.collect(Collectors.toList()));
97+
}
98+
if (mlModelProperties.hasDownstreamJobs()) {
99+
lineageInfo.setDownstreamJobs(
100+
mlModelProperties.getDownstreamJobs().stream()
101+
.map(urn -> urn.toString())
102+
.collect(Collectors.toList()));
103+
}
104+
result.setMlModelLineageInfo(lineageInfo);
90105

91106
return result;
92107
}

‎datahub-graphql-core/src/main/resources/lineage.graphql

+29
Original file line numberDiff line numberDiff line change
@@ -25,3 +25,32 @@ input LineageEdge {
2525
"""
2626
upstreamUrn: String!
2727
}
28+
29+
"""
30+
Represents lineage information for ML entities.
31+
"""
32+
type MLModelLineageInfo {
33+
"""
34+
List of jobs or processes used to train the model.
35+
"""
36+
trainingJobs: [String!]
37+
38+
"""
39+
List of jobs or processes that use this model.
40+
"""
41+
downstreamJobs: [String!]
42+
}
43+
44+
extend type MLModelProperties {
45+
"""
46+
Information related to lineage to this model group
47+
"""
48+
mlModelLineageInfo: MLModelLineageInfo
49+
}
50+
51+
extend type MLModelGroupProperties {
52+
"""
53+
Information related to lineage to this model group
54+
"""
55+
mlModelLineageInfo: MLModelLineageInfo
56+
}

‎datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelGroupPropertiesMapperTest.java

+68
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
2+
3+
import static org.testng.Assert.assertEquals;
4+
import static org.testng.Assert.assertNotNull;
5+
import static org.testng.Assert.assertNull;
6+
7+
import com.linkedin.common.urn.Urn;
8+
import com.linkedin.ml.metadata.MLModelGroupProperties;
9+
import java.net.URISyntaxException;
10+
import org.testng.annotations.Test;
11+
12+
public class MLModelGroupPropertiesMapperTest {
13+
14+
@Test
15+
public void testMapMLModelGroupProperties() throws URISyntaxException {
16+
// Create backend ML Model Group Properties
17+
MLModelGroupProperties input = new MLModelGroupProperties();
18+
19+
// Set description
20+
input.setDescription("a ml trust model group");
21+
22+
// Set Name
23+
input.setName("ML trust model group");
24+
25+
// Create URN
26+
Urn groupUrn =
27+
Urn.createFromString(
28+
"urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,another-group,PROD)");
29+
30+
// Map the properties
31+
com.linkedin.datahub.graphql.generated.MLModelGroupProperties result =
32+
MLModelGroupPropertiesMapper.map(null, input, groupUrn);
33+
34+
// Verify mapped properties
35+
assertNotNull(result);
36+
assertEquals(result.getDescription(), "a ml trust model group");
37+
assertEquals(result.getName(), "ML trust model group");
38+
39+
// Verify lineage info is null as in the mock data
40+
assertNotNull(result.getMlModelLineageInfo());
41+
assertNull(result.getMlModelLineageInfo().getTrainingJobs());
42+
assertNull(result.getMlModelLineageInfo().getDownstreamJobs());
43+
}
44+
45+
@Test
46+
public void testMapWithMinimalProperties() throws URISyntaxException {
47+
// Create backend ML Model Group Properties with minimal information
48+
MLModelGroupProperties input = new MLModelGroupProperties();
49+
50+
// Create URN
51+
Urn groupUrn =
52+
Urn.createFromString(
53+
"urn:li:mlModelGroup:(urn:li:dataPlatform:sagemaker,another-group,PROD)");
54+
55+
// Map the properties
56+
com.linkedin.datahub.graphql.generated.MLModelGroupProperties result =
57+
MLModelGroupPropertiesMapper.map(null, input, groupUrn);
58+
59+
// Verify basic mapping with minimal properties
60+
assertNotNull(result);
61+
assertNull(result.getDescription());
62+
63+
// Verify lineage info is null
64+
assertNotNull(result.getMlModelLineageInfo());
65+
assertNull(result.getMlModelLineageInfo().getTrainingJobs());
66+
assertNull(result.getMlModelLineageInfo().getDownstreamJobs());
67+
}
68+
}

‎datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/types/mlmodel/mappers/MLModelPropertiesMapperTest.java

+187
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
package com.linkedin.datahub.graphql.types.mlmodel.mappers;
2+
3+
import static org.testng.Assert.assertEquals;
4+
import static org.testng.Assert.assertNotNull;
5+
import static org.testng.Assert.assertNull;
6+
7+
import com.linkedin.common.MLFeatureUrnArray;
8+
import com.linkedin.common.TimeStamp;
9+
import com.linkedin.common.VersionTag;
10+
import com.linkedin.common.url.Url;
11+
import com.linkedin.common.urn.MLFeatureUrn;
12+
import com.linkedin.common.urn.MLModelUrn;
13+
import com.linkedin.common.urn.Urn;
14+
import com.linkedin.data.template.StringArray;
15+
import com.linkedin.data.template.StringMap;
16+
import com.linkedin.ml.metadata.MLHyperParam;
17+
import com.linkedin.ml.metadata.MLHyperParamArray;
18+
import com.linkedin.ml.metadata.MLMetric;
19+
import com.linkedin.ml.metadata.MLMetricArray;
20+
import com.linkedin.ml.metadata.MLModelProperties;
21+
import java.net.URISyntaxException;
22+
import org.testng.annotations.Test;
23+
24+
public class MLModelPropertiesMapperTest {
25+
26+
@Test
27+
public void testMapMLModelProperties() throws URISyntaxException {
28+
MLModelProperties input = new MLModelProperties();
29+
30+
// Set basic properties
31+
input.setName("TestModel");
32+
input.setDescription("A test ML model");
33+
input.setType("Classification");
34+
35+
// Set version
36+
VersionTag versionTag = new VersionTag();
37+
versionTag.setVersionTag("1.0.0");
38+
input.setVersion(versionTag);
39+
40+
// Set external URL
41+
Url externalUrl = new Url("https://example.com/model");
42+
input.setExternalUrl(externalUrl);
43+
44+
// Set created and last modified timestamps
45+
TimeStamp createdTimeStamp = new TimeStamp();
46+
createdTimeStamp.setTime(1000L);
47+
Urn userUrn = Urn.createFromString("urn:li:corpuser:test");
48+
createdTimeStamp.setActor(userUrn);
49+
input.setCreated(createdTimeStamp);
50+
51+
TimeStamp lastModifiedTimeStamp = new TimeStamp();
52+
lastModifiedTimeStamp.setTime(2000L);
53+
lastModifiedTimeStamp.setActor(userUrn);
54+
input.setLastModified(lastModifiedTimeStamp);
55+
56+
// Set custom properties
57+
StringMap customProps = new StringMap();
58+
customProps.put("key1", "value1");
59+
customProps.put("key2", "value2");
60+
input.setCustomProperties(customProps);
61+
62+
// Set hyper parameters
63+
MLHyperParamArray hyperParams = new MLHyperParamArray();
64+
MLHyperParam hyperParam1 = new MLHyperParam();
65+
hyperParam1.setName("learning_rate");
66+
hyperParam1.setValue("0.01");
67+
hyperParams.add(hyperParam1);
68+
input.setHyperParams(hyperParams);
69+
70+
// Set training metrics
71+
MLMetricArray trainingMetrics = new MLMetricArray();
72+
MLMetric metric1 = new MLMetric();
73+
metric1.setName("accuracy");
74+
metric1.setValue("0.95");
75+
trainingMetrics.add(metric1);
76+
input.setTrainingMetrics(trainingMetrics);
77+
78+
// Set ML features
79+
MLFeatureUrnArray mlFeatures = new MLFeatureUrnArray();
80+
MLFeatureUrn featureUrn = MLFeatureUrn.createFromString("urn:li:mlFeature:(dataset,feature)");
81+
mlFeatures.add(featureUrn);
82+
input.setMlFeatures(mlFeatures);
83+
84+
// Set tags
85+
StringArray tags = new StringArray();
86+
tags.add("tag1");
87+
tags.add("tag2");
88+
input.setTags(tags);
89+
90+
// Set training and downstream jobs
91+
input.setTrainingJobs(
92+
new com.linkedin.common.UrnArray(Urn.createFromString("urn:li:dataJob:train")));
93+
input.setDownstreamJobs(
94+
new com.linkedin.common.UrnArray(Urn.createFromString("urn:li:dataJob:predict")));
95+
96+
// Create ML Model URN
97+
MLModelUrn modelUrn =
98+
MLModelUrn.createFromString(
99+
"urn:li:mlModel:(urn:li:dataPlatform:sagemaker,unittestmodel,PROD)");
100+
101+
// Map the properties
102+
com.linkedin.datahub.graphql.generated.MLModelProperties result =
103+
MLModelPropertiesMapper.map(null, input, modelUrn);
104+
105+
// Verify mapped properties
106+
assertNotNull(result);
107+
assertEquals(result.getName(), "TestModel");
108+
assertEquals(result.getDescription(), "A test ML model");
109+
assertEquals(result.getType(), "Classification");
110+
assertEquals(result.getVersion(), "1.0.0");
111+
assertEquals(result.getExternalUrl(), "https://example.com/model");
112+
113+
// Verify audit stamps
114+
assertNotNull(result.getCreated());
115+
assertEquals(result.getCreated().getTime().longValue(), 1000L);
116+
assertEquals(result.getCreated().getActor(), userUrn.toString());
117+
118+
assertNotNull(result.getLastModified());
119+
assertEquals(result.getLastModified().getTime().longValue(), 2000L);
120+
assertEquals(result.getLastModified().getActor(), userUrn.toString());
121+
122+
// Verify custom properties
123+
assertNotNull(result.getCustomProperties());
124+
125+
// Verify hyper parameters
126+
assertNotNull(result.getHyperParams());
127+
assertEquals(result.getHyperParams().size(), 1);
128+
assertEquals(result.getHyperParams().get(0).getName(), "learning_rate");
129+
assertEquals(result.getHyperParams().get(0).getValue(), "0.01");
130+
131+
// Verify training metrics
132+
assertNotNull(result.getTrainingMetrics());
133+
assertEquals(result.getTrainingMetrics().size(), 1);
134+
assertEquals(result.getTrainingMetrics().get(0).getName(), "accuracy");
135+
assertEquals(result.getTrainingMetrics().get(0).getValue(), "0.95");
136+
137+
// Verify ML features
138+
assertNotNull(result.getMlFeatures());
139+
assertEquals(result.getMlFeatures().size(), 1);
140+
assertEquals(result.getMlFeatures().get(0), featureUrn.toString());
141+
142+
// Verify tags
143+
assertNotNull(result.getTags());
144+
assertEquals(result.getTags().get(0), "tag1");
145+
assertEquals(result.getTags().get(1), "tag2");
146+
147+
// Verify lineage info
148+
assertNotNull(result.getMlModelLineageInfo());
149+
assertEquals(result.getMlModelLineageInfo().getTrainingJobs().size(), 1);
150+
assertEquals(result.getMlModelLineageInfo().getTrainingJobs().get(0), "urn:li:dataJob:train");
151+
assertEquals(result.getMlModelLineageInfo().getDownstreamJobs().size(), 1);
152+
assertEquals(
153+
result.getMlModelLineageInfo().getDownstreamJobs().get(0), "urn:li:dataJob:predict");
154+
}
155+
156+
@Test
157+
public void testMapWithMissingName() throws URISyntaxException {
158+
MLModelProperties input = new MLModelProperties();
159+
MLModelUrn modelUrn =
160+
MLModelUrn.createFromString(
161+
"urn:li:mlModel:(urn:li:dataPlatform:sagemaker,missingnamemodel,PROD)");
162+
163+
com.linkedin.datahub.graphql.generated.MLModelProperties result =
164+
MLModelPropertiesMapper.map(null, input, modelUrn);
165+
166+
// Verify that name is extracted from URN when not present in input
167+
assertEquals(result.getName(), "missingnamemodel");
168+
}
169+
170+
@Test
171+
public void testMapWithMinimalProperties() throws URISyntaxException {
172+
MLModelProperties input = new MLModelProperties();
173+
MLModelUrn modelUrn =
174+
MLModelUrn.createFromString(
175+
"urn:li:mlModel:(urn:li:dataPlatform:sagemaker,minimalmodel,PROD)");
176+
177+
com.linkedin.datahub.graphql.generated.MLModelProperties result =
178+
MLModelPropertiesMapper.map(null, input, modelUrn);
179+
180+
// Verify basic mapping with minimal properties
181+
assertNotNull(result);
182+
assertEquals(result.getName(), "minimalmodel");
183+
assertNull(result.getDescription());
184+
assertNull(result.getType());
185+
assertNull(result.getVersion());
186+
}
187+
}

‎datahub-web-react/src/app/entity/EntityPage.tsx

+1
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ export const EntityPage = ({ entityType }: Props) => {
6666
entityType === EntityType.MlfeatureTable ||
6767
entityType === EntityType.MlmodelGroup ||
6868
entityType === EntityType.GlossaryTerm ||
69+
entityType === EntityType.DataProcessInstance ||
6970
entityType === EntityType.GlossaryNode;
7071

7172
return (

0 commit comments

Comments
 (0)