From 50a3882400df738b576c0b3cbbf59373f299985a Mon Sep 17 00:00:00 2001 From: aldbr Date: Wed, 4 Dec 2024 15:03:33 +0100 Subject: [PATCH] feat(framework): make logging with context var more generic --- .../gLogger/gLogger/Advanced/index.rst | 62 ++++++++++++++++ src/DIRAC/FrameworkSystem/Client/Logger.py | 4 ++ .../private/standardLogging/LoggingContext.py | 16 +++++ .../test/Test_Logging_ContextVars.py | 71 +++++++++++++++++++ .../Client/Matcher.py | 4 +- .../WorkloadManagementSystem/DB/JobDB.py | 4 +- .../DB/JobLoggingDB.py | 4 +- .../DB/PilotAgentsDB.py | 4 +- .../DB/TaskQueueDB.py | 4 +- .../Service/MatcherHandler.py | 2 +- .../Utilities/ContextVars.py | 16 ----- 11 files changed, 164 insertions(+), 27 deletions(-) create mode 100644 src/DIRAC/FrameworkSystem/private/standardLogging/LoggingContext.py create mode 100644 src/DIRAC/FrameworkSystem/private/standardLogging/test/Test_Logging_ContextVars.py delete mode 100644 src/DIRAC/WorkloadManagementSystem/Utilities/ContextVars.py diff --git a/docs/source/DeveloperGuide/AddingNewComponents/Utilities/gLogger/gLogger/Advanced/index.rst b/docs/source/DeveloperGuide/AddingNewComponents/Utilities/gLogger/gLogger/Advanced/index.rst index 538ec77cb68..c754041aff8 100644 --- a/docs/source/DeveloperGuide/AddingNewComponents/Utilities/gLogger/gLogger/Advanced/index.rst +++ b/docs/source/DeveloperGuide/AddingNewComponents/Utilities/gLogger/gLogger/Advanced/index.rst @@ -217,6 +217,68 @@ This option can not be modified in the children of *gLogger*, even by *gLogger* itself after the configuration, so the children receive the *gLogger* configuration. +Add variables to different *Logging* objects depending on the context +--------------------------------------------------------------------- + +In complex cases, it can be useful to have loggers that change depending on +the execution context, without having to pass logger instances explicitly +through multiple layers of function calls. + +Python's `contextvars` module provides context-local storage, which can be used +to store and retrieve context-specific data, such as logger instances. + +gLogger supports the use of context variables to manage loggers in a flexible way. + +Provide a Context Logger +~~~~~~~~~~~~~~~~~~~~~~~~ + +When you have a *Logging* instance that you want to use in a specific context, +you can set it in the context variable: + +:: + + # Create a logger instance + logger = gLogger.getSubLogger("MyContextLogger") + + # Set it in the context variable + contextLogger.set(logger) + +Then, the instances within the context block will use the shared *Logging* object +set in the context variable: + +:: + + with setContextLogger(contextualLogger): + # Any logging within this block will use contextualLogger + obj = MyClass() + obj.do_something() # This will use contextualLogger + +Consume a Context Logger +~~~~~~~~~~~~~~~~~~~~~~~~ + +In functions or classes that need to log messages, you can retrieve the logger +from the context variable: + +:: + + class MyClass: + def __init__(self): + # Get the default logger if no context logger is set + self._defaultLogger = gLogger.getSubLogger("MyClass") + + @property + def log(self): + # Return the context logger if set, otherwise the default logger + return contextLogger.get() or self._defaultLogger + + @log.setter + def log(self, value): + # Optionally, allow setting a new default logger + self._defaultLogger = value + + def do_something(self): + self.log.notice("Doing something") + Some examples and summaries --------------------------- diff --git a/src/DIRAC/FrameworkSystem/Client/Logger.py b/src/DIRAC/FrameworkSystem/Client/Logger.py index 864883ac866..7b48d9b49ee 100755 --- a/src/DIRAC/FrameworkSystem/Client/Logger.py +++ b/src/DIRAC/FrameworkSystem/Client/Logger.py @@ -1,3 +1,4 @@ +from DIRAC.FrameworkSystem.private.standardLogging.LoggingContext import contextLogger, setContextLogger from DIRAC.FrameworkSystem.private.standardLogging.LoggingRoot import LoggingRoot gLogger = LoggingRoot() @@ -5,3 +6,6 @@ def getLogger(): return gLogger + + +__all__ = ["contextLogger", "setContextLogger", "getLogger"] diff --git a/src/DIRAC/FrameworkSystem/private/standardLogging/LoggingContext.py b/src/DIRAC/FrameworkSystem/private/standardLogging/LoggingContext.py new file mode 100644 index 00000000000..1e16ffdb3df --- /dev/null +++ b/src/DIRAC/FrameworkSystem/private/standardLogging/LoggingContext.py @@ -0,0 +1,16 @@ +""" Logging context module""" + +# Context variable for the logger (adapted to the request of the pilot reference) +import contextvars +from contextlib import contextmanager + +contextLogger = contextvars.ContextVar("Logger", default=None) + + +@contextmanager +def setContextLogger(logger_name): + token = contextLogger.set(logger_name) + try: + yield + finally: + contextLogger.reset(token) diff --git a/src/DIRAC/FrameworkSystem/private/standardLogging/test/Test_Logging_ContextVars.py b/src/DIRAC/FrameworkSystem/private/standardLogging/test/Test_Logging_ContextVars.py new file mode 100644 index 00000000000..c6dd5811e66 --- /dev/null +++ b/src/DIRAC/FrameworkSystem/private/standardLogging/test/Test_Logging_ContextVars.py @@ -0,0 +1,71 @@ +""" Test the context variable logger """ + +from DIRAC import gLogger +from DIRAC.FrameworkSystem.private.standardLogging.Logging import Logging +from DIRAC.FrameworkSystem.private.standardLogging.test.TestLogUtilities import gLoggerReset +from DIRAC.FrameworkSystem.private.standardLogging.LoggingContext import contextLogger, setContextLogger + + +class A: + def __init__(self): + # Get the logger from the context variable + self._defaultLogger = gLogger.getSubLogger("A") + + # Use a property to get and set the logger, this is necessary to use the context variable + @property + def log(self): + return contextLogger.get() or self._defaultLogger + + @log.setter + def log(self, value: Logging): + self._defaultLogger = value + + def do_something(self): + self.log.notice("A is doing something") + + +class B: + def __init__(self, a: A, pilotRef: str = None): + self.a = A() + + # Get the logger from the context variable + if pilotRef: + self.log = gLogger.getLocalSubLogger(f"[{pilotRef}]B") + contextLogger.set(self.log) + else: + self.log = gLogger.getSubLogger("B") + + def do_something_else(self): + with setContextLogger(self.log): + self.a.do_something() + self.log.notice("B is doing something else") + + +def test_contextvar_logger(): + capturedBackend, log, sublog = gLoggerReset() + + # Create an instance of A + a = A() + + # Create an instance of B and call its method without setting the pilotRef + # Log signature coming from A and B should be different + b1 = B(a) + b1.do_something_else() + assert "Framework/B NOTICE: A is doing something" in capturedBackend.getvalue() + assert "Framework/B NOTICE: B is doing something else" in capturedBackend.getvalue() + + # Create an instance of B and call its method with setting the pilotRef + # Log signature coming from A and B should be similar because of the pilotRef + capturedBackend.truncate(0) + + b2 = B(a, "pilotRef") + b2.do_something_else() + assert "Framework/[pilotRef]B NOTICE: A is doing something" in capturedBackend.getvalue() + assert "Framework/[pilotRef]B NOTICE: B is doing something else" in capturedBackend.getvalue() + + # Now we check that the logger of b1 is not the same as the logger of b2 (b1 should still use its own logger) + capturedBackend.truncate(0) + + b1.do_something_else() + assert "Framework/B NOTICE: A is doing something" in capturedBackend.getvalue() + assert "Framework/B NOTICE: B is doing something else" in capturedBackend.getvalue() diff --git a/src/DIRAC/WorkloadManagementSystem/Client/Matcher.py b/src/DIRAC/WorkloadManagementSystem/Client/Matcher.py index 087847e1644..0f2443f66e8 100644 --- a/src/DIRAC/WorkloadManagementSystem/Client/Matcher.py +++ b/src/DIRAC/WorkloadManagementSystem/Client/Matcher.py @@ -9,6 +9,7 @@ from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations from DIRAC.Core.Security import Properties from DIRAC.Core.Utilities.PrettyPrint import printDict +from DIRAC.FrameworkSystem.Client.Logger import setContextLogger from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus from DIRAC.WorkloadManagementSystem.Client import JobStatus, PilotStatus from DIRAC.WorkloadManagementSystem.Client.Limiter import Limiter @@ -16,7 +17,6 @@ from DIRAC.WorkloadManagementSystem.DB.JobLoggingDB import JobLoggingDB from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB from DIRAC.WorkloadManagementSystem.DB.TaskQueueDB import TaskQueueDB, multiValueMatchFields, singleValueDefFields -from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import setPilotRefLogger class PilotVersionError(Exception): @@ -61,7 +61,7 @@ def __init__(self, pilotAgentsDB=None, jobDB=None, tqDB=None, jlDB=None, opsHelp def selectJob(self, resourceDescription, credDict): """Main job selection function to find the highest priority job matching the resource capacity""" - with setPilotRefLogger(self.log): + with setContextLogger(self.log): startTime = time.time() resourceDict = self._getResourceDict(resourceDescription, credDict) diff --git a/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py b/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py index 08091fe031f..180c75a9412 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/JobDB.py @@ -21,6 +21,7 @@ from DIRAC.Core.Utilities.ClassAd.ClassAdLight import ClassAd from DIRAC.Core.Utilities.DErrno import EWMSJMAN, EWMSSUBM, cmpError from DIRAC.Core.Utilities.ReturnValues import S_ERROR, S_OK +from DIRAC.FrameworkSystem.Client.Logger import contextLogger from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus from DIRAC.WorkloadManagementSystem.Client import JobMinorStatus, JobStatus from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient @@ -32,7 +33,6 @@ extractJDL, fixJDL, ) -from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger class JobDB(DB): @@ -69,7 +69,7 @@ def __init__(self, parentLogger=None): @property def log(self): - return pilotRefLogger.get() or self._defaultLogger + return contextLogger.get() or self._defaultLogger @log.setter def log(self, value): diff --git a/src/DIRAC/WorkloadManagementSystem/DB/JobLoggingDB.py b/src/DIRAC/WorkloadManagementSystem/DB/JobLoggingDB.py index 62f5f8e3a58..b00e3ec3af2 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/JobLoggingDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/JobLoggingDB.py @@ -11,7 +11,7 @@ from DIRAC import S_ERROR, S_OK from DIRAC.Core.Base.DB import DB from DIRAC.Core.Utilities import TimeUtilities -from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger +from DIRAC.FrameworkSystem.Client.Logger import contextLogger MAGIC_EPOC_NUMBER = 1270000000 @@ -29,7 +29,7 @@ def __init__(self, parentLogger=None): @property def log(self): - return pilotRefLogger.get() or self._defaultLogger + return contextLogger.get() or self._defaultLogger @log.setter def log(self, value): diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index 6a64c07f5ea..0942b4d1fae 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py @@ -30,9 +30,9 @@ from DIRAC.Core.Base.DB import DB from DIRAC.Core.Utilities import DErrno from DIRAC.Core.Utilities.MySQL import _quotedList +from DIRAC.FrameworkSystem.Client.Logger import contextLogger from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus from DIRAC.WorkloadManagementSystem.Client import PilotStatus -from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger class PilotAgentsDB(DB): @@ -43,7 +43,7 @@ def __init__(self, parentLogger=None): @property def log(self): - return pilotRefLogger.get() or self._defaultLogger + return contextLogger.get() or self._defaultLogger @log.setter def log(self, value): diff --git a/src/DIRAC/WorkloadManagementSystem/DB/TaskQueueDB.py b/src/DIRAC/WorkloadManagementSystem/DB/TaskQueueDB.py index d2d426c92be..c9f97259f06 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/TaskQueueDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/TaskQueueDB.py @@ -11,8 +11,8 @@ from DIRAC.Core.Security import Properties from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations from DIRAC.ConfigurationSystem.Client.Helpers import Registry +from DIRAC.FrameworkSystem.Client.Logger import contextLogger from DIRAC.WorkloadManagementSystem.private.SharesCorrector import SharesCorrector -from DIRAC.WorkloadManagementSystem.Utilities.ContextVars import pilotRefLogger DEFAULT_GROUP_SHARE = 1000 TQ_MIN_SHARE = 0.001 @@ -68,7 +68,7 @@ def __init__(self, parentLogger=None): @property def log(self): - return pilotRefLogger.get() or self._defaultLogger + return contextLogger.get() or self._defaultLogger @log.setter def log(self, value): diff --git a/src/DIRAC/WorkloadManagementSystem/Service/MatcherHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/MatcherHandler.py index cd3dc992c48..7d7bf4428b6 100755 --- a/src/DIRAC/WorkloadManagementSystem/Service/MatcherHandler.py +++ b/src/DIRAC/WorkloadManagementSystem/Service/MatcherHandler.py @@ -56,7 +56,7 @@ def export_requestJob(self, resourceDescription): resourceDescription["Setup"] = self.serviceInfoDict["clientSetup"] credDict = self.getRemoteCredentials() - pilotRef = resourceDescription.get("PilotReference") + pilotRef = resourceDescription.get("PilotReference", "Unknown") try: opsHelper = Operations(group=credDict["group"]) diff --git a/src/DIRAC/WorkloadManagementSystem/Utilities/ContextVars.py b/src/DIRAC/WorkloadManagementSystem/Utilities/ContextVars.py deleted file mode 100644 index 1044c453764..00000000000 --- a/src/DIRAC/WorkloadManagementSystem/Utilities/ContextVars.py +++ /dev/null @@ -1,16 +0,0 @@ -""" Context variables for the Workload Management System """ - -# Context variable for the logger (adapted to the request of the pilot reference) -import contextvars -from contextlib import contextmanager - -pilotRefLogger = contextvars.ContextVar("PilotRefLogger", default=None) - - -@contextmanager -def setPilotRefLogger(logger_name): - token = pilotRefLogger.set(logger_name) - try: - yield - finally: - pilotRefLogger.reset(token)