snowflake-ml-python 1.16.0__py3-none-any.whl → 1.17.0__py3-none-any.whl

snowflake/ml/jobs/_interop/results.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+from typing import Any, Optional
+
+
+@dataclass(frozen=True)
+class ExecutionResult:
+    """
+    A result of a job execution.
+
+    Args:
+        success: Whether the execution was successful.
+        value: The value of the execution.
+    """
+
+    success: bool
+    value: Any
+
+    def get_value(self, wrap_exceptions: bool = True) -> Any:
+        if not self.success:
+            assert isinstance(self.value, BaseException), "Unexpected non-exception value for failed result"
+            self._raise_exception(self.value, wrap_exceptions)
+        return self.value
+
+    def _raise_exception(self, exception: BaseException, wrap_exceptions: bool) -> None:
+        if wrap_exceptions:
+            raise RuntimeError(f"Job execution failed with error: {exception!r}") from exception
+        else:
+            raise exception
+
+
+@dataclass(frozen=True)
+class LoadedExecutionResult(ExecutionResult):
+    """
+    A result of a job execution that has been loaded from a file.
+    """
+
+    load_error: Optional[Exception] = None
+    result_metadata: Optional[dict[str, Any]] = None
+
+    def get_value(self, wrap_exceptions: bool = True) -> Any:
+        if not self.success:
+            # Raise the original exception if available, otherwise raise the load error
+            ex = self.value
+            if not isinstance(ex, BaseException):
+                ex = RuntimeError(f"Unknown error {ex or ''}")
+                ex.__cause__ = self.load_error
+            self._raise_exception(ex, wrap_exceptions)
+        else:
+            if self.load_error:
+                raise ValueError("Job execution succeeded but result retrieval failed") from self.load_error
+            return self.value

snowflake/ml/jobs/_interop/utils.py

