snowflake-ml-python 1.8.4__py3-none-any.whl → 1.8.6__py3-none-any.whl

snowflake/ml/jobs/_utils/payload_utils.py

@@ -12,10 +12,17 @@ import cloudpickle as cp
 from packaging import version
 
 from snowflake import snowpark
-from snowflake.ml.jobs._utils import constants, types
+from snowflake.ml.jobs._utils import (
+    constants,
+    function_payload_utils,
+    stage_utils,
+    types,
+)
 from snowflake.snowpark import exceptions as sp_exceptions
 from snowflake.snowpark._internal import code_generation
 
+cp.register_pickle_by_value(function_payload_utils)
+
 _SUPPORTED_ARG_TYPES = {str, int, float}
 _SUPPORTED_ENTRYPOINT_EXTENSIONS = {".py"}
 _ENTRYPOINT_FUNC_NAME = "func"
@@ -100,6 +107,11 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
             # 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"
+                echo "{constants.LOG_START_MSG}"
+            fi
+
             # Use the parsed variables
             echo "Head Instance Index: $head_index"
             echo "Head Instance IP: $head_ip"
@@ -117,9 +129,7 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
             exit 1
         fi
 
-        if [ "$SNOWFLAKE_JOB_INDEX" -ne "$head_index" ]; then
-            NODE_TYPE="worker"
-        fi
+
     fi
 
     # Common parameters for both head and worker nodes
@@ -168,6 +178,10 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
         # Start Ray on a worker node - run in background
         ray start "${{common_params[@]}}" "${{worker_params[@]}}" -v --block &
 
+        echo "Worker node started on address $eth0Ip. See more logs in the head node."
+
+        echo "{constants.LOG_END_MSG}"
+
         # Start the worker shutdown listener in the background
         echo "Starting worker shutdown listener..."
         python worker_shutdown_listener.py
@@ -189,15 +203,16 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
 
         # Start Ray on the head node
         ray start "${{common_params[@]}}" "${{head_params[@]}}" -v
+
         ##### End Ray configuration #####
 
         # TODO: Monitor MLRS and handle process crashes
         python -m web.ml_runtime_grpc_server &
 
         # TODO: Launch worker service(s) using SQL if Ray and MLRS successfully started
+        echo Running command: python "$@"
 
         # Run user's Python entrypoint
-        echo Running command: python "$@"
         python "$@"
 
         # After the user's job completes, signal workers to shut down
@@ -209,20 +224,23 @@ _STARTUP_SCRIPT_CODE = textwrap.dedent(
 ).strip()
 
 
-def resolve_source(source: Union[Path, Callable[..., Any]]) -> Union[Path, Callable[..., Any]]:
+def resolve_source(
+    source: Union[Path, stage_utils.StagePath, Callable[..., Any]]
+) -> Union[Path, stage_utils.StagePath, Callable[..., Any]]:
     if callable(source):
         return source
-    elif isinstance(source, Path):
-        # Validate source
-        source = source
+    elif isinstance(source, (Path, stage_utils.StagePath)):
         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.")
