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

snowflake/ml/jobs/_utils/interop_utils.py

@@ -0,0 +1,442 @@
+import builtins
+import functools
+import importlib
+import json
+import os
+import pickle
+import re
+import sys
+import traceback
+from collections import namedtuple
+from dataclasses import dataclass
+from types import TracebackType
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union, cast
+
+from snowflake import snowpark
+from snowflake.snowpark import exceptions as sp_exceptions
+
+_TRACEBACK_ENTRY_PATTERN = re.compile(
+    r'File "(?P<filename>[^"]+)", line (?P<lineno>\d+), in (?P<name>[^\n]+)(?:\n(?!^\s*File)^\s*(?P<line>[^\n]+))?\n',
+    flags=re.MULTILINE,
+)
+_REMOTE_ERROR_ATTR_NAME = "_remote_error"
+
+RemoteError = namedtuple("RemoteError", ["exc_type", "exc_msg", "exc_tb"])
+
+
+@dataclass(frozen=True)
+class ExecutionResult:
+    result: 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):
+            exc_type = type(self.exception)
+            return {
+                "success": False,
+                "exc_type": f"{exc_type.__module__}.{exc_type.__name__}",
+                "exc_value": self.exception,
+                "exc_tb": "".join(traceback.format_tb(self.exception.__traceback__)),
+            }
+        return {
+            "success": True,
+            "result_type": type(self.result).__qualname__,
+            "result": self.result,
+        }
+
+    @classmethod
+    def from_dict(cls, result_dict: Dict[str, Any]) -> "ExecutionResult":
+        if not isinstance(result_dict.get("success"), bool):
+            raise ValueError("Invalid result dictionary")
+
+        if result_dict["success"]:
+            # Load successful result
+            return cls(result=result_dict.get("result"))
+
+        # Load exception
+        exc_type = result_dict.get("exc_type", "RuntimeError")
+        exc_value = result_dict.get("exc_value", "Unknown error")
+        exc_tb = result_dict.get("exc_tb", "")
+        return cls(exception=load_exception(exc_type, exc_value, exc_tb))
+
+
+def fetch_result(session: snowpark.Session, result_path: str) -> ExecutionResult:
+    """
+    Fetch the serialized result from the specified path.
+
+    Args:
+        session: Snowpark Session to use for file operations.
+        result_path: The path to the serialized result file.
+
+    Returns:
+        A dictionary containing the execution result if available, None otherwise.
+    """
+    try:
+        # TODO: Check if file exists
+        with session.file.get_stream(result_path) as result_stream:
+            return ExecutionResult.from_dict(pickle.load(result_stream))
+    except (sp_exceptions.SnowparkSQLException, TypeError, pickle.UnpicklingError):
+        # Fall back to JSON result if loading pickled result fails for any reason
+        result_json_path = os.path.splitext(result_path)[0] + ".json"
+        with session.file.get_stream(result_json_path) as result_stream:
+            return ExecutionResult.from_dict(json.load(result_stream))
+
+
+def load_exception(exc_type_name: str, exc_value: Union[Exception, str], exc_tb: str) -> Exception:
+    """
+    Create an exception with a string-formatted traceback.
+
+    When this exception is raised and not caught, it will display the original traceback.
+    When caught, it behaves like a regular exception without showing the traceback.
+
+    Args:
+        exc_type_name: Name of the exception type (e.g., 'ValueError', 'RuntimeError')
+        exc_value: The deserialized exception value or exception string (i.e. message)
+        exc_tb: String representation of the traceback
+
+    Returns:
+        An exception object with the original traceback information
+
+    # noqa: DAR401
+    """
+    if isinstance(exc_value, Exception):
+        exception = exc_value
+    else:
+        # Try to load the original exception type if possible
+        try:
+            # First check built-in exceptions
+            exc_type = getattr(builtins, exc_type_name, None)
+            if exc_type is None and "." in exc_type_name:
+                # Try to import from module path if it's a qualified name
+                module_path, class_name = exc_type_name.rsplit(".", 1)
+                module = importlib.import_module(module_path)
+                exc_type = getattr(module, class_name)
+            if exc_type is None or not issubclass(exc_type, Exception):
+                raise TypeError(f"{exc_type_name} is not a known exception type")
+            # Create the exception instance
+            exception = exc_type(exc_value)
+        except (ImportError, AttributeError, TypeError):
+            # Fall back to a generic exception
+            exception = RuntimeError(
+                f"Exception deserialization failed, original exception: {exc_type_name}: {exc_value}"
+            )
+
+    # Attach the traceback information to the exception
+    return _attach_remote_error_info(exception, exc_type_name, str(exc_value), exc_tb)
+
+
+def _attach_remote_error_info(ex: Exception, exc_type: str, exc_msg: str, traceback_str: str) -> Exception:
+    """
+    Attach a string-formatted traceback to an exception.
+
+    When the exception is raised and not caught, it will display the original traceback.
+    When caught, it behaves like a regular exception without showing the traceback.
+
+    Args:
+        ex: The exception object to modify
+        exc_type: The original exception type name
+        exc_msg: The original exception message
+        traceback_str: String representation of the traceback
+
+    Returns:
+        An exception object with the original traceback information
+    """
+    # Store the traceback information
+    exc_type = exc_type.rsplit(".", 1)[-1]  # Remove module path
+    setattr(ex, _REMOTE_ERROR_ATTR_NAME, RemoteError(exc_type=exc_type, exc_msg=exc_msg, exc_tb=traceback_str))
+    return ex
+
+
+def _retrieve_remote_error_info(ex: Optional[BaseException]) -> Optional[RemoteError]:
+    """
+    Retrieve the string-formatted traceback from an exception if it exists.
+
+    Args:
+        ex: The exception to retrieve the traceback from
+
+    Returns:
+        The remote error tuple if it exists, None otherwise
+    """
+    if not ex:
+        return None
+    return getattr(ex, _REMOTE_ERROR_ATTR_NAME, None)
+
+
+# ###############################################################################
+# ------------------------------- !!! NOTE !!! -------------------------------- #
+# ###############################################################################
+# Job execution results (including uncaught exceptions) are serialized to file(s)
+# in mljob_launcher.py. When the job is executed remotely, the serialized results
+# are fetched and deserialized in the local environment. If the result contains
+# an exception the original traceback is reconstructed and displayed to the user.
+#
+# It's currently impossible to recreate the original traceback object, so the
+# following overrides are necessary to attach and display the deserialized
+# traceback during exception handling.
+#
+# The following code implements the necessary overrides including sys.excepthook
+# modifications and IPython traceback formatting. The hooks are applied on init
+# and will be active for the duration of the process. The hooks are designed to
+# self-uninstall in the event of an error in case of future compatibility issues.
+# ###############################################################################
+
+
+def _revert_func_wrapper(
+    patched_func: Callable[..., Any],
+    original_func: Callable[..., Any],
+    uninstall_func: Callable[[], None],
+) -> Callable[..., Any]:
+    """
+    Create a wrapper function that uninstalls the original function if an error occurs during execution.
+
+    This wrapper provides a fallback mechanism where if the patched function fails, it will:
+    1. Uninstall the patched function using the provided uninstall_func, reverting back to using the original function
+    2. Re-execute the current call using the original (unpatched) function with the same arguments
+
+    Args:
+        patched_func: The patched function to call.
+        original_func: The original function to call if patched_func fails.
+        uninstall_func: The function to call to uninstall the patched function.
+
+    Returns:
+        A wrapped function that calls patched_func and uninstalls on failure.
+    """
+
+    @functools.wraps(patched_func)
+    def wrapped(*args: Any, **kwargs: Any) -> Any:
+        try:
+            return patched_func(*args, **kwargs)
+        except Exception:
+            # Uninstall and revert to original on failure
+            uninstall_func()
+            return original_func(*args, **kwargs)
+
+    return wrapped
+
+
+def _install_sys_excepthook() -> None:
+    """
+    Install a custom sys.excepthook to handle remote exception tracebacks.
+
+    sys.excepthook is the global hook that Python calls when an unhandled exception occurs.
+    By default it prints the exception type, message and traceback to stderr.
+
+    We override sys.excepthook to intercept exceptions that contain our special RemoteError
+    attribute. These exceptions come from deserialized remote execution results and contain
+    the original traceback information from where they occurred.
+
+    When such an exception is detected, we format and display the original remote traceback
+    instead of the local one, which provides better debugging context by showing where the
+    error actually happened during remote execution.
+
+    The custom hook maintains proper exception chaining for both __cause__ (from raise from)
+    and __context__ (from implicit exception chaining).
+    """
+    # Attach the custom excepthook for standard Python scripts if not already attached
+    if not hasattr(sys, "_original_excepthook"):
+        original_excepthook = sys.excepthook
+
+        def custom_excepthook(
+            exc_type: Type[BaseException],
+            exc_value: BaseException,
+            exc_tb: Optional[TracebackType],
+            *,
+            seen_exc_ids: Optional[Set[int]] = None,
+        ) -> None:
+            if seen_exc_ids is None:
+                seen_exc_ids = set()
+            seen_exc_ids.add(id(exc_value))
+
+            cause = getattr(exc_value, "__cause__", None)
+            context = getattr(exc_value, "__context__", None)
+            if cause:
+                # Handle cause-chained exceptions
+                custom_excepthook(type(cause), cause, cause.__traceback__, seen_exc_ids=seen_exc_ids)
+                print(  # noqa: T201
+                    "\nThe above exception was the direct cause of the following exception:\n", file=sys.stderr
+                )
+            elif context and not getattr(exc_value, "__suppress_context__", False):
+                # Handle context-chained exceptions
+                # Only process context if it's different from cause to avoid double printing
+                custom_excepthook(type(context), context, context.__traceback__, seen_exc_ids=seen_exc_ids)
+                print(  # noqa: T201
+                    "\nDuring handling of the above exception, another exception occurred:\n", file=sys.stderr
+                )
+
+            if (remote_err := _retrieve_remote_error_info(exc_value)) and isinstance(remote_err, RemoteError):
+                # Display stored traceback for deserialized exceptions
+                print("Traceback (from remote execution):", file=sys.stderr)  # noqa: T201
+                print(remote_err.exc_tb, end="", file=sys.stderr)  # noqa: T201
+                print(f"{remote_err.exc_type}: {remote_err.exc_msg}", file=sys.stderr)  # noqa: T201
+            else:
+                # Fall back to the original excepthook
+                traceback.print_exception(exc_type, exc_value, exc_tb, file=sys.stderr, chain=False)
+
+        sys._original_excepthook = original_excepthook  # type: ignore[attr-defined]
+        sys.excepthook = _revert_func_wrapper(custom_excepthook, original_excepthook, _uninstall_sys_excepthook)
+
+
+def _uninstall_sys_excepthook() -> None:
+    """
+    Restore the original excepthook for the current process.
+
+    This is useful when we want to revert to the default behavior after installing a custom excepthook.
+    """
+    if hasattr(sys, "_original_excepthook"):
+        sys.excepthook = sys._original_excepthook
+        del sys._original_excepthook
+
+
+def _install_ipython_hook() -> bool:
+    """Install IPython-specific exception handling hook to improve remote error reporting.
+
+    This function enhances IPython's error formatting capabilities by intercepting and customizing
+    how remote execution errors are displayed. It modifies two key IPython traceback formatters:
+
+    1. VerboseTB.format_exception_as_a_whole: Customizes the full traceback formatting for remote
+        errors by:
+        - Adding a "(from remote execution)" header instead of "(most recent call last)"
+        - Properly formatting the remote traceback entries
+        - Maintaining original behavior for non-remote errors
+
+    2. ListTB.structured_traceback: Modifies the structured traceback output by:
+        - Parsing and formatting remote tracebacks appropriately
+        - Adding remote execution context to the output
+        - Preserving original functionality for local errors
+
+    The modifications are needed because IPython's default error handling doesn't properly display
+    remote execution errors that occur in Snowpark/Snowflake operations. The custom formatters
+    ensure that error messages from remote executions are properly captured, formatted and displayed
+    with the correct context and traceback information.
+
+    Returns:
+         bool: True if IPython hooks were successfully installed, False if IPython is not available
+         or not in an IPython environment.
+
+    Note:
+         This function maintains the ability to revert changes through _uninstall_ipython_hook by
+         storing original implementations before applying modifications.
+    """
+    try:
+        from IPython.core.getipython import get_ipython
+        from IPython.core.ultratb import ListTB, VerboseTB
+
+        if get_ipython() is None:
+            return False
+    except ImportError:
+        return False
+
+    def parse_traceback_str(traceback_str: str) -> List[Tuple[str, int, str, str]]:
+        return [
+            (m.group("filename"), int(m.group("lineno")), m.group("name"), m.group("line"))
+            for m in re.finditer(_TRACEBACK_ENTRY_PATTERN, traceback_str)
+        ]
+
+    if not hasattr(VerboseTB, "_original_format_exception_as_a_whole"):
+        original_format_exception_as_a_whole = VerboseTB.format_exception_as_a_whole
+
+        def custom_format_exception_as_a_whole(
+            self: VerboseTB,
+            etype: Type[BaseException],
+            evalue: Optional[BaseException],
+            etb: Optional[TracebackType],
+            number_of_lines_of_context: int,
+            tb_offset: Optional[int],
+            **kwargs: Any,
+        ) -> List[List[str]]:
+            if (remote_err := _retrieve_remote_error_info(evalue)) and isinstance(remote_err, RemoteError):
+                # Implementation forked from IPython.core.ultratb.VerboseTB.format_exception_as_a_whole
+                head = self.prepare_header(remote_err.exc_type, long_version=False).replace(
+                    "(most recent call last)",
+                    "(from remote execution)",
+                )
+
+                frames = ListTB._format_list(
+                    self,
+                    parse_traceback_str(remote_err.exc_tb),
+                )
+                formatted_exception = self.format_exception(remote_err.exc_type, remote_err.exc_msg)
+
+                return [[head] + frames + formatted_exception]
+            return original_format_exception_as_a_whole(  # type: ignore[no-any-return]
+                self,
+                etype=etype,
+                evalue=evalue,
+                etb=etb,
+                number_of_lines_of_context=number_of_lines_of_context,
+                tb_offset=tb_offset,
+                **kwargs,
+            )
+
+        VerboseTB._original_format_exception_as_a_whole = original_format_exception_as_a_whole
+        VerboseTB.format_exception_as_a_whole = _revert_func_wrapper(
+            custom_format_exception_as_a_whole, original_format_exception_as_a_whole, _uninstall_ipython_hook
+        )
+
+    if not hasattr(ListTB, "_original_structured_traceback"):
+        original_structured_traceback = ListTB.structured_traceback
+
+        def structured_traceback(
+            self: ListTB,
+            etype: type,
+            evalue: Optional[BaseException],
+            etb: Optional[TracebackType],
+            tb_offset: Optional[int] = None,
+            **kwargs: Any,
+        ) -> List[str]:
+            if (remote_err := _retrieve_remote_error_info(evalue)) and isinstance(remote_err, RemoteError):
+                tb_list = [
+                    (m.group("filename"), m.group("lineno"), m.group("name"), m.group("line"))
+                    for m in re.finditer(_TRACEBACK_ENTRY_PATTERN, remote_err.exc_tb or "")
+                ]
+                out_list = original_structured_traceback(self, etype, evalue, tb_list, tb_offset, **kwargs)
+                if out_list:
+                    out_list[0] = out_list[0].replace(
+                        "(most recent call last)",
+                        "(from remote execution)",
+                    )
+                return cast(List[str], out_list)
+            return original_structured_traceback(  # type: ignore[no-any-return]
+                self, etype, evalue, etb, tb_offset, **kwargs
+            )
+
+        ListTB._original_structured_traceback = original_structured_traceback
+        ListTB.structured_traceback = _revert_func_wrapper(
+            structured_traceback, original_structured_traceback, _uninstall_ipython_hook
+        )
+
+    return True
+
+
+def _uninstall_ipython_hook() -> None:
+    """
+    Restore the original IPython traceback formatting if it was modified.
+
+    This is useful when we want to revert to the default behavior after installing a custom hook.
+    """
+    try:
+        from IPython.core.ultratb import ListTB, VerboseTB
+
+        if hasattr(VerboseTB, "_original_format_exception_as_a_whole"):
+            VerboseTB.format_exception_as_a_whole = VerboseTB._original_format_exception_as_a_whole
+            del VerboseTB._original_format_exception_as_a_whole
+
+        if hasattr(ListTB, "_original_structured_traceback"):
+            ListTB.structured_traceback = ListTB._original_structured_traceback
+            del ListTB._original_structured_traceback
+    except ImportError:
+        pass
+
+
+def install_exception_display_hooks() -> None:
+    if not _install_ipython_hook():
+        _install_sys_excepthook()
+
+
+# ------ Install the custom traceback hooks by default ------ #
+install_exception_display_hooks()

