SIENTIAPDE-1350: Refactor: Improve logging, configuration, and resource management. Includes .env updates, client initialization logging, and resource closing.

model_manager/activities/activities.py

@@ -152,3 +152,5 @@ class Activities(ExperimentTracking, Training, Cleanup):
         Prefer calling this method explicitly rather than relying on __del__.
         """
         ExperimentTracking.close(self)
+        self.info('Postgres client closed')
+        self.storage_repository.close()

model_manager/activities/experiment_tracking.py

@@ -90,6 +90,8 @@ class ExperimentTracking(Postgres):
             metrics_controller=metrics_controller,
         )
 
+        self.info(f'Postgres client initialized at {host}:{port}')
+
     def __del__(self):
         """
         Destructor to safely handle cleanup during garbage collection.

model_manager/utils/connectors_config.py

@@ -84,6 +84,7 @@ def build_mongodb_config() -> dict[str, Any]:
         'connection_string': connection_string,
         'database_name': getenv('MONGODB_DATABASE_NAME', 'sientia'),
         'ttl_index_seconds': int(getenv('MONGODB_TTL_INDEX_HOURS', '1')) * 3600,
+        'uri': uri,
     }
 
 

model_manager/utils/repository/model_repository.py

@@ -31,8 +31,8 @@ warnings.filterwarnings('ignore', category=FutureWarning, message=".*'squared' i
 class ModelRepository:
     def __init__(self, url, username, password, logger: Logger):
         self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
-
         self.logger = logger
+        self.logger.info(f'MLFlow client initialized at {url}')
 
     def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
         """

model_manager/utils/repository/storage_repository.py

@@ -86,7 +86,11 @@ class StorageRepository:
             use_ssl=use_ssl,
         )
 
-        self.logger.info(f'MinIO client initialized successfully: {endpoint_url}')
+        self.logger.info(f'MinIO client initialized at {endpoint_url}')
+
+    def close(self) -> None:
+        self.minio_client.close()
+        self.logger.info('MinIO client closed')
 
     def fetch_file(self, bucket_name: str, file_name: str) -> BytesIO:
         """

model_manager/worker/worker.py

@@ -73,15 +73,14 @@ async def main():
         SystemExit: On graceful shutdown or error conditions
     """
     host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
+    use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
     logger = get_logger(__name__)
 
     metadata = {
         'pod_id': POD_ID,
     }
 
-    logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
     start_prometheus_server(logger, metadata)
-    logger.custom_info('Starting Notification Handler...', metadata)
     mongo_config = build_mongodb_config()
 
     notification_handler = NotificationHandler(
@@ -91,7 +90,7 @@ async def main():
         project_name=os.getenv('PROJECT_NAME', 'model-manager'),
     )
 
-    logger.custom_info('Starting Activities...', metadata)
+    logger.custom_info(f'MongoDB client initialized at {mongo_config["uri"]}', metadata)
 
     activities = Activities(
         postgres_config=build_postgres_config(),
@@ -101,23 +100,22 @@ async def main():
         notification_handler=notification_handler,
     )
 
-    logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
-
     new_runtime = Runtime(
         telemetry=TelemetryConfig(
             metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
         )
     )
 
-    logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
+    logger.custom_info(f'SDK metrics server initialized on port {SDK_METRICS_PORT}', metadata)
 
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'),
         runtime=new_runtime,
+        tls=use_tls,
     )
 
-    logger.custom_info('Starting Workers...', metadata)
+    logger.custom_info(f'Temporal client initialized at {host}', metadata)
 
     workers = [
         Worker(
@@ -159,7 +157,7 @@ async def main():
     for w in workers:
         handlers.append(w.run())
 
-    logger.custom_info('Workers started successfully', metadata)
+    logger.custom_info('Model manager workers initialized', metadata)
 
     try:
         # This will run the workers and wait for them to complete.
@@ -169,6 +167,7 @@ async def main():
         logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
     finally:
         notification_handler.shutdown()
+        logger.custom_info('MongoDB client closed', metadata)
         await activities.shutdown()
         # Exit with a non-zero status code to indicate failure to Kubernetes
         metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
@@ -195,7 +194,7 @@ def start_prometheus_server(logger: SientiaLogger, metadata: dict[str, str | Non
     try:
         port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        logger.custom_info(f'Prometheus server started on port {port}.', metadata)
+        logger.custom_info(f'Prometheus server initialized on port {port}.', metadata)
         metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
     except Exception as e:  # noqa: BLE001
         logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)