snowflake-ml-python 1.16.0__py3-none-any.whl → 1.18.0__py3-none-any.whl

snowflake/ml/jobs/_utils/constants.py

@@ -12,6 +12,9 @@ PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
 RESULT_PATH_ENV_VAR = "MLRS_RESULT_PATH"
 MIN_INSTANCES_ENV_VAR = "MLRS_MIN_INSTANCES"
 TARGET_INSTANCES_ENV_VAR = "SNOWFLAKE_JOBS_COUNT"
+INSTANCES_MIN_WAIT_ENV_VAR = "MLRS_INSTANCES_MIN_WAIT"
+INSTANCES_TIMEOUT_ENV_VAR = "MLRS_INSTANCES_TIMEOUT"
+INSTANCES_CHECK_INTERVAL_ENV_VAR = "MLRS_INSTANCES_CHECK_INTERVAL"
 RUNTIME_IMAGE_TAG_ENV_VAR = "MLRS_CONTAINER_IMAGE_TAG"
 
 # Stage mount paths
@@ -19,7 +22,7 @@ STAGE_VOLUME_MOUNT_PATH = "/mnt/job_stage"
 APP_STAGE_SUBPATH = "app"
 SYSTEM_STAGE_SUBPATH = "system"
 OUTPUT_STAGE_SUBPATH = "output"
-RESULT_PATH_DEFAULT_VALUE = f"{OUTPUT_STAGE_SUBPATH}/mljob_result.pkl"
+RESULT_PATH_DEFAULT_VALUE = f"{OUTPUT_STAGE_SUBPATH}/mljob_result"
 
 # Default container image information
 DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
@@ -53,8 +56,9 @@ ENABLE_HEALTH_CHECKS_ENV_VAR = "ENABLE_HEALTH_CHECKS"
 ENABLE_HEALTH_CHECKS = "false"
 
 # Job status polling constants
-JOB_POLL_INITIAL_DELAY_SECONDS = 0.1
+JOB_POLL_INITIAL_DELAY_SECONDS = 5
 JOB_POLL_MAX_DELAY_SECONDS = 30
+JOB_SPCS_TIMEOUT_SECONDS = 30
 
 # Log start and end messages
 LOG_START_MSG = "--------------------------------\nML job started\n--------------------------------"
@@ -70,6 +74,7 @@ COMMON_INSTANCE_FAMILIES = {
     "CPU_X64_XS": ComputeResources(cpu=1, memory=6),
     "CPU_X64_S": ComputeResources(cpu=3, memory=13),
     "CPU_X64_M": ComputeResources(cpu=6, memory=28),
+    "CPU_X64_SL": ComputeResources(cpu=14, memory=54),
     "CPU_X64_L": ComputeResources(cpu=28, memory=116),
     "HIGHMEM_X64_S": ComputeResources(cpu=6, memory=58),
 }
@@ -82,6 +87,7 @@ AWS_INSTANCE_FAMILIES = {
 }
 AZURE_INSTANCE_FAMILIES = {
     "HIGHMEM_X64_M": ComputeResources(cpu=28, memory=244),
+    "HIGHMEM_X64_SL": ComputeResources(cpu=92, memory=654),
     "HIGHMEM_X64_L": ComputeResources(cpu=92, memory=654),
     "GPU_NV_XS": ComputeResources(cpu=3, memory=26, gpu=1, gpu_type="T4"),
     "GPU_NV_SM": ComputeResources(cpu=32, memory=424, gpu=1, gpu_type="A10"),
@@ -89,7 +95,15 @@ AZURE_INSTANCE_FAMILIES = {
     "GPU_NV_3M": ComputeResources(cpu=44, memory=424, gpu=2, gpu_type="A100"),
     "GPU_NV_SL": ComputeResources(cpu=92, memory=858, gpu=4, gpu_type="A100"),
 }
+GCP_INSTANCE_FAMILIES = {
+    "HIGHMEM_X64_M": ComputeResources(cpu=28, memory=244),
+    "HIGHMEM_X64_SL": ComputeResources(cpu=92, memory=654),
+    "GPU_GCP_NV_L4_1_24G": ComputeResources(cpu=6, memory=28, gpu=1, gpu_type="L4"),
+    "GPU_GCP_NV_L4_4_24G": ComputeResources(cpu=44, memory=178, gpu=4, gpu_type="L4"),
+    "GPU_GCP_NV_A100_8_40G": ComputeResources(cpu=92, memory=654, gpu=8, gpu_type="A100"),
+}
 CLOUD_INSTANCE_FAMILIES = {
     SnowflakeCloudType.AWS: AWS_INSTANCE_FAMILIES,
     SnowflakeCloudType.AZURE: AZURE_INSTANCE_FAMILIES,
+    SnowflakeCloudType.GCP: GCP_INSTANCE_FAMILIES,
 }

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("."):
@@ -488,10 +488,13 @@ class JobPayload:
                 " comment = 'Created by snowflake.ml.jobs Python API'",
                 params=[stage_name],
             )
