fslinstaller.py 48.2 KB
Newer Older
1
#!/usr/bin/env python
2
3
4
#
# FSL installer script.
#
5
6


7
from __future__ import print_function, division, unicode_literals
8

9
10
11
import functools      as ft
import os.path        as op
import subprocess     as sp
12
import textwrap       as tw
13
14
15
import                   argparse
import                   contextlib
import                   getpass
Paul McCarthy's avatar
Paul McCarthy committed
16
import                   hashlib
17
18
import                   json
import                   logging
19
import                   os
20
import                   platform
21
import                   re
22
import                   readline
23
24
import                   shlex
import                   shutil
25
import                   sys
26
import                   tempfile
27
28
import                   threading
import                   time
29

30
# TODO check py2/3
31
32
33
34
35
36
37
38
try:
    import                   urllib
    import urllib.parse   as urlparse
    import urllib.request as urlrequest
except ImportError:
    import urllib2 as urllib
    import urllib2 as urlparse
    import urllib2 as urlrequest
39

40
41
42
43
44
45
try:                import queue
except ImportError: import Queue as queue


PY2 = sys.version[0] == '2'

46

47
48
49
log = logging.getLogger(__name__)


50
51
52
__absfile__ = op.abspath(__file__)


Paul McCarthy's avatar
Paul McCarthy committed
53
54
55
56
57
58
59
60
61
62
63
64
__version__ = '0.0.0'
"""Installer version number. This is automatically updated in release versions
whenever a new version is released.
"""


DEFAULT_INSTALLATION_DIRECTORY = '/usr/local/fsl'
"""Default FSL installation directory. """


FSL_INSTALLER_MANIFEST = 'http://18.133.213.73/installer/manifest,json'
"""URL to download the FSL installer manifest file from. The installer
65
66
67
68
69
manifest file is a JSON file which contains information about available FSL
versions.

See the Context.download_manifest function, and an example manifest file
in test/data/manifest.json, for more details.
Paul McCarthy's avatar
Paul McCarthy committed
70
71
72
"""


73
74
75
76
77
78
SUPPORTED_CUDAS = ['9.2', '10.2', '11.1']
"""Versions of CUDA that CUDA-capable FSL packages are built for. Used
by Context.identify_cuda. Must be in increasing order.
"""


79
80
81
82
83
84
FIRST_FSL_CONDA_RELEASE = '6.0.6'
"""Oldest conda-based FSL version that can be updated in-place by this
installer script.
"""


Paul McCarthy's avatar
Paul McCarthy committed
85
86
87
88
89
90
class InvalidPassword(Exception):
    """Exception raised by Context.get_admin_password if the user gives an
    incorrect password.
    """


91
92
93
94
95
96
97
98
99
100
101
102
class UnknownVersion(Exception):
    """Exception raised by Context.build if the user has requested a FSL
    version that does not exist.
    """


class BuildNotAvailable(Exception):
    """Exception raised by Context.build if there is no available FSL version
    that matches the target platform and/or requested CUDA version.
    """


103
104
105
106
107
108
class UnsupportedPlatform(Exception):
    """Exception raised by the identify_platform function if FSL is not
    available on this platform.
    """


109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
@ft.total_ordering
class Version(object):
    """Class to represent and compare version strings.  Accepted version
    strings are of the form X.Y.Z, where X, Y, and Z are all integers.
    """
    def __init__(self, verstr):
        major, minor, patch = verstr.split('.')[:3]
        self.verstr         = verstr
        self.major          = int(major)
        self.minor          = int(minor)
        self.patch          = int(patch)

    def __str__(self):
        return self.verstr

    def __eq__(self, other):
        return all((self.major == other.major,
                    self.minor == other.minor,
                    self.patch == other.patch))

    def __lt__(self, other):
        for p1, p2 in zip((self.major,  self.minor,  self.patch),
                          (other.major, other.minor, other.patch)):
            if p1 < p2: return True
            if p1 > p2: return False
        return False


Paul McCarthy's avatar
Paul McCarthy committed
137
138
139
class Context(object):
    """Bag of information and settings created in main, and passed around
    this script.
140
141
142

    Several settings are lazily evaluated on first access, but once evaluated,
    their values are immutable.
Paul McCarthy's avatar
Paul McCarthy committed
143
144
145
    """

    def __init__(self, args):
146
147
148
        """Create the context with the argparse.Namespace object containing
        parsed command-line arguments.
        """
149

150
151
        self.args  = args
        self.shell = op.basename(os.environ['SHELL'])
Paul McCarthy's avatar
Paul McCarthy committed
152
153

        # These attributes are updated on-demand via
154
155
156
157
158
159
160
        # the property accessors defined below, or all
        # all updated via the finalise-settings method.
        self.__platform       = None
        self.__cuda           = None
        self.__manifest       = None
        self.__build          = None
        self.__destdir        = None
Paul McCarthy's avatar
Paul McCarthy committed
161
162
        self.__need_admin     = None
        self.__admin_password = None
163
164
165


    def finalise_settings(self):
