#ENG-1292 TEE host api deployment and logging

.gitignore

@@ -12,3 +12,4 @@ verified-inference/enclave/.DS_Store
 **/**/node_modules
 SDKs/JS/dist
 **node_modules
+*.log

tee-host-api/api/api_logger.py

@@ -0,0 +1,48 @@
+import logging
+import os
+from typing import Any
+from typing import Optional
+
+from pythonjsonlogger import jsonlogger
+
+LOGGING_MESSAGE_FORMAT = "%(asctime)s %(name)-12s %(levelname)s %(message)s"
+LOG_FILE_PATH = "logs/tee-host-api.log"
+
+logger: Optional[Any] = None
+
+
+# pylint: disable=W0603
+def get():
+    global logger
+    if logger:
+        return logger
+    name = "tee-host-api"
+    file_handler = get_file_logger()
+    console_handler = get_console_logger()
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.DEBUG)
+
+    logger.addHandler(console_handler)
+    logger.addHandler(file_handler)
+    apply_default_formatter(file_handler)
+    apply_default_formatter(console_handler)
+
+    return logger
+
+
+def get_file_logger() -> logging.FileHandler:
+    os.makedirs(os.path.dirname(LOG_FILE_PATH), exist_ok=True)
+    file_handler = logging.FileHandler(LOG_FILE_PATH)
+    file_handler.setLevel(logging.DEBUG)
+    return file_handler
+
+
+def get_console_logger() -> logging.StreamHandler:
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.DEBUG)
+    return console_handler
+
+
+def apply_default_formatter(handler: logging.Handler):
+    formatter = jsonlogger.JsonFormatter(LOGGING_MESSAGE_FORMAT)
+    handler.setFormatter(formatter)

verified-inference/host/api/domain/tee/deploy_enclave_use_case.py → tee-host-api/api/domain/tee/deploy_enclave_use_case.py

@@ -1,13 +1,18 @@
 from typing import Any
 from typing import Dict
+
+import api_logger
+
 from repository import nitro_cli_repository
 
+logger = api_logger.get()
+
 
 async def execute(name: str, docker_hub_image: str, env_vars: Dict[str, Any]) -> Dict:
     build_result = nitro_cli_repository.build_enclave(name, docker_hub_image)
-    print(f"Enclave built: {build_result}")
+    logger.info(f"Enclave built: {build_result}")
 
     running_result = nitro_cli_repository.run_enclave(name)
-    print(f"Enclave running: {running_result}")
+    logger.info(f"Enclave running: {running_result}")
 
     return running_result

verified-inference/host/api/domain/tee/get_all_enclaves_use_case.py → tee-host-api/api/domain/tee/get_all_enclaves_use_case.py

@@ -1,10 +1,16 @@
 from typing import Dict
+
+import api_logger
+
 from repository import nitro_cli_repository
 from domain.tee.entities import Enclave, EnclaveState
 
+logger = api_logger.get()
+
 
 async def execute() -> Dict[str, Enclave]:
     enclaves = nitro_cli_repository.describe_enclaves()
+    logger.info(f"Enclaves: {enclaves}")
     # Return a dictionary of Enclave objects with the enclave name as the key
     return {
         enclave["EnclaveName"]: Enclave(

verified-inference/host/api/domain/tee/inject_env_vars_use_case.py → tee-host-api/api/domain/tee/inject_env_vars_use_case.py

@@ -3,6 +3,10 @@
 import socket
 from typing import Dict, Any
 
+import api_logger
+
+logger = api_logger.get()
+
 ENCLAVE_ENV_VAR_RECEIVER_PORT = 3000
 CONNECTION_TIMEOUT = 30.0
 
@@ -11,17 +15,19 @@ async def execute(enclave_cid: int, env_vars: Dict[str, Any]) -> None:
     for i in range(10):
         try:
             response = _send_env_vars(enclave_cid, env_vars)
-            print(f"Enclave CID: {enclave_cid} - Send env vars response: {response}")
+            logger.info(
+                f"Enclave CID: {enclave_cid} - Send env vars response: {response}"
+            )
             return
         except socket.timeout:
-            print(f"Enclave CID: {enclave_cid} - Connection timed out.")
+            logger.error(f"Enclave CID: {enclave_cid} - Connection timed out.")
         except socket.error as e:
-            print(f"Enclave CID: {enclave_cid} - Socket error: {e}")
+            logger.error(f"Enclave CID: {enclave_cid} - Socket error: {e}")
         except Exception as e:
-            print(f"Enclave CID: {enclave_cid} - Unexpected error: {e}")
+            logger.error(f"Enclave CID: {enclave_cid} - Unexpected error: {e}")
         await asyncio.sleep(5)
-        print(f"Enclave CID: {enclave_cid} - Retrying..")
-    print(f"Enclave CID: {enclave_cid} - Failed to send env vars")
+        logger.warning(f"Enclave CID: {enclave_cid} - Retrying..")
+    logger.error(f"Enclave CID: {enclave_cid} - Failed to send env vars")
 
 
 def _send_env_vars(enclave_cid: int, env_vars: Dict[str, Any]) -> str:

verified-inference/host/api/domain/tee/request_attestation_from_enclave_use_case.py → tee-host-api/api/domain/tee/request_attestation_from_enclave_use_case.py

@@ -1,14 +1,19 @@
+import api_logger
+
 from repository.vsock_repository import request_attestation
 from domain.tee import get_all_enclaves_use_case
 
+logger = api_logger.get()
+
 
 async def execute(enclave_name: str) -> str:
     enclaves = await get_all_enclaves_use_case.execute()
     if enclave_name not in enclaves:
         raise Exception(f"Enclave {enclave_name} not found")
 
     enclave = enclaves[enclave_name]
-    print(f"Requesting attestation from enclave: {enclave_name}")
+    logger.info(f"Requesting attestation from enclave: {enclave_name}")
     response = await request_attestation(enclave.cid)
+    logger.info(f"Attestation response: {response}")
 
     return response

verified-inference/host/api/domain/tee/shutdown_enclave_use_case.py → tee-host-api/api/domain/tee/shutdown_enclave_use_case.py

@@ -1,14 +1,18 @@
+import api_logger
+
 from domain.tee import get_all_enclaves_use_case
 from repository.vsock_repository import shutdown
 
+logger = api_logger.get()
+
 
 async def execute(enclave_name: str) -> bool:
     enclaves = await get_all_enclaves_use_case.execute()
     if enclave_name not in enclaves:
         raise Exception(f"Enclave {enclave_name} not found")
 
     enclave = enclaves[enclave_name]
-    print(f"Requesting shutdown from enclave: {enclave_name}")
+    logger.info(f"Requesting shutdown from enclave: {enclave_name}")
     shutdown_result = await shutdown(enclave.cid)
-    print(f"Enclave shut down: {shutdown_result}")
+    logger.info(f"Enclave shut down: {shutdown_result}")
     return shutdown_result

verified-inference/host/api/domain/tee/terminate_enclave_use_case.py → tee-host-api/api/domain/tee/terminate_enclave_use_case.py

@@ -1,9 +1,12 @@
+import api_logger
+
 from typing import Dict
 from repository import nitro_cli_repository
 
+logger = api_logger.get()
+
 
 async def execute(name: str) -> Dict:
     terminate_result = nitro_cli_repository.terminate_enclave(name)
-    print(f"Enclave terminated: {terminate_result}")
-
+    logger.info(f"Enclave terminated: {terminate_result}")
     return terminate_result

verified-inference/host/api/repository/nitro_cli_repository.py → tee-host-api/api/repository/nitro_cli_repository.py

@@ -4,6 +4,9 @@
 import re
 import time
 
+import api_logger
+
+logger = api_logger.get()
 
 TEE_CPU_COUNT = 2
 TEE_MEMORY_IN_MB_LOW = 4096
@@ -39,9 +42,10 @@ def _run_command(command: list) -> str:
             check=True,
         )
         if result.stderr:
-            print(f"Error logs: {result.stderr.strip()}")
+            logger.error(f"Error logs: {result.stderr.strip()}")
         return result.stdout.strip()
     except subprocess.CalledProcessError as e:
+        logger.error(f"Error running command: {command}\nError: {e}")
         error_message = e.stderr.strip()
         error_code = _parse_error_code(error_message)
         raise NitroCLIError(message=error_message, code=error_code) from e
@@ -78,7 +82,7 @@ def _create_run_enclave_command(enclave_name: str, memory_mb: int) -> List[str]:
 def _handle_insufficient_memory_error(enclave_name: str, error: NitroCLIError) -> Dict:
     min_memory = _parse_min_memory(error.message)
     if min_memory:
-        print(f"Retrying with minimum memory: {min_memory} MB")
+        logger.warning(f"Retrying with minimum memory: {min_memory} MB")
         time.sleep(1)  # add a small delay to avoid immediate retry
         return run_enclave(enclave_name, memory_mb=min_memory + 10, retry=False)
     raise error
@@ -108,6 +112,7 @@ def run_enclave(
         res = _run_command(command)
         return json.loads(res)
     except NitroCLIError as e:
+        logger.error(f"Error running enclave: {e}")
         if e.code == "E26" and retry:
             return _handle_insufficient_memory_error(enclave_name, e)
         raise
@@ -131,10 +136,3 @@ def describe_enclaves() -> List[Dict]:
     ]
     res = _run_command(command)
     return json.loads(res)
-
-
-if __name__ == "__main__":
-    try:
-        print(run_enclave("067988c3-6ede-7553-8000-e7d0807256bb", memory_mb=10000))
-    except NitroCLIError as e:
-        print(f"Caught an error: {e}")

verified-inference/host/api/repository/vsock_repository.py → tee-host-api/api/repository/vsock_repository.py

@@ -3,6 +3,10 @@
 from enum import Enum
 from typing import Optional
 
+import api_logger
+
+logger = api_logger.get()
+
 VSOCK_PORT = 8000
 TIMEOUT_IN_SECONDS = 5
 BUFFER_SIZE = 8192
@@ -18,6 +22,7 @@ async def request_attestation(enclave_cid: int) -> str:
         enclave_cid, VsockRequestType.GET_ATTESTATION, receive_response=True
     )
     if response is None:
