#1677 Fix config of user facing execution parameters in spawning elastic tasks Pull request opened by fg91 TL;DR When using

@task(task_config=flytekitplugins.kfpytorch.Elastic())

, the task function is started in a number of worker processes using torch

elastic_launch

(

torchrun

). The processes can be created using fork or spawn which is controlled by the arg

Elastic(start_method=...)

. When using

fork

, the child process inherits a copy of the parent process' stack including the flyte context and the user facing execution parameters

ctx = flytekit.current_context()

. When spawning, however, fresh processes are started and the flyte context and the execution parameters are not transferred to the child process currently. This means that within a task with ``@task(task_config=Elastic(start_method="spawn"))` the execution id and the checkpoint cannot be accessed from the execution parameters. This PR fixes this by setting up the flyte context in the spawned worker processes. Type ☑︎ Bug Fix ☐ Feature ☐ Plugin Are all requirements met? ☑︎ Code completed ☐ Smoke tested ☑︎ Unit tests added ☐ Code documentation added ☑︎ Any pending items have an associated Issue Complete description In the spawned worker processes I call

flytekit.bin.entrypoint.setup_execution

which sets up the flyte context the same way as when a normal python task is started. Raw data prefix and checkpoint pathes are transferred from the parent process. Tracking Issue NA Follow-up issue NA flyteorg/flytekit ✅ All checks have passed 30/30 successful checks