wehub-resource-sync
30 lines
936 B
Python
30 lines
936 B
Python
import urllib.parse
|
|
|
|
from mlflow.environment_variables import MLFLOW_UC_OSS_TOKEN
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.utils.databricks_utils import get_databricks_host_creds
|
|
from mlflow.utils.rest_utils import MlflowHostCreds
|
|
from mlflow.utils.uri import (
|
|
_DATABRICKS_UNITY_CATALOG_SCHEME,
|
|
)
|
|
|
|
|
|
def get_oss_host_creds(server_uri=None):
|
|
"""
|
|
Retrieve the host credentials for the OSS server.
|
|
|
|
Args:
|
|
server_uri (str): The URI of the server.
|
|
|
|
Returns:
|
|
MlflowHostCreds: The host credentials for the OSS server.
|
|
"""
|
|
parsed_uri = urllib.parse.urlparse(server_uri)
|
|
|
|
if parsed_uri.scheme != "uc":
|
|
raise MlflowException("The scheme of the server_uri should be 'uc'")
|
|
|
|
if parsed_uri.path == _DATABRICKS_UNITY_CATALOG_SCHEME:
|
|
return get_databricks_host_creds(parsed_uri.path)
|
|
return MlflowHostCreds(host=parsed_uri.path, token=MLFLOW_UC_OSS_TOKEN.get())
|