<gh_stars>10-100
from datetime import datetime
import json
import os
import platform
import requests
import signal
from subprocess import Popen, CREATE_NEW_CONSOLE
import sys
import time
from typing import Any, Dict, List, Union
import urllib.parse
import yaml


"""
Handles command line arguments.
TODO use argparse.ArgumentParser

:raises ValueError: if no config yaml path was specified
"""
class ArgParser():
    def __init__(self):
        args: List[str] = sys.argv
        del args[0]

        value: Any = self.get_value(args, '.yaml')
        if value:
            self.config_path: str = str(value)
        else:
            raise ValueError('No config yaml path specified.')

        value = self.get_value(args, 'run-id')
        self.run_id: str = str(value) if value else 'run'

        value = self.get_value(args, 'num-envs')
        self.num_envs: int = int(value) if value else 1

        value = self.get_value(args, 'base-port')
        self.base_port: int = int(value) if value else 5005

        self.env_args: List[str] = args

    """
    Returns argument value for a specified search string.

    :param List[str] args: arguments list
    :param str search: search string
    :return: argument value
    :rtype: Any
    """
    def get_value(self, args: List[str], search: str) -> Any:
        value: Any = None
        l: List[int] = [n for n, s in enumerate(args) if search in s]
        if l:
            arg: str = args[l[0]]
            value = arg.split('=')[1] if '=' in arg else arg
            del args[l[0]]
        return value

    """
    Returns a run ID with # suffix

    :param int n: run count
    :return: run ID
    :rtype: List[str]
    """
    def get_run_id(self, n: int) -> str:
        return f'{self.run_id}-{n}'

    """
    Returns the arguments for an mlagents-learn subprocess.

    :param int n: run count
    :param int i: slot index
    :param str config_path: config path for training run
    :return: arguments list
    :rtype: List[str]
    """
    def get_process_args(self, n: int, i: int, config_path: str) -> List[str]:
        args: List[str] = ['mlagents-learn', config_path, f'--run-id={self.get_run_id(n)}', f'--base-port={self.base_port + i}']
        args.extend(self.env_args)
        return args

    def __str__(self) -> str:
        return f'config_path: {self.config_path}, num_envs: {str(self.num_envs)}, env_args: {", ".join(self.env_args)}'

"""
Stores value options for config param.
"""
class ValueOption():
    """
    :param str key: name of config param
    :param List[Any] values: list of possible values
    """
    def __init__(self, key: str, values: List[Any]):
        self.key: str = key
        self.values: List[Any] = values
        log(f'Found config param option - {self}')

    def __str__(self) -> str:
        return f'{self.key}: {", ".join(map(str, self.values))}'


"""
Stop condition for training runs.
A run can stop prematurely if the latest scalar value for a 
specified metric (tag) is outside of specified min/max limits.
"""
class StopCondition():
    tb_api = 'http://localhost:6006/data/plugin/scalars/scalars?'

    """
    :param Dict[str, Any] params: tag/step/min/max params
    """
    def __init__(self, params: Dict[str, Any]):
        assert 'tag' in params, 'No tag found in stop condition.'
        self.tag: str = params['tag']
        self.min: float = float(params['min'] if 'min' in params else -999999999)
        self.max: float = float(params['max'] if 'max' in params else  999999999)
        self.max = max(self.max, self.min)
        self.step: int = int(params['step'] if 'step' in params else 0)
        log(f'Found stop condition - {self}')

    """
    Calls TensorBoard HTTP API for specified run and checks whether 
    the latest scalar value for {tag} is outside of min/max limits.

    :param str run_id: run id
    :return: true if value is out of min/max limits
    :return: message if value is out of min/max limits 
    :rtype: bool
    """
    def evaluate(self, run_id: str) -> Union[bool, str]:
        args: Dict[str, str] = {'run': run_id, 'tag': self.tag}
        url: str = StopCondition.tb_api + urllib.parse.urlencode(args)
        try:
            r: Response = requests.get(url=url, verify=False, timeout=5)
            if r.status_code == requests.codes.ok:
                data: List[List[float]] = json.loads(r.text)
                step: int = data[-1][1] # Latest step
                value: float = data[-1][2] # Latest scalar
                if step >= self.step:
                    if value < self.min:
                        return True, f'{self.tag}: {value} < {self.min} [step: {step}]'
                    elif value > self.max:
                        return True, f'{self.tag}: {value} > {self.max} [step: {step}]'
            #else:
                # log(r.text)
                # No scalar data yet.
        except:
            log('Could not connect to TensorBoard.')

        return False, None

    def __eq__(self, other):
        return self.tag == other.tag

    def __str__(self) -> str:
        return f'tag: {self.tag}, step: {self.step}, min: {str(self.min)}, max: {str(self.max)}'


