snowflake-ml-python 1.24.0__py3-none-any.whl → 1.25.0__py3-none-any.whl

snowflake/ml/jobs/_utils/feature_flags.py

@@ -1,6 +1,11 @@
 import os
-from enum import Enum
-from typing import Optional
+from typing import Callable, Optional, Union
+
+from snowflake.ml._internal.utils.snowflake_env import SnowflakeCloudType
+from snowflake.snowpark import context as sp_context
+
+# Default value type: can be a bool or a callable that returns a bool
+DefaultValue = Union[bool, Callable[[], bool]]
 
 
 def parse_bool_env_value(value: Optional[str], default: bool = False) -> bool:
@@ -28,22 +33,101 @@ def parse_bool_env_value(value: Optional[str], default: bool = False) -> bool:
         return default
 
 
-class FeatureFlags(Enum):
-    USE_SUBMIT_JOB_V2 = "MLRS_USE_SUBMIT_JOB_V2"
-    ENABLE_RUNTIME_VERSIONS = "MLRS_ENABLE_RUNTIME_VERSIONS"
-    ENABLE_STAGE_MOUNT_V2 = "MLRS_ENABLE_STAGE_MOUNT_V2"
+def _enabled_in_clouds(*clouds: SnowflakeCloudType) -> Callable[[], bool]:
+    """Create a callable that checks if the current environment is in any of the specified clouds.
+
+    This factory function returns a callable that can be used as a dynamic default
+    for feature flags. The returned callable will check if the current Snowflake
+    session is connected to a region in any of the specified cloud providers.
+
+    Args:
+        *clouds: One or more SnowflakeCloudType values to check against.
+
+    Returns:
+        A callable that returns True if running in any of the specified clouds,
+        False otherwise (including when no session is available).
+
+    Example:
+        >>> # Enable feature only in GCP
+        >>> default=_enabled_in_clouds(SnowflakeCloudType.GCP)
+        >>>
+        >>> # Enable feature in both GCP and Azure
+        >>> default=_enabled_in_clouds(SnowflakeCloudType.GCP, SnowflakeCloudType.AZURE)
+    """
+    cloud_set = frozenset(clouds)
+
+    def check() -> bool:
+        try:
+            from snowflake.ml._internal.utils.snowflake_env import get_current_cloud
+
+            session = sp_context.get_active_session()
+            current_cloud = get_current_cloud(session, default=SnowflakeCloudType.AWS)
+            return current_cloud in cloud_set
+        except Exception:
+            # If we can't determine the cloud (no session, SQL error, etc.),
+            # default to False for safety
+            return False
+
+    return check
 
-    def is_enabled(self, default: bool = False) -> bool:
-        """Check if the feature flag is enabled.
+
+class _FeatureFlag:
+    """A feature flag backed by an environment variable with a configurable default.
+
+    The default value can be a constant boolean or a callable that dynamically
+    determines the default based on runtime context (e.g., cloud provider).
+    """
+
+    def __init__(self, env_var: str, default: DefaultValue = False) -> None:
+        """Initialize a feature flag.
 
         Args:
-            default: The default value to return if the environment variable is not set.
+            env_var: The environment variable name that controls this flag.
+            default: The default value when the env var is not set. Can be:
+                - A boolean constant (True/False)
+                - A callable that returns a boolean (evaluated at check time)
+        """
+        self._env_var = env_var
+        self._default = default
+
+    @property
+    def value(self) -> str:
+        """Return the environment variable name (for compatibility with Enum-style access)."""
+        return self._env_var
+
+    def _get_default(self) -> bool:
+        """Get the default value, calling it if it's a callable."""
+        if callable(self._default):
+            return self._default()
+        return self._default
+
+    def is_enabled(self) -> bool:
+        """Check if the feature flag is enabled.
+
+        First checks the environment variable. If not set or unrecognized,
+        falls back to the configured default value.
 
         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.
+            True if the feature is enabled, False otherwise.
         """
-        return parse_bool_env_value(os.getenv(self.value), default)
+        env_value = os.getenv(self._env_var)
+        if env_value is not None:
+            # Environment variable is set, parse it
+            result = parse_bool_env_value(env_value, default=self._get_default())
+            return result
+        else:
+            # Environment variable not set, use the default
+            return self._get_default()
 
     def __str__(self) -> str:
-        return self.value
+        return self._env_var
+
+
+class FeatureFlags:
+    """Collection of feature flags for ML Jobs."""
+
+    ENABLE_RUNTIME_VERSIONS = _FeatureFlag("MLRS_ENABLE_RUNTIME_VERSIONS", default=True)
+    ENABLE_STAGE_MOUNT_V2 = _FeatureFlag(
+        "MLRS_ENABLE_STAGE_MOUNT_V2",
+        default=_enabled_in_clouds(SnowflakeCloudType.GCP),
+    )

snowflake/ml/jobs/_utils/payload_utils.py

@@ -17,20 +17,12 @@ import cloudpickle as cp
 from packaging import version
 
 from snowflake import snowpark
-from snowflake.ml.jobs._utils import (
-    constants,
-    function_payload_utils,
-    query_helper,
-    stage_utils,
-    types,
-)
+from snowflake.ml.jobs._utils import constants, query_helper, stage_utils, types
 from snowflake.snowpark import exceptions as sp_exceptions
 from snowflake.snowpark._internal import code_generation
 from snowflake.snowpark._internal.utils import zip_file_or_directory_to_stream
 
 logger = logging.getLogger(__name__)
-
-cp.register_pickle_by_value(function_payload_utils)
 ImportType = Union[str, Path, ModuleType]
 
 _SUPPORTED_ARG_TYPES = {str, int, float}
