AWS Batch

This contains the TorchX AWS Batch scheduler which can be used to run TorchX components directly on AWS Batch.

This scheduler is in prototype stage and may change without notice.

Prerequisites

You’ll need to create an AWS Batch queue configured for multi-node parallel jobs.

See https://docs.aws.amazon.com/batch/latest/userguide/Batch_GetStarted.html for how to setup a job queue and compute environment. It needs to be backed by EC2 for multi-node parallel jobs.

See https://docs.aws.amazon.com/batch/latest/userguide/multi-node-parallel-jobs.html for more information on distributed jobs.

If you want to use workspaces and container patching you’ll also need to configure a docker registry to store the patched containers with your changes such as AWS ECR.

See https://docs.aws.amazon.com/AmazonECR/latest/userguide/getting-started-cli.html#cli-create-repository for how to create a image repository.

class torchx.schedulers.aws_batch_scheduler.AWSBatchScheduler(session_name: str, client: Any | None = None, log_client: Any | None = None, docker_client: DockerClient | None = None)

Bases: DockerWorkspaceMixin, Scheduler[Opts]

AWSBatchScheduler is a TorchX scheduling interface to AWS Batch.

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

Authentication is loaded from the environment using the boto3 credential handling.

Config Options

    usage:
        queue=QUEUE,[user=USER],[image_repo=IMAGE_REPO],[privileged=PRIVILEGED],[share_id=SHARE_ID],[priority=PRIORITY],[job_role_arn=JOB_ROLE_ARN],[execution_role_arn=EXECUTION_ROLE_ARN],[ulimits=ULIMITS],[quiet=QUIET]

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

    optional arguments:
        user=USER (str, runner)
            The username to tag the job with. `getpass.getuser()` if not specified.
        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
        privileged=PRIVILEGED (bool, False)
            If true runs the container with elevated permissions. Equivalent to running with `docker run --privileged`.
        share_id=SHARE_ID (str, None)
            The share identifier for the job. This must be set if and only if the job queue has a scheduling policy.
        priority=PRIORITY (int, 0)
            The scheduling priority for the job within the context of share_id. Higher number (between 0 and 9999) means higher priority. This will only take effect if the job queue has a scheduling policy.
        job_role_arn=JOB_ROLE_ARN (str, None)
            The Amazon Resource Name (ARN) of the IAM role that the container can assume for AWS permissions.
        execution_role_arn=EXECUTION_ROLE_ARN (str, None)
            The Amazon Resource Name (ARN) of the IAM role that the ECS agent can assume for AWS permissions.
        ulimits=ULIMITS (list, None)
            Ulimit settings in format: name:softLimit:hardLimit (multiple separated by commas).
        quiet=QUIET (bool, False)
            whether to suppress verbose output for image building. Defaults to ``False``.

Mounts

This class supports bind mounting host directories, efs volumes and host devices.

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

• efs volume: type=volume,src=<efs id>,dst=<container path>[,readonly]

• devices: type=device,src=/dev/infiniband/uverbs0,[dst=<container path>][,perm=rwm]

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

For other filesystems such as FSx you can mount them onto the host and bind mount them into your job: https://repost.aws/knowledge-center/batch-fsx-lustre-file-system-mount

For Elastic Fabric Adapter (EFA) you’ll need to use a device mount to mount them into the container: https://docs.aws.amazon.com/batch/latest/userguide/efa.html

Compatibility

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

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[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[BatchJob]) → str

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

class torchx.schedulers.aws_batch_scheduler.BatchJob(name: str, queue: str, share_id: str | None, job_def: dict[str, object], images_to_push: dict[str, tuple[str, str]])

Reference

torchx.schedulers.aws_batch_scheduler.create_scheduler(session_name: str, client: Any | None = None, log_client: Any | None = None, docker_client: DockerClient | None = None, **kwargs: object) → AWSBatchScheduler