def test_resolve_tracking_uri_with_no_param():
with mock.patch(
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
) as get_tracking_uri_mock:
default_uri = "databricks://tracking_zlkjdas"
get_tracking_uri_mock.return_value = default_uri
assert _resolve_tracking_uri() == default_uri
def test_resolve_tracking_uri_with_param():
with mock.patch(
"mlflow.tracking._tracking_service.utils.get_tracking_uri"
) as get_tracking_uri_mock:
get_tracking_uri_mock.return_value = "databricks://tracking_qoeirj"
overriding_uri = "databricks://tracking_poiwerow"
assert _resolve_tracking_uri(overriding_uri) == overriding_uri
def __init__(self, tracking_uri=None, registry_uri=None):
"""
:param tracking_uri: Address of local or remote tracking server. If not provided, defaults
to the service set by ``mlflow.tracking.set_tracking_uri``. See
`Where Runs Get Recorded <../tracking.html#where-runs-get-recorded>`_
for more info.
:param registry_uri: Address of local or remote model registry server. If not provided,
defaults to the service set by ``mlflow.tracking.set_registry_uri``. If
no such service was set, defaults to the tracking uri of the client.
"""
final_tracking_uri = utils._resolve_tracking_uri(tracking_uri)
self._registry_uri = registry_utils._resolve_registry_uri(registry_uri, tracking_uri)
self._tracking_client = TrackingServiceClient(final_tracking_uri)
def get_store(self, store_uri=None, artifact_uri=None):
"""Get a store from the registry based on the scheme of store_uri
:param store_uri: The store URI. If None, it will be inferred from the environment. This URI
is used to select which tracking store implementation to instantiate and
is passed to the constructor of the implementation.
:param artifact_uri: Artifact repository URI. Passed through to the tracking store
implementation.
:return: An instance of `mlflow.store.tracking.AbstractStore` that fulfills the store URI
requirements.
"""
from mlflow.tracking._tracking_service import utils
store_uri = utils._resolve_tracking_uri(store_uri)
builder = self.get_store_builder(store_uri)
return builder(store_uri=store_uri, artifact_uri=artifact_uri)
def _resolve_registry_uri(registry_uri=None, tracking_uri=None):
return registry_uri or _get_registry_uri_from_context(
) or _resolve_tracking_uri(tracking_uri)
