<commit_before>import os
import posixpath
import tempfile
from contextlib import contextmanager
import urllib.parse

from mlflow.entities import FileInfo
from mlflow.exceptions import MlflowException
from mlflow.store.artifact.artifact_repo import ArtifactRepository
from mlflow.utils.file_utils import mkdir, relative_path_to_artifact_path


class HdfsArtifactRepository(ArtifactRepository):
    """
    Stores artifacts on HDFS.

    This repository is used with URIs of the form ``hdfs:/<path>``. The repository can only be used
    together with the RestStore.
    """

    def __init__(self, artifact_uri):
        self.scheme, self.host, self.port, self.path = _resolve_connection_params(artifact_uri)
        super().__init__(artifact_uri)

    def log_artifact(self, local_file, artifact_path=None):
        """
            Log artifact in hdfs.
        :param local_file: source file path
        :param artifact_path: when specified will attempt to write under artifact_uri/artifact_path
        """
        hdfs_base_path = _resolve_base_path(self.path, artifact_path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
            _, file_name = os.path.split(local_file)
            destination = posixpath.join(hdfs_base_path, file_name)
            with hdfs.open(destination, "wb") as output:
                output.write(open(local_file, "rb").read())

    def log_artifacts(self, local_dir, artifact_path=None):
        """
            Log artifacts in hdfs.
            Missing remote sub-directories will be created if needed.
        :param local_dir: source dir path
        :param artifact_path: when specified will attempt to write under artifact_uri/artifact_path
        """
        hdfs_base_path = _resolve_base_path(self.path, artifact_path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:

            if not hdfs.exists(hdfs_base_path):
                hdfs.mkdir(hdfs_base_path)

            for subdir_path, _, files in os.walk(local_dir):

                relative_path = _relative_path_local(local_dir, subdir_path)

                hdfs_subdir_path = (
                    posixpath.join(hdfs_base_path, relative_path)
                    if relative_path
                    else hdfs_base_path
                )

                if not hdfs.exists(hdfs_subdir_path):
                    hdfs.mkdir(hdfs_subdir_path)

                for each_file in files:
                    source = os.path.join(subdir_path, each_file)
                    destination = posixpath.join(hdfs_subdir_path, each_file)
                    with hdfs.open(destination, "wb") as output_stream:
                        output_stream.write(open(source, "rb").read())

    def list_artifacts(self, path=None):
        """
        Lists files and directories under artifacts directory for the current run_id.
        (self.path contains the base path - hdfs:/some/path/run_id/artifacts)

        :param path: Relative source path. Possible subdirectory existing under
                     hdfs:/some/path/run_id/artifacts
        :return: List of FileInfos under given path
        """
        hdfs_base_path = _resolve_base_path(self.path, path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
            paths = []
            if hdfs.exists(hdfs_base_path):
                for file_detail in hdfs.ls(hdfs_base_path, detail=True):
                    file_name = file_detail.get("name")
                    # Strip off anything that comes before the artifact root e.g. hdfs://name
                    offset = file_name.index(self.path)
                    rel_path = _relative_path_remote(self.path, file_name[offset:])
                    is_dir = file_detail.get("kind") == "directory"
                    size = file_detail.get("size")
                    paths.append(FileInfo(rel_path, is_dir, size))
            return sorted(paths, key=lambda f: paths)

    def _walk_path(self, hdfs, hdfs_path):
        if hdfs.exists(hdfs_path):
            if hdfs.isdir(hdfs_path):
                for subdir, _, files in hdfs.walk(hdfs_path):
                    if subdir != hdfs_path:
                        yield subdir, hdfs.isdir(subdir), hdfs.info(subdir).get("size")
                    for f in files:
                        file_path = posixpath.join(subdir, f)
                        yield file_path, hdfs.isdir(file_path), hdfs.info(file_path).get("size")
            else:
                yield hdfs_path, False, hdfs.info(hdfs_path).get("size")

    def download_artifacts(self, artifact_path, dst_path=None):
        """
        Download an artifact file or directory to a local directory/file if applicable, and
        return a local path for it.
        The caller is responsible for managing the lifecycle of the downloaded artifacts.

        (self.path contains the base path - hdfs:/some/path/run_id/artifacts)

        :param artifact_path: Relative source path to the desired artifacts file or directory.
        :param dst_path: Absolute path of the local filesystem destination directory to which
                         to download the specified artifacts. This directory must already
                         exist. If unspecified, the artifacts will be downloaded to a new,
                         uniquely-named
                         directory on the local filesystem.

        :return: Absolute path of the local filesystem location containing the downloaded
        artifacts - file/directory.
        """

        hdfs_base_path = _resolve_base_path(self.path, artifact_path)
        local_dir = _tmp_dir(dst_path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:

            if not hdfs.isdir(hdfs_base_path):
                local_path = os.path.join(local_dir, os.path.normpath(artifact_path))
                _download_hdfs_file(hdfs, hdfs_base_path, local_path)
                return local_path

            for path, is_dir, _ in self._walk_path(hdfs, hdfs_base_path):

                relative_path = _relative_path_remote(hdfs_base_path, path)
                local_path = os.path.join(local_dir, relative_path) if relative_path else local_dir

                if is_dir:
                    mkdir(local_path)
                else:
                    _download_hdfs_file(hdfs, path, local_path)
            return local_dir

    def _download_file(self, remote_file_path, local_path):
        raise MlflowException("This is not implemented. Should never be called.")

    def delete_artifacts(self, artifact_path=None):
        path = posixpath.join(self.path, artifact_path) if artifact_path else self.path
        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
            hdfs.delete(path, recursive=True)


@contextmanager
def hdfs_system(scheme, host, port):
    """<commit_msg>Uniquely-named directory was always created for hdfs artifact repo (#5210)<commit_after>import os
import posixpath
import tempfile
from contextlib import contextmanager
import urllib.parse

from mlflow.entities import FileInfo
from mlflow.exceptions import MlflowException
from mlflow.store.artifact.artifact_repo import ArtifactRepository
from mlflow.utils.file_utils import mkdir, relative_path_to_artifact_path


class HdfsArtifactRepository(ArtifactRepository):
    """
    Stores artifacts on HDFS.

    This repository is used with URIs of the form ``hdfs:/<path>``. The repository can only be used
    together with the RestStore.
    """

    def __init__(self, artifact_uri):
        self.scheme, self.host, self.port, self.path = _resolve_connection_params(artifact_uri)
        super().__init__(artifact_uri)

    def log_artifact(self, local_file, artifact_path=None):
        """
            Log artifact in hdfs.
        :param local_file: source file path
        :param artifact_path: when specified will attempt to write under artifact_uri/artifact_path
        """
        hdfs_base_path = _resolve_base_path(self.path, artifact_path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
            _, file_name = os.path.split(local_file)
            destination = posixpath.join(hdfs_base_path, file_name)
            with hdfs.open(destination, "wb") as output:
                output.write(open(local_file, "rb").read())

    def log_artifacts(self, local_dir, artifact_path=None):
        """
            Log artifacts in hdfs.
            Missing remote sub-directories will be created if needed.
        :param local_dir: source dir path
        :param artifact_path: when specified will attempt to write under artifact_uri/artifact_path
        """
        hdfs_base_path = _resolve_base_path(self.path, artifact_path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:

            if not hdfs.exists(hdfs_base_path):
                hdfs.mkdir(hdfs_base_path)

            for subdir_path, _, files in os.walk(local_dir):

                relative_path = _relative_path_local(local_dir, subdir_path)

                hdfs_subdir_path = (
                    posixpath.join(hdfs_base_path, relative_path)
                    if relative_path
                    else hdfs_base_path
                )

                if not hdfs.exists(hdfs_subdir_path):
                    hdfs.mkdir(hdfs_subdir_path)

                for each_file in files:
                    source = os.path.join(subdir_path, each_file)
                    destination = posixpath.join(hdfs_subdir_path, each_file)
                    with hdfs.open(destination, "wb") as output_stream:
                        output_stream.write(open(source, "rb").read())

    def list_artifacts(self, path=None):
        """
        Lists files and directories under artifacts directory for the current run_id.
        (self.path contains the base path - hdfs:/some/path/run_id/artifacts)

        :param path: Relative source path. Possible subdirectory existing under
                     hdfs:/some/path/run_id/artifacts
        :return: List of FileInfos under given path
        """
        hdfs_base_path = _resolve_base_path(self.path, path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
            paths = []
            if hdfs.exists(hdfs_base_path):
                for file_detail in hdfs.ls(hdfs_base_path, detail=True):
                    file_name = file_detail.get("name")
                    # Strip off anything that comes before the artifact root e.g. hdfs://name
                    offset = file_name.index(self.path)
                    rel_path = _relative_path_remote(self.path, file_name[offset:])
                    is_dir = file_detail.get("kind") == "directory"
                    size = file_detail.get("size")
                    paths.append(FileInfo(rel_path, is_dir, size))
            return sorted(paths, key=lambda f: paths)

    def _walk_path(self, hdfs, hdfs_path):
        if hdfs.exists(hdfs_path):
            if hdfs.isdir(hdfs_path):
                for subdir, _, files in hdfs.walk(hdfs_path):
                    if subdir != hdfs_path:
                        yield subdir, hdfs.isdir(subdir), hdfs.info(subdir).get("size")
                    for f in files:
                        file_path = posixpath.join(subdir, f)
                        yield file_path, hdfs.isdir(file_path), hdfs.info(file_path).get("size")
            else:
                yield hdfs_path, False, hdfs.info(hdfs_path).get("size")

    def download_artifacts(self, artifact_path, dst_path=None):
        """
        Download an artifact file or directory to a local directory/file if applicable, and
        return a local path for it.
        The caller is responsible for managing the lifecycle of the downloaded artifacts.

        (self.path contains the base path - hdfs:/some/path/run_id/artifacts)

        :param artifact_path: Relative source path to the desired artifacts file or directory.
        :param dst_path: Absolute path of the local filesystem destination directory to which
                         to download the specified artifacts. This directory must already
                         exist. If unspecified, the artifacts will be downloaded to a new,
                         uniquely-named
                         directory on the local filesystem.

        :return: Absolute path of the local filesystem location containing the downloaded
        artifacts - file/directory.
        """

        hdfs_base_path = _resolve_base_path(self.path, artifact_path)
        if dst_path and os.path.exists(dst_path):
            local_dir = os.path.abspath(dst_path)
        else:
            local_dir = _tmp_dir(dst_path)

        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:

            if not hdfs.isdir(hdfs_base_path):
                local_path = os.path.join(local_dir, os.path.normpath(artifact_path))
                _download_hdfs_file(hdfs, hdfs_base_path, local_path)
                return local_path

            for path, is_dir, _ in self._walk_path(hdfs, hdfs_base_path):

                relative_path = _relative_path_remote(hdfs_base_path, path)
                local_path = os.path.join(local_dir, relative_path) if relative_path else local_dir

                if is_dir:
                    mkdir(local_path)
                else:
                    _download_hdfs_file(hdfs, path, local_path)
            return local_dir

    def _download_file(self, remote_file_path, local_path):
        raise MlflowException("This is not implemented. Should never be called.")

    def delete_artifacts(self, artifact_path=None):
        path = posixpath.join(self.path, artifact_path) if artifact_path else self.path
        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
            hdfs.delete(path, recursive=True)


@contextmanager
def hdfs_system(scheme, host, port):
    """
{'fraction_non_alphanumeric': 0.07798589844028203, 'fraction_numerical': 0.0002848799943024001, 'mean_word_length': 2.263304671159656, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4322048', 'n_tokens_mistral': 4216, 'n_tokens_neox': 3922, 'n_words': 1008}