pytorch

1
#!/usr/bin/env python3
2

3
"""Generates a matrix for docker releases through github actions
4

5
Will output a condensed version of the matrix. Will include fllowing:
6
    * CUDA version short
7
    * CUDA full version
8
    * CUDNN version short
9
    * Image type either runtime or devel
10
    * Platform linux/arm64,linux/amd64
11

12
"""
13

14
import json
15
from typing import Dict, List
16

17
import generate_binary_build_matrix
18

19

20
DOCKER_IMAGE_TYPES = ["runtime", "devel"]
21

22

23
def generate_docker_matrix() -> Dict[str, List[Dict[str, str]]]:
24
    ret: List[Dict[str, str]] = []
25
    # CUDA amd64 Docker images are available as both runtime and devel while
26
    # CPU arm64 image is only available as runtime.
27
    for cuda, version in generate_binary_build_matrix.CUDA_ARCHES_FULL_VERSION.items():
28
        for image in DOCKER_IMAGE_TYPES:
29
            ret.append(
30
                {
31
                    "cuda": cuda,
32
                    "cuda_full_version": version,
33
                    "cudnn_version": generate_binary_build_matrix.CUDA_ARCHES_CUDNN_VERSION[
34
                        cuda
35
                    ],
36
                    "image_type": image,
37
                    "platform": "linux/amd64",
38
                }
39
            )
40
    ret.append(
41
        {
42
            "cuda": "cpu",
43
            "cuda_full_version": "",
44
            "cudnn_version": "",
45
            "image_type": "runtime",
46
            "platform": "linux/arm64",
47
        }
48
    )
49

50
    return {"include": ret}
51

52

53
if __name__ == "__main__":
54
    build_matrix = generate_docker_matrix()
55
    print(json.dumps(build_matrix))
56