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

snowflake/ml/jobs/manager.py

@@ -1,18 +1,20 @@
 import logging
 import pathlib
 import textwrap
-from typing import Any, Callable, Literal, Optional, TypeVar, Union, cast, overload
+from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
 from uuid import uuid4
 
+import pandas as pd
 import yaml
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml.jobs import job as jb
-from snowflake.ml.jobs._utils import payload_utils, spec_utils
+from snowflake.ml.jobs._utils import payload_utils, query_helper, spec_utils
 from snowflake.snowpark.context import get_active_session
 from snowflake.snowpark.exceptions import SnowparkSQLException
+from snowflake.snowpark.functions import coalesce, col, lit, when
 
 logger = logging.getLogger(__name__)
 
@@ -25,39 +27,127 @@ T = TypeVar("T")
 @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["limit", "scope"])
 def list_jobs(
     limit: int = 10,
-    scope: Union[Literal["account", "database", "schema"], str, None] = None,
+    database: Optional[str] = None,
+    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
-) -> snowpark.DataFrame:
+) -> pd.DataFrame:
     """
-    Returns a Snowpark DataFrame with the list of jobs in the current session.
+    Returns a Pandas DataFrame with the list of jobs in the current session.
 
     Args:
         limit: The maximum number of jobs to return. Non-positive values are treated as no limit.
-        scope: The scope to list jobs from, such as "schema" or "compute pool <pool_name>".
+        database: The database to use. If not specified, uses the current database.
+        schema: The schema to use. If not specified, uses the current schema.
         session: The Snowpark session to use. If none specified, uses active session.
 
     Returns:
         A DataFrame with the list of jobs.
 
+    Raises:
+        SnowparkSQLException: if there is an error retrieving the job history.
+
     Examples:
         >>> from snowflake.ml.jobs import list_jobs
-        >>> list_jobs(limit=5).show()
+        >>> list_jobs(limit=5)
     """
     session = session or get_active_session()
+    try:
+        df = _get_job_history_spcs(
+            session,
+            limit=limit,
+            database=database,
+            schema=schema,
+        )
+        return df.to_pandas()
+    except SnowparkSQLException as spcs_error:
+        if spcs_error.sql_error_code == 2143:
+            logger.debug("Job history is not enabled. Please enable it to use this feature.")
+            df = _get_job_services(session, limit=limit, database=database, schema=schema)
+            return df.to_pandas()
+        raise
+
+
+def _get_job_services(
+    session: snowpark.Session, limit: int = 10, database: Optional[str] = None, schema: Optional[str] = None
+) -> snowpark.DataFrame:
     query = "SHOW JOB SERVICES"
     query += f" LIKE '{JOB_ID_PREFIX}%'"
