Skip to content
Snippets Groups Projects
my_spawner.py 9.1 KiB
Newer Older
# The spawner is based on PRP's https://gitlab.nrp-nautilus.io/vpeddu/jupyterlab-west.git
# This file is managed by [youngsu.kim|dung.vu]@csusb.edu

from kubespawner import KubeSpawner

class MySpawner(KubeSpawner):
    #Read custom HTML for spawner from my_spanwer.html
    with open('/etc/jupyterhub/custom/my_spawner.html', 'r') as f:
        profile_form_template = f.readlines()
    profile_form_template = "".join(profile_form_template)

    def options_from_form(self, formdata):
        # Uncomment these lines to add pvc to a class         
        # course_pvc_test = [
        #     "000065181@csusb.edu",
        #     "dung.vu@csusb.edu",
        #     "dvu@csusb.edu",
        #     "youngsu.kim@csusb.edu",
        # ]
        
        # cephfs_pvc_users = {
        #     "000065181@csusb.edu":"csusb-dungvu-share", 
        #     "dung.vu@csusb.edu":"csusb-dungvu-share", 
        #     "dvu@csusb.edu":"csusb-dungvu-share",
        #     "youngsu.kim@csusb.edu":"csusb-dungvu-share",
        # }

        if not self.profile_list or not hasattr(self, '_profile_list'):
            return formdata
        selected_profile = int(formdata.get('profile', [0])[0])
        options = self._profile_list[selected_profile]
        self.log.debug("Applying KubeSpawner override for profile '%s'", options['display_name'])
        kubespawner_override = options.get('kubespawner_override', {})
        for k, v in kubespawner_override.items():
            if callable(v):
                v = v(self)
                self.log.debug(".. overriding KubeSpawner value %s=%s (callable result)", k, v)
            else:
                self.log.debug(".. overriding KubeSpawner value %s=%s", k, v)
            setattr(self, k, v)

        gpus = int(formdata.get('gpus', [0])[0])
        
        #Add at least one gpu for certain profiles
        # if "Datascience" in options['display_name'] and gpus == int(0):
        #     gpus = int(1)
        
        #These two lines to avoid overriding the GPU assigned to some stacks by default
        if "nvidia.com/gpu" in options["kubespawner_override"]["extra_resource_limits"] and gpus == 0:
            pass
        else 
            setattr(self, "extra_resource_limits", {"nvidia.com/gpu": gpus})
            
        # setattr(self, "extra_resource_limits", {"nvidia.com/gpu": gpus})
        setattr(self, "mem_guarantee", formdata.get('ram', [0])[0]+"G")
        setattr(self, "cpu_guarantee", float(formdata.get('cores', [0])[0]))
        setattr(self, "mem_limit", formdata.get('ram', [0])[0]+"G")
        setattr(self, "cpu_limit", float(formdata.get('cores', [0])[0]))

        nodeSelectorTermsExpressions = [{
            'key': 'topology.kubernetes.io/region',
            'operator': 'In',
            'values': ["us-west"]
            }]

        if formdata.get('gputype', [0])[0]:
            nodeSelectorTermsExpressions.append({
                'key': 'nvidia.com/gpu.product',
                'operator': 'In',
                'values': formdata.get('gputype', [0])
            })

        setattr(self, 'extra_pod_config', {
            'affinity': {
                'nodeAffinity': {
                    'requiredDuringSchedulingIgnoredDuringExecution': {
                        'nodeSelectorTerms': [{
                            'matchExpressions': nodeSelectorTermsExpressions,
                        }],
                    },
                },
            }
            })

        self.volume_mounts = [
            {
                'name': 'volume-{username}',
                'mountPath': '/home/jovyan',
            }
            ]
        self.volumes = [
            {
                'name': 'volume-{username}',
                'persistentVolumeClaim': {
                    'claimName': 'claim-{username}'
                    }
            }
            ]

        # if formdata.get('shm', [0])[0]:
        #     self.volume_mounts.append({
        #         'name': 'dshm',
        #         'mountPath': '/dev/shm',
        #         })
        #     self.volumes.append({
        #         'name': 'dshm',
        #         'emptyDir': {'medium': 'Memory'}
        #         })

        # The following is based on PRP's JupyterHub West 
        # We use course based shared folders 
        # if formdata.get('cephfs', [0])[0] and self.user.name in cephfs_pvc_users:
        #     self.volume_mounts.append({
        #         'name': 'cephfs',
        #         'mountPath': '/home/jovyan/cephfs',
        #     })
        #     self.volumes.append({
        #         'name': 'cephfs',
        #         'persistentVolumeClaim': {
        #         'claimName': cephfs_pvc_users[self.user.name]
        #         }
        #     })

        # Uncomment these lines to add pvc to a class 
        # if formdata.get('cephfscourse', [0])[0] and self.user.name in course_pvc_test:
        #     self.volume_mounts.append({
        #         'name': 'volumn-course-pvc-test',
        #         'mountPath': '/home/jovyan/course-pvc-test',
        #         })
        #     self.volumes.append({ 
        #         'name': 'volumn-course-pvc-test',
        #         'persistentVolumeClaim': {
        #             'claimName': 'course-pvc-test',
        #             }
        #         })

        return options

    profile_list = [
        {
            "display_name": "Stack Minimal",
            "default": true,
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/minimal"
            }
        },
        {
            "display_name": "Stack Minimal + Desktop GUI + Relion",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/relion-desktop"
            }
        },
        {
            "display_name": "Stack Datascience",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/datascience",
                "extra_resource_limits": {
                    "nvidia.com/gpu": "2"
                }
            }
        },
        {
            "display_name": "Stack Tensorflow [2GPUs]",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/prp",
                "extra_resource_limits": {
                    "nvidia.com/gpu": "2"
                }
            }
        },
        {
            "display_name": "Stack Tensorflow [4GPUs]",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/prp",
                "extra_resource_limits": {
                    "nvidia.com/gpu": "4"
                }
            }
        },
        {
            "display_name": "Stack R-Studio",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/r-studio"
            }
        },
        {
            "display_name": "Stack Pyspark [2GPUs]",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/pyspark",
                "extra_resource_limits": {
                    "nvidia.com/gpu": "2"
                }
            }
        },
        {
            "display_name": "Stack Pyspark [4GPUs]",
            "kubespawner_override": {
                "image_spec": "localhost:30081/prp/jupyter-stack/pyspark",
                "extra_resource_limits": {
                    "nvidia.com/gpu": "4"
                }
            }
        },
        {
            "display_name": "Stack SageMath",
            "kubespawner_override": {
                "image_spec": "gitlab-registry.nrp-nautilus.io/youngsu_kim/csusb-jupyter-stack/stack-sagemath:b353d66b"
            }
        }
    ]

    # profile_list = [
    #     {
    #         'display_name': 'Stack Minimal',
    #         'kubespawner_override': {
    #             'image_spec': 'localhost:30081/prp/jupyter-stack/minimal',
    #         }
    #     },
    #     {
    #         'display_name': 'Stack Minimal + Desktop GUI + Relion',
    #         'kubespawner_override': {
    #             'image_spec': 'localhost:30081/prp/jupyter-stack/relion-desktop',
    #         }
    #     },
    #     {
    #         'display_name': 'R Studio Server',
    #         'kubespawner_override': {
    #             'image_spec': 'localhost:30081/prp/jupyter-stack/r-studio',
    #         }
    #     },
    #     {
    #         'display_name': 'Stack Tensorflow',
    #         'kubespawner_override': {
    #             'image_spec': 'localhost:30081/prp/jupyter-stack/tensorflow',
    #         }
    #     },
    #     {
    #         'display_name': 'Stack Tensorflow + PRP added libs',
    #         # 'default': 'true',
    #     },
    #     {
    #         'display_name': 'Stack Datascience',
    #         'default': 'true',
    #         'kubespawner_override': {
    #             'image_spec': 'localhost:30081/prp/jupyter-stack/datascience',
    #         }
    #     },
    #     {
    #         'display_name': 'Stack All Spark',
    #         'kubespawner_override': {
    #             'image_spec': 'localhost:30081/prp/jupyter-stack/all-spark',
    #         }
    #     }
    # ]

c.JupyterHub.spawner_class = MySpawner
# c.JupyterHub.template_vars = {'announcement_login': 'announcements here'}