def test_get_store_file_store_from_env(tmp_wkdir, uri):
env = {_TRACKING_URI_ENV_VAR: uri}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, FileStore)
assert os.path.abspath(
store.root_directory) == os.path.abspath("other/path")
def test_get_store_basic_rest_store():
env = {_TRACKING_URI_ENV_VAR: "https://my-tracking-server:5050"}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
assert store.get_host_creds().host == "https://my-tracking-server:5050"
assert store.get_host_creds().token is None
def test_get_store_file_store_from_arg(tmp_wkdir):
env = {}
with mock.patch.dict(os.environ, env):
store = _get_store("other/path")
assert isinstance(store, FileStore)
assert os.path.abspath(
store.root_directory) == os.path.abspath("other/path")
def _get_host_creds_from_default_store():
store = utils._get_store()
if not isinstance(store, RestStore):
raise MlflowException(
'Failed to get credentials for DBFS; they are read from the ' +
'Databricks CLI credentials or MLFLOW_TRACKING* environment ' +
'variables.')
return store.get_host_creds
def test_get_store_rest_store_with_no_insecure():
env = {
_TRACKING_URI_ENV_VAR: "https://my-tracking-server:5050",
_TRACKING_INSECURE_TLS_ENV_VAR: "false",
}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
assert not store.get_host_creds().ignore_tls_verification
# By default, should not ignore verification.
env = {
_TRACKING_URI_ENV_VAR: "https://my-tracking-server:5050",
}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
assert not store.get_host_creds().ignore_tls_verification
def test_get_store_rest_store_with_token():
env = {
_TRACKING_URI_ENV_VAR: "https://my-tracking-server:5050",
_TRACKING_TOKEN_ENV_VAR: "my-token",
}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
assert store.get_host_creds().token == "my-token"
def test_get_store_databricks():
env = {
_TRACKING_URI_ENV_VAR: "databricks",
'DATABRICKS_HOST': "https://my-tracking-server",
'DATABRICKS_TOKEN': "abcdef",
}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
assert store.get_host_creds().host == "https://my-tracking-server"
assert store.get_host_creds().token == "abcdef"
def test_get_store_rest_store_with_password():
env = {
_TRACKING_URI_ENV_VAR: "https://my-tracking-server:5050",
_TRACKING_USERNAME_ENV_VAR: "Bob",
_TRACKING_PASSWORD_ENV_VAR: "Ross",
}
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
assert store.get_host_creds().host == "https://my-tracking-server:5050"
assert store.get_host_creds().username == "Bob"
assert store.get_host_creds().password == "Ross"
def test_get_store_databricks_profile():
env = {
_TRACKING_URI_ENV_VAR: "databricks://mycoolprofile",
}
# It's kind of annoying to setup a profile, and we're not really trying to test
# that anyway, so just check if we raise a relevant exception.
with mock.patch.dict(os.environ, env):
store = _get_store()
assert isinstance(store, RestStore)
with pytest.raises(Exception) as e_info:
store.get_host_creds()
assert 'mycoolprofile' in str(e_info.value)
def test_get_store_sqlalchemy_store(tmp_wkdir, db_type):
patch_create_engine = mock.patch("sqlalchemy.create_engine")
uri = "{}://hostname/database".format(db_type)
env = {_TRACKING_URI_ENV_VAR: uri}
with mock.patch.dict(os.environ, env), patch_create_engine as mock_create_engine,\
mock.patch("mlflow.store.db.utils._verify_schema"), \
mock.patch("mlflow.store.db.utils._initialize_tables"):
store = _get_store()
assert isinstance(store, SqlAlchemyStore)
assert store.db_uri == uri
assert store.artifact_root_uri == "./mlruns"
mock_create_engine.assert_called_once_with(uri, pool_pre_ping=True)
def __init__(self, tracking_uri):
"""
:param tracking_uri: Address of local or remote tracking server.
"""
self.tracking_uri = tracking_uri
self.store = utils._get_store(self.tracking_uri)