From badb3909fa4de8cbb1796553a8208dd7ebabbf86 Mon Sep 17 00:00:00 2001 From: Renan Souza Date: Wed, 18 Dec 2024 15:37:55 -0500 Subject: [PATCH] Code reformat --- src/flowcept/commons/daos/docdb_dao/docdb_dao_base.py | 2 +- src/flowcept/configs.py | 4 ++++ src/flowcept/instrumentation/flowcept_task.py | 4 +++- src/flowcept/instrumentation/flowcept_torch.py | 4 +++- 4 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/flowcept/commons/daos/docdb_dao/docdb_dao_base.py b/src/flowcept/commons/daos/docdb_dao/docdb_dao_base.py index cdb9ab23..872d84e5 100644 --- a/src/flowcept/commons/daos/docdb_dao/docdb_dao_base.py +++ b/src/flowcept/commons/daos/docdb_dao/docdb_dao_base.py @@ -67,7 +67,7 @@ def get_instance(*args, **kwargs) -> "DocumentDBDAO": DocumentDBDAO._instance = LMDBDAO() else: - raise NotImplementedError + raise Exception("All dbs are disabled. You can't use this.") # TODO: revise, this below may be better in subclasses DocumentDBDAO._instance._initialized = True return DocumentDBDAO._instance diff --git a/src/flowcept/configs.py b/src/flowcept/configs.py index 115ce1fb..dba19ed2 100644 --- a/src/flowcept/configs.py +++ b/src/flowcept/configs.py @@ -104,6 +104,10 @@ else: LMDB_ENABLED = _lmdb_settings.get("enabled", False) +if not LMDB_ENABLED and not MONGO_ENABLED: + # At least one of these variables need to be enabled. + LMDB_ENABLED = True + ########################## # Buffer Settings # ########################## diff --git a/src/flowcept/instrumentation/flowcept_task.py b/src/flowcept/instrumentation/flowcept_task.py index 3d361ae0..d1e06629 100644 --- a/src/flowcept/instrumentation/flowcept_task.py +++ b/src/flowcept/instrumentation/flowcept_task.py @@ -1,4 +1,5 @@ """Task module.""" + import threading from time import time from functools import wraps @@ -18,6 +19,7 @@ _thread_local = threading.local() + # TODO: :code-reorg: consider moving it to utils and reusing it in dask interceptor def default_args_handler(task_message: TaskObject, *args, **kwargs): """Get default arguments.""" @@ -157,4 +159,4 @@ def wrapper(*args, **kwargs): def get_current_context_task_id(): """Retrieve the current task object from thread-local storage.""" - return getattr(_thread_local, '_flowcept_current_context_task_id', None) + return getattr(_thread_local, "_flowcept_current_context_task_id", None) diff --git a/src/flowcept/instrumentation/flowcept_torch.py b/src/flowcept/instrumentation/flowcept_torch.py index d3bfac95..713ab00c 100644 --- a/src/flowcept/instrumentation/flowcept_torch.py +++ b/src/flowcept/instrumentation/flowcept_torch.py @@ -142,7 +142,9 @@ def __init__(self, *args, **kwargs): self._parent_task_id = kwargs.get( "parent_task_id", get_current_context_task_id() ) # to be used by forward layers - self._parent_workflow_id = kwargs.get("parent_workflow_id", Flowcept.current_workflow_id) + self._parent_workflow_id = kwargs.get( + "parent_workflow_id", Flowcept.current_workflow_id + ) self._campaign_id = kwargs.get("campaign_id", Flowcept.campaign_id) if kwargs.get("save_workflow", True): self._workflow_id = self._register_as_workflow()