fslinstaller.py 48.8 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
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):
596

597
        value = self.fmt(value)
598
599
600
601

        if self.label is None: line = '{} ...'.format(value)
        else:                  line = '{}{} ...'.format(value, self.label)

602
603
604
605
        printmsg(line, end='\r')

    def progress(self, value, total):

606
        value = min(value, total)
607
608
609
610
611
612

        # 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

613
614
615
616
617
618
        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
619
        #     - space+spin at the end
620
        width     = width - (len(suffix) + 5)
621
622
        completed = int(round(width * (value  / total)))
        remaining = width - completed
623
624
625
        progress  = '[{}{}] {}'.format('#' * completed,
                                       ' ' * remaining,
                                       suffix)
626

627
        printmsg(progress, end='')
628
629
        printmsg(' ', end='')
        self.spin()
630
        printmsg(end='\r')
631
632
633
634
635
636
637


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

643
        try:
644
            result = Process.check_output('tput cols', log_output=False)
645
            return int(result.strip())
646
        except Exception:
647
648
649
            return fallback


650
@contextlib.contextmanager
651
652
653
654
655
656
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.
657
658
    """

659
660
661
    if override_dir is None: tmpdir = tempfile.mkdtemp()
    else:                    tmpdir = override_dir

662
663
664
    prevdir = os.getcwd()

    try:
665
666
        os.chdir(tmpdir)
        yield tmpdir
667
668
669

    finally:
        os.chdir(prevdir)
670
671
        if override_dir is None:
            shutil.rmtree(tmpdir)
672
673
674
675
676
677
678
679
680
681
682
683
684
685


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
686
687
688
689
class ChecksumError(Exception):
    """Exception raised by the sha256 function if a file checksume does
    not match the expected checksum.
    """
690
691


692
def sha256(filename, check_against=None, blocksize=1048576):
Paul McCarthy's avatar
Paul McCarthy committed
693
694
695
    """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.
696
697
    """

Paul McCarthy's avatar
Paul McCarthy committed
698
    hashobj = hashlib.sha256()
699

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

Paul McCarthy's avatar
Paul McCarthy committed
707
708
    checksum = hashobj.hexdigest()

709
    if check_against is not None:
Paul McCarthy's avatar
Paul McCarthy committed
710
        if checksum != check_against:
711
712
713
714
            raise ChecksumError('File {} does not match expected checksum '
                                '({})'.format(filename, check_against))

    return checksum
715
716


Paul McCarthy's avatar
Paul McCarthy committed
717
718
719
720
class DownloadFailed(Exception):
    """Exception type raised by the download_file function if a
    download fails for some reason.
    """
721
722


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

726
727
    def default_progress(downloaded, total):
        pass
Paul McCarthy's avatar
Paul McCarthy committed
728

729
730
731
732
    if progress is None:
        progress = default_progress

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

734
    req = None
Paul McCarthy's avatar
Paul McCarthy committed
735
    try:
736
737
738
739
        # 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
740
741
742
743
744
745

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

            downloaded = 0

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

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


763
class Process(object):
764
    """Container for a subprocess.Popen object, allowing non-blocking
765
766
767
768
769
770
771
    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
772
773
    """

774

775
    def __init__(self, cmd, admin=False, ctx=None, log_output=True, **kwargs):
776
777
        """Run the specified command. Starts threads to capture stdout and
        stderr.
778

779
780
781
782
783
784
785
        :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
786
        """
787

788
789
790
791
792
793
794
795
796
797
        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)

798
        self.popen = Process.popen(self.cmd, self.admin, self.ctx, **kwargs)
799

800
        # threads for consuming stdout/stderr
801
802
        self.stdout_thread = threading.Thread(
            target=Process.forward_stream,
803
            args=(self.popen.stdout, self.stdoutq, cmd, 'stdout', log_output))
804
805
        self.stderr_thread = threading.Thread(
            target=Process.forward_stream,
806
            args=(self.popen.stderr, self.stderrq, cmd, 'stderr', log_output))
807

808
809
810
811
        self.stdout_thread.daemon = True
        self.stderr_thread.daemon = True
        self.stdout_thread.start()
        self.stderr_thread.start()
812
813


814
815
816
817
818
819
820
821
822
    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()


823
824
825
826
827
828
829
830
831
832
833
    @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


834
    @staticmethod
835
    def check_output(cmd, *args, **kwargs):
836
837
838
        """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.
