yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-12 18:47:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ingest/mlflow): add mlflow auth config (#12984)

Browse Source
This commit is contained in:
Hyejin Yoon 2025-03-27 11:30:30 +09:00 committed by GitHub
parent fbeedac497
commit 55cf62b49f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 19 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
metadata-ingestion/src/datahub/ingestion/source/mlflow.py
View File

@ -1,4 +1,5 @@
import json
import os
import time
from dataclasses import dataclass
from typing import Any, Callable, Iterable, List, Optional, Tuple, TypeVar, Union
@ -115,6 +116,13 @@ class MLflowConfig(StatefulIngestionConfigBase, EnvConfigMixin):
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"
)
@dataclass
class MLflowRegisteredModelStageInfo:
@ -161,7 +169,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,
)