"""
Stores behavior specific info.
"""
class Behavior():
    """
    :param str name: behavior name
    :param str run_id: run id
    :param Dict[str, Any] config: behavior config settings
    :param Dict[str, Any] defaults: default settings if available
    """
    def __init__(self, name: str, run_id: str, config: Dict[str, Any], defaults: Dict[str, Any]):
        self.name: str = name
        # Objects generated from opt_values and opt_stop fields
        self.value_options: List[ValueOption] = []
        self.stop_conditions: List[StopCondition] = []

        # Run IDs for each value combination,
        # verbose run IDs contain behavior name: RunID-#/BehaviorName
        # that's how they are listed in TensorBoard
        self.verbose_run_ids: List[str] = []
        # Info strings for each value combination
        self.value_infos: List[List[str]] = []
        # Config settings for each value combination
        self.mod_configs: List[Dict[str, Any]] = []

        if defaults is not None:
            self.copy_defaults(config, defaults)
        parsed: Dict[str, Any] = self.parse_config(config)

        if self.value_options:
            # List of param names we have optional values for
            param_names: List[str] = [x.key for x in self.value_options]
            # Create all possible combinations of optional values 
            value_combos: List[List[Any]] = self.get_value_combinations()
            # Create specific config settings for each value combination
            for i, value_combo in enumerate(value_combos):
                mod_config: Dict[str, Any] = self.insert_values(parsed, param_names, value_combo)
                self.mod_configs.append(mod_config)
                self.verbose_run_ids.append(f'{run_id}-{str(i)}\{name}')
                # Info for value options:
                # RunID-#
                # - Behavior name
                #   - param1 name: param1 value
                #   - param2 name: param2 value
                # ...
                value_info: List[str] = [f'\n{run_id}-{str(i)}\n- {name}\n']
                for j, param_name in enumerate(param_names):
                    value_info.append(f'  - {param_name}: {str(value_combo[j])}\n')
                self.value_infos.append(value_info)
        else:
            # No value options, keep behavior as is
            self.mod_configs.append(parsed)
            self.verbose_run_ids.append(f'{run_id}-0\{name}')
            self.value_infos.append([f'- {name}\n', '  - no value options\n'])

    def __str__(self) -> str:
        return f'run_ids: {", ".join(self.run_ids)}'

    """
    Parses the 'opt_' params in config yaml.
    Returns a config settings copy without those params.
    Generates ValueOption and StopCondition objects.
    
    :param Dict[str, Any] config: behavior config settings
    :param str key: config param name
    :return: updated config settings copy
    :rtype: Dict[str, Any]
    """
    def parse_config(self, config: Dict[str, Any], key: str = None) -> Dict[str, Any]:
        result: Dict[str, Any] = {}
        for k, v in config.items():
            if 'opt_values' in k:
                self.value_options.append(ValueOption(key, v))
            elif 'opt_stop' in k:
                self.stop_conditions.append(StopCondition(v))
            if isinstance(v, dict):
                v = self.parse_config(v, k)
            if 'opt_' not in k:
                result[k] = v
        return result

    """
    Copies default settings into behavior config settings.
    The resulting run config won't contain a default_settings section.
    If there are opt_values set in the original defaults, then they will be 
    applied to each individual behavior. Note that this can blow up the total
    number of runs, because every behavior permutation is combined with every
    other behavior permutation.

    :param Dict[str, Any] config: config settings
    :param Dict[str, Any] defaults: default settings
    :rtype: None
    """
    def copy_defaults(self, config: Dict[str, Any], defaults: Dict[str, Any]) -> None:
        for k, v in defaults.items():
            if k not in config:
                config[k] = v
            if isinstance(v, dict):
                self.copy_defaults(config[k], v)

    """
    Generates option value combinations.

    :param List[List[Any]] result: value combination list
    :param List[Any] tmp: temp. value list
    :param int i: value option index
    :return: list of all value combinations
    :rtype: List[List[Any]]
    """
    def get_value_combinations(self, result: List[List[Any]] = None, tmp: List[Any] = None, i: int = 0) -> List[List[Any]]:
        n: int = len(self.value_options)
        result: List[List[Any]] = [] if result is None else result
        tmp: List[Any] = [None] * n if tmp is None else tmp
        for v in self.value_options[i].values:
            tmp[i] = v # Current value by option index
            if i is n - 1:
                result.append(tmp.copy())
            else:
                self.get_value_combinations(result, tmp, i + 1)
        return result

    """
    Inserts value combination in config settings.
    Returns a config settings copy with modified params.

    :param Dict[str, Any] config: config settings
    :param List[str] param_names: list of param names
    :param List[Any] values: list of possible values
    :return: updated config settings copy
    :rtype: Dict[str, Any]
    """
    def insert_values(self, config: Dict[str, Any], param_names: List[str], values: List[Any]) -> Dict[str, Any]:
        result: Dict[str, Any] = {}
        for k, v in config.items():
            if isinstance(v, dict):
                v = self.insert_values(v, param_names, values)
            result[k] = values[param_names.index(k)] if k in param_names else v
        return result


