gen_json.py 2.71 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25
# Copyright (c) 2016, Daniele Venzano
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Jupyter Zoe application description generator."""

import json
import sys
import os

REPOSITORY = os.getenv("REPOSITORY", default="zapps")
VERSION = os.getenv("VERSION", default="latest")

UPSTREAM_NOTEBOOKS = [
Daniele Venzano's avatar
Daniele Venzano committed
26
    ("pytorch-notebook", REPOSITORY + "/pytorch:" + VERSION, []),
Daniele Venzano's avatar
Daniele Venzano committed
27
    ("pytorch-nb-gpu", REPOSITORY + "/pytorch:" + VERSION, ['gpu'])
28 29 30 31 32 33 34 35 36 37 38 39 40 41 42
]

ZOE_APPLICATION_DESCRIPTION_VERSION = 3

options = {
    'core_limit': {
        'value': 4,
        'description': 'Notebook core limit'
    },
    'memory_limit': {
        'value': 4 * (1024**3),
        'description': 'Notebook memory limit (bytes)'
    }
}

43
def jupyter_service(memory_limit, core_limit, image, labels):
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
    """
    :rtype: dict
    """
    service = {
        'name': "jupyter",
        'image': image,
        'monitor': True,
        'resources': {
            "memory": {
                "min": memory_limit,
                "max": memory_limit
            },
            "cores": {
                "min": core_limit,
                "max": core_limit
            }
        },
        'ports': [
            {
                'name': 'Jupyter Notebook interface',
                'protocol': 'tcp',
                'port_number': 8888,
                'url_template': 'http://{ip_port}/'
            }
        ],
        'environment': [
            ['NB_UID', '1000'],
            ['HOME', '/mnt/workspace']
        ],
        'volumes': [],
        'command': None,
        'work_dir': '/mnt/workspace',
        'total_count': 1,
        'essential_count': 1,
        'replicas': 1,
        'startup_order': 0
    }
81 82
    if len(labels) > 0:
        service['labels'] = labels
83 84 85 86 87

    return service


if __name__ == '__main__':
88
    for app_name, image, labels in UPSTREAM_NOTEBOOKS:
89 90 91 92 93 94
        app = {
            'name': app_name,
            'version': ZOE_APPLICATION_DESCRIPTION_VERSION,
            'will_end': False,
            'size': 512,
            'services': [
95
                jupyter_service(options["memory_limit"]["value"], options["core_limit"]["value"], image, labels)
96 97 98 99 100 101 102
            ]
        }

        json.dump(app, open(app_name + ".json", "w"), sort_keys=True, indent=4)

    print("ZApps written")