Use spot instances for PRs (#457)

mozilla · Feb 23, 2024 · 9ae9d9a · 9ae9d9a
1 parent 3f135aa
commit 9ae9d9a
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 3 deletions.
diff --git a/taskcluster/kinds/train-backwards/kind.yml b/taskcluster/kinds/train-backwards/kind.yml
@@ -52,7 +52,10 @@ tasks:
                 - dependencies
                 - attributes
                 - run.command
-        worker-type: b-largegpu-largedisk-standard
+        worker-type:
+            by-tasks-for:
+                github-pull-request: b-largegpu
+                default: b-largegpu-largedisk-standard
         worker:
             max-run-time: 2592000
             env:

diff --git a/taskcluster/kinds/train-student/kind.yml b/taskcluster/kinds/train-student/kind.yml
@@ -48,7 +48,10 @@ tasks:
                     - pipeline/train/train.sh
                 from-parameters:
                     marian_args: training_config.marian-args.training-student
-        worker-type: b-largegpu-xlargedisk-standard
+        worker-type:
+            by-tasks-for:
+                github-pull-request: b-largegpu-largedisk
+                default: b-largegpu-xlargedisk-standard
         worker:
             max-run-time: 2592000
             env:

diff --git a/taskcluster/kinds/train-teacher/kind.yml b/taskcluster/kinds/train-teacher/kind.yml
@@ -71,7 +71,10 @@ tasks:
                     marian_args: training_config.marian-args.training-teacher
                     teacher-ensemble: training_config.experiment.teacher-ensemble
                     pretrained_teacher: training_config.experiment.pretrained-models.train-teacher
-        worker-type: b-largegpu-xlargedisk-standard
+        worker-type:
+            by-tasks-for:
+                github-pull-request: b-largegpu
+                default: b-largegpu-xlargedisk-standard
         worker:
             max-run-time: 2592000
             env:

diff --git a/taskcluster/translations_taskgraph/transforms/worker_env.py b/taskcluster/translations_taskgraph/transforms/worker_env.py
@@ -8,10 +8,24 @@
 # runtime, or adjust in kinds when changing worker types.
 
 from taskgraph.transforms.base import TransformSequence
+from taskgraph.util.schema import resolve_keyed_by
 
 transforms = TransformSequence()
 
 
+@transforms.add
+def evaluate_keyed_by(config, jobs):
+    for job in jobs:
+        resolve_keyed_by(
+            job,
+            "worker-type",
+            item_name=job["description"],
+            **{"tasks-for": config.params["tasks_for"]},
+        )
+
+        yield job
+
+
 @transforms.add
 def inject_worker_env(config, jobs):
     for job in jobs: