OLMo-Bitnet-1B / util.py

update inference code

2010c83 2 months ago

No virus

23.4 kB

	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))