[Core feature] Improve flytekitplugins-kfpytorch user experience with default pod template and other reasonable defaults about flyte HOT 2 OPEN

Fully agree that this should be simplified.

Questions to discuss:

• Shared memory:

  • Do we need to specify an amount? We've had this volume configured in our default pod template and never had any issues:

     volumeMounts:
        - mountPath: /dev/shm
          name: dshm
      volumes:
        - name: dshm
          emptyDir:
            medium: Memory

  • Do we try to merge this into the pod template a user might have provided to the task or should the shared memory volume only be added if the user doesn’t provide a pod template?

• Timeouts:
  • For the join timeout I feel we should consider the scenario that some workers have a hot start (node is up and image is cached) while other workers have a cold start, i.e. node needs to be scaled up and image has to be pulled. I feel 15 minutes, as you specified, is a good value here. Are there other opinions?
  • Clarify whether the timeout in the rdzv config is the same timeout as in torch.distributed.init_process_group and decide on a reasonable default value.

from flyte.

Just to circle back to this: we opted to:

1. Initialize the Elastic task config with a default pod template:

PodTemplate(
    primary_container_name="pytorch",
    pod_spec=V1PodSpec(
        containers=[
            V1Container(
                name="pytorch",
                volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm")]
            )
        ],
        volumes=[V1Volume(name="dshm", empty_dir=V1EmptyDirVolumeSource(medium="Memory"))]
    ),
)

This would not be exposed to the end user, but they could still override this by specifying pod_template in the @task decorator.

2. Set the default rdvz_configs join_timeout to 900 (15 minutes). Digging into the pytorch docs/code, it looks like timeout and join_timeout are the same, I think timeout is a legacy argument for the `EtcdRendezvousHandler:
   • https://pytorch.org/docs/stable/_modules/torch/distributed/elastic/rendezvous/dynamic_rendezvous.html#create_handler
     

• https://pytorch.org/docs/stable/elastic/rendezvous.html#etcd-rendezvous-legacy
  

from flyte.