From 4266618ba57346f94c164fa8a7581507713b05d0 Mon Sep 17 00:00:00 2001 From: Marek Wolan Date: Fri, 5 Jan 2024 11:25:57 +0000 Subject: [PATCH] Make pcap logs and sys logs optional --- CHANGELOG.md | 2 +- .../config/_package_data/example_config.yaml | 2 ++ src/primaite/session/io.py | 9 +++++++-- src/primaite/session/session.py | 2 +- src/primaite/simulator/__init__.py | 6 ++++-- .../simulator/system/core/packet_capture.py | 8 ++++++-- src/primaite/simulator/system/core/sys_log.py | 18 +++++++++++++----- 7 files changed, 34 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ccaa411a..9e44efe3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] - +- Made packet capture and system logging optional ### Added diff --git a/src/primaite/config/_package_data/example_config.yaml b/src/primaite/config/_package_data/example_config.yaml index a764703b..81c9643e 100644 --- a/src/primaite/config/_package_data/example_config.yaml +++ b/src/primaite/config/_package_data/example_config.yaml @@ -14,6 +14,8 @@ io_settings: save_checkpoints: true checkpoint_interval: 5 save_step_metadata: false + save_pcap_logs: true + save_sys_logs: true game: diff --git a/src/primaite/session/io.py b/src/primaite/session/io.py index 0d80a385..b4b740e9 100644 --- a/src/primaite/session/io.py +++ b/src/primaite/session/io.py @@ -24,9 +24,13 @@ class SessionIOSettings(BaseModel): save_transactions: bool = True """Whether to save transactions, If true, the session path will have a transactions folder.""" save_tensorboard_logs: bool = False - """Whether to save tensorboard logs. If true, the session path will have a tenorboard_logs folder.""" + """Whether to save tensorboard logs. If true, the session path will have a tensorboard_logs folder.""" save_step_metadata: bool = False """Whether to save the RL agents' action, environment state, and other data at every single step.""" + save_pcap_logs: bool = False + """Whether to save PCAP logs.""" + save_sys_logs: bool = False + """Whether to save system logs.""" class SessionIO: @@ -39,9 +43,10 @@ class SessionIO: def __init__(self, settings: SessionIOSettings = SessionIOSettings()) -> None: self.settings: SessionIOSettings = settings self.session_path: Path = self.generate_session_path() - # set global SIM_OUTPUT path SIM_OUTPUT.path = self.session_path / "simulation_output" + SIM_OUTPUT.save_pcap_logs = self.settings.save_pcap_logs + SIM_OUTPUT.save_sys_logs = self.settings.save_sys_logs # warning TODO: must be careful not to re-initialise sessionIO because it will create a new path each time it's # possible refactor needed diff --git a/src/primaite/session/session.py b/src/primaite/session/session.py index ef462d83..0197ac9d 100644 --- a/src/primaite/session/session.py +++ b/src/primaite/session/session.py @@ -54,7 +54,7 @@ class PrimaiteSession: self.policy: PolicyABC """The reinforcement learning policy.""" - self.io_manager = SessionIO() + self.io_manager: Optional["SessionIO"] = None """IO manager for the session.""" self.game: PrimaiteGame = game diff --git a/src/primaite/simulator/__init__.py b/src/primaite/simulator/__init__.py index 19c86e28..aebd77cf 100644 --- a/src/primaite/simulator/__init__.py +++ b/src/primaite/simulator/__init__.py @@ -7,11 +7,13 @@ from primaite import _PRIMAITE_ROOT __all__ = ["SIM_OUTPUT"] -class __SimOutput: +class _SimOutput: def __init__(self): self._path: Path = ( _PRIMAITE_ROOT.parent.parent / "simulation_output" / datetime.now().strftime("%Y-%m-%d_%H-%M-%S") ) + self.save_pcap_logs: bool = False + self.save_sys_logs: bool = False @property def path(self) -> Path: @@ -23,4 +25,4 @@ class __SimOutput: self._path.mkdir(exist_ok=True, parents=True) -SIM_OUTPUT = __SimOutput() +SIM_OUTPUT = _SimOutput() diff --git a/src/primaite/simulator/system/core/packet_capture.py b/src/primaite/simulator/system/core/packet_capture.py index 1539e024..bfb6a055 100644 --- a/src/primaite/simulator/system/core/packet_capture.py +++ b/src/primaite/simulator/system/core/packet_capture.py @@ -41,6 +41,9 @@ class PacketCapture: def setup_logger(self): """Set up the logger configuration.""" + if not SIM_OUTPUT.save_pcap_logs: + return + log_path = self._get_log_path() file_handler = logging.FileHandler(filename=log_path) @@ -88,5 +91,6 @@ class PacketCapture: :param frame: The PCAP frame to capture. """ - msg = frame.model_dump_json() - self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL + if SIM_OUTPUT.save_pcap_logs: + msg = frame.model_dump_json() + self.logger.log(level=60, msg=msg) # Log at custom log level > CRITICAL diff --git a/src/primaite/simulator/system/core/sys_log.py b/src/primaite/simulator/system/core/sys_log.py index 41ce8fee..00e6920b 100644 --- a/src/primaite/simulator/system/core/sys_log.py +++ b/src/primaite/simulator/system/core/sys_log.py @@ -41,6 +41,9 @@ class SysLog: The logger is set to the DEBUG level, and is equipped with a handler that writes to a file and filters out JSON-like messages. """ + if not SIM_OUTPUT.save_sys_logs: + return + log_path = self._get_log_path() file_handler = logging.FileHandler(filename=log_path) file_handler.setLevel(logging.DEBUG) @@ -91,7 +94,8 @@ class SysLog: :param msg: The message to be logged. """ - self.logger.debug(msg) + if SIM_OUTPUT.save_sys_logs: + self.logger.debug(msg) def info(self, msg: str): """ @@ -99,7 +103,8 @@ class SysLog: :param msg: The message to be logged. """ - self.logger.info(msg) + if SIM_OUTPUT.save_sys_logs: + self.logger.info(msg) def warning(self, msg: str): """ @@ -107,7 +112,8 @@ class SysLog: :param msg: The message to be logged. """ - self.logger.warning(msg) + if SIM_OUTPUT.save_sys_logs: + self.logger.warning(msg) def error(self, msg: str): """ @@ -115,7 +121,8 @@ class SysLog: :param msg: The message to be logged. """ - self.logger.error(msg) + if SIM_OUTPUT.save_sys_logs: + self.logger.error(msg) def critical(self, msg: str): """ @@ -123,4 +130,5 @@ class SysLog: :param msg: The message to be logged. """ - self.logger.critical(msg) + if SIM_OUTPUT.save_sys_logs: + self.logger.critical(msg)