# Buggy code
def exec_cluster(
    config_file,
    cmd=None,
    docker=False,
    screen=False,
    tmux=False,
    stop=False,
    start=False,
    override_cluster_name=None,
    port_forward=None,
    with_output=False,
):
    """Runs a command on the specified cluster.

    Arguments:
        config_file: path to the cluster yaml
        cmd: command to run
        docker: whether to run command in docker container of config
        screen: whether to run in a screen
        tmux: whether to run in a tmux session
        stop: whether to stop the cluster after command run
        start: whether to start the cluster if it isn't up
        override_cluster_name: set the name of the cluster
        port_forward (int or list[int]): port(s) to forward
    """
    assert not (screen and tmux), "Can specify only one of `screen` or `tmux`."

    config = yaml.safe_load(open(config_file).read())
    if override_cluster_name is not None:
        config["cluster_name"] = override_cluster_name
    config = _bootstrap_config(config)

    head_node = _get_head_node(
        config, config_file, override_cluster_name, create_if_needed=start
    )

    provider = get_node_provider(config["provider"], config["cluster_name"])
    try:
        updater = NodeUpdaterThread(
            node_id=head_node,
            provider_config=config["provider"],
            provider=provider,
            auth_config=config["auth"],
            cluster_name=config["cluster_name"],
            file_mounts=config["file_mounts"],
            initialization_commands=[],
            setup_commands=[],
            ray_start_commands=[],
            runtime_hash="",
        )

        def wrap_docker(command):
            container_name = config["docker"]["container_name"]
            if not container_name:
                raise ValueError("Docker container not specified in config.")
            return with_docker_exec([command], container_name=container_name)[0]

        if cmd:
            cmd = wrap_docker(cmd) if docker else cmd

            if stop:
                shutdown_cmd = (
                    "ray stop; ray teardown ~/ray_bootstrap_config.yaml "
                    "--yes --workers-only"
                )
                if docker:
                    shutdown_cmd = wrap_docker(shutdown_cmd)
                cmd += "; {}; sudo shutdown -h now".format(shutdown_cmd)

        result = _exec(
            updater,
            cmd,
            screen,
            tmux,
            port_forward=port_forward,
            with_output=with_output,
        )

        if tmux or screen:
            attach_command_parts = ["ray attach", config_file]
            if override_cluster_name is not None:
                attach_command_parts.append(
                    "--cluster-name={}".format(override_cluster_name)
                )
            if tmux:
                attach_command_parts.append("--tmux")
            elif screen:
                attach_command_parts.append("--screen")

            attach_command = " ".join(attach_command_parts)
            attach_info = "Use `{}` to check on command status.".format(attach_command)
            logger.info(attach_info)
        return result
    finally:
        provider.cleanup()