-    if scope:
-        query += f" IN {scope}"
+    database = database or session.get_current_database()
+    schema = schema or session.get_current_schema()
+    if database is None and schema is None:
+        query += "IN account"
+    elif not schema:
+        query += f" IN DATABASE {database}"
+    else:
+        query += f" IN {database}.{schema}"
     if limit > 0:
         query += f" LIMIT {limit}"
     df = session.sql(query)
     df = df.select(
         df['"name"'],
-        df['"owner"'],
         df['"status"'],
-        df['"created_on"'],
+        lit(None).alias('"message"'),
+        df['"database_name"'],
+        df['"schema_name"'],
+        df['"owner"'],
         df['"compute_pool"'],
-    ).order_by('"created_on"', ascending=False)
+        df['"target_instances"'],
+        df['"created_on"'].alias('"created_time"'),
+        when(col('"status"').isin(jb.TERMINAL_JOB_STATUSES), col('"updated_on"'))
+        .otherwise(lit(None))
+        .alias('"completed_time"'),
+    ).order_by('"created_time"', ascending=False)
+    return df
+
+
+def _get_job_history_spcs(
+    session: snowpark.Session,
+    limit: int = 10,
+    database: Optional[str] = None,
+    schema: Optional[str] = None,
+    include_deleted: bool = False,
+    created_time_start: Optional[str] = None,
+    created_time_end: Optional[str] = None,
+) -> snowpark.DataFrame:
+    query = ["select * from table(snowflake.spcs.get_job_history("]
+    query_params = []
+    if created_time_start:
+        query_params.append(f"created_time_start => TO_TIMESTAMP_LTZ('{created_time_start}')")
+    if created_time_end:
+        query_params.append(f"created_time_end => TO_TIMESTAMP_LTZ('{created_time_end}')")
+    query.append(",".join(query_params))
+    query.append("))")
+    condition = []
+    database = database or session.get_current_database()
+    schema = schema or session.get_current_schema()
+
+    # format database and schema identifiers
+    if database:
+        condition.append(f"DATABASE_NAME = '{identifier.resolve_identifier(database)}'")
+
+    if schema:
+        condition.append(f"SCHEMA_NAME = '{identifier.resolve_identifier(schema)}'")
+
+    if not include_deleted:
+        condition.append("DELETED_TIME IS NULL")
+
+    if len(condition) > 0:
+        query.append("WHERE " + " AND ".join(condition))
+    if limit > 0:
+        query.append(f"LIMIT {limit}")
+    df = session.sql("\n".join(query))
+    df = df.select(
+        df["NAME"].alias('"name"'),
+        df["STATUS"].alias('"status"'),
+        df["MESSAGE"].alias('"message"'),
+        df["DATABASE_NAME"].alias('"database_name"'),
+        df["SCHEMA_NAME"].alias('"schema_name"'),
+        df["OWNER"].alias('"owner"'),
+        df["COMPUTE_POOL_NAME"].alias('"compute_pool"'),
+        coalesce(df["PARAMETERS"]["REPLICAS"], lit(1)).alias('"target_instances"'),
+        df["CREATED_TIME"].alias('"created_time"'),
+        df["COMPLETED_TIME"].alias('"completed_time"'),
+    )
     return df
 
 
@@ -74,9 +164,9 @@ def get_job(job_id: str, session: Optional[snowpark.Session] = None) -> jb.MLJob
 
     job_id = f"{database}.{schema}.{job_name}"
     try:
-        # Validate that job exists by doing a status check
+        # Validate that job exists by doing a spec lookup
         job = jb.MLJob[Any](job_id, session=session)
-        _ = job.status
+        _ = job._service_spec
         return job
     except SnowparkSQLException as e:
         if "does not exist" in e.message:
@@ -95,7 +185,7 @@ def delete_job(job: Union[str, jb.MLJob[Any]], session: Optional[snowpark.Sessio
         logger.info(f"Successfully cleaned up stage files for job {job.id} at {stage_path}")
     except Exception as e:
         logger.warning(f"Failed to clean up stage files for job {job.id}: {e}")
-    session.sql("DROP SERVICE IDENTIFIER(?)", params=(job.id,)).collect()
+    query_helper.run_query(session, "DROP SERVICE IDENTIFIER(?)", params=(job.id,))
 
 
 @telemetry.send_api_usage_telemetry(project=_PROJECT)
@@ -105,17 +195,11 @@ def submit_file(
     *,
     stage_name: str,
     args: Optional[list[str]] = None,
-    env_vars: Optional[dict[str, str]] = None,
+    target_instances: int = 1,
     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,
-    target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> jb.MLJob[None]:
     """
     Submit a Python file as a job to the compute pool.
@@ -125,18 +209,20 @@ def submit_file(
         compute_pool: The compute pool to use for the job.
         stage_name: The name of the stage where the job payload will be uploaded.
         args: A list of arguments to pass to the job.
-        env_vars: Environment variables to set in container
+        target_instances: The number of nodes in the job. If none specified, create a single node job.
         pip_requirements: A list of pip requirements for the job.
         external_access_integrations: A list of external access integrations.
-        query_warehouse: The query warehouse to use. Defaults to session warehouse.
-        spec_overrides: Custom service specification overrides to apply.
-        target_instances: The number of instances to use for the job. If none specified, single node job is created.
-        min_instances: The minimum number of nodes required to start the job. If none specified,
-            defaults to target_instances. If set, the job will not start until the minimum number of nodes is available.
-        enable_metrics: Whether to enable metrics publishing for the job.
-        database: The database to use.
-        schema: The schema to use.
         session: The Snowpark session to use. If none specified, uses active session.
+        kwargs: Additional keyword arguments. Supported arguments:
+            database (str): The database to use for the job.
+            schema (str): The schema to use for the job.
+            min_instances (int): The minimum number of nodes required to start the job.
+                If none specified, defaults to target_instances. If set, the job
+                will not start until the minimum number of nodes is available.
+            env_vars (dict): Environment variables to set in container.
+            enable_metrics (bool): Whether to enable metrics publishing for the job.
+            query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
+            spec_overrides (dict): A dictionary of overrides for the service spec.
 
     Returns:
         An object representing the submitted job.
@@ -146,17 +232,11 @@ def submit_file(
         args=args,
         compute_pool=compute_pool,
         stage_name=stage_name,
-        env_vars=env_vars,
+        target_instances=target_instances,
         pip_requirements=pip_requirements,
         external_access_integrations=external_access_integrations,
-        query_warehouse=query_warehouse,
-        spec_overrides=spec_overrides,
-        target_instances=target_instances,
-        min_instances=min_instances,
-        enable_metrics=enable_metrics,
-        database=database,
-        schema=schema,
         session=session,
+        **kwargs,
     )
 
 
@@ -168,17 +248,11 @@ def submit_directory(
     entrypoint: str,
     stage_name: str,
     args: Optional[list[str]] = None,
-    env_vars: Optional[dict[str, str]] = None,
+    target_instances: int = 1,
     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,
-    target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> jb.MLJob[None]:
     """
     Submit a directory containing Python script(s) as a job to the compute pool.
@@ -189,18 +263,20 @@ def submit_directory(
         entrypoint: The relative path to the entry point script inside the source directory.
         stage_name: The name of the stage where the job payload will be uploaded.
         args: A list of arguments to pass to the job.
-        env_vars: Environment variables to set in container
+        target_instances: The number of nodes in the job. If none specified, create a single node job.
         pip_requirements: A list of pip requirements for the job.
         external_access_integrations: A list of external access integrations.
-        query_warehouse: The query warehouse to use. Defaults to session warehouse.
-        spec_overrides: Custom service specification overrides to apply.
-        target_instances: The number of instances to use for the job. If none specified, single node job is created.
-        min_instances: The minimum number of nodes required to start the job. If none specified,
-            defaults to target_instances. If set, the job will not start until the minimum number of nodes is available.
-        enable_metrics: Whether to enable metrics publishing for the job.
-        database: The database to use.
-        schema: The schema to use.
         session: The Snowpark session to use. If none specified, uses active session.
+        kwargs: Additional keyword arguments. Supported arguments:
+            database (str): The database to use for the job.
+            schema (str): The schema to use for the job.
+            min_instances (int): The minimum number of nodes required to start the job.
+                If none specified, defaults to target_instances. If set, the job
+                will not start until the minimum number of nodes is available.
+            env_vars (dict): Environment variables to set in container.
+            enable_metrics (bool): Whether to enable metrics publishing for the job.
+            query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
+            spec_overrides (dict): A dictionary of overrides for the service spec.
 
     Returns:
         An object representing the submitted job.
@@ -211,17 +287,11 @@ def submit_directory(
         args=args,
         compute_pool=compute_pool,
         stage_name=stage_name,
-        env_vars=env_vars,
+        target_instances=target_instances,
         pip_requirements=pip_requirements,
         external_access_integrations=external_access_integrations,
-        query_warehouse=query_warehouse,
-        spec_overrides=spec_overrides,
-        target_instances=target_instances,
-        min_instances=min_instances,
-        enable_metrics=enable_metrics,
-        database=database,
-        schema=schema,
         session=session,
+        **kwargs,
     )
 
 
@@ -233,17 +303,11 @@ def submit_from_stage(
     entrypoint: str,
     stage_name: str,
     args: Optional[list[str]] = None,
-    env_vars: Optional[dict[str, str]] = None,
+    target_instances: int = 1,
     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,
-    target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> jb.MLJob[None]:
     """
     Submit a directory containing Python script(s) as a job to the compute pool.
@@ -254,19 +318,20 @@ def submit_from_stage(
         entrypoint: a stage path containing the entry point script inside the source directory.
         stage_name: The name of the stage where the job payload will be uploaded.
         args: A list of arguments to pass to the job.
-        env_vars: Environment variables to set in container
+        target_instances: The number of nodes in the job. If none specified, create a single node job.
         pip_requirements: A list of pip requirements for the job.
         external_access_integrations: A list of external access integrations.
-        query_warehouse: The query warehouse to use. Defaults to session warehouse.
-        spec_overrides: Custom service specification overrides to apply.
-        target_instances: The number of instances to use for the job. If none specified, single node job is created.
-        min_instances: The minimum number of nodes required to start the job. If none specified,
-            defaults to target_instances. If set, the job will not start until the minimum number of nodes is available.
-        enable_metrics: Whether to enable metrics publishing for the job.
-        database: The database to use.
-        schema: The schema to use.
         session: The Snowpark session to use. If none specified, uses active session.
-
+        kwargs: Additional keyword arguments. Supported arguments:
+            database (str): The database to use for the job.
+            schema (str): The schema to use for the job.
+            min_instances (int): The minimum number of nodes required to start the job.
+                If none specified, defaults to target_instances. If set, the job
+                will not start until the minimum number of nodes is available.
+            env_vars (dict): Environment variables to set in container.
+            enable_metrics (bool): Whether to enable metrics publishing for the job.
+            query_warehouse (str): The query warehouse to use. Defaults to session warehouse.
+            spec_overrides (dict): A dictionary of overrides for the service spec.
 
     Returns:
         An object representing the submitted job.
@@ -277,17 +342,11 @@ def submit_from_stage(
         args=args,
         compute_pool=compute_pool,
         stage_name=stage_name,
-        env_vars=env_vars,
+        target_instances=target_instances,
         pip_requirements=pip_requirements,
         external_access_integrations=external_access_integrations,
-        query_warehouse=query_warehouse,
-        spec_overrides=spec_overrides,
-        target_instances=target_instances,
-        min_instances=min_instances,
-        enable_metrics=enable_metrics,
-        database=database,
-        schema=schema,
         session=session,
+        **kwargs,
     )
 
 
@@ -299,17 +358,11 @@ def _submit_job(
     stage_name: str,
     entrypoint: Optional[str] = None,
     args: Optional[list[str]] = None,
-    env_vars: Optional[dict[str, str]] = None,
+    target_instances: int = 1,
     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,
-    target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> jb.MLJob[None]:
     ...
 
@@ -322,17 +375,11 @@ def _submit_job(
     stage_name: str,
     entrypoint: Optional[str] = None,
     args: Optional[list[str]] = None,
-    env_vars: Optional[dict[str, str]] = None,
+    target_instances: int = 1,
     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,
-    target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> jb.MLJob[T]:
     ...
 
@@ -345,8 +392,9 @@ def _submit_job(
         # TODO: Log lengths of args, env_vars, and spec_overrides values
         "pip_requirements",
         "external_access_integrations",
+        "num_instances",  # deprecated
         "target_instances",
-        "enable_metrics",
+        "min_instances",
     ],
 )
 def _submit_job(
@@ -356,17 +404,9 @@ def _submit_job(
     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,
     target_instances: int = 1,
-    min_instances: Optional[int] = None,
-    enable_metrics: bool = False,
-    database: Optional[str] = None,
-    schema: Optional[str] = None,
     session: Optional[snowpark.Session] = None,
+    **kwargs: Any,
 ) -> jb.MLJob[T]:
     """
     Submit a job to the compute pool.
@@ -377,37 +417,48 @@ def _submit_job(
         stage_name: The name of the stage where the job payload will be uploaded.
         entrypoint: The entry point for the job execution. Required if source is a directory.
         args: A list of arguments to pass to the job.
-        env_vars: Environment variables to set in container
-        pip_requirements: A list of pip requirements for the job.
-        external_access_integrations: A list of external access integrations.
-        query_warehouse: The query warehouse to use. Defaults to session warehouse.
-        spec_overrides: Custom service specification overrides to apply.
         target_instances: The number of instances to use for the job. If none specified, single node job is created.
-        min_instances: The minimum number of nodes required to start the job. If none specified,
-            defaults to target_instances. If set, the job will not start until the minimum number of nodes is available.
-        enable_metrics: Whether to enable metrics publishing for the job.
-        database: The database to use.
-        schema: The schema to use.
         session: The Snowpark session to use. If none specified, uses active session.
+        kwargs: Additional keyword arguments.
 
     Returns:
         An object representing the submitted job.
 
     Raises:
-        RuntimeError: If required Snowflake features are not enabled.
         ValueError: If database or schema value(s) are invalid
+        SnowparkSQLException: If there is an error submitting the job.
     """
+    session = session or get_active_session()
+
+    # Check for deprecated args
+    if "num_instances" in kwargs:
+        logger.warning(
+            "'num_instances' is deprecated and will be removed in a future release. Use 'target_instances' instead."
+        )
+        target_instances = max(target_instances, kwargs.pop("num_instances"))
+
+    # Use kwargs for less common optional parameters
+    database = kwargs.pop("database", None)
+    schema = kwargs.pop("schema", None)
+    min_instances = kwargs.pop("min_instances", target_instances)
+    pip_requirements = kwargs.pop("pip_requirements", None)
+    external_access_integrations = kwargs.pop("external_access_integrations", None)
+    env_vars = kwargs.pop("env_vars", None)
+    spec_overrides = kwargs.pop("spec_overrides", None)
+    enable_metrics = kwargs.pop("enable_metrics", True)
+    query_warehouse = kwargs.pop("query_warehouse", None)
+
+    # Warn if there are unknown kwargs
+    if kwargs:
+        logger.warning(f"Ignoring unknown kwargs: {kwargs.keys()}")
+
+    # Validate parameters
     if database and not schema:
         raise ValueError("Schema must be specified if database is specified.")
     if target_instances < 1:
         raise ValueError("target_instances must be greater than 0.")
-
-    min_instances = target_instances if min_instances is None else min_instances
     if not (0 < min_instances <= target_instances):
         raise ValueError("min_instances must be greater than 0 and less than or equal to target_instances.")
-
-    session = session or get_active_session()
-
     if min_instances > 1:
         # Validate min_instances against compute pool max_nodes
         pool_info = jb._get_compute_pool_info(session, compute_pool)
@@ -418,13 +469,8 @@ def _submit_job(
                 f"of compute pool '{compute_pool}'. Reduce min_instances or increase max_nodes."
             )
 
-    # Validate database and schema identifiers on client side since
-    # SQL parser for EXECUTE JOB SERVICE seems to struggle with this
-    database = identifier.resolve_identifier(cast(str, database or session.get_current_database()))
-    schema = identifier.resolve_identifier(cast(str, schema or session.get_current_schema()))
-
     job_name = f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"
-    job_id = f"{database}.{schema}.{job_name}"
+    job_id = identifier.get_schema_level_object_identifier(database, schema, job_name)
     stage_path_parts = identifier.parse_snowflake_stage_path(stage_name.lstrip("@"))
     stage_name = f"@{'.'.join(filter(None, stage_path_parts[:3]))}"
     stage_path = pathlib.PurePosixPath(f"{stage_name}{stage_path_parts[-1].rstrip('/')}/{job_name}")
@@ -453,7 +499,48 @@ def _submit_job(
     if spec_overrides:
         spec = spec_utils.merge_patch(spec, spec_overrides, display_name="spec_overrides")
 
-    # Generate SQL command for job submission
+    query_text, params = _generate_submission_query(
+        spec, external_access_integrations, query_warehouse, target_instances, session, compute_pool, job_id
+    )
+    try:
+        _ = query_helper.run_query(session, query_text, params=params)
+    except SnowparkSQLException as e:
+        if "Invalid spec: unknown option 'resourceManagement' for 'spec'." in e.message:
+            logger.warning("Dropping 'resourceManagement' from spec because control policy is not enabled.")
+            spec["spec"].pop("resourceManagement", None)
+            query_text, params = _generate_submission_query(
+                spec, external_access_integrations, query_warehouse, target_instances, session, compute_pool, job_id
+            )
+            _ = query_helper.run_query(session, query_text, params=params)
+        else:
+            raise
+    return get_job(job_id, session=session)
+
+
+def _generate_submission_query(
+    spec: dict[str, Any],
+    external_access_integrations: list[str],
+    query_warehouse: Optional[str],
+    target_instances: int,
+    session: snowpark.Session,
+    compute_pool: str,
+    job_id: str,
+) -> tuple[str, list[Any]]:
+    """
+    Generate the SQL query for job submission.
+
+    Args:
+        spec: The service spec for the job.
+        external_access_integrations: The external access integrations for the job.
+        query_warehouse: The query warehouse for the job.
+        target_instances: The number of instances for the job.
+        session: The Snowpark session to use.
+        compute_pool: The compute pool to use for the job.
+        job_id: The ID of the job.
+
+    Returns:
+        A tuple containing the SQL query text and the parameters for the query.
+    """
     query_template = textwrap.dedent(
         """\
         EXECUTE JOB SERVICE
@@ -477,17 +564,5 @@ def _submit_job(
     if target_instances > 1:
         query.append("REPLICAS = ?")
         params.append(target_instances)
-
-    # Submit job
     query_text = "\n".join(line for line in query if line)
-
-    try:
-        _ = session.sql(query_text, params=params).collect()
-    except SnowparkSQLException as e:
-        if "invalid property 'ASYNC'" in e.message:
-            raise RuntimeError(
-                "SPCS Async Jobs not enabled. Set parameter `ENABLE_SNOWSERVICES_ASYNC_JOBS = TRUE` to enable."
-            ) from e
-        raise
-
-    return jb.MLJob(job_id, service_spec=spec, session=session)
+    return query_text, params

snowflake/ml/lineage/lineage_node.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Literal, Optional, Union
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
-from snowflake.ml._internal.utils import identifier
+from snowflake.ml._internal.utils import identifier, mixins
 
 if TYPE_CHECKING:
     from snowflake.ml import dataset
@@ -15,7 +15,7 @@ _PROJECT = "LINEAGE"
 DOMAIN_LINEAGE_REGISTRY: dict[str, type["LineageNode"]] = {}
 
 
-class LineageNode:
+class LineageNode(mixins.SerializableSessionMixin):
     """
     Represents a node in a lineage graph and serves as the base class for all machine learning objects.
     """

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

@@ -1,4 +1,5 @@
 import enum
+import logging
 import pathlib
 import tempfile
 import warnings
@@ -10,7 +11,7 @@ from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.lineage import lineage_node
-from snowflake.ml.model import type_hints as model_types
+from snowflake.ml.model import task, type_hints
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
@@ -401,7 +402,7 @@ class ModelVersion(lineage_node.LineageNode):
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
     )
-    def get_model_task(self) -> model_types.Task:
+    def get_model_task(self) -> task.Task:
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
             subproject=_TELEMETRY_SUBPROJECT,
@@ -607,8 +608,8 @@ class ModelVersion(lineage_node.LineageNode):
         self,
         *,
         force: bool = False,
-        options: Optional[model_types.ModelLoadOption] = None,
-    ) -> model_types.SupportedModelType:
+        options: Optional[type_hints.ModelLoadOption] = None,
+    ) -> type_hints.SupportedModelType:
         """Load the underlying original Python object back from a model.
             This operation requires to have the exact the same environment as the one when logging the model, otherwise,
             the model might be not functional or some other problems might occur.
@@ -889,6 +890,17 @@ class ModelVersion(lineage_node.LineageNode):
             project=_TELEMETRY_PROJECT,
             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. "