"""
Loads, parses and saves config settings.
"""
class Config():
    """
    :param ArgParser args: ArgParser instance
    """
    def __init__(self, args: ArgParser):
        file_path: str = args.config_path
        name: str = os.path.basename(file_path).split('.')[0]
        dir: str = os.path.dirname(file_path)
        
        config: Dict[str, Any] = self.load_config(file_path)
        defaults: Dict[str, Any] = config['default_settings'] if 'default_settings' in config else None

        log(f'Parsing {name}...')
        self.behaviors: List[Behavior] = []
        self.stop_conditions: List[StopCondition] = []

        for k, v in config['behaviors'].items():
            b: Behavior = Behavior(k, args.run_id, v, defaults)
            self.behaviors.append(b) 
            # Stop conditions are global in the sense that if they are set for
            # only one of multiple behaviors, the corresponding training run will 
            # stop when they are met - regardless of the other behaviors' metrics.
            for cond in b.stop_conditions:
                if cond not in self.stop_conditions:
                    self.stop_conditions.append(cond)

        # Create all possible run ID combinations
        # Outer list: by combo, inner list: verbose run ids
        self.verbose_run_id_combos: List[List[str]] = self.get_run_id_combinations()
        self.num_runs: int = len(self.verbose_run_id_combos)

        # Gather config settings for each run id combination
        # Outer list: by combo, inner list: by run id, dict: config settings
        run_configs: List[List[Dict[str, Any]]] = []

        config_names: List[str] = []
        config_info: List[str] = []

        for i, run_id_combo in enumerate(self.verbose_run_id_combos):
            config_name: str = name + '-' + str(i)
            config_names.append(config_name)
            # Configs by run id
            configs: List[Dict[str, Any]] = []
            for run_id in run_id_combo:
                # Find matching behavior
                for behavior in self.behaviors:
                    if run_id in behavior.verbose_run_ids:
                        b: int = behavior.verbose_run_ids.index(run_id)
                        configs.append(behavior.mod_configs[b])
                        config_info.extend(behavior.value_infos[b])
                        break
            run_configs.append(configs)

        n: int = len(run_configs)
        assert n is self.num_runs, f'Wrong number of configs {n}, should be {self.num_runs}'
        self.save_info(config_info, dir)

        # Build and save combined config settings
        self.config_paths: List[str] = []
        for i in range(self.num_runs):
            save_config: Dict[str, Any] = {'behaviors': {}}
            for j, b in enumerate(self.behaviors):
                save_config['behaviors'][b.name] = run_configs[i][j]
            file_path: str = self.save_config(save_config, dir, config_names[i])
            self.config_paths.append(file_path)
        
        log(f'{self.num_runs} training runs queued. See config_info.txt for details.')

    """
    Generates run id combinations.

    :param List[List[str]] result: run id combination list
    :param List[str] tmp: temp. run id list
    :param int i: behavior index
    :return: list of all run id combinations
    :rytpe: List[List[str]]
    """
    def get_run_id_combinations(self, result: List[List[str]] = None, tmp: List[str] = None, i: int = 0) -> List[List[str]]:
        n: int = len(self.behaviors)
        result: List[List[str]] = [] if result is None else result
        tmp: List[str] = [None] * n if tmp is None else tmp
        for id in self.behaviors[i].verbose_run_ids:
            tmp[i] = id # Current run id by behavior index
            if i is n - 1:
                result.append(tmp.copy())
            else:
                self.get_run_id_combinations(result, tmp, i + 1)
        return result

    """
    Loads config settings from yaml file.

    :param str path: path to config yaml file
    :return: config settings
    :rytpe: Dict[str, Any]
    """
    def load_config(self, path: str) -> Dict[str, Any]:
        try:
            with open(path) as f:
                config = yaml.load(f, Loader=yaml.FullLoader)
            return config
        except FileNotFoundError:
            log(f'Could not load configuration from {path}.')

    """
    Saves config settings to yaml file.

    :param Dict[str, Any] config: config settings
    :param str dir: file directory
    :param str name: file name
    :return: path to config file
    :rytpe: str
    """
    def save_config(self, config: Dict[str, Any], dir: str, name: str) -> str:
        path: str = os.path.join(dir, name + ".yaml")  
        try:
            with open(path, 'w') as f:
                try:
                    yaml.dump(config, f, sort_keys=False)
                except TypeError:  # Older versions of pyyaml don't support sort_keys
                    yaml.dump(dir, f)
            return path
        except FileNotFoundError:
            log(f'Could not save configuration to {path}.')

    """
    Saves config info to text file.

    :param List[str] info: lines of output text
    :param str dir: file directory
    :rytpe: None
    """
    def save_info(self, info: List[str], dir: str) -> None:
        path: str = os.path.join(dir, "config_info.txt")
        try:
            with open(path, "w") as f:
                for line in info:
                    f.write(line)
        except FileNotFoundError:
            log(f'Could not save configuration info to {path}.')

    def __str__(self):
        return self.name
        

