snowflake-ml-python 1.7.2__py3-none-any.whl → 1.7.4__py3-none-any.whl

snowflake/ml/jobs/_utils/payload_utils.py

@@ -0,0 +1,352 @@
+import inspect
+import io
+import sys
+import textwrap
+from pathlib import Path, PurePath
+from typing import (
+    Any,
+    Callable,
+    List,
+    Optional,
+    Type,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+)
+
+import cloudpickle as cp
+
+from snowflake import snowpark
+from snowflake.ml.jobs._utils import constants, types
+from snowflake.snowpark._internal import code_generation
+
+_SUPPORTED_ARG_TYPES = {str, int, float}
+_STARTUP_SCRIPT_PATH = PurePath("startup.sh")
+_STARTUP_SCRIPT_CODE = textwrap.dedent(
+    f"""
+    #!/bin/bash
+
+    ##### 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.PAYLOAD_DIR_ENV_VAR}
+    fi
+
+    ##### Set up Python environment #####
+    export PYTHONPATH=/opt/env/site-packages/
+    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)
+
+    # Configure IP address and logging directory
+    eth0Ip=$(ifconfig eth0 | sed -En -e 's/.*inet ([0-9.]+).*/\1/p')
+    log_dir="/tmp/ray"
+
+    # Check if eth0Ip is empty and set default if necessary
+    if [ -z "$eth0Ip" ]; then
+        # This should never happen, but just in case ethOIp is not set, we should default to localhost
+        eth0Ip="127.0.0.1"
+    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"
+    )
+
+    # 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}}"      # Listening port 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 on
+        "--dashboard-port=${{DASHBOARD_PORT}}"                  # Port to bind the dashboard server for local debugging
+        "--resources={{\\"node_tag:head\\":1}}"                   # Resource tag for selecting head as coordinator
+    )
+
+    # Start Ray on the head node
+    ray start "${{common_params[@]}}" "${{head_params[@]}}" &
+    ##### 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
+
+    # Run user's Python entrypoint
+    echo Running command: python "$@"
+    python "$@"
+    """
+).strip()
+
+
+class JobPayload:
+    def __init__(
+        self,
+        source: Union[str, Path, Callable[..., Any]],
+        entrypoint: Optional[Union[str, Path]] = None,
+        *,
+        pip_requirements: Optional[List[str]] = None,
+    ) -> None:
+        self.source = Path(source) if isinstance(source, str) else source
+        self.entrypoint = Path(entrypoint) if isinstance(entrypoint, str) else entrypoint
+        self.pip_requirements = pip_requirements
+
+    def validate(self) -> None:
+        if callable(self.source):
+            # Any entrypoint value is OK for callable payloads (including None aka default)
+            # since we will generate the file from the serialized callable
+            pass
+        elif isinstance(self.source, Path):
+            # Validate self.source and self.entrypoint for files
+            if not self.source.exists():
+                raise FileNotFoundError(f"{self.source} does not exist")
+            if self.entrypoint is None:
+                if self.source.is_file():
+                    self.entrypoint = self.source
+                else:
+                    raise ValueError("entrypoint must be provided when source is a directory")
+            if not self.entrypoint.is_file():
+                # Check if self.entrypoint is a valid relative path
+                self.entrypoint = self.source.joinpath(self.entrypoint)
+                if not self.entrypoint.is_file():
+                    raise FileNotFoundError(f"File {self.entrypoint} does not exist")
+            if not self.entrypoint.is_relative_to(self.source):
+                raise ValueError(f"{self.entrypoint} must be a subpath of {self.source}")
+            if self.entrypoint.suffix != ".py":
+                raise NotImplementedError("Only Python entrypoints are supported currently")
+        else:
+            raise ValueError("Unsupported source type. Source must be a file, directory, or callable.")
+
+    def upload(self, session: snowpark.Session, stage_path: Union[str, PurePath]) -> types.UploadedPayload:
+        # Validate payload
+        self.validate()
+
+        # Prepare local variables
+        stage_path = PurePath(stage_path) if isinstance(stage_path, str) else stage_path
+        source = self.source
+        entrypoint = self.entrypoint or Path(constants.DEFAULT_ENTRYPOINT_PATH)
+
+        # Create stage if necessary
+        stage_name = stage_path.parts[0]
+        session.sql(
+            f"create stage if not exists {stage_name.lstrip('@')}"
+            " encryption = ( type = 'SNOWFLAKE_SSE' )"
+            " comment = 'Created by snowflake.ml.jobs Python API'"
+        ).collect()
+
+        # Upload payload to stage
+        if not isinstance(source, Path):
+            source_code = generate_python_code(source, source_code_display=True)
+            _ = session.file.put_stream(
+                io.BytesIO(source_code.encode()),
+                stage_location=stage_path.joinpath(entrypoint).as_posix(),
+                auto_compress=False,
+                overwrite=True,
+            )
+            source = entrypoint.parent
+        elif source.is_dir():
+            # Manually traverse the directory and upload each file, since Snowflake PUT
+            # can't handle directories. Reduce the number of PUT operations by using
+            # wildcard patterns to batch upload files with the same extension.
+            for path in {
+                p.parent.joinpath(f"*{p.suffix}") if p.suffix else p for p in source.resolve().rglob("*") if p.is_file()
+            }:
+                session.file.put(
+                    str(path),
+                    stage_path.joinpath(path.parent.relative_to(source)).as_posix(),
+                    overwrite=True,
+                    auto_compress=False,
+                )
+        else:
+            session.file.put(
+                str(source.resolve()),
+                stage_path.as_posix(),
+                overwrite=True,
+                auto_compress=False,
+            )
+            source = source.parent
+
+        # Upload requirements
+        # TODO: Check if payload includes both a requirements.txt file and pip_requirements
+        if self.pip_requirements:
+            # Upload requirements.txt to stage
+            session.file.put_stream(
+                io.BytesIO("\n".join(self.pip_requirements).encode()),
+                stage_location=stage_path.joinpath("requirements.txt").as_posix(),
+                auto_compress=False,
+                overwrite=True,
+            )
+
+        # Upload startup script
+        # TODO: Make sure payload does not include file with same name
+        session.file.put_stream(
+            io.BytesIO(_STARTUP_SCRIPT_CODE.encode()),
+            stage_location=stage_path.joinpath(_STARTUP_SCRIPT_PATH).as_posix(),
+            auto_compress=False,
+            overwrite=False,  # FIXME
+        )
+
+        return types.UploadedPayload(
+            stage_path=stage_path,
+            entrypoint=[
+                "bash",
+                _STARTUP_SCRIPT_PATH,
+                entrypoint.relative_to(source),
+            ],
+        )
+
+
+def get_parameter_type(param: inspect.Parameter) -> Optional[Type[object]]:
+    # Unwrap Optional type annotations
+    param_type = param.annotation
+    if get_origin(param_type) is Union and len(get_args(param_type)) == 2 and type(None) in get_args(param_type):
+        param_type = next(t for t in get_args(param_type) if t is not type(None))
+
+    # Return None for empty type annotations
+    if param_type == inspect.Parameter.empty:
+        return None
+    return cast(Type[object], param_type)
+
+
+def validate_parameter_type(param_type: Type[object], param_name: str) -> None:
+    # Validate param_type is a supported type
+    if param_type not in _SUPPORTED_ARG_TYPES:
+        raise ValueError(
+            f"Unsupported argument type {param_type} for '{param_name}'."
+            f" Supported types: {', '.join(t.__name__ for t in _SUPPORTED_ARG_TYPES)}"
+        )
+
+
+def generate_python_code(func: Callable[..., Any], source_code_display: bool = False) -> str:
+    signature = inspect.signature(func)
+    if any(
+        p.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
+        for p in signature.parameters.values()
+    ):
+        raise NotImplementedError("Function must not have unpacking arguments (* or **)")
+
+    # Mirrored from Snowpark generate_python_code() function
+    # https://github.com/snowflakedb/snowpark-python/blob/main/src/snowflake/snowpark/_internal/udf_utils.py
+    try:
+        source_code_comment = (
+            code_generation.generate_source_code(func) if source_code_display else ""  # type: ignore[arg-type]
+        )
+    except Exception as exc:
+        error_msg = f"Source code comment could not be generated for {func} due to error {exc}."
+        source_code_comment = code_generation.comment_source_code(error_msg)
+
+    func_name = "func"
+    func_code = f"""
+{source_code_comment}
+
+import pickle
+{func_name} = pickle.loads(bytes.fromhex('{cp.dumps(func).hex()}'))
+"""
+
+    # Generate argparse logic for argument handling (type coercion, default values, etc)
+    argparse_code = ["import argparse", "", "parser = argparse.ArgumentParser()"]
+    argparse_postproc = []
+    for name, param in signature.parameters.items():
+        opts = {}
+
+        param_type = get_parameter_type(param)
+        if param_type is not None:
+            validate_parameter_type(param_type, name)
+            opts["type"] = param_type.__name__
+
+        if param.default != inspect.Parameter.empty:
+            opts["default"] = f"'{param.default}'" if isinstance(param.default, str) else param.default
+
+        if param.kind == inspect.Parameter.KEYWORD_ONLY:
+            # Keyword argument
+            argparse_code.append(
+                f"parser.add_argument('--{name}', required={'default' not in opts},"
+                f" {', '.join(f'{k}={v}' for k, v in opts.items())})"
+            )
+        else:
+            # Positional argument. Use `argparse.add_mutually_exclusive_group()`
+            # to allow passing positional args by name as well
+            group_name = f"{name}_group"
+            argparse_code.append(
+                f"{group_name} = parser.add_mutually_exclusive_group(required={'default' not in opts})"
+            )
+            argparse_code.append(
+                f"{group_name}.add_argument('pos-{name}', metavar='{name}', nargs='?',"
+                f" {', '.join(f'{k}={v}' for k, v in opts.items() if k != 'default')})"
+            )
+            argparse_code.append(
+                f"{group_name}.add_argument('--{name}', {', '.join(f'{k}={v}' for k, v in opts.items())})"
+            )
+            argparse_code.append("")  # Add newline for readability
+            argparse_postproc.append(
+                f"args.{name} = {name} if ({name} := args.__dict__.pop('pos-{name}')) is not None else args.{name}"
+            )
+    argparse_code.append("args = parser.parse_args()")
+    param_code = "\n".join(argparse_code + argparse_postproc)
+
+    return f"""
+### Version guard to check compatibility across Python versions ###
+import sys
+import warnings
+
+if sys.version_info.major != {sys.version_info.major} or sys.version_info.minor != {sys.version_info.minor}:
+    warnings.warn(
+        "Python version mismatch: job was created using"
+        " python{sys.version_info.major}.{sys.version_info.minor}"
+        f" but runtime environment uses python{{sys.version_info.major}}.{{sys.version_info.minor}}."
+        " Compatibility across Python versions is not guaranteed and may result in unexpected behavior."
+        " This will be fixed in a future release; for now, please use Python version"
+        f" {{sys.version_info.major}}.{{sys.version_info.minor}}.",
+        RuntimeWarning,
+        stacklevel=0,
+    )
+### End version guard ###
+
+{func_code.strip()}
+
+if __name__ == '__main__':
+{textwrap.indent(param_code, '    ')}
+
+    {func_name}(**vars(args))
+"""

snowflake/ml/jobs/_utils/spec_utils.py

@@ -0,0 +1,298 @@
+import logging
+from math import ceil
+from pathlib import PurePath
+from typing import Any, Dict, List, Optional, Union
+
+from snowflake import snowpark
+from snowflake.ml._internal.utils import snowflake_env
+from snowflake.ml.jobs._utils import constants, 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 = session.sql(f"show compute pools like '{compute_pool}'").collect()
+    if not rows:
+        raise ValueError(f"Compute pool '{compute_pool}' not found")
+    instance_family: str = rows[0]["instance_family"]
+
+    # Get the cloud we're using (AWS, Azure, etc)
+    region = snowflake_env.get_regions(session)[snowflake_env.get_current_region_id(session)]
+    cloud = region["cloud"]
+
+    return (
+        constants.COMMON_INSTANCE_FAMILIES.get(instance_family)
+        or constants.CLOUD_INSTANCE_FAMILIES[cloud][instance_family]
+    )
+
+
+def _get_image_spec(session: snowpark.Session, compute_pool: str) -> types.ImageSpec:
+    # Retrieve compute pool node resources
+    resources = _get_node_resources(session, compute_pool=compute_pool)
+
+    # Use MLRuntime image
+    image_repo = constants.DEFAULT_IMAGE_REPO
+    image_name = constants.DEFAULT_IMAGE_GPU if resources.gpu > 0 else constants.DEFAULT_IMAGE_CPU
+    image_tag = constants.DEFAULT_IMAGE_TAG
+
+    # Try to pull latest image tag from server side if possible
+    query_result = session.sql("SHOW PARAMETERS LIKE 'constants.RUNTIME_BASE_IMAGE_TAG' IN ACCOUNT").collect()
+    if query_result:
+        image_tag = query_result[0]["value"]
+
+    # TODO: Should each instance consume the entire pod?
+    return types.ImageSpec(
+        repo=image_repo,
+        image_name=image_name,
+        image_tag=image_tag,
+        resource_requests=resources,
+        resource_limits=resources,
+    )
+
+
+def generate_spec_overrides(
+    environment_vars: Optional[Dict[str, str]] = None,
+    custom_overrides: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Generate a dictionary of service specification overrides.
+
+    Args:
+        environment_vars: Environment variables to set in primary container
+        custom_overrides: Custom service specification overrides
+
+    Returns:
+        Resulting service specifiation patch dict. Empty if no overrides were supplied.
+    """
+    # Generate container level overrides
+    container_spec: Dict[str, Any] = {
+        "name": constants.DEFAULT_CONTAINER_NAME,
+    }
+    if environment_vars:
+        # TODO: Validate environment variables
+        container_spec["env"] = environment_vars
+
+    # Build container override spec only if any overrides were supplied
+    spec = {}
+    if len(container_spec) > 1:
+        spec = {
+            "spec": {
+                "containers": [container_spec],
+            }
+        }
+
+    # Apply custom overrides
+    if custom_overrides:
+        spec = merge_patch(spec, custom_overrides, display_name="custom_overrides")
+
+    return spec
+
+
+def generate_service_spec(
+    session: snowpark.Session,
+    compute_pool: str,
+    payload: types.UploadedPayload,
+    args: Optional[List[str]] = None,
+) -> Dict[str, Any]:
+    """
+    Generate a service specification for a job.
+
+    Args:
+        session: Snowflake session
+        compute_pool: Compute pool for job execution
+        payload: Uploaded job payload
+        args: Arguments to pass to entrypoint script
+
+    Returns:
+        Job service specification
+    """
+    # Set resource requests/limits, including nvidia.com/gpu quantity if applicable
+    image_spec = _get_image_spec(session, compute_pool)
+    resource_requests: Dict[str, Union[str, int]] = {
+        "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
+        "memory": f"{image_spec.resource_limits.memory}Gi",
+    }
+    resource_limits: Dict[str, Union[str, int]] = {
+        "cpu": f"{int(image_spec.resource_requests.cpu * 1000)}m",
+        "memory": f"{image_spec.resource_limits.memory}Gi",
+    }
+    if image_spec.resource_limits.gpu > 0:
+        resource_requests["nvidia.com/gpu"] = image_spec.resource_requests.gpu
+        resource_limits["nvidia.com/gpu"] = image_spec.resource_limits.gpu
+
+    # Add local volumes for ephemeral logs and artifacts
+    volumes: List[Dict[str, str]] = []
+    volume_mounts: List[Dict[str, str]] = []
+    for volume_name, mount_path in [
+        ("system-logs", "/var/log/managedservices/system/mlrs"),
+        ("user-logs", "/var/log/managedservices/user/mlrs"),
+    ]:
+        volume_mounts.append(
+            {
+                "name": volume_name,
+                "mountPath": mount_path,
+            }
+        )
+        volumes.append(
+            {
+                "name": volume_name,
+                "source": "local",
+            }
+        )
+
+    # Mount 30% of memory limit as a memory-backed volume
+    memory_volume_name = "dshm"
+    memory_volume_size = min(
+        ceil(image_spec.resource_limits.memory * constants.MEMORY_VOLUME_SIZE),
+        image_spec.resource_requests.memory,
+    )
+    volume_mounts.append(
+        {
+            "name": memory_volume_name,
+            "mountPath": "/dev/shm",
+        }
+    )
+    volumes.append(
+        {
+            "name": memory_volume_name,
+            "source": "memory",
+            "size": f"{memory_volume_size}Gi",
+        }
+    )
+
+    # Mount payload as volume
+    stage_mount = PurePath("/opt/app")
+    stage_volume_name = "stage-volume"
+    volume_mounts.append(
+        {
+            "name": stage_volume_name,
+            "mountPath": stage_mount.as_posix(),
+        }
+    )
+    volumes.append(
+        {
+            "name": stage_volume_name,
+            "source": payload.stage_path.as_posix(),
+        }
+    )
+
+    # TODO: Add hooks for endpoints for integration with TensorBoard etc
+
+    # Assemble into service specification dict
+    spec = {
+        "spec": {
+            "containers": [
+                {
+                    "name": constants.DEFAULT_CONTAINER_NAME,
+                    "image": image_spec.full_name,
+                    "command": ["/usr/local/bin/_entrypoint.sh"],
+                    "args": [
+                        stage_mount.joinpath(v).as_posix() if isinstance(v, PurePath) else v for v in payload.entrypoint
+                    ]
+                    + (args or []),
+                    "env": {
+                        constants.PAYLOAD_DIR_ENV_VAR: stage_mount.as_posix(),
+                    },
+                    "volumeMounts": volume_mounts,
+                    "resources": {
+                        "requests": resource_requests,
+                        "limits": resource_limits,
+                    },
+                },
+            ],
+            "volumes": volumes,
+        }
+    }
+
+    return spec
+
+
+def merge_patch(base: Any, patch: Any, display_name: str = "") -> Any:
+    """
+    Implements a modified RFC7386 JSON Merge Patch
+    https://datatracker.ietf.org/doc/html/rfc7386
+
+    Behavior differs from the RFC in the following ways:
+      1. Empty nested dictionaries resulting from the patch are treated as None and are pruned
+      2. Attempts to merge lists of dicts using a merge key (default "name").
+         See _merge_lists_of_dicts for details on list merge behavior.
+
+    Args:
+        base: The base object to patch.
+        patch: The patch object.
+        display_name: The name of the patch object for logging purposes.
+
+    Returns:
+        The patched object.
+    """
+    if not type(base) is type(patch):
+        if base is not None:
+            logging.warning(f"Type mismatch while merging {display_name} (base={type(base)}, patch={type(patch)})")
+        return patch
+    elif isinstance(patch, list) and all(isinstance(v, dict) for v in base + patch):
+        # TODO: Should we prune empty lists?
+        return _merge_lists_of_dicts(base, patch, display_name=display_name)
+    elif not isinstance(patch, dict) or len(patch) == 0:
+        return patch
+
+    result = dict(base)  # Shallow copy
+    for key, value in patch.items():
+        if value is None:
+            result.pop(key, None)
+        else:
+            merge_result = merge_patch(result.get(key, None), value, display_name=f"{display_name}.{key}")
+            if isinstance(merge_result, dict) and len(merge_result) == 0:
+                result.pop(key, None)
+            else:
+                result[key] = merge_result
+
+    return result
+
+
+def _merge_lists_of_dicts(
+    base: List[Dict[str, Any]], patch: List[Dict[str, Any]], merge_key: str = "name", display_name: str = ""
+) -> List[Dict[str, Any]]:
+    """
+    Attempts to merge lists of dicts by matching on a merge key (default "name").
+    - If the merge key is missing, the behavior falls back to overwriting the list.
+    - If the merge key is present, the behavior is to match the list elements based on the
+        merge key and preserving any unmatched elements from the base list.
+    - Matched entries may be dropped in the following way(s):
+        1. The matching patch entry has a None key entry, e.g. { "name": "foo", None: None }.
+
+    Args:
+        base: The base list of dicts.
+        patch: The patch list of dicts.
+        merge_key: The key to use for merging.
+        display_name: The name of the patch object for logging purposes.
+
+    Returns:
+        The merged list of dicts if merging successful, else returns the patch list.
+    """
+    if any(merge_key not in d for d in base + patch):
+        logging.warning(f"Missing merge key {merge_key} in {display_name}. Falling back to overwrite behavior.")
+        return patch
+
+    # Build mapping of merge key values to list elements for the base list
+    result = {d[merge_key]: d for d in base}
+    if len(result) != len(base):
+        logging.warning(f"Duplicate merge key {merge_key} in {display_name}. Falling back to overwrite behavior.")
+        return patch
+
+    # Apply patches
+    for d in patch:
+        key = d[merge_key]
+
+        # Removal case 1: `None` key in patch entry
+        if None in d:
+            result.pop(key, None)
+            continue
+
+        # Apply patch
+        if key in result:
+            d = merge_patch(result[key], d, display_name=f"{display_name}[{merge_key}={d[merge_key]}]")
+            # TODO: Should we drop the item if the patch result is empty save for the merge key?
+            #       Can check `d.keys() <= {merge_key}`
+        result[key] = d
+
+    return list(result.values())

snowflake/ml/jobs/_utils/types.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+from pathlib import PurePath
+from typing import List, Literal, Optional, Union
+
+JOB_STATUS = Literal[
+    "PENDING",
+    "RUNNING",
+    "FAILED",
+    "DONE",
+    "INTERNAL_ERROR",
+]
+
+
+@dataclass(frozen=True)
+class UploadedPayload:
+    # TODO: Include manifest of payload files for validation
+    stage_path: PurePath
+    entrypoint: List[Union[str, PurePath]]
+
+
+@dataclass(frozen=True)
+class ComputeResources:
+    cpu: float  # Number of vCPU cores
+    memory: float  # Memory in GiB
+    gpu: int = 0  # Number of GPUs
+    gpu_type: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class ImageSpec:
+    repo: str
+    image_name: str
+    image_tag: str
+    resource_requests: ComputeResources
+    resource_limits: ComputeResources
+
+    @property
+    def full_name(self) -> str:
+        return f"{self.repo}/{self.image_name}:{self.image_tag}"