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

snowflake/ml/feature_store/feature_view.py

@@ -22,6 +22,7 @@ from snowflake.ml.feature_store import feature_store
 from snowflake.ml.feature_store.entity import Entity
 from snowflake.ml.lineage import lineage_node
 from snowflake.snowpark import DataFrame, Session
+from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.types import (
     DateType,
     StructType,
@@ -167,6 +168,7 @@ class FeatureView(lineage_node.LineageNode):
         refresh_freq: Optional[str] = None,
         desc: str = "",
         warehouse: Optional[str] = None,
+        initialize: str = "ON_CREATE",
         **_kwargs: Any,
     ) -> None:
         """
@@ -190,6 +192,10 @@ class FeatureView(lineage_node.LineageNode):
             warehouse: warehouse to refresh feature view. Not needed for static feature view (refresh_freq is None).
                 For managed feature view, this warehouse will overwrite the default warehouse of Feature Store if it is
                 specified, otherwise the default warehouse will be used.
+            initialize: Specifies the behavior of the initial refresh of feature view. This property cannot be altered
+                after you register the feature view. It supports ON_CREATE (default) or ON_SCHEDULE. ON_CREATE refreshes
+                the feature view synchronously at creation. ON_SCHEDULE refreshes the feature view at the next scheduled
+                refresh. It is only effective when refresh_freq is not None.
             _kwargs: reserved kwargs for system generated args. NOTE: DO NOT USE.
 
         Example::
@@ -227,10 +233,14 @@ class FeatureView(lineage_node.LineageNode):
         self._query: str = self._get_query()
         self._version: Optional[FeatureViewVersion] = None
         self._status: FeatureViewStatus = FeatureViewStatus.DRAFT
-        self._feature_desc: OrderedDict[SqlIdentifier, str] = OrderedDict((f, "") for f in self._get_feature_names())
+        feature_names = self._get_feature_names()
+        self._feature_desc: Optional[OrderedDict[SqlIdentifier, str]] = (
+            OrderedDict((f, "") for f in feature_names) if feature_names is not None else None
+        )
         self._refresh_freq: Optional[str] = refresh_freq
         self._database: Optional[SqlIdentifier] = None
         self._schema: Optional[SqlIdentifier] = None
+        self._initialize: str = initialize
         self._warehouse: Optional[SqlIdentifier] = SqlIdentifier(warehouse) if warehouse is not None else None
         self._refresh_mode: Optional[str] = _kwargs.get("refresh_mode", "AUTO")
         self._refresh_mode_reason: Optional[str] = None
@@ -345,6 +355,15 @@ class FeatureView(lineage_node.LineageNode):
                 ('START_STATION_LATITUDE', 'Latitude of the start station.')])
 
         """
+        if self._feature_desc is None:
+            warnings.warn(
+                "Failed to read feature view schema. Probably feature view is not refreshed yet. "
+                "Schema will be available after initial refresh.",
+                stacklevel=2,
+                category=UserWarning,
+            )
+            return self
+
         for f, d in descs.items():
             f = SqlIdentifier(f)
             if f not in self._feature_desc:
@@ -424,10 +443,10 @@ class FeatureView(lineage_node.LineageNode):
 
     @property
     def feature_names(self) -> List[SqlIdentifier]:
-        return list(self._feature_desc.keys())
+        return list(self._feature_desc.keys()) if self._feature_desc is not None else []
 
     @property
-    def feature_descs(self) -> Dict[SqlIdentifier, str]:
+    def feature_descs(self) -> Optional[Dict[SqlIdentifier, str]]:
         return self._feature_desc
 
     def list_columns(self) -> DataFrame:
