oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'origin/dev' into feature/1709-base-simulation-class

Browse Source
This commit is contained in:
Marek Wolan
2023-07-28 14:49:56 +01:00
parent b129c4fc97 fd28475e83
commit c1bcc372b1
18 changed files with 1485 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/primaite/agents/agent_abc.py
View File

@@ -52,6 +52,8 @@ class AgentSessionABC(ABC):
training_config_path: Optional[Union[str, Path]] = None,
lay_down_config_path: Optional[Union[str, Path]] = None,
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
Initialise an agent session from config files, or load a previous session.
@@ -64,6 +66,10 @@ class AgentSessionABC(ABC):
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param session_path: directory path of the session to load
"""
# initialise variables
@@ -72,6 +78,8 @@ class AgentSessionABC(ABC):
self._can_learn: bool = False
self._can_evaluate: bool = False
self.is_eval = False
self.legacy_training_config = legacy_training_config
self.legacy_lay_down_config = legacy_lay_down_config
self.session_timestamp: datetime = datetime.now()
@@ -91,12 +99,14 @@ class AgentSessionABC(ABC):
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Union[Path, str] = training_config_path
self._training_config: TrainingConfig = training_config.load(self._training_config_path)
self._training_config: TrainingConfig = training_config.load(
self._training_config_path, legacy_file=legacy_training_config
)
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Union[Path, str] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path)
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path, legacy_lay_down_config)
self.sb3_output_verbose_level = self._training_config.sb3_output_verbose_level
# set random UUID for session

12
src/primaite/agents/sb3.py
View File

@@ -26,6 +26,8 @@ class SB3Agent(AgentSessionABC):
training_config_path: Optional[Union[str, Path]] = None,
lay_down_config_path: Optional[Union[str, Path]] = None,
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
Initialise the SB3 Agent training session.
@@ -35,11 +37,17 @@ class SB3Agent(AgentSessionABC):
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:raises ValueError: If the training config contains an unexpected value for agent_framework (should be "SB3")
:raises ValueError: If the training config contains an unexpected value for agent_identifies (should be `PPO`
or `A2C`)
"""
super().__init__(training_config_path, lay_down_config_path, session_path)
super().__init__(
training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config
)
if not self._training_config.agent_framework == AgentFramework.SB3:
msg = f"Expected SB3 agent_framework, " f"got {self._training_config.agent_framework}"
_LOGGER.error(msg)
@@ -75,6 +83,8 @@ class SB3Agent(AgentSessionABC):
lay_down_config_path=self._lay_down_config_path,
session_path=self.session_path,
timestamp_str=self.timestamp_str,
legacy_training_config=self.legacy_training_config,
legacy_lay_down_config=self.legacy_lay_down_config,
)
# check if there is a zip file that needs to be loaded

23
src/primaite/cli.py
View File

