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.
- kumoai/__init__.py +41 -35
- kumoai/_version.py +1 -1
- kumoai/client/client.py +15 -13
- kumoai/client/endpoints.py +1 -0
- kumoai/client/jobs.py +24 -0
- kumoai/client/pquery.py +6 -2
- kumoai/client/rfm.py +35 -7
- kumoai/connector/utils.py +23 -2
- kumoai/experimental/rfm/__init__.py +191 -48
- kumoai/experimental/rfm/authenticate.py +3 -4
- kumoai/experimental/rfm/backend/__init__.py +0 -0
- kumoai/experimental/rfm/backend/local/__init__.py +42 -0
- kumoai/experimental/rfm/{local_graph_store.py → backend/local/graph_store.py} +65 -127
- kumoai/experimental/rfm/backend/local/sampler.py +312 -0
- kumoai/experimental/rfm/backend/local/table.py +113 -0
- kumoai/experimental/rfm/backend/snow/__init__.py +37 -0
- kumoai/experimental/rfm/backend/snow/sampler.py +297 -0
- kumoai/experimental/rfm/backend/snow/table.py +242 -0
- kumoai/experimental/rfm/backend/sqlite/__init__.py +32 -0
- kumoai/experimental/rfm/backend/sqlite/sampler.py +398 -0
- kumoai/experimental/rfm/backend/sqlite/table.py +184 -0
- kumoai/experimental/rfm/base/__init__.py +30 -0
- kumoai/experimental/rfm/base/column.py +152 -0
- kumoai/experimental/rfm/base/expression.py +44 -0
- kumoai/experimental/rfm/base/sampler.py +761 -0
- kumoai/experimental/rfm/base/source.py +19 -0
- kumoai/experimental/rfm/base/sql_sampler.py +143 -0
- kumoai/experimental/rfm/base/table.py +735 -0
- kumoai/experimental/rfm/graph.py +1237 -0
- kumoai/experimental/rfm/infer/__init__.py +8 -0
- kumoai/experimental/rfm/infer/dtype.py +82 -0
- kumoai/experimental/rfm/infer/multicategorical.py +1 -1
- kumoai/experimental/rfm/infer/pkey.py +128 -0
- kumoai/experimental/rfm/infer/stype.py +35 -0
- kumoai/experimental/rfm/infer/time_col.py +61 -0
- kumoai/experimental/rfm/pquery/__init__.py +0 -4
- kumoai/experimental/rfm/pquery/executor.py +27 -27
- kumoai/experimental/rfm/pquery/pandas_executor.py +64 -40
- kumoai/experimental/rfm/relbench.py +76 -0
- kumoai/experimental/rfm/rfm.py +386 -276
- kumoai/experimental/rfm/sagemaker.py +138 -0
- kumoai/kumolib.cp311-win_amd64.pyd +0 -0
- kumoai/pquery/predictive_query.py +10 -6
- kumoai/spcs.py +1 -3
- kumoai/testing/decorators.py +1 -1
- kumoai/testing/snow.py +50 -0
- kumoai/trainer/distilled_trainer.py +175 -0
- kumoai/trainer/trainer.py +9 -10
- kumoai/utils/__init__.py +3 -2
- kumoai/utils/display.py +51 -0
- kumoai/utils/progress_logger.py +188 -16
- kumoai/utils/sql.py +3 -0
- {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/METADATA +13 -2
- {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/RECORD +57 -36
- kumoai/experimental/rfm/local_graph.py +0 -810
- kumoai/experimental/rfm/local_graph_sampler.py +0 -184
- kumoai/experimental/rfm/local_pquery_driver.py +0 -494
- kumoai/experimental/rfm/local_table.py +0 -545
- kumoai/experimental/rfm/pquery/backend.py +0 -136
- kumoai/experimental/rfm/pquery/pandas_backend.py +0 -478
- kumoai/experimental/rfm/utils.py +0 -344
- {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/WHEEL +0 -0
- {kumoai-2.12.0.dev202510231830.dist-info → kumoai-2.14.0.dev202512311733.dist-info}/licenses/LICENSE +0 -0
- {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
|
-
|
|
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
|
-
|
|
127
|
-
|
|
128
|
-
|
|
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
|
-
"
|
|
142
|
-
"are specified. If running from a
|
|
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
|
-
"
|
|
148
|
-
"without an active
|
|
149
|
-
"a
|
|
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
|
-
"
|
|
159
|
-
"credentials provided. Please either set the
|
|
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
|
|
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
|
-
|
|
178
|
-
|
|
179
|
-
|
|
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
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
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
|
-
|
|
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.
|
|
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) ->
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
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()
|
kumoai/client/endpoints.py
CHANGED
|
@@ -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
|
-
|
|
180
|
-
|
|
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,
|
|
32
|
-
|
|
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(
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
-
|
|
2
|
-
|
|
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 .
|
|
38
|
-
from .
|
|
39
|
-
|
|
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:
|
|
45
|
-
api_key:
|
|
46
|
-
snowflake_credentials:
|
|
47
|
-
snowflake_application:
|
|
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
|
-
|
|
51
|
-
|
|
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
|
-
|
|
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:
|
|
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
|
|
67
|
+
from typing import Any
|
|
69
68
|
|
|
70
69
|
logger = logging.getLogger('kumoai')
|
|
71
70
|
|
|
72
|
-
token_status:
|
|
71
|
+
token_status: dict[str, Any] = {
|
|
73
72
|
'token': None,
|
|
74
73
|
'token_name': None,
|
|
75
74
|
'failed': False
|
|
File without changes
|