refactored update-version.sh to handle new branching strategy

ci/release/update-version.sh

@@ -1,16 +1,86 @@
 #!/bin/bash
-# Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
 ###########################
 # ci-imgs Version Updater #
 ###########################
 
 ## Usage
-# bash update-version.sh <new_version>
+# Primary interface:   bash update-version.sh <new_version> [--run-context=main|release]
+# Fallback interface:  [RAPIDS_RUN_CONTEXT=main|release] bash update-version.sh <new_version>
+# CLI arguments take precedence over environment variables
+# Defaults to main when no run-context is specified
 
-set -e
+set -euo pipefail
+
+# Verify we're running from the repository root
+if [[ ! -d ".git" ]]; then
+    echo "Error: This script must be run from the repository root directory."
+    echo "Expected to find: .git/"
+    echo "Current directory: $(pwd)"
+    exit 1
+fi
+
+# Parse command line arguments
+CLI_RUN_CONTEXT=""
+VERSION_ARG=""
+
+for arg in "$@"; do
+    case $arg in
+        --run-context=*)
+            CLI_RUN_CONTEXT="${arg#*=}"
+            shift
+            ;;
+        *)
+            if [[ -z "$VERSION_ARG" ]]; then
+                VERSION_ARG="$arg"
+            fi
+            ;;
+    esac
+done
 
 # Format is YY.MM.PP - no leading 'v' or trailing 'a'
-NEXT_FULL_TAG=$1
+NEXT_FULL_TAG="$VERSION_ARG"
+
+# Determine RUN_CONTEXT with CLI precedence over environment variable, defaulting to main
+if [[ -n "$CLI_RUN_CONTEXT" ]]; then
+    RUN_CONTEXT="$CLI_RUN_CONTEXT"
+    echo "Using run-context from CLI: $RUN_CONTEXT"
+elif [[ -n "${RAPIDS_RUN_CONTEXT}" ]]; then
+    RUN_CONTEXT="$RAPIDS_RUN_CONTEXT"
+    echo "Using run-context from environment: $RUN_CONTEXT"
+else
+    RUN_CONTEXT="main"
+    echo "No run-context provided, defaulting to: $RUN_CONTEXT"
+fi
+
+# Validate RUN_CONTEXT value
+if [[ "${RUN_CONTEXT}" != "main" && "${RUN_CONTEXT}" != "release" ]]; then
+    echo "Error: Invalid run-context value '${RUN_CONTEXT}'"
+    echo "Valid values: main, release"
+    exit 1
+fi
+
+# Validate version argument
+if [[ -z "$NEXT_FULL_TAG" ]]; then
+    echo "Error: Version argument is required"
+    echo "Usage: $0 <new_version> [--run-context=<context>]"
+    echo "   or: [RAPIDS_RUN_CONTEXT=<context>] $0 <new_version>"
+    echo "Note: Defaults to main when run-context is not specified"
+    exit 1
+fi
+
+# Get <major>.<minor> for next version
+NEXT_MAJOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[1]}')
+NEXT_MINOR=$(echo $NEXT_FULL_TAG | awk '{split($0, a, "."); print a[2]}')
+NEXT_SHORT_TAG=${NEXT_MAJOR}.${NEXT_MINOR}
+
+# Log update context
+if [[ "${RUN_CONTEXT}" == "main" ]]; then
+    echo "Preparing development branch update for $NEXT_FULL_TAG (targeting main branch)"
+elif [[ "${RUN_CONTEXT}" == "release" ]]; then
+    echo "Preparing release branch update for $NEXT_FULL_TAG (targeting release/${NEXT_SHORT_TAG} branch)"
+fi
 
 echo "Updating to $NEXT_FULL_TAG"
 echo "${NEXT_FULL_TAG}" > VERSION