Custom Components

This is a guide on how to build a simple app and custom component spec and launch it via two different schedulers.

See the Quickstart Guide for installation and basic usage.

Hello World

Lets start off with writing a simple “Hello World” python app. This is just a normal python program and can contain anything you’d like.

Note

This example uses Jupyter Notebook %%writefile to create local files for example purposes. Under normal usage you would have these as standalone files.

%%writefile 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)

Writing my_app.py

Now that we have an app we can write the component file for it. This function allows us to reuse and share our app in a user friendly way.

We can use this component from the torchx cli or programmatically as part of a pipeline.

%%writefile 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,
                ],
            )
        ],
    )

Writing my_component.py

We can execute our component via torchx run. The local_cwd scheduler executes the component relative to the current directory.

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

torchx 2022-03-28 19:31:45 INFO     Log files located in: /tmp/torchx_lehw4e9s/torchx/hello_world-t7kj2b3gpgp3w/greeter/0
torchx 2022-03-28 19:31:46 INFO     Waiting for the app to finish...
greeter/0 Hello, your name!
torchx 2022-03-28 19:31:47 INFO     Job finished: SUCCEEDED

local_cwd://torchx/hello_world-t7kj2b3gpgp3w

If we want to run in other environments, we can build a Docker container so we can run our component in Docker enabled environments such as Kubernetes or via the local Docker scheduler.

Note

This requires Docker installed and won’t work in environments such as Google Colab. If you have not done so already follow the install instructions on: https://docs.docker.com/get-docker/

%%writefile Dockerfile.custom

FROM ghcr.io/pytorch/torchx:0.1.0rc1

ADD my_app.py .

Writing Dockerfile.custom

Once we have the Dockerfile created we can create our docker image.

%%sh
docker build -t my_app:latest -f Dockerfile.custom .

Step 1/2 : FROM ghcr.io/pytorch/torchx:0.1.0rc1
 ---> 3dbec59e8049
Step 2/2 : ADD my_app.py .
 ---> Using cache
 ---> fa8c0e244754
Successfully built fa8c0e244754
Successfully tagged my_app:latest

We can then launch it on the local scheduler.

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

torchx 2022-03-28 19:31:48 INFO     Building workspace: file:///home/tristanr/Developer/torchx-0.1.2/docs/source for role[0]: greeter, image: my_app:latest
torchx 2022-03-28 19:31:50 WARNING  failed to pull image my_app:latest, falling back to local: 404 Client Error for http+docker://localhost/v1.41/images/create?tag=latest&fromImage=my_app: Not Found ("pull access denied for my_app, repository does not exist or may require 'docker login': denied: requested access to the resource is denied")
torchx 2022-03-28 19:31:57 INFO     Done building workspace
torchx 2022-03-28 19:31:57 INFO     New image: sha256:7fdf685edad448a965a5ce479727b45b5eeef168d324d282e9a88e6af3b590f4 built from workspace
torchx 2022-03-28 19:31:58 INFO     Waiting for the app to finish...
greeter/0 Hello, your name!
torchx 2022-03-28 19:31:59 INFO     Job finished: SUCCEEDED

local_docker://torchx/hello_world-rgzwt93vjpgm7

If you have a Kubernetes cluster you can use the Kubernetes scheduler to launch this on the cluster instead.

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

Builtins

TorchX also provides a number of builtin components with premade images. You can discover them via:

%%sh
torchx builtins

Found 10 builtin components:
  1. utils.binary
  2. utils.booth
  3. utils.copy
  4. utils.echo
  5. utils.python
  6. utils.sh
  7. utils.touch
  8. serve.torchserve
  9. metrics.tensorboard
 10. dist.ddp

You can use these either from the CLI, from a pipeline or programmatically like you would any other component.

%%sh
torchx run utils.echo --msg "Hello :)"

torchx 2022-03-28 19:32:02 INFO     Building workspace: file:///home/tristanr/Developer/torchx-0.1.2/docs/source for role[0]: echo, image: ghcr.io/pytorch/torchx:0.1.2
torchx 2022-03-28 19:32:09 INFO     Done building workspace
torchx 2022-03-28 19:32:09 INFO     New image: sha256:bee1c5460939f1b51f14d96572586357f8fc59562403ea82ddd2c3688a6b4c60 built from workspace
torchx 2022-03-28 19:32:09 INFO     Waiting for the app to finish...
torchx 2022-03-28 19:32:10 INFO     Job finished: SUCCEEDED
echo/0 Hello :)

local_docker://torchx/echo-x2lkw7lsv14vpc