acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/source/metadata/lineage.py

@@ -36,9 +36,10 @@ from datahub.ingestion.api.source_helpers import (
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import get_default_graph
-from datahub.metadata.com.linkedin.pegasus2avro.dataset import (
-    FineGrainedLineageDownstreamType,
-    FineGrainedLineageUpstreamType,
+from datahub.ingestion.graph.config import ClientMode
+from datahub.metadata.schema_classes import (
+    FineGrainedLineageDownstreamTypeClass,
+    FineGrainedLineageUpstreamTypeClass,
 )
 
 logger = logging.getLogger(__name__)
@@ -48,7 +49,7 @@ class EntityConfig(EnvConfigMixin):
     name: str
     type: str
     platform: str
-    platform_instance: Optional[str]
+    platform_instance: Optional[str] = None
 
     @validator("type")
     def type_must_be_supported(cls, v: str) -> str:
@@ -79,9 +80,9 @@ class FineGrainedLineageConfig(ConfigModel):
     @validator("upstreamType")
     def upstream_type_must_be_supported(cls, v: str) -> str:
         allowed_types = [
-            FineGrainedLineageUpstreamType.FIELD_SET,
-            FineGrainedLineageUpstreamType.DATASET,
-            FineGrainedLineageUpstreamType.NONE,
+            FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+            FineGrainedLineageUpstreamTypeClass.DATASET,
+            FineGrainedLineageUpstreamTypeClass.NONE,
         ]
         if v not in allowed_types:
             raise ValueError(
@@ -92,8 +93,8 @@ class FineGrainedLineageConfig(ConfigModel):
     @validator("downstreamType")
     def downstream_type_must_be_supported(cls, v: str) -> str:
         allowed_types = [
-            FineGrainedLineageDownstreamType.FIELD_SET,
-            FineGrainedLineageDownstreamType.FIELD,
+            FineGrainedLineageDownstreamTypeClass.FIELD_SET,
+            FineGrainedLineageDownstreamTypeClass.FIELD,
         ]
         if v not in allowed_types:
             raise ValueError(
@@ -210,7 +211,7 @@ def _get_lineage_mcp(
 
     # extract the old lineage and save it for the new mcp
     if preserve_upstream:
-        client = get_default_graph()
+        client = get_default_graph(ClientMode.INGESTION)
 
         old_upstream_lineage = get_aspects_for_entity(
             client._session,

datahub/ingestion/source/mlflow.py

@@ -1,10 +1,13 @@
+import json
+import os
 import time
 from dataclasses import dataclass
-from typing import Any, Callable, Iterable, List, Optional, TypeVar, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, TypeVar, Union
 
 from mlflow import MlflowClient
-from mlflow.entities import Experiment, Run
+from mlflow.entities import Dataset as MlflowDataset, Experiment, Run
 from mlflow.entities.model_registry import ModelVersion, RegisteredModel
+from mlflow.exceptions import MlflowException
 from mlflow.store.entities import PagedList
 from pydantic.fields import Field
 
@@ -14,7 +17,7 @@ from datahub.api.entities.dataprocess.dataprocess_instance import (
 )
 from datahub.configuration.source_common import EnvConfigMixin
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.emitter.mcp_builder import ContainerKey
+from datahub.emitter.mcp_builder import ExperimentKey
 from datahub.ingestion.api.common import PipelineContext
 from datahub.ingestion.api.decorators import (
     SupportStatus,
@@ -29,10 +32,15 @@ from datahub.ingestion.api.source import (
     SourceReport,
 )
 from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.source.common.subtypes import MLAssetSubTypes
+from datahub.ingestion.source.common.data_platforms import KNOWN_VALID_PLATFORM_NAMES
+from datahub.ingestion.source.common.subtypes import (
+    MLAssetSubTypes,
+    SourceCapabilityModifier,
+)
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
     StaleEntityRemovalSourceReport,
+    StatefulStaleMetadataRemovalConfig,
 )
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
@@ -42,6 +50,7 @@ from datahub.metadata.schema_classes import (
     AuditStampClass,
     ContainerClass,
     DataPlatformInstanceClass,
+    DataProcessInstanceInputClass,
     DataProcessInstanceOutputClass,
     DataProcessInstancePropertiesClass,
     DataProcessInstanceRunEventClass,
@@ -60,24 +69,19 @@ from datahub.metadata.schema_classes import (
     TagAssociationClass,
     TagPropertiesClass,
     TimeStampClass,
+    UpstreamClass,
+    UpstreamLineageClass,
     VersionPropertiesClass,
     VersionTagClass,
     _Aspect,
 )
-from datahub.metadata.urns import (
-    DataPlatformUrn,
-    MlModelUrn,
-    VersionSetUrn,
-)
+from datahub.metadata.urns import DataPlatformUrn, DatasetUrn, MlModelUrn, VersionSetUrn
 from datahub.sdk.container import Container
+from datahub.sdk.dataset import Dataset
 
 T = TypeVar("T")
 
 
-class ContainerKeyWithId(ContainerKey):
-    id: str
-
-
 class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
     tracking_uri: Optional[str] = Field(
         default=None,
@@ -105,6 +109,22 @@ class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
             " If neither is set, external URLs are not generated."
         ),
     )
+    materialize_dataset_inputs: Optional[bool] = Field(
+        default=False,
+        description="Whether to materialize dataset inputs for each run",
+    )
+    source_mapping_to_platform: Optional[dict] = Field(
+        default=None, description="Mapping of source type to datahub platform"
+    )
+
+    username: Optional[str] = Field(
+        default=None, description="Username for MLflow authentication"
+    )
+    password: Optional[str] = Field(
+        default=None, description="Password for MLflow authentication"
+    )
+
+    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
 
 
 @dataclass
@@ -116,11 +136,18 @@ class MLflowRegisteredModelStageInfo:
 
 @platform_name("MLflow")
 @config_class(MLflowConfig)
-@support_status(SupportStatus.TESTING)
+@support_status(SupportStatus.INCUBATING)
 @capability(
     SourceCapability.DESCRIPTIONS,
     "Extract descriptions for MLflow Registered Models and Model Versions",
 )
+@capability(
+    SourceCapability.CONTAINERS,
+    "Extract ML experiments",
+    subtype_modifier=[
+        SourceCapabilityModifier.MLFLOW_EXPERIMENT,
+    ],
+)
 @capability(SourceCapability.TAGS, "Extract tags for MLflow Registered Model Stages")
 class MLflowSource(StatefulIngestionSourceBase):
     platform = "mlflow"
@@ -152,7 +179,17 @@ class MLflowSource(StatefulIngestionSourceBase):
         self.ctx = ctx
         self.config = config
         self.report = StaleEntityRemovalSourceReport()
-        self.client = MlflowClient(
+        self.client = self._configure_client()
+
+    def _configure_client(self) -> MlflowClient:
+        if bool(self.config.username) != bool(self.config.password):
+            raise ValueError("Both username and password must be set together")
+
+        if self.config.username and self.config.password:
+            os.environ["MLFLOW_TRACKING_USERNAME"] = self.config.username
+            os.environ["MLFLOW_TRACKING_PASSWORD"] = self.config.password
+
+        return MlflowClient(
             tracking_uri=self.config.tracking_uri,
             registry_uri=self.config.registry_uri,
         )
@@ -213,6 +250,7 @@ class MLflowSource(StatefulIngestionSourceBase):
             if runs:
                 for run in runs:
                     yield from self._get_run_workunits(experiment, run)
+                    yield from self._get_dataset_input_workunits(run)
 
     def _get_experiment_custom_properties(self, experiment):
         experiment_custom_props = getattr(experiment, "tags", {}) or {}
@@ -224,7 +262,7 @@ class MLflowSource(StatefulIngestionSourceBase):
         self, experiment: Experiment
     ) -> Iterable[MetadataWorkUnit]:
         experiment_container = Container(
-            container_key=ContainerKeyWithId(
+            container_key=ExperimentKey(
                 platform=str(DataPlatformUrn(platform_name=self.platform)),
                 id=experiment.name,
             ),
@@ -262,10 +300,187 @@ class MLflowSource(StatefulIngestionSourceBase):
                 type="SKIPPED", nativeResultType=self.platform
             )
 
+    def _get_dataset_schema(
+        self, dataset: MlflowDataset
+    ) -> Optional[List[Tuple[str, str]]]:
+        try:
+            schema_dict = json.loads(dataset.schema)
+        except json.JSONDecodeError:
+            self.report.warning(
+                title="Failed to load dataset schema",
+                message="Schema metadata will be missing due to a JSON parsing error.",
+                context=f"Dataset: {dataset.name}, Schema: {dataset.schema}",
+            )
+            return None
+
+        if "mlflow_colspec" in schema_dict:
+            try:
+                return [
+                    (field["name"], field["type"])
+                    for field in schema_dict["mlflow_colspec"]
+                ]
+            except (KeyError, TypeError):
+                return None
+        # If the schema is not formatted, return None
+        return None
+
+    def _get_external_dataset_urn(self, platform: str, dataset_name: str) -> str:
+        """
+        Get the URN for an external dataset.
+        Args:
+            platform: The platform of the external dataset (e.g., 's3', 'bigquery')
+            dataset: The MLflow dataset
+        Returns:
+            str: The URN of the external dataset
+        """
+        return str(DatasetUrn(platform=platform, name=dataset_name))
+
+    def _get_dataset_input_workunits(self, run: Run) -> Iterable[MetadataWorkUnit]:
+        """
+        Generate workunits for dataset inputs in a run.
+
+        For each dataset input:
+        1. If source type is 'local' or 'code':
+           - Create a local dataset reference
+        2. Otherwise:
+           - If materialization is enabled:
+             - Create a hosted dataset and a dataset reference with upstream
+           - If materialization is not enabled:
+             - Create a dataset reference and add upstream if dataset exists
+        3. Add all dataset references as upstreams for the run
+        """
+        run_urn = DataProcessInstance(
+            id=run.info.run_id,
+            orchestrator=self.platform,
+        ).urn
+
+        dataset_reference_urns = []
+
+        for dataset_input in run.inputs.dataset_inputs:
+            dataset = dataset_input.dataset
+            source_type = dataset.source_type
+            dataset_tags = {k[1]: v[1] for k, v in dataset_input.tags}
+
+            # Prepare dataset properties
+            custom_properties = dataset_tags
+            formatted_schema = self._get_dataset_schema(dataset)
+            if formatted_schema is None:
+                custom_properties["schema"] = dataset.schema
+
+            # Handle local/code datasets
+            if source_type in ("local", "code"):
+                local_dataset = Dataset(
+                    platform=self.platform,
+                    name=dataset.name,
+                    schema=formatted_schema,
+                    custom_properties=custom_properties,
+                )
+                yield from local_dataset.as_workunits()
+                dataset_reference_urns.append(local_dataset.urn)
+                continue
+
+            # Handle hosted datasets
+            formatted_platform = self._get_dataset_platform_from_source_type(
+                source_type
+            )
+
+            # Validate platform if materialization is enabled
+            if self.config.materialize_dataset_inputs:
+                if not formatted_platform:
+                    self.report.failure(
+                        title="Unable to materialize dataset inputs",
+                        message=f"No mapping dataPlatform found for dataset input source type '{source_type}'",
+                        context=f"please add `materialize_dataset_inputs.source_mapping_to_platform` in config "
+                        f"(e.g. '{source_type}': 'snowflake')",
+                    )
+                    continue
+                # Create hosted dataset
+                hosted_dataset = Dataset(
+                    platform=formatted_platform,
+                    name=dataset.name,
+                    schema=formatted_schema,
+                    custom_properties=dataset_tags,
+                )
+                yield from hosted_dataset.as_workunits()
+
+            # Create dataset reference with upstream
+            hosted_dataset_reference = Dataset(
+                platform=self.platform,
+                name=dataset.name,
+                schema=formatted_schema,
+                custom_properties=dataset_tags,
+                upstreams=UpstreamLineageClass(
+                    upstreams=[
+                        UpstreamClass(
+                            self._get_external_dataset_urn(
+                                formatted_platform, dataset.name
+                            ),
+                            type="COPY",
+                        )
+                    ]
+                )
+                if formatted_platform
+                else None,
+            )
+            dataset_reference_urns.append(hosted_dataset_reference.urn)
+            yield from hosted_dataset_reference.as_workunits()
+
+        # Add dataset references as upstreams for the run
+        if dataset_reference_urns:
+            input_edges = [
+                EdgeClass(destinationUrn=str(dataset_ref_urn))
+                for dataset_ref_urn in dataset_reference_urns
+            ]
+            yield MetadataChangeProposalWrapper(
+                entityUrn=str(run_urn),
+                aspect=DataProcessInstanceInputClass(inputs=[], inputEdges=input_edges),
+            ).as_workunit()
+
+    def _get_dataset_platform_from_source_type(self, source_type: str) -> Optional[str]:
+        """
+        Map MLflow source type to DataHub platform.
+
+        Priority:
+        1. User-provided mapping in config
+        2. Internal mapping
+        3. Direct platform match from list of supported platforms
+        """
+        source_type = source_type.lower()
+
+        # User-provided mapping
+        platform = self._get_platform_from_user_mapping(source_type)
+        if platform:
+            return platform
+
+        # Internal mapping
+        if source_type == "gs":
+            return "gcs"
+
+        # Check direct platform match
+        if self._is_valid_platform(source_type):
+            return source_type
+
+        return None
+
+    def _get_platform_from_user_mapping(self, source_type: str) -> Optional[str]:
+        """
+        Get platform from user-provided mapping in config.
+        Returns None if mapping is invalid or platform is not supported.
+        """
+        source_mapping = self.config.source_mapping_to_platform
+        if not source_mapping:
+            return None
+
+        platform = source_mapping.get(source_type)
+        if not platform:
+            return None
+
+        return platform
+
     def _get_run_workunits(
         self, experiment: Experiment, run: Run
     ) -> Iterable[MetadataWorkUnit]:
-        experiment_key = ContainerKeyWithId(
+        experiment_key = ExperimentKey(
             platform=str(DataPlatformUrn(self.platform)), id=experiment.name
         )
 
@@ -385,8 +600,8 @@ class MLflowSource(StatefulIngestionSourceBase):
         )
         return runs
 
-    @staticmethod
     def _traverse_mlflow_search_func(
+        self,
         search_func: Callable[..., PagedList[T]],
         **kwargs: Any,
     ) -> Iterable[T]:
@@ -394,12 +609,24 @@ class MLflowSource(StatefulIngestionSourceBase):
         Utility to traverse an MLflow search_* functions which return PagedList.
         """
         next_page_token = None
-        while True:
-            paged_list = search_func(page_token=next_page_token, **kwargs)
-            yield from paged_list.to_list()
-            next_page_token = paged_list.token
-            if not next_page_token:
+        try:
+            while True:
+                paged_list = search_func(page_token=next_page_token, **kwargs)
+                yield from paged_list.to_list()
+                next_page_token = paged_list.token
+                if not next_page_token:
+                    return
+        except MlflowException as e:
+            if e.error_code == "ENDPOINT_NOT_FOUND":
+                self.report.warning(
+                    title="MLflow API Endpoint Not Found for Experiments.",
+                    message="Please upgrade to version 1.28.0 or higher to ensure compatibility. Skipping ingestion for experiments and runs.",
+                    context=None,
+                    exc=e,
+                )
                 return
+            else:
+                raise  # Only re-raise other exceptions
 
     def _get_latest_version(self, registered_model: RegisteredModel) -> Optional[str]:
         return (
@@ -659,6 +886,10 @@ class MLflowSource(StatefulIngestionSourceBase):
         )
         return wu
 
+    def _is_valid_platform(self, platform: Optional[str]) -> bool:
+        """Check if platform is registered as a source plugin"""
+        return platform in KNOWN_VALID_PLATFORM_NAMES
+
     @classmethod
     def create(cls, config_dict: dict, ctx: PipelineContext) -> "MLflowSource":
         config = MLflowConfig.parse_obj(config_dict)