snowflake-ml-python 1.8.1__py3-none-any.whl → 1.8.3__py3-none-any.whl

snowflake/ml/jobs/job.py

@@ -1,20 +1,32 @@
 import time
-from typing import Any, List, Optional, cast
+from typing import Any, Generic, Literal, Optional, TypeVar, Union, cast, overload
+
+import yaml
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
-from snowflake.ml.jobs._utils import constants, types
+from snowflake.ml.jobs._utils import constants, interop_utils, types
 from snowflake.snowpark import context as sp_context
 
 _PROJECT = "MLJob"
 TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "INTERNAL_ERROR"}
 
+T = TypeVar("T")
+
 
-class MLJob:
-    def __init__(self, id: str, session: Optional[snowpark.Session] = None) -> None:
+class MLJob(Generic[T]):
+    def __init__(
+        self,
+        id: str,
+        service_spec: Optional[dict[str, Any]] = None,
+        session: Optional[snowpark.Session] = None,
+    ) -> None:
         self._id = id
+        self._service_spec_cached: Optional[dict[str, Any]] = service_spec
         self._session = session or sp_context.get_active_session()
+
         self._status: types.JOB_STATUS = "PENDING"
+        self._result: Optional[interop_utils.ExecutionResult] = None
 
     @property
     def id(self) -> str:
@@ -29,33 +41,76 @@ class MLJob:
             self._status = _get_status(self._session, self.id)
         return self._status
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
-    def get_logs(self, limit: int = -1) -> str:
+    @property
+    def _service_spec(self) -> dict[str, Any]:
+        """Get the job's service spec."""
+        if not self._service_spec_cached:
+            self._service_spec_cached = _get_service_spec(self._session, self.id)
+        return self._service_spec_cached
+
+    @property
+    def _container_spec(self) -> dict[str, Any]:
+        """Get the job's main container spec."""
+        containers = self._service_spec["spec"]["containers"]
+        container_spec = next(c for c in containers if c["name"] == constants.DEFAULT_CONTAINER_NAME)
+        return cast(dict[str, Any], container_spec)
+
+    @property
+    def _stage_path(self) -> str:
+        """Get the job's artifact storage stage location."""
+        volumes = self._service_spec["spec"]["volumes"]
+        stage_path = next(v for v in volumes if v["name"] == constants.STAGE_VOLUME_NAME)["source"]
+        return cast(str, stage_path)
+
+    @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:
+            raise RuntimeError(f"Job {self.id} doesn't have a result path configured")
+        return f"{self._stage_path}/{result_path}"
+
+    @overload
+    def get_logs(self, limit: int = -1, instance_id: Optional[int] = None, *, as_list: Literal[True]) -> list[str]:
+        ...
+
+    @overload
+    def get_logs(self, limit: int = -1, instance_id: Optional[int] = None, *, as_list: Literal[False] = False) -> str:
+        ...
+
+    def get_logs(
+        self, limit: int = -1, instance_id: Optional[int] = None, *, as_list: bool = False
+    ) -> Union[str, list[str]]:
         """
         Return the job's execution logs.
 
         Args:
             limit: The maximum number of lines to return. Negative values are treated as no limit.
+            instance_id: Optional instance ID to get logs from a specific instance.
+                         If not provided, returns logs from the head node.
+            as_list: If True, returns logs as a list of lines. Otherwise, returns logs as a single string.
 
         Returns:
             The job's execution logs.
         """
-        logs = _get_logs(self._session, self.id, limit)
+        logs = _get_logs(self._session, self.id, limit, instance_id)
         assert isinstance(logs, str)  # mypy
+        if as_list:
+            return logs.splitlines()
         return logs
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
-    def show_logs(self, limit: int = -1) -> None:
+    def show_logs(self, limit: int = -1, instance_id: Optional[int] = None) -> None:
         """
         Display the job's execution logs.
 
         Args:
             limit: The maximum number of lines to display. Negative values are treated as no limit.
+            instance_id: Optional instance ID to get logs from a specific instance.
+                         If not provided, displays logs from the head node.
         """
-        print(self.get_logs(limit))  # noqa: T201: we need to print here.
+        print(self.get_logs(limit, instance_id, as_list=False))  # noqa: T201: we need to print here.
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
-    @telemetry.send_api_usage_telemetry(project=_PROJECT)
+    @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["timeout"])
     def wait(self, timeout: float = -1) -> types.JOB_STATUS:
         """
         Block until completion. Returns completion status.
@@ -78,20 +133,58 @@ class MLJob:
             delay = min(delay * 2, constants.JOB_POLL_MAX_DELAY_SECONDS)  # Exponential backoff
         return self.status
 
+    @snowpark._internal.utils.private_preview(version="1.8.2")
+    @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["timeout"])
+    def result(self, timeout: float = -1) -> T:
+        """
+        Block until completion. Returns job execution result.
+
+        Args:
+            timeout: The maximum time to wait in seconds. Negative values are treated as no timeout.
+
+        Returns:
+            T: The deserialized job result.  # noqa: DAR401
+
+        Raises:
+            RuntimeError: If the job failed or if the job doesn't have a result to retrieve.
+            TimeoutError: If the job does not complete within the specified timeout.  # noqa: DAR402
+        """
+        if self._result is None:
+            self.wait(timeout)
+            try:
+                self._result = interop_utils.fetch_result(self._session, self._result_path)
+            except Exception as e:
+                raise RuntimeError(f"Failed to retrieve result for job (id={self.id})") from e
+
+        if self._result.success:
+            return cast(T, self._result.result)
+        raise RuntimeError(f"Job execution failed (id={self.id})") from self._result.exception
+
+
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "instance_id"])
+def _get_status(session: snowpark.Session, job_id: str, instance_id: Optional[int] = None) -> types.JOB_STATUS:
+    """Retrieve job or job instance execution status."""
+    if instance_id is not None:
+        # Get specific instance status
+        rows = session.sql("SHOW SERVICE INSTANCES IN SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
+        for row in rows:
+            if row["instance_id"] == str(instance_id):
+                return cast(types.JOB_STATUS, row["status"])
+        raise ValueError(f"Instance {instance_id} not found in job {job_id}")
+    else:
+        (row,) = session.sql("DESCRIBE SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
+        return cast(types.JOB_STATUS, row["status"])
+
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
-def _get_status(session: snowpark.Session, job_id: str) -> types.JOB_STATUS:
-    """Retrieve job execution status."""
-    # TODO: snowflake-snowpark-python<1.24.0 shows spurious error messages on
-    #       `DESCRIBE` queries with bind variables
-    #       Switch to use bind variables instead of client side formatting after
-    #       updating to snowflake-snowpark-python>=1.24.0
-    (row,) = session.sql(f"DESCRIBE SERVICE {job_id}").collect()
-    return cast(types.JOB_STATUS, row["status"])
-
-
-@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "limit"])
-def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1) -> str:
+def _get_service_spec(session: snowpark.Session, job_id: str) -> dict[str, Any]:
+    """Retrieve job execution service spec."""
+    (row,) = session.sql("DESCRIBE SERVICE IDENTIFIER(?)", params=[job_id]).collect()
+    return cast(dict[str, Any], yaml.safe_load(row["spec"]))
+
+
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "limit", "instance_id"])
+def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1, instance_id: Optional[int] = None) -> str:
     """
     Retrieve the job's execution logs.
 
@@ -99,15 +192,54 @@ def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1) -> str:
         job_id: The job ID.
         limit: The maximum number of lines to return. Negative values are treated as no limit.
         session: The Snowpark session to use. If none specified, uses active session.
