generate_environment_files.py 3.22 KB
Newer Older
Paul McCarthy's avatar
Paul McCarthy committed
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
98
99
100
101
102
103
104
#!/usr/bin/env python
#
# Generate FSL conda environment files from fsl-release.yml


import itertools  as it
import os.path    as op
import               os
import               re
import               sys

from manifest_rules.utils import (sprun,
                                  load_release_info,
                                  generate_environment_file_name,
                                  generate_development_version_identifier)


def filter_packages(release_info, platform, cudaver):
    """Return a list of packages that should be included in the environment
    file for the platform and CUDA version.
    """
    packages  = release_info['packages']
    cudapat   = r'cuda-[\d]+\.[\d]+'
    cudapkgs  = [p for p in packages if re.search(cudapat, p) is not None]
    otherpkgs = [p for p in packages if p not in cudapkgs]

    if cudaver is None:
        return otherpkgs

    else:
        include  = f'cuda-{cudaver}'
        cudapkgs = [p for p in cudapkgs if include in p]

        return otherpkgs + cudapkgs


def generate_environment(release_info, version, platform, cudaver, outfile):
    """Genereate an environment file for the platform and CUDA version.
    """
    channels = list(release_info['channels'])
    packages = filter_packages(release_info, platform, cudaver)

    # Dev release
    if version is None:
        channels = [release_info['internal_channel']] + channels

    with open(outfile, 'wt') as f:
        f.write('name: FSL\n')
        f.write('channels:\n')
        for channel in channels:
            f.write(f' - {channel}\n')
        f.write('dependencies:\n')
        for package in packages:
            f.write(f' - {package}\n')


def generate_variants(release_info):
    """Generate a list of (platform, cudaver) pairs for which an environment
    file should be generated. For non-CUDA environments, cudaver will be None.
    """
    cudas     = [None] + list(release_info['cuda'])
    platforms = list(release_info['miniconda'].keys())
    variants  = list(it.product(platforms, cudas))

    # CUDA builds are only supported
    # on linux-64 at this time
    variants = [(plat, cuda) for plat, cuda in variants
                if not (('linux-64' not in plat) and cuda is not None)]

    return variants


def main():
    # Save generated environment files here
    release_file = op.abspath(sys.argv[1])
    outdir       = op.abspath(sys.argv[2])

    # Tags on fsl/conda/manifest denote a
    # public FSL release.  All other
    # commits denote an internal release.
    version      = os.environ.get('CI_COMMIT_TAG', None)
    release_info = load_release_info(release_file)

    # Generate environment files
    # for all platforms+CUDA versions
    variants = generate_variants(release_info)

    if not op.exists(outdir):
        os.mkdir(outdir)

    generated = []

    for platform, cuda in variants:
        filename = generate_environment_file_name(version, platform, cuda)
        filename = op.join(outdir, filename)
        generate_environment(release_info, version, platform, cuda, filename)
        generated.append((platform, cuda, op.basename(filename)))

    for platform, cuda, filename in generated:
        print(f'Platform {platform} [CUDA {cuda}]: {filename}')


if __name__ == '__main__':
    sys.exit(main())