
from kubernetes.client.models import V1Tolerationvertex_settings = VertexOrchestratorSettings(

pod_settings={

"affinity": {

"nodeAffinity": {

"requiredDuringSchedulingIgnoredDuringExecution": {

"nodeSelectorTerms": [

"matchExpressions": [

"key": "node.kubernetes.io/name",

"operator": "In",

"values": ["my_powerful_node_group"],

},

"tolerations": [

V1Toleration(

key="node.kubernetes.io/name",

operator="Equal",

value="",

effect="NoSchedule"

If your pipelines steps have certain hardware requirements, you can specify them as ResourceSettings:

resource_settings = ResourceSettings(cpu_count=8, memory="16GB")

These settings can then be specified on either pipeline-level or step-level:

# Either specify on pipeline-level

@pipeline(

settings={

"orchestrator.vertex": vertex_settings,

"resources": resource_settings,

def my_pipeline():

...

# OR specify settings on step-level

@step(

settings={

"orchestrator.vertex": vertex_settings,

"resources": resource_settings,

def my_step():

...

Check out the SDK docs for a full list of available attributes and this docs page for more information on how to specify settings.

For more information and a full list of configurable attributes of the Vertex orchestrator, check out the SDK Docs .

Enabling CUDA for GPU-backed hardware

Note that if you wish to use this orchestrator to run steps on a GPU, you will need to follow the instructions on this page to ensure that it works. It requires adding some extra settings customization and is essential to enable CUDA for the GPU to give its full acceleration.

PreviousKubernetes Orchestrator

NextAWS Sagemaker Orchestrator

Last updated 18 days ago