Merge remote-tracking branch 'origin/4.0.0a1-dev' into feature/2913_Extensible_Rewards

This commit is contained in:
Marek Wolan
2025-01-02 15:11:25 +00:00
7 changed files with 3790 additions and 0 deletions

View File

@@ -1 +1,259 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict, Iterable, List, Optional
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
class FileObservation(AbstractObservation, identifier="FILE"):
"""File observation, provides status information about a file within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for FileObservation."""
file_name: str
"""Name of the file, used for querying simulation state dictionary."""
include_num_access: Optional[bool] = None
"""Whether to include the number of accesses to the file in the observation."""
file_system_requires_scan: Optional[bool] = None
"""If True, the file must be scanned to update the health state. Tf False, the true state is always shown."""
def __init__(self, where: WhereType, include_num_access: bool, file_system_requires_scan: bool) -> None:
"""
Initialise a file observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this file.
A typical location for a file might be
['network', 'nodes', <node_hostname>, 'file_system', 'folder', <folder_name>, 'files', <file_name>].
:type where: WhereType
:param include_num_access: Whether to include the number of accesses to the file in the observation.
:type include_num_access: bool
:param file_system_requires_scan: If True, the file must be scanned to update the health state. Tf False,
the true state is always shown.
:type file_system_requires_scan: bool
"""
self.where: WhereType = where
self.include_num_access: bool = include_num_access
self.file_system_requires_scan: bool = file_system_requires_scan
self.default_observation: ObsType = {"health_status": 0}
if self.include_num_access:
self.default_observation["num_access"] = 0
# TODO: allow these to be configured in yaml
self.high_threshold = 10
self.med_threshold = 5
self.low_threshold = 0
def _categorise_num_access(self, num_access: int) -> int:
"""
Represent number of file accesses as a categorical variable.
:param num_access: Number of file accesses.
:return: Bin number corresponding to the number of accesses.
"""
if num_access > self.high_threshold:
return 3
elif num_access > self.med_threshold:
return 2
elif num_access > self.low_threshold:
return 1
return 0
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the health status of the file and optionally the number of accesses.
:rtype: ObsType
"""
file_state = access_from_nested_dict(state, self.where)
if file_state is NOT_PRESENT_IN_STATE:
return self.default_observation
if self.file_system_requires_scan:
health_status = file_state["visible_status"]
else:
health_status = file_state["health_status"]
obs = {"health_status": health_status}
if self.include_num_access:
obs["num_access"] = self._categorise_num_access(file_state["num_access"])
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for file status.
:rtype: spaces.Space
"""
space = {"health_status": spaces.Discrete(6)}
if self.include_num_access:
space["num_access"] = spaces.Discrete(4)
return spaces.Dict(space)
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FileObservation:
"""
Create a file observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the file observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this file's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed file observation instance.
:rtype: FileObservation
:param file_system_requires_scan: If True, the folder must be scanned to update the health state. Tf False,
the true state is always shown.
:type file_system_requires_scan: bool
"""
return cls(
where=parent_where + ["files", config.file_name],
include_num_access=config.include_num_access,
file_system_requires_scan=config.file_system_requires_scan,
)
class FolderObservation(AbstractObservation, identifier="FOLDER"):
"""Folder observation, provides status information about a folder within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for FolderObservation."""
folder_name: str
"""Name of the folder, used for querying simulation state dictionary."""
files: List[FileObservation.ConfigSchema] = []
"""List of file configurations within the folder."""
num_files: Optional[int] = None
"""Number of spaces for file observations in this folder."""
include_num_access: Optional[bool] = None
"""Whether files in this folder should include the number of accesses in their observation."""
file_system_requires_scan: Optional[bool] = None
"""If True, the folder must be scanned to update the health state. Tf False, the true state is always shown."""
def __init__(
self,
where: WhereType,
files: Iterable[FileObservation],
num_files: int,
include_num_access: bool,
file_system_requires_scan: bool,
) -> None:
"""
Initialise a folder observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this folder.
A typical location for a folder might be ['network', 'nodes', <node_hostname>, 'folders', <folder_name>].
:type where: WhereType
:param files: List of file observation instances within the folder.
:type files: Iterable[FileObservation]
:param num_files: Number of files expected in the folder.
:type num_files: int
:param include_num_access: Whether to include the number of accesses to files in the observation.
:type include_num_access: bool
:param file_system_requires_scan: If True, the folder must be scanned to update the health state. Tf False,
the true state is always shown.
:type file_system_requires_scan: bool
"""
self.where: WhereType = where
self.file_system_requires_scan: bool = file_system_requires_scan
self.files: List[FileObservation] = files
while len(self.files) < num_files:
self.files.append(
FileObservation(
where=None,
include_num_access=include_num_access,
file_system_requires_scan=self.file_system_requires_scan,
)
)
while len(self.files) > num_files:
truncated_file = self.files.pop()
msg = f"Too many files in folder observation. Truncating file {truncated_file}"
_LOGGER.warning(msg)
self.default_observation = {
"health_status": 0,
}
if self.files:
self.default_observation["FILES"] = {i + 1: f.default_observation for i, f in enumerate(self.files)}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the health status of the folder and status of files within the folder.
:rtype: ObsType
"""
folder_state = access_from_nested_dict(state, self.where)
if folder_state is NOT_PRESENT_IN_STATE:
return self.default_observation
if self.file_system_requires_scan:
health_status = folder_state["visible_status"]
else:
health_status = folder_state["health_status"]
obs = {}
obs["health_status"] = health_status
if self.files:
obs["FILES"] = {i + 1: file.observe(state) for i, file in enumerate(self.files)}
return obs
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for folder status.
:rtype: spaces.Space
"""
shape = {"health_status": spaces.Discrete(6)}
if self.files:
shape["FILES"] = spaces.Dict({i + 1: f.space for i, f in enumerate(self.files)})
return spaces.Dict(shape)
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> FolderObservation:
"""
Create a folder observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the folder observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this folder's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed folder observation instance.
:rtype: FolderObservation
"""
where = parent_where + ["file_system", "folders", config.folder_name]
# pass down shared/common config items
for file_config in config.files:
file_config.include_num_access = config.include_num_access
file_config.file_system_requires_scan = config.file_system_requires_scan
files = [FileObservation.from_config(config=f, parent_where=where) for f in config.files]
return cls(
where=where,
files=files,
num_files=config.num_files,
include_num_access=config.include_num_access,
file_system_requires_scan=config.file_system_requires_scan,
)