166
167
        """Finalise values for all information and settings in the Context.
        """
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
        self.platform
        self.cuda
        self.manifest
        self.build
        self.destdir
        self.need_admin
        self.admin_password


    @property
    def platform(self):
        """The platform we are running on, e.g. "linux-64", "macos-64". """
        if self.__platform is None:
            self.__platform = Context.identify_platform()
        return self.__platform


    @property
    def cuda(self):
        """The available CUDA version, or a CUDA version requested by the user.
        """
        if self.__cuda is not None:
            return self.__cuda
        if self.args.cuda is not None:
            self.__cuda = self.args.cuda
        if self.__cuda is None:
            self.__cuda = Context.identify_cuda()
        return self.__cuda


    @property
    def build(self):
        """Returns a suitable FSL build (a dictionary entry from the FSL
201
202
        installer manifest) for the target platform and requested FSL/CUDA
        versions.
203
204
205
206
207
208
209
210
211
212

        The returned dictionary has the following elements:
          - 'version'      FSL version.
          - 'platform':    Platform identifier (e.g. 'linux-64')
          - 'environment': Environment file to download
          - 'sha256':      Checksum of environment file
          - 'output':      Number of lines of expected output, for reporting
                           progress
          - 'cuda':        X.Y CUDA version, if a CUDA-enabled version of FSL
                           is to be installed.
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
        """

        fslversion = self.args.fslversion

        if fslversion not in self.manifest['versions']:
            raise UnknownVersion(
                'FSL version {} is not available'.format(args.fslversion))

        if fslversion == 'latest':
            fslversion = self.manifest['versions']['latest']

        match = None

        for build in self.manifest['versions'][fslversion]:
            if build['platform']       == self.platform and \
               build.get('cuda', None) == self.cuda:
                match = build
                break
        else:
            raise BuildNotAvailable(
                'Cannot find a version of FSL matching platform '
                '{} and CUDA {}'.format(self.platform, self.cuda))

236
        match['version'] = fslversion
237
        return match
Paul McCarthy's avatar
Paul McCarthy committed
238
239
240
241


    @property
    def destdir(self):
242
243
244
        """Installation directory. If not specified at the command line, the
        user is prompted to enter a directory.
        """
245
246
247
248
249
250
251
252
253

        if self.__destdir is not None:
            return self.__destdir

        # The loop below validates the destination directory
        # both when specified at commmand line or
        # interactively.  In either case, if invalid, the
        # user is re-prompted to enter a new destination.
        destdir = None
254
        if self.args.dest is not None:
255
256
257
258
259
260
261
            response = self.args.dest
        else:
            response = None

        while destdir is None:

            if response is None:
262
263
                printmsg('Where do you want to install FSL?',
                         IMPORTANT, EMPHASIS)
264
265
266
267
                printmsg('Press enter to install to the default location '
                         '[{}]\n'.format(DEFAULT_INSTALLATION_DIRECTORY), INFO)
                response = prompt('FSL installation directory [{}]:'.format(
                    DEFAULT_INSTALLATION_DIRECTORY), QUESTION, EMPHASIS)
268
                response = response.rstrip(op.sep)
269

270
271
                if response == '':
                    response = DEFAULT_INSTALLATION_DIRECTORY
272

273
            response  = op.expanduser(op.expandvars(response))
274
275
276
277
278
279
280
281
282
283
            response  = op.abspath(response)
            parentdir = op.dirname(response)
            if op.exists(parentdir):
                destdir = response
            else:
                printmsg('Destination directory {} does not '
                         'exist!'.format(parentdir), ERROR)
                response = None

        self.__destdir = destdir
Paul McCarthy's avatar
Paul McCarthy committed
284
285
286
287
288
        return self.__destdir


    @property
    def need_admin(self):
289
290
291
        """Returns True if administrator privileges will be needed to install
        FSL.
        """
Paul McCarthy's avatar
Paul McCarthy committed
292
293
        if self.__need_admin is not None:
            return self.__need_admin
294
295
296
        parentdir = op.dirname(self.destdir)
        self.__need_admin = Context.check_need_admin(parentdir)
        return self.__need_admin
Paul McCarthy's avatar
Paul McCarthy committed
297
298
299
300


    @property
    def admin_password(self):
301
302
        """Returns the user's administrator password, prompting them if needed.
        """
Paul McCarthy's avatar
Paul McCarthy committed
303
304
305
306
307
308
309
310
311
312
313
        if self.__admin_password is not None:
            return self.__admin_password
        if self.__need_admin == False:
            return None
        if self.__destdir is None:
            raise RuntimeError('Destination directory has not been set')
        self.__admin_password = Context.get_admin_password()


    @property
    def manifest(self):
314
        """Returns the FSL installer manifest as a dictionary. """
Paul McCarthy's avatar
Paul McCarthy committed
315
        if self.__manifest is None:
316
317
            self.__manifest = Context.download_manifest(self.args.manifest,
                                                        self.args.workdir)
Paul McCarthy's avatar
Paul McCarthy committed
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
        return self.__manifest


    @staticmethod
    def identify_platform():
        """Figures out what platform we are running on. Returns a platform
        identifier string - one of:

          - "linux-64" (Linux, x86_64)
          - "macos-64" (macOS, x86_64)
        """

        platforms = {
            ('linux',  'x86_64') : 'linux-64',
            ('darwin', 'x86_64') : 'linux-64',
333
334
335

            # M1 builds (and possbily ARM for Linux)
            # will be added in the future
Paul McCarthy's avatar
Paul McCarthy committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
            ('darwin', 'arm64')  : 'macos-64',
        }

        system = platform.system().lower()
        cpu    = platform.machine()
        key    = (system, cpu)

        if key not in platforms:
            raise UnsupportedPlatform()

        return platforms[key]


    @staticmethod
    def identify_cuda():
        """Identifies the CUDA version supported on the platform. Returns a
352
353
        string containing the 'X.Y' CUDA version, or None if CUDA is not
        supported.
Paul McCarthy's avatar
Paul McCarthy committed
354
355
        """

356
357
358
        if getattr(Context.identify_cuda, 'no_cuda', False):
            return None

Paul McCarthy's avatar
Paul McCarthy committed
359
        try:
360
            output = Process.check_output('nvidia-smi')
361
        except Exception:
362
            Context.identify_cuda.no_cuda = True
Paul McCarthy's avatar
Paul McCarthy committed
363
364
            return None

365
366
367
368
369
370
371
372
        pat   = r'CUDA Version: (\S+)'
        lines = output.split('\n')
        for line in lines:
            match = re.search(pat, line)
            if match:
                cudaver = match.group(1)
                break
        else:
373
            Context.identify_cuda.no_cuda = True
374
            return None
Paul McCarthy's avatar
Paul McCarthy committed
375
376
377

        # Return the most suitable CUDA
        # version that we have a build for
378
379
380
        for supported in reversed(SUPPORTED_CUDAS):
            if float(cudaver) >= float(supported):
                return supported
