Skip to content

Commit 7f85849

Browse files
chtruong814claude
andauthored
ci: shard tests to run more in parallel (#2345)
Signed-off-by: Charlie Truong <chtruong@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 9b39c07 commit 7f85849

54 files changed

Lines changed: 1907 additions & 480 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/actions/test-template/action.yml

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -41,19 +41,6 @@ inputs:
4141
description: "Run tests on CPU only"
4242
required: false
4343
default: "false"
44-
azure-client-id:
45-
description: "Azure Client ID"
46-
required: true
47-
azure-tenant-id:
48-
description: "Azure Tenant ID"
49-
required: true
50-
azure-subscription-id:
51-
description: "Azure Subscription ID"
52-
required: true
53-
has-azure-credentials:
54-
description: "Has Azure credentials"
55-
required: false
56-
default: "false"
5744
is_fork_pr:
5845
description: "Whether this is a pull request from a fork"
5946
required: false
@@ -77,31 +64,16 @@ inputs:
7764
runs:
7865
using: "composite"
7966
steps:
80-
- name: Install Azure CLI
81-
if: ${{ inputs.has-azure-credentials == 'true' }}
82-
shell: bash
83-
run: |
84-
for i in 1 2 3; do
85-
curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash && break
86-
echo "Attempt $i failed, retrying in 10s..."
87-
sleep 10
88-
done
89-
9067
- name: Install uuidgen
9168
shell: bash -x -e -u -o pipefail {0}
92-
if: ${{ contains(inputs.runner, 'gcp') }}
69+
if: ${{ contains(inputs.runner, 'aws') || contains(inputs.runner, 'gcp') }}
9370
run: |
9471
for i in 1 2 3; do
9572
apt-get update && apt-get install -y uuid-runtime && break
9673
echo "Attempt $i failed, retrying in 10s..."
9774
sleep 10
9875
done
9976
100-
- name: Docker system cleanup
101-
shell: bash
102-
run: |
103-
docker system prune -af --filter "until=48h" --force || true
104-
10577
- name: Docker pull image
10678
shell: bash
10779
run: |
@@ -138,6 +110,7 @@ runs:
138110
docker run --rm -u root --runtime=nvidia --gpus all \
139111
--shm-size=64g \
140112
--env TRANSFORMERS_OFFLINE=0 \
113+
--env GHA_RUNNER=${{ inputs.runner }} \
141114
--env HYDRA_FULL_ERROR=1 \
142115
--env HF_HOME=/home/TestData/nemo-rl/hf_home \
143116
--env HF_DATASETS_CACHE=/home/TestData/nemo-rl/hf_datasets_cache \
Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
name: Build container
15+
16+
on:
17+
workflow_call:
18+
inputs:
19+
build-ref:
20+
required: false
21+
default: ${{ github.sha }}
22+
description: Ref, branch, or SHA to build.
23+
type: string
24+
image-name:
25+
required: true
26+
description: Name of the image to build and push.
27+
type: string
28+
build-args:
29+
required: false
30+
default: ""
31+
description: Additional Docker build args.
32+
type: string
33+
build-contexts:
34+
required: false
35+
default: ""
36+
description: Additional Docker build contexts.
37+
type: string
38+
dockerfile:
39+
required: true
40+
description: Path to the Dockerfile.
41+
type: string
42+
platform:
43+
required: true
44+
description: Docker build platform.
45+
type: string
46+
runner:
47+
required: true
48+
description: Runner to use for the build.
49+
type: string
50+
registry:
51+
required: true
52+
description: Container registry to push to.
53+
type: string
54+
target:
55+
required: false
56+
default: ""
57+
description: Dockerfile stage to build.
58+
type: string
59+
60+
permissions:
61+
contents: read
62+
pull-requests: read
63+
64+
defaults:
65+
run:
66+
shell: bash -x -e -u -o pipefail {0}
67+
68+
jobs:
69+
build:
70+
runs-on: ${{ inputs.runner }}
71+
env:
72+
REGISTRY: ${{ inputs.registry }}
73+
IMAGE_NAME: ${{ inputs.image-name }}
74+
GH_REF: ${{ github.ref }}
75+
RUN_ID: ${{ github.run_id }}
76+
steps:
77+
- name: Checkout repository
78+
uses: actions/checkout@v6
79+
with:
80+
ref: ${{ inputs.build-ref }}
81+
submodules: recursive
82+
83+
- name: Set up Docker Buildx
84+
uses: docker/setup-buildx-action@v3
85+
86+
- name: Get recently merged PR cache refs
87+
id: recent_pr_cache_refs
88+
uses: actions/github-script@v8
89+
env:
90+
REGISTRY: ${{ inputs.registry }}
91+
IMAGE_NAME: ${{ inputs.image-name }}
92+
with:
93+
script: |
94+
const [owner, repo] = process.env.GITHUB_REPOSITORY.split("/");
95+
const result = await github.graphql(`
96+
query($owner: String!, $repo: String!) {
97+
repository(owner: $owner, name: $repo) {
98+
pullRequests(states: MERGED, first: 100, orderBy: {field: UPDATED_AT, direction: DESC}) {
99+
nodes {
100+
number
101+
}
102+
}
103+
}
104+
}
105+
`, { owner, repo });
106+
107+
const refs = result.repository.pullRequests.nodes
108+
.map(({ number }) => `type=registry,ref=${process.env.REGISTRY}/${process.env.IMAGE_NAME}:${number}-buildcache,mode=max`)
109+
.join("\n");
110+
111+
core.setOutput("cache-from", refs);
112+
core.info(`Found ${result.repository.pullRequests.nodes.length} recently merged PR cache refs.`);
113+
114+
- name: Compute build metadata
115+
id: build_meta
116+
shell: bash
117+
run: |
118+
set -euo pipefail
119+
120+
PR_NUMBER=""
121+
if [[ "$GH_REF" =~ refs/heads/pull-request/([0-9]+) ]]; then
122+
PR_NUMBER="${BASH_REMATCH[1]}"
123+
fi
124+
125+
TAGS=("$REGISTRY/$IMAGE_NAME:$RUN_ID")
126+
if [[ "$GH_REF" == "refs/heads/main" ]]; then
127+
CACHE_KEY="main"
128+
TAGS+=("$REGISTRY/$IMAGE_NAME:main")
129+
elif [[ -n "$PR_NUMBER" ]]; then
130+
CACHE_KEY="$PR_NUMBER"
131+
TAGS+=("$REGISTRY/$IMAGE_NAME:$PR_NUMBER")
132+
else
133+
CACHE_KEY=$(printf '%s' "${GITHUB_REF_NAME:-$RUN_ID}" | tr '/' '-' | tr -cd '[:alnum:]._-')
134+
if [[ -z "$CACHE_KEY" ]]; then
135+
CACHE_KEY="$RUN_ID"
136+
fi
137+
fi
138+
139+
CACHE_FROM=(
140+
"type=registry,ref=$REGISTRY/$IMAGE_NAME:main-buildcache,mode=max"
141+
)
142+
if [[ "$CACHE_KEY" != "main" ]]; then
143+
CACHE_FROM+=("type=registry,ref=$REGISTRY/$IMAGE_NAME:$CACHE_KEY-buildcache,mode=max")
144+
fi
145+
146+
{
147+
echo "tags<<EOF"
148+
printf '%s\n' "${TAGS[@]}"
149+
echo "EOF"
150+
echo "cache-from<<EOF"
151+
printf '%s\n' "${CACHE_FROM[@]}"
152+
echo "EOF"
153+
echo "cache-to=type=registry,ref=$REGISTRY/$IMAGE_NAME:$CACHE_KEY-buildcache,mode=max"
154+
} >> "$GITHUB_OUTPUT"
155+
156+
- name: Build and push
157+
uses: docker/build-push-action@v5
158+
with:
159+
file: ${{ inputs.dockerfile }}
160+
push: true
161+
context: .
162+
platforms: ${{ inputs.platform }}
163+
build-contexts: ${{ inputs.build-contexts }}
164+
build-args: ${{ inputs.build-args }}
165+
cache-from: |
166+
${{ steps.build_meta.outputs.cache-from }}
167+
${{ steps.recent_pr_cache_refs.outputs.cache-from }}
168+
cache-to: ${{ steps.build_meta.outputs.cache-to }}
169+
no-cache: false
170+
tags: |
171+
${{ steps.build_meta.outputs.tags }}
172+
target: ${{ inputs.target }}
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
name: Approve Test Queue
16+
17+
on:
18+
schedule:
19+
- cron: "*/5 * * * *"
20+
workflow_dispatch:
21+
22+
jobs:
23+
approve-test-queue:
24+
if: github.repository == 'NVIDIA-NeMo/RL'
25+
uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_test_approval_queue.yml@v1.3.0
26+
with:
27+
workflow_name: CICD NeMo RL
28+
max_concurrency_internal: ${{ fromJSON(vars.MAX_CONCURRENCY || '3') }}
29+
max_concurrency_external: ${{ fromJSON(vars.MAX_CONCURRENCY_EXTERNAL || '3') }}
30+
secrets:
31+
PAT: ${{ secrets.PAT }}
32+
NVIDIA_MANAGEMENT_ORG_PAT: ${{ secrets.NVIDIA_MANAGEMENT_ORG_PAT }}
33+
SLACK_CI_CHANNEL_WEBHOOK: ${{ secrets.SLACK_GITHUB_CI_WEBHOOK }}
34+
SLACK_TEAM_GROUP_ID: ${{ secrets.SLACK_TEAM_GROUP_ID }}

0 commit comments

Comments
 (0)