hyperion-sdk 0.2.0.dev1741815359__py3-none-any.whl

hyperion/dateutils.py

@@ -0,0 +1,238 @@
+import datetime
+import re
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Literal, TypeAlias, cast
+
+from dateutil.relativedelta import relativedelta
+
+from hyperion.logging import get_logger
+
+TIME_UNITS = ["s", "m", "h", "d", "w", "M", "y"]
+PATT_TIME_RESOLUTION = re.compile(rf"(?P<value>\d+)(?P<unit>[{''.join(TIME_UNITS)}])")
+TimeResolutionUnit: TypeAlias = Literal["s", "m", "h", "d", "w", "M", "y"]
+
+logger = get_logger("hyperion-dateutils")
+
+
+@dataclass(frozen=True, eq=True)
+class TimeResolution:
+    value: int
+    unit: TimeResolutionUnit
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.value, int):
+            super().__setattr__("value", int(self.value))
+        if self.unit not in TIME_UNITS:
+            raise ValueError(f"Unknown time unit {self.unit!r}. Pick one of {', '.join(TIME_UNITS)}")
+
+    def __repr__(self) -> str:
+        return f"{self.value}{self.unit}"
+
+    @staticmethod
+    def from_str(string: str) -> "TimeResolution":
+        if (rematch := PATT_TIME_RESOLUTION.match(string)) is None:
+            raise ValueError(f"Invalid time resolution specification {string!r}. Use expressions such as 1d, 5s or 3M.")
+        value = int(rematch.group("value"))
+        unit = cast(TimeResolutionUnit, rematch.group("unit"))
+        return TimeResolution(value=value, unit=unit)
+
+    @property
+    def delta(self) -> datetime.timedelta | relativedelta:
+        match self.unit:
+            case "s":
+                return datetime.timedelta(seconds=self.value)
+            case "m":
+                return datetime.timedelta(minutes=self.value)
+            case "h":
+                return datetime.timedelta(hours=self.value)
+            case "d":
+                return datetime.timedelta(days=self.value)
+            case "w":
+                return datetime.timedelta(days=self.value * 7)
+            case "M":
+                return relativedelta(months=self.value)
+            case "y":
+                return relativedelta(years=self.value)
+            case _:
+                raise ValueError(f"Unsupported time unit {self.unit!r}.")
+
+
+def truncate_datetime(base: datetime.datetime | datetime.date, unit: TimeResolutionUnit) -> datetime.datetime:
+    """Truncate datetime to the specified unit (set all smaller units to zero)."""
+    if not isinstance(base, datetime.datetime):
+        base = datetime.datetime(base.year, base.month, base.day, tzinfo=datetime.timezone.utc)
+    match unit:
+        case "s":
+            return base.replace(microsecond=0)
+        case "m":
+            return base.replace(second=0, microsecond=0)
+        case "h":
+            return base.replace(minute=0, second=0, microsecond=0)
+        case "d":
+            return base.replace(hour=0, minute=0, second=0, microsecond=0)
+        case "w":
+            return base.replace(hour=0, minute=0, second=0, microsecond=0) - datetime.timedelta(days=base.weekday())
+        case "M":
+            return base.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
+        case "y":
+            return base.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0)
+    raise ValueError(f"Unknown time unit {unit!r}. Pick one of {', '.join(TIME_UNITS)}")
+
+
+def iter_dates_between(
+    start_date: datetime.datetime | datetime.date,
+    end_date: datetime.datetime | datetime.date,
+    granularity: TimeResolutionUnit,
+) -> Iterator[datetime.datetime]:
+    """
+    Iterate over datetimes between start_date and end_date with steps based on the given granularity.
+    Includes the start_date and may include end_date (if end date is reachable from start date with given granularity).
+
+    :param start_date: The starting datetime.
+    :param end_date: The ending datetime.
+    :param granularity: The granularity for steps (e.g., "d" for days, "M" for months).
+    :return: An iterator of datetime objects.
+    """
+    start_date = assure_timezone(start_date)
+    end_date = assure_timezone(end_date)
+    logger.debug(
+        "Generating dates between two points.", start_date=start_date.isoformat(), end_date=end_date.isoformat()
+    )
+    if start_date > end_date:
+        raise ValueError("Start date cannot be later than end date.")
+
+    current = start_date
+
+    delta: datetime.timedelta | relativedelta
+
+    match granularity:
+        case "s":
+            delta = datetime.timedelta(seconds=1)
+        case "m":
+            delta = datetime.timedelta(minutes=1)
+        case "h":
+            delta = datetime.timedelta(hours=1)
+        case "d":
+            delta = datetime.timedelta(days=1)
+        case "w":
+            delta = datetime.timedelta(weeks=1)
+        case "M":
+            delta = relativedelta(months=1)
+        case "y":
+            delta = relativedelta(years=1)
+        case default:
+            raise ValueError(f"Unsupported granularity {default!r}.")
+
+    while current <= end_date:
+        yield current
+        current += delta
+
+
+def quantize_datetime(base: datetime.datetime, resolution: TimeResolution | str) -> datetime.datetime:
+    """
+    Quantize a datetime to the next interval based on the specified resolution.
+
+    This function aligns a given datetime to the next moment in time defined
+    by the resolution. The resolution is expressed as a unit (seconds, minutes,
+    hours, or days) and a value (e.g., 5 seconds, 15 minutes).
+
+    **Important Notes:**
+    - The resolution is always calculated relative to the higher unit, which may lead
+      to overlapping intervals for non-standard values. For example:
+        - A 7-second resolution could result in intervals ending at 12:50:56 and 12:51:03,
+          with another interval starting at 12:51:00, causing overlaps.
+    - To avoid such overlaps, it is recommended to use resolutions that are divisors of
+      the higher unit (e.g., 60 for seconds and minutes).
+
+    Parameters:
+    - base (datetime.datetime): The datetime to quantize.
+    - resolution (TimeResolution | str): The resolution for quantization.
+      If a string is provided, it should follow the format "{value}{unit}"
+      (e.g., "5s", "15m", "2h").
+
+    Returns:
+    - datetime.datetime: The quantized datetime.
+
+    Raises:
+    - ValueError: If the resolution unit is unsupported.
+    """
+    resolution = resolution if isinstance(resolution, TimeResolution) else TimeResolution.from_str(resolution)
+    base_truncated = truncate_datetime(base, resolution.unit)
+
+    def _get_shift(value: int) -> int:
+        return resolution.value - (value % resolution.value)
+
+    match resolution.unit:
+        case "s":
+            seconds_shift = _get_shift(base.second)
+            return base_truncated + datetime.timedelta(seconds=seconds_shift)
+        case "m":
+            minutes_shift = _get_shift(base.minute)
+            return base_truncated + datetime.timedelta(minutes=minutes_shift)
+        case "h":
+            hours_shift = _get_shift(base.hour)
+            return base_truncated + datetime.timedelta(hours=hours_shift)
+        case "d":
+            days_shift = _get_shift(base.day)
+            return base_truncated + datetime.timedelta(days=days_shift)
+        case "w":
+            days_shift = resolution.value * 7 - (base.weekday() % resolution.value)
+            return base_truncated + datetime.timedelta(days=days_shift)
+        case "M":
+            months_shift = _get_shift(base.month)
+            return base_truncated + relativedelta(months=months_shift)
+        case "y":
+            years_shift = _get_shift(base.year)
+            return base_truncated + relativedelta(years=years_shift)
+    raise ValueError(f"Unsupported resolution unit {resolution.unit!r} for quantization.")  # pragma: no cover
+
+
+def assure_timezone(
+    base: datetime.datetime | datetime.date, tz: datetime.timezone = datetime.timezone.utc
+) -> datetime.datetime:
+    """Assure datetime has a datetime and return it timezone-aware if not."""
+    if not isinstance(base, datetime.datetime):
+        return datetime.datetime(base.year, base.month, base.day, tzinfo=tz)
+    if base.tzinfo is not None:
+        if base.tzinfo == tz:
+            return base
+        return base.astimezone(tz)
+    logger.warning(f"A timezone-unaware timestamp was given, assuming {tz!r}.")
+    return base.replace(tzinfo=tz)
+
+
+def get_date_pattern(date: datetime.datetime, unit: TimeResolutionUnit) -> str:
+    """Get a date pattern string up to the specified unit level.
+
+    Examples:
+        >>> get_date_pattern(datetime.datetime(2025, 1, 12, 12), "h")
+        "2025-01-12T12"
+        >>> get_date_pattern(datetime.datetime(2025, 1, 12, 12), "d")
+        "2025-01-12"
+        >>> get_date_pattern(datetime.datetime(2025, 1, 12, 12), "M")
+        "2025-01"
+        >>> get_date_pattern(datetime.datetime(2025, 1, 12, 12), "y")
+        "2025"
+    """
+    truncated = truncate_datetime(date, unit)
+    match unit:
+        case "s":
+            return truncated.strftime("%Y-%m-%dT%H:%M:%S")
+        case "m":
+            return truncated.strftime("%Y-%m-%dT%H:%M")
+        case "h":
+            return truncated.strftime("%Y-%m-%dT%H")
+        case "d":
+            return truncated.strftime("%Y-%m-%d")
+        case "w":
+            return truncated.strftime("%Y-%m-%d")  # Week is special, keep the day
+        case "M":
+            return truncated.strftime("%Y-%m")
+        case "y":
+            return truncated.strftime("%Y")
+    raise ValueError(f"Unsupported time unit {unit!r}.")
+
+
+def utcnow() -> datetime.datetime:
+    return datetime.datetime.now(tz=datetime.timezone.utc)

