renaming modules and adding unit tests

Author: SebastianScherer88

pyproject.toml

@@ -250,7 +250,7 @@ dev = [
     "darglint",
     "pytest-randomly",
     "pytest-mock",
-    "pytest-clarity",
+    #"pytest-clarity",
     "pytest-instafail",
     "pytest-rerunfailures",
     "pytest-split",

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

@@ -29,7 +29,7 @@
     SagemakerStepOperatorConfig,
     SagemakerStepOperatorFlavor,
 )
-from zenml.integrations.aws.flavors.batch_step_operator_flavor import (
+from zenml.integrations.aws.flavors.aws_batch_step_operator_flavor import (
     AWSBatchStepOperatorConfig,
     AWSBatchStepOperatorFlavor
 )

src/zenml/integrations/aws/flavors/batch_step_operator_flavor.py renamed to src/zenml/integrations/aws/flavors/aws_batch_step_operator_flavor.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Amazon SageMaker step operator flavor."""
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
+from typing import TYPE_CHECKING, Dict, Optional, Type
 
 from pydantic import Field
 
@@ -27,7 +27,6 @@
     BaseStepOperatorConfig,
     BaseStepOperatorFlavor,
 )
-from zenml.utils import deprecation_utils
 
 if TYPE_CHECKING:
     from zenml.integrations.aws.step_operators import AWSBatchStepOperator
@@ -36,13 +35,13 @@
 class AWSBatchStepOperatorSettings(BaseSettings):
     """Settings for the Sagemaker step operator."""
 
-    instance_type: Optional[str] = Field(
-        'optimal',
+    instance_type: str = Field(
+        default='optimal',
         description="The instance type for AWS Batch to use for the step" \
         " execution. Example: 'm5.xlarge'",
     )
     environment: Dict[str, str] = Field(
-        default_factory=dict,
+        default={},
         description="Environment variables to pass to the container during " \
             "execution. Example: {'LOG_LEVEL': 'INFO', 'DEBUG_MODE': 'False'}",
     )

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

@@ -16,7 +16,7 @@
 from zenml.integrations.aws.step_operators.sagemaker_step_operator import (  # noqa: F401
     SagemakerStepOperator,
 )
-from zenml.integrations.aws.step_operators.batch_step_operator import (  # noqa: F401
+from zenml.integrations.aws.step_operators.aws_batch_step_operator import (  # noqa: F401
     AWSBatchStepOperator,
 )
 __all__ = ["SagemakerStepOperator","AWSBatchStepOperator"]

src/zenml/integrations/aws/step_operators/batch_step_operator.py renamed to src/zenml/integrations/aws/step_operators/aws_batch_step_operator.py

@@ -14,33 +14,28 @@
 """Implementation of the Sagemaker Step Operator."""
 
 import time
+import math
 from typing import (
     TYPE_CHECKING,
     Dict,
     List,
     Optional,
     Tuple,
     Type,
-    Union,
     Literal,
     cast,
 )
 from pydantic import BaseModel
-
 import boto3
-from sagemaker.estimator import Estimator
-from sagemaker.inputs import TrainingInput
-from sagemaker.session import Session
 
 from zenml.client import Client
 from zenml.config.build_configuration import BuildConfiguration
 from zenml.enums import StackComponentType
-from zenml.integrations.aws.flavors.batch_step_operator_flavor import (
+from zenml.integrations.aws.flavors.aws_batch_step_operator_flavor import (
     AWSBatchStepOperatorConfig,
     AWSBatchStepOperatorSettings,
 )
-from zenml.integrations.aws.step_operators.batch_step_operator_entrypoint_config import (
-    BATCH_STEP_ENV_VAR_SIZE_LIMIT,
+from zenml.integrations.aws.step_operators.aws_batch_step_operator_entrypoint_config import (
     AWSBatchEntrypointConfiguration,
 )
 from zenml.logger import get_logger
@@ -49,7 +44,6 @@
 from zenml.step_operators.step_operator_entrypoint_configuration import (
     StepOperatorEntrypointConfiguration,
 )
-from zenml.utils.env_utils import split_environment_variables
 from zenml.utils.string_utils import random_str
 
 if TYPE_CHECKING:
@@ -60,7 +54,7 @@
 
 logger = get_logger(__name__)
 
-BATCH_DOCKER_IMAGE_KEY = "batch_step_operator"
+BATCH_DOCKER_IMAGE_KEY = "aws_batch_step_operator"
 _ENTRYPOINT_ENV_VARIABLE = "__ZENML_ENTRYPOINT"
 
 class AWSBatchJobDefinitionContainerProperties(BaseModel):
@@ -232,11 +226,11 @@ def map_environment(environment: Dict[str,str]) -> List[Dict[str,str]]:
         """
 
         return [
-            {"name":k,"value":v} for k,v in environment
+            {"name":k,"value":v} for k,v in environment.items()
         ]
 
     @staticmethod
-    def map_resource_settings(resource_settings: ResourceSettings) -> List[Dict[str,str]]:
+    def map_resource_settings(resource_settings: "ResourceSettings") -> List[Dict[str,str]]:
         """Utility to map the resource_settings to the resource convention used
         in the AWS Batch Job definition spec.
 
@@ -253,25 +247,31 @@ def map_resource_settings(resource_settings: ResourceSettings) -> List[Dict[str,
         else:
 
             if resource_settings.cpu_count is not None:
+
+                cpu_count_int = math.ceil(resource_settings.cpu_count)
+
+                if cpu_count_int != resource_settings.cpu_count:
+                    logger.info(f"AWS Batch only accepts int type cpu resource requirements. Converted {resource_settings.cpu_count} to {cpu_count_int}")
+
                 mapped_resource_settings.append(
                     {
-                        "value": resource_settings.cpu_count,
+                        "value": str(cpu_count_int),
                         "type": 'VCPU'
                     }
                 )
 
             if resource_settings.gpu_count is not None:
                 mapped_resource_settings.append(
                     {
-                        "value": resource_settings.gpu_count,
+                        "value": str(resource_settings.gpu_count),
                         "type": 'GPU'
                     }
                 )
 
             if resource_settings.get_memory() is not None:
                 mapped_resource_settings.append(
                     {
-                        "value": resource_settings.get_memory(),
+                        "value": str(int(resource_settings.get_memory(unit="MiB"))),
                         "type": 'MEMORY'
                     }
                 )

src/zenml/integrations/aws/step_operators/batch_step_operator_entrypoint_config.py renamed to src/zenml/integrations/aws/step_operators/aws_batch_step_operator_entrypoint_config.py

@@ -16,24 +16,7 @@
 from zenml.step_operators.step_operator_entrypoint_configuration import (
     StepOperatorEntrypointConfiguration,
 )
-from zenml.utils.env_utils import reconstruct_environment_variables
-
-BATCH_STEP_ENV_VAR_SIZE_LIMIT = 512
 
 
 class AWSBatchEntrypointConfiguration(StepOperatorEntrypointConfiguration):
-    """Entrypoint configuration for ZenML Batch step operator.
-
-    The only purpose of this entrypoint configuration is to reconstruct the
-    environment variables that exceed the maximum length of 512 characters
-    allowed for Batch steps from their individual components.
-    """
-
-    def run(self) -> None:
-        """Runs the step."""
-        # Reconstruct the environment variables that exceed the maximum length
-        # of 512 characters from their individual chunks
-        reconstruct_environment_variables()
-
-        # Run the step
-        super().run()
+    """Entrypoint configuration for ZenML Batch step operator."""

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

@@ -0,0 +1,13 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.

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

@@ -0,0 +1,86 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+
+
+import pytest
+
+from zenml.config.resource_settings import ResourceSettings
+from zenml.integrations.aws.step_operators.aws_batch_step_operator import AWSBatchStepOperator
+
+def test_aws_batch_step_operator_map_environment():
+    test_environment = {'key_1':'value_1','key_2':'value_2'}
+    expected = [
+        {
+            "name": "key_1",
+            "value": "value_1"
+        },
+        {
+            "name": "key_2",
+            "value": "value_2"
+        }
+    ]
+
+    assert AWSBatchStepOperator.map_environment(test_environment) == expected
+
+@pytest.mark.parametrize(
+        "test_resource_settings,expected",
+        [
+            (
+                ResourceSettings(),
+                []
+            ),
+            (
+                ResourceSettings(cpu_count=0.4,gpu_count=1,memory="10MiB"),
+                [
+                    {
+                        "value":"1",
+                        "type":"VCPU"
+                    },
+                    {
+                        "value":"1",
+                        "type":"GPU"
+                    },
+                    {
+                        "value": "10",
+                        "type": "MEMORY"
+                    }
+                ]
+            ),
+            (
+                ResourceSettings(cpu_count=1,gpu_count=1),
+                [
+                    {
+                        "value":"1",
+                        "type":"VCPU"
+                    },
+                    {
+                        "value":"1",
+                        "type":"GPU"
+                    },
+                ]
+            ),
+            (
+                ResourceSettings(memory="1GiB"),
+                [
+                    {
+                        "value": "1024",
+                        "type": "MEMORY"
+                    }
+                ]
+            ),
+
+        ]
+)
+def test_aws_batch_step_operator_map_resource_settings(test_resource_settings,expected):
+    assert AWSBatchStepOperator.map_resource_settings(test_resource_settings) == expected

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

@@ -0,0 +1,8 @@
+from zenml.integrations.aws.flavors.aws_batch_step_operator_flavor import AWSBatchStepOperatorSettings
+
+def test_aws_batch_step_operator_settings():
+    AWSBatchStepOperatorSettings(
+        instance_type="g4dn.xlarge",
+        environment={"key_1":"value_1","key_2":"value_2"},
+        timeout_seconds=60
+    )