snowflake-ml-python 1.7.2__py3-none-any.whl → 1.7.4__py3-none-any.whl

snowflake/ml/_internal/utils/identifier.py

@@ -158,8 +158,10 @@ def parse_schema_level_object_identifier(
     res = _SF_SCHEMA_LEVEL_OBJECT_RE.fullmatch(object_name)
     if not res:
         raise ValueError(
-            "Invalid identifier because it does not follow the pattern. "
-            f"It should start with [[database.]schema.]object. Getting {object_name}"
+            f"Invalid object name `{object_name}` cannot be parsed as a SQL identifier. "
+            "Alphanumeric characters and underscores are permitted. "
+            "See https://docs.snowflake.com/en/sql-reference/identifiers-syntax for "
+            "more information."
         )
     return (
         res.group("db"),

snowflake/ml/data/__init__.py

@@ -1,5 +1,8 @@
+from pkgutil import extend_path
+
 from .data_connector import DataConnector
 from .data_ingestor import DataIngestor, DataIngestorType
 from .data_source import DataFrameInfo, DatasetInfo, DataSource
 
 __all__ = ["DataConnector", "DataSource", "DataFrameInfo", "DatasetInfo", "DataIngestor", "DataIngestorType"]
+__path__ = extend_path(__path__, __name__)

snowflake/ml/data/_internal/arrow_ingestor.py

@@ -2,7 +2,7 @@ import collections
 import logging
 import os
 import time
-from typing import Any, Deque, Dict, Iterator, List, Optional, Union
+from typing import Any, Deque, Dict, Iterator, List, Optional, Sequence, Union
 
 import numpy as np
 import numpy.typing as npt
@@ -47,7 +47,7 @@ class ArrowIngestor(data_ingestor.DataIngestor):
     def __init__(
         self,
         session: snowpark.Session,
-        data_sources: List[data_source.DataSource],
+        data_sources: Sequence[data_source.DataSource],
         format: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
@@ -60,14 +60,14 @@ class ArrowIngestor(data_ingestor.DataIngestor):
             kwargs: Miscellaneous arguments passed to underlying PyArrow Dataset initializer.
         """
         self._session = session
-        self._data_sources = data_sources
+        self._data_sources = list(data_sources)
         self._format = format
         self._kwargs = kwargs
 
         self._schema: Optional[pa.Schema] = None
 
     @classmethod
-    def from_sources(cls, session: snowpark.Session, sources: List[data_source.DataSource]) -> "ArrowIngestor":
+    def from_sources(cls, session: snowpark.Session, sources: Sequence[data_source.DataSource]) -> "ArrowIngestor":
         return cls(session, sources)
 
     @property

snowflake/ml/data/data_connector.py

@@ -1,5 +1,16 @@
 import os
-from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Type, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Sequence,
+    Type,
+    TypeVar,
+    cast,
+)
 
 import numpy.typing as npt
 from typing_extensions import deprecated
@@ -12,6 +23,7 @@ from snowflake.ml.modeling._internal.constants import (
     IN_ML_RUNTIME_ENV_VAR,
     USE_OPTIMIZED_DATA_INGESTOR,
 )
+from snowflake.snowpark import context as sf_context
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -35,8 +47,10 @@ class DataConnector:
     def __init__(
         self,
         ingestor: data_ingestor.DataIngestor,
+        **kwargs: Any,
     ) -> None:
         self._ingestor = ingestor
+        self._kwargs = kwargs
 
     @classmethod
     @snowpark._internal.utils.private_preview(version="1.6.0")
@@ -44,20 +58,34 @@ class DataConnector:
         cls: Type[DataConnectorType],
         df: snowpark.DataFrame,
         ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> DataConnectorType:
         if len(df.queries["queries"]) != 1 or len(df.queries["post_actions"]) != 0:
             raise ValueError("DataFrames with multiple queries and/or post-actions not supported")
-        source = data_source.DataFrameInfo(df.queries["queries"][0])
-        assert df._session is not None
-        return cls.from_sources(df._session, [source], ingestor_class=ingestor_class, **kwargs)
+        return cast(
+            DataConnectorType,
+            cls.from_sql(df.queries["queries"][0], session=df._session, ingestor_class=ingestor_class, **kwargs),
+        )
+
+    @classmethod
+    @snowpark._internal.utils.private_preview(version="1.7.3")
+    def from_sql(
+        cls: Type[DataConnectorType],
+        query: str,
+        session: Optional[snowpark.Session] = None,
+        ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
+        **kwargs: Any,
+    ) -> DataConnectorType:
+        session = session or sf_context.get_active_session()
+        source = data_source.DataFrameInfo(query)
+        return cls.from_sources(session, [source], ingestor_class=ingestor_class, **kwargs)
 
     @classmethod
     def from_dataset(
         cls: Type[DataConnectorType],
         ds: "dataset.Dataset",
         ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> DataConnectorType:
         dsv = ds.selected_version
         assert dsv is not None
@@ -75,9 +103,9 @@ class DataConnector:
     def from_sources(
         cls: Type[DataConnectorType],
         session: snowpark.Session,
-        sources: List[data_source.DataSource],
+        sources: Sequence[data_source.DataSource],
         ingestor_class: Optional[Type[data_ingestor.DataIngestor]] = None,
-        **kwargs: Any
+        **kwargs: Any,
     ) -> DataConnectorType:
         ingestor_class = ingestor_class or cls.DEFAULT_INGESTOR_CLASS
         ingestor = ingestor_class.from_sources(session, sources)
@@ -130,7 +158,11 @@ class DataConnector:
         func_params_to_log=["batch_size", "shuffle", "drop_last_batch"],
     )
     def to_torch_datapipe(
-        self, *, batch_size: int, shuffle: bool = False, drop_last_batch: bool = True
+        self,
+        *,
+        batch_size: int,
+        shuffle: bool = False,
+        drop_last_batch: bool = True,
     ) -> "torch_data.IterDataPipe":  # type: ignore[type-arg]
         """Transform the Snowflake data into a ready-to-use Pytorch datapipe.
 
@@ -149,8 +181,13 @@ class DataConnector:
         """
         from snowflake.ml.data import torch_utils
 
+        expand_dims = self._kwargs.get("expand_dims", True)
         return torch_utils.TorchDataPipeWrapper(
-            self._ingestor, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last_batch
+            self._ingestor,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last_batch,
+            expand_dims=expand_dims,
         )
 
     @telemetry.send_api_usage_telemetry(
@@ -179,8 +216,13 @@ class DataConnector:
         """
         from snowflake.ml.data import torch_utils
 
+        expand_dims = self._kwargs.get("expand_dims", True)
         return torch_utils.TorchDatasetWrapper(
-            self._ingestor, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last_batch
+            self._ingestor,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            drop_last=drop_last_batch,
+            expand_dims=expand_dims,
         )
 
     @telemetry.send_api_usage_telemetry(

snowflake/ml/data/data_ingestor.py

@@ -6,6 +6,7 @@ from typing import (
     List,
     Optional,
     Protocol,
+    Sequence,
     Type,
     TypeVar,
 )
@@ -25,7 +26,7 @@ DataIngestorType = TypeVar("DataIngestorType", bound="DataIngestor")
 class DataIngestor(Protocol):
     @classmethod
     def from_sources(
-        cls: Type[DataIngestorType], session: snowpark.Session, sources: List[data_source.DataSource]
+        cls: Type[DataIngestorType], session: snowpark.Session, sources: Sequence[data_source.DataSource]
     ) -> DataIngestorType:
         raise NotImplementedError
 

snowflake/ml/data/torch_utils.py

@@ -17,6 +17,7 @@ class TorchDatasetWrapper(torch.utils.data.IterableDataset[Dict[str, Any]]):
         batch_size: Optional[int],
         shuffle: bool = False,
         drop_last: bool = False,
+        expand_dims: bool = True,
     ) -> None:
         """Not intended for direct usage. Use DataConnector.to_torch_dataset() instead"""
         squeeze = False
@@ -29,6 +30,7 @@ class TorchDatasetWrapper(torch.utils.data.IterableDataset[Dict[str, Any]]):
         self._shuffle = shuffle
         self._drop_last = drop_last
         self._squeeze_outputs = squeeze
+        self._expand_dims = expand_dims
 
     def __iter__(self) -> Iterator[Dict[str, Union[npt.NDArray[Any], List[Any]]]]:
         max_idx = 0
@@ -47,7 +49,10 @@ class TorchDatasetWrapper(torch.utils.data.IterableDataset[Dict[str, Any]]):
         ):
             # Skip indices during multi-process data loading to prevent data duplication
             if counter == filter_idx:
-                yield {k: _preprocess_array(v, squeeze=self._squeeze_outputs) for k, v in batch.items()}
+                yield {
+                    k: _preprocess_array(v, squeeze=self._squeeze_outputs, expand_dims=self._expand_dims)
+                    for k, v in batch.items()
+                }
             if counter < max_idx:
                 counter += 1
             else:
@@ -58,13 +63,21 @@ class TorchDataPipeWrapper(TorchDatasetWrapper, torch.utils.data.IterDataPipe[Di
     """Wrap a DataIngestor into a PyTorch IterDataPipe"""
 
     def __init__(
-        self, ingestor: data_ingestor.DataIngestor, *, batch_size: int, shuffle: bool = False, drop_last: bool = False
+        self,
+        ingestor: data_ingestor.DataIngestor,
+        *,
+        batch_size: int,
+        shuffle: bool = False,
+        drop_last: bool = False,
+        expand_dims: bool = True,
     ) -> None:
         """Not intended for direct usage. Use DataConnector.to_torch_datapipe() instead"""
-        super().__init__(ingestor, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last)
+        super().__init__(ingestor, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, expand_dims=expand_dims)
 
 
-def _preprocess_array(arr: npt.NDArray[Any], squeeze: bool = False) -> Union[npt.NDArray[Any], List[np.object_]]:
+def _preprocess_array(
+    arr: npt.NDArray[Any], squeeze: bool = False, expand_dims: bool = True
+) -> Union[npt.NDArray[Any], List[np.object_]]:
     """Preprocesses batch column values."""
     single_dimensional = arr.ndim < 2 and not arr.dtype == np.object_
 
@@ -73,7 +86,7 @@ def _preprocess_array(arr: npt.NDArray[Any], squeeze: bool = False) -> Union[npt
         arr = arr.squeeze(axis=0)
 
     # For single dimensional data,
-    if single_dimensional:
+    if single_dimensional and expand_dims:
         axis = 0 if arr.ndim == 0 else 1
         arr = np.expand_dims(arr, axis=axis)
 

snowflake/ml/dataset/dataset.py

@@ -419,7 +419,6 @@ class Dataset(lineage_node.LineageNode):
     @telemetry.send_api_usage_telemetry(project=_PROJECT)
     def delete(self) -> None:
         """Delete Dataset and all contained versions"""
-        # TODO: Check and warn if any versions exist
         self._session.sql(f"DROP DATASET {self.fully_qualified_name}").collect(
             statement_params=_TELEMETRY_STATEMENT_PARAMS
         )

snowflake/ml/feature_store/examples/example_helper.py

@@ -45,8 +45,9 @@ class ExampleHelper:
         """Return a dataframe object about descriptions of all examples."""
         root_dir = Path(__file__).parent
         rows = []
+        hide_folders = ["citibike_trip_features", "source_data"]
         for f_name in os.listdir(root_dir):
-            if os.path.isdir(os.path.join(root_dir, f_name)) and f_name[0].isalpha() and f_name != "source_data":
+            if os.path.isdir(os.path.join(root_dir, f_name)) and f_name[0].isalpha() and f_name not in hide_folders:
                 source_file_path = root_dir.joinpath(f"{f_name}/source.yaml")
                 source_dict = self._read_yaml(str(source_file_path))
                 rows.append((f_name, source_dict["model_category"], source_dict["desc"], source_dict["label_columns"]))

snowflake/ml/fileset/fileset.py

@@ -2,6 +2,8 @@ import functools
 import inspect
 from typing import Any, Callable, List, Optional
 
+from typing_extensions import deprecated
+
 from snowflake import snowpark
 from snowflake.connector import connection
 from snowflake.ml._internal import telemetry
@@ -11,11 +13,9 @@ from snowflake.ml._internal.exceptions import (
     fileset_error_messages,
     fileset_errors,
 )
-from snowflake.ml._internal.utils import (
-    identifier,
-    import_utils,
-    snowpark_dataframe_utils,
-)
+from snowflake.ml._internal.utils import identifier, snowpark_dataframe_utils
+from snowflake.ml.data import data_connector
+from snowflake.ml.data._internal import arrow_ingestor
 from snowflake.ml.fileset import sfcfs
 from snowflake.snowpark import exceptions as snowpark_exceptions, functions
 
@@ -44,6 +44,10 @@ def _raise_if_deleted(func: Callable[..., Any]) -> Callable[..., Any]:
     return raise_if_deleted_helper
 
 
+@deprecated(
+    "FileSet is deprecated and will be removed in a future release."
+    " Use snowflake.ml.dataset.Dataset and snowflake.ml.data.DataConnector instead"
+)
 class FileSet:
     """A FileSet represents an immutable snapshot of the result of a query in the form of files."""
 
@@ -285,6 +289,16 @@ class FileSet:
         """Get the Snowflake absolute path to this FileSet directory."""
         return _fileset_absolute_path(self._target_stage_loc, self.name)
 
+    def _to_data_connector(self) -> data_connector.DataConnector:
+        self._fs.optimize_read(self._list_files())
+        ingester = arrow_ingestor.ArrowIngestor(
+            self._snowpark_session,
+            self._list_files(),
+            format="parquet",
+            filesystem=self._fs,
+        )
+        return data_connector.DataConnector(ingester, expand_dims=False)
+
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
     )
@@ -362,13 +376,9 @@ class FileSet:
         ----
         {'_COL_1':[10]}
         """
-        IterableWrapper, _ = import_utils.import_or_get_dummy("torchdata.datapipes.iter.IterableWrapper")
-        torch_datapipe_module, _ = import_utils.import_or_get_dummy("snowflake.ml.fileset.torch_datapipe")
-
-        self._fs.optimize_read(self._list_files())
-
-        input_dp = IterableWrapper(self._list_files())
-        return torch_datapipe_module.ReadAndParseParquet(input_dp, self._fs, batch_size, shuffle, drop_last_batch)
+        return self._to_data_connector().to_torch_datapipe(
+            batch_size=batch_size, shuffle=shuffle, drop_last_batch=drop_last_batch
+        )
 
     @telemetry.send_api_usage_telemetry(
         project=_PROJECT,
@@ -402,12 +412,8 @@ class FileSet:
         ----
         {'_COL_1': <tf.Tensor: shape=(1,), dtype=int64, numpy=[10]>}
         """
-        tf_dataset_module, _ = import_utils.import_or_get_dummy("snowflake.ml.fileset.tf_dataset")
-
-        self._fs.optimize_read(self._list_files())
-
-        return tf_dataset_module.read_and_parse_parquet(
-            self._list_files(), self._fs, batch_size, shuffle, drop_last_batch
+        return self._to_data_connector().to_tf_dataset(
+            batch_size=batch_size, shuffle=shuffle, drop_last_batch=drop_last_batch
         )
 
     @telemetry.send_api_usage_telemetry(

snowflake/ml/jobs/__init__.py

@@ -0,0 +1,21 @@
+from snowflake.ml.jobs._utils.types import JOB_STATUS
+from snowflake.ml.jobs.decorators import remote
+from snowflake.ml.jobs.job import MLJob
+from snowflake.ml.jobs.manager import (
+    delete_job,
+    get_job,
+    list_jobs,
+    submit_directory,
+    submit_file,
+)
+
+__all__ = [
+    "remote",
+    "submit_file",
+    "submit_directory",
+    "list_jobs",
+    "get_job",
+    "delete_job",
+    "MLJob",
+    "JOB_STATUS",
+]

snowflake/ml/jobs/_utils/constants.py

@@ -0,0 +1,51 @@
+from snowflake.ml._internal.utils.snowflake_env import SnowflakeCloudType
+from snowflake.ml.jobs._utils.types import ComputeResources
+
+# SPCS specification constants
+DEFAULT_CONTAINER_NAME = "main"
+PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
+
+# Default container image information
+DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
+DEFAULT_IMAGE_CPU = "st_plat/runtime/x86/runtime_image/snowbooks"
+DEFAULT_IMAGE_GPU = "st_plat/runtime/x86/generic_gpu/runtime_image/snowbooks"
+DEFAULT_IMAGE_TAG = "0.8.0"
+DEFAULT_ENTRYPOINT_PATH = "func.py"
+
+# Percent of container memory to allocate for /dev/shm volume
+MEMORY_VOLUME_SIZE = 0.3
+
+# Job status polling constants
+JOB_POLL_INITIAL_DELAY_SECONDS = 0.1
+JOB_POLL_MAX_DELAY_SECONDS = 1
+
+# Compute pool resource information
+# TODO: Query Snowflake for resource information instead of relying on this hardcoded
+#       table from https://docs.snowflake.com/en/sql-reference/sql/create-compute-pool
+COMMON_INSTANCE_FAMILIES = {
+    "CPU_X64_XS": ComputeResources(cpu=1, memory=6),
+    "CPU_X64_S": ComputeResources(cpu=3, memory=13),
+    "CPU_X64_M": ComputeResources(cpu=6, memory=28),
+    "CPU_X64_L": ComputeResources(cpu=28, memory=116),
+    "HIGHMEM_X64_S": ComputeResources(cpu=6, memory=58),
+}
+AWS_INSTANCE_FAMILIES = {
+    "HIGHMEM_X64_M": ComputeResources(cpu=28, memory=240),
+    "HIGHMEM_X64_L": ComputeResources(cpu=124, memory=984),
+    "GPU_NV_S": ComputeResources(cpu=6, memory=27, gpu=1, gpu_type="A10G"),
+    "GPU_NV_M": ComputeResources(cpu=44, memory=178, gpu=4, gpu_type="A10G"),
+    "GPU_NV_L": ComputeResources(cpu=92, memory=1112, gpu=8, gpu_type="A100"),
+}
+AZURE_INSTANCE_FAMILIES = {
+    "HIGHMEM_X64_M": ComputeResources(cpu=28, memory=244),
+    "HIGHMEM_X64_L": ComputeResources(cpu=92, memory=654),
+    "GPU_NV_XS": ComputeResources(cpu=3, memory=26, gpu=1, gpu_type="T4"),
+    "GPU_NV_SM": ComputeResources(cpu=32, memory=424, gpu=1, gpu_type="A10"),
+    "GPU_NV_2M": ComputeResources(cpu=68, memory=858, gpu=2, gpu_type="A10"),
+    "GPU_NV_3M": ComputeResources(cpu=44, memory=424, gpu=2, gpu_type="A100"),
+    "GPU_NV_SL": ComputeResources(cpu=92, memory=858, gpu=4, gpu_type="A100"),
+}
+CLOUD_INSTANCE_FAMILIES = {
+    SnowflakeCloudType.AWS: AWS_INSTANCE_FAMILIES,
+    SnowflakeCloudType.AZURE: AZURE_INSTANCE_FAMILIES,
+}