added support for multinode aws batch job type

Author: SebastianScherer88

src/zenml/integrations/aws/flavors/aws_batch_step_operator_flavor.py

@@ -46,6 +46,7 @@ class AWSBatchStepOperatorSettings(BaseSettings):
             "execution. Example: {'LOG_LEVEL': 'INFO', 'DEBUG_MODE': 'False'}",
     )
     node_count: PositiveInt = Field(
+        default=1,
         description="The number of AWS Batch nodes to run the step on. If > 1," \
         "an AWS Batch multinode job will be run, with the network connectivity" \
         "between the nodes provided by AWS Batch. See https://docs.aws.amazon.com/batch/latest/userguide/multi-node-parallel-jobs.html" \

src/zenml/integrations/aws/step_operators/__init__.py

@@ -18,5 +18,6 @@
 )
 from zenml.integrations.aws.step_operators.aws_batch_step_operator import (  # noqa: F401
     AWSBatchStepOperator,
+    get_aws_batch_context
 )
-__all__ = ["SagemakerStepOperator","AWSBatchStepOperator"]
+__all__ = ["SagemakerStepOperator","AWSBatchStepOperator","get_aws_batch_context"]

src/zenml/integrations/aws/step_operators/aws_batch_step_operator.py

@@ -25,7 +25,8 @@
     Literal,
     cast,
 )
-from pydantic import BaseModel, PositiveInt
+from pydantic import BaseModel, PositiveInt, Field
+from pydantic_settings import BaseSettings as PydanticBaseSettings
 import boto3
 
 from zenml.client import Client
@@ -57,6 +58,43 @@
 BATCH_DOCKER_IMAGE_KEY = "aws_batch_step_operator"
 _ENTRYPOINT_ENV_VARIABLE = "__ZENML_ENTRYPOINT"
 
+class AWSBatchContext(PydanticBaseSettings):
+    """A utility to access the AWS Batch job context environment variables."""
+    main_node_index: int = Field(
+        description="This variable is set to the index number of the job's "
+            "main node. Your application code can compare the "
+            "AWS_BATCH_JOB_MAIN_NODE_INDEX to the AWS_BATCH_JOB_NODE_INDEX on "
+            "an individual node to determine if it's the main node.",
+        alias="AWS_BATCH_JOB_MAIN_NODE_INDEX"
+    )
+
+    main_node_address: str = Field(
+        description="This variable is only set in multi-node parallel job "
+            "child nodes. This variable isn't present on the main node. This "
+            "variable is set to the private IPv4 address of the job's main "
+            "node. Your child node's application code can use this address to "
+            "communicate with the main node.",
+        alias="AWS_BATCH_JOB_MAIN_NODE_PRIVATE_IPV4_ADDRESS"
+    )
+
+    node_index: int = Field(
+        description="This variable is set to the node index number of the "
+            "node. The node index begins at 0, and each node receives a unique"
+            " index number. For example, a multi-node parallel job with 10 "
+            "children has index values of 0-9.",
+        alias="AWS_BATCH_JOB_NODE_INDEX"
+    )
+
+    num_nodes: int = Field(
+        description="This variable is set to the number of nodes that you have"
+            "requested for your multi-node parallel job.",
+        alias="AWS_BATCH_JOB_NUM_NODES"
+    )
+
+def get_aws_batch_context() -> AWSBatchContext:
+    """Utility to retrieve the AWS Batch runtime context."""
+    return AWSBatchContext()
+
 class AWSBatchJobDefinitionContainerProperties(BaseModel):
     """An AWS Batch job subconfiguration model for a container specification."""
     image: str

tests/integration/integrations/aws/step_operators/test_aws_batch_step_operator.py

@@ -16,9 +16,25 @@
 import pytest
 
 from zenml.config.resource_settings import ResourceSettings
-from zenml.integrations.aws.step_operators.aws_batch_step_operator import AWSBatchStepOperator
+from zenml.integrations.aws.step_operators.aws_batch_step_operator import AWSBatchStepOperator, get_aws_batch_context
+
+def test_aws_batch_context(monkeypatch):
+    """Tests the AWSBatchContext class."""
+
+    monkeypatch.setenv('AWS_BATCH_JOB_MAIN_NODE_INDEX',0)
+    monkeypatch.setenv('AWS_BATCH_JOB_MAIN_NODE_PRIVATE_IPV4_ADDRESS','test-address')
+    monkeypatch.setenv('AWS_BATCH_JOB_NODE_INDEX',1)
+    monkeypatch.setenv('AWS_BATCH_JOB_NUM_NODES',2)
+
+    test_aws_batch_context = get_aws_batch_context()
+    assert test_aws_batch_context.main_node_index == 0
+    assert test_aws_batch_context.main_node_address == 'test-address'
+    assert test_aws_batch_context.node_index == 1
+    assert test_aws_batch_context.num_nodes == 2
 
 def test_aws_batch_step_operator_map_environment():
+    """Tests the AWSBatchStepOperator's map_environment class method."""
+
     test_environment = {'key_1':'value_1','key_2':'value_2'}
     expected = [
         {
@@ -83,4 +99,6 @@ def test_aws_batch_step_operator_map_environment():
         ]
 )
 def test_aws_batch_step_operator_map_resource_settings(test_resource_settings,expected):
+    """Tests the AWSBatchStepOperator's map_resource_settings class method."""
+
     assert AWSBatchStepOperator.map_resource_settings(test_resource_settings) == expected