@@ -561,7 +553,6 @@ class JobPayload:
         env_vars = {
             constants.STAGE_MOUNT_PATH_ENV_VAR: constants.STAGE_VOLUME_MOUNT_PATH,
             constants.PAYLOAD_DIR_ENV_VAR: constants.APP_STAGE_SUBPATH,
-            constants.RESULT_PATH_ENV_VAR: constants.RESULT_PATH_DEFAULT_VALUE,
         }
 
         return types.UploadedPayload(
@@ -691,14 +682,9 @@ def _generate_param_handler_code(signature: inspect.Signature, output_name: str
     return param_code
 
 
-def generate_python_code(payload: Callable[..., Any], source_code_display: bool = False) -> str:
+def generate_python_code(function: Callable[..., Any], source_code_display: bool = False) -> str:
     """Generate an entrypoint script from a Python function."""
 
-    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}
@@ -711,7 +697,7 @@ def generate_python_code(payload: Callable[..., Any], source_code_display: bool
     source_code_comment = _generate_source_code_comment(function) if source_code_display else ""
 
     arg_dict_name = "kwargs"
-    if isinstance(payload, function_payload_utils.FunctionPayload):
+    if getattr(function, constants.IS_MLJOB_REMOTE_ATTR, None):
         param_code = f"{arg_dict_name} = {{}}"
     else:
         param_code = _generate_param_handler_code(signature, arg_dict_name)
@@ -721,7 +707,7 @@ import pickle
 
 try:
     {textwrap.indent(source_code_comment, '    ')}
-    {_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(payload).hex()}'))
+    {_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(function).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(
@@ -747,26 +733,6 @@ if __name__ == '__main__':
 """
 
 
-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
-
-
 def get_payload_name(source: Union[str, Callable[..., Any]], entrypoint: Optional[Union[str, list[str]]] = None) -> str:
 
     if entrypoint and isinstance(entrypoint, (list, tuple)):
@@ -775,7 +741,7 @@ def get_payload_name(source: Union[str, Callable[..., Any]], entrypoint: Optiona
         return f"{PurePath(entrypoint).stem}"
     elif source and not callable(source):
         return f"{PurePath(source).stem}"
-    elif isinstance(source, function_payload_utils.FunctionPayload):
-        return f"{source.function.__name__}"
+    elif callable(source):
+        return f"{source.__name__}"
     else:
         return f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"

snowflake/ml/jobs/_utils/runtime_env_utils.py

@@ -1,117 +1,18 @@
-import datetime
-import logging
-from typing import Any, Literal, Optional, Union
-
-from packaging.version import Version
-from pydantic import BaseModel, Field, RootModel, field_validator
+from typing import Optional, cast
 
 from snowflake import snowpark
-from snowflake.ml.jobs._utils import constants, query_helper
-
-
-class SpcsContainerRuntime(BaseModel):
-    python_version: Version = Field(alias="pythonVersion")
-    hardware_type: str = Field(alias="hardwareType")
-    runtime_container_image: str = Field(alias="runtimeContainerImage")
-
-    @field_validator("python_version", mode="before")
-    @classmethod
-    def validate_python_version(cls, v: Union[str, Version]) -> Version:
-        if isinstance(v, Version):
-            return v
-        try:
-            return Version(v)
-        except Exception:
-            raise ValueError(f"Invalid Python version format: {v}")
-
-    class Config:
-        frozen = True
-        extra = "allow"
-        arbitrary_types_allowed = True
-
-
-class RuntimeEnvironmentEntry(BaseModel):
-    spcs_container_runtime: Optional[SpcsContainerRuntime] = Field(alias="spcsContainerRuntime", default=None)
-    created_on: datetime.datetime = Field(alias="createdOn")
-    id: Optional[str] = Field(alias="id")
-
-    class Config:
-        extra = "allow"
-        frozen = True
-
-
-class RuntimeEnvironmentsDict(RootModel[dict[str, RuntimeEnvironmentEntry]]):
-    @field_validator("root", mode="before")
-    @classmethod
-    def _filter_to_dict_entries(cls, data: Any) -> dict[str, dict[str, Any]]:
-        """
-        Pre-validation hook: keep only those items at the root level
-        whose values are dicts. Non-dict values will be dropped.
+from snowflake.ml.jobs._utils import query_helper
 
-        Args:
-            data: The input data to filter, expected to be a dictionary.
 
-        Returns:
-            A dictionary containing only the key-value pairs where values are dictionaries.
-
-        Raises:
-            ValueError: If input data is not a dictionary.
-        """
-        # If the entire root is not a dict, raise error immediately
-        if not isinstance(data, dict):
-            raise ValueError(f"Expected dictionary data, but got {type(data).__name__}: {data}")
-
-        # Filter out any key whose value is not a dict
-        return {key: value for key, value in data.items() if isinstance(value, dict)}
-
-    def get_spcs_container_runtimes(
-        self,
-        *,
-        hardware_type: Optional[str] = None,
-        python_version: Optional[Version] = None,
-    ) -> list[SpcsContainerRuntime]:
-        # TODO(SNOW-2682000): parse version from NRE in a safer way, like relying on the label,id or image tag.
-        entries: list[RuntimeEnvironmentEntry] = [
-            entry
-            for entry in self.root.values()
-            if entry.spcs_container_runtime is not None
-            and (hardware_type is None or entry.spcs_container_runtime.hardware_type.lower() == hardware_type.lower())
-            and (
-                python_version is None
-                or (
-                    entry.spcs_container_runtime.python_version.major == python_version.major
-                    and entry.spcs_container_runtime.python_version.minor == python_version.minor
-                )
-            )
-        ]
-        entries.sort(key=lambda e: e.created_on, reverse=True)
-
-        return [entry.spcs_container_runtime for entry in entries if entry.spcs_container_runtime is not None]
-
-
-def _extract_image_tag(image_url: str) -> Optional[str]:
-    image_tag = image_url.rsplit(":", 1)[-1]
-    return image_tag
-
-
-def find_runtime_image(
-    session: snowpark.Session, target_hardware: Literal["CPU", "GPU"], target_python_version: Optional[str] = None
+def get_runtime_image(
+    session: snowpark.Session, compute_pool: str, runtime_environment: Optional[str] = None
 ) -> Optional[str]:
-    python_version = (
-        Version(target_python_version) if target_python_version else Version(constants.DEFAULT_PYTHON_VERSION)
-    )
-    rows = query_helper.run_query(session, "CALL SYSTEM$NOTEBOOKS_FIND_LABELED_RUNTIMES()")
+    runtime_environment = runtime_environment if runtime_environment else ""
+    rows = query_helper.run_query(session, f"CALL SYSTEM$GET_ML_JOB_RUNTIME('{compute_pool}', '{runtime_environment}')")
     if not rows:
-        return None
-    try:
-        runtime_envs = RuntimeEnvironmentsDict.model_validate_json(rows[0][0])
-        spcs_container_runtimes = runtime_envs.get_spcs_container_runtimes(
-            hardware_type=target_hardware,
-            python_version=python_version,
-        )
-    except Exception as e:
-        logging.warning(f"Failed to parse runtime image name from {rows[0][0]}, error: {e}")
-        return None
-
-    selected_runtime = spcs_container_runtimes[0] if spcs_container_runtimes else None
-    return selected_runtime.runtime_container_image if selected_runtime else None
+        raise ValueError("Failed to get any available runtime image")
+    image = rows[0][0]
+    url, tag = image.rsplit(":", 1)
+    if url is None or tag is None:
+        raise ValueError(f"image {image} is not a valid runtime image")
+    return cast(str, image) if image else None

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

@@ -1,6 +1,7 @@
 import argparse
 import copy
 import importlib.util
+import io
 import json
 import logging
 import math
@@ -12,15 +13,22 @@ import sys
 import time
 import traceback
 import zipfile
-from pathlib import Path
-from typing import Any, Optional
+from pathlib import Path, PurePosixPath
+from typing import Any, Callable, Optional
 
 # 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")
+STAGE_RESULT_PATH = os.environ.get("MLRS_STAGE_RESULT_PATH")
+# Updated MLRS_RESULT_PATH to use unique stage mounts for each ML Job.
+# To prevent output collisions between jobs sharing the same definition,
+# the server-side mount now dynamically includes the job_name.
+# Format: @payload_stage/{job_definition_name}/{job_name}/mljob_result
+JOB_RESULT_PATH = os.environ.get("MLRS_RESULT_PATH", "mljob_result")
+if STAGE_RESULT_PATH:
+    JOB_RESULT_PATH = os.path.join(STAGE_RESULT_PATH, JOB_RESULT_PATH)
 PAYLOAD_PATH = os.environ.get("MLRS_PAYLOAD_DIR")
 
 if PAYLOAD_PATH and not os.path.isabs(PAYLOAD_PATH):
@@ -347,24 +355,156 @@ def wait_for_instances(
     )
 
 
-def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = None) -> Any:
+def _load_dto_fallback(function_args: str, path_transform: Callable[[str], str]) -> Any:
+    from snowflake.ml.jobs._interop import data_utils
+    from snowflake.ml.jobs._interop.utils import DEFAULT_CODEC, DEFAULT_PROTOCOL
+    from snowflake.snowpark import exceptions as sp_exceptions
+
+    try:
+        with data_utils.open_stream(function_args, "r") as stream:
+            # Load the DTO as a dict for easy fallback to legacy loading if necessary
+            data = DEFAULT_CODEC.decode(stream, as_dict=True)
+    # the exception could be OSError or BlockingIOError(the file name is too long)
+    except OSError as e:
+        # path_or_data might be inline data
+        try:
+            data = DEFAULT_CODEC.decode(io.StringIO(function_args), as_dict=True)
+        except Exception:
+            raise e
+
+    if data["protocol"] is not None:
+        try:
+            from snowflake.ml.jobs._interop.dto_schema import ProtocolInfo
+
+            protocol_info = ProtocolInfo.model_validate(data["protocol"])
+            logger.debug(f"Loading result value with protocol {protocol_info}")
+            result_value = DEFAULT_PROTOCOL.load(protocol_info, session=None, path_transform=path_transform)
+        except sp_exceptions.SnowparkSQLException:
+            raise
+    else:
+        result_value = None
+
+    return data["value"] or result_value
+
+
+def _unpack_obj_fallback(obj: Any, session: Optional[snowflake.snowpark.Session]) -> Any:
+    SESSION_KEY_PREFIX = "session@"
+
+    if not isinstance(obj, dict):
+        return obj
+    elif len(obj) == 1 and SESSION_KEY_PREFIX in obj:
+        return session
+    else:
+        type = obj.get("type@", None)
+        # If type is None, we are unpacking a dict
+        if type is None:
+            result_dict = {}
+            for k, v in obj.items():
+                if k.startswith(SESSION_KEY_PREFIX):
+                    result_key = k[len(SESSION_KEY_PREFIX) :]
+                    result_dict[result_key] = session
+                else:
+                    result_dict[k] = _unpack_obj_fallback(v, session)
+            return result_dict
+        # If type is not None, we are unpacking a tuple or list
+        else:
+            indexes = []
+            for k, _ in obj.items():
+                if "#" in k:
+                    indexes.append(int(k.split("#")[-1]))
+
+            if not indexes:
+                return tuple() if type is tuple else []
+            result_list: list[Any] = [None] * (max(indexes) + 1)
+
+            for k, v in obj.items():
+                if k == "type@":
+                    continue
+                idx = int(k.split("#")[-1])
+                if k.startswith(SESSION_KEY_PREFIX):
+                    result_list[idx] = session
+                else:
+                    result_list[idx] = _unpack_obj_fallback(v, session)
+            return tuple(result_list) if type is tuple else result_list
+
+
+def _load_function_args(
+    session: snowflake.snowpark.Session,
+    function_args: Optional[str] = None,
+) -> tuple[tuple[Any, ...], dict[str, Any]]:
+    """Load and deserialize function arguments.
+
+    Args:
+        function_args: Inline serialized function arguments or path to serialized file.
+        session: Optional Snowpark session for stage access if needed.
+
+    Returns:
+        A tuple of (positional_args, keyword_args)
+
+    """
+    if not function_args:
+        return (), {}
+
+    def path_transform(stage_path: str) -> str:
+        if not PAYLOAD_PATH:
+            return stage_path
+
+        payload_path = PurePosixPath(PAYLOAD_PATH)
+        payload_dir_name = payload_path.name  # e.g., "app"
+
+        # Parse stage path and find the payload directory
+        stage_parts = PurePosixPath(stage_path.lstrip("@")).parts
+
+        try:
+            # Find index of payload directory (e.g., "app") in stage path
+            idx = stage_parts.index(payload_dir_name)
+            # Get relative path after the payload directory
+            relative_parts = stage_parts[idx + 1 :]
+            return str(payload_path.joinpath(*relative_parts))
+        except (ValueError, IndexError):
+            # Fallback to just the filename
+            return str(payload_path / PurePosixPath(stage_path).name)
+
+    try:
+        from snowflake.ml.jobs._interop import utils as interop_utils
+
+        args, kwargs = interop_utils.load(
+            function_args,
+            session=session,
+            path_transform=path_transform,
+        )
+        return args, kwargs
+    except (AttributeError, ImportError):
+        # Backwards compatibility: load may not exist in older SnowML versions
+        packed = _load_dto_fallback(function_args, path_transform)
+        args, kwargs = _unpack_obj_fallback(packed, session)
+        return args, kwargs
+
+
+def run_script(
+    script_path: str,
+    payload_args: Optional[tuple[Any, ...]] = None,
+    payload_kwargs: Optional[dict[str, Any]] = None,
+    main_func: Optional[str] = None,
+) -> Any:
     """
     Execute a Python script and return its result.
 
     Args:
-        script_path: Path to the Python script
-        script_args: Arguments to pass to the script
-        main_func: The name of the function to call in the script (if any)
+        script_path: Path to the Python script.
+        payload_args: Positional arguments to pass to the script or entrypoint.
+        payload_kwargs: Keyword arguments to pass to the script or entrypoint.
+        main_func: The name of the function to call in the script (if any).
 
     Returns:
         Result from script execution, either from the main function or the script's __return__ value
 
     Raises:
         RuntimeError: If the specified main_func is not found or not callable
+        ValueError: If payload_kwargs is provided for runpy execution.
     """
     # Save original sys.argv and modify it for the script (applies to runpy execution only)
     original_argv = sys.argv
-    sys.argv = [script_path, *script_args]
 
     try:
         if main_func:
@@ -381,10 +521,13 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
                 raise RuntimeError(f"Function '{main_func}' not a valid entrypoint for {script_path}")
 
             # Call main function
-            result = func(*script_args)
+            result = func(*(payload_args or ()), **(payload_kwargs or {}))
             return result
         else:
-            # Use runpy for other scripts
+            if payload_kwargs:
+                raise ValueError("payload_kwargs is not supported for runpy execution; use payload_args instead")
+            # Save original sys.argv and modify it for the script.
+            sys.argv = [script_path, *(payload_args or ())]
             globals_dict = runpy.run_path(script_path, run_name="__main__")
             result = globals_dict.get("__return__", None)
             return result
@@ -393,24 +536,28 @@ def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = N
         sys.argv = original_argv
 
 
-def main(entrypoint: str, *script_args: Any, script_main_func: Optional[str] = None) -> Any:
+def main(
+    entrypoint: str,
+    session: snowflake.snowpark.Session,
+    payload_args: Optional[tuple[Any, ...]] = None,
+    payload_kwargs: Optional[dict[str, Any]] = None,
+    script_main_func: Optional[str] = None,
+) -> Any:
     """Executes a Python script and serializes the result to JOB_RESULT_PATH.
 
     Args:
         entrypoint (str): The job payload entrypoint to execute.
-        script_args (Any): Arguments to pass to the script.
+        payload_args (tuple[Any, ...], optional): Positional args to pass to the script or entrypoint.
+        payload_kwargs (dict[str, Any], optional): Keyword args to pass to the script or entrypoint.
         script_main_func (str, optional): The name of the function to call in the script (if any).
+        session (snowflake.snowpark.Session, optional): Snowpark session for stage access if needed.
 
     Returns:
         Any: The result of the script execution.
 
     Raises:
-        Exception: Re-raises any exception caught during script execution.
+        ValueError: If payload_kwargs is provided for runpy execution.
     """
-    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:
@@ -420,12 +567,6 @@ def main(entrypoint: str, *script_args: Any, script_main_func: Optional[str] = N
     except ModuleNotFoundError:
         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 = snowflake.snowpark.Session.builder.configs(config).create()  # noqa: F841
-
     execution_result_is_error = False
     execution_result_value = None
     try:
@@ -446,10 +587,21 @@ def main(entrypoint: str, *script_args: Any, script_main_func: Optional[str] = N
 
         if is_python:
             # Run as Python script
-            execution_result_value = run_script(resolved_entrypoint, *script_args, main_func=script_main_func)
+            execution_result_value = run_script(
+                resolved_entrypoint,
+                payload_args=payload_args,
+                payload_kwargs=payload_kwargs,
+                main_func=script_main_func,
+            )
         else:
             # Run as subprocess
-            run_command(resolved_entrypoint, *script_args)
+            if payload_kwargs:
+                raise ValueError("payload_kwargs is not supported for subprocesses")
+
+            run_command(
+                resolved_entrypoint,
+                *(payload_args or ()),
+            )
 
         # Log end marker for user script execution
         print(LOG_END_MSG)  # noqa: T201
@@ -487,11 +639,36 @@ if __name__ == "__main__":
     parser.add_argument(
         "--script_main_func", required=False, help="The name of the main function to call in the script"
     )
+    parser.add_argument(
+        "--function_args",
+        required=False,
+        help="Serialized function arguments or path to serialized function arguments file",
+    )
     args, unknown_args = parser.parse_known_args()
 
+    try:
+        from snowflake.ml._internal.utils.connection_params import SnowflakeLoginOptions
+    except ImportError:
+        from snowflake.ml.utils.connection_params import SnowflakeLoginOptions
+
+    # Create a Snowpark session before starting
+    # Session can be retrieved from using snowflake.snowpark.context.get_active_session()
+    # _load_function_args will use the session to load the function arguments
+    config = SnowflakeLoginOptions()
+    config["client_session_keep_alive"] = "True"
+    session = snowflake.snowpark.Session.builder.configs(config).create()  # noqa: F841
+
+    if args.function_args:
+        if args.script_args or unknown_args:
+            raise ValueError("Only one of function_args and script_args can be provided")
+        payload_args, payload_kwargs = _load_function_args(session, args.function_args)
+    else:
+        payload_args, payload_kwargs = (args.script_args + unknown_args), {}
+
     main(
         args.entrypoint,
-        *args.script_args,
-        *unknown_args,
+        session=session,
+        payload_args=payload_args,
+        payload_kwargs=payload_kwargs,
         script_main_func=args.script_main_func,
     )