# Traceback (AssertionError)
ray up config/example_full.yaml
2020-06-08 01:13:10,526 INFO config.py:143 -- _configure_iam_role: Role not specified for head node, using arn:aws:iam::<redacted>:instance-profile/ray-autoscaler-v1
2020-06-08 01:13:11,089 INFO config.py:194 -- _configure_key_pair: KeyName not specified for nodes, using ray-autoscaler_us-west-2
2020-06-08 01:13:11,365 INFO config.py:235 -- _configure_subnet: SubnetIds not specified for head node, using [('subnet-<redacted>', 'us-west-2b'), ('subnet-<redacted>', 'us-west-2a')]
2020-06-08 01:13:11,366 INFO config.py:241 -- _configure_subnet: SubnetId not specified for workers, using [('subnet-<redacted>', 'us-west-2b'), ('subnet-<redacted>', 'us-west-2a')]
2020-06-08 01:13:11,725 INFO config.py:261 -- _configure_security_group: SecurityGroupIds not specified for head node, using ray-autoscaler-default (sg-<redacted>)
2020-06-08 01:13:11,725 INFO config.py:268 -- _configure_security_group: SecurityGroupIds not specified for workers, using ray-autoscaler-default (sg-<redacted>)
This will restart cluster services [y/N]: y
2020-06-08 01:13:15,214 INFO commands.py:238 -- get_or_create_head_node: Updating files on head node...
2020-06-08 01:13:15,215 INFO updater.py:379 -- NodeUpdater: i-<redacted>: Updating to <redacted>
2020-06-08 01:13:15,216 INFO updater.py:423 -- NodeUpdater: i-<redacted>: Waiting for remote shell...
2020-06-08 01:13:15,422 INFO log_timer.py:22 -- AWSNodeProvider: Set tag ray-node-status=waiting-for-ssh on ['i-<redacted>
Exception in thread Thread-2:
Traceback (most recent call last):
File "/usr/lib/python3.6/threading.py", line 916, in _bootstrap_inner
self.run()
File "/home/richard/improbable/ray/python/ray/autoscaler/updater.py", line 383, in run
self.do_update()
File "/home/richard/improbable/ray/python/ray/autoscaler/updater.py", line 450, in do_update
self.wait_ready(deadline)
File "/home/richard/improbable/ray/python/ray/autoscaler/updater.py", line 444, in wait_ready
assert False, "Unable to connect to node"
AssertionError: Unable to connect to node

2020-06-08 01:11:59,679 ERROR commands.py:304 -- get_or_create_head_node: Updating <redacted> failed
2020-06-08 01:11:59,701 INFO log_timer.py:22 -- AWSNodeProvider: Set tag ray-node-status=update-failed on ['i-<redacted>']  [LogTimer=333ms]

# Fixed code
def exec_cluster(
    config_file,
    cmd=None,
    docker=False,
    screen=False,
    tmux=False,
    stop=False,
    start=False,
    override_cluster_name=None,
    port_forward=None,
    with_output=False,
):
    """Runs a command on the specified cluster.

    Arguments:
        config_file: path to the cluster yaml
        cmd: command to run
        docker: whether to run command in docker container of config
        screen: whether to run in a screen
        tmux: whether to run in a tmux session
        stop: whether to stop the cluster after command run
        start: whether to start the cluster if it isn't up
        override_cluster_name: set the name of the cluster
        port_forward (int or list[int]): port(s) to forward
    """
    assert not (screen and tmux), "Can specify only one of `screen` or `tmux`."

    config = yaml.safe_load(open(config_file).read())
    if override_cluster_name is not None:
        config["cluster_name"] = override_cluster_name
    config = _bootstrap_config(config)

    head_node = _get_head_node(
        config, config_file, override_cluster_name, create_if_needed=start
    )

    provider = get_node_provider(config["provider"], config["cluster_name"])
    try:
        updater = NodeUpdaterThread(
            node_id=head_node,
            provider_config=config["provider"],
            provider=provider,
            auth_config=config["auth"],
            cluster_name=config["cluster_name"],
            file_mounts=config["file_mounts"],
            initialization_commands=[],
            setup_commands=[],
            ray_start_commands=[],
            runtime_hash="",
            docker_config=config["docker"],
        )

        def wrap_docker(command):
            container_name = config["docker"]["container_name"]
            if not container_name:
                raise ValueError("Docker container not specified in config.")
            return with_docker_exec([command], container_name=container_name)[0]

        if cmd:
            cmd = wrap_docker(cmd) if docker else cmd

            if stop:
                shutdown_cmd = (
                    "ray stop; ray teardown ~/ray_bootstrap_config.yaml "
                    "--yes --workers-only"
                )
                if docker:
                    shutdown_cmd = wrap_docker(shutdown_cmd)
                cmd += "; {}; sudo shutdown -h now".format(shutdown_cmd)

        result = _exec(
            updater,
            cmd,
            screen,
            tmux,
            port_forward=port_forward,
            with_output=with_output,
        )

        if tmux or screen:
            attach_command_parts = ["ray attach", config_file]
            if override_cluster_name is not None:
                attach_command_parts.append(
                    "--cluster-name={}".format(override_cluster_name)
                )
            if tmux:
                attach_command_parts.append("--tmux")
            elif screen:
                attach_command_parts.append("--screen")

            attach_command = " ".join(attach_command_parts)
            attach_info = "Use `{}` to check on command status.".format(attach_command)
            logger.info(attach_info)
        return result
    finally:
        provider.cleanup()