From 54e8d2a9de3320ef320f46e7207b42a5092e642e Mon Sep 17 00:00:00 2001
From: Mathieu Guillame-Bert <gbm@google.com>
Date: Thu, 1 Oct 2026 05:49:26 -0700
Subject: [PATCH] Fix the GitHub CI caching (do not save empty caches, do not
 cancel the runs on main), and fix a flaky timeseries node prediction test.

PiperOrigin-RevId: 991619945
---
 .github/workflows/ci.yml                      | 34 +++++++++++++------
 .../ten_lines/node_prediction_test.py         |  7 ++--
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 53373f2..c89cb0f 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -8,7 +8,10 @@ on:
   workflow_dispatch:
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
+  # Only the runs of a pull request are cancelled by a newer commit. The runs on
+  # main are not (each one has its own group): each commit is tested, and the
+  # runs save the caches restored by the pull request runs.
+  group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.ref || github.run_id }}
   cancel-in-progress: true
 
 permissions:
@@ -45,10 +48,11 @@ jobs:
 
       # Cache of the build container's /root/.cache: Bazel output base and
       # repository cache (compiled C++ dependencies, test results), and
-      # bazelisk. The pip cache is not saved (see below). Restored from the latest run for this Python version,
-      # even if the dependencies changed since: Bazel and pip only reuse what
-      # is still valid, and script/test.sh invalidates the cached test results
-      # when the installed Python packages change.
+      # bazelisk. The pip cache is not saved (see below). Restored from the
+      # latest run for this Python version, even if the dependencies changed
+      # since: Bazel and pip only reuse what is still valid, and script/test.sh
+      # invalidates the cached test results when the installed Python packages
+      # change.
       - name: Restore build cache
         id: restore-cache
         uses: actions/cache/restore@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
@@ -73,21 +77,29 @@ jobs:
       # the runner user to be saved. The pip wheels (~3GB) are not saved: they
       # are quick to download, and the repository cache quota (10GB) must fit
       # the caches of the three Python versions.
+      # Nothing is saved if there is no Bazel cache (e.g. the run was cancelled
+      # before the build started): an empty cache would be restored by the next
+      # run instead of a complete one.
       - name: Prepare build cache for saving
+        id: prepare-cache
         if: always()
         env:
           CACHE_DIR: ${{ runner.temp }}/dgf-cache
         run: |
-          if [ -d "${CACHE_DIR}" ]; then
-            sudo rm -rf "${CACHE_DIR}/pip"
-            sudo chown -R "$(id -u):$(id -g)" "${CACHE_DIR}"
-            sudo du -sh "${CACHE_DIR}"/* "${CACHE_DIR}"
+          if [ ! -d "${CACHE_DIR}/bazel" ]; then
+            echo "No Bazel cache to save."
+            exit 0
           fi
+          sudo rm -rf "${CACHE_DIR}/pip"
+          sudo chown -R "$(id -u):$(id -g)" "${CACHE_DIR}"
+          du -sh "${CACHE_DIR}"/* "${CACHE_DIR}"
+          echo "save=true" >> "${GITHUB_OUTPUT}"
 
-      # Saved even if the build failed, so that the next run does not have to
+      # Saved even if the build failed or was cancelled (Bazel's cache remains
+      # valid if it is interrupted), so that the next run does not have to
       # recompile everything.
       - name: Save build cache
-        if: always()
+        if: ${{ always() && steps.prepare-cache.outputs.save == 'true' }}
         uses: actions/cache/save@55cc8345863c7cc4c66a329aec7e433d2d1c52a9 # v6.1.0
         with:
           path: ${{ runner.temp }}/dgf-cache
diff --git a/dgf/src/learning/ten_lines/node_prediction_test.py b/dgf/src/learning/ten_lines/node_prediction_test.py
index 54dd52a..0c8ebcb 100644
--- a/dgf/src/learning/ten_lines/node_prediction_test.py
+++ b/dgf/src/learning/ten_lines/node_prediction_test.py
@@ -1299,6 +1299,9 @@ def _create_timeseries_graph_and_schema(
   def test_train_and_predict_with_timeseries(self):
     graph, schema = self._create_timeseries_graph_and_schema()
 
+    # The training is not seeded. With a larger learning rate, it sometimes
+    # collapses to predicting the mean label (1.5) for both queries (~5% of the
+    # runs with learning_rate=0.01 and 100 steps, 0/40 with these values).
     model = node_prediction_lib.train_node_model(
         graph=graph,
         valid_graph=graph,
@@ -1306,13 +1309,13 @@ def test_train_and_predict_with_timeseries(self):
         target_nodeset="query",
         target_column="label",
         time_aware=True,
-        num_train_steps=100,
+        num_train_steps=300,
         batch_size=2,
         sampling_width=2,
         num_sampling_hops=1,
         node_embedding_dim=16,
         num_layers=1,
-        learning_rate=0.01,
+        learning_rate=0.005,
         verbose=0,
     )