839

840
        :arg cmd: The command to run, as a string
841
842
        """

843
        proc = Process(cmd, *args, **kwargs)
844
        proc.wait()
845

846
        if proc.returncode != 0:
847
848
            raise RuntimeError(cmd)

849
850
851
852
853
854
855
856
        stdout = ''
        while True:
            try:
                stdout += proc.stdoutq.get_nowait()
            except queue.Empty:
                break

        return stdout
857
858
859


    @staticmethod
860
    def check_call(cmd, *args, **kwargs):
861
862
863
864
        """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.

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


    @staticmethod
874
    def monitor_progress(cmd, total=None, *args, **kwargs):
875
876
877
878
879
880
881
882
883
884
        """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:

885
            proc   = Process(cmd, *args, **kwargs)
886
            nlines = 0 if total else None
887
888
889

            prog.update(nlines, total)

890
            while proc.returncode is None:
891
                try:
892
                    line    = proc.stdoutq.get(timeout=0.5)
893
894
                    nlines  = (nlines + 1) if total else None

895
                except queue.Empty:
896
                    pass
897
898

                prog.update(nlines, total)
899
                proc.popen.poll()
900

901
902
903
904
905
906
            # force progress bar to 100% when finished
            if proc.returncode == 0:
                prog.update(total, total)
            else:
                raise RuntimeError(cmd)

907
908

    @staticmethod
909
    def forward_stream(stream, queue, cmd, streamname, log_output):
910
911
912
        """Reads lines from stream and pushes them onto queue until popen
        is finished. Logs every line.

913
        :arg stream:     stream to forward
914
915
916
        :arg queue:      queue.Queue to push lines onto
        :arg cmd:        string - the command that is running
        :arg streamname: string - 'stdout' or 'stderr'
917
        :arg log_output: If True, log all stdout/stderr.
918
919
        """

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


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

935
936
937
        :arg cmd:    The command to run, as a string

        :arg admin:  Whether to run with administrative privileges
938

939
940
        :arg ctx:    The installer Context object. Only required if admin is
                     True.
941

942
943
        :arg kwargs: Passed to subprocess.Popen. stdin, stdout, and stderr
                     will be silently clobbered
944

945
        :returns:    The subprocess.Popen object.
946
947
948
949
950
951
952
        """

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

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

953
954
955
956
        cmd              = shlex.split(cmd)
        kwargs['stdin']  = sp.PIPE
        kwargs['stdout'] = sp.PIPE
        kwargs['stderr'] = sp.PIPE
957
958

        if admin: proc = Process.sudo_popen(cmd, password, **kwargs)
959
        else:     proc = sp.Popen(          cmd,           **kwargs)
960
961
962
963
964

        return proc


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

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


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.
995
996

    This function assumes that it is run within a temporary/scratch directory.