Paul McCarthy's avatar
Paul McCarthy committed
381

382
        Context.identify_cuda.no_cuda = True
383
        return None
Paul McCarthy's avatar
Paul McCarthy committed
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399


    @staticmethod
    def check_need_admin(dirname):
        """Returns True if dirname needs administrator privileges to write to,
        False otherwise.
        """
        # TODO os.supports_effective_ids added in python 3.3
        return not os.access(dirname, os.W_OK | os.X_OK)


    @staticmethod
    def get_admin_password():
        """Prompt the user for their administrator password."""

        def validate_admin_password(password):
400
401
402
            proc = sudo_popen(['true'], password)
            proc.communicate()
            return proc.returncode == 0
Paul McCarthy's avatar
Paul McCarthy committed
403

Paul McCarthy's avatar
Paul McCarthy committed
404
405
406
407
408
409
410
        for attempt in range(3):
            if attempt == 0:
                msg = 'Your administrator password is needed to ' \
                      'install FSL: '
            else:
                msg = 'Your administrator password is needed to ' \
                      'install FSL [attempt {} of 3]:'.format(attempt + 1)
411
            printmsg(msg, IMPORTANT, end='')
Paul McCarthy's avatar
Paul McCarthy committed
412
413
414
            password = getpass.getpass('')
            valid    = validate_admin_password(password)

Paul McCarthy's avatar
Paul McCarthy committed
415
            if valid:
416
                printmsg('Password accepted', INFO)
Paul McCarthy's avatar
Paul McCarthy committed
417
418
419
                break
            else:
                printmsg('Incorrect password', WARNING)
Paul McCarthy's avatar
Paul McCarthy committed
420
421
422
423
424
425
426
427

        if not valid:
            raise InvalidPassword()

        return password


    @staticmethod
428
    def download_manifest(url, workdir=None):
Paul McCarthy's avatar
Paul McCarthy committed
429
430
431
432
        """Downloads the installer manifest file, which contains information
        about available FSL vesrions, and the most recent version number of the
        installer (this script).

433
434
435
        The manifest file is a JSON file. Lines beginning
        with a double-forward-slash are ignored. See test/data/manifes.json
        for an example.
Paul McCarthy's avatar
Paul McCarthy committed
436
437
438
439
        """

        log.debug('Downloading FSL installer manifest from %s', url)

440
        with tempdir(workdir):
Paul McCarthy's avatar
Paul McCarthy committed
441
442
            download_file(url, 'manifest.json')
            with open('manifest.json') as f:
443
444
445
446
                lines = f.readlines()

        # Drop comments
        lines = [l for l in lines if not l.lstrip().startswith('//')]
Paul McCarthy's avatar
Paul McCarthy committed
447

448
        return json.loads('\n'.join(lines))
Paul McCarthy's avatar
Paul McCarthy committed
449
450


451
452
453
454
455
456
457
458
459
460
461
462
# List of modifiers which can be used to change how
# a message is printed by the printmsg function.
INFO      = 1
IMPORTANT = 2
QUESTION  = 3
PROMPT    = 4
WARNING   = 5
ERROR     = 6
EMPHASIS  = 7
UNDERLINE = 8
RESET     = 9
ANSICODES = {
463
464
465
466
467
468
469
470
471
    INFO      : '\033[37m',         # Light grey
    IMPORTANT : '\033[92m',         # Green
    QUESTION  : '\033[36m\033[4m',  # Blue+underline
    PROMPT    : '\033[36m\033[1m',  # Bright blue+bold
    WARNING   : '\033[93m',         # Yellow
    ERROR     : '\033[91m',         # Red
    EMPHASIS  : '\033[1m',          # White+bold
    UNDERLINE : '\033[4m',          # Underline
    RESET     : '\033[0m',          # Used internally
472
473
474
}


475
def printmsg(msg='', *msgtypes, **kwargs):
476
477
478
479
480
481
    """Prints msg according to the ANSI codes provided in msgtypes.
    All other keyword arguments are passed through to the print function.
    """
    msgcodes = [ANSICODES[t] for t in msgtypes]
    msgcodes = ''.join(msgcodes)
    print('{}{}{}'.format(msgcodes, msg, ANSICODES[RESET]), **kwargs)
482
    sys.stdout.flush()
483
484
485
486
487
488


def prompt(prompt, *msgtypes, **kwargs):
    """Prompts the user for some input. msgtypes and kwargs are passed
    through to the printmsg function.
    """
489
    printmsg(prompt, *msgtypes, end='', **kwargs)
490
491
492

    if PY2: return raw_input(' ').strip()
    else:   return input(    ' ').strip()
493
494


495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
class Progress(object):
    """Simple progress reporter. Displays one of the following:

       - If both a value and total are provided, a progress bar is shown
       - If only a value is provided, a cumulative count is shown
       - If nothing is provided, a spinner is shown.

    Use as a context manager, and call the update method to report progress,
    e,g:

        with Progress('%') as p:
            for i in range(100):
                p.update(i + 1, 100)
    """

    def __init__(self,
                 label='',
                 transform=None,
                 fmt='{:.1f}',
                 total=None,
                 width=None):
        """Create a Progress reporter.

        :arg label:     Units (e.g. "MB", "%",)

        :arg transform: Function to transform values (see e.g.
                        Progress.bytes_to_mb)

        :arg fmt:       Template string used to format value / total.

        :arg total:     Maximum value - overrides the total value passed to
                        the update method.

        :arg width:     Maximum width, if a progress bar is displayed. Default
                        is to automatically infer the terminal width (see
                        Progress.get_terminal_width).
        """

        if transform is None:
            transform = Progress.default_transform

        self.width     = width
        self.fmt       = fmt.format
        self.total     = total
        self.label     = label
        self.transform = transform

        # used by the spin function
        self.__last_spin = None

    @staticmethod
    def default_transform(val, total):
        return val, total

    @staticmethod
    def bytes_to_mb(val, total):
        if val   is not None: val   = val   / 1048576
        if total is not None: total = total / 1048576
        return val, total

    @staticmethod
    def percent(val, total):
        if val is None or total is None:
            return val, total
        return 100 * (val / total), 100

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
        printmsg()

    def update(self, value=None, total=None):

        if total is None:
            total = self.total

        value, total = self.transform(value, total)

        if value is None and total is None:
            self.spin()
        elif value is not None and total is None:
            self.count(value)
        elif value is not None and total is not None:
            self.progress(value, total)

    def spin(self):

        symbols = ['|', '/', '-',  '\\']

        if self.__last_spin is not None: last = self.__last_spin
        else:                            last = symbols[-1]

        idx  = symbols.index(last)
        idx  = (idx + 1) % len(symbols)
        this = symbols[idx]

        printmsg(this, end='\r')
        self.__last_spin = this

    def count(self, value):
        value = self.fmt(value)
        line  = '{}{} ...'.format(value, self.label)
        printmsg(line, end='\r')

    def progress(self, value, total):

