snowflake-ml-python 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

snowflake/ml/model/_deploy_client/utils/snowservice_client.py

@@ -0,0 +1,178 @@
+import json
+import logging
+import time
+from typing import Optional
+
+from snowflake.ml.model._deploy_client.utils import constants
+from snowflake.snowpark import Session
+
+
+class SnowServiceClient:
+    """
+    SnowService client implementation: a Python wrapper for SnowService SQL queries.
+    """
+
+    def __init__(self, session: Session) -> None:
+        """Initialization
+
+        Args:
+            session: Snowpark session
+        """
+        self.session = session
+
+    def create_or_replace_service(
+        self,
+        service_name: str,
+        compute_pool: str,
+        spec_stage_location: str,
+        *,
+        min_instances: int = 1,
+        max_instances: int = 1,
+    ) -> None:
+        """Create or replace service. Since SnowService doesn't support the CREATE OR REPLACE service syntax, we will
+        first attempt to drop the service if it exists, and then create the service. Please note that this approach may
+        have side effects due to the lack of transaction support.
+
+        Args:
+            service_name: Name of the service.
+            min_instances: Minimum number of service replicas.
+            max_instances: Maximum number of service replicas.
+            compute_pool: Name of the compute pool.
+            spec_stage_location: Stage path for the service spec.
+        """
+        self._drop_service_if_exists(service_name)
+        sql = f"""
+             CREATE SERVICE {service_name}
+                 MIN_INSTANCES={min_instances}
+                 MAX_INSTANCES={max_instances}
+                 COMPUTE_POOL={compute_pool}
+                 SPEC=@{spec_stage_location}
+         """
+        logging.info(f"Create service with SQL: \n {sql}")
+        self.session.sql(sql).collect()
+
+    def _drop_service_if_exists(self, service_name: str) -> None:
+        """Drop service if it already exists.
+
+        Args:
+            service_name: Name of the service.
+        """
+        self.session.sql(f"DROP SERVICE IF EXISTS {service_name}").collect()
+
+    def create_or_replace_service_function(
+        self,
+        service_func_name: str,
+        service_name: str,
+        *,
+        endpoint_name: str = constants.PREDICT,
+        path_at_service_endpoint: str = constants.PREDICT,
+    ) -> None:
+        """Create or replace service function.
+
+        Args:
+            service_func_name: Name of the service function.
+            service_name: Name of the service.
+            endpoint_name: Name the service endpoint, declared in the service spec, indicating the listening port.
+            path_at_service_endpoint: Specify the path/route at the service endpoint. Multiple paths can exist for a
+                given endpoint. For example, an inference server listening on port 5000 may have paths like "/predict"
+                and "/monitoring
+
+        """
+        sql = f"""
+            CREATE OR REPLACE FUNCTION {service_func_name}(input OBJECT)
+                RETURNS OBJECT
+                SERVICE={service_name}
+                ENDPOINT={endpoint_name}
+                AS '/{path_at_service_endpoint}'
+            """
+        logging.info(f"Create service function with SQL: \n {sql}")
+        self.session.sql(sql).collect()
+
+    def block_until_resource_is_ready(
+        self,
+        resource_name: str,
+        resource_type: constants.ResourceType,
+        *,
+        max_retries: int = 60,
+        retry_interval_secs: int = 5,
+    ) -> None:
+        """Blocks execution until the specified resource is ready.
+        Note that this is a best-effort approach because when launching a service, it's possible for it to initially
+        fail due to a system error. However, SnowService may automatically retry and recover the service, leading to
+        potential false-negative information.
+
+        Args:
+            resource_name: Name of the resource.
+            resource_type: Type of the resource.
+            max_retries: The maximum number of retries to check the resource readiness (default: 60).
+            retry_interval_secs: The number of seconds to wait between each retry (default: 5).
+
+        Raises:
+            RuntimeError: If the resource received the following status [failed, not_found, internal_error, deleting]
+            RuntimeError: If the resource does not reach the ready/done state within the specified number of retries.
+        """
+        for _ in range(max_retries):
+            status = self.get_resource_status(resource_name=resource_name, resource_type=resource_type)
+            if status in [constants.ResourceStatus.READY, constants.ResourceStatus.DONE]:
+                return
+            elif status in [
+                constants.ResourceStatus.FAILED,
+                constants.ResourceStatus.NOT_FOUND,
+                constants.ResourceStatus.INTERNAL_ERROR,
+                constants.ResourceStatus.DELETING,
+            ]:
+                error_log = self.get_resource_log(
+                    resource_name=resource_name,
+                    resource_type=resource_type,
+                    container_name=constants.INFERENCE_SERVER_CONTAINER,
+                )
+                raise RuntimeError(f"{resource_type} {resource_name} failed. \n {error_log if error_log else ''}")
+            time.sleep(retry_interval_secs)
+
+        raise RuntimeError("Resource never reached the ready/done state.")
+
+    def get_resource_log(
+        self, resource_name: str, resource_type: constants.ResourceType, container_name: str
+    ) -> Optional[str]:
+        if resource_type != constants.ResourceType.SERVICE:
+            raise NotImplementedError(f"{resource_type.name} is not yet supported in get_resource_log function")
+        try:
+            row = self.session.sql(
+                f"CALL SYSTEM$GET_SNOWSERVICE_LOGS('{resource_name}', '0', '{container_name}')"
+            ).collect()
+            return str(row[0]["SYSTEM$GET_SNOWSERVICE_LOGS"])
+        except Exception:
+            return None
+
+    def get_resource_status(
+        self, resource_name: str, resource_type: constants.ResourceType
+    ) -> Optional[constants.ResourceStatus]:
+        """Get resource status.
+
+        Args:
+            resource_name: Name of the resource.
+            resource_type: Type of the resource.
+
+        Raises:
+            ValueError: If resource type does not have a corresponding system function for querying status.
+            RuntimeError: If corresponding status call failed.
+
+        Returns:
+            Optional[constants.ResourceStatus]: The status of the resource, or None if the resource status is empty.
+        """
+        if resource_type not in constants.RESOURCE_TO_STATUS_FUNCTION_MAPPING:
+            raise ValueError(f"Status querying is not supported for resources of type '{resource_type}'.")
+        status_func = constants.RESOURCE_TO_STATUS_FUNCTION_MAPPING[resource_type]
+        try:
+            row = self.session.sql(f"CALL {status_func}('{resource_name}');").collect()
+        except Exception as e:
+            raise RuntimeError(f"Error while querying the {resource_type} {resource_name} status: {str(e)}")
+        resource_metadata = json.loads(row[0][status_func])[0]
+        logging.info(f"Resource status metadata: {resource_metadata}")
+        if resource_metadata and resource_metadata["status"]:
+            try:
+                status = resource_metadata["status"]
+                return constants.ResourceStatus(status)
+            except ValueError:
+                logging.warning(f"Unknown status returned: {status}")
+        return None