"""
Handles training runs.
"""
class Runner():
    """
    :param ArgParser args: ArgParser instance
    """
    def __init__(self, args: ArgParser):
        self.args: ArgParser = args
        self.config: Config = Config(args)
        # Each slot can store a subprocess
        num_slots: int = args.num_envs
        self.slots: List[Popen] = [None] * num_slots
        # Short run IDs don't contain behavior names, these 
        # are the IDs passed as arguments to the subprocesses
        self.short_run_ids: List[str] = [None] * num_slots
        # Need to store verbose run IDs for the subprocesses too,
        # in order to evaluate stop conditions for each behavior
        self.verbose_run_ids: List[List[str]] = [None] * num_slots
        self.run_count: int = 0
        self.run_controller()

    """
    Starts training runs and checks for stop conditions.

    :param int interval: state check interval
    :rytpe: None
    """
    def run_controller(self, interval: int = 60) -> None:
        interrupt: bool = False

        while self.has_active_runs() or self.has_pending_runs():
            while True:
                i: int = self.get_free_slot()
                if i > -1 and self.has_pending_runs():
                    self.start_process(i)
                else:
                    break
            try:
                while True:
                    time.sleep(interval)
                    exit: bool = False
                    done: List[int] = []
                    for i, slot in enumerate(self.slots):
                        if slot:
                            id: str = self.short_run_ids[i]
                            log(f'Checking {id} progress...')
                            if slot.poll() is not None:
                                done.append(i)
                            elif self.must_stop(i):
                                is_windows: bool = self.stop_process(i)
                                if is_windows:
                                    # Process was terminated, free up slot right away
                                    self.slots[i] = None
                                    exit = True
                                # else: Clean process shut down triggered via signal  
                                # should by caught by next poll call (not tested)

                    for i in done:
                        id: str = self.short_run_ids[i]
                        code: int = self.slots[i].returncode
                        if code is 0:
                            log(f'{id} complete.')
                        else:
                            log(f'An error occurred in {id}: {code}.')
                        self.slots[i].kill() # TODO do we need this?
                        self.slots[i] = None
                        exit = True

                    if exit:
                        break

            except KeyboardInterrupt:
                interrupt = True
                for i, slot in enumerate(self.slots):
                    if slot:
                        self.stop_process(i)
                break

        if interrupt:
            log('Training was interrupted.')
        else:
            log('All training runs complete.')
            
    
    """
    Starts a training run / launches a subprocess.

    :param int i: slot index where the process will be stored
    :rytpe: None
    """
    def start_process(self, i: int) -> None:
        n: int = self.run_count
        self.run_count += 1

        args: List[str] = self.args.get_process_args(n, i, self.config.config_paths[n])
        self.slots[i] = Popen(args, creationflags=CREATE_NEW_CONSOLE)

        run_id: str = self.args.get_run_id(n)
        self.short_run_ids[i] = run_id
        self.verbose_run_ids[i] = self.config.verbose_run_id_combos[n]
        log(f'{run_id} started.')        
            
    """
    Stops specified process.
    TODO Not sure how to get mlagents subprocess to shut down
    gracefully on Windows. mlagents-learn / trainer_controller.py 
    https://github.com/Unity-Technologies/ml-agents/blob/main/ml-agents/mlagents/trainers/trainer_controller.py#L181
    doesn't handle CTRL_BREAK_EVENT, https://stackoverflow.com/a/47314406
    For now, we just terminate the process on Windows, therefore
    it won't be able to exit cleanly and store its latest model.
    This should only be an issue if we were to use stop conditions  
    for catching *successful* training runs though. If we merely 
    stop runs because they don't produce high enough rewards after
    some number of steps for instance, then we probably wouldn't 
    need their models anyway.

    :param int i: process slot index
    :return: true if platform is windows
    :rytpe: bool
    """
    def stop_process(self, i : int):
        if platform.system() is 'Windows':
            self.slots[i].terminate()
            return True
        else:
            # TODO not tested yet, rather use SIGINT?
            self.slots[i].send_signal(signal.CTRL_C_EVENT)
            return False

    """
    Whether any stop condition was met for a specified process.

    :param int i: process slot index
    :return: true if process must stop.
    :rytpe: bool
    """
    def must_stop(self, i: int) -> bool:
        for cond in self.config.stop_conditions:
            for id in self.verbose_run_ids[i]:
                stop, reason = cond.evaluate(id)
                if stop:
                    log(f'Stopping {self.short_run_ids[i]} because {reason}')
                    return True
        return False

    """
    Whether there are any pending runs.

    :return: true if there are any pending runs
    :rytpe: bool
    """
    def has_pending_runs(self) -> bool:
        return self.run_count < self.config.num_runs

    """
    Whether there are any active runs.

    :return: true if there are any active runs
    :rytpe: bool
    """
    def has_active_runs(self) -> bool:
        for slot in self.slots:
            if slot is not None:
                return True
        return False

    """
    Returns next available process slot index or -1.

    :return: next available process slot index or -1
    :rytpe: int
    """
    def get_free_slot(self) -> int:
        for i, slot in enumerate(self.slots):
            if slot is None:
                return  i
        return -1


def log(msg):
    now: datetime = datetime.now()
    current_time: str = now.strftime("%H:%M:%S")
    print(f'[{current_time}] {msg}')

def main():
    runner = Runner(ArgParser())

if __name__ == "__main__":
    main()