wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/lib/file_stream_utils.py

@@ -0,0 +1,118 @@
+#
+from typing import Any, Dict, Iterable
+
+
+def split_files(
+    files: Dict[str, Any], max_bytes: int = 10 * 1024 * 1024
+) -> Iterable[Dict[str, Dict]]:
+    """Split a file's dict (see `files` arg) into smaller dicts.
+
+    Each smaller dict will have at most `MAX_BYTES` size.
+
+    This method is used in `FileStreamAPI._send()` to limit the size of post requests
+    sent to wandb server.
+
+    Arguments:
+    files (dict): `dict` of form {file_name: {'content': ".....", 'offset': 0}}
+        The key `file_name` can also be mapped to a List [{"offset": int, "content": str}]
+    `max_bytes`: max size for chunk in bytes
+    """
+    current_volume: Dict[str, Dict] = {}
+    current_size = 0
+
+    def _str_size(x):
+        return len(x) if isinstance(x, bytes) else len(x.encode("utf-8"))
+
+    def _file_size(file):
+        size = file.get("_size")
+        if size is None:
+            size = sum(map(_str_size, file["content"]))
+            file["_size"] = size
+        return size
+
+    def _split_file(file, num_lines):
+        offset = file["offset"]
+        content = file["content"]
+        name = file["name"]
+        f1 = {"offset": offset, "content": content[:num_lines], "name": name}
+        f2 = {
+            "offset": offset + num_lines,
+            "content": content[num_lines:],
+            "name": name,
+        }
+        return f1, f2
+
+    def _num_lines_from_num_bytes(file, num_bytes):
+        size = 0
+        num_lines = 0
+        content = file["content"]
+        while num_lines < len(content):
+            size += _str_size(content[num_lines])
+            if size > num_bytes:
+                break
+            num_lines += 1
+        return num_lines
+
+    files_stack = []
+    for k, v in files.items():
+        if isinstance(v, list):
+            for item in v:
+                files_stack.append(
+                    {"name": k, "offset": item["offset"], "content": item["content"]}
+                )
+        else:
+            files_stack.append(
+                {"name": k, "offset": v["offset"], "content": v["content"]}
+            )
+
+    while files_stack:
+        f = files_stack.pop()
+        if f["name"] in current_volume:
+            files_stack.append(f)
+            yield current_volume
+            current_volume = {}
+            current_size = 0
+            continue
+        # For each file, we have to do 1 of 4 things:
+        # - Add the file as such to the current volume if possible.
+        # - Split the file and add the first part to the current volume and push the second part back onto the stack.
+        # - If that's not possible, check if current volume is empty:
+        # - If empty, add first line of file to current volume and push rest onto stack (This volume will exceed MAX_MB).
+        # - If not, push file back to stack and yield current volume.
+        fsize = _file_size(f)
+        rem = max_bytes - current_size
+        if fsize <= rem:
+            current_volume[f["name"]] = {
+                "offset": f["offset"],
+                "content": f["content"],
+            }
+            current_size += fsize
+        else:
+            num_lines = _num_lines_from_num_bytes(f, rem)
+            if not num_lines and not current_volume:
+                num_lines = 1
+            if num_lines:
+                f1, f2 = _split_file(f, num_lines)
+                current_volume[f1["name"]] = {
+                    "offset": f1["offset"],
+                    "content": f1["content"],
+                }
+                files_stack.append(f2)
+                yield current_volume
+                current_volume = {}
+                current_size = 0
+                continue
+            else:
+                files_stack.append(f)
+                yield current_volume
+                current_volume = {}
+                current_size = 0
+                continue
+        if current_size >= max_bytes:
+            yield current_volume
+            current_volume = {}
+            current_size = 0
+            continue
+
+    if current_volume:
+        yield current_volume

wandb/sdk/lib/filenames.py

