snowflake-ml-python 1.9.1__py3-none-any.whl → 1.10.0__py3-none-any.whl

snowflake/ml/jobs/_utils/stage_utils.py

@@ -14,7 +14,10 @@ _SNOWURL_PATH_RE = re.compile(
     r"(?P<path>versions(?:/(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)$"
 )
 
-_STAGEF_PATH_RE = re.compile(r"^@(?P<stage>~|%?\w+)(?:/(?P<relpath>[\w\-./]*))?$")
+# Break long regex into two main parts
+_STAGE_PATTERN = rf"~|%?(?:(?:{identifier._SF_IDENTIFIER}\.?){{,2}}{identifier._SF_IDENTIFIER})"
+_RELPATH_PATTERN = r"[\w\-./]*"
+_STAGEF_PATH_RE = re.compile(rf"^@(?P<stage>{_STAGE_PATTERN})(?:/(?P<relpath>{_RELPATH_PATTERN}))?$")
 
 
 class StagePath:
@@ -29,6 +32,14 @@ class StagePath:
         self._root = self._raw_path[0:start].rstrip("/") if relpath else self._raw_path.rstrip("/")
         self._path = Path(relpath or "")
 
+    @property
+    def parts(self) -> tuple[str, ...]:
+        return self._path.parts
+
+    @property
+    def name(self) -> str:
+        return self._path.name
+
     @property
     def parent(self) -> "StagePath":
         if self._path.parent == Path(""):
@@ -51,18 +62,28 @@ class StagePath:
         else:
             return f"{self.root}/{path}"
 
-    def is_relative_to(self, path: Union[str, PathLike[str], "StagePath"]) -> bool:
+    def is_relative_to(self, *other: Union[str, os.PathLike[str]]) -> bool:
+        if not other:
+            raise TypeError("is_relative_to() requires at least one argument")
+        # For now, we only support a single argument, like pathlib.Path in Python < 3.12
+        path = other[0]
         stage_path = path if isinstance(path, StagePath) else StagePath(os.fspath(path))
         if stage_path.root == self.root:
             return self._path.is_relative_to(stage_path._path)
         else:
             return False
 
-    def relative_to(self, path: Union[str, PathLike[str], "StagePath"]) -> PurePath:
+    def relative_to(self, *other: Union[str, os.PathLike[str]]) -> PurePath:
+        if not other:
+            raise TypeError("relative_to() requires at least one argument")
+        if not self.is_relative_to(*other):
+            raise ValueError(f"{other} does not start with {self._raw_path}")
+        path = other[0]
         stage_path = path if isinstance(path, StagePath) else StagePath(os.fspath(path))
         if self.root == stage_path.root:
             return self._path.relative_to(stage_path._path)
-        raise ValueError(f"{self._raw_path} does not start with {stage_path._raw_path}")
+        else:
+            raise ValueError(f"{self._raw_path} does not start with {stage_path._raw_path}")
 
     def absolute(self) -> "StagePath":
         return self
@@ -88,6 +109,9 @@ class StagePath:
     def __str__(self) -> str:
         return self.as_posix()
 
+    def __repr__(self) -> str:
+        return f"StagePath('{self.as_posix()}')"
+
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, StagePath):
             raise NotImplementedError
@@ -96,24 +120,16 @@ class StagePath:
     def __fspath__(self) -> str:
         return self._compose_path(self._path)
 
-    def joinpath(self, *args: Union[str, PathLike[str], "StagePath"]) -> "StagePath":
+    def joinpath(self, *args: Union[str, PathLike[str]]) -> "StagePath":
         path = self
         for arg in args:
             path = path._make_child(arg)
         return path
 
-    def _make_child(self, path: Union[str, PathLike[str], "StagePath"]) -> "StagePath":
+    def _make_child(self, path: Union[str, PathLike[str]]) -> "StagePath":
         stage_path = path if isinstance(path, StagePath) else StagePath(os.fspath(path))
         if self.root == stage_path.root:
             child_path = self._path.joinpath(stage_path._path)
             return StagePath(self._compose_path(child_path))
         else:
             return stage_path
