Custom Components

Tip

This guide covers writing a custom component (a Python function returning an AppDef) and launching it on both local and container-based schedulers. Check the builtins first – you may not need a custom component.

Prerequisites: Quickstart (installation) and Basic Concepts (core concepts).

Builtins

Discover available builtins with torchx builtins:

$ torchx run utils.echo --msg "Hello :)"

Hello World

Create my_app.py:

import sys
import argparse

def main(user: str) -> None:
    print(f"Hello, {user}!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Hello world app"
    )
    parser.add_argument(
        "--user",
        type=str,
        help="the person to greet",
        required=True,
    )
    args = parser.parse_args(sys.argv[1:])

    main(args.user)

Next, write a component – a factory function that returns an AppDef.

Create my_component.py:

import torchx.specs as specs

def greet(user: str, image: str = "my_app:latest") -> specs.AppDef:
    return specs.AppDef(
        name="hello_world",
        roles=[
            specs.Role(
                name="greeter",
                image=image,
                entrypoint="python",
                args=[
                    "-m", "my_app",
                    "--user", user,
                ],
            )
        ],
    )

Run the component with the local_cwd scheduler (executes in the current directory):

$ torchx run --scheduler local_cwd my_component.py:greet --user "your name"

The same launch works from Python using the Runner:

from torchx.runner import get_runner

with get_runner() as runner:
    # reference the file:function component, same as the CLI
    app_handle = runner.run_component(
        "my_component.py:greet",
        ["--user", "your name"],
        scheduler="local_cwd",
    )

    # alternatively, call the component function directly
    from my_component import greet

    app = greet(user="your name")
    app_handle = runner.run(app, scheduler="local_cwd")

For container-based schedulers, build a Docker image:

Note

Requires Docker: https://docs.docker.com/get-docker/

Create Dockerfile.custom:

FROM ghcr.io/pytorch/torchx:latest

ADD my_app.py .

Build the image:

$ docker build -t my_app:latest -f Dockerfile.custom .

Launch on the local Docker scheduler:

$ torchx run --scheduler local_docker my_component.py:greet --image "my_app:latest" --user "your name"

Or push and launch on a Kubernetes cluster:

$ docker push my_app:latest
$ torchx run --scheduler kubernetes my_component.py:greet --image "my_app:latest" --user "your name"

See also

Quick Reference

Single-page reference with imports, types, and copy-pasteable recipes.

Component Best Practices

Best practices for entrypoints, simplicity, named resources, and testing.

Advanced Usage

Registering custom components as CLI builtins (via entry points; @register migration pending for components).

Overview

Browse the builtin component library.