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

Fix trying to init obs before building network

Browse Source
This commit is contained in:
Marek Wolan
2023-06-01 17:42:35 +01:00
parent 2b25573378
commit 7041b79d2a
2 changed files with 61 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
src/primaite/environment/observations.py
View File

@@ -1,17 +1,22 @@
"""Module for handling configurable observation spaces in PrimAITE.""" """Module for handling configurable observation spaces in PrimAITE."""
import logging import logging
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from enum import Enum from typing import TYPE_CHECKING, List, Tuple
from typing import List, Tuple
import numpy as np import numpy as np
from gym import spaces from gym import spaces
from primaite.common.enums import FileSystemState, HardwareState, SoftwareState from primaite.common.enums import FileSystemState, HardwareState, SoftwareState
from primaite.environment.primaite_env import Primaite
from primaite.nodes.active_node import ActiveNode from primaite.nodes.active_node import ActiveNode
from primaite.nodes.service_node import ServiceNode 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__) _LOGGER = logging.getLogger(__name__)
@@ -19,9 +24,9 @@ class AbstractObservationComponent(ABC):
"""Represents a part of the PrimAITE observation space.""" """Represents a part of the PrimAITE observation space."""
@abstractmethod @abstractmethod
def __init__(self, env: Primaite): def __init__(self, env: "Primaite"):
_LOGGER.info(f"Initialising {self} observation component") _LOGGER.info(f"Initialising {self} observation component")
self.env: Primaite = env self.env: "Primaite" = env
self.space: spaces.Space self.space: spaces.Space
self.current_observation: np.ndarray # type might be too restrictive? self.current_observation: np.ndarray # type might be too restrictive?
return NotImplemented return NotImplemented
@@ -51,7 +56,7 @@ class NodeLinkTable(AbstractObservationComponent):
For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be For example if the environment has 5 nodes, 7 links, and 3 services, the observation space shape will be
``(12, 7)`` ``(12, 7)``
#todo: clean up description #TODO: clean up description
""" """
@@ -59,7 +64,7 @@ class NodeLinkTable(AbstractObservationComponent):
_MAX_VAL = 1_000_000 _MAX_VAL = 1_000_000
_DATA_TYPE = np.int64 _DATA_TYPE = np.int64
def __init__(self, env: Primaite): def __init__(self, env: "Primaite"):
super().__init__(env) super().__init__(env)
# 1. Define the shape of your observation space component # 1. Define the shape of your observation space component
@@ -76,16 +81,16 @@ class NodeLinkTable(AbstractObservationComponent):
) )
# 3. Initialise Observation with zeroes # 3. Initialise Observation with zeroes
self.current_observation = np.zeroes(observation_shape, dtype=self._DATA_TYPE) self.current_observation = np.zeros(observation_shape, dtype=self._DATA_TYPE)
def update_obs(self): def update(self):
"""Update the observation. """Update the observation.
Update the environment's observation state based on the current status of nodes and links. 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` 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. This function can only be called if the observation space setting is set to BOX.
todo: complete description.. TODO: complete description..
""" """
item_index = 0 item_index = 0
nodes = self.env.nodes nodes = self.env.nodes
@@ -136,7 +141,7 @@ class NodeLinkTable(AbstractObservationComponent):
class NodeStatuses(AbstractObservationComponent): class NodeStatuses(AbstractObservationComponent):
"""todo: complete description. """TODO: complete description.
This will create the observation space with node observations followed by link observations. 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: Each node has 3 elements in the observation space plus 1 per service, more specifically:
@@ -148,7 +153,7 @@ class NodeStatuses(AbstractObservationComponent):
_DATA_TYPE = np.int64 _DATA_TYPE = np.int64
def __init__(self, env: Primaite): def __init__(self, env: "Primaite"):
super().__init__(env) super().__init__(env)
# 1. Define the shape of your observation space component # 1. Define the shape of your observation space component
@@ -166,8 +171,8 @@ class NodeStatuses(AbstractObservationComponent):
# 3. Initialise observation with zeroes # 3. Initialise observation with zeroes
self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)
def update_obs(self): def update(self):
"""todo: complete description. """TODO: complete description.
Update the environment's observation state based on the current status of nodes and links. Update the environment's observation state based on the current status of nodes and links.
@@ -196,7 +201,7 @@ class NodeStatuses(AbstractObservationComponent):
class LinkTrafficLevels(AbstractObservationComponent): class LinkTrafficLevels(AbstractObservationComponent):
"""todo: complete description. """TODO: complete description.
Each link has one element in the observation space, corresponding to the traffic load, Each link has one element in the observation space, corresponding to the traffic load,
it can take the following values: it can take the following values:
@@ -211,7 +216,7 @@ class LinkTrafficLevels(AbstractObservationComponent):
def __init__( def __init__(
self, self,
env: Primaite, env: "Primaite",
combine_service_traffic: bool = False, combine_service_traffic: bool = False,
quantisation_levels: int = 5, quantisation_levels: int = 5,
): ):
@@ -234,8 +239,8 @@ class LinkTrafficLevels(AbstractObservationComponent):
# 3. Initialise observation with zeroes # 3. Initialise observation with zeroes
self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE) self.current_observation = np.zeros(len(shape), dtype=self._DATA_TYPE)
def update_obs(self): def update(self):
"""todo: complete description.""" """TODO: complete description."""
obs = [] obs = []
for _, link in self.env.links.items(): for _, link in self.env.links.items():
bandwidth = link.bandwidth bandwidth = link.bandwidth
@@ -262,15 +267,14 @@ class LinkTrafficLevels(AbstractObservationComponent):
class ObservationsHandler: class ObservationsHandler:
"""Component-based observation space handler.""" """Component-based observation space handler."""
class registry(Enum): registry = {
"""todo: complete description.""" "NODE_LINK_TABLE": NodeLinkTable,
"NODE_STATUSES": NodeStatuses,
NODE_LINK_TABLE: NodeLinkTable "LINK_TRAFFIC_LEVELS": LinkTrafficLevels,
NODE_STATUSES: NodeStatuses }
LINK_TRAFFIC_LEVELS: LinkTrafficLevels
def __init__(self): def __init__(self):
"""todo: complete description.""" """TODO: complete description."""
"""Initialise the handler without any components yet. They""" """Initialise the handler without any components yet. They"""
self.registered_obs_components: List[AbstractObservationComponent] = [] self.registered_obs_components: List[AbstractObservationComponent] = []
self.space: spaces.Space self.space: spaces.Space
@@ -279,33 +283,33 @@ class ObservationsHandler:
# self.registry.LINK_TRAFFIC_LEVELS # self.registry.LINK_TRAFFIC_LEVELS
def update_obs(self): def update_obs(self):
"""todo: complete description.""" """TODO: complete description."""
current_obs = [] current_obs = []
for obs in self.registered_obs_components: for obs in self.registered_obs_components:
obs.update_obs() obs.update()
current_obs.append(obs.current_observation) current_obs.append(obs.current_observation)
self.current_observation = tuple(current_obs) self.current_observation = tuple(current_obs)
def register(self, obs_component: AbstractObservationComponent): def register(self, obs_component: AbstractObservationComponent):
"""todo: complete description.""" """TODO: complete description."""
self.registered_obs_components.append(obs_component) self.registered_obs_components.append(obs_component)
self.update_space() self.update_space()
def deregister(self, obs_component: AbstractObservationComponent): def deregister(self, obs_component: AbstractObservationComponent):
"""todo: complete description.""" """TODO: complete description."""
self.registered_obs_components.remove(obs_component) self.registered_obs_components.remove(obs_component)
self.update_space() self.update_space()
def update_space(self): def update_space(self):
"""todo: complete description.""" """TODO: complete description."""
component_spaces = [] component_spaces = []
for obs_comp in self.registered_obs_components: for obs_comp in self.registered_obs_components:
component_spaces.append(obs_comp.space) component_spaces.append(obs_comp.space)
self.space = spaces.Tuple(component_spaces) self.space = spaces.Tuple(component_spaces)
@classmethod @classmethod
def from_config(cls, obs_space_config): def from_config(cls, env: "Primaite", obs_space_config: dict):
"""todo: complete description. """TODO: complete description.
This method parses config items related to the observation space, then This method parses config items related to the observation space, then
creates the necessary components and adds them to the observation handler. creates the necessary components and adds them to the observation handler.
@@ -316,11 +320,13 @@ class ObservationsHandler:
for component_cfg in obs_space_config["components"]: for component_cfg in obs_space_config["components"]:
# Figure out which class can instantiate the desired component # Figure out which class can instantiate the desired component
comp_type = component_cfg["name"] comp_type = component_cfg["name"]
comp_class = cls.registry[comp_type].value comp_class = cls.registry[comp_type]
# Create the component with options from the YAML # Create the component with options from the YAML
component = comp_class(**component_cfg["options"]) options = component_cfg.get("options") or {}
component = comp_class(env, **options)
handler.register(component) handler.register(component)
handler.update_obs()
return handler return handler

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

