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

kumoai/__init__.py

@@ -0,0 +1,300 @@
+import warnings
+import os
+import sys
+import threading
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+import logging
+
+from kumoapi.typing import Dtype, Stype
+
+from kumoai.client.client import KumoClient
+from kumoai._logging import initialize_logging, _ENV_KUMO_LOG
+from kumoai._singleton import Singleton
+from kumoai.futures import create_future, initialize_event_loop
+from kumoai.spcs import (
+    _get_active_session,
+    _get_spcs_token,
+    _run_refresh_spcs_token,
+)
+
+initialize_logging()
+initialize_event_loop()
+
+
+@dataclass
+class GlobalState(metaclass=Singleton):
+    r"""Global storage of the state needed to create a Kumo client object. A
+    singleton so its initialized state can be referenced elsewhere for free.
+    """
+
+    # NOTE fork semantics: CoW on Linux, and re-execed on Windows. So this will
+    # likely not work on Windows unless we have special handling for the shared
+    # state:
+    _url: Optional[str] = None
+    _api_key: Optional[str] = None
+    _snowflake_credentials: Optional[Dict[str, Any]] = None
+    _spcs_token: Optional[str] = None
+    _snowpark_session: Optional[Any] = None
+
+    thread_local: threading.local = threading.local()
+
+    def clear(self) -> None:
+        if hasattr(self.thread_local, '_client'):
+            del self.thread_local._client
+        self._url = None
+        self._api_key = None
+        self._snowflake_credentials = None
+        self._spcs_token = None
+
+    def set_spcs_token(self, spcs_token: str) -> None:
+        # Set the spcs token in the global state. This will be picked up the
+        # next time client() is accessed.
+        self._spcs_token = spcs_token
+
+    @property
+    def initialized(self) -> bool:
+        return self._url is not None and (
+            self._api_key is not None or self._snowflake_credentials
+            is not None or self._snowpark_session is not None)
+
+    @property
+    def client(self) -> KumoClient:
+        r"""Accessor for the Kumo client. Note that clients are stored as
+        thread-local variables as the requests Session library is not
+        guaranteed to be thread-safe.
+
+        For more information, see https://github.com/psf/requests/issues/1871.
+        """
+        if self._url is None or (self._api_key is None
+                                 and self._spcs_token is None
+                                 and self._snowpark_session is None):
+            raise ValueError("Client creation or authentication failed. "
+                             "Please re-create your client before proceeding.")
+
+        if hasattr(self.thread_local, '_client'):
+            # Set the spcs token in the client to ensure it has the latest.
+            self.thread_local._client.set_spcs_token(self._spcs_token)
+            return self.thread_local._client
+
+        client = KumoClient(self._url, self._api_key, self._spcs_token)
+        self.thread_local._client = client
+        return client
+
+    @property
+    def is_spcs(self) -> bool:
+        return (self._api_key is None
+                and (self._snowflake_credentials is not None
+                     or self._snowpark_session is not None))
+
+
+global_state: GlobalState = GlobalState()
+
+
+def init(
+    url: Optional[str] = None,
+    api_key: Optional[str] = None,
+    snowflake_credentials: Optional[Dict[str, str]] = None,
+    snowflake_application: Optional[str] = None,
+    log_level: str = "INFO",
+) -> None:
+    r"""Initializes and authenticates the API key against the Kumo service.
+    Successful authentication is required to use the SDK.
+
+    Example:
+        >>> import kumoai
+        >>> kumoai.init(url="<api_url>", api_key="<api_key>")  # doctest: +SKIP
+
+    Args:
+        url: The Kumo API endpoint. Can also be provided via the
+            ``KUMO_API_ENDPOINT`` envronment variable. Will be inferred from
+            the provided API key, if not provided.
+        api_key: The Kumo API key. Can also be provided via the
+            ``KUMO_API_KEY`` environment variable.
+        snowflake_credentials: The Snowflake credentials to authenticate
+            against the Kumo service. The dictionary should contain the keys
+            ``"user"``, ``"password"``, and ``"account"``. This should only be
+            provided for SPCS.
+        snowflake_application: The Snowflake application.
+        log_level: The logging level that Kumo operates under. Defaults to
+            INFO; for more information, please see
+            :class:`~kumoai.set_log_level`. Can also be set with the
+            environment variable ``KUMOAI_LOG``.
+    """  # noqa
+    # Avoid mutations to the global state after it is set:
+    if global_state.initialized:
+        warnings.warn("Kumo SDK already initialized. To re-initialize the "
+                      "SDK, please start a new interpreter. No changes will "
+                      "be made to the current session.")
+        return
+
+    set_log_level(os.getenv(_ENV_KUMO_LOG, log_level))
+
+    # Get API key:
+    api_key = api_key or os.getenv("KUMO_API_KEY")
+
+    snowpark_session = None
+    if snowflake_application:
+        if url is not None:
+            raise ValueError(
+                "Kumo SDK initialization failed. Both 'snowflake_application' "
+                "and 'url' are specified. If running from a Snowflake "
+                "notebook, specify only 'snowflake_application'.")
+        snowpark_session = _get_active_session()
+        if not snowpark_session:
+            raise ValueError(
+                "Kumo SDK initialization failed. 'snowflake_application' is "
+                "specified without an active Snowpark session. If running "
+                "outside a Snowflake notebook, specify a URL and credentials.")
+        description = snowpark_session.sql(
+            f"DESCRIBE SERVICE {snowflake_application}."
+            "USER_SCHEMA.KUMO_SERVICE").collect()[0]
+        url = f"http://{description.dns_name}:8888/public_api"
+
+    if api_key is None and not snowflake_application:
+        if snowflake_credentials is None:
+            raise ValueError(
+                "Kumo SDK initialization failed. Neither an API key nor "
+                "Snowflake credentials provided. Please either set the "
+                "'KUMO_API_KEY' or explicitly call `kumoai.init(...)`.")
+        if (set(snowflake_credentials.keys())
+                != {'user', 'password', 'account'}):
+            raise ValueError(
+                f"Provided Snowflake credentials should be a dictionary with "
+                f"keys 'user', 'password', and 'account'. Only "
+                f"{set(snowflake_credentials.keys())} were provided.")
+
+    # Get or infer URL:
+    url = url or os.getenv("KUMO_API_ENDPOINT")
+    try:
+        if api_key:
+            url = url or f"http://{api_key.split(':')[0]}.kumoai.cloud/api"
+    except KeyError:
+        pass
+    if url is None:
+        raise ValueError("Kumo SDK initialization failed since no endpoint "
+                         "URL was provided. Please either set the "
+                         "'KUMO_API_ENDPOINT' environment variable or "
+                         "explicitly call `kumoai.init(...)`.")
+
+    # Assign global state after verification that client can be created and
+    # authenticated successfully:
+    spcs_token = _get_spcs_token(
+        snowflake_credentials
+    ) if not api_key and snowflake_credentials else None
+    client = KumoClient(url=url, api_key=api_key, spcs_token=spcs_token)
+    client.authenticate()
+    global_state._url = client._url
+    global_state._api_key = client._api_key
+    global_state._snowflake_credentials = snowflake_credentials
+    global_state._spcs_token = client._spcs_token
+    global_state._snowpark_session = snowpark_session
+
+    if not api_key and snowflake_credentials:
+        # Refresh token every 10 minutes (expires in 1 hour):
+        create_future(_run_refresh_spcs_token(minutes=10))
+
+    logger = logging.getLogger('kumoai')
+    log_level = logging.getLevelName(logger.getEffectiveLevel())
+
+    logger.info(f"Initialized Kumo SDK v{__version__} against deployment "
+                f"'{url}'")
+
+
+def set_log_level(level: str) -> None:
+    r"""Sets the Kumo logging level, which defines the amount of output that
+    methods produce.
+
+    Example:
+        >>> import kumoai
+        >>> kumoai.set_log_level("INFO")  # doctest: +SKIP
+
+    Args:
+        level: the logging level. Can be one of (in order of lowest to highest
+            log output) :obj:`DEBUG`, :obj:`INFO`, :obj:`WARNING`,
+            :obj:`ERROR`, :obj:`FATAL`, :obj:`CRITICAL`.
+    """
+    # logging library will ensure `level` is a valid string, and raise a
+    # warning if not:
+    logging.getLogger('kumoai').setLevel(level)
+
+
+# Try to initialize purely with environment variables:
+if ("pytest" not in sys.modules and "KUMO_API_KEY" in os.environ
+        and "KUMO_API_ENDPOINT" in os.environ):
+    init()
+
+import kumoai.connector  # noqa
+import kumoai.encoder  # noqa
+import kumoai.graph  # noqa
+import kumoai.pquery  # noqa
+import kumoai.trainer  # noqa
+import kumoai.utils  # noqa
+import kumoai.databricks  # noqa
+
+from kumoai.connector import (  # noqa
+    SourceTable, SourceTableFuture, S3Connector, SnowflakeConnector,
+    DatabricksConnector, BigQueryConnector, FileUploadConnector, GlueConnector)
+from kumoai.graph import Column, Edge, Graph, Table  # noqa
+from kumoai.pquery import (  # noqa
+    PredictionTableGenerationPlan, PredictiveQuery,
+    TrainingTableGenerationPlan, TrainingTable, TrainingTableJob,
+    PredictionTable, PredictionTableJob)
+from kumoai.trainer import (  # noqa
+    ModelPlan, Trainer, TrainingJobResult, TrainingJob,
+    BatchPredictionJobResult, BatchPredictionJob)
+from kumoai._version import __version__  # noqa
+
+__all__ = [
+    'Dtype',
+    'Stype',
+    'SourceTable',
+    'SourceTableFuture',
+    'S3Connector',
+    'SnowflakeConnector',
+    'DatabricksConnector',
+    'BigQueryConnector',
+    'FileUploadConnector',
+    'GlueConnector',
+    'Column',
+    'Table',
+    'Graph',
+    'Edge',
+    'PredictiveQuery',
+    'TrainingTable',
+    'TrainingTableJob',
+    'TrainingTableGenerationPlan',
+    'PredictionTable',
+    'PredictionTableJob',
+    'PredictionTableGenerationPlan',
+    'Trainer',
+    'TrainingJobResult',
+    'TrainingJob',
+    'BatchPredictionJobResult',
+    'BatchPredictionJob',
+    'ModelPlan',
+    '__version__',
+]
+
+
+def in_snowflake_notebook() -> bool:
+    try:
+        from snowflake.snowpark.context import get_active_session
+        import streamlit  # noqa: F401
+        get_active_session()
+        return True
+    except Exception:
+        return False
+
+
+def in_notebook() -> bool:
+    if in_snowflake_notebook():
+        return True
+    try:
+        from IPython import get_ipython
+        shell = get_ipython()
+        if 'google.colab' in str(shell.__class__):
+            return True
+        return shell.__class__.__name__ == 'ZMQInteractiveShell'
+    except Exception:
+        return False

