dd-trace-py/.github/workflows/pytorch_gpu_tests.yml at 35c1f5e56dc8fe080ae6fdc9e70198d804538be5 · DataDog/dd-trace-py · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
name: Pytorch Unit Tests (with GPU)

on:
  push:
    branches:
      - 'main'
    paths:
      - ddtrace/internal/datadog/profiling/**
      - ddtrace/profiling/**
      - scripts/download-s3-wheels.sh
      - .github/workflows/pytorch_gpu_tests.yml
  pull_request:
    paths:
      - ddtrace/internal/datadog/profiling/**
      - ddtrace/profiling/**
      - scripts/download-s3-wheels.sh
      - .github/workflows/pytorch_gpu_tests.yml
  workflow_dispatch:

jobs:
   unit-tests:
    runs-on: APM-4-CORE-GPU-LINUX
    steps:
      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
        # Include all history and tags
        with:
          persist-credentials: false
          fetch-depth: 0

      - name: Setup uv
        uses: astral-sh/setup-uv@cec208311dfd045dd5311c1add060b2062131d57 # v8.0.0
        with:
          version: "0.10.4"

      - name: Show system info
        run: uname -a

      - name: Resolve commit SHA
        id: resolve-sha
        run: |
          if [ "${{ github.event_name }}" = "pull_request" ]; then
            echo "sha=${{ github.event.pull_request.head.sha }}" >> "$GITHUB_OUTPUT"
          else
            echo "sha=${{ github.sha }}" >> "$GITHUB_OUTPUT"
          fi

      - name: Download ddtrace wheels from S3
        run: |
          bash scripts/download-s3-wheels.sh \
            "${{ steps.resolve-sha.outputs.sha }}" \
            ./wheels \
            manylinux2014_x86_64

      - name: Run tests
        env:
          DD_PROFILING_ENABLED: "true"
          DD_PROFILING_PYTORCH_ENABLED: "true"
        run: |
          set -euo pipefail

          REPO_ROOT="$(pwd)"
          WHEEL_PATH=$(find "$REPO_ROOT/wheels" -maxdepth 1 -type f -name "ddtrace-*-cp312-cp312-*x86_64*.whl" | sort | head -n 1)
          if [ -z "$WHEEL_PATH" ]; then
            echo "Could not find a cp312 x86_64 wheel in ./wheels" >&2
            ls -la ./wheels >&2
            exit 1
          fi
          echo "Using wheel: $WHEEL_PATH"

          ISOLATED_DIR="$(mktemp -d)"
          TEST_ROOT="$ISOLATED_DIR/tests"

          mkdir -p "$TEST_ROOT/profiling"
          cp "$REPO_ROOT/tests/__init__.py" "$TEST_ROOT/"
          cp "$REPO_ROOT/tests/utils.py" "$TEST_ROOT/"
          cp "$REPO_ROOT/tests/subprocesstest.py" "$TEST_ROOT/"
          cp "$REPO_ROOT/tests/profiling/__init__.py" "$TEST_ROOT/profiling/"
          cp "$REPO_ROOT/tests/profiling/test_pytorch.py" "$TEST_ROOT/profiling/"
          cp "$REPO_ROOT/tests/profiling/simple_program_pytorch_gpu.py" "$TEST_ROOT/profiling/"
          cp "$REPO_ROOT/tests/profiling/simple_program_pytorch_cpu.py" "$TEST_ROOT/profiling/"
          cp -R "$REPO_ROOT/tests/profiling/collector" "$TEST_ROOT/profiling/"

          cd "$ISOLATED_DIR"

          uv run \
            --no-project \
            --with "$WHEEL_PATH" \
            --with protobuf \
            --with pytest \
            --with pytest-cov \
            --with requests \
            --with hypothesis \
            --python 3.12 \
            --with "torch>=1.8.1,<2.7" \
            --with torchvision \
            --with zstandard \
            pytest tests/profiling/test_pytorch.py -vvv --capture=tee-sys