snowflake-ml-python 2.4.0__py3-none-any.whl → 2.6.0__py3-none-any.whl

snowflake/cortex/_classify_text.py

@@ -6,24 +6,23 @@ from typing_extensions import deprecated
 from snowflake import snowpark
 from snowflake.cortex._util import (
     CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
-    call_sql_function,
+    SnowflakeAuthenticationException,
     warn_cortex_deprecated,
 )
 from snowflake.ml._internal import telemetry
-from snowflake.snowpark import functions
+from snowflake.snowpark import context, functions
 
 _AI_CLASSIFY_FUNCTION_NAME = "AI_CLASSIFY"
-# Legacy CLASSIFY_TEXT always returned exactly one label. AI_CLASSIFY defaults to single-label mode, but
-# we set output_mode explicitly so the single-label contract can't drift and multi-label results can
-# never leak through.
-_AI_CLASSIFY_SINGLE_LABEL_CONFIG = {"output_mode": "single"}
+# AI_CLASSIFY requires its config argument to be a constant object literal. A Python dict is rendered by
+# Snowpark as PARSE_JSON('...')::OBJECT, which the server rejects ("... is not a constant"), so emit the
+# config as a SQL object constant via sql_expr -- this mirrors snowflake.snowpark.functions.ai_classify.
+# Legacy CLASSIFY_TEXT was always single-label, so set output_mode explicitly to keep that contract.
+_AI_CLASSIFY_SINGLE_LABEL_CONFIG_SQL = "{'output_mode': 'single'}"
 
 
 def _adapt_ai_classify_result(result: Any) -> str:
     """Map an AI_CLASSIFY result to the legacy CLASSIFY_TEXT {"label": ...} JSON string."""
-    # Legacy classify_text returned the raw CLASSIFY_TEXT OBJECT, which the immediate (non-Column) path
-    # collects as a JSON string such as '{"label": "positive"}'. We reproduce that exact shape so
-    # existing callers that json.loads the result and read ["label"] keep working.
+    # Used for non-Column results (e.g. a dict/JSON string): reproduce the legacy single-label object shape.
     parsed: Any = result
     if isinstance(result, str):
         try:
@@ -40,6 +39,20 @@ def _adapt_ai_classify_result(result: Any) -> str:
     return json.dumps({"label": label})
 
 
+def _ai_classify_label_column(
+    str_input: str | snowpark.Column,
+    categories: list[str] | snowpark.Column,
+) -> snowpark.Column:
+    """Call AI_CLASSIFY in single-label mode and re-wrap to the legacy CLASSIFY_TEXT {"label": ...} OBJECT."""
+    result = functions.builtin(_AI_CLASSIFY_FUNCTION_NAME)(
+        str_input, categories, functions.sql_expr(_AI_CLASSIFY_SINGLE_LABEL_CONFIG_SQL)
+    )
+    # AI_CLASSIFY returns {"labels": [...]}; collapse to the legacy single {"label": ...}. Subscript access
+    # (not functions.get with a bare string, which Snowpark treats as a column identifier) uses literal keys.
+    label = functions.coalesce(result["labels"][0], result["label"])
+    return functions.object_construct(functions.lit("label"), label)
+
+
 @telemetry.send_api_usage_telemetry(
     project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
 )