View File

@@ -1 +1,153 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Any, Dict, List
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite import getLogger
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
_LOGGER = getLogger(__name__)
class LinkObservation(AbstractObservation, identifier="LINK"):
"""Link observation, providing information about a specific link within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for LinkObservation."""
link_reference: str
"""Reference identifier for the link."""
def __init__(self, where: WhereType) -> None:
"""
Initialise a link observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this link.
A typical location for a link might be ['network', 'links', <link_reference>].
:type where: WhereType
"""
self.where = where
self.default_observation: ObsType = {"PROTOCOLS": {"ALL": 0}}
def observe(self, state: Dict) -> Any:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing information about the link.
:rtype: Any
"""
link_state = access_from_nested_dict(state, self.where)
if link_state is NOT_PRESENT_IN_STATE:
self.where[-1] = "<->".join(self.where[-1].split("<->")[::-1]) # try swapping endpoint A and B
link_state = access_from_nested_dict(state, self.where)
if link_state is NOT_PRESENT_IN_STATE:
return self.default_observation
bandwidth = link_state["bandwidth"]
load = link_state["current_load"]
if load == 0:
utilisation_category = 0
else:
utilisation_fraction = load / bandwidth
utilisation_category = int(utilisation_fraction * 9) + 1
return {"PROTOCOLS": {"ALL": min(utilisation_category, 10)}}
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for link status.
:rtype: spaces.Space
"""
return spaces.Dict({"PROTOCOLS": spaces.Dict({"ALL": spaces.Discrete(11)})})
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> LinkObservation:
"""
Create a link observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the link observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this link.
A typical location might be ['network', 'links', <link_reference>].
:type parent_where: WhereType, optional
:return: Constructed link observation instance.
:rtype: LinkObservation
"""
link_reference = config.link_reference
if parent_where == []:
where = ["network", "links", link_reference]
else:
where = parent_where + ["links", link_reference]
return cls(where=where)
class LinksObservation(AbstractObservation, identifier="LINKS"):
"""Collection of link observations representing multiple links within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for LinksObservation."""
link_references: List[str]
"""List of reference identifiers for the links."""
def __init__(self, where: WhereType, links: List[LinkObservation]) -> None:
"""
Initialise a links observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for these links.
A typical location for links might be ['network', 'links'].
:type where: WhereType
:param links: List of link observations.
:type links: List[LinkObservation]
"""
self.where: WhereType = where
self.links: List[LinkObservation] = links
self.default_observation: ObsType = {i + 1: l.default_observation for i, l in enumerate(self.links)}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing information about multiple links.
:rtype: ObsType
"""
return {i + 1: l.observe(state) for i, l in enumerate(self.links)}
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for multiple links.
:rtype: spaces.Space
"""
return spaces.Dict({i + 1: l.space for i, l in enumerate(self.links)})
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> LinksObservation:
"""
Create a links observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the links observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about these links.
A typical location might be ['network'].
:type parent_where: WhereType, optional
:return: Constructed links observation instance.
:rtype: LinksObservation
"""
where = parent_where + ["network"]
link_cfgs = [LinkObservation.ConfigSchema(link_reference=ref) for ref in config.link_references]
links = [LinkObservation.from_config(c, parent_where=where) for c in link_cfgs]
return cls(where=where, links=links)