+        logger.error(f"Enclave CID: {enclave_cid} - No response received from attestation request")
         raise Exception("No response received from attestation request")
     return response
 
@@ -31,8 +36,10 @@ async def shutdown(enclave_cid: int) -> bool:
         if e.errno == errno.ENOTCONN:
             # The enclave closed the connection, which is expected
             return True
+        logger.error(f"Enclave CID: {enclave_cid} - Error shutting down enclave: {e}")
         return False
     if response is None:
+        logger.error(f"Enclave CID: {enclave_cid} - No response received from shutdown request")
         raise Exception("No response received from shutdown request")
     return True
 
@@ -56,7 +63,7 @@ def _send_request(
             return response.decode()
         return None
     except Exception as e:
-        print(f"Error in vsock communication: {e}")
+        logger.error(f"Error in vsock communication: {e}")
         raise e
     finally:
         vsock.close()

verified-inference/host/api/routers/routes/tee_router.py → tee-host-api/api/routers/routes/tee_router.py

@@ -33,7 +33,6 @@
     response_model=TeeDeploymentResponse,
 )
 async def deploy(
-    api_request: Request,
     request: TeeDeploymentRequest,
 ):
     return await tee_deploy_service.execute(
@@ -51,7 +50,6 @@ async def deploy(
     response_model=TeeTerminateResponse,
 )
 async def terminate(
-    api_request: Request,
     request: TeeTerminateRequest,
 ):
     return await tee_terminate_service.execute(request.enclave_name)
@@ -64,9 +62,7 @@ async def terminate(
     response_description="Returns all enclaves",
     response_model=TeeGetEnclavesResponse,
 )
-async def enclaves(
-    api_request: Request,
-):
+async def enclaves():
     return await tee_get_all_enclaves_service.execute()
 
 
@@ -78,7 +74,6 @@ async def enclaves(
     response_model=TeeGetEnclaveResponse,
 )
 async def enclave(
-    api_request: Request,
     enclave_name: str,
 ):
     return await tee_get_enclave_service.execute(enclave_name)
@@ -92,7 +87,6 @@ async def enclave(
     response_model=TeeAttestationResponse,
 )
 async def attestation(
-    api_request: Request,
     enclave_name: str,
 ):
     return await tee_attestation_service.execute(enclave_name)

verified-inference/host/api/service/tee/tee_attestation_service.py → tee-host-api/api/service/tee/tee_attestation_service.py

@@ -1,13 +1,19 @@
 from fastapi import HTTPException
+
+import api_logger
+
 from service.tee.entities import TeeAttestationResponse
 from domain.tee import request_attestation_from_enclave_use_case
 
+logger = api_logger.get()
+
 
 async def execute(name: str) -> TeeAttestationResponse:
     try:
         result = await request_attestation_from_enclave_use_case.execute(name)
         return TeeAttestationResponse(attestation=result)
     except Exception as e:
+        logger.error(f"Error attesting enclave: {str(e)}")
         raise HTTPException(
             status_code=500, detail=f"Error attesting enclave: {str(e)}"
         )

verified-inference/host/api/service/tee/tee_chat_service.py → tee-host-api/api/service/tee/tee_chat_service.py

@@ -1,13 +1,14 @@
 import asyncio
 import socket
-import httpx
 from typing import AsyncGenerator
-from fastapi import HTTPException
 from fastapi import Request
-from fastapi.responses import StreamingResponse
+
+import api_logger
 
 from domain.tee import get_all_enclaves_use_case
 
+logger = api_logger.get()
+
 VSOCK_PORT = 8001
 BUFFER_SIZE = 4096
 
@@ -134,7 +135,7 @@ async def _vsock_stream_generator(
             yield chunk_data
 
     except Exception as e:
-        print(f"Error streaming from VSOCK socket: {e}")
+        logger.error(f"Enclave CID: {cid} - Error streaming from VSOCK socket: {e}")
 
     finally:
         writer.close()

verified-inference/host/api/service/tee/tee_deploy_service.py → tee-host-api/api/service/tee/tee_deploy_service.py

@@ -1,10 +1,14 @@
 from typing import Any
 from typing import Dict
 from fastapi import HTTPException
+
+import api_logger
+
 from service.tee.entities import TeeDeploymentResponse
 from domain.tee import deploy_enclave_use_case
 from domain.tee import inject_env_vars_use_case
 
+logger = api_logger.get()
 
 async def execute(
     name: str, docker_hub_image: str, env_vars: Dict[str, Any]
@@ -14,6 +18,7 @@ async def execute(
         await inject_env_vars_use_case.execute(result["EnclaveCID"], env_vars)
         return TeeDeploymentResponse(result=result)
     except Exception as e:
+        logger.error(f"Enclave name: {name} - Error deploying enclave: {str(e)}")
         raise HTTPException(
             status_code=500, detail=f"Error deploying enclave: {str(e)}"
         )