snowflake/ml/jobs/_utils/payload_utils.py

@@ -27,6 +27,7 @@ from snowflake.snowpark._internal import code_generation
 
 _SUPPORTED_ARG_TYPES = {str, int, float}
 _SUPPORTED_ENTRYPOINT_EXTENSIONS = {".py"}
+_ENTRYPOINT_FUNC_NAME = "func"
 _STARTUP_SCRIPT_PATH = PurePath("startup.sh")
 _STARTUP_SCRIPT_CODE = textwrap.dedent(
     f"""
@@ -73,14 +74,14 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
     ##### Ray configuration #####
     shm_size=$(df --output=size --block-size=1 /dev/shm | tail -n 1)
 
-    # Check if the instance ip retrieval module exists, which is a prerequisite for multi node jobs
+    # Check if the local get_instance_ip.py script exists
     HELPER_EXISTS=$(
-        python3 -c "import snowflake.runtime.utils.get_instance_ip" 2>/dev/null && echo "true" || echo "false"
+        [ -f "get_instance_ip.py" ] && echo "true" || echo "false"
     )
 
     # Configure IP address and logging directory
     if [ "$HELPER_EXISTS" = "true" ]; then
-        eth0Ip=$(python3 -m snowflake.runtime.utils.get_instance_ip "$SNOWFLAKE_SERVICE_NAME" --instance-index=-1)
+        eth0Ip=$(python3 get_instance_ip.py "$SNOWFLAKE_SERVICE_NAME" --instance-index=-1)
     else
         eth0Ip=$(ifconfig eth0 2>/dev/null | sed -En -e 's/.*inet ([0-9.]+).*/\1/p')
     fi
@@ -103,7 +104,7 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
 
     # Determine if it should be a worker or a head node for batch jobs
     if [[ "$SNOWFLAKE_JOBS_COUNT" -gt 1 && "$HELPER_EXISTS" = "true" ]]; then
-        head_info=$(python3 -m snowflake.runtime.utils.get_instance_ip "$SNOWFLAKE_SERVICE_NAME" --head)
+        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"
@@ -166,10 +167,17 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
             "--object-store-memory=${{shm_size}}"
         )
 
-        # Start Ray on a worker node
-        ray start "${{common_params[@]}}" "${{worker_params[@]}}" -v --block
-    else
+        # Start Ray on a worker node - run in background
+        ray start "${{common_params[@]}}" "${{worker_params[@]}}" -v --block &
+
+        # Start the worker shutdown listener in the background
+        echo "Starting worker shutdown listener..."
+        python worker_shutdown_listener.py
+        WORKER_EXIT_CODE=$?
 
+        echo "Worker shutdown listener exited with code $WORKER_EXIT_CODE"
+        exit $WORKER_EXIT_CODE
+    else
         # Additional head-specific parameters
         head_params=(
             "--head"
@@ -193,13 +201,39 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
         # Run user's Python entrypoint
         echo Running command: python "$@"
         python "$@"
+
+        # After the user's job completes, signal workers to shut down
+        echo "User job completed. Signaling workers to shut down..."
+        python signal_workers.py --wait-time 15
+        echo "Head node job completed. Exiting."
     fi
     """
 ).strip()
 
 
