kumoai 2.13.0.dev202511191731__cp310-cp310-macosx_11_0_arm64.whl → 2.14.0rc2__cp310-cp310-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.
- kumoai/__init__.py +35 -26
- kumoai/_version.py +1 -1
- kumoai/client/client.py +6 -0
- kumoai/client/jobs.py +26 -0
- kumoai/client/pquery.py +6 -2
- kumoai/connector/utils.py +44 -9
- kumoai/experimental/rfm/__init__.py +70 -68
- 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 +366 -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 +454 -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/mapper.py +67 -0
- kumoai/experimental/rfm/base/sampler.py +782 -0
- kumoai/experimental/rfm/base/source.py +19 -0
- kumoai/experimental/rfm/base/sql_sampler.py +366 -0
- kumoai/experimental/rfm/base/table.py +741 -0
- kumoai/experimental/rfm/{local_graph.py → graph.py} +581 -154
- 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/executor.py +27 -27
- kumoai/experimental/rfm/pquery/pandas_executor.py +30 -32
- kumoai/experimental/rfm/relbench.py +76 -0
- kumoai/experimental/rfm/rfm.py +775 -481
- kumoai/experimental/rfm/sagemaker.py +15 -7
- kumoai/experimental/rfm/task_table.py +292 -0
- kumoai/pquery/predictive_query.py +10 -6
- kumoai/pquery/training_table.py +16 -2
- kumoai/testing/decorators.py +1 -1
- kumoai/testing/snow.py +50 -0
- kumoai/trainer/distilled_trainer.py +175 -0
- kumoai/utils/__init__.py +3 -2
- kumoai/utils/display.py +87 -0
- kumoai/utils/progress_logger.py +190 -12
- kumoai/utils/sql.py +3 -0
- {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/METADATA +10 -8
- {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/RECORD +54 -30
- kumoai/experimental/rfm/local_graph_sampler.py +0 -182
- kumoai/experimental/rfm/local_pquery_driver.py +0 -689
- kumoai/experimental/rfm/local_table.py +0 -545
- kumoai/experimental/rfm/utils.py +0 -344
- {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/WHEEL +0 -0
- {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/licenses/LICENSE +0 -0
- {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.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:
|
|
@@ -198,10 +197,8 @@ def init(
|
|
|
198
197
|
logger = logging.getLogger('kumoai')
|
|
199
198
|
log_level = logging.getLevelName(logger.getEffectiveLevel())
|
|
200
199
|
|
|
201
|
-
logger.info(
|
|
202
|
-
|
|
203
|
-
f"against deployment {url}, with "
|
|
204
|
-
f"log level {log_level}.")
|
|
200
|
+
logger.info(f"Initialized Kumo SDK v{__version__} against deployment "
|
|
201
|
+
f"'{url}'")
|
|
205
202
|
|
|
206
203
|
|
|
207
204
|
def set_log_level(level: str) -> None:
|
|
@@ -280,7 +277,19 @@ __all__ = [
|
|
|
280
277
|
]
|
|
281
278
|
|
|
282
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
|
+
|
|
283
290
|
def in_notebook() -> bool:
|
|
291
|
+
if in_snowflake_notebook():
|
|
292
|
+
return True
|
|
284
293
|
try:
|
|
285
294
|
from IPython import get_ipython
|
|
286
295
|
shell = get_ipython()
|
kumoai/_version.py
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
__version__ = '2.
|
|
1
|
+
__version__ = '2.14.0rc2'
|
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,
|
|
@@ -132,6 +133,11 @@ class KumoClient:
|
|
|
132
133
|
from kumoai.client.jobs import TrainingJobAPI
|
|
133
134
|
return TrainingJobAPI(self)
|
|
134
135
|
|
|
136
|
+
@property
|
|
137
|
+
def distillation_job_api(self) -> 'DistillationJobAPI':
|
|
138
|
+
from kumoai.client.jobs import DistillationJobAPI
|
|
139
|
+
return DistillationJobAPI(self)
|
|
140
|
+
|
|
135
141
|
@property
|
|
136
142
|
def batch_prediction_job_api(self) -> 'BatchPredictionJobAPI':
|
|
137
143
|
from kumoai.client.jobs import BatchPredictionJobAPI
|
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."""
|
|
@@ -320,12 +344,14 @@ class GenerateTrainTableJobAPI(CommonJobAPI[GenerateTrainTableRequest,
|
|
|
320
344
|
id: str,
|
|
321
345
|
source_table_type: SourceTableType,
|
|
322
346
|
train_table_mod: TrainingTableSpec,
|
|
347
|
+
extensive_validation: bool,
|
|
323
348
|
) -> ValidationResponse:
|
|
324
349
|
response = self._client._post(
|
|
325
350
|
f'{self._base_endpoint}/{id}/validate_custom_train_table',
|
|
326
351
|
json=to_json_dict({
|
|
327
352
|
'custom_table': source_table_type,
|
|
328
353
|
'train_table_mod': train_table_mod,
|
|
354
|
+
'extensive_validation': extensive_validation,
|
|
329
355
|
}),
|
|
330
356
|
)
|
|
331
357
|
return parse_response(ValidationResponse, response)
|
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/connector/utils.py
CHANGED
|
@@ -1,10 +1,10 @@
|
|
|
1
1
|
import asyncio
|
|
2
2
|
import csv
|
|
3
|
-
import gc
|
|
4
3
|
import io
|
|
5
4
|
import math
|
|
6
5
|
import os
|
|
7
6
|
import re
|
|
7
|
+
import sys
|
|
8
8
|
import tempfile
|
|
9
9
|
import threading
|
|
10
10
|
import time
|
|
@@ -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
|
|
|
@@ -899,7 +920,10 @@ def _read_remote_file_with_progress(
|
|
|
899
920
|
if capture_first_line and not seen_nl:
|
|
900
921
|
header_line = bytes(header_acc)
|
|
901
922
|
|
|
902
|
-
|
|
923
|
+
if sys.version_info >= (3, 13):
|
|
924
|
+
mv = memoryview(buf.getvalue())
|
|
925
|
+
else:
|
|
926
|
+
mv = buf.getbuffer() # zero-copy view of BytesIO internal buffer
|
|
903
927
|
return buf, mv, header_line
|
|
904
928
|
|
|
905
929
|
|
|
@@ -978,7 +1002,10 @@ def _iter_mv_chunks(mv: memoryview,
|
|
|
978
1002
|
n = mv.nbytes
|
|
979
1003
|
while pos < n:
|
|
980
1004
|
nxt = min(n, pos + part_size)
|
|
981
|
-
|
|
1005
|
+
if sys.version_info >= (3, 13):
|
|
1006
|
+
yield mv[pos:nxt].tobytes()
|
|
1007
|
+
else:
|
|
1008
|
+
yield mv[pos:nxt] # zero-copy slice
|
|
982
1009
|
pos = nxt
|
|
983
1010
|
|
|
984
1011
|
|
|
@@ -1168,7 +1195,7 @@ def _detect_and_validate_csv(head_bytes: bytes) -> str:
|
|
|
1168
1195
|
- Re-serializes those rows and validates with pandas (small nrows) to catch
|
|
1169
1196
|
malformed inputs.
|
|
1170
1197
|
- Raises ValueError on empty input or if parsing fails with the chosen
|
|
1171
|
-
|
|
1198
|
+
delimiter.
|
|
1172
1199
|
"""
|
|
1173
1200
|
if not head_bytes:
|
|
1174
1201
|
raise ValueError("Could not auto-detect a delimiter: file is empty.")
|
|
@@ -1452,13 +1479,17 @@ def _remote_upload_file(name: str, fs: Filesystem, url: str, info: dict,
|
|
|
1452
1479
|
if renamed_cols_msg:
|
|
1453
1480
|
logger.info(renamed_cols_msg)
|
|
1454
1481
|
|
|
1482
|
+
try:
|
|
1483
|
+
if isinstance(data_mv, memoryview):
|
|
1484
|
+
data_mv.release()
|
|
1485
|
+
except Exception:
|
|
1486
|
+
pass
|
|
1487
|
+
|
|
1455
1488
|
try:
|
|
1456
1489
|
if buf:
|
|
1457
1490
|
buf.close()
|
|
1458
1491
|
except Exception:
|
|
1459
1492
|
pass
|
|
1460
|
-
del buf, data_mv, header_line
|
|
1461
|
-
gc.collect()
|
|
1462
1493
|
|
|
1463
1494
|
logger.info("Upload complete. Validated table %s.", name)
|
|
1464
1495
|
|
|
@@ -1698,13 +1729,17 @@ def _remote_upload_directory(
|
|
|
1698
1729
|
else:
|
|
1699
1730
|
break
|
|
1700
1731
|
|
|
1732
|
+
try:
|
|
1733
|
+
if isinstance(data_mv, memoryview):
|
|
1734
|
+
data_mv.release()
|
|
1735
|
+
except Exception:
|
|
1736
|
+
pass
|
|
1737
|
+
|
|
1701
1738
|
try:
|
|
1702
1739
|
if buf:
|
|
1703
1740
|
buf.close()
|
|
1704
1741
|
except Exception:
|
|
1705
1742
|
pass
|
|
1706
|
-
del buf, data_mv, header_line
|
|
1707
|
-
gc.collect()
|
|
1708
1743
|
|
|
1709
1744
|
_safe_bar_update(file_bar, 1)
|
|
1710
1745
|
_merge_status_update(fpath)
|
|
@@ -1,54 +1,27 @@
|
|
|
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 dataclasses import dataclass
|
|
35
|
-
from enum import Enum
|
|
36
1
|
import ipaddress
|
|
37
2
|
import logging
|
|
3
|
+
import os
|
|
38
4
|
import re
|
|
39
5
|
import socket
|
|
40
6
|
import threading
|
|
41
|
-
from
|
|
42
|
-
import
|
|
7
|
+
from dataclasses import dataclass
|
|
8
|
+
from enum import Enum
|
|
43
9
|
from urllib.parse import urlparse
|
|
10
|
+
|
|
44
11
|
import kumoai
|
|
45
12
|
from kumoai.client.client import KumoClient
|
|
46
|
-
from .
|
|
47
|
-
|
|
48
|
-
from .local_table import LocalTable
|
|
49
|
-
from .local_graph import LocalGraph
|
|
50
|
-
from .rfm import ExplainConfig, Explanation, KumoRFM
|
|
13
|
+
from kumoai.spcs import _get_active_session
|
|
14
|
+
|
|
51
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 .task_table import TaskTable
|
|
24
|
+
from .rfm import ExplainConfig, Explanation, KumoRFM
|
|
52
25
|
|
|
53
26
|
logger = logging.getLogger('kumoai_rfm')
|
|
54
27
|
|
|
@@ -77,7 +50,8 @@ class InferenceBackend(str, Enum):
|
|
|
77
50
|
|
|
78
51
|
|
|
79
52
|
def _detect_backend(
|
|
80
|
-
url: str
|
|
53
|
+
url: str, #
|
|
54
|
+
) -> tuple[InferenceBackend, str | None, str | None]:
|
|
81
55
|
parsed = urlparse(url)
|
|
82
56
|
|
|
83
57
|
# Remote SageMaker
|
|
@@ -101,12 +75,27 @@ def _detect_backend(
|
|
|
101
75
|
return InferenceBackend.REST, None, None
|
|
102
76
|
|
|
103
77
|
|
|
78
|
+
def _get_snowflake_url(snowflake_application: str) -> str:
|
|
79
|
+
snowpark_session = _get_active_session()
|
|
80
|
+
if not snowpark_session:
|
|
81
|
+
raise ValueError(
|
|
82
|
+
"KumoRFM initialization failed. 'snowflake_application' is "
|
|
83
|
+
"specified without an active Snowpark session. If running outside "
|
|
84
|
+
"a Snowflake notebook, specify a URL and credentials.")
|
|
85
|
+
with snowpark_session.connection.cursor() as cur:
|
|
86
|
+
cur.execute(
|
|
87
|
+
f"DESCRIBE SERVICE {snowflake_application}.user_schema.rfm_service"
|
|
88
|
+
f" ->> SELECT \"dns_name\" from $1")
|
|
89
|
+
dns_name: str = cur.fetchone()[0]
|
|
90
|
+
return f"http://{dns_name}:8000/api"
|
|
91
|
+
|
|
92
|
+
|
|
104
93
|
@dataclass
|
|
105
94
|
class RfmGlobalState:
|
|
106
95
|
_url: str = '__url_not_provided__'
|
|
107
96
|
_backend: InferenceBackend = InferenceBackend.UNKNOWN
|
|
108
|
-
_region:
|
|
109
|
-
_endpoint_name:
|
|
97
|
+
_region: str | None = None
|
|
98
|
+
_endpoint_name: str | None = None
|
|
110
99
|
_thread_local = threading.local()
|
|
111
100
|
|
|
112
101
|
# Thread-safe init-once.
|
|
@@ -115,6 +104,9 @@ class RfmGlobalState:
|
|
|
115
104
|
|
|
116
105
|
@property
|
|
117
106
|
def client(self) -> KumoClient:
|
|
107
|
+
if self._backend == InferenceBackend.UNKNOWN:
|
|
108
|
+
raise RuntimeError("KumoRFM is not yet initialized")
|
|
109
|
+
|
|
118
110
|
if self._backend == InferenceBackend.REST:
|
|
119
111
|
return kumoai.global_state.client
|
|
120
112
|
|
|
@@ -149,60 +141,70 @@ global_state = RfmGlobalState()
|
|
|
149
141
|
|
|
150
142
|
|
|
151
143
|
def init(
|
|
152
|
-
url:
|
|
153
|
-
api_key:
|
|
154
|
-
snowflake_credentials:
|
|
155
|
-
snowflake_application:
|
|
144
|
+
url: str | None = None,
|
|
145
|
+
api_key: str | None = None,
|
|
146
|
+
snowflake_credentials: dict[str, str] | None = None,
|
|
147
|
+
snowflake_application: str | None = None,
|
|
156
148
|
log_level: str = "INFO",
|
|
157
149
|
) -> None:
|
|
158
150
|
with global_state._lock:
|
|
159
151
|
if global_state._initialized:
|
|
160
152
|
if url != global_state._url:
|
|
161
|
-
raise
|
|
162
|
-
"
|
|
163
|
-
"URL. Re-initialization with a different URL is not "
|
|
153
|
+
raise RuntimeError(
|
|
154
|
+
"KumoRFM has already been initialized with a different "
|
|
155
|
+
"API URL. Re-initialization with a different URL is not "
|
|
164
156
|
"supported.")
|
|
165
157
|
return
|
|
166
158
|
|
|
159
|
+
if snowflake_application:
|
|
160
|
+
if url is not None:
|
|
161
|
+
raise ValueError(
|
|
162
|
+
"KumoRFM initialization failed. Both "
|
|
163
|
+
"'snowflake_application' and 'url' are specified. If "
|
|
164
|
+
"running from a Snowflake notebook, specify only "
|
|
165
|
+
"'snowflake_application'.")
|
|
166
|
+
url = _get_snowflake_url(snowflake_application)
|
|
167
|
+
api_key = "test:DISABLED"
|
|
168
|
+
|
|
167
169
|
if url is None:
|
|
168
170
|
url = os.getenv("RFM_API_URL", "https://kumorfm.ai/api")
|
|
169
171
|
|
|
170
172
|
backend, region, endpoint_name = _detect_backend(url)
|
|
171
173
|
if backend == InferenceBackend.REST:
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
kumoai.init(url=url, api_key=api_key,
|
|
180
|
-
snowflake_credentials=snowflake_credentials,
|
|
181
|
-
snowflake_application=snowflake_application,
|
|
182
|
-
log_level=log_level)
|
|
174
|
+
kumoai.init(
|
|
175
|
+
url=url,
|
|
176
|
+
api_key=api_key,
|
|
177
|
+
snowflake_credentials=snowflake_credentials,
|
|
178
|
+
snowflake_application=snowflake_application,
|
|
179
|
+
log_level=log_level,
|
|
180
|
+
)
|
|
183
181
|
elif backend == InferenceBackend.AWS_SAGEMAKER:
|
|
184
182
|
assert region
|
|
185
183
|
assert endpoint_name
|
|
186
184
|
KumoClient_SageMakerAdapter(region, endpoint_name).authenticate()
|
|
185
|
+
logger.info("KumoRFM initialized in AWS SageMaker")
|
|
187
186
|
else:
|
|
188
187
|
assert backend == InferenceBackend.LOCAL_SAGEMAKER
|
|
189
188
|
KumoClient_SageMakerProxy_Local(url).authenticate()
|
|
189
|
+
logger.info(f"KumoRFM initialized in local SageMaker at '{url}'")
|
|
190
190
|
|
|
191
191
|
global_state._url = url
|
|
192
192
|
global_state._backend = backend
|
|
193
193
|
global_state._region = region
|
|
194
194
|
global_state._endpoint_name = endpoint_name
|
|
195
195
|
global_state._initialized = True
|
|
196
|
-
logger.info("Kumo RFM initialized with backend: %s, url: %s", backend,
|
|
197
|
-
url)
|
|
198
196
|
|
|
199
197
|
|
|
198
|
+
LocalGraph = Graph # NOTE Backward compatibility - do not use anymore.
|
|
199
|
+
|
|
200
200
|
__all__ = [
|
|
201
|
+
'authenticate',
|
|
202
|
+
'init',
|
|
203
|
+
'Table',
|
|
201
204
|
'LocalTable',
|
|
202
|
-
'
|
|
205
|
+
'Graph',
|
|
206
|
+
'TaskTable',
|
|
203
207
|
'KumoRFM',
|
|
204
208
|
'ExplainConfig',
|
|
205
209
|
'Explanation',
|
|
206
|
-
'authenticate',
|
|
207
|
-
'init',
|
|
208
210
|
]
|
|
@@ -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
|
|
@@ -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
|
+
]
|