From c2fe579cccb11ecd3b2096404267cce66b896153 Mon Sep 17 00:00:00 2001
From: Eugene Rakhmatulin <eugr@eugr.com>
Date: Wed, 25 Mar 2026 23:16:56 -0700
Subject: [PATCH] Enhance .env file handling and validation in scripts

---
 autodiscover.sh   |  8 +++++
 build-and-copy.sh | 82 +++++++++++++++++------------------------------
 launch-cluster.sh |  5 ++-
 3 files changed, 42 insertions(+), 53 deletions(-)

diff --git a/autodiscover.sh b/autodiscover.sh
index 54ee4e0..11d771d 100644
--- a/autodiscover.sh
+++ b/autodiscover.sh
@@ -4,11 +4,19 @@
 # This is called early so that DOTENV_* variables are available to all functions
 load_env_if_exists() {
     local env_file="${CONFIG_FILE:-}"
+    local config_explicit="${CONFIG_FILE_SET:-false}"
     
     # If CONFIG_FILE is not set, check default location
     if [[ -z "$env_file" ]]; then
         local script_dir="$(dirname "$(realpath "${BASH_SOURCE[0]}")")"
         env_file="$script_dir/.env"
+        config_explicit="false"
+    fi
+    
+    # Validate config file exists if explicitly specified
+    if [[ "$config_explicit" == "true" ]] && [[ ! -f "$env_file" ]]; then
+        echo "Error: Config file not found: $env_file"
+        exit 1
     fi
     
     if [[ -f "$env_file" ]]; then
diff --git a/build-and-copy.sh b/build-and-copy.sh
index 1aa3628..f841c6f 100755
--- a/build-and-copy.sh
+++ b/build-and-copy.sh
@@ -286,27 +286,8 @@ usage() {
     exit 1
 }
 
-# Set default CONFIG_FILE
-SCRIPT_DIR="$(dirname "$(realpath "$0")")"
-export CONFIG_FILE="$SCRIPT_DIR/.env"
-
-# Parse --config argument first
-i=1
-while [[ $i -le $# ]]; do
-    arg="${!i}"
-    if [[ "$arg" == "--config" ]]; then
-        next_i=$((i+1))
-        CONFIG_FILE="${!next_i}"
-        export CONFIG_FILE
-        break
-    fi
-    i=$((i+1))
-done
-
-# Source autodiscover.sh to load .env file
-source "$(dirname "$0")/autodiscover.sh"
-
-# Now parse all arguments normally
+# Parse all arguments
+CONFIG_FILE_SET=false
 while [[ "$#" -gt 0 ]]; do
     case $1 in
         -t|--tag) IMAGE_TAG="$2"; shift ;;
@@ -320,34 +301,6 @@ while [[ "$#" -gt 0 ]]; do
                 add_copy_hosts "$1"
                 shift
             done
-
-            if [ "${#COPY_HOSTS[@]}" -eq 0 ]; then
-                # Try to use COPY_HOSTS from .env first
-                if [[ -n "$DOTENV_COPY_HOSTS" ]]; then
-                    echo "Using COPY_HOSTS from .env: $DOTENV_COPY_HOSTS"
-                    IFS=',' read -ra HOSTS_FROM_ENV <<< "$DOTENV_COPY_HOSTS"
-                    COPY_HOSTS=("${HOSTS_FROM_ENV[@]}")
-                else
-                    echo "No hosts specified. Using autodiscovery..."
-                    source "$(dirname "$0")/autodiscover.sh"
-
-                    detect_nodes
-                    if [ $? -ne 0 ]; then
-                        echo "Error: Autodiscovery failed."
-                        exit 1
-                    fi
-
-                    if [ ${#PEER_NODES[@]} -gt 0 ]; then
-                        COPY_HOSTS=("${PEER_NODES[@]}")
-                    fi
-
-                    if [ "${#COPY_HOSTS[@]}" -eq 0 ]; then
-                         echo "Error: Autodiscovery found no other nodes."
-                         exit 1
-                    fi
-                    echo "Autodiscovered hosts: ${COPY_HOSTS[*]}"
-                fi
-            fi
             continue
             ;;
         -j|--build-jobs) BUILD_JOBS="$2"; shift ;;
@@ -380,17 +333,42 @@ while [[ "$#" -gt 0 ]]; do
                 exit 1
             fi
             ;;
-        --config) CONFIG_FILE="$2"; shift ;;
+        --config) CONFIG_FILE="$2"; CONFIG_FILE_SET=true; shift ;;
         -h|--help) usage ;;
         *) echo "Unknown parameter passed: $1"; usage ;;
     esac
     shift
 done
 
-# Set CONFIG_FILE and source autodiscover.sh to load .env
-export CONFIG_FILE
+# Source autodiscover.sh to load .env file
 source "$(dirname "$0")/autodiscover.sh"
 
+# Handle COPY_HOSTS from .env or autodiscovery if not specified via arguments
+if [ "${#COPY_HOSTS[@]}" -eq 0 ]; then
+    if [[ -n "$DOTENV_COPY_HOSTS" ]]; then
+        echo "Using COPY_HOSTS from .env: $DOTENV_COPY_HOSTS"
+        IFS=',' read -ra HOSTS_FROM_ENV <<< "$DOTENV_COPY_HOSTS"
+        COPY_HOSTS=("${HOSTS_FROM_ENV[@]}")
+    else
+        echo "No hosts specified. Using autodiscovery..."
+        detect_nodes
+        if [ $? -ne 0 ]; then
+            echo "Error: Autodiscovery failed."
+            exit 1
+        fi
+
+        if [ ${#PEER_NODES[@]} -gt 0 ]; then
+            COPY_HOSTS=("${PEER_NODES[@]}")
+        fi
+
+        if [ "${#COPY_HOSTS[@]}" -eq 0 ]; then
+              echo "Error: Autodiscovery found no other nodes."
+              exit 1
+        fi
+        echo "Autodiscovered hosts: ${COPY_HOSTS[*]}"
+    fi
+fi
+
 # Validate flag combinations
 if [ -n "$VLLM_PRS" ]; then
     if [ "$EXP_MXFP4" = true ]; then echo "Error: --apply-vllm-pr is incompatible with --exp-mxfp4"; exit 1; fi
diff --git a/launch-cluster.sh b/launch-cluster.sh
index 24b6a42..f7dc1cd 100755
--- a/launch-cluster.sh
+++ b/launch-cluster.sh
@@ -159,9 +159,12 @@ done
 # Set .env file path (use default if not specified)
 if [[ -z "$CONFIG_FILE" ]]; then
     CONFIG_FILE="$SCRIPT_DIR/.env"
+    CONFIG_FILE_SET=false
+else
+    CONFIG_FILE_SET=true
 fi
 
-# Load .env file if exists
+# Load .env file
 if [[ -f "$CONFIG_FILE" ]]; then
     echo "Loading configuration from .env file..."