dockerrun

#!/usr/bin/env python3
import os
import grp
import argparse
import random
from subprocess import call, check_output


# set requested GPUs
try:
    os.environ['NV_GPU'] = os.environ['GPU']
    GPU = os.environ['GPU']
except KeyError:
    GPU = ''
if GPU == '':
    GPU = 'cpu'
    os.environ['NV_GPU'] = ','  # empty NV_GPU leads to all GPUs being mounted into the container


# find other containers using the same GPU
def check_other_containers(gpu):
    if (gpu == "") or (gpu == "cpu"):
        return()
    
    # get container ids and container names
    a = check_output(['docker', 'ps'], text=True).split('\n')
    containerIds = []
    containerNames = []
    for l in a[1:]:
        containerId = l.strip().split(" ")[0]
        if containerId.strip() != "":
            containerIds.append(containerId)
        containerName = l.strip().split(" ")[-1]
        if containerName.strip() != "":
            containerNames.append(containerName)

    # check if any container uses the specified GPU
    for containerId, containerName in zip(containerIds, containerNames):
        try:
            usedGpu = check_output(['docker', 'exec', containerId, 'sh', '-c', "echo $GPU"]).strip()
            if usedGpu == gpu:
                print(">>> WARNING the following container use the requested GPU:")
                print("CONTAINER ID: {}, CONTAINER NAME: {}".format(containerId, containerName))
        except:
            print(">>> WARNING cannot probe container it is unclear if it uses a GPU:")
            print("CONTAINER ID: {}, CONTAINER NAME: {}".format(containerId, containerName))


# parse custom arguments
parser = argparse.ArgumentParser()

parser.add_argument('-u', '--user', metavar='USER_NAME', type=str,
                    default=os.environ['USER'],
					help='Specify an alternate user name (default: $USER)')

parser.add_argument('--uid', metavar='USER_ID', type=int,
                    default=os.getuid(),
					help='Specify an alternative user ID (default: $UID)') 

parser.add_argument('--jupyterport', metavar='JUPYTER_PORT', type=int,
                    default=random.randint(600,800),
					help='Specify a jupyter port (default: randomly chosen in port range [600, 800])')
parser.add_argument('--sshport', metavar='SSH_PORT', type=int,
					help='Specify an SSH port (default: no ssh port)')

parser.add_argument('--name', metavar='CONTAINER_NAME', type=str,
                    default="",
					help='Name of container (default: $USER-[CPU/$GPU])')

parser.add_argument('--local-datasets', default=False, action='store_true')

(args, extras) = parser.parse_known_args()


# Set user group and directories
# All user groups available
user_groups=[grp.getgrgid(g).gr_name for g in os.getgroups()]

# Set primary group name
group_name='ECKERLAB'
group_id=47162

if group_name not in user_groups:
    group_id = os.getegid()
    group_name = grp.getgrgid(group_id).gr_name
    home_dir = '/home/' + args.user
    bind_dir = ['-v', '/srv/pcpool/' + home_dir + ':' + home_dir]
else:
    home_dir = '/usr/users/'
    bind_dir = ['-v', home_dir + args.user + ':' + home_dir + args.user,
                '-v', home_dir + 'agecker:' + home_dir + 'agecker']

    if args.local_datasets:
        bind_dir += ['-v', '/user/' + args.user + '/local_datasets' + ':/local_datasets']


# container name
gpu_name = '-gpu' + GPU.translate(str.maketrans('', '',  ',')) if GPU != "cpu" else '-cpu'
container_name = os.environ['USER'] + gpu_name
if args.name != "":
    container_name = container_name + "-" + args.name
if os.environ.get('CONTAINER_POSTFIX'):
    container_name = container_name + "-" + os.environ['CONTAINER_POSTFIX']

# put together the nvidia-docker command
command = ['docker', 'run', '-d']

if args.sshport:
    command.extend(['-p', str(args.sshport) + ':22'])  # set SSH port

if args.jupyterport:
    command.extend(['-p', str(args.jupyterport) + ':8888'])  # set JUPYTER port

if '--env-file' not in extras:
    print('Setting standard value JUPYTER_TOKEN =token')
    command += ['-e', 'JUPYTER_TOKEN=token']


# mount directories
command.extend(bind_dir)

command += [
		   '-e', 'USER=' + args.user,                  # set user-name
		   '-e', 'NB_UID=' + str(args.uid),
		   '-e', 'NB_USER=' + args.user,
		   '-e', 'NB_GID=' + str(group_id),
		   '-e', 'NB_GROUP=' + group_name,
		   '-e', 'GPU=' + GPU,                         # set GPU
		   '-w', home_dir,                         
		   '-e', 'JUPYTER_ENABLE_LAB=yes',
		   '-e', 'GRANT_SUDO=yes',
		   '-e', 'CHOWN_HOME=yes',
		   '--name', container_name,                   # set container name
		   '--shm-size', '128G',                       # set size of /dev/shm
		   '--user', 'root',                           # set user
		   '--rm',
		   ] + extras


# print directions
print("Setting Notebook port binding to: {} (to set manually add --jupyterport <port> as flag)".format(args.jupyterport))
print("")
print("You can now open the notebook on the host machine by directing your browser to")
print("")
print("    http://localhost:%i" % args.jupyterport)
print("")

check_other_containers(GPU)
call(command)