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

#915 - Refactored documentation and included APi docs, dependencies.

Browse Source 
- make files now re-build autosummary and deps file.
- Added typer and platformdirs to deps in pyproject.toml.
- Made root_is_pure = True in setup.py as platform/python specific wheels don't need to be built but the option is there should we need to.
-
Added an e2e test for primaite.main.run func.
This commit is contained in:
Chris McCarthy
2023-06-08 15:57:38 +01:00
parent 92334133c9
commit 5558750f13
26 changed files with 923 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/VERSION
View File

@@ -1 +1 @@
1.2.1
2.0.0dev0

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

@@ -2,7 +2,7 @@
from pathlib import Path
from typing import Final
from primaite import getLogger, USERS_CONFIG_DIR
from primaite import USERS_CONFIG_DIR, getLogger
_LOGGER = getLogger(__name__)
@@ -17,7 +17,7 @@ def ddos_basic_one_config_path() -> Path:
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_1_DDOS_basic.yaml"
if not path.exists():
msg = f"Example config not found. Please run 'primaite setup'"
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
@@ -32,7 +32,7 @@ def ddos_basic_two_config_path() -> Path:
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_2_DDOS_basic.yaml"
if not path.exists():
msg = f"Example config not found. Please run 'primaite setup'"
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
@@ -47,7 +47,7 @@ def dos_very_basic_config_path() -> Path:
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_3_DOS_very_basic.yaml"
if not path.exists():
msg = f"Example config not found. Please run 'primaite setup'"
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
@@ -62,7 +62,7 @@ def data_manipulation_config_path() -> Path:
"""
path = _EXAMPLE_LAY_DOWN / "lay_down_config_5_data_manipulation.yaml"
if not path.exists():
msg = f"Example config not found. Please run 'primaite setup'"
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)

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

@@ -1,17 +1,16 @@
# Crown Copyright (C) Dstl 2022. DEFCON 703. Shared in confidence.
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Union, Final
from typing import Any, Dict, Final, Union
import yaml
from primaite import getLogger, USERS_CONFIG_DIR
from primaite import USERS_CONFIG_DIR, getLogger
from primaite.common.enums import ActionType
_LOGGER = getLogger(__name__)
_EXAMPLE_TRAINING: Final[
Path] = USERS_CONFIG_DIR / "example_config" / "training"
_EXAMPLE_TRAINING: Final[Path] = USERS_CONFIG_DIR / "example_config" / "training"
@dataclass()
@@ -99,25 +98,37 @@ class TrainingConfig:
file_system_restoring_limit: int # The time take to restore a file
file_system_scanning_limit: int # The time taken to scan the file system
def to_dict(self, json_serializable: bool = True):
"""
Serialise the ``TrainingConfig`` as dict.
:param json_serializable: If True, Enums are converted to their
string name.
:return: The ``TrainingConfig`` as a dict.
"""
data = self.__dict__
if json_serializable:
data["action_type"] = self.action_type.value
return data
def main_training_config_path() -> Path:
"""
The path to the example training_config_main.yaml file
The path to the example training_config_main.yaml file.
:return: The file path.
"""
path = _EXAMPLE_TRAINING / "training_config_main.yaml"
if not path.exists():
msg = f"Example config not found. Please run 'primaite setup'"
msg = "Example config not found. Please run 'primaite setup'"
_LOGGER.critical(msg)
raise FileNotFoundError(msg)
return path
def load(file_path: Union[str, Path],
legacy_file: bool = False) -> TrainingConfig:
def load(file_path: Union[str, Path], legacy_file: bool = False) -> TrainingConfig:
"""
Read in a training config yaml file.
@@ -162,9 +173,7 @@ def load(file_path: Union[str, Path],
def convert_legacy_training_config_dict(
legacy_config_dict: Dict[str, Any],
num_steps: int = 256,
action_type: str = "ANY"
legacy_config_dict: Dict[str, Any], num_steps: int = 256, action_type: str = "ANY"
) -> Dict[str, Any]:
"""
Convert a legacy training config dict to the new format.

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

@@ -4,7 +4,6 @@
import copy
import csv
import logging
import os.path
from datetime import datetime
from pathlib import Path
from typing import Dict, Tuple, Union
@@ -80,10 +79,12 @@ class Primaite(Env):
self._training_config_path = training_config_path
self._lay_down_config_path = lay_down_config_path
self.config_values: TrainingConfig = training_config.load(training_config_path)
self.training_config: TrainingConfig = training_config.load(
training_config_path
)
# Number of steps in an episode
self.episode_steps = self.config_values.num_steps
self.episode_steps = self.training_config.num_steps
super(Primaite, self).__init__()
@@ -91,7 +92,7 @@ class Primaite(Env):
self.transaction_list = transaction_list
# The agent in use
self.agent_identifier = self.config_values.agent_identifier
self.agent_identifier = self.training_config.agent_identifier
# Create a dictionary to hold all the nodes
self.nodes: Dict[str, NodeUnion] = {}
@@ -135,6 +136,9 @@ class Primaite(Env):
# Create step count
self.step_count = 0
self.total_step_count: int = 0
"""The total number of time steps completed."""
# Create step info dictionary
self.step_info = {}
@@ -168,7 +172,7 @@ class Primaite(Env):
# 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.config_data = yaml.safe_load(file)
self.lay_down_config = yaml.safe_load(file)
self.load_lay_down_config()
# Store the node objects as node attributes
@@ -187,7 +191,6 @@ class Primaite(Env):
plt.tight_layout()
nx.draw_networkx(self.network, with_labels=True)
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
file_path = session_path / f"network_{timestamp_str}.png"
plt.savefig(file_path, format="PNG")
@@ -238,8 +241,6 @@ class Primaite(Env):
# Set up a csv to store the results of the training
try:
now = datetime.now() # current date and time
time = now.strftime("%Y%m%d_%H%M%S")
header = ["Episode", "Average Reward"]
file_name = f"average_reward_per_episode_{timestamp_str}.csv"
@@ -303,7 +304,8 @@ class Primaite(Env):
done = False
self.step_count += 1
# print("Episode step: " + str(self.stepCount))
self.total_step_count += 1
# print("Episode step: " + str(self.step_count))
# Need to clear traffic on all links first
for link_key, link_value in self.links.items():
@@ -375,13 +377,13 @@ class Primaite(Env):
self.green_iers,
self.red_iers,
self.step_count,
self.config_values,
self.training_config,
)
# print(f" Step {self.step_count} Reward: {str(reward)}")
self.total_reward += reward
if self.step_count == self.episode_steps:
self.average_reward = self.total_reward / self.step_count
if self.config_values.session_type == "EVALUATION":
if self.training_config.session_type == "EVALUATION":
# For evaluation, need to trigger the done value = True when
# step count is reached in order to prevent neverending episode
done = True
@@ -891,7 +893,7 @@ class Primaite(Env):
def load_lay_down_config(self):
"""Loads config data in order to build the environment configuration."""
for item in self.config_data:
for item in self.lay_down_config:
if item["itemType"] == "NODE":
# Create a node
self.create_node(item)
@@ -951,7 +953,7 @@ class Primaite(Env):
node_type,
node_priority,
node_hardware_state,
self.config_values,
self.training_config,
)
elif node_class == "ACTIVE":
# Active nodes have IP address, Software State and file system state
@@ -967,7 +969,7 @@ class Primaite(Env):
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
self.training_config,
)
elif node_class == "SERVICE":
# Service nodes have IP address, Software State, file system state and list of services
@@ -983,7 +985,7 @@ class Primaite(Env):
node_ip_address,
node_software_state,
node_file_system_state,
self.config_values,
self.training_config,
)
node_services = item["services"]
for service in node_services:
@@ -1257,7 +1259,7 @@ class Primaite(Env):
Uses config data config data in order to build the environment
configuration.
"""
for item in self.config_data:
for item in self.lay_down_config:
if item["itemType"] == "NODE":
# Reset a node's state (normal and reference)
self.reset_node(item)

147
src/primaite/main.py
View File

@@ -6,11 +6,12 @@ TODO: This will eventually be refactored out into a proper Session class.
TODO: The passing about of session_dir and timestamp_str is temporary and
will be cleaned up once we move to a proper Session class.
"""
import argparse
import json
import time
from datetime import datetime
from pathlib import Path
from typing import Final, Union
from uuid import uuid4
from stable_baselines3 import A2C, PPO
from stable_baselines3.common.evaluation import evaluate_policy
@@ -19,8 +20,7 @@ from stable_baselines3.ppo import MlpPolicy as PPOMlp
from primaite import SESSIONS_DIR, getLogger
from primaite.config.lay_down_config import data_manipulation_config_path
from primaite.config.training_config import TrainingConfig, \
main_training_config_path
from primaite.config.training_config import TrainingConfig, main_training_config_path
from primaite.environment.primaite_env import Primaite
from primaite.transactions.transactions_to_file import write_transaction_to_file
@@ -60,10 +60,7 @@ def run_generic(env: Primaite, config_values: TrainingConfig):
def run_stable_baselines3_ppo(
env: Primaite,
config_values: TrainingConfig,
session_path: Path,
timestamp_str: str
env: Primaite, config_values: TrainingConfig, session_path: Path, timestamp_str: str
):
"""
Run against a stable_baselines3 PPO agent.
@@ -98,8 +95,8 @@ def run_stable_baselines3_ppo(
# We're in a training session
print("Starting training session...")
_LOGGER.debug("Starting training session...")
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
for episode in range(config_values.num_episodes):
agent.learn(total_timesteps=config_values.num_steps)
_save_agent(agent, session_path, timestamp_str)
else:
# Default to being in an evaluation session
@@ -111,9 +108,7 @@ def run_stable_baselines3_ppo(
def run_stable_baselines3_a2c(
env: Primaite,
config_values: TrainingConfig,
session_path: Path, timestamp_str: str
env: Primaite, config_values: TrainingConfig, session_path: Path, timestamp_str: str
):
"""
Run against a stable_baselines3 A2C agent.
@@ -148,8 +143,8 @@ def run_stable_baselines3_a2c(
# We're in a training session
print("Starting training session...")
_LOGGER.debug("Starting training session...")
for episode in range(0, config_values.num_episodes):
agent.learn(total_timesteps=1)
for episode in range(config_values.num_episodes):
agent.learn(total_timesteps=config_values.num_steps)
_save_agent(agent, session_path, timestamp_str)
else:
# Default to being in an evaluation session
@@ -160,6 +155,68 @@ def run_stable_baselines3_a2c(
env.close()
def _write_session_metadata_file(
session_dir: Path, uuid: str, session_timestamp: datetime, env: Primaite
):
"""
Write the ``session_metadata.json`` file.
Creates a ``session_metadata.json`` in the ``session_dir`` directory
and adds the following key/value pairs:
- uuid: The UUID assigned to the session upon instantiation.
- start_datetime: The date & time the session started in iso format.
- end_datetime: NULL.
- total_episodes: NULL.
- total_time_steps: NULL.
- env:
- training_config:
- All training config items
- lay_down_config:
- All lay down config items
"""
metadata_dict = {
"uuid": uuid,
"start_datetime": session_timestamp.isoformat(),
"end_datetime": None,
"total_episodes": None,
"total_time_steps": None,
"env": {
"training_config": env.training_config.to_dict(json_serializable=True),
"lay_down_config": env.lay_down_config,
},
}
filepath = session_dir / "session_metadata.json"
_LOGGER.debug(f"Writing Session Metadata file: {filepath}")
with open(filepath, "w") as file:
json.dump(metadata_dict, file)
def _update_session_metadata_file(session_dir: Path, env: Primaite):
"""
Update the ``session_metadata.json`` file.
Updates the `session_metadata.json`` in the ``session_dir`` directory
with the following key/value pairs:
- end_datetime: NULL.
- total_episodes: NULL.
- total_time_steps: NULL.
"""
with open(session_dir / "session_metadata.json", "r") as file:
metadata_dict = json.load(file)
metadata_dict["end_datetime"] = datetime.now().isoformat()
metadata_dict["total_episodes"] = env.episode_count
metadata_dict["total_time_steps"] = env.total_step_count
filepath = session_dir / "session_metadata.json"
_LOGGER.debug(f"Updating Session Metadata file: {filepath}")
with open(filepath, "w") as file:
json.dump(metadata_dict, file)
def _save_agent(agent: OnPolicyAlgorithm, session_path: Path, timestamp_str: str):
"""
Persist an agent.
@@ -197,10 +254,7 @@ def _get_session_path(session_timestamp: datetime) -> Path:
return session_path
def run(
training_config_path: Union[str, Path],
lay_down_config_path: Union[str, Path]
):
def run(training_config_path: Union[str, Path], lay_down_config_path: Union[str, Path]):
"""Run the PrimAITE Session.
:param training_config_path: The training config filepath.
@@ -208,12 +262,12 @@ def run(
"""
# Welcome message
print("Welcome to the Primary-level AI Training Environment (PrimAITE)")
uuid = str(uuid4())
session_timestamp: Final[datetime] = datetime.now()
session_path = _get_session_path(session_timestamp)
session_dir = _get_session_path(session_timestamp)
timestamp_str = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S")
print(f"The output directory for this session is: {session_path}")
print(f"The output directory for this session is: {session_dir}")
# Create a list of transactions
# A transaction is an object holding the:
@@ -230,11 +284,17 @@ def run(
training_config_path=training_config_path,
lay_down_config_path=lay_down_config_path,
transaction_list=transaction_list,
session_path=session_path,
session_path=session_dir,
timestamp_str=timestamp_str,
)
config_values = env.config_values
print("Writing Session Metadata file...")
_write_session_metadata_file(
session_dir=session_dir, uuid=uuid, session_timestamp=session_timestamp, env=env
)
config_values = env.training_config
# Get the number of steps (which is stored in the child config file)
config_values.num_steps = env.episode_steps
@@ -246,14 +306,14 @@ def run(
run_stable_baselines3_ppo(
env=env,
config_values=config_values,
session_path=session_path,
session_path=session_dir,
timestamp_str=timestamp_str,
)
elif config_values.agent_identifier == "STABLE_BASELINES3_A2C":
run_stable_baselines3_a2c(
env=env,
config_values=config_values,
session_path=session_path,
session_path=session_dir,
timestamp_str=timestamp_str,
)
@@ -261,29 +321,32 @@ def run(
_LOGGER.debug("Session finished")
print("Saving transaction logs...")
_LOGGER.debug("Saving transaction logs...")
write_transaction_to_file(
transaction_list=transaction_list,
session_path=session_path,
session_path=session_dir,
timestamp_str=timestamp_str,
)
print("Updating Session Metadata file...")
_update_session_metadata_file(session_dir=session_dir, env=env)
print("Finished")
_LOGGER.debug("Finished")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--tc")
parser.add_argument("--ldc")
args = parser.parse_args()
if not args.tc:
_LOGGER.error(
"Please provide a training config file using the --tc " "argument"
)
if not args.ldc:
_LOGGER.error(
"Please provide a lay down config file using the --ldc " "argument"
)
run(training_config_path=args.tc, lay_down_config_path=args.ldc)
# if __name__ == "__main__":
# parser = argparse.ArgumentParser()
# parser.add_argument("--tc")
# parser.add_argument("--ldc")
# args = parser.parse_args()
# if not args.tc:
# _LOGGER.error(
# "Please provide a training config file using the --tc " "argument"
# )
# if not args.ldc:
# _LOGGER.error(
# "Please provide a lay down config file using the --ldc " "argument"
# )
# run(training_config_path=args.tc, lay_down_config_path=args.ldc)
run(main_training_config_path(), data_manipulation_config_path())

2
src/primaite/transactions/transactions_to_file.py
View File

@@ -82,9 +82,9 @@ def write_transaction_to_file(transaction_list, session_path: Path, timestamp_st
header = ["Timestamp", "Episode", "Step", "Reward"]
header = header + action_header + obs_header_initial + obs_header_new
try:
filename = session_path / f"all_transactions_{timestamp_str}.csv"
_LOGGER.debug(f"Saving transaction logs: {filename}")
csv_file = open(filename, "w", encoding="UTF8", newline="")
csv_writer = csv.writer(csv_file)
csv_writer.writerow(header)