+        raise ValueError("Unsupported source type. Source must be a stage, file, directory, or callable.")
 
 
-def resolve_entrypoint(source: Union[Path, Callable[..., Any]], entrypoint: Optional[Path]) -> types.PayloadEntrypoint:
+def resolve_entrypoint(
+    source: Union[Path, stage_utils.StagePath, Callable[..., Any]],
+    entrypoint: Optional[Union[stage_utils.StagePath, Path]],
+) -> types.PayloadEntrypoint:
     if callable(source):
         # Entrypoint is generated for callable payloads
         return types.PayloadEntrypoint(
@@ -237,11 +255,11 @@ def resolve_entrypoint(source: Union[Path, Callable[..., Any]], entrypoint: Opti
             # Infer entrypoint from source
             entrypoint = parent
         else:
-            raise ValueError("entrypoint must be provided when source is a directory")
+            raise ValueError("Entrypoint must be provided when source is a directory")
     elif entrypoint.is_absolute():
         # Absolute path - validate it's a subpath of source dir
         if not entrypoint.is_relative_to(parent):
-            raise ValueError(f"Entrypoint must be a subpath of {parent}, got: {entrypoint})")
+            raise ValueError(f"Entrypoint must be a subpath of {parent}, got: {entrypoint}")
     else:
         # Relative path
         if (abs_entrypoint := entrypoint.absolute()).is_relative_to(parent) and abs_entrypoint.is_file():
@@ -257,6 +275,7 @@ def resolve_entrypoint(source: Union[Path, Callable[..., Any]], entrypoint: Opti
             "Entrypoint not found. Ensure the entrypoint is a valid file and is under"
             f" the source directory (source={parent}, entrypoint={entrypoint})"
         )
+
     if entrypoint.suffix not in _SUPPORTED_ENTRYPOINT_EXTENSIONS:
         raise ValueError(
             "Unsupported entrypoint type:"
@@ -277,8 +296,9 @@ class JobPayload:
         *,
         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
+        # for stage path like snow://domain....., Path(path) will remove duplicate /, it will become snow:/ domain...
+        self.source = stage_utils.identify_stage_path(source) if isinstance(source, str) else source
+        self.entrypoint = stage_utils.identify_stage_path(entrypoint) if isinstance(entrypoint, str) else entrypoint
         self.pip_requirements = pip_requirements
 
     def upload(self, session: snowpark.Session, stage_path: Union[str, PurePath]) -> types.UploadedPayload:
@@ -302,7 +322,7 @@ class JobPayload:
             ).collect()
 
         # Upload payload to stage
-        if not isinstance(source, Path):
+        if not isinstance(source, (Path, stage_utils.StagePath)):
             source_code = generate_python_code(source, source_code_display=True)
             _ = session.file.put_stream(
                 io.BytesIO(source_code.encode()),
@@ -313,27 +333,38 @@ class JobPayload:
             source = Path(entrypoint.file_path.parent)
             if not any(r.startswith("cloudpickle") for r in pip_requirements):
                 pip_requirements.append(f"cloudpickle~={version.parse(cp.__version__).major}.0")
-        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
-            # wildcard patterns to batch upload files with the same extension.
-            for path in {
-                p.parent.joinpath(f"*{p.suffix}") if p.suffix else p for p in source.resolve().rglob("*") if p.is_file()
-            }:
+
+        elif isinstance(source, stage_utils.StagePath):
+            # copy payload to stage
+            if source == entrypoint.file_path:
+                source = source.parent
+            source_path = source.as_posix() + "/"
+            session.sql(f"copy files into {stage_path}/ from {source_path}").collect()
+
+        elif isinstance(source, Path):
+            if 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
+                # wildcard patterns to batch upload files with the same extension.
+                for path in {
+                    p.parent.joinpath(f"*{p.suffix}") if p.suffix else p
+                    for p in source.resolve().rglob("*")
+                    if p.is_file()
+                }:
+                    session.file.put(
+                        str(path),
+                        stage_path.joinpath(path.parent.relative_to(source)).as_posix(),
+                        overwrite=True,
+                        auto_compress=False,
+                    )
+            else:
                 session.file.put(
-                    str(path),
-                    stage_path.joinpath(path.parent.relative_to(source)).as_posix(),
+                    str(source.resolve()),
+                    stage_path.as_posix(),
                     overwrite=True,
                     auto_compress=False,
                 )
-        else:
-            session.file.put(
-                str(source.resolve()),
-                stage_path.as_posix(),
-                overwrite=True,
-                auto_compress=False,
-            )
-            source = source.parent
+                source = source.parent
 
         # Upload requirements
         # TODO: Check if payload includes both a requirements.txt file and pip_requirements
@@ -494,9 +525,15 @@ def _generate_param_handler_code(signature: inspect.Signature, output_name: str
     return param_code
 
 
-def generate_python_code(func: Callable[..., Any], source_code_display: bool = False) -> str:
+def generate_python_code(payload: Callable[..., Any], source_code_display: bool = False) -> str:
     """Generate an entrypoint script from a Python function."""
-    signature = inspect.signature(func)
+
+    if isinstance(payload, function_payload_utils.FunctionPayload):
+        function = payload.function
+    else:
+        function = payload
+
+    signature = inspect.signature(function)
     if any(
         p.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
         for p in signature.parameters.values()
@@ -505,21 +542,20 @@ def generate_python_code(func: Callable[..., Any], source_code_display: bool = F
 
     # Mirrored from Snowpark generate_python_code() function
     # 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 ""
+    source_code_comment = _generate_source_code_comment(function) if source_code_display else ""
 
     arg_dict_name = "kwargs"
-    if getattr(func, constants.IS_MLJOB_REMOTE_ATTR, None):
+    if isinstance(payload, function_payload_utils.FunctionPayload):
         param_code = f"{arg_dict_name} = {{}}"
     else:
         param_code = _generate_param_handler_code(signature, arg_dict_name)
-
     return f"""
 import sys
 import pickle
 
 try:
     {textwrap.indent(source_code_comment, '    ')}
-    {_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(func).hex()}'))
+    {_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(payload).hex()}'))
 except (TypeError, pickle.PickleError):
     if sys.version_info.major != {sys.version_info.major} or sys.version_info.minor != {sys.version_info.minor}:
         raise RuntimeError(
@@ -543,3 +579,23 @@ if __name__ == '__main__':
 
     __return__ = {_ENTRYPOINT_FUNC_NAME}(**{arg_dict_name})
 """
+
+
+def create_function_payload(
+    func: Callable[..., Any], *args: Any, **kwargs: Any
+) -> function_payload_utils.FunctionPayload:
+    signature = inspect.signature(func)
+    bound = signature.bind(*args, **kwargs)
+    bound.apply_defaults()
+    session_argument = ""
+    session = None
+    for name, val in list(bound.arguments.items()):
+        if isinstance(val, snowpark.Session):
+            if session:
+                raise TypeError(f"Expected only one Session-type argument, but got both {session_argument} and {name}.")
+            session = val
+            session_argument = name
+            del bound.arguments[name]
+    payload = function_payload_utils.FunctionPayload(func, session, session_argument, *bound.args, **bound.kwargs)
+
+    return payload

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

@@ -1,4 +1,26 @@
+from snowflake.ml.jobs._utils import constants as mljob_constants
+
 # Constants defining the shutdown signal actor configuration.
 SHUTDOWN_ACTOR_NAME = "ShutdownSignal"
 SHUTDOWN_ACTOR_NAMESPACE = "default"
 SHUTDOWN_RPC_TIMEOUT_SECONDS = 5.0
+
+
+# The followings are Inherited from snowflake.ml.jobs._utils.constants
+# We need to copy them here since snowml package on the server side does
+# not have the latest version of the code
+
+# Log start and end messages
+LOG_START_MSG = getattr(
+    mljob_constants,
+    "LOG_START_MSG",
+    "--------------------------------\nML job started\n--------------------------------",
+)
+LOG_END_MSG = getattr(
+    mljob_constants,
+    "LOG_END_MSG",
+    "--------------------------------\nML job finished\n--------------------------------",
+)
+
+# min_instances environment variable name
+MIN_INSTANCES_ENV_VAR = getattr(mljob_constants, "MIN_INSTANCES_ENV_VAR", "MLRS_MIN_INSTANCES")

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

@@ -2,25 +2,35 @@ import argparse
 import copy
 import importlib.util
 import json
+import logging
 import os
 import runpy
 import sys
+import time
 import traceback
 import warnings
 from pathlib import Path
 from typing import Any, Optional
 
 import cloudpickle
+from constants import LOG_END_MSG, LOG_START_MSG, MIN_INSTANCES_ENV_VAR
 
 from snowflake.ml.jobs._utils import constants
 from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
 from snowflake.snowpark import Session
 
+# Configure logging
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
 # 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")
 
+# Constants for the wait_for_min_instances function
+CHECK_INTERVAL = 10  # seconds
+TIMEOUT = 720  # seconds
+
 
 try:
     from snowflake.ml.jobs._utils.interop_utils import ExecutionResult
@@ -62,6 +72,48 @@ class SimpleJSONEncoder(json.JSONEncoder):
             return f"Unserializable object: {repr(obj)}"
 
 
+def wait_for_min_instances(min_instances: int) -> None:
+    """
+    Wait until the specified minimum number of instances are available in the Ray cluster.
+
+    Args:
+        min_instances: Minimum number of instances required
+
+    Raises:
+        TimeoutError: If failed to connect to Ray or if minimum instances are not available within timeout
+    """
+    if min_instances <= 1:
+        logger.debug("Minimum instances is 1 or less, no need to wait for additional instances")
+        return
+
+    # mljob_launcher runs inside the CR where mlruntime libraries are available, so we can import common_util directly
+    from common_utils import common_util as mlrs_util
+
+    start_time = time.time()
+    timeout = os.getenv("JOB_MIN_INSTANCES_TIMEOUT", TIMEOUT)
+    check_interval = os.getenv("JOB_MIN_INSTANCES_CHECK_INTERVAL", CHECK_INTERVAL)
+    logger.debug(f"Waiting for at least {min_instances} instances to be ready (timeout: {timeout}s)")
+
+    while time.time() - start_time < timeout:
+        total_nodes = mlrs_util.get_num_ray_nodes()
+
+        if total_nodes >= min_instances:
+            elapsed = time.time() - start_time
+            logger.info(f"Minimum instance requirement met: {total_nodes} instances available after {elapsed:.1f}s")
+            return
+
+        logger.debug(
+            f"Waiting for instances: {total_nodes}/{min_instances} available "
+            f"(elapsed: {time.time() - start_time:.1f}s)"
+        )
+        time.sleep(check_interval)
+
+    raise TimeoutError(
+        f"Timed out after {timeout}s waiting for {min_instances} instances, only "
+        f"{mlrs_util.get_num_ray_nodes()} available"
+    )
+
+
 def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = None) -> Any:
     """
     Execute a Python script and return its result.
@@ -86,6 +138,7 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
     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
@@ -126,9 +179,21 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
     Raises:
         Exception: Re-raises any exception caught during script execution.
     """
-    # Run the script with the specified arguments
     try:
+        # Wait for minimum required instances if specified
+        min_instances_str = os.environ.get(MIN_INSTANCES_ENV_VAR) or "1"
+        if min_instances_str and int(min_instances_str) > 1:
+            wait_for_min_instances(int(min_instances_str))
+
+        # Log start marker for 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)
+
+        # Log end marker for user script execution
+        print(LOG_END_MSG)  # noqa: T201
+
         result_obj = ExecutionResult(result=result)
         return result_obj
     except Exception as e:

snowflake/ml/jobs/_utils/spec_utils.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from math import ceil
 from pathlib import PurePath
 from typing import Any, Optional, Union
@@ -30,7 +31,7 @@ def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.Image
     # Use MLRuntime image
     image_repo = constants.DEFAULT_IMAGE_REPO
     image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
-    image_tag = constants.DEFAULT_IMAGE_TAG
+    image_tag = _get_runtime_image_tag()
 
     # TODO: Should each instance consume the entire pod?
     return types.ImageSpec(
@@ -85,7 +86,8 @@ def generate_service_spec(
     compute_pool: str,
     payload: types.UploadedPayload,
     args: Optional[list[str]] = None,
-    num_instances: Optional[int] = None,
+    target_instances: int = 1,
+    min_instances: int = 1,
     enable_metrics: bool = False,
 ) -> dict[str, Any]:
     """
@@ -96,13 +98,13 @@ def generate_service_spec(
         compute_pool: Compute pool for job execution
         payload: Uploaded job payload
         args: Arguments to pass to entrypoint script
-        num_instances: Number of instances for multi-node job
+        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
 
     Returns:
         Job service specification
     """
-    is_multi_node = num_instances is not None and num_instances > 1
     image_spec = _get_image_spec(session, compute_pool)
 
     # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
@@ -180,10 +182,11 @@ def generate_service_spec(
     }
     endpoints = []
 
-    if is_multi_node:
+    if target_instances > 1:
         # Update environment variables for multi-node job
         env_vars.update(constants.RAY_PORTS)
-        env_vars["ENABLE_HEALTH_CHECKS"] = constants.ENABLE_HEALTH_CHECKS
+        env_vars[constants.ENABLE_HEALTH_CHECKS_ENV_VAR] = constants.ENABLE_HEALTH_CHECKS
+        env_vars[constants.MIN_INSTANCES_ENV_VAR] = str(min_instances)
 
         # Define Ray endpoints for intra-service instance communication
         ray_endpoints = [
@@ -344,3 +347,24 @@ def _merge_lists_of_dicts(
         result[key] = d
 
     return list(result.values())
+
+
+def _get_runtime_image_tag() -> str:
+    """
+    Detect runtime image tag from container environment.
+
+    Checks in order:
+    1. Environment variable MLRS_CONTAINER_IMAGE_TAG
+    2. Falls back to hardcoded default
+
+    Returns:
+        str: The runtime image tag to use for job containers
+    """
+    env_tag = os.environ.get(constants.RUNTIME_IMAGE_TAG_ENV_VAR)
+    if env_tag:
+        logging.debug(f"Using runtime image tag from environment: {env_tag}")
+        return env_tag
+
+    # Fall back to default
+    logging.debug(f"Using default runtime image tag: {constants.DEFAULT_IMAGE_TAG}")
+    return constants.DEFAULT_IMAGE_TAG

snowflake/ml/jobs/_utils/stage_utils.py

@@ -0,0 +1,119 @@
+import os
+import re
+from os import PathLike
+from pathlib import Path, PurePath
+from typing import Union
+
+from snowflake.ml._internal.utils import identifier
+
+PROTOCOL_NAME = "snow"
+_SNOWURL_PATH_RE = re.compile(
+    rf"^(?:(?:{PROTOCOL_NAME}://)?"
+    r"(?<!@)(?P<domain>\w+)/"
+    rf"(?P<name>(?:{identifier._SF_IDENTIFIER}\.){{,2}}{identifier._SF_IDENTIFIER})/)?"
+    r"(?P<path>versions(?:/(?P<version>[^/]+)(?:/(?P<relpath>.*))?)?)$"
+)
+
+_STAGEF_PATH_RE = re.compile(r"^@(?P<stage>~|%?\w+)(?:/(?P<relpath>[\w\-./]*))?$")
+
+
+class StagePath:
+    def __init__(self, path: str) -> None:
+        stage_match = _SNOWURL_PATH_RE.fullmatch(path) or _STAGEF_PATH_RE.fullmatch(path)
+        if not stage_match:
+            raise ValueError(f"{path} is not a valid stage path")
+        path = path.strip()
+        self._raw_path = path
+        relpath = stage_match.group("relpath")
+        start, _ = stage_match.span("relpath")
+        self._root = self._raw_path[0:start].rstrip("/") if relpath else self._raw_path.rstrip("/")
+        self._path = Path(relpath or "")
+
+    @property
+    def parent(self) -> "StagePath":
+        if self._path.parent == Path(""):
+            return StagePath(self._root)
+        else:
+            return StagePath(f"{self._root}/{self._path.parent}")
+
+    @property
+    def root(self) -> str:
+        return self._root
+
+    @property
+    def suffix(self) -> str:
+        return self._path.suffix
+
+    def _compose_path(self, path: Path) -> str:
+        # in pathlib, Path("") = "."
+        if path == Path(""):
+            return self.root
+        else:
+            return f"{self.root}/{path}"
+
+    def is_relative_to(self, path: Union[str, PathLike[str], "StagePath"]) -> bool:
+        stage_path = path if isinstance(path, StagePath) else StagePath(os.fspath(path))
+        if stage_path.root == self.root:
+            return self._path.is_relative_to(stage_path._path)
+        else:
+            return False
+
+    def relative_to(self, path: Union[str, PathLike[str], "StagePath"]) -> PurePath:
+        stage_path = path if isinstance(path, StagePath) else StagePath(os.fspath(path))
+        if self.root == stage_path.root:
+            return self._path.relative_to(stage_path._path)
+        raise ValueError(f"{self._raw_path} does not start with {stage_path._raw_path}")
+
+    def absolute(self) -> "StagePath":
+        return self
+
+    def as_posix(self) -> str:
+        return self._compose_path(self._path)
+
+    # TODO Add actual implementation https://snowflakecomputing.atlassian.net/browse/SNOW-2112795
+    def exists(self) -> bool:
+        return True
+
+    # TODO Add actual implementation https://snowflakecomputing.atlassian.net/browse/SNOW-2112795
+    def is_file(self) -> bool:
+        return True
+
+    # TODO Add actual implementation https://snowflakecomputing.atlassian.net/browse/SNOW-2112795
+    def is_dir(self) -> bool:
+        return True
+
+    def is_absolute(self) -> bool:
+        return True
+
+    def __str__(self) -> str:
+        return self.as_posix()
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, StagePath):
+            raise NotImplementedError
+        return bool(self.root == other.root and self._path == other._path)
+
+    def __fspath__(self) -> str:
+        return self._compose_path(self._path)
+
+    def joinpath(self, *args: Union[str, PathLike[str], "StagePath"]) -> "StagePath":
+        path = self
+        for arg in args:
+            path = path._make_child(arg)
+        return path
+
+    def _make_child(self, path: Union[str, PathLike[str], "StagePath"]) -> "StagePath":
+        stage_path = path if isinstance(path, StagePath) else StagePath(os.fspath(path))
+        if self.root == stage_path.root:
+            child_path = self._path.joinpath(stage_path._path)
+            return StagePath(self._compose_path(child_path))
+        else:
+            return stage_path
+
+
+def identify_stage_path(path: str) -> Union[StagePath, Path]:
+    try:
+        stage_path = StagePath(path)
+    except ValueError:
+        return Path(path)
+    return stage_path

snowflake/ml/jobs/_utils/types.py

@@ -2,18 +2,22 @@ from dataclasses import dataclass
 from pathlib import PurePath
 from typing import Literal, Optional, Union
 
+from snowflake.ml.jobs._utils import stage_utils
+
 JOB_STATUS = Literal[
     "PENDING",
     "RUNNING",
     "FAILED",
     "DONE",
+    "CANCELLING",
+    "CANCELLED",
     "INTERNAL_ERROR",
 ]
 
 
 @dataclass(frozen=True)
 class PayloadEntrypoint:
-    file_path: PurePath
+    file_path: Union[PurePath, stage_utils.StagePath]
     main_func: Optional[str]
 
 

snowflake/ml/jobs/decorators.py

@@ -7,7 +7,7 @@ from typing_extensions import ParamSpec
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml.jobs import job as jb, manager as jm
-from snowflake.ml.jobs._utils import constants
+from snowflake.ml.jobs._utils import payload_utils
 
 _PROJECT = "MLJob"
 
@@ -24,7 +24,8 @@ def remote(
     external_access_integrations: Optional[list[str]] = None,
     query_warehouse: Optional[str] = None,
     env_vars: Optional[dict[str, str]] = None,
-    num_instances: Optional[int] = None,
+    target_instances: int = 1,
+    min_instances: Optional[int] = None,
     enable_metrics: bool = False,
     database: Optional[str] = None,
     schema: Optional[str] = None,
@@ -40,7 +41,9 @@ def remote(
         external_access_integrations: A list of external access integrations.
         query_warehouse: The query warehouse to use. Defaults to session warehouse.
         env_vars: Environment variables to set in container
-        num_instances: The number of nodes in the job. If none specified, create a single node job.
+        target_instances: The number of nodes in the job. If none specified, create a single node job.
+        min_instances: The minimum number of nodes required to start the job. If none specified,
+            defaults to target_instances. If set, the job will not start until the minimum number of nodes is available.
         enable_metrics: Whether to enable metrics publishing for the job.
         database: The database to use for the job.
         schema: The schema to use for the job.
@@ -59,8 +62,7 @@ def remote(
 
         @functools.wraps(func)
         def wrapper(*args: _Args.args, **kwargs: _Args.kwargs) -> jb.MLJob[_ReturnValue]:
-            payload = functools.partial(func, *args, **kwargs)
-            setattr(payload, constants.IS_MLJOB_REMOTE_ATTR, True)
+            payload = payload_utils.create_function_payload(func, *args, **kwargs)
             job = jm._submit_job(
                 source=payload,
                 stage_name=stage_name,
@@ -69,11 +71,12 @@ def remote(
                 external_access_integrations=external_access_integrations,
                 query_warehouse=query_warehouse,
                 env_vars=env_vars,
-                num_instances=num_instances,
+                target_instances=target_instances,
+                min_instances=min_instances,
                 enable_metrics=enable_metrics,
                 database=database,
                 schema=schema,
-                session=session,
+                session=payload.session or session,
             )
             assert isinstance(job, jb.MLJob), f"Unexpected job type: {type(job)}"
             return job