snowflake-ml-python 1.20.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

snowflake/ml/jobs/_utils/payload_utils.py

@@ -3,14 +3,14 @@ import importlib
 import inspect
 import io
 import itertools
-import keyword
 import logging
 import pickle
 import sys
 import textwrap
 from importlib.abc import Traversable
 from pathlib import Path, PurePath
-from typing import Any, Callable, Optional, Union, cast, get_args, get_origin
+from types import ModuleType
+from typing import IO, Any, Callable, Optional, Union, cast, get_args, get_origin
 
 import cloudpickle as cp
 from packaging import version
@@ -25,279 +25,109 @@ from snowflake.ml.jobs._utils import (
 )
 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}
 _SUPPORTED_ENTRYPOINT_EXTENSIONS = {".py"}
 _ENTRYPOINT_FUNC_NAME = "func"
 _STARTUP_SCRIPT_PATH = PurePath("startup.sh")
-_STARTUP_SCRIPT_CODE = textwrap.dedent(
-    f"""
-    #!/bin/bash
-
-    ##### Get system scripts directory #####
-    SYSTEM_DIR=$(cd "$(dirname "$0")" && pwd)
-
-    ##### Perform common set up steps #####
-    set -e # exit if a command fails
-
-    echo "Creating log directories..."
-    mkdir -p /var/log/managedservices/user/mlrs
-    mkdir -p /var/log/managedservices/system/mlrs
-    mkdir -p /var/log/managedservices/system/ray
-
-    echo "*/1 * * * * root /etc/ray_copy_cron.sh" >> /etc/cron.d/ray_copy_cron
-    echo "" >> /etc/cron.d/ray_copy_cron
-    chmod 744 /etc/cron.d/ray_copy_cron
-
-    service cron start
-
-    mkdir -p /tmp/prometheus-multi-dir
-
-    # Change directory to user payload directory
-    if [ -n "${constants.PAYLOAD_DIR_ENV_VAR}" ]; then
-        cd ${constants.STAGE_MOUNT_PATH_ENV_VAR}/${constants.PAYLOAD_DIR_ENV_VAR}
-    fi
-
-    ##### Set up Python environment #####
-    export PYTHONPATH=/opt/env/site-packages/
-    MLRS_SYSTEM_REQUIREMENTS_FILE=${{MLRS_SYSTEM_REQUIREMENTS_FILE:-"${{SYSTEM_DIR}}/requirements.txt"}}
-
-    if [ -f "${{MLRS_SYSTEM_REQUIREMENTS_FILE}}" ]; then
-        echo "Installing packages from $MLRS_SYSTEM_REQUIREMENTS_FILE"
-        if ! pip install --no-index -r $MLRS_SYSTEM_REQUIREMENTS_FILE; then
-            echo "Offline install failed, falling back to regular pip install"
-            pip install -r $MLRS_SYSTEM_REQUIREMENTS_FILE
-        fi
-    fi
-
-    MLRS_REQUIREMENTS_FILE=${{MLRS_REQUIREMENTS_FILE:-"requirements.txt"}}
-    if [ -f "${{MLRS_REQUIREMENTS_FILE}}" ]; then
-        # TODO: Prevent collisions with MLRS packages using virtualenvs
-        echo "Installing packages from $MLRS_REQUIREMENTS_FILE"
-        pip install -r $MLRS_REQUIREMENTS_FILE
-    fi
-
-    MLRS_CONDA_ENV_FILE=${{MLRS_CONDA_ENV_FILE:-"environment.yml"}}
-    if [ -f "${{MLRS_CONDA_ENV_FILE}}" ]; then
-        # TODO: Handle conda environment
-        echo "Custom conda environments not currently supported"
-        exit 1
-    fi
-    ##### End Python environment setup #####
-
-    ##### Ray configuration #####
-    shm_size=$(df --output=size --block-size=1 /dev/shm | tail -n 1)
-
-    # Check if the local get_instance_ip.py script exists
-    HELPER_EXISTS=$(
-        [ -f "${{SYSTEM_DIR}}/get_instance_ip.py" ] && echo "true" || echo "false"
-    )
-
 
-    # Configure IP address and logging directory
-    if [ "$HELPER_EXISTS" = "true" ]; then
-        eth0Ip=$(python3 "${{SYSTEM_DIR}}/get_instance_ip.py" \
-            "$SNOWFLAKE_SERVICE_NAME" --instance-index=-1)
-    else
-        eth0Ip=$(ifconfig eth0 2>/dev/null | sed -En -e 's/.*inet ([0-9.]+).*/\1/p')
-    fi
-    log_dir="/tmp/ray"
-
-    # Check if eth0Ip is a valid IP address and fall back to default if necessary
-    if [[ ! $eth0Ip =~ ^[0-9]+\\.[0-9]+\\.[0-9]+\\.[0-9]+$ ]]; then
-        eth0Ip="127.0.0.1"
-    fi
-
-    # Get the environment values of SNOWFLAKE_JOBS_COUNT and SNOWFLAKE_JOB_INDEX for batch jobs
-    # These variables don't exist for non-batch jobs, so set defaults
-    if [ -z "$SNOWFLAKE_JOBS_COUNT" ]; then
-        SNOWFLAKE_JOBS_COUNT=1
-    fi
-
-    if [ -z "$SNOWFLAKE_JOB_INDEX" ]; then
-        SNOWFLAKE_JOB_INDEX=0
-    fi
-
-    # Determine if it should be a worker or a head node for batch jobs
-    if [[ "$SNOWFLAKE_JOBS_COUNT" -gt 1 && "$HELPER_EXISTS" = "true" ]]; then
-        head_info=$(python3 "${{SYSTEM_DIR}}/get_instance_ip.py" "$SNOWFLAKE_SERVICE_NAME" --head)
-        if [ $? -eq 0 ]; then
-            # Parse the output using read
-            read head_index head_ip head_status<<< "$head_info"
-
-            if [ "$SNOWFLAKE_JOB_INDEX" -ne "$head_index" ]; then
-                NODE_TYPE="worker"
-                echo "{constants.LOG_START_MSG}"
-            fi
-
-            # Use the parsed variables
-            echo "Head Instance Index: $head_index"
-            echo "Head Instance IP: $head_ip"
-            echo "Head Instance Status: $head_status"
-
-            # If the head status is not "READY" or "PENDING", exit early
-            if [ "$head_status" != "READY" ] && [ "$head_status" != "PENDING" ]; then
-                echo "Head instance status is not READY or PENDING. Exiting."
-                exit 0
-            fi
-
-        else
-            echo "Error: Failed to get head instance information."
-            echo "$head_info" # Print the error message
-            exit 1
-        fi
-
-
-    fi
-
-    # Common parameters for both head and worker nodes
-    common_params=(
-        "--node-ip-address=$eth0Ip"
-        "--object-manager-port=${{RAY_OBJECT_MANAGER_PORT:-12011}}"
-        "--node-manager-port=${{RAY_NODE_MANAGER_PORT:-12012}}"
-        "--runtime-env-agent-port=${{RAY_RUNTIME_ENV_AGENT_PORT:-12013}}"
-        "--dashboard-agent-grpc-port=${{RAY_DASHBOARD_AGENT_GRPC_PORT:-12014}}"
-        "--dashboard-agent-listen-port=${{RAY_DASHBOARD_AGENT_LISTEN_PORT:-12015}}"
-        "--min-worker-port=${{RAY_MIN_WORKER_PORT:-12031}}"
-        "--max-worker-port=${{RAY_MAX_WORKER_PORT:-13000}}"
-        "--metrics-export-port=11502"
-        "--temp-dir=$log_dir"
-        "--disable-usage-stats"
-    )
-
-    if [ "$NODE_TYPE" = "worker" ]; then
-        # Use head_ip as head address if it exists
-        if [ ! -z "$head_ip" ]; then
-            RAY_HEAD_ADDRESS="$head_ip"
-        fi
-
-        # If RAY_HEAD_ADDRESS is still empty, exit with an error
-        if [ -z "$RAY_HEAD_ADDRESS" ]; then
-            echo "Error: Failed to determine head node address using default instance-index=0"
-            exit 1
-        fi
-
-        if [ -z "$SERVICE_NAME" ]; then
-            SERVICE_NAME="$SNOWFLAKE_SERVICE_NAME"
-        fi
-
-        if [ -z "$RAY_HEAD_ADDRESS" ] || [ -z "$SERVICE_NAME" ]; then
-            echo "Error: RAY_HEAD_ADDRESS and SERVICE_NAME must be set."
-            exit 1
-        fi
-
-        # Additional worker-specific parameters
-        worker_params=(
-            "--address=${{RAY_HEAD_ADDRESS}}:12001"       # Connect to head node
-            "--resources={{\\"${{SERVICE_NAME}}\\":1, \\"node_tag:worker\\":1}}"  # Tag for node identification
-            "--object-store-memory=${{shm_size}}"
-        )
 
-        # Start Ray on a worker node - run in background
-        ray start "${{common_params[@]}}" "${{worker_params[@]}}" -v --block &
-
-        echo "Worker node started on address $eth0Ip. See more logs in the head node."
-
-        echo "{constants.LOG_END_MSG}"
-
-        # Start the worker shutdown listener in the background
-        echo "Starting worker shutdown listener..."
-        python "${{SYSTEM_DIR}}/worker_shutdown_listener.py"
-        WORKER_EXIT_CODE=$?
-
-        echo "Worker shutdown listener exited with code $WORKER_EXIT_CODE"
-        exit $WORKER_EXIT_CODE
-    else
-        # Additional head-specific parameters
-        head_params=(
-            "--head"
-            "--port=${{RAY_HEAD_GCS_PORT:-12001}}"                                  # Port of Ray (GCS server)
-            "--ray-client-server-port=${{RAY_HEAD_CLIENT_SERVER_PORT:-10001}}"      # Rort for Ray Client Server
-            "--dashboard-host=${{NODE_IP_ADDRESS}}"                                 # Host to bind the dashboard server
-            "--dashboard-grpc-port=${{RAY_HEAD_DASHBOARD_GRPC_PORT:-12002}}"        # Dashboard head to listen for grpc
-            "--dashboard-port=${{DASHBOARD_PORT}}"                  # Port to bind the dashboard server for debugging
-            "--resources={{\\"node_tag:head\\":1}}"                   # Resource tag for selecting head as coordinator
+def _compress_and_upload_file(
+    session: snowpark.Session, source_path: Path, stage_path: PurePath, import_path: Optional[str] = None
+) -> None:
+    absolute_source_path = source_path.absolute()
+    leading_path = absolute_source_path.as_posix()[: -len(import_path)] if import_path else None
+    filename = f"{source_path.name}.zip" if source_path.is_dir() or source_path.suffix == ".py" else source_path.name
+    with zip_file_or_directory_to_stream(source_path.absolute().as_posix(), leading_path) as stream:
+        session.file.put_stream(
+            cast(IO[bytes], stream),
+            stage_path.joinpath(filename).as_posix(),
+            auto_compress=False,
+            overwrite=True,
         )
 
-        # Start Ray on the head node
-        ray start "${{common_params[@]}}" "${{head_params[@]}}" -v
-
-        ##### End Ray configuration #####
-
-        # TODO: Monitor MLRS and handle process crashes
-        python -m web.ml_runtime_grpc_server &
-
-        # TODO: Launch worker service(s) using SQL if Ray and MLRS successfully started
-        echo Running command: python "$@"
-
-        # Run user's Python entrypoint
-        python "$@"
-
-        # After the user's job completes, signal workers to shut down
-        echo "User job completed. Signaling workers to shut down..."
-        python "${{SYSTEM_DIR}}/signal_workers.py" --wait-time 15
-        echo "Head node job completed. Exiting."
-    fi
-    """
-).strip()
 