@@ -149,7 +149,8 @@ class Primaite(Env):
# The action type # The action type
self.action_type = 0 self.action_type = 0
# todo: proper description here # TODO: proper description here
self.obs_config: dict
self.obs_handler: ObservationsHandler self.obs_handler: ObservationsHandler
# Open the config file and build the environment laydown # Open the config file and build the environment laydown
@@ -161,10 +162,6 @@ class Primaite(Env):
_LOGGER.error("Could not load the environment configuration") _LOGGER.error("Could not load the environment configuration")
_LOGGER.error("Exception occured", exc_info=True) _LOGGER.error("Exception occured", exc_info=True)
# If it doesn't exist after parsing config, create default obs space.
if self.get("obs_handler") is None:
self.configure_obs_space()
# Store the node objects as node attributes # Store the node objects as node attributes
# (This is so we can access them as objects) # (This is so we can access them as objects)
for node in self.network: for node in self.network:
@@ -195,6 +192,10 @@ class Primaite(Env):
_LOGGER.error("Exception occured", exc_info=True) _LOGGER.error("Exception occured", exc_info=True)
print("Could not save network diagram") print("Could not save network diagram")
# # If it doesn't exist after parsing config, create default obs space.
# if getattr(self, "obs_handler", None) is None:
# self.configure_obs_space()
# Initiate observation space # Initiate observation space
self.observation_space, self.env_obs = self.init_observations() self.observation_space, self.env_obs = self.init_observations()
@@ -646,13 +647,22 @@ class Primaite(Env):
pass pass
def init_observations(self) -> Tuple[spaces.Space, np.ndarray]: def init_observations(self) -> Tuple[spaces.Space, np.ndarray]:
"""todo: write docstring.""" """TODO: write docstring."""
if getattr(self, "obs_config", None) is None:
self.obs_config = {
"components": [
{"name": "NODE_LINK_TABLE"},
]
}
self.obs_handler = ObservationsHandler.from_config(self, self.obs_config)
return self.obs_handler.space, self.obs_handler.current_observation return self.obs_handler.space, self.obs_handler.current_observation
def update_environent_obs(self): def update_environent_obs(self):
"""Updates the observation space based on the node and link status. """Updates the observation space based on the node and link status.
todo: better docstring TODO: better docstring
""" """
self.obs_handler.update_obs() self.obs_handler.update_obs()
self.env_obs = self.obs_handler.current_observation self.env_obs = self.obs_handler.current_observation
@@ -692,7 +702,7 @@ class Primaite(Env):
self.get_action_info(item) self.get_action_info(item)
elif item["itemType"] == "OBSERVATION_SPACE": elif item["itemType"] == "OBSERVATION_SPACE":
# Get the observation information # Get the observation information
self.configure_obs_space(item) self.save_obs_config(item)
elif item["itemType"] == "STEPS": elif item["itemType"] == "STEPS":
# Get the steps information # Get the steps information
self.get_steps_info(item) self.get_steps_info(item)
@@ -1025,16 +1035,9 @@ class Primaite(Env):
""" """
self.action_type = ActionType[action_info["type"]] self.action_type = ActionType[action_info["type"]]
def configure_obs_space(self, observation_config: Optional[Dict] = None): def save_obs_config(self, obs_config: Optional[Dict] = None):
"""todo: better docstring.""" """TODO: better docstring."""
if observation_config is None: self.obs_config = obs_config
observation_config = {
"components": [
{"name": "NODE_LINK_TABLE"},
]
}
self.obs_handler = ObservationsHandler[observation_config]
def get_steps_info(self, steps_info): def get_steps_info(self, steps_info):
""" """