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

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,178 @@
+import argparse
+import importlib.util
+import json
+import os
+import runpy
+import sys
+import traceback
+import warnings
+from pathlib import Path
+from typing import Any, Dict, 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:
+        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
+        result_obj = ExecutionResult(exception=e.with_traceback(tb))
+        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,
+    )

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

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+# This file is part of the Ray-based distributed job system for Snowflake ML.
+# Architecture overview:
+# - Head node creates a ShutdownSignal actor and signals workers when job completes
+# - Worker nodes listen for this signal and gracefully shut down
+# - This ensures clean termination of distributed Ray jobs
+import argparse
+import logging
+import socket
+import sys
+import time
+from typing import Any, Dict, List, Set
+
+import ray
+from constants import (
+    SHUTDOWN_ACTOR_NAME,
+    SHUTDOWN_ACTOR_NAMESPACE,
+    SHUTDOWN_RPC_TIMEOUT_SECONDS,
+)
+from ray.actor import ActorHandle
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+
+@ray.remote
+class ShutdownSignal:
+    """A simple Ray actor that workers can check to determine if they should shutdown"""
+
+    def __init__(self) -> None:
+        self.shutdown_requested = False
+        self.timestamp = None
+        self.hostname = socket.gethostname()
+        self.acknowledged_workers = set()
+        logging.info(f"ShutdownSignal actor created on {self.hostname}")
+
+    def request_shutdown(self) -> Dict[str, Any]:
+        """Signal workers to shut down"""
+        self.shutdown_requested = True
+        self.timestamp = time.time()
+        logging.info(f"Shutdown requested by head node at {self.timestamp}")
+        return {"status": "shutdown_requested", "timestamp": self.timestamp, "host": self.hostname}
+
+    def should_shutdown(self) -> Dict[str, Any]:
+        """Check if shutdown has been requested"""
+        return {"shutdown": self.shutdown_requested, "timestamp": self.timestamp, "host": self.hostname}
+
+    def ping(self) -> Dict[str, Any]:
+        """Simple method to test connectivity"""
+        return {"status": "alive", "host": self.hostname}
+
+    def acknowledge_shutdown(self, worker_id: str) -> Dict[str, Any]:
+        """Worker acknowledges it has received the shutdown signal and is terminating"""
+        self.acknowledged_workers.add(worker_id)
+        logging.info(f"Worker {worker_id} acknowledged shutdown. Total acknowledged: {len(self.acknowledged_workers)}")
+
+        return {"status": "acknowledged", "worker_id": worker_id, "acknowledged_count": len(self.acknowledged_workers)}
+
+    def get_acknowledgment_workers(self) -> Set[str]:
+        """Get the set of workers who have acknowledged shutdown"""
+        return self.acknowledged_workers
+
+
+def get_worker_node_ids() -> List[str]:
+    """Get the IDs of all active worker nodes.
+
+    Returns:
+        List[str]: List of worker node IDs. Empty list if no workers are present.
+    """
+    worker_nodes = [
+        node for node in ray.nodes() if node.get("Alive") and node.get("Resources", {}).get("node_tag:worker", 0) > 0
+    ]
+
+    worker_node_ids = [node.get("NodeName") for node in worker_nodes]
+
+    if worker_node_ids:
+        logging.info(f"Found {len(worker_node_ids)} worker nodes")
+    else:
+        logging.info("No active worker nodes found")
+
+    return worker_node_ids
+
+
+def get_or_create_shutdown_signal() -> ActorHandle:
+    """Get existing shutdown signal actor or create a new one.
+
+    Returns:
+        ActorHandle: Reference to shutdown signal actor
+    """
+    try:
+        # Try to get existing actor
+        shutdown_signal = ray.get_actor(SHUTDOWN_ACTOR_NAME, namespace=SHUTDOWN_ACTOR_NAMESPACE)
+        logging.info("Found existing shutdown signal actor")
+    except (ValueError, ray.exceptions.RayActorError) as e:
+        logging.info(f"Creating new shutdown signal actor: {e}")
+        # Create new actor if it doesn't exist
+        shutdown_signal = ShutdownSignal.options(
+            name=SHUTDOWN_ACTOR_NAME,
+            namespace=SHUTDOWN_ACTOR_NAMESPACE,
+            lifetime="detached",  # Ensure actor survives client disconnect
+            resources={"node_tag:head": 0.001},  # Resource constraint to ensure it runs on head node
+        ).remote()
+
+        # Verify actor is created and accessible
+        ping_result = ray.get(shutdown_signal.ping.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+        logging.debug(f"New actor ping response: {ping_result}")
+
+    return shutdown_signal
+
+
+def request_shutdown(shutdown_signal: ActorHandle) -> None:
+    """Request workers to shut down.
+
+    Args:
+        shutdown_signal: Reference to the shutdown signal actor
+    """
+    response = ray.get(shutdown_signal.request_shutdown.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+    logging.info(f"Shutdown requested: {response}")
+
+
+def verify_shutdown(shutdown_signal: ActorHandle) -> None:
+    """Verify that shutdown was properly signaled.
+
+    Args:
+        shutdown_signal: Reference to the shutdown signal actor
+    """
+    check = ray.get(shutdown_signal.should_shutdown.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS)
+    logging.debug(f"Shutdown status check: {check}")
+
+
+def wait_for_acknowledgments(shutdown_signal: ActorHandle, worker_node_ids: List[str], wait_time: int) -> None:
+    """Wait for workers to acknowledge shutdown.
+
+    Args:
+        shutdown_signal: Reference to the shutdown signal actor
+        worker_node_ids: List of worker node IDs
+        wait_time: Time in seconds to wait for acknowledgments
+
+    Raises:
+        TimeoutError: When workers don't acknowledge within the wait time or if actor communication times out
+    """
+    if not worker_node_ids:
+        return
+
+    logging.info(f"Waiting up to {wait_time}s for workers to acknowledge shutdown signal...")
+    start_time = time.time()
+    check_interval = 1.0
+
+    while time.time() - start_time < wait_time:
+        try:
+            ack_workers = ray.get(
+                shutdown_signal.get_acknowledgment_workers.remote(), timeout=SHUTDOWN_RPC_TIMEOUT_SECONDS
+            )
+            if ack_workers and ack_workers == set(worker_node_ids):
+                logging.info(
+                    f"All {len(worker_node_ids)} workers acknowledged shutdown. "
+                    f"Completed in {time.time() - start_time:.2f}s"
+                )
+                return
+            else:
+                logging.debug(f"Waiting for acknowledgments: {len(ack_workers)}/{len(worker_node_ids)} workers")
+        except Exception as e:
+            logging.warning(f"Error checking acknowledgment status: {e}")
+
+        time.sleep(check_interval)
+
+    raise TimeoutError(
+        f"Timed out waiting for {len(worker_node_ids)} workers to acknowledge shutdown after {wait_time}s"
+    )
+
+
+def signal_workers(wait_time: int = 10) -> int:
+    """
+    Signal worker nodes to shut down by creating a shutdown signal actor.
+
+    Args:
+        wait_time: Time in seconds to wait for workers to receive the message
+
+    Returns:
+        0 for success, 1 for failure
+    """
+    ray.init(address="auto", ignore_reinit_error=True)
+
+    worker_node_ids = get_worker_node_ids()
+
+    if worker_node_ids:
+        shutdown_signal = get_or_create_shutdown_signal()
+        request_shutdown(shutdown_signal)
+        verify_shutdown(shutdown_signal)
+        wait_for_acknowledgments(shutdown_signal, worker_node_ids, wait_time)
+    else:
+        logging.info("No active worker nodes found to signal.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Signal Ray workers to shutdown")
+    parser.add_argument(
+        "--wait-time", type=int, default=10, help="Time in seconds to wait for workers to receive the signal"
+    )
+    args = parser.parse_args()
+
+    sys.exit(signal_workers(args.wait_time))