Canary

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "Node-Pod-Orchestration"
-version = "0.5.1"
+version = "0.5.5"
 description = ""
 authors = ["Alexander Röhl <alexander.roehl@uni-tuebingen.de>", "David Hieber <david.hieber@uni-tuebingen.de>"]
 license = "Apache 2.0"

src/api/api.py

@@ -1,11 +1,12 @@
+import time
 import uvicorn
-import os
+from typing import Optional
 import threading
 from fastapi import APIRouter, FastAPI, Depends, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 
-from src.utils.hub_client import init_hub_client_with_client, get_node_id_by_client
+from src.utils.hub_client import init_hub_client, get_node_id_by_client
 from src.utils.other import extract_hub_envs
 from src.api.oauth import valid_access_token
 from src.resources.database.entity import Database
@@ -24,6 +25,7 @@
 
 logger = get_logger()
 
+
 class PodOrchestrationAPI:
     """FastAPI application exposing the Pod Orchestration REST endpoints.
 
@@ -49,16 +51,11 @@ def __init__(self, database: Database, namespace: str = 'default'):
         """
         self.database = database
 
-        client_id, client_secret, hub_url_core, hub_auth, enable_hub_logging, http_proxy, https_proxy = extract_hub_envs()
+        self.node_id = None
+        self.hub_client = None
+        self.enable_hub_logging = None
+        self._set_node_id_and_hub_client(max_attempts=100)
 