View File

@@ -1 +1,164 @@
# © Crown-owned copyright 2025, Defence Science and Technology Laboratory UK
from __future__ import annotations
from typing import Dict
from gymnasium import spaces
from gymnasium.core import ObsType
from primaite.game.agent.observations.observations import AbstractObservation, WhereType
from primaite.game.agent.utils import access_from_nested_dict, NOT_PRESENT_IN_STATE
class ServiceObservation(AbstractObservation, identifier="SERVICE"):
"""Service observation, shows status of a service in the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for ServiceObservation."""
service_name: str
"""Name of the service, used for querying simulation state dictionary"""
def __init__(self, where: WhereType) -> None:
"""
Initialise a service observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this service.
A typical location for a service might be ['network', 'nodes', <node_hostname>, 'services', <service_name>].
:type where: WhereType
"""
self.where = where
self.default_observation = {"operating_status": 0, "health_status": 0}
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Observation containing the operating status and health status of the service.
:rtype: ObsType
"""
service_state = access_from_nested_dict(state, self.where)
if service_state is NOT_PRESENT_IN_STATE:
return self.default_observation
return {
"operating_status": service_state["operating_state"],
"health_status": service_state["health_state_visible"],
}
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for service status.
:rtype: spaces.Space
"""
return spaces.Dict({"operating_status": spaces.Discrete(7), "health_status": spaces.Discrete(5)})
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ServiceObservation:
"""
Create a service observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the service observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this service's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed service observation instance.
:rtype: ServiceObservation
"""
return cls(where=parent_where + ["services", config.service_name])
class ApplicationObservation(AbstractObservation, identifier="APPLICATION"):
"""Application observation, shows the status of an application within the simulation environment."""
class ConfigSchema(AbstractObservation.ConfigSchema):
"""Configuration schema for ApplicationObservation."""
application_name: str
"""Name of the application, used for querying simulation state dictionary"""
def __init__(self, where: WhereType) -> None:
"""
Initialise an application observation instance.
:param where: Where in the simulation state dictionary to find the relevant information for this application.
A typical location for an application might be
['network', 'nodes', <node_hostname>, 'applications', <application_name>].
:type where: WhereType
"""
self.where = where
self.default_observation = {"operating_status": 0, "health_status": 0, "num_executions": 0}
# TODO: allow these to be configured in yaml
self.high_threshold = 10
self.med_threshold = 5
self.low_threshold = 0
def _categorise_num_executions(self, num_executions: int) -> int:
"""
Represent number of file accesses as a categorical variable.
:param num_access: Number of file accesses.
:return: Bin number corresponding to the number of accesses.
"""
if num_executions > self.high_threshold:
return 3
elif num_executions > self.med_threshold:
return 2
elif num_executions > self.low_threshold:
return 1
return 0
def observe(self, state: Dict) -> ObsType:
"""
Generate observation based on the current state of the simulation.
:param state: Simulation state dictionary.
:type state: Dict
:return: Obs containing the operating status, health status, and number of executions of the application.
:rtype: ObsType
"""
application_state = access_from_nested_dict(state, self.where)
if application_state is NOT_PRESENT_IN_STATE:
return self.default_observation
return {
"operating_status": application_state["operating_state"],
"health_status": application_state["health_state_visible"],
"num_executions": self._categorise_num_executions(application_state["num_executions"]),
}
@property
def space(self) -> spaces.Space:
"""
Gymnasium space object describing the observation space shape.
:return: Gymnasium space representing the observation space for application status.
:rtype: spaces.Space
"""
return spaces.Dict(
{
"operating_status": spaces.Discrete(7),
"health_status": spaces.Discrete(5),
"num_executions": spaces.Discrete(4),
}
)
@classmethod
def from_config(cls, config: ConfigSchema, parent_where: WhereType = []) -> ApplicationObservation:
"""
Create an application observation from a configuration schema.
:param config: Configuration schema containing the necessary information for the application observation.
:type config: ConfigSchema
:param parent_where: Where in the simulation state dictionary to find the information about this application's
parent node. A typical location for a node might be ['network', 'nodes', <node_hostname>].
:type parent_where: WhereType, optional
:return: Constructed application observation instance.
:rtype: ApplicationObservation
"""
return cls(where=parent_where + ["applications", config.application_name])