@@ -0,0 +1,144 @@
+import logging
+import os
+import traceback
+from pathlib import PurePath
+from typing import Any, Callable, Optional
+
+import pydantic
+
+from snowflake import snowpark
+from snowflake.ml.jobs._interop import data_utils, exception_utils, legacy, protocols
+from snowflake.ml.jobs._interop.dto_schema import (
+    ExceptionMetadata,
+    ResultDTO,
+    ResultMetadata,
+)
+from snowflake.ml.jobs._interop.results import ExecutionResult, LoadedExecutionResult
+from snowflake.snowpark import exceptions as sp_exceptions
+
+DEFAULT_CODEC = data_utils.JsonDtoCodec
+DEFAULT_PROTOCOL = protocols.AutoProtocol()
+DEFAULT_PROTOCOL.try_register_protocol(protocols.CloudPickleProtocol)
+DEFAULT_PROTOCOL.try_register_protocol(protocols.ArrowTableProtocol)
+DEFAULT_PROTOCOL.try_register_protocol(protocols.PandasDataFrameProtocol)
+DEFAULT_PROTOCOL.try_register_protocol(protocols.NumpyArrayProtocol)
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_result(result: ExecutionResult, path: str, session: Optional[snowpark.Session] = None) -> None:
+    """
+    Save the result to a file.
+    """
+    result_dto = ResultDTO(
+        success=result.success,
+        value=result.value,
+    )
+
+    try:
+        # Try to encode result directly
+        payload = DEFAULT_CODEC.encode(result_dto)
+    except TypeError:
+        result_dto.value = None  # Remove raw value to avoid serialization error
+        result_dto.metadata = _get_metadata(result.value)  # Add metadata for client fallback on protocol mismatch
+        try:
+            path_dir = PurePath(path).parent.as_posix()
+            protocol_info = DEFAULT_PROTOCOL.save(result.value, path_dir, session=session)
+            result_dto.protocol = protocol_info
+
+        except Exception as e:
+            logger.warning(f"Error dumping result value: {repr(e)}")
+            result_dto.serialize_error = repr(e)
+
+        # Encode the modified result DTO
+        payload = DEFAULT_CODEC.encode(result_dto)
+
+    with data_utils.open_stream(path, "wb", session=session) as stream:
+        stream.write(payload)
+
+
+def load_result(
+    path: str, session: Optional[snowpark.Session] = None, path_transform: Optional[Callable[[str], str]] = None
+) -> ExecutionResult:
+    """Load the result from a file on a Snowflake stage."""
+    try:
+        with data_utils.open_stream(path, "r", session=session) as stream:
+            # Load the DTO as a dict for easy fallback to legacy loading if necessary
+            dto_dict = DEFAULT_CODEC.decode(stream, as_dict=True)
+    except UnicodeDecodeError:
+        # Path may be a legacy result file (cloudpickle)
+        # TODO: Re-use the stream
+        assert session is not None
+        return legacy.load_legacy_result(session, path)
+
+    try:
+        dto = ResultDTO.model_validate(dto_dict)
+    except pydantic.ValidationError as e:
+        if "success" in dto_dict:
+            assert session is not None
+            if path.endswith(".json"):
+                path = os.path.splitext(path)[0] + ".pkl"
+            return legacy.load_legacy_result(session, path, result_json=dto_dict)
+        raise ValueError("Invalid result schema") from e
+
+    # Try loading data from file using the protocol info
+    result_value = None
+    data_load_error = None
+    if dto.protocol is not None:
+        try:
+            logger.debug(f"Loading result value with protocol {dto.protocol}")
+            result_value = DEFAULT_PROTOCOL.load(dto.protocol, session=session, path_transform=path_transform)
+        except sp_exceptions.SnowparkSQLException:
+            raise  # Data retrieval errors should be bubbled up
+        except Exception as e:
+            logger.debug(f"Error loading result value with protocol {dto.protocol}: {repr(e)}")
+            data_load_error = e
+
+    # Wrap serialize_error in a TypeError
+    if dto.serialize_error:
+        serialize_error = TypeError("Original result serialization failed with error: " + dto.serialize_error)
+        if data_load_error:
+            data_load_error.__context__ = serialize_error
+        else:
+            data_load_error = serialize_error
+
+    # Prepare to assemble the final result
+    result_value = result_value if result_value is not None else dto.value
+    if not dto.success and result_value is None:
+        # Try to reconstruct exception from metadata if available
+        if isinstance(dto.metadata, ExceptionMetadata):
+            logger.debug(f"Reconstructing exception from metadata {dto.metadata}")
+            result_value = exception_utils.build_exception(
+                type_str=dto.metadata.type,
+                message=dto.metadata.message,
+                traceback=dto.metadata.traceback,
+                original_repr=dto.metadata.repr,
+            )
+
+        # Generate a generic error if we still don't have a value,
+        # attaching the data load error if any
+        if result_value is None:
+            result_value = exception_utils.RemoteError("Unknown remote error")
+            result_value.__cause__ = data_load_error
+
+    return LoadedExecutionResult(
+        success=dto.success,
+        value=result_value,
+        load_error=data_load_error,
+    )
+
+
+def _get_metadata(value: Any) -> ResultMetadata:
+    type_name = f"{type(value).__module__}.{type(value).__name__}"
+    if isinstance(value, BaseException):
+        return ExceptionMetadata(
+            type=type_name,
+            repr=repr(value),
+            message=str(value),
+            traceback="".join(traceback.format_tb(value.__traceback__)),
+        )
+    return ResultMetadata(
+        type=type_name,
+        repr=repr(value),
+    )

snowflake/ml/jobs/_utils/constants.py

@@ -12,6 +12,9 @@ PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
 RESULT_PATH_ENV_VAR = "MLRS_RESULT_PATH"
 MIN_INSTANCES_ENV_VAR = "MLRS_MIN_INSTANCES"
 TARGET_INSTANCES_ENV_VAR = "SNOWFLAKE_JOBS_COUNT"
+INSTANCES_MIN_WAIT_ENV_VAR = "MLRS_INSTANCES_MIN_WAIT"
+INSTANCES_TIMEOUT_ENV_VAR = "MLRS_INSTANCES_TIMEOUT"
+INSTANCES_CHECK_INTERVAL_ENV_VAR = "MLRS_INSTANCES_CHECK_INTERVAL"
 RUNTIME_IMAGE_TAG_ENV_VAR = "MLRS_CONTAINER_IMAGE_TAG"
 
 # Stage mount paths
