kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/experimental/rfm/sagemaker.py

@@ -0,0 +1,138 @@
+import base64
+import json
+from typing import Any
+
+import requests
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import Endpoint, HTTPMethod
+from kumoai.exceptions import HTTPException
+
+try:
+    # isort: off
+    from mypy_boto3_sagemaker_runtime.client import SageMakerRuntimeClient
+    from mypy_boto3_sagemaker_runtime.type_defs import (
+        InvokeEndpointOutputTypeDef, )
+    # isort: on
+except ImportError:
+    SageMakerRuntimeClient = Any
+    InvokeEndpointOutputTypeDef = Any
+
+
+class SageMakerResponseAdapter(requests.Response):
+    def __init__(self, sm_response: InvokeEndpointOutputTypeDef):
+        super().__init__()
+        # Read the body bytes
+        self._content = sm_response['Body'].read()
+        self.status_code = 200
+        self.headers['Content-Type'] = sm_response.get('ContentType',
+                                                       'application/json')
+        # Optionally, you can store original sm_response for debugging
+        self.sm_response = sm_response
+
+    @property
+    def text(self) -> str:
+        assert isinstance(self._content, bytes)
+        return self._content.decode('utf-8')
+
+    def json(self, **kwargs) -> dict[str, Any]:  # type: ignore
+        return json.loads(self.text, **kwargs)
+
+
+class KumoClient_SageMakerAdapter(KumoClient):
+    def __init__(self, region: str, endpoint_name: str):
+        import boto3
+        self._client: SageMakerRuntimeClient = boto3.client(
+            service_name="sagemaker-runtime", region_name=region)
+        self._endpoint_name = endpoint_name
+
+        # Recording buffers.
+        self._recording_active = False
+        self._recorded_reqs: list[dict[str, Any]] = []
+        self._recorded_resps: list[dict[str, Any]] = []
+
+    def authenticate(self) -> None:
+        # TODO(siyang): call /ping to verify?
+        pass
+
+    def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
+        assert endpoint.method == HTTPMethod.POST
+        if 'json' in kwargs:
+            payload = json.dumps(kwargs.pop('json'))
+        elif 'data' in kwargs:
+            raw_payload = kwargs.pop('data')
+            assert isinstance(raw_payload, bytes)
+            payload = base64.b64encode(raw_payload).decode()
+        else:
+            raise HTTPException(400, 'Unable to send data to KumoRFM.')
+
+        request = {
+            'method': endpoint.get_path().rsplit('/')[-1],
+            'payload': payload,
+        }
+        response: InvokeEndpointOutputTypeDef = self._client.invoke_endpoint(
+            EndpointName=self._endpoint_name,
+            ContentType="application/json",
+            Body=json.dumps(request),
+        )
+
+        adapted_response = SageMakerResponseAdapter(response)
+
+        # If validation is active, store input/output
+        if self._recording_active:
+            self._recorded_reqs.append(request)
+            self._recorded_resps.append(adapted_response.json())
+
+        return adapted_response
+
+    def start_recording(self) -> None:
+        """Start recording requests/responses to/from sagemaker endpoint."""
+        assert not self._recording_active
+        self._recording_active = True
+        self._recorded_reqs.clear()
+        self._recorded_resps.clear()
+
+    def end_recording(self) -> list[tuple[dict[str, Any], dict[str, Any]]]:
+        """Stop recording and return recorded requests/responses."""
+        assert self._recording_active
+        self._recording_active = False
+        recorded = list(zip(self._recorded_reqs, self._recorded_resps))
+        self._recorded_reqs.clear()
+        self._recorded_resps.clear()
+        return recorded
+
+
+class KumoClient_SageMakerProxy_Local(KumoClient):
+    def __init__(self, url: str):
+        self._client = KumoClient(url, api_key=None)
+        self._client._api_url = self._client._url
+        self._endpoint = Endpoint('/invocations', HTTPMethod.POST)
+
+    def authenticate(self) -> None:
+        try:
+            self._client._session.get(
+                self._url + '/ping',
+                verify=self._verify_ssl).raise_for_status()
+        except Exception:
+            raise ValueError(
+                "Client authentication failed. Please check if you "
+                "have a valid API key/credentials.")
+
+    def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
+        assert endpoint.method == HTTPMethod.POST
+        if 'json' in kwargs:
+            payload = json.dumps(kwargs.pop('json'))
+        elif 'data' in kwargs:
+            raw_payload = kwargs.pop('data')
+            assert isinstance(raw_payload, bytes)
+            payload = base64.b64encode(raw_payload).decode()
+        else:
+            raise HTTPException(400, 'Unable to send data to KumoRFM.')
+        return self._client._request(
+            self._endpoint,
+            json={
+                'method': endpoint.get_path().rsplit('/')[-1],
+                'payload': payload,
+            },
+            **kwargs,
+        )

