Add -e/--env passthrough to run-recipe.py

Fixes #81. Allows passing environment variables (e.g. HF_TOKEN)
through to the container when launching via recipes, mirroring
the existing -e flag in launch-cluster.sh.

Usage: ./run-recipe.sh glm-4.7-flash-awq --solo -e HF_TOKEN=$HF_TOKEN
This commit is contained in:
mariosaladino
2026-03-06 21:50:29 +01:00
parent eb8abcca7f
commit f95beba566
2 changed files with 50 additions and 0 deletions

View File

@@ -818,6 +818,7 @@ Examples:
launch_group.add_argument("-d", "--daemon", action="store_true", help="Run in daemon mode")
launch_group.add_argument("-t", "--container", dest="container_override", help="Override container image from recipe")
launch_group.add_argument("--nccl-debug", choices=["VERSION", "WARN", "INFO", "TRACE"], help="NCCL debug level")
launch_group.add_argument("-e", "--env", action="append", dest="env_vars", default=[], metavar="VAR=VALUE", help="Environment variable to pass to container (e.g. -e HF_TOKEN=xxx). Can be used multiple times.")
# Cluster discovery options
discover_group = parser.add_argument_group("Cluster discovery")
@@ -1118,6 +1119,8 @@ Examples:
cmd_parts.extend(["-n", ",".join(nodes)])
if args.nccl_debug:
cmd_parts.extend(["--nccl-debug", args.nccl_debug])
for env_var in args.env_vars:
cmd_parts.extend(["-e", env_var])
cmd_parts.extend(["\\", "\n --launch-script", "/tmp/tmpXXXXXX.sh"])
print(" ".join(cmd_parts))
print()
@@ -1159,6 +1162,9 @@ Examples:
if args.nccl_debug:
cmd.extend(["--nccl-debug", args.nccl_debug])
for env_var in args.env_vars:
cmd.extend(["-e", env_var])
# Add launch script
cmd.extend(["--launch-script", temp_script])

View File

@@ -728,6 +728,48 @@ test_launch_cmd_no_solo_in_cluster() {
fi
}
# Test: -e / --env passthrough to launch-cluster.sh
test_launch_cmd_env_passthrough() {
log_test "Launch command includes -e env vars"
recipe_name=$(find_solo_recipe)
if [[ -z "$recipe_name" ]]; then
log_skip "No solo-capable recipes found"
return
fi
output=$("$PROJECT_DIR/run-recipe.py" "$recipe_name" --dry-run --solo -e HF_TOKEN=test123 -e MY_VAR=hello 2>&1)
launch_cmd=$(extract_launch_cmd "$output")
if echo "$launch_cmd" | grep -q "\-e HF_TOKEN=test123" && echo "$launch_cmd" | grep -q "\-e MY_VAR=hello"; then
log_pass "Launch command includes -e env vars"
else
log_fail "-e env vars not found in launch command"
log_verbose "Launch cmd: $launch_cmd"
fi
}
# Test: no -e flags when none specified
test_launch_cmd_no_env_by_default() {
log_test "Launch command omits -e when no env vars specified"
recipe_name=$(find_solo_recipe)
if [[ -z "$recipe_name" ]]; then
log_skip "No solo-capable recipes found"
return
fi
output=$("$PROJECT_DIR/run-recipe.py" "$recipe_name" --dry-run --solo 2>&1)
launch_cmd=$(extract_launch_cmd "$output")
if echo "$launch_cmd" | grep -q " -e "; then
log_fail "Unexpected -e flag in launch command"
log_verbose "Launch cmd: $launch_cmd"
else
log_pass "Launch command correctly omits -e when none specified"
fi
}
# ==============================================================================
# README Documentation Verification Tests
# ==============================================================================
@@ -1203,6 +1245,8 @@ main() {
test_launch_cmd_launch_script
test_launch_cmd_container_override
test_launch_cmd_no_solo_in_cluster
test_launch_cmd_env_passthrough
test_launch_cmd_no_env_by_default
echo ""
# README documentation verification tests