hyperion/entities/catalog.py

@@ -0,0 +1,190 @@
+import datetime
+import json
+from dataclasses import dataclass, field
+from typing import ClassVar, Literal, Protocol
+
+from pydantic import BaseModel
+
+from hyperion.dateutils import TimeResolution, assure_timezone
+
+AssetType = Literal["data_lake", "feature", "persistent_store"]
+
+
+def get_prefixed_path(path: str, prefix: str = "") -> str:
+    """Get the path with the given prefix.
+
+    Args:
+        path (str): The path.
+        prefix (str): The prefix.
+
+    Returns:
+        str: The path with the prefix.
+    """
+    prefix = prefix.strip("/")
+    if prefix:
+        prefix = f"{prefix}/"
+    return prefix + path
+
+
+class AssetProtocol(Protocol):
+    """Protocol for assets in the catalog.
+
+    This protocol defines the interface for assets in the catalog.
+
+    Attributes:
+        asset_type (ClassVar[AssetType]): The type of the asset.
+        name (str): The name of the asset.
+        schema_version (int): The schema version of the asset.
+    """
+
+    asset_type: ClassVar[AssetType]
+
+    @property
+    def name(self) -> str: ...
+
+    @property
+    def schema_version(self) -> int:
+        """The schema version of the asset."""
+
+    def get_path(self, prefix: str = "") -> str:
+        """Get the path for the asset with the given prefix.
+
+        Args:
+            prefix (str): The prefix for the path.
+
+        Returns:
+            str: The path for the asset.
+        """
+
+    def to_metadata(self) -> dict[str, str]:
+        """Get the metadata for the asset.
+
+        Returns:
+            dict[str, str]: The metadata for the asset.
+        """
+
+
+@dataclass(frozen=True, eq=True)
+class DataLakeAsset:
+    """Data lake asset.
+
+    Attributes:
+        asset_type (ClassVar[AssetType]): The type of the asset.
+        name (str): The name of the asset.
+        date (datetime.datetime): The date of the asset.
+        schema_version (int): The schema version of the asset.
+    """
+
+    asset_type: ClassVar[AssetType] = "data_lake"
+    name: str
+    date: datetime.datetime
+    schema_version: int = 1
+
+    def get_path(self, prefix: str = "") -> str:
+        """Get the path for the asset with the given prefix."""
+        date = assure_timezone(self.date).isoformat()
+        return get_prefixed_path(f"{self.name}/date={date}/v{self.schema_version}.avro", prefix)
+
+    def to_metadata(self) -> dict[str, str]:
+        """Get the metadata for the asset."""
+        return {"name": self.name, "date": self.date.isoformat(), "schema_version": str(self.schema_version)}
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.get_path()!r})"
+
+
+@dataclass(frozen=True, eq=True)
+class PersistentStoreAsset:
+    """Persistent store asset.
+
+    Attributes:
+        asset_type (ClassVar[AssetType]): The type of the asset.
+        name (str): The name of the asset.
+        schema_version (int): The schema version of the asset.
+    """
+
+    asset_type: ClassVar[AssetType] = "persistent_store"
+    name: str
+    schema_version: int = 1
+
+    def get_path(self, prefix: str = "") -> str:
+        """Get the path for the asset with the given prefix."""
+        return get_prefixed_path(f"{self.name}/v{self.schema_version}.avro", prefix)
+
+    def to_metadata(self) -> dict[str, str]:
+        """Get the metadata for the asset."""
+        return {"name": self.name, "schema_version": str(self.schema_version)}
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.get_path()!r})"
+
+
+@dataclass(frozen=True, eq=True)
+class FeatureAsset:
+    """Feature asset.
+
+    Attributes:
+        asset_type (ClassVar[AssetType]): The type of the asset.
+        name (str): The name of the asset.
+        partition_date (datetime.datetime): The partition timestamp of the asset.
+        resolution (TimeResolution | str): The resolution of the asset.
+        schema_version (int): The schema version of the asset.
+        partition_keys (dict[str, str]): The partition keys of the asset.
+    """
+
+    asset_type: ClassVar[AssetType] = "feature"
+    name: str
+    partition_date: datetime.datetime
+    resolution: TimeResolution | str
+    schema_version: int = 1
+    partition_keys: dict[str, str] = field(default_factory=dict)
+
+    @property
+    def time_resolution(self) -> TimeResolution:
+        """The time resolution of the feature."""
+        if isinstance(self.resolution, TimeResolution):
+            return self.resolution
+        return TimeResolution.from_str(self.resolution)
+
+    @property
+    def feature_name(self) -> str:
+        """The name of the feature including the time resolution."""
+        return f"{self.name}.{self.time_resolution!r}"
+
+    def _get_partition_keys_prefix(self) -> str:
+        key_names = sorted(self.partition_keys.keys())
+        return ("/".join(f"{key}={self.partition_keys[key]}" for key in key_names)).strip("/")
+
+    def get_path(self, prefix: str = "") -> str:
+        """Get the path for the asset with the given prefix."""
+        partition_date = assure_timezone(self.partition_date).isoformat()
+        keys_prefix = self._get_partition_keys_prefix()
+        if keys_prefix:
+            keys_prefix = keys_prefix + "/"
+        return get_prefixed_path(
+            f"{self.feature_name}/{keys_prefix}partition_date={partition_date}/v{self.schema_version}.avro", prefix
+        )
+
+    def to_metadata(self) -> dict[str, str]:
+        """Get the metadata for the asset."""
+        return {
+            "name": self.name,
+            "partition_date": self.partition_date.isoformat(),
+            "schema_version": str(self.schema_version),
+            "partition_keys": json.dumps(self.partition_keys),
+        }
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.get_path()!r})"
+
+
+class FeatureModel(BaseModel):
+    """A base class for feature models.
+
+    You may use this base class (along with pydantic's BaseModel) to define type-safe feature models.
+    Use with "AssetCollection" to make powerful typed feature collections.
+    """
+
+    asset_name: ClassVar[str] = NotImplemented
+    resolution: ClassVar[TimeResolution] = NotImplemented
+    schema_version: ClassVar[int] = 1