kumoai/experimental/rfm/task_table.py

@@ -0,0 +1,231 @@
+from collections.abc import Sequence
+
+import pandas as pd
+from kumoapi.task import TaskType
+from kumoapi.typing import Dtype, Stype
+
+from kumoai.experimental.rfm.base import Column
+from kumoai.experimental.rfm.infer import contains_timestamp, infer_dtype
+
+
+class TaskTable:
+    r"""A :class:`TaskTable` fully specifies the task, *i.e.* its context and
+    prediction examples with entity IDs, targets and timestamps.
+
+    Args:
+        task_type: The task type.
+        context_df: The data frame holding context examples.
+        pred_df: The data frame holding prediction examples.
+        entity_table_name: The entity table to predict for. For link prediction
+            tasks, needs to hold both entity and target table names.
+        entity_column: The name of the entity column.
+        target_column: The name of the target column.
+        time_column: The name of the time column, if it exists.
+    """
+    def __init__(
+        self,
+        task_type: TaskType,
+        context_df: pd.DataFrame,
+        pred_df: pd.DataFrame,
+        entity_table_name: str | Sequence[str],
+        entity_column: str,
+        target_column: str,
+        time_column: str | None = None,
+    ) -> None:
+
+        task_type = TaskType(task_type)
+        if task_type not in {  # Currently supported task types:
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+                TaskType.REGRESSION,
+                TaskType.TEMPORAL_LINK_PREDICTION,
+        }:
+            raise ValueError  # TODO
+        self._task_type = task_type
+
+        # TODO Check dfs (unify from local table)
+        self._context_df = context_df.copy(deep=False)
+        self._pred_df = pred_df.copy(deep=False)
+
+        self._dtype_dict: dict[str, Dtype] = {
+            column_name: infer_dtype(context_df[column_name])
+            for column_name in context_df.columns
+        }
+
+        self._entity_table_names: tuple[str] | tuple[str, str]
+        if isinstance(entity_table_name, str):
+            self._entity_table_names = (entity_table_name, )
+        elif len(entity_table_name) == 1:
+            self._entity_table_names = (entity_table_name[0], )
+        elif len(entity_table_name) == 2:
+            self._entity_table_names = (
+                entity_table_name[0],
+                entity_table_name[1],
+            )
+        else:
+            raise ValueError  # TODO
+
+        self._entity_column: str = ''
+        self._target_column: str = ''
+        self._time_column: str | None = None
+
+        self.entity_column = entity_column
+        self.target_column = target_column
+        if time_column is not None:
+            self.time_column = time_column
+
+    @property
+    def task_type(self) -> TaskType:
+        return self._task_type
+
+    # Entity column ###########################################################
+
+    @property
+    def entity_table_name(self) -> str:
+        return self._entity_table_names[0]
+
+    @property
+    def entity_table_names(self) -> tuple[str] | tuple[str, str]:
+        return self._entity_table_names
+
+    @property
+    def entity_column(self) -> Column:
+        return Column(
+            name=self._entity_column,
+            expr=None,
+            dtype=self._dtype_dict[self._entity_column],
+            stype=Stype.ID,
+        )
+
+    @entity_column.setter
+    def entity_column(self, name: str) -> None:
+        if name in {self._target_column, self._time_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not Stype.ID.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._entity_column = name
+
+    # Target column ###########################################################
+
+    @property
+    def _target_stype(self) -> Stype:
+        if self.task_type in {
+                TaskType.BINARY_CLASSIFICATION,
+                TaskType.MULTICLASS_CLASSIFICATION,
+        }:
+            return Stype.categorical
+        if self.task_type in {TaskType.REGRESSION}:
+            return Stype.numerical
+        if self.task_type.is_link_pred:
+            return Stype.multicategorical
+        raise ValueError
+
+    @property
+    def target_column(self) -> Column:
+        return Column(
+            name=self._target_column,
+            expr=None,
+            dtype=self._dtype_dict[self._target_column],
+            stype=self._target_stype,
+        )
+
+    @target_column.setter
+    def target_column(self, name: str) -> None:
+        if name in {self._entity_column, self._time_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if not self._target_stype.supports_dtype(self._dtype_dict[name]):
+            raise ValueError  # TODO
+
+        self._target_column = name
+
+    # Time column #############################################################
+
+    def has_time_column(self) -> bool:
+        r"""Returns ``True`` if this task has a time column; ``False``
+        otherwise.
+        """
+        return self._time_column is not None
+
+    @property
+    def time_column(self) -> Column | None:
+        r"""The time column of this task.
+
+        The getter returns the time column of this task, or ``None`` if no
+        such time column is present.
+
+        The setter sets a column as a time column for this task, and raises a
+        :class:`ValueError` if the time column has a non-timestamp compatible
+        data type or if the column name does not match a column in the data
+        frame.
+        """
+        if self._time_column is None:
+            return None
+        return Column(
+            name=self._time_column,
+            expr=None,
+            dtype=self._dtype_dict[self._time_column],
+            stype=Stype.timestamp,
+        )
+
+    @time_column.setter
+    def time_column(self, name: str | None) -> None:
+        if name is None:
+            self._time_column = None
+            return
+
+        if name in {self._entity_column, self._target_column}:
+            raise ValueError  # TODO
+        if name not in self._context_df:
+            raise ValueError  # TODO
+        if name not in self._pred_df:
+            raise ValueError  # TODO
+        if not contains_timestamp(
+                ser=self._context_df[name],
+                column_name=name,
+                dtype=self._dtype_dict[name],
+        ):
+            raise ValueError  # TODO
+
+        self._time_column = name
+
+    # Metadata ################################################################
+
+    @property
+    def metadata(self) -> pd.DataFrame:
+        raise NotImplementedError
+
+    def print_metadata(self) -> None:
+        raise NotImplementedError
+
+    # Python builtins #########################################################
+
+    def __hash__(self) -> int:
+        return hash((
+            self.task_type,
+            self.entity_table_names,
+            self._entity_column,
+            self._target_column,
+            self._time_column,
+        ))
+
+    def __repr__(self) -> str:
+        if self.task_type.is_link_pred:
+            entity_table_repr = f'entity_table_names={self.entity_table_names}'
+        else:
+            entity_table_repr = f'entity_table_name={self.entity_table_name}'
+        return (f'{self.__class__.__name__}(\n'
+                f'  task_type={self.task_type},\n'
+                f'  num_context_examples={len(self._context_df)},\n'
+                f'  num_prediction_examples={len(self._pred_df)},\n'
+                f'  {entity_table_repr},\n'
+                f'  entity_column={self._entity_column},\n'
+                f'  target_column={self._target_column},\n'
+                f'  time_column={self._time_column},\n'
+                f')')

kumoai/formatting.py

@@ -0,0 +1,30 @@
+from kumoapi.jobs import ErrorDetails
+
+
+def pretty_print_error_details(error_details: ErrorDetails) -> str:
+    """Pretty prints the ErrorDetails combining all the individual items.
+    If there are CTAs, they are also displayed after creating
+    corresponding hyperlinks.
+
+    Arguments:
+    error_details (ErrorDetails): Standard ErrorDetails response from
+        get_errors APIs.
+    """
+    out = ""
+    ctr = None
+    if len(error_details.items) != 1:
+        out += "Encountered multiple errors:\n"
+        ctr = 1
+    for error_detail in error_details.items:
+        if ctr is not None:
+            out += f'{ctr}.'
+            ctr += 1
+        if error_detail.title is not None:
+            out += f'{error_detail.title}: '
+        out += error_detail.description
+        if error_detail.cta is not None:
+            out += 'Follow the link for potential resolution:'
+            f' {error_detail.cta.url}'
+        out += '\n'
+
+    return out

kumoai/futures.py

@@ -0,0 +1,99 @@
+import asyncio
+import concurrent
+import logging
+import threading
+from abc import ABC, abstractmethod
+from asyncio.events import AbstractEventLoop
+from typing import Any, Awaitable, Coroutine, Generic, TypeVar
+
+logger = logging.getLogger(__name__)
+
+CoroFuncType = Awaitable[Any]
+
+# Kumo global event loop (our implementation of green threads for pollers and
+# other interactions with the Kumo backend that require long-running tasks).
+# Since the caller may have their own event loop that we do not want to
+# mess with, _do not_ ever call `set_event_loop` here!! Instead, be extra
+# cautious to pass this loop everywhere.
+_KUMO_EVENT_LOOP: AbstractEventLoop = asyncio.new_event_loop()
+
+
+def initialize_event_loop() -> None:
+    def _run_background_loop(loop: AbstractEventLoop) -> None:
+        asyncio.set_event_loop(loop)
+        loop.run_forever()
+
+    t = threading.Thread(target=_run_background_loop,
+                         args=(_KUMO_EVENT_LOOP, ), daemon=True)
+    t.start()
+
+
+def create_future(coro: Coroutine[Any, Any, Any]) -> concurrent.futures.Future:
+    r"""Creates a future to execute in the Kumo event loop."""
+    # NOTE this function creates a future, chains it to the output of the
+    # execution of `coro` in the Kumo event loop, and handles exceptions
+    # before scheduling to run in the loop:
+    return asyncio.run_coroutine_threadsafe(coro, _KUMO_EVENT_LOOP)
+
+
+T = TypeVar("T")
+
+
+class KumoFuture(ABC, Generic[T]):
+    r"""Abstract base class for a Kumo future object."""
+
+    # We cannot use Python future implementations (`asyncio.Future` or
+    # `concurrent.futures.Future`) as they are native to the Python
+    # implementation of asyncio and threading, and thus not easily extensible.
+    # Python additionally recommends not exposing low-level Future objects in
+    # user facing APIs.
+    @abstractmethod
+    def result(self) -> T:
+        r"""Returns the resolved state of the future.
+
+        Raises:
+            Exception:
+                If the future is complete but in a failed state due to an
+                exception being raised, this method will raise the same
+                exception.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def future(self) -> 'concurrent.futures.Future[T]':
+        r"""Returns the :obj:`concurrent.futures.Future` object wrapped by
+        this future. It is not recommended to access this object directly.
+        """
+        raise NotImplementedError
+
+    def done(self) -> bool:
+        r"""Returns :obj:`True` if this future has been resolved with
+        ``result()``, or :obj:`False` if this future is still
+        in-progress. Note that this method will return :obj:`True` if the
+        future is complete, but in a failed state, and that this method will
+        return :obj:`False` if the job is complete, but the future has not
+        been awaited.
+        """
+        return self.future().done()
+
+
+class KumoProgressFuture(KumoFuture[T]):
+    @abstractmethod
+    def _attach_internal(self, interval_s: float = 4.0) -> T:
+        raise NotImplementedError
+
+    def attach(self, interval_s: float = 4.0) -> T:
+        r"""Allows a user to attach to a running job and view its progress.
+
+        Args:
+            interval_s (float): Time interval (seconds) between polls, minimum
+                value allowed is 4 seconds.
+        """
+        try:
+            return self._attach_internal(interval_s=interval_s)
+        except Exception:
+            logger.warning(
+                "Detailed job tracking has become temporarily unavailable. "
+                "The job is continuing to proceed on the Kumo server, "
+                "and this call will complete when the job has finished.")
+            return self.result()

kumoai/graph/__init__.py

@@ -0,0 +1,12 @@
+from .column import Column, TimestampUnit
+from .table import Table
+from .graph import Graph, Edge, GraphHealthStats
+
+__all__ = [
+    'TimestampUnit',
+    'Column',
+    'Table',
+    'Graph',
+    'Edge',
+    'GraphHealthStats',
+]

kumoai/graph/column.py

@@ -0,0 +1,106 @@
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+from kumoapi.table import TimestampUnit
+from kumoapi.typing import Dtype, Stype
+
+from kumoai.mixin import CastMixin
+
+
+@dataclass(init=False)
+class Column(CastMixin):
+    r"""A column represents metadata information for a column in a Kumo
+    :class:`~kumoai.graph.Table`. Columns can be created independent of
+    a table, or can be fetched from a table with the
+    :meth:`~kumoai.graph.Table.column` method.
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Fetch a column from a `kumoai.Table`:
+        table = kumoai.Table(...)
+
+        column = table.column("col_name")
+        column = table["col_name"]  # equivalent to the above.
+
+        # Edit a column's data type:
+        print("Existing dtype: ", column.dtype)
+        column.dtype = "int"
+
+        # Edit a column's semantic type:
+        print("Existing stype: ", column.stype)
+        column.stype = "ID"
+
+    Args:
+        name: The name of this column.
+        stype: The semantic type of this column. Semantic types can be
+            specified as strings: the list of possible semantic types
+            is located at :class:`~kumoai.Stype`.
+        dtype: The data type of this column. Data types can be specified
+            as strings: the list of possible data types is located at
+            :class:`~kumoai.Dtype`.
+        timestamp_format: If this column represents a timestamp, the format
+            that the timestamp should be parsed in. The format can either be
+            a :class:`~kumoapi.table.TimestampUnit` for integer columns or a
+            string with a format identifier described
+            `here <https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html>`__
+            for a SaaS Kumo deployment and
+            `here <https://docs.snowflake.com/en/sql-reference/date-time-input-output#about-the-elements-used-in-input-and-output-formats>`__
+            for a Snowpark Container Services Kumo deployment. If left empty,
+            will be intelligently inferred by Kumo.
+    """  # noqa: E501
+    name: str
+    stype: Optional[Stype] = None
+    dtype: Optional[Dtype] = None
+    timestamp_format: Optional[Union[str, TimestampUnit]] = None
+
+    def __init__(
+        self,
+        name: str,
+        stype: Optional[Union[Stype, str]] = None,
+        dtype: Optional[Union[Dtype, str]] = None,
+        timestamp_format: Optional[Union[str, TimestampUnit]] = None,
+    ) -> None:
+        self.name = name
+        self.stype = Stype(stype) if stype is not None else None
+        self.dtype = Dtype(dtype) if dtype is not None else None
+        try:
+            self.timestamp_format = TimestampUnit(timestamp_format)
+        except ValueError:
+            self.timestamp_format = timestamp_format
+
+    def __hash__(self) -> int:
+        return hash((self.name, self.stype, self.dtype, self.timestamp_format))
+
+    def __setattr__(self, key: Any, value: Any) -> None:
+        if key == 'name' and value != getattr(self, key, value):
+            raise AttributeError("Attribute 'name' is read-only")
+        elif key == 'stype' and isinstance(value, str):
+            value = Stype(value)
+        elif key == 'dtype' and isinstance(value, str):
+            value = Dtype(value)
+        elif key == 'timestamp_format' and isinstance(value, str):
+            try:
+                value = TimestampUnit(value)
+            except ValueError:
+                pass
+        super().__setattr__(key, value)
+
+    def update(self, obj: 'Column', override: bool = True) -> 'Column':
+        for key in self.__dict__:
+            if key[0] == '_':  # Skip private attributes:
+                continue
+            value = getattr(obj, key, None)
+            if value is not None:
+                if override or getattr(self, key, None) is None:
+                    setattr(self, key, value)
+        return self
+
+    def __repr__(self) -> str:
+        out = (f"Column(name=\"{self.name}\", stype=\"{self.stype}\", "
+               f"dtype=\"{self.dtype}\"")
+        if self.timestamp_format is not None:
+            out += f", timestamp_format=\"{self.timestamp_format}\""
+        out += ")"
+        return out