@@ -56,32 +69,25 @@ def classify_text(
         session: The snowpark session to use. Will be inferred by context if not specified.
 
     Returns:
-        A column of classification responses.
+        A column of classification responses (legacy CLASSIFY_TEXT {"label": ...} shape).
+
+    Raises:
+        SnowflakeAuthenticationException: if no session is provided and none is active.
     """
     warn_cortex_deprecated("classify_text", "ai_classify")
-    result = _classify_text_impl(
-        _AI_CLASSIFY_FUNCTION_NAME, str_input, categories, _AI_CLASSIFY_SINGLE_LABEL_CONFIG, session=session
-    )
-    if isinstance(result, snowpark.Column):
-        # AI_CLASSIFY returns {"labels": [...]}; re-wrap as the legacy CLASSIFY_TEXT {"label": ...} OBJECT
-        # so downstream SQL that reads col["label"] is unchanged. Subscript access (not functions.get with
-        # a bare string, which Snowpark treats as a column identifier) extracts the field by literal key.
-        label = functions.coalesce(result["labels"][0], result["label"])
-        return functions.object_construct(functions.lit("label"), label)
-    return _adapt_ai_classify_result(result)
-
+    column = _ai_classify_label_column(str_input, categories)
+    if isinstance(str_input, snowpark.Column) or isinstance(categories, snowpark.Column):
+        return column
 
-def _classify_text_impl(
-    function: str,
-    str_input: str | snowpark.Column,
-    categories: list[str] | snowpark.Column,
-    config: dict[str, Any] | None = None,
-    session: snowpark.Session | None = None,
-) -> str | snowpark.Column:
-    args: list[Any] = [str_input, categories]
-    if config is not None:
-        args.append(config)
-    return cast(str | snowpark.Column, call_sql_function(function, session, *args))
+    # Immediate (literal) path: evaluate and return the collected OBJECT as a JSON string, exactly as the
+    # legacy CLASSIFY_TEXT immediate path did.
+    session = session or context.get_active_session()
+    if session is None:
+        raise SnowflakeAuthenticationException(
+            """Session required. Provide the session through a session=... argument or ensure an active session is
+            available in your environment."""
+        )
+    return cast(str, session.create_dataframe([snowpark.Row()]).select(column).collect()[0][0])
 
 
 ClassifyText = deprecated(

snowflake/ml/_internal/env_utils.py

@@ -90,6 +90,9 @@ def _validate_pip_requirement_string(req_str: str) -> requirements.Requirement:
     Returns:
         A requirements.Requirement object containing the requirement information.
     """
+    # A requirement string must not carry surrounding whitespace: packaging>=26 accepts only a
+    # requirement that ends at the end of the string, so a trailing newline is a parse error.
+    req_str = req_str.strip()
     try:
         r = requirements.Requirement(req_str)
 
@@ -780,8 +783,9 @@ def load_requirements_file(path: pathlib.Path) -> list[requirements.Requirement]
     with open(path, encoding="utf-8") as f:
         reqs = f.readlines()
 
-    # Filter out pip options
-    reqs = [line for line in reqs if line.strip() and not line.strip().startswith(("-", "#"))]
+    # "-" lines are pip options (e.g. --extra-index-url) and "#" lines are comments; neither is a
+    # requirement. Strip as well, since a requirement string must not keep readlines()' newline.
+    reqs = [stripped for line in reqs if (stripped := line.strip()) and not stripped.startswith(("-", "#"))]
 
     return validate_pip_requirement_string_list(reqs)
 

snowflake/ml/_internal/file_utils.py

@@ -120,7 +120,21 @@ def make_archive(
     )
 
 
-def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py_file: bool = True) -> None:
+def zip_python_package(
+    zipfile_path: str,
+    package_name: str,
+    ignore_generated_py_file: bool = True,
+    *,
+    exclude_dir_names: tuple[str, ...] = (),
+) -> None:
+    """Zip an imported Python package.
+
+    Args:
+        zipfile_path: Destination zip path.
+        package_name: Dotted package name to zip.
+        ignore_generated_py_file: When True, omit generated bytecode and stub files.
+        exclude_dir_names: Directory names to omit anywhere in the package. The directory entry itself is not written.
+    """
     import importlib_resources
     from importlib_resources import abc as importlib_resources_abc
 
@@ -131,6 +145,7 @@ def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py
 
     base_dirs = package_name.split(".")
     assert len(base_dirs) >= 1, "Invalid package name."