-
-
-def identify_stage_path(path: str) -> Union[StagePath, Path]:
-    try:
-        stage_path = StagePath(path)
-    except ValueError:
-        return Path(path)
-    return stage_path

snowflake/ml/jobs/_utils/types.py

@@ -1,8 +1,7 @@
+import os
 from dataclasses import dataclass
 from pathlib import PurePath
-from typing import Literal, Optional, Union
-
-from snowflake.ml.jobs._utils import stage_utils
+from typing import Iterator, Literal, Optional, Protocol, Union, runtime_checkable
 
 JOB_STATUS = Literal[
     "PENDING",
@@ -15,9 +14,70 @@ JOB_STATUS = Literal[
 ]
 
 
+@runtime_checkable
+class PayloadPath(Protocol):
+    """A protocol for path-like objects used in this module, covering methods from pathlib.Path and StagePath."""
+
+    @property
+    def name(self) -> str:
+        ...
+
+    @property
+    def suffix(self) -> str:
+        ...
+
+    @property
+    def parent(self) -> "PayloadPath":
+        ...
+
+    def exists(self) -> bool:
+        ...
+
+    def is_file(self) -> bool:
+        ...
+
+    def is_absolute(self) -> bool:
+        ...
+
+    def absolute(self) -> "PayloadPath":
+        ...
+
+    def joinpath(self, *other: Union[str, os.PathLike[str]]) -> "PayloadPath":
+        ...
+
+    def as_posix(self) -> str:
+        ...
+
+    def is_relative_to(self, *other: Union[str, os.PathLike[str]]) -> bool:
+        ...
+
+    def relative_to(self, *other: Union[str, os.PathLike[str]]) -> PurePath:
+        ...
+
+    def __fspath__(self) -> str:
+        ...
+
+    def __str__(self) -> str:
+        ...
+
+    def __repr__(self) -> str:
+        ...
+
+
+@dataclass
+class PayloadSpec:
+    """Represents a payload item to be uploaded."""
+
+    source_path: PayloadPath
+    remote_relative_path: Optional[PurePath] = None
+
+    def __iter__(self) -> Iterator[Union[PayloadPath, Optional[PurePath]]]:
+        return iter((self.source_path, self.remote_relative_path))
+
+
 @dataclass(frozen=True)
 class PayloadEntrypoint:
-    file_path: Union[PurePath, stage_utils.StagePath]
+    file_path: PayloadPath
     main_func: Optional[str]
 
 

snowflake/ml/jobs/job.py

@@ -3,6 +3,7 @@ import logging
 import os
 import time
 from functools import cached_property
+from pathlib import Path
 from typing import Any, Generic, Literal, Optional, TypeVar, Union, cast, overload
 
 import yaml
@@ -95,10 +96,24 @@ class MLJob(Generic[T], SerializableSessionMixin):
     @property
     def _result_path(self) -> str:
         """Get the job's result file location."""
-        result_path = self._container_spec["env"].get(constants.RESULT_PATH_ENV_VAR)
-        if result_path is None:
+        result_path_str = self._container_spec["env"].get(constants.RESULT_PATH_ENV_VAR)
+        if result_path_str is None:
             raise RuntimeError(f"Job {self.name} doesn't have a result path configured")
-        return f"{self._stage_path}/{result_path}"
+        volume_mounts = self._container_spec["volumeMounts"]
+        stage_mount_str = next(v for v in volume_mounts if v.get("name") == constants.STAGE_VOLUME_NAME)["mountPath"]
+
+        result_path = Path(result_path_str)
+        stage_mount = Path(stage_mount_str)
+        try:
+            relative_path = result_path.relative_to(stage_mount)
+        except ValueError:
+            if result_path.is_absolute():
+                raise ValueError(
+                    f"Result path {result_path} is absolute, but should be relative to stage mount {stage_mount}"
+                )
+            relative_path = result_path
+
+        return f"{self._stage_path}/{relative_path.as_posix()}"
 
     @overload
     def get_logs(
@@ -181,7 +196,10 @@ class MLJob(Generic[T], SerializableSessionMixin):
         start_time = time.monotonic()
         warning_shown = False
         while (status := self.status) not in TERMINAL_JOB_STATUSES:
-            if status == "PENDING" and not warning_shown:
+            elapsed = time.monotonic() - start_time
+            if elapsed >= timeout >= 0:
+                raise TimeoutError(f"Job {self.name} did not complete within {timeout} seconds")
+            elif status == "PENDING" and not warning_shown and elapsed >= 2:  # Only show warning after 2s
                 pool_info = _get_compute_pool_info(self._session, self._compute_pool)
                 if (pool_info.max_nodes - pool_info.active_nodes) < self.min_instances:
                     logger.warning(
@@ -189,8 +207,6 @@ class MLJob(Generic[T], SerializableSessionMixin):
                         f"{self.min_instances} nodes required). Job execution may be delayed."
                     )
                     warning_shown = True
-            if timeout >= 0 and (elapsed := time.monotonic() - start_time) >= timeout:
-                raise TimeoutError(f"Job {self.name} did not complete within {elapsed} seconds")
             time.sleep(delay)
             delay = min(delay * 1.2, constants.JOB_POLL_MAX_DELAY_SECONDS)  # Exponential backoff
         return self.status

snowflake/ml/jobs/manager.py

@@ -447,6 +447,10 @@ def _submit_job(
     spec_overrides = kwargs.pop("spec_overrides", None)
     enable_metrics = kwargs.pop("enable_metrics", True)
     query_warehouse = kwargs.pop("query_warehouse", None)
+    additional_payloads = kwargs.pop("additional_payloads", None)
+
+    if additional_payloads:
+        logger.warning("'additional_payloads' is in private preview since 1.9.1. Do not use it in production.")
 
     # Warn if there are unknown kwargs
     if kwargs:
@@ -477,9 +481,7 @@ def _submit_job(
 
     # Upload payload
     uploaded_payload = payload_utils.JobPayload(
-        source,
-        entrypoint=entrypoint,
-        pip_requirements=pip_requirements,
+        source, entrypoint=entrypoint, pip_requirements=pip_requirements, additional_payloads=additional_payloads
     ).upload(session, stage_path)
 
     # Generate service spec

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

@@ -1,5 +1,4 @@
 import enum
-import logging
 import pathlib
 import tempfile
 import warnings
@@ -881,6 +880,7 @@ class ModelVersion(lineage_node.LineageNode):
 
         Raises:
             ValueError: Illegal external access integration arguments.
+            exceptions.SnowparkSQLException: if service already exists.
 
         Returns:
             If `block=True`, return result information about service creation from server.
@@ -891,16 +891,6 @@ class ModelVersion(lineage_node.LineageNode):
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
-        # Check root logger level and emit warning if needed
-        root_logger = logging.getLogger()
-        if root_logger.level in (logging.WARNING, logging.ERROR):
-            warnings.warn(
-                "Suppressing service logs. Set the log level to INFO if you would like "
-                "verbose service logs (e.g., logging.getLogger().setLevel(logging.INFO)).",
-                UserWarning,
-                stacklevel=2,
-            )
-
         if build_external_access_integration is not None:
             msg = (
                 "`build_external_access_integration` is deprecated. "
@@ -917,39 +907,60 @@ class ModelVersion(lineage_node.LineageNode):
 
         service_db_id, service_schema_id, service_id = sql_identifier.parse_fully_qualified_name(service_name)
         image_repo_db_id, image_repo_schema_id, image_repo_id = sql_identifier.parse_fully_qualified_name(image_repo)
-        return self._service_ops.create_service(
-            database_name=None,
-            schema_name=None,
-            model_name=self._model_name,
-            version_name=self._version_name,
-            service_database_name=service_db_id,
-            service_schema_name=service_schema_id,
-            service_name=service_id,
-            image_build_compute_pool_name=(
-                sql_identifier.SqlIdentifier(image_build_compute_pool)
-                if image_build_compute_pool
-                else sql_identifier.SqlIdentifier(service_compute_pool)
-            ),
-            service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
-            image_repo_database_name=image_repo_db_id,
-            image_repo_schema_name=image_repo_schema_id,
-            image_repo_name=image_repo_id,
-            ingress_enabled=ingress_enabled,
-            max_instances=max_instances,
-            cpu_requests=cpu_requests,
-            memory_requests=memory_requests,
-            gpu_requests=gpu_requests,
-            num_workers=num_workers,
-            max_batch_rows=max_batch_rows,
-            force_rebuild=force_rebuild,
-            build_external_access_integrations=(
-                None
-                if build_external_access_integrations is None
-                else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
-            ),
-            block=block,
-            statement_params=statement_params,
-        )
+
+        from snowflake.ml.model import event_handler
+        from snowflake.snowpark import exceptions
+
+        model_event_handler = event_handler.ModelEventHandler()
+
+        with model_event_handler.status("Creating model inference service", total=6, block=block) as status:
+            try:
+                result = self._service_ops.create_service(
+                    database_name=None,
+                    schema_name=None,
+                    model_name=self._model_name,
+                    version_name=self._version_name,
+                    service_database_name=service_db_id,
+                    service_schema_name=service_schema_id,
+                    service_name=service_id,
+                    image_build_compute_pool_name=(
+                        sql_identifier.SqlIdentifier(image_build_compute_pool)
+                        if image_build_compute_pool
+                        else sql_identifier.SqlIdentifier(service_compute_pool)
+                    ),
+                    service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
+                    image_repo=image_repo,
+                    ingress_enabled=ingress_enabled,
+                    max_instances=max_instances,
+                    cpu_requests=cpu_requests,
+                    memory_requests=memory_requests,
+                    gpu_requests=gpu_requests,
+                    num_workers=num_workers,
+                    max_batch_rows=max_batch_rows,
+                    force_rebuild=force_rebuild,
+                    build_external_access_integrations=(
+                        None
+                        if build_external_access_integrations is None
+                        else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
+                    ),
+                    block=block,
+                    statement_params=statement_params,
+                    progress_status=status,
+                )
+                status.update(label="Model service created successfully", state="complete", expanded=False)
+                return result
+            except exceptions.SnowparkSQLException as e:
+                # Check if the error is because the service already exists
+                if "already exists" in str(e).lower() or "100132" in str(
+                    e
+                ):  # 100132 is Snowflake error code for object already exists
+                    status.update("service already exists")
+                    status.complete()
+                    status.update(label="Service already exists", state="error", expanded=False)
+                    raise
+                else:
+                    status.update(label="Service creation failed", state="error", expanded=False)
+                    raise
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
@@ -1045,7 +1056,6 @@ class ModelVersion(lineage_node.LineageNode):
         )
         target_function_info = self._get_function_info(function_name=function_name)
         job_db_id, job_schema_id, job_id = sql_identifier.parse_fully_qualified_name(job_name)
-        image_repo_db_id, image_repo_schema_id, image_repo_id = sql_identifier.parse_fully_qualified_name(image_repo)
         output_table_db_id, output_table_schema_id, output_table_id = sql_identifier.parse_fully_qualified_name(
             output_table_name
         )
@@ -1064,9 +1074,7 @@ class ModelVersion(lineage_node.LineageNode):
             job_name=job_id,
             compute_pool_name=sql_identifier.SqlIdentifier(compute_pool),
             warehouse_name=sql_identifier.SqlIdentifier(warehouse),
-            image_repo_database_name=image_repo_db_id,
-            image_repo_schema_name=image_repo_schema_id,
-            image_repo_name=image_repo_id,
+            image_repo=image_repo,
             output_table_database_name=output_table_db_id,
             output_table_schema_name=output_table_schema_id,
             output_table_name=output_table_id,