hyperion/infrastructure/aws.py

@@ -0,0 +1,220 @@
+"""AWS helpers and methods."""
+
+import datetime
+import logging
+from asyncio import Semaphore
+from collections.abc import AsyncIterator, Iterator
+from contextlib import ExitStack, asynccontextmanager
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import IO, BinaryIO, TypeVar, cast
+
+import aioboto3
+import boto3
+import botocore.exceptions
+
+from hyperion.config import storage_config
+from hyperion.logging import get_logger
+
+PathOrIOBinary = str | Path | BinaryIO | IO[bytes]
+
+T = TypeVar("T")
+
+logger = get_logger("aws")
+
+logging.getLogger("botocore.endpoint").setLevel("WARNING")
+logging.getLogger("botocore").setLevel("WARNING")
+
+
+class S3StorageClass(str, Enum):
+    """S3 storage classes."""
+
+    STANDARD = "STANDARD"
+    REDUCED_REDUNDANCY = "REDUCED_REDUNDANCY"
+    STANDARD_IA = "STANDARD_IA"
+    ONEZONE_IA = "ONEZONE_IA"
+    INTELLIGENT_TIERING = "INTELLIGENT_TIERING"
+    GLACIER = "GLACIER"
+    DEEP_ARCHIVE = "DEEP_ARCHIVE"
+    OUTPOSTS = "OUTPOSTS"
+    GLACIER_IR = "GLACIER_IR"
+    SNOW = "SNOW"
+    EXPRESS_ONEZONE = "EXPRESS_ONEZONE"
+
+
+@dataclass
+class S3ObjectAttributes:
+    """Attributes of an S3 object."""
+
+    last_modified: datetime.datetime
+    etag: str
+    storage_class: S3StorageClass
+    object_size: int
+
+
+class S3Client:
+    """S3 client."""
+
+    _storage_semaphore: Semaphore | None = None
+
+    @classmethod
+    @asynccontextmanager
+    async def semaphore(cls) -> AsyncIterator[None]:
+        """Semaphore for asynchronous storage operations."""
+        # Lazily initialize the semaphore instance
+        if cls._storage_semaphore is None:
+            logger.debug(
+                "Initializing new storage operations semaphore.", max_concurrency=storage_config.max_concurrency
+            )
+            cls._storage_semaphore = Semaphore(storage_config.max_concurrency)
+        logger.debug("Attempting to acquire the storage operations semaphore.")
+        async with cls._storage_semaphore:
+            logger.debug("Green light, semaphore open.")
+            yield
+
+    def __init__(self) -> None:
+        self._client = boto3.client("s3")
+        self._aio_session = aioboto3.Session()
+
+    async def upload_async(self, file: PathOrIOBinary, bucket: str, name: str) -> None:
+        """Upload a file to S3 asynchronously.
+
+        Args:
+            file (PathOrIOBinary): The file to upload.
+            bucket (str): The bucket to upload the file to.
+            name (str): The name of the file in the bucket.
+        """
+        with ExitStack() as file_context:
+            if isinstance(file, str | Path):
+                path = Path(file)
+                logger.debug("Uploading from path.", path=path.as_posix(), bucket=bucket, name=name)
+                file = file_context.enter_context(path.open("rb"))
+            async with self.semaphore(), self._aio_session.client("s3") as s3:
+                try:
+                    await s3.upload_fileobj(file, bucket, name)
+                except botocore.exceptions.ClientError:
+                    logger.error("Error when uploading file to S3.", bucket=bucket, name=name)
+                    raise
+
+    def iter_objects(self, bucket: str, prefix: str) -> Iterator[str]:
+        """Iterate over objects in an S3 bucket.
+
+        Args:
+            bucket (str): The bucket to list objects in.
+            prefix (str): The prefix to filter objects by.
+
+        Yields:
+            Iterator[str]: The keys of the objects in the bucket.
+        """
+        paginator = self._client.get_paginator("list_objects_v2")
+        pagination_config = {"StartingToken": None}
+        logger.debug("Listing contents of a bucket.", bucket=bucket, prefix=prefix)
+        response_iterator = paginator.paginate(Bucket=bucket, Prefix=prefix, PaginationConfig=pagination_config)
+        for response in response_iterator:
+            logger.debug("Received a response from S3.", response=response)
+            if "Contents" not in response:
+                continue
+            yield from (s3_object["Key"] for s3_object in response["Contents"])
+
+    def upload(self, file: PathOrIOBinary, bucket: str, name: str) -> None:
+        """Upload a file to S3.
+
+        Args:
+            file (PathOrIOBinary): The file to upload.
+            bucket (str): The bucket to upload the file to.
+            name (str): The name of the file in the bucket.
+        """
+        if isinstance(file, str | Path):
+            file = Path(file)
+            logger.debug("Uploading from path.", path=file.as_posix(), bucket=bucket, name=name)
+            try:
+                self._client.upload_file(file, bucket, name)
+            except botocore.exceptions.ClientError:
+                logger.error("Error when uploading file to S3.", file=file.as_posix(), bucket=bucket, name=name)
+                raise
+            return
+        logger.debug("Uploading from an open file stream.", bucket=bucket, name=name)
+        try:
+            self._client.upload_fileobj(file, bucket, name)
+        except botocore.exceptions.ClientError:
+            logger.error("Error when uploading file to S3.", bucket=bucket, name=name)
+            raise
+
+    def get_object_attributes(self, bucket: str, name: str) -> S3ObjectAttributes:
+        """Get the attributes of an object in S3.
+
+        Args:
+            bucket (str): The bucket containing the object.
+            name (str): The name of the object.
+
+        Returns:
+            S3ObjectAttributes: The attributes of the object.
+        """
+        response = self._client.get_object_attributes(
+            Bucket=bucket,
+            Key=name,
+            ObjectAttributes=[
+                "ETag",
+                "Checksum",
+                "ObjectParts",
+                "StorageClass",
+                "ObjectSize",
+            ],
+        )
+        try:
+            return S3ObjectAttributes(
+                last_modified=response["LastModified"],
+                etag=response["ETag"],
+                storage_class=S3StorageClass(response["StorageClass"]),
+                object_size=int(response["ObjectSize"]),
+            )
+        except Exception:
+            logger.error(
+                "Failed to get object attributes, the response is probably invalid.",
+                bucket=bucket,
+                object_name=name,
+                response=response,
+            )
+            raise
+
+    def download(self, bucket: str, name: str, file: PathOrIOBinary) -> None:
+        """Download a file from S3.
+
+        Args:
+            bucket (str): The bucket containing the file.
+            name (str): The name of the file.
+            file (PathOrIOBinary): The file to download to.
+        """
+        if isinstance(file, str | Path):
+            file = Path(file)
+            logger.debug("Downloading into a path.", path=file.as_posix(), bucket=bucket, name=name)
+            try:
+                self._client.download_file(bucket, name, file)
+            except botocore.exceptions.ClientError:
+                logger.error("Error when downloading file from S3.", bucket=bucket, name=name, path=file.as_posix())
+                raise
+            return
+        logger.debug("Downloading into a file stream.", bucket=bucket, name=name)
+        try:
+            self._client.download_fileobj(bucket, name, file)
+        except botocore.exceptions.ClientError:
+            logger.error("Error when downloading file from S3.", bucket=bucket, name=name)
+            raise
+
+    def download_as_string(self, bucket: str, name: str) -> str:
+        """Download an object from S3 as a string.
+
+        Args:
+            bucket (str): The bucket containing the object.
+            name (str): The name of the object.
+
+        Returns:
+            str: The object as a string.
+        """
+        logger.debug("Downloading object as a string.", bucket=bucket, name=name)
+        try:
+            return cast(str, self._client.get_object(Bucket=bucket, Key=name)["Body"].read().decode("utf-8"))
+        except botocore.exceptions.ClientError:
+            logger.error("Error when downloading file from S3.", bucket=bucket, name=name)
+            raise