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

snowflake/ml/fileset/snowfs.py

@@ -1,10 +1,9 @@
 import collections
 import logging
 import re
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import fsspec
-import packaging.version as pkg_version
 
 from snowflake import snowpark
 from snowflake.connector import connection
@@ -12,7 +11,7 @@ from snowflake.ml._internal.exceptions import (
     error_codes,
     exceptions as snowml_exceptions,
 )
-from snowflake.ml._internal.utils import identifier, snowflake_env
+from snowflake.ml._internal.utils import identifier
 from snowflake.ml.fileset import embedded_stage_fs, sfcfs
 
 PROTOCOL_NAME = "snow"
@@ -28,10 +27,6 @@ _SNOWURL_PATTERN = re.compile(
     r"(?P<path>versions/(?:(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)"
 )
 
-# FIXME(dhung): Temporary fix for bug in GS version 8.17
-_BUG_VERSION_MIN = pkg_version.Version("8.17")  # Inclusive minimum version with bugged behavior
-_BUG_VERSION_MAX = pkg_version.Version("8.18")  # Exclusive maximum version with bugged behavior
-
 
 class SnowFileSystem(sfcfs.SFFileSystem):
     """A filesystem that allows user to access Snowflake embedded stage files with valid Snowflake locations.
@@ -54,21 +49,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
     ) -> None:
         super().__init__(sf_connection=sf_connection, snowpark_session=snowpark_session, **kwargs)
 
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        if SnowFileSystem._IS_BUGGED_VERSION is None:
-            try:
-                sf_version = snowflake_env.get_current_snowflake_version(self._session)
-                SnowFileSystem._IS_BUGGED_VERSION = _BUG_VERSION_MIN <= sf_version < _BUG_VERSION_MAX
-            except Exception:
-                SnowFileSystem._IS_BUGGED_VERSION = False
-
-    def info(self, path: str, **kwargs: Any) -> Dict[str, Any]:
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        res: Dict[str, Any] = super().info(path, **kwargs)
-        if res.get("type") == "directory" and not res["name"].endswith("/"):
-            res["name"] += "/"
-        return res
-
     def _get_stage_fs(
         self, sf_file_path: _SFFileEntityPath  # type: ignore[override]
     ) -> embedded_stage_fs.SFEmbeddedStageFileSystem:
@@ -100,11 +80,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
         if stage_name.startswith(protocol):
             stage_name = stage_name[len(protocol) :]
         abs_path = stage_name + "/" + path
-        # FIXME(dhung): Temporary fix for bug in GS version 8.17
-        if self._IS_BUGGED_VERSION:
-            match = _SNOWURL_PATTERN.fullmatch(abs_path)
-            assert match is not None
-            abs_path = abs_path.replace(match.group("relpath"), match.group("relpath").lstrip("/"))
         return abs_path
 
     @classmethod
@@ -143,9 +118,6 @@ class SnowFileSystem(sfcfs.SFFileSystem):
             version = snowurl_match.group("version")
             relative_path = snowurl_match.group("relpath") or ""
             logging.debug(f"Parsed snow URL: {snowurl_match.groups()}")
-            # FIXME(dhung): Temporary fix for bug in GS version 8.17
-            if cls._IS_BUGGED_VERSION:
-                filepath = filepath.replace(f"{version}/", f"{version}//")
             return _SFFileEntityPath(
                 domain=domain, name=name, version=version, relative_path=relative_path, filepath=filepath
             )

snowflake/ml/fileset/stage_fs.py

@@ -2,13 +2,13 @@ import inspect
 import logging
 import time
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 import fsspec
 from fsspec.implementations import http as httpfs
 
 from snowflake import snowpark
-from snowflake.connector import connection, errorcode
+from snowflake.connector import connection, errorcode, errors as snowpark_errors
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.exceptions import (
     error_codes,
@@ -18,6 +18,7 @@ from snowflake.ml._internal.exceptions import (
 )
 from snowflake.snowpark import exceptions as snowpark_exceptions
 from snowflake.snowpark._internal import utils as snowpark_utils
+from snowflake.snowpark._internal.analyzer import snowflake_plan
 
 # The default length of how long a presigned url stays active in seconds.
 # Presigned url here is used to fetch file objects from Snowflake when SFStageFileSystem.open() is called.
@@ -167,7 +168,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         try:
             loc = self.stage_name
             path = path.lstrip("/")
-            objects = self._session.sql(f"LIST '{loc}/{path}'").collect()
+            async_job: snowpark.AsyncJob = self._session.sql(f"LIST '{loc}/{path}'").collect(block=False)
+            objects: List[snowpark.Row] = _resolve_async_job(async_job)
         except snowpark_exceptions.SnowparkClientException as e:
             if e.message.startswith(fileset_errors.ERRNO_DOMAIN_NOT_EXIST):
                 raise snowml_exceptions.SnowflakeMLException(
@@ -289,9 +291,7 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
                     original_exception=e,
                 )
 
-    def _parse_list_result(
-        self, list_result: List[Tuple[str, int, str, str]], search_path: str
-    ) -> List[Dict[str, Any]]:
+    def _parse_list_result(self, list_result: List[snowpark.Row], search_path: str) -> List[Dict[str, Any]]:
         """Convert the result from LIST query to the expected format of fsspec ls() method.
 
         Note that Snowflake LIST query has different behavior with ls(). LIST query will return all the stage files
