feat(haproxy-load-balancer): add haproxy load balancer

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/common/helpers.py

@@ -236,7 +236,15 @@ def apply_url_override(url: str) -> str:
     else:
         # Local executor - use localhost
         endpoint_uri = cfg.deployment.endpoints[endpoint_type]
-        endpoint_url = f"http://127.0.0.1:{cfg.deployment.port}{endpoint_uri}"
+
+        # Use HAProxy port if multiple_instances is enabled
+        if cfg.deployment.get("multiple_instances", False):
+            proxy_config = cfg.execution.get("proxy", {}).get("config", {})
+            port = proxy_config.get("haproxy_port", 5009)
+        else:
+            port = cfg.deployment.port
+
+        endpoint_url = f"http://127.0.0.1:{port}{endpoint_uri}"
         return endpoint_url
 
 

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/configs/deployment/nim.yaml

@@ -17,6 +17,7 @@ type: nim
 image: ??? # e.g., nvcr.io/nim/meta/llama-3.1-8b-instruct:1.8.6
 served_model_name: ???
 port: 8000
+multiple_instances: false # If true, deploy across multiple nodes with --nodes $num_nodes --ntasks $num_nodes
 
 # NIM containers use default entrypoint - no custom command needed
 # Configuration is done via environment variables in lepton_config

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/configs/deployment/sglang.yaml

@@ -22,6 +22,7 @@ tensor_parallel_size: 8
 data_parallel_size: 1
 extra_args: ""
 env_vars: {} # {name: value} dict
+multiple_instances: false # If true, deploy across multiple nodes with --nodes $num_nodes --ntasks $num_nodes
 
 endpoints:
   chat: /v1/chat/completions

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/configs/deployment/vllm.yaml

@@ -24,6 +24,7 @@ data_parallel_size: 1
 gpu_memory_utilization: 0.95
 extra_args: ""
 env_vars: {} # {name: value} dict
+multiple_instances: false # If true, deploy across multiple nodes with --nodes $num_nodes --ntasks $num_nodes
 
 endpoints:
   chat: /v1/chat/completions

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/configs/execution/slurm/default.yaml

@@ -32,3 +32,11 @@ mounts:
   deployment: {}
   evaluation: {}
   mount_home: true
+
+proxy:
+  type: haproxy
+  image: haproxy:latest
+  config:
+    haproxy_port: 5009
+    health_check_path: /health
+    health_check_status: 200

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/executors/slurm/executor.py

@@ -30,6 +30,7 @@
 from typing import Dict, List, Optional
 
 import yaml
+from jinja2 import Environment, FileSystemLoader
 from omegaconf import DictConfig, OmegaConf
 
 from nemo_evaluator_launcher.common.execdb import (
@@ -40,11 +41,8 @@
 )
 from nemo_evaluator_launcher.common.helpers import (
     get_api_key_name,
-    get_endpoint_url,
     get_eval_factory_command,
-    get_eval_factory_config,
     get_eval_factory_dataset_size_from_run_config,
-    get_health_url,
     get_timestamp_string,
 )
 from nemo_evaluator_launcher.common.mapping import (
@@ -123,6 +121,18 @@ def execute_eval(cfg: DictConfig, dry_run: bool = False) -> str:
                     invocation_id=invocation_id,
                     job_id=job_id,
                 )
+
+                # Create HAProxy config file with placeholder IPs only if multiple_instances is true
+                if cfg.deployment.get("multiple_instances", False):
+                    haproxy_config = _generate_haproxy_config_with_placeholders(cfg)
+                    # Save both template and working config
+                    haproxy_template_path = local_task_subdir / "haproxy.cfg.template"
+                    haproxy_config_path = local_task_subdir / "haproxy.cfg"
+                    with open(haproxy_template_path, "w") as f:
+                        f.write(haproxy_config)
+                    with open(haproxy_config_path, "w") as f:
+                        f.write(haproxy_config)
+
                 local_runsub_path = local_task_subdir / "run.sub"
                 remote_runsub_path = remote_task_subdir / "run.sub"
                 with open(local_runsub_path, "w") as f:
@@ -455,15 +465,6 @@ def _create_slurm_sbatch_script(
     tasks_mapping = load_tasks_mapping()
     task_definition = get_task_from_mapping(task.name, tasks_mapping)
 
-    # Create merged config for get_endpoint_url
-    merged_nemo_evaluator_config = get_eval_factory_config(cfg, task)
-    health_url = get_health_url(
-        cfg,
-        get_endpoint_url(
-            cfg, merged_nemo_evaluator_config, task_definition["endpoint_type"]
-        ),
-    )
-
     # TODO(public release): convert to template
     s = "#!/bin/bash\n"
 
@@ -560,36 +561,30 @@ def _create_slurm_sbatch_script(
             deployment_mounts_list.append(f"{source_mnt}:{target_mnt}")
 
         # add deployment srun command
-        s += "# deployment server\n"
-        s += "srun --mpi pmix --overlap "
-        s += "--container-image {} ".format(cfg.deployment.image)
-        if deployment_mounts_list:
-            s += "--container-mounts {} ".format(",".join(deployment_mounts_list))
-        if not cfg.execution.get("mounts", {}).get("mount_home", True):
-            s += "--no-container-mount-home "
-        s += "--output {} ".format(remote_task_subdir / "logs" / "server-%A.out")
-        deployment_env_var_names = list(
-            cfg.execution.get("env_vars", {}).get("deployment", {})
-        )
-        if cfg.deployment.get("env_vars"):
-            warnings.warn(
-                "cfg.deployment.env_vars will be deprecated in future versions. "
-                "Use cfg.execution.env_vars.deployment instead.",
-                category=DeprecationWarning,
-                stacklevel=2,
-            )
-            deployment_env_var_names.extend(list(cfg.deployment["env_vars"]))
-        if deployment_env_var_names:
-            s += f"--container-env {','.join(deployment_env_var_names)} "
-        s += "{} &\n\n".format(cfg.deployment.command)  # run asynchronously
-        s += (
-            "SERVER_PID=$!  # capture the PID of the server background srun process\n\n"
+        s += _generate_deployment_srun_command(
+            cfg, deployment_mounts_list, remote_task_subdir
         )
 
         # wait for the server to initialize
-        s += _WAIT_FOR_SERVER_HANDLER.format(health_url=health_url)
+        health_path = cfg.deployment.get("health_check_path", "/health")
+        # For multi-instance check all node IPs, for single instance check localhost
+        if cfg.deployment.get("multiple_instances", False):
+            ip_list = '"${NODES_IPS_ARRAY[@]}"'
+        else:
+            ip_list = '"127.0.0.1"'
+        s += _get_wait_for_server_handler(
+            ip_list,
+            cfg.deployment.port,
+            health_path,
+            "server",
+            check_pid=True,
+        )
         s += "\n\n"
 
+        # add HAProxy load balancer only if multiple_instances is true
+        if cfg.deployment.get("multiple_instances", False):
+            s += _get_proxy_server_srun_command(cfg, remote_task_subdir)
+
     # prepare evaluation mounts
     evaluation_mounts_list = [
         "{}:/results".format(remote_task_subdir / "artifacts"),
@@ -615,6 +610,7 @@ def _create_slurm_sbatch_script(
 
     s += "# evaluation client\n"
     s += "srun --mpi pmix --overlap "
+    s += "--nodes 1 --ntasks 1 "  # Client always runs on single node
     s += "--container-image {} ".format(eval_image)
     evaluation_env_var_names = list(
         cfg.execution.get("env_vars", {}).get("evaluation", {})
@@ -632,7 +628,10 @@ def _create_slurm_sbatch_script(
 
     # terminate the server after all evaluation clients finish
     if cfg.deployment.type != "none":
-        s += "kill $SERVER_PID  # terminate the server to finish gracefully\n\n"
+        s += "kill $SERVER_PID  # terminate the server to finish gracefully\n"
+        if cfg.deployment.get("multiple_instances", False):
+            s += "kill $HAPROXY_PID  # terminate HAProxy to finish gracefully\n"
+        s += "\n"
 
     # auto-export
     ae_cfg = cfg.execution.get("auto_export")
@@ -1094,9 +1093,180 @@ def _get_progress(
 """.strip()
 
 
-_WAIT_FOR_SERVER_HANDLER = """
-date
-# wait for the server to initialize
-bash -c 'while [[ "$(curl -s -o /dev/null -w "%{{http_code}}" {health_url})" != "200" ]]; do kill -0 '"$SERVER_PID"' 2>/dev/null || {{ echo "Server process '"$SERVER_PID"' died"; exit 1; }}; sleep 5; done'
+def _generate_haproxy_config_with_placeholders(cfg):
+    """Generate HAProxy configuration with placeholder IPs using Jinja template."""
+    # Set up Jinja environment
+    template_dir = Path(__file__).parent
+    template_path = template_dir / "haproxy.cfg.template"
+
+    if not template_path.exists():
+        raise FileNotFoundError(f"HAProxy template not found: {template_path}")
+
+    env = Environment(loader=FileSystemLoader(template_dir))
+    template = env.get_template("haproxy.cfg.template")
+
+    # Prepare template data with placeholder IPs - use actual number of nodes
+    num_nodes = cfg.execution.num_nodes
+    nodes = []
+    for i in range(num_nodes):
+        nodes.append({"ip": f"{{IP_{i}}}", "port": cfg.deployment.port})
+
+    # Get health check parameters from execution config
+    proxy_config = cfg.execution.get("proxy", {}).get("config", {})
+    health_check_path = proxy_config.get("health_check_path", "/health")
+    health_check_status = proxy_config.get("health_check_status", 200)
+    haproxy_port = proxy_config.get("haproxy_port", 5009)
+
+    # Render template
+    config = template.render(
+        haproxy_port=haproxy_port,
+        health_check_path=health_check_path,
+        health_check_status=health_check_status,
+        nodes=nodes,
+    )
+
+    return config
+
+
+def _generate_haproxy_config(cfg, nodes_ips):
+    """Generate HAProxy configuration using Jinja template."""
+    # Set up Jinja environment
+    template_dir = Path(__file__).parent
+    template_path = template_dir / "haproxy.cfg.template"
+
+    if not template_path.exists():
+        raise FileNotFoundError(f"HAProxy template not found: {template_path}")
+
+    env = Environment(loader=FileSystemLoader(template_dir))
+    template = env.get_template("haproxy.cfg.template")
+
+    # Prepare template data
+    nodes = []
+    for i, ip in enumerate(nodes_ips, 1):
+        nodes.append(
+            {"ip": ip, "port": cfg.deployment.port}  # All nodes use the same port
+        )
+
+    # Get health check parameters from deployment config
+    health_check_path = cfg.deployment.get("health_check_path", "/health")
+    health_check_status = cfg.deployment.get("health_check_status", 200)
+    haproxy_port = cfg.deployment.get("haproxy_port", 5009)
+
+    # Render template
+    config = template.render(
+        haproxy_port=haproxy_port,
+        health_check_path=health_check_path,
+        health_check_status=health_check_status,
+        nodes=nodes,
+    )
+
+    return config
+
+
+def _generate_deployment_srun_command(
+    cfg, deployment_mounts_list, remote_task_subdir, instance_id: int = 0
+):
+    """Generate the deployment srun command with proper node/ntask configuration."""
+    s = ""
+    s += "# deployment server\n"
+    s += "# Get node IPs\n"
+    s += "nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )\n"
+    s += 'nodes_array=("${nodes[@]}")  # Ensure nodes are stored properly\n'
+    s += 'export NODES_IPS_ARRAY=($(for node in "${nodes_array[@]}"; do srun --nodelist=$node --ntasks=1 --nodes=1 hostname --ip-address; done))\n'
+    s += 'echo "Node IPs: ${NODES_IPS_ARRAY[@]}"\n'
+    s += "# Export MASTER_IP as the first node IP\n"
+    s += "export MASTER_IP=${NODES_IPS_ARRAY[0]}\n"
+    s += 'echo "MASTER_IP: $MASTER_IP"\n'
+    s += "srun --mpi pmix --overlap "
+    s += f"--nodes {cfg.execution.num_nodes} --ntasks {cfg.execution.num_nodes} "
+    s += "--container-image {} ".format(cfg.deployment.image)
+    if deployment_mounts_list:
+        s += "--container-mounts {} ".format(",".join(deployment_mounts_list))
+    if not cfg.execution.get("mounts", {}).get("mount_home", True):
+        s += "--no-container-mount-home "
+    s += "--output {} ".format(remote_task_subdir / "logs" / "server-%A-%t.out")
+
+    deployment_env_var_names = list(
+        cfg.execution.get("env_vars", {}).get("deployment", {})
+    )
+    if cfg.deployment.get("env_vars"):
+        warnings.warn(
+            "cfg.deployment.env_vars will be deprecated in future versions. "
+            "Use cfg.execution.env_vars.deployment instead.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        deployment_env_var_names.extend(list(cfg.deployment["env_vars"]))
+
+    # Always add MASTER_IP to the environment variables
+    if "MASTER_IP" not in deployment_env_var_names:
+        deployment_env_var_names.append("MASTER_IP")
+
+    if deployment_env_var_names:
+        s += f"--container-env {','.join(deployment_env_var_names)} "
+    s += "{} &\n\n".format(cfg.deployment.command)  # run asynchronously
+    s += "SERVER_PID=$!  # capture the PID of the server background srun process\n\n"
+
+    return s
+
+
+def _get_wait_for_server_handler(
+    ip_list: str,
+    port: int,
+    health_check_path: str,
+    service_name: str = "server",
+    check_pid: bool = False,
+):
+    """Generate wait for server handler that takes a list of IPs."""
+    pid_check = ""
+    if check_pid:
+        pid_check = 'kill -0 "$SERVER_PID" 2>/dev/null || { echo "Server process $SERVER_PID died"; exit 1; }'
+
+    handler = f"""date
+# wait for the {service_name} to initialize
+for ip in {ip_list}; do
+  echo "Waiting for {service_name} on $ip..."
+  while [[ "$(curl -s -o /dev/null -w "%{{http_code}}" http://$ip:{port}{health_check_path})" != "200" ]]; do
+    {pid_check}
+    sleep 5
+  done
+  echo "{service_name} ready on $ip!"
+done
 date
 """.strip()
+
+    return handler
+
+
+def _get_proxy_server_srun_command(cfg, remote_task_subdir):
+    """Generate HAProxy proxy server srun command using template-based config."""
+    s = ""
+    s += "# HAProxy load balancer\n"
+    s += "# Copy template to config file (important for restarts)\n"
+    s += f"cp {remote_task_subdir}/haproxy.cfg.template {remote_task_subdir}/haproxy.cfg\n"
+    s += "# Replace placeholder IPs with actual node IPs\n"
+    s += f"haproxy_config_file={remote_task_subdir}/haproxy.cfg\n"
+    s += 'for i in "${!NODES_IPS_ARRAY[@]}"; do\n'
+    s += '    ip="${NODES_IPS_ARRAY[$i]}"\n'
+    s += '    sed -i "s/{IP_$i}/$ip/g" "$haproxy_config_file"\n'
+    s += "done\n"
+    s += "\n"
+    s += "srun --mpi pmix --overlap "
+    s += "--nodes 1 --ntasks 1 "
+    s += f"--container-image {cfg.execution.get('proxy', {}).get('image', 'haproxy:latest')} "
+    s += f"--container-mounts {remote_task_subdir}/haproxy.cfg:/usr/local/etc/haproxy/haproxy.cfg:ro "
+    s += f"--output {remote_task_subdir}/logs/haproxy-%A.out "
+    s += "haproxy -f /usr/local/etc/haproxy/haproxy.cfg &\n"
+    s += "HAPROXY_PID=$!  # capture the PID of the HAProxy background srun process\n"
+    s += 'echo "HAProxy started with PID: $HAPROXY_PID"\n\n'
+
+    # Wait for HAProxy to be ready on localhost
+    proxy_config = cfg.execution.get("proxy", {}).get("config", {})
+    haproxy_port = proxy_config.get("haproxy_port", 5009)
+    health_path = proxy_config.get("health_check_path", "/health")
+    s += _get_wait_for_server_handler(
+        "127.0.0.1", haproxy_port, health_path, "HAProxy", check_pid=False
+    )
+    s += "\n"
+
+    return s

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/executors/slurm/haproxy.cfg.template

@@ -0,0 +1,26 @@
+global
+    log stdout format raw local0
+    maxconn 4096
+
+defaults
+    log     global
+    mode    http
+    option  httplog
+    timeout connect 10s
+    timeout client  100000s
+    timeout server  100000s
+
+frontend service_frontend
+    bind *:{{ haproxy_port }}
+    default_backend service_backend
+
+backend service_backend
+    mode http
+    option httpchk GET {{ health_check_path }}
+    http-check expect status {{ health_check_status }}
+    option http-server-close
+    balance leastconn
+{% for node in nodes %}
+    server node{{ loop.index }} {{ node.ip }}:{{ node.port }} check
+{% endfor %}
+