snowflake-ml-python 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

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

@@ -70,13 +70,16 @@ class SnowServiceClient:
         logger.debug(f"Create service with SQL: \n {sql}")
         self.session.sql(sql).collect()
 
-    def create_job(self, compute_pool: str, spec_stage_location: str, external_access_integrations: List[str]) -> None:
+    def create_job(
+        self, job_name: str, compute_pool: str, spec_stage_location: str, external_access_integrations: List[str]
+    ) -> None:
         """Execute the job creation SQL command. Note that the job creation is synchronous, hence we execute it in a
         async way so that we can query the log in the meantime.
 
         Upon job failure, full job container log will be logged.
 
         Args:
+            job_name: name of the job
             compute_pool: name of the compute pool
             spec_stage_location: path to the stage location where the spec is located at.
             external_access_integrations: EAIs for network connection.
@@ -84,19 +87,18 @@ class SnowServiceClient:
         stage, path = uri.get_stage_and_path(spec_stage_location)
         sql = textwrap.dedent(
             f"""
-            EXECUTE SERVICE
+            EXECUTE JOB SERVICE
             IN COMPUTE POOL {compute_pool}
             FROM {stage}
-            SPEC = '{path}'
+            SPECIFICATION_FILE = '{path}'
+            NAME = {job_name}
             EXTERNAL_ACCESS_INTEGRATIONS = ({', '.join(external_access_integrations)})
             """
         )
         logger.debug(f"Create job with SQL: \n {sql}")
-        cur = self.session._conn._conn.cursor()
-        cur.execute_async(sql)
-        job_id = cur._sfqid
+        self.session.sql(sql).collect_nowait()
         self.block_until_resource_is_ready(
-            resource_name=str(job_id),
+            resource_name=job_name,
             resource_type=constants.ResourceType.JOB,
             container_name=constants.KANIKO_CONTAINER_NAME,
             max_retries=240,
@@ -182,10 +184,7 @@ class SnowServiceClient:
         """
         assert resource_type == constants.ResourceType.SERVICE or resource_type == constants.ResourceType.JOB
         query_command = ""
-        if resource_type == constants.ResourceType.SERVICE:
-            query_command = f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
-        elif resource_type == constants.ResourceType.JOB:
-            query_command = f"CALL SYSTEM$GET_JOB_LOGS('{resource_name}', '{container_name}')"
+        query_command = f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
         logger.warning(
             f"Best-effort log streaming from SPCS will be enabled when python logging level is set to INFO."
             f"Alternatively, you can also query the logs by running the query '{query_command}'"
@@ -201,7 +200,7 @@ class SnowServiceClient:
                 )
                 lsp.process_new_logs(resource_log, log_level=logging.INFO)
 
-            status = self.get_resource_status(resource_name=resource_name, resource_type=resource_type)
+            status = self.get_resource_status(resource_name=resource_name)
 
             if resource_type == constants.ResourceType.JOB and status == constants.ResourceStatus.DONE:
                 return
@@ -246,52 +245,24 @@ class SnowServiceClient:
     def get_resource_log(
         self, resource_name: str, resource_type: constants.ResourceType, container_name: str
     ) -> Optional[str]:
