hi guys is there a way to set the task s pod template at run Flyte #ask-the-community

hi guys, is there a way to set the task’s pod temp...

Brian Tang

01/29/2024, 4:13 AM

hi guys, is there a way to set the task’s pod template at runtime? i went through this, but it’s not exactly what we want. our use-case is similar to `with_overrides`: we want to target a different nodegroup for the task to run on using k8s labels. we currently have a dynamic workflow that computes the dataset size and estimated training workload, and then uses

with_overrides

to specify mem, cpu and gpu. we want to use this estimation to target a diff nodegroup w better GPUs

Brian Tang

01/29/2024, 8:20 AM

we went w this approach, involving creating different tasks with diff pod templates:

Copy code

def task():
    ...

@task(pod_template=with_gpu)
def task_wrapper_gpu():
    task()

@task(pod_template=with_big_gpu)
def task_wrapper_big_gpu():
    task()

then in our dynamic workflow, we select the task based on the runtime params:

Copy code

@dynamic
def wf():
    resource = sizing()
    
    if resource == "big":
        t = task_wrapper_big_gpu
    else:
        t = task_wrapper_gpu()
    
    out = t()

Dan Rammer (hamersaw)

01/29/2024, 3:47 PM

PodTemplate should probably be supported as an override - would you mind filing an issue?

Dan Rammer (hamersaw)

01/29/2024, 3:49 PM

Also in this scenario have you tried using a conditional from a flyte workflow? This would negate the need to stand up a new pod for a dynamic workflow just to do this.

2 Views

Open in Slack

Previous Next