snowflake-ml-python 1.25.0__tar.gz → 1.25.1__tar.gz

{snowflake_ml_python-1.25.0 → snowflake_ml_python-1.25.1}/CHANGELOG.md

@@ -1,5 +1,11 @@
 # Release History
 
+## 1.25.1
+
+### Bug Fixes
+
+* ML Job: Reverted changes related to the introduction of ML Job Definitions.
+
 ## 1.25.0
 
 ### New Features

{snowflake_ml_python-1.25.0/snowflake_ml_python.egg-info → snowflake_ml_python-1.25.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowflake-ml-python
-Version: 1.25.0
+Version: 1.25.1
 Summary: The machine learning client library that is used for interacting with Snowflake to build machine learning solutions.
 Author-email: "Snowflake, Inc" <support@snowflake.com>
 License: 
@@ -417,6 +417,12 @@ NOTE: Version 1.7.0 is used as example here. Please choose the the latest versio
 
 # Release History
 
+## 1.25.1
+
+### Bug Fixes
+
+* ML Job: Reverted changes related to the introduction of ML Job Definitions.
+
 ## 1.25.0
 
 ### New Features

{snowflake_ml_python-1.25.0 → snowflake_ml_python-1.25.1}/snowflake/ml/_internal/utils/mixins.py

@@ -9,7 +9,6 @@ _SESSION_ACCOUNT_KEY = "session$account"
 _SESSION_ROLE_KEY = "session$role"
 _SESSION_DATABASE_KEY = "session$database"
 _SESSION_SCHEMA_KEY = "session$schema"
-_SESSION_STATE_ATTR = "_session_state"
 
 
 def _identifiers_match(saved: Optional[str], current: Optional[str]) -> bool:
@@ -62,7 +61,7 @@ class SerializableSessionMixin:
         else:
             self.__dict__.update(state)
 
-        setattr(self, _SESSION_STATE_ATTR, session_state)
+        self._set_session(session_state)
 
     def _set_session(self, session_state: _SessionState) -> None:
 
@@ -87,27 +86,3 @@ class SerializableSessionMixin:
                     ),
                 ),
             )
-
-    @property
-    def session(self) -> Optional[snowpark_session.Session]:
-        if _SESSION_KEY not in self.__dict__:
-            session_state = getattr(self, _SESSION_STATE_ATTR, None)
-            if session_state is not None:
-                self._set_session(session_state)
-        return self.__dict__.get(_SESSION_KEY)
-
-    @session.setter
-    def session(self, value: Optional[snowpark_session.Session]) -> None:
-        self.__dict__[_SESSION_KEY] = value
-
-    # _getattr__ is only called when an attribute is NOT found through normal lookup.
-    # 1. Data descriptors (like @property with setter) from the class hierarchy
-    # 2. Instance __dict__ (e.g., self.x = 10)
-    # 3. Non-data descriptors (methods, `@property without setter) from the class hierarchy
-    # __getattr__ — only called if steps 1-3 all fail
-    def __getattr__(self, name: str) -> Any:
-        if name == _SESSION_KEY:
-            return self.session
-        if hasattr(super(), "__getattr__"):
-            return super().__getattr__(name)  # type: ignore[misc]
-        raise AttributeError(f"{type(self).__name__!s} object has no attribute {name!r}")

{snowflake_ml_python-1.25.0 → snowflake_ml_python-1.25.1}/snowflake/ml/jobs/_interop/data_utils.py

@@ -31,7 +31,7 @@ class StageFileWriter(io.IOBase):
             # Only upload if buffer has content and no exception occurred
             if write_contents and self._buffer.tell() > 0:
                 self._buffer.seek(0)
-                self._session.file.put_stream(self._buffer, self._path, auto_compress=False)
+                self._session.file.put_stream(self._buffer, self._path)
             self._buffer.close()
             self._closed = True
 
@@ -84,15 +84,15 @@ class DtoCodec(Protocol):
 
     @overload
     @staticmethod
-    def decode(stream: io.IOBase, as_dict: Literal[False] = False) -> dto_schema.PayloadDTO:
+    def decode(stream: io.IOBase, as_dict: Literal[False] = False) -> dto_schema.ResultDTO:
         ...
 
     @staticmethod