-
+        payload_name = None
         # Upload payload to stage - organize into app/ subdirectory
         app_stage_path = stage_path.joinpath(constants.APP_STAGE_SUBPATH)
         if not isinstance(source, types.PayloadPath):
+            if isinstance(source, function_payload_utils.FunctionPayload):
+                payload_name = source.function.__name__
+
             source_code = generate_python_code(source, source_code_display=True)
             _ = session.file.put_stream(
                 io.BytesIO(source_code.encode()),
@@ -502,12 +505,14 @@ class JobPayload:
             source = Path(entrypoint.file_path.parent)
 
         elif isinstance(source, stage_utils.StagePath):
+            payload_name = entrypoint.file_path.stem
             # copy payload to stage
             if source == entrypoint.file_path:
                 source = source.parent
             upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
 
         elif isinstance(source, Path):
+            payload_name = entrypoint.file_path.stem
             upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
             if source.is_file():
                 source = source.parent
@@ -562,6 +567,7 @@ class JobPayload:
                 *python_entrypoint,
             ],
             env_vars=env_vars,
+            payload_name=payload_name,
         )
 
 

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

@@ -104,7 +104,7 @@ def _get_image_spec(
             image_tag = runtime_environment
         else:
             container_image = runtime_environment
-    elif feature_flags.FeatureFlags.ENABLE_IMAGE_VERSION_ENV_VAR.is_enabled():
+    elif feature_flags.FeatureFlags.ENABLE_RUNTIME_VERSIONS.is_enabled():
         container_image = _get_runtime_image(session, hardware)  # type: ignore[arg-type]
 
     container_image = container_image or f"{image_repo}/{image_name}:{image_tag}"
@@ -266,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"},

snowflake/ml/jobs/_utils/stage_utils.py

@@ -32,6 +32,10 @@ class StagePath:
         self._root = self._raw_path[0:start].rstrip("/") if relpath else self._raw_path.rstrip("/")
         self._path = Path(relpath or "")
 
+    @property
+    def stem(self) -> str:
+        return self._path.stem
+
     @property
     def parts(self) -> tuple[str, ...]:
         return self._path.parts

snowflake/ml/jobs/_utils/types.py

@@ -11,6 +11,7 @@ JOB_STATUS = Literal[
     "CANCELLING",
     "CANCELLED",
     "INTERNAL_ERROR",
+    "DELETED",
 ]
 
 
@@ -22,6 +23,10 @@ class PayloadPath(Protocol):
     def name(self) -> str:
         ...
 
+    @property
+    def stem(self) -> str:
+        ...
+
     @property
     def suffix(self) -> str:
         ...
@@ -91,6 +96,7 @@ class UploadedPayload:
     stage_path: PurePath
     entrypoint: list[Union[str, PurePath]]
     env_vars: dict[str, str] = field(default_factory=dict)
+    payload_name: Optional[str] = None
 
 
 @dataclass(frozen=True)
@@ -106,3 +112,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