handle single node use case

Netflix · Sep 22, 2023 · e172d98 · e172d98
1 parent 50fa111
commit e172d98
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/metaflow/plugins/aws/batch/batch_client.py b/metaflow/plugins/aws/batch/batch_client.py
@@ -320,14 +320,23 @@ def _register_job_definition(
                 )
             else:
                 job_definition["containerProperties"]["linuxParameters"]["devices"] = []
-                for i in range(int(efa)):
+                if (num_parallel or 0) > 1:
+                    for i in range(int(efa)):
+                        job_definition["containerProperties"]["linuxParameters"][
+                            "devices"
+                        ].append(
+                            {
+                                "hostPath": "/dev/infiniband/uverbs{}".format(i),
+                                "containerPath": "/dev/infiniband/uverbs{}".format(i),
+                                "permissions": ["READ", "WRITE", "MKNOD"]
+                            }
+                        )
+                else:
                     job_definition["containerProperties"]["linuxParameters"][
                         "devices"
                     ].append(
                         {
-                            "hostPath": "/dev/infiniband/uverbs{}".format(i),
-                            "containerPath": "/dev/infiniband/uverbs{}".format(i),
-                            "permissions": ["READ", "WRITE", "MKNOD"]
+                            "hostPath": "/dev/infiniband/uverbs0"
                         }
                     )