snowflake-ml-python 1.8.2__py3-none-any.whl → 1.8.4__py3-none-any.whl

snowflake/ml/jobs/_utils/spec_utils.py

@@ -1,7 +1,7 @@
 import logging
 from math import ceil
 from pathlib import PurePath
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import snowflake_env
@@ -11,14 +11,11 @@ from snowflake.ml.jobs._utils import constants, types
 def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.ComputeResources:
     """Extract resource information for the specified compute pool"""
     # Get the instance family
-    rows = session.sql(f"show compute pools like '{compute_pool}'").collect()
+    rows = session.sql("show compute pools like ?", params=[compute_pool]).collect()
     if not rows:
         raise ValueError(f"Compute pool '{compute_pool}' not found")
     instance_family: str = rows[0]["instance_family"]
-
-    # Get the cloud we're using (AWS, Azure, etc)
-    region = snowflake_env.get_regions(session)[snowflake_env.get_current_region_id(session)]
-    cloud = region["cloud"]
+    cloud = snowflake_env.get_current_cloud(session, default=snowflake_env.SnowflakeCloudType.AWS)
 
     return (
         constants.COMMON_INSTANCE_FAMILIES.get(instance_family)
@@ -26,22 +23,14 @@ def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.C
     )
 
 
-def _get_image_spec(session: snowpark.Session, compute_pool: str, image_tag: Optional[str] = None) -> types.ImageSpec:
+def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.ImageSpec:
     # Retrieve compute pool node resources
     resources = _get_node_resources(session, compute_pool=compute_pool)
 
     # Use MLRuntime image
     image_repo = constants.DEFAULT_IMAGE_REPO
     image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
