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

snowflake/cortex/__init__.py

@@ -1,5 +1,10 @@
 from snowflake.cortex._classify_text import ClassifyText, classify_text
-from snowflake.cortex._complete import Complete, CompleteOptions, complete
+from snowflake.cortex._complete import (
+    Complete,
+    CompleteOptions,
+    ConversationMessage,
+    complete,
+)
 from snowflake.cortex._embed_text_768 import EmbedText768, embed_text_768
 from snowflake.cortex._embed_text_1024 import EmbedText1024, embed_text_1024
 from snowflake.cortex._extract_answer import ExtractAnswer, extract_answer
@@ -14,6 +19,7 @@ __all__ = [
     "Complete",
     "complete",
     "CompleteOptions",
+    "ConversationMessage",
     "EmbedText768",
     "embed_text_768",
     "EmbedText1024",

snowflake/ml/_internal/platform_capabilities.py

@@ -11,6 +11,9 @@ from snowflake.snowpark import (
     session as snowpark_session,
 )
 
+LIVE_COMMIT_PARAMETER = "ENABLE_LIVE_VERSION_IN_SDK"
+INLINE_DEPLOYMENT_SPEC_PARAMETER = "ENABLE_INLINE_DEPLOYMENT_SPEC"
+
 
 class PlatformCapabilities:
     """Class that retrieves platform feature values for the currently running server.
@@ -18,12 +21,12 @@ class PlatformCapabilities:
     Example usage:
     ```
     pc = PlatformCapabilities.get_instance(session)
-    if pc.is_nested_function_enabled():
-        # Nested functions are enabled.
-        print("Nested functions are enabled.")
+    if pc.is_inlined_deployment_spec_enabled():
+        # Inline deployment spec is enabled.
+        print("Inline deployment spec is enabled.")
     else:
-        # Nested functions are disabled.
-        print("Nested functions are disabled or not supported.")
+        # Inline deployment spec is disabled.
+        print("Inline deployment spec is disabled or not supported.")
     ```
     """
 
@@ -50,9 +53,11 @@ class PlatformCapabilities:
 
     # For contextmanager, we need to have return type Iterator[Never]. However, Never type is introduced only in
     # Python 3.11. So, we are ignoring the type for this method.
+    _dummy_features: dict[str, Any] = {"dummy": "dummy"}
+
     @classmethod  # type: ignore[arg-type]
     @contextmanager
-    def mock_features(cls, features: dict[str, Any]) -> None:  # type: ignore[misc]
+    def mock_features(cls, features: dict[str, Any] = _dummy_features) -> None:  # type: ignore[misc]
         logging.debug(f"Setting mock features: {features}")
         cls.set_mock_features(features)
         try:
@@ -61,14 +66,11 @@ class PlatformCapabilities:
             logging.debug(f"Clearing mock features: {features}")
             cls.clear_mock_features()
 
-    def is_nested_function_enabled(self) -> bool:
-        return self._get_bool_feature("SPCS_MODEL_ENABLE_EMBEDDED_SERVICE_FUNCTIONS", False)
-
     def is_inlined_deployment_spec_enabled(self) -> bool:
-        return self._get_bool_feature("ENABLE_INLINE_DEPLOYMENT_SPEC", False)
+        return self._get_bool_feature(INLINE_DEPLOYMENT_SPEC_PARAMETER, False)
 
     def is_live_commit_enabled(self) -> bool:
-        return self._get_bool_feature("ENABLE_BUNDLE_MODULE_CHECKOUT", False)
+        return self._get_bool_feature(LIVE_COMMIT_PARAMETER, False)
 
     @staticmethod
     def _get_features(session: snowpark_session.Session) -> dict[str, Any]:

snowflake/ml/_internal/utils/identifier.py

@@ -12,7 +12,7 @@ SF_IDENTIFIER_RE = re.compile(_SF_IDENTIFIER)
 _SF_SCHEMA_LEVEL_OBJECT = (
     rf"(?:(?:(?P<db>{_SF_IDENTIFIER})\.)?(?P<schema>{_SF_IDENTIFIER})\.)?(?P<object>{_SF_IDENTIFIER})"
 )
-_SF_STAGE_PATH = rf"{_SF_SCHEMA_LEVEL_OBJECT}(?P<path>.*)"
+_SF_STAGE_PATH = rf"@?{_SF_SCHEMA_LEVEL_OBJECT}(?P<path>/.*)?"
 _SF_SCHEMA_LEVEL_OBJECT_RE = re.compile(_SF_SCHEMA_LEVEL_OBJECT)
 _SF_STAGE_PATH_RE = re.compile(_SF_STAGE_PATH)
 
@@ -197,7 +197,7 @@ def parse_snowflake_stage_path(
         res.group("db"),
         res.group("schema"),
         res.group("object"),
-        res.group("path"),
+        res.group("path") or "",
     )
 
 

snowflake/ml/jobs/_utils/constants.py

@@ -13,7 +13,7 @@ STAGE_VOLUME_MOUNT_PATH = "/mnt/app"
 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 = "1.0.1"
+DEFAULT_IMAGE_TAG = "1.2.3"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
 
 # Percent of container memory to allocate for /dev/shm volume

snowflake/ml/jobs/_utils/payload_utils.py

@@ -9,6 +9,7 @@ from pathlib import Path, PurePath
 from typing import Any, Callable, Optional, Union, cast, get_args, get_origin
 
 import cloudpickle as cp
+from packaging import version
 
 from snowflake import snowpark
 from snowflake.ml.jobs._utils import constants, types
@@ -97,11 +98,18 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
         head_info=$(python3 get_instance_ip.py "$SNOWFLAKE_SERVICE_NAME" --head)
         if [ $? -eq 0 ]; then
             # Parse the output using read
-            read head_index head_ip <<< "$head_info"
+            read head_index head_ip head_status<<< "$head_info"
 
             # Use the parsed variables
             echo "Head Instance Index: $head_index"
             echo "Head Instance IP: $head_ip"
+            echo "Head Instance Status: $head_status"
+
+            # If the head status is not "READY" or "PENDING", exit early
+            if [ "$head_status" != "READY" ] && [ "$head_status" != "PENDING" ]; then
+                echo "Head instance status is not READY or PENDING. Exiting."
+                exit 0
+            fi
 
         else
             echo "Error: Failed to get head instance information."
@@ -278,17 +286,19 @@ class JobPayload:
         stage_path = PurePath(stage_path) if isinstance(stage_path, str) else stage_path
         source = resolve_source(self.source)
         entrypoint = resolve_entrypoint(source, self.entrypoint)
+        pip_requirements = self.pip_requirements or []
 
         # Create stage if necessary
         stage_name = stage_path.parts[0].lstrip("@")
         # Explicitly check if stage exists first since we may not have CREATE STAGE privilege
         try:
-            session.sql(f"describe stage {stage_name}").collect()
+            session.sql("describe stage identifier(?)", params=[stage_name]).collect()
         except sp_exceptions.SnowparkSQLException:
             session.sql(
-                f"create stage if not exists {stage_name}"
+                "create stage if not exists identifier(?)"
                 " encryption = ( type = 'SNOWFLAKE_SSE' )"
-                " comment = 'Created by snowflake.ml.jobs Python API'"
+                " comment = 'Created by snowflake.ml.jobs Python API'",
+                params=[stage_name],
             ).collect()
 
         # Upload payload to stage
@@ -301,6 +311,8 @@ class JobPayload:
                 overwrite=True,
             )
             source = Path(entrypoint.file_path.parent)