@@ -19,7 +22,7 @@ STAGE_VOLUME_MOUNT_PATH = "/mnt/job_stage"
 APP_STAGE_SUBPATH = "app"
 SYSTEM_STAGE_SUBPATH = "system"
 OUTPUT_STAGE_SUBPATH = "output"
-RESULT_PATH_DEFAULT_VALUE = f"{OUTPUT_STAGE_SUBPATH}/mljob_result.pkl"
+RESULT_PATH_DEFAULT_VALUE = f"{OUTPUT_STAGE_SUBPATH}/mljob_result"
 
 # Default container image information
 DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"

snowflake/ml/jobs/_utils/feature_flags.py

@@ -1,16 +1,48 @@
 import os
 from enum import Enum
+from typing import Optional
+
+
+def parse_bool_env_value(value: Optional[str], default: bool = False) -> bool:
+    """Parse a boolean value from an environment variable string.
+
+    Args:
+        value: The environment variable value to parse (may be None).
+        default: The default value to return if the value is None or unrecognized.
+
+    Returns:
+        True if the value is a truthy string (true, 1, yes, on - case insensitive),
+        False if the value is a falsy string (false, 0, no, off - case insensitive),
+        or the default value if the value is None or unrecognized.
+    """
+    if value is None:
+        return default
+
+    normalized_value = value.strip().lower()
+    if normalized_value in ("true", "1", "yes", "on"):
+        return True
+    elif normalized_value in ("false", "0", "no", "off"):
+        return False
+    else:
+        # For unrecognized values, return the default
+        return default
 
 
 class FeatureFlags(Enum):
     USE_SUBMIT_JOB_V2 = "MLRS_USE_SUBMIT_JOB_V2"
-    ENABLE_IMAGE_VERSION_ENV_VAR = "MLRS_ENABLE_RUNTIME_VERSIONS"
+    ENABLE_RUNTIME_VERSIONS = "MLRS_ENABLE_RUNTIME_VERSIONS"
+
+    def is_enabled(self, default: bool = False) -> bool:
+        """Check if the feature flag is enabled.
 
-    def is_enabled(self) -> bool:
-        return os.getenv(self.value, "false").lower() == "true"
+        Args:
+            default: The default value to return if the environment variable is not set.
 
-    def is_disabled(self) -> bool:
-        return not self.is_enabled()
+        Returns:
+            True if the environment variable is set to a truthy value,
+            False if set to a falsy value, or the default value if not set.
+        """
+        return parse_bool_env_value(os.getenv(self.value), default)
 
     def __str__(self) -> str:
         return self.value

snowflake/ml/jobs/_utils/payload_utils.py

