https://flyte.org logo
#ask-the-community
Title
# ask-the-community
b

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
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()
d

Dan Rammer (hamersaw)

01/29/2024, 3:47 PM
PodTemplate should probably be supported as an override - would you mind filing an issue?
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.