-        if resource_type == constants.ResourceType.SERVICE:
-            try:
-                row = self.session.sql(
-                    f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
-                ).collect()
-                return str(row[0]["SYSTEM$GET_SERVICE_LOGS"])
-            except Exception:
-                return None
-        elif resource_type == constants.ResourceType.JOB:
-            try:
-                row = self.session.sql(f"CALL SYSTEM$GET_JOB_LOGS('{resource_name}', '{container_name}')").collect()
-                return str(row[0]["SYSTEM$GET_JOB_LOGS"])
-            except Exception:
-                return None
-        else:
-            raise snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.NOT_IMPLEMENTED,
-                original_exception=NotImplementedError(
-                    f"{resource_type.name} is not yet supported in get_resource_log function"
-                ),
-            )
-
-    def get_resource_status(
-        self, resource_name: str, resource_type: constants.ResourceType
-    ) -> Optional[constants.ResourceStatus]:
+        try:
+            row = self.session.sql(
+                f"CALL SYSTEM$GET_SERVICE_LOGS('{resource_name}', '0', '{container_name}')"
+            ).collect()
+            return str(row[0]["SYSTEM$GET_SERVICE_LOGS"])
+        except Exception:
+            return None
+
+    def get_resource_status(self, resource_name: str) -> Optional[constants.ResourceStatus]:
         """Get resource status.
 
         Args:
             resource_name: Name of the resource.
-            resource_type: Type of the resource.
-
-        Raises:
-            SnowflakeMLException: If resource type does not have a corresponding system function for querying status.
-            SnowflakeMLException: 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 snowml_exceptions.SnowflakeMLException(
-                error_code=error_codes.INVALID_ARGUMENT,
-                original_exception=ValueError(
-                    f"Status querying is not supported for resources of type '{resource_type}'."
-                ),
-            )
-        status_func = constants.RESOURCE_TO_STATUS_FUNCTION_MAPPING[resource_type]
+        status_func = "SYSTEM$GET_SERVICE_STATUS"
         try:
             row = self.session.sql(f"CALL {status_func}('{resource_name}');").collect()
         except Exception:

snowflake/ml/model/_model_composer/model_composer.py

@@ -8,8 +8,10 @@ from typing import Any, Dict, List, Optional
 
 from absl import logging
 from packaging import requirements
+from typing_extensions import deprecated
 
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml._internal.lineage import data_source
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._model_composer.model_manifest import model_manifest
 from snowflake.ml.model._packager import model_packager
@@ -134,6 +136,7 @@ class ModelComposer:
             model_meta=self.packager.meta,
             model_file_rel_path=pathlib.PurePosixPath(self.model_file_rel_path),
             options=options,
+            data_sources=self._get_data_sources(model),
         )
 
         file_utils.upload_directory_to_stage(
@@ -143,7 +146,8 @@ class ModelComposer:
             statement_params=self._statement_params,
         )
 
-    def load(
+    @deprecated("Only used by PrPr model registry. Use static method version of load instead.")
+    def legacy_load(
         self,
         *,
         meta_only: bool = False,
@@ -163,3 +167,20 @@ class ModelComposer:
         with zipfile.ZipFile(self.model_local_path, mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
             zf.extractall(path=self._packager_workspace_path)
         self.packager.load(meta_only=meta_only, options=options)
+
+    @staticmethod
+    def load(
+        workspace_path: pathlib.Path,
+        *,
+        meta_only: bool = False,
+        options: Optional[model_types.ModelLoadOption] = None,
+    ) -> model_packager.ModelPackager:
+        mp = model_packager.ModelPackager(str(workspace_path / ModelComposer.MODEL_DIR_REL_PATH))
+        mp.load(meta_only=meta_only, options=options)
+        return mp
+
+    def _get_data_sources(self, model: model_types.SupportedModelType) -> Optional[List[data_source.DataSource]]:
+        data_sources = getattr(model, "_data_sources", None)
+        if isinstance(data_sources, list) and all(isinstance(item, data_source.DataSource) for item in data_sources):
+            return data_sources
+        return None

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -1,16 +1,17 @@
 import collections
+import copy
 import pathlib
-from typing import Any, Dict, List, Optional, cast
+from typing import List, Optional, cast
 
 import yaml
 
+from snowflake.ml._internal.lineage import data_source
 from snowflake.ml.model import type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._model_composer.model_method import (
     function_generator,
     model_method,
 )
-from snowflake.ml.model._model_composer.model_runtime import model_runtime
 from snowflake.ml.model._packager.model_meta import model_meta as model_meta_api
 from snowflake.snowpark import Session
 
@@ -36,24 +37,23 @@ class ModelManifest:
         model_meta: model_meta_api.ModelMetadata,
         model_file_rel_path: pathlib.PurePosixPath,
         options: Optional[type_hints.ModelSaveOption] = None,
+        data_sources: Optional[List[data_source.DataSource]] = None,
     ) -> None:
         if options is None:
             options = {}
-        self.runtimes = [
-            model_runtime.ModelRuntime(
-                session=session,
-                name=ModelManifest._DEFAULT_RUNTIME_NAME,
-                model_meta=model_meta,
-                imports=[model_file_rel_path],
-            )
-        ]
+
+        runtime_to_use = copy.deepcopy(model_meta.runtimes["cpu"])
+        runtime_to_use.name = self._DEFAULT_RUNTIME_NAME
+        runtime_to_use.imports.append(model_file_rel_path)
+        runtime_dict = runtime_to_use.save(self.workspace_path)
+
         self.function_generator = function_generator.FunctionGenerator(model_file_rel_path=model_file_rel_path)
         self.methods: List[model_method.ModelMethod] = []
         for target_method in model_meta.signatures.keys():
             method = model_method.ModelMethod(
                 model_meta=model_meta,
                 target_method=target_method,
-                runtime_name=self.runtimes[0].name,
+                runtime_name=self._DEFAULT_RUNTIME_NAME,
                 function_generator=self.function_generator,
                 options=model_method.get_model_method_options_from_options(options, target_method),
             )
@@ -71,7 +71,16 @@ class ModelManifest:
 
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
-            runtimes={runtime.name: runtime.save(self.workspace_path) for runtime in self.runtimes},
+            runtimes={
+                self._DEFAULT_RUNTIME_NAME: model_manifest_schema.ModelRuntimeDict(
+                    language="PYTHON",
+                    version=runtime_to_use.runtime_env.python_version,
+                    imports=runtime_dict["imports"],
+                    dependencies=model_manifest_schema.ModelRuntimeDependenciesDict(
+                        conda=runtime_dict["dependencies"]["conda"]
+                    ),
+                )
+            },
             methods=[
                 method.save(
                     self.workspace_path,
@@ -83,7 +92,9 @@ class ModelManifest:
             ],
         )
 
-        manifest_dict["user_data"] = self.generate_user_data_with_client_data(model_meta)
+        lineage_sources = self._extract_lineage_info(data_sources)
+        if lineage_sources:
+            manifest_dict["lineage_sources"] = lineage_sources
 
         with (self.workspace_path / ModelManifest.MANIFEST_FILE_REL_PATH).open("w", encoding="utf-8") as f:
             # Anchors are not supported in the server, avoid that.
@@ -104,42 +115,18 @@ class ModelManifest:
 
         return res
 
-    def generate_user_data_with_client_data(self, model_meta: model_meta_api.ModelMetadata) -> Dict[str, Any]:
-        client_data = model_manifest_schema.SnowparkMLDataDict(
-            schema_version=model_manifest_schema.MANIFEST_CLIENT_DATA_SCHEMA_VERSION,
-            functions=[
-                model_manifest_schema.ModelFunctionInfoDict(
-                    name=method.method_name.identifier(),
-                    target_method=method.target_method,
-                    signature=model_meta.signatures[method.target_method].to_dict(),
+    def _extract_lineage_info(
+        self, data_sources: Optional[List[data_source.DataSource]]
+    ) -> List[model_manifest_schema.LineageSourceDict]:
+        result = []
+        if data_sources:
+            for source in data_sources:
+                result.append(
+                    model_manifest_schema.LineageSourceDict(
+                        # Currently, we only support lineage from Dataset.
+                        type=model_manifest_schema.LineageSourceTypes.DATASET.value,
+                        entity=source.fully_qualified_name,
+                        version=source.version,
+                    )
                 )
-                for method in self.methods
-            ],
-        )
-        return {model_manifest_schema.MANIFEST_CLIENT_DATA_KEY_NAME: client_data}
-
-    @staticmethod
-    def parse_client_data_from_user_data(raw_user_data: Dict[str, Any]) -> model_manifest_schema.SnowparkMLDataDict:
-        raw_client_data = raw_user_data.get(model_manifest_schema.MANIFEST_CLIENT_DATA_KEY_NAME, {})
-        if not isinstance(raw_client_data, dict) or "schema_version" not in raw_client_data:
-            raise ValueError(f"Ill-formatted client data {raw_client_data} in user data found.")
-        loaded_client_data_schema_version = raw_client_data["schema_version"]
-        if (
-            not isinstance(loaded_client_data_schema_version, str)
-            or loaded_client_data_schema_version != model_manifest_schema.MANIFEST_CLIENT_DATA_SCHEMA_VERSION
-        ):
-            raise ValueError(f"Unsupported client data schema version {loaded_client_data_schema_version} confronted.")
-
-        return_functions_info: List[model_manifest_schema.ModelFunctionInfoDict] = []
-        loaded_functions_info = raw_client_data.get("functions", [])
-        for func in loaded_functions_info:
-            fi = model_manifest_schema.ModelFunctionInfoDict(
-                name=func["name"],
-                target_method=func["target_method"],
-                signature=func["signature"],
-            )
-            return_functions_info.append(fi)
-
-        return model_manifest_schema.SnowparkMLDataDict(
-            schema_version=loaded_client_data_schema_version, functions=return_functions_info
-        )
+        return result

snowflake/ml/model/_model_composer/model_manifest/model_manifest_schema.py

@@ -1,5 +1,5 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
-
+import enum
 from typing import Any, Dict, List, Literal, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
@@ -12,6 +12,11 @@ MANIFEST_CLIENT_DATA_KEY_NAME = "snowpark_ml_data"
 MANIFEST_CLIENT_DATA_SCHEMA_VERSION = "2024-02-01"
 
 
+class ModelMethodFunctionTypes(enum.Enum):
+    FUNCTION = "FUNCTION"
+    TABLE_FUNCTION = "TABLE_FUNCTION"
+
+
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
 
@@ -49,11 +54,13 @@ class ModelFunctionInfo(TypedDict):
     Attributes:
         name: Name of the function to be called via SQL.
         target_method: actual target method name to be called.
+        target_method_function_type: target method function type (FUNCTION or TABLE_FUNCTION).
         signature: The signature of the model method.
     """
 
     name: Required[str]
     target_method: Required[str]
