feat(back): fluidattacks#973.1 tags support

- add tag support for aws jobs
danmur97 · Oct 24, 2022 · 36eafce · 36eafce
1 parent 22c92c3
commit 36eafce
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git a/src/args/compute-on-aws-batch/default.nix b/src/args/compute-on-aws-batch/default.nix
@@ -19,6 +19,7 @@
   queue,
   name,
   setup,
+  tags,
   vcpus,
 }:
 makeScript {
@@ -65,6 +66,13 @@ makeScript {
     __argName__ = name;
     __argParallel__ = parallel;
     __argQueue__ = queue;
+    __argTags__ = let
+      tag_names = builtins.attrNames tags;
+      encode_tag = key: "${key}=${tags."${key}"}";
+      encoded = map encode_tag tag_names;
+      encoded_tags = builtins.concatStringsSep "," encoded;
+    in
+      encoded_tags;
   };
   searchPaths = {
     bin = [

diff --git a/src/args/compute-on-aws-batch/entrypoint.sh b/src/args/compute-on-aws-batch/entrypoint.sh
@@ -22,6 +22,7 @@ function main {
   local name="__argName__"
   local queue="__argQueue__"
   local parallel="__argParallel__"
+  local tags="__argTags__"
   local submit_job_args
 
   : \
@@ -76,6 +77,7 @@ function main {
       --job-definition "${definition}"
       --retry-strategy "attempts=${attempts}"
       --timeout "attemptDurationSeconds=${attempt_duration_seconds}"
+      --tags "${tags}"
     ) \
     && if [ "${parallel}" -gt "1" ]; then
       submit_job_args+=(--array-properties "size=${parallel}")

diff --git a/src/evaluator/modules/compute-on-aws-batch/default.nix b/src/evaluator/modules/compute-on-aws-batch/default.nix
@@ -73,6 +73,10 @@ in {
             default = [];
             type = lib.types.listOf lib.types.package;
           };
+          tags = lib.mkOption {
+            default = {};
+            type = lib.types.attrsOf lib.types.str;
+          };
           vcpus = lib.mkOption {
             type = lib.types.ints.positive;
           };