Hi all, I am trying to use ktensorflow.TFJob wit a custom image. However I'm not sure what exactly gets executed. So here is a rough outline of the code:

@task(
    task_config=TfJob(
        chief=Chief(replicas=1, image=<some_image>),
        ps=PS(replicas=0, image=<some_image>),
        worker=Worker(replicas=1, image=<some_image>),
    ),
)
def func(some_arg: int):
    do somethign...

does the image <some_image> get run on each pod

yes, all the pods will use same image worker and master will run the same code https://www.tensorflow.org/guide/keras/distributed_training

Thanks for your responses. so my current setup builds the image for workers/chief etc and then spawns a TFJob so that those images get executed and TF_CONFIG is set properly by flyte. This is one step of my workflow. The problem Im running into is that the image that is being run for chief/worker 's entrypoint expects cmdline args. How can I pass these in via the definition of the TFJob task?

yes, you can’t override the entrypoint

ok trying to figure out how to do this in my toolchain (bazel docker io)

what is absl? what do you want to run? download data?

absl is a command line parser

you want to use it in the pod running TF training?

yes

maybe create an init container that run absl for your task. https://docs.flyte.org/en/latest/deployment/configuration/general.html#compile-time-podtemplates

would this init container then run my custom image?

yes, you can use any image and args

can I alternatively not specify a custom image but make the TFJob decorate a container task