Add execution_role_arn to aws_batch_scheduler

torchx/schedulers/aws_batch_scheduler.py

@@ -174,6 +174,7 @@ def _role_to_node_properties(
     start_idx: int,
     privileged: bool = False,
     job_role_arn: Optional[str] = None,
+    execution_role_arn: Optional[str] = None,
 ) -> Dict[str, object]:
     role.mounts += get_device_mounts(role.resource.devices)
 
@@ -250,6 +251,8 @@ def _role_to_node_properties(
     }
     if job_role_arn:
         container["jobRoleArn"] = job_role_arn
+    if execution_role_arn:
+        container["executionRoleArn"] = execution_role_arn
     if role.num_replicas > 1:
         instance_type = instance_type_from_resource(role.resource)
         if instance_type is not None:
@@ -355,6 +358,7 @@ class AWSBatchOpts(TypedDict, total=False):
     share_id: Optional[str]
     priority: int
     job_role_arn: Optional[str]
+    execution_role_arn: Optional[str]
 
 
 class AWSBatchScheduler(DockerWorkspaceMixin, Scheduler[AWSBatchOpts]):
@@ -505,6 +509,7 @@ def _submit_dryrun(self, app: AppDef, cfg: AWSBatchOpts) -> AppDryRunInfo[BatchJ
                     start_idx=node_idx,
                     privileged=cfg["privileged"],
                     job_role_arn=cfg.get("job_role_arn"),
+                    execution_role_arn=cfg.get("execution_role_arn"),
                 )
             )
             node_idx += role.num_replicas
@@ -585,6 +590,11 @@ def _run_opts(self) -> runopts:
             type_=str,
             help="The Amazon Resource Name (ARN) of the IAM role that the container can assume for AWS permissions.",
         )
+        opts.add(
+            "execution_role_arn",
+            type_=str,
+            help="The Amazon Resource Name (ARN) of the IAM role that the ECS agent can assume for AWS permissions.",
+        )
         return opts
 
     def _get_job_id(self, app_id: str) -> Optional[str]:

torchx/schedulers/test/aws_batch_scheduler_test.py

@@ -161,6 +161,15 @@ def test_submit_dryrun_job_role_arn(self) -> None:
             cfg["job_role_arn"], node_groups[0]["container"]["jobRoleArn"]
         )
 
+    def test_submit_dryrun_execution_role_arn(self) -> None:
+        cfg = AWSBatchOpts({"queue": "ignored_in_test", "execution_role_arn": "veryexecutive"})
+        info = create_scheduler("test").submit_dryrun(_test_app(), cfg)
+        node_groups = info.request.job_def["nodeProperties"]["nodeRangeProperties"]
+        self.assertEqual(1, len(node_groups))
+        self.assertEqual(
+            cfg["execution_role_arn"], node_groups[0]["container"]["executionRoleArn"]
+        )
+
     def test_submit_dryrun_privileged(self) -> None:
         cfg = AWSBatchOpts({"queue": "ignored_in_test", "privileged": True})
         info = create_scheduler("test").submit_dryrun(_test_app(), cfg)