remove unneccesary code & update scripts to reflect recent changes

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)

metadata-ingestion/src/datahub/ingestion/source/mlflow.py

@@ -42,7 +42,6 @@
     AuditStampClass,
     ContainerClass,
     DataPlatformInstanceClass,
-    DataProcessInstanceOutputClass,
     DataProcessInstancePropertiesClass,
     DataProcessInstanceRunEventClass,
     DataProcessInstanceRunResultClass,
@@ -307,13 +306,13 @@ def _get_run_workunits(
             aspect=ContainerClass(container=experiment_key.as_urn()),
         ).as_workunit()
 
-        model_versions = self.get_mlflow_model_versions_from_run(run.info.run_id)
-        if model_versions:
-            model_version_urn = self._make_ml_model_urn(model_versions[0])
-            yield MetadataChangeProposalWrapper(
-                entityUrn=str(data_process_instance.urn),
-                aspect=DataProcessInstanceOutputClass(outputs=[model_version_urn]),
-            ).as_workunit()
+        # model_versions = self.get_mlflow_model_versions_from_run(run.info.run_id)
+        # if model_versions:
+        #     model_version_urn = self._make_ml_model_urn(model_versions[0])
+        #     yield MetadataChangeProposalWrapper(
+        #         entityUrn=str(data_process_instance.urn),
+        #         aspect=DataProcessInstanceOutputClass(outputs=[model_version_urn]),
+        #     ).as_workunit()  # TODO: if model.trainingJob has job, shouldn't this not needed?
 
         metrics = self._get_run_metrics(run)
         hyperparams = self._get_run_params(run)

metadata-ingestion/tests/integration/mlflow/mlflow_mcps_golden.json

@@ -301,24 +301,6 @@
         "lastRunId": "no-run-id-provided"
     }
 },
-{
-    "entityType": "dataProcessInstance",
-    "entityUrn": "urn:li:dataProcessInstance:2666299269d6ebea994d5ec0c29e3aca",
-    "changeType": "UPSERT",
-    "aspectName": "dataProcessInstanceOutput",
-    "aspect": {
-        "json": {
-            "outputs": [
-                "urn:li:mlModel:(urn:li:dataPlatform:mlflow,test-model_1,PROD)"
-            ]
-        }
-    },
-    "systemMetadata": {
-        "lastObserved": 1615443388097,
-        "runId": "mlflow-source-test",
-        "lastRunId": "no-run-id-provided"
-    }
-},
 {
     "entityType": "dataProcessInstance",
     "entityUrn": "urn:li:dataProcessInstance:2666299269d6ebea994d5ec0c29e3aca",