+        instance_id: Optional instance ID to get logs from a specific instance.
 
     Returns:
         The job's execution logs.
     """
-    params: List[Any] = [job_id]
+    # If instance_id is not specified, try to get the head instance ID
+    if instance_id is None:
+        instance_id = _get_head_instance_id(session, job_id)
+
+    # Assemble params: [job_id, instance_id, container_name, (optional) limit]
+    params: list[Any] = [
+        job_id,
+        0 if instance_id is None else instance_id,
+        constants.DEFAULT_CONTAINER_NAME,
+    ]
     if limit > 0:
         params.append(limit)
+
     (row,) = session.sql(
-        f"SELECT SYSTEM$GET_SERVICE_LOGS(?, 0, '{constants.DEFAULT_CONTAINER_NAME}'{f', ?' if limit > 0 else ''})",
+        f"SELECT SYSTEM$GET_SERVICE_LOGS(?, ?, ?{f', ?' if limit > 0 else ''})",
         params=params,
     ).collect()
     return str(row[0])
+
+
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
+def _get_head_instance_id(session: snowpark.Session, job_id: str) -> Optional[int]:
+    """
+    Retrieve the head instance ID of a job.
+
+    Args:
+        session: The Snowpark session to use.
+        job_id: The job ID.
+
+    Returns:
+        The head instance ID of the job. Returns None if the head instance has not started yet.
+    """
+    rows = session.sql("SHOW SERVICE INSTANCES IN SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
+    if not rows:
+        return None
+
+    # Sort by start_time first, then by instance_id
+    sorted_instances = sorted(rows, key=lambda x: (x["start_time"], int(x["instance_id"])))
+    head_instance = sorted_instances[0]
+    if not head_instance["start_time"]:
+        # If head instance hasn't started yet, return None
+        return None
+    try:
+        return int(head_instance["instance_id"])
+    except (ValueError, TypeError):
+        return 0

snowflake/ml/jobs/manager.py

@@ -1,6 +1,7 @@
+import logging
 import pathlib
 import textwrap
-from typing import Any, Callable, Dict, List, Literal, Optional, Union
+from typing import Any, Callable, Literal, Optional, TypeVar, Union, overload
 from uuid import uuid4
 
 import yaml
@@ -13,11 +14,14 @@ from snowflake.ml.jobs._utils import payload_utils, spec_utils
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
 
+logger = logging.getLogger(__name__)
+
 _PROJECT = "MLJob"
 JOB_ID_PREFIX = "MLJOB_"
 
+T = TypeVar("T")
+
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["limit", "scope"])
 def list_jobs(
     limit: int = 10,
@@ -57,9 +61,8 @@ def list_jobs(
     return df
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
-def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob:
+def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob[Any]:
     """Retrieve a job service from the backend."""
     session = session or get_active_session()
 
@@ -71,7 +74,8 @@ def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob
 
     try:
         # Validate that job exists by doing a status check
-        job = jb.MLJob(job_id, session=session)
+        # FIXME: Retrieve return path
+        job = jb.MLJob[Any](job_id, session=session)
         _ = job.status
         return job
     except SnowparkSQLException as e:
@@ -80,9 +84,8 @@ def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob
         raise
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
-def delete_job(job: Union[str, jb.MLJob], session: Optional[snowpark.Session] = None) -> None:
+def delete_job(job: Union[str, jb.MLJob[Any]], session: Optional[snowpark.Session] = None) -> None:
     """Delete a job service from the backend. Status and logs will be lost."""
     if isinstance(job, jb.MLJob):
         job_id = job.id
@@ -93,23 +96,22 @@ def delete_job(job: Union[str, jb.MLJob], session: Optional[snowpark.Session] =
     session.sql("DROP SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def submit_file(
     file_path: str,
     compute_pool: str,
     *,
     stage_name: str,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
-) -> jb.MLJob:
+) -> jb.MLJob[None]:
     """
     Submit a Python file as a job to the compute pool.
 