-    def decode(stream: io.IOBase, as_dict: bool = False) -> Union[dto_schema.PayloadDTO, dict[str, Any]]:
+    def decode(stream: io.IOBase, as_dict: bool = False) -> Union[dto_schema.ResultDTO, dict[str, Any]]:
         pass
 
     @staticmethod
-    def encode(dto: dto_schema.PayloadDTO) -> bytes:
+    def encode(dto: dto_schema.ResultDTO) -> bytes:
         pass
 
 
@@ -104,18 +104,18 @@ class JsonDtoCodec(DtoCodec):
 
     @overload
     @staticmethod
-    def decode(stream: io.IOBase, as_dict: Literal[False] = False) -> dto_schema.PayloadDTO:
+    def decode(stream: io.IOBase, as_dict: Literal[False] = False) -> dto_schema.ResultDTO:
         ...
 
     @staticmethod
-    def decode(stream: io.IOBase, as_dict: bool = False) -> Union[dto_schema.PayloadDTO, dict[str, Any]]:
+    def decode(stream: io.IOBase, as_dict: bool = False) -> Union[dto_schema.ResultDTO, dict[str, Any]]:
         data = cast(dict[str, Any], json.load(stream))
         if as_dict:
             return data
-        return dto_schema.ResultDTOAdapter.validate_python(data)
+        return dto_schema.ResultDTO.model_validate(data)
 
     @staticmethod
-    def encode(dto: dto_schema.PayloadDTO) -> bytes:
+    def encode(dto: dto_schema.ResultDTO) -> bytes:
         # Temporarily extract the value to avoid accidentally applying model_dump() on it
         result_value = dto.value
         dto.value = None  # Clear value to avoid serializing it in the model_dump

{snowflake_ml_python-1.25.0 → snowflake_ml_python-1.25.1}/snowflake/ml/jobs/_interop/dto_schema.py

@@ -1,7 +1,7 @@
-from typing import Any, Literal, Optional, Union
+from typing import Any, Optional, Union
 
-from pydantic import BaseModel, Discriminator, Tag, TypeAdapter, model_validator
-from typing_extensions import Annotated, NotRequired, TypedDict
+from pydantic import BaseModel, model_validator
+from typing_extensions import NotRequired, TypedDict
 
 
 class BinaryManifest(TypedDict):
@@ -67,47 +67,22 @@ class ExceptionMetadata(ResultMetadata):
     traceback: str
 
 
-class PayloadDTO(BaseModel):
-    """
-    Base class for serializable payloads.
-
-    Args:
-        kind: Discriminator field for DTO type dispatch.
-        value: The payload value (if JSON-serializable).
-        protocol: The protocol used to serialize the payload (if not JSON-serializable).
-    """
-
-    kind: Literal["base"] = "base"
-    value: Optional[Any] = None
-    protocol: Optional[ProtocolInfo] = None
-    serialize_error: Optional[str] = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_fields(cls, data: Any) -> Any:
-        """Ensure at least one of value or protocol keys is specified."""
-        if cls is PayloadDTO and isinstance(data, dict):
-            required_fields = {"value", "protocol"}
-            if not any(field in data for field in required_fields):
-                raise ValueError("At least one of 'value' or 'protocol' must be specified")
-        return data
-
-
-class ResultDTO(PayloadDTO):
+class ResultDTO(BaseModel):
     """
     A JSON representation of an execution result.
 
     Args:
-        kind: Discriminator field for DTO type dispatch.
         success: Whether the execution was successful.
         value: The value of the execution or the exception if the execution failed.
         protocol: The protocol used to serialize the result.
         metadata: The metadata of the result.
     """
 
-    kind: Literal["result"] = "result"  # type: ignore[assignment]
     success: bool
+    value: Optional[Any] = None
+    protocol: Optional[ProtocolInfo] = None
     metadata: Optional[Union[ResultMetadata, ExceptionMetadata]] = None
+    serialize_error: Optional[str] = None
 
     @model_validator(mode="before")
     @classmethod
@@ -118,23 +93,3 @@ class ResultDTO(PayloadDTO):
             if not any(field in data for field in required_fields):
                 raise ValueError("At least one of 'value', 'protocol', or 'metadata' must be specified")
         return data
-
-
-def _get_dto_kind(data: Any) -> str:
-    """Extract the 'kind' discriminator from input, defaulting to 'result' for backward compatibility."""
-    if isinstance(data, dict):
-        kind = data.get("kind", "result")
-    else:
-        kind = getattr(data, "kind", "result")
-    return str(kind)
-
-
-AnyResultDTO = Annotated[
-    Union[
-        Annotated[ResultDTO, Tag("result")],
-        Annotated[PayloadDTO, Tag("base")],
-    ],
-    Discriminator(_get_dto_kind),
-]
-
-ResultDTOAdapter: TypeAdapter[AnyResultDTO] = TypeAdapter(AnyResultDTO)

{snowflake_ml_python-1.25.0 → snowflake_ml_python-1.25.1}/snowflake/ml/jobs/_interop/protocols.py

@@ -17,8 +17,6 @@ Condition = Union[type, tuple[type, ...], Callable[[Any], bool], None]
 
 logger = logging.getLogger(__name__)
 
-SESSION_KEY_PREFIX = "session@"
-
 
 class SerializationError(TypeError):
     """Exception raised when a serialization protocol fails."""
@@ -138,10 +136,9 @@ class CloudPickleProtocol(SerializationProtocol):
 
     def save(self, obj: Any, dest_dir: str, session: Optional[snowpark.Session] = None) -> ProtocolInfo:
         """Save the object to the destination directory."""
-        replaced_obj = self._pack_obj(obj)
         result_path = posixpath.join(dest_dir, self.DEFAULT_PATH)
         with data_utils.open_stream(result_path, "wb", session=session) as f:
-            self._backend.dump(replaced_obj, f)
+            self._backend.dump(obj, f)
         manifest: BinaryManifest = {"path": result_path}
         return self.protocol_info.with_manifest(manifest)
 
@@ -160,15 +157,12 @@ class CloudPickleProtocol(SerializationProtocol):
         payload_manifest = cast(BinaryManifest, payload_info.manifest)
         try:
             if payload_bytes := payload_manifest.get("bytes"):
-                result = self._backend.loads(payload_bytes)
-            elif payload_b64 := payload_manifest.get("base64"):
-                result = self._backend.loads(base64.b64decode(payload_b64))
-            else:
-                result_path = path_transform(payload_manifest["path"]) if path_transform else payload_manifest["path"]
-                with data_utils.open_stream(result_path, "rb", session=session) as f:
-                    result = self._backend.load(f)
-
-            return self._unpack_obj(result, session=session)
+                return self._backend.loads(payload_bytes)
+            if payload_b64 := payload_manifest.get("base64"):
+                return self._backend.loads(base64.b64decode(payload_b64))
+            result_path = path_transform(payload_manifest["path"]) if path_transform else payload_manifest["path"]
+            with data_utils.open_stream(result_path, "rb", session=session) as f:
+                return self._backend.load(f)
         except (
             pickle.UnpicklingError,
             TypeError,
@@ -179,117 +173,6 @@ class CloudPickleProtocol(SerializationProtocol):
                 raise error from pickle_error
             raise
 
-    def _pack_obj(self, obj: Any) -> Any:
-        """Pack objects into JSON-safe dicts using reserved marker keys.
-
-        Markers:
-        - "type@": container type for list/tuple (list or tuple)
-        - "#<i>": positional element for list/tuple at index i
-        - "session@": placeholder for snowpark.Session values
-          - "session@#<i>" for list/tuple entries
-          - "session@<key>" for dict entries
-          - {"session@": None} for a bare Session object
-
-        Example:
-            obj = {"x": [1, session], "s": session}
-            packed = {
-                "x": {"type@": list, "#0": 1, "session@#1": None},
-                "session@s": None,
-            }
-            _unpack_obj(packed, session) == obj
-
-        Args:
-            obj: Object to pack into JSON-safe marker dictionaries.
-
-        Returns:
-            Packed representation with markers for session references.
-        """
-        arguments: dict[str, Any] = {}
-        if isinstance(obj, tuple) or isinstance(obj, list):
-            arguments = {"type@": type(obj)}
-            for i, arg in enumerate(obj):
-                if isinstance(arg, snowpark.Session):
-                    arguments[f"{SESSION_KEY_PREFIX}#{i}"] = None
-                else:
-                    arguments[f"#{i}"] = self._pack_obj(arg)
-            return arguments
-        elif isinstance(obj, dict):
-            for k, v in obj.items():
-                if isinstance(v, snowpark.Session):
-                    arguments[f"{SESSION_KEY_PREFIX}{k}"] = None
-                else:
-                    arguments[k] = self._pack_obj(v)
-            return arguments
-        elif isinstance(obj, snowpark.Session):
-            # Box session into a dict marker so we can distinguish it from other plain objects.
-            arguments[f"{SESSION_KEY_PREFIX}"] = None
-            return arguments
-        else:
-            return obj
-
-    def _unpack_obj(self, obj: Any, session: Optional[snowpark.Session] = None) -> Any:
-        """Unpack dict markers back into containers and Session references.
-
-        Markers:
-        - "type@": container type for list/tuple (list or tuple)
-        - "#<i>": positional element for list/tuple at index i
-        - "session@": placeholder for snowpark.Session values
-          - "session@#<i>" for list/tuple entries
-          - "session@<key>" for dict entries
-          - {"session@": None} for a bare Session object
-
-        Example:
-            packed = {
-                "x": {"type@": list, "#0": 1, "session@#1": None},
-                "session@s": None,
-            }
-            obj = _unpack_obj(packed, session)
-            # obj == {"x": [1, session], "s": session}
-
-        Args:
-            obj: Packed object with marker dictionaries.
-            session: Session to inject for session markers.
-
-        Returns:
-            Unpacked object with session references restored.
-        """
-        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] = self._unpack_obj(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] = self._unpack_obj(v, session)
-                return tuple(result_list) if type is tuple else result_list
-
 
 class ArrowTableProtocol(SerializationProtocol):
     """

snowflake_ml_python-1.25.1/snowflake/ml/jobs/_interop/utils.py

@@ -0,0 +1,144 @@
+import logging
+import os
+import traceback
+from pathlib import PurePath
+from typing import Any, Callable, Optional
+
+import pydantic
+
+from snowflake import snowpark
+from snowflake.ml.jobs._interop import data_utils, exception_utils, legacy, protocols
+from snowflake.ml.jobs._interop.dto_schema import (
+    ExceptionMetadata,
+    ResultDTO,
+    ResultMetadata,
+)
+from snowflake.ml.jobs._interop.results import ExecutionResult, LoadedExecutionResult
+from snowflake.snowpark import exceptions as sp_exceptions
+
+DEFAULT_CODEC = data_utils.JsonDtoCodec
+DEFAULT_PROTOCOL = protocols.AutoProtocol()
+DEFAULT_PROTOCOL.try_register_protocol(protocols.CloudPickleProtocol)
+DEFAULT_PROTOCOL.try_register_protocol(protocols.ArrowTableProtocol)
+DEFAULT_PROTOCOL.try_register_protocol(protocols.PandasDataFrameProtocol)
+DEFAULT_PROTOCOL.try_register_protocol(protocols.NumpyArrayProtocol)
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_result(result: ExecutionResult, path: str, session: Optional[snowpark.Session] = None) -> None:
+    """
+    Save the result to a file.
+    """
+    result_dto = ResultDTO(
+        success=result.success,
+        value=result.value,
+    )
+
+    try:
+        # Try to encode result directly
+        payload = DEFAULT_CODEC.encode(result_dto)
+    except TypeError:
+        result_dto.value = None  # Remove raw value to avoid serialization error
+        result_dto.metadata = _get_metadata(result.value)  # Add metadata for client fallback on protocol mismatch
+        try:
+            path_dir = PurePath(path).parent.as_posix()
+            protocol_info = DEFAULT_PROTOCOL.save(result.value, path_dir, session=session)
+            result_dto.protocol = protocol_info
+
+        except Exception as e:
+            logger.warning(f"Error dumping result value: {repr(e)}")
+            result_dto.serialize_error = repr(e)
+
+        # Encode the modified result DTO
+        payload = DEFAULT_CODEC.encode(result_dto)
+
+    with data_utils.open_stream(path, "wb", session=session) as stream:
+        stream.write(payload)
+
+
+def load_result(
+    path: str, session: Optional[snowpark.Session] = None, path_transform: Optional[Callable[[str], str]] = None
+) -> ExecutionResult:
+    """Load the result from a file on a Snowflake stage."""
+    try:
+        with data_utils.open_stream(path, "r", session=session) as stream:
+            # Load the DTO as a dict for easy fallback to legacy loading if necessary
+            dto_dict = DEFAULT_CODEC.decode(stream, as_dict=True)
+    except UnicodeDecodeError:
+        # Path may be a legacy result file (cloudpickle)
+        # TODO: Re-use the stream
+        assert session is not None
+        return legacy.load_legacy_result(session, path)
+
+    try:
+        dto = ResultDTO.model_validate(dto_dict)
+    except pydantic.ValidationError as e:
+        if "success" in dto_dict:
+            assert session is not None
+            if path.endswith(".json"):
+                path = os.path.splitext(path)[0] + ".pkl"
+            return legacy.load_legacy_result(session, path, result_json=dto_dict)
+        raise ValueError("Invalid result schema") from e
+
+    # Try loading data from file using the protocol info
+    result_value = None
+    data_load_error = None
+    if dto.protocol is not None:
+        try:
+            logger.debug(f"Loading result value with protocol {dto.protocol}")
+            result_value = DEFAULT_PROTOCOL.load(dto.protocol, session=session, path_transform=path_transform)
+        except sp_exceptions.SnowparkSQLException:
+            raise  # Data retrieval errors should be bubbled up
+        except Exception as e:
+            logger.debug(f"Error loading result value with protocol {dto.protocol}: {repr(e)}")
+            data_load_error = e
+
+    # Wrap serialize_error in a TypeError
+    if dto.serialize_error:
+        serialize_error = TypeError("Original result serialization failed with error: " + dto.serialize_error)
+        if data_load_error:
+            data_load_error.__context__ = serialize_error
+        else:
+            data_load_error = serialize_error
+
+    # Prepare to assemble the final result
+    result_value = result_value if result_value is not None else dto.value
+    if not dto.success and result_value is None:
+        # Try to reconstruct exception from metadata if available
+        if isinstance(dto.metadata, ExceptionMetadata):
+            logger.debug(f"Reconstructing exception from metadata {dto.metadata}")
+            result_value = exception_utils.build_exception(
+                type_str=dto.metadata.type,
+                message=dto.metadata.message,
+                traceback=dto.metadata.traceback,
+                original_repr=dto.metadata.repr,
+            )
+
+        # Generate a generic error if we still don't have a value,
+        # attaching the data load error if any
+        if result_value is None:
+            result_value = exception_utils.RemoteError("Unknown remote error")
+            result_value.__cause__ = data_load_error
+
+    return LoadedExecutionResult(
+        success=dto.success,
+        value=result_value,
+        load_error=data_load_error,
+    )
+
+
+def _get_metadata(value: Any) -> ResultMetadata:
+    type_name = f"{type(value).__module__}.{type(value).__name__}"
+    if isinstance(value, BaseException):
+        return ExceptionMetadata(
+            type=type_name,
+            repr=repr(value),
+            message=str(value),
+            traceback="".join(traceback.format_tb(value.__traceback__)),
+        )
+    return ResultMetadata(
+        type=type_name,
+        repr=repr(value),
+    )

{snowflake_ml_python-1.25.0 → snowflake_ml_python-1.25.1}/snowflake/ml/jobs/_utils/constants.py

@@ -5,7 +5,6 @@ from snowflake.ml.jobs._utils.types import ComputeResources
 DEFAULT_CONTAINER_NAME = "main"
 MEMORY_VOLUME_NAME = "dshm"
 STAGE_VOLUME_NAME = "stage-volume"
-RESULT_VOLUME_NAME = "result-volume"
 DEFAULT_PYTHON_VERSION = "3.10"
 
 # Environment variables
@@ -110,6 +109,3 @@ CLOUD_INSTANCE_FAMILIES = {
     SnowflakeCloudType.AZURE: AZURE_INSTANCE_FAMILIES,
     SnowflakeCloudType.GCP: GCP_INSTANCE_FAMILIES,
 }
-
-# Magic attributes
-IS_MLJOB_REMOTE_ATTR = "_is_mljob_remote_callable"

snowflake_ml_python-1.25.1/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_python-1.25.0 → snowflake_ml_python-1.25.1}/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()}"