997
998
    """

999
    metadata = ctx.manifest['miniconda'][ctx.platform]
1000
1001
1002
    url      = metadata['url']
    checksum = metadata['sha256']
    output   = metadata.get('output', '').strip()
1003

1004
1005
    if output == '': output = None
    else:            output = int(output)
Paul McCarthy's avatar
Paul McCarthy committed
1006

1007
1008
    # Download
    printmsg('Downloading miniconda from {}...'.format(url))
1009
1010
1011
1012
    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)
1013
1014

    # Install
1015
    printmsg('Installing miniconda at {}...'.format(ctx.destdir))
1016
    cmd = 'sh miniforge.sh -b -p {}'.format(ctx.destdir)
1017
    Process.monitor_progress(cmd, output, ctx.need_admin, ctx)
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
1050
1051
1052

    # 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)

1053
1054
    Process.check_call('cp .condarc {}'.format(ctx.destdir),
                       ctx.need_admin, ctx)
1055
1056


1057
def install_fsl(ctx, update_from=None):
1058
    """Install FSL into ctx.destdir (which is assumed to be a miniconda
1059
    installation.
1060

1061
1062
    :arg update_from: FSL version string of existing installation that is
                      being updated, or None if this is a new installation.
1063
1064
    """

1065
1066
1067
1068
    build    = ctx.build
    url      = build['environment']
    checksum = build['sha256']

1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
    # expected number of output lines
    # for new install or upgrade, used
    # for progress reporting
    if update_from is None:
        output = build.get('output', {}).get('install', None)
    else:
        output = build.get('output', {}).get(update_from, None)

    if output in ('', None): output = None
    else:                    output = int(output)
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088

    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))

Paul McCarthy's avatar
Paul McCarthy committed
1089
1090
1091
    # post-link scripts call $FSLDIR/share/fsl/sbin/createFSLWrapper
    # (part of fsl/base), which will only do its thing if the following
    # env vars are set
1092
1093
    env = os.environ.copy()
    env['FSL_CREATE_WRAPPER_SCRIPTS'] = '1'
Paul McCarthy's avatar
Paul McCarthy committed
1094
    env['FSLDIR']                     = ctx.destdir
1095

1096
    Process.monitor_progress(cmd, output, ctx.need_admin, ctx, env=env)
1097

1098

1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
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))


1118
1119
1120
1121
1122
1123
def post_install_cleanup(ctx):
    """Cleans up the FSL directory after installation. """

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

1124
    Process.check_call(cmd, ctx.need_admin, ctx)
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141


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()]
1142
    else:
1143
        lines = []
1144

1145
1146
1147
1148
    # replace block
    try:
        idx   = lines.index(searchline)
        lines = lines[:idx] + content + lines[idx + numlines:]
1149

1150
1151
1152
    # append to end
    except ValueError:
        lines = lines + [''] + content
1153

1154
1155
    with open(filename, 'wt') as f:
        f.write('\n'.join(lines))
1156
1157


1158
def configure_shell(ctx):
1159
    """Configures the user's shell environment (e.g. ~/.bash_profile). """
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
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244

    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)
1245
1246


Paul McCarthy's avatar
Paul McCarthy committed
1247
def self_update(ctx):
1248
    """Checks to see if a newer version of the installer (this script) is
1249
    available and if so, downloads it to a temporary file, and runs it in
Paul McCarthy's avatar
Paul McCarthy committed
1250
    place of this script.
1251
1252
1253
    """

    thisver   = Version(__version__)
Paul McCarthy's avatar
Paul McCarthy committed
1254
    latestver = Version(ctx.manifest['installer']['version'])
1255
1256
1257
1258
1259
1260

    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
1261
1262
    log.debug('New version of installer is available '
              '(%s) - self-updating', latestver)
1263

1264
1265
    tmpf = tempfile.NamedTemporaryFile(
        prefix='new_fslinstaller', delete=False, dir=ctx.args.workdir)
Paul McCarthy's avatar
Paul McCarthy committed
1266
1267
    tmpf.close()
    tmpf = tmpf.name
1268

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

1271
    if not ctx.args.no_checksum:
1272
        sha256(tmpf, ctx.manifest['installer']['sha256'])
1273

Paul McCarthy's avatar
Paul McCarthy committed
1274
1275
1276
    cmd = [sys.executable, tmpf] + sys.argv[1:]
    log.debug('Running new installer: %s', cmd)
    os.execv(sys.executable, cmd)
1277
1278


1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
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


1294
1295
1296
1297
1298
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.

1299
1300
    Returns the old FSL version string if the existing FSL installation
    should be updated, or None if it should be overwritten.
1301
1302
1303
1304
1305
1306
    """

    installed = read_fslversion(ctx.destdir)

    # Cannot detect a FSL installation
    if installed is None:
1307
        return None
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320

    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)
1321
        return None
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345

    # 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
1346
            return None
1347
1348
1349
1350
1351
1352
        else:
            printmsg('Aborting installation', ERROR, EMPHASIS)
            sys.exit(1)

    # --update -> don't prompt
    if ctx.args.update:
1353
        return str(installed)
1354
1355
1356
1357

    response = prompt('Would you like to upgrade from version {} to '
                      'version {} [y/N]?'.format(installed, requested),
                      QUESTION, EMPHASIS)
1358
    if response.lower() in ('y', 'yes'):