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

snowflake/ml/jobs/_utils/payload_utils.py

@@ -6,17 +6,7 @@ import pickle
 import sys
 import textwrap
 from pathlib import Path, PurePath
-from typing import (
-    Any,
-    Callable,
-    List,
-    Optional,
-    Type,
-    Union,
-    cast,
-    get_args,
-    get_origin,
-)
+from typing import Any, Callable, Optional, Union, cast, get_args, get_origin
 
 import cloudpickle as cp
 
@@ -27,6 +17,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 +64,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 +94,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 +157,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 +191,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 +242,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:
@@ -232,46 +267,17 @@ class JobPayload:
         source: Union[str, Path, Callable[..., Any]],
         entrypoint: Optional[Union[str, Path]] = None,
         *,
-        pip_requirements: Optional[List[str]] = None,
+        pip_requirements: Optional[list[str]] = None,
     ) -> None:
         self.source = Path(source) if isinstance(source, str) else source
         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 +296,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,17 +343,35 @@ 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,
             ],
         )
 
 
-def _get_parameter_type(param: inspect.Parameter) -> Optional[Type[object]]:
+def _get_parameter_type(param: inspect.Parameter) -> Optional[type[object]]:
     # Unwrap Optional type annotations
     param_type = param.annotation
     if get_origin(param_type) is Union and len(get_args(param_type)) == 2 and type(None) in get_args(param_type):
@@ -356,10 +380,10 @@ def _get_parameter_type(param: inspect.Parameter) -> Optional[Type[object]]:
     # Return None for empty type annotations
     if param_type == inspect.Parameter.empty:
         return None