+    target_method_function_type: Required[str]
     signature: Required[model_signature.ModelSignature]
 
 
@@ -68,8 +75,19 @@ class SnowparkMLDataDict(TypedDict):
     functions: Required[List[ModelFunctionInfoDict]]
 
 
+class LineageSourceTypes(enum.Enum):
+    DATASET = "DATASET"
+
+
+class LineageSourceDict(TypedDict):
+    type: Required[str]
+    entity: Required[str]
+    version: NotRequired[str]
+
+
 class ModelManifestDict(TypedDict):
     manifest_version: Required[str]
     runtimes: Required[Dict[str, ModelRuntimeDict]]
     methods: Required[List[ModelMethodDict]]
     user_data: NotRequired[Dict[str, Any]]
+    lineage_sources: NotRequired[List[LineageSourceDict]]

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -1,5 +1,4 @@
 import collections
-import enum
 import pathlib
 from typing import List, Optional, TypedDict, Union
 
@@ -13,11 +12,6 @@ from snowflake.ml.model._packager.model_meta import model_meta as model_meta_api
 from snowflake.snowpark._internal import type_utils
 
 
-class ModelMethodFunctionTypes(enum.Enum):
-    FUNCTION = "FUNCTION"
-    TABLE_FUNCTION = "TABLE_FUNCTION"
-
-
 class ModelMethodOptions(TypedDict):
     """Options when creating model method.
 
@@ -33,9 +27,9 @@ def get_model_method_options_from_options(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
     method_option = options.get("method_options", {}).get(target_method, {})
-    global_function_type = options.get("function_type", ModelMethodFunctionTypes.FUNCTION.value)
+    global_function_type = options.get("function_type", model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value)
     function_type = method_option.get("function_type", global_function_type)
-    if function_type not in [function_type.value for function_type in ModelMethodFunctionTypes]:
+    if function_type not in [function_type.value for function_type in model_manifest_schema.ModelMethodFunctionTypes]:
         raise NotImplementedError
 
     # TODO(TH): enforce minimum snowflake version
@@ -89,7 +83,9 @@ class ModelMethod:
         if self.target_method not in self.model_meta.signatures.keys():
             raise ValueError(f"Target method {self.target_method} is not available in the signatures of the model.")
 
-        self.function_type = self.options.get("function_type", ModelMethodFunctionTypes.FUNCTION.value)
+        self.function_type = self.options.get(
+            "function_type", model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+        )
 
     @staticmethod
     def _get_method_arg_from_feature(
@@ -134,7 +130,7 @@ class ModelMethod:
             List[model_manifest_schema.ModelMethodSignatureField],
             List[model_manifest_schema.ModelMethodSignatureFieldWithName],
         ]
-        if self.function_type == ModelMethodFunctionTypes.TABLE_FUNCTION.value:
+        if self.function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
             outputs = [
                 ModelMethod._get_method_arg_from_feature(ft, case_sensitive=self.options.get("case_sensitive", False))
                 for ft in self.model_meta.signatures[self.target_method].outputs

snowflake/ml/model/_packager/model_env/model_env.py

@@ -284,6 +284,7 @@ class ModelEnv:
                         " This may prevent model deploying to Snowflake Warehouse."
                     ),
                     category=UserWarning,
+                    stacklevel=2,
                 )
             if len(channel_dependencies) == 0 and channel not in self._conda_dependencies:
                 warnings.warn(
@@ -292,6 +293,7 @@ class ModelEnv:
                         " This may prevent model deploying to Snowflake Warehouse."
                     ),
                     category=UserWarning,
+                    stacklevel=2,
                 )
                 self._conda_dependencies[channel] = []
 
@@ -307,6 +309,7 @@ class ModelEnv:
                             " This may be unintentional."
                         ),
                         category=UserWarning,
+                        stacklevel=2,
                     )
 
         if pip_requirements_list:
@@ -316,6 +319,7 @@ class ModelEnv:
                     " This may prevent model deploying to Snowflake Warehouse."
                 ),
                 category=UserWarning,
+                stacklevel=2,
             )
             for pip_dependency in pip_requirements_list:
                 if any(
@@ -338,6 +342,7 @@ class ModelEnv:
                     " This may prevent model deploying to Snowflake Warehouse."
                 ),
                 category=UserWarning,
+                stacklevel=2,
             )
             for pip_dependency in pip_requirements_list:
                 if any(
@@ -372,3 +377,39 @@ class ModelEnv:
             "cuda_version": self.cuda_version,
             "snowpark_ml_version": self.snowpark_ml_version,
         }
+
+    def validate_with_local_env(
+        self, check_snowpark_ml_version: bool = False
+    ) -> List[env_utils.IncorrectLocalEnvironmentError]:
+        errors = []
+        try:
+            env_utils.validate_py_runtime_version(str(self._python_version))
+        except env_utils.IncorrectLocalEnvironmentError as e:
+            errors.append(e)
+
+        for conda_reqs in self._conda_dependencies.values():
+            for conda_req in conda_reqs:
+                try:
+                    env_utils.validate_local_installed_version_of_pip_package(
+                        env_utils.try_convert_conda_requirement_to_pip(conda_req)
+                    )
+                except env_utils.IncorrectLocalEnvironmentError as e:
+                    errors.append(e)
+
+        for pip_req in self._pip_requirements:
+            try:
+                env_utils.validate_local_installed_version_of_pip_package(pip_req)
+            except env_utils.IncorrectLocalEnvironmentError as e:
+                errors.append(e)
+
+        if check_snowpark_ml_version:
+            # For Modeling model
+            if self._snowpark_ml_version.base_version != snowml_env.VERSION:
+                errors.append(
+                    env_utils.IncorrectLocalEnvironmentError(
+                        f"The local installed version of Snowpark ML library is {snowml_env.VERSION} "
+                        f"which differs from required version {self.snowpark_ml_version}."
+                    )
+                )
+
+        return errors