fix overwrite of neuron devices when efa devices are also specified

Netflix · Jan 23, 2024 · c0a858c · c0a858c
1 parent a46f4a3
commit c0a858c
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/metaflow/plugins/aws/batch/batch_client.py b/metaflow/plugins/aws/batch/batch_client.py
@@ -344,7 +344,10 @@ def _register_job_definition(
  "Invalid efa value: ({}) (should be 0 or greater)".format(efa)
  )
  else:
- job_definition["containerProperties"]["linuxParameters"]["devices"] = []
+ if "linuxParameters" not in job_definition["containerProperties"]:
+ job_definition["containerProperties"]["linuxParameters"] = {}
+ if "devices" not in job_definition["containerProperties"]["linuxParameters"]:
+ job_definition["containerProperties"]["linuxParameters"]["devices"] = []
  if (num_parallel or 0) > 1:
  # Multi-node parallel jobs require the container path and permissions explicitly specified in Job definition
  for i in range(int(efa)):