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

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

@@ -1,5 +1,6 @@
 import logging
 import os
+import re
 import sys
 from math import ceil
 from pathlib import PurePath
@@ -10,6 +11,8 @@ from snowflake.ml._internal.utils import snowflake_env
 from snowflake.ml.jobs._utils import constants, feature_flags, query_helper, types
 from snowflake.ml.jobs._utils.runtime_env_utils import RuntimeEnvironmentsDict
 
+_OCI_TAG_REGEX = re.compile("^[a-zA-Z0-9._-]{1,128}$")
+
 
 def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.ComputeResources:
     """Extract resource information for the specified compute pool"""
@@ -56,22 +59,55 @@ def _get_runtime_image(session: snowpark.Session, target_hardware: Literal["CPU"
     return selected_runtime.runtime_container_image if selected_runtime else None
 
 
-def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.ImageSpec:
+def _check_image_tag_valid(tag: Optional[str]) -> bool:
+    if tag is None:
+        return False
+
+    return _OCI_TAG_REGEX.fullmatch(tag) is not None
+
+
+def _get_image_spec(
+    session: snowpark.Session, compute_pool: str, runtime_environment: Optional[str] = None
+) -> types.ImageSpec:
+    """
+    Resolve image specification (container image and resources) for the job.
+
+    Behavior:
+    - If `runtime_environment` is empty or the feature flag is disabled, use the
+      default image tag and image name.
+    - If `runtime_environment` is a valid image tag, use that tag with the default
+      repository/name.
+    - If `runtime_environment` is a full image URL, use it directly.
+    - If the feature flag is enabled and `runtime_environment` is not provided,
+      select an ML Runtime image matching the local Python major.minor
+    - When multiple inputs are provided, `runtime_environment` takes priority.
+
+    Args:
+        session: Snowflake session.
+        compute_pool: Compute pool used to infer CPU/GPU resources.
+        runtime_environment: Optional image tag or full image URL to override.
+
+    Returns:
+        Image spec including container image and resource requests/limits.
+    """
     # Retrieve compute pool node resources
     resources = _get_node_resources(session, compute_pool=compute_pool)
+    hardware = "GPU" if resources.gpu > 0 else "CPU"
+    image_tag = _get_runtime_image_tag()
+    image_repo = constants.DEFAULT_IMAGE_REPO
+    image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
 
     # Use MLRuntime image
-    hardware = "GPU" if resources.gpu > 0 else "CPU"
     container_image = None
-    if feature_flags.FeatureFlags.ENABLE_IMAGE_VERSION_ENV_VAR.is_enabled():
+    if runtime_environment:
+        if _check_image_tag_valid(runtime_environment):
+            image_tag = runtime_environment
+        else:
+            container_image = runtime_environment
+    elif feature_flags.FeatureFlags.ENABLE_RUNTIME_VERSIONS.is_enabled():
         container_image = _get_runtime_image(session, hardware)  # type: ignore[arg-type]
 
-    if not container_image:
-        image_repo = constants.DEFAULT_IMAGE_REPO
-        image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
-        image_tag = _get_runtime_image_tag()
-        container_image = f"{image_repo}/{image_name}:{image_tag}"
-
+    container_image = container_image or f"{image_repo}/{image_name}:{image_tag}"
     # TODO: Should each instance consume the entire pod?
     return types.ImageSpec(
         resource_requests=resources,
@@ -127,6 +163,7 @@ def generate_service_spec(
     target_instances: int = 1,
     min_instances: int = 1,
     enable_metrics: bool = False,
+    runtime_environment: Optional[str] = None,
 ) -> dict[str, Any]:
     """
     Generate a service specification for a job.
@@ -139,11 +176,12 @@ def generate_service_spec(
         target_instances: Number of instances for multi-node job
         enable_metrics: Enable platform metrics for the job
         min_instances: Minimum number of instances required to start the job
+        runtime_environment: The runtime image to use. Only support image tag or full image URL.
 
     Returns:
         Job service specification
     """
-    image_spec = _get_image_spec(session, compute_pool)
+    image_spec = _get_image_spec(session, compute_pool, runtime_environment)
 
     # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
     resource_requests: dict[str, Union[str, int]] = {
@@ -228,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"},
@@ -317,7 +356,7 @@ def merge_patch(base: Any, patch: Any, display_name: str = "") -> Any:
     Returns:
         The patched object.
     """
-    if not type(base) is type(patch):
+    if type(base) is not type(patch):
         if base is not None:
             logging.warning(f"Type mismatch while merging {display_name} (base={type(base)}, patch={type(patch)})")
         return patch

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