Delete util.py

util.py

@@ -1,655 +0,0 @@
-import logging
-import os
-import re
-import socket
-import sys
-import time
-import warnings
-from datetime import datetime
-from enum import Enum
-from itertools import cycle, islice
-from pathlib import Path
-from queue import Queue
-from threading import Thread
-from typing import Any, Callable, Dict, Optional, Union
-
-import boto3
-import botocore.exceptions as boto_exceptions
-import rich
-from botocore.config import Config
-from rich.console import Console, ConsoleRenderable
-from rich.highlighter import NullHighlighter
-from rich.progress import Progress
-from rich.text import Text
-from rich.traceback import Traceback
-
-from .aliases import PathOrStr
-from .exceptions import (
-    OLMoCliError,
-    OLMoEnvironmentError,
-    OLMoError,
-    OLMoNetworkError,
-    OLMoThreadError,
-)
-from .torch_util import get_global_rank, get_local_rank, get_node_rank, is_distributed
-
-try:
-    from functools import cache
-except ImportError:
-    from functools import lru_cache as cache
-
-
-class StrEnum(str, Enum):
-    """
-    This is equivalent to Python's :class:`enum.StrEnum` since version 3.11.
-    We include this here for compatibility with older version of Python.
-    """
-
-    def __str__(self) -> str:
-        return self.value
-
-    def __repr__(self) -> str:
-        return f"'{str(self)}'"
-
-
-_log_extra_fields: Dict[str, Any] = {}
-log = logging.getLogger(__name__)
-
-
-class LogFilterType(StrEnum):
-    rank0_only = "rank0_only"
-    local_rank0_only = "local_rank0_only"
-    all_ranks = "all_ranks"
-
-
-def log_extra_field(field_name: str, field_value: Any) -> None:
-    global _log_extra_fields
-    if field_value is None:
-        if field_name in _log_extra_fields:
-            del _log_extra_fields[field_name]
-    else:
-        _log_extra_fields[field_name] = field_value
-
-
-def setup_logging(log_filter_type: LogFilterType = LogFilterType.rank0_only) -> None:
-    """
-    :param rank0_only: INFO and below messages will only be emitted on the rank0 process.
-    """
-    log_extra_field("hostname", socket.gethostname())
-    if is_distributed():
-        log_extra_field("node_rank", get_node_rank())
-        log_extra_field("local_rank", get_local_rank())
-        log_extra_field("global_rank", get_global_rank())
-    else:
-        log_extra_field("node_rank", 0)
-        log_extra_field("local_rank", 0)
-        log_extra_field("global_rank", 0)
-
-    old_log_record_factory = logging.getLogRecordFactory()
-
-    def log_record_factory(*args, **kwargs) -> logging.LogRecord:
-        record = old_log_record_factory(*args, **kwargs)
-        for field_name, field_value in _log_extra_fields.items():
-            setattr(record, field_name, field_value)
-        return record
-
-    logging.setLogRecordFactory(log_record_factory)
-
-    handler: logging.Handler
-    if (
-        os.environ.get("OLMo_NONINTERACTIVE", False)
-        or os.environ.get("DEBIAN_FRONTEND", None) == "noninteractive"
-        or not sys.stdout.isatty()
-    ):
-        handler = logging.StreamHandler(sys.stdout)
-        formatter = logging.Formatter(
-            "%(asctime)s\t%(hostname)s:%(local_rank)s\t%(name)s:%(lineno)s\t%(levelname)s\t%(message)s"
-        )
-        formatter.default_time_format = "%Y-%m-%d %H:%M:%S"
-        formatter.default_msec_format = "%s.%03d"
-        handler.setFormatter(formatter)
-    else:
-        handler = RichHandler()
-
-    def rank0_filter(record: logging.LogRecord) -> int:
-        if record.levelno > logging.INFO:
-            return 1
-        if getattr(record, "global_rank", 0) == 0:
-            return 1
-        else:
-            return 0
-
-    def local_rank0_filter(record: logging.LogRecord) -> int:
-        if record.levelno > logging.INFO:
-            return 1
-        if getattr(record, "local_rank", 0) == 0:
-            return 1
-        else:
-            return 0
-
-    if log_filter_type == LogFilterType.rank0_only:
-        filter = rank0_filter
-    elif log_filter_type == LogFilterType.local_rank0_only:
-        filter = local_rank0_filter  # type: ignore
-    elif log_filter_type == LogFilterType.all_ranks:
-        filter = None
-    else:
-        raise ValueError(log_filter_type)
-
-    if filter is not None:
-        handler.addFilter(filter)  # type: ignore
-    logging.basicConfig(handlers=[handler], level=logging.INFO)
-
-    logging.captureWarnings(True)
-    logging.getLogger("urllib3").setLevel(logging.ERROR)
-
-
-def excepthook(exctype, value, traceback):
-    """
-    Used to patch `sys.excepthook` in order to log exceptions.
-    """
-    if issubclass(exctype, KeyboardInterrupt):
-        sys.__excepthook__(exctype, value, traceback)
-    elif issubclass(exctype, OLMoCliError):
-        rich.get_console().print(f"[yellow]{value}[/]", highlight=False)
-    elif issubclass(exctype, OLMoError):
-        rich.get_console().print(Text(f"{exctype.__name__}:", style="red"), value, highlight=False)
-    else:
-        log.critical("Uncaught %s: %s", exctype.__name__, value, exc_info=(exctype, value, traceback))
-
-
-def install_excepthook():
-    sys.excepthook = excepthook
-
-
-def filter_warnings():
-    # Filter internal deprecation warnings from torch
-    warnings.filterwarnings(
-        action="ignore",
-        category=UserWarning,
-        message="torch.distributed.*_base is a private function and will be deprecated.*",
-    )
-    warnings.filterwarnings(
-        action="ignore",
-        category=UserWarning,
-        message="TypedStorage is deprecated.*",
-    )
-    warnings.filterwarnings(
-        action="ignore",
-        category=UserWarning,
-        message="Please use DTensor instead.*",
-    )
-    # Torchvision warnings. We don't actually use torchvision.
-    warnings.filterwarnings(
-        action="ignore",
-        message="failed to load.*",
-        module="torchvision.io.image",
-    )
-
-
-def set_env_variables():
-    os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-
-def prepare_cli_environment(log_filter_type: Optional[LogFilterType] = None):
-    if log_filter_type is None:
-        log_filter_type = LogFilterType(os.environ.get("LOG_FILTER_TYPE", "rank0_only"))
-    rich.reconfigure(width=max(rich.get_console().width, 180), soft_wrap=True)
-    setup_logging(log_filter_type=log_filter_type)
-    install_excepthook()
-    filter_warnings()
-    set_env_variables()
-
-
-def clean_opt(arg: str) -> str:
-    if "=" not in arg:
-        arg = f"{arg}=True"
-    name, val = arg.split("=", 1)
-    name = name.strip("-").replace("-", "_")
-    return f"{name}={val}"
-
-
-class RichHandler(logging.Handler):
-    """
-    A simplified version of rich.logging.RichHandler from
-    https://github.com/Textualize/rich/blob/master/rich/logging.py
-    """
-
-    def __init__(
-        self,
-        *,
-        level: Union[int, str] = logging.NOTSET,
-        console: Optional[Console] = None,
-        markup: bool = False,
-    ) -> None:
-        super().__init__(level=level)
-        self.console = console or rich.get_console()
-        self.highlighter = NullHighlighter()
-        self.markup = markup
-
-    def emit(self, record: logging.LogRecord) -> None:
-        try:
-            if hasattr(record.msg, "__rich__") or hasattr(record.msg, "__rich_console__"):
-                self.console.print(record.msg)
-            else:
-                msg: Any = record.msg
-                if isinstance(record.msg, str):
-                    msg = self.render_message(record=record, message=record.getMessage())
-                renderables = [
-                    self.get_time_text(record),
-                    self.get_level_text(record),
-                    self.get_location_text(record),
-                    msg,
-                ]
-                if record.exc_info is not None:
-                    tb = Traceback.from_exception(*record.exc_info)  # type: ignore
-                    renderables.append(tb)
-                self.console.print(*renderables)
-        except Exception:
-            self.handleError(record)
-
-    def render_message(self, *, record: logging.LogRecord, message: str) -> ConsoleRenderable:
-        use_markup = getattr(record, "markup", self.markup)
-        message_text = Text.from_markup(message) if use_markup else Text(message)
-
-        highlighter = getattr(record, "highlighter", self.highlighter)
-        if highlighter:
-            message_text = highlighter(message_text)
-
-        return message_text
-
-    def get_time_text(self, record: logging.LogRecord) -> Text:
-        log_time = datetime.fromtimestamp(record.created)
-        time_str = log_time.strftime("[%Y-%m-%d %X]")
-        return Text(time_str, style="log.time", end=" ")
-
-    def get_level_text(self, record: logging.LogRecord) -> Text:
-        level_name = record.levelname
-        level_text = Text.styled(level_name.ljust(8), f"logging.level.{level_name.lower()}")
-        level_text.style = "log.level"
-        level_text.end = " "
-        return level_text
-
-    def get_location_text(self, record: logging.LogRecord) -> Text:
-        name_and_line = f"{record.name}:{record.lineno}" if record.name != "root" else "root"
-        text = f"[{name_and_line}, rank={record.local_rank}]"  # type: ignore
-        return Text(text, style="log.path")
-
-
-def wait_for(condition: Callable[[], bool], description: str, timeout: float = 10.0):
-    """Wait for the condition function to return True."""
-    start_time = time.monotonic()
-    while not condition():
-        time.sleep(0.5)
-        if time.monotonic() - start_time > timeout:
-            raise TimeoutError(f"{description} timed out")
-
-
-def is_url(path: PathOrStr) -> bool:
-    return re.match(r"[a-z0-9]+://.*", str(path)) is not None
-
-
-def dir_is_empty(dir: PathOrStr) -> bool:
-    dir = Path(dir)
-    if not dir.is_dir():
-        return True
-    try:
-        next(dir.glob("*"))
-        return False
-    except StopIteration:
-        return True
-
-
-def get_progress_bar() -> Progress:
-    from cached_path import get_download_progress
-
-    return get_download_progress()
-
-
-def resource_path(
-    folder: PathOrStr, fname: str, local_cache: Optional[PathOrStr] = None, progress: Optional[Progress] = None
-) -> Path:
-    if local_cache is not None and (local_path := Path(local_cache) / fname).is_file():
-        log.info(f"Found local cache of {fname} at {local_path}")
-        return local_path
-    else:
-        from cached_path import cached_path
-
-        return cached_path(f"{str(folder).rstrip('/')}/{fname}", progress=progress)
-
-
-def file_size(path: PathOrStr) -> int:
-    """
-    Get the size of a local or remote file in bytes.
-    """
-    if is_url(path):
-        from urllib.parse import urlparse
-
-        parsed = urlparse(str(path))
-        if parsed.scheme == "gs":
-            return _gcs_file_size(parsed.netloc, parsed.path.strip("/"))
-        elif parsed.scheme in ("s3", "r2"):
-            return _s3_file_size(parsed.scheme, parsed.netloc, parsed.path.strip("/"))
-        elif parsed.scheme == "file":
-            return file_size(str(path).replace("file://", "", 1))
-        else:
-            raise NotImplementedError(f"file size not implemented for '{parsed.scheme}' files")
-    else:
-        return os.stat(path).st_size
-
-
-def upload(source: PathOrStr, target: str, save_overwrite: bool = False):
-    """Upload source file to a target location on GCS or S3."""
-    from urllib.parse import urlparse
-
-    source = Path(source)
-    assert source.is_file()
-    parsed = urlparse(target)
-    if parsed.scheme == "gs":
-        _gcs_upload(source, parsed.netloc, parsed.path.strip("/"), save_overwrite=save_overwrite)
-    elif parsed.scheme in ("s3", "r2"):
-        _s3_upload(source, parsed.scheme, parsed.netloc, parsed.path.strip("/"), save_overwrite=save_overwrite)
-    else:
-        raise NotImplementedError(f"Upload not implemented for '{parsed.scheme}' scheme")
-
-
-def get_bytes_range(source: PathOrStr, bytes_start: int, num_bytes: int) -> bytes:
-    if is_url(source):
-        from urllib.parse import urlparse
-
-        parsed = urlparse(str(source))
-        if parsed.scheme == "gs":
-            return _gcs_get_bytes_range(parsed.netloc, parsed.path.strip("/"), bytes_start, num_bytes)
-        elif parsed.scheme in ("s3", "r2"):
-            return _s3_get_bytes_range(
-                parsed.scheme, parsed.netloc, parsed.path.strip("/"), bytes_start, num_bytes
-            )
-        elif parsed.scheme == "file":
-            return get_bytes_range(str(source).replace("file://", "", 1), bytes_start, num_bytes)
-        else:
-            raise NotImplementedError(f"file size not implemented for '{parsed.scheme}' files")
-    else:
-        with open(source, "rb") as f:
-            f.seek(bytes_start)
-            return f.read(num_bytes)
-
-
-def find_latest_checkpoint(dir: PathOrStr) -> Optional[PathOrStr]:
-    if is_url(dir):
-        from urllib.parse import urlparse
-
-        parsed = urlparse(str(dir))
-        if parsed.scheme == "gs":
-            raise NotImplementedError
-        elif parsed.scheme in ("s3", "r2"):
-            return _s3_find_latest_checkpoint(parsed.scheme, parsed.netloc, parsed.path.strip("/"))
-        elif parsed.scheme == "file":
-            return find_latest_checkpoint(str(dir).replace("file://", "", 1))
-        else:
-            raise NotImplementedError(f"find_latest_checkpoint not implemented for '{parsed.scheme}' files")
-    else:
-        latest_step = 0
-        latest_checkpoint: Optional[Path] = None
-        for path in Path(dir).glob("step*"):
-            if path.is_dir():
-                try:
-                    step = int(path.name.replace("step", "").replace("-unsharded", ""))
-                except ValueError:
-                    continue
-                # We prioritize sharded checkpoints over unsharded checkpoints.
-                if step > latest_step or (step == latest_step and not path.name.endswith("-unsharded")):
-                    latest_step = step
-                    latest_checkpoint = path
-        return latest_checkpoint
-
-
-def _gcs_upload(source: Path, bucket_name: str, key: str, save_overwrite: bool = False):
-    from google.cloud import storage as gcs
-
-    storage_client = gcs.Client()
-    bucket = storage_client.bucket(bucket_name)
-    blob = bucket.blob(key)
-    if not save_overwrite and blob.exists():
-        raise FileExistsError(f"gs://{bucket_name}/{key} already exists. Use save_overwrite to overwrite it.")
-    blob.upload_from_filename(source)
-
-
-def _gcs_file_size(bucket_name: str, key: str) -> int:
-    from google.api_core.exceptions import NotFound
-    from google.cloud import storage as gcs
-
-    storage_client = gcs.Client()
-    bucket = storage_client.bucket(bucket_name)
-    blob = bucket.blob(key)
-    try:
-        blob.reload()
-    except NotFound:
-        raise FileNotFoundError(f"gs://{bucket_name}/{key}")
-    assert blob.size is not None
-    return blob.size
-
-
-def _gcs_get_bytes_range(bucket_name: str, key: str, bytes_start: int, num_bytes: int) -> bytes:
-    from google.api_core.exceptions import NotFound
-    from google.cloud import storage as gcs
-
-    storage_client = gcs.Client()
-    bucket = storage_client.bucket(bucket_name)
-    blob = bucket.blob(key)
-    try:
-        blob.reload()
-    except NotFound:
-        raise FileNotFoundError(f"gs://{bucket_name}/{key}")
-    return blob.download_as_bytes(start=bytes_start, end=bytes_start + num_bytes - 1)
-
-
-def _get_s3_profile_name(scheme: str) -> Optional[str]:
-    if scheme == "s3":
-        # For backwards compatibility, we assume S3 uses the default profile if S3_PROFILE is not set.
-        return os.environ.get("S3_PROFILE")
-    if scheme == "r2":
-        profile_name = os.environ.get("R2_PROFILE")
-        if profile_name is None:
-            raise OLMoEnvironmentError(
-                "R2 profile name is not set. Did you forget to set the 'R2_PROFILE' env var?"
-            )
-
-        return profile_name
-
-    raise NotImplementedError(f"Cannot get profile name for scheme {scheme}")
-
-
-def _get_s3_endpoint_url(scheme: str) -> Optional[str]:
-    if scheme == "s3":
-        return None
-    if scheme == "r2":
-        r2_endpoint_url = os.environ.get("R2_ENDPOINT_URL")
-        if r2_endpoint_url is None:
-            raise OLMoEnvironmentError(
-                "R2 endpoint url is not set. Did you forget to set the 'R2_ENDPOINT_URL' env var?"
-            )
-
-        return r2_endpoint_url
-
-    raise NotImplementedError(f"Cannot get endpoint url for scheme {scheme}")
-
-
-@cache
-def _get_s3_client(scheme: str):
-    session = boto3.Session(profile_name=_get_s3_profile_name(scheme))
-    return session.client(
-        "s3",
-        endpoint_url=_get_s3_endpoint_url(scheme),
-        config=Config(retries={"max_attempts": 10, "mode": "standard"}),
-        use_ssl=not int(os.environ.get("OLMO_NO_SSL", "0")),
-    )
-
-
-def _wait_before_retry(attempt: int):
-    time.sleep(min(0.5 * 2**attempt, 3.0))
-
-
-def _s3_upload(
-    source: Path, scheme: str, bucket_name: str, key: str, save_overwrite: bool = False, max_attempts: int = 3
-):
-    err: Optional[Exception] = None
-    if not save_overwrite:
-        for attempt in range(1, max_attempts + 1):
-            try:
-                _get_s3_client(scheme).head_object(Bucket=bucket_name, Key=key)
-                raise FileExistsError(
-                    f"s3://{bucket_name}/{key} already exists. Use save_overwrite to overwrite it."
-                )
-            except boto_exceptions.ClientError as e:
-                if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
-                    err = None
-                    break
-                err = e
-
-            if attempt < max_attempts:
-                log.warning("%s failed attempt %d with retriable error: %s", _s3_upload.__name__, attempt, err)
-                _wait_before_retry(attempt)
-
-        if err is not None:
-            raise OLMoNetworkError(f"Failed to check object existence during {scheme} upload") from err
-
-    try:
-        _get_s3_client(scheme).upload_file(source, bucket_name, key)
-    except boto_exceptions.ClientError as e:
-        raise OLMoNetworkError(f"Failed to upload to {scheme}") from e
-
-
-def _s3_file_size(scheme: str, bucket_name: str, key: str, max_attempts: int = 3) -> int:
-    err: Optional[Exception] = None
-    for attempt in range(1, max_attempts + 1):
-        try:
-            return _get_s3_client(scheme).head_object(Bucket=bucket_name, Key=key)["ContentLength"]
-        except boto_exceptions.ClientError as e:
-            if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
-                raise FileNotFoundError(f"s3://{bucket_name}/{key}") from e
-            err = e
-
-        if attempt < max_attempts:
-            log.warning("%s failed attempt %d with retriable error: %s", _s3_file_size.__name__, attempt, err)
-            _wait_before_retry(attempt)
-
-    raise OLMoNetworkError(f"Failed to get {scheme} file size") from err
-
-
-def _s3_get_bytes_range(
-    scheme: str, bucket_name: str, key: str, bytes_start: int, num_bytes: int, max_attempts: int = 3
-) -> bytes:
-    err: Optional[Exception] = None
-    for attempt in range(1, max_attempts + 1):
-        try:
-            return (
-                _get_s3_client(scheme)
-                .get_object(
-                    Bucket=bucket_name, Key=key, Range=f"bytes={bytes_start}-{bytes_start + num_bytes - 1}"
-                )["Body"]
-                .read()
-            )
-        except boto_exceptions.ClientError as e:
-            if e.response["ResponseMetadata"]["HTTPStatusCode"] == 404:
-                raise FileNotFoundError(f"{scheme}://{bucket_name}/{key}") from e
-            err = e
-        except (boto_exceptions.HTTPClientError, boto_exceptions.ConnectionError) as e:
-            # ResponseStreamingError (subclass of HTTPClientError) can happen as
-            # a result of a failed read from the stream (http.client.IncompleteRead).
-            # Retrying can help in this case.
-            err = e
-
-        if attempt < max_attempts:
-            log.warning(
-                "%s failed attempt %d with retriable error: %s", _s3_get_bytes_range.__name__, attempt, err
-            )
-            _wait_before_retry(attempt)
-
-    # When torch's DataLoader intercepts exceptions, it may try to re-raise them
-    # by recalling their constructor with a single message arg. Torch has some
-    # logic to deal with the absence of a single-parameter constructor, but it
-    # doesn't gracefully handle other possible failures in calling such a constructor
-    # This can cause an irrelevant exception (e.g. KeyError: 'error'), resulting
-    # in us losing the true exception info. To avoid this, we change the exception
-    # to a type that has a single-parameter constructor.
-    raise OLMoNetworkError(f"Failed to get bytes range from {scheme}") from err
-
-
-def _s3_find_latest_checkpoint(scheme: str, bucket_name: str, prefix: str) -> Optional[str]:
-    if not prefix.endswith("/"):
-        prefix = f"{prefix}/"
-    response = _get_s3_client(scheme).list_objects(Bucket=bucket_name, Prefix=prefix, Delimiter="/")
-    assert not response["IsTruncated"]  # need to handle this if it happens
-    latest_step = 0
-    latest_checkpoint: Optional[str] = None
-    for item in response["CommonPrefixes"]:
-        prefix = item["Prefix"].strip("/")
-        checkpoint_name = os.path.split(prefix)[-1]
-        if not checkpoint_name.startswith("step"):
-            continue
-        try:
-            step = int(checkpoint_name.replace("step", "").replace("-unsharded", ""))
-        except ValueError:
-            continue
-        # Make sure the checkpoint dir contains a config, otherwise the checkpoint is incomplete
-        # (upload might have have failed part way through).
-        try:
-            _s3_file_size(scheme, bucket_name, f"{prefix}/config.yaml")
-        except FileNotFoundError:
-            continue
-        # We prioritize sharded checkpoints over unsharded ones.
-        if step > latest_step or (step == latest_step and not checkpoint_name.endswith("-unsharded")):
-            latest_step = step
-            latest_checkpoint = f"{scheme}://ai2-llm/{prefix}"
-    return latest_checkpoint
-
-
-def default_thread_count() -> int:
-    return int(os.environ.get("OLMO_NUM_THREADS") or min(32, (os.cpu_count() or 1) + 4))
-
-
-def pass_through_fn(fn, *args, **kwargs):
-    return fn(*args, **kwargs)
-
-
-def threaded_generator(g, maxsize: int = 16, thread_name: Optional[str] = None):
-    q: Queue = Queue(maxsize=maxsize)
-
-    sentinel = object()
-
-    def fill_queue():
-        try:
-            for value in g:
-                q.put(value)
-        except Exception as e:
-            q.put(e)
-        finally:
-            q.put(sentinel)
-
-    thread_name = thread_name or repr(g)
-    thread = Thread(name=thread_name, target=fill_queue, daemon=True)
-    thread.start()
-
-    for x in iter(q.get, sentinel):
-        if isinstance(x, Exception):
-            raise OLMoThreadError(f"generator thread {thread_name} failed") from x
-        else:
-            yield x
-
-
-def roundrobin(*iterables):
-    """
-    Call the given iterables in a round-robin fashion. For example:
-    ``roundrobin('ABC', 'D', 'EF') --> A D E B F C``
-    """
-    # Adapted from https://docs.python.org/3/library/itertools.html#itertools-recipes
-    num_active = len(iterables)
-    nexts = cycle(iter(it).__next__ for it in iterables)
-    while num_active:
-        try:
-            for next in nexts:
-                yield next()
-        except StopIteration:
-            # Remove the iterator we just exhausted from the cycle.
-            num_active -= 1
-            nexts = cycle(islice(nexts, num_active))