feat(docs/mlflow): update sample scripts to be compatible with edges/versioning (#12878)

Author: yoonhyejin

metadata-ingestion/examples/ai/dh_ai_client.py

@@ -14,6 +14,7 @@
     ChangeTypeClass,
     DataProcessInstanceRunResultClass,
     DataProcessRunStatusClass,
+    EdgeClass,
 )
 from datahub.metadata.urns import (
     ContainerUrn,
@@ -255,7 +256,7 @@ def create_model(
         version_props = {
             "version": version_tag,
             "versionSet": str(version_set_urn),
-            "sortId": "AAAAAAAA",
+            "sortId": str(version_tag).zfill(10),
         }
 
         # Add alias if provided
@@ -266,22 +267,10 @@ def create_model(
             models.VersionPropertiesClass, version_props
         )
 
-        # Create version set properties
-        version_set_properties = models.VersionSetPropertiesClass(
-            latest=str(model_urn),
-            versioningScheme="ALPHANUMERIC_GENERATED_BY_DATAHUB",
-        )
-
         mcps = [
             self._create_mcp(
                 str(model_urn), properties, "mlModel", "mlModelProperties"
             ),
-            self._create_mcp(
-                str(version_set_urn),
-                version_set_properties,
-                "versionSet",
-                "versionSetProperties",
-            ),
             self._create_mcp(
                 str(model_urn), version_properties, "mlModel", "versionProperties"
             ),
@@ -429,7 +418,13 @@ def add_input_datasets_to_run(self, run_urn: str, dataset_urns: List[str]) -> No
             entity_urn=run_urn,
             entity_type="dataProcessInstance",
             aspect_name="dataProcessInstanceInput",
-            aspect=DataProcessInstanceInput(inputs=dataset_urns),
+            aspect=DataProcessInstanceInput(
+                inputs=[],
+                inputEdges=[
+                    EdgeClass(destinationUrn=str(dataset_urn))
+                    for dataset_urn in dataset_urns
+                ],
+            ),
         )
         self._emit_mcps(mcp)
         logger.info(f"Added input datasets to run {run_urn}")
@@ -440,7 +435,13 @@ def add_output_datasets_to_run(self, run_urn: str, dataset_urns: List[str]) -> N
             entity_urn=run_urn,
             entity_type="dataProcessInstance",
             aspect_name="dataProcessInstanceOutput",
-            aspect=DataProcessInstanceOutput(outputs=dataset_urns),
+            aspect=DataProcessInstanceOutput(
+                outputEdges=[
+                    EdgeClass(destinationUrn=str(dataset_urn))
+                    for dataset_urn in dataset_urns
+                ],
+                outputs=[],
+            ),
         )
         self._emit_mcps(mcp)
         logger.info(f"Added output datasets to run {run_urn}")

metadata-ingestion/examples/ai/dh_ai_client_sample.py

@@ -96,15 +96,19 @@
         end_timestamp=1628580001000,
     )
     # Create datasets
-    input_dataset_urn = client.create_dataset(
-        platform="snowflake",
-        name="iris_input",
-    )
+    input_dataset_urns = [
+        client.create_dataset(
+            platform="snowflake",
+            name="iris_input",
+        )
+    ]
 
-    output_dataset_urn = client.create_dataset(
-        platform="snowflake",
-        name="iris_ouptut",
-    )
+    output_dataset_urns = [
+        client.create_dataset(
+            platform="snowflake",
+            name="iris_ouptut",
+        )
+    ]
 
     # Add run to experiment
     client.add_run_to_experiment(run_urn=run_urn, experiment_urn=experiment_urn)
@@ -125,10 +129,6 @@
     )
 
     # Add input and output datasets to run
-    client.add_input_datasets_to_run(
-        run_urn=run_urn, dataset_urns=[str(input_dataset_urn)]
-    )
+    client.add_input_datasets_to_run(run_urn=run_urn, dataset_urns=input_dataset_urns)
 
-    client.add_output_datasets_to_run(
-        run_urn=run_urn, dataset_urns=[str(output_dataset_urn)]
-    )
+    client.add_output_datasets_to_run(run_urn=run_urn, dataset_urns=output_dataset_urns)