train = ShellTask(
⚠ name="train",
      debug=True,
      script="""
      set -ex
      NPROC_PER_NODE=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
      python -m torch.distributed.launch --nproc_per_node $NPROC_PER_NODE pocs/wav2vec2/flyte.py --cfg {inp>
      """,
      inputs=kwtypes(cfg=str, save_directory=str),
      task_config=Pod(pod_spec=generate_gpu_pod_spec_for_task(), primary_container_name="primary"),
  )

⚠ @workflow
  def wf(cfg: str, save_directory: str) -> None:
      train(cfg=cfg, save_directory=save_directory)

@task(task_config=Pod(…))

k8sPod

container

k8sPod