-def _resolve_entrypoint(parent: Path, entrypoint: Optional[Path]) -> Path:
-    parent = parent.absolute()
+def resolve_source(source: Union[Path, Callable[..., Any]]) -> Union[Path, Callable[..., Any]]:
+    if callable(source):
+        return source
+    elif isinstance(source, Path):
+        # Validate source
+        source = source
+        if not source.exists():
+            raise FileNotFoundError(f"{source} does not exist")
+        return source.absolute()
+    else:
+        raise ValueError("Unsupported source type. Source must be a file, directory, or callable.")
+
+
+def resolve_entrypoint(source: Union[Path, Callable[..., Any]], entrypoint: Optional[Path]) -> types.PayloadEntrypoint:
+    if callable(source):
+        # Entrypoint is generated for callable payloads
+        return types.PayloadEntrypoint(
+            file_path=entrypoint or Path(constants.DEFAULT_ENTRYPOINT_PATH),
+            main_func=_ENTRYPOINT_FUNC_NAME,
+        )
+
+    # Resolve entrypoint path for file-based payloads
+    parent = source.absolute()
     if entrypoint is None:
         if parent.is_file():
             # Infer entrypoint from source
@@ -218,12 +252,23 @@ def _resolve_entrypoint(parent: Path, entrypoint: Optional[Path]) -> Path:
         else:
             # Relative to source dir
             entrypoint = parent.joinpath(entrypoint)
