Dynamic CPU/GPU option

CyberDataLab · Aug 8, 2024 · bc5a45b · bc5a45b
1 parent 1deb0af
commit bc5a45b
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 2 deletions.
diff --git a/nebula/controller.py b/nebula/controller.py
@@ -80,6 +80,7 @@ def __init__(self, args):
         self.n_nodes = 0
         self.mender = None if self.simulation else Mender()
         self.use_blockchain = args.use_blockchain if hasattr(args, "use_blockchain") else False
+        self.gpu_available = False
 
     def start(self):
         banner = """
@@ -334,6 +335,7 @@ def run_frontend(self):
                     - ./config/nebula:/etc/nginx/sites-available/default
                 environment:
                     - NEBULA_PRODUCTION={production}
+                    - NEBULA_GPU_AVAILABLE={gpu_available}
                     - NEBULA_ADVANCED_ANALYTICS={advanced_analytics}
                     - SERVER_LOG=/nebula/app/logs/server.log
                     - NEBULA_LOGS_DIR=/nebula/app/logs/
@@ -377,10 +379,16 @@ def run_frontend(self):
                     external: true
         """
         )
+
+        try:
+            subprocess.check_call(["nvidia-smi"])
+            self.gpu_available = True
+        except Exception as e:
+            logging.info("No GPU available for the frontend, nodes will be deploy in CPU mode")
 
         # Generate the Docker Compose file dynamically
         services = ""
-        services += frontend_template.format(production=self.production, advanced_analytics=self.advanced_analytics, path=self.root_path, gw="192.168.10.1", ip="192.168.10.100", frontend_port=self.frontend_port, statistics_port=self.statistics_port)
+        services += frontend_template.format(production=self.production, gpu_available=self.gpu_available, advanced_analytics=self.advanced_analytics, path=self.root_path, gw="192.168.10.1", ip="192.168.10.100", frontend_port=self.frontend_port, statistics_port=self.statistics_port)
         docker_compose_file = docker_compose_template.format(services)
 
         if self.production:

diff --git a/nebula/frontend/app.py b/nebula/frontend/app.py
@@ -57,6 +57,7 @@
 
 class Settings:
     production: bool = os.environ.get("NEBULA_PRODUCTION", "False") == "True"
+    gpu_available: bool = os.environ.get("NEBULA_GPU_AVAILABLE", "False") == "True"
     advanced_analytics: bool = os.environ.get("NEBULA_ADVANCED_ANALYTICS", "False") == "True"
     log_dir: str = os.environ.get("NEBULA_LOGS_DIR")
     config_dir: str = os.environ.get("NEBULA_CONFIG_DIR")
@@ -933,7 +934,7 @@ async def nebula_dashboard_download_logs_metrics(scenario_name: str, request: Re
 @app.get("/nebula/dashboard/deployment/", response_class=HTMLResponse)
 async def nebula_dashboard_deployment(request: Request, session: Dict = Depends(get_session)):
     scenario_running = get_running_scenario()
-    return templates.TemplateResponse("deployment.html", {"request": request, "scenario_running": scenario_running, "user_logged_in": session.get("user")})
+    return templates.TemplateResponse("deployment.html", {"request": request, "scenario_running": scenario_running, "user_logged_in": session.get("user"), "gpu_available": settings.gpu_available})
 
 
 def attack_node_assign(

diff --git a/nebula/frontend/templates/deployment.html b/nebula/frontend/templates/deployment.html
@@ -450,8 +450,12 @@ <h5 class="step-title">Accelerator definition</h5>
                     <div class="form-check">
                         <select class="form-control" id="resourceUsage" name="resources"
                             style="display: inline; width: 20%">
+                            {% if gpu_available %}
                             <option value="gpu" selected>GPU</option>
                             <option value="cpu">CPU</option>
+                            {% else %}
+                            <option value="cpu" selected>CPU</option>
+                            {% endif %}
                         </select>
                     </div>
                 </div>