Kubernetes

This contains the TorchX Kubernetes scheduler which can be used to run TorchX components on a Kubernetes cluster.

Prerequisites

The TorchX Kubernetes scheduler depends on Volcano. If you’re trying to do an upgrade you’ll need to completely remove all non-Job Volcano resources and recreate.

Install Volcano:

kubectl apply -f https://raw.githubusercontent.com/volcano-sh/volcano/v1.6.0/installer/volcano-development.yaml

See the Volcano Quickstart for more information.

Pod Overlay

You can overlay arbitrary Kubernetes Pod fields on generated pods by setting the kubernetes metadata on your role. The value can be:

• A dict with the overlay structure

• A resource URI pointing to a YAML file (e.g. file://, s3://, gs://)

Merge semantics: - dict: recursive merge (upsert) - list: append by default, replace if tuple (Python) or !!python/tuple tag (YAML) - primitives: replace

from torchx.specs import Role

# Dict overlay - lists append, tuples replace
role = Role(
    name="trainer",
    image="my-image:latest",
    entrypoint="train.py",
    metadata={
        "kubernetes": {
            "spec": {
                "nodeSelector": {"gpu": "true"},
                "tolerations": [{"key": "nvidia.com/gpu", "operator": "Exists"}],  # appends
                "volumes": ({"name": "my-volume", "emptyDir": {}},)  # replaces
            }
        }
    }
)

# File URI overlay
role = Role(
    name="trainer",
    image="my-image:latest",
    entrypoint="train.py",
    metadata={
        "kubernetes": "file:///path/to/pod_overlay.yaml"
    }
)

CLI usage with builtin components:

$ torchx run --scheduler kubernetes dist.ddp \
    --metadata kubernetes=file:///path/to/pod_overlay.yaml \
    --script train.py

Example pod_overlay.yaml:

spec:
  nodeSelector:
    node.kubernetes.io/instance-type: p4d.24xlarge
  tolerations:
    - key: nvidia.com/gpu
      operator: Exists
      effect: NoSchedule
  volumes: !!python/tuple
    - name: my-volume
      emptyDir: {}

The overlay is deep-merged with the generated pod, preserving existing fields and adding or overriding specified ones.

class torchx.schedulers.kubernetes_scheduler.KubernetesScheduler(session_name: str, client: ApiClient | None = None, docker_client: DockerClient | None = None)

Bases: DockerWorkspaceMixin, Scheduler[Opts]

KubernetesScheduler is a TorchX scheduling interface to Kubernetes.

Important: Volcano is required to be installed on the Kubernetes cluster. TorchX requires gang scheduling for multi-replica/multi-role execution and Volcano is currently the only supported scheduler with Kubernetes. For installation instructions see: https://github.com/volcano-sh/volcano

This has been confirmed to work with Volcano v1.3.0 and Kubernetes versions v1.18-1.21. See https://github.com/meta-pytorch/torchx/issues/120 which is tracking Volcano support for Kubernetes v1.22.

Note

AppDefs that have more than 0 retries may not be displayed as pods if they failed. This occurs due to known bug in Volcano(as per 1.4.0 release): https://github.com/volcano-sh/volcano/issues/1651

$ pip install torchx[kubernetes]
$ torchx run --scheduler kubernetes --scheduler_args namespace=default,queue=test utils.echo --image alpine:latest --msg hello
kubernetes://torchx_user/1234
$ torchx status kubernetes://torchx_user/1234
...

Cancellation

Canceling a job aborts it while preserving the job spec for inspection and cloning via kubectl apply. Use the delete command to remove the job entirely:

$ torchx cancel kubernetes://namespace/jobname  # abort, preserves spec
$ torchx delete kubernetes://namespace/jobname  # delete completely

Config Options

    usage:
        queue=QUEUE,[namespace=NAMESPACE],[service_account=SERVICE_ACCOUNT],[priority_class=PRIORITY_CLASS],[validate_spec=VALIDATE_SPEC],[reserved_millicpu=RESERVED_MILLICPU],[reserved_memmb=RESERVED_MEMMB],[image_repo=IMAGE_REPO],[efa_device_count=EFA_DEVICE_COUNT],[quiet=QUIET]

    required arguments:
        queue=QUEUE (str)
            Volcano queue to schedule job in.

    optional arguments:
        namespace=NAMESPACE (str, default)
            Kubernetes namespace to schedule job in.
        service_account=SERVICE_ACCOUNT (str, None)
            The service account name to set on the pod specs.
        priority_class=PRIORITY_CLASS (str, None)
            The name of the PriorityClass to set on the job specs.
        validate_spec=VALIDATE_SPEC (bool, True)
            Validate job spec using Kubernetes API dry-run before submission.
        reserved_millicpu=RESERVED_MILLICPU (int, 100)
            Amount of CPU in millicores to reserve for Kubernetes system overhead (default: 100).
        reserved_memmb=RESERVED_MEMMB (int, 1024)
            Amount of memory in MB to reserve for Kubernetes system overhead (default: 1024).
        image_repo=IMAGE_REPO (str, None)
            (remote jobs) the image repository to use when pushing patched images, must have push access. Ex: example.com/your/container
        efa_device_count=EFA_DEVICE_COUNT (int, None)
            EFA device count override: None/unset=use resource spec, 0=remove EFA, N>0=set EFA count to N.
        quiet=QUIET (bool, False)
            whether to suppress verbose output for image building. Defaults to ``False``.

Mounts

Mounting external filesystems/volumes is via the HostPath and PersistentVolumeClaim support.

• hostPath volumes: type=bind,src=<host path>,dst=<container path>[,readonly]

• PersistentVolumeClaim: type=volume,src=<claim>,dst=<container path>[,readonly]

• host devices: type=device,src=/dev/foo[,dst=<container path>][,perm=rwm] If you specify a host device the job will run in privileged mode since Kubernetes doesn’t expose a way to pass –device to the underlying container runtime. Users should prefer to use device plugins.

See torchx.specs.parse_mounts() for more info.

External docs: https://kubernetes.io/docs/concepts/storage/persistent-volumes/

Resources / Allocation

To select a specific machine type you can add a capability to your resources with node.kubernetes.io/instance-type which will constrain the launched jobs to nodes of that instance type.

>>> from torchx import specs
>>> specs.Resource(
...     cpu=4,
...     memMB=16000,
...     gpu=2,
...     capabilities={
...         "node.kubernetes.io/instance-type": "<cloud instance type>",
...     },
... )
Resource(...)

Kubernetes may reserve some memory for the host. TorchX assumes you’re scheduling on whole hosts and thus will automatically reduce the resource request by a small amount to account for the node reserved CPU and memory. If you run into scheduling issues you may need to reduce the requested CPU and memory from the host values.

Compatibility

Feature	Scheduler Support
Fetch Logs	✔️
Distributed Jobs	✔️
Cancel Job	✔️
Describe Job	Partial support. KubernetesScheduler will return job and replica status but does not provide the complete original AppSpec.
Workspaces / Patching	✔️
Mounts	✔️
Elasticity	Requires Volcano >1.6

describe(app_id: str) → torchx.schedulers.api.DescribeAppResponse | None

Returns app description, or None if it no longer exists.

list(cfg: Optional[Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None]] = None) → list[torchx.schedulers.api.ListAppResponse]