+            if not any(r.startswith("cloudpickle") for r in pip_requirements):
+                pip_requirements.append(f"cloudpickle~={version.parse(cp.__version__).major}.0")
         elif source.is_dir():
             # Manually traverse the directory and upload each file, since Snowflake PUT
             # can't handle directories. Reduce the number of PUT operations by using
@@ -325,10 +337,10 @@ class JobPayload:
 
         # Upload requirements
         # TODO: Check if payload includes both a requirements.txt file and pip_requirements
-        if self.pip_requirements:
+        if pip_requirements:
             # Upload requirements.txt to stage
             session.file.put_stream(
-                io.BytesIO("\n".join(self.pip_requirements).encode()),
+                io.BytesIO("\n".join(pip_requirements).encode()),
                 stage_location=stage_path.joinpath("requirements.txt").as_posix(),
                 auto_compress=False,
                 overwrite=True,
@@ -495,13 +507,6 @@ def generate_python_code(func: Callable[..., Any], source_code_display: bool = F
     # https://github.com/snowflakedb/snowpark-python/blob/main/src/snowflake/snowpark/_internal/udf_utils.py
     source_code_comment = _generate_source_code_comment(func) if source_code_display else ""
 
-    func_code = f"""
-{source_code_comment}
-
-import pickle
-{_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(func).hex()}'))
-"""
-
     arg_dict_name = "kwargs"
     if getattr(func, constants.IS_MLJOB_REMOTE_ATTR, None):
         param_code = f"{arg_dict_name} = {{}}"
@@ -509,25 +514,29 @@ import pickle
         param_code = _generate_param_handler_code(signature, arg_dict_name)
 
     return f"""
-### Version guard to check compatibility across Python versions ###
-import os
 import sys
-import warnings
-
-if sys.version_info.major != {sys.version_info.major} or sys.version_info.minor != {sys.version_info.minor}:
-    warnings.warn(
-        "Python version mismatch: job was created using"
-        " python{sys.version_info.major}.{sys.version_info.minor}"
-        f" but runtime environment uses python{{sys.version_info.major}}.{{sys.version_info.minor}}."
-        " Compatibility across Python versions is not guaranteed and may result in unexpected behavior."
-        " This will be fixed in a future release; for now, please use Python version"
-        f" {{sys.version_info.major}}.{{sys.version_info.minor}}.",
-        RuntimeWarning,
-        stacklevel=0,
-    )
-### End version guard ###
+import pickle
 
-{func_code.strip()}
+try:
+    {textwrap.indent(source_code_comment, '    ')}
+    {_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(func).hex()}'))
+except (TypeError, pickle.PickleError):
+    if sys.version_info.major != {sys.version_info.major} or sys.version_info.minor != {sys.version_info.minor}:
+        raise RuntimeError(
+            "Failed to deserialize function due to Python version mismatch."
+            f" Runtime environment is Python {{sys.version_info.major}}.{{sys.version_info.minor}}"
+            " but function was serialized using Python {sys.version_info.major}.{sys.version_info.minor}."
+        ) from None
+    raise
+except AttributeError as e:
+    if 'cloudpickle' in str(e):
+        import cloudpickle as cp
+        raise RuntimeError(
+            "Failed to deserialize function due to cloudpickle version mismatch."
+            f" Runtime environment uses cloudpickle=={{cp.__version__}}"
+            " but job was serialized using cloudpickle=={cp.__version__}."
+        ) from e
+    raise
 
 if __name__ == '__main__':
 {textwrap.indent(param_code, '    ')}

snowflake/ml/jobs/_utils/scripts/get_instance_ip.py

@@ -29,7 +29,7 @@ def get_self_ip() -> Optional[str]:
         return None
 
 
-def get_first_instance(service_name: str) -> Optional[tuple[str, str]]:
+def get_first_instance(service_name: str) -> Optional[tuple[str, str, str]]:
     """Get the first instance of a batch job based on start time and instance ID.
 
     Args:
@@ -42,7 +42,7 @@ def get_first_instance(service_name: str) -> Optional[tuple[str, str]]:
 
     session = session_utils.get_session()
     df = session.sql(f"show service instances in service {service_name}")
-    result = df.select('"instance_id"', '"ip_address"', '"start_time"').collect()
+    result = df.select('"instance_id"', '"ip_address"', '"start_time"', '"status"').collect()
 
     if not result:
         return None
@@ -57,7 +57,7 @@ def get_first_instance(service_name: str) -> Optional[tuple[str, str]]:
     ip_address = head_instance["ip_address"]
     try:
         socket.inet_aton(ip_address)  # Validate IPv4 address
-        return (head_instance["instance_id"], ip_address)
+        return (head_instance["instance_id"], ip_address, head_instance["status"])
     except OSError:
         logger.error(f"Error: Invalid IP address format: {ip_address}")
         return None
@@ -110,7 +110,7 @@ def main():
             head_info = get_first_instance(args.service_name)
             if head_info:
                 # Print to stdout to allow capture but don't use logger
-                sys.stdout.write(f"{head_info[0]} {head_info[1]}\n")
+                sys.stdout.write(" ".join(head_info) + "\n")
                 sys.exit(0)
             time.sleep(args.retry_interval)
         # If we get here, we've timed out

snowflake/ml/jobs/_utils/scripts/mljob_launcher.py

@@ -59,7 +59,7 @@ class SimpleJSONEncoder(json.JSONEncoder):
         try:
             return super().default(obj)
         except TypeError:
-            return str(obj)
+            return f"Unserializable object: {repr(obj)}"
 
 
 def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = None) -> Any:

snowflake/ml/jobs/_utils/spec_utils.py

@@ -11,7 +11,7 @@ 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"]

snowflake/ml/jobs/decorators.py

@@ -26,6 +26,8 @@ def remote(
     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]]]:
     """
@@ -40,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:
@@ -67,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 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"}
@@ -28,6 +31,14 @@ class MLJob(Generic[T]):
         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"""
@@ -67,7 +78,7 @@ 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}"
 
     @overload
@@ -128,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
@@ -154,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"])
@@ -172,14 +183,14 @@ 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]:
     """Retrieve job execution service spec."""
-    (row,) = session.sql("DESCRIBE SERVICE IDENTIFIER(?)", params=[job_id]).collect()
+    row = _get_service_info(session, job_id)
     return cast(dict[str, Any], yaml.safe_load(row["spec"]))
 
 
@@ -196,10 +207,21 @@ 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] = [
@@ -210,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])
 
 
@@ -223,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
@@ -243,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