602
        value = min(value, total)
603
604
605
606
607
608

        # arbitrary fallback of 50 columns if
        # terminal width cannot be determined
        if self.width is None: width = Progress.get_terminal_width(50)
        else:                  width = self.width

609
610
611
612
613
614
        fvalue = self.fmt(value)
        ftotal = self.fmt(total)
        suffix = '{} / {} {}'.format(fvalue, ftotal, self.label).rstrip()

        # +5: - square brackets around bar
        #     - space between bar and tally
615
        #     - space+spin at the end
616
        width     = width - (len(suffix) + 5)
617
618
        completed = int(round(width * (value  / total)))
        remaining = width - completed
619
620
621
        progress  = '[{}{}] {}'.format('#' * completed,
                                       ' ' * remaining,
                                       suffix)
622

623
        printmsg(progress, end='')
624
625
        printmsg(' ', end='')
        self.spin()
626
        printmsg(end='\r')
627
628
629
630
631
632
633


    @staticmethod
    def get_terminal_width(fallback=None):
        """Return the number of columns in the current terminal, or fallback
        if it cannot be determined.
        """
634
        # os.get_terminal_size added in python
635
636
637
638
        # 3.3, so we try it but fall back to tput
        if hasattr(os, 'get_terminal_size'):
            return os.get_terminal_size()[0]

639
        try:
640
            result = Process.check_output('tput cols', log_output=False)
641
            return int(result.strip())
642
        except Exception:
643
644
645
            return fallback


646
@contextlib.contextmanager
647
648
649
650
651
652
def tempdir(override_dir=None):
    """Returns a context manager which creates, changes into, and returns a
    temporary directory, and then deletes it on exit.

    If override_dir is not None, instead of creating and changing into a
    temporary directory, this function just changes into override_dir.
653
654
    """

655
656
657
    if override_dir is None: tmpdir = tempfile.mkdtemp()
    else:                    tmpdir = override_dir

658
659
660
    prevdir = os.getcwd()

    try:
661
662
        os.chdir(tmpdir)
        yield tmpdir
663
664
665

    finally:
        os.chdir(prevdir)
666
667
        if override_dir is None:
            shutil.rmtree(tmpdir)
668
669
670
671
672
673
674
675
676
677
678
679
680
681


def memoize(f):
    """Decorator to memoize a function. """

    cache = f.cache = {}
    def g(*args, **kwargs):
        key = (f, tuple(args), frozenset(kwargs.items()))
        if key not in cache:
            cache[key] = f(*args, **kwargs)
        return cache[key]
    return g


Paul McCarthy's avatar
Paul McCarthy committed
682
683
684
685
class ChecksumError(Exception):
    """Exception raised by the sha256 function if a file checksume does
    not match the expected checksum.
    """
686
687


688
def sha256(filename, check_against=None, blocksize=1048576):
Paul McCarthy's avatar
Paul McCarthy committed
689
690
691
    """Calculate the SHA256 checksum of the given file. If check_against
    is provided, it is compared against the calculated checksum, and an
    error is raised if they are not the same.
692
693
    """

Paul McCarthy's avatar
Paul McCarthy committed
694
    hashobj = hashlib.sha256()
695

Paul McCarthy's avatar
Paul McCarthy committed
696
697
698
699
700
701
    with open(filename, 'rb') as f:
        while True:
            block = f.read(blocksize)
            if len(block) == 0:
                break
            hashobj.update(block)
702

Paul McCarthy's avatar
Paul McCarthy committed
703
704
    checksum = hashobj.hexdigest()

705
    if check_against is not None:
Paul McCarthy's avatar
Paul McCarthy committed
706
        if checksum != check_against:
707
708
709
710
            raise ChecksumError('File {} does not match expected checksum '
                                '({})'.format(filename, check_against))

    return checksum
711
712


Paul McCarthy's avatar
Paul McCarthy committed
713
714
715
716
class DownloadFailed(Exception):
    """Exception type raised by the download_file function if a
    download fails for some reason.
    """
717
718


719
def download_file(url, destination, progress=None, blocksize=131072):
Paul McCarthy's avatar
Paul McCarthy committed
720
721
    """Download a file from url, saving it to destination. """

722
723
    def default_progress(downloaded, total):
        pass
Paul McCarthy's avatar
Paul McCarthy committed
724

725
726
727
728
    if progress is None:
        progress = default_progress

    log.debug('Downloading %s ...', url)
Paul McCarthy's avatar
Paul McCarthy committed
729

730
    req = None
Paul McCarthy's avatar
Paul McCarthy committed
731
    try:
732
733
734
735
        # py2: urlopen result cannot be
        # used as a context manager
        req = urlrequest.urlopen(url)
        with open(destination, 'wb') as outf:
Paul McCarthy's avatar
Paul McCarthy committed
736
737
738
739
740
741

            try:             total = int(req.headers['content-length'])
            except KeyError: total = None

            downloaded = 0

