def test_list_artifacts(ftp_mock): artifact_root_path = "/experiment_id/run_id/" repo = FTPArtifactRepository("ftp://test_ftp" + artifact_root_path) repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) # mocked file structure # |- file # |- model # |- model.pb file_path = "file" file_size = 678 dir_path = "model" ftp_mock.cwd = MagicMock(side_effect=[None, ftplib.error_perm, None]) ftp_mock.nlst = MagicMock(return_value=[file_path, dir_path]) ftp_mock.size = MagicMock(return_value=file_size) artifacts = repo.list_artifacts(path=None) ftp_mock.nlst.assert_called_once_with(artifact_root_path) ftp_mock.size.assert_called_once_with(artifact_root_path + file_path) assert len(artifacts) == 2 assert artifacts[0].path == file_path assert artifacts[0].is_dir is False assert artifacts[0].file_size == file_size assert artifacts[1].path == dir_path assert artifacts[1].is_dir is True assert artifacts[1].file_size is None
def from_artifact_uri(artifact_uri, store): """ Given an artifact URI for an Experiment Run (e.g., /local/file/path or s3://my/bucket), returns an ArtifactReposistory instance capable of logging and downloading artifacts on behalf of this URI. :param store: An instance of AbstractStore which the artifacts are registered in. """ if artifact_uri.startswith("s3:/"): # Import these locally to avoid creating a circular import loop from mlflow.store.s3_artifact_repo import S3ArtifactRepository return S3ArtifactRepository(artifact_uri) elif artifact_uri.startswith("gs:/"): from mlflow.store.gcs_artifact_repo import GCSArtifactRepository return GCSArtifactRepository(artifact_uri) elif artifact_uri.startswith("wasbs:/"): from mlflow.store.azure_blob_artifact_repo import AzureBlobArtifactRepository return AzureBlobArtifactRepository(artifact_uri) elif artifact_uri.startswith("ftp:/"): from mlflow.store.ftp_artifact_repo import FTPArtifactRepository return FTPArtifactRepository(artifact_uri) elif artifact_uri.startswith("sftp:/"): from mlflow.store.sftp_artifact_repo import SFTPArtifactRepository return SFTPArtifactRepository(artifact_uri) elif artifact_uri.startswith("dbfs:/"): from mlflow.store.dbfs_artifact_repo import DbfsArtifactRepository if not isinstance(store, RestStore): raise MlflowException( '`store` must be an instance of RestStore.') return DbfsArtifactRepository(artifact_uri, store.get_host_creds) else: from mlflow.store.local_artifact_repo import LocalArtifactRepository return LocalArtifactRepository(artifact_uri)
def test_download_artifacts(ftp_mock): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) # mocked file structure # |- model # |- model.pb # |- variables # |- test.txt dir_name = 'model' # list artifacts at sub directory level file_path = 'model.pb' subdir_name = 'variables' subfile_name = 'test.txt' ftp_dir_name = '/some/path/' + dir_name ftp_mock.cwd = MagicMock(side_effect=[None, None, ftplib.error_perm, None, None, ftplib.error_perm]) ftp_mock.nlst = MagicMock(side_effect=[[file_path, subdir_name], [subfile_name]]) repo.download_artifacts("model") cwd_call_args = sorted([ftp_mock.cwd.call_args_list[i][0][0] for i in range(6)]) expect_args = [ftp_dir_name, ftp_dir_name, os.path.join(ftp_dir_name, file_path), os.path.join(ftp_dir_name, subdir_name), os.path.join(ftp_dir_name, subdir_name), os.path.join(ftp_dir_name, subdir_name, subfile_name)] assert cwd_call_args == expect_args assert ftp_mock.nlst.call_count == 2 assert ftp_mock.retrbinary.call_args_list[0][0][0] == 'RETR ' + file_path assert ftp_mock.retrbinary.call_args_list[1][0][0] == 'RETR ' + subfile_name
def test_list_artifacts_empty(ftp_mock): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) ftp_mock.nlst = MagicMock(return_value=[]) assert repo.list_artifacts() == [] ftp_mock.nlst.assert_called_once_with("/some/path")
def test_download_artifacts(ftp_mock): artifact_root_path = "/some/path" repo = FTPArtifactRepository("ftp://test_ftp" + artifact_root_path) repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) # mocked file structure # |- model # |- model.pb # |- variables # |- test.txt dir_path = os.path.join(artifact_root_path, 'model') # list artifacts at sub directory level model_file_path_sub = 'model.pb' model_file_path_full = os.path.join(dir_path, model_file_path_sub) subdir_name = 'variables' subdir_path_full = os.path.join(dir_path, subdir_name) subfile_name = 'test.txt' subfile_path_full = os.path.join(artifact_root_path, subdir_path_full, subfile_name) is_dir_mapping = { dir_path: True, model_file_path_full: False, subdir_path_full: True, subfile_path_full: False, } is_dir_call_args = [ dir_path, model_file_path_full, subdir_path_full, model_file_path_full, subdir_path_full, subfile_path_full, subfile_path_full ] cwd_side_effect = [ None if is_dir_mapping[call_arg] else ftplib.error_perm for call_arg in is_dir_call_args ] ftp_mock.cwd = MagicMock(side_effect=cwd_side_effect) ftp_mock.nlst = MagicMock( side_effect=[[model_file_path_sub, subdir_name], [subfile_name]]) repo.download_artifacts("model") cwd_call_args = [ arg_entry[0][0] for arg_entry in ftp_mock.cwd.call_args_list ] assert cwd_call_args == is_dir_call_args assert ftp_mock.nlst.call_count == 2 assert ftp_mock.retrbinary.call_args_list[0][0][ 0] == 'RETR ' + model_file_path_full assert ftp_mock.retrbinary.call_args_list[1][0][ 0] == 'RETR ' + subfile_path_full
def test_download_artifacts_single(ftp_mock): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) ftp_mock.cwd = MagicMock(side_effect=ftplib.error_perm) repo.download_artifacts("test.txt") ftp_mock.retrbinary.assert_called_once() assert ftp_mock.retrbinary.call_args_list[0][0][0] == 'RETR /some/path/test.txt'
def test_log_artifact_reuse_ftp_client(ftp_mock, tmpdir): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) d = tmpdir.mkdir("data") file = d.join("test.txt") file.write("hello world!") fpath = file.strpath repo.log_artifact(fpath) repo.log_artifact(fpath, "subdir1/subdir2") repo.log_artifact(fpath, "subdir3") assert repo.get_ftp_client.call_count == 3
def test_log_artifact_multiple_calls(ftp_mock, tmpdir): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) d = tmpdir.mkdir("data") file1 = d.join("test1.txt") file1.write("hello world!") fpath1 = d + '/test1.txt' fpath1 = fpath1.strpath file2 = d.join("test2.txt") file2.write("hello world!") fpath2 = d + '/test2.txt' fpath2 = fpath2.strpath ftp_mock.cwd = MagicMock(side_effect=[ ftplib.error_perm, None, ftplib.error_perm, None, None, None ]) repo.log_artifact(fpath1) ftp_mock.mkd.assert_called_once_with('/some/path') ftp_mock.cwd.assert_called_with('/some/path') ftp_mock.storbinary.assert_called() assert ftp_mock.storbinary.call_args_list[0][0][0] == 'STOR test1.txt' ftp_mock.reset_mock() repo.log_artifact(fpath1, "subdir") ftp_mock.mkd.assert_called_once_with('/some/path/subdir') ftp_mock.cwd.assert_called_with('/some/path/subdir') ftp_mock.storbinary.assert_called() assert ftp_mock.storbinary.call_args_list[0][0][0] == 'STOR test1.txt' ftp_mock.reset_mock() repo.log_artifact(fpath2) ftp_mock.mkd.assert_not_called() ftp_mock.cwd.assert_called_with('/some/path') ftp_mock.storbinary.assert_called() assert ftp_mock.storbinary.call_args_list[0][0][0] == 'STOR test2.txt'
def test_log_artifact(ftp_mock, tmpdir): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) d = tmpdir.mkdir("data") f = d.join("test.txt") f.write("hello world!") fpath = d + '/test.txt' fpath = fpath.strpath repo.log_artifact(fpath) ftp_mock.mkd.assert_called_once_with('/some/path') ftp_mock.cwd.assert_called_once_with('/some/path') ftp_mock.storbinary.assert_called_once() assert ftp_mock.storbinary.call_args_list[0][0][0] == 'STOR test.txt'
def test_log_artifacts(ftp_mock, tmpdir): repo = FTPArtifactRepository("ftp://test_ftp/some/path") repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) subd = tmpdir.mkdir("data").mkdir("subdir") subd.join("a.txt").write("A") subd.join("b.txt").write("B") subd.join("c.txt").write("C") repo.log_artifacts(subd.strpath) ftp_mock.mkd.assert_any_call('/some/path/subdir') ftp_mock.cwd.assert_any_call('/some/path/subdir') assert ftp_mock.storbinary.call_count == 3 storbinary_call_args = sorted([ftp_mock.storbinary.call_args_list[i][0][0] for i in range(3)]) assert storbinary_call_args == ['STOR a.txt', 'STOR b.txt', 'STOR c.txt']
def test_list_artifacts_with_subdir(ftp_mock): artifact_root_path = "/experiment_id/run_id/" repo = FTPArtifactRepository("sftp://test_sftp" + artifact_root_path) repo.get_ftp_client = MagicMock() call_mock = MagicMock(return_value=ftp_mock) repo.get_ftp_client.return_value = MagicMock(__enter__=call_mock) # mocked file structure # |- model # |- model.pb # |- variables dir_name = 'model' # list artifacts at sub directory level file_path = 'model.pb' file_size = 345 subdir_name = 'variables' ftp_mock.nlst = MagicMock(return_value=[file_path, subdir_name]) ftp_mock.cwd = MagicMock(side_effect=[None, ftplib.error_perm, None]) ftp_mock.size = MagicMock(return_value=file_size) artifacts = repo.list_artifacts(path=dir_name) ftp_mock.nlst.assert_called_once_with(artifact_root_path + dir_name) ftp_mock.size.assert_called_once_with(artifact_root_path + dir_name + '/' + file_path) assert len(artifacts) == 2 assert artifacts[0].path == dir_name + '/' + file_path assert artifacts[0].is_dir is False assert artifacts[0].file_size == file_size assert artifacts[1].path == dir_name + '/' + subdir_name assert artifacts[1].is_dir is True assert artifacts[1].file_size is None