+
+    # Validate resolved entrypoint file
     if not entrypoint.is_file():
         raise FileNotFoundError(
             "Entrypoint not found. Ensure the entrypoint is a valid file and is under"
             f" the source directory (source={parent}, entrypoint={entrypoint})"
         )
-    return entrypoint
+    if entrypoint.suffix not in _SUPPORTED_ENTRYPOINT_EXTENSIONS:
+        raise ValueError(
+            "Unsupported entrypoint type:"
+            f" supported={','.join(_SUPPORTED_ENTRYPOINT_EXTENSIONS)} got={entrypoint.suffix}"
+        )
+
+    return types.PayloadEntrypoint(
+        file_path=entrypoint,  # entrypoint is an absolute path at this point
+        main_func=None,
+    )
 
 
 class JobPayload:
@@ -238,40 +283,11 @@ class JobPayload:
         self.entrypoint = Path(entrypoint) if isinstance(entrypoint, str) else entrypoint
         self.pip_requirements = pip_requirements
 
-    def validate(self) -> None:
-        if callable(self.source):
-            # Any entrypoint value is OK for callable payloads (including None aka default)
-            # since we will generate the file from the serialized callable
-            pass
-        elif isinstance(self.source, Path):
-            # Validate source
-            source = self.source
-            if not source.exists():
-                raise FileNotFoundError(f"{source} does not exist")
-            source = source.absolute()
-
-            # Validate entrypoint
-            entrypoint = _resolve_entrypoint(source, self.entrypoint)
-            if entrypoint.suffix not in _SUPPORTED_ENTRYPOINT_EXTENSIONS:
-                raise ValueError(
-                    "Unsupported entrypoint type:"
-                    f" supported={','.join(_SUPPORTED_ENTRYPOINT_EXTENSIONS)} got={entrypoint.suffix}"
-                )
-
-            # Update fields with normalized values
-            self.source = source
-            self.entrypoint = entrypoint
-        else:
-            raise ValueError("Unsupported source type. Source must be a file, directory, or callable.")
-
     def upload(self, session: snowpark.Session, stage_path: Union[str, PurePath]) -> types.UploadedPayload:
