snowflake-ml-python 1.30.0__tar.gz → 1.31.0__tar.gz

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/CHANGELOG.md

@@ -1,5 +1,36 @@
 # Release History
 
+## 1.31.0
+
+### New Features
+
+### Bug Fixes
+
+* Registry: Fixed custom model handler to explicitly persist `PARTITIONED=False` in model metadata for
+  `@inference_api` methods. Previously, custom models using `@inference_api` with `TABLE_FUNCTION` type
+  were incorrectly reported as partitioned.
+* Registry: Fixed sklearn and SnowML model handlers to explicitly persist `PARTITIONED=False` metadata for
+  the `explain` method, preventing it from being incorrectly reported as partitioned at read time.
+* Registry: Fixed inconsistent `infer_signature` for columns containing NaN values. Previously,
+  the inferred dtype (`DOUBLE` vs `INT64`) depended on whether NaN rows survived truncation, causing
+  downstream validation failures. Now, the original pandas column dtype is respected during `infer_signature`,
+  so columns with mix of integer and `np.nan` or `None` are consistently inferred as `DOUBLE` regardless of NaN position.
+* Registry: Fixed a bug where invalid parameter types (e.g., passing a string to an integer parameter)
+  were silently coerced instead of raising an error. Parameters are now strictly validated against their
+  declared types before inference. This may break code that relied on the silent coercion behavior
+  (e.g., `model.run(max_tokens="100")` or `model.run(some_int_param=True)`).
+* ML Job: Fixed explicitly named `MLJobDefinition` objects so they can be overwritten in place and invoked
+  repeatedly without job name collisions; the `generate_suffix` argument has been removed.
+* Experiment Tracking live logging (PrPr): Fixed a bug that generated too many lines of log that consist entirely of
+  whitespace.
+
+### Behavior Changes
+
+* Registry: Warehouse partitioned model inference no longer includes non-partition input columns (which were always
+  NULL) in the output. The output now contains only the model's output columns and the partition column.
+
+### Deprecations
+
 ## 1.30.0
 
 ### New Features