-    return cast(Type[object], param_type)
+    return cast(type[object], param_type)
 
 
-def _validate_parameter_type(param_type: Type[object], param_name: str) -> None:
+def _validate_parameter_type(param_type: type[object], param_name: str) -> None:
     # Validate param_type is a supported type
     if param_type not in _SUPPORTED_ARG_TYPES:
         raise ValueError(
@@ -471,12 +495,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 +510,7 @@ import pickle
 
     return f"""
 ### Version guard to check compatibility across Python versions ###
+import os
 import sys
 import warnings
 
@@ -508,5 +532,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

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

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+# This file is modified from mlruntime/service/snowflake/runtime/utils
+import argparse
+import logging
+import socket
+import sys
+import time
+from typing import Optional
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def get_self_ip() -> Optional[str]:
+    """Get the IP address of the current service instance.
+    References:
+    - https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-services#general-guidelines-related-to-service-to-service-communications # noqa: E501
+
+    Returns:
+        Optional[str]: The IP address of the current service instance, or None if unable to retrieve.
+    """
+    try:
+        hostname = socket.gethostname()
+        instance_ip = socket.gethostbyname(hostname)
+        return instance_ip
+    except OSError as e:
+        logger.error(f"Error: Unable to get IP address via socket. {e}")
+        return None
+
+
+def get_first_instance(service_name: str) -> Optional[tuple[str, str]]:
+    """Get the first instance of a batch job based on start time and instance ID.
+
+    Args:
+        service_name (str): The name of the service to query.
+
+    Returns:
+        tuple[str, str]: A tuple containing (instance_id, ip_address) of the head instance.
+    """
+    from snowflake.runtime.utils import session_utils
+
+    session = session_utils.get_session()
+    df = session.sql(f"show service instances in service {service_name}")
+    result = df.select('"instance_id"', '"ip_address"', '"start_time"').collect()
+
+    if not result:
+        return None
+
+    # Sort by start_time first, then by instance_id
+    sorted_instances = sorted(result, key=lambda x: (x["start_time"], int(x["instance_id"])))
+    head_instance = sorted_instances[0]
+    if not head_instance["instance_id"] or not head_instance["ip_address"]:
+        return None
+
+    # Validate head instance IP
+    ip_address = head_instance["ip_address"]
+    try:
+        socket.inet_aton(ip_address)  # Validate IPv4 address
+        return (head_instance["instance_id"], ip_address)
+    except OSError:
+        logger.error(f"Error: Invalid IP address format: {ip_address}")
+        return None
+
+
+def main():
+    """Retrieves the IP address of a specified service instance or the current service.
+    Args:
+        service_name (str,required) Name of the service to query
+        --instance-index (int, optional) Index of the service instance to query. Default: -1
+            Currently only supports -1 to get the IP address of the current service instance.
+        --head (bool, optional) Get the head instance information using show services.
+            If set, instance-index will be ignored, and the script will return the index and IP address of
+              the head instance, split by a space. Default: False.
+        --timeout (int, optional) Maximum time to wait for IP address retrieval in seconds. Default: 720 seconds
+        --retry-interval (int, optional) Time to wait between retry attempts in seconds. Default: 10 seconds
+    Usage Examples:
+        python get_instance_ip.py myservice --instance-index=1 --retry-interval=5
+    Returns:
+        Prints the IP address to stdout if successful. Exits with status code 0 on success, 1 on failure
+    """
+
+    parser = argparse.ArgumentParser(description="Get IP address of a service instance")
+    group = parser.add_mutually_exclusive_group()
+    parser.add_argument("service_name", help="Name of the service")
+    group.add_argument(
+        "--instance-index",
+        type=int,
+        default=-1,
+        help="Index of service instance (default: -1 for self instance)",
+    )
+    group.add_argument(
+        "--head",
+        action="store_true",
+        help="Get head instance information using show services",
+    )
+    parser.add_argument("--timeout", type=int, default=720, help="Timeout in seconds (default: 720)")
+    parser.add_argument(
+        "--retry-interval",
+        type=int,
+        default=10,
+        help="Retry interval in seconds (default: 10)",
+    )
+
+    args = parser.parse_args()
+    start_time = time.time()
+
+    if args.head:
+        while time.time() - start_time < args.timeout:
+            head_info = get_first_instance(args.service_name)
+            if head_info:
+                # Print to stdout to allow capture but don't use logger
+                sys.stdout.write(f"{head_info[0]} {head_info[1]}\n")
+                sys.exit(0)
+            time.sleep(args.retry_interval)
+        # If we get here, we've timed out
+        logger.error("Error: Unable to retrieve head IP address")
+        sys.exit(1)
+
+    # If the index is -1, use get_self_ip to get the IP address of the current service
+    if args.instance_index == -1:
+        ip_address = get_self_ip()
+        if ip_address:
+            sys.stdout.write(f"{ip_address}\n")
+            sys.exit(0)
+        else:
+            logger.error("Error: Unable to retrieve self IP address")
+            sys.exit(1)
+    else:
+        # We don't support querying a specific instance index other than -1
+        logger.error("Error: Invalid arguments. Only --instance-index=-1 is supported for now.")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

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

@@ -0,0 +1,181 @@
+import argparse
+import copy
+import importlib.util
+import json
+import os
+import runpy
+import sys
+import traceback
+import warnings
+from pathlib import Path
+from typing import Any, Optional
+
+import cloudpickle
+
+from snowflake.ml.jobs._utils import constants
+from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
+from snowflake.snowpark import Session
+
+# Fallbacks in case of SnowML version mismatch
+RESULT_PATH_ENV_VAR = getattr(constants, "RESULT_PATH_ENV_VAR", "MLRS_RESULT_PATH")
+
+JOB_RESULT_PATH = os.environ.get(RESULT_PATH_ENV_VAR, "mljob_result.pkl")
+
+
+try:
+    from snowflake.ml.jobs._utils.interop_utils import ExecutionResult
+except ImportError:
+    from dataclasses import dataclass
+
+    @dataclass(frozen=True)
+    class ExecutionResult:  # type: ignore[no-redef]
+        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):
+                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,
+            }
+
+
+# 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 str(obj)
+
+
+def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = None) -> Any:
+    """
+    Execute a Python script and return its result.
+
+    Args:
+        script_path: Path to the Python script
+        script_args: Arguments to pass to the script
+        main_func: The name of the function to call in the script (if any)
+
+    Returns:
+        Result from script execution, either from the main function or the script's __return__ value
+
+    Raises:
+        RuntimeError: If the specified main_func is not found or not callable
+    """
+    # Save original sys.argv and modify it for the script (applies to runpy execution only)
+    original_argv = sys.argv
+    sys.argv = [script_path, *script_args]
+
+    # Create a Snowpark session before running the script
+    # Session can be retrieved from using snowflake.snowpark.context.get_active_session()
+    session = Session.builder.configs(SnowflakeLoginOptions()).create()  # noqa: F841
+
+    try:
+        if main_func:
+            # Use importlib for scripts with a main function defined
+            module_name = Path(script_path).stem
+            spec = importlib.util.spec_from_file_location(module_name, script_path)
+            assert spec is not None
+            assert spec.loader is not None
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+
+            # Validate main function
+            if not (func := getattr(module, main_func, None)) or not callable(func):
+                raise RuntimeError(f"Function '{main_func}' not a valid entrypoint for {script_path}")
+
+            # Call main function
+            result = func(*script_args)
+            return result
+        else:
+            # Use runpy for other scripts
+            globals_dict = runpy.run_path(script_path, run_name="__main__")
+            result = globals_dict.get("__return__", None)
+            return result
+    finally:
+        # Restore original sys.argv
+        sys.argv = original_argv
+
+
+def main(script_path: str, *script_args: Any, script_main_func: Optional[str] = None) -> ExecutionResult:
+    """Executes a Python script and serializes the result to JOB_RESULT_PATH.
+
+    Args:
+        script_path (str): Path to the Python script to execute.
+        script_args (Any): Arguments to pass to the script.
+        script_main_func (str, optional): The name of the function to call in the script (if any).
+
+    Returns:
+        ExecutionResult: Object containing execution results.
+
+    Raises:
+        Exception: Re-raises any exception caught during script execution.
+    """
+    # Run the script with the specified arguments
+    try:
+        result = run_script(script_path, *script_args, main_func=script_main_func)
+        result_obj = ExecutionResult(result=result)
+        return result_obj
+    except Exception as e:
+        tb = e.__traceback__
+        skip_files = {__file__, runpy.__file__}
+        while tb and tb.tb_frame.f_code.co_filename in skip_files:
+            # Skip any frames preceding user script execution
+            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)
+        raise
+    finally:
+        result_dict = result_obj.to_dict()
+        try:
+            # Serialize result using cloudpickle
+            result_pickle_path = JOB_RESULT_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(JOB_RESULT_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
+            )
+
+
+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")
+    parser.add_argument(
+        "--script_main_func", required=False, help="The name of the main function to call in the script"
+    )
+    args, unknown_args = parser.parse_known_args()
+
+    main(
+        args.script_path,
+        *args.script_args,
+        *unknown_args,
+        script_main_func=args.script_main_func,
+    )