Start data science environment container

start_data_science_env.py

from __future__ import print_function

import os, time, argparse
from subprocess import call, check_call, check_output


def stop_cmd(name):
    return ['docker', 'stop', name], False


def rm_cmd(name):
    return ['docker', 'rm', name], False


def pull_cmd(name):
    return ['docker', 'pull', name], True


def data_only_container_run_cmd(name, dest, source=None, check=True):
    if source:
        volume_args = ['-v', '{s}:{d}'.format(s=source, d=dest)]
    else:
        volume_args = ['-v', '{d}'.format(d=dest)]

    args = ['docker', 'run', '-d']
    args += ['--name', name]
    args += volume_args
    args += ['busybox']
    args += ['echo', name]

    return args, check


def mkdir_cmd(name, dest, source=None, check=True):
    if source is not None:
        return ['mkdir', '-p', source], False


def stateless_container_run_cmd(name, image, environments=None, volumes_from=None, ports=None, links=None):
    env_args, vol_args, port_args, link_args = [[],]*4
    if environments is not None:
        env_args = [['-e', '{0}={1}'.format(*pair)] for pair in environments]
    if volumes_from is not None:
        vol_args = [['--volumes-from', vol] for vol in volumes_from]
    if ports is not None:
        port_args = [['-p', port] for port in ports]
    if links is not None:
        link_args = [['--link', '{}:{}'.format(*link)] for link in links]

    args = ['docker', 'run', '-d']
    args += ['--name', name]
    args += [item for sublist in  env_args for item in sublist]
    args += [item for sublist in  vol_args for item in sublist]
    args += [item for sublist in  port_args for item in sublist]
    args += [item for sublist in  link_args for item in sublist]
    args += [image]

    return args, True


def main(notebook_dir, data_dir, password="YourPassword", launch_browser=True, docker_namespace="calvingiles", postgres=False):
    run_stack = []
    stop_stack = []
    rm_stack = []
    
    # data only containers
    data_only_containers = [('notebooks', '/notebooks', notebook_dir),
                            ('data', '/data', data_dir),
                            ]
    
    for data_only_args in data_only_containers:
        stop_stack.append(stop_cmd(data_only_args[0]))
        rm_stack.append(rm_cmd(data_only_args[0]))
        run_stack.append(mkdir_cmd(*data_only_args))
        run_stack.append(data_only_container_run_cmd(*data_only_args))


    # database containers
    database_containers = []
    if postgres is True:
        # postgres database
        data_containers = [('dev-postgres-data', '/var/lib/postgresql/data')]
        for data_only_args in data_containers:
            # stop_stack.append(stop_cmd(data_only_args[0]))
            # rm_stack.append(rm_cmd(data_only_args[0]))
            run_stack.append(data_only_container_run_cmd(*data_only_args, check=False))
        
        volumes = [data_only_args[0] for data_only_args in data_containers]
        ports = ['5433:5432']  # Expose on local port
        name = 'dev-postgres'
        image = 'postgres'

        stop_stack.append(stop_cmd(name))
        rm_stack.append(rm_cmd(name))
        run_stack.append(pull_cmd(image))
        run_stack.append(stateless_container_run_cmd(name=name, image=image, volumes_from=volumes, ports=ports))


    # notebook server container
    environments = [('PASSWORD', password)]
    volumes = [data_only_args[0] for data_only_args in data_only_containers]
    notebook_port = 443
    ports = ['{}:8888'.format(notebook_port)]

    name = 'ds-env'
    image = '{}/data-science-environment'.format(docker_namespace)
    if postgres is True:
        links = [('dev-postgres', 'postgres')]
    else:
        links = []

    stop_stack.append(stop_cmd(name))
    rm_stack.append(rm_cmd(name))
    run_stack.append(pull_cmd(image))
    run_stack.append(stateless_container_run_cmd(name=name, image=image, environments=environments, volumes_from=volumes, ports=ports, links=links))


    # run the commands
    exec_cmd(['boot2docker', 'start'])
    for cmd, check in stop_stack[::-1] + rm_stack[::-1] + run_stack:
        exec_cmd(cmd, check=check)
    print("Notebook server setup succesfully.")
    
    ip_str = exec_cmd(['boot2docker', 'ip'], result=True)
    b2d_ip = ip_str.strip().split()[-1]
    notebook_addr = "https://{ip}:{port}".format(ip=b2d_ip, port=notebook_port)
    if launch_browser:
        print("Launching notebook at {}".format(notebook_addr))
        time.sleep(1.0)
        exec_cmd(['open', notebook_addr])

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-R", "--dry-run", action="store_true", default=False, help="Just output the commands.")
    parser.add_argument("-n", "--notebook-dir", default="~/notebooks", help="Specify a non-default location for your notebooks.")
    parser.add_argument("-d", "--data-dir", default="~/data", help="Specify a non-default location for your data.")
    parser.add_argument("-L", "--no-launch", default=False, action="store_true", help="Don't launch the notebook in the browser.")
    parser.add_argument("-p", "--password", default="YourPassword", help="Password for the notebook server.")
    parser.add_argument("--postgres", default=False, action="store_true", help="Start a postgres database container.")
    parser.add_argument("--docker-namespace", default="calvingiles", help="Set your own docker username if you have customised the Dockerfile for yourself.")
    args = parser.parse_args()
    if args.dry_run:
        def exec_cmd(args, check=True, result=False):
            print(' '.join(args))
            call(['echo', '[CMD]'] + args)
    else:
        def exec_cmd(args, check=True, result=False):
            call(['echo', '[CMD]'] + args)
            if result:
                return check_output(args)
            elif check:
                check_call(args)
            else:
                call(args)
                
    main(notebook_dir=os.path.expanduser(args.notebook_dir),
         data_dir=os.path.expanduser(args.data_dir),
         password=args.password,
         launch_browser=not args.no_launch,
         docker_namespace=args.docker_namespace,
         postgres=args.postgres,
         )