Skip to content

Commit 1928c5d

Browse files
mydatasciencemydatascience
committed
Some initial tests
Signed-off-by: Vladimir Suvorov <suvorovv@google.com>
1 parent 5c35ab8 commit 1928c5d

File tree

3 files changed

+337
-74
lines changed

3 files changed

+337
-74
lines changed

‎.github/workflows/build_upload_internal.yml‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# Copyright 2025 Google LLC
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+
# https://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: Build and Upload Internal
16+
17+
on:
18+
workflow_call:
19+
inputs:
20+
device_type:
21+
required: true
22+
type: string
23+
device_name:
24+
required: true
25+
type: string
26+
cloud_runner:
27+
required: true
28+
type: string
29+
build_mode:
30+
required: true
31+
type: string
32+
base_image:
33+
required: true
34+
type: string
35+
36+
jobs:
37+
build_upload:
38+
runs-on: [self-hosted, ${{ inputs.cloud_runner }}]
39+
steps:
40+
- name: Checkout code
41+
uses: actions/checkout@v4
42+
with:
43+
fetch-depth: 0
44+
45+
- name: Authenticate gcloud
46+
run: |
47+
gcloud auth configure-docker --quiet
48+
gcloud auth configure-docker us-docker.pkg.dev --quiet
49+
50+
- name: Build Docker image
51+
run: |
52+
# Build tunix image based on the specified base image
53+
docker build \
54+
--build-arg BASE_IMAGE="${{ inputs.base_image }}" \
55+
--build-arg DEVICE_TYPE="${{ inputs.device_type }}" \
56+
-t "gcr.io/tpu-prod-env-multipod/tunix_${{ github.run_id }}:${{ inputs.device_type }}" \
57+
-f Dockerfile.tpu \
58+
.
59+
60+
- name: Push Docker image
61+
run: |
62+
docker push "gcr.io/tpu-prod-env-multipod/tunix_${{ github.run_id }}:${{ inputs.device_type }}"

‎.github/workflows/run_tests_internal.yml‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# Copyright 2025 Google LLC
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+
# https://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: Run Tests Internal
16+
17+
on:
18+
workflow_call:
19+
inputs:
20+
device_type:
21+
required: true
22+
type: string
23+
device_name:
24+
required: true
25+
type: string
26+
cloud_runner:
27+
required: true
28+
type: string
29+
image_type:
30+
required: false
31+
type: string
32+
default: ""
33+
pytest_marker:
34+
required: true
35+
type: string
36+
xla_python_client_mem_fraction:
37+
required: false
38+
type: string
39+
default: "0.75"
40+
tf_force_gpu_allow_growth:
41+
required: false
42+
type: string
43+
default: "false"
44+
container_resource_option:
45+
required: false
46+
type: string
47+
default: ""
48+
is_scheduled_run:
49+
required: false
50+
type: boolean
51+
default: false
52+
53+
jobs:
54+
run_tests:
55+
runs-on: [self-hosted, ${{ inputs.cloud_runner }}]
56+
steps:
57+
- name: Checkout code
58+
uses: actions/checkout@v4
59+
with:
60+
fetch-depth: 0
61+
62+
- name: Authenticate gcloud
63+
run: |
64+
gcloud auth configure-docker --quiet
65+
gcloud auth configure-docker us-docker.pkg.dev --quiet
66+
67+
- name: Pull Docker image
68+
run: |
69+
docker pull "gcr.io/tpu-prod-env-multipod/tunix_${{ github.run_id }}:${{ inputs.device_type }}"
70+
71+
- name: Run tests in Docker container
72+
run: |
73+
docker run \
74+
--rm \
75+
${{ inputs.container_resource_option }} \
76+
-e XLA_PYTHON_CLIENT_MEM_FRACTION="${{ inputs.xla_python_client_mem_fraction }}" \
77+
-e TF_FORCE_GPU_ALLOW_GROWTH="${{ inputs.tf_force_gpu_allow_growth }}" \
78+
-e PYTEST_MARKER="${{ inputs.pytest_marker }}" \
79+
-e IS_SCHEDULED_RUN="${{ inputs.is_scheduled_run }}" \
80+
-v "$(pwd):/workspace" \
81+
-w /workspace \
82+
"gcr.io/tpu-prod-env-multipod/tunix_${{ github.run_id }}:${{ inputs.device_type }}" \
83+
bash -c "
84+
# Install tunix in development mode
85+
pip install -e .
86+
pip install pytest pytest-xdist
87+
88+
# Verify TPU availability
89+
python -c \"
90+
import jax
91+
print(f'JAX version: {jax.__version__}')
92+
print(f'JAX devices: {jax.devices()}')
93+
print(f'TPU available: {len(jax.devices()) > 0}')
94+
\"
95+
96+
# Run tests based on marker
97+
if [ \"\$IS_SCHEDULED_RUN\" = \"true\" ]; then
98+
echo 'Running scheduled test suite...'
99+
python -m pytest tests/ -v --tb=short -m \"\$PYTEST_MARKER\" --timeout=600
100+
else
101+
echo 'Running PR test suite...'
102+
python -m pytest tests/ -v --tb=short -m \"\$PYTEST_MARKER\" --timeout=300
103+
fi
104+
105+
# Test tunix imports
106+
python -c \"
107+
import tunix
108+
import tunix.models
109+
import tunix.generate
110+
import tunix.sft
111+
import tunix.distillation
112+
import tunix.rl
113+
print('All tunix modules imported successfully')
114+
\"
115+
"

0 commit comments

Comments
 (0)