from __future__ import annotations
from collections.abc import Mapping
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, Union, Optional, Dict
import signal
import logging
from uuid import uuid4
import aiomultiprocess
from palaestrai.core import EventStateMachine as ESM
from palaestrai.core import BasicState, RuntimeConfig
from palaestrai.core.protocol import (
EnvironmentSetupRequest,
EnvironmentSetupResponse,
ShutdownRequest,
ShutdownResponse,
AttachDescriptor,
AttachServiceResponse,
)
from .environment import Environment
from .loading import instantiate_environment
from palaestrai.util.dynaloader import load_with_params
if TYPE_CHECKING:
import multiprocessing
LOG = logging.getLogger(__name__)
@dataclass(frozen=True)
class ProviderProvenance:
"""Provider provenance owned by the environment conductor."""
provider: str
scenario_name: str
resolution_path: str
input_keys: tuple[str, ...]
ignored_keys: tuple[str, ...]
_INLINE_PROVIDER_KEYS = {
"scenario_provider",
"scenario",
"definition",
"config",
"skip_download",
"silent",
"meta_params",
}
[docs]
@ESM.monitor(is_mdp_worker=True)
class EnvironmentConductor:
"""The environment conductor creates new environment instances.
There could be multiple simulation runs and each would need a
separate environment. The environment conductor controls the
creation of those new environment instances.
Parameters
----------
env_cfg : dict
Dictionary with parameters needed by the environment
seed : uuid4
Random seed for recreation
uid : uuid4
Unique identifier
"""
def __init__(self, env_cfg, seed: int, uid=None):
self._uid: str = uid if uid else "EnvironmentConductor-%s" % uuid4()
self._seed: int = seed
self._environment_configuration = deepcopy(env_cfg)
self._state = BasicState.PRISTINE
self.provider_provenance: Optional[ProviderProvenance] = None
self._environments: Dict[str, Environment] = {}
self._environment_processes: Dict[str, aiomultiprocess.Process] = {}
LOG.debug("%s created.", self)
@property
def uid(self) -> str:
return str(self._uid)
@property
def seed(self) -> int:
return int(self._seed)
def _load_environment(self, configuration=None) -> Environment:
"""Loads the ::`Environment` and necessary dependent classes."""
configuration = (
self._environment_configuration
if configuration is None
else configuration
)
env_name = None
env_params = {}
try:
env_name = configuration["environment"].get("uid", "Unknown")
env_clazz = configuration["environment"]["name"]
raw_params = deepcopy(configuration["environment"]["params"])
nested_params = raw_params.get("params")
provider_params = (
deepcopy(nested_params)
if isinstance(nested_params, Mapping)
else raw_params
)
env_params = raw_params
env_uid = "%s.%s-%s" % (self.uid, env_name, str(uuid4())[-6:])
except KeyError:
LOG.critical(
"%s could not load environment: Configuration not present. "
'Key "environment" is missing in environment configuration. '
"The configuration currently contains: %s",
self,
configuration,
)
raise
provider = provider_params.get("scenario_provider", "file")
ignored_keys = []
if provider not in ("file", "inline"):
LOG.warning(
"Unknown scenario_provider '%s', falling back to file mode. "
"Supported providers: file, inline",
provider,
)
provider = "file"
ignored_keys.append("scenario_provider")
if provider == "inline" and "config" in provider_params:
LOG.warning(
"Ignoring `config` because inline scenario_provider does not "
"use a scenario file."
)
ignored_keys.append("config")
if provider == "inline":
for key in provider_params:
if key not in _INLINE_PROVIDER_KEYS:
LOG.warning(
"Ignoring unsupported inline provider key '%s'.",
key,
)
ignored_keys.append(key)
scenario = provider_params.get("scenario", {})
if isinstance(scenario, Mapping):
scenario_name = scenario.get(
"name", provider_params.get("name", "")
)
else:
scenario_name = provider_params.get("name", "")
resolution_path = "inline"
if provider == "file":
configured_path = provider_params.get("config", "file")
if isinstance(configured_path, (list, tuple)):
configured_path = (
configured_path[0] if configured_path else "file"
)
resolution_path = str(configured_path)
self.provider_provenance = ProviderProvenance(
provider=provider,
scenario_name=str(scenario_name),
resolution_path=resolution_path,
input_keys=tuple(provider_params.keys()),
ignored_keys=tuple(ignored_keys),
)
LOG.debug(
"%s loading Environment '%s' (UID %s).",
self,
env_name,
env_uid,
)
try:
environment = instantiate_environment(
configuration,
uid=env_uid,
seed=self.seed + len(self._environments),
broker_uri=RuntimeConfig().broker_uri,
)
environment._name = env_name
environment.provider_provenance = deepcopy(
self.provider_provenance
)
except ValueError as e:
LOG.critical(
"%s could not load environment '%s': %s. Perhaps a typo in "
"your configuration? %s",
self,
env_name,
e,
configuration["environment"],
)
raise e
if "state_transformer" in configuration:
environment._state_transformer = load_with_params(
configuration["state_transformer"]["name"],
configuration["state_transformer"]["params"],
)
LOG.debug(
"%s loaded %s for %s",
self,
environment._state_transformer,
environment,
)
if "reward" in configuration:
environment.reward = load_with_params(
configuration["reward"]["name"],
configuration["reward"]["params"],
)
LOG.debug(
"%s loaded %s for %s",
self,
environment.reward,
environment,
)
return environment
def _configuration_for_setup(
self, request: EnvironmentSetupRequest
) -> dict:
base = deepcopy(self._environment_configuration)
logical_uid = base["environment"].get("uid")
override = request.environment_configuration_overrides.get(logical_uid)
if override is None:
return base
candidate = deepcopy(override)
if candidate.get("environment", {}).get("uid") != logical_uid:
raise ValueError(
f"environment override UID does not match {logical_uid!r}"
)
return candidate
@ESM.attaches
def _init_environment(self, env: Environment, target_service_uid: str):
"""Initialize a new environment.
Creates a new environment instance with its own UID.
Returns
-------
str
The unique identifier of the new environment
"""
return AttachDescriptor(
target_service_uid=target_service_uid,
instance=env,
)
@ESM.on(AttachServiceResponse)
def _handle_attach_service_response(self, response: AttachServiceResponse):
if not response.success:
LOG.critical(
"%s encountered a fatal error while executing %s: %s. "
"Judgement day is nigh!",
self,
self._environments,
response.error,
)
raise RuntimeError(response.error)
def setup(self):
self._state = BasicState.RUNNING
self.mdp_service = self.uid
LOG.info("%s commencing run: creating better worlds.", self.uid)
@ESM.on(EnvironmentSetupRequest)
def handle_environment_setup_request(self, request):
LOG.debug(
"%s received EnvironmentSetupRequest(experiment_run_id=%s).",
self,
request.experiment_run_id,
)
configuration = self._configuration_for_setup(request)
env = self._load_environment(configuration)
assert env.uid not in self._environments
self._environments[env.uid] = env
LOG.info('Loaded environment "%s", starting subprocess...', env.name)
self._init_environment(env, target_service_uid=request.sender)
return EnvironmentSetupResponse(
sender_environment_conductor=request.receiver,
receiver_simulation_controller=request.sender,
experiment_run_id=request.experiment_run_id,
experiment_run_instance_id=request.experiment_run_instance_id,
experiment_run_phase=request.experiment_run_phase,
environment_id=env.uid,
environment_name=env.name,
environment_type=configuration["environment"]["name"],
environment_parameters=deepcopy(
configuration["environment"].get("params", dict())
),
mode=request.mode,
)
@ESM.on(ShutdownRequest)
def handle_shutdown_request(self, request: ShutdownRequest):
self._state = BasicState.STOPPING
self.stop() # type: ignore[attr-defined]
return ShutdownResponse(
sender=self.uid,
receiver=request.sender,
experiment_run_id=request.experiment_run_id,
experiment_run_instance_id=request.experiment_run_instance_id,
experiment_run_phase=request.experiment_run_phase,
)
@ESM.on(signal.SIGCHLD)
def _handle_child(
self, process: Union[aiomultiprocess.Process, multiprocessing.Process]
):
LOG.debug(
"%s: Environment child process %s endet.", self, process.name
)
if process.exitcode is not None and process.exitcode > 0:
self._state = BasicState.ERROR
LOG.critical(
"Call Arthur Dent! "
"The vogons demolished environment process %s "
"(exited prematurely with rc %s)",
process,
process.exitcode,
)
self.stop( # type: ignore[attr-defined]
RuntimeError(f"Environment process {process.name} died.")
)
def teardown(self):
if self._state != BasicState.ERROR:
self._state = BasicState.FINISHED
self._environment_processes.clear()
LOG.info("%s completed shutdown.", self.uid)
def __str__(self):
return f"{self.__class__.__name__}(id=0x{id(self):x}, uid={self.uid})"