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

snowflake/ml/jobs/_utils/payload_utils.py

@@ -17,12 +17,20 @@ import cloudpickle as cp
 from packaging import version
 
 from snowflake import snowpark
-from snowflake.ml.jobs._utils import constants, query_helper, stage_utils, types
+from snowflake.ml.jobs._utils import (
+    constants,
+    function_payload_utils,
+    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}
@@ -553,6 +561,7 @@ 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(
@@ -682,9 +691,14 @@ def _generate_param_handler_code(signature: inspect.Signature, output_name: str
     return param_code
 
 
-def generate_python_code(function: 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."""
 
+    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}
@@ -697,7 +711,7 @@ def generate_python_code(function: Callable[..., Any], source_code_display: bool
     source_code_comment = _generate_source_code_comment(function) if source_code_display else ""
 
     arg_dict_name = "kwargs"
-    if getattr(function, 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)
@@ -707,7 +721,7 @@ import pickle
 
 try:
     {textwrap.indent(source_code_comment, '    ')}
-    {_ENTRYPOINT_FUNC_NAME} = pickle.loads(bytes.fromhex('{_serialize_callable(function).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(
@@ -733,6 +747,26 @@ 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)):
@@ -741,7 +775,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 callable(source):
-        return f"{source.__name__}"
+    elif isinstance(source, function_payload_utils.FunctionPayload):
+        return f"{source.function.__name__}"
     else:
         return f"{JOB_ID_PREFIX}{str(uuid4()).replace('-', '_').upper()}"

snowflake/ml/jobs/_utils/runtime_env_utils.py

@@ -1,18 +1,117 @@
-from typing import Optional, cast
+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 snowflake import snowpark
-from snowflake.ml.jobs._utils import query_helper
+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.
 
+        Args:
+            data: The input data to filter, expected to be a dictionary.
 
-def get_runtime_image(
-    session: snowpark.Session, compute_pool: str, runtime_environment: Optional[str] = None
+        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
 ) -> Optional[str]:
-    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}')")
+    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()")
     if not rows:
-        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
+        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

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

@@ -1,7 +1,6 @@
 import argparse
 import copy
 import importlib.util
-import io
 import json
 import logging
 import math
@@ -13,22 +12,15 @@ import sys
 import time
 import traceback
 import zipfile
-from pathlib import Path, PurePosixPath
-from typing import Any, Callable, Optional
+from pathlib import Path
+from typing import Any, 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")
-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)
+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):
@@ -355,156 +347,24 @@ def wait_for_instances(
     )
 
 
-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:
+def run_script(script_path: str, *script_args: Any, main_func: Optional[str] = None) -> Any:
     """
     Execute a Python script and return its result.
 
     Args:
-        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).
+        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)
 
     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:
@@ -521,13 +381,10 @@ def run_script(
                 raise RuntimeError(f"Function '{main_func}' not a valid entrypoint for {script_path}")
 
             # Call main function
-            result = func(*(payload_args or ()), **(payload_kwargs or {}))
+            result = func(*script_args)
             return result
         else:
-            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 ())]
+            # Use runpy for other scripts
             globals_dict = runpy.run_path(script_path, run_name="__main__")
             result = globals_dict.get("__return__", None)
             return result
@@ -536,28 +393,24 @@ def run_script(
         sys.argv = original_argv
 
 
-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:
+def main(entrypoint: str, *script_args: Any, 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.
-        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_args (Any): Arguments to pass to the script.
         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:
-        ValueError: If payload_kwargs is provided for runpy execution.
+        Exception: Re-raises any exception caught during script 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:
@@ -567,6 +420,12 @@ def main(
     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:
@@ -587,21 +446,10 @@ def main(
 
         if is_python:
             # Run as Python script
-            execution_result_value = run_script(
-                resolved_entrypoint,
-                payload_args=payload_args,
-                payload_kwargs=payload_kwargs,
-                main_func=script_main_func,
-            )
+            execution_result_value = run_script(resolved_entrypoint, *script_args, main_func=script_main_func)
         else:
             # Run as subprocess
-            if payload_kwargs:
-                raise ValueError("payload_kwargs is not supported for subprocesses")
-
-            run_command(
-                resolved_entrypoint,
-                *(payload_args or ()),
-            )
+            run_command(resolved_entrypoint, *script_args)
 
         # Log end marker for user script execution
         print(LOG_END_MSG)  # noqa: T201
@@ -639,36 +487,11 @@ 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,
-        session=session,
-        payload_args=payload_args,
-        payload_kwargs=payload_kwargs,
+        *args.script_args,
+        *unknown_args,
         script_main_func=args.script_main_func,
     )

snowflake/ml/jobs/_utils/spec_utils.py

@@ -0,0 +1,22 @@
+from snowflake import snowpark
+from snowflake.ml._internal.utils import snowflake_env
+from snowflake.ml.jobs._utils import constants, query_helper, types
+
+
+def _get_node_resources(session: snowpark.Session, compute_pool: str) -> types.ComputeResources:
+    """Extract resource information for the specified compute pool"""
+    # Get the instance family
+    rows = query_helper.run_query(
+        session,
+        "show compute pools like ?",
+        params=[compute_pool],
+    )
+    if not rows:
+        raise ValueError(f"Compute pool '{compute_pool}' not found")
+    instance_family: str = rows[0]["instance_family"]
+    cloud = snowflake_env.get_current_cloud(session, default=snowflake_env.SnowflakeCloudType.AWS)
+
+    return (
+        constants.COMMON_INSTANCE_FAMILIES.get(instance_family)
+        or constants.CLOUD_INSTANCE_FAMILIES[cloud][instance_family]
+    )

snowflake/ml/jobs/decorators.py

@@ -1,12 +1,13 @@
 import copy
+import functools
 from typing import Any, Callable, Optional, TypeVar
 
 from typing_extensions import ParamSpec
 
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
-from snowflake.ml.jobs import job_definition as jd
-from snowflake.ml.jobs._utils import arg_protocol, constants
+from snowflake.ml.jobs import job as jb, manager as jm
+from snowflake.ml.jobs._utils import payload_utils
 
 _PROJECT = "MLJob"
 
@@ -24,7 +25,7 @@ def remote(
     external_access_integrations: Optional[list[str]] = None,
     session: Optional[snowpark.Session] = None,
     **kwargs: Any,
-) -> Callable[[Callable[_Args, _ReturnValue]], jd.MLJobDefinition[_Args, _ReturnValue]]:
+) -> Callable[[Callable[_Args, _ReturnValue]], Callable[_Args, jb.MLJob[_ReturnValue]]]:
     """
     Submit a job to the compute pool.
 
@@ -50,25 +51,29 @@ def remote(
         Decorator that dispatches invocations of the decorated function as remote jobs.
     """
 
-    def decorator(func: Callable[_Args, _ReturnValue]) -> jd.MLJobDefinition[_Args, _ReturnValue]:
+    def decorator(func: Callable[_Args, _ReturnValue]) -> Callable[_Args, jb.MLJob[_ReturnValue]]:
         # Copy the function to avoid modifying the original
         # We need to modify the line number of the function to exclude the
         # decorator from the copied source code
         wrapped_func = copy.copy(func)
         wrapped_func.__code__ = wrapped_func.__code__.replace(co_firstlineno=func.__code__.co_firstlineno + 1)
 
-        setattr(wrapped_func, constants.IS_MLJOB_REMOTE_ATTR, True)
-        return jd.MLJobDefinition.register(
-            source=wrapped_func,
-            compute_pool=compute_pool,
-            stage_name=stage_name,
-            target_instances=target_instances,
-            pip_requirements=pip_requirements,
-            external_access_integrations=external_access_integrations,
-            session=session or snowpark.context.get_active_session(),
-            arg_protocol=arg_protocol.ArgProtocol.PICKLE,
-            generate_suffix=True,
-            **kwargs,
-        )
+        @functools.wraps(func)
+        def wrapper(*_args: _Args.args, **_kwargs: _Args.kwargs) -> jb.MLJob[_ReturnValue]:
+            payload = payload_utils.create_function_payload(func, *_args, **_kwargs)
+            job = jm._submit_job(
+                source=payload,
+                stage_name=stage_name,
+                compute_pool=compute_pool,
+                target_instances=target_instances,
+                pip_requirements=pip_requirements,
+                external_access_integrations=external_access_integrations,
+                session=payload.session or session,
+                **kwargs,
+            )
+            assert isinstance(job, jb.MLJob), f"Unexpected job type: {type(job)}"
+            return job
+
+        return wrapper
 
     return decorator