Encoding workflow dependencies in AWS Batch
Most users of HPC or Batch systems need to analyze data with multiple operations to get meaningful results. That’s really driven by the nature of scientific research or engineering processes – it’s rare that a single task generates the insight you need.
AWS Batch customers are no exception, of course, which is why Batch supports nested dependencies to create relationships between jobs. In today’s post, we’ll walk you through how to encode job dependencies in Batch, using a simple machine learning pipeline as an example.
Let’s take the following high-level diagram depicting a simple machine learning pipeline.
The diagram shows a straight-forward set of serial steps (or jobs) that transform data into a trained model that can be used for inference within some other system. Even for this simple example workflow there are a few things to note: (1) each step can have different requirements for CPU, memory, storage, or even a GPU, so it makes sense to encapsulate each as a separate job definition; (2) steps further in the chain depend on data from an earlier step; and because of that, (3) if an earlier step does not succeed, then it makes no sense to provision resources for the dependent processes that follow.
When using AWS Batch, you have a choice for how to encode the dependencies between jobs:
- Use the SubmitJob API to define the dependencies as you are submitting the job requests. This job dependencies are directly encoded within, and managed by, Batch.
- Define the job dependencies outside of Batch using a workflow framework, such as Apache Airflow or AWS Step Functions.
I’ll describe both of these methods, starting with directly encoding job dependencies in the Batch API. I’ll use AWS Step Functions, which integrates natively with AWS Batch, to describe encoding dependencies outside of Batch.
Encoding job dependencies at runtime with AWS Batch
When you submit a job request to AWS Batch, you have the option of defining a dependency on a previously submitted job. This is what we mean by “runtime” – that some API requests have been made and you can refer to these in subsequent requests in an iterative manner and that job dependencies are not formally defined before submitting requests are acknowledged by the Batch service.
The following example shows how to submit a job using the AWS CLI and then using the returned job ID in a subsequent job request to define the dependency.
As the following animation (Figure 2) shows, the Batch scheduler will keep the second job (Job B) in the PENDING state until the first job (Job A) completes successfully. After Job A succeeds, Job B progresses from PENDING to RUNNABLE and eventually to the SUCCEEDED state.
If Job A fails, Job B also immediately transitions from PENDING to FAILED, as shown in Figure 2.
Defining dependencies for array jobs
The previous example shows how dependencies work for basic jobs, which are individual units of work. AWS Batch also allows you to submit array jobs – which allow you to submit multiple units of work that share common parameters such as the application, vCPUs, and memory in a single API request.
When you submit an array job request with an array size of 1000, a single parent job gets created that spawns 1000 child basic jobs. Each child job has the same base job ID as the parent, but adds the array index as a suffix. The array index is also passed to the child job as the environment variable AWS_BATCH_JOB_ARRAY_INDEX.You can define two types of job dependencies for array jobs, SEQUENTIAL or N_TO_N. A SEQUENTIAL dependency type for array jobs will result in the basic child jobs to be dependent on the previously indexed job to succeed before it can start work. The job with index 1 will not start until index 0 succeeds. This is a way to define a sequence of operations within the elements of the array job. Figure 3 is an illustration of the SEQUENTIAL array job dependency. Failure of a parent will cascade down the line and also transition the parent array job to FAILED.
The second array job dependency type, N_TO_N, is between array jobs. Given two array jobs, A and B, and given that B depends on A, then each child job of B must wait for the corresponding index of child jobs in A to complete before it can begin. For example, Job B:1 can’t start until job A:1 finishes. Figure 4 illustrates the dependency definition A and B, and the child job dependencies.
If you cancel or shut down a parent array job, all of the child jobs are cancelled or shut down with it. You can cancel or shut down individual child jobs (which moves them to the FAILED status) without affecting the other child jobs. However, if a child array job enters a FAILED status, the parent job also fails.
Finally, dependencies between basic and array jobs act as they would for dependencies between basic jobs, where a dependency from a basic on an array job parent will only start the basic job if all the array child-jobs complete successfully.
Defining job dependencies outside of AWS Batch
Another way to define dependencies between jobs is to do so outside of AWS Batch using a workflow framework such as Apache Airflow or AWS Step Functions (and there are many more – Batch is widely supported by the community).
The main advantages of leveraging a workflow framework are:
- Workflow check-pointing for manual reviews, or taking advantage of flexible compute options like EC2 Spot Instances
- Restarting workflows from a checkpoint so that you don’t waste compute resources running jobs that already have results
- Leveraging multiple compute resources or services depending on the task requirement
- Advanced retry strategies, such as adding more job resources (e.g., more memory) or switching resource pools
- Multiple error-handling routines that are triggered based on the type of error encountered
AWS Step Functions allows you to define a workflow as a state machine using the Amazon States Language – a JSON-based, structured language. The diagram on the right side of this page shows the visual representation of a state machine of the previous machine learning example. Each of the task states integrate with AWS Batch to perform work. By default, when a state reports an error, AWS Step Functions causes the execution to fail entirely and future tasks are not sent to AWS Batch at all. However, Step Functions has some advanced error handling features that you can take advantage of to deal with different types of errors, such as runtime parameters, timeouts, and permissions errors. Being able to define fallback states and advanced retry strategies are the main advantage of working with full-featured workflow frameworks such as Step Functions.
You can read about common workflow design patterns in HPC, and how to encode those patterns using Step Functions and Batch, in the Orchestrating high performance computing with AWS Step Functions and AWS Batch blog post by AWS solutions architects, Dan Fox and Sabha Parameswaran.
We discussed two methods for encoding dependencies between jobs submitted to AWS Batch. The first was leveraging the features of the AWS Batch API for basic and array jobs as you submit work to the job queue. This has the advantage of working within a single service for all your workflow needs.
The second method described encoding dependencies outside of AWS Batch using AWS Step Functions as an example. Use external dependency management when you want to take advantage of advanced error handling capabilities, have to support multiple job schedulers and resources, or are already have experience working within a framework that fits into your current systems.