265 lines
11 KiB
Python
265 lines
11 KiB
Python
# © 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, discriminator="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, discriminator="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)}
|
|
|
|
self.cached_obs: Optional[ObsType] = self.default_observation
|
|
|
|
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:
|
|
if not folder_state["scanned_this_step"]:
|
|
health_status = self.cached_obs["health_status"]
|
|
else:
|
|
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,
|
|
)
|