wehub-resource-sync
987 lines
32 KiB
Python
987 lines
32 KiB
Python
import atexit
|
|
import codecs
|
|
import errno
|
|
import fnmatch
|
|
import gzip
|
|
import importlib.util
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import pathlib
|
|
import posixpath
|
|
import shutil
|
|
import stat
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
from concurrent.futures import as_completed
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from subprocess import CalledProcessError, TimeoutExpired
|
|
from types import TracebackType
|
|
from typing import Any
|
|
from urllib.parse import unquote
|
|
from urllib.request import pathname2url
|
|
|
|
from mlflow.entities import FileInfo
|
|
from mlflow.environment_variables import (
|
|
_MLFLOW_MPD_NUM_RETRIES,
|
|
_MLFLOW_MPD_RETRY_INTERVAL_SECONDS,
|
|
MLFLOW_DOWNLOAD_CHUNK_TIMEOUT,
|
|
MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_artifacts_pb2 import ArtifactCredentialType
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.utils import download_cloud_file_chunk
|
|
from mlflow.utils.databricks_utils import (
|
|
get_databricks_local_temp_dir,
|
|
get_databricks_nfs_temp_dir,
|
|
)
|
|
from mlflow.utils.os import is_windows
|
|
from mlflow.utils.process import cache_return_value_per_process
|
|
from mlflow.utils.request_utils import cloud_storage_http_request, download_chunk
|
|
from mlflow.utils.rest_utils import augmented_raise_for_status
|
|
|
|
ENCODING = "utf-8"
|
|
_PROGRESS_BAR_DISPLAY_THRESHOLD = 500_000_000 # 500 MB
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
# This is for backward compatibility with databricks-feature-engineering<=0.10.2
|
|
if importlib.util.find_spec("yaml") is not None:
|
|
try:
|
|
from yaml import CSafeDumper as YamlSafeDumper
|
|
except ImportError:
|
|
from yaml import SafeDumper as YamlSafeDumper # noqa: F401
|
|
|
|
|
|
class ArtifactProgressBar:
|
|
def __init__(self, desc, total, step, **kwargs) -> None:
|
|
self.desc = desc
|
|
self.total = total
|
|
self.step = step
|
|
self.pbar = None
|
|
self.progress = 0
|
|
self.kwargs = kwargs
|
|
|
|
def set_pbar(self):
|
|
if MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR.get():
|
|
try:
|
|
from tqdm.auto import tqdm
|
|
|
|
self.pbar = tqdm(total=self.total, desc=self.desc, **self.kwargs)
|
|
except ImportError:
|
|
pass
|
|
|
|
@classmethod
|
|
def chunks(cls, file_size, desc, chunk_size):
|
|
bar = cls(
|
|
desc,
|
|
total=file_size,
|
|
step=chunk_size,
|
|
unit="iB",
|
|
unit_scale=True,
|
|
unit_divisor=1024,
|
|
miniters=1,
|
|
)
|
|
if file_size >= _PROGRESS_BAR_DISPLAY_THRESHOLD:
|
|
bar.set_pbar()
|
|
return bar
|
|
|
|
@classmethod
|
|
def files(cls, desc, total):
|
|
bar = cls(desc, total=total, step=1)
|
|
bar.set_pbar()
|
|
return bar
|
|
|
|
def update(self):
|
|
if self.pbar:
|
|
update_step = min(self.total - self.progress, self.step)
|
|
self.pbar.update(update_step)
|
|
self.pbar.refresh()
|
|
self.progress += update_step
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args):
|
|
if self.pbar:
|
|
self.pbar.close()
|
|
|
|
|
|
def is_directory(name):
|
|
return os.path.isdir(name)
|
|
|
|
|
|
def is_file(name):
|
|
return os.path.isfile(name)
|
|
|
|
|
|
def exists(name):
|
|
return os.path.exists(name)
|
|
|
|
|
|
def list_all(root, filter_func=lambda x: True, full_path=False):
|
|
"""List all entities directly under 'dir_name' that satisfy 'filter_func'
|
|
|
|
Args:
|
|
root: Name of directory to start search.
|
|
filter_func: function or lambda that takes path.
|
|
full_path: If True will return results as full path including `root`.
|
|
|
|
Returns:
|
|
list of all files or directories that satisfy the criteria.
|
|
|
|
"""
|
|
if not is_directory(root):
|
|
raise Exception(f"Invalid parent directory '{root}'")
|
|
matches = [x for x in os.listdir(root) if filter_func(os.path.join(root, x))]
|
|
return [os.path.join(root, m) for m in matches] if full_path else matches
|
|
|
|
|
|
def list_subdirs(dir_name, full_path=False):
|
|
"""
|
|
Equivalent to UNIX command:
|
|
``find $dir_name -depth 1 -type d``
|
|
|
|
Args:
|
|
dir_name: Name of directory to start search.
|
|
full_path: If True will return results as full path including `root`.
|
|
|
|
Returns:
|
|
list of all directories directly under 'dir_name'.
|
|
"""
|
|
return list_all(dir_name, os.path.isdir, full_path)
|
|
|
|
|
|
def list_files(dir_name, full_path=False):
|
|
"""
|
|
Equivalent to UNIX command:
|
|
``find $dir_name -depth 1 -type f``
|
|
|
|
Args:
|
|
dir_name: Name of directory to start search.
|
|
full_path: If True will return results as full path including `root`.
|
|
|
|
Returns:
|
|
list of all files directly under 'dir_name'.
|
|
"""
|
|
return list_all(dir_name, os.path.isfile, full_path)
|
|
|
|
|
|
def find(root, name, full_path=False):
|
|
"""Search for a file in a root directory. Equivalent to:
|
|
``find $root -name "$name" -depth 1``
|
|
|
|
Args:
|
|
root: Name of root directory for find.
|
|
name: Name of file or directory to find directly under root directory.
|
|
full_path: If True will return results as full path including `root`.
|
|
|
|
Returns:
|
|
list of matching files or directories.
|
|
"""
|
|
path_name = os.path.join(root, name)
|
|
return list_all(root, lambda x: x == path_name, full_path)
|
|
|
|
|
|
def mkdir(root, name=None):
|
|
"""Make directory with name "root/name", or just "root" if name is None.
|
|
|
|
Args:
|
|
root: Name of parent directory.
|
|
name: Optional name of leaf directory.
|
|
|
|
Returns:
|
|
Path to created directory.
|
|
"""
|
|
target = os.path.join(root, name) if name is not None else root
|
|
try:
|
|
os.makedirs(target, exist_ok=True)
|
|
except OSError as e:
|
|
if e.errno != errno.EEXIST or not os.path.isdir(target):
|
|
raise e
|
|
return target
|
|
|
|
|
|
def make_containing_dirs(path):
|
|
"""
|
|
Create the base directory for a given file path if it does not exist; also creates parent
|
|
directories.
|
|
"""
|
|
dir_name = os.path.dirname(path)
|
|
if not os.path.exists(dir_name):
|
|
os.makedirs(dir_name)
|
|
|
|
|
|
class TempDir:
|
|
def __init__(self, chdr=False, remove_on_exit=True):
|
|
self._dir = None
|
|
self._path = None
|
|
self._chdr = chdr
|
|
self._remove = remove_on_exit
|
|
|
|
def __enter__(self):
|
|
self._path = os.path.abspath(create_tmp_dir())
|
|
assert os.path.exists(self._path)
|
|
if self._chdr:
|
|
self._dir = os.path.abspath(os.getcwd())
|
|
os.chdir(self._path)
|
|
return self
|
|
|
|
def __exit__(self, tp, val, traceback):
|
|
if self._chdr and self._dir:
|
|
os.chdir(self._dir)
|
|
self._dir = None
|
|
if self._remove and os.path.exists(self._path):
|
|
shutil.rmtree(self._path)
|
|
|
|
assert not self._remove or not os.path.exists(self._path)
|
|
|
|
def path(self, *path):
|
|
return os.path.join("./", *path) if self._chdr else os.path.join(self._path, *path)
|
|
|
|
|
|
def read_file_lines(parent_path, file_name):
|
|
"""Return the contents of the file as an array where each element is a separate line.
|
|
|
|
Args:
|
|
parent_path: Full path to the directory that contains the file.
|
|
file_name: Leaf file name.
|
|
|
|
Returns:
|
|
All lines in the file as an array.
|
|
|
|
"""
|
|
file_path = os.path.join(parent_path, file_name)
|
|
with codecs.open(file_path, mode="r", encoding=ENCODING) as f:
|
|
return f.readlines()
|
|
|
|
|
|
def read_file(parent_path, file_name):
|
|
"""Return the contents of the file.
|
|
|
|
Args:
|
|
parent_path: Full path to the directory that contains the file.
|
|
file_name: Leaf file name.
|
|
|
|
Returns:
|
|
The contents of the file.
|
|
|
|
"""
|
|
file_path = os.path.join(parent_path, file_name)
|
|
with codecs.open(file_path, mode="r", encoding=ENCODING) as f:
|
|
return f.read()
|
|
|
|
|
|
def get_file_info(path, rel_path):
|
|
"""Returns file meta data : location, size, ... etc
|
|
|
|
Args:
|
|
path: Path to artifact.
|
|
rel_path: Relative path.
|
|
|
|
Returns:
|
|
`FileInfo` object
|
|
"""
|
|
if is_directory(path):
|
|
return FileInfo(rel_path, True, None)
|
|
else:
|
|
return FileInfo(rel_path, False, os.path.getsize(path))
|
|
|
|
|
|
def mv(target, new_parent):
|
|
shutil.move(target, new_parent)
|
|
|
|
|
|
def write_to(filename, data):
|
|
with codecs.open(filename, mode="w", encoding=ENCODING) as handle:
|
|
handle.write(data)
|
|
|
|
|
|
def append_to(filename, data):
|
|
with open(filename, "a") as handle:
|
|
handle.write(data)
|
|
|
|
|
|
def make_tarfile(output_filename, source_dir, archive_name, custom_filter=None):
|
|
# Helper for filtering out modification timestamps
|
|
def _filter_timestamps(tar_info):
|
|
tar_info.mtime = 0
|
|
return tar_info if custom_filter is None else custom_filter(tar_info)
|
|
|
|
unzipped_file_handle, unzipped_filename = tempfile.mkstemp()
|
|
try:
|
|
with tarfile.open(unzipped_filename, "w") as tar:
|
|
tar.add(source_dir, arcname=archive_name, filter=_filter_timestamps)
|
|
# When gzipping the tar, don't include the tar's filename or modification time in the
|
|
# zipped archive (see https://docs.python.org/3/library/gzip.html#gzip.GzipFile)
|
|
with (
|
|
gzip.GzipFile(
|
|
filename="", fileobj=open(output_filename, "wb"), mode="wb", mtime=0
|
|
) as gzipped_tar,
|
|
open(unzipped_filename, "rb") as tar,
|
|
):
|
|
gzipped_tar.write(tar.read())
|
|
finally:
|
|
os.close(unzipped_file_handle)
|
|
|
|
|
|
def _copy_project(src_path, dst_path=""):
|
|
"""Internal function used to copy MLflow project during development.
|
|
|
|
Copies the content of the whole directory tree except patterns defined in .dockerignore.
|
|
The MLflow is assumed to be accessible as a local directory in this case.
|
|
|
|
Args:
|
|
src_path: Path to the original MLflow project
|
|
dst_path: MLflow will be copied here
|
|
|
|
Returns:
|
|
Name of the MLflow project directory.
|
|
"""
|
|
|
|
def _docker_ignore(mlflow_root):
|
|
docker_ignore = os.path.join(mlflow_root, ".dockerignore")
|
|
patterns = []
|
|
if os.path.exists(docker_ignore):
|
|
with open(docker_ignore) as f:
|
|
patterns = [x.strip() for x in f]
|
|
|
|
def ignore(_, names):
|
|
res = set()
|
|
for p in patterns:
|
|
res.update(set(fnmatch.filter(names, p)))
|
|
return list(res)
|
|
|
|
return ignore if patterns else None
|
|
|
|
mlflow_dir = "mlflow-project"
|
|
# check if we have project root
|
|
assert os.path.isfile(os.path.join(src_path, "pyproject.toml")), "file not found " + str(
|
|
os.path.abspath(os.path.join(src_path, "pyproject.toml"))
|
|
)
|
|
shutil.copytree(src_path, os.path.join(dst_path, mlflow_dir), ignore=_docker_ignore(src_path))
|
|
return mlflow_dir
|
|
|
|
|
|
def _copy_file_or_tree(src, dst, dst_dir=None):
|
|
"""
|
|
Returns:
|
|
The path to the copied artifacts, relative to `dst`.
|
|
"""
|
|
dst_subpath = os.path.basename(os.path.abspath(src))
|
|
if dst_dir is not None:
|
|
dst_subpath = os.path.join(dst_dir, dst_subpath)
|
|
dst_path = os.path.join(dst, dst_subpath)
|
|
if os.path.isfile(src):
|
|
dst_dirpath = os.path.dirname(dst_path)
|
|
if not os.path.exists(dst_dirpath):
|
|
os.makedirs(dst_dirpath)
|
|
shutil.copy(src=src, dst=dst_path)
|
|
else:
|
|
shutil.copytree(src=src, dst=dst_path, ignore=shutil.ignore_patterns("__pycache__"))
|
|
return dst_subpath
|
|
|
|
|
|
def _get_local_project_dir_size(project_path):
|
|
"""Internal function for reporting the size of a local project directory before copying to
|
|
destination for cli logging reporting to stdout.
|
|
|
|
Args:
|
|
project_path: local path of the project directory
|
|
|
|
Returns:
|
|
directory file sizes in KB, rounded to single decimal point for legibility
|
|
"""
|
|
|
|
total_size = 0
|
|
for root, _, files in os.walk(project_path):
|
|
for f in files:
|
|
path = os.path.join(root, f)
|
|
total_size += os.path.getsize(path)
|
|
return round(total_size / 1024.0, 1)
|
|
|
|
|
|
def _get_local_file_size(file):
|
|
"""
|
|
Get the size of a local file in KB
|
|
"""
|
|
return round(os.path.getsize(file) / 1024.0, 1)
|
|
|
|
|
|
def get_parent_dir(path):
|
|
return os.path.abspath(os.path.join(path, os.pardir))
|
|
|
|
|
|
def relative_path_to_artifact_path(path):
|
|
if os.path == posixpath:
|
|
return path
|
|
if os.path.abspath(path) == path:
|
|
raise Exception("This method only works with relative paths.")
|
|
return unquote(pathname2url(path))
|
|
|
|
|
|
def path_to_local_file_uri(path):
|
|
"""
|
|
Convert local filesystem path to local file uri.
|
|
"""
|
|
return pathlib.Path(os.path.abspath(path)).as_uri()
|
|
|
|
|
|
def path_to_local_sqlite_uri(path):
|
|
"""
|
|
Convert local filesystem path to sqlite uri.
|
|
"""
|
|
path = posixpath.abspath(pathname2url(os.path.abspath(path)))
|
|
prefix = "sqlite://" if sys.platform == "win32" else "sqlite:///"
|
|
return prefix + path
|
|
|
|
|
|
def local_file_uri_to_path(uri):
|
|
"""
|
|
Convert URI to local filesystem path.
|
|
No-op if the uri does not have the expected scheme.
|
|
"""
|
|
path = uri
|
|
if uri.startswith("file:"):
|
|
parsed_path = urllib.parse.urlparse(uri)
|
|
path = parsed_path.path
|
|
# Fix for retaining server name in UNC path.
|
|
if is_windows() and parsed_path.netloc:
|
|
return urllib.request.url2pathname(rf"\\{parsed_path.netloc}{path}")
|
|
return urllib.request.url2pathname(path)
|
|
|
|
|
|
def get_local_path_or_none(path_or_uri):
|
|
"""Check if the argument is a local path (no scheme or file:///) and return local path if true,
|
|
None otherwise.
|
|
"""
|
|
parsed_uri = urllib.parse.urlparse(path_or_uri)
|
|
if len(parsed_uri.scheme) == 0 or parsed_uri.scheme == "file" and len(parsed_uri.netloc) == 0:
|
|
return local_file_uri_to_path(path_or_uri)
|
|
else:
|
|
return None
|
|
|
|
|
|
def download_file_using_http_uri(http_uri, download_path, chunk_size=100000000, headers=None):
|
|
"""
|
|
Downloads a file specified using the `http_uri` to a local `download_path`. This function
|
|
uses a `chunk_size` to ensure an OOM error is not raised a large file is downloaded.
|
|
|
|
Note : This function is meant to download files using presigned urls from various cloud
|
|
providers.
|
|
"""
|
|
if headers is None:
|
|
headers = {}
|
|
with cloud_storage_http_request("get", http_uri, stream=True, headers=headers) as response:
|
|
augmented_raise_for_status(response)
|
|
with open(download_path, "wb") as output_file:
|
|
for chunk in response.iter_content(chunk_size=chunk_size):
|
|
if not chunk:
|
|
break
|
|
output_file.write(chunk)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Chunk:
|
|
index: int
|
|
start: int
|
|
end: int
|
|
path: str
|
|
|
|
|
|
def _yield_chunks(path, file_size, chunk_size):
|
|
num_requests = int(math.ceil(file_size / float(chunk_size)))
|
|
for i in range(num_requests):
|
|
range_start = i * chunk_size
|
|
range_end = min(range_start + chunk_size - 1, file_size - 1)
|
|
yield _Chunk(i, range_start, range_end, path)
|
|
|
|
|
|
def parallelized_download_file_using_http_uri(
|
|
thread_pool_executor,
|
|
http_uri,
|
|
download_path,
|
|
remote_file_path,
|
|
file_size,
|
|
uri_type,
|
|
chunk_size,
|
|
env,
|
|
headers=None,
|
|
):
|
|
"""
|
|
Downloads a file specified using the `http_uri` to a local `download_path`. This function
|
|
sends multiple requests in parallel each specifying its own desired byte range as a header,
|
|
then reconstructs the file from the downloaded chunks. This allows for downloads of large files
|
|
without OOM risk.
|
|
|
|
Note : This function is meant to download files using presigned urls from various cloud
|
|
providers.
|
|
Returns a dict of chunk index : exception, if one was thrown for that index.
|
|
"""
|
|
|
|
def run_download(chunk: _Chunk):
|
|
try:
|
|
subprocess.run(
|
|
[
|
|
sys.executable,
|
|
download_cloud_file_chunk.__file__,
|
|
"--range-start",
|
|
str(chunk.start),
|
|
"--range-end",
|
|
str(chunk.end),
|
|
"--headers",
|
|
json.dumps(headers or {}),
|
|
"--download-path",
|
|
download_path,
|
|
"--http-uri",
|
|
http_uri,
|
|
],
|
|
text=True,
|
|
check=True,
|
|
capture_output=True,
|
|
timeout=MLFLOW_DOWNLOAD_CHUNK_TIMEOUT.get(),
|
|
env=env,
|
|
)
|
|
except (TimeoutExpired, CalledProcessError) as e:
|
|
raise MlflowException(
|
|
f"""
|
|
----- stdout -----
|
|
{e.stdout.strip()}
|
|
|
|
----- stderr -----
|
|
{e.stderr.strip()}
|
|
"""
|
|
) from e
|
|
|
|
chunks = _yield_chunks(remote_file_path, file_size, chunk_size)
|
|
# Create file if it doesn't exist or erase the contents if it does. We should do this here
|
|
# before sending to the workers so they can each individually seek to their respective positions
|
|
# and write chunks without overwriting.
|
|
with open(download_path, "w"):
|
|
pass
|
|
if uri_type == ArtifactCredentialType.GCP_SIGNED_URL or uri_type is None:
|
|
chunk = next(chunks)
|
|
# GCP files could be transcoded, in which case the range header is ignored.
|
|
# Test if this is the case by downloading one chunk and seeing if it's larger than the
|
|
# requested size. If yes, let that be the file; if not, continue downloading more chunks.
|
|
download_chunk(
|
|
range_start=chunk.start,
|
|
range_end=chunk.end,
|
|
headers=headers,
|
|
download_path=download_path,
|
|
http_uri=http_uri,
|
|
)
|
|
downloaded_size = os.path.getsize(download_path)
|
|
# If downloaded size was equal to the chunk size it would have been downloaded serially,
|
|
# so we don't need to consider this here
|
|
if downloaded_size > chunk_size:
|
|
return {}
|
|
|
|
futures = {thread_pool_executor.submit(run_download, chunk): chunk for chunk in chunks}
|
|
failed_downloads = {}
|
|
with ArtifactProgressBar.chunks(file_size, f"Downloading {download_path}", chunk_size) as pbar:
|
|
for future in as_completed(futures):
|
|
chunk = futures[future]
|
|
try:
|
|
future.result()
|
|
except Exception as e:
|
|
_logger.debug(
|
|
f"Failed to download chunk {chunk.index} for {chunk.path}: {e}. "
|
|
f"The download of this chunk will be retried later."
|
|
)
|
|
failed_downloads[chunk] = future.exception()
|
|
else:
|
|
pbar.update()
|
|
|
|
return failed_downloads
|
|
|
|
|
|
def download_chunk_retries(*, chunks, http_uri, headers, download_path):
|
|
num_retries = _MLFLOW_MPD_NUM_RETRIES.get()
|
|
interval = _MLFLOW_MPD_RETRY_INTERVAL_SECONDS.get()
|
|
for chunk in chunks:
|
|
_logger.info(f"Retrying download of chunk {chunk.index} for {chunk.path}")
|
|
for retry in range(num_retries):
|
|
try:
|
|
download_chunk(
|
|
range_start=chunk.start,
|
|
range_end=chunk.end,
|
|
headers=headers,
|
|
download_path=download_path,
|
|
http_uri=http_uri,
|
|
)
|
|
_logger.info(f"Successfully downloaded chunk {chunk.index} for {chunk.path}")
|
|
break
|
|
except Exception:
|
|
if retry == num_retries - 1:
|
|
raise
|
|
time.sleep(interval)
|
|
|
|
|
|
def _handle_readonly_on_windows(func, path, exc_info):
|
|
"""
|
|
This function should not be called directly but should be passed to `onerror` of
|
|
`shutil.rmtree` in order to reattempt the removal of a read-only file after making
|
|
it writable on Windows.
|
|
|
|
References:
|
|
- https://bugs.python.org/issue19643
|
|
- https://bugs.python.org/issue43657
|
|
"""
|
|
exc_type, exc_value = exc_info[:2]
|
|
should_reattempt = (
|
|
is_windows()
|
|
and func in (os.unlink, os.rmdir)
|
|
and issubclass(exc_type, PermissionError)
|
|
and exc_value.winerror == 5
|
|
)
|
|
if not should_reattempt:
|
|
raise exc_value
|
|
os.chmod(path, stat.S_IWRITE)
|
|
func(path)
|
|
|
|
|
|
def _get_tmp_dir():
|
|
from mlflow.utils.databricks_utils import get_repl_id, is_in_databricks_runtime
|
|
|
|
if is_in_databricks_runtime():
|
|
try:
|
|
return get_databricks_local_temp_dir()
|
|
except Exception:
|
|
pass
|
|
|
|
if repl_id := get_repl_id():
|
|
return os.path.join("/tmp", "repl_tmp_data", repl_id)
|
|
|
|
return None
|
|
|
|
|
|
def create_tmp_dir():
|
|
if directory := _get_tmp_dir():
|
|
os.makedirs(directory, exist_ok=True)
|
|
return tempfile.mkdtemp(dir=directory)
|
|
|
|
return tempfile.mkdtemp()
|
|
|
|
|
|
@cache_return_value_per_process
|
|
def get_or_create_tmp_dir():
|
|
"""
|
|
Get or create a temporary directory which will be removed once python process exit.
|
|
"""
|
|
from mlflow.utils.databricks_utils import get_repl_id, is_in_databricks_runtime
|
|
|
|
if is_in_databricks_runtime() and get_repl_id() is not None:
|
|
# Note: For python process attached to databricks notebook, atexit does not work.
|
|
# The directory returned by `get_databricks_local_tmp_dir`
|
|
# will be removed once databricks notebook detaches.
|
|
# The temp directory is designed to be used by all kinds of applications,
|
|
# so create a child directory "mlflow" for storing mlflow temp data.
|
|
try:
|
|
repl_local_tmp_dir = get_databricks_local_temp_dir()
|
|
except Exception:
|
|
repl_local_tmp_dir = os.path.join("/tmp", "repl_tmp_data", get_repl_id())
|
|
|
|
tmp_dir = os.path.join(repl_local_tmp_dir, "mlflow")
|
|
os.makedirs(tmp_dir, exist_ok=True)
|
|
else:
|
|
tmp_dir = tempfile.mkdtemp()
|
|
# mkdtemp creates a directory with permission 0o700
|
|
# For Spark UDFs, we need to make it accessible to other processes
|
|
# Use 0o750 (owner: rwx, group: r-x, others: None) instead of 0o777
|
|
# This allows read/execute but not write for group and others
|
|
os.chmod(tmp_dir, 0o750)
|
|
atexit.register(shutil.rmtree, tmp_dir, ignore_errors=True)
|
|
|
|
return tmp_dir
|
|
|
|
|
|
@cache_return_value_per_process
|
|
def get_or_create_nfs_tmp_dir():
|
|
"""
|
|
Get or create a temporary NFS directory which will be removed once python process exit.
|
|
"""
|
|
from mlflow.utils.databricks_utils import get_repl_id, is_in_databricks_runtime
|
|
from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir
|
|
|
|
nfs_root_dir = get_nfs_cache_root_dir()
|
|
|
|
if is_in_databricks_runtime() and get_repl_id() is not None:
|
|
# Note: In databricks, atexit hook does not work.
|
|
# The directory returned by `get_databricks_nfs_tmp_dir`
|
|
# will be removed once databricks notebook detaches.
|
|
# The temp directory is designed to be used by all kinds of applications,
|
|
# so create a child directory "mlflow" for storing mlflow temp data.
|
|
try:
|
|
repl_nfs_tmp_dir = get_databricks_nfs_temp_dir()
|
|
except Exception:
|
|
repl_nfs_tmp_dir = os.path.join(nfs_root_dir, "repl_tmp_data", get_repl_id())
|
|
|
|
tmp_nfs_dir = os.path.join(repl_nfs_tmp_dir, "mlflow")
|
|
os.makedirs(tmp_nfs_dir, exist_ok=True)
|
|
else:
|
|
tmp_nfs_dir = tempfile.mkdtemp(dir=nfs_root_dir)
|
|
# mkdtemp creates a directory with permission 0o700
|
|
# For Spark UDFs, we need to make it accessible to other processes
|
|
# Use 0o750 (owner: rwx, group: r-x, others: None) instead of 0o777
|
|
os.chmod(tmp_nfs_dir, 0o750)
|
|
atexit.register(shutil.rmtree, tmp_nfs_dir, ignore_errors=True)
|
|
|
|
return tmp_nfs_dir
|
|
|
|
|
|
def shutil_copytree_without_file_permissions(src_dir, dst_dir):
|
|
"""
|
|
Copies the directory src_dir into dst_dir, without preserving filesystem permissions
|
|
"""
|
|
for dirpath, dirnames, filenames in os.walk(src_dir):
|
|
for dirname in dirnames:
|
|
relative_dir_path = os.path.relpath(os.path.join(dirpath, dirname), src_dir)
|
|
# For each directory <dirname> immediately under <dirpath>, create an equivalently-named
|
|
# directory under the destination directory
|
|
abs_dir_path = os.path.join(dst_dir, relative_dir_path)
|
|
if not os.path.exists(abs_dir_path):
|
|
os.mkdir(abs_dir_path)
|
|
for filename in filenames:
|
|
# For each file with name <filename> immediately under <dirpath>, copy that file to
|
|
# the appropriate location in the destination directory
|
|
file_path = os.path.join(dirpath, filename)
|
|
relative_file_path = os.path.relpath(file_path, src_dir)
|
|
abs_file_path = os.path.join(dst_dir, relative_file_path)
|
|
shutil.copy2(file_path, abs_file_path)
|
|
|
|
|
|
def contains_path_separator(path):
|
|
"""
|
|
Returns True if a path contains a path separator, False otherwise.
|
|
"""
|
|
return any((sep in path) for sep in (os.path.sep, os.path.altsep) if sep is not None)
|
|
|
|
|
|
def contains_percent(path):
|
|
"""
|
|
Returns True if a path contains a percent character, False otherwise.
|
|
"""
|
|
return "%" in path
|
|
|
|
|
|
def read_chunk(path: os.PathLike, size: int, start_byte: int = 0) -> bytes:
|
|
"""Read a chunk of bytes from a file.
|
|
|
|
Args:
|
|
path: Path to the file.
|
|
size: The size of the chunk.
|
|
start_byte: The start byte of the chunk.
|
|
|
|
Returns:
|
|
The chunk of bytes.
|
|
|
|
"""
|
|
with open(path, "rb") as f:
|
|
if start_byte > 0:
|
|
f.seek(start_byte)
|
|
return f.read(size)
|
|
|
|
|
|
@contextmanager
|
|
def remove_on_error(path: os.PathLike, onerror=None):
|
|
"""A context manager that removes a file or directory if an exception is raised during
|
|
execution.
|
|
|
|
Args:
|
|
path: Path to the file or directory.
|
|
onerror: A callback function that will be called with the captured exception before
|
|
the file or directory is removed. For example, you can use this callback to
|
|
log the exception.
|
|
|
|
"""
|
|
try:
|
|
yield
|
|
except Exception as e:
|
|
if onerror:
|
|
onerror(e)
|
|
if os.path.exists(path):
|
|
if os.path.isfile(path):
|
|
os.remove(path)
|
|
elif os.path.isdir(path):
|
|
shutil.rmtree(path)
|
|
_logger.warning(
|
|
f"Failed to remove {path}" if os.path.exists(path) else f"Successfully removed {path}"
|
|
)
|
|
raise
|
|
|
|
|
|
def get_total_file_size(path: str | pathlib.Path) -> int | None:
|
|
"""Return the size of all files under given path, including files in subdirectories.
|
|
|
|
Args:
|
|
path: The absolute path of a local directory.
|
|
|
|
Returns:
|
|
size in bytes.
|
|
|
|
"""
|
|
try:
|
|
if isinstance(path, pathlib.Path):
|
|
path = str(path)
|
|
if not os.path.exists(path):
|
|
raise MlflowException(
|
|
message=f"The given {path} does not exist.", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
if not os.path.isdir(path):
|
|
raise MlflowException(
|
|
message=f"The given {path} is not a directory.", error_code=INVALID_PARAMETER_VALUE
|
|
)
|
|
|
|
total_size = 0
|
|
for cur_path, dirs, files in os.walk(path):
|
|
full_paths = [os.path.join(cur_path, file) for file in files]
|
|
total_size += sum(map(os.path.getsize, full_paths))
|
|
return total_size
|
|
except Exception as e:
|
|
_logger.info(f"Failed to get the total size of {path} because of error :{e}")
|
|
return None
|
|
|
|
|
|
def write_yaml(
|
|
root: str,
|
|
file_name: str,
|
|
data: dict[str, Any],
|
|
overwrite: bool = False,
|
|
sort_keys: bool = True,
|
|
ensure_yaml_extension: bool = True,
|
|
) -> None:
|
|
"""
|
|
NEVER TOUCH THIS FUNCTION. KEPT FOR BACKWARD COMPATIBILITY with
|
|
databricks-feature-engineering<=0.10.2
|
|
"""
|
|
import yaml
|
|
|
|
with open(os.path.join(root, file_name), "w") as f:
|
|
yaml.safe_dump(
|
|
data,
|
|
f,
|
|
default_flow_style=False,
|
|
allow_unicode=True,
|
|
sort_keys=sort_keys,
|
|
)
|
|
|
|
|
|
def read_yaml(root: str, file_name: str) -> dict[str, Any]:
|
|
"""
|
|
NEVER TOUCH THIS FUNCTION. KEPT FOR BACKWARD COMPATIBILITY with
|
|
databricks-feature-engineering<=0.10.2
|
|
"""
|
|
import yaml
|
|
|
|
with open(os.path.join(root, file_name)) as f:
|
|
return yaml.safe_load(f)
|
|
|
|
|
|
class ExclusiveFileLock:
|
|
"""
|
|
Exclusive file lock (only works on Unix system)
|
|
"""
|
|
|
|
def __init__(self, path: str):
|
|
if os.name == "nt":
|
|
raise MlflowException("ExclusiveFileLock class does not support Windows system.")
|
|
self.path = path
|
|
self.fd = None
|
|
|
|
def __enter__(self) -> None:
|
|
# Python on Windows does not have `fcntl` module, so importing it lazily.
|
|
import fcntl # clint: disable=lazy-import
|
|
|
|
# Open file (create if missing)
|
|
self.fd = open(self.path, "w")
|
|
# Acquire exclusive lock (blocking)
|
|
fcntl.flock(self.fd, fcntl.LOCK_EX)
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
):
|
|
# Python on Windows does not have `fcntl` module, so importing it lazily.
|
|
import fcntl # clint: disable=lazy-import
|
|
|
|
# Release lock
|
|
fcntl.flock(self.fd, fcntl.LOCK_UN)
|
|
self.fd.close()
|
|
|
|
|
|
def check_tarfile_security(archive_path: str) -> None:
|
|
"""
|
|
Check the tar file content.
|
|
If its members contain any of the following paths:
|
|
* An absolute path.
|
|
* A relative path that escapes the extraction directory.
|
|
* A relative path that goes through a symlink.
|
|
then raise an error.
|
|
"""
|
|
with tarfile.open(archive_path, "r") as tar:
|
|
symlink_set = set()
|
|
for m in tar.getmembers():
|
|
# Normalize backslashes to forward slashes before path validation to prevent
|
|
# bypass on Windows where backslashes are treated as directory separators.
|
|
path = posixpath.normpath(m.name.replace("\\", "/"))
|
|
_check_path_is_safe(path)
|
|
if m.issym():
|
|
symlink_set.add(path)
|
|
elif m.islnk():
|
|
symlink_set.add(path)
|
|
# Hard link targets are dangerous: tar.extract creates an actual hard
|
|
# link to the target path, so validate they don't escape.
|
|
link_target = posixpath.normpath(m.linkname.replace("\\", "/"))
|
|
_check_path_is_safe(link_target, context=f"hard link target of {path}")
|
|
link_parent = posixpath.dirname(path)
|
|
resolved = posixpath.normpath(posixpath.join(link_parent, link_target))
|
|
if resolved == ".." or resolved.startswith("../"):
|
|
raise MlflowException.invalid_parameter_value(
|
|
"Hard link target that escapes the extraction directory is not "
|
|
f"allowed, but got {path} -> {link_target}."
|
|
)
|
|
for m in tar.getmembers():
|
|
if not m.issym() and not m.islnk():
|
|
path = posixpath.normpath(m.name.replace("\\", "/"))
|
|
path_parts = path.split("/")
|
|
for prefix_len in range(1, len(path_parts) + 1):
|
|
prefix_path = "/".join(path_parts[:prefix_len])
|
|
if prefix_path in symlink_set:
|
|
raise MlflowException.invalid_parameter_value(
|
|
"Destination path in the archive file can not go through a symlink, "
|
|
f"but got path {path}."
|
|
)
|
|
|
|
|
|
def _check_path_is_safe(path: str, context: str = "") -> None:
|
|
label = f" ({context})" if context else ""
|
|
# Reject Unix absolute paths
|
|
if path.startswith("/"):
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Absolute path destination in the archive file is not allowed, "
|
|
f"but got path {path}{label}."
|
|
)
|
|
# Reject Windows drive-letter absolute paths (e.g., C:/...) and UNC paths (//server/...)
|
|
if len(path) >= 3 and path[0].isalpha() and path[1:3] == ":/":
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Absolute path destination in the archive file is not allowed, "
|
|
f"but got path {path}{label}."
|
|
)
|
|
path_parts = path.split("/")
|
|
if path_parts[0] == "..":
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Escaped path destination in the archive file is not allowed, "
|
|
f"but got path {path}{label}."
|
|
)
|