@@ -146,7 +148,6 @@ def submit_file(
     )
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def submit_directory(
     dir_path: str,
@@ -154,16 +155,16 @@ def submit_directory(
     *,
     entrypoint: str,
     stage_name: str,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
-) -> jb.MLJob:
+) -> jb.MLJob[None]:
     """
     Submit a directory containing Python script(s) as a job to the compute pool.
 
@@ -202,6 +203,46 @@ def submit_directory(
     )
 
 
+@overload
+def _submit_job(
+    source: str,
+    compute_pool: str,
+    *,
+    stage_name: str,
+    entrypoint: Optional[str] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
+    query_warehouse: Optional[str] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
+    num_instances: Optional[int] = None,
+    enable_metrics: bool = False,
+    session: Optional[snowpark.Session] = None,
+) -> jb.MLJob[None]:
+    ...
+
+
+@overload
+def _submit_job(
+    source: Callable[..., T],
+    compute_pool: str,
+    *,
+    stage_name: str,
+    entrypoint: Optional[str] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
+    query_warehouse: Optional[str] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
+    num_instances: Optional[int] = None,
+    enable_metrics: bool = False,
+    session: Optional[snowpark.Session] = None,
+) -> jb.MLJob[T]:
+    ...
+
+
 @telemetry.send_api_usage_telemetry(
     project=_PROJECT,
     func_params_to_log=[
@@ -210,24 +251,26 @@ def submit_directory(
         # TODO: Log lengths of args, env_vars, and spec_overrides values
         "pip_requirements",
         "external_access_integrations",
+        "num_instances",
+        "enable_metrics",
     ],
 )
 def _submit_job(
-    source: Union[str, Callable[..., Any]],
+    source: Union[str, Callable[..., T]],
     compute_pool: str,
     *,
     stage_name: str,
     entrypoint: Optional[str] = None,
-    args: Optional[List[str]] = None,
-    env_vars: Optional[Dict[str, str]] = None,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    spec_overrides: Optional[Dict[str, Any]] = None,
+    spec_overrides: Optional[dict[str, Any]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
     session: Optional[snowpark.Session] = None,
-) -> jb.MLJob:
+) -> jb.MLJob[T]:
     """
     Submit a job to the compute pool.
 
@@ -252,6 +295,12 @@ def _submit_job(
     Raises:
         RuntimeError: If required Snowflake features are not enabled.
     """
+    # Display warning about PrPr parameters
+    if num_instances is not None:
+        logger.warning(
+            "_submit_job() parameter 'num_instances' is in private preview since 1.8.2. Do not use it in production.",
+        )
+
     session = session or get_active_session()
     job_id = f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"
     stage_name = "@" + stage_name.lstrip("@").rstrip("/")
@@ -314,5 +363,4 @@ def _submit_job(
             ) from e
         raise
 
-    # TODO: Wrap snowflake.core.service.JobService object
-    return jb.MLJob(job_id, session=session)
+    return jb.MLJob(job_id, service_spec=spec, session=session)

snowflake/ml/lineage/lineage_node.py

@@ -1,6 +1,6 @@
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Set, Type, Union
+from typing import TYPE_CHECKING, Literal, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
@@ -12,7 +12,7 @@ if TYPE_CHECKING:
     from snowflake.ml.model._client.model import model_version_impl
 
 _PROJECT = "LINEAGE"
-DOMAIN_LINEAGE_REGISTRY: Dict[str, Type["LineageNode"]] = {}
+DOMAIN_LINEAGE_REGISTRY: dict[str, type["LineageNode"]] = {}
 
 
 class LineageNode:
@@ -87,8 +87,8 @@ class LineageNode:
     def lineage(
         self,
         direction: Literal["upstream", "downstream"] = "downstream",
-        domain_filter: Optional[Set[Literal["feature_view", "dataset", "model", "table", "view"]]] = None,
-    ) -> List[Union["feature_view.FeatureView", "dataset.Dataset", "model_version_impl.ModelVersion", "LineageNode"]]:
+        domain_filter: Optional[set[Literal["feature_view", "dataset", "model", "table", "view"]]] = None,
+    ) -> list[Union["feature_view.FeatureView", "dataset.Dataset", "model_version_impl.ModelVersion", "LineageNode"]]:
         """
         Retrieves the lineage nodes connected to this node.
 
@@ -109,7 +109,7 @@ class LineageNode:
         if domain_filter is not None:
             domain_filter = {d.lower() for d in domain_filter}  # type: ignore[misc]
 
-        lineage_nodes: List["LineageNode"] = []
+        lineage_nodes: list["LineageNode"] = []
         for row in df.collect():
             lineage_object = (
                 json.loads(row["TARGET_OBJECT"])

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

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union
+from typing import Optional, Union
 
 import pandas as pd
 
@@ -224,7 +224,7 @@ class Model:
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def versions(self) -> List[model_version_impl.ModelVersion]:
+    def versions(self) -> list[model_version_impl.ModelVersion]:
         """Get all versions in the model.
 
         Returns:
@@ -298,7 +298,7 @@ class Model:
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def show_tags(self) -> Dict[str, str]:
+    def show_tags(self) -> dict[str, str]:
         """Get a dictionary showing the tag and its value attached to the model.
 
         Returns: