"""Module for handling configurable observation spaces in PrimAITE.""" import logging from abc import ABC, abstractmethod from typing import TYPE_CHECKING, List, Tuple import numpy as np from gym import spaces from primaite.common.enums import FileSystemState, HardwareState, SoftwareState from primaite.nodes.active_node import ActiveNode from primaite.nodes.service_node import ServiceNode # This dependency is only needed for type hints, # TYPE_CHECKING is False at runtime and True when typecheckers are performing typechecking # Therefore, this avoids circular dependency problem. if TYPE_CHECKING: from primaite.environment.primaite_env import Primaite _LOGGER = logging.getLogger(__name__) class AbstractObservationComponent(ABC): """Represents a part of the PrimAITE observation space.""" @abstractmethod def __init__(self, env: "Primaite"): _LOGGER.info(f"Initialising {self} observation component") self.env: "Primaite" = env self.space: spaces.Space self.current_observation: np.ndarray # type might be too restrictive? return NotImplemented @abstractmethod def update(self): """Look at the environment and update the current observation value.""" self.current_observation = NotImplemented class NodeLinkTable(AbstractObservationComponent): """Table with nodes/links as rows and hardware/software status as cols. Initialise the observation space with the BOX option chosen. This will create the observation space formatted as a table of integers. There is one row per node, followed by one row per link. Columns are as follows: * node/link ID * node hardware status / 0 for links * node operating system status (if active/service) / 0 for links * node file system status (active/service only) / 0 for links * node service1 status / traffic load from that service for links * node service2 status / traffic load from that service for links * ... * node serviceN status / traffic load from that service for links For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be ``(12, 7)`` #TODO: clean up description """ _FIXED_PARAMETERS = 4 _MAX_VAL = 1_000_000 _DATA_TYPE = np.int64 def __init__(self, env: "Primaite"): super().__init__(env) # 1. Define the shape of your observation space component num_items = self.env.num_links + self.env.num_nodes num_columns = self.env.num_services + self._FIXED_PARAMETERS observation_shape = (num_items, num_columns) # 2. Create Observation space self.space = spaces.Box( low=0, high=self._MAX_VAL, shape=observation_shape, dtype=self._DATA_TYPE, ) # 3. Initialise Observation with zeroes self.current_observation = np.zeros(observation_shape, dtype=self._DATA_TYPE) def update(self): """Update the observation. Update the environment's observation state based on the current status of nodes and links. The structure of the observation space is described in :func:`~_init_box_observations` This function can only be called if the observation space setting is set to BOX. TODO: complete description.. """ item_index = 0 nodes = self.env.nodes links = self.env.links # Do nodes first for _, node in nodes.items(): self.current_observation[item_index][0] = int(node.node_id) self.current_observation[item_index][1] = node.hardware_state.value if isinstance(node, ActiveNode) or isinstance(node, ServiceNode): self.current_observation[item_index][2] = node.software_state.value self.current_observation[item_index][ 3 ] = node.file_system_state_observed.value else: self.current_observation[item_index][2] = 0 self.current_observation[item_index][3] = 0 service_index = 4 if isinstance(node, ServiceNode): for service in self.env.services_list: if node.has_service(service): self.current_observation[item_index][ service_index ] = node.get_service_state(service).value else: self.current_observation[item_index][service_index] = 0 service_index += 1 else: # Not a service node for service in self.env.services_list: self.current_observation[item_index][service_index] = 0 service_index += 1 item_index += 1 # Now do links for _, link in links.items(): self.current_observation[item_index][0] = int(link.get_id()) self.current_observation[item_index][1] = 0 self.current_observation[item_index][2] = 0 self.current_observation[item_index][3] = 0 protocol_list = link.get_protocol_list() protocol_index = 0 for protocol in protocol_list: self.current_observation[item_index][ protocol_index + 4 ] = protocol.get_load() protocol_index += 1 item_index += 1 class NodeStatuses(AbstractObservationComponent): """TODO: complete description. This will create the observation space with node observations followed by link observations. Each node has 3 elements in the observation space plus 1 per service, more specifically: * hardware state * operating system state * file system state * service states (one per service) """ _DATA_TYPE = np.int64 def __init__(self, env: "Primaite"): super().__init__(env) # 1. Define the shape of your observation space component shape = [ len(HardwareState) + 1, len(SoftwareState) + 1, len(FileSystemState) + 1, ] services_shape = [len(SoftwareState) + 1] * self.env.num_services shape = shape + services_shape # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) def update(self): """TODO: complete description. Update the environment's observation state based on the current status of nodes and links. The structure of the observation space is described in :func:`~_init_multidiscrete_observations` This function can only be called if the observation space setting is set to MULTIDISCRETE. """ obs = [] for _, node in self.env.nodes.items(): hardware_state = node.hardware_state.value software_state = 0 file_system_state = 0 service_states = [0] * self.env.num_services if isinstance(node, ActiveNode): software_state = node.software_state.value file_system_state = node.file_system_state_observed.value if isinstance(node, ServiceNode): for i, service in enumerate(self.env.services_list): if node.has_service(service): service_states[i] = node.get_service_state(service).value obs.extend([hardware_state, software_state, file_system_state, *service_states]) self.current_observation[:] = obs class LinkTrafficLevels(AbstractObservationComponent): """TODO: complete description. Each link has one element in the observation space, corresponding to the traffic load, it can take the following values: 0 = No traffic (0% of bandwidth) 1 = No traffic (0%-33% of bandwidth) 2 = No traffic (33%-66% of bandwidth) 3 = No traffic (66%-100% of bandwidth) 4 = No traffic (100% of bandwidth) """ _DATA_TYPE = np.int64 def __init__( self, env: "Primaite", combine_service_traffic: bool = False, quantisation_levels: int = 5, ): super().__init__(env) self._combine_service_traffic: bool = combine_service_traffic self._quantisation_levels: int = quantisation_levels self._entries_per_link: int = 1 if not self._combine_service_traffic: self._entries_per_link = self.env.num_services # 1. Define the shape of your observation space component shape = ( [self._quantisation_levels] * self.env.num_links * self._entries_per_link ) # 2. Create Observation space self.space = spaces.MultiDiscrete(shape) # 3. Initialise observation with zeroes self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) def update(self): """TODO: complete description.""" obs = [] for _, link in self.env.links.items(): bandwidth = link.bandwidth if self._combine_service_traffic: loads = [link.get_current_load()] else: loads = [protocol.get_load() for protocol in link.protocol_list] for load in loads: if load <= 0: traffic_level = 0 elif load >= bandwidth: traffic_level = self._quantisation_levels - 1 else: traffic_level = (load / bandwidth) // ( 1 / (self._quantisation_levels - 1) ) + 1 obs.append(int(traffic_level)) self.current_observation[:] = obs class ObservationsHandler: """Component-based observation space handler.""" registry = { "NODE_LINK_TABLE": NodeLinkTable, "NODE_STATUSES": NodeStatuses, "LINK_TRAFFIC_LEVELS": LinkTrafficLevels, } def __init__(self): """TODO: complete description.""" """Initialise the handler without any components yet. They""" self.registered_obs_components: List[AbstractObservationComponent] = [] self.space: spaces.Space self.current_observation: Tuple[np.ndarray] # i can access the registry items like this: # self.registry.LINK_TRAFFIC_LEVELS def update_obs(self): """TODO: complete description.""" current_obs = [] for obs in self.registered_obs_components: obs.update() current_obs.append(obs.current_observation) self.current_observation = tuple(current_obs) def register(self, obs_component: AbstractObservationComponent): """TODO: complete description.""" self.registered_obs_components.append(obs_component) self.update_space() def deregister(self, obs_component: AbstractObservationComponent): """TODO: complete description.""" self.registered_obs_components.remove(obs_component) self.update_space() def update_space(self): """TODO: complete description.""" component_spaces = [] for obs_comp in self.registered_obs_components: component_spaces.append(obs_comp.space) self.space = spaces.Tuple(component_spaces) @classmethod def from_config(cls, env: "Primaite", obs_space_config: dict): """TODO: complete description. This method parses config items related to the observation space, then creates the necessary components and adds them to the observation handler. """ # Instantiate the handler handler = cls() for component_cfg in obs_space_config["components"]: # Figure out which class can instantiate the desired component comp_type = component_cfg["name"] comp_class = cls.registry[comp_type] # Create the component with options from the YAML options = component_cfg.get("options") or {} component = comp_class(env, **options) handler.register(component) handler.update_obs() return handler