-        self.enable_hub_logging = enable_hub_logging
-        self.hub_client = init_hub_client_with_client(client_id,
-                                                          client_secret,
-                                                          hub_url_core,
-                                                          hub_auth,
-                                                          http_proxy,
-                                                          https_proxy)
-        self.node_id = get_node_id_by_client(self.hub_client, client_id) if self.hub_client else None
         self.namespace = namespace
         app = FastAPI(title="FLAME PO",
                       docs_url="/api/docs",
@@ -329,12 +326,16 @@ def stop_all_analysis_call(self):
         """
         try:
             response = stop_analysis('all', self.database)
-            for analysis_id in self.database.get_analysis_ids():
-                stream_logs(AnalysisStoppedLog(analysis_id),
-                            self.node_id,
-                            self.enable_hub_logging,
-                            self.database,
-                            self.hub_client)
+            self._set_node_id_and_hub_client(max_attempts=5)
+            if self.node_id is not None:
+                for analysis_id in self.database.get_analysis_ids():
+                    stream_logs(AnalysisStoppedLog(analysis_id),
+                                self.node_id,
+                                self.enable_hub_logging,
+                                self.database,
+                                self.hub_client)
+            else:
+                logger.warning(f"Couldn't forward logs for stopped analyses.")
             return response
         except Exception as e:
             logger.error(f"Error stopping ALL analyzes: {repr(e)}")
@@ -354,11 +355,15 @@ def stop_analysis_call(self, analysis_id: str):
         """
         try:
             response = stop_analysis(analysis_id, self.database)
-            stream_logs(AnalysisStoppedLog(analysis_id),
-                        self.node_id,
-                        self.enable_hub_logging,
-                        self.database,
-                        self.hub_client)
+            self._set_node_id_and_hub_client(max_attempts=5)
+            if self.node_id is not None:
+                stream_logs(AnalysisStoppedLog(analysis_id),
+                            self.node_id,
+                            self.enable_hub_logging,
+                            self.database,
+                            self.hub_client)
+            else:
+                logger.warning(f"Couldn't forward logs for stopped analysis.")
             return response
         except Exception as e:
             logger.error(f"Error stopping analysis: {repr(e)}")
@@ -414,7 +419,11 @@ def cleanup_call(self, cleanup_type: str):
             HTTPException: 500 on any downstream failure (details in logs).
         """
         try:
-            return cleanup(cleanup_type, self.database, self.namespace)
+            self._set_node_id_and_hub_client(max_attempts=5)
+            return cleanup(cleanup_type,
+                           self.database,
+                           self.hub_client,
+                           self.namespace)
         except Exception as e:
             logger.error(f"Error cleaning up: {repr(e)}")
             raise HTTPException(status_code=500, detail=f"Error cleaning up (see po logs).")
@@ -433,9 +442,14 @@ def stream_logs_call(self, body: CreateLogEntity):
             HTTPException: 500 on any downstream failure (details in logs).
         """
         try:
-            return stream_logs(body, self.node_id, self.enable_hub_logging, self.database, self.hub_client)
-        except Exception as e:
-            logger.error(f"Error streaming logs: {repr(e)}")
+            self._set_node_id_and_hub_client(max_attempts=5)
+            if self.node_id is not None:
+                return stream_logs(body, self.node_id, self.enable_hub_logging, self.database, self.hub_client)
+            else:
+                raise Exception
+        except Exception:
+            logger.error(f"Error streaming logs: node_id={self.node_id}, "
+                         f"hub_client={self.hub_client}.")
             raise HTTPException(status_code=500, detail=f"Error streaming logs (see po logs).")
 
     def health_call(self):
@@ -452,3 +466,27 @@ def health_call(self):
             raise RuntimeError("Main thread is not alive.")
         else:
             return {'status': "ok"}
+
+    def _set_node_id_and_hub_client(self, max_attempts: Optional[int] = None) -> None:
+        current_attempt = 1
+        client_id, client_secret, hub_url_core, hub_auth, enable_hub_logging, http_proxy, https_proxy = extract_hub_envs()
+        if (self.hub_client is None) or (get_node_id_by_client(self.hub_client, client_id) is None):
+            self.node_id = None
+            while self.node_id is None:
+                self.enable_hub_logging = enable_hub_logging
+                self.hub_client = init_hub_client(client_id,
+                                                  client_secret,
+                                                  hub_url_core,
+                                                  hub_auth,
+                                                  http_proxy,
+                                                  https_proxy)
+                self.node_id = get_node_id_by_client(self.hub_client, client_id) if self.hub_client else None
+
+                if (max_attempts is not None) and (current_attempt >= max_attempts):
+                    break
+                else:
+                    current_attempt += 1
+                    if current_attempt % 60 == 59:
+                        logger.warning(f"Unable to connect to Hub, attempt {current_attempt}"
+                                       f"{'/' + str(max_attempts) if max_attempts is not None else ''}.")
+                    time.sleep(1)

src/main.py

@@ -1,10 +1,6 @@
 import os
 from threading import Thread
-
 from dotenv import load_dotenv, find_dotenv
-# load env
-load_dotenv(find_dotenv())
-
 
 from src.resources.database.entity import Database
 from src.api.api import PodOrchestrationAPI
@@ -13,7 +9,12 @@
 from src.utils.po_logging import get_logger
 
 
+# connect logger
 logger = get_logger()
+# load cluster config
+load_cluster_config()
+# load env
+load_dotenv(find_dotenv())
 
 
 def main():
@@ -23,9 +24,6 @@ def main():
     spawns the FastAPI server in a background thread, and starts the blocking
     status monitoring loop on the main thread.
     """
-    # load cluster config
-    load_cluster_config()
-
     if not os.getenv('NGINX_IMAGE'):
         logger.warning("Environment variable 'NGINX_IMAGE' is not set, defaulting to 'nginx:1.29.8'.")
 

src/resources/analysis/entity.py

@@ -8,6 +8,7 @@
 from src.resources.database.db_models import AnalysisDB
 from src.resources.database.entity import Database
 from src.status.constants import AnalysisStatus
+from src.utils.mb_client import delete_subscription
 
 
 class Analysis(BaseModel):
@@ -55,15 +56,10 @@ def start(self, database: Database, namespace: str = 'default') -> None:
         self.analysis_config['PROJECT_ID'] = self.project_id
         self.analysis_config['DEPLOYMENT_NAME'] = self.deployment_name
         self.namespace = namespace
-        self.pod_ids = create_analysis_deployment(name=self.deployment_name,
-                                                  image=self.image_url,
-                                                  env=self.analysis_config,
-                                                  namespace=namespace)
-
         database.create_analysis(analysis_id=self.analysis_id,
                                  deployment_name=self.deployment_name,
                                  project_id=self.project_id,
-                                 pod_ids=self.pod_ids,
+                                 pod_ids=None,
                                  status=self.status,
                                  log=self.log,
                                  registry_url=self.registry_url,
@@ -74,6 +70,12 @@ def start(self, database: Database, namespace: str = 'default') -> None:
                                  kong_token=self.kong_token,
                                  restart_counter=self.restart_counter,
                                  progress=self.progress)
+        self.pod_ids = create_analysis_deployment(name=self.deployment_name,
+                                                  image=self.image_url,
+                                                  env=self.analysis_config,
+                                                  namespace=namespace)
+
+        database.update_deployment(self.deployment_name, pod_ids=json.dumps(self.pod_ids))
 
     def stop(self,
              database: Database,
@@ -94,6 +96,9 @@ def stop(self,
         # Update the database
         database.update_deployment(self.deployment_name, status=self.status)
         database.update_deployment(self.deployment_name, log=self.log)
+        # end message-broker subscription
+        self.tokens = create_analysis_tokens(kong_token=self.kong_token, analysis_id=self.analysis_id)
+        delete_subscription(self.analysis_id, self.tokens['KEYCLOAK_TOKEN'])
 
 
 def read_db_analysis(analysis: AnalysisDB) -> Analysis:

src/resources/log/entity.py

@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel
 
-from src.status.constants import _MAX_RESTARTS
+from src.status.constants import _MAX_RESTARTS, AnalysisStatus
 
 
 class LogEntity(BaseModel):
@@ -99,4 +99,8 @@ def __init__(self, analysis_id: str) -> None:
         log = (f"[flame -- POAPI: ANALYSISSTOPPED -- "
                f"{time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}] "
                f"Info: The analysis was stopped either locally, or externally on another node.")
-        super().__init__(log=log, log_type="info", analysis_id=analysis_id, status="stopped", progress=0)
+        super().__init__(log=log,
+                         log_type="info",
+                         analysis_id=analysis_id,
+                         status=AnalysisStatus.STOPPED.value,
+                         progress=0)