snowflake-ml-python 1.5.4__py3-none-any.whl → 1.6.1__py3-none-any.whl

snowflake/cortex/__init__.py

@@ -1,3 +1,4 @@
+from snowflake.cortex._classify_text import ClassifyText
 from snowflake.cortex._complete import Complete, CompleteOptions
 from snowflake.cortex._extract_answer import ExtractAnswer
 from snowflake.cortex._sentiment import Sentiment
@@ -5,6 +6,7 @@ from snowflake.cortex._summarize import Summarize
 from snowflake.cortex._translate import Translate
 
 __all__ = [
+    "ClassifyText",
     "Complete",
     "CompleteOptions",
     "ExtractAnswer",

snowflake/cortex/_classify_text.py

@@ -0,0 +1,36 @@
+from typing import List, Optional, Union
+
+from snowflake import snowpark
+from snowflake.cortex._util import CORTEX_FUNCTIONS_TELEMETRY_PROJECT, call_sql_function
+from snowflake.ml._internal import telemetry
+
+
+@telemetry.send_api_usage_telemetry(
+    project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+)
+def ClassifyText(
+    str_input: Union[str, snowpark.Column],
+    categories: Union[List[str], snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[str, snowpark.Column]:
+    """Use the LLM inference service to classify the INPUT text into one of the target CATEGORIES.
+
+    Args:
+        str_input: A Column of strings to classify.
+        categories: A list of candidate categories to classify the INPUT text into.
+        session: The snowpark session to use. Will be inferred by context if not specified.
+
+    Returns:
+        A column of classification responses.
+    """
+
+    return _classify_text_impl("snowflake.cortex.classify_text", str_input, categories, session=session)
+
+
+def _classify_text_impl(
+    function: str,
+    str_input: Union[str, snowpark.Column],
+    categories: Union[List[str], snowpark.Column],
+    session: Optional[snowpark.Session] = None,
+) -> Union[str, snowpark.Column]:
+    return call_sql_function(function, session, str_input, categories)

snowflake/cortex/_complete.py

@@ -1,6 +1,7 @@
 import json
 import logging
-from typing import Iterator, List, Optional, TypedDict, Union, cast
+import time
+from typing import Any, Callable, Iterator, List, Optional, TypedDict, Union, cast
 from urllib.parse import urlunparse
 
 import requests
@@ -52,12 +53,43 @@ class ResponseParseException(Exception):
     pass
 
 
+_MAX_RETRY_SECONDS = 30
+
+
+def retry(func: Callable[..., requests.Response]) -> Callable[..., requests.Response]:
+    def inner(*args: Any, **kwargs: Any) -> requests.Response:
+        deadline = cast(Optional[float], kwargs["deadline"])
+        kwargs = {key: value for key, value in kwargs.items() if key != "deadline"}
+        expRetrySeconds = 0.5
+        while True:
+            if deadline is not None and time.time() > deadline:
+                raise TimeoutError()
+            response = func(*args, **kwargs)
+            if response.status_code >= 200 and response.status_code < 300:
+                return response
+            retry_status_codes = [429, 503, 504]
+            if response.status_code not in retry_status_codes:
+                response.raise_for_status()
+            logger.debug(f"request failed with status code {response.status_code}, retrying")
+
+            # Formula: delay(i) = max(RetryAfterHeader, min(2^i, _MAX_RETRY_SECONDS)).
+            expRetrySeconds = min(2 * expRetrySeconds, _MAX_RETRY_SECONDS)
+            retrySeconds = expRetrySeconds
+            retryAfterHeader = response.headers.get("retry-after")
+            if retryAfterHeader is not None:
+                retrySeconds = max(retrySeconds, int(retryAfterHeader))
+            logger.debug(f"sleeping for {retrySeconds}s before retrying")
+            time.sleep(retrySeconds)
+
+    return inner
+
+
+@retry
 def _call_complete_rest(
     model: str,
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
-    stream: bool = False,
 ) -> requests.Response:
     session = session or context.get_active_session()
     if session is None:
@@ -78,7 +110,7 @@ def _call_complete_rest(
     scheme = "https"
     if hasattr(session.connection, "scheme"):
         scheme = session.connection.scheme
-    url = urlunparse((scheme, session.connection.host, "api/v2/cortex/inference/complete", "", "", ""))
+    url = urlunparse((scheme, session.connection.host, "api/v2/cortex/inference:complete", "", "", ""))
 
     headers = {
         "Content-Type": "application/json",
@@ -88,7 +120,7 @@ def _call_complete_rest(
 
     data = {
         "model": model,
-        "stream": stream,
+        "stream": True,
     }
     if isinstance(prompt, List):
         data["messages"] = prompt
@@ -104,33 +136,20 @@ def _call_complete_rest(
         if "top_p" in options:
             data["top_p"] = options["top_p"]
 
-    logger.debug(f"making POST request to {url} (model={model}, stream={stream})")
-    response = requests.post(
+    logger.debug(f"making POST request to {url} (model={model})")
+    return requests.post(
         url,
         json=data,
         headers=headers,
-        stream=stream,
+        stream=True,
     )
-    response.raise_for_status()
-    return response
 
 
-def _process_rest_response(response: requests.Response, stream: bool = False) -> Union[str, Iterator[str]]:
-    if stream:
-        return _return_stream_response(response)
-
-    try:
-        content = response.json()["choices"][0]["message"]["content"]
-        assert isinstance(content, str)
-        return content
-    except (KeyError, IndexError, AssertionError) as e:
-        # Unlike the streaming case, errors are not ignored because a message must be returned.
-        raise ResponseParseException("Failed to parse message from response.") from e
-
-
-def _return_stream_response(response: requests.Response) -> Iterator[str]:
+def _return_stream_response(response: requests.Response, deadline: Optional[float]) -> Iterator[str]:
     client = SSEClient(response)
     for event in client.events():
+        if deadline is not None and time.time() > deadline:
+            raise TimeoutError()
         try:
             yield json.loads(event.data)["choices"][0]["delta"]["content"]
         except (json.JSONDecodeError, KeyError, IndexError):
@@ -206,19 +225,23 @@ def _complete_impl(
     prompt: Union[str, List[ConversationMessage], snowpark.Column],
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
-    use_rest_api_experimental: bool = False,
     stream: bool = False,
     function: str = "snowflake.cortex.complete",
+    timeout: Optional[float] = None,
+    deadline: Optional[float] = None,
 ) -> Union[str, Iterator[str], snowpark.Column]:
-    if use_rest_api_experimental:
+    if timeout is not None and deadline is not None:
+        raise ValueError('only one of "timeout" and "deadline" must be set')
+    if timeout is not None:
+        deadline = time.time() + timeout
+    if stream:
         if not isinstance(model, str):
             raise ValueError("in REST mode, 'model' must be a string")
         if not isinstance(prompt, str) and not isinstance(prompt, List):
             raise ValueError("in REST mode, 'prompt' must be a string or a list of ConversationMessage")
-        response = _call_complete_rest(model, prompt, options, session=session, stream=stream)
-        return _process_rest_response(response, stream=stream)
-    if stream is True:
-        raise ValueError("streaming can only be enabled in REST mode, set use_rest_api_experimental=True")
+        response = _call_complete_rest(model, prompt, options, session=session, deadline=deadline)
+        assert response.status_code >= 200 and response.status_code < 300
+        return _return_stream_response(response, deadline)
     return _complete_sql_impl(function, model, prompt, options, session)
 
 
@@ -231,8 +254,9 @@ def Complete(
     *,
     options: Optional[CompleteOptions] = None,
     session: Optional[snowpark.Session] = None,
-    use_rest_api_experimental: bool = False,
     stream: bool = False,
+    timeout: Optional[float] = None,
+    deadline: Optional[float] = None,
 ) -> Union[str, Iterator[str], snowpark.Column]:
     """Complete calls into the LLM inference service to perform completion.
 
@@ -241,19 +265,26 @@ def Complete(
         prompt: A Column of prompts to send to the LLM.
         options: A instance of snowflake.cortex.CompleteOptions
         session: The snowpark session to use. Will be inferred by context if not specified.
-        use_rest_api_experimental (bool): Toggles between the use of SQL and REST implementation. This feature is
-            experimental and can be removed at any time.
         stream (bool): Enables streaming. When enabled, a generator function is returned that provides the streaming
             output as it is received. Each update is a string containing the new text content since the previous update.
-            The use of streaming requires the experimental use_rest_api_experimental flag to be enabled.
+        timeout (float): Timeout in seconds to retry failed REST requests.
+        deadline (float): Time in seconds since the epoch (as returned by time.time()) to retry failed REST requests.
 
     Raises:
-        ValueError: If `stream` is set to True and `use_rest_api_experimental` is set to False.
+        ValueError: incorrect argument.
 
     Returns:
         A column of string responses.
     """
     try:
-        return _complete_impl(model, prompt, options, session, use_rest_api_experimental, stream)
+        return _complete_impl(
+            model,
+            prompt,
+            options=options,
+            session=session,
+            stream=stream,
+            timeout=timeout,
+            deadline=deadline,
+        )
     except ValueError as err:
         raise err

snowflake/cortex/_util.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Union, cast
+from typing import Dict, List, Optional, Union, cast
 
 from snowflake import snowpark
 from snowflake.snowpark import context, functions
@@ -23,7 +23,7 @@ class SnowflakeConfigurationException(Exception):
 def call_sql_function(
     function: str,
     session: Optional[snowpark.Session],
-    *args: Union[str, snowpark.Column, Dict[str, Union[int, float]]],
+    *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
 ) -> Union[str, snowpark.Column]:
     handle_as_column = False
 
@@ -40,7 +40,7 @@ def call_sql_function(
 
 
 def _call_sql_function_column(
-    function: str, *args: Union[str, snowpark.Column, Dict[str, Union[int, float]]]
+    function: str, *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]]
 ) -> snowpark.Column:
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
@@ -48,7 +48,7 @@ def _call_sql_function_column(
 def _call_sql_function_immediate(
     function: str,
     session: Optional[snowpark.Session],
-    *args: Union[str, snowpark.Column, Dict[str, Union[int, float]]],
+    *args: Union[str, List[str], snowpark.Column, Dict[str, Union[int, float]]],
 ) -> str:
     session = session or context.get_active_session()
     if session is None:

snowflake/ml/_internal/env_utils.py

@@ -27,7 +27,6 @@ class CONDA_OS(Enum):
     NO_ARCH = "noarch"
 
 
-_SNOWFLAKE_CONDA_CHANNEL_URL = "https://repo.anaconda.com/pkgs/snowflake"
 _NODEFAULTS = "nodefaults"
 _SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
 _SNOWFLAKE_CONDA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
@@ -36,6 +35,7 @@ _SUPPORTED_PACKAGE_SPEC_OPS = ["==", ">=", "<=", ">", "<"]
 DEFAULT_CHANNEL_NAME = ""
 SNOWML_SPROC_ENV = "IN_SNOWML_SPROC"
 SNOWPARK_ML_PKG_NAME = "snowflake-ml-python"
+SNOWFLAKE_CONDA_CHANNEL_URL = "https://repo.anaconda.com/pkgs/snowflake"
 
 
 def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
@@ -370,7 +370,7 @@ def get_matched_package_versions_in_snowflake_conda_channel(
 
     assert not snowpark_utils.is_in_stored_procedure()  # type: ignore[no-untyped-call]
 
-    url = f"{_SNOWFLAKE_CONDA_CHANNEL_URL}/{conda_os.value}/repodata.json"
+    url = f"{SNOWFLAKE_CONDA_CHANNEL_URL}/{conda_os.value}/repodata.json"
 
     if req.name not in _SNOWFLAKE_CONDA_PACKAGE_CACHE:
         try:
@@ -477,6 +477,7 @@ def save_conda_env_file(
     path: pathlib.Path,
     conda_chan_deps: DefaultDict[str, List[requirements.Requirement]],
     python_version: str,
+    default_channel_override: str = SNOWFLAKE_CONDA_CHANNEL_URL,
 ) -> None:
     """Generate conda.yml file given a dict of dependencies after validation.
     The channels part of conda.yml file will contains Snowflake Anaconda Channel, nodefaults and all channel names
@@ -489,6 +490,7 @@ def save_conda_env_file(
         path: Path to the conda.yml file.
         conda_chan_deps: Dict of conda dependencies after validated.
         python_version: A string 'major.minor' showing python version relate to model.
+        default_channel_override: The default channel to be put in the first place of the channels section.
     """
     assert path.suffix in [".yml", ".yaml"], "Conda environment file should have extension of yml or yaml."
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -499,7 +501,11 @@ def save_conda_env_file(
     channels = list(dict(sorted(conda_chan_deps.items(), key=lambda item: len(item[1]), reverse=True)).keys())
     if DEFAULT_CHANNEL_NAME in channels:
         channels.remove(DEFAULT_CHANNEL_NAME)
-    env["channels"] = [_SNOWFLAKE_CONDA_CHANNEL_URL] + channels + [_NODEFAULTS]
+
+    if default_channel_override in channels:
+        channels.remove(default_channel_override)
+
+    env["channels"] = [default_channel_override] + channels + [_NODEFAULTS]
     env["dependencies"] = [f"python=={python_version}.*"]
     for chan, reqs in conda_chan_deps.items():
         env["dependencies"].extend(
@@ -567,8 +573,8 @@ def load_conda_env_file(
     python_version = None
 
     channels = env.get("channels", [])
-    if _SNOWFLAKE_CONDA_CHANNEL_URL in channels:
-        channels.remove(_SNOWFLAKE_CONDA_CHANNEL_URL)
+    if len(channels) >= 1:
+        channels = channels[1:]  # Skip the first channel which is the default channel
     if _NODEFAULTS in channels:
         channels.remove(_NODEFAULTS)
 

snowflake/ml/_internal/exceptions/modeling_error_messages.py

@@ -4,7 +4,10 @@ ATTRIBUTE_NOT_SET = (
     "-differences."
 )
 SIZE_MISMATCH = "Size mismatch: {}={}, {}={}."
-INVALID_MODEL_PARAM = "Invalid parameter {} for model {}. Valid parameters: {}."
+INVALID_MODEL_PARAM = (
+    "Invalid parameter {} for model {}. Valid parameters: {}."
+    "Note: Scikit learn params cannot be set until the model has been fit."
+)
 UNSUPPORTED_MODEL_CONVERSION = "Object doesn't support {}. Please use {}."
 INCOMPATIBLE_NEW_SKLEARN_PARAM = "Incompatible scikit-learn version: {} requires scikit-learn>={}. Installed: {}."
 REMOVED_SKLEARN_PARAM = "Incompatible scikit-learn version: {} is removed in scikit-learn>={}. Installed: {}."

snowflake/ml/_internal/lineage/lineage_utils.py

@@ -1,9 +1,9 @@
 import copy
 import functools
-from typing import Any, Callable, List, Optional
+from typing import Any, Callable, List, Optional, get_args
 
 from snowflake import snowpark
-from snowflake.ml._internal.lineage import data_source
+from snowflake.ml.data import data_source
 
 _DATA_SOURCES_ATTR = "_data_sources"
 
@@ -39,7 +39,7 @@ def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
     result: Optional[List[data_source.DataSource]] = None
     for arg in args:
         srcs = getattr(arg, _DATA_SOURCES_ATTR, None)
-        if isinstance(srcs, list) and all(isinstance(s, data_source.DataSource) for s in srcs):
+        if isinstance(srcs, list) and all(isinstance(s, get_args(data_source.DataSource)) for s in srcs):
             if result is None:
                 result = []
             result += srcs
@@ -49,7 +49,7 @@ def get_data_sources(*args: Any) -> Optional[List[data_source.DataSource]]:
 def set_data_sources(obj: Any, data_sources: Optional[List[data_source.DataSource]]) -> None:
     """Helper method for attaching data sources to an object"""
     if data_sources:
-        assert all(isinstance(ds, data_source.DataSource) for ds in data_sources)
+        assert all(isinstance(ds, get_args(data_source.DataSource)) for ds in data_sources)
     setattr(obj, _DATA_SOURCES_ATTR, data_sources)
 
 

snowflake/ml/_internal/telemetry.py

@@ -44,6 +44,20 @@ _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
 
+@enum.unique
+class TelemetryProject(enum.Enum):
+    MLOPS = "MLOps"
+    MODELING = "ModelDevelopment"
+    # TODO: Update with remaining projects.
+
+
+@enum.unique
+class TelemetrySubProject(enum.Enum):
+    MONITORING = "Monitoring"
+    REGISTRY = "ModelManagement"
+    # TODO: Update with remaining subprojects.
+
+
 @enum.unique
 class TelemetryField(enum.Enum):
     # constants
@@ -277,6 +291,7 @@ def send_api_usage_telemetry(
         ]
     ] = None,
     sfqids_extractor: Optional[Callable[..., List[str]]] = None,
+    subproject_extractor: Optional[Callable[[Any], str]] = None,
     custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, _ReturnValue]]:
     """
@@ -290,6 +305,7 @@ def send_api_usage_telemetry(
         conn_attr_name: Name of the SnowflakeConnection attribute in `self`.
         api_calls_extractor: Extract API calls from `self`.
         sfqids_extractor: Extract sfqids from `self`.
+        subproject_extractor: Extract subproject at runtime from `self`.
         custom_tags: Custom tags.
 
     Returns:
@@ -297,10 +313,14 @@ def send_api_usage_telemetry(
 
     Raises:
         TypeError: If `conn_attr_name` is provided but the conn attribute is not of type SnowflakeConnection.
+        ValueError: If both `subproject` and `subproject_extractor` are provided
 
     # noqa: DAR402
     """
 
+    if subproject is not None and subproject_extractor is not None:
+        raise ValueError("Specifying both subproject and subproject_extractor is not allowed")
+
     def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, _ReturnValue]:
         @functools.wraps(func)
         def wrap(*args: Any, **kwargs: Any) -> _ReturnValue:
@@ -322,9 +342,13 @@ def send_api_usage_telemetry(
             if sfqids_extractor:
                 sfqids = sfqids_extractor(args[0])
 
+            subproject_name = subproject
+            if subproject_extractor is not None:
+                subproject_name = subproject_extractor(args[0])
+
             statement_params = get_function_usage_statement_params(
                 project=project,
-                subproject=subproject,
+                subproject=subproject_name,
                 function_category=TelemetryField.FUNC_CAT_USAGE.value,
                 function_name=_get_full_func_name(func),
                 function_parameters=params,
@@ -381,7 +405,7 @@ def send_api_usage_telemetry(
                         raise e.original_exception from e
 
             # TODO(hayu): [SNOW-750287] Optimize telemetry client to a singleton.
-            telemetry = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject)
+            telemetry = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject_name)
             telemetry_args = dict(
                 func_name=_get_full_func_name(func),
                 function_category=TelemetryField.FUNC_CAT_USAGE.value,

snowflake/ml/_internal/utils/pkg_version_utils.py

@@ -26,30 +26,11 @@ def get_valid_pkg_versions_supported_in_snowflake_conda_channel(
     pkg_versions: List[str], session: Session, subproject: Optional[str] = None
 ) -> List[str]:
     if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
-        return _get_valid_pkg_versions_supported_in_snowflake_conda_channel_sync(pkg_versions, session, subproject)
+        return pkg_versions
     else:
         return _get_valid_pkg_versions_supported_in_snowflake_conda_channel_async(pkg_versions, session, subproject)
 
 
-def _get_valid_pkg_versions_supported_in_snowflake_conda_channel_sync(
-    pkg_versions: List[str], session: Session, subproject: Optional[str] = None
-) -> List[str]:
-    for pkg_version in pkg_versions:
-        if pkg_version not in cache:
-            pkg_version_list = _query_pkg_version_supported_in_snowflake_conda_channel(
-                pkg_version=pkg_version, session=session, block=True, subproject=subproject
-            )
-            assert isinstance(pkg_version_list, list)  # keep mypy happy
-            try:
-                cache[pkg_version] = pkg_version_list[0]["VERSION"]
-            except IndexError:
-                cache[pkg_version] = None
-
-    pkg_version_conda_list = _get_conda_packages_and_emit_warnings(pkg_versions)
-
-    return pkg_version_conda_list
-
-
 def _get_valid_pkg_versions_supported_in_snowflake_conda_channel_async(
     pkg_versions: List[str], session: Session, subproject: Optional[str] = None
 ) -> List[str]:
@@ -60,7 +41,11 @@ def _get_valid_pkg_versions_supported_in_snowflake_conda_channel_async(
             async_job = _query_pkg_version_supported_in_snowflake_conda_channel(
                 pkg_version=pkg_version, session=session, block=False, subproject=subproject
             )
-            assert isinstance(async_job, AsyncJob)
+            if isinstance(async_job, list):
+                raise RuntimeError(
+                    "Async job was expected, executed query was returned. Please contact Snowflake support."
+                )
+
             pkg_version_async_job_list.append((pkg_version, async_job))
 
     # Populate the cache.
@@ -143,7 +128,8 @@ def _get_conda_packages_and_emit_warnings(pkg_versions: List[str]) -> List[str]:
         warnings.warn(
             f"Package {', '.join([pkg[0] for pkg in pkg_version_warning_list])} is not supported "
             f"in snowflake conda channel for python runtime "
-            f"{', '.join([pkg[1] for pkg in pkg_version_warning_list])}."
+            f"{', '.join([pkg[1] for pkg in pkg_version_warning_list])}.",
+            stacklevel=1,
         )
 
     return pkg_version_conda_list