kumoai 2.8.0.dev202508221830__cp312-cp312-win_amd64.whl → 2.13.0.dev202512041141__cp312-cp312-win_amd64.whl

kumoai/experimental/rfm/sagemaker.py

@@ -0,0 +1,138 @@
+import base64
+import json
+from typing import Any, Dict, List, Tuple
+
+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/jobs.py

@@ -2,6 +2,7 @@ from abc import ABC, abstractmethod
 from typing import Generic, Mapping, Optional, TypeVar
 
 from kumoapi.jobs import JobStatusReport
+from typing_extensions import Self
 
 from kumoai.client.jobs import CommonJobAPI, JobRequestType, JobResourceType
 
@@ -9,12 +10,37 @@ IDType = TypeVar('IDType', bound=str)
 
 
 class JobInterface(ABC, Generic[IDType, JobRequestType, JobResourceType]):
-    r"""Defines a standard interface for job objects"""
+    r"""Defines a standard interface for job objects."""
     @staticmethod
     @abstractmethod
     def _api() -> CommonJobAPI[JobRequestType, JobResourceType]:
         pass
 
+    @classmethod
+    def search_by_tags(cls, tags: Mapping[str, str],
+                       limit: int = 10) -> list[Self]:
+        r"""Returns a list of job instances from a set of job tags.
+
+        Args:
+            tags (Mapping[str, str]): Tags by which to search.
+            limit (int): Max number of jobs to list, default 10.
+
+        Example:
+            >>> # doctest: +SKIP
+            >>> tags = {'pquery_name': 'my_pquery_name'}
+            >>> jobs = BatchPredictionJob.search_by_tags(tags)
+            Search limited to 10 results based on the `limit` parameter.
+            Found 2 jobs.
+        """
+        print(f"Search limited to {limit} results based on the `limit` "
+              "parameter.")
+
+        jobs = cls._api().list(limit=limit, additional_tags=tags)
+
+        print(f"Found {len(jobs)} jobs.")
+
+        return [cls(j.job_id) for j in jobs]  # type: ignore
+
     @property
     @abstractmethod
     def id(self) -> IDType:

kumoai/pquery/prediction_table.py

@@ -4,6 +4,7 @@ import asyncio
 import logging
 from concurrent.futures import Future
 from datetime import datetime
+from functools import cached_property
 from typing import List, Optional, Union
 
 import pandas as pd
