kumoai 2.10.0.dev202509231831__cp313-cp313-macosx_11_0_arm64.whl → 2.14.0.dev202512161731__cp313-cp313-macosx_11_0_arm64.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.
Potentially problematic release.
This version of kumoai might be problematic. Click here for more details.
- kumoai/__init__.py +22 -11
- kumoai/_version.py +1 -1
- kumoai/client/client.py +17 -16
- kumoai/client/endpoints.py +1 -0
- kumoai/client/pquery.py +6 -2
- kumoai/client/rfm.py +37 -8
- kumoai/connector/utils.py +23 -2
- kumoai/experimental/rfm/__init__.py +164 -46
- 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} +49 -86
- kumoai/experimental/rfm/backend/local/sampler.py +315 -0
- kumoai/experimental/rfm/backend/local/table.py +119 -0
- kumoai/experimental/rfm/backend/snow/__init__.py +37 -0
- kumoai/experimental/rfm/backend/snow/sampler.py +274 -0
- kumoai/experimental/rfm/backend/snow/table.py +135 -0
- kumoai/experimental/rfm/backend/sqlite/__init__.py +32 -0
- kumoai/experimental/rfm/backend/sqlite/sampler.py +353 -0
- kumoai/experimental/rfm/backend/sqlite/table.py +126 -0
- kumoai/experimental/rfm/base/__init__.py +25 -0
- kumoai/experimental/rfm/base/column.py +66 -0
- kumoai/experimental/rfm/base/sampler.py +773 -0
- kumoai/experimental/rfm/base/source.py +19 -0
- kumoai/experimental/rfm/base/sql_sampler.py +60 -0
- kumoai/experimental/rfm/{local_table.py → base/table.py} +245 -156
- kumoai/experimental/rfm/{local_graph.py → graph.py} +425 -137
- kumoai/experimental/rfm/infer/__init__.py +6 -0
- kumoai/experimental/rfm/infer/dtype.py +79 -0
- kumoai/experimental/rfm/infer/pkey.py +126 -0
- kumoai/experimental/rfm/infer/time_col.py +62 -0
- kumoai/experimental/rfm/infer/timestamp.py +7 -4
- kumoai/experimental/rfm/pquery/__init__.py +4 -4
- kumoai/experimental/rfm/pquery/{backend.py → executor.py} +24 -58
- kumoai/experimental/rfm/pquery/{pandas_backend.py → pandas_executor.py} +278 -224
- kumoai/experimental/rfm/rfm.py +669 -246
- kumoai/experimental/rfm/sagemaker.py +138 -0
- kumoai/jobs.py +1 -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/trainer.py +12 -10
- kumoai/utils/__init__.py +3 -2
- kumoai/utils/progress_logger.py +239 -4
- kumoai/utils/sql.py +3 -0
- {kumoai-2.10.0.dev202509231831.dist-info → kumoai-2.14.0.dev202512161731.dist-info}/METADATA +15 -5
- {kumoai-2.10.0.dev202509231831.dist-info → kumoai-2.14.0.dev202512161731.dist-info}/RECORD +50 -32
- kumoai/experimental/rfm/local_graph_sampler.py +0 -176
- kumoai/experimental/rfm/local_pquery_driver.py +0 -404
- kumoai/experimental/rfm/utils.py +0 -344
- {kumoai-2.10.0.dev202509231831.dist-info → kumoai-2.14.0.dev202512161731.dist-info}/WHEEL +0 -0
- {kumoai-2.10.0.dev202509231831.dist-info → kumoai-2.14.0.dev202512161731.dist-info}/licenses/LICENSE +0 -0
- {kumoai-2.10.0.dev202509231831.dist-info → kumoai-2.14.0.dev202512161731.dist-info}/top_level.txt +0 -0
kumoai/__init__.py
CHANGED
|
@@ -184,15 +184,12 @@ def init(
|
|
|
184
184
|
snowflake_credentials
|
|
185
185
|
) if not api_key and snowflake_credentials else None
|
|
186
186
|
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.")
|
|
187
|
+
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
|
|
196
193
|
|
|
197
194
|
if not api_key and snowflake_credentials:
|
|
198
195
|
# Refresh token every 10 minutes (expires in 1 hour):
|
|
@@ -200,9 +197,11 @@ def init(
|
|
|
200
197
|
|
|
201
198
|
logger = logging.getLogger('kumoai')
|
|
202
199
|
log_level = logging.getLevelName(logger.getEffectiveLevel())
|
|
200
|
+
|
|
203
201
|
logger.info(
|
|
204
|
-
"Successfully initialized the Kumo SDK
|
|
205
|
-
"
|
|
202
|
+
f"Successfully initialized the Kumo SDK (version {__version__}) "
|
|
203
|
+
f"against deployment {url}, with "
|
|
204
|
+
f"log level {log_level}.")
|
|
206
205
|
|
|
207
206
|
|
|
208
207
|
def set_log_level(level: str) -> None:
|
|
@@ -281,7 +280,19 @@ __all__ = [
|
|
|
281
280
|
]
|
|
282
281
|
|
|
283
282
|
|
|
283
|
+
def in_snowflake_notebook() -> bool:
|
|
284
|
+
try:
|
|
285
|
+
from snowflake.snowpark.context import get_active_session
|
|
286
|
+
import streamlit # noqa: F401
|
|
287
|
+
get_active_session()
|
|
288
|
+
return True
|
|
289
|
+
except Exception:
|
|
290
|
+
return False
|
|
291
|
+
|
|
292
|
+
|
|
284
293
|
def in_notebook() -> bool:
|
|
294
|
+
if in_snowflake_notebook():
|
|
295
|
+
return True
|
|
285
296
|
try:
|
|
286
297
|
from IPython import get_ipython
|
|
287
298
|
shell = get_ipython()
|
kumoai/_version.py
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
__version__ = '2.
|
|
1
|
+
__version__ = '2.14.0.dev202512161731'
|
kumoai/client/client.py
CHANGED
|
@@ -20,7 +20,6 @@ if TYPE_CHECKING:
|
|
|
20
20
|
)
|
|
21
21
|
from kumoai.client.online import OnlineServingEndpointAPI
|
|
22
22
|
from kumoai.client.pquery import PQueryAPI
|
|
23
|
-
from kumoai.client.rfm import RFMAPI
|
|
24
23
|
from kumoai.client.source_table import SourceTableAPI
|
|
25
24
|
from kumoai.client.table import TableAPI
|
|
26
25
|
|
|
@@ -33,6 +32,7 @@ class KumoClient:
|
|
|
33
32
|
url: str,
|
|
34
33
|
api_key: Optional[str],
|
|
35
34
|
spcs_token: Optional[str] = None,
|
|
35
|
+
verify_ssl: bool = True,
|
|
36
36
|
) -> None:
|
|
37
37
|
r"""Creates a client against the Kumo public API, provided a URL of
|
|
38
38
|
the endpoint and an authentication token.
|
|
@@ -42,11 +42,14 @@ class KumoClient:
|
|
|
42
42
|
api_key: the public API authentication token.
|
|
43
43
|
spcs_token: the SPCS token used for authentication to access the
|
|
44
44
|
Kumo API endpoint.
|
|
45
|
+
verify_ssl: whether to verify SSL certificates. Set to False to
|
|
46
|
+
skip SSL certificate verification (equivalent to curl -k).
|
|
45
47
|
"""
|
|
46
48
|
self._url = url
|
|
47
49
|
self._api_url = f"{url}/{API_VERSION}"
|
|
48
50
|
self._api_key = api_key
|
|
49
51
|
self._spcs_token = spcs_token
|
|
52
|
+
self._verify_ssl = verify_ssl
|
|
50
53
|
|
|
51
54
|
retry_strategy = Retry(
|
|
52
55
|
total=10, # Maximum number of retries
|
|
@@ -69,11 +72,15 @@ class KumoClient:
|
|
|
69
72
|
self._session.headers.update(
|
|
70
73
|
{'Authorization': f'Snowflake Token={self._spcs_token}'})
|
|
71
74
|
|
|
72
|
-
def authenticate(self) ->
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
75
|
+
def authenticate(self) -> None:
|
|
76
|
+
"""Raises an exception if authentication fails."""
|
|
77
|
+
try:
|
|
78
|
+
self._session.get(self._url + '/v1/connectors',
|
|
79
|
+
verify=self._verify_ssl).raise_for_status()
|
|
80
|
+
except Exception:
|
|
81
|
+
raise ValueError(
|
|
82
|
+
"Client authentication failed. Please check if you "
|
|
83
|
+
"have a valid API key/credentials.")
|
|
77
84
|
|
|
78
85
|
def set_spcs_token(self, spcs_token: str) -> None:
|
|
79
86
|
r"""Sets the SPCS token for the client and updates the session
|
|
@@ -158,12 +165,6 @@ class KumoClient:
|
|
|
158
165
|
from kumoai.client.online import OnlineServingEndpointAPI
|
|
159
166
|
return OnlineServingEndpointAPI(self)
|
|
160
167
|
|
|
161
|
-
@property
|
|
162
|
-
def rfm_api(self) -> 'RFMAPI':
|
|
163
|
-
r"""Returns the typed RFM API."""
|
|
164
|
-
from kumoai.client.rfm import RFMAPI
|
|
165
|
-
return RFMAPI(self)
|
|
166
|
-
|
|
167
168
|
def _request(self, endpoint: Endpoint, **kwargs: Any) -> requests.Response:
|
|
168
169
|
r"""Send a HTTP request to the specified endpoint."""
|
|
169
170
|
endpoint_str = endpoint.get_path()
|
|
@@ -184,7 +185,7 @@ class KumoClient:
|
|
|
184
185
|
:meth:`requests.Session.get`.
|
|
185
186
|
"""
|
|
186
187
|
url = self._format_endpoint_url(endpoint)
|
|
187
|
-
return self._session.get(url=url, **kwargs)
|
|
188
|
+
return self._session.get(url=url, verify=self._verify_ssl, **kwargs)
|
|
188
189
|
|
|
189
190
|
def _post(self, endpoint: str, **kwargs: Any) -> requests.Response:
|
|
190
191
|
r"""Send a POST request to the specified endpoint, with keyword
|
|
@@ -192,7 +193,7 @@ class KumoClient:
|
|
|
192
193
|
:meth:`requests.Session.post`.
|
|
193
194
|
"""
|
|
194
195
|
url = self._format_endpoint_url(endpoint)
|
|
195
|
-
return self._session.post(url=url, **kwargs)
|
|
196
|
+
return self._session.post(url=url, verify=self._verify_ssl, **kwargs)
|
|
196
197
|
|
|
197
198
|
def _patch(self, endpoint: str, **kwargs: Any) -> requests.Response:
|
|
198
199
|
r"""Send a PATCH request to the specified endpoint, with keyword
|
|
@@ -200,7 +201,7 @@ class KumoClient:
|
|
|
200
201
|
:meth:`requests.Session.patch`.
|
|
201
202
|
"""
|
|
202
203
|
url = self._format_endpoint_url(endpoint)
|
|
203
|
-
return self._session.patch(url=url, **kwargs)
|
|
204
|
+
return self._session.patch(url=url, verify=self._verify_ssl, **kwargs)
|
|
204
205
|
|
|
205
206
|
def _delete(self, endpoint: str, **kwargs: Any) -> requests.Response:
|
|
206
207
|
r"""Send a DELETE request to the specified endpoint, with keyword
|
|
@@ -208,7 +209,7 @@ class KumoClient:
|
|
|
208
209
|
:meth:`requests.Session.delete`.
|
|
209
210
|
"""
|
|
210
211
|
url = self._format_endpoint_url(endpoint)
|
|
211
|
-
return self._session.delete(url=url, **kwargs)
|
|
212
|
+
return self._session.delete(url=url, verify=self._verify_ssl, **kwargs)
|
|
212
213
|
|
|
213
214
|
def _format_endpoint_url(self, endpoint: str) -> str:
|
|
214
215
|
if endpoint[0] == "/":
|
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/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,6 +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,
|
|
6
|
+
RFMExplanationResponse,
|
|
7
|
+
RFMParseQueryRequest,
|
|
8
|
+
RFMParseQueryResponse,
|
|
4
9
|
RFMPredictResponse,
|
|
5
10
|
RFMValidateQueryRequest,
|
|
6
11
|
RFMValidateQueryResponse,
|
|
@@ -25,28 +30,35 @@ class RFMAPI:
|
|
|
25
30
|
Returns:
|
|
26
31
|
RFMPredictResponse containing the predictions
|
|
27
32
|
"""
|
|
28
|
-
# Send binary data to the predict endpoint
|
|
29
33
|
response = self._client._request(
|
|
30
|
-
RFMEndpoints.predict,
|
|
31
|
-
|
|
34
|
+
RFMEndpoints.predict,
|
|
35
|
+
data=request,
|
|
36
|
+
headers={'Content-Type': 'application/x-protobuf'},
|
|
37
|
+
)
|
|
32
38
|
raise_on_error(response)
|
|
33
39
|
return parse_response(RFMPredictResponse, response)
|
|
34
40
|
|
|
35
|
-
def explain(
|
|
41
|
+
def explain(
|
|
42
|
+
self,
|
|
43
|
+
request: bytes,
|
|
44
|
+
skip_summary: bool = False,
|
|
45
|
+
) -> RFMExplanationResponse:
|
|
36
46
|
"""Explain the RFM model on the given context.
|
|
37
47
|
|
|
38
48
|
Args:
|
|
39
49
|
request: The predict request as serialized protobuf.
|
|
50
|
+
skip_summary: Whether to skip generating a human-readable summary
|
|
51
|
+
of the explanation.
|
|
40
52
|
|
|
41
53
|
Returns:
|
|
42
54
|
RFMPredictResponse containing the explanations
|
|
43
55
|
"""
|
|
44
|
-
|
|
56
|
+
params: dict[str, Any] = {'generate_summary': not skip_summary}
|
|
45
57
|
response = self._client._request(
|
|
46
|
-
RFMEndpoints.explain, data=request,
|
|
58
|
+
RFMEndpoints.explain, data=request, params=params,
|
|
47
59
|
headers={'Content-Type': 'application/x-protobuf'})
|
|
48
60
|
raise_on_error(response)
|
|
49
|
-
return parse_response(
|
|
61
|
+
return parse_response(RFMExplanationResponse, response)
|
|
50
62
|
|
|
51
63
|
def evaluate(self, request: bytes) -> RFMEvaluateResponse:
|
|
52
64
|
"""Evaluate the RFM model on the given context.
|
|
@@ -57,7 +69,6 @@ class RFMAPI:
|
|
|
57
69
|
Returns:
|
|
58
70
|
RFMEvaluateResponse containing the computed metrics
|
|
59
71
|
"""
|
|
60
|
-
# Send binary data to the evaluate endpoint
|
|
61
72
|
response = self._client._request(
|
|
62
73
|
RFMEndpoints.evaluate, data=request,
|
|
63
74
|
headers={'Content-Type': 'application/x-protobuf'})
|
|
@@ -81,3 +92,21 @@ class RFMAPI:
|
|
|
81
92
|
json=to_json_dict(request))
|
|
82
93
|
raise_on_error(response)
|
|
83
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,43 +1,123 @@
|
|
|
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.9, 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.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 typing import Dict, Optional, Tuple
|
|
10
|
+
from urllib.parse import urlparse
|
|
11
|
+
|
|
36
12
|
import kumoai
|
|
37
|
-
from .
|
|
38
|
-
|
|
39
|
-
from .rfm import KumoRFM
|
|
13
|
+
from kumoai.client.client import KumoClient
|
|
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) -> Tuple[InferenceBackend, Optional[str], Optional[str]]:
|
|
53
|
+
parsed = urlparse(url)
|
|
54
|
+
|
|
55
|
+
# Remote SageMaker
|
|
56
|
+
if ("runtime.sagemaker" in parsed.netloc
|
|
57
|
+
and parsed.path.endswith("/invocations")):
|
|
58
|
+
# Example: https://runtime.sagemaker.us-west-2.amazonaws.com/
|
|
59
|
+
# endpoints/Name/invocations
|
|
60
|
+
match = re.search(r"runtime\.sagemaker\.([a-z0-9-]+)\.amazonaws\.com",
|
|
61
|
+
parsed.netloc)
|
|
62
|
+
region = match.group(1) if match else None
|
|
63
|
+
m = re.search(r"/endpoints/([^/]+)/invocations", parsed.path)
|
|
64
|
+
endpoint_name = m.group(1) if m else None
|
|
65
|
+
return InferenceBackend.AWS_SAGEMAKER, region, endpoint_name
|
|
66
|
+
|
|
67
|
+
# Local SageMaker
|
|
68
|
+
if parsed.port == 8080 and parsed.path.endswith(
|
|
69
|
+
"/invocations") and _is_local_address(parsed.hostname):
|
|
70
|
+
return InferenceBackend.LOCAL_SAGEMAKER, None, None
|
|
71
|
+
|
|
72
|
+
# Default: regular REST
|
|
73
|
+
return InferenceBackend.REST, None, None
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
@dataclass
|
|
77
|
+
class RfmGlobalState:
|
|
78
|
+
_url: str = '__url_not_provided__'
|
|
79
|
+
_backend: InferenceBackend = InferenceBackend.UNKNOWN
|
|
80
|
+
_region: Optional[str] = None
|
|
81
|
+
_endpoint_name: Optional[str] = None
|
|
82
|
+
_thread_local = threading.local()
|
|
83
|
+
|
|
84
|
+
# Thread-safe init-once.
|
|
85
|
+
_initialized: bool = False
|
|
86
|
+
_lock: threading.Lock = threading.Lock()
|
|
87
|
+
|
|
88
|
+
@property
|
|
89
|
+
def client(self) -> KumoClient:
|
|
90
|
+
if self._backend == InferenceBackend.REST:
|
|
91
|
+
return kumoai.global_state.client
|
|
92
|
+
|
|
93
|
+
if hasattr(self._thread_local, '_sagemaker'):
|
|
94
|
+
# Set the spcs token in the client to ensure it has the latest.
|
|
95
|
+
return self._thread_local._sagemaker
|
|
96
|
+
|
|
97
|
+
sagemaker_client: KumoClient
|
|
98
|
+
if self._backend == InferenceBackend.LOCAL_SAGEMAKER:
|
|
99
|
+
sagemaker_client = KumoClient_SageMakerProxy_Local(self._url)
|
|
100
|
+
else:
|
|
101
|
+
assert self._backend == InferenceBackend.AWS_SAGEMAKER
|
|
102
|
+
assert self._region
|
|
103
|
+
assert self._endpoint_name
|
|
104
|
+
sagemaker_client = KumoClient_SageMakerAdapter(
|
|
105
|
+
self._region, self._endpoint_name)
|
|
106
|
+
|
|
107
|
+
self._thread_local._sagemaker = sagemaker_client
|
|
108
|
+
return sagemaker_client
|
|
109
|
+
|
|
110
|
+
def reset(self) -> None: # For testing only.
|
|
111
|
+
with self._lock:
|
|
112
|
+
self._initialized = False
|
|
113
|
+
self._url = '__url_not_provided__'
|
|
114
|
+
self._backend = InferenceBackend.UNKNOWN
|
|
115
|
+
self._region = None
|
|
116
|
+
self._endpoint_name = None
|
|
117
|
+
self._thread_local = threading.local()
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
global_state = RfmGlobalState()
|
|
41
121
|
|
|
42
122
|
|
|
43
123
|
def init(
|
|
@@ -47,19 +127,57 @@ def init(
|
|
|
47
127
|
snowflake_application: Optional[str] = None,
|
|
48
128
|
log_level: str = "INFO",
|
|
49
129
|
) -> None:
|
|
50
|
-
|
|
51
|
-
|
|
130
|
+
with global_state._lock:
|
|
131
|
+
if global_state._initialized:
|
|
132
|
+
if url != global_state._url:
|
|
133
|
+
raise ValueError(
|
|
134
|
+
"Kumo RFM has already been initialized with a different "
|
|
135
|
+
"URL. Re-initialization with a different URL is not "
|
|
136
|
+
"supported.")
|
|
137
|
+
return
|
|
52
138
|
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
snowflake_application=snowflake_application,
|
|
56
|
-
log_level=log_level)
|
|
139
|
+
if url is None:
|
|
140
|
+
url = os.getenv("RFM_API_URL", "https://kumorfm.ai/api")
|
|
57
141
|
|
|
142
|
+
backend, region, endpoint_name = _detect_backend(url)
|
|
143
|
+
if backend == InferenceBackend.REST:
|
|
144
|
+
# Initialize kumoai.global_state
|
|
145
|
+
if (kumoai.global_state.initialized
|
|
146
|
+
and kumoai.global_state._url != url):
|
|
147
|
+
raise ValueError(
|
|
148
|
+
"Kumo AI SDK has already been initialized with different "
|
|
149
|
+
"API URL. Please restart Python interpreter and "
|
|
150
|
+
"initialize via kumoai.rfm.init()")
|
|
151
|
+
kumoai.init(url=url, api_key=api_key,
|
|
152
|
+
snowflake_credentials=snowflake_credentials,
|
|
153
|
+
snowflake_application=snowflake_application,
|
|
154
|
+
log_level=log_level)
|
|
155
|
+
elif backend == InferenceBackend.AWS_SAGEMAKER:
|
|
156
|
+
assert region
|
|
157
|
+
assert endpoint_name
|
|
158
|
+
KumoClient_SageMakerAdapter(region, endpoint_name).authenticate()
|
|
159
|
+
else:
|
|
160
|
+
assert backend == InferenceBackend.LOCAL_SAGEMAKER
|
|
161
|
+
KumoClient_SageMakerProxy_Local(url).authenticate()
|
|
162
|
+
|
|
163
|
+
global_state._url = url
|
|
164
|
+
global_state._backend = backend
|
|
165
|
+
global_state._region = region
|
|
166
|
+
global_state._endpoint_name = endpoint_name
|
|
167
|
+
global_state._initialized = True
|
|
168
|
+
logger.info("Kumo RFM initialized with backend: %s, url: %s", backend,
|
|
169
|
+
url)
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
LocalGraph = Graph # NOTE Backward compatibility - do not use anymore.
|
|
58
173
|
|
|
59
174
|
__all__ = [
|
|
60
|
-
'LocalTable',
|
|
61
|
-
'LocalGraph',
|
|
62
|
-
'KumoRFM',
|
|
63
175
|
'authenticate',
|
|
64
176
|
'init',
|
|
177
|
+
'Table',
|
|
178
|
+
'LocalTable',
|
|
179
|
+
'Graph',
|
|
180
|
+
'KumoRFM',
|
|
181
|
+
'ExplainConfig',
|
|
182
|
+
'Explanation',
|
|
65
183
|
]
|
|
File without changes
|
|
@@ -0,0 +1,42 @@
|
|
|
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 .table import LocalTable
|
|
35
|
+
from .graph_store import LocalGraphStore
|
|
36
|
+
from .sampler import LocalSampler
|
|
37
|
+
|
|
38
|
+
__all__ = [
|
|
39
|
+
'LocalTable',
|
|
40
|
+
'LocalGraphStore',
|
|
41
|
+
'LocalSampler',
|
|
42
|
+
]
|