def _init_server(backend_uri, root_artifact_uri):
"""
Launch a new REST server using the tracking store specified by backend_uri and root artifact
directory specified by root_artifact_uri.
:returns A tuple (url, process) containing the string URL of the server and a handle to the
server process (a multiprocessing.Process object).
"""
kiwi.set_tracking_uri(None)
server_port = get_safe_port()
env = {
BACKEND_STORE_URI_ENV_VAR:
backend_uri,
ARTIFACT_ROOT_ENV_VAR:
path_to_local_file_uri(
tempfile.mkdtemp(dir=local_file_uri_to_path(root_artifact_uri))),
}
with mock.patch.dict(os.environ, env):
cmd = [
"python", "-c",
'from mlflow.server import app; app.run("{hostname}", {port})'.
format(hostname=LOCALHOST, port=server_port)
]
process = Popen(cmd)
_await_server_up_or_die(server_port)
url = "http://{hostname}:{port}".format(hostname=LOCALHOST,
port=server_port)
print("Launching tracking server against backend URI %s. Server URL: %s" %
(backend_uri, url))
return url, process
def test_log_artifact_with_dirs(tmpdir):
# Test log artifact with a directory
art_dir = tmpdir.mkdir("parent")
file0 = art_dir.join("file0")
file0.write("something")
file1 = art_dir.join("file1")
file1.write("something")
sub_dir = art_dir.mkdir("child")
with start_run():
artifact_uri = kiwi.get_artifact_uri()
run_artifact_dir = local_file_uri_to_path(artifact_uri)
kiwi.log_artifact(str(art_dir))
base = os.path.basename(str(art_dir))
assert os.listdir(run_artifact_dir) == [base]
assert set(os.listdir(os.path.join(run_artifact_dir, base))) == \
{'child', 'file0', 'file1'}
with open(os.path.join(run_artifact_dir, base, "file0")) as f:
assert f.read() == "something"
# Test log artifact with directory and specified parent folder
art_dir = tmpdir.mkdir("dir")
with start_run():
artifact_uri = kiwi.get_artifact_uri()
run_artifact_dir = local_file_uri_to_path(artifact_uri)
kiwi.log_artifact(str(art_dir), "some_parent")
assert os.listdir(run_artifact_dir) == [
os.path.basename("some_parent")
]
assert os.listdir(os.path.join(run_artifact_dir, "some_parent")) == \
[os.path.basename(str(art_dir))]
sub_dir = art_dir.mkdir("another_dir")
with start_run():
artifact_uri = kiwi.get_artifact_uri()
run_artifact_dir = local_file_uri_to_path(artifact_uri)
kiwi.log_artifact(str(art_dir), "parent/and_child")
assert os.listdir(os.path.join(run_artifact_dir, "parent", "and_child")) == \
[os.path.basename(str(art_dir))]
assert os.listdir(os.path.join(run_artifact_dir,
"parent", "and_child",
os.path.basename(str(art_dir)))) == \
[os.path.basename(str(sub_dir))]
def test_log_artifact():
artifact_src_dir = tempfile.mkdtemp()
# Create artifacts
_, path0 = tempfile.mkstemp(dir=artifact_src_dir)
_, path1 = tempfile.mkstemp(dir=artifact_src_dir)
for i, path in enumerate([path0, path1]):
with open(path, "w") as handle:
handle.write("%s" % str(i))
# Log an artifact, verify it exists in the directory returned by get_artifact_uri
# after the run finishes
artifact_parent_dirs = ["some_parent_dir", None]
for parent_dir in artifact_parent_dirs:
with start_run():
artifact_uri = kiwi.get_artifact_uri()
run_artifact_dir = local_file_uri_to_path(artifact_uri)
kiwi.log_artifact(path0, parent_dir)
expected_dir = os.path.join(run_artifact_dir, parent_dir) \
if parent_dir is not None else run_artifact_dir
assert os.listdir(expected_dir) == [os.path.basename(path0)]
logged_artifact_path = os.path.join(expected_dir, path0)
assert filecmp.cmp(logged_artifact_path, path0, shallow=False)
# Log multiple artifacts, verify they exist in the directory returned by get_artifact_uri
for parent_dir in artifact_parent_dirs:
with start_run():
artifact_uri = kiwi.get_artifact_uri()
run_artifact_dir = local_file_uri_to_path(artifact_uri)
kiwi.log_artifacts(artifact_src_dir, parent_dir)
# Check that the logged artifacts match
expected_artifact_output_dir = os.path.join(run_artifact_dir, parent_dir) \
if parent_dir is not None else run_artifact_dir
dir_comparison = filecmp.dircmp(artifact_src_dir,
expected_artifact_output_dir)
assert len(dir_comparison.left_only) == 0
assert len(dir_comparison.right_only) == 0
assert len(dir_comparison.diff_files) == 0
assert len(dir_comparison.funny_files) == 0
def __init__(self, db_uri, default_artifact_root):
"""
Create a database backed store.
:param db_uri: The SQLAlchemy database URI string to connect to the database. See
the `SQLAlchemy docs
<https://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls>`_
for format specifications. Mlflow supports the dialects ``mysql``,
``mssql``, ``sqlite``, and ``postgresql``.
:param default_artifact_root: Path/URI to location suitable for large data (such as a blob
store object, DBFS path, or shared NFS file system).
"""
super(SqlAlchemyStore, self).__init__()
self.db_uri = db_uri
self.db_type = extract_db_type_from_uri(db_uri)
self.artifact_root_uri = default_artifact_root
self.engine = kiwi.store.db.utils.create_sqlalchemy_engine(db_uri)
# On a completely fresh MLflow installation against an empty database (verify database
# emptiness by checking that 'experiments' etc aren't in the list of table names), run all
# DB migrations
expected_tables = [
SqlExperiment.__tablename__,
SqlRun.__tablename__,
SqlMetric.__tablename__,
SqlParam.__tablename__,
SqlTag.__tablename__,
SqlExperimentTag.__tablename__,
SqlLatestMetric.__tablename__,
]
inspected_tables = set(
sqlalchemy.inspect(self.engine).get_table_names())
if any([table not in inspected_tables for table in expected_tables]):
kiwi.store.db.utils._initialize_tables(self.engine)
Base.metadata.bind = self.engine
SessionMaker = sqlalchemy.orm.sessionmaker(bind=self.engine)
self.ManagedSessionMaker = kiwi.store.db.utils._get_managed_session_maker(
SessionMaker, self.db_type)
kiwi.store.db.utils._verify_schema(self.engine)
if is_local_uri(default_artifact_root):
mkdir(local_file_uri_to_path(default_artifact_root))
if len(self.list_experiments()) == 0:
with self.ManagedSessionMaker() as session:
self._create_default_experiment(session)
def __init__(self, root_directory=None, artifact_root_uri=None):
"""
Create a new FileStore with the given root directory and a given default artifact root URI.
"""
super(FileStore, self).__init__()
self.root_directory = local_file_uri_to_path(root_directory
or _default_root_dir())
self.artifact_root_uri = artifact_root_uri or path_to_local_file_uri(
self.root_directory)
self.trash_folder = os.path.join(self.root_directory,
FileStore.TRASH_FOLDER_NAME)
# Create root directory if needed
if not exists(self.root_directory):
mkdir(self.root_directory)
self._create_experiment_with_id(
name=Experiment.DEFAULT_EXPERIMENT_NAME,
experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
artifact_uri=None)
# Create trash folder if needed
if not exists(self.trash_folder):
mkdir(self.trash_folder)
def __init__(self, *args, **kwargs):
super(LocalArtifactRepository, self).__init__(*args, **kwargs)
self._artifact_dir = local_file_uri_to_path(self.artifact_uri)
def delete_artifacts(self, artifact_path=None):
artifact_path = os.path.join(self._artifact_dir, artifact_path) if artifact_path \
else self._artifact_dir
shutil.rmtree(local_file_uri_to_path(artifact_path))