@@ -217,9 +218,10 @@ class PredictionTableJob(JobInterface[GeneratePredictionTableJobID,
     ) -> None:
         self.job_id = job_id
         self.job: Optional[GeneratePredictionTableJobResource] = None
-        # A training table holds a reference to the future that tracks the
-        # execution of its generation.
-        self._fut: Future = create_future(self._poll())
+
+    @cached_property
+    def _fut(self) -> Future:
+        return create_future(self._poll())
 
     @override
     @property

kumoai/pquery/training_table.py

@@ -5,6 +5,7 @@ import logging
 import os
 import time
 from concurrent.futures import Future
+from functools import cached_property
 from typing import List, Optional, Tuple, Union
 
 import pandas as pd
@@ -308,9 +309,10 @@ class TrainingTableJob(JobInterface[GenerateTrainTableJobID,
         job_id: GenerateTrainTableJobID,
     ) -> None:
         self.job_id = job_id
-        # A training table holds a reference to the future that tracks the
-        # execution of its generation.
-        self._fut: Future[TrainingTable] = create_future(_poll(job_id))
+
+    @cached_property
+    def _fut(self) -> Future[TrainingTable]:
+        return create_future(_poll(self.job_id))
 
     @override
     @property

kumoai/spcs.py

@@ -54,9 +54,7 @@ def _refresh_spcs_token() -> None:
         api_key=global_state._api_key,
         spcs_token=spcs_token,
     )
-    if not client.authenticate():
-        raise ValueError("Client authentication failed. Please check if you "
-                         "have a valid API key.")
+    client.authenticate()
 
     # Update state:
     global_state.set_spcs_token(spcs_token)

kumoai/testing/decorators.py

@@ -25,7 +25,7 @@ def onlyFullTest(func: Callable) -> Callable:
 def has_package(package: str) -> bool:
     r"""Returns ``True`` in case ``package`` is installed."""
     req = Requirement(package)
-    if importlib.util.find_spec(req.name) is None:
+    if importlib.util.find_spec(req.name) is None:  # type: ignore
         return False
 
     try:

kumoai/trainer/job.py

@@ -4,7 +4,7 @@ import concurrent.futures
 import time
 from datetime import datetime, timezone
 from functools import cached_property
-from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 from urllib.parse import urlparse, urlunparse
 
 import pandas as pd
@@ -600,8 +600,10 @@ class TrainingJob(JobInterface[TrainingJobID, TrainingJobRequest,
 
     def __init__(self, job_id: TrainingJobID) -> None:
         self.job_id = job_id
-        self._fut: concurrent.futures.Future = create_future(
-            _poll_training(job_id))
+
+    @cached_property
+    def _fut(self) -> concurrent.futures.Future:
+        return create_future(_poll_training(self.job_id))
 
     @override
     @property
@@ -1002,31 +1004,6 @@ class BatchPredictionJob(JobInterface[BatchPredictionJobID,
         """
         return self._api().get_config(self.job_id)
 
-    @classmethod
-    def search_by_tags(cls, tags: Mapping[str, str],
-                       limit: int = 10) -> list['BatchPredictionJob']:
-        r"""Returns a list of :class:`~kumoai.trainer.job.BatchPredictionJob`
-        instances from a set of job tags.
-
-        Args:
-            tags (Mapping[str, str]): Tags by which to search.
-            limit (int): Max number of jobs to list, default 10.
-
-        Example:
-            >>> tags = {'pquery_name': 'my_pquery_name'}
-            >>> jobs = BatchPredictionJob.search_by_tags(tags)
-            Search limited to 10 results based on the `limit` parameter.
-            Found 2 jobs.
-        """
-        print(f"Search limited to {limit} results based on the `limit` "
-              "parameter.")
-
-        jobs = cls._api().list(limit=limit, additional_tags=tags)
-
-        print(f"Found {len(jobs)} jobs.")
-
-        return [cls(j.job_id) for j in jobs]
-
 
 def _get_batch_prediction_job(job_id: str) -> BatchPredictionJobResource:
     api = global_state.client.batch_prediction_job_api
@@ -1097,8 +1074,10 @@ class BaselineJob(JobInterface[BaselineJobID, BaselineJobRequest,
 
     def __init__(self, job_id: BaselineJobID) -> None:
         self.job_id = job_id
-        self._fut: concurrent.futures.Future = create_future(
-            _poll_baseline(job_id))
+
+    @cached_property
+    def _fut(self) -> concurrent.futures.Future:
+        return create_future(_poll_baseline(self.job_id))
 
     @override
     @property

kumoai/trainer/trainer.py

@@ -20,7 +20,6 @@ from kumoapi.jobs import (
     TrainingJobResource,
 )
 from kumoapi.model_plan import ModelPlan
-from kumoapi.task import TaskType
 
 from kumoai import global_state
 from kumoai.artifact_export.config import OutputConfig
@@ -190,6 +189,7 @@ class Trainer:
         *,
         non_blocking: bool = False,
         custom_tags: Mapping[str, str] = {},
+        warm_start_job_id: Optional[TrainingJobID] = None,
     ) -> Union[TrainingJob, TrainingJobResult]:
         r"""Fits a model to the specified graph and training table, with the
         strategy defined by this :class:`Trainer`'s :obj:`model_plan`.
@@ -207,6 +207,11 @@ class Trainer:
             custom_tags: Additional, customer defined k-v tags to be associated
                 with the job to be launched. Job tags are useful for grouping
                 and searching jobs.
+            warm_start_job_id: Optional job ID of a completed training job to
+                warm start from. Initializes the new model with the best
+                weights from the specified job, using its model
+                architecture, column processing, and neighbor sampling
+                configurations.
 
         Returns:
             Union[TrainingJobResult, TrainingJob]:
@@ -241,6 +246,7 @@ class Trainer:
                 graph_snapshot_id=graph.snapshot(non_blocking=non_blocking),
                 train_table_job_id=job_id,
                 custom_train_table=custom_table,
+                warm_start_job_id=warm_start_job_id,
             ))
 
         out = TrainingJob(job_id=self._training_job_id)
@@ -353,6 +359,9 @@ class Trainer:
                 'deprecated. Please use output_config to specify these '
                 'parameters.')
         assert output_config is not None
+        # Be able to pass output_config as a dictionary
+        if isinstance(output_config, dict):
+            output_config = OutputConfig(**output_config)
         output_table_name = to_db_table_name(output_config.output_table_name)
         validate_output_arguments(
             output_config.output_types,
@@ -395,15 +404,15 @@ class Trainer:
                 pred_table_data_path = prediction_table.table_data_uri
 
         api = global_state.client.batch_prediction_job_api
-
-        from kumoai.pquery.predictive_query import PredictiveQuery
-        pquery = PredictiveQuery.load_from_training_job(training_job_id)
-        if pquery.get_task_type() == TaskType.BINARY_CLASSIFICATION:
-            if binary_classification_threshold is None:
-                logger.warning("No binary classification threshold provided. "
-                               "Using default threshold of 0.5.")
-                binary_classification_threshold = 0.5
-
+        # Remove to resolve https://github.com/kumo-ai/kumo/issues/24250
+        # from kumoai.pquery.predictive_query import PredictiveQuery
+        # pquery = PredictiveQuery.load_from_training_job(training_job_id)
+        # if pquery.get_task_type() == TaskType.BINARY_CLASSIFICATION:
+        #     if binary_classification_threshold is None:
+        #         logger.warning(
+        # "No binary classification threshold provided. "
+        # "Using default threshold of 0.5.")
+        #         binary_classification_threshold = 0.5
         job_id, response = api.maybe_create(
             BatchPredictionRequest(
                 dict(custom_tags),

kumoai/utils/__init__.py

@@ -1,9 +1,10 @@
-from .progress_logger import ProgressLogger
+from .progress_logger import ProgressLogger, InteractiveProgressLogger
 from .forecasting import ForecastVisualizer
 from .datasets import from_relbench
 
 __all__ = [
     'ProgressLogger',
+    'InteractiveProgressLogger',
     'ForecastVisualizer',
     'from_relbench',
 ]

kumoai/utils/progress_logger.py

@@ -1,9 +1,18 @@
+import sys
 import time
 from typing import Any, List, Optional, Union
 
 from rich.console import Console, ConsoleOptions, RenderResult
 from rich.live import Live
 from rich.padding import Padding
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    Task,
+    TextColumn,
+    TimeRemainingColumn,
+)
 from rich.spinner import Spinner
 from rich.table import Table
 from rich.text import Text
@@ -11,27 +20,13 @@ from typing_extensions import Self
 
 
 class ProgressLogger:
-    def __init__(
-        self,
-        msg: str,
-        verbose: bool = True,
-        refresh_per_second: int = 10,
-    ) -> None:
-
+    def __init__(self, msg: str) -> None:
         self.msg = msg
-        self.verbose = verbose
-        self.refresh_per_second = refresh_per_second
         self.logs: List[str] = []
 
         self.start_time: Optional[float] = None
         self.end_time: Optional[float] = None
 
-        self._live: Optional[Live] = None
-        self._exception: bool = False
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}({self.msg})'
-
     @property
     def duration(self) -> float:
         assert self.start_time is not None
@@ -44,6 +39,77 @@ class ProgressLogger:
 
     def __enter__(self) -> Self:
         self.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.end_time = time.perf_counter()
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.msg})'
+
+
+class ColoredMofNCompleteColumn(MofNCompleteColumn):
+    def __init__(self, style: str = 'green') -> None:
+        super().__init__()
+        self.style = style
+
+    def render(self, task: Task) -> Text:
+        return Text(str(super().render(task)), style=self.style)
+
+
+class ColoredTimeRemainingColumn(TimeRemainingColumn):
+    def __init__(self, style: str = 'cyan') -> None:
+        super().__init__()
+        self.style = style
+
+    def render(self, task: Task) -> Text:
+        return Text(str(super().render(task)), style=self.style)
+
+
+class InteractiveProgressLogger(ProgressLogger):
+    def __init__(
+        self,
+        msg: str,
+        verbose: bool = True,
+        refresh_per_second: int = 10,
+    ) -> None:
+        super().__init__(msg=msg)
+
+        self.verbose = verbose
+        self.refresh_per_second = refresh_per_second
+
+        self._progress: Optional[Progress] = None
+        self._task: Optional[int] = None
+
+        self._live: Optional[Live] = None
+        self._exception: bool = False
+
+    def init_progress(self, total: int, description: str) -> None:
+        assert self._progress is None
+        if self.verbose:
+            self._progress = Progress(
+                TextColumn(f'   ↳ {description}', style='dim'),
+                BarColumn(bar_width=None),
+                ColoredMofNCompleteColumn(style='dim'),
+                TextColumn('•', style='dim'),
+                ColoredTimeRemainingColumn(style='dim'),
+            )
+            self._task = self._progress.add_task("Progress", total=total)
+
+    def step(self) -> None:
+        if self.verbose:
+            assert self._progress is not None
+            assert self._task is not None
+            self._progress.update(self._task, advance=1)  # type: ignore
+
+    def __enter__(self) -> Self:
+        from kumoai import in_notebook
+
+        super().__enter__()
+
+        if not in_notebook():  # Render progress bar in TUI.
+            sys.stdout.write("\x1b]9;4;3\x07")
+            sys.stdout.flush()
 
         if self.verbose:
             self._live = Live(
@@ -56,16 +122,27 @@ class ProgressLogger:
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        self.end_time = time.perf_counter()
+        from kumoai import in_notebook
+
+        super().__exit__(exc_type, exc_val, exc_tb)
 
         if exc_type is not None:
             self._exception = True
 
+        if self._progress is not None:
+            self._progress.stop()
+            self._progress = None
+            self._task = None
+
         if self._live is not None:
             self._live.update(self, refresh=True)
             self._live.stop()
             self._live = None
 
+        if not in_notebook():
+            sys.stdout.write("\x1b]9;4;0\x07")
+            sys.stdout.flush()
+
     def __rich_console__(
         self,
         console: Console,
@@ -95,3 +172,6 @@ class ProgressLogger:
             table.add_row('', Text(f'↳ {log}', style='dim'))
 
         yield table
+
+        if self.verbose and self._progress is not None:
+            yield self._progress.get_renderable()

{kumoai-2.8.0.dev202508221830.dist-info → kumoai-2.13.0.dev202512041141.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kumoai
-Version: 2.8.0.dev202508221830
+Version: 2.13.0.dev202512041141
 Summary: AI on the Modern Data Stack
 Author-email: "Kumo.AI" <hello@kumo.ai>
 License-Expression: MIT
@@ -9,13 +9,12 @@ Project-URL: documentation, https://kumo.ai/docs
 Keywords: deep-learning,graph-neural-networks,cloud-data-warehouse
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Programming Language :: Python
-Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.9
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pandas
@@ -24,7 +23,7 @@ Requires-Dist: requests>=2.28.2
 Requires-Dist: urllib3
 Requires-Dist: plotly
 Requires-Dist: typing_extensions>=4.5.0
-Requires-Dist: kumo-api==0.31.0
+Requires-Dist: kumo-api==0.48.0
 Requires-Dist: tqdm>=4.66.0
 Requires-Dist: aiohttp>=3.10.0
 Requires-Dist: pydantic>=1.10.21
@@ -39,6 +38,16 @@ Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-mock; extra == "test"
 Requires-Dist: requests-mock; extra == "test"
+Provides-Extra: sqlite
+Requires-Dist: adbc_driver_sqlite; extra == "sqlite"
+Provides-Extra: snowflake
+Requires-Dist: snowflake-connector-python; extra == "snowflake"
+Requires-Dist: pyyaml; extra == "snowflake"
+Provides-Extra: sagemaker
+Requires-Dist: boto3<2.0,>=1.30.0; extra == "sagemaker"
+Requires-Dist: mypy-boto3-sagemaker-runtime<2.0,>=1.34.0; extra == "sagemaker"
+Provides-Extra: test-sagemaker
+Requires-Dist: sagemaker<3.0; extra == "test-sagemaker"
 Dynamic: license-file
 Dynamic: requires-dist
 
@@ -54,7 +63,7 @@ interact with the Kumo machine learning platform
 
 ## Installation
 
-The Kumo SDK is available for Python 3.9 to Python 3.13. To install, simply run
+The Kumo SDK is available for Python 3.10 to Python 3.13. To install, simply run
 
 ```
 pip install kumoai