Repository: allenai/OLMo
File: olmo/util.py (Python)

Code before:

def _s3_upload(source: Path, bucket_name: str, key: str, save_overwrite: bool = False):
    from botocore.exceptions import ClientError

    if not save_overwrite:
        try:
            s3_client.head_object(Bucket=bucket_name, Key=key)
            raise FileExistsError(f"s3://{bucket_name}/{key} already exists. Use save_overwrite to overwrite it.")
        except ClientError as e:
            if int(e.response["Error"]["Code"]) != 404:
                raise
    s3_client.upload_file(source, bucket_name, key)


def _s3_file_size(bucket_name: str, key: str) -> int:
    from botocore.exceptions import ClientError

    try:
        return s3_client.head_object(Bucket=bucket_name, Key=key)["ContentLength"]
    except ClientError as e:
        if int(e.response["Error"]["Code"]) != 404:
            raise
        raise FileNotFoundError(f"s3://{bucket_name}/{key}")


def _s3_get_bytes_range(bucket_name: str, key: str, bytes_start: int, num_bytes: int, max_attempts: int = 2) -> bytes:
    from botocore.exceptions import ClientError

    err = None
    for _ in range(max_attempts):
        try:
            return s3_client.get_object(
                Bucket=bucket_name, Key=key, Range=f"bytes={bytes_start}-{bytes_start + num_bytes - 1}"
            )["Body"].read()
        except ClientError as e:
            if int(e.response["Error"]["Code"]) == 404:
                raise FileNotFoundError(f"s3://{bucket_name}/{key}")
            err = e
        except ResponseStreamingError as e:
            # ResponseStreamingError can happen as a result of a failed read from
            # the stream (http.client.IncompleteRead) Retrying can help in this case.
            err = e

    # 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.e. 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('Failed to get bytes range from s3') from err

Review comment (suggestion):
I feel like 3 retries is a good default.
```suggestion
def _s3_get_bytes_range(bucket_name: str, key: str, bytes_start: int, num_bytes: int, max_attempts: int = 3) -> bytes:
```

Code after:

def _wait_before_retry(attempt: int):
    time.sleep(min(0.5 * 2**attempt, 3.0))


def _s3_upload(source: Path, 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:
                s3_client.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 int(e.response["Error"]["Code"]) == 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("Failed to check object existence during s3 upload") from err

    try:
        s3_client.upload_file(source, bucket_name, key)
    except boto_exceptions.ClientError as e:
        raise OlmoNetworkError("Failed to upload to s3") from e


def _s3_file_size(bucket_name: str, key: str, max_attempts: int = 3) -> int:
    err: Optional[Exception] = None
    for attempt in range(1, max_attempts + 1):
        try:
            return s3_client.head_object(Bucket=bucket_name, Key=key)["ContentLength"]
        except boto_exceptions.ClientError as e:
            if int(e.response["Error"]["Code"]) == 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("Failed to get s3 file size") from err


def _s3_get_bytes_range(