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

snowflake/ml/_internal/telemetry.py

@@ -411,16 +411,13 @@ def send_custom_usage(
     **kwargs: Any,
 ) -> None:
     conn = _get_snowflake_connection()
-    if conn is None:
-        raise ValueError(
-            """Snowflake connection is required to send custom telemetry. This means there
-            must be at least one active session, or that telemetry is being sent from within an SPCS service."""
-        )
 
-    client = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject)
-    common_metrics = client._create_basic_telemetry_data(telemetry_type=telemetry_type)
-    data = {**common_metrics, TelemetryField.KEY_DATA.value: data, **kwargs}
-    client._send(msg=data)
+    # Send telemetry if Snowflake connection is available.
+    if conn is not None:
+        client = _SourceTelemetryClient(conn=conn, project=project, subproject=subproject)
+        common_metrics = client._create_basic_telemetry_data(telemetry_type=telemetry_type)
+        data = {**common_metrics, TelemetryField.KEY_DATA.value: data, **kwargs}
+        client._send(msg=data)
 
 
 def send_api_usage_telemetry(

snowflake/ml/_internal/utils/connection_params.py

@@ -0,0 +1,196 @@
+import configparser
+import os
+from typing import Optional, Union
+
+from absl import logging
+from cryptography.hazmat import backends
+from cryptography.hazmat.primitives import serialization
+
+_DEFAULT_CONNECTION_FILE = "~/.snowsql/config"
+
+
+def _read_token(token_file: str = "") -> str:
+    """
+    Reads token from environment or file provided.
+
+    First tries to read the token from environment variable
+    (`SNOWFLAKE_TOKEN`) followed by the token file.
+    Both the options are tried out in SnowServices.
+
+    Args:
+        token_file: File from which token needs to be read. Optional.
+
+    Returns:
+        the token.
+    """
+    token = os.getenv("SNOWFLAKE_TOKEN", "")
+    if token:
+        return token
+    if token_file and os.path.exists(token_file):
+        with open(token_file) as f:
+            token = f.read()
+    return token
+
+
+_ENCRYPTED_PKCS8_PK_HEADER = b"-----BEGIN ENCRYPTED PRIVATE KEY-----"
+_UNENCRYPTED_PKCS8_PK_HEADER = b"-----BEGIN PRIVATE KEY-----"
+
+
+def _load_pem_to_der(private_key_path: str) -> bytes:
+    """Given a private key file path (in PEM format), decode key data into DER format."""
+    with open(private_key_path, "rb") as f:
+        private_key_pem = f.read()
+    private_key_passphrase: Optional[str] = os.getenv("SNOWFLAKE_PRIVATE_KEY_PASSPHRASE", None)
+
+    # Only PKCS#8 format key will be accepted. However, openssl
+    # transparently handle PKCS#8 and PKCS#1 format (by some fallback
+    # logic) and their is no function to distinguish between them. By
+    # reading openssl source code, apparently they also relies on header
+    # to determine if give bytes is PKCS#8 format or not
+    if not private_key_pem.startswith(_ENCRYPTED_PKCS8_PK_HEADER) and not private_key_pem.startswith(
+        _UNENCRYPTED_PKCS8_PK_HEADER
+    ):
+        raise Exception("Private key provided is not in PKCS#8 format. Please use correct format.")
+
+    if private_key_pem.startswith(_ENCRYPTED_PKCS8_PK_HEADER) and private_key_passphrase is None:
+        raise Exception(
+            "Private key is encrypted but passphrase could not be found. "
+            "Please set SNOWFLAKE_PRIVATE_KEY_PASSPHRASE env variable."
+        )
+
+    if private_key_pem.startswith(_UNENCRYPTED_PKCS8_PK_HEADER):
+        private_key_passphrase = None
+
+    private_key = serialization.load_pem_private_key(
+        private_key_pem,
+        str.encode(private_key_passphrase) if private_key_passphrase is not None else private_key_passphrase,
+        backends.default_backend(),
+    )
+
+    return private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+
+def _connection_properties_from_env() -> dict[str, str]:
+    """Returns a dict with all possible login related env variables."""
+    sf_conn_prop = {
+        # Mandatory fields
+        "account": os.environ["SNOWFLAKE_ACCOUNT"],
+        "database": os.environ["SNOWFLAKE_DATABASE"],
+        # With a default value
+        "token_file": os.getenv("SNOWFLAKE_TOKEN_FILE", "/snowflake/session/token"),
+        "ssl": os.getenv("SNOWFLAKE_SSL", "on"),
+        "protocol": os.getenv("SNOWFLAKE_PROTOCOL", "https"),
+    }
+    # With empty default value
+    for key, env_var in {
+        "user": "SNOWFLAKE_USER",
+        "authenticator": "SNOWFLAKE_AUTHENTICATOR",
+        "password": "SNOWFLAKE_PASSWORD",
+        "host": "SNOWFLAKE_HOST",
+        "port": "SNOWFLAKE_PORT",
+        "schema": "SNOWFLAKE_SCHEMA",
+        "warehouse": "SNOWFLAKE_WAREHOUSE",
+        "private_key_path": "SNOWFLAKE_PRIVATE_KEY_PATH",
+    }.items():
+        value = os.getenv(env_var, "")
+        if value:
+            sf_conn_prop[key] = value
+    return sf_conn_prop
+
+
+def _load_from_snowsql_config_file(connection_name: str, login_file: str = "") -> dict[str, str]:
+    """Loads the dictionary from snowsql config file."""
+    snowsql_config_file = login_file if login_file else os.path.expanduser(_DEFAULT_CONNECTION_FILE)
+    if not os.path.exists(snowsql_config_file):
+        logging.error(f"Connection name given but snowsql config file is not found at: {snowsql_config_file}")
+        raise Exception("Snowflake SnowSQL config not found.")
+
+    config = configparser.ConfigParser(inline_comment_prefixes="#")
+
+    snowflake_connection_name = os.getenv("SNOWFLAKE_CONNECTION_NAME")
+    if snowflake_connection_name is not None:
+        connection_name = snowflake_connection_name
+
+    if connection_name:
+        if not connection_name.startswith("connections."):
+            connection_name = "connections." + connection_name
+    else:
+        # See https://docs.snowflake.com/en/user-guide/snowsql-start.html#configuring-default-connection-settings
+        connection_name = "connections"
+
+    logging.info(f"Reading {snowsql_config_file} for connection parameters defined as {connection_name}")
+    config.read(snowsql_config_file)
+    conn_params = dict(config[connection_name])
+    # Remap names to appropriate args in Python Connector API
+    # Note: "dbname" should become "database"
+    conn_params = {k.replace("name", ""): v.strip('"') for k, v in conn_params.items()}
+    if "db" in conn_params:
+        conn_params["database"] = conn_params["db"]
+        del conn_params["db"]
+    return conn_params
+
+
+def SnowflakeLoginOptions(connection_name: str = "", login_file: Optional[str] = None) -> dict[str, Union[str, bytes]]:
+    """Returns a dict that can be used directly into snowflake python connector or Snowpark session config.
+
+    NOTE: Token/Auth information is sideloaded in all cases above, if provided in following order:
+      1. If SNOWFLAKE_TOKEN is defined in the environment, it will be used.
+      2. If SNOWFLAKE_TOKEN_FILE is defined in the environment and file matching the value found, content of the file
+         will be used.
+
+    If token is found, username, password will be reset and 'authenticator' will be set to 'oauth'.
+
+    Python Connector:
+    >> ctx = snowflake.connector.connect(**(SnowflakeLoginOptions()))
+
+    Snowpark Session:
+    >> session = Session.builder.configs(SnowflakeLoginOptions()).create()
+
+    Usage Note:
+      Ideally one should have a snowsql config file. Read more here:
+      https://docs.snowflake.com/en/user-guide/snowsql-start.html#configuring-default-connection-settings
+
+      If snowsql config file does not exist, it tries auth from env variables.
+
+    Args:
+        connection_name: Name of the connection to look for inside the config file. If environment variable
+            SNOWFLAKE_CONNECTION_NAME is provided, it will override the input connection_name.
+        login_file: If provided, this is used as config file instead of default one (_DEFAULT_CONNECTION_FILE).
+
+    Returns:
+        A dict with connection parameters.
+
+    Raises:
+        Exception: if none of config file and environment variable are present.
+    """
+    conn_prop: dict[str, Union[str, bytes]] = {}
+    login_file = login_file or os.path.expanduser(_DEFAULT_CONNECTION_FILE)
+    # If login file exists, use this exclusively.
+    if os.path.exists(login_file):
+        conn_prop = {**(_load_from_snowsql_config_file(connection_name, login_file))}
+    else:
+        # If environment exists for SNOWFLAKE_ACCOUNT, assume everything
+        # comes from environment. Mixing it not allowed.
+        account = os.getenv("SNOWFLAKE_ACCOUNT", "")
+        if account:
+            conn_prop = {**_connection_properties_from_env()}
+        else:
+            raise Exception("Snowflake credential is neither set in env nor a login file was provided.")
+
+    # Token, if specified, is always side-loaded in all cases.
+    token = _read_token(str(conn_prop["token_file"]) if "token_file" in conn_prop else "")
+    if token:
+        conn_prop["token"] = token
+        if "authenticator" not in conn_prop or conn_prop["authenticator"]:
+            conn_prop["authenticator"] = "oauth"
+    elif "private_key_path" in conn_prop and "private_key" not in conn_prop:
+        conn_prop["private_key"] = _load_pem_to_der(str(conn_prop["private_key_path"]))
+
+    if "ssl" in conn_prop and conn_prop["ssl"].lower() == "off":
+        conn_prop["protocol"] = "http"
+
+    return conn_prop

snowflake/ml/jobs/__init__.py

@@ -7,6 +7,7 @@ from snowflake.ml.jobs.manager import (
     list_jobs,
     submit_directory,
     submit_file,
+    submit_from_stage,
 )
 
 __all__ = [
@@ -18,4 +19,5 @@ __all__ = [
     "delete_job",
     "MLJob",
     "JOB_STATUS",
+    "submit_from_stage",
 ]

snowflake/ml/jobs/_utils/constants.py

@@ -6,6 +6,7 @@ DEFAULT_CONTAINER_NAME = "main"
 PAYLOAD_DIR_ENV_VAR = "MLRS_PAYLOAD_DIR"
 RESULT_PATH_ENV_VAR = "MLRS_RESULT_PATH"
 MIN_INSTANCES_ENV_VAR = "MLRS_MIN_INSTANCES"
+RUNTIME_IMAGE_TAG_ENV_VAR = "MLRS_CONTAINER_IMAGE_TAG"
 MEMORY_VOLUME_NAME = "dshm"
 STAGE_VOLUME_NAME = "stage-volume"
 STAGE_VOLUME_MOUNT_PATH = "/mnt/app"
@@ -14,7 +15,7 @@ STAGE_VOLUME_MOUNT_PATH = "/mnt/app"
 DEFAULT_IMAGE_REPO = "/snowflake/images/snowflake_images"
 DEFAULT_IMAGE_CPU = "st_plat/runtime/x86/runtime_image/snowbooks"
 DEFAULT_IMAGE_GPU = "st_plat/runtime/x86/generic_gpu/runtime_image/snowbooks"
-DEFAULT_IMAGE_TAG = "1.2.3"
+DEFAULT_IMAGE_TAG = "1.4.2"
 DEFAULT_ENTRYPOINT_PATH = "func.py"
 
 # Percent of container memory to allocate for /dev/shm volume
@@ -43,7 +44,7 @@ ENABLE_HEALTH_CHECKS = "false"
 
 # Job status polling constants
 JOB_POLL_INITIAL_DELAY_SECONDS = 0.1
-JOB_POLL_MAX_DELAY_SECONDS = 1
+JOB_POLL_MAX_DELAY_SECONDS = 30
 
 # Magic attributes
 IS_MLJOB_REMOTE_ATTR = "_is_mljob_remote_callable"

snowflake/ml/jobs/_utils/function_payload_utils.py

@@ -0,0 +1,43 @@
+import inspect
+from typing import Any, Callable, Optional
+
+from snowflake import snowpark
+from snowflake.snowpark import context as sp_context
+
+
+class FunctionPayload:
+    def __init__(
+        self,
+        func: Callable[..., Any],
+        session: Optional[snowpark.Session] = None,
+        session_argument: str = "",
+        *args: Any,
+        **kwargs: Any
+    ) -> None:
+        self.function = func
+        self.args = args
+        self.kwargs = kwargs
+        self._session = session
+        self._session_argument = session_argument
+
+    @property
+    def session(self) -> Optional[snowpark.Session]:
+        return self._session
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Customize pickling to exclude session."""
+        state = self.__dict__.copy()
+        state["_session"] = None
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore session from context during unpickling."""
+        self.__dict__.update(state)
+        self._session = sp_context.get_active_session()
+
+    def __call__(self) -> Any:
+        sig = inspect.signature(self.function)
+        bound = sig.bind_partial(*self.args, **self.kwargs)
+        bound.arguments[self._session_argument] = self._session
+
+        return self.function(*bound.args, **bound.kwargs)

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"
@@ -217,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(
@@ -245,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():
@@ -265,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:"
@@ -285,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:
@@ -310,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()),
@@ -321,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
@@ -502,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()
@@ -513,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(
@@ -551,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,10 +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
-# Inherited from snowflake.ml.jobs._utils.constants
-LOG_START_MSG = "--------------------------------\nML job started\n--------------------------------"
-LOG_END_MSG = "--------------------------------\nML job finished\n--------------------------------"
+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

@@ -13,7 +13,7 @@ from pathlib import Path
 from typing import Any, Optional
 
 import cloudpickle
-from constants import LOG_END_MSG, LOG_START_MSG
+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
@@ -72,28 +72,6 @@ class SimpleJSONEncoder(json.JSONEncoder):
             return f"Unserializable object: {repr(obj)}"
 
 
-def get_active_node_count() -> int:
-    """
-    Count the number of active nodes in the Ray cluster.
-
-    Returns:
-        int: Total count of active nodes
-    """
-    import ray
-
-    if not ray.is_initialized():
-        ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
-    try:
-        nodes = [node for node in ray.nodes() if node.get("Alive")]
-        total_active = len(nodes)
-
-        logger.info(f"Active nodes: {total_active}")
-        return total_active
-    except Exception as e:
-        logger.warning(f"Error getting active node count: {e}")
-        return 0
-
-
 def wait_for_min_instances(min_instances: int) -> None:
     """
     Wait until the specified minimum number of instances are available in the Ray cluster.
@@ -108,13 +86,16 @@ def wait_for_min_instances(min_instances: int) -> None:
         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 = get_active_node_count()
+        total_nodes = mlrs_util.get_num_ray_nodes()
 
         if total_nodes >= min_instances:
             elapsed = time.time() - start_time
@@ -128,7 +109,8 @@ def wait_for_min_instances(min_instances: int) -> None:
         time.sleep(check_interval)
 
     raise TimeoutError(
-        f"Timed out after {timeout}s waiting for {min_instances} instances, only {get_active_node_count()} available"
+        f"Timed out after {timeout}s waiting for {min_instances} instances, only "
+        f"{mlrs_util.get_num_ray_nodes()} available"
     )
 
 
@@ -199,7 +181,7 @@ def main(script_path: str, *script_args: Any, script_main_func: Optional[str] =
     """
     try:
         # Wait for minimum required instances if specified
-        min_instances_str = os.environ.get("JOB_MIN_INSTANCES", 1)
+        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))
 

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(
@@ -346,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