@@ -0,0 +1,64 @@
+import os
+from typing import Callable, Generator, Union
+
+WANDB_DIRS = ("wandb", ".wandb")
+
+CONFIG_FNAME = "config.yaml"
+OUTPUT_FNAME = "output.log"
+DIFF_FNAME = "diff.patch"
+SUMMARY_FNAME = "wandb-summary.json"
+METADATA_FNAME = "wandb-metadata.json"
+REQUIREMENTS_FNAME = "requirements.txt"
+HISTORY_FNAME = "wandb-history.jsonl"
+EVENTS_FNAME = "wandb-events.jsonl"
+JOBSPEC_FNAME = "wandb-jobspec.json"
+CONDA_ENVIRONMENTS_FNAME = "conda-environment.yaml"
+
+
+def is_wandb_file(name: str) -> bool:
+    return (
+        name.startswith("wandb")
+        or name == METADATA_FNAME
+        or name == CONFIG_FNAME
+        or name == REQUIREMENTS_FNAME
+        or name == OUTPUT_FNAME
+        or name == DIFF_FNAME
+        or name == CONDA_ENVIRONMENTS_FNAME
+    )
+
+
+def filtered_dir(
+    root: str,
+    include_fn: Union[Callable[[str, str], bool], Callable[[str], bool]],
+    exclude_fn: Union[Callable[[str, str], bool], Callable[[str], bool]],
+) -> Generator[str, None, None]:
+    """Simple generator to walk a directory."""
+    import inspect
+
+    # compatibility with old API, which didn't pass root
+    def _include_fn(path: str, root: str) -> bool:
+        return (
+            include_fn(path, root)  # type: ignore
+            if len(inspect.signature(include_fn).parameters) == 2
+            else include_fn(path)  # type: ignore
+        )
+
+    def _exclude_fn(path: str, root: str) -> bool:
+        return (
+            exclude_fn(path, root)  # type: ignore
+            if len(inspect.signature(exclude_fn).parameters) == 2
+            else exclude_fn(path)  # type: ignore
+        )
+
+    for dirpath, _, files in os.walk(root):
+        for fname in files:
+            file_path = os.path.join(dirpath, fname)
+            if _include_fn(file_path, root) and not _exclude_fn(file_path, root):
+                yield file_path
+
+
+def exclude_wandb_fn(path: str, root: str) -> bool:
+    return any(
+        os.path.relpath(path, root).startswith(wandb_dir + os.sep)
+        for wandb_dir in WANDB_DIRS
+    )

wandb/sdk/lib/filesystem.py

