Add support for job_role_arn in the demand execution lambda functions

pyproject.toml

@@ -18,8 +18,8 @@ dynamic = [
     "version",
 ]
 dependencies = [
-    "aibs-informatics-aws-utils>=0.0.8,<1",
-    "aibs-informatics-core>=0.2.4,<1",
+    "aibs-informatics-aws-utils>=0.0.12,<1",
+    "aibs-informatics-core>=0.2.6,<1",
     "aws-lambda-powertools ~= 2.35",
     "pydantic >= 2.0.3, < 3",
     "aws-lambda-typing",

src/aibs_informatics_aws_lambda/handlers/batch/create.py

@@ -63,6 +63,7 @@ def handle(
             mount_points=request.mount_points,
             volumes=request.volumes,
             privileged=request.privileged,
+            job_role_arn=request.job_role_arn,
         )
 
         response = register_job_definition(

src/aibs_informatics_aws_lambda/handlers/batch/model.py

@@ -29,6 +29,7 @@ class CreateDefinitionAndPrepareArgsRequest(SchemaModel):
     image: str = custom_field()
     job_definition_name: str = custom_field()
     job_queue_name: str = custom_field()
+    job_role_arn: Optional[str] = custom_field(default=None)
     job_name: Optional[str] = custom_field(default=None)
     command: List[str] = custom_field(default_factory=list)
     environment: Dict[str, str] = custom_field(default_factory=dict)

src/aibs_informatics_aws_lambda/handlers/demand/context_manager.py

@@ -636,6 +636,7 @@ def generate_batch_job_builder(  # noqa: C901
     if tmp_mount_point:
         vol_configurations.append(BatchEFSConfiguration(tmp_mount_point, read_only=False))
     logger.info("Constructing BatchJobBuilder instance...")
+    assert demand_execution.execution_platform.aws_batch is not None
     return BatchJobBuilder(
         image=demand_execution.execution_image,
         job_definition_name=env_base.get_job_name(
@@ -657,6 +658,9 @@ def generate_batch_job_builder(  # noqa: C901
         env_base=env_base,
         # TODO: need to make this configurable
         privileged=True,
+        job_role_arn=demand_execution.execution_platform.aws_batch.job_role
+        if demand_execution.execution_platform.aws_batch
+        else None,
     )
 
 

src/aibs_informatics_aws_lambda/handlers/demand/model.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import List, Optional, Union
 
 from aibs_informatics_core.models.aws.s3 import S3Path
 from aibs_informatics_core.models.base import (
@@ -17,15 +17,6 @@
 from aibs_informatics_aws_lambda.handlers.data_sync.model import RemoveDataPathsRequest
 
 
-@dataclass
-class CreateDefinitionAndPrepareArgsResponse(SchemaModel):
-    job_name: str = custom_field()
-    job_definition_arn: Optional[str] = custom_field()
-    job_queue_arn: str = custom_field()
-    parameters: Dict[str, Any] = custom_field()
-    container_overrides: Dict[str, Any] = custom_field()
-
-
 @dataclass
 class FileSystemConfiguration(SchemaModel):
     file_system: Optional[str] = None

src/aibs_informatics_aws_lambda/handlers/demand/scaffolding.py

@@ -101,6 +101,7 @@ def handle(self, request: PrepareDemandScaffoldingRequest) -> PrepareDemandScaff
                 volumes=batch_job_builder.volumes,
                 retry_strategy=build_retry_strategy(num_retries=5),
                 privileged=batch_job_builder.privileged,
+                job_role_arn=batch_job_builder.job_role_arn,
             ),
         )
 

test/aibs_informatics_aws_lambda/handlers/demand/test_scaffolding.py

@@ -203,6 +203,7 @@ def test__handle__simple_case(self, mock_construct_batch_efs_configuration) -> N
         batch_job_builder.mount_points = []
         batch_job_builder.volumes = []
         batch_job_builder.privileged = False
+        batch_job_builder.job_role_arn = None
 
         expected = PrepareDemandScaffoldingResponse(
             demand_execution=self.demand_execution,
@@ -299,6 +300,7 @@ def test__handle__file_system_overrides(self, mock_construct_batch_efs_configura
         batch_job_builder.mount_points = []
         batch_job_builder.volumes = []
         batch_job_builder.privileged = False
+        batch_job_builder.job_role_arn = None
 
         expected = PrepareDemandScaffoldingResponse(
             demand_execution=self.demand_execution,