@@ -18,9 +18,9 @@ app = typer.Typer()
@app.command()
def build_dirs() -> None:
"""Build the PrimAITE app directories."""
from primaite.setup import setup_app_dirs
from primaite import PRIMAITE_PATHS
setup_app_dirs.run()
PRIMAITE_PATHS.mkdirs()
@app.command()
@@ -137,7 +137,13 @@ def setup(overwrite_existing: bool = True) -> None:
@app.command()
def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[str] = None) -> None:
def session(
tc: Optional[str] = None,
ldc: Optional[str] = None,
load: Optional[str] = None,
legacy_tc: bool = False,
legacy_ldc: bool = False,
) -> None:
"""
Run a PrimAITE session.
@@ -153,6 +159,10 @@ def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[
will use the default training config and laydown config. Inversely, if a training config and laydown config
is passed while a session directory is passed, PrimAITE will load the session and ignore the training config
and laydown config.
legacy_tc: If the training config file is a legacy file from PrimAITE < 2.0.
legacy_ldf: If the lay down config file is a legacy file from PrimAITE < 2.0.
"""
from primaite.config.lay_down_config import dos_very_basic_config_path
from primaite.config.training_config import main_training_config_path
@@ -170,7 +180,12 @@ def session(tc: Optional[str] = None, ldc: Optional[str] = None, load: Optional[
if not ldc:
ldc = dos_very_basic_config_path()
run(training_config_path=tc, lay_down_config_path=ldc)
run(
training_config_path=tc,
lay_down_config_path=ldc,
legacy_training_config=legacy_tc,
legacy_lay_down_config=legacy_ldc,
)
@app.command()

43
src/primaite/config/lay_down_config.py
View File

@@ -1,7 +1,7 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
from logging import Logger
from pathlib import Path
from typing import Any, Dict, Final, Union
from typing import Any, Dict, Final, List, Union
import yaml
@@ -12,14 +12,43 @@ _LOGGER: Logger = getLogger(__name__)
_EXAMPLE_LAY_DOWN: Final[Path] = PRIMAITE_PATHS.user_config_path / "example_config" / "lay_down"
def convert_legacy_lay_down_config_dict(legacy_config_dict: Dict[str, Any]) -> Dict[str, Any]:
def convert_legacy_lay_down_config(legacy_config: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""
Convert a legacy lay down config dict to the new format.
Convert a legacy lay down config to the new format.
:param legacy_config_dict: A legacy lay down config dict.
:param legacy_config: A legacy lay down config.
"""
_LOGGER.warning("Legacy lay down config conversion not yet implemented")
return legacy_config_dict
field_conversion_map = {
"itemType": "item_type",
"portsList": "ports_list",
"serviceList": "service_list",
"baseType": "node_class",
"nodeType": "node_type",
"hardwareState": "hardware_state",
"softwareState": "software_state",
"startStep": "start_step",
"endStep": "end_step",
"fileSystemState": "file_system_state",
"ipAddress": "ip_address",
"missionCriticality": "mission_criticality",
}
new_config = []
for item in legacy_config:
if "itemType" in item:
if item["itemType"] in ["ACTIONS", "STEPS"]:
continue
new_dict = {}
for key in item.keys():
conversion_key = field_conversion_map.get(key)
if key == "id" and "itemType" in item:
if item["itemType"] == "NODE":
conversion_key = "node_id"
if conversion_key:
new_dict[conversion_key] = item[key]
else:
new_dict[key] = item[key]
new_config.append(new_dict)
return new_config
def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict:
@@ -39,7 +68,7 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> Dict:
_LOGGER.debug(f"Loading lay down config file: {file_path}")
if legacy_file:
try:
config = convert_legacy_lay_down_config_dict(config)
config = convert_legacy_lay_down_config(config)
except KeyError:
msg = (
f"Failed to convert lay down config file {file_path} "

18
src/primaite/config/training_config.py
View File

@@ -291,12 +291,14 @@ def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConf
if legacy_file:
try:
config = convert_legacy_training_config_dict(config)
except KeyError:
except KeyError as e:
msg = (
f"Failed to convert training config file {file_path} "
f"from legacy format. Attempting to use file as is."
)
_LOGGER.error(msg)
raise e
try:
return TrainingConfig.from_dict(config)
except TypeError as e:
@@ -314,6 +316,9 @@ def convert_legacy_training_config_dict(
agent_identifier: AgentIdentifier = AgentIdentifier.PPO,
action_type: ActionType = ActionType.ANY,
num_train_steps: int = 256,
num_eval_steps: int = 256,
num_train_episodes: int = 10,
num_eval_episodes: int = 1,
) -> Dict[str, Any]:
"""
Convert a legacy training config dict to the new format.
@@ -325,8 +330,14 @@ def convert_legacy_training_config_dict(
training configs don't have agent_identifier values.
:param action_type: The action space type to set as legacy training configs
don't have action_type values.
:param num_train_steps: The number of steps to set as legacy training configs
:param num_train_steps: The number of train steps to set as legacy training configs
don't have num_train_steps values.
:param num_eval_steps: The number of eval steps to set as legacy training configs
don't have num_eval_steps values.
:param num_train_episodes: The number of train episodes to set as legacy training configs
don't have num_train_episodes values.
:param num_eval_episodes: The number of eval episodes to set as legacy training configs
don't have num_eval_episodes values.
:return: The converted training config dict.
"""
config_dict = {
@@ -334,6 +345,9 @@ def convert_legacy_training_config_dict(
"agent_identifier": agent_identifier.name,
"action_type": action_type.name,
"num_train_steps": num_train_steps,
"num_eval_steps": num_eval_steps,
"num_train_episodes": num_train_episodes,
"num_eval_episodes": num_eval_episodes,
"sb3_output_verbose_level": SB3OutputVerboseLevel.INFO.name,
}
session_type_map = {"TRAINING": "TRAIN", "EVALUATION": "EVAL"}

21
src/primaite/environment/primaite_env.py
View File

@@ -10,7 +10,6 @@ from typing import Any, Dict, Final, List, Tuple, Union
import networkx as nx
import numpy as np
import yaml
from gym import Env, spaces
from matplotlib import pyplot as plt
@@ -34,6 +33,7 @@ from primaite.common.enums import (
)
from primaite.common.service import Service
from primaite.config import training_config
from primaite.config.lay_down_config import load
from primaite.config.training_config import TrainingConfig
from primaite.environment.observations import ObservationsHandler
from primaite.environment.reward import calculate_reward_function
@@ -68,6 +68,8 @@ class Primaite(Env):
lay_down_config_path: Union[str, Path],
session_path: Path,
timestamp_str: str,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
The Primaite constructor.
@@ -76,13 +78,19 @@ class Primaite(Env):
:param lay_down_config_path: The lay down config filepath.
:param session_path: The directory path the session is writing to.
:param timestamp_str: The session timestamp in the format: <yyyy-mm-dd>_<hh-mm- ss>.
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
self.session_path: Final[Path] = session_path
self.timestamp_str: Final[str] = timestamp_str
self._training_config_path: Union[str, Path] = training_config_path
self._lay_down_config_path: Union[str, Path] = lay_down_config_path
self.legacy_training_config = legacy_training_config
self.legacy_lay_down_config = legacy_lay_down_config
self.training_config: TrainingConfig = training_config.load(training_config_path)
self.training_config: TrainingConfig = training_config.load(training_config_path, self.legacy_training_config)
_LOGGER.info(f"Using: {str(self.training_config)}")
# Number of steps in an episode
@@ -191,11 +199,8 @@ class Primaite(Env):
self._obs_space_description: List[str] = None
"The env observation space description for transactions writing"
# Open the config file and build the environment laydown
with open(self._lay_down_config_path, "r") as file:
# Open the config file and build the environment laydown
self.lay_down_config = yaml.safe_load(file)
self.load_lay_down_config()
self.lay_down_config = load(self._lay_down_config_path, self.legacy_lay_down_config)
self.load_lay_down_config()
# Store the node objects as node attributes
# (This is so we can access them as objects)
@@ -1027,7 +1032,7 @@ class Primaite(Env):
acl_rule_destination = item["destination"]
acl_rule_protocol = item["protocol"]
acl_rule_port = item["port"]
acl_rule_position = item["position"]
acl_rule_position = item.get("position")
self.acl.add_rule(
acl_rule_permission,

18
src/primaite/main.py
View File

@@ -14,18 +14,26 @@ def run(
training_config_path: Optional[Union[str, Path]] = "",
lay_down_config_path: Optional[Union[str, Path]] = "",
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
Run the PrimAITE Session.
:param training_config_path: YAML file containing configurable items defined in
`primaite.config.training_config.TrainingConfig`
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:type training_config_path: Union[path, str]
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
session = PrimaiteSession(training_config_path, lay_down_config_path, session_path)
session = PrimaiteSession(
training_config_path, lay_down_config_path, session_path, legacy_training_config, legacy_lay_down_config
)
session.setup()
session.learn()

22
src/primaite/primaite_session.py
View File

@@ -34,6 +34,8 @@ class PrimaiteSession:
training_config_path: Optional[Union[str, Path]] = "",
lay_down_config_path: Optional[Union[str, Path]] = "",
session_path: Optional[Union[str, Path]] = None,
legacy_training_config: bool = False,
legacy_lay_down_config: bool = False,
) -> None:
"""
The PrimaiteSession constructor.
@@ -44,12 +46,18 @@ class PrimaiteSession:
:param lay_down_config_path: YAML file containing configurable items for generating network laydown.
:type lay_down_config_path: Union[path, str]
:param session_path: directory path of the session to load
:param legacy_training_config: True if the training config file is a legacy file from PrimAITE < 2.0,
otherwise False.
:param legacy_lay_down_config: True if the lay_down config file is a legacy file from PrimAITE < 2.0,
otherwise False.
"""
self._agent_session: AgentSessionABC = None # noqa
self.session_path: Path = session_path # noqa
self.timestamp_str: str = None # noqa
self.learning_path: Path = None # noqa
self.evaluation_path: Path = None # noqa
self.legacy_training_config = legacy_training_config
self.legacy_lay_down_config = legacy_lay_down_config
# check if session path is provided
if session_path is not None:
@@ -67,12 +75,14 @@ class PrimaiteSession:
if not isinstance(training_config_path, Path):
training_config_path = Path(training_config_path)
self._training_config_path: Final[Union[Path, str]] = training_config_path
self._training_config: Final[TrainingConfig] = training_config.load(self._training_config_path)
self._training_config: Final[TrainingConfig] = training_config.load(
self._training_config_path, legacy_training_config
)
if not isinstance(lay_down_config_path, Path):
lay_down_config_path = Path(lay_down_config_path)
self._lay_down_config_path: Final[Union[Path, str]] = lay_down_config_path
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path) # noqa
self._lay_down_config: Dict = lay_down_config.load(self._lay_down_config_path, legacy_lay_down_config) # noqa
def setup(self) -> None:
"""Performs the session setup."""
@@ -139,7 +149,13 @@ class PrimaiteSession:
elif self._training_config.agent_framework == AgentFramework.SB3:
_LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.SB3}")
# Stable Baselines3 Agent
self._agent_session = SB3Agent(self._training_config_path, self._lay_down_config_path, self.session_path)
self._agent_session = SB3Agent(
self._training_config_path,
self._lay_down_config_path,
self.session_path,
self.legacy_training_config,
self.legacy_lay_down_config,
)
elif self._training_config.agent_framework == AgentFramework.RLLIB:
_LOGGER.debug(f"PrimaiteSession Setup: Agent Framework = {AgentFramework.RLLIB}")