Paul McCarthy's avatar
Paul McCarthy committed
742
            progress(downloaded, total)
Paul McCarthy's avatar
Paul McCarthy committed
743
744
745
746
747
748
            while True:
                block = req.read(blocksize)
                if len(block) == 0:
                    break
                downloaded += len(block)
                outf.write(block)
749
                progress(downloaded, total)
Paul McCarthy's avatar
Paul McCarthy committed
750

751
752
753
754
    except Exception:
        raise DownloadFailed('A network error has occurred while '
                             'trying to download {}'.format(url))
    finally:
755
756
        if req:
            req.close()
Paul McCarthy's avatar
Paul McCarthy committed
757
758


759
class Process(object):
760
    """Container for a subprocess.Popen object, allowing non-blocking
761
762
763
764
765
766
767
    line-based access to its standard output and error streams via separate
    queues, while logging all outputs.

    Don't create a Process directly - use one of the following static methods:
     - Process.check_output
     - Process.check_call
     - Process.monitor_progress
768
769
    """

770

771
    def __init__(self, cmd, admin=False, ctx=None, log_output=True, **kwargs):
772
773
        """Run the specified command. Starts threads to capture stdout and
        stderr.
774

775
776
777
778
779
780
781
        :arg cmd:        Command to run - passed directly to subprocess.Popen
        :arg admin:      Run the command with administrative privileges
        :arg ctx:        The installer Context. Only used for admin password -
                         can be None if admin is False.
        :arg log_output: If True, the command and all of its stdout/stderr are
                         logged.
        :arg kwargs:     Passed to subprocess.Popen
782
        """
783

784
785
786
787
788
789
790
791
792
793
        self.ctx        = ctx
        self.cmd        = cmd
        self.admin      = admin
        self.log_output = log_output
        self.stdoutq    = queue.Queue()
        self.stderrq    = queue.Queue()

        if log_output:
            log.debug('Running %s [as admin: %s]', cmd, admin)

794
        self.popen = Process.popen(self.cmd, self.admin, self.ctx, **kwargs)
795

796
        # threads for consuming stdout/stderr
797
798
        self.stdout_thread = threading.Thread(
            target=Process.forward_stream,
799
            args=(self.popen, self.stdoutq, cmd, 'stdout', log_output))
800
801
        self.stderr_thread = threading.Thread(
            target=Process.forward_stream,
802
            args=(self.popen, self.stderrq, cmd, 'stderr', log_output))
803

804
805
806
807
        self.stdout_thread.daemon = True
        self.stderr_thread.daemon = True
        self.stdout_thread.start()
        self.stderr_thread.start()
808
809


810
811
812
813
814
815
816
817
818
    def wait(self):
        """Waits for the process to terminate, then waits for the stdout
        and stderr consumer threads to finish.
        """
        self.popen.wait()
        self.stdout_thread.join()
        self.stderr_thread.join()


819
820
821
822
823
824
825
826
827
828
829
    @property
    def returncode(self):
        """Process return code. Returns None until the process has terminated,
        and the stdout/stderr consumer threads have finished.
        """
        if self.popen.returncode is None: return None
        if self.stdout_thread.is_alive(): return None
        if self.stderr_thread.is_alive(): return None
        return self.popen.returncode


830
    @staticmethod
831
    def check_output(cmd, *args, **kwargs):
832
833
834
        """Behaves like subprocess.check_output. Runs the given command, then
        waits until it finishes, and return its standard output. An error
        is raised if the process returns a non-zero exit code.
835

836
        :arg cmd: The command to run, as a string
837
838
        """

839
        proc = Process(cmd, *args, **kwargs)
840
        proc.wait()
841

842
        if proc.returncode != 0:
843
844
            raise RuntimeError(cmd)

845
846
847
848
849
850
851
852
        stdout = ''
        while True:
            try:
                stdout += proc.stdoutq.get_nowait()
            except queue.Empty:
                break

        return stdout
853
854
855


    @staticmethod
856
    def check_call(cmd, *args, **kwargs):
857
858
859
860
        """Behaves like subprocess.check_call. Runs the given command, then
        waits until it finishes. An error is raised if the process returns a
        non-zero exit code.

861
        :arg cmd: The command to run, as a string
862
        """
863
        proc = Process(cmd, *args, **kwargs)
864
        proc.wait()
865
        if proc.returncode != 0:
866
867
868
869
            raise RuntimeError(cmd)


    @staticmethod
870
    def monitor_progress(cmd, total=None, *args, **kwargs):
871
872
873
874
875
876
877
878
879
880
        """Runs the given command, and shows a progress bar under the
        assumption that cmd will produce "total" number of lines of output.
        """
        if total is None: label = None
        else:             label = '%'

        with Progress(label=label,
                      fmt='{:.0f}',
                      transform=Progress.percent) as prog:

881
            proc   = Process(cmd, *args, **kwargs)
882
883
884
885
            nlines = 0

            prog.update(nlines, total)

886
            while proc.returncode is None:
887
                try:
888
889
                    line    = proc.stdoutq.get(timeout=0.5)
                    nlines += 1
890
                except queue.Empty:
891
                    pass
892
893
894

                prog.update(nlines, total)

895
896
897
898
899
900
            # force progress bar to 100% when finished
            if proc.returncode == 0:
                prog.update(total, total)
            else:
                raise RuntimeError(cmd)

901
902

    @staticmethod
903
    def forward_stream(popen, queue, cmd, streamname, log_output):
904
905
906
907
908
909
910
        """Reads lines from stream and pushes them onto queue until popen
        is finished. Logs every line.

        :arg popen:      subprocess.Popen object
        :arg queue:      queue.Queue to push lines onto
        :arg cmd:        string - the command that is running
        :arg streamname: string - 'stdout' or 'stderr'
911
        :arg log_output: If True, log all stdout/stderr.
912
913
914
915
916
        """

        if streamname == 'stdout': stream = popen.stdout
        else:                      stream = popen.stderr

917
        while True:
