fix: Pass instance type into aws batch job definition even when num_replicas = 1

Author: dataproblems

torchx/schedulers/aws_batch_scheduler.py

@@ -255,7 +255,7 @@ def _role_to_node_properties(
         container["jobRoleArn"] = job_role_arn
     if execution_role_arn:
         container["executionRoleArn"] = execution_role_arn
-    if role.num_replicas > 1:
+    if role.num_replicas > 0:
         instance_type = instance_type_from_resource(role.resource)
         if instance_type is not None:
             container["instanceType"] = instance_type

torchx/schedulers/test/aws_batch_scheduler_test.py

@@ -195,15 +195,15 @@ def test_submit_dryrun_instance_type_multinode(self) -> None:
             node_groups[0]["container"]["instanceType"],
         )
 
-    def test_submit_dryrun_no_instance_type_singlenode(self) -> None:
+    def test_submit_dryrun_instance_type_singlenode(self) -> None:
         cfg = AWSBatchOpts({"queue": "ignored_in_test", "privileged": True})
         resource = specs.named_resources_aws.aws_p3dn_24xlarge()
         app = _test_app(num_replicas=1, resource=resource)
         info = create_scheduler("test").submit_dryrun(app, cfg)
         # pyre-ignore[16]
         node_groups = info.request.job_def["nodeProperties"]["nodeRangeProperties"]
         self.assertEqual(1, len(node_groups))
-        self.assertTrue("instanceType" not in node_groups[0]["container"])
+        self.assertTrue("instanceType" in node_groups[0]["container"])
 
     def test_submit_dryrun_no_instance_type_non_aws(self) -> None:
         cfg = AWSBatchOpts({"queue": "ignored_in_test", "privileged": True})