-        # Validate payload
-        self.validate()
-
         # Prepare local variables
         stage_path = PurePath(stage_path) if isinstance(stage_path, str) else stage_path
-        source = self.source
-        entrypoint = self.entrypoint or Path(constants.DEFAULT_ENTRYPOINT_PATH)
+        source = resolve_source(self.source)
+        entrypoint = resolve_entrypoint(source, self.entrypoint)
 
         # Create stage if necessary
         stage_name = stage_path.parts[0].lstrip("@")
@@ -290,11 +306,11 @@ class JobPayload:
             source_code = generate_python_code(source, source_code_display=True)
             _ = session.file.put_stream(
                 io.BytesIO(source_code.encode()),
-                stage_location=stage_path.joinpath(entrypoint).as_posix(),
+                stage_location=stage_path.joinpath(entrypoint.file_path).as_posix(),
                 auto_compress=False,
                 overwrite=True,
             )
-            source = entrypoint.parent
+            source = Path(entrypoint.file_path.parent)
         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
@@ -337,12 +353,30 @@ class JobPayload:
             overwrite=False,  # FIXME
         )
 
+        # Upload system scripts
+        scripts_dir = Path(__file__).parent.joinpath("scripts")
+        for script_file in scripts_dir.glob("*"):
+            if script_file.is_file():
+                session.file.put(
+                    script_file.as_posix(),
+                    stage_path.as_posix(),
+                    overwrite=True,
+                    auto_compress=False,
+                )
+
+        python_entrypoint: List[Union[str, PurePath]] = [
+            PurePath("mljob_launcher.py"),
+            entrypoint.file_path.relative_to(source),
+        ]
+        if entrypoint.main_func:
+            python_entrypoint += ["--script_main_func", entrypoint.main_func]
+
         return types.UploadedPayload(
             stage_path=stage_path,
             entrypoint=[
                 "bash",
                 _STARTUP_SCRIPT_PATH,
-                entrypoint.relative_to(source),
+                *python_entrypoint,
             ],
         )
 
@@ -471,12 +505,11 @@ 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_name = "func"
     func_code = f"""
 {source_code_comment}
 
 import pickle
-{func_name} = pickle.loads(bytes.fromhex('{_serialize_callable(func).hex()}'))
+{_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(func).hex()}'))
 """
 
     arg_dict_name = "kwargs"
@@ -487,6 +520,7 @@ import pickle
 
     return f"""
 ### Version guard to check compatibility across Python versions ###
+import os
 import sys
 import warnings
 
@@ -508,5 +542,5 @@ if sys.version_info.major != {sys.version_info.major} or sys.version_info.minor
 if __name__ == '__main__':
 {textwrap.indent(param_code, '    ')}
 
-    {func_name}(**{arg_dict_name})
+    __return__ = {_ENTRYPOINT_FUNC_NAME}(**{arg_dict_name})
 """

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

@@ -0,0 +1,4 @@
+# Constants defining the shutdown signal actor configuration.
+SHUTDOWN_ACTOR_NAME = "ShutdownSignal"
+SHUTDOWN_ACTOR_NAMESPACE = "default"
+SHUTDOWN_RPC_TIMEOUT_SECONDS = 5.0