918
919
920
921
922
            line = stream.readline().decode('utf-8')
            if line == '':
                break
            else:
                queue.put(line)
923
                if log_output:
924
                    log.debug(' [%s]: %s', streamname, line.rstrip())
925
926
927


    @staticmethod
928
    def popen(cmd, admin=False, ctx=None, **kwargs):
929
930
        """Runs the given command via subprocess.Popen, as administrator if
        requested.
931

932
933
934
        :arg cmd:    The command to run, as a string

        :arg admin:  Whether to run with administrative privileges
935

936
937
        :arg ctx:    The installer Context object. Only required if admin is
                     True.
938

939
940
        :arg kwargs: Passed to subprocess.Popen. stdin, stdout, and stderr
                     will be silently clobbered
941

942
        :returns:    The subprocess.Popen object.
943
944
945
946
947
948
949
        """

        admin = admin and os.getuid() != 0

        if admin: password = ctx.password
        else:     password = None

950
951
952
953
        cmd              = shlex.split(cmd)
        kwargs['stdin']  = sp.PIPE
        kwargs['stdout'] = sp.PIPE
        kwargs['stderr'] = sp.PIPE
954
955

        if admin: proc = Process.sudo_popen(cmd, password, **kwargs)
956
        else:     proc = sp.Popen(          cmd,           **kwargs)
957
958
959
960
961

        return proc


    @staticmethod
962
963
964
    def sudo_popen(cmd, password, **kwargs):
        """Runs "sudo cmd" using subprocess.Popen. Used by Process.popen.
        Assumes that kwargs contains stdin=sp.PIPE
965
        """
966

967
968
969
970
        cmd  = ['sudo', '-S', '-k'] + cmd
        proc = sp.Popen(cmd, **kwargs)
        proc.stdin.write('{}\n'.format(password))
        return proc
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991


def list_available_versions(ctx):
    """Lists available FSL versions. """
    printmsg('Available FSL versions:', EMPHASIS)
    for version in ctx.manifest['versions']:
        if version == 'latest':
            continue
        printmsg(version, IMPORTANT, EMPHASIS)
        for build in ctx.manifest['versions'][version]:
            if build.get('cuda', '').strip() != '':
                template = '  {platform} [CUDA {cuda}]'
            else:
                template = '  {platform}'
            printmsg(template.format(**build), EMPHASIS, end=' ')
            printmsg(build['environment'], INFO)


def install_miniconda(ctx):
    """Downloads the miniconda/miniforge installer, and installs it to the
    destination directory.
992
993

    This function assumes that it is run within a temporary/scratch directory.
994
995
    """

996
    metadata = ctx.manifest['miniconda'][ctx.platform]
997
998
999
    url      = metadata['url']
    checksum = metadata['sha256']
    output   = metadata.get('output', '').strip()
1000

1001
1002
    if output == '': output = None
    else:            output = int(output)
Paul McCarthy's avatar
Paul McCarthy committed
1003

1004
1005
    # Download
    printmsg('Downloading miniconda from {}...'.format(url))
1006
1007
1008
1009
    with Progress('MB', transform=Progress.bytes_to_mb) as prog:
        download_file(url, 'miniforge.sh', prog.update)
    if not ctx.args.no_checksum:
        sha256('miniforge.sh', checksum)
1010
1011

    # Install
1012
    printmsg('Installing miniconda at {}...'.format(ctx.destdir))
1013
    cmd = 'sh miniforge.sh -b -p {}'.format(ctx.destdir)
1014
    Process.monitor_progress(cmd, output, ctx.need_admin, ctx)
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049

    # Create .condarc config file
    condarc = tw.dedent("""
    # Putting a .condarc file into the root environment
    # directory will override ~/.condarc if it exists,
    # but will not override a system condarc (e.g. at
    # /etc/condarc/condarc). There is currently no
    # workaround for this - see:
    #  - https://github.com/conda/conda/issues/8599
    #  - https://github.com/conda/conda/issues/8804

    # Try and make package downloads more robust
    remote_read_timeout_secs:    240
    remote_connect_timeout_secs: 20
    remote_max_retries:          10
    remote_backoff_factor:       5
    safety_checks:               warn

    # Channel priority is important. In older versions
    # of FSL we placed the FSL conda channel at the
    # bottom (lowest priority) for legacy reasons (to
    # ensure that conda-forge versions of e.g. VTK were
    # preferred over legacy FSL conda versions).
    #
    # https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-channels.html
    channel_priority: strict
    channels:
      - http://18.133.213.73/production/
      - conda-forge
      - defaults
    """)

    with open('.condarc', 'wt') as f:
        f.write(condarc)

1050
1051
    Process.check_call('cp .condarc {}'.format(ctx.destdir),
                       ctx.need_admin, ctx)
1052
1053
1054


def install_fsl(ctx):
1055
    """Install FSL into ctx.destdir (which is assumed to be a miniconda
1056
    installation.
1057
1058

    This function assumes that it is run within a temporary/scratch directory.
1059
1060
    """

1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
    build    = ctx.build
    url      = build['environment']
    checksum = build['sha256']
    output   = build.get('output', None)

    if output == '': output = None
    else:            output = int(output)

    printmsg('Downloading FSL environment specification '
             'from {}...'.format(url))
    download_file(url, 'environment.yml')

    conda = op.join(ctx.destdir, 'bin', 'conda')
    cmd   = conda + ' env update -n base -f environment.yml'

    printmsg('Installing FSL into {}...'.format(ctx.destdir))

    env = os.environ.copy()
    env['FSL_CREATE_WRAPPER_SCRIPTS'] = '1'

1081
    Process.monitor_progress(cmd, output, ctx.need_admin, ctx, env=env)
1082

1083