@@ -463,7 +482,17 @@ class FeatureView(lineage_node.LineageNode):
 
         """
         session = self._feature_df.session
-        rows = []
+        rows = []  # type: ignore[var-annotated]
+
+        if self.feature_descs is None:
+            warnings.warn(
+                "Failed to read feature view schema. Probably feature view is not refreshed yet. "
+                "Schema will be available after initial refresh.",
+                stacklevel=2,
+                category=UserWarning,
+            )
+            return session.create_dataframe(rows, schema=["name", "category", "dtype", "desc"])
+
         for name, type in self._feature_df.dtypes:
             if SqlIdentifier(name) in self.feature_descs:
                 desc = self.feature_descs[SqlIdentifier(name)]
@@ -565,6 +594,10 @@ class FeatureView(lineage_node.LineageNode):
         )
         self._warehouse = SqlIdentifier(new_value)
 
+    @property
+    def initialize(self) -> str:
+        return self._initialize
+
     @property
     def output_schema(self) -> StructType:
         return self._infer_schema_df.schema
@@ -601,33 +634,49 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
                 f"FeatureView name `{self._name}` contains invalid character `{_FEATURE_VIEW_NAME_DELIMITER}`."
             )
 
-        unescaped_df_cols = to_sql_identifiers(self._infer_schema_df.columns)
-        for e in self._entities:
-            for k in e.join_keys:
-                if k not in unescaped_df_cols:
-                    raise ValueError(
-                        f"join_key {k} in Entity {e.name} is not found in input dataframe: {unescaped_df_cols}"
-                    )
-
-        if self._timestamp_col is not None:
-            ts_col = self._timestamp_col
-            if ts_col == SqlIdentifier(_TIMESTAMP_COL_PLACEHOLDER):
-                raise ValueError(f"Invalid timestamp_col name, cannot be {_TIMESTAMP_COL_PLACEHOLDER}.")
-            if ts_col not in to_sql_identifiers(self._infer_schema_df.columns):
-                raise ValueError(f"timestamp_col {ts_col} is not found in input dataframe.")
-
-            col_type = self._infer_schema_df.schema[ts_col].datatype
-            if not isinstance(col_type, (DateType, TimeType, TimestampType, _NumericType)):
-                raise ValueError(f"Invalid data type for timestamp_col {ts_col}: {col_type}.")
+        df_cols = self._get_column_names()
+        if df_cols is not None:
+            for e in self._entities:
+                for k in e.join_keys:
+                    if k not in df_cols:
+                        raise ValueError(f"join_key {k} in Entity {e.name} is not found in input dataframe: {df_cols}")
+
+            if self._timestamp_col is not None:
+                ts_col = self._timestamp_col
+                if ts_col == SqlIdentifier(_TIMESTAMP_COL_PLACEHOLDER):
+                    raise ValueError(f"Invalid timestamp_col name, cannot be {_TIMESTAMP_COL_PLACEHOLDER}.")
+                if ts_col not in df_cols:
+                    raise ValueError(f"timestamp_col {ts_col} is not found in input dataframe.")
+
+                col_type = self._infer_schema_df.schema[ts_col].datatype
+                if not isinstance(col_type, (DateType, TimeType, TimestampType, _NumericType)):
+                    raise ValueError(f"Invalid data type for timestamp_col {ts_col}: {col_type}.")
 
         if re.match(_RESULT_SCAN_QUERY_PATTERN, self._query) is not None:
             raise ValueError(f"feature_df should not be reading from RESULT_SCAN. Invalid query: {self._query}")
 
-    def _get_feature_names(self) -> List[SqlIdentifier]:
+        if self._initialize not in ["ON_CREATE", "ON_SCHEDULE"]:
+            raise ValueError("'initialize' only supports ON_CREATE or ON_SCHEDULE.")
+
+    def _get_column_names(self) -> Optional[List[SqlIdentifier]]:
+        try:
+            return to_sql_identifiers(self._infer_schema_df.columns)
+        except SnowparkSQLException as e:
+            warnings.warn(
+                "Failed to read feature view schema. Probably feature view is not refreshed yet. "
+                f"Schema will be available after initial refresh. Original exception: {e}",
+                stacklevel=2,
+                category=UserWarning,
+            )
+            return None
+
+    def _get_feature_names(self) -> Optional[List[SqlIdentifier]]:
         join_keys = [k for e in self._entities for k in e.join_keys]
         ts_col = [self._timestamp_col] if self._timestamp_col is not None else []
-        feature_names = to_sql_identifiers(self._infer_schema_df.columns, case_sensitive=False)
-        return [c for c in feature_names if c not in join_keys + ts_col]
+        feature_names = self._get_column_names()
+        if feature_names is not None:
+            return [c for c in feature_names if c not in join_keys + ts_col]
+        return None
 
     def __repr__(self) -> str:
         states = (f"{k}={v}" for k, v in vars(self).items())
@@ -670,11 +719,13 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         fv_dict["_schema"] = str(self._schema) if self._schema is not None else None
         fv_dict["_warehouse"] = str(self._warehouse) if self._warehouse is not None else None
         fv_dict["_timestamp_col"] = str(self._timestamp_col) if self._timestamp_col is not None else None
+        fv_dict["_initialize"] = str(self._initialize)
 
         feature_desc_dict = {}
-        for k, v in self._feature_desc.items():
-            feature_desc_dict[k.identifier()] = v
-        fv_dict["_feature_desc"] = feature_desc_dict
+        if self._feature_desc is not None:
+            for k, v in self._feature_desc.items():
+                feature_desc_dict[k.identifier()] = v
+            fv_dict["_feature_desc"] = feature_desc_dict
 
         lineage_node_keys = [key for key in fv_dict if key.startswith("_node") or key == "_session"]
 
@@ -706,7 +757,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             >>> ).attach_feature_desc({"AGE": "my age", "TITLE": '"my title"'})
             >>> fv = fs.register_feature_view(draft_fv, '1.0')
             <BLANKLINE>
-            fv.to_df().show()
+            >>> fv.to_df().show()
             ----------------------------------------------------------------...
             |"NAME"  |"ENTITIES"                |"TIMESTAMP_COL"  |"DESC"  |
             ----------------------------------------------------------------...
@@ -760,6 +811,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
             warehouse=json_dict["_warehouse"],
             refresh_mode=json_dict["_refresh_mode"],
             refresh_mode_reason=json_dict["_refresh_mode_reason"],
+            initialize=json_dict["_initialize"],
             owner=json_dict["_owner"],
             infer_schema_df=session.sql(json_dict.get("_infer_schema_query", None)),
             session=session,
@@ -801,7 +853,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
 
     @staticmethod
     def _load_from_lineage_node(session: Session, name: str, version: str) -> FeatureView:
-        db_name, feature_store_name, feature_view_name, _ = identifier.parse_schema_level_object_identifier(name)
+        db_name, feature_store_name, feature_view_name = identifier.parse_schema_level_object_identifier(name)
 
         session_warehouse = session.get_current_warehouse()
 
@@ -830,6 +882,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         warehouse: Optional[str],
         refresh_mode: Optional[str],
         refresh_mode_reason: Optional[str],
+        initialize: str,
         owner: Optional[str],
         infer_schema_df: Optional[DataFrame],
         session: Session,
@@ -850,6 +903,7 @@ Got {len(self._feature_df.queries['queries'])}: {self._feature_df.queries['queri
         fv._warehouse = SqlIdentifier(warehouse) if warehouse is not None else None
         fv._refresh_mode = refresh_mode
         fv._refresh_mode_reason = refresh_mode_reason
+        fv._initialize = initialize
         fv._owner = owner
         fv.attach_feature_desc(feature_descs)
 

snowflake/ml/fileset/embedded_stage_fs.py

@@ -35,7 +35,7 @@ class SFEmbeddedStageFileSystem(stage_fs.SFStageFileSystem):
         **kwargs: Any,
     ) -> None:
 
-        (db, schema, object_name, _) = identifier.parse_schema_level_object_identifier(name)
+        (db, schema, object_name) = identifier.parse_schema_level_object_identifier(name)
         self._name = name  # TODO: Require or resolve FQN
         self._domain = domain
 

snowflake/ml/fileset/fileset.py

@@ -538,7 +538,7 @@ def _validate_target_stage_loc(snowpark_session: snowpark.Session, target_stage_
             original_exception=fileset_errors.FileSetLocationError('FileSet location should start with "@".'),
         )
     try:
-        db, schema, stage, _ = identifier.parse_schema_level_object_identifier(target_stage_loc[1:])
+        db, schema, stage, _ = identifier.parse_snowflake_stage_path(target_stage_loc[1:])
         if db is None or schema is None:
             raise ValueError("The stage path should be in the form '@<database>.<schema>.<stage>/*'")
         df_stages = snowpark_session.sql(f"Show stages like '{stage}' in SCHEMA {db}.{schema}")

snowflake/ml/fileset/sfcfs.py

@@ -15,6 +15,7 @@ from snowflake.ml._internal.exceptions import (
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.fileset import stage_fs
 from snowflake.ml.utils import connection_params
+from snowflake.snowpark import context, exceptions as snowpark_exceptions
 
 PROTOCOL_NAME = "sfc"
 
@@ -84,7 +85,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
         """
         if kwargs.get(_RECREATE_FROM_SERIALIZED):
             try:
-                snowpark_session = self._create_default_session()
+                snowpark_session = self._get_default_session()
             except Exception as e:
                 raise snowml_exceptions.SnowflakeMLException(
                     error_code=error_codes.SNOWML_DESERIALIZATION_FAILED,
@@ -103,7 +104,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
 
         super().__init__(**kwargs)
 
-    def _create_default_session(self) -> snowpark.Session:
+    def _get_default_session(self) -> snowpark.Session:
         """Create a Snowpark Session from default login options.
 
         Returns:
@@ -114,6 +115,11 @@ class SFFileSystem(fsspec.AbstractFileSystem):
             ValueError: Snowflake Connection could not be created.
 
         """
+        try:
+            return context.get_active_session()
+        except snowpark_exceptions.SnowparkSessionException:
+            pass
+
         try:
             snowflake_config = connection_params.SnowflakeLoginOptions()
         except Exception as e:
@@ -328,7 +334,7 @@ class SFFileSystem(fsspec.AbstractFileSystem):
                 ),
             )
         try:
-            res = identifier.parse_schema_level_object_identifier(path[1:])
+            res = identifier.parse_snowflake_stage_path(path[1:])
             if res[1] is None or res[0] is None or (res[3] and not res[3].startswith("/")):
                 raise ValueError("Invalid path. Missing database or schema identifier.")
             logging.debug(f"Parsed path: {res}")