@@ -268,7 +268,7 @@ def upload_payloads(session: snowpark.Session, stage_path: PurePath, *payload_sp
                 # can't handle directories. Reduce the number of PUT operations by using
                 # wildcard patterns to batch upload files with the same extension.
                 upload_path_patterns = set()
-                for p in source_path.resolve().rglob("*"):
+                for p in source_path.rglob("*"):
                     if p.is_dir():
                         continue
                     if p.name.startswith("."):

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

@@ -9,19 +9,23 @@ import runpy
 import sys
 import time
 import traceback
-import warnings
-from pathlib import Path
 from typing import Any, Optional
 
-import cloudpickle
-
-from snowflake.ml.jobs._utils import constants
-from snowflake.snowpark import Session
-
-try:
-    from snowflake.ml._internal.utils.connection_params import SnowflakeLoginOptions
-except ImportError:
-    from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
+# Ensure payload directory is in sys.path for module imports before importing other modules
+# This is needed to support relative imports in user scripts and to allow overriding
+# modules using modules in the payload directory
+# TODO: Inject the environment variable names at job submission time
+STAGE_MOUNT_PATH = os.environ.get("MLRS_STAGE_MOUNT_PATH", "/mnt/job_stage")
+JOB_RESULT_PATH = os.environ.get("MLRS_RESULT_PATH", "output/mljob_result.pkl")
+PAYLOAD_PATH = os.environ.get("MLRS_PAYLOAD_DIR")
+if PAYLOAD_PATH and not os.path.isabs(PAYLOAD_PATH):
+    PAYLOAD_PATH = os.path.join(STAGE_MOUNT_PATH, PAYLOAD_PATH)
+if PAYLOAD_PATH and PAYLOAD_PATH not in sys.path:
+    sys.path.insert(0, PAYLOAD_PATH)
+
+# Imports below must come after sys.path modification to support module overrides
+import snowflake.ml.jobs._utils.constants  # noqa: E402
+import snowflake.snowpark  # noqa: E402
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -33,48 +37,74 @@ logger = logging.getLogger(__name__)
 # not have the latest version of the code
 # Log start and end messages
 LOG_START_MSG = getattr(
-    constants,
+    snowflake.ml.jobs._utils.constants,
     "LOG_START_MSG",
     "--------------------------------\nML job started\n--------------------------------",
 )
 LOG_END_MSG = getattr(
-    constants,
+    snowflake.ml.jobs._utils.constants,
     "LOG_END_MSG",
     "--------------------------------\nML job finished\n--------------------------------",
 )
+MIN_INSTANCES_ENV_VAR = getattr(
+    snowflake.ml.jobs._utils.constants,
+    "MIN_INSTANCES_ENV_VAR",
+    "MLRS_MIN_INSTANCES",
+)
+TARGET_INSTANCES_ENV_VAR = getattr(
+    snowflake.ml.jobs._utils.constants,
+    "TARGET_INSTANCES_ENV_VAR",
+    "SNOWFLAKE_JOBS_COUNT",
+)
+INSTANCES_MIN_WAIT_ENV_VAR = getattr(
+    snowflake.ml.jobs._utils.constants,
+    "INSTANCES_MIN_WAIT_ENV_VAR",
+    "MLRS_INSTANCES_MIN_WAIT",
+)
+INSTANCES_TIMEOUT_ENV_VAR = getattr(
+    snowflake.ml.jobs._utils.constants,
+    "INSTANCES_TIMEOUT_ENV_VAR",
+    "MLRS_INSTANCES_TIMEOUT",
+)
+INSTANCES_CHECK_INTERVAL_ENV_VAR = getattr(
+    snowflake.ml.jobs._utils.constants,
+    "INSTANCES_CHECK_INTERVAL_ENV_VAR",
+    "MLRS_INSTANCES_CHECK_INTERVAL",
+)
 
-# min_instances environment variable name
-MIN_INSTANCES_ENV_VAR = getattr(constants, "MIN_INSTANCES_ENV_VAR", "MLRS_MIN_INSTANCES")
-TARGET_INSTANCES_ENV_VAR = getattr(constants, "TARGET_INSTANCES_ENV_VAR", "SNOWFLAKE_JOBS_COUNT")
-
-# Fallbacks in case of SnowML version mismatch
-STAGE_MOUNT_PATH_ENV_VAR = getattr(constants, "STAGE_MOUNT_PATH_ENV_VAR", "MLRS_STAGE_MOUNT_PATH")
-RESULT_PATH_ENV_VAR = getattr(constants, "RESULT_PATH_ENV_VAR", "MLRS_RESULT_PATH")
-PAYLOAD_DIR_ENV_VAR = getattr(constants, "PAYLOAD_DIR_ENV_VAR", "MLRS_PAYLOAD_DIR")
 
 # Constants for the wait_for_instances function
-MIN_WAIT_TIME = float(os.getenv("MLRS_INSTANCES_MIN_WAIT") or -1)  # seconds
-TIMEOUT = float(os.getenv("MLRS_INSTANCES_TIMEOUT") or 720)  # seconds
-CHECK_INTERVAL = float(os.getenv("MLRS_INSTANCES_CHECK_INTERVAL") or 10)  # seconds
+MIN_INSTANCES = int(os.environ.get(MIN_INSTANCES_ENV_VAR) or "1")
+TARGET_INSTANCES = int(os.environ.get(TARGET_INSTANCES_ENV_VAR) or MIN_INSTANCES)
+MIN_WAIT_TIME = float(os.getenv(INSTANCES_MIN_WAIT_ENV_VAR) or -1)  # seconds
+TIMEOUT = float(os.getenv(INSTANCES_TIMEOUT_ENV_VAR) or 720)  # seconds
+CHECK_INTERVAL = float(os.getenv(INSTANCES_CHECK_INTERVAL_ENV_VAR) or 10)  # seconds
 
-STAGE_MOUNT_PATH = os.environ.get(STAGE_MOUNT_PATH_ENV_VAR, "/mnt/job_stage")
-JOB_RESULT_PATH = os.environ.get(RESULT_PATH_ENV_VAR, "output/mljob_result.pkl")
 
+def save_mljob_result_v2(value: Any, is_error: bool, path: str) -> None:
+    from snowflake.ml.jobs._interop import (
+        results as interop_result,
+        utils as interop_utils,
+    )
+
+    result_obj = interop_result.ExecutionResult(success=not is_error, value=value)
+    interop_utils.save_result(result_obj, path)
 
-try:
-    from snowflake.ml.jobs._utils.interop_utils import ExecutionResult
-except ImportError:
+
+def save_mljob_result_v1(value: Any, is_error: bool, path: str) -> None:
     from dataclasses import dataclass
 
+    import cloudpickle
+
+    # Directly in-line the ExecutionResult class since the legacy type
+    # instead of attempting to import the to-be-deprecated
+    # snowflake.ml.jobs._utils.interop module
+    # Eventually, this entire function will be removed in favor of v2
     @dataclass(frozen=True)
-    class ExecutionResult:  # type: ignore[no-redef]
+    class ExecutionResult:
         result: Optional[Any] = None
         exception: Optional[BaseException] = None
 
-        @property
-        def success(self) -> bool:
-            return self.exception is None
-
         def to_dict(self) -> dict[str, Any]:
             """Return the serializable dictionary."""
             if isinstance(self.exception, BaseException):
@@ -91,14 +121,45 @@ except ImportError:
                 "result": self.result,
             }
 