1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
def finalise_installation(ctx):
    """Performs some finalisation tasks. Includes:
      - Saving the installed version to $FSLDIR/etc/fslversion
      - Saving this installer script and the environment file to
        $FSLDIR/share/fsl/installer/
    """
    with open('fslversion', 'wt') as f:
        f.write(ctx.build['version'])

    call    = ft.partial(Process.check_call, admin=ctx.need_admin, ctx=ctx)
    etcdir  = op.join(ctx.destdir, 'etc')
    instdir = op.join(ctx.destdir, 'share', 'fsl', 'installer')

    call('cp fslversion {}'     .format(etcdir))
    call('mkdir -p {}'          .format(instdir))
    call('cp environment.yml {}'.format(instdir))
    call('cp {} {}'             .format(__absfile__, instdir))


1103
1104
1105
1106
1107
1108
def post_install_cleanup(ctx):
    """Cleans up the FSL directory after installation. """

    conda = op.join(ctx.destdir, 'bin', 'conda')
    cmd   = conda + ' clean -y --all'

1109
    Process.check_call(cmd, ctx.need_admin, ctx)
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126


def patch_file(filename, searchline, numlines, content):
    """Used by configure_shell and configure_matlab. Adds to, modifies,
    or creates the specified file.

    If a line matching searchline is found in the file, numlines (starting
    from searchline) are replaced with content.

    Otherwise, content is appended to the end of the file.
    """

    content = content.split('\n')

    if op.isfile(filename):
        with open(filename) as f:
            lines = [l.strip() for l in f.readlines()]
1127
    else:
1128
        lines = []
1129

1130
1131
1132
1133
    # replace block
    try:
        idx   = lines.index(searchline)
        lines = lines[:idx] + content + lines[idx + numlines:]
1134

1135
1136
1137
    # append to end
    except ValueError:
        lines = lines + [''] + content
1138

1139
1140
    with open(filename, 'wt') as f:
        f.write('\n'.join(lines))
1141
1142


1143
def configure_shell(ctx):
1144
    """Configures the user's shell environment (e.g. ~/.bash_profile). """
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229

    bourne_shells  = ['sh', 'bash', 'zsh', 'dash']
    csh_shells     = ['csh', 'tcsh']

    # we edit the first file that exists
    # in the list of candidate profile files,
    shell_profiles = {'sh'   : ['.profile'],
                      'bash' : ['.bash_profile', '.profile'],
                      'dash' : ['.bash_profile', '.profile'],
                      'zsh'  : ['.zprofile'],
                      'csh'  : ['.cshrc'],
                      'tcsh' : ['.tcshrc']}

    # Do not change the format of these configurations -
    # they are kept exactly as-is for compatibility with
    # legacy FSL installations, i.e. so we can modify
    # profiles with an existing configuration from older
    # FSL versions
    bourne_cfg = tw.dedent("""
    # FSL Setup
    FSLDIR={fsldir}
    PATH=${{FSLDIR}}/share/fsl/bin:${{PATH}}
    export FSLDIR PATH
    . ${{FSLDIR}}/etc/fslconf/fsl.sh
    """).format(fsldir=ctx.destdir)

    csh_cfg = tw.dedent("""
    # FSL Setup
    setenv FSLDIR {fsldir}
    setenv PATH ${{FSLDIR}}/share/fsl/bin:${{PATH}}
    source ${{FSLDIR}}/etc/fslconf/fsl.csh
    """).format(fsldir=ctx.destdir)

    if ctx.shell not in bourne_shells + csh_shells:
        printmsg('Shell {} not recognised - skipping environment '
                 'setup'.format(ctx.shell), WARNING, EMPHASIS)
        return

    if ctx.shell in bourne_shells: cfg = bourne_cfg
    else:                          cfg = csh_cfg

    # find the profile file to edit
    profile    = None
    candidates = [op.join(ctx.args.homedir, p)
                  for p in shell_profiles[ctx.shell]]
    for candidate in candidates:
        if op.isfile(candidate):
            profile = candidate
            break

    # if no candidate profile files
    # exist, fall back to the first one
    if profile is None:
        profile = candidates[0]

    printmsg('Adding FSL configuration to {}'.format(profile))

    patch_file(profile, '# FSL Setup', len(cfg.split('\n')), cfg)


def configure_matlab(ctx):
    """Creates/appends FSL configuration code to ~/Documents/MATLAB/startup.m.
    """

    # Do not change the format of this configuration -
    # see in-line comments in configure_shell.
    cfg = tw.dedent("""
    % FSL Setup
    setenv( 'FSLDIR', '{fsldir}' );
    setenv('FSLOUTPUTTYPE', 'NIFTI_GZ');
    fsldir = getenv('FSLDIR');
    fsldirmpath = sprintf('%s/etc/matlab',fsldir);
    path(path, fsldirmpath);
    clear fsldir fsldirmpath;
    """).format(fsldir=ctx.destdir)

    matlab_dir = op.expanduser(op.join(ctx.args.homedir, 'Documents', 'MATLAB'))
    startup_m  = op.join(matlab_dir, 'startup.m')

    if not op.exists(matlab_dir):
        os.makedirs(matlab_dir)

    printmsg('Adding FSL configuration to {}'.format(startup_m))

    patch_file(startup_m, '% FSL Setup', len(cfg.split('\n')), cfg)
1230
1231


Paul McCarthy's avatar
Paul McCarthy committed
1232
def self_update(ctx):
1233
    """Checks to see if a newer version of the installer (this script) is
1234
    available and if so, downloads it to a temporary file, and runs it in
Paul McCarthy's avatar
Paul McCarthy committed
1235
    place of this script.
1236
1237
1238
    """

    thisver   = Version(__version__)
Paul McCarthy's avatar
Paul McCarthy committed
1239
    latestver = Version(ctx.manifest['installer']['version'])
1240
1241
1242
1243
1244
1245

    if latestver <= thisver:
        log.debug('Installer is up to date (this vesrion: %s, '
                  'latest version: %s)', thisver, latestver)
        return

Paul McCarthy's avatar
Paul McCarthy committed
1246
1247
    log.debug('New version of installer is available '
              '(%s) - self-updating', latestver)
1248