-
-    # Try to pull latest image tag from server side if possible
-    if not image_tag:
-        query_result = session.sql("SHOW PARAMETERS LIKE 'constants.RUNTIME_BASE_IMAGE_TAG' IN ACCOUNT").collect()
-        if query_result:
-            image_tag = query_result[0]["value"]
-
-    if image_tag is None:
-        image_tag = constants.DEFAULT_IMAGE_TAG
+    image_tag = constants.DEFAULT_IMAGE_TAG
 
     # TODO: Should each instance consume the entire pod?
     return types.ImageSpec(
@@ -54,9 +43,9 @@ def _get_image_spec(session: snowpark.Session, compute_pool: str, image_tag: Opt
 
 
 def generate_spec_overrides(
-    environment_vars: Optional[Dict[str, str]] = None,
-    custom_overrides: Optional[Dict[str, Any]] = None,
-) -> Dict[str, Any]:
+    environment_vars: Optional[dict[str, str]] = None,
+    custom_overrides: Optional[dict[str, Any]] = None,
+) -> dict[str, Any]:
     """
     Generate a dictionary of service specification overrides.
 
@@ -68,7 +57,7 @@ def generate_spec_overrides(
         Resulting service specifiation patch dict. Empty if no overrides were supplied.
     """
     # Generate container level overrides
-    container_spec: Dict[str, Any] = {
+    container_spec: dict[str, Any] = {
         "name": constants.DEFAULT_CONTAINER_NAME,
     }
     if environment_vars:
@@ -95,10 +84,10 @@ def generate_service_spec(
     session: snowpark.Session,
     compute_pool: str,
     payload: types.UploadedPayload,
-    args: Optional[List[str]] = None,
+    args: Optional[list[str]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
-) -> Dict[str, Any]:
+) -> dict[str, Any]:
     """
     Generate a service specification for a job.
 
@@ -117,11 +106,11 @@ def generate_service_spec(
     image_spec = _get_image_spec(session, compute_pool)
 
     # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
-    resource_requests: Dict[str, Union[str, int]] = {
+    resource_requests: dict[str, Union[str, int]] = {
         "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
         "memory": f"{image_spec.resource_limits.memory}Gi",
     }
-    resource_limits: Dict[str, Union[str, int]] = {
+    resource_limits: dict[str, Union[str, int]] = {
         "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
         "memory": f"{image_spec.resource_limits.memory}Gi",
     }
@@ -130,8 +119,8 @@ def generate_service_spec(
         resource_limits["nvidia.com/gpu"] = image_spec.resource_limits.gpu
 
     # Add local volumes for ephemeral logs and artifacts
-    volumes: List[Dict[str, str]] = []
-    volume_mounts: List[Dict[str, str]] = []
+    volumes: list[dict[str, str]] = []
+    volume_mounts: list[dict[str, str]] = []
     for volume_name, mount_path in [
         ("system-logs", "/var/log/managedservices/system/mlrs"),
         ("user-logs", "/var/log/managedservices/user/mlrs"),
@@ -302,11 +291,11 @@ def merge_patch(base: Any, patch: Any, display_name: str = "") -> Any:
 
 
 def _merge_lists_of_dicts(
-    base: List[Dict[str, Any]],
-    patch: List[Dict[str, Any]],
+    base: list[dict[str, Any]],
+    patch: list[dict[str, Any]],
     merge_key: str = "name",
     display_name: str = "",
-) -> List[Dict[str, Any]]:
+) -> list[dict[str, Any]]:
     """
     Attempts to merge lists of dicts by matching on a merge key (default "name").
     - If the merge key is missing, the behavior falls back to overwriting the list.

snowflake/ml/jobs/_utils/types.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from pathlib import PurePath
-from typing import List, Literal, Optional, Union
+from typing import Literal, Optional, Union
 
 JOB_STATUS = Literal[
     "PENDING",
@@ -21,7 +21,7 @@ class PayloadEntrypoint:
 class UploadedPayload:
     # TODO: Include manifest of payload files for validation
     stage_path: PurePath
-    entrypoint: List[Union[str, PurePath]]
+    entrypoint: list[Union[str, PurePath]]
 
 
 @dataclass(frozen=True)

snowflake/ml/jobs/decorators.py

@@ -1,6 +1,6 @@
 import copy
 import functools
-from typing import Callable, Dict, List, Optional, TypeVar
+from typing import Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
@@ -15,18 +15,19 @@ _Args = ParamSpec("_Args")
 _ReturnValue = TypeVar("_ReturnValue")
 
 
-@snowpark._internal.utils.private_preview(version="1.7.4")
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
 def remote(
     compute_pool: str,
     *,
     stage_name: str,
-    pip_requirements: Optional[List[str]] = None,
-    external_access_integrations: Optional[List[str]] = None,
+    pip_requirements: Optional[list[str]] = None,
+    external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
-    env_vars: Optional[Dict[str, str]] = None,
+    env_vars: Optional[dict[str, str]] = None,
     num_instances: Optional[int] = None,
     enable_metrics: bool = False,
+    database: Optional[str] = None,
+    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
 ) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob[_ReturnValue]]]:
     """
@@ -41,6 +42,8 @@ def remote(
         env_vars: Environment variables to set in container
         num_instances: The number of nodes in the job. If none specified, create a single node job.
         enable_metrics: Whether to enable metrics publishing for the job.
+        database: The database to use for the job.
+        schema: The schema to use for the job.
         session: The Snowpark session to use. If none specified, uses active session.
 
     Returns:
@@ -68,6 +71,8 @@ def remote(
                 env_vars=env_vars,
                 num_instances=num_instances,
                 enable_metrics=enable_metrics,
+                database=database,
+                schema=schema,
                 session=session,
             )
             assert isinstance(job, jb.MLJob), f"Unexpected job type: {type(job)}"

snowflake/ml/jobs/job.py

@@ -1,12 +1,15 @@
 import time
-from typing import Any, Dict, Generic, List, Optional, TypeVar, cast
+from functools import cached_property
+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._internal.utils import identifier
 from snowflake.ml.jobs._utils import constants, interop_utils, types
-from snowflake.snowpark import context as sp_context
+from snowflake.snowpark import Row, context as sp_context
+from snowflake.snowpark.exceptions import SnowparkSQLException
 
 _PROJECT = "MLJob"
 TERMINAL_JOB_STATUSES = {"FAILED", "DONE", "INTERNAL_ERROR"}
@@ -18,16 +21,24 @@ class MLJob(Generic[T]):
     def __init__(
         self,
         id: str,
-        service_spec: Optional[Dict[str, Any]] = None,
+        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._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
 
+    @cached_property
+    def name(self) -> str:
+        return identifier.parse_schema_level_object_identifier(self.id)[-1]
+
+    @cached_property
+    def num_instances(self) -> int:
+        return _get_num_instances(self._session, self.id)
+
     @property
     def id(self) -> str:
         """Get the unique job ID"""
@@ -42,18 +53,18 @@ class MLJob(Generic[T]):
         return self._status
 
     @property
-    def _service_spec(self) -> Dict[str, Any]:
+    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]:
+    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)
+        return cast(dict[str, Any], container_spec)
 
     @property
     def _stage_path(self) -> str:
@@ -67,11 +78,20 @@ class MLJob(Generic[T]):
         """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")
+            raise RuntimeError(f"Job {self.name} doesn't have a result path configured")
         return f"{self._stage_path}/{result_path}"
 
-    @snowpark._internal.utils.private_preview(version="1.7.4")
-    def get_logs(self, limit: int = -1, instance_id: Optional[int] = None) -> str:
+    @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.
 
@@ -79,15 +99,17 @@ class MLJob(Generic[T]):
             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, 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, instance_id: Optional[int] = None) -> None:
         """
         Display the job's execution logs.
@@ -97,9 +119,8 @@ class MLJob(Generic[T]):
             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, instance_id))  # 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, func_params_to_log=["timeout"])
     def wait(self, timeout: float = -1) -> types.JOB_STATUS:
         """