+def _upload_directory(session: snowpark.Session, source_path: Path, payload_stage_path: PurePath) -> None:
+    # 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.
+    upload_path_patterns = set()
+    for p in source_path.rglob("*"):
+        if p.is_dir():
+            continue
+        # Skip python cache files
+        if "__pycache__" in p.parts or p.suffix == ".pyc":
+            continue
+        if p.name.startswith("."):
+            # Hidden files: use .* pattern for batch upload
+            if p.suffix:
+                upload_path_patterns.add(p.parent.joinpath(f".*{p.suffix}"))
+            else:
+                upload_path_patterns.add(p.parent.joinpath(".*"))
+        else:
+            # Regular files: use * pattern for batch upload
+            if p.suffix:
+                upload_path_patterns.add(p.parent.joinpath(f"*{p.suffix}"))
+            else:
+                upload_path_patterns.add(p)
 
-def resolve_path(path: str) -> types.PayloadPath:
-    try:
-        stage_path = stage_utils.StagePath(path)
-    except ValueError:
-        return Path(path)
-    return stage_path
+    for path in upload_path_patterns:
+        session.file.put(
+            str(path),
+            payload_stage_path.joinpath(path.parent.relative_to(source_path)).as_posix(),
+            overwrite=True,
+            auto_compress=False,
+        )
 
 
 def upload_payloads(session: snowpark.Session, stage_path: PurePath, *payload_specs: types.PayloadSpec) -> None:
-    for source_path, remote_relative_path in payload_specs:
+    for spec in payload_specs:
+        source_path = spec.source_path
+        remote_relative_path = spec.remote_relative_path
+        compress = spec.compress
         payload_stage_path = stage_path.joinpath(remote_relative_path) if remote_relative_path else stage_path
         if isinstance(source_path, stage_utils.StagePath):
             # only copy files into one stage directory from another stage directory, not from stage file
             # due to incomplete of StagePath functionality
-            session.sql(f"copy files into {payload_stage_path.as_posix()}/ from {source_path.as_posix()}/").collect()
+            if source_path.as_posix().endswith(".py"):
+                session.sql(f"copy files into {stage_path.as_posix()}/ from {source_path.as_posix()}").collect()
+            else:
+                session.sql(
+                    f"copy files into {payload_stage_path.as_posix()}/ from {source_path.as_posix()}/"
+                ).collect()
         elif isinstance(source_path, Path):
             if source_path.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.
-                upload_path_patterns = set()
-                for p in source_path.rglob("*"):
-                    if p.is_dir():
-                        continue
-                    if p.name.startswith("."):
-                        # Hidden files: use .* pattern for batch upload
-                        if p.suffix:
-                            upload_path_patterns.add(p.parent.joinpath(f".*{p.suffix}"))
-                        else:
-                            upload_path_patterns.add(p.parent.joinpath(".*"))
-                    else:
-                        # Regular files: use * pattern for batch upload
-                        if p.suffix:
-                            upload_path_patterns.add(p.parent.joinpath(f"*{p.suffix}"))
-                        else:
-                            upload_path_patterns.add(p)
-
-                for path in upload_path_patterns:
+                if compress:
+                    _compress_and_upload_file(
+                        session,
+                        source_path,
+                        stage_path,
+                        remote_relative_path.as_posix() if remote_relative_path else None,
+                    )
+                else:
+                    _upload_directory(session, source_path, payload_stage_path)
+
+            elif source_path.is_file():
+                if compress and source_path.suffix == ".py":
+                    _compress_and_upload_file(
+                        session,
+                        source_path,
+                        stage_path,
+                        remote_relative_path.as_posix() if remote_relative_path else None,
+                    )
+                else:
                     session.file.put(
-                        str(path),
-                        payload_stage_path.joinpath(path.parent.relative_to(source_path)).as_posix(),
+                        str(source_path.resolve()),
+                        payload_stage_path.as_posix(),
                         overwrite=True,
                         auto_compress=False,
                     )
-            else:
-                session.file.put(
-                    str(source_path.resolve()),
-                    payload_stage_path.as_posix(),
-                    overwrite=True,
-                    auto_compress=False,
-                )
 
 
 def upload_system_resources(session: snowpark.Session, stage_path: PurePath) -> None:
@@ -336,8 +166,30 @@ def resolve_source(
 
 def resolve_entrypoint(
     source: Union[types.PayloadPath, Callable[..., Any]],
-    entrypoint: Optional[types.PayloadPath],
-) -> types.PayloadEntrypoint:
+    entrypoint: Optional[Union[types.PayloadPath, list[str]]],
+) -> Union[types.PayloadEntrypoint, list[str]]:
+    """Resolve and validate the entrypoint for a job payload.
+
+    Args:
+        source: The source path or callable for the job payload.
+        entrypoint: The entrypoint specification. Can be:
+            - A path (str or Path) to a Python script file
+            - A list of strings representing a custom command (passed through as-is)
+            - None (inferred from source if source is a file)
+
+    Returns:
+        Either a PayloadEntrypoint object for file-based entrypoints, or the list
+        of strings passed through unchanged for custom command entrypoints.
+
+    Raises:
+        ValueError: If the entrypoint is invalid or cannot be resolved.
+        FileNotFoundError: If the entrypoint file does not exist.
+    """
+    # If entrypoint is a list, pass it through without resolution/validation
+    # This allows users to specify custom entrypoints (e.g., installed CLI tools)
+    if isinstance(entrypoint, (list, tuple)):
+        return entrypoint
+
     if callable(source):
         # Entrypoint is generated for callable payloads
         return types.PayloadEntrypoint(
@@ -385,87 +237,223 @@ def resolve_entrypoint(
     )
 
 
-def resolve_additional_payloads(
-    additional_payloads: Optional[list[Union[str, tuple[str, str]]]]
-) -> list[types.PayloadSpec]:
+def get_zip_file_from_path(path: types.PayloadPath) -> types.PayloadPath:
+    """Finds the path of the outermost zip archive from a given file path.
+
+    Examples:
+        >>> get_zip_file_from_path("/path/to/archive.zip/nested_file.py")
+        "/path/to/archive.zip"
+        >>> get_zip_file_from_path("/path/to/archive.zip")
+        "/path/to/archive.zip"
+        >>> get_zip_file_from_path("/path/to/regular_file.py")
+        "/path/to/regular_file.py"
+
+    Args:
+        path: The file path to inspect.
+
+    Returns:
+        str: The path to the outermost zip file, or the original path if
+            none is found.
     """
-    Determine how to stage local packages so that imports continue to work.
+
+    path_str = path.as_posix()
+
+    index = path_str.rfind(".zip/")
+    if index != -1:
+        return stage_utils.resolve_path(path_str[: index + 4])
+    return path
+
+
+def _finalize_payload_pair(
+    p: types.PayloadPath, base_import_path: Optional[str]
+) -> tuple[types.PayloadPath, Optional[str]]:
+    """Finalize the `(payload_path, import_path)` pair based on source type.
+
+    - Zip file: ignore import path (returns `(p, None)`).
+    - Python file: if `base_import_path` is provided, append ".py"; otherwise None.
+    - Directory: preserve `base_import_path` as-is.
+    - Stage file: use `base_import_path` as-is since we do not compress stage files.
+    - Other files: ignore import path (None).
 
     Args:
-        additional_payloads: A list of directory paths, each optionally paired with a dot-separated
-            import path
-            e.g. [("proj/src/utils", "src.utils"), "proj/src/helper"]
-            if there is no import path, the last part of path will be considered as import path
-            e.g. the import path of "proj/src/helper" is "helper"
+        p (types.PayloadPath): The resolved source path
+        base_import_path (Optional[str]): Slash-separated import path
 
     Returns:
-        A list of payloadSpec for additional payloads.
+        tuple[types.PayloadPath, Optional[str]]: `(p, final_import_path)` where:
+            - `final_import_path` is None for zip archives and non-Python files.
+            - `final_import_path` is `base_import_path + ".py"` for Python files when
+              `base_import_path` is provided; otherwise None.
+            - `final_import_path` is `base_import_path` for directories.
+
+    """
+    if p.suffix == ".zip":
+        final_import_path = None
+    elif isinstance(p, stage_utils.StagePath):
+        final_import_path = base_import_path
+    elif p.is_file():
+        if p.suffix == ".py":
+            final_import_path = (base_import_path + ".py") if base_import_path else None
+        else:
+            final_import_path = None
+    else:
+        final_import_path = base_import_path
+
+    validate_import_path(p, final_import_path)
+    return (p, None) if p.suffix == ".zip" else (p, final_import_path)
 
-    Raises:
-        FileNotFoundError: If any specified package path does not exist.
-        ValueError: If the format of local_packages is invalid.
 
+def resolve_import_path(
+    path: Union[types.PayloadPath, ModuleType],
+    import_path: Optional[str] = None,
+) -> list[tuple[types.PayloadPath, Optional[str]]]:
     """
-    if not additional_payloads:
-        return []
+    Resolve and normalize the import path for modules, Python files, or zip payloads.
 
-    logger.warning(
-        "When providing a stage path as an additional payload, "
-        "please ensure it points to a directory. "
-        "Files are not currently supported."
-    )
+    Args:
+        path (Union[types.PayloadPath, ModuleType]): The source path or module to resolve.
+            - If a directory is provided, it is compressed as a zip archive preserving its structure.
+            - If a single Python file is provided, the file itself is zipped.
+            - If a module is provided, it is treated as a directory or Python file.
+            - If a zip file is provided, it is uploaded as it is.
+            - If a stage file is provided, we only support stage file when the import path is provided
+        import_path (Optional[str], optional): Explicit import path to use. If None,
+            the function infers it from `path`.
 
-    additional_payloads_paths = []
-    for pkg in additional_payloads:
-        if isinstance(pkg, str):
-            source_path = resolve_path(pkg).absolute()
-            module_path = source_path.name
-        elif isinstance(pkg, tuple):
-            try:
-                source_path_str, module_path = pkg
-            except ValueError:
-                raise ValueError(
-                    f"Invalid format in `additional_payloads`. "
-                    f"Expected a tuple of (source_path, module_path). Got {pkg}"
-                )
-            source_path = resolve_path(source_path_str).absolute()
+    Returns:
+        list[tuple[types.PayloadPath, Optional[str]]]: A list of tuples where each tuple
+        contains the resolved payload path and its corresponding import path (if any).
+
+    Raises:
+        FileNotFoundError: If the provided `path` does not exist.
+        NotImplementedError: If the stage file is provided without an import path.
+        ValueError: If the import path cannot be resolved or is invalid.
+    """
+    if import_path is None:
+        import_path = path.stem if isinstance(path, types.PayloadPath) else path.__name__
+    import_path = import_path.strip().replace(".", "/") if import_path else None
+    if isinstance(path, Path):
+        if not path.exists():
+            raise FileNotFoundError(f"{path} is not found")
+        return [_finalize_payload_pair(path.absolute(), import_path)]
+    elif isinstance(path, stage_utils.StagePath):
+        if import_path:
+            return [_finalize_payload_pair(path.absolute(), import_path)]
+        raise NotImplementedError("We only support stage file when the import path is provided")
+    elif isinstance(path, ModuleType):
+        if hasattr(path, "__path__"):
+            paths = [get_zip_file_from_path(stage_utils.resolve_path(p).absolute()) for p in path.__path__]
+            return [_finalize_payload_pair(p, import_path) for p in paths]
+        elif hasattr(path, "__file__") and path.__file__:
+            p = get_zip_file_from_path(Path(path.__file__).absolute())
+            return [_finalize_payload_pair(p, import_path)]
         else:
-            raise ValueError("the format of additional payload is not correct")
+            raise ValueError(f"Module {path} is not a valid module")
+    else:
+        raise ValueError(f"Module {path} is not a valid imports")
 
-        if not source_path.exists():
-            raise FileNotFoundError(f"{source_path} does not exist")
 
-        if isinstance(source_path, Path):
-            if source_path.is_file():
-                raise ValueError(f"file is not supported for additional payloads: {source_path}")
+def validate_import_path(source: Union[str, types.PayloadPath], import_path: Optional[str]) -> None:
+    """Validate the import path for local python file or directory."""
+    if import_path is None:
+        return
 
-        module_parts = module_path.split(".")
-        for part in module_parts:
-            if not part.isidentifier() or keyword.iskeyword(part):
-                raise ValueError(
-                    f"Invalid module import path '{module_path}'. "
-                    f"'{part}' is not a valid Python identifier or is a keyword."
-                )
+    source_path = stage_utils.resolve_path(source) if isinstance(source, str) else source
+    if isinstance(source_path, stage_utils.StagePath):
+        if not source_path.as_posix().endswith(import_path + ".py"):
+            raise ValueError(f"Import path {import_path} must end with the source name {source_path}")
+    elif (source_path.is_file() and source_path.suffix == ".py") or source_path.is_dir():
+        if not source_path.as_posix().endswith(import_path):
+            raise ValueError(f"Import path {import_path} must end with the source name {source_path}")
 
-        dest_path = PurePath(*module_parts)
-        additional_payloads_paths.append(types.PayloadSpec(source_path, dest_path))
-    return additional_payloads_paths
+
+def upload_imports(
+    imports: Optional[list[Union[str, Path, ModuleType, tuple[Union[str, Path, ModuleType], Optional[str]]]]],
+    session: snowpark.Session,
+    stage_path: PurePath,
+) -> None:
+    """Resolve paths and upload imports for ML Jobs.
+
+    Args:
+        imports: Optional list of paths/modules, or tuples of
+            ``(path_or_module, import_path)``. The path can be a local
+            directory, a local ``.py`` file, a local ``.zip`` file, or a stage
+            path (for example, ``@stage/path``). If a tuple is provided and the
+            first element is a local directory or ``.py`` file, the second
+            element denotes the Python import path (dot or slash separated) to
+            which the content should be mounted. If not provided for local
+            sources, it defaults to the stem of the path/module. For stage
+            paths or non-Python local files, the import path is ignored.
+        session: Active Snowpark session used to upload files.
+        stage_path: Destination stage subpath where payloads will be uploaded.
+
+    Raises:
+        ValueError: If a import has an invalid format or the
+            provided import path is incompatible with the source.
+
+    """
+    if not imports:
+        return
+    for additional_payload in imports:
+        if isinstance(additional_payload, tuple):
+            source, import_path = additional_payload
+        elif isinstance(additional_payload, str) or isinstance(additional_payload, ModuleType):
+            source = additional_payload
+            import_path = None
+        else:
+            raise ValueError(f"Invalid import format: {additional_payload}")
+        resolved_imports = resolve_import_path(
+            stage_utils.resolve_path(source) if not isinstance(source, ModuleType) else source, import_path
+        )
+        for source_path, import_path in resolved_imports:
+            # TODO(SNOW-2467038): support import path for stage files or directories
+            if isinstance(source_path, stage_utils.StagePath):
+                remote = None
+                compress = False
+            elif source_path.as_posix().endswith(".zip"):
+                remote = None
+                compress = False
+            elif source_path.is_dir() or source_path.suffix == ".py":
+                remote = PurePath(import_path) if import_path else None
+                compress = True
+            else:
+                # if the file is not a python file, ignore the import path
+                remote = None
+                compress = False
+
+            upload_payloads(session, stage_path, types.PayloadSpec(source_path, remote, compress=compress))
 
 
 class JobPayload:
     def __init__(
         self,
         source: Union[str, Path, Callable[..., Any]],
-        entrypoint: Optional[Union[str, Path]] = None,
+        entrypoint: Optional[Union[str, Path, list[str]]] = None,
         *,
         pip_requirements: Optional[list[str]] = None,
-        additional_payloads: Optional[list[Union[str, tuple[str, str]]]] = None,
+        imports: Optional[list[Union[ImportType, tuple[ImportType, Optional[str]]]]] = None,
     ) -> None:
+        """Initialize a job payload.
+
+        Args:
+            source: The source for the job payload. Can be a file path, directory path,
+                stage path, or a callable.
+            entrypoint: The entrypoint for job execution. Can be:
+                - A path (str or Path) to a Python script file
+                - A list of strings representing a custom command (e.g., ["arctic_training"])
+                  which is passed through as-is without resolution or validation
+                - None (inferred from source if source is a file)
+            pip_requirements: A list of pip requirements for the job.
+            imports: A list of additional imports/payloads for the job.
+        """
         # for stage path like snow://domain....., Path(path) will remove duplicate /, it will become snow:/ domain...
-        self.source = resolve_path(source) if isinstance(source, str) else source
-        self.entrypoint = resolve_path(entrypoint) if isinstance(entrypoint, str) else entrypoint
+        self.source = stage_utils.resolve_path(source) if isinstance(source, str) else source
+        if isinstance(entrypoint, list):
+            self.entrypoint: Optional[Union[types.PayloadPath, list[str]]] = entrypoint
+        else:
+            self.entrypoint = stage_utils.resolve_path(entrypoint) if isinstance(entrypoint, str) else entrypoint
         self.pip_requirements = pip_requirements
-        self.additional_payloads = additional_payloads
+        self.imports = imports
 
     def upload(self, session: snowpark.Session, stage_path: Union[str, PurePath]) -> types.UploadedPayload:
         # Prepare local variables
@@ -473,7 +461,6 @@ class JobPayload:
         source = resolve_source(self.source)
         entrypoint = resolve_entrypoint(source, self.entrypoint)
         pip_requirements = self.pip_requirements or []
-        additional_payload_specs = resolve_additional_payloads(self.additional_payloads)
 
         # Create stage if necessary
         stage_name = stage_path.parts[0].lstrip("@")
@@ -488,70 +475,81 @@ class JobPayload:
                 " comment = 'Created by snowflake.ml.jobs Python API'",
                 params=[stage_name],
             )
-        payload_name = None
+
         # Upload payload to stage - organize into app/ subdirectory
         app_stage_path = stage_path.joinpath(constants.APP_STAGE_SUBPATH)
-        if not isinstance(source, types.PayloadPath):
-            if isinstance(source, function_payload_utils.FunctionPayload):
-                payload_name = source.function.__name__
-
-            source_code = generate_python_code(source, source_code_display=True)
-            _ = session.file.put_stream(
-                io.BytesIO(source_code.encode()),
-                stage_location=app_stage_path.joinpath(entrypoint.file_path).as_posix(),
+        upload_imports(self.imports, session, app_stage_path)
+
+        # Handle list entrypoints (custom commands like ["arctic_training"])
+        if isinstance(entrypoint, (list, tuple)):
+            payload_name = entrypoint[0] if entrypoint else None
+            # For list entrypoints, still upload source if it's a path
+            if isinstance(source, Path):
+                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
+            elif isinstance(source, stage_utils.StagePath):
+                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
+            python_entrypoint: list[Union[str, PurePath]] = list(entrypoint)
+        else:
+            # Standard file-based entrypoint handling
+            payload_name = None
+            if not isinstance(source, types.PayloadPath):
+                if isinstance(source, function_payload_utils.FunctionPayload):
+                    payload_name = source.function.__name__
+
+                source_code = generate_python_code(source, source_code_display=True)
+                _ = session.file.put_stream(
+                    io.BytesIO(source_code.encode()),
+                    stage_location=app_stage_path.joinpath(entrypoint.file_path).as_posix(),
+                    auto_compress=False,
+                    overwrite=True,
+                )
+                source = Path(entrypoint.file_path.parent)
+
+            elif isinstance(source, stage_utils.StagePath):
+                payload_name = entrypoint.file_path.stem
+                # copy payload to stage
+                if source == entrypoint.file_path:
+                    source = source.parent
+                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
+
+            elif isinstance(source, Path):
+                payload_name = entrypoint.file_path.stem
+                upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
+                if source.is_file():
+                    source = source.parent
+
+            python_entrypoint = [
+                PurePath(
+                    constants.STAGE_VOLUME_MOUNT_PATH,
+                    constants.APP_STAGE_SUBPATH,
+                    entrypoint.file_path.relative_to(source).as_posix(),
+                ),
+            ]
+            if entrypoint.main_func:
+                python_entrypoint += ["--script_main_func", entrypoint.main_func]
+
+        if pip_requirements:
+            session.file.put_stream(
+                io.BytesIO("\n".join(pip_requirements).encode()),
+                stage_location=app_stage_path.joinpath("requirements.txt").as_posix(),
                 auto_compress=False,
                 overwrite=True,
             )
-            source = Path(entrypoint.file_path.parent)
-
-        elif isinstance(source, stage_utils.StagePath):
-            payload_name = entrypoint.file_path.stem
-            # copy payload to stage
-            if source == entrypoint.file_path:
-                source = source.parent
-            upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
 
-        elif isinstance(source, Path):
-            payload_name = entrypoint.file_path.stem
-            upload_payloads(session, app_stage_path, types.PayloadSpec(source, None))
-            if source.is_file():
-                source = source.parent
-
-        upload_payloads(session, app_stage_path, *additional_payload_specs)
-
-        if not any(r.startswith("cloudpickle") for r in pip_requirements):
-            pip_requirements.append(f"cloudpickle~={version.parse(cp.__version__).major}.0")
-
-        #  Upload system scripts and requirements.txt generated by pip_requirements to system/ directory
+        # Upload system scripts and other assets to system/ directory
         system_stage_path = stage_path.joinpath(constants.SYSTEM_STAGE_SUBPATH)
-        if pip_requirements:
-            # Upload requirements.txt to stage
+        system_pip_requirements = []
+        if not any(r.startswith("cloudpickle") for r in pip_requirements):
+            system_pip_requirements.append(f"cloudpickle~={version.parse(cp.__version__).major}.0")
+        if system_pip_requirements:
+            # Upload requirements.txt to system path in stage
             session.file.put_stream(
-                io.BytesIO("\n".join(pip_requirements).encode()),
+                io.BytesIO("\n".join(system_pip_requirements).encode()),
                 stage_location=system_stage_path.joinpath("requirements.txt").as_posix(),
                 auto_compress=False,
                 overwrite=True,
             )
-
-        # TODO: Make sure payload does not include file with same name
-        session.file.put_stream(
-            io.BytesIO(_STARTUP_SCRIPT_CODE.encode()),
-            stage_location=system_stage_path.joinpath(_STARTUP_SCRIPT_PATH).as_posix(),
-            auto_compress=False,
-            overwrite=False,  # FIXME
-        )
-
         upload_system_resources(session, system_stage_path)
-        python_entrypoint: list[Union[str, PurePath]] = [
-            PurePath(constants.STAGE_VOLUME_MOUNT_PATH, constants.SYSTEM_STAGE_SUBPATH, "mljob_launcher.py"),
-            PurePath(
-                constants.STAGE_VOLUME_MOUNT_PATH,
-                constants.APP_STAGE_SUBPATH,
-                entrypoint.file_path.relative_to(source).as_posix(),
-            ),
-        ]
-        if entrypoint.main_func:
-            python_entrypoint += ["--script_main_func", entrypoint.main_func]
 
         env_vars = {
             constants.STAGE_MOUNT_PATH_ENV_VAR: constants.STAGE_VOLUME_MOUNT_PATH,