@@ -312,7 +312,8 @@ class SFStageFileSystem(fsspec.AbstractFileSystem):
         """
         files: Dict[str, Dict[str, Any]] = {}
         search_path = search_path.strip("/")
-        for name, size, md5, last_modified in list_result:
+        for row in list_result:
+            name, size, md5, last_modified = row["name"], row["size"], row["md5"], row["last_modified"]
             obj_path = self._stage_path_to_relative_path(name)
             if obj_path == search_path:
                 # If there is a exact match, then the matched object will always be a file object.
@@ -408,3 +409,20 @@ def _match_error_code(ex: snowpark_exceptions.SnowparkSQLException, error_code:
     # Snowpark writes error code to message instead of populating e.error_code
     error_code_str = str(error_code)
     return ex.error_code == error_code_str or error_code_str in ex.message
+
+
+@snowflake_plan.SnowflakePlan.Decorator.wrap_exception  # type: ignore[misc]
+def _resolve_async_job(async_job: snowpark.AsyncJob) -> List[snowpark.Row]:
+    # Make sure Snowpark exceptions are properly caught and converted by wrap_exception wrapper
+    try:
+        query_result = cast(List[snowpark.Row], async_job.result("row"))
+        return query_result
+    except snowpark_errors.DatabaseError as e:
+        # HACK: Snowpark surfaces a generic exception if query doesn't complete immediately
+        # assume it's due to FileNotFound
+        if type(e) is snowpark_errors.DatabaseError and "results are unavailable" in str(e):
+            raise snowml_exceptions.SnowflakeMLException(
+                error_code=error_codes.SNOWML_NOT_FOUND,
+                original_exception=fileset_errors.StageNotFoundError("Query failed."),
+            ) from e
+        raise

snowflake/ml/model/_client/model/model_impl.py

@@ -1,9 +1,9 @@
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Union
 
 import pandas as pd
 
 from snowflake.ml._internal import telemetry
-from snowflake.ml._internal.utils import identifier, sql_identifier
+from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.model._client.model import model_version_impl
 from snowflake.ml.model._client.ops import model_ops
 
@@ -45,7 +45,7 @@ class Model:
     @property
     def fully_qualified_name(self) -> str:
         """Return the fully qualified name of the model that can be used to refer to it in SQL."""
-        return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
+        return self._model_ops._model_version_client.fully_qualified_object_name(None, None, self._model_name)
 
     @property
     @telemetry.send_api_usage_telemetry(
@@ -76,6 +76,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.get_comment(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -92,6 +94,8 @@ class Model:
         )
         return self._model_ops.set_comment(
             comment=comment,
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -109,7 +113,7 @@ class Model:
             class_name=self.__class__.__name__,
         )
         default_version_name = self._model_ops.get_default_version(
-            model_name=self._model_name, statement_params=statement_params
+            database_name=None, schema_name=None, model_name=self._model_name, statement_params=statement_params
         )
         return self.version(default_version_name)
 
@@ -129,7 +133,11 @@ class Model:
         else:
             version_name = version._version_name
         self._model_ops.set_default_version(
-            model_name=self._model_name, version_name=version_name, statement_params=statement_params
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
 
     @telemetry.send_api_usage_telemetry(
@@ -155,6 +163,8 @@ class Model:
         )
         version_id = sql_identifier.SqlIdentifier(version_name)
         if self._model_ops.validate_existence(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=version_id,
             statement_params=statement_params,
@@ -184,6 +194,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         version_names = self._model_ops.list_models_or_versions(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -211,6 +223,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         rows = self._model_ops.show_models_or_versions(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             statement_params=statement_params,
         )
@@ -231,6 +245,8 @@ class Model:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         self._model_ops.delete_model_or_version(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=sql_identifier.SqlIdentifier(version_name),
             statement_params=statement_params,
@@ -250,29 +266,9 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        return self._model_ops.show_tags(model_name=self._model_name, statement_params=statement_params)
-
-    def _parse_tag_name(
-        self,
-        tag_name: str,
-    ) -> Tuple[sql_identifier.SqlIdentifier, sql_identifier.SqlIdentifier, sql_identifier.SqlIdentifier]:
-        _tag_db, _tag_schema, _tag_name, _ = identifier.parse_schema_level_object_identifier(tag_name)
-        if _tag_db is None:
-            tag_db_id = self._model_ops._model_client._database_name
-        else:
-            tag_db_id = sql_identifier.SqlIdentifier(_tag_db)
-
-        if _tag_schema is None:
-            tag_schema_id = self._model_ops._model_client._schema_name
-        else:
-            tag_schema_id = sql_identifier.SqlIdentifier(_tag_schema)
-
-        if _tag_name is None:
-            raise ValueError(f"Unable parse the tag name `{tag_name}` you input.")
-
-        tag_name_id = sql_identifier.SqlIdentifier(_tag_name)
-
-        return tag_db_id, tag_schema_id, tag_name_id
+        return self._model_ops.show_tags(
+            database_name=None, schema_name=None, model_name=self._model_name, statement_params=statement_params
+        )
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
@@ -292,8 +288,10 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        tag_db_id, tag_schema_id, tag_name_id = self._parse_tag_name(tag_name)
+        tag_db_id, tag_schema_id, tag_name_id = sql_identifier.parse_fully_qualified_name(tag_name)
         return self._model_ops.get_tag_value(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             tag_database_name=tag_db_id,
             tag_schema_name=tag_schema_id,
@@ -317,8 +315,10 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        tag_db_id, tag_schema_id, tag_name_id = self._parse_tag_name(tag_name)
+        tag_db_id, tag_schema_id, tag_name_id = sql_identifier.parse_fully_qualified_name(tag_name)
         self._model_ops.set_tag(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             tag_database_name=tag_db_id,
             tag_schema_name=tag_schema_id,
@@ -342,8 +342,10 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        tag_db_id, tag_schema_id, tag_name_id = self._parse_tag_name(tag_name)
+        tag_db_id, tag_schema_id, tag_name_id = sql_identifier.parse_fully_qualified_name(tag_name)
         self._model_ops.unset_tag(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             tag_database_name=tag_db_id,
             tag_schema_name=tag_schema_id,
@@ -365,15 +367,20 @@ class Model:
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
         )
-        db, schema, model, _ = identifier.parse_schema_level_object_identifier(model_name)
-        new_model_db = sql_identifier.SqlIdentifier(db) if db else None
-        new_model_schema = sql_identifier.SqlIdentifier(schema) if schema else None
-        new_model_id = sql_identifier.SqlIdentifier(model)
+        new_db, new_schema, new_model = sql_identifier.parse_fully_qualified_name(model_name)
+
         self._model_ops.rename(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
-            new_model_db=new_model_db,
-            new_model_schema=new_model_schema,
-            new_model_name=new_model_id,
+            new_model_db=new_db,
+            new_model_schema=new_schema,
+            new_model_name=new_model,
             statement_params=statement_params,
         )
-        self._model_name = new_model_id
+        self._model_ops = model_ops.ModelOperator(
+            self._model_ops._session,
+            database_name=new_db or self._model_ops._model_client._database_name,
+            schema_name=new_schema or self._model_ops._model_client._schema_name,
+        )
+        self._model_name = new_model

snowflake/ml/model/_client/model/model_version_impl.py

@@ -72,7 +72,7 @@ class ModelVersion:
     @property
     def fully_qualified_model_name(self) -> str:
         """Return the fully qualified name of the model to which the model version belongs."""
-        return self._model_ops._model_version_client.fully_qualified_model_name(self._model_name)
+        return self._model_ops._model_version_client.fully_qualified_object_name(None, None, self._model_name)
 
     @property
     @telemetry.send_api_usage_telemetry(
@@ -103,6 +103,8 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.get_comment(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -120,6 +122,8 @@ class ModelVersion:
         )
         return self._model_ops.set_comment(
             comment=comment,
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -140,7 +144,11 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops._metadata_ops.load(
-            model_name=self._model_name, version_name=self._version_name, statement_params=statement_params
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
         )["metrics"]
 
     @telemetry.send_api_usage_telemetry(
@@ -183,6 +191,8 @@ class ModelVersion:
         metrics[metric_name] = value
         self._model_ops._metadata_ops.save(
             metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -211,6 +221,8 @@ class ModelVersion:
         del metrics[metric_name]
         self._model_ops._metadata_ops.save(
             metadata_ops.ModelVersionMetadataSchema(metrics=metrics),
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -222,6 +234,8 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         return self._model_ops.get_functions(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             statement_params=statement_params,
@@ -309,6 +323,8 @@ class ModelVersion:
             method_function_type=target_function_info["target_method_function_type"],
             signature=target_function_info["signature"],
             X=X,
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             strict_input_validation=strict_input_validation,
@@ -341,6 +357,8 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
         self._model_ops.download_files(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             target_path=target_local_path,
@@ -380,6 +398,8 @@ class ModelVersion:
             with tempfile.TemporaryDirectory() as tmp_workspace_for_validation:
                 ws_path_for_validation = pathlib.Path(tmp_workspace_for_validation)
                 self._model_ops.download_files(
+                    database_name=None,
+                    schema_name=None,
                     model_name=self._model_name,
                     version_name=self._version_name,
                     target_path=ws_path_for_validation,
@@ -417,6 +437,8 @@ class ModelVersion:
         # We need the folder to be existed.
         workspace = pathlib.Path(tempfile.mkdtemp())
         self._model_ops.download_files(
+            database_name=None,
+            schema_name=None,
             model_name=self._model_name,
             version_name=self._version_name,
             target_path=workspace,

snowflake/ml/model/_client/ops/metadata_ops.py

@@ -61,12 +61,18 @@ class MetadataOperator:
     def _get_current_metadata_dict(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         version_info_list = self._model_client.show_versions(
-            model_name=model_name, version_name=version_name, statement_params=statement_params
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
         metadata_str = version_info_list[0][self._model_client.MODEL_VERSION_METADATA_COL_NAME]
         if not metadata_str:
@@ -79,12 +85,18 @@ class MetadataOperator:
     def load(
         self,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> ModelVersionMetadataSchema:
         metadata_dict = self._get_current_metadata_dict(
-            model_name=model_name, version_name=version_name, statement_params=statement_params
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
         return MetadataOperator._parse(metadata_dict)
 
@@ -92,14 +104,25 @@ class MetadataOperator:
         self,
         metadata: ModelVersionMetadataSchema,
         *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> None:
         metadata_dict = self._get_current_metadata_dict(
-            model_name=model_name, version_name=version_name, statement_params=statement_params
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )
         metadata_dict.update({**metadata, "snowpark_ml_schema_version": MODEL_VERSION_METADATA_SCHEMA_VERSION})
         self._model_version_client.set_metadata(
-            metadata_dict, model_name=model_name, version_name=version_name, statement_params=statement_params
+            metadata_dict,
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
         )