@@ -118,7 +139,7 @@ class MLJob(Generic[T]):
         start_time = time.monotonic()
         while self.status not in TERMINAL_JOB_STATUSES:
             if timeout >= 0 and (elapsed := time.monotonic() - start_time) >= timeout:
-                raise TimeoutError(f"Job {self.id} did not complete within {elapsed} seconds")
+                raise TimeoutError(f"Job {self.name} did not complete within {elapsed} seconds")
             time.sleep(delay)
             delay = min(delay * 2, constants.JOB_POLL_MAX_DELAY_SECONDS)  # Exponential backoff
         return self.status
@@ -144,11 +165,11 @@ class MLJob(Generic[T]):
             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
+                raise RuntimeError(f"Failed to retrieve result for job (id={self.name})") 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
+        raise RuntimeError(f"Job execution failed (id={self.name})") from self._result.exception
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id", "instance_id"])
@@ -162,15 +183,15 @@ def _get_status(session: snowpark.Session, job_id: str, instance_id: Optional[in
                 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()
+        row = _get_service_info(session, job_id)
         return cast(types.JOB_STATUS, row["status"])
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
-def _get_service_spec(session: snowpark.Session, job_id: str) -> Dict[str, Any]:
+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"]))
+    row = _get_service_info(session, job_id)
+    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"])
@@ -186,13 +207,24 @@ def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1, instance_
 
     Returns:
         The job's execution logs.
+
+    Raises:
+        SnowparkSQLException: if the container is pending
+        RuntimeError: if failed to get head instance_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)
+        try:
+            instance_id = _get_head_instance_id(session, job_id)
+        except RuntimeError:
+            raise RuntimeError(
+                "Failed to retrieve job logs. "
+                "Logs may be inaccessible due to job expiration and can be retrieved from Event Table instead."
+            )
 
     # Assemble params: [job_id, instance_id, container_name, (optional) limit]
-    params: List[Any] = [
+    params: list[Any] = [
         job_id,
         0 if instance_id is None else instance_id,
         constants.DEFAULT_CONTAINER_NAME,
@@ -200,10 +232,15 @@ def _get_logs(session: snowpark.Session, job_id: str, limit: int = -1, instance_
     if limit > 0:
         params.append(limit)
 
-    (row,) = session.sql(
-        f"SELECT SYSTEM$GET_SERVICE_LOGS(?, ?, ?{f', ?' if limit > 0 else ''})",
-        params=params,
-    ).collect()
+    try:
+        (row,) = session.sql(
+            f"SELECT SYSTEM$GET_SERVICE_LOGS(?, ?, ?{f', ?' if limit > 0 else ''})",
+            params=params,
+        ).collect()
+    except SnowparkSQLException as e:
+        if "Container Status: PENDING" in e.message:
+            return "Warning: Waiting for container to start. Logs will be shown when available."
+        raise
     return str(row[0])
 
 
@@ -213,18 +250,27 @@ def _get_head_instance_id(session: snowpark.Session, job_id: str) -> Optional[in
     Retrieve the head instance ID of a job.
 
     Args:
-        session: The Snowpark session to use.
-        job_id: The job ID.
+        session (Session): The Snowpark session to use.
+        job_id (str): The job ID.
 
     Returns:
-        The head instance ID of the job. Returns None if the head instance has not started yet.
+        Optional[int]: The head instance ID of the job, or None if the head instance has not started yet.
+
+    Raises:
+        RuntimeError: If the instances died or if some instances disappeared.
     """
     rows = session.sql("SHOW SERVICE INSTANCES IN SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
     if not rows:
         return None
+    if _get_num_instances(session, job_id) > len(rows):
+        raise RuntimeError("Couldn’t retrieve head instance due to missing instances.")
 
     # Sort by start_time first, then by instance_id
-    sorted_instances = sorted(rows, key=lambda x: (x["start_time"], int(x["instance_id"])))
+    try:
+        sorted_instances = sorted(rows, key=lambda x: (x["start_time"], int(x["instance_id"])))
+    except TypeError:
+        raise RuntimeError("Job instance information unavailable.")
+
     head_instance = sorted_instances[0]
     if not head_instance["start_time"]:
         # If head instance hasn't started yet, return None
@@ -233,3 +279,14 @@ def _get_head_instance_id(session: snowpark.Session, job_id: str) -> Optional[in
         return int(head_instance["instance_id"])
     except (ValueError, TypeError):
         return 0
+
+
+def _get_service_info(session: snowpark.Session, job_id: str) -> Row:
+    (row,) = session.sql("DESCRIBE SERVICE IDENTIFIER(?)", params=(job_id,)).collect()
+    return row
+
+
+@telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["job_id"])
+def _get_num_instances(session: snowpark.Session, job_id: str) -> int:
+    row = _get_service_info(session, job_id)
+    return int(row["target_instances"]) if row["target_instances"] else 0