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

Merge branch 'dev' into feature/2041_2042-Add-NTP-Services

Browse Source
This commit is contained in:
Nick Todd
2023-12-07 14:55:01 +00:00
parent 44ada941e6 e93fb85672
commit c3015b7f0e
10 changed files with 78 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
.azure/azure-ci-build-pipeline.yaml
View File

@@ -18,36 +18,43 @@ parameters:
py: '3.8'
img: 'ubuntu-latest'
every_time: false
publish_coverage: false
- job_name: 'UbuntuPython310'
py: '3.10'
img: 'ubuntu-latest'
every_time: true
publish_coverage: true
- job_name: 'WindowsPython38'
py: '3.8'
img: 'windows-latest'
every_time: false
publish_coverage: false
- job_name: 'WindowsPython310'
py: '3.10'
img: 'windows-latest'
every_time: false
publish_coverage: false
- job_name: 'MacOSPython38'
py: '3.8'
img: 'macOS-latest'
every_time: false
publish_coverage: false
- job_name: 'MacOSPython310'
py: '3.10'
img: 'macOS-latest'
every_time: false
publish_coverage: false
stages:
- stage: Test
jobs:
- ${{ each item in parameters.matrix }}:
- job: ${{ item.job_name }}
timeoutInMinutes: 90
cancelTimeoutInMinutes: 1
pool:
vmImage: ${{ item.img }}
condition: or( eq(variables['Build.Reason'], 'PullRequest'), ${{ item.every_time }} )
condition: and(succeeded(), or( eq(variables['Build.Reason'], 'PullRequest'), ${{ item.every_time }} ))
steps:
- task: UsePythonVersion@0
@@ -109,12 +116,12 @@ stages:
- publish: $(System.DefaultWorkingDirectory)/htmlcov/
# publish the html report - so we can debug the coverage if needed
condition: ${{ item.every_time }} # should only be run once
condition: ${{ item.publish_coverage }} # should only be run once
artifact: coverage_report
- task: PublishCodeCoverageResults@2
# publish the code coverage so it can be viewed in the run coverage page
condition: ${{ item.every_time }} # should only be run once
condition: ${{ item.publish_coverage }} # should only be run once
inputs:
codeCoverageTool: Cobertura
summaryFileLocation: '$(System.DefaultWorkingDirectory)/**/coverage.xml'

4
docs/source/primaite_session.rst
View File

@@ -31,4 +31,6 @@ Outputs
Running a session creates a session output directory in your user data folder. The filepath looks like this:
``~/primaite/3.0.0/sessions/YYYY-MM-DD/HH-MM-SS/``. This folder contains the simulation sys logs generated by each node,
the saved agent checkpoints, and final model.
the saved agent checkpoints, and final model. The folder also contains a .json file for each episode step that
contains the action, reward, and simulation state. These can be found in
``~/primaite/3.0.0/sessions/YYYY-MM-DD/HH-MM-SS/simulation_output/episode_<n>/step_metadata/step_<n>.json``

9
src/primaite/__init__.py
View File

@@ -1,4 +1,5 @@
# © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
import datetime as datetime
import logging
import logging.config
import shutil
@@ -38,6 +39,7 @@ class _PrimaitePaths:
self.app_config_file_path = self.generate_app_config_file_path()
self.app_log_dir_path = self.generate_app_log_dir_path()
self.app_log_file_path = self.generate_app_log_file_path()
self.episode_log_file_path = self.generate_episode_log_file_path()
def _get_dirs_properties(self) -> List[str]:
class_items = self.__class__.__dict__.items()
@@ -105,6 +107,13 @@ class _PrimaitePaths:
"""The PrimAITE app log file path."""
return self.app_log_dir_path / "primaite.log"
def generate_episode_log_file_path(self) -> Path:
"""The PrimAITE app episode step log file path."""
date_string = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
self.episode_log_dir_path = self.app_log_dir_path / date_string
self.episode_log_dir_path.mkdir(exist_ok=True, parents=True)
return self.episode_log_dir_path / "episode.log"
def __repr__(self) -> str:
properties_str = ", ".join([f"{p}='{getattr(self, p)}'" for p in self._get_dirs_properties()])
return f"{self.__class__.__name__}({properties_str})"

1
src/primaite/config/_package_data/example_config.yaml
View File

@@ -13,6 +13,7 @@ training_config:
io_settings:
save_checkpoints: true
checkpoint_interval: 5
save_step_metadata: false
game:

1
src/primaite/config/_package_data/example_config_2_rl_agents.yaml
View File

@@ -9,6 +9,7 @@ training_config:
io_settings:
save_checkpoints: true
checkpoint_interval: 5
save_step_metadata: false
game:

9
src/primaite/game/game.py
View File

@@ -10,6 +10,7 @@ from primaite.game.agent.data_manipulation_bot import DataManipulationAgent
from primaite.game.agent.interface import AbstractAgent, AgentSettings, ProxyAgent, RandomAgent
from primaite.game.agent.observations import ObservationManager
from primaite.game.agent.rewards import RewardFunction
from primaite.session.io import SessionIO, SessionIOSettings
from primaite.simulator.network.hardware.base import NIC, NodeOperatingState
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.router import ACLAction, Router
@@ -86,6 +87,9 @@ class PrimaiteGame:
self.ref_map_links: Dict[str, str] = {}
"""Mapping from human-readable link reference to link object. Used when parsing config files."""
self.save_step_metadata: bool = False
"""Whether to save the RL agents' action, environment state, and other data at every single step."""
def step(self):
"""
Perform one step of the simulation/agent loop.
@@ -182,8 +186,13 @@ class PrimaiteGame:
:return: A PrimaiteGame object.
:rtype: PrimaiteGame
"""
io_settings = cfg.get("io_settings", {})
_ = SessionIO(SessionIOSettings(**io_settings))
# Instantiating this ensures that the game saves to the correct output dir even without being part of a session
game = cls()
game.options = PrimaiteGameOptions(**cfg["game"])
game.save_step_metadata = cfg.get("io_settings", {}).get("save_step_metadata") or False
# 1. create simulation
sim = game.simulation

39
src/primaite/session/environment.py
View File

@@ -1,3 +1,4 @@
import json
from typing import Any, Dict, Final, Optional, SupportsFloat, Tuple
import gymnasium
@@ -6,6 +7,7 @@ from ray.rllib.env.multi_agent_env import MultiAgentEnv
from primaite.game.agent.interface import ProxyAgent
from primaite.game.game import PrimaiteGame
from primaite.simulator import SIM_OUTPUT
class PrimaiteGymEnv(gymnasium.Env):
@@ -30,6 +32,7 @@ class PrimaiteGymEnv(gymnasium.Env):
self.game.apply_agent_actions()
self.game.advance_timestep()
state = self.game.get_sim_state()
self.game.update_agents(state)
next_obs = self._get_obs()
@@ -37,8 +40,26 @@ class PrimaiteGymEnv(gymnasium.Env):
terminated = False
truncated = self.game.calculate_truncated()
info = {}
if self.game.save_step_metadata:
self._write_step_metadata_json(action, state, reward)
return next_obs, reward, terminated, truncated, info
def _write_step_metadata_json(self, action: int, state: Dict, reward: int):
output_dir = SIM_OUTPUT.path / f"episode_{self.game.episode_counter}" / "step_metadata"
output_dir.mkdir(parents=True, exist_ok=True)
path = output_dir / f"step_{self.game.step_counter}.json"
data = {
"episode": self.game.episode_counter,
"step": self.game.step_counter,
"action": int(action),
"reward": int(reward),
"state": state,
}
with open(path, "w") as file:
json.dump(data, file)
def reset(self, seed: Optional[int] = None) -> Tuple[ObsType, Dict[str, Any]]:
"""Reset the environment."""
print(
@@ -162,8 +183,26 @@ class PrimaiteRayMARLEnv(MultiAgentEnv):
infos = {}
terminateds["__all__"] = len(self.terminateds) == len(self.agents)
truncateds["__all__"] = self.game.calculate_truncated()
if self.game.save_step_metadata:
self._write_step_metadata_json(actions, state, rewards)
return next_obs, rewards, terminateds, truncateds, infos
def _write_step_metadata_json(self, actions: Dict, state: Dict, rewards: Dict):
output_dir = SIM_OUTPUT.path / f"episode_{self.game.episode_counter}" / "step_metadata"
output_dir.mkdir(parents=True, exist_ok=True)
path = output_dir / f"step_{self.game.step_counter}.json"
data = {
"episode": self.game.episode_counter,
"step": self.game.step_counter,
"actions": {agent_name: int(action) for agent_name, action in actions.items()},
"reward": rewards,
"state": state,
}
with open(path, "w") as file:
json.dump(data, file)
def _get_obs(self) -> Dict[str, ObsType]:
"""Return the current observation."""
obs = {}

2
src/primaite/session/io.py
View File

@@ -25,6 +25,8 @@ class SessionIOSettings(BaseModel):
"""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."""
save_step_metadata: bool = False
"""Whether to save the RL agents' action, environment state, and other data at every single step."""
class SessionIO:

4
src/primaite/simulator/network/hardware/nodes/router.py
View File

@@ -66,9 +66,9 @@ class ACLRule(SimComponent):
state = super().describe_state()
state["action"] = self.action.value
state["protocol"] = self.protocol.value if self.protocol else None
state["src_ip_address"] = self.src_ip_address if self.src_ip_address else None
state["src_ip_address"] = str(self.src_ip_address) if self.src_ip_address else None
state["src_port"] = self.src_port.value if self.src_port else None
state["dst_ip_address"] = self.dst_ip_address if self.dst_ip_address else None
state["dst_ip_address"] = str(self.dst_ip_address) if self.dst_ip_address else None
state["dst_port"] = self.dst_port.value if self.dst_port else None
return state

2
src/primaite/simulator/network/hardware/nodes/switch.py
View File

@@ -57,7 +57,7 @@ class Switch(Node):
state = super().describe_state()
state["ports"] = {port_num: port.describe_state() for port_num, port in self.switch_ports.items()}
state["num_ports"] = self.num_ports # redundant?
state["mac_address_table"] = {mac: port for mac, port in self.mac_address_table.items()}
state["mac_address_table"] = {mac: port.port_num for mac, port in self.mac_address_table.items()}
return state
def _add_mac_table_entry(self, mac_address: str, switch_port: SwitchPort):