+    # Create a custom JSON encoder that converts non-serializable types to strings
+    class SimpleJSONEncoder(json.JSONEncoder):
+        def default(self, obj: Any) -> Any:
+            try:
+                return super().default(obj)
+            except TypeError:
+                return f"Unserializable object: {repr(obj)}"
+
+    result_obj = ExecutionResult(result=None if is_error else value, exception=value if is_error else None)
+    result_dict = result_obj.to_dict()
+    try:
+        # Serialize result using cloudpickle
+        result_pickle_path = path
+        with open(result_pickle_path, "wb") as f:
+            cloudpickle.dump(result_dict, f)  # Pickle dictionary form for compatibility
+    except Exception as pkl_exc:
+        logger.warning(f"Failed to pickle result to {result_pickle_path}: {pkl_exc}")
 
-# Create a custom JSON encoder that converts non-serializable types to strings
-class SimpleJSONEncoder(json.JSONEncoder):
-    def default(self, obj: Any) -> Any:
-        try:
-            return super().default(obj)
-        except TypeError:
-            return f"Unserializable object: {repr(obj)}"
+    try:
+        # Serialize result to JSON as fallback path in case of cross version incompatibility
+        result_json_path = os.path.splitext(path)[0] + ".json"
+        with open(result_json_path, "w") as f:
+            json.dump(result_dict, f, indent=2, cls=SimpleJSONEncoder)
+    except Exception as json_exc:
+        logger.warning(f"Failed to serialize JSON result to {result_json_path}: {json_exc}")
+
+
+def save_mljob_result(result_obj: Any, is_error: bool, path: str) -> None:
+    """Saves the result or error message to a file in the stage mount path.
+
+    Args:
+        result_obj: The result object to save, either the return value or the exception.
+        is_error: Whether the result_obj is a raised exception.
+        path: The file path to save the result to.
+    """
+    try:
+        save_mljob_result_v2(result_obj, is_error, path)
+    except ImportError:
+        save_mljob_result_v1(result_obj, is_error, path)
 
 
 def wait_for_instances(
@@ -225,20 +286,10 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
     original_argv = sys.argv
     sys.argv = [script_path, *script_args]
 
-    # Ensure payload directory is in sys.path for module imports
-    # This is needed because mljob_launcher.py is now in /mnt/job_stage/system
-    # but user scripts are in the payload directory and may import from each other
-    payload_dir = os.environ.get(PAYLOAD_DIR_ENV_VAR)
-    if payload_dir and not os.path.isabs(payload_dir):
-        payload_dir = os.path.join(STAGE_MOUNT_PATH, payload_dir)
-    if payload_dir and payload_dir not in sys.path:
-        sys.path.insert(0, payload_dir)
-
     try:
-
         if main_func:
             # Use importlib for scripts with a main function defined
-            module_name = Path(script_path).stem
+            module_name = os.path.splitext(os.path.basename(script_path))[0]
             spec = importlib.util.spec_from_file_location(module_name, script_path)
             assert spec is not None
             assert spec.loader is not None
@@ -262,7 +313,7 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
         sys.argv = original_argv
 
 
-def main(script_path: str, *script_args: Any, script_main_func: Optional[str] = None) -> ExecutionResult:
+def main(script_path: str, *script_args: Any, script_main_func: Optional[str] = None) -> Any:
     """Executes a Python script and serializes the result to JOB_RESULT_PATH.
 
     Args:
@@ -271,55 +322,53 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
         script_main_func (str, optional): The name of the function to call in the script (if any).
 
     Returns:
-        ExecutionResult: Object containing execution results.
+        Any: The result of the script execution.
 
     Raises:
         Exception: Re-raises any exception caught during script execution.
     """
-    # Ensure the output directory exists before trying to write result files.
-    result_abs_path = (
-        JOB_RESULT_PATH if os.path.isabs(JOB_RESULT_PATH) else os.path.join(STAGE_MOUNT_PATH, JOB_RESULT_PATH)
-    )
-    output_dir = os.path.dirname(result_abs_path)
-    os.makedirs(output_dir, exist_ok=True)
+    try:
+        from snowflake.ml._internal.utils.connection_params import SnowflakeLoginOptions
+    except ImportError:
+        from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
 
+    # Initialize Ray if available
     try:
         import ray
 
         ray.init(address="auto")
     except ModuleNotFoundError:
-        warnings.warn("Ray is not installed, skipping Ray initialization", ImportWarning, stacklevel=1)
+        logger.debug("Ray is not installed, skipping Ray initialization")
 
     # Create a Snowpark session before starting
     # Session can be retrieved from using snowflake.snowpark.context.get_active_session()
     config = SnowflakeLoginOptions()
     config["client_session_keep_alive"] = "True"
-    session = Session.builder.configs(config).create()  # noqa: F841
+    session = snowflake.snowpark.Session.builder.configs(config).create()  # noqa: F841
 
+    execution_result_is_error = False
+    execution_result_value = None
     try:
-        # Wait for minimum required instances if specified
-        min_instances_str = os.environ.get(MIN_INSTANCES_ENV_VAR) or "1"
-        target_instances_str = os.environ.get(TARGET_INSTANCES_ENV_VAR) or min_instances_str
-        if target_instances_str and int(target_instances_str) > 1:
-            wait_for_instances(
-                int(min_instances_str),
-                int(target_instances_str),
-                min_wait_time=MIN_WAIT_TIME,
-                timeout=TIMEOUT,
-                check_interval=CHECK_INTERVAL,
-            )
-
-        # Log start marker for user script execution
+        # Wait for minimum required instances before starting user script execution
+        wait_for_instances(
+            MIN_INSTANCES,
+            TARGET_INSTANCES,
+            min_wait_time=MIN_WAIT_TIME,
+            timeout=TIMEOUT,
+            check_interval=CHECK_INTERVAL,
+        )
+
+        # Log start marker before starting user script execution
         print(LOG_START_MSG)  # noqa: T201
 
-        # Run the script with the specified arguments
-        result = run_script(script_path, *script_args, main_func=script_main_func)
+        # Run the user script
+        execution_result_value = run_script(script_path, *script_args, main_func=script_main_func)
 
         # Log end marker for user script execution
         print(LOG_END_MSG)  # noqa: T201
 
-        result_obj = ExecutionResult(result=result)
-        return result_obj
+        return execution_result_value
+
     except Exception as e:
         tb = e.__traceback__
         skip_files = {__file__, runpy.__file__}
@@ -328,35 +377,23 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
             tb = tb.tb_next
         cleaned_ex = copy.copy(e)  # Need to create a mutable copy of exception to set __traceback__
         cleaned_ex = cleaned_ex.with_traceback(tb)
-        result_obj = ExecutionResult(exception=cleaned_ex)
+        execution_result_value = cleaned_ex
+        execution_result_is_error = True
         raise
     finally:
-        result_dict = result_obj.to_dict()
-        try:
-            # Serialize result using cloudpickle
-            result_pickle_path = result_abs_path
-            with open(result_pickle_path, "wb") as f:
-                cloudpickle.dump(result_dict, f)  # Pickle dictionary form for compatibility
-        except Exception as pkl_exc:
-            warnings.warn(f"Failed to pickle result to {result_pickle_path}: {pkl_exc}", RuntimeWarning, stacklevel=1)
-
-        try:
-            # Serialize result to JSON as fallback path in case of cross version incompatibility
-            # TODO: Manually convert non-serializable types to strings
-            result_json_path = os.path.splitext(result_abs_path)[0] + ".json"
-            with open(result_json_path, "w") as f:
-                json.dump(result_dict, f, indent=2, cls=SimpleJSONEncoder)
-        except Exception as json_exc:
-            warnings.warn(
-                f"Failed to serialize JSON result to {result_json_path}: {json_exc}", RuntimeWarning, stacklevel=1
-            )
-
-        # Close the session after serializing the result
+        # Ensure the output directory exists before trying to write result files.
+        result_abs_path = (
+            JOB_RESULT_PATH if os.path.isabs(JOB_RESULT_PATH) else os.path.join(STAGE_MOUNT_PATH, JOB_RESULT_PATH)
+        )
+        output_dir = os.path.dirname(result_abs_path)
+        os.makedirs(output_dir, exist_ok=True)
+
+        # Save the result before closing the session
+        save_mljob_result(execution_result_value, execution_result_is_error, result_abs_path)
         session.close()
 
 
 if __name__ == "__main__":
-    # Parse command line arguments
     parser = argparse.ArgumentParser(description="Launch a Python script and save the result")
     parser.add_argument("script_path", help="Path to the Python script to execute")
     parser.add_argument("script_args", nargs="*", help="Arguments to pass to the script")

snowflake/ml/jobs/_utils/spec_utils.py

@@ -104,7 +104,7 @@ def _get_image_spec(
             image_tag = runtime_environment
         else:
             container_image = runtime_environment
-    elif feature_flags.FeatureFlags.ENABLE_IMAGE_VERSION_ENV_VAR.is_enabled():
+    elif feature_flags.FeatureFlags.ENABLE_RUNTIME_VERSIONS.is_enabled():
         container_image = _get_runtime_image(session, hardware)  # type: ignore[arg-type]
 
     container_image = container_image or f"{image_repo}/{image_name}:{image_tag}"
@@ -266,6 +266,7 @@ def generate_service_spec(
             {"name": "ray-client-server-endpoint", "port": 10001, "protocol": "TCP"},
             {"name": "ray-gcs-endpoint", "port": 12001, "protocol": "TCP"},
             {"name": "ray-dashboard-grpc-endpoint", "port": 12002, "protocol": "TCP"},
+            {"name": "ray-dashboard-endpoint", "port": 12003, "protocol": "TCP"},
             {"name": "ray-object-manager-endpoint", "port": 12011, "protocol": "TCP"},
             {"name": "ray-node-manager-endpoint", "port": 12012, "protocol": "TCP"},
             {"name": "ray-runtime-agent-endpoint", "port": 12013, "protocol": "TCP"},

snowflake/ml/jobs/_utils/types.py

@@ -11,6 +11,7 @@ JOB_STATUS = Literal[
     "CANCELLING",
     "CANCELLED",
     "INTERNAL_ERROR",
+    "DELETED",
 ]
 
 
@@ -106,3 +107,12 @@ class ImageSpec:
     resource_requests: ComputeResources
     resource_limits: ComputeResources
     container_image: str
+
+
+@dataclass(frozen=True)
+class ServiceInfo:
+    database_name: str
+    schema_name: str
+    status: str
+    compute_pool: str
+    target_instances: int