kumoai/_logging.py

@@ -0,0 +1,29 @@
+import logging
+import os
+
+_ENV_KUMO_LOG = "KUMO_LOG"
+
+
+def initialize_logging() -> None:
+    r"""Initializes Kumo logging."""
+    logger: logging.Logger = logging.getLogger('kumoai')
+
+    # From openai-python/blob/main/src/openai/_utils/_logs.py#L4
+    logging.basicConfig(
+        format=(
+            "[%(asctime)s - %(name)s:%(lineno)d - %(levelname)s] %(message)s"),
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    default_level = os.getenv(_ENV_KUMO_LOG, "INFO")
+    try:
+        logger.setLevel(default_level)
+    except (TypeError, ValueError):
+        logger.setLevel(logging.INFO)
+        logger.warning(
+            "Logging level %s could not be properly parsed. "
+            "Defaulting to INFO log level.", default_level)
+
+    for name in ["matplotlib", "urllib3", "snowflake"]:
+        # TODO(dm) required for spcs
+        logging.getLogger(name).setLevel(logging.ERROR)

kumoai/_singleton.py

@@ -0,0 +1,25 @@
+from abc import ABCMeta
+from typing import Any, Dict
+
+
+class Singleton(ABCMeta):
+    r"""A per-process singleton definition."""
+    _instances: Dict[type, Any] = {}
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        if cls not in cls._instances:
+            # Calls the `__init__` method of the subclass and returns a
+            # reference, which is stored to prevent multiple instantiations.
+            instance = super(Singleton, cls).__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+            return instance
+        return cls._instances[cls]
+
+    def clear(cls) -> None:
+        r"""Clears the singleton class instance, so the next construction
+        will re-initialize the clas.
+        """
+        try:
+            del Singleton._instances[cls]
+        except KeyError:
+            pass

kumoai/_version.py

@@ -0,0 +1 @@
+__version__ = '2.14.0.dev202601011731'

kumoai/artifact_export/__init__.py

@@ -0,0 +1,9 @@
+from .job import ArtifactExportJob, ArtifactExportResult
+from .config import OutputConfig, TrainingTableExportConfig
+
+__all__ = [
+    "ArtifactExportJob",
+    "ArtifactExportResult",
+    "OutputConfig",
+    "TrainingTableExportConfig",
+]

kumoai/artifact_export/config.py

@@ -0,0 +1,209 @@
+import functools
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+
+from kumoapi.jobs import MetadataField, WriteMode
+from kumoapi.typing import WITH_PYDANTIC_V2
+from pydantic.dataclasses import dataclass as pydantic_dataclass
+
+if WITH_PYDANTIC_V2:
+    from pydantic import field_validator, model_validator  # type: ignore
+else:
+    from pydantic import root_validator, validator
+
+from kumoai.connector.base import Connector
+
+
+def compatible_field_validator(field_name: str):  # type: ignore
+    """Decorator factory that creates a field validator compatible with both
+    Pydantic v1 and v2.
+
+    Usage:
+        @compatible_field_validator('field_name')
+        def validate_field(cls, v, values_or_info):
+            # Your validation logic here
+            return v
+    """
+    def decorator(func):  # type: ignore
+        if WITH_PYDANTIC_V2:
+
+            @field_validator(field_name)
+            @classmethod
+            @functools.wraps(func)
+            def wrapper(cls, v, info):  # type: ignore
+                # Convert info to values dict for compatibility
+                values = info.data if hasattr(info, 'data') else {}
+                return func(cls, v, values)
+
+            return wrapper
+        else:
+
+            @validator(field_name)
+            @functools.wraps(func)
+            def wrapper(cls, v, values):  # type: ignore
+                return func(cls, v, values)
+
+            return wrapper
+
+    return decorator
+
+
+# TODO: probably will need to be removed b/c using __post_init__ instead
+def compatible_model_validator(mode='before'):  # type: ignore
+    """Decorator factory that creates a model validator compatible with both
+    Pydantic v1 and v2.
+
+    Usage:
+        @compatible_model_validator()
+        def validate_model(cls, values):
+            # Your validation logic here
+            return values
+    """
+    def decorator(func):  # type: ignore
+        if WITH_PYDANTIC_V2:
+
+            @model_validator(mode=mode)
+            @classmethod
+            @functools.wraps(func)
+            def wrapper(cls, values):  # type: ignore
+                return func(cls, values)
+
+            return wrapper
+        else:
+
+            @root_validator
+            @functools.wraps(func)
+            def wrapper(cls, values):  # type: ignore
+                return func(cls, values)
+
+            return wrapper
+
+    return decorator
+
+
+@dataclass(frozen=True)
+class QueryConnectorConfig:
+    # If using OVERWRITE, big query connector will first write to a staging
+    # table followed by overwriting to the destination table.
+    # When using APPEND, it is strongly recommended to use
+    # MetadataField.JOB_TIMESTAMP to indicate the timestamp of the job.
+    write_mode: WriteMode = WriteMode.APPEND
+
+
+@dataclass(frozen=True)
+class BigQueryOutputConfig(QueryConnectorConfig):
+    pass
+
+
+@dataclass(frozen=True)
+class SnowflakeConnectorConfig(QueryConnectorConfig):
+    pass
+
+
+CONNECTOR_CONFIG_MAPPING = {
+    'BigQueryConnector': BigQueryOutputConfig,
+    'SnowflakeConnector': SnowflakeConnectorConfig,
+    # 'DatabricksConnector': DatabricksOutputConfig,
+    # 'S3Connector': S3OutputConfig,
+}
+
+
+@pydantic_dataclass(frozen=True, config={'arbitrary_types_allowed': True})
+class OutputConfig:
+    """Output configuration associated with a Batch Prediction Job.
+    Specifies the output types and optionally output data source
+    configuration.
+
+    Args:
+        output_types(`Set[str]`): The types of outputs that should be produced
+            by the prediction job. Can include either ``'predictions'``,
+            ``'embeddings'``, or both.
+        output_connector(`Connector` or None): The output data source that Kumo
+                should write batch predictions to, if it is None,
+                produce local download output only.
+        output_table_name(`str` or `Tuple[str, str]` or None): The name of the
+            table in the output data source
+            that Kumo should write batch predictions to. In the case of
+            a Databricks connector, this should be a tuple of two strings:
+            the schema name and the output prediction table name.
+        output_metadata_fields(`List[MetadataField]` or None): Any additional
+            metadata fields to include as new columns in the produced
+            ``'predictions'`` output. Currently, allowed options are
+            ``JOB_TIMESTAMP`` and ``ANCHOR_TIMESTAMP``.
+        connector_specific_config(`QueryConnectorConfig` or None): The custom
+            connector specific output config for predictions, for
+            example whether to append or overwrite existing table.
+    """
+    output_types: Set[str]
+    output_connector: Optional[Connector] = None
+    output_table_name: Optional[Union[str, Tuple]] = None
+    output_metadata_fields: Optional[List[MetadataField]] = None
+    connector_specific_config: Optional[Union[
+        BigQueryOutputConfig,
+        SnowflakeConnectorConfig,
+    ]] = None
+
+    @compatible_field_validator('connector_specific_config')
+    def validate_connector_config(cls, v: Any, values: Dict) -> Any:
+        """Validate the connector specific output config. Raises ValueError if
+        there is a mismatch between the connector type and the config type.
+        """
+        # Skip validation if no connector or no specific config
+        if values.get('output_connector') is None or v is None:
+            return v
+
+        connector_type = type(values['output_connector']).__name__
+        expected_config_type = CONNECTOR_CONFIG_MAPPING.get(connector_type)
+
+        # If we don't have a mapping for this connector type, it doesn't
+        # support specific configs yet
+        if expected_config_type is None:
+            raise ValueError(
+                f"Connector type '{connector_type}' does not support "
+                f"specific output configurations")
+
+        # Check if the provided config is of the correct type
+        if not isinstance(v, expected_config_type):
+            raise ValueError(
+                f"Connector type '{connector_type}' requires output "
+                f"config of type '{expected_config_type.__name__}', but "
+                f"got '{type(v).__name__}'")
+
+        return v
+
+
+@pydantic_dataclass(frozen=True, config={'arbitrary_types_allowed': True})
+class TrainingTableExportConfig(OutputConfig):
+    """Export configuration associated with a Training Table.
+
+    Args:
+        output_types(`Set[str]`): The artifact to export from the training
+            table job. Currently only `'training_table'` is supported.
+            Which exports the full training table to the output connector.
+        output_connector(`Connector`): The output data source that Kumo should
+            write training table artifacts to.
+        output_table_name(str): The name of the table in the output data source
+            that Kumo should write batch predictions to. In the case of
+            a Databricks connector, this should be a tuple of two strings:
+            the schema name and the output prediction table name.
+        connector_specific_config(QueryConnectorConfig or None):
+            Defines custom connector specific output
+            for example whether to append or overwrite
+            existing table. This is currently only supported for BigQuery and
+            Snowflake.
+    """
+    output_connector: Connector
+    output_table_name: str
+
+    def __post_init__(self) -> None:
+        if self.output_types != {'training_table'}:
+            raise ValueError("output_type must be set(['training_table'])"
+                             f" (got {self.output_types})")
+        if self.output_connector is None:
+            raise ValueError("output_connector is required")
+        if self.output_table_name is None:
+            raise ValueError("output_table_name is required")
+        if self.output_metadata_fields is not None:
+            raise ValueError(
+                "output_metadata_fields is not supported for training"
+                "table export")

kumoai/artifact_export/job.py

@@ -0,0 +1,108 @@
+import asyncio
+import concurrent
+import concurrent.futures
+import time
+
+from kumoapi.common import JobStatus
+from typing_extensions import override
+
+from kumoai import global_state
+from kumoai.futures import KumoProgressFuture, create_future
+
+
+class ArtifactExportResult:
+    r"""Represents a completed artifact export job."""
+    def __init__(self, job_id: str) -> None:
+        self.job_id = job_id
+
+    def tracking_url(self) -> str:
+        r"""Returns a tracking URL pointing to the UI display of
+        this prediction export job.
+        """
+        raise NotImplementedError
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(job_id={self.job_id})"
+
+
+class ArtifactExportJob(KumoProgressFuture[ArtifactExportResult]):
+    """Represents an in-progress artifact export job."""
+    def __init__(self, job_id: str) -> None:
+        self.job_id = job_id
+        self._fut: concurrent.futures.Future = create_future(
+            _poll_export(job_id))
+
+    @property
+    def id(self) -> str:
+        """The unique ID of this export job."""
+        return self.job_id
+
+    @override
+    def result(self) -> ArtifactExportResult:
+        return self._fut.result()
+
+    @override
+    def future(self) -> 'concurrent.futures.Future[ArtifactExportResult]':
+        return self._fut
+
+    @override
+    def _attach_internal(
+        self,
+        interval_s: float = 20.0,
+    ) -> ArtifactExportResult:
+        """Allows a user to attach to a running export job and view
+        its progress.
+
+        Args:
+            interval_s (float): Time interval (seconds) between polls, minimum
+                value allowed is 4 seconds.
+        """
+        assert interval_s >= 4.0
+        print(f"Attaching to export job {self.job_id}. To detach from "
+              f"this job, please enter Ctrl+C (the job will continue to run, "
+              f"and you can re-attach anytime).")
+
+        # TODO improve print statements.
+        # Will require changes to status to return
+        # JobStatusReport instead of JobStatus.
+        while not self.done():
+            status = self.status()
+            print(f"Export job {self.job_id} status: {status}")
+            time.sleep(interval_s)
+
+        return self.result()
+
+    def status(self) -> JobStatus:
+        """Returns the status of a running export job."""
+        return get_export_status(self.job_id)
+
+    def cancel(self) -> bool:
+        """Cancels a running export job.
+
+        Returns:
+            bool: True if the job is in a terminal state.
+        """
+        api = global_state.client.artifact_export_api
+        status = api.cancel(self.job_id)
+        if status == JobStatus.CANCELLED:
+            return True
+        return False
+
+
+def get_export_status(job_id: str) -> JobStatus:
+    api = global_state.client.artifact_export_api
+    resource = api.get(job_id)
+    return resource
+
+
+async def _poll_export(job_id: str) -> ArtifactExportResult:
+    status = get_export_status(job_id)
+    while not status.is_terminal:
+        await asyncio.sleep(10)
+        status = get_export_status(job_id)
+
+    if status != JobStatus.DONE:
+        raise RuntimeError(f"Export job {job_id} failed "
+                           f"with job status {status}.")
+
+    return ArtifactExportResult(job_id=job_id)

kumoai/client/__init__.py

@@ -0,0 +1,5 @@
+from .client import KumoClient
+
+__all__ = [
+    'KumoClient',
+]