1249
1250
    tmpf = tempfile.NamedTemporaryFile(
        prefix='new_fslinstaller', delete=False, dir=ctx.args.workdir)
Paul McCarthy's avatar
Paul McCarthy committed
1251
1252
    tmpf.close()
    tmpf = tmpf.name
1253

Paul McCarthy's avatar
Paul McCarthy committed
1254
    download_file(ctx.manifest['installer']['url'], tmpf)
1255

1256
    if not ctx.args.no_checksum:
1257
        sha256(tmpf, ctx.manifest['installer']['sha256'])
1258

Paul McCarthy's avatar
Paul McCarthy committed
1259
1260
1261
    cmd = [sys.executable, tmpf] + sys.argv[1:]
    log.debug('Running new installer: %s', cmd)
    os.execv(sys.executable, cmd)
1262
1263


1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
def read_fslversion(destdir):
    """Reads the FSL version from an existing FSL installation. Returns the
    version string, or None if it can't be read.
    """
    fslversion = op.join(destdir, 'etc', 'fslversion')
    if not op.exists(fslversion):
        return None
    try:
        with open(fslversion, 'rt') as f:
            fslversion = f.readline().split(':')[0]
    except:
        return None
    return fslversion


1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
def update_destdir(ctx):
    """Called by main. Checks if the destination directory is an FSL
    installation, and determines / asks the user whether they want to update
    it.

    Returns True if the existing FSL installation should be updated, False
    if it should be overwritten.
    """

    installed = read_fslversion(ctx.destdir)

    # Cannot detect a FSL installation
    if installed is None:
        return False

    printmsg()
    printmsg('Existing FSL installation [version {}] detected '
             'at {}'.format(installed, ctx.destdir), INFO)

    installed  = Version(installed)
    requested  = Version(ctx.build['version'])
    updateable = Version(FIRST_FSL_CONDA_RELEASE)

    # Too old (pre-conda)
    if installed < updateable:
        printmsg('FSL version {} is too old to update - you will '
                 'need to overwrite it'.format(installed), INFO)
        return False

    # Existing install is equal to
    # or newer than requested
    if installed >= requested:
        if installed == requested:
            msg       = '\nFSL version {installed} is already installed!'
            promptmsg = 'Do you want to re-install FSL {installed} [y/N]?'
        else:
            msg       = '\nInstalled version [{installed}] is newer than ' \
                        'the requested version [{requested}]!'
            promptmsg = 'Do you want to replace your existing version ' \
                        '[{installed}] with an older version [{requested}] ' \
                        '[y/N]?'

        msg       = msg      .format(installed=installed, requested=requested)
        promptmsg = promptmsg.format(installed=installed, requested=requested)

        printmsg(msg, WARNING, EMPHASIS)
        response = prompt(promptmsg, QUESTION, EMPHASIS)

        # Overwrite/re-install - don't ask user
        # again if they want to overwrite destdir
        if response.lower() in ('y', 'yes'):
            ctx.args.overwrite = True
            return False
        else:
            printmsg('Aborting installation', ERROR, EMPHASIS)
            sys.exit(1)

    # --update -> don't prompt
    if ctx.args.update:
        return True

    response = prompt('Would you like to upgrade from version {} to '
                      'version {} [y/N]?'.format(installed, requested),
                      QUESTION, EMPHASIS)
1343
1344
1345
1346
1347
1348
1349
1350
1351
    if response.lower() in ('y', 'yes'):

        # We don't maintain expected #lines for
        # upgrades, so can't report progress
        ctx.build.pop('output', None)
        return True
    else:
        printmsg('Aborting installation', ERROR, EMPHASIS)
        sys.exit(1)
1352
1353


1354
def overwrite_destdir(ctx):
1355
1356
1357
    """Called by main if the destination directory already exists. Asks the
    user if they want to overwrite it. If they do, or if the --overwrite
    option was specified, the directory is deleted
1358
    """
1359
1360

    if not ctx.args.overwrite:
1361
        printmsg()
1362
1363
        printmsg('Destination directory [{}] already exists!'
                 .format(ctx.destdir), WARNING, EMPHASIS)
1364
1365
        response = prompt('Do you want to overwrite it [y/N]?',
                          QUESTION, EMPHASIS)
1366
1367
1368
        if response.lower() not in ('y', 'yes'):
            printmsg('Aborting installation', ERROR, EMPHASIS)
            sys.exit(1)
1369
1370

    printmsg('Deleting directory {}'.format(ctx.destdir), IMPORTANT)
1371
    Process.check_call('rm -r {}'.format(ctx.destdir), ctx.need_admin, ctx)
1372
1373


1374
1375
1376
1377
def parse_args(argv=None):
    """Parse command-line arguments, returns an argparse.Namespace object. """

    helps = {
1378
1379
1380
1381
1382

        'version'      : 'Print installer version number and exit',
        'listversions' : 'List available FSL versions and exit',
        'dest'         : 'Install FSL into this folder (default: '
                         '{})'.format(DEFAULT_INSTALLATION_DIRECTORY),
1383
1384
1385
1386
        'update'       : 'Update existing FSL installation if possible, '
                         'without asking',
        'overwrite'    : 'Delete existing destination directory if it exists, '
                         'without asking',
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
        'no_env'       : 'Do not modify your shell or MATLAB configuration '
                         'implies --no_shell and --no_matlab)',
        'no_shell'     : 'Do not modify your shell configuration',
        'no_matlab'    : 'Do not modify your MATLAB configuration',
        'fslversion'   : 'Install this specific version of FSL',
        'cuda'         : 'Install FSL for this CUDA version (default: '
                         'automatically detected)',

        # Do not automatically update the installer script,
        'no_self_update' : argparse.SUPPRESS,
Paul McCarthy's avatar
Paul McCarthy committed
1397
1398

        # Path to local installer manifest file
1399
        'manifest'       : argparse.SUPPRESS,
Paul McCarthy's avatar
Paul McCarthy committed
1400
1401

        # Print debugging messages
1402
        'debug'          : argparse