+    excluded_dirs = set(exclude_dir_names)
 
     with zipfile.ZipFile(zipfile_path, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
         for i in range(len(base_dirs)):
@@ -141,6 +156,8 @@ def zip_python_package(zipfile_path: str, package_name: str, ignore_generated_py
         def _add_to_zip(
             zf: zipfile.ZipFile, path_info: importlib_resources_abc.Traversable, base_arcname: pathlib.PurePosixPath
         ) -> None:
+            if path_info.name in excluded_dirs:
+                return
             if path_info.is_file():
                 if ignore_generated_py_file and path_info.name.endswith(GENERATED_PY_FILE_EXT):
                     return

snowflake/ml/_internal/platform_capabilities.py

@@ -93,7 +93,7 @@ class PlatformCapabilities:
         return self._get_bool_feature(ENABLE_MODEL_METHOD_SIGNATURE_PARAMETERS, False)
 
     def is_lora_adapters_enabled(self) -> bool:
-        return False
+        return self._get_bool_feature(ENABLE_LORA_ADAPTERS, False)
 
     def is_single_output_native_type_enabled(self) -> bool:
         """Whether a single, non-OBJECT model FUNCTION output is registered with its native SQL type.

snowflake/ml/dataset/dataset.py

@@ -341,7 +341,7 @@ class Dataset(lineage_node.LineageNode):
         metadata = dataset_metadata.DatasetMetadata(
             source_query=source_query,
             owner=self._session.sql("SELECT CURRENT_USER()").collect(statement_params=_TELEMETRY_STATEMENT_PARAMS)[0][
-                "CURRENT_USER()"
+                0
             ],
             exclude_cols=exclude_cols,
             label_cols=label_cols,

snowflake/ml/experiment/_entities/run.py

@@ -2,9 +2,11 @@ import types
 import warnings
 from typing import TYPE_CHECKING
 
-from snowflake.ml._internal.utils import sql_identifier
+from snowflake import snowpark
+from snowflake.ml._internal.utils import identifier, sql_identifier
 from snowflake.ml.experiment import _experiment_info as experiment_info
 from snowflake.ml.experiment._entities import run_metadata
+from snowflake.ml.lineage import lineage_node
 
 if TYPE_CHECKING:
     from snowflake.ml.experiment import experiment_tracking
@@ -12,13 +14,24 @@ if TYPE_CHECKING:
 METADATA_SIZE_WARNING_MESSAGE = "It is likely that no further metrics or parameters will be logged for this run."
 
 
-class Run:
+class Run(lineage_node.LineageNode):
+    """An experiment run.
+
+    When created by ``ExperimentTracking.start_run``, the run is an active context manager that
+    patches model logging and ends the run on exit. When returned by the lineage system
+    (``model_version.lineage(direction="upstream", domain_filter={"experiment"})``), it is a
+    read-only handle exposing the identifiers needed to look up metrics and parameters.
+    """
+
     def __init__(
         self,
-        experiment_tracking: "experiment_tracking.ExperimentTracking",
+        experiment_tracking: "experiment_tracking.ExperimentTracking | None" = None,
         *,
         experiment_name: sql_identifier.SqlIdentifier,
         run_name: sql_identifier.SqlIdentifier,
+        session: snowpark.Session | None = None,
+        database_name: str | None = None,
+        schema_name: str | None = None,
     ) -> None:
         self._experiment_tracking = experiment_tracking
         self.experiment_name = experiment_name
@@ -27,10 +40,31 @@ class Run:
         # Whether we've already shown the user a warning about exceeding the run metadata size limit.
         self._warned_about_metadata_size = False
 
+        if experiment_tracking is not None:
+            _session = experiment_tracking._session
+            self.database_name = str(experiment_tracking._database_name)
+            self.schema_name = str(experiment_tracking._schema_name)
+        else:
+            if session is None or database_name is None or schema_name is None:
+                raise ValueError(
+                    "session, database_name, and schema_name are required when experiment_tracking is None"
+                )
+            _session = session
+            self.database_name = database_name
+            self.schema_name = schema_name
+
         self._patcher = experiment_info.ExperimentInfoPatcher(
             experiment_info=self._get_experiment_info(),
         )
 
+        # A run is modelled in lineage as a version of its experiment, so the node name is the experiment's.
+        super().__init__(
+            session=_session,
+            name=self._fully_qualified_experiment_name(),
+            domain="experiment",
+            version=str(run_name),
+        )
+
     def __enter__(self) -> "Run":
         self._patcher.__enter__()
         return self
@@ -42,18 +76,19 @@ class Run:
         tb: types.TracebackType | None,
     ) -> None:
         self._patcher.__exit__(exc_type, exc_value, tb)
-        if self._experiment_tracking._run is self:
+        if self._experiment_tracking is not None and self._experiment_tracking._run is self:
             if exc_type is not None:
                 self._experiment_tracking._try_log_exception(exc_type, exc_value, tb)
                 self._experiment_tracking.end_run(status=run_metadata.RunStatus.FAILED.value)
             else:
                 self._experiment_tracking.end_run()
 
+    def _fully_qualified_experiment_name(self) -> str:
+        return sql_identifier.get_fully_qualified_name(self.database_name, self.schema_name, self.experiment_name)
+
     def _get_experiment_info(self) -> experiment_info.ExperimentInfo:
         return experiment_info.ExperimentInfo(
-            fully_qualified_name=self._experiment_tracking._sql_client.fully_qualified_object_name(
-                self._experiment_tracking._database_name, self._experiment_tracking._schema_name, self.experiment_name
-            ),
+            fully_qualified_name=self._fully_qualified_experiment_name(),
             run_name=self.name.identifier(),
         )
 
@@ -67,9 +102,27 @@ class Run:
             RuntimeWarning: If the experiment tracking instance was created with warnings_as_errors=True.
         """
         message = f"{sql_error_msg}. {METADATA_SIZE_WARNING_MESSAGE}"
-        if self._experiment_tracking._warnings_as_errors:
+        if self._experiment_tracking is not None and self._experiment_tracking._warnings_as_errors:
             raise RuntimeWarning(message)
         # Once the limit is hit every subsequent log attempt fails the same way, so warn only once.
         if not self._warned_about_metadata_size:
             warnings.warn(message, RuntimeWarning, stacklevel=2)
             self._warned_about_metadata_size = True
+
+    @staticmethod
+    def _load_from_lineage_node(
+        session: snowpark.Session,
+        name: str,
+        version: str,
+    ) -> "Run":
+        database_name, schema_name, experiment_name = identifier.parse_schema_level_object_identifier(name)
+        return Run(
+            experiment_name=sql_identifier.SqlIdentifier(experiment_name),
+            run_name=sql_identifier.SqlIdentifier(version),
+            session=session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
+
+
+lineage_node.DOMAIN_LINEAGE_REGISTRY["experiment"] = Run

snowflake/ml/experiment/experiment_tracking.py

@@ -25,7 +25,6 @@ from snowflake.ml.experiment._client import (
     experiment_tracking_sql_client as sql_client,
 )
 from snowflake.ml.experiment._entities import run_metadata
-from snowflake.ml.lineage import lineage_node
 from snowflake.ml.model import type_hints
 from snowflake.ml.utils import sql_client as sql_client_utils
 from snowflake.snowpark import functions
@@ -328,15 +327,14 @@ class ExperimentTracking:
         else:
             raise RuntimeError("No run is active. Please provide a run_name or start a run.")
 
-        experiment_fqn = self._sql_client.fully_qualified_object_name(
-            self._database_name, self._schema_name, self._experiment.name
-        )
-        run_node = lineage_node.LineageNode(
-            session=self._session,
-            name=experiment_fqn,
-            domain="experiment",
-            version=resolved_run_name.identifier(),
-        )
+        if self._run and self._run.name == resolved_run_name:
+            run_node = self._run
+        else:
+            run_node = entities.Run(
+                experiment_tracking=self,
+                experiment_name=self._experiment.name,
+                run_name=resolved_run_name,
+            )
         logged_models = run_node.lineage(direction="downstream", domain_filter={"model"})
         return [model for model in logged_models if isinstance(model, ml_model.ModelVersion)]
 

snowflake/ml/feature_store/decl/__init__.py

@@ -0,0 +1,168 @@
+"""Public API for the declarative feature store library.
+
+Import the key functions and types from this package:
+
+    from snowflake.ml.feature_store.decl import (
+        load_specs,
+        validate_specs,
+        generate_plan,
+        fetch_applied_state,
+        SpecBatch,
+        AppliedState,
+        Plan,
+        PlanOptions,
+        ValidationResult,
+        PlanOp,
+        AppliedObject,
+        SpecLoadError,
+        ValidationError,
+        DependencyError,
+        StateDriftError,
+        FSBaseType,
+        FeatureViewKind,
+        SourceType,
+        FeatureAggregationMethod,
+        OpKind,
+        TYPE_ALIASES,
+        normalize_type,
+        FSColumn,
+        SpecBase,
+        Entity,
+        StreamingSource,
+        BatchSource,
+        SourceRef,
+        UDF,
+        Feature,
+        FeatureView,
+        FeatureViewRef,
+        FeatureGroup,
+    )
+"""
+
+try:
+    from snowflake.ml.feature_store.decl.api import (
+        discover_project,
+        fetch_applied_state,
+        generate_plan,
+        load_manifest,
+        load_project,
+        load_specs,
+        resolve_datasource_columns,
+        resolve_target,
+        validate_specs,
+    )
+except ModuleNotFoundError:  # lower-layer BUILD targets omit api.py
+    pass
+from snowflake.ml.feature_store.decl.enums import (
+    TYPE_ALIASES,
+    FeatureAggregationMethod,
+    FeatureViewKind,
+    FSBaseType,
+    OpKind,
+    SourceType,
+    normalize_type,
+)
+from snowflake.ml.feature_store.decl.errors import (
+    DependencyError,
+    SpecLoadError,
+    StateDriftError,
+    ValidationError,
+)
+from snowflake.ml.feature_store.decl.manifest import (
+    FSManifest,
+    FSTarget,
+    FSTemplating,
+    InvalidManifestError,
+    ManifestConfigurationError,
+    ManifestNotFoundError,
+    TargetContext,
+)
+from snowflake.ml.feature_store.decl.project_paths import FSProjectPaths
+from snowflake.ml.feature_store.decl.spec_models import (
+    UDF,
+    Backfill,
+    BatchFeatureView,
+    BatchSource,
+    Entity,
+    Feature,
+    FeatureGroup,
+    FeatureView,
+    FeatureViewRef,
+    FSColumn,
+    RealtimeFeatureView,
+    SourceRef,
+    SpecBase,
+    StorageConfig,
+    StreamingFeatureView,
+    StreamingSource,
+)
+from snowflake.ml.feature_store.decl.types import (
+    AppliedObject,
+    AppliedState,
+    Plan,
+    PlanOp,
+    PlanOptions,
+    SpecBatch,
+    ValidationResult,
+)
+
+__all__ = [
+    # api.py
+    "load_specs",
+    "validate_specs",
+    "generate_plan",
+    "fetch_applied_state",
+    "resolve_datasource_columns",
+    "load_manifest",
+    "discover_project",
+    "resolve_target",
+    "load_project",
+    # enums.py
+    "FSBaseType",
+    "FeatureViewKind",
+    "SourceType",
+    "FeatureAggregationMethod",
+    "OpKind",
+    "TYPE_ALIASES",
+    "normalize_type",
+    # errors.py
+    "SpecLoadError",
+    "ValidationError",
+    "DependencyError",
+    "StateDriftError",
+    # manifest.py
+    "FSManifest",
+    "FSTarget",
+    "FSTemplating",
+    "TargetContext",
+    "ManifestNotFoundError",
+    "InvalidManifestError",
+    "ManifestConfigurationError",
+    # project_paths.py
+    "FSProjectPaths",
+    # spec_models.py
+    "FSColumn",
+    "SpecBase",
+    "Entity",
+    "StreamingSource",
+    "BatchSource",
+    "SourceRef",
+    "UDF",
+    "Feature",
+    "FeatureView",
+    "StreamingFeatureView",
+    "BatchFeatureView",
+    "RealtimeFeatureView",
+    "FeatureViewRef",
+    "FeatureGroup",
+    "Backfill",
+    "StorageConfig",
+    # types.py
+    "AppliedObject",
+    "AppliedState",
+    "PlanOp",
+    "Plan",
+    "ValidationResult",
+    "PlanOptions",
+    "SpecBatch",
+]