snowflake/ml/fileset/stage_fs.py

@@ -431,4 +431,5 @@ def _resolve_async_job(async_job: snowpark.AsyncJob) -> List[snowpark.Row]:
                 error_code=error_codes.SNOWML_NOT_FOUND,
                 original_exception=fileset_errors.StageNotFoundError("Query failed."),
             ) from e
-        raise
+        assert e.msg is not None
+        raise snowpark_exceptions.SnowparkSQLException(e.msg, conn_error=e) from e

snowflake/ml/lineage/lineage_node.py

@@ -118,16 +118,21 @@ class LineageNode:
             )
             domain = lineage_object["domain"].lower()
             if domain_filter is None or domain in domain_filter:
+                obj_name = ".".join(
+                    identifier.rename_to_valid_snowflake_identifier(s)
+                    for s in identifier.parse_schema_level_object_identifier(lineage_object["name"])
+                    if s is not None
+                )
                 if domain in DOMAIN_LINEAGE_REGISTRY and lineage_object["status"] == "ACTIVE":
                     lineage_nodes.append(
                         DOMAIN_LINEAGE_REGISTRY[domain]._load_from_lineage_node(
-                            self._session, lineage_object["name"], lineage_object.get("version")
+                            self._session, obj_name, lineage_object.get("version")
                         )
                     )
                 else:
                     lineage_nodes.append(
                         LineageNode(
-                            name=lineage_object["name"],
+                            name=obj_name,
                             version=lineage_object.get("version"),
                             domain=domain,
                             status=lineage_object["status"],

snowflake/ml/model/__init__.py

@@ -1,6 +1,5 @@
 from snowflake.ml.model._client.model.model_impl import Model
 from snowflake.ml.model._client.model.model_version_impl import ExportMode, ModelVersion
 from snowflake.ml.model.models.huggingface_pipeline import HuggingFacePipelineModel
-from snowflake.ml.model.models.llm import LLM, LLMOptions
 
-__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel", "LLM", "LLMOptions"]
+__all__ = ["Model", "ModelVersion", "ExportMode", "HuggingFacePipelineModel"]

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

@@ -306,6 +306,23 @@ class ModelVersion(lineage_node.LineageNode):
             statement_params=statement_params,
         )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def get_model_task(self) -> model_types.Task:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops.get_model_task(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -406,6 +423,7 @@ class ModelVersion(lineage_node.LineageNode):
             partition_column = sql_identifier.SqlIdentifier(partition_column)
 
         functions: List[model_manifest_schema.ModelFunctionInfo] = self._functions
+
         if function_name:
             req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
             find_method: Callable[[model_manifest_schema.ModelFunctionInfo], bool] = (
@@ -606,8 +624,9 @@ class ModelVersion(lineage_node.LineageNode):
             "image_repo_database",
             "image_repo_schema",
             "image_repo",
-            "image_name",
             "gpu_requests",
+            "num_workers",
+            "max_batch_rows",
         ],
     )
     def create_service(
@@ -617,11 +636,11 @@ class ModelVersion(lineage_node.LineageNode):
         image_build_compute_pool: Optional[str] = None,
         service_compute_pool: str,
         image_repo: str,
-        image_name: Optional[str] = None,
         ingress_enabled: bool = False,
-        min_instances: int = 1,
         max_instances: int = 1,
         gpu_requests: Optional[str] = None,
+        num_workers: Optional[int] = None,
+        max_batch_rows: Optional[int] = None,
         force_rebuild: bool = False,
         build_external_access_integration: str,
     ) -> str:
@@ -630,22 +649,27 @@ class ModelVersion(lineage_node.LineageNode):
         Args:
             service_name: The name of the service, can be fully qualified. If not fully qualified, the database or
                 schema of the model will be used.
-            image_build_compute_pool: The name of the compute pool used to build the model inference image. Use
+            image_build_compute_pool: The name of the compute pool used to build the model inference image. It uses
                 the service compute pool if None.
             service_compute_pool: The name of the compute pool used to run the inference service.
             image_repo: The name of the image repository, can be fully qualified. If not fully qualified, the database
                 or schema of the model will be used.
-            image_name: The name of the model inference image. Use a generated name if None.
-            ingress_enabled: Whether to enable ingress.
-            min_instances: The minimum number of inference service instances to run.
-            max_instances: The maximum number of inference service instances to run.
+            ingress_enabled: If true, creates an service endpoint associated with the service. User must have
+                BIND SERVICE ENDPOINT privilege on the account.
+            max_instances: The maximum number of inference service instances to run. The same value it set to
+                MIN_INSTANCES property of the service.
             gpu_requests: The gpu limit for GPU based inference. Can be integer, fractional or string values. Use CPU
                 if None.
+            num_workers: The number of workers to run the inference service for handling requests in parallel within an
+                instance of the service. By default, it is set to 2*vCPU+1 of the node for CPU based inference and 1 for
+                GPU based inference. For GPU based inference, please see best practices before playing with this value.
+            max_batch_rows: The maximum number of rows to batch for inference. Auto determined if None. Minimum 32.
             force_rebuild: Whether to force a model inference image rebuild.
-            build_external_access_integration: The external access integration for image build.
+            build_external_access_integration: The external access integration for image build. This is usually
+                permitting access to conda & PyPI repositories.
 
         Returns:
-            The service name.
+            Result information about service creation from server.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
@@ -670,15 +694,75 @@ class ModelVersion(lineage_node.LineageNode):
             image_repo_database_name=image_repo_db_id,
             image_repo_schema_name=image_repo_schema_id,
             image_repo_name=image_repo_id,
-            image_name=sql_identifier.SqlIdentifier(image_name) if image_name else None,
             ingress_enabled=ingress_enabled,
-            min_instances=min_instances,
             max_instances=max_instances,
             gpu_requests=gpu_requests,
+            num_workers=num_workers,
+            max_batch_rows=max_batch_rows,
             force_rebuild=force_rebuild,
             build_external_access_integration=sql_identifier.SqlIdentifier(build_external_access_integration),
             statement_params=statement_params,
         )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def list_services(
+        self,
+    ) -> List[str]:
+        """List all the service names using this model version.
+
+        Returns:
+            List of service_names: The name of the service, can be fully qualified. If not fully qualified, the database
+                or schema of the model will be used.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+
+        return self._model_ops.list_inference_services(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def delete_service(
+        self,
+        service_name: str,
+    ) -> None:
+        """Drops the given service.
+
+        Args:
+            service_name: The name of the service, can be fully qualified. If not fully qualified, the database or
+                schema of the model will be used.
+
+        Raises:
+            ValueError: If the service does not exist or operation is not permitted by user or service does not belong
+                to this model.
+        """
+        if not service_name:
+            raise ValueError("service_name cannot be empty.")
+
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        self._model_ops.delete_service(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            service_name=service_name,
+            statement_params=statement_params,
+        )
+
 
 lineage_node.DOMAIN_LINEAGE_REGISTRY["model"] = ModelVersion