Shell tasks

Run bash scripts in Flyte workflows using ShellTask.

Use ShellTask to execute bash scripts within a Flyte workflow. Shell tasks are useful when you need to run system commands, invoke external tools, or orchestrate scripts that are not written in Python.

Imports

shell_task.py

import os
from pathlib import Path
from typing import Tuple

from flytekit import kwtypes, task, workflow
from flytekit.extras.tasks.shell import OutputLocation, ShellTask
from flytekit.types.directory import FlyteDirectory
from flytekit.types.file import FlyteFile

Define a shell task

Create a ShellTask by providing a name, the bash script to execute, and optional inputs and outputs:

shell_task.py

t1 = ShellTask(
    name="task_1",
    debug=True,
    script="""
    set -ex
    echo "hey it's me {inputs.x}" >> {outputs.j}
    """,
    inputs=kwtypes(x=str),
    output_locs=[
        OutputLocation(var="j", var_type=FlyteFile, location="{inputs.x}.txt")
    ],
)

The ShellTask constructor accepts these key parameters:

Parameter	Description
name	Unique name for the task
script	The bash script to run. Use {inputs.<name>} and {outputs.<name>} as placeholders
inputs	Input type definitions using kwtypes
output_locs	Output file locations — a list of OutputLocation objects
debug	When True, prints the rendered script before execution

Inputs and outputs

Reference inputs and outputs inside the script using {inputs.<name>} and {outputs.<name>} syntax. Flyte replaces these placeholders with the actual values at execution time.

shell_task.py

t2 = ShellTask(
    name="task_2",
    debug=True,
    script="""
    set -ex
    cp {inputs.x} {outputs.j}
    """,
    inputs=kwtypes(x=FlyteFile),
    output_locs=[
        OutputLocation(var="j", var_type=FlyteFile, location="/tmp/output.txt")
    ],
)


t3 = ShellTask(
    name="task_3",
    debug=True,
    script="""
    set -ex
    cp {inputs.x} {outputs.j}
    tar -zcvf {outputs.k} {inputs.y}
    """,
    inputs=kwtypes(x=FlyteFile, y=FlyteDirectory),
    output_locs=[
        OutputLocation(var="j", var_type=FlyteFile, location="/tmp/output.txt"),
        OutputLocation(var="k", var_type=FlyteFile, location="/tmp/output.tar.gz"),
    ],
)

Create FlyteFile and FlyteDirectory inputs

Define a Python task to create the FlyteFile and FlyteDirectory inputs needed by the shell tasks. A .gitkeep placeholder file ensures the directory exists:

shell_task.py

@task
def create_entities() -> Tuple[FlyteFile, FlyteDirectory]:
    working_dir = Path(os.getcwd())
    # Create a FlyteFile
    ff = working_dir / "test.txt"
    ff.write_text("hello flyte")
    # Create a FlyteDirectory with a placeholder file
    fd = working_dir / "test_dir"
    fd.mkdir(exist_ok=True)
    (fd / ".gitkeep").touch()
    return FlyteFile(path=str(ff)), FlyteDirectory(path=str(fd))

Wire tasks into a workflow

shell_task.py

@workflow
def shell_task_wf() -> FlyteFile:
    ff, fd = create_entities()
    t1_out = t1(x="foo")
    t2_out = t2(x=t1_out)
    t3_out = t3(x=t2_out, y=fd)
    return t3_out

Run locally

shell_task.py

if __name__ == "__main__":
    print(shell_task_wf())

Or use pyflyte run:

pyflyte run shell_task.py shell_task_wf

Complete example

shell_task.py

import os
from pathlib import Path
from typing import Tuple

from flytekit import kwtypes, task, workflow
from flytekit.extras.tasks.shell import OutputLocation, ShellTask
from flytekit.types.directory import FlyteDirectory
from flytekit.types.file import FlyteFile


t1 = ShellTask(
    name="task_1",
    debug=True,
    script="""
    set -ex
    echo "hey it's me {inputs.x}" >> {outputs.j}
    """,
    inputs=kwtypes(x=str),
    output_locs=[
        OutputLocation(var="j", var_type=FlyteFile, location="{inputs.x}.txt")
    ],
)


t2 = ShellTask(
    name="task_2",
    debug=True,
    script="""
    set -ex
    cp {inputs.x} {outputs.j}
    """,
    inputs=kwtypes(x=FlyteFile),
    output_locs=[
        OutputLocation(var="j", var_type=FlyteFile, location="/tmp/output.txt")
    ],
)


t3 = ShellTask(
    name="task_3",
    debug=True,
    script="""
    set -ex
    cp {inputs.x} {outputs.j}
    tar -zcvf {outputs.k} {inputs.y}
    """,
    inputs=kwtypes(x=FlyteFile, y=FlyteDirectory),
    output_locs=[
        OutputLocation(var="j", var_type=FlyteFile, location="/tmp/output.txt"),
        OutputLocation(var="k", var_type=FlyteFile, location="/tmp/output.tar.gz"),
    ],
)


@task
def create_entities() -> Tuple[FlyteFile, FlyteDirectory]:
    working_dir = Path(os.getcwd())
    ff = working_dir / "test.txt"
    ff.write_text("hello flyte")
    fd = working_dir / "test_dir"
    fd.mkdir(exist_ok=True)
    (fd / ".gitkeep").touch()
    return FlyteFile(path=str(ff)), FlyteDirectory(path=str(fd))


@workflow
def shell_task_wf() -> FlyteFile:
    ff, fd = create_entities()
    t1_out = t1(x="foo")
    t2_out = t2(x=t1_out)
    t3_out = t3(x=t2_out, y=fd)
    return t3_out


if __name__ == "__main__":
    print(shell_task_wf())

The full source for this example is available in the Flytesnacks repository.