Merge branch 'ko3n1g/ci/fix-ut-notifications' into 'main'

ci: Add notifications for unit tests

See merge request ADLR/megatron-lm!2331

.gitlab/stages/01.test.yml

@@ -46,7 +46,7 @@ test:build_image:
       
         ADDITIONAL_PARAMS=()
 
-        if [[ "$CI_COMMIT_BRANCH" == "ci-rebuild-mcore-nemo-image" ]]; then
+        if [[ "$CI_COMMIT_BRANCH" == "ci-rebuild-mcore-nemo-image" || "$CI_COMMIT_BRANCH" == "main" ]]; then
           ADDITIONAL_PARAMS+=("--pull")
           ADDITIONAL_PARAMS+=("--cache-to type=registry,ref=${IMAGE}-buildcache:main")
         fi

tests/functional_tests/shell_test_utils/notify_unit_tests.sh

@@ -11,7 +11,7 @@ collect_jobs () {
                   -s \
                   --globoff \
                   --header "PRIVATE-TOKEN: $RO_API_TOKEN" \
-                  "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/pipelines/${DOWNSTREAM_PIPELINE_ID}/jobs?page=$PAGE&per_page=$PER_PAGE"
+                  "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/pipelines/${CI_PIPELINE_ID}/jobs?page=$PAGE&per_page=$PER_PAGE"
               )
     # Combine the results
     RESULTS=$(jq -s '.[0] + .[1]' <<< "$RESULTS $RESPONSE")
@@ -36,20 +36,13 @@ CONTEXT="unit-tests-extended"
 
 # Fetch Elastic logs
 set +x
-PIPELINE_JSON=$(curl \
-                  --fail \
-                  --silent \
-                  --header "PRIVATE-TOKEN: ${RO_API_TOKEN}" \
-                  "https://${GITLAB_ENDPOINT}/api/v4/projects/${CI_PROJECT_ID}/pipelines/${CI_PIPELINE_ID}/jobs"
-                ) || ret_code=$?
+UNIT_TESTS_JOBS=$(collect_jobs | jq '[.[] | select(.name | startswith("test:pyt"))]')
 set -x
 if [[ ${ret_code:-0} -ne 0 ]]; then
     echo CI_PIPELINE_ID=$CI_PIPELINE_ID does not exist
     exit 1
 fi
 
-UNIT_TESTS_JOBS=$(echo -E $PIPELINE_JSON | jq '[.[] | select(.name | startswith("test:pyt"))]')
-
 if [[ $UNIT_TESTS_JOBS == null ]]; then
     FAILED_JOBS=$(curl \
                     --fail \