#2374 Remove primaite session
This commit is contained in:
Marek Wolan
@@ -1,33 +1,28 @@
|
||||
import pytest
|
||||
import ray
|
||||
import yaml
|
||||
from ray import air, tune
|
||||
from ray.rllib.algorithms.ppo import PPOConfig
|
||||
|
||||
from primaite.config.load import data_manipulation_config_path
|
||||
from primaite.game.game import PrimaiteGame
|
||||
from primaite.session.environment import PrimaiteRayMARLEnv
|
||||
from tests import TEST_ASSETS_ROOT
|
||||
|
||||
MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Slow, reenable later")
|
||||
def test_rllib_multi_agent_compatibility():
|
||||
"""Test that the PrimaiteRayEnv class can be used with a multi agent RLLIB system."""
|
||||
|
||||
with open(data_manipulation_config_path(), "r") as f:
|
||||
with open(MULTI_AGENT_PATH, "r") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
|
||||
game = PrimaiteGame.from_config(cfg)
|
||||
|
||||
ray.shutdown()
|
||||
ray.init()
|
||||
|
||||
env_config = {"game": game}
|
||||
config = (
|
||||
PPOConfig()
|
||||
.environment(env=PrimaiteRayMARLEnv, env_config={"game": game})
|
||||
.environment(env=PrimaiteRayMARLEnv, env_config=cfg)
|
||||
.rollouts(num_rollout_workers=0)
|
||||
.multi_agent(
|
||||
policies={agent.agent_name for agent in game.rl_agents},
|
||||
policies={agent["ref"] for agent in cfg["agents"]},
|
||||
policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,
|
||||
)
|
||||
.training(train_batch_size=128)
|
||||
|
||||
91
tests/e2e_integration_tests/test_environment.py
Normal file
91
tests/e2e_integration_tests/test_environment.py
Normal file
@@ -0,0 +1,91 @@
|
||||
import pydantic
|
||||
import pytest
|
||||
import yaml
|
||||
from gymnasium.core import ObsType
|
||||
from numpy import ndarray
|
||||
|
||||
from primaite.session.environment import PrimaiteGymEnv, PrimaiteRayMARLEnv
|
||||
from primaite.simulator.network.hardware.nodes.host.server import Printer
|
||||
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
|
||||
from tests import TEST_ASSETS_ROOT
|
||||
|
||||
CFG_PATH = TEST_ASSETS_ROOT / "configs/test_primaite_session.yaml"
|
||||
TRAINING_ONLY_PATH = TEST_ASSETS_ROOT / "configs/train_only_primaite_session.yaml"
|
||||
EVAL_ONLY_PATH = TEST_ASSETS_ROOT / "configs/eval_only_primaite_session.yaml"
|
||||
MISCONFIGURED_PATH = TEST_ASSETS_ROOT / "configs/bad_primaite_session.yaml"
|
||||
MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
|
||||
|
||||
|
||||
class TestPrimaiteEnvironment:
|
||||
def test_creating_env(self):
|
||||
"""Check that environment loads correctly from config and it can be reset."""
|
||||
with open(CFG_PATH, "r") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
env = PrimaiteGymEnv(game_config=cfg)
|
||||
|
||||
def env_checks():
|
||||
assert env is not None
|
||||
assert env.game.simulation
|
||||
assert len(env.game.agents) == 3
|
||||
assert len(env.game.rl_agents) == 1
|
||||
|
||||
assert env.game.simulation.network
|
||||
assert len(env.game.simulation.network.nodes) == 12
|
||||
wireless = env.game.simulation.network.get_node_by_hostname("router_2")
|
||||
assert isinstance(wireless, WirelessRouter)
|
||||
printer = env.game.simulation.network.get_node_by_hostname("HP_LaserJet_Pro_4102fdn_printer")
|
||||
assert isinstance(printer, Printer)
|
||||
|
||||
env_checks()
|
||||
env.reset()
|
||||
env_checks()
|
||||
|
||||
def test_step_env(self):
|
||||
"""Make sure you can go all the way through the session without errors."""
|
||||
with open(CFG_PATH, "r") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
env = PrimaiteGymEnv(game_config=cfg)
|
||||
|
||||
assert (num_actions := len(env.agent.action_manager.action_map)) == 54
|
||||
# run every action and make sure there's no crash
|
||||
for act in range(num_actions):
|
||||
env.step(act)
|
||||
# try running action number outside the action map to check that it fails.
|
||||
with pytest.raises(KeyError):
|
||||
env.step(num_actions)
|
||||
|
||||
obs, rew, trunc, term, info = env.step(0)
|
||||
assert isinstance(obs, ndarray)
|
||||
|
||||
def test_multi_agent_env(self):
|
||||
"""Check that we can run a training session with a multi agent system."""
|
||||
with open(MULTI_AGENT_PATH, "r") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
env = PrimaiteRayMARLEnv(env_config=cfg)
|
||||
|
||||
assert set(env._agent_ids) == {"defender1", "defender2"}
|
||||
|
||||
assert len(env.agents) == 2
|
||||
defender1 = env.agents["defender1"]
|
||||
defender2 = env.agents["defender2"]
|
||||
assert (num_actions_1 := len(defender1.action_manager.action_map)) == 54
|
||||
assert (num_actions_2 := len(defender2.action_manager.action_map)) == 38
|
||||
|
||||
# ensure we can run all valid actions without error
|
||||
for act_1 in range(num_actions_1):
|
||||
env.step({"defender1": act_1, "defender2": 0})
|
||||
for act_2 in range(num_actions_2):
|
||||
env.step({"defender1": 0, "defender2": act_2})
|
||||
|
||||
# ensure we get error when taking an invalid action
|
||||
with pytest.raises(KeyError):
|
||||
env.step({"defender1": num_actions_1, "defender2": 0})
|
||||
with pytest.raises(KeyError):
|
||||
env.step({"defender1": 0, "defender2": num_actions_2})
|
||||
|
||||
def test_error_thrown_on_bad_configuration(self):
|
||||
"""Make sure we throw an error when the config is bad."""
|
||||
with open(MISCONFIGURED_PATH, "r") as f:
|
||||
cfg = yaml.safe_load(f)
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
env = PrimaiteGymEnv(game_config=cfg)
|
||||
@@ -1,109 +0,0 @@
|
||||
import pydantic
|
||||
import pytest
|
||||
|
||||
from primaite.simulator.network.hardware.nodes.host.server import Printer
|
||||
from primaite.simulator.network.hardware.nodes.network.wireless_router import WirelessRouter
|
||||
from tests import TEST_ASSETS_ROOT
|
||||
from tests.conftest import TempPrimaiteSession
|
||||
|
||||
CFG_PATH = TEST_ASSETS_ROOT / "configs/test_primaite_session.yaml"
|
||||
TRAINING_ONLY_PATH = TEST_ASSETS_ROOT / "configs/train_only_primaite_session.yaml"
|
||||
EVAL_ONLY_PATH = TEST_ASSETS_ROOT / "configs/eval_only_primaite_session.yaml"
|
||||
MISCONFIGURED_PATH = TEST_ASSETS_ROOT / "configs/bad_primaite_session.yaml"
|
||||
MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
class TestPrimaiteSession:
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
@pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
|
||||
def test_creating_session(self, temp_primaite_session):
|
||||
"""Check that creating a session from config works."""
|
||||
with temp_primaite_session as session:
|
||||
if not isinstance(session, TempPrimaiteSession):
|
||||
raise AssertionError
|
||||
|
||||
assert session is not None
|
||||
assert session.env.game.simulation
|
||||
assert len(session.env.game.agents) == 3
|
||||
assert len(session.env.game.rl_agents) == 1
|
||||
|
||||
assert session.policy
|
||||
assert session.env
|
||||
|
||||
assert session.env.game.simulation.network
|
||||
assert len(session.env.game.simulation.network.nodes) == 12
|
||||
wireless = session.env.game.simulation.network.get_node_by_hostname("router_2")
|
||||
assert isinstance(wireless, WirelessRouter)
|
||||
printer = session.env.game.simulation.network.get_node_by_hostname("HP_LaserJet_Pro_4102fdn_printer")
|
||||
assert isinstance(printer, Printer)
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
@pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
|
||||
def test_start_session(self, temp_primaite_session):
|
||||
"""Make sure you can go all the way through the session without errors."""
|
||||
with temp_primaite_session as session:
|
||||
session: TempPrimaiteSession
|
||||
session.start_session()
|
||||
|
||||
session_path = session.io_manager.session_path
|
||||
assert session_path.exists()
|
||||
print(list(session_path.glob("*")))
|
||||
checkpoint_dir = session_path / "checkpoints" / "sb3_final"
|
||||
assert checkpoint_dir.exists()
|
||||
checkpoint_1 = checkpoint_dir / "sb3_model_640_steps.zip"
|
||||
checkpoint_2 = checkpoint_dir / "sb3_model_1280_steps.zip"
|
||||
checkpoint_3 = checkpoint_dir / "sb3_model_1920_steps.zip"
|
||||
assert checkpoint_1.exists()
|
||||
assert checkpoint_2.exists()
|
||||
assert not checkpoint_3.exists()
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
@pytest.mark.parametrize("temp_primaite_session", [[TRAINING_ONLY_PATH]], indirect=True)
|
||||
def test_training_only_session(self, temp_primaite_session):
|
||||
"""Check that you can run a training-only session."""
|
||||
with temp_primaite_session as session:
|
||||
session: TempPrimaiteSession
|
||||
session.start_session()
|
||||
# TODO: include checks that the model was trained, e.g. that the loss changed and checkpoints were saved?
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
@pytest.mark.parametrize("temp_primaite_session", [[EVAL_ONLY_PATH]], indirect=True)
|
||||
def test_eval_only_session(self, temp_primaite_session):
|
||||
"""Check that you can load a model and run an eval-only session."""
|
||||
with temp_primaite_session as session:
|
||||
session: TempPrimaiteSession
|
||||
session.start_session()
|
||||
# TODO: include checks that the model was loaded and that the eval-only session ran
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
@pytest.mark.skip(reason="Slow, reenable later")
|
||||
@pytest.mark.parametrize("temp_primaite_session", [[MULTI_AGENT_PATH]], indirect=True)
|
||||
def test_multi_agent_session(self, temp_primaite_session):
|
||||
"""Check that we can run a training session with a multi agent system."""
|
||||
with temp_primaite_session as session:
|
||||
session.start_session()
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
def test_error_thrown_on_bad_configuration(self):
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
session = TempPrimaiteSession.from_config(MISCONFIGURED_PATH)
|
||||
|
||||
@pytest.mark.skip(reason="Session is not being maintained and will be removed in the subsequent beta release.")
|
||||
@pytest.mark.skip(
|
||||
reason="Currently software cannot be dynamically created/destroyed during simulation. Therefore, "
|
||||
"reset doesn't implement software restore."
|
||||
)
|
||||
@pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
|
||||
def test_session_sim_reset(self, temp_primaite_session):
|
||||
with temp_primaite_session as session:
|
||||
session: TempPrimaiteSession
|
||||
client_1 = session.game.simulation.network.get_node_by_hostname("client_1")
|
||||
client_1.software_manager.uninstall("DataManipulationBot")
|
||||
|
||||
assert "DataManipulationBot" not in client_1.software_manager.software
|
||||
|
||||
session.game.reset()
|
||||
client_1 = session.game.simulation.network.get_node_by_hostname("client_1")
|
||||
|
||||
assert "DataManipulationBot" in client_1.software_manager.software
|
||||
Reference in New Issue
Block a user