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

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

        if attempt < max_attempts:
            _wait_before_retry(attempt)

    raise OlmoNetworkError('Failed to get s3 file size') from err


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

    err: Optional[Exception] = None
    for attempt in range(1, max_attempts + 1):
        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

        if attempt < max_attempts:
            _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('Failed to get bytes range from s3') from err


def is_weight_decay_module(module: nn.Module) -> bool:
    """Returns true if the module should use weight decay."""
    from .model import LayerNormBase

    return not isinstance(module, (LayerNormBase, nn.LayerNorm, nn.Embedding))


Review comment (suggestion):
I was just reading a little more about boto3 error handling and I think we can make this a little more robust. Instead of `ResponseStreamingError` I think we should catch a combination of `HTTPClientError` (the base class for `ResponseStreamingError`) and `ConnectionError` (from botocore, not the Python built-in exception) - the base class for a lot of other retry-able errors.
```suggestion
        except (HTTPClientError, ConnectionError) as e:
```

Code after:
        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(
    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 s3_client.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 int(e.response["Error"]["Code"]) == 404:
                raise FileNotFoundError(f"s3://{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("Failed to get bytes range from s3") from err


def is_weight_decay_module(module: nn.Module) -> bool:
    """Returns true if the module should use weight decay."""
    from .model import LayerNormBase