Repository: ray-project/ray
File: python/ray/autoscaler/_private/local/config.py (Python)

Code before:
import copy
import os
from typing import Any, Dict, Optional

from ray._common.utils import get_default_ray_temp_dir
from ray.autoscaler._private.cli_logger import cli_logger

unsupported_field_message = "The field {} is not supported for on-premise clusters."

LOCAL_CLUSTER_NODE_TYPE = "local.cluster.node"


def prepare_local(config: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """
    Prepare local cluster config for ingestion by cluster launcher and
    autoscaler.
    """
    config = copy.deepcopy(config)
    for field in "head_node", "worker_nodes", "available_node_types":
        if config.get(field):
            # If the config already contains the internal node type, it's been prepared via ray up already hence return as-is.
            if (
                field == "available_node_types"
                and LOCAL_CLUSTER_NODE_TYPE in config.get(field, {})
            ):
                return None
            err_msg = unsupported_field_message.format(field)
            cli_logger.abort(err_msg)
    # We use a config with a single node type for on-prem clusters.
    # Resources internally detected by Ray are not overridden by the autoscaler
    # (see NodeProvider.do_update)
    config["available_node_types"] = {
        LOCAL_CLUSTER_NODE_TYPE: {"node_config": {}, "resources": {}}
    }
    config["head_node_type"] = LOCAL_CLUSTER_NODE_TYPE
    if "coordinator_address" in config["provider"]:
        config = prepare_coordinator(config)
    else:
        config = prepare_manual(config)
    return config


def prepare_coordinator(config: Dict[str, Any]) -> Dict[str, Any]:
    config = copy.deepcopy(config)
    # User should explicitly set the max number of workers for the coordinator
    # to allocate.
    if "max_workers" not in config:
        cli_logger.abort(
            "The field `max_workers` is required when using an "
            "automatically managed on-premise cluster."
        )

Review comment (suggestion):
Let's tweak the signature to return a tuple of the config and if we edited it or not, idk if I like that we'll just return None because of this policy.

Code after:
import copy
import os
from typing import Any, Dict, Tuple

from ray._common.utils import get_default_ray_temp_dir
from ray.autoscaler._private.cli_logger import cli_logger

unsupported_field_message = "The field {} is not supported for on-premise clusters."

LOCAL_CLUSTER_NODE_TYPE = "local.cluster.node"


def prepare_local(config: Dict[str, Any]) -> Tuple[Dict[str, Any], bool]:
    """
    Prepare local cluster config for ingestion by cluster launcher and
    autoscaler.
    """
    config = copy.deepcopy(config)
    for field in "head_node", "worker_nodes", "available_node_types":
        if config.get(field):
            # If the config already contains the internal node type, it's been prepared via ray up already hence return as-is.
            if (
                field == "available_node_types"
                and LOCAL_CLUSTER_NODE_TYPE in config.get(field, {})
            ):
                return config, False
            err_msg = unsupported_field_message.format(field)
            cli_logger.abort(err_msg)
    # We use a config with a single node type for on-prem clusters.
    # Resources internally detected by Ray are not overridden by the autoscaler
    # (see NodeProvider.do_update)
    config["available_node_types"] = {
        LOCAL_CLUSTER_NODE_TYPE: {"node_config": {}, "resources": {}}
    }
    config["head_node_type"] = LOCAL_CLUSTER_NODE_TYPE
    if "coordinator_address" in config["provider"]:
        config = prepare_coordinator(config)
    else:
        config = prepare_manual(config)
    return config, True


def prepare_coordinator(config: Dict[str, Any]) -> Dict[str, Any]:
    config = copy.deepcopy(config)
    # User should explicitly set the max number of workers for the coordinator
    # to allocate.
    if "max_workers" not in config:
        cli_logger.abort(
            "The field `max_workers` is required when using an "
            "automatically managed on-premise cluster."
        )