@@ -0,0 +1,372 @@
+import contextlib
+import ctypes
+import errno
+import logging
+import os
+import platform
+import re
+import shutil
+import tempfile
+import threading
+from pathlib import Path
+from typing import IO, Any, BinaryIO, Generator, Optional
+
+from wandb.sdk.lib.paths import StrPath
+
+logger = logging.getLogger(__name__)
+
+# https://en.wikipedia.org/wiki/Filename#Comparison_of_filename_limitations
+PROBLEMATIC_PATH_CHARS = "".join(chr(i) for i in range(0, 32)) + ':"*<>?|'
+
+
+def mkdir_exists_ok(dir_name: StrPath) -> None:
+    """Create `dir_name` and any parent directories if they don't exist.
+
+    Raises:
+        FileExistsError: if `dir_name` exists and is not a directory.
+        PermissionError: if `dir_name` is not writable.
+    """
+    try:
+        os.makedirs(dir_name, exist_ok=True)
+    except FileExistsError as e:
+        raise FileExistsError(f"{dir_name!s} exists and is not a directory") from e
+    except PermissionError as e:
+        raise PermissionError(f"{dir_name!s} is not writable") from e
+
+
+def path_fallbacks(path: StrPath) -> Generator[str, None, None]:
+    """Yield variations of `path` that may exist on the filesystem.
+
+    Return a sequence of paths that should be checked in order for existence or
+    create-ability. Essentially, keep replacing "suspect" characters until we run out.
+    """
+    path = str(path)
+    root, tail = os.path.splitdrive(path)
+    yield os.path.join(root, tail)
+    for char in PROBLEMATIC_PATH_CHARS:
+        if char in tail:
+            tail = tail.replace(char, "-")
+            yield os.path.join(root, tail)
+
+
+def mkdir_allow_fallback(dir_name: StrPath) -> StrPath:
+    """Create `dir_name`, removing invalid path characters if necessary.
+
+    Returns:
+        The path to the created directory, which may not be the original path.
+    """
+    for new_name in path_fallbacks(dir_name):
+        try:
+            os.makedirs(new_name, exist_ok=True)
+            if Path(new_name) != Path(dir_name):
+                logger.warning(f"Creating '{new_name}' instead of '{dir_name}'")
+            return Path(new_name) if isinstance(dir_name, Path) else new_name
+        except (ValueError, NotADirectoryError):
+            pass
+        except OSError as e:
+            if e.errno != 22:
+                raise
+
+    raise OSError(f"Unable to create directory '{dir_name}'")
+
+
+def files_in(path: StrPath) -> Generator[os.DirEntry, None, None]:
+    """Yield a directory entry for each file under a given path (recursive)."""
+    if not os.path.isdir(path):
+        return
+    for entry in os.scandir(path):
+        if entry.is_dir():
+            yield from files_in(entry.path)
+        else:
+            yield entry
+
+
+class WriteSerializingFile:
+    """Wrapper for a file object that serializes writes."""
+
+    def __init__(self, f: BinaryIO) -> None:
+        self.lock = threading.Lock()
+        self.f = f
+
+    def write(self, *args, **kargs) -> None:  # type: ignore
+        self.lock.acquire()
+        try:
+            self.f.write(*args, **kargs)
+            self.f.flush()
+        finally:
+            self.lock.release()
+
+    def close(self) -> None:
+        self.lock.acquire()  # wait for pending writes
+        try:
+            self.f.close()
+        finally:
+            self.lock.release()
+
+
+class CRDedupedFile(WriteSerializingFile):
+    def __init__(self, f: BinaryIO) -> None:
+        super().__init__(f=f)
+        self._buff = b""
+
+    def write(self, data) -> None:  # type: ignore
+        lines = re.split(b"\r\n|\n", data)
+        ret = []  # type: ignore
+        for line in lines:
+            if line[:1] == b"\r":
+                if ret:
+                    ret.pop()
+                elif self._buff:
+                    self._buff = b""
+            line = line.split(b"\r")[-1]
+            if line:
+                ret.append(line)
+        if self._buff:
+            ret.insert(0, self._buff)
+        if ret:
+            self._buff = ret.pop()
+        super().write(b"\n".join(ret) + b"\n")
+
+    def close(self) -> None:
+        if self._buff:
+            super().write(self._buff)
+        super().close()
+
+
+def copy_or_overwrite_changed(source_path: StrPath, target_path: StrPath) -> StrPath:
+    """Copy source_path to target_path, unless it already exists with the same mtime.
+
+    We liberally add write permissions to deal with the case of multiple users needing
+    to share the same cache or run directory.
+
+    Args:
+        source_path: The path to the file to copy.
+        target_path: The path to copy the file to.
+
+    Returns:
+        The path to the copied file (which may be different from target_path).
+    """
+    return_type = type(target_path)
+
+    target_path = system_preferred_path(target_path, warn=True)
+
+    need_copy = (
+        not os.path.isfile(target_path)
+        or os.stat(source_path).st_mtime != os.stat(target_path).st_mtime
+    )
+
+    permissions_plus_write = os.stat(source_path).st_mode
+    if need_copy:
+        dir_name, file_name = os.path.split(target_path)
+        target_path = os.path.join(mkdir_allow_fallback(dir_name), file_name)
+        try:
+            # Use copy2 to preserve file metadata (including modified time).
+            shutil.copy2(source_path, target_path)
+        except PermissionError:
+            # If the file is read-only try to make it writable.
+            try:
+                os.chmod(target_path, permissions_plus_write)
+                shutil.copy2(source_path, target_path)
+            except PermissionError as e:
+                raise PermissionError("Unable to overwrite '{target_path!s}'") from e
+        # Prevent future permissions issues by universal write permissions now.
+        os.chmod(target_path, permissions_plus_write)
+
+    return return_type(target_path)  # type: ignore  # 'os.PathLike' is abstract.
+
+
+@contextlib.contextmanager
+def safe_open(
+    path: StrPath, mode: str = "r", *args: Any, **kwargs: Any
+) -> Generator[IO, None, None]:
+    """Open a file, ensuring any changes only apply atomically after close.
+
+    This context manager ensures that even unsuccessful writes will not leave a "dirty"
+    file or overwrite good data, and that all temp data is cleaned up.
+
+    The semantics and behavior are intended to be nearly identical to the built-in
+    open() function. Differences:
+        - It creates any parent directories that don't exist, rather than raising.
+        - In 'x' mode, it checks at the beginning AND end of the write and fails if the
+            file exists either time.
+    """
+    path = Path(path).resolve()
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if "x" in mode and path.exists():
+        raise FileExistsError(f"{path!s} already exists")
+
+    if "r" in mode and "+" not in mode:
+        # This is read-only, so we can just open the original file.
+        # TODO (hugh): create a reflink and read from that.
+        with path.open(mode, *args, **kwargs) as f:
+            yield f
+        return
+
+    with tempfile.TemporaryDirectory(dir=path.parent) as tmp_dir:
+        tmp_path = Path(tmp_dir) / path.name
+
+        if ("r" in mode or "a" in mode) and path.exists():
+            # We need to copy the original file in order to support reads and appends.
+            # TODO (hugh): use reflinks to avoid the copy on platforms that support it.
+            shutil.copy2(path, tmp_path)
+
+        with tmp_path.open(mode, *args, **kwargs) as f:
+            yield f
+            f.flush()
+            os.fsync(f.fileno())
+
+        if "x" in mode:
+            # Ensure that if another process has beaten us to writing the file we raise
+            # rather than overwrite. os.link() atomically creates a hard link to the
+            # target file and will raise FileExistsError if the target already exists.
+            os.link(tmp_path, path)
+            os.unlink(tmp_path)
+        else:
+            tmp_path.replace(path)
+
+
+def safe_copy(source_path: StrPath, target_path: StrPath) -> StrPath:
+    """Copy a file atomically.
+
+    Copying is not usually atomic, and on operating systems that allow multiple
+    writers to the same file, the result can get corrupted. If two writers copy
+    to the same file, the contents can become interleaved.
+
+    We mitigate the issue somewhat by copying to a temporary file first and
+    then renaming. Renaming is atomic: if process 1 renames file A to X and
+    process 2 renames file B to X, then X will either contain the contents
+    of A or the contents of B, not some mixture of both.
+    """
+    # TODO (hugh): check that there is enough free space.
+    output_path = Path(target_path).resolve()
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with tempfile.TemporaryDirectory(dir=output_path.parent) as tmp_dir:
+        tmp_path = (Path(tmp_dir) / Path(source_path).name).with_suffix(".tmp")
+        shutil.copy2(source_path, tmp_path)
+        tmp_path.replace(output_path)
+    return target_path
+
+
+def _reflink_linux(existing_path: Path, new_path: Path) -> None:
+    """Create a reflink to `existing_path` at `new_path` on Linux."""
+    import fcntl
+
+    FICLONE = 0x40049409  # magic number from <linux/fs.h>  # noqa: N806
+    with open(existing_path, "rb") as t_f, open(new_path, "wb+") as l_f:
+        fcntl.ioctl(l_f.fileno(), FICLONE, t_f.fileno())
+
+
+def _reflink_macos(existing_path: Path, new_path: Path) -> None:
+    try:
+        clib = ctypes.CDLL("libc.dylib", use_errno=True)
+    except (FileNotFoundError, OSError) as e:
+        if ctypes.get_errno() != errno.ENOENT and not isinstance(e, FileNotFoundError):
+            raise
+        # Before macOS 11 (<Nov 2020) clib was in libSystem.dylib, so we can try there.
+        clib = ctypes.CDLL("/usr/lib/libSystem.dylib", use_errno=True)
+
+    try:
+        clonefile = clib.clonefile
+    except AttributeError:
+        raise OSError(errno.ENOTSUP, "'clonefile' is not available on this system")
+
+    clonefile.argtypes = (ctypes.c_char_p, ctypes.c_char_p, ctypes.c_int)
+    clonefile.restype = ctypes.c_int
+
+    if clonefile(os.fsencode(existing_path), os.fsencode(new_path), ctypes.c_int(0)):
+        # Anything other than 0 is an error.
+        err = ctypes.get_errno()
+        raise OSError(err, os.strerror(err), existing_path)
+
+
+def reflink(existing_path: StrPath, new_path: StrPath, overwrite: bool = False) -> None:
+    """Create a reflink to `existing_path` at `new_path`.
+
+    A reflink (reflective link) is a copy-on-write reference to a file. Once linked, the
+    file and link are both "real" files (not symbolic or hard links) and each can be
+    modified independently without affecting the other; however, they share the same
+    underlying data blocks on disk so until one is modified they are "zero-cost" copies.
+
+    Reflinks have all the functionality of copies, so we should use them wherever they
+    are supported if we would otherwise copy a file. (This is not particularly radical--
+    GNU `cp` defaults to `reflink=auto`, using it whenever available) However, support
+    for them is limited to a small number of filesystems. They should work on:
+    - Linux with a Btrfs or XFS filesystem (NOT ext4)
+    - macOS 10.13 or later with an APFS filesystem (called clone files)
+
+    Reflinks are also supported on Solaris and Windows with ReFSv2, but we haven't
+    implemented support for them.
+
+    Like hard links, a reflink can only be created on the same filesystem as the target.
+    """
+    if platform.system() == "Linux":
+        link_fn = _reflink_linux
+    elif platform.system() == "Darwin":
+        link_fn = _reflink_macos
+    else:
+        raise OSError(
+            errno.ENOTSUP, f"reflinks are not supported on {platform.system()}"
+        )
+
+    new_path = Path(new_path).resolve()
+    existing_path = Path(existing_path).resolve()
+    if new_path.exists():
+        if not overwrite:
+            raise FileExistsError(f"{new_path} already exists")
+        logger.warning(f"Overwriting existing file {new_path}.")
+        new_path.unlink()
+
+    # Create any missing parent directories.
+    new_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        link_fn(existing_path, new_path)
+    except OSError as e:
+        base_msg = f"failed to create reflink from {existing_path} to {new_path}."
+        if e.errno in (errno.EPERM, errno.EACCES):
+            raise PermissionError(f"Insufficient permissions; {base_msg}") from e
+        if e.errno == errno.ENOENT:
+            raise FileNotFoundError(f"File not found; {base_msg}") from e
+        if e.errno == errno.EXDEV:
+            raise ValueError(f"Cannot link across filesystems; {base_msg}") from e
+        if e.errno == errno.EISDIR:
+            raise IsADirectoryError(f"Cannot reflink a directory; {base_msg}") from e
+        if e.errno in (errno.EOPNOTSUPP, errno.ENOTSUP):
+            raise OSError(
+                errno.ENOTSUP,
+                f"Filesystem does not support reflinks; {base_msg}",
+            ) from e
+        if e.errno == errno.EINVAL:
+            raise ValueError(f"Cannot link file ranges; {base_msg}") from e
+        raise
+
+
+def check_exists(path: StrPath) -> Optional[StrPath]:
+    """Look for variations of `path` and return the first found.
+
+    This exists to support former behavior around system-dependent paths; we used to use
+    ':' in Artifact paths unless we were on Windows, but this has issues when e.g. a
+    Linux machine is accessing an NTFS filesystem; we might need to look for the
+    alternate path. This checks all the possible directories we would consider creating.
+    """
+    for dest in path_fallbacks(path):
+        if os.path.exists(dest):
+            return Path(dest) if isinstance(path, Path) else dest
+    return None
+
+
+def system_preferred_path(path: StrPath, warn: bool = False) -> StrPath:
+    """Replace ':' with '-' in paths on Windows.
+
+    Args:
+        path: The path to convert.
+        warn: Whether to warn if ':' is replaced.
+    """
+    if platform.system() != "Windows":
+        return path
+    head, tail = os.path.splitdrive(path)
+    if warn and ":" in tail:
+        logger.warning(f"Replacing ':' in {tail} with '-'")
+    new_path = head + tail.replace(":", "-")
+    return Path(new_path) if isinstance(path, Path) else new_path