Lists jobs on this scheduler.

log_iter(app_id: str, role_name: str, k: int = 0, regex: str | None = None, since: datetime.datetime | None = None, until: datetime.datetime | None = None, should_tail: bool = False, streams: torchx.schedulers.api.Stream | None = None) → Iterable[str]

Returns an iterator over log lines for the k-th replica of role_name.

Important

Not all schedulers support log iteration, tailing, or time-based cursors. Check the specific scheduler docs.

Lines include trailing whitespace (\n). When should_tail=True, the iterator blocks until the app reaches a terminal state.

Parameters:

• k – replica (node) index

• regex – optional filter pattern

• since – start cursor (scheduler-dependent)

• until – end cursor (scheduler-dependent)

• should_tail – if True, follow output like tail -f

• streams – stdout, stderr, or combined

Raises:

NotImplementedError – if the scheduler does not support log iteration

schedule(dryrun_info: AppDryRunInfo[KubernetesJob]) → str

Submits a previously dry-run request. Returns the app_id.

class torchx.schedulers.kubernetes_scheduler.KubernetesJob(images_to_push: dict[str, tuple[str, str]], resource: dict[str, Any])

Reference

torchx.schedulers.kubernetes_scheduler.create_scheduler(session_name: str, client: ApiClient | None = None, docker_client: DockerClient | None = None, **kwargs: Any) → KubernetesScheduler

torchx.schedulers.kubernetes_scheduler.app_to_resource(app: AppDef, queue: str, service_account: str | None, priority_class: str | None = None, reserved_millicpu: int = 100, reserved_memmb: int = 1024, efa_device_count: int | None = None) → dict[str, Any]

app_to_resource creates a volcano job kubernetes resource definition from the provided AppDef. The resource definition can be used to launch the app on Kubernetes.

To support macros we generate one task per replica instead of using the volcano replicas field since macros change the arguments on a per replica basis.

Volcano has two levels of retries: one at the task level and one at the job level. When using the APPLICATION retry policy, the job level retry count is set to the minimum of the max_retries of the roles.

torchx.schedulers.kubernetes_scheduler.pod_labels(app: AppDef, role_idx: int, role: Role, replica_id: int, app_id: str) → dict[str, str]

torchx.schedulers.kubernetes_scheduler.role_to_pod(name: str, role: Role, service_account: str | None, reserved_millicpu: int = 100, reserved_memmb: int = 1024, efa_device_count: int | None = None) → V1Pod

torchx.schedulers.kubernetes_scheduler.sanitize_for_serialization(obj: object) → object