snowflake/ml/model/_deploy_client/warehouse/deploy.py

@@ -1,4 +1,5 @@
 import os
+import posixpath
 import tempfile
 import warnings
 from types import ModuleType
@@ -6,7 +7,8 @@ from typing import IO, List, Optional, Tuple, TypedDict, Union
 
 from typing_extensions import Unpack
 
-from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml._internal import env_utils, file_utils
+from snowflake.ml._internal.utils import identifier
 from snowflake.ml.model import (
     _env as model_env,
     _model,
@@ -37,6 +39,7 @@ def _deploy_to_warehouse(
         **kwargs: Options that control some features in generated udf code.
 
     Raises:
+        ValueError: Raised when model file name is unable to encoded using ASCII.
         ValueError: Raised when incompatible model.
         ValueError: Raised when target method does not exist in model.
         ValueError: Raised when confronting invalid stage location.
@@ -44,14 +47,20 @@ def _deploy_to_warehouse(
     Returns:
         The metadata of the model deployed.
     """
+    # TODO(SNOW-862576): Should remove check on ASCII encoding after SNOW-862576 fixed.
     if model_dir_path:
         model_dir_path = os.path.normpath(model_dir_path)
         model_dir_name = os.path.basename(model_dir_path)
+        if not file_utils._able_ascii_encode(model_dir_name):
+            raise ValueError(f"Model file name {model_dir_name} cannot be encoded using ASCII. Please rename.")
         extract_model_code = infer_template._EXTRACT_LOCAL_MODEL_CODE.format(model_dir_name=model_dir_name)
         meta = _model.load_model(model_dir_path=model_dir_path, meta_only=True)
     else:
         assert model_stage_file_path is not None, "Unreachable assertion error."
-        model_stage_file_name = os.path.basename(model_stage_file_path)
+        model_stage_file_name = posixpath.basename(model_stage_file_path)
+        if not file_utils._able_ascii_encode(model_stage_file_name):
+            raise ValueError(f"Model file name {model_stage_file_name} cannot be encoded using ASCII. Please rename.")
+
         extract_model_code = infer_template._EXTRACT_STAGE_MODEL_CODE.format(
             model_stage_file_name=model_stage_file_name
         )
@@ -59,32 +68,26 @@ def _deploy_to_warehouse(
 
     relax_version = kwargs.get("relax_version", False)
 
+    disable_local_conda_resolver = kwargs.get("disable_local_conda_resolver", False)
+
     if target_method not in meta.signatures.keys():
         raise ValueError(f"Target method {target_method} does not exist in model.")
 
-    _use_local_snowml = kwargs.get("_use_local_snowml", False)
-
     final_packages = _get_model_final_packages(
-        meta, session, relax_version=relax_version, _use_local_snowml=_use_local_snowml
+        meta, session, relax_version=relax_version, disable_local_conda_resolver=disable_local_conda_resolver
     )
 
     stage_location = kwargs.get("permanent_udf_stage_location", None)
     if stage_location:
-        stage_location = stage_location.strip().rstrip("/")
+        stage_location = posixpath.normpath(stage_location.strip())
         if not stage_location.startswith("@"):
             raise ValueError(f"Invalid stage location {stage_location}.")
 
-    _snowml_wheel_path = None
-    if _use_local_snowml:
-        _snowml_wheel_path = file_utils.upload_snowml(session, stage_location=stage_location)
-
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
         _write_UDF_py_file(f.file, extract_model_code, target_method, **kwargs)
         print(f"Generated UDF file is persisted at: {f.name}")
-        imports = (
-            ([model_dir_path] if model_dir_path else [])
-            + ([model_stage_file_path] if model_stage_file_path else [])
-            + ([_snowml_wheel_path] if _snowml_wheel_path else [])
+        imports = ([model_dir_path] if model_dir_path else []) + (
+            [model_stage_file_path] if model_stage_file_path else []
         )
 
         class _UDFParams(TypedDict):
@@ -99,7 +102,7 @@ def _deploy_to_warehouse(
         params = _UDFParams(
             file_path=f.name,
             func_name="infer",
-            name=f"{udf_name}",
+            name=identifier.get_inferred_name(udf_name),
             return_type=st.PandasSeriesType(st.MapType(st.StringType(), st.VariantType())),
             input_types=[st.PandasDataFrameType([st.MapType()])],
             imports=list(imports),
@@ -139,6 +142,7 @@ def _write_UDF_py_file(
         extract_model_code=extract_model_code,
         keep_order_code=infer_template._KEEP_ORDER_CODE_TEMPLATE if keep_order else "",
         target_method=target_method,
+        code_dir_name=_model_meta.ModelMetadata.MODEL_CODE_DIR,
     )
     f.write(udf_code)
     f.flush()
@@ -148,7 +152,7 @@ def _get_model_final_packages(
     meta: _model_meta.ModelMetadata,
     session: snowpark_session.Session,
     relax_version: Optional[bool] = False,
-    _use_local_snowml: Optional[bool] = False,
+    disable_local_conda_resolver: Optional[bool] = False,
 ) -> List[str]:
     """Generate final packages list of dependency of a model to be deployed to warehouse.
 
@@ -157,7 +161,8 @@ def _get_model_final_packages(
         session: Snowpark connection session.
         relax_version: Whether or not relax the version restriction when fail to resolve dependencies.
             Defaults to False.
-        _use_local_snowml: Flag to indicate if using local SnowML code as execution library
+        disable_local_conda_resolver: Set to disable use local conda resolver to do pre-check on environment and rely on
+            the information schema only. Defaults to False.
 
     Raises:
         RuntimeError: Raised when PIP requirements and dependencies from non-Snowflake anaconda channel found.
@@ -174,18 +179,10 @@ def _get_model_final_packages(
         raise RuntimeError("PIP requirements and dependencies from non-Snowflake anaconda channel is not supported.")
 
     deps = meta._conda_dependencies[""]
-    if _use_local_snowml:
-        local_snowml_version = snowml_env.VERSION
-        snowml_dept = next((dep for dep in deps if dep.name == env_utils._SNOWML_PKG_NAME), None)
-        if snowml_dept:
-            if not snowml_dept.specifier.contains(local_snowml_version) and not relax_version:
-                raise RuntimeError(
-                    "Incompatible snowflake-ml-python-version is found. "
-                    + f"Require {snowml_dept.specifier}, got {local_snowml_version}."
-                )
-            deps.remove(snowml_dept)
 
     try:
+        if disable_local_conda_resolver:
+            raise ImportError("Raise to disable local conda resolver. Should be captured.")
         final_packages = env_utils.resolve_conda_environment(
             deps, [model_env._SNOWFLAKE_CONDA_CHANNEL_URL], python_version=meta.python_version
         )

snowflake/ml/model/_deploy_client/warehouse/infer_template.py

@@ -48,18 +48,21 @@ class FileLock:
 IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
 import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]
 
-from snowflake.ml.model import _model
-
 {extract_model_code}
 
+sys.path.insert(0, os.path.join(extracted_model_dir_path, "{code_dir_name}"))
+from snowflake.ml.model import _model
 model, meta = _model._load_model_for_deploy(extracted_model_dir_path)
 
+features = meta.signatures["{target_method}"].inputs
+input_cols = [feature.name for feature in features]
+dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
+
 # TODO(halu): Wire `max_batch_size`.
 # TODO(halu): Avoid per batch async detection branching.
 @vectorized(input=pd.DataFrame, max_batch_size=10)
 def infer(df):
-    input_cols = [spec.name for spec in meta.signatures["{target_method}"].inputs]
-    input_df = pd.io.json.json_normalize(df[0])
+    input_df = pd.io.json.json_normalize(df[0]).astype(dtype=dtype_map)
     if inspect.iscoroutinefunction(model.{target_method}):
         predictions_df = anyio.run(model.{target_method}, input_df[input_cols])
     else:

snowflake/ml/model/_deployer.py

@@ -1,9 +1,7 @@
-import json
 import traceback
 from enum import Enum
 from typing import Optional, TypedDict, Union, overload
 
-import numpy as np
 import pandas as pd
 from typing_extensions import Required
 
@@ -184,7 +182,6 @@ def predict(
 
     Raises:
         ValueError: Raised when the input is too large to use keep_order option.
-        NotImplementedError: FeatureGroupSpec is not supported.
 
     Returns:
         The output dataframe.
@@ -199,19 +196,19 @@ def predict(
     # Validate and prepare input
     if not isinstance(X, SnowparkDataFrame):
         df = model_signature._convert_and_validate_local_data(X, sig.inputs)
-        s_df = session.create_dataframe(df)
+        s_df = model_signature._SnowparkDataFrameHandler.convert_from_df(session, df, keep_order=keep_order)
     else:
         model_signature._validate_snowpark_data(X, sig.inputs)
         s_df = X
 
-    if keep_order:
-        # ID is UINT64 type, this we should limit.
-        if s_df.count() > 2**64:
-            raise ValueError("Unable to keep order of a DataFrame with more than 2 ** 64 rows.")
-        s_df = s_df.with_column(
-            infer_template._KEEP_ORDER_COL_NAME,
-            F.monotonically_increasing_id(),
-        )
+        if keep_order:
+            # ID is UINT64 type, this we should limit.
+            if s_df.count() > 2**64:
+                raise ValueError("Unable to keep order of a DataFrame with more than 2 ** 64 rows.")
+            s_df = s_df.with_column(
+                infer_template._KEEP_ORDER_COL_NAME,
+                F.monotonically_increasing_id(),
+            )
 
     # Infer and get intermediate result
     input_cols = []
@@ -223,7 +220,9 @@ def predict(
                 F.col(col_name),
             ]
         )
-    output_obj = F.call_udf(deployment["name"], F.object_construct(*input_cols))  # type:ignore[arg-type]
+    output_obj = F.call_udf(
+        identifier.get_inferred_name(deployment["name"]), F.object_construct(*input_cols)  # type:ignore[arg-type]
+    )
     if output_with_input_features:
         df_res = s_df.with_column(INTERMEDIATE_OBJ_NAME, output_obj)
     else:
@@ -243,24 +242,12 @@ def predict(
         output_cols.append(F.col(INTERMEDIATE_OBJ_NAME)[output_feature.name].astype(output_feature.as_snowpark_type()))
 
     df_res = df_res.with_columns(
-        [identifier.quote_name_without_upper_casing(output_feature.name) for output_feature in sig.outputs],
+        [identifier.get_inferred_name(output_feature.name) for output_feature in sig.outputs],
         output_cols,
     ).drop(INTERMEDIATE_OBJ_NAME)
 
     # Get final result
     if not isinstance(X, SnowparkDataFrame):
-        dtype_map = {}
-        for feature in sig.outputs:
-            if isinstance(feature, model_signature.FeatureGroupSpec):
-                raise NotImplementedError("FeatureGroupSpec is not supported.")
-            assert isinstance(feature, model_signature.FeatureSpec), "Invalid feature kind."
-            dtype_map[feature.name] = feature.as_dtype()
-        df_local = df_res.to_pandas()
-        # This is because Array and object will generate variant type and requires an additional loads to
-        # get correct data otherwise it would be string.
-        for col_name in [col_name for col_name, col_dtype in dtype_map.items() if col_dtype == np.object0]:
-            df_local[col_name] = df_local[col_name].map(json.loads)
-        df_local = df_local.astype(dtype=dtype_map)
-        return pd.DataFrame(df_local)
+        return model_signature._SnowparkDataFrameHandler.convert_to_df(df_res, features=sig.outputs)
     else:
         return df_res

snowflake/ml/model/_env.py

@@ -36,7 +36,7 @@ def save_conda_env_file(
     for chan, reqs in deps.items():
         env["dependencies"].extend([f"{chan}::{str(req)}" if chan else str(req) for req in reqs])
 
-    with open(path, "w") as f:
+    with open(path, "w", encoding="utf-8") as f:
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
     return path
@@ -54,7 +54,7 @@ def save_requirements_file(dir_path: str, pip_deps: List[requirements.Requiremen
     """
     requirements = "\n".join(map(str, pip_deps))
     path = os.path.join(dir_path, _REQUIREMENTS_FILE_NAME)
-    with open(path, "w") as out:
+    with open(path, "w", encoding="utf-8") as out:
         out.write(requirements)
 
     return path
@@ -69,7 +69,7 @@ def load_conda_env_file(path: str) -> Tuple[DefaultDict[str, List[requirements.R
     Returns:
         A tuple of Dict of conda dependencies after validated and a string 'major.minor.patchlevel' of python version.
     """
-    with open(path) as f:
+    with open(path, encoding="utf-8") as f:
         env = yaml.safe_load(stream=f)
 
     assert isinstance(env, dict)
@@ -99,7 +99,7 @@ def load_requirements_file(path: str) -> List[requirements.Requirement]:
     Returns:
         List of dependencies string after validated.
     """
-    with open(path) as f:
+    with open(path, encoding="utf-8") as f:
         reqs = f.readlines()
 
     return env_utils.validate_pip_requirement_string_list(reqs)

snowflake/ml/model/_handlers/_base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Generic, Optional
 
-from typing_extensions import TypeGuard
+from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml.model import _model_meta, type_hints as model_types
 
@@ -43,6 +43,7 @@ class _ModelHandler(ABC, Generic[model_types._ModelType]):
         model_blobs_dir_path: str,
         sample_input: Optional[model_types.SupportedDataType] = None,
         is_sub_model: Optional[bool] = False,
+        **kwargs: Unpack[model_types.ModelSaveOption],
     ) -> None:
         """Save the model.
 
@@ -53,6 +54,7 @@ class _ModelHandler(ABC, Generic[model_types._ModelType]):
             model_blobs_dir_path: Directory path to the model.
             sample_input: Sample input to infer the signatures from.
             is_sub_model: Flag to show if it is a sub model, a sub model does not need signature.
+            kwargs: Additional saving options.
         """
         ...
 

snowflake/ml/model/_handlers/custom.py

@@ -1,16 +1,19 @@
 import inspect
 import os
+import pathlib
 import sys
 from typing import TYPE_CHECKING, Dict, Optional
 
 import anyio
 import cloudpickle
+import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import file_utils, type_utils
 from snowflake.ml.model import (
     _model_handler,
     _model_meta as model_meta_api,
+    model_signature,
     type_hints as model_types,
 )
 from snowflake.ml.model._handlers import _base
@@ -55,6 +58,10 @@ class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
             target_method = getattr(model, target_method_name, None)
             assert callable(target_method) and inspect.ismethod(target_method)
             target_method = target_method.__func__
+
+            if not isinstance(sample_input, pd.DataFrame):
+                sample_input = model_signature._convert_local_data_to_df(sample_input)
+
             if inspect.iscoroutinefunction(target_method):
                 with anyio.start_blocking_portal() as portal:
                     predictions_df = portal.call(target_method, model, sample_input)
@@ -102,7 +109,9 @@ class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
             model_type=_CustomModelHandler.handler_type,
             path=_CustomModelHandler.MODEL_BLOB_FILE,
             artifacts={
-                name: os.path.join(_CustomModelHandler.MODEL_ARTIFACTS_DIR, os.path.basename(os.path.normpath(uri)))
+                name: pathlib.Path(
+                    os.path.join(_CustomModelHandler.MODEL_ARTIFACTS_DIR, os.path.basename(os.path.normpath(path=uri)))
+                ).as_posix()
                 for name, uri in model.context.artifacts.items()
             },
         )
@@ -129,7 +138,10 @@ class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
         assert issubclass(ModelClass, custom_model.CustomModel)
 
         artifacts_meta = model_blob_metadata.artifacts
-        artifacts = {name: os.path.join(model_blob_path, rel_path) for name, rel_path in artifacts_meta.items()}
+        artifacts = {
+            name: str(pathlib.PurePath(model_blob_path) / pathlib.PurePosixPath(rel_path))
+            for name, rel_path in artifacts_meta.items()
+        }
         models: Dict[str, model_types.SupportedModelType] = dict()
         for sub_model_name, _ref in m.context.model_refs.items():
             model_type = model_meta.models[sub_model_name].model_type