@samitouri / QOSami-colab-cli-2
Code Commits Issues Pulls Sessions CI/CD
main
main
sh 132 lines 5.22 KB
Raw
1 #!/bin/bash
2 # Copyright 2026 Google LLC
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 # http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15
16 # Integration Test: `colab run <script.py> [args...]`
17 #
18 # Verifies the shebang-friendly one-shot execution flow:
19 # 1. `colab run` allocates a CPU VM, runs the script, releases the VM.
20 # 2. `sys.argv` and `__name__ == "__main__"` are honored.
21 # 3. After the run finishes, no orphan VMs remain.
22 # 4. `colab run --keep` leaves the session alive; `colab stop` clears it.
23
24 # Don't `set -e` so we can capture failures and clean up explicitly.
25
26 # ---------- Auth detection (mirrors integration/repro_keep_alive/test.sh) ----
27 if [ -f "$HOME/.config/colab-cli/token.json" ]; then
28 AUTH_FLAGS="--auth=oauth2"
29 elif command -v gcloud > /dev/null && gcloud auth application-default print-access-token > /dev/null 2>&1; then
30 ADC_TOKEN=$(gcloud auth application-default print-access-token 2>/dev/null)
31 ADC_SCOPES=$(curl -s "https://www.googleapis.com/oauth2/v3/tokeninfo?access_token=$ADC_TOKEN" | python3 -c "import json,sys; print(json.load(sys.stdin).get('scope',''))" 2>/dev/null)
32 if echo "$ADC_SCOPES" | grep -q "colaboratory" && echo "$ADC_SCOPES" | grep -q "userinfo.email"; then
33 AUTH_FLAGS="--auth=adc"
34 else
35 echo "Error: ADC token lacks the required scopes (colaboratory + userinfo.email)."
36 echo "Re-issue ADC creds with all required scopes:"
37 echo " gcloud auth application-default login \\"
38 echo " --scopes=openid,\\"
39 echo " https://www.googleapis.com/auth/cloud-platform,\\"
40 echo " https://www.googleapis.com/auth/userinfo.email,\\"
41 echo " https://www.googleapis.com/auth/colaboratory"
42 exit 1
43 fi
44 else
45 echo "Error: No usable auth provider found."
46 exit 1
47 fi
48 echo "[*] Using $AUTH_FLAGS"
49
50 # ---------- Isolated session state -------------------------------------------
51 TMP_DIR=$(mktemp -d)
52 SESSION_FILE="$TMP_DIR/sessions.json"
53 SCRIPT_PATH="$TMP_DIR/script.py"
54 KEEP_SESSION_NAME="repro-run-keep-$(date +%s)"
55
56 cleanup() {
57 echo "[*] Cleaning up..."
58 uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$KEEP_SESSION_NAME" 2>/dev/null || true
59 rm -rf "$TMP_DIR"
60 }
61 trap cleanup EXIT
62
63 cat > "$SCRIPT_PATH" <<'PYEOF'
64 import sys
65 print(f"argv={sys.argv}")
66 print(f"is_main={__name__ == '__main__'}")
67 PYEOF
68 SCRIPT_BASENAME=$(basename "$SCRIPT_PATH")
69
70 # ---------- Phase 1: basic run + auto-cleanup --------------------------------
71 echo "[*] Phase 1: colab run <script.py> hello world"
72 OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" run "$SCRIPT_PATH" hello world 2>&1)
73 RC=$?
74 echo "$OUTPUT"
75
76 if [ $RC -ne 0 ]; then
77 echo "[FAILURE] colab run exited $RC"
78 exit 1
79 fi
80 if ! echo "$OUTPUT" | grep -q "argv=\['$SCRIPT_BASENAME', 'hello', 'world'\]"; then
81 echo "[FAILURE] argv was not propagated as expected."
82 echo " Wanted substring: argv=['$SCRIPT_BASENAME', 'hello', 'world']"
83 exit 1
84 fi
85 if ! echo "$OUTPUT" | grep -q "is_main=True"; then
86 echo "[FAILURE] __name__ was not set to '__main__'."
87 exit 1
88 fi
89
90 # Verify cleanup actually happened — no orphan assignments remain.
91 SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1)
92 echo "$SESSIONS_OUT"
93 if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then
94 echo "[FAILURE] After auto-cleanup, server still reports active sessions."
95 echo " (Possible orphan VM — investigate.)"
96 exit 1
97 fi
98 echo "[SUCCESS] Phase 1 passed: argv passthrough, __main__, and auto-cleanup."
99
100 # ---------- Phase 2: --keep leaves the session alive -------------------------
101 echo ""
102 echo "[*] Phase 2: colab run --keep -s $KEEP_SESSION_NAME <script.py>"
103 OUTPUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" run --keep -s "$KEEP_SESSION_NAME" "$SCRIPT_PATH" keep_arg 2>&1)
104 RC=$?
105 echo "$OUTPUT"
106
107 if [ $RC -ne 0 ]; then
108 echo "[FAILURE] colab run --keep exited $RC"
109 exit 1
110 fi
111 if ! echo "$OUTPUT" | grep -q "argv=\['$SCRIPT_BASENAME', 'keep_arg'\]"; then
112 echo "[FAILURE] --keep run did not produce the expected argv output."
113 exit 1
114 fi
115
116 SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1)
117 echo "$SESSIONS_OUT"
118 if ! echo "$SESSIONS_OUT" | grep -q "\[$KEEP_SESSION_NAME\]"; then
119 echo "[FAILURE] --keep session $KEEP_SESSION_NAME not found in colab sessions."
120 exit 1
121 fi
122
123 uv run colab $AUTH_FLAGS --config "$SESSION_FILE" stop -s "$KEEP_SESSION_NAME"
124 SESSIONS_OUT=$(uv run colab $AUTH_FLAGS --config "$SESSION_FILE" sessions 2>&1)
125 if ! echo "$SESSIONS_OUT" | grep -q "No active sessions found on server."; then
126 echo "[FAILURE] After manual stop of $KEEP_SESSION_NAME, sessions remain."
127 exit 1
128 fi
129
130 echo "[SUCCESS] Phase 2 passed: --keep persists the session, manual stop clears it."
131 echo "[SUCCESS] All phases passed."
132 exit 0