Assert that GPUs are available and restart tasks if they are not (#1112)

* Assert that GPUs are available and output information about them

* Add the GPU assertion to pipeline/translate/translate.py

* Fix linting error

* Add the USE_CPU environment

Author: gregtatum

pipeline/bicleaner/bicleaner.sh

@@ -65,7 +65,7 @@ else
                # The GPU devices have failed to be found, and bicleaner AI falls back
                # to operate on the CPU very slowly. To guard against this wasting expensive
                # GPU time, always check that it can find GPUs.
-               python3 -c "import tensorflow; exit(0) if tensorflow.config.list_physical_devices('GPU') else exit(9001)"
+               python3 -c "import tensorflow; exit(0) if tensorflow.config.list_physical_devices('GPU') else exit(75)"
                bicleaner-ai-classify --disable_hardrules --scol ${scol} --tcol ${tcol} - - $1
        }
        export -f biclean

pipeline/common/marian.py

@@ -2,6 +2,10 @@
 Common utilities related to working with Marian.
 """
 
+from logging import Logger
+import os
+import sys
+import subprocess
 from pathlib import Path
 from typing import Union
 
@@ -51,3 +55,68 @@ def marian_args_to_dict(extra_marian_args: list[str]) -> dict[str, Union[str, bo
             decoder_config[previous_key] = [prev_value, arg]
 
     return decoder_config
+
+
+def assert_gpus_available(logger: Logger) -> None:
+    """
+    Sometimes the GPUs aren't available when running tasks on GPU machines in the cloud.
+    This function reports on the GPUs available, and exits the task with an
+    EX_TEMPFAIL (75) exit code when the machines are not available. Taskcluster can
+    restart the tasks via the `retry-exit-status` property.
+    """
+
+    if "USE_CPU" in os.environ or "COMET_CPU" in os.environ:
+        return
+
+    query = {
+        "name": "Name",
+        "driver_version": "Driver Version",
+        "vbios_version": "GPU BIOS",
+        "memory.total": "Memory Total",
+        "memory.free": "Memory Free",
+        "compute_cap": "Compute Capability (https://developer.nvidia.com/cuda-gpus)",
+        "temperature.gpu": "GPU temperature (Celsius)",
+    }
+
+    fields = list(query.keys())
+
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                f"--query-gpu={','.join(fields)}",
+                "--format=csv,noheader,nounits",
+            ],
+            check=False,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+    except FileNotFoundError:
+        raise Exception(
+            "nvidia-smi not found. Ensure NVIDIA drivers are installed and nvidia-smi is in PATH."
+        )
+
+    if result.returncode != 0:
+        stdout = result.stdout.strip()
+        stderr = result.stderr.strip()
+        logger.error(f"nvidia-smi failed with return code {result.returncode}")
+        for line in stdout.splitlines():
+            logger.error(f"[nvidia-smi(stdout)] {line}")
+        for line in stderr.splitlines():
+            logger.error(f"[nvidia-smi(stderr)] {line}")
+        logger.error("No GPUs were found available on this machine. Exiting with EX_TEMPFAIL (75)")
+        sys.exit(75)
+
+    output = result.stdout.strip()
+    if not output:
+        logger.info("No GPUs found by nvidia-smi, exiting EX_TEMPFAIL (75).")
+        sys.exit(75)
+
+    logger.info("CUDA-capable GPU(s) detected.")
+    for idx, line in enumerate(output.splitlines()):
+        values = [v.strip() for v in line.split(",")]
+        logger.info(f"GPU {idx}:")
+        for key, value in zip(query.values(), values):
+            logger.info(f"  {key}: {value}")
+        logger.info("")

pipeline/eval/eval.py

@@ -53,6 +53,7 @@
 
 from pipeline.common.downloads import decompress_file
 from pipeline.common.logging import get_logger
+from pipeline.common.marian import assert_gpus_available
 
 logger = get_logger("eval")
 try:
@@ -212,6 +213,8 @@ def main(args_list: Optional[list[str]] = None) -> None:
 
     logger.info("Save the original target sentences to the artifacts")
 
+    assert_gpus_available(logger)
+
     decompress_file(target_file_compressed, keep_original=False, decompressed_path=target_ref_file)
 
     run_bash_oneliner(

pipeline/train/train.py

@@ -16,6 +16,7 @@
 from pipeline.common.downloads import read_lines, write_lines
 from pipeline.common.logging import get_logger
 from pipeline.common.command_runner import apply_command_args, run_command_pipeline
+from pipeline.common.marian import assert_gpus_available
 
 logger = get_logger(__file__)
 train_dir = Path(__file__).parent
@@ -478,6 +479,8 @@ def main() -> None:
         help="Additional parameters for the training script",
     )
 
+    assert_gpus_available(logger)
+
     with tempfile.TemporaryDirectory() as temp_dir:
         train_cli = TrainCLI(parser.parse_args(), Path(temp_dir))
         train_cli.log_config()

pipeline/translate/translate.py

@@ -7,7 +7,6 @@
 from glob import glob
 import os
 from pathlib import Path
-import sys
 import tempfile
 
 from pipeline.common.command_runner import apply_command_args, run_command
@@ -22,6 +21,7 @@
 )
 from pipeline.common.marian import get_combined_config
 from pipeline.translate.translate_ctranslate2 import translate_with_ctranslate2
+from pipeline.common.marian import assert_gpus_available
 
 logger = get_logger(__file__)
 
@@ -187,6 +187,8 @@ def main() -> None:
             pass
         return
 
+    assert_gpus_available(logger)
+
     if decoder == Decoder.ctranslate2:
         translate_with_ctranslate2(
             input_zst=input_zst,
@@ -257,12 +259,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    try:
-        main()
-    except RuntimeError as e:
-        # On GCP instances, we occasionally find that a GPU is not found even
-        # when it has been requested. Exiting with a unique error code in these
-        # cases allows us to automatically retry such tasks in Taskcluster.
-        if len(e.args) > 0 and "no CUDA-capable device is detected" in e.args[0]:
-            logger.exception("couldn't find GPU, exiting with 9002")
-            sys.exit(9002)
+    main()

taskcluster/kinds/backtranslations-mono-trg-translate/kind.yml

@@ -111,9 +111,9 @@ tasks:
                 CUDA_DIR: fetches/cuda-toolkit
                 CUDNN_DIR: fetches/cuda-toolkit
                 MARIAN: $MOZ_FETCHES_DIR
+            # 75 - EX_TEMPFAIL, used for when the GPUs aren't available on the machine.
             # 128 happens when cloning this repository fails
-            # 9002 happens if no GPU is attached
-            retry-exit-status: [128, 9002]
+            retry-exit-status: [75, 128]
 
         # Don't run unless explicitly scheduled
         run-on-tasks-for: []

taskcluster/kinds/backtranslations-train-backwards-model/kind.yml

@@ -74,7 +74,9 @@ tasks:
             docker-image: {"in-tree": "train"}
             max-run-time: 2592000
             # train_taskcluster.py exits with 17 if a request to Taskcluster fails
-            retry-exit-status: [17]
+            # 75 - EX_TEMPFAIL, used for when the GPUs aren't available on the machine.
+            # 128 happens when cloning this repository fails
+            retry-exit-status: [17, 75, 128]
             env:
                 # Weight & Biases trigger
                 WANDB_PUBLICATION: "{wandb_publication}"
@@ -91,8 +93,6 @@ tasks:
 
             # Taskcluster proxy is required to read secrets
             taskcluster-proxy: true
-            # 128 happens when cloning this repository fails
-            retry-exit-status: [128]
 
         # The task needs to be able to read that secret to publish on Weight & Biases
         scopes:

taskcluster/kinds/corpus-clean-parallel-bicleaner-ai/kind.yml

@@ -79,10 +79,10 @@ tasks:
                 SRC: "{src_locale}"
                 TRG: "{trg_locale}"
             # 128 happens when cloning this repository fails
-            # 9001 is the code for when tensorflow fails to find GPUs on the system,
-            # and biclenaer reverts to CPU time. Rather than waste time, we should
-            # restart the task.
-            retry-exit-status: [128,9001]
+            # 75 is the unix code EX_TEMPFAIL, which indicates a temporary failure.
+            # This is used when the GPUs can't be accessed. Bicleaner reverts to CPU
+            # time in this case, which is a waste of time. The task should be restarted.
+            retry-exit-status: [128,75]
 
         # Don't run unless explicitly scheduled
         run-on-tasks-for: []

taskcluster/kinds/distillation-mono-src-translate/kind.yml

@@ -112,9 +112,9 @@ tasks:
                 CUDA_DIR: fetches/cuda-toolkit
                 CUDNN_DIR: fetches/cuda-toolkit
                 MARIAN: $MOZ_FETCHES_DIR
+            # 75 - EX_TEMPFAIL, used for when the GPUs aren't available on the machine.
             # 128 happens when cloning this repository fails
-            # 9002 happens if no GPU is attached
-            retry-exit-status: [128, 9002]
+            retry-exit-status: [75, 128]
 
         marian-args:
             from-parameters: training_config.marian-args.decoding-teacher

taskcluster/kinds/distillation-parallel-src-translate/kind.yml

@@ -113,9 +113,9 @@ tasks:
                 CUDA_DIR: fetches/cuda-toolkit
                 CUDNN_DIR: fetches/cuda-toolkit
                 MARIAN: $MOZ_FETCHES_DIR
+            # 75 - EX_TEMPFAIL, used for when the GPUs aren't available on the machine.
             # 128 happens when cloning this repository fails
-            # 9002 happens if no GPU is attached
-            retry-exit-status: [128, 9002]
+            retry-exit-status: [75, 128]
 
         run:
             using: run-task