From c62dde964dd9b4b9bd6bd109244f3ba517cac694 Mon Sep 17 00:00:00 2001 From: davidhieber Date: Tue, 14 Nov 2023 14:49:30 +0100 Subject: [PATCH] feat: add the option to use a seingel gpu --- airflow/dags/DAG_PHT_run_train.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/airflow/dags/DAG_PHT_run_train.py b/airflow/dags/DAG_PHT_run_train.py index 31e4d3e..333c066 100644 --- a/airflow/dags/DAG_PHT_run_train.py +++ b/airflow/dags/DAG_PHT_run_train.py @@ -273,6 +273,10 @@ def execute_container(train_state): raise ValueError( f"Invalid gpu configuration: {gpu_config}. Must be a list of integers or 'all'" ) + elif isinstance(gpu_config, int): + device_request = docker.types.DeviceRequest( + device_ids=[str(gpu_config)], capabilities=[["gpu"]] + ) else: raise ValueError( f"Invalid gpu configuration: {gpu_config}. Must be a list of integers or 'all'"