Official AWS Batch step operator

[SebastianScherer88]

Contact Details [Optional]

scherersebastian@yahoo.de

Feature Description

It would be great to get an official version of this interesting AWS Batch step operator into the main zenml library.

Problem or Use Case

I think it's quite common for people to have

• AWS infra and
• heterogenuous compute requirements in their pipeline steps - not everything needs to run on sagemaker

running a local (docker) orchestrator that can push individual components to powerful remote execution engines like AWS Batch sounds super useful to me - its definitely something i would be interested in (I work in ML as an ops engineer).

Happy to contribute based on the linked reference plugin implementation provided by you guys

Proposed Solution

A hardened, more configurable version of the linked plugin implementation that

• allows for step resource configuration that get mapped canonically (where possible) onto AWS Batch resource specs
• default to AWS infra settings that are compatible with current terraform setup utils (where possible)
• integrates with every orchestrator that honours the canonical steplauncher appoach (i.e. not the LocalDockerOrchestrator)

Alternatives Considered

The official Sagemaker step operator. AWS Batch would be a cheaper (no ml.... instance sagemaker type $ markup) and more flexible way of launching scalable custom compute jobs

Additional Context

Implementation draft (unofficial AWS Batch step operator plugin)

Priority

Low - Nice to have

Code of Conduct

• I agree to follow this project's Code of Conduct

[htahir1]

@SebastianScherer88 thanks for raising this — a step operator for this use-case would be very useful.

If you’re up for contributing it, a good starting point is the small example plugin that shows the skeleton for a batch-style step operator:
👉 https://github.com/zenml-io/zenml-plugins/tree/main/step_operator_batch

For structure and conventions, these existing step operators are good references:

• SageMaker: https://github.com/zenml-io/zenml/blob/main/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py
• Vertex AI: https://github.com/zenml-io/zenml/blob/main/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py
• Azure ML: https://github.com/zenml-io/zenml/blob/main/src/zenml/integrations/azure/step_operators/azureml_step_operator.py
• Kubernetes: https://github.com/zenml-io/zenml/blob/main/src/zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py
• Spark: https://github.com/zenml-io/zenml/blob/main/src/zenml/integrations/spark/step_operators/spark_step_operator.py

A few implementation notes that typically matter:

• Implement launch(...) to submit the step to your target backend and handle env/config transfer.
• Implement get_docker_builds(...) if the operator needs a custom image per step.
• Provide a validator to enforce required remote components (e.g., remote artifact store / container registry).
• Expose settings_class and (if needed) an entrypoint_config_class for runtime config.
• The abstract interface lives here: BaseStepOperator.

If you run into design questions, happy to discuss/review an outline or early PR. And once there’s a working PR, we’ll sort you out with some ZenML swag as a small thank-you. 🙂

[SebastianScherer88]

cool i'll have a look at those resources. i'll let you guys know when i got PR draft ready content and take it from there? i dont think i can push feature branches or create PRs atm 🤔

[SebastianScherer88]

@htahir1 heya 👋 , i got a first draft and some passing unit tests. a proper e2e test would require some non-trivial AWS resource configuration 🤔 do you guys have something on your side? i could stand something up in pulumi if you have an AWS playground env for functional testing purposes as well 🏗️ happy to discuss how to proceed on this

[htahir1]

@SebastianScherer88 Epic - mind joining our Slack and we can figure that out? https://zenml.io/slack (DM me when joined, on Monday I can start the process to give you some access). Can you also push the code so I can take an initial look? Thanks!

[SebastianScherer88]

just pinged you on slack 🔔

sadly i can't push branches to this github repo:

(.zenml) C:\Users\scher\Repositories\zenml>git push --set-upstream origin feature/aws-step-operator
remote: Permission to zenml-io/zenml.git denied to SebastianScherer88.
fatal: unable to access 'https://github.com/zenml-io/zenml.git/': The requested URL returned error: 403