snowflake-ml-python 1.20.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

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

@@ -6,9 +6,13 @@ import logging
 import math
 import os
 import runpy
+import shutil
+import subprocess
 import sys
 import time
 import traceback
+import zipfile
+from pathlib import Path
 from typing import Any, Optional
 
 # Ensure payload directory is in sys.path for module imports before importing other modules
@@ -18,11 +22,17 @@ from typing import Any, Optional
 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)
 
+if PAYLOAD_PATH:
+    if PAYLOAD_PATH not in sys.path:
+        sys.path.insert(0, PAYLOAD_PATH)
+    for zip_file in Path(PAYLOAD_PATH).rglob("*.zip"):
+        fpath = str(zip_file)
+        if fpath not in sys.path and zipfile.is_zipfile(fpath):
+            sys.path.insert(0, fpath)
 # 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
@@ -81,6 +91,76 @@ TIMEOUT = float(os.getenv(INSTANCES_TIMEOUT_ENV_VAR) or 720)  # seconds
 CHECK_INTERVAL = float(os.getenv(INSTANCES_CHECK_INTERVAL_ENV_VAR) or 10)  # seconds
 
 
+def is_python_script(file_path: str) -> bool:
+    """Check if a file is a Python script by examining its shebang.
+
+    Args:
+        file_path: Path to the file to check.
+
+    Returns:
+        True if the file has a shebang line containing 'python', False otherwise.
+    """
+    try:
+        with open(file_path, "rb") as f:
+            first_line = f.readline()
+            if first_line.startswith(b"#!"):
+                shebang = first_line.decode("utf-8", errors="ignore").lower()
+                return "python" in shebang
+    except OSError:
+        pass
+    return False
+
+
+def resolve_entrypoint(entrypoint: str) -> tuple[str, bool]:
+    """Resolve the entrypoint to determine how to execute it.
+
+    Args:
+        entrypoint: The entrypoint string (file path or command name).
+
+    Returns:
+        A tuple of (resolved_path, is_python):
+        - resolved_path: The path to the executable/script.
+        - is_python: True if this should be run as a Python script.
+    """
+    # Check if entrypoint is an existing file
+    if os.path.isfile(entrypoint):
+        # Always run as Python script for backward compatibility
+        return entrypoint, True
+
+    # Try to resolve as a command using shutil.which
+    resolved_path = shutil.which(entrypoint)
+    if resolved_path:
+        if is_python_script(resolved_path):
+            return resolved_path, True
+        else:
+            # Assume it's meant to be used as a command and not a Python script
+            return entrypoint, False
+
+    # If we can't resolve it, assume it's meant to be a Python script path
+    # (this preserves backwards compatibility and will fail with a clear error)
+    return entrypoint, True
+
+
+def run_command(command: str, *args: Any) -> None:
+    """Execute a command as a subprocess, streaming output and raising an exception if it fails.
+
+    Args:
+        command: Path to the executable.
+        args: Arguments to pass to the command.
+
+    Raises:
+        CalledProcessError: If the subprocess exits with a non-zero return code.
+    """
+    cmd = [command, *[str(arg) for arg in args]]
+    logger.debug(f"Running subprocess: {' '.join(cmd)}")
+
+    # Run subprocess without capturing output - let stdout/stderr flow directly to console
+    result = subprocess.run(cmd)
+
+    if result.returncode != 0:
+        raise subprocess.CalledProcessError(result.returncode, cmd)
+
+
 def save_mljob_result_v2(value: Any, is_error: bool, path: str) -> None:
     from snowflake.ml.jobs._interop import (
         results as interop_result,
@@ -313,11 +393,11 @@ 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) -> Any:
+def main(entrypoint: str, *script_args: Any, script_main_func: Optional[str] = None) -> Any:
     """Executes a Python script and serializes the result to JOB_RESULT_PATH.
 
     Args:
-        script_path (str): Path to the Python script to execute.
+        entrypoint (str): The job payload entrypoint 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).
 
@@ -361,8 +441,15 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
         # Log start marker before starting user script execution
         print(LOG_START_MSG)  # noqa: T201
 
-        # Run the user script
-        execution_result_value = run_script(script_path, *script_args, main_func=script_main_func)
+        # Resolve entrypoint to determine execution method
+        resolved_entrypoint, is_python = resolve_entrypoint(entrypoint)
+
+        if is_python:
+            # Run as Python script
+            execution_result_value = run_script(resolved_entrypoint, *script_args, main_func=script_main_func)
+        else:
+            # Run as subprocess
+            run_command(resolved_entrypoint, *script_args)
 
         # Log end marker for user script execution
         print(LOG_END_MSG)  # noqa: T201
@@ -395,7 +482,7 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
 
 if __name__ == "__main__":
     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("entrypoint", help="The job payload entrypoint 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"
@@ -403,7 +490,7 @@ if __name__ == "__main__":
     args, unknown_args = parser.parse_known_args()
 
     main(
-        args.script_path,
+        args.entrypoint,
         *args.script_args,
         *unknown_args,
         script_main_func=args.script_main_func,

snowflake/ml/jobs/_utils/scripts/start_mlruntime.sh

@@ -0,0 +1,92 @@
+#!/bin/bash
+
+set -e # exit if a command fails
+
+echo "Creating log directories..."
+mkdir -p /var/log/managedservices/user/mlrs
+mkdir -p /var/log/managedservices/system/mlrs
+mkdir -p /var/log/managedservices/system/ray
+
+echo "*/1 * * * * root /etc/ray_copy_cron.sh" >> /etc/cron.d/ray_copy_cron
+echo "" >> /etc/cron.d/ray_copy_cron
+chmod 744 /etc/cron.d/ray_copy_cron
+
+service cron start
+
+mkdir -p /tmp/prometheus-multi-dir
+
+# Configure IP address and logging directory
+eth0Ip=$(ifconfig eth0 | sed -En -e 's/.*inet ([0-9.]+).*/\1/p')
+log_dir="/tmp/ray"
+
+# Check if eth0Ip is empty and set default if necessary
+if [ -z "$eth0Ip" ]; then
+    # This should never happen, but just in case ethOIp is not set, we should default to localhost
+    eth0Ip="127.0.0.1"
+fi
+
+shm_size=$(df --output=size --block-size=1 /dev/shm | tail -n 1)
+total_memory_size=$(awk '/MemTotal/ {print int($2/1024)}' /proc/meminfo)
+
+# Determine if dashboard should be enabled based on total memory size
+# Enable dashboard only if total memory size >= 8GB (i.e. not on XS compute pool)
+# TODO (SNOW-2860029): use a environment variable to determine the node type
+total_memory_threshold=8192
+if [ "$total_memory_size" -ge "$total_memory_threshold" ]; then
+    enable_dashboard="true"
+else
+    enable_dashboard="false"
+fi
+
+echo "Shared memory size: $shm_size bytes"
+echo "Dashboard enabled: $enable_dashboard"
+
+# Common parameters for both head and worker nodes
+common_params=(
+    "--node-ip-address=$eth0Ip"
+    "--object-manager-port=${RAY_OBJECT_MANAGER_PORT:-12011}"
+    "--node-manager-port=${RAY_NODE_MANAGER_PORT:-12012}"
+    "--runtime-env-agent-port=${RAY_RUNTIME_ENV_AGENT_PORT:-12013}"
+    "--dashboard-agent-grpc-port=${RAY_DASHBOARD_AGENT_GRPC_PORT:-12014}"
+    "--dashboard-agent-listen-port=${RAY_DASHBOARD_AGENT_LISTEN_PORT:-12015}"
+    "--min-worker-port=${RAY_MIN_WORKER_PORT:-12031}"
+    "--max-worker-port=${RAY_MAX_WORKER_PORT:-13000}"
+    "--metrics-export-port=11502"
+    "--temp-dir=$log_dir"
+    "--disable-usage-stats"
+)
+
+# Specific parameters for head and worker nodes
+if [ "$NODE_TYPE" = "worker" ]; then
+    # Check mandatory environment variables for worker
+    if [ -z "$RAY_HEAD_ADDRESS" ] || [ -z "$SERVICE_NAME" ]; then
+        echo "Error: RAY_HEAD_ADDRESS and SERVICE_NAME must be set."
+        exit 1
+    fi
+
+    # Additional worker-specific parameters
+    worker_params=(
+        "--address=${RAY_HEAD_ADDRESS}:${RAY_HEAD_GCS_PORT:-12001}"   # Connect to head node
+        "--resources={\"${SERVICE_NAME}\":1, \"node_tag:worker\":1}"  # Custom resource for node identification
+        "--object-store-memory=${shm_size}"
+    )
+
+    # Start Ray on a worker node
+    ray start "${common_params[@]}" "${worker_params[@]}" "$@" -v
+else
+    # Additional head-specific parameters
+   head_params=(
+        "--head"
+        "--include-dashboard=$enable_dashboard"
+        "--disable-usage-stats"
+        "--port=${RAY_HEAD_GCS_PORT:-12001}"                                  # Port of Ray (GCS server)
+        "--ray-client-server-port=${RAY_HEAD_CLIENT_SERVER_PORT:-10001}"      # Listening port for Ray Client Server
+        "--dashboard-host=${NODE_IP_ADDRESS}"                                            # Host to bind the dashboard server
+        "--dashboard-grpc-port=${RAY_HEAD_DASHBOARD_GRPC_PORT:-12002}"        # Dashboard head to listen for grpc on
+        "--dashboard-port=${DASHBOARD_PORT}"                  # Port to bind the dashboard server for local debugging
+        "--resources={\"node_tag:head\":1}"                   # Resource tag for selecting head as coordinator
+    )
+
+    # Start Ray
+    ray start "${common_params[@]}" "${head_params[@]}" "$@"
+fi

snowflake/ml/jobs/_utils/scripts/startup.sh

@@ -0,0 +1,112 @@
+#!/bin/bash
+
+set -e # exit if a command fails
+
+# Get and change to system scripts directory
+SYSTEM_DIR=$(cd "$(dirname "$0")" && pwd)
+
+# Change directory to user payload directory
+if [ -n "${MLRS_PAYLOAD_DIR}" ]; then
+    cd ${MLRS_STAGE_MOUNT_PATH}/${MLRS_PAYLOAD_DIR}
+fi
+
+##### Set up Python environment #####
+export PYTHONPATH=/opt/env/site-packages/
+MLRS_SYSTEM_REQUIREMENTS_FILE=${MLRS_SYSTEM_REQUIREMENTS_FILE:-"${SYSTEM_DIR}/requirements.txt"}
+if [ -f "${MLRS_SYSTEM_REQUIREMENTS_FILE}" ]; then
+    echo "Installing packages from $MLRS_SYSTEM_REQUIREMENTS_FILE"
+    if ! pip install --no-index -r $MLRS_SYSTEM_REQUIREMENTS_FILE; then
+        echo "Offline install failed, falling back to regular pip install"
+        pip install -r $MLRS_SYSTEM_REQUIREMENTS_FILE
+    fi
+fi
+
+MLRS_REQUIREMENTS_FILE=${MLRS_REQUIREMENTS_FILE:-"requirements.txt"}
+if [ -f "${MLRS_REQUIREMENTS_FILE}" ]; then
+    # TODO: Prevent collisions with MLRS packages using virtualenvs
+    echo "Installing packages from $MLRS_REQUIREMENTS_FILE"
+    pip install -r $MLRS_REQUIREMENTS_FILE
+fi
+
+MLRS_CONDA_ENV_FILE=${MLRS_CONDA_ENV_FILE:-"environment.yml"}
+if [ -f "${MLRS_CONDA_ENV_FILE}" ]; then
+    # TODO: Handle conda environment
+    echo "Custom conda environments not currently supported"
+    exit 1
+fi
+##### End Python environment setup #####
+
+##### Set up multi-node configuration #####
+# Configure IP address
+if [ -f "${SYSTEM_DIR}/get_instance_ip.py" ]; then
+    eth0Ip=$(python3 "${SYSTEM_DIR}/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
+
+# Check if eth0Ip is a valid IP address and fall back to default if necessary
+if [[ ! $eth0Ip =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
+    eth0Ip="127.0.0.1"
+fi
+
+# Set default values for job environment variables if they don't exist
+# (e.g. some are only populate by SPCS for batch jobs, others just may not be set at all)
+export SNOWFLAKE_JOBS_COUNT=${SNOWFLAKE_JOBS_COUNT:-1}
+export SNOWFLAKE_JOB_INDEX=${SNOWFLAKE_JOB_INDEX:-0}
+export SERVICE_NAME="${SERVICE_NAME:-$SNOWFLAKE_SERVICE_NAME}"
+
+##### Ray configuration #####
+
+# Determine if it should be a worker or a head node for batch jobs
+if [[ "$SNOWFLAKE_JOBS_COUNT" -gt 1 ]]; then
+    head_info=$(python3 "${SYSTEM_DIR}/get_instance_ip.py" "$SNOWFLAKE_SERVICE_NAME" --head)
+    if [ $? -eq 0 ]; then
+        # Parse the output using read
+        read head_index head_ip head_status<<< "$head_info"
+
+        if [ "$SNOWFLAKE_JOB_INDEX" -ne "$head_index" ]; then
+            NODE_TYPE="worker"
+        fi
+
+        # Use the parsed variables
+        echo "Head Instance Index: $head_index"
+        echo "Head Instance IP: $head_ip"
+        echo "Head Instance Status: $head_status"
+
+        # If the head status is not "READY" or "PENDING", exit early
+        if [ "$head_status" != "READY" ] && [ "$head_status" != "PENDING" ]; then
+            echo "Head instance status is not READY or PENDING. Exiting."
+            exit 0
+        fi
+
+    else
+        echo "Error: Failed to get head instance information."
+        echo "$head_info" # Print the error message
+        exit 1
+    fi
+fi
+
+# Start ML Runtime (non-blocking call)
+NODE_TYPE=$NODE_TYPE RAY_HEAD_ADDRESS="$head_ip" bash ${SYSTEM_DIR}/start_mlruntime.sh
+
+if [ "$NODE_TYPE" = "worker" ]; then
+    echo "Worker node started on address $eth0Ip. See more logs in the head node."
+
+    # Start the worker shutdown listener in the background
+    echo "Starting worker shutdown listener..."
+    python "${SYSTEM_DIR}/worker_shutdown_listener.py"
+    WORKER_EXIT_CODE=$?
+
+    echo "Worker shutdown listener exited with code $WORKER_EXIT_CODE"
+    exit $WORKER_EXIT_CODE
+else
+    # Run user's Python entrypoint via mljob_launcher
+    echo Running command: python "${SYSTEM_DIR}/mljob_launcher.py" "$@"
+    python "${SYSTEM_DIR}/mljob_launcher.py" "$@"
+
+    # After the user's job completes, signal workers to shut down
+    echo "User job completed. Signaling workers to shut down..."
+    python "${SYSTEM_DIR}/signal_workers.py" --wait-time 15
+    echo "Head node job completed. Exiting."
+fi