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.
Files changed (58) hide show
  1. kumoai/__init__.py +35 -26
  2. kumoai/_version.py +1 -1
  3. kumoai/client/client.py +6 -0
  4. kumoai/client/jobs.py +26 -0
  5. kumoai/client/pquery.py +6 -2
  6. kumoai/connector/utils.py +44 -9
  7. kumoai/experimental/rfm/__init__.py +70 -68
  8. kumoai/experimental/rfm/authenticate.py +3 -4
  9. kumoai/experimental/rfm/backend/__init__.py +0 -0
  10. kumoai/experimental/rfm/backend/local/__init__.py +42 -0
  11. kumoai/experimental/rfm/{local_graph_store.py → backend/local/graph_store.py} +65 -127
  12. kumoai/experimental/rfm/backend/local/sampler.py +312 -0
  13. kumoai/experimental/rfm/backend/local/table.py +113 -0
  14. kumoai/experimental/rfm/backend/snow/__init__.py +37 -0
  15. kumoai/experimental/rfm/backend/snow/sampler.py +366 -0
  16. kumoai/experimental/rfm/backend/snow/table.py +242 -0
  17. kumoai/experimental/rfm/backend/sqlite/__init__.py +32 -0
  18. kumoai/experimental/rfm/backend/sqlite/sampler.py +454 -0
  19. kumoai/experimental/rfm/backend/sqlite/table.py +184 -0
  20. kumoai/experimental/rfm/base/__init__.py +30 -0
  21. kumoai/experimental/rfm/base/column.py +152 -0
  22. kumoai/experimental/rfm/base/expression.py +44 -0
  23. kumoai/experimental/rfm/base/mapper.py +67 -0
  24. kumoai/experimental/rfm/base/sampler.py +782 -0
  25. kumoai/experimental/rfm/base/source.py +19 -0
  26. kumoai/experimental/rfm/base/sql_sampler.py +366 -0
  27. kumoai/experimental/rfm/base/table.py +741 -0
  28. kumoai/experimental/rfm/{local_graph.py → graph.py} +581 -154
  29. kumoai/experimental/rfm/infer/__init__.py +8 -0
  30. kumoai/experimental/rfm/infer/dtype.py +82 -0
  31. kumoai/experimental/rfm/infer/multicategorical.py +1 -1
  32. kumoai/experimental/rfm/infer/pkey.py +128 -0
  33. kumoai/experimental/rfm/infer/stype.py +35 -0
  34. kumoai/experimental/rfm/infer/time_col.py +61 -0
  35. kumoai/experimental/rfm/pquery/executor.py +27 -27
  36. kumoai/experimental/rfm/pquery/pandas_executor.py +30 -32
  37. kumoai/experimental/rfm/relbench.py +76 -0
  38. kumoai/experimental/rfm/rfm.py +775 -481
  39. kumoai/experimental/rfm/sagemaker.py +15 -7
  40. kumoai/experimental/rfm/task_table.py +292 -0
  41. kumoai/pquery/predictive_query.py +10 -6
  42. kumoai/pquery/training_table.py +16 -2
  43. kumoai/testing/decorators.py +1 -1
  44. kumoai/testing/snow.py +50 -0
  45. kumoai/trainer/distilled_trainer.py +175 -0
  46. kumoai/utils/__init__.py +3 -2
  47. kumoai/utils/display.py +87 -0
  48. kumoai/utils/progress_logger.py +190 -12
  49. kumoai/utils/sql.py +3 -0
  50. {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/METADATA +10 -8
  51. {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/RECORD +54 -30
  52. kumoai/experimental/rfm/local_graph_sampler.py +0 -182
  53. kumoai/experimental/rfm/local_pquery_driver.py +0 -689
  54. kumoai/experimental/rfm/local_table.py +0 -545
  55. kumoai/experimental/rfm/utils.py +0 -344
  56. {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/WHEEL +0 -0
  57. {kumoai-2.13.0.dev202511191731.dist-info → kumoai-2.14.0rc2.dist-info}/licenses/LICENSE +0 -0
  58. {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
- "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:
@@ -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
- f"Successfully initialized the Kumo SDK (version {__version__}) "
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.13.0.dev202511191731'
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
- 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/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
- _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
 
@@ -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
- mv = buf.getbuffer() # zero-copy view of BytesIO internal buffer
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
- yield mv[pos:nxt] # zero-copy slice
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
- delimiter.
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 typing import Optional, Dict, Tuple
42
- import os
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 .sagemaker import (KumoClient_SageMakerAdapter,
47
- KumoClient_SageMakerProxy_Local)
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) -> Tuple[InferenceBackend, Optional[str], Optional[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: Optional[str] = None
109
- _endpoint_name: Optional[str] = None
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: Optional[str] = None,
153
- api_key: Optional[str] = None,
154
- snowflake_credentials: Optional[Dict[str, str]] = None,
155
- snowflake_application: Optional[str] = None,
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 ValueError(
162
- "Kumo RFM has already been initialized with a different "
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
- # Initialize kumoai.global_state
173
- if (kumoai.global_state.initialized
174
- and kumoai.global_state._url != url):
175
- raise ValueError(
176
- "Kumo AI SDK has already been initialized with different "
177
- "API URL. Please restart Python interpreter and "
178
- "initialize via kumoai.rfm.init()")
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
- 'LocalGraph',
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: 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
@@ -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
+ ]