@@ -319,6 +350,10 @@ mv = registry.log_model(
 * Registry: Added support for `image-text-to-text` task type in `huggingface.TransformersPipeline`.
   Note: Requires vLLM inference engine while creating the service.
 
+* ML Job: Added support for custom command entrypoints. The `entrypoint` parameter now accepts a
+  `list[str]` (e.g., `["arctic_training", "run_causal.yml"]`), which is executed directly as a
+  command rather than being treated as a Python script.
+
 ### Bug Fixes
 
 ### Behavior Changes

{snowflake_ml_python-1.30.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.31.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.30.0
+Version: 1.31.0
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -416,6 +416,37 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.31.0
+
+### New Features
+
+### Bug Fixes
+
+* Registry: Fixed custom model handler to explicitly persist `PARTITIONED=False` in model metadata for
+  `@inference_api` methods. Previously, custom models using `@inference_api` with `TABLE_FUNCTION` type
+  were incorrectly reported as partitioned.
+* Registry: Fixed sklearn and SnowML model handlers to explicitly persist `PARTITIONED=False` metadata for
+  the `explain` method, preventing it from being incorrectly reported as partitioned at read time.
+* Registry: Fixed inconsistent `infer_signature` for columns containing NaN values. Previously,
+  the inferred dtype (`DOUBLE` vs `INT64`) depended on whether NaN rows survived truncation, causing
+  downstream validation failures. Now, the original pandas column dtype is respected during `infer_signature`,
+  so columns with mix of integer and `np.nan` or `None` are consistently inferred as `DOUBLE` regardless of NaN position.
+* Registry: Fixed a bug where invalid parameter types (e.g., passing a string to an integer parameter)
+  were silently coerced instead of raising an error. Parameters are now strictly validated against their
+  declared types before inference. This may break code that relied on the silent coercion behavior
+  (e.g., `model.run(max_tokens="100")` or `model.run(some_int_param=True)`).
+* ML Job: Fixed explicitly named `MLJobDefinition` objects so they can be overwritten in place and invoked
+  repeatedly without job name collisions; the `generate_suffix` argument has been removed.
+* Experiment Tracking live logging (PrPr): Fixed a bug that generated too many lines of log that consist entirely of
+  whitespace.
+
+### Behavior Changes
+
+* Registry: Warehouse partitioned model inference no longer includes non-partition input columns (which were always
+  NULL) in the output. The output now contains only the model's output columns and the partition column.
+
+### Deprecations
+
 ## 1.30.0
 
 ### New Features
@@ -735,6 +766,10 @@ mv = registry.log_model(
 * Registry: Added support for `image-text-to-text` task type in `huggingface.TransformersPipeline`.
   Note: Requires vLLM inference engine while creating the service.
 
+* ML Job: Added support for custom command entrypoints. The `entrypoint` parameter now accepts a
+  `list[str]` (e.g., `["arctic_training", "run_causal.yml"]`), which is executed directly as a
+  command rather than being treated as a Python script.
+
 ### Bug Fixes
 
 ### Behavior Changes

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/experiment/_logging/experiment_logger.py

@@ -20,14 +20,31 @@ class ExperimentLogger(TextIO):
         self.exp_id = exp_id
         self.run_id = run_id
         self.stream = stream
+        self._buffer = ""  # Buffer to store incomplete lines
+
+    def _write_line(self, line: str) -> None:
+        log_message = {
+            "body": line,
+            "attributes": {
+                "snow.experiment.id": self.exp_id,
+                "snow.experiment.run.id": self.run_id,
+                "snow.experiment.stream": self.stream,
+            },
+        }
+        json_data = json.dumps(log_message)
+        self.file.write(json_data + "\n")
 
     def close(self) -> None:
+        self.flush()
         self.file.close()
 
     def fileno(self) -> int:
         return self.file.fileno()
 
     def flush(self) -> None:
+        if self._buffer:
+            self._write_line(self._buffer)
+            self._buffer = ""
         self.file.flush()
 
     def isatty(self) -> bool:
@@ -61,17 +78,11 @@ class ExperimentLogger(TextIO):
         return True
 
     def write(self, data: str) -> int:
-        log_message = {
-            "body": data,
-            "attributes": {
-                "snow.experiment.id": self.exp_id,
-                "snow.experiment.run.id": self.run_id,
-                "snow.experiment.stream": self.stream,
-            },
-        }
-        json_data = json.dumps(log_message)
-        self.file.write(json_data + "\n")
-        return len(data)  # Return length of input, not JSON output
+        self._buffer += data
+        while "\n" in self._buffer:
+            line, self._buffer = self._buffer.split("\n", 1)
+            self._write_line(line)
+        return len(data)
 
     def writelines(self, lines: Iterable[str]) -> None:
         for line in lines:

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/experiment/experiment_tracking.py

@@ -128,10 +128,17 @@ class ExperimentTracking:
         state["_session"] = None
         state["_sql_client"] = None
         state["_registry"] = None
+        # Save logger buffer state before clearing _logging_context
+        if self._logging_context:
+            state["_stdout_buffer"] = self._logging_context.stdout_logger._buffer
+            state["_stderr_buffer"] = self._logging_context.stderr_logger._buffer
         state["_logging_context"] = None
         return state
 
     def __setstate__(self, state: dict[str, Any]) -> None:
+        stdout_buffer = state.pop("_stdout_buffer", "")
+        stderr_buffer = state.pop("_stderr_buffer", "")
+
         if hasattr(super(), "__setstate__"):
             super().__setstate__(state)  # type: ignore[misc]
         else:
@@ -153,6 +160,10 @@ class ExperimentTracking:
         )
         if self._run:
             self._patch_stdout_and_stderr()
+            # _patch_stdout_and_stderr sets _logging_context if live logging is enabled
+            if self._logging_context:
+                self._logging_context.stdout_logger._buffer = stdout_buffer
+                self._logging_context.stderr_logger._buffer = stderr_buffer
 
     def set_experiment(
         self,

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/jobs/__init__.py

@@ -1,5 +1,5 @@
 from snowflake.ml.jobs._interop.exception_utils import install_exception_display_hooks
-from snowflake.ml.jobs._utils.types import JOB_STATUS
+from snowflake.ml.jobs._utils.type_utils import JOB_STATUS
 from snowflake.ml.jobs.decorators import remote
 from snowflake.ml.jobs.job import MLJob
 from snowflake.ml.jobs.job_definition import MLJobDefinition

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/jobs/_utils/constants.py

@@ -1,5 +1,5 @@
 from snowflake.ml._internal.utils.snowflake_env import SnowflakeCloudType
-from snowflake.ml.jobs._utils.types import ComputeResources
+from snowflake.ml.jobs._utils.type_utils import ComputeResources
 
 # SPCS specification constants
 DEFAULT_CONTAINER_NAME = "main"

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/jobs/_utils/payload_utils.py

@@ -17,7 +17,7 @@ import cloudpickle as cp
 from packaging import version
 
 from snowflake import snowpark
-from snowflake.ml.jobs._utils import constants, query_helper, stage_utils, types
+from snowflake.ml.jobs._utils import constants, query_helper, stage_utils, type_utils
 from snowflake.snowpark import exceptions as sp_exceptions
 from snowflake.snowpark._internal import code_generation
 from snowflake.snowpark._internal.utils import zip_file_or_directory_to_stream
@@ -87,7 +87,7 @@ def _upload_directory(
 
 
 def upload_payloads(
-    session: snowpark.Session, stage_path: PurePath, *payload_specs: types.PayloadSpec, overwrite: bool = True
+    session: snowpark.Session, stage_path: PurePath, *payload_specs: type_utils.PayloadSpec, overwrite: bool = True
 ) -> None:
     for spec in payload_specs:
         source_path = spec.source_path
@@ -156,11 +156,11 @@ def upload_system_resources(session: snowpark.Session, stage_path: PurePath) ->
 
 
 def resolve_source(
-    source: Union[types.PayloadPath, Callable[..., Any]]
-) -> Union[types.PayloadPath, Callable[..., Any]]:
+    source: Union[type_utils.PayloadPath, Callable[..., Any]]
+) -> Union[type_utils.PayloadPath, Callable[..., Any]]:
     if callable(source):
         return source
-    elif isinstance(source, types.PayloadPath):
+    elif isinstance(source, type_utils.PayloadPath):
         if not source.exists():
             raise FileNotFoundError(f"{source} does not exist")
         return source.absolute()
@@ -169,9 +169,9 @@ def resolve_source(
 
 
 def resolve_entrypoint(
-    source: Union[types.PayloadPath, Callable[..., Any]],
-    entrypoint: Optional[Union[types.PayloadPath, list[str]]],
-) -> Union[types.PayloadEntrypoint, list[str]]:
+    source: Union[type_utils.PayloadPath, Callable[..., Any]],
+    entrypoint: Optional[Union[type_utils.PayloadPath, list[str]]],
+) -> Union[type_utils.PayloadEntrypoint, list[str]]:
     """Resolve and validate the entrypoint for a job payload.
 
     Args:
@@ -196,7 +196,7 @@ def resolve_entrypoint(
 
     if callable(source):
         # Entrypoint is generated for callable payloads
-        return types.PayloadEntrypoint(
+        return type_utils.PayloadEntrypoint(
             file_path=entrypoint or Path(constants.DEFAULT_ENTRYPOINT_PATH),
             main_func=_ENTRYPOINT_FUNC_NAME,
         )
@@ -235,13 +235,13 @@ def resolve_entrypoint(
             f" supported={','.join(_SUPPORTED_ENTRYPOINT_EXTENSIONS)} got={entrypoint.suffix}"
         )
 
-    return types.PayloadEntrypoint(
+    return type_utils.PayloadEntrypoint(
         file_path=entrypoint,  # entrypoint is an absolute path at this point
         main_func=None,
     )
 
 
-def get_zip_file_from_path(path: types.PayloadPath) -> types.PayloadPath:
+def get_zip_file_from_path(path: type_utils.PayloadPath) -> type_utils.PayloadPath:
     """Finds the path of the outermost zip archive from a given file path.
 
     Examples:
@@ -269,8 +269,8 @@ def get_zip_file_from_path(path: types.PayloadPath) -> types.PayloadPath:
 
 
 def _finalize_payload_pair(
-    p: types.PayloadPath, base_import_path: Optional[str]
-) -> tuple[types.PayloadPath, Optional[str]]:
+    p: type_utils.PayloadPath, base_import_path: Optional[str]
+) -> tuple[type_utils.PayloadPath, Optional[str]]:
     """Finalize the `(payload_path, import_path)` pair based on source type.
 
     - Zip file: ignore import path (returns `(p, None)`).
@@ -280,11 +280,11 @@ def _finalize_payload_pair(
     - Other files: ignore import path (None).
 
     Args:
-        p (types.PayloadPath): The resolved source path
+        p (type_utils.PayloadPath): The resolved source path
         base_import_path (Optional[str]): Slash-separated import path
 
     Returns:
-        tuple[types.PayloadPath, Optional[str]]: `(p, final_import_path)` where:
+        tuple[type_utils.PayloadPath, Optional[str]]: `(p, final_import_path)` where:
             - `final_import_path` is None for zip archives and non-Python files.
             - `final_import_path` is `base_import_path + ".py"` for Python files when
               `base_import_path` is provided; otherwise None.
@@ -308,14 +308,14 @@ def _finalize_payload_pair(
 
 
 def resolve_import_path(
-    path: Union[types.PayloadPath, ModuleType],
+    path: Union[type_utils.PayloadPath, ModuleType],
     import_path: Optional[str] = None,
-) -> list[tuple[types.PayloadPath, Optional[str]]]:
+) -> list[tuple[type_utils.PayloadPath, Optional[str]]]:
     """
     Resolve and normalize the import path for modules, Python files, or zip payloads.
 
     Args:
-        path (Union[types.PayloadPath, ModuleType]): The source path or module to resolve.
+        path (Union[type_utils.PayloadPath, ModuleType]): The source path or module to resolve.
             - If a directory is provided, it is compressed as a zip archive preserving its structure.
             - If a single Python file is provided, the file itself is zipped.
             - If a module is provided, it is treated as a directory or Python file.
@@ -325,7 +325,7 @@ def resolve_import_path(
             the function infers it from `path`.
 
     Returns:
-        list[tuple[types.PayloadPath, Optional[str]]]: A list of tuples where each tuple
+        list[tuple[type_utils.PayloadPath, Optional[str]]]: A list of tuples where each tuple
         contains the resolved payload path and its corresponding import path (if any).
 
     Raises:
@@ -334,7 +334,7 @@ def resolve_import_path(
         ValueError: If the import path cannot be resolved or is invalid.
     """
     if import_path is None:
-        import_path = path.stem if isinstance(path, types.PayloadPath) else path.__name__
+        import_path = path.stem if isinstance(path, type_utils.PayloadPath) else path.__name__
     import_path = import_path.strip().replace(".", "/") if import_path else None
     if isinstance(path, Path):
         if not path.exists():
@@ -357,7 +357,7 @@ def resolve_import_path(
         raise ValueError(f"Module {path} is not a valid imports")
 
 
-def validate_import_path(source: Union[str, types.PayloadPath], import_path: Optional[str]) -> None:
+def validate_import_path(source: Union[str, type_utils.PayloadPath], import_path: Optional[str]) -> None:
     """Validate the import path for local python file or directory."""
     if import_path is None:
         return
@@ -425,7 +425,7 @@ def upload_imports(
                 remote = None
                 compress = False
 
-            upload_payloads(session, stage_path, types.PayloadSpec(source_path, remote, compress=compress))
+            upload_payloads(session, stage_path, type_utils.PayloadSpec(source_path, remote, compress=compress))
 
 
 class JobPayload:
@@ -453,7 +453,7 @@ class JobPayload:
         # for stage path like snow://domain....., Path(path) will remove duplicate /, it will become snow:/ domain...
         self.source = stage_utils.resolve_path(source) if isinstance(source, str) else source
         if isinstance(entrypoint, list):
-            self.entrypoint: Optional[Union[types.PayloadPath, list[str]]] = entrypoint
+            self.entrypoint: Optional[Union[type_utils.PayloadPath, list[str]]] = entrypoint
         else:
             self.entrypoint = stage_utils.resolve_path(entrypoint) if isinstance(entrypoint, str) else entrypoint
         self.pip_requirements = pip_requirements
@@ -461,7 +461,7 @@ class JobPayload:
 
     def upload(
         self, session: snowpark.Session, stage_path: Union[str, PurePath], overwrite: bool = False
-    ) -> types.UploadedPayload:
+    ) -> type_utils.UploadedPayload:
         # Prepare local variables
         stage_path = PurePath(stage_path) if isinstance(stage_path, str) else stage_path
         source = resolve_source(self.source)
@@ -490,13 +490,13 @@ class JobPayload:
         if isinstance(entrypoint, (list, tuple)):
             # For list entrypoints, still upload source if it's a path
             if isinstance(source, Path):
-                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
+                upload_payloads(session, app_stage_path, type_utils.PayloadSpec(source, None))
             elif isinstance(source, stage_utils.StagePath):
-                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
+                upload_payloads(session, app_stage_path, type_utils.PayloadSpec(source, None))
             python_entrypoint: list[Union[str, PurePath]] = list(entrypoint)
         else:
             # Standard file-based entrypoint handling
-            if not isinstance(source, types.PayloadPath):
+            if not isinstance(source, type_utils.PayloadPath):
                 source_code = generate_python_code(source, source_code_display=True)
                 _ = session.file.put_stream(
                     io.BytesIO(source_code.encode()),
@@ -510,10 +510,10 @@ class JobPayload:
                 # copy payload to stage
                 if source == entrypoint.file_path:
                     source = source.parent
-                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None), overwrite=overwrite)
+                upload_payloads(session, app_stage_path, type_utils.PayloadSpec(source, None), overwrite=overwrite)
 
             elif isinstance(source, Path):
-                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None), overwrite=overwrite)
+                upload_payloads(session, app_stage_path, type_utils.PayloadSpec(source, None), overwrite=overwrite)
                 if source.is_file():
                     source = source.parent
 
@@ -555,7 +555,7 @@ class JobPayload:
             constants.PAYLOAD_DIR_ENV_VAR: constants.APP_STAGE_SUBPATH,
         }
 
-        return types.UploadedPayload(
+        return type_utils.UploadedPayload(
             stage_path=stage_path,
             entrypoint=[
                 "bash",
@@ -733,15 +733,20 @@ if __name__ == '__main__':
 """
 
 
-def get_payload_name(source: Union[str, Callable[..., Any]], entrypoint: Optional[Union[str, list[str]]] = None) -> str:
+def get_payload_name(
+    source: Union[str, Callable[..., Any]],
+    entrypoint: Optional[Union[str, list[str]]] = None,
+    generate_suffix: bool = True,
+) -> str:
+    suffix = f"_{str(uuid4().hex)[:8]}" if generate_suffix else ""
 
     if entrypoint and isinstance(entrypoint, (list, tuple)):
-        return entrypoint[0]
+        return entrypoint[0] + suffix
     elif entrypoint and isinstance(entrypoint, str):
-        return f"{PurePath(entrypoint).stem}"
+        return PurePath(entrypoint).stem + suffix
     elif source and not callable(source):
-        return f"{PurePath(source).stem}"
+        return PurePath(source).stem + suffix
     elif callable(source):
-        return f"{source.__name__}"
+        return source.__name__ + suffix
     else:
-        return f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"
+        return f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"  # NOTE: Already includes a suffix

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/jobs/_utils/stage_utils.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Union
 from snowflake.ml._internal.utils import identifier
 
 if TYPE_CHECKING:
-    from snowflake.ml.jobs._utils import types
+    from snowflake.ml.jobs._utils import type_utils
 
 PROTOCOL_NAME = "snow"
 _SNOWURL_PATH_RE = re.compile(
@@ -155,7 +155,7 @@ class StagePath:
         return path
 
 
-def resolve_path(path: Union[str, Path]) -> "types.PayloadPath":
+def resolve_path(path: Union[str, Path]) -> "type_utils.PayloadPath":
     """
     Resolve a path to either a StagePath or a local Path.
 

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/jobs/decorators.py

@@ -73,7 +73,6 @@ def remote(
             external_access_integrations=external_access_integrations,
             session=session,
             arg_protocol=arg_protocol.ArgProtocol.PICKLE,
-            generate_suffix=True,
             **kwargs,
         )
 

{snowflake_ml_python-1.30.0 → snowflake_ml_python-1.31.0}/snowflake/ml/jobs/job.py

@@ -13,7 +13,7 @@ from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import identifier
 from snowflake.ml._internal.utils.mixins import SerializableSessionMixin
 from snowflake.ml.jobs._interop import results as interop_result, utils as interop_utils
-from snowflake.ml.jobs._utils import constants, query_helper, stage_utils, types
+from snowflake.ml.jobs._utils import constants, query_helper, stage_utils, type_utils
 from snowflake.snowpark import Row, context as sp_context
 from snowflake.snowpark.exceptions import SnowparkSQLException
 
@@ -38,11 +38,11 @@ class MLJob(Generic[T], SerializableSessionMixin):
         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._status: type_utils.JOB_STATUS = "PENDING"
         self._result: Optional[interop_result.ExecutionResult] = None
 
     @cached_property
-    def _service_info(self) -> types.ServiceInfo:
+    def _service_info(self) -> type_utils.ServiceInfo:
         """Get the job's service info."""
         return _resolve_service_info(self.id, self._session)
 
@@ -67,7 +67,7 @@ class MLJob(Generic[T], SerializableSessionMixin):
         return self._id
 
     @property
-    def status(self) -> types.JOB_STATUS:
+    def status(self) -> type_utils.JOB_STATUS:
         """Get the job's execution status."""
         if self._status not in TERMINAL_JOB_STATUSES:
             # Query backend for job status if not in terminal state
@@ -232,7 +232,7 @@ class MLJob(Generic[T], SerializableSessionMixin):
         print(self.get_logs(limit, instance_id, as_list=False, verbose=verbose))  # noqa: T201: we need to print here.
 
     @telemetry.send_api_usage_telemetry(project=_PROJECT, func_params_to_log=["timeout"])
-    def wait(self, timeout: float = -1) -> types.JOB_STATUS:
+    def wait(self, timeout: float = -1) -> type_utils.JOB_STATUS:
         """
         Block until completion. Returns completion status.
 
@@ -282,6 +282,10 @@ class MLJob(Generic[T], SerializableSessionMixin):
     def get_ray_dashboard_url(self) -> Optional[str]:
         """
         Get the Ray dashboard URL for the job.
+
+        Returns:
+            Optional[str]: The Ray dashboard URL if the job is running and has a Ray dashboard endpoint,
+                None otherwise.
         """
         if self.status != JOB_RUNNING_STATUS:
             logger.warning("Ray dashboard is not available for non-running jobs")
@@ -334,7 +338,7 @@ class MLJob(Generic[T], SerializableSessionMixin):
 
 
 @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:
+def _get_status(session: snowpark.Session, job_id: str, instance_id: Optional[int] = None) -> type_utils.JOB_STATUS:
     """Retrieve job or job instance execution status."""
     try:
         if instance_id is not None:
@@ -342,15 +346,15 @@ def _get_status(session: snowpark.Session, job_id: str, instance_id: Optional[in
             rows = query_helper.run_query(session, "SHOW SERVICE INSTANCES IN SERVICE IDENTIFIER(?)", params=(job_id,))
             for row in rows:
                 if row["instance_id"] == str(instance_id):
-                    return cast(types.JOB_STATUS, row["status"])
+                    return cast(type_utils.JOB_STATUS, row["status"])
             raise ValueError(f"Instance {instance_id} not found in job {job_id}")
         else:
             row = _get_service_info(session, job_id)
-            return cast(types.JOB_STATUS, row["status"])
+            return cast(type_utils.JOB_STATUS, row["status"])
     except SnowparkSQLException as e:
         if e.sql_error_code == 2003:
             row = _get_service_info_spcs(session, job_id)
-            return cast(types.JOB_STATUS, row["STATUS"])
+            return cast(type_utils.JOB_STATUS, row["STATUS"])
         raise
 
 
@@ -708,7 +712,7 @@ def _get_service_info_spcs(session: snowpark.Session, job_id: str) -> Any:
         raise SnowparkSQLException(f"Job {job_id} does not exist or could not be retrieved", sql_error_code=2003)
 
 
-def _resolve_service_info(id: str, session: snowpark.Session) -> types.ServiceInfo:
+def _resolve_service_info(id: str, session: snowpark.Session) -> type_utils.ServiceInfo:
     try:
         row = _get_service_info(session, id)
     except SnowparkSQLException as e:
@@ -751,10 +755,10 @@ def _resolve_service_info(id: str, session: snowpark.Session) -> types.ServiceIn
     database_name = row["database_name"] if "database_name" in row else row["DATABASE_NAME"]
     schema_name = row["schema_name"] if "schema_name" in row else row["SCHEMA_NAME"]
 
-    return types.ServiceInfo(
+    return type_utils.ServiceInfo(
         database_name=database_name,
         schema_name=schema_name,
-        status=cast(types.JOB_STATUS, status),
+        status=cast(type_utils.JOB_STATUS, status),
         compute_pool=cast(str, compute_pool),
         target_instances=target_instances,
     )