# Buggy code
def get_agent_class(alg):
    """Returns the class of an known agent given its name."""

    if alg == "PPO":
        from ray.rllib import ppo

        return ppo.PPOAgent
    elif alg == "ES":
        from ray.rllib import es

        return es.ESAgent
    elif alg == "DQN":
        from ray.rllib import dqn

        return dqn.DQNAgent
    elif alg == "APEX":
        from ray.rllib import dqn

        return dqn.ApexAgent
    elif alg == "A3C":
        from ray.rllib import a3c

        return a3c.A3CAgent
    elif alg == "BC":
        from ray.rllib import bc

        return bc.BCAgent
    elif alg == "PG":
        from ray.rllib import pg

        return pg.PGAgent
    elif alg == "script":
        from ray.tune import script_runner

        return script_runner.ScriptRunner
    elif alg == "__fake":
        return _MockAgent
    elif alg == "__sigmoid_fake_data":
        return _SigmoidFakeData
    elif alg == "__parameter_tuning":
        return _ParameterTuningAgent
    else:
        raise Exception(("Unknown algorithm {}.").format(alg))


# Traceback (_frozen_importlib._DeadlockError)
Traceback (most recent call last):
File "/home/ubuntu/ray/python/ray/worker.py", line 1720, in fetch_and_execute_function_to_run
function = pickle.loads(serialized_function)
File "/home/ubuntu/ray/python/ray/rllib/__init__.py", line 17, in <module>
_register_all()
File "/home/ubuntu/ray/python/ray/rllib/__init__.py", line 14, in _register_all
register_trainable(key, get_agent_class(key))
File "/home/ubuntu/ray/python/ray/rllib/agent.py", line 229, in get_agent_class
from ray.rllib import es
File "/home/ubuntu/ray/python/ray/rllib/es/__init__.py", line 1, in <module>
from ray.rllib.es.es import (ESAgent, DEFAULT_CONFIG)
File "/home/ubuntu/ray/python/ray/rllib/es/es.py", line 19, in <module>
from ray.rllib.es import policies
File "<frozen importlib._bootstrap>", line 968, in _find_and_load
File "<frozen importlib._bootstrap>", line 168, in __enter__
File "<frozen importlib._bootstrap>", line 110, in acquire
_frozen_importlib._DeadlockError: deadlock detected by _ModuleLock('ray.rllib.es.policies') at 139937598221224

# Fixed code
def get_agent_class(alg):
    """Returns the class of a known agent given its name."""

    if alg == "PPO":
        from ray.rllib import ppo

        return ppo.PPOAgent
    elif alg == "ES":
        from ray.rllib import es

        return es.ESAgent
    elif alg == "DQN":
        from ray.rllib import dqn

        return dqn.DQNAgent
    elif alg == "APEX":
        from ray.rllib import dqn

        return dqn.ApexAgent
    elif alg == "A3C":
        from ray.rllib import a3c

        return a3c.A3CAgent
    elif alg == "BC":
        from ray.rllib import bc

        return bc.BCAgent
    elif alg == "PG":
        from ray.rllib import pg

        return pg.PGAgent
    elif alg == "script":
        from ray.tune import script_runner

        return script_runner.ScriptRunner
    elif alg == "__fake":
        return _MockAgent
    elif alg == "__sigmoid_fake_data":
        return _SigmoidFakeData
    elif alg == "__parameter_tuning":
        return _ParameterTuningAgent
    else:
        raise Exception(("Unknown algorithm {}.").format(alg))