import konduktor
# Describe the work to run on each node.
task = konduktor.Task(
name="python-api-demo",
run="python train.py --epochs 10",
workdir=".",
envs={"myenv": "foo"},
)
# Specify compute requirements and image
resources = konduktor.Resources(
cpus=4,
memory=12,
accelerators="H100",
image_id="docker.io/ryanattrainy/pytorch-mnist:cpu",
labels={"mylabel": "bar"},
job_config={
"max_restarts": 3,
"completions": 2,
},
)
task.set_resources(resources)
job_id = konduktor.launch(task)
print(f"Submitted job: {job_id}")