kumoai 2.12.0.dev202510231830__cp311-cp311-win_amd64.whl → 2.14.0.dev202512311733__cp311-cp311-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (64) hide show
  1. kumoai/__init__.py +41 -35
  2. kumoai/_version.py +1 -1
  3. kumoai/client/client.py +15 -13
  4. kumoai/client/endpoints.py +1 -0
  5. kumoai/client/jobs.py +24 -0
  6. kumoai/client/pquery.py +6 -2
  7. kumoai/client/rfm.py +35 -7
  8. kumoai/connector/utils.py +23 -2
  9. kumoai/experimental/rfm/__init__.py +191 -48
  10. kumoai/experimental/rfm/authenticate.py +3 -4
  11. kumoai/experimental/rfm/backend/__init__.py +0 -0
  12. kumoai/experimental/rfm/backend/local/__init__.py +42 -0
  13. kumoai/experimental/rfm/{local_graph_store.py → backend/local/graph_store.py} +65 -127
  14. kumoai/experimental/rfm/backend/local/sampler.py +312 -0
  15. kumoai/experimental/rfm/backend/local/table.py +113 -0
  16. kumoai/experimental/rfm/backend/snow/__init__.py +37 -0
  17. kumoai/experimental/rfm/backend/snow/sampler.py +297 -0
  18. kumoai/experimental/rfm/backend/snow/table.py +242 -0
  19. kumoai/experimental/rfm/backend/sqlite/__init__.py +32 -0
  20. kumoai/experimental/rfm/backend/sqlite/sampler.py +398 -0
  21. kumoai/experimental/rfm/backend/sqlite/table.py +184 -0
  22. kumoai/experimental/rfm/base/__init__.py +30 -0
  23. kumoai/experimental/rfm/base/column.py +152 -0
  24. kumoai/experimental/rfm/base/expression.py +44 -0
  25. kumoai/experimental/rfm/base/sampler.py +761 -0
  26. kumoai/experimental/rfm/base/source.py +19 -0
  27. kumoai/experimental/rfm/base/sql_sampler.py +143 -0
  28. kumoai/experimental/rfm/base/table.py +735 -0
  29. kumoai/experimental/rfm/graph.py +1237 -0
  30. kumoai/experimental/rfm/infer/__init__.py +8 -0
  31. kumoai/experimental/rfm/infer/dtype.py +82 -0
  32. kumoai/experimental/rfm/infer/multicategorical.py +1 -1
  33. kumoai/experimental/rfm/infer/pkey.py +128 -0
  34. kumoai/experimental/rfm/infer/stype.py +35 -0
  35. kumoai/experimental/rfm/infer/time_col.py +61 -0
  36. kumoai/experimental/rfm/pquery/__init__.py +0 -4
  37. kumoai/experimental/rfm/pquery/executor.py +27 -27
  38. kumoai/experimental/rfm/pquery/pandas_executor.py +64 -40
  39. kumoai/experimental/rfm/relbench.py +76 -0
  40. kumoai/experimental/rfm/rfm.py +386 -276
  41. kumoai/experimental/rfm/sagemaker.py +138 -0
  42. kumoai/kumolib.cp311-win_amd64.pyd +0 -0
  43. kumoai/pquery/predictive_query.py +10 -6
  44. kumoai/spcs.py +1 -3
  45. kumoai/testing/decorators.py +1 -1
  46. kumoai/testing/snow.py +50 -0
  47. kumoai/trainer/distilled_trainer.py +175 -0
  48. kumoai/trainer/trainer.py +9 -10
  49. kumoai/utils/__init__.py +3 -2
  50. kumoai/utils/display.py +51 -0
  51. kumoai/utils/progress_logger.py +188 -16
  52. kumoai/utils/sql.py +3 -0
  53. {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/METADATA +13 -2
  54. {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/RECORD +57 -36
  55. kumoai/experimental/rfm/local_graph.py +0 -810
  56. kumoai/experimental/rfm/local_graph_sampler.py +0 -184
  57. kumoai/experimental/rfm/local_pquery_driver.py +0 -494
  58. kumoai/experimental/rfm/local_table.py +0 -545
  59. kumoai/experimental/rfm/pquery/backend.py +0 -136
  60. kumoai/experimental/rfm/pquery/pandas_backend.py +0 -478
  61. kumoai/experimental/rfm/utils.py +0 -344
  62. {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/WHEEL +0 -0
  63. {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/licenses/LICENSE +0 -0
  64. {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/top_level.txt +0 -0
kumoai/__init__.py CHANGED
@@ -1,3 +1,4 @@
1
+ import warnings
1
2
  import os
2
3
  import sys
3
4
  import threading
@@ -68,9 +69,8 @@ class GlobalState(metaclass=Singleton):
68
69
  if self._url is None or (self._api_key is None
69
70
  and self._spcs_token is None
70
71
  and self._snowpark_session is None):
71
- raise ValueError(
72
- "Client creation or authentication failed; please re-create "
73
- "your client before proceeding.")
72
+ raise ValueError("Client creation or authentication failed. "
73
+ "Please re-create your client before proceeding.")
74
74
 
75
75
  if hasattr(self.thread_local, '_client'):
76
76
  # Set the spcs token in the client to ensure it has the latest.
@@ -123,10 +123,9 @@ def init(
123
123
  """ # noqa
124
124
  # Avoid mutations to the global state after it is set:
125
125
  if global_state.initialized:
126
- print(
127
- "Client has already been created. To re-initialize Kumo, please "
128
- "start a new interpreter. No changes will be made to the current "
129
- "session.")
126
+ warnings.warn("Kumo SDK already initialized. To re-initialize the "
127
+ "SDK, please start a new interpreter. No changes will "
128
+ "be made to the current session.")
130
129
  return
131
130
 
132
131
  set_log_level(os.getenv(_ENV_KUMO_LOG, log_level))
@@ -138,15 +137,15 @@ def init(
138
137
  if snowflake_application:
139
138
  if url is not None:
140
139
  raise ValueError(
141
- "Client creation failed: both snowflake_application and url "
142
- "are specified. If running from a snowflake notebook, specify"
143
- "only snowflake_application.")
140
+ "Kumo SDK initialization failed. Both 'snowflake_application' "
141
+ "and 'url' are specified. If running from a Snowflake "
142
+ "notebook, specify only 'snowflake_application'.")
144
143
  snowpark_session = _get_active_session()
145
144
  if not snowpark_session:
146
145
  raise ValueError(
147
- "Client creation failed: snowflake_application is specified "
148
- "without an active snowpark session. If running outside "
149
- "a snowflake notebook, specify a URL and credentials.")
146
+ "Kumo SDK initialization failed. 'snowflake_application' is "
147
+ "specified without an active Snowpark session. If running "
148
+ "outside a Snowflake notebook, specify a URL and credentials.")
150
149
  description = snowpark_session.sql(
151
150
  f"DESCRIBE SERVICE {snowflake_application}."
152
151
  "USER_SCHEMA.KUMO_SERVICE").collect()[0]
@@ -155,14 +154,14 @@ def init(
155
154
  if api_key is None and not snowflake_application:
156
155
  if snowflake_credentials is None:
157
156
  raise ValueError(
158
- "Client creation failed: Neither API key nor snowflake "
159
- "credentials provided. Please either set the 'KUMO_API_KEY' "
160
- "or explicitly call `kumoai.init(...)`.")
157
+ "Kumo SDK initialization failed. Neither an API key nor "
158
+ "Snowflake credentials provided. Please either set the "
159
+ "'KUMO_API_KEY' or explicitly call `kumoai.init(...)`.")
161
160
  if (set(snowflake_credentials.keys())
162
161
  != {'user', 'password', 'account'}):
163
162
  raise ValueError(
164
- f"Provided credentials should be a dictionary with keys "
165
- f"'user', 'password', and 'account'. Only "
163
+ f"Provided Snowflake credentials should be a dictionary with "
164
+ f"keys 'user', 'password', and 'account'. Only "
166
165
  f"{set(snowflake_credentials.keys())} were provided.")
167
166
 
168
167
  # Get or infer URL:
@@ -173,10 +172,10 @@ def init(
173
172
  except KeyError:
174
173
  pass
175
174
  if url is None:
176
- raise ValueError(
177
- "Client creation failed: endpoint URL not provided. Please "
178
- "either set the 'KUMO_API_ENDPOINT' environment variable or "
179
- "explicitly call `kumoai.init(...)`.")
175
+ raise ValueError("Kumo SDK initialization failed since no endpoint "
176
+ "URL was provided. Please either set the "
177
+ "'KUMO_API_ENDPOINT' environment variable or "
178
+ "explicitly call `kumoai.init(...)`.")
180
179
 
181
180
  # Assign global state after verification that client can be created and
182
181
  # authenticated successfully:
@@ -184,15 +183,12 @@ def init(
184
183
  snowflake_credentials
185
184
  ) if not api_key and snowflake_credentials else None
186
185
  client = KumoClient(url=url, api_key=api_key, spcs_token=spcs_token)
187
- if client.authenticate():
188
- global_state._url = client._url
189
- global_state._api_key = client._api_key
190
- global_state._snowflake_credentials = snowflake_credentials
191
- global_state._spcs_token = client._spcs_token
192
- global_state._snowpark_session = snowpark_session
193
- else:
194
- raise ValueError("Client authentication failed. Please check if you "
195
- "have a valid API key.")
186
+ client.authenticate()
187
+ global_state._url = client._url
188
+ global_state._api_key = client._api_key
189
+ global_state._snowflake_credentials = snowflake_credentials
190
+ global_state._spcs_token = client._spcs_token
191
+ global_state._snowpark_session = snowpark_session
196
192
 
197
193
  if not api_key and snowflake_credentials:
198
194
  # Refresh token every 10 minutes (expires in 1 hour):
@@ -201,10 +197,8 @@ def init(
201
197
  logger = logging.getLogger('kumoai')
202
198
  log_level = logging.getLevelName(logger.getEffectiveLevel())
203
199
 
204
- logger.info(
205
- f"Successfully initialized the Kumo SDK (version {__version__}) "
206
- f"against deployment {url}, with "
207
- f"log level {log_level}.")
200
+ logger.info(f"Initialized Kumo SDK v{__version__} against deployment "
201
+ f"'{url}'")
208
202
 
209
203
 
210
204
  def set_log_level(level: str) -> None:
@@ -283,7 +277,19 @@ __all__ = [
283
277
  ]
284
278
 
285
279
 
280
+ def in_snowflake_notebook() -> bool:
281
+ try:
282
+ from snowflake.snowpark.context import get_active_session
283
+ import streamlit # noqa: F401
284
+ get_active_session()
285
+ return True
286
+ except Exception:
287
+ return False
288
+
289
+
286
290
  def in_notebook() -> bool:
291
+ if in_snowflake_notebook():
292
+ return True
287
293
  try:
288
294
  from IPython import get_ipython
289
295
  shell = get_ipython()
kumoai/_version.py CHANGED
@@ -1 +1 @@
1
- __version__ = '2.12.0.dev202510231830'
1
+ __version__ = '2.14.0.dev202512311733'
kumoai/client/client.py CHANGED
@@ -13,6 +13,7 @@ if TYPE_CHECKING:
13
13
  ArtifactExportJobAPI,
14
14
  BaselineJobAPI,
15
15
  BatchPredictionJobAPI,
16
+ DistillationJobAPI,
16
17
  GeneratePredictionTableJobAPI,
17
18
  GenerateTrainTableJobAPI,
18
19
  LLMJobAPI,
@@ -20,7 +21,6 @@ if TYPE_CHECKING:
20
21
  )
21
22
  from kumoai.client.online import OnlineServingEndpointAPI
22
23
  from kumoai.client.pquery import PQueryAPI
23
- from kumoai.client.rfm import RFMAPI
24
24
  from kumoai.client.source_table import SourceTableAPI
25
25
  from kumoai.client.table import TableAPI
26
26
 
@@ -73,12 +73,15 @@ class KumoClient:
73
73
  self._session.headers.update(
74
74
  {'Authorization': f'Snowflake Token={self._spcs_token}'})
75
75
 
76
- def authenticate(self) -> bool:
77
- r"""Raises an exception if authentication fails. Succeeds if the
78
- client is properly formed.
79
- """
80
- return self._session.get(f"{self._url}/v1/connectors",
81
- verify=self._verify_ssl).ok
76
+ def authenticate(self) -> None:
77
+ """Raises an exception if authentication fails."""
78
+ try:
79
+ self._session.get(self._url + '/v1/connectors',
80
+ verify=self._verify_ssl).raise_for_status()
81
+ except Exception:
82
+ raise ValueError(
83
+ "Client authentication failed. Please check if you "
84
+ "have a valid API key/credentials.")
82
85
 
83
86
  def set_spcs_token(self, spcs_token: str) -> None:
84
87
  r"""Sets the SPCS token for the client and updates the session
@@ -130,6 +133,11 @@ class KumoClient:
130
133
  from kumoai.client.jobs import TrainingJobAPI
131
134
  return TrainingJobAPI(self)
132
135
 
136
+ @property
137
+ def distillation_job_api(self) -> 'DistillationJobAPI':
138
+ from kumoai.client.jobs import DistillationJobAPI
139
+ return DistillationJobAPI(self)
140
+
133
141
  @property
134
142
  def batch_prediction_job_api(self) -> 'BatchPredictionJobAPI':
135
143
  from kumoai.client.jobs import BatchPredictionJobAPI
@@ -163,12 +171,6 @@ class KumoClient:
163
171
  from kumoai.client.online import OnlineServingEndpointAPI
164
172
  return OnlineServingEndpointAPI(self)
165
173
 
166
- @property
167
- def rfm_api(self) -> 'RFMAPI':
168
- r"""Returns the typed RFM API."""
169
- from kumoai.client.rfm import RFMAPI
170
- return RFMAPI(self)
171
-
172
174
  def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
173
175
  r"""Send a HTTP request to the specified endpoint."""
174
176
  endpoint_str = endpoint.get_path()
@@ -147,3 +147,4 @@ class RFMEndpoints:
147
147
  explain = Endpoint(f"{BASE}/explain", HTTPMethod.POST)
148
148
  evaluate = Endpoint(f"{BASE}/evaluate", HTTPMethod.POST)
149
149
  validate_query = Endpoint(f"{BASE}/validate_query", HTTPMethod.POST)
150
+ parse_query = Endpoint(f"{BASE}/parse_query", HTTPMethod.POST)
kumoai/client/jobs.py CHANGED
@@ -22,6 +22,8 @@ from kumoapi.jobs import (
22
22
  BatchPredictionRequest,
23
23
  CancelBatchPredictionJobResponse,
24
24
  CancelTrainingJobResponse,
25
+ DistillationJobRequest,
26
+ DistillationJobResource,
25
27
  ErrorDetails,
26
28
  GeneratePredictionTableJobResource,
27
29
  GeneratePredictionTableRequest,
@@ -171,6 +173,28 @@ class TrainingJobAPI(CommonJobAPI[TrainingJobRequest, TrainingJobResource]):
171
173
  return resource.config
172
174
 
173
175
 
176
+ class DistillationJobAPI(CommonJobAPI[DistillationJobRequest,
177
+ DistillationJobResource]):
178
+ r"""Typed API definition for the distillation job resource."""
179
+ def __init__(self, client: KumoClient) -> None:
180
+ super().__init__(client, '/training_jobs/distilled_training_job',
181
+ DistillationJobResource)
182
+
183
+ def get_config(self, job_id: str) -> DistillationJobRequest:
184
+ raise NotImplementedError(
185
+ "Getting the configuration for a distillation job is "
186
+ "not implemented yet.")
187
+
188
+ def get_progress(self, id: str) -> AutoTrainerProgress:
189
+ raise NotImplementedError(
190
+ "Getting the progress for a distillation job is not "
191
+ "implemented yet.")
192
+
193
+ def cancel(self, id: str) -> CancelTrainingJobResponse:
194
+ raise NotImplementedError(
195
+ "Cancelling a distillation job is not implemented yet.")
196
+
197
+
174
198
  class BatchPredictionJobAPI(CommonJobAPI[BatchPredictionRequest,
175
199
  BatchPredictionJobResource]):
176
200
  r"""Typed API definition for the prediction job resource."""
kumoai/client/pquery.py CHANGED
@@ -176,8 +176,12 @@ def filter_model_plan(
176
176
  # Undefined
177
177
  pass
178
178
 
179
- new_opt_fields.append((field.name, _type, default))
180
- new_opts.append(getattr(section, field.name))
179
+ # Forward compatibility - Remove any newly introduced arguments not
180
+ # returned yet by the backend:
181
+ value = getattr(section, field.name)
182
+ if value != MissingType.VALUE:
183
+ new_opt_fields.append((field.name, _type, default))
184
+ new_opts.append(value)
181
185
 
182
186
  Section = dataclass(
183
187
  config=dict(validate_assignment=True),
kumoai/client/rfm.py CHANGED
@@ -1,7 +1,11 @@
1
+ from typing import Any
2
+
1
3
  from kumoapi.json_serde import to_json_dict
2
4
  from kumoapi.rfm import (
3
5
  RFMEvaluateResponse,
4
6
  RFMExplanationResponse,
7
+ RFMParseQueryRequest,
8
+ RFMParseQueryResponse,
5
9
  RFMPredictResponse,
6
10
  RFMValidateQueryRequest,
7
11
  RFMValidateQueryResponse,
@@ -26,25 +30,32 @@ class RFMAPI:
26
30
  Returns:
27
31
  RFMPredictResponse containing the predictions
28
32
  """
29
- # Send binary data to the predict endpoint
30
33
  response = self._client._request(
31
- RFMEndpoints.predict, data=request,
32
- headers={'Content-Type': 'application/x-protobuf'})
34
+ RFMEndpoints.predict,
35
+ data=request,
36
+ headers={'Content-Type': 'application/x-protobuf'},
37
+ )
33
38
  raise_on_error(response)
34
39
  return parse_response(RFMPredictResponse, response)
35
40
 
36
- def explain(self, request: bytes) -> RFMExplanationResponse:
41
+ def explain(
42
+ self,
43
+ request: bytes,
44
+ skip_summary: bool = False,
45
+ ) -> RFMExplanationResponse:
37
46
  """Explain the RFM model on the given context.
38
47
 
39
48
  Args:
40
49
  request: The predict request as serialized protobuf.
50
+ skip_summary: Whether to skip generating a human-readable summary
51
+ of the explanation.
41
52
 
42
53
  Returns:
43
54
  RFMPredictResponse containing the explanations
44
55
  """
45
- # Send binary data to the explain endpoint
56
+ params: dict[str, Any] = {'generate_summary': not skip_summary}
46
57
  response = self._client._request(
47
- RFMEndpoints.explain, data=request,
58
+ RFMEndpoints.explain, data=request, params=params,
48
59
  headers={'Content-Type': 'application/x-protobuf'})
49
60
  raise_on_error(response)
50
61
  return parse_response(RFMExplanationResponse, response)
@@ -58,7 +69,6 @@ class RFMAPI:
58
69
  Returns:
59
70
  RFMEvaluateResponse containing the computed metrics
60
71
  """
61
- # Send binary data to the evaluate endpoint
62
72
  response = self._client._request(
63
73
  RFMEndpoints.evaluate, data=request,
64
74
  headers={'Content-Type': 'application/x-protobuf'})
@@ -82,3 +92,21 @@ class RFMAPI:
82
92
  json=to_json_dict(request))
83
93
  raise_on_error(response)
84
94
  return parse_response(RFMValidateQueryResponse, response)
95
+
96
+ def parse_query(
97
+ self,
98
+ request: RFMParseQueryRequest,
99
+ ) -> RFMParseQueryResponse:
100
+ """Validate a predictive query against a graph.
101
+
102
+ Args:
103
+ request: The request object containing
104
+ the query and graph definition
105
+
106
+ Returns:
107
+ RFMParseQueryResponse containing the QueryDefinition
108
+ """
109
+ response = self._client._request(RFMEndpoints.parse_query,
110
+ json=to_json_dict(request))
111
+ raise_on_error(response)
112
+ return parse_response(RFMParseQueryResponse, response)
kumoai/connector/utils.py CHANGED
@@ -381,8 +381,29 @@ def _handle_duplicate_names(names: List[str]) -> List[str]:
381
381
 
382
382
 
383
383
  def _sanitize_columns(names: List[str]) -> Tuple[List[str], bool]:
384
- _SAN_RE = re.compile(r"[^0-9A-Za-z]+")
384
+ """Normalize column names in a CSV or Parquet file.
385
+
386
+ Rules:
387
+ - Replace any non-alphanumeric character with "_"
388
+ - Strip leading/trailing underscores
389
+ - Ensure uniqueness by appending suffixes: _1, _2, ...
390
+ - Auto-name empty columns as auto_named_<n>
391
+
392
+ Returns:
393
+ (new_column_names, changed)
394
+ """
395
+ _SAN_RE = re.compile(r"[^0-9A-Za-z,\t]")
396
+ # 1) Replace non-alphanumeric sequences with underscore
385
397
  new = [_SAN_RE.sub("_", n).strip("_") for n in names]
398
+
399
+ # 2) Auto-name any empty column names to match UI behavior
400
+ unnamed_counter = 0
401
+ for i, n in enumerate(new):
402
+ if not n:
403
+ new[i] = f"auto_named_{unnamed_counter}"
404
+ unnamed_counter += 1
405
+
406
+ # 3) Ensure uniqueness (append suffixes where needed)
386
407
  new = _handle_duplicate_names(new)
387
408
  return new, new != names
388
409
 
@@ -1168,7 +1189,7 @@ def _detect_and_validate_csv(head_bytes: bytes) -> str:
1168
1189
  - Re-serializes those rows and validates with pandas (small nrows) to catch
1169
1190
  malformed inputs.
1170
1191
  - Raises ValueError on empty input or if parsing fails with the chosen
1171
- delimiter.
1192
+ delimiter.
1172
1193
  """
1173
1194
  if not head_bytes:
1174
1195
  raise ValueError("Could not auto-detect a delimiter: file is empty.")
@@ -1,65 +1,208 @@
1
- try:
2
- import kumoai.kumolib # noqa: F401
3
- except Exception as e:
4
- import platform
5
-
6
- _msg = f"""RFM is not supported in your environment.
7
-
8
- 💻 Your Environment:
9
- Python version: {platform.python_version()}
10
- Operating system: {platform.system()}
11
- CPU architecture: {platform.machine()}
12
- glibc version: {platform.libc_ver()[1]}
13
-
14
- ✅ Supported Environments:
15
- * Python versions: 3.10, 3.11, 3.12, 3.13
16
- * Operating systems and CPU architectures:
17
- * Linux (x86_64)
18
- * macOS (arm64)
19
- * Windows (x86_64)
20
- * glibc versions: >=2.28
21
-
22
- ❌ Unsupported Environments:
23
- * Python versions: 3.8, 3.9, 3.14
24
- * Operating systems and CPU architectures:
25
- * Linux (arm64)
26
- * macOS (x86_64)
27
- * Windows (arm64)
28
- * glibc versions: <2.28
29
-
30
- Please create a feature request at 'https://github.com/kumo-ai/kumo-rfm'."""
31
-
32
- raise RuntimeError(_msg) from e
33
-
34
- from typing import Optional, Dict
1
+ import ipaddress
2
+ import logging
35
3
  import os
4
+ import re
5
+ import socket
6
+ import threading
7
+ from dataclasses import dataclass
8
+ from enum import Enum
9
+ from urllib.parse import urlparse
10
+
36
11
  import kumoai
37
- from .local_table import LocalTable
38
- from .local_graph import LocalGraph
39
- from .rfm import KumoRFM
12
+ from kumoai.client.client import KumoClient
13
+ from kumoai.spcs import _get_active_session
14
+
40
15
  from .authenticate import authenticate
16
+ from .sagemaker import (
17
+ KumoClient_SageMakerAdapter,
18
+ KumoClient_SageMakerProxy_Local,
19
+ )
20
+ from .base import Table
21
+ from .backend.local import LocalTable
22
+ from .graph import Graph
23
+ from .rfm import ExplainConfig, Explanation, KumoRFM
24
+
25
+ logger = logging.getLogger('kumoai_rfm')
26
+
27
+
28
+ def _is_local_address(host: str | None) -> bool:
29
+ """Return True if the hostname/IP refers to the local machine."""
30
+ if not host:
31
+ return False
32
+ try:
33
+ infos = socket.getaddrinfo(host, None)
34
+ for _, _, _, _, sockaddr in infos:
35
+ ip = sockaddr[0]
36
+ ip_obj = ipaddress.ip_address(ip)
37
+ if ip_obj.is_loopback or ip_obj.is_unspecified:
38
+ return True
39
+ return False
40
+ except Exception:
41
+ return False
42
+
43
+
44
+ class InferenceBackend(str, Enum):
45
+ REST = "REST"
46
+ LOCAL_SAGEMAKER = "LOCAL_SAGEMAKER"
47
+ AWS_SAGEMAKER = "AWS_SAGEMAKER"
48
+ UNKNOWN = "UNKNOWN"
49
+
50
+
51
+ def _detect_backend(
52
+ url: str, #
53
+ ) -> tuple[InferenceBackend, str | None, str | None]:
54
+ parsed = urlparse(url)
55
+
56
+ # Remote SageMaker
57
+ if ("runtime.sagemaker" in parsed.netloc
58
+ and parsed.path.endswith("/invocations")):
59
+ # Example: https://runtime.sagemaker.us-west-2.amazonaws.com/
60
+ # endpoints/Name/invocations
61
+ match = re.search(r"runtime\.sagemaker\.([a-z0-9-]+)\.amazonaws\.com",
62
+ parsed.netloc)
63
+ region = match.group(1) if match else None
64
+ m = re.search(r"/endpoints/([^/]+)/invocations", parsed.path)
65
+ endpoint_name = m.group(1) if m else None
66
+ return InferenceBackend.AWS_SAGEMAKER, region, endpoint_name
67
+
68
+ # Local SageMaker
69
+ if parsed.port == 8080 and parsed.path.endswith(
70
+ "/invocations") and _is_local_address(parsed.hostname):
71
+ return InferenceBackend.LOCAL_SAGEMAKER, None, None
72
+
73
+ # Default: regular REST
74
+ return InferenceBackend.REST, None, None
75
+
76
+
77
+ def _get_snowflake_url(snowflake_application: str) -> str:
78
+ snowpark_session = _get_active_session()
79
+ if not snowpark_session:
80
+ raise ValueError(
81
+ "KumoRFM initialization failed. 'snowflake_application' is "
82
+ "specified without an active Snowpark session. If running outside "
83
+ "a Snowflake notebook, specify a URL and credentials.")
84
+ with snowpark_session.connection.cursor() as cur:
85
+ cur.execute(
86
+ f"DESCRIBE SERVICE {snowflake_application}.user_schema.rfm_service"
87
+ f" ->> SELECT \"dns_name\" from $1")
88
+ dns_name: str = cur.fetchone()[0]
89
+ return f"http://{dns_name}:8000/api"
90
+
91
+
92
+ @dataclass
93
+ class RfmGlobalState:
94
+ _url: str = '__url_not_provided__'
95
+ _backend: InferenceBackend = InferenceBackend.UNKNOWN
96
+ _region: str | None = None
97
+ _endpoint_name: str | None = None
98
+ _thread_local = threading.local()
99
+
100
+ # Thread-safe init-once.
101
+ _initialized: bool = False
102
+ _lock: threading.Lock = threading.Lock()
103
+
104
+ @property
105
+ def client(self) -> KumoClient:
106
+ if self._backend == InferenceBackend.UNKNOWN:
107
+ raise RuntimeError("KumoRFM is not yet initialized")
108
+
109
+ if self._backend == InferenceBackend.REST:
110
+ return kumoai.global_state.client
111
+
112
+ if hasattr(self._thread_local, '_sagemaker'):
113
+ # Set the spcs token in the client to ensure it has the latest.
114
+ return self._thread_local._sagemaker
115
+
116
+ sagemaker_client: KumoClient
117
+ if self._backend == InferenceBackend.LOCAL_SAGEMAKER:
118
+ sagemaker_client = KumoClient_SageMakerProxy_Local(self._url)
119
+ else:
120
+ assert self._backend == InferenceBackend.AWS_SAGEMAKER
121
+ assert self._region
122
+ assert self._endpoint_name
123
+ sagemaker_client = KumoClient_SageMakerAdapter(
124
+ self._region, self._endpoint_name)
125
+
126
+ self._thread_local._sagemaker = sagemaker_client
127
+ return sagemaker_client
128
+
129
+ def reset(self) -> None: # For testing only.
130
+ with self._lock:
131
+ self._initialized = False
132
+ self._url = '__url_not_provided__'
133
+ self._backend = InferenceBackend.UNKNOWN
134
+ self._region = None
135
+ self._endpoint_name = None
136
+ self._thread_local = threading.local()
137
+
138
+
139
+ global_state = RfmGlobalState()
41
140
 
42
141
 
43
142
  def init(
44
- url: Optional[str] = None,
45
- api_key: Optional[str] = None,
46
- snowflake_credentials: Optional[Dict[str, str]] = None,
47
- snowflake_application: Optional[str] = None,
143
+ url: str | None = None,
144
+ api_key: str | None = None,
145
+ snowflake_credentials: dict[str, str] | None = None,
146
+ snowflake_application: str | None = None,
48
147
  log_level: str = "INFO",
49
148
  ) -> None:
50
- if url is None:
51
- url = os.getenv("KUMO_API_URL", "https://kumorfm.ai/api")
149
+ with global_state._lock:
150
+ if global_state._initialized:
151
+ if url != global_state._url:
152
+ raise RuntimeError(
153
+ "KumoRFM has already been initialized with a different "
154
+ "API URL. Re-initialization with a different URL is not "
155
+ "supported.")
156
+ return
52
157
 
53
- kumoai.init(url=url, api_key=api_key,
158
+ if snowflake_application:
159
+ if url is not None:
160
+ raise ValueError(
161
+ "KumoRFM initialization failed. Both "
162
+ "'snowflake_application' and 'url' are specified. If "
163
+ "running from a Snowflake notebook, specify only "
164
+ "'snowflake_application'.")
165
+ url = _get_snowflake_url(snowflake_application)
166
+ api_key = "test:DISABLED"
167
+
168
+ if url is None:
169
+ url = os.getenv("RFM_API_URL", "https://kumorfm.ai/api")
170
+
171
+ backend, region, endpoint_name = _detect_backend(url)
172
+ if backend == InferenceBackend.REST:
173
+ kumoai.init(
174
+ url=url,
175
+ api_key=api_key,
54
176
  snowflake_credentials=snowflake_credentials,
55
177
  snowflake_application=snowflake_application,
56
- log_level=log_level)
178
+ log_level=log_level,
179
+ )
180
+ elif backend == InferenceBackend.AWS_SAGEMAKER:
181
+ assert region
182
+ assert endpoint_name
183
+ KumoClient_SageMakerAdapter(region, endpoint_name).authenticate()
184
+ logger.info("KumoRFM initialized in AWS SageMaker")
185
+ else:
186
+ assert backend == InferenceBackend.LOCAL_SAGEMAKER
187
+ KumoClient_SageMakerProxy_Local(url).authenticate()
188
+ logger.info(f"KumoRFM initialized in local SageMaker at '{url}'")
189
+
190
+ global_state._url = url
191
+ global_state._backend = backend
192
+ global_state._region = region
193
+ global_state._endpoint_name = endpoint_name
194
+ global_state._initialized = True
195
+
57
196
 
197
+ LocalGraph = Graph # NOTE Backward compatibility - do not use anymore.
58
198
 
59
199
  __all__ = [
60
- 'LocalTable',
61
- 'LocalGraph',
62
- 'KumoRFM',
63
200
  'authenticate',
64
201
  'init',
202
+ 'Table',
203
+ 'LocalTable',
204
+ 'Graph',
205
+ 'KumoRFM',
206
+ 'ExplainConfig',
207
+ 'Explanation',
65
208
  ]
@@ -2,12 +2,11 @@ import logging
2
2
  import os
3
3
  import platform
4
4
  from datetime import datetime
5
- from typing import Optional
6
5
 
7
6
  from kumoai import in_notebook
8
7
 
9
8
 
10
- def authenticate(api_url: Optional[str] = None) -> None:
9
+ def authenticate(api_url: str | None = None) -> None:
11
10
  """Authenticates the user and sets the Kumo API key for the SDK.
12
11
 
13
12
  This function detects the current environment and launches the appropriate
@@ -65,11 +64,11 @@ def _authenticate_local(api_url: str, redirect_port: int = 8765) -> None:
65
64
  import webbrowser
66
65
  from getpass import getpass
67
66
  from socketserver import TCPServer
68
- from typing import Any, Dict
67
+ from typing import Any
69
68
 
70
69
  logger = logging.getLogger('kumoai')
71
70
 
72
- token_status: Dict[str, Any] = {
71
+ token_status: dict[str, Any] = {
73
72
  'token': None,
74
73
  'token_name': None,
75
74
  'failed': False
File without changes