# Buggy code
def get_command_runner(
    self,
    log_prefix,
    node_id,
    auth_config,
    cluster_name,
    process_runner,
    use_internal_ip,
):
    """Returns the CommandRunner class used to perform SSH commands.

    Args:
    log_prefix(str): stores "NodeUpdater: {}: ".format(<node_id>). Used
        to print progress in the CommandRunner.
    node_id(str): the node ID.
    auth_config(dict): the authentication configs from the autoscaler
        yaml file.
    cluster_name(str): the name of the cluster.
    process_runner(module): the module to use to run the commands
        in the CommandRunner. E.g., subprocess.
    use_internal_ip(bool): whether the node_id belongs to an internal ip
        or external ip.
    """

    return SSHCommandRunner(
        log_prefix,
        node_id,
        self,
        auth_config,
        cluster_name,
        process_runner,
        use_internal_ip,
    )


# 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 get_command_runner(
    self,
    log_prefix,
    node_id,
    auth_config,
    cluster_name,
    process_runner,
    use_internal_ip,
    docker_config=None,
):
    """Returns the CommandRunner class used to perform SSH commands.

    Args:
    log_prefix(str): stores "NodeUpdater: {}: ".format(<node_id>). Used
        to print progress in the CommandRunner.
    node_id(str): the node ID.
    auth_config(dict): the authentication configs from the autoscaler
        yaml file.
    cluster_name(str): the name of the cluster.
    process_runner(module): the module to use to run the commands
        in the CommandRunner. E.g., subprocess.
    use_internal_ip(bool): whether the node_id belongs to an internal ip
        or external ip.
    docker_config(dict): If set, the docker information of the docker
        container that commands should be run on.
    """
    common_args = {
        "log_prefix": log_prefix,
        "node_id": node_id,
        "provider": self,
        "auth_config": auth_config,
        "cluster_name": cluster_name,
        "process_runner": process_runner,
        "use_internal_ip": use_internal_ip,
    }
    if docker_config and docker_config["container_name"] != "":
        return DockerCommandRunner(docker_config, **common_args)
    else:
        return SSHCommandRunner(**common_args)