snowflake.ml.model.JobSpec

class snowflake.ml.model.JobSpec(*, image_repo: Optional[str] = None, job_name: Optional[str] = None, num_workers: Optional[int] = None, function_name: Optional[str] = None, force_rebuild: bool = False, max_batch_rows: int = 1024, warehouse: Optional[str] = None, cpu_requests: Optional[str] = None, memory_requests: Optional[str] = None, gpu_requests: Optional[str] = None, replicas: Optional[int] = None)

Bases: BaseModel

Specification for batch inference job execution.

Defines the compute resources, job settings, and execution parameters for running batch inference jobs in Snowflake.

image_repo

Container image repository for the inference job. If not specified, uses the default repository.

Type:

Optional[str]

job_name

Custom name for the batch inference job. If not provided, a name will be auto-generated in the form of “BATCH_INFERENCE_<UUID>”.

Type:

Optional[str]

num_workers

The number of workers to run the inference service for handling requests in parallel within an instance of the service. By default, it is set to 2*vCPU+1 of the node for CPU based inference and 1 for GPU based inference. For GPU based inference, please see best practices before playing with this value.

Type:

Optional[int]

function_name

Name of the specific function to call for inference. Required when the model has multiple inference functions.

Type:

Optional[str]

force_rebuild

Whether to force rebuilding the container image even if it already exists. Defaults to False.

Type:

bool

max_batch_rows

Maximum number of rows to process in a single batch. Defaults to 1024. Larger values may improve throughput.

Type:

int

warehouse

Snowflake warehouse to use for the batch inference job. If not specified, uses the session’s current warehouse.

Type:

Optional[str]

cpu_requests

The cpu limit for CPU based inference. Can be an integer, fractional or string values. If None, we attempt to utilize all the vCPU of the node.

Type:

Optional[str]

memory_requests

The memory limit for inference. Can be an integer or a fractional value, but requires a unit (GiB, MiB). If None, we attempt to utilize all the memory of the node.

Type:

Optional[str]

gpu_requests

The gpu limit for GPU based inference. Can be integer or string values. Use CPU if None.

Type:

Optional[str]

replicas

Number of SPCS job nodes used for distributed inference. If not specified, defaults to 1 replica.

Type:

Optional[int]

Example

>>> job_spec = JobSpec(
...     job_name="my_inference_job",
...     num_workers=4,
...     cpu_requests="2",
...     memory_requests="8Gi",
...     max_batch_rows=2048
... )
Copy

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Attributes

model_computed_fields: ClassVar[dict[str, ComputedFieldInfo]] = {}

A dictionary of computed field names and their corresponding ComputedFieldInfo objects.

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

model_extra

Get extra fields set during validation.

Returns:

A dictionary of extra fields, or None if config.extra is not set to “allow”.

model_fields: ClassVar[dict[str, FieldInfo]] = {'cpu_requests': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'force_rebuild': FieldInfo(annotation=bool, required=False, default=False), 'function_name': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'gpu_requests': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'image_repo': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'job_name': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'max_batch_rows': FieldInfo(annotation=int, required=False, default=1024), 'memory_requests': FieldInfo(annotation=Union[str, NoneType], required=False, default=None), 'num_workers': FieldInfo(annotation=Union[int, NoneType], required=False, default=None), 'replicas': FieldInfo(annotation=Union[int, NoneType], required=False, default=None), 'warehouse': FieldInfo(annotation=Union[str, NoneType], required=False, default=None)}

Metadata about the fields defined on the model, mapping of field names to [FieldInfo][pydantic.fields.FieldInfo].

This replaces Model.__fields__ from Pydantic V1.

model_fields_set

Returns the set of fields that have been explicitly set on this model instance.

Returns:

A set of strings representing the fields that have been set,

i.e. that were not filled from defaults.

image_repo: Optional[str]
job_name: Optional[str]
num_workers: Optional[int]
function_name: Optional[str]
force_rebuild: bool
max_batch_rows: int
warehouse: Optional[str]
cpu_requests: Optional[str]
memory_requests: Optional[str]
gpu_requests: Optional[str]
replicas: Optional[int]