snowflake-ml-python 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

snowflake/cortex/_complete.py

@@ -23,7 +23,7 @@ def Complete(
         A column of string responses.
     """
 
-    return _complete_impl("snowflake.ml.complete", model, prompt, session=session)
+    return _complete_impl("snowflake.cortex.complete", model, prompt, session=session)
 
 
 def _complete_impl(

snowflake/cortex/_extract_answer.py

@@ -25,7 +25,7 @@ def ExtractAnswer(
         A column of strings containing answers.
     """
 
-    return _extract_answer_impl("snowflake.ml.extract_answer", from_text, question, session=session)
+    return _extract_answer_impl("snowflake.cortex.extract_answer", from_text, question, session=session)
 
 
 def _extract_answer_impl(

snowflake/cortex/_sentiment.py

@@ -22,7 +22,7 @@ def Sentiment(
         A column of floats. 1 represents positive sentiment, -1 represents negative sentiment.
     """
 
-    return _sentiment_impl("snowflake.ml.sentiment", text, session=session)
+    return _sentiment_impl("snowflake.cortex.sentiment", text, session=session)
 
 
 def _sentiment_impl(

snowflake/cortex/_summarize.py

@@ -23,7 +23,7 @@ def Summarize(
         A column of string summaries.
     """
 
-    return _summarize_impl("snowflake.ml.summarize", text, session=session)
+    return _summarize_impl("snowflake.cortex.summarize", text, session=session)
 
 
 def _summarize_impl(

snowflake/cortex/_translate.py

@@ -27,7 +27,7 @@ def Translate(
         A column of string translations.
     """
 
-    return _translate_impl("snowflake.ml.translate", text, from_language, to_language, session=session)
+    return _translate_impl("snowflake.cortex.translate", text, from_language, to_language, session=session)
 
 
 def _translate_impl(

snowflake/ml/_internal/env_utils.py

@@ -4,6 +4,7 @@ import pathlib
 import re
 import textwrap
 import warnings
+from enum import Enum
 from importlib import metadata as importlib_metadata
 from typing import Any, DefaultDict, Dict, List, Optional, Tuple
 
@@ -18,10 +19,22 @@ from snowflake.ml._internal.exceptions import (
 )
 from snowflake.ml._internal.utils import query_result_checker
 from snowflake.snowpark import session
+from snowflake.snowpark._internal import utils as snowpark_utils
+
+
+class CONDA_OS(Enum):
+    LINUX_64 = "linux-64"
+    LINUX_AARCH64 = "linux-aarch64"
+    OSX_64 = "osx-64"
+    OSX_ARM64 = "osx-arm64"
+    WIN_64 = "win-64"
+    NO_ARCH = "noarch"
+
 
 _SNOWFLAKE_CONDA_CHANNEL_URL = "https://repo.anaconda.com/pkgs/snowflake"
 _NODEFAULTS = "nodefaults"
 _INFO_SCHEMA_PACKAGES_HAS_RUNTIME_VERSION: Optional[bool] = None
+_SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
 _SNOWFLAKE_CONDA_PACKAGE_CACHE: Dict[str, List[version.Version]] = {}
 
 DEFAULT_CHANNEL_NAME = ""
@@ -217,6 +230,7 @@ def get_local_installed_version_of_pip_package(pip_req: requirements.Requirement
         warnings.warn(
             f"Package requirement {str(pip_req)} specified, while version {local_dist_version} is installed. "
             "Local version will be ignored to conform to package requirement.",
+            stacklevel=2,
             category=UserWarning,
         )
         return pip_req
@@ -265,10 +279,58 @@ def _check_runtime_version_column_existence(session: session.Session) -> bool:
     return result == 1
 
 
-def validate_requirements_in_snowflake_conda_channel(
+def get_matched_package_versions_in_snowflake_conda_channel(
+    req: requirements.Requirement,
+    python_version: str = snowml_env.PYTHON_VERSION,
+    conda_os: CONDA_OS = CONDA_OS.LINUX_64,
+) -> List[version.Version]:
+    """Search the snowflake anaconda channel for packages that matches the specifier. Note that this will be the
+    source of truth for checking whether a package indeed exists in Snowflake conda channel.
+
+    Given that a package comes in different architectures, we only check for the Linux x86_64 architecture and assume
+    the package exists in other architectures. If such an assumption does not hold true for a certain package, the
+    caller should specify the architecture to search for.
+
+    Args:
+        req: Requirement specifier.
+        python_version: A string of python version where model is run.
+        conda_os: Specified platform to search availability of the package.
+
+    Returns:
+        List of package versions that meet the requirement specifier.
+    """
+    # Move the retryable_http import here as when UDF import this file, it won't have the "requests" dependency.
+    from snowflake.ml._internal.utils import retryable_http
+
+    assert not snowpark_utils.is_in_stored_procedure()  # type: ignore[no-untyped-call]
+
+    url = f"{_SNOWFLAKE_CONDA_CHANNEL_URL}/{conda_os.value}/repodata.json"
+
+    if req.name not in _SNOWFLAKE_CONDA_PACKAGE_CACHE:
+        http_client = retryable_http.get_http_client()
+        parsed_python_version = version.Version(python_version)
+        python_version_build_str = f"py{parsed_python_version.major}{parsed_python_version.minor}"
+        repodata = http_client.get(url).json()
+        assert isinstance(repodata, dict)
+        packages_info = repodata["packages"]
+        assert isinstance(packages_info, dict)
+        version_list = [
+            version.parse(package_info["version"])
+            for package_info in packages_info.values()
+            if package_info["name"] == req.name and python_version_build_str in package_info["build"]
+        ]
+        _SNOWFLAKE_CONDA_PACKAGE_CACHE[req.name] = version_list
+
+    matched_versions = list(req.specifier.filter(set(_SNOWFLAKE_CONDA_PACKAGE_CACHE.get(req.name, []))))
+    return matched_versions
+
+
+def validate_requirements_in_information_schema(
     session: session.Session, reqs: List[requirements.Requirement], python_version: str
 ) -> Optional[List[str]]:
-    """Search the snowflake anaconda channel for packages with version meet the specifier.
+    """Look up the information_schema table to check if a package with the specified specifier exists in the Snowflake
+    Conda channel. Note that this is not the source of truth due to the potential delay caused by a package that might
+    exist in the information_schema table but has not yet become available in the Snowflake Conda channel.
 
     Args:
         session: Snowflake connection session.
@@ -285,7 +347,7 @@ def validate_requirements_in_snowflake_conda_channel(
     ret_list = []
     reqs_to_request = []
     for req in reqs:
-        if req.name not in _SNOWFLAKE_CONDA_PACKAGE_CACHE:
+        if req.name not in _SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE:
             reqs_to_request.append(req)
     if reqs_to_request:
         pkg_names_str = " OR ".join(
@@ -326,13 +388,13 @@ def validate_requirements_in_snowflake_conda_channel(
             for row in result:
                 req_name = row["PACKAGE_NAME"]
                 req_ver = version.parse(row["VERSION"])
-                cached_req_ver_list = _SNOWFLAKE_CONDA_PACKAGE_CACHE.get(req_name, [])
+                cached_req_ver_list = _SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE.get(req_name, [])
                 cached_req_ver_list.append(req_ver)
-                _SNOWFLAKE_CONDA_PACKAGE_CACHE[req_name] = cached_req_ver_list
+                _SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE[req_name] = cached_req_ver_list
         except snowflake.connector.DataError:
             return None
     for req in reqs:
-        available_versions = list(req.specifier.filter(set(_SNOWFLAKE_CONDA_PACKAGE_CACHE.get(req.name, []))))
+        available_versions = list(req.specifier.filter(set(_SNOWFLAKE_INFO_SCHEMA_PACKAGE_CACHE.get(req.name, []))))
         if not available_versions:
             return None
         else:

snowflake/ml/_internal/file_utils.py

@@ -28,6 +28,7 @@ import cloudpickle
 
 from snowflake import snowpark
 from snowflake.ml._internal.exceptions import exceptions
+from snowflake.snowpark import exceptions as snowpark_exceptions
 
 GENERATED_PY_FILE_EXT = (".pyc", ".pyo", ".pyd", ".pyi")
 
@@ -286,8 +287,16 @@ def stage_file_exists(
         return False
 
 
+def _retry_on_sql_error(exception: Exception) -> bool:
+    return isinstance(exception, snowpark_exceptions.SnowparkSQLException)
+
+
 def upload_directory_to_stage(
-    session: snowpark.Session, local_path: pathlib.Path, stage_path: pathlib.PurePosixPath
+    session: snowpark.Session,
+    local_path: pathlib.Path,
+    stage_path: pathlib.PurePosixPath,
+    *,
+    statement_params: Optional[Dict[str, Any]] = None,
 ) -> None:
     """Upload a local folder recursively to a stage and keep the structure.
 
@@ -295,7 +304,10 @@ def upload_directory_to_stage(
         session: Snowpark Session.
         local_path: Local path to upload.
         stage_path: Base path in the stage.
+        statement_params: Statement Params.
     """
+    import retrying
+
     file_operation = snowpark.FileOperation(session=session)
 
     for root, _, filenames in os.walk(local_path):
@@ -305,16 +317,26 @@ def upload_directory_to_stage(
             stage_dir_path = (
                 stage_path / pathlib.PurePosixPath(local_file_path.relative_to(local_path).as_posix()).parent
             )
-            file_operation.put(
+            retrying.retry(
+                retry_on_exception=_retry_on_sql_error,
+                stop_max_attempt_number=5,
+                wait_exponential_multiplier=100,
+                wait_exponential_max=10000,
+            )(file_operation.put)(
                 str(local_file_path),
                 str(stage_dir_path),
                 auto_compress=False,
                 overwrite=False,
+                statement_params=statement_params,
             )
 
 
 def download_directory_from_stage(
-    session: snowpark.Session, stage_path: pathlib.PurePosixPath, local_path: pathlib.Path
+    session: snowpark.Session,
+    stage_path: pathlib.PurePosixPath,
+    local_path: pathlib.Path,
+    *,
+    statement_params: Optional[Dict[str, Any]] = None,
 ) -> None:
     """Upload a folder in stage recursively to a folder in local and keep the structure.
 
@@ -322,7 +344,10 @@ def download_directory_from_stage(
         session: Snowpark Session.
         stage_path: Stage path to download from.
         local_path: Local path as the base of destination.
+        statement_params: Statement Params.
     """
+    import retrying
+
     file_operation = file_operation = snowpark.FileOperation(session=session)
     file_list = [
         pathlib.PurePosixPath(stage_path.parts[0], *pathlib.PurePosixPath(row.name).parts[1:])
@@ -331,4 +356,9 @@ def download_directory_from_stage(
     for stage_file_path in file_list:
         local_file_dir = local_path / stage_file_path.relative_to(stage_path).parent
         local_file_dir.mkdir(parents=True, exist_ok=True)
-        file_operation.get(str(stage_file_path), str(local_file_dir))
+        retrying.retry(
+            retry_on_exception=_retry_on_sql_error,
+            stop_max_attempt_number=5,
+            wait_exponential_multiplier=100,
+            wait_exponential_max=10000,
+        )(file_operation.get)(str(stage_file_path), str(local_file_dir), statement_params=statement_params)

snowflake/ml/_internal/telemetry.py

@@ -42,6 +42,7 @@ class TelemetryField(enum.Enum):
     NAME = "name"
     # types of telemetry
     TYPE_FUNCTION_USAGE = "function_usage"
+    TYPE_SNOWML_SPCS_USAGE = "snowml_spcs_usage"
     # message keys for telemetry
     KEY_PROJECT = "project"
     KEY_SUBPROJECT = "subproject"
@@ -207,6 +208,23 @@ def suppress_exceptions(func: Callable[..., Any]) -> Callable[..., Any]:
     return wrapper
 
 
+def send_custom_usage(
+    project: str,
+    *,
+    telemetry_type: str,
+    subproject: Optional[str] = None,
+    data: Optional[Dict[str, Any]] = None,
+    **kwargs: Any,
+) -> None:
+    active_session = next(iter(session._get_active_sessions()))
+    assert active_session, "Missing active session object"
+
+    client = _SourceTelemetryClient(conn=active_session._conn._conn, project=project, subproject=subproject)
+    common_metrics = client._create_basic_telemetry_data(telemetry_type=telemetry_type)
+    data = {**common_metrics, TelemetryField.KEY_DATA.value: data, **kwargs}
+    client._send(msg=data)
+
+
 def send_api_usage_telemetry(
     project: str,
     subproject: Optional[str] = None,
@@ -228,7 +246,8 @@ def send_api_usage_telemetry(
     custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, _ReturnValue]]:
     """
-    Decorator that sends API usage telemetry.
+    Decorator that sends API usage telemetry and adds function usage statement parameters to the dataframe returned by
+    the function.
 
     Args:
         project: Project.
@@ -253,6 +272,51 @@ def send_api_usage_telemetry(
         def wrap(*args: Any, **kwargs: Any) -> _ReturnValue:
             params = _get_func_params(func, func_params_to_log, args, kwargs) if func_params_to_log else None
 
+            api_calls: List[Union[Dict[str, Union[Callable[..., Any], str]], Callable[..., Any], str]] = []
+            if api_calls_extractor:
+                extracted_api_calls = api_calls_extractor(args[0])
+                for api_call in extracted_api_calls:
+                    if isinstance(api_call, str):
+                        api_calls.append({TelemetryField.NAME.value: api_call})
+                    elif callable(api_call):
+                        api_calls.append({TelemetryField.NAME.value: _get_full_func_name(api_call)})
+                    else:
+                        api_calls.append(api_call)
+            api_calls.append({TelemetryField.NAME.value: _get_full_func_name(func)})
+
+            sfqids = None
+            if sfqids_extractor:
+                sfqids = sfqids_extractor(args[0])
+
+            statement_params = get_function_usage_statement_params(
+                project=project,
+                subproject=subproject,
+                function_category=TelemetryField.FUNC_CAT_USAGE.value,
+                function_name=_get_full_func_name(func),
+                function_parameters=params,
+                api_calls=api_calls,
+                custom_tags=custom_tags,
+            )
+
+            def update_stmt_params_if_snowpark_df(obj: _ReturnValue, statement_params: Dict[str, Any]) -> _ReturnValue:
+                """
+                Update SnowML function usage statement parameters to the object if it is a Snowpark DataFrame.
+                Used to track APIs returning a Snowpark DataFrame.
+
+                Args:
+                    obj: Object to check and update.
+                    statement_params: Statement parameters.
+
+                Returns:
+                    Updated object.
+                """
+                if isinstance(obj, dataframe.DataFrame):
+                    if hasattr(obj, "_statement_params") and obj._statement_params:
+                        obj._statement_params.update(statement_params)
+                    else:
+                        obj._statement_params = statement_params  # type: ignore[assignment]
+                return obj
+
             # prioritize `conn_attr_name` over the active session
             if conn_attr_name:
                 # raise AttributeError if conn attribute does not exist in `self`
@@ -266,36 +330,20 @@ def send_api_usage_telemetry(
                 # server no default session
                 except snowpark_exceptions.SnowparkSessionException:
                     try:
-                        return func(*args, **kwargs)
+                        return update_stmt_params_if_snowpark_df(func(*args, **kwargs), statement_params)
                     except Exception as e:
                         if isinstance(e, snowml_exceptions.SnowflakeMLException):
-                            e = e.original_exception
+                            raise e.original_exception.with_traceback(e.__traceback__) from None
                         # suppress SnowparkSessionException from telemetry in the stack trace
                         raise e from None
 
                 conn = active_session._conn._conn
                 if (not active_session.telemetry_enabled) or (conn is None):
                     try:
-                        return func(*args, **kwargs)
+                        return update_stmt_params_if_snowpark_df(func(*args, **kwargs), statement_params)
                     except snowml_exceptions.SnowflakeMLException as e:
                         raise e.original_exception from e
 
-            api_calls: List[Dict[str, Any]] = []
-            if api_calls_extractor:
-                extracted_api_calls = api_calls_extractor(args[0])
-                for api_call in extracted_api_calls:
-                    if isinstance(api_call, str):
-                        api_calls.append({TelemetryField.NAME.value: api_call})
-                    elif callable(api_call):
-                        api_calls.append({TelemetryField.NAME.value: _get_full_func_name(api_call)})
-                    else:
-                        api_calls.append(api_call)
-            api_calls.append({TelemetryField.NAME.value: _get_full_func_name(func)})
-
-            sfqids = None
-            if sfqids_extractor:
-                sfqids = sfqids_extractor(args[0])
-
             # TODO(hayu): [SNOW-750287] Optimize telemetry client to a singleton.
             telemetry = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject)
             telemetry_args = dict(
@@ -314,22 +362,24 @@ def send_api_usage_telemetry(
                     if hasattr(e, "_snowflake_ml_handled") and e._snowflake_ml_handled:
                         raise e
                     if isinstance(e, snowpark_exceptions.SnowparkClientException):
-                        e = snowml_exceptions.SnowflakeMLException(
+                        me = snowml_exceptions.SnowflakeMLException(
                             error_code=error_codes.INTERNAL_SNOWPARK_ERROR, original_exception=e
                         )
                     else:
-                        e = snowml_exceptions.SnowflakeMLException(
+                        me = snowml_exceptions.SnowflakeMLException(
                             error_code=error_codes.UNDEFINED, original_exception=e
                         )
-                telemetry_args["error"] = repr(e)
-                telemetry_args["error_code"] = e.error_code
-                e.original_exception._snowflake_ml_handled = True  # type: ignore[attr-defined]
-                if e.suppress_source_trace:
-                    raise e.original_exception from None
                 else:
-                    raise e.original_exception from e
+                    me = e
+                telemetry_args["error"] = repr(me)
+                telemetry_args["error_code"] = me.error_code
+                me.original_exception._snowflake_ml_handled = True  # type: ignore[attr-defined]
+                if me.suppress_source_trace:
+                    raise me.original_exception from None
+                else:
+                    raise me.original_exception from e
             else:
-                return res
+                return update_stmt_params_if_snowpark_df(res, statement_params)
             finally:
                 telemetry.send_function_usage_telemetry(**telemetry_args)
                 global _log_counter
@@ -343,68 +393,6 @@ def send_api_usage_telemetry(
     return decorator
 
 
-def add_stmt_params_to_df(
-    project: str,
-    subproject: Optional[str] = None,
-    *,
-    function_category: str = TelemetryField.FUNC_CAT_USAGE.value,
-    func_params_to_log: Optional[Iterable[str]] = None,
-    api_calls: Optional[
-        List[
-            Union[
-                Dict[str, Union[Callable[..., Any], str]],
-                Union[Callable[..., Any], str],
-            ]
-        ]
-    ] = None,
-    custom_tags: Optional[Dict[str, Union[bool, int, str, float]]] = None,
-) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, _ReturnValue]]:
-    """
-    Decorator that adds function usage statement parameters to the dataframe returned by the function.
-
-    Args:
-        project: Project.
-        subproject: Subproject.
-        function_category: Function category.
-        func_params_to_log: Function parameters to log.
-        api_calls: API calls in the function.
-        custom_tags: Custom tags.
-
-    Returns:
-        Decorator that adds function usage statement parameters to the dataframe returned by the decorated function.
-    """
-
-    def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, _ReturnValue]:
-        @functools.wraps(func)
-        def wrap(*args: Any, **kwargs: Any) -> _ReturnValue:
-            params = _get_func_params(func, func_params_to_log, args, kwargs) if func_params_to_log else None
-            statement_params = get_function_usage_statement_params(
-                project=project,
-                subproject=subproject,
-                function_category=function_category,
-                function_name=_get_full_func_name(func),
-                function_parameters=params,
-                api_calls=api_calls,
-                custom_tags=custom_tags,
-            )
-
-            try:
-                res = func(*args, **kwargs)
-                if isinstance(res, dataframe.DataFrame):
-                    if hasattr(res, "_statement_params") and res._statement_params:
-                        res._statement_params.update(statement_params)
-                    else:
-                        res._statement_params = statement_params  # type: ignore[assignment]
-            except Exception:
-                raise
-            else:
-                return res
-
-        return cast(Callable[_Args, _ReturnValue], wrap)
-
-    return decorator
-
-
 def _get_full_func_name(func: Callable[..., Any]) -> str:
     """
     Get the full function name with module and qualname.