Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

start jupyterhub-singleuser within docker spawer #95

Merged
merged 5 commits into from
Sep 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
143 changes: 143 additions & 0 deletions src/jupyterhub_config/hub_singleuser.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,143 @@
import fcntl
import grp
import logging
import os
import pwd
import subprocess
import tempfile
from pathlib import Path


class SingleUserEnvManager:
"""
Manages the environment and user-specific settings for a single user in a JupyterHub server.
"""

def __init__(self,
username,
groupname: str = 'jupyterhub', ):
"""
Initializes the environment manager for the specified user.
:param username: The name of the user for whom the environment is set up.
:param groupname: The name of the OS system group to which the user belongs.
"""
self.username = username # OS system username
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't understand why the juipyterhub username and the linux username need to be the same

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

somewhere down the line (I think it's auth) will need to check that system user exist with the same jupyterhub username.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't quite get that - why would the linux username need to match the auth username?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe for some spawner it's not required. But I don't see it concerns us.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yarn is working fine with different users.

Not if it's bypassing fair share it isn't

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will explore some yarn setting (spark.yarn.queue) on the client's side.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you just use the same user for the containers you don't have to worry about this

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I created another issue for this. #97

It's really not how Jupyterhub works as far as I understand.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you're saying that JH requires the linux username to match the JH username then yeah, that won't work and we'll have to look into modifying yarn or how the submit works

self.groupname = groupname # OS system group

Check warning on line 25 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L24-L25

Added lines #L24 - L25 were not covered by tests

# Inherit the current environment variables
self.environment = dict(os.environ)

Check warning on line 28 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L28

Added line #L28 was not covered by tests
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This still makes me nervous...

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will get a whitelist later once system is up and running.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍


self.global_home = Path(os.environ['JUPYTERHUB_USER_HOME'])
self.user_dir = self.global_home / username
self.log = self._init_logger()

Check warning on line 32 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L30-L32

Added lines #L30 - L32 were not covered by tests

def setup_environment(self):
"""
Set up the user's environment by ensuring the system user exists, and the workspace
permissions are correctly configured.
"""
self.log.info(f"Setting up environment for user: {self.username}")
self._ensure_system_user()
self._ensure_workspace_permission()

Check warning on line 41 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L39-L41

Added lines #L39 - L41 were not covered by tests

def start_single_user_server(self):
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

copied from customer_spawner.py

"""
Starts the single-user Jupyter server for the user, inheriting the configured environment.
"""
self.log.info(f'Starting single-user server for {self.username} with environment: {self.environment}')
env_vars = [f'{key}={value}' for key, value in self.environment.items()]

Check warning on line 48 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L47-L48

Added lines #L47 - L48 were not covered by tests

cmd = ['sudo', '-E', '-u', self.username, 'env'] + env_vars + [

Check warning on line 50 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L50

Added line #L50 was not covered by tests
os.path.join(os.environ['JUPYTERHUB_CONFIG_DIR'], 'spawn_notebook.sh')]

self.log.info(f'Executing command: {" ".join(cmd)}')
subprocess.run(cmd, check=True)

Check warning on line 54 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L53-L54

Added lines #L53 - L54 were not covered by tests

def _ensure_system_user(self):
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

mostly copied from customer_spawner.py. But using pwd and grp to check existing user and group.

"""
Create a system user with the given username if it does not already exist.
Ensure the group exists before creating the user.
Use a file lock to prevent race conditions.
"""

lock_file = os.path.join(tempfile.gettempdir(), f'user_creation_{self.username}.lock')

Check warning on line 63 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L63

Added line #L63 was not covered by tests

with open(lock_file, 'w') as lock:
fcntl.flock(lock, fcntl.LOCK_EX)
try:

Check warning on line 67 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L65-L67

Added lines #L65 - L67 were not covered by tests
# Check if user already exists
try:
pwd.getpwnam(self.username)
self.log.info(f'User {self.username} already exists')
return
except KeyError:

Check warning on line 73 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L69-L73

Added lines #L69 - L73 were not covered by tests

# Create the user
self.log.info(f'Creating system user: {self.username}')
useradd_cmd = ['sudo', 'useradd', '-r']

Check warning on line 77 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L76-L77

Added lines #L76 - L77 were not covered by tests

if self.groupname:

Check warning on line 79 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L79

Added line #L79 was not covered by tests
# Check if the group exists, create if necessary
try:
grp.getgrnam(self.groupname)
self.log.info(f'Group {self.groupname} already exists')
except KeyError:
self.log.info(f'Group {self.groupname} does not exist, creating it.')
subprocess.run(['sudo', 'groupadd', self.groupname], check=True)

Check warning on line 86 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L81-L86

Added lines #L81 - L86 were not covered by tests

useradd_cmd.extend(['-g', self.groupname])

Check warning on line 88 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L88

Added line #L88 was not covered by tests

useradd_cmd.append(self.username)

Check warning on line 90 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L90

Added line #L90 was not covered by tests

self.log.info(f'Creating system user: {self.username}')
subprocess.run(useradd_cmd, check=True)

Check warning on line 93 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L92-L93

Added lines #L92 - L93 were not covered by tests

except subprocess.CalledProcessError as e:
raise ValueError(f'Failed to create system user: {e}')

Check warning on line 96 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L95-L96

Added lines #L95 - L96 were not covered by tests

finally:
fcntl.flock(lock, fcntl.LOCK_UN)

Check warning on line 99 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L99

Added line #L99 was not covered by tests

def _ensure_workspace_permission(self):
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

copied from custom_spawner.py

"""
Ensure the user's workspace has the correct permissions.
"""
try:
user_info = pwd.getpwnam(self.username)
except KeyError:
raise ValueError(f'System user {self.username} does not exist')
gid = user_info.pw_gid
group_name = grp.getgrgid(gid).gr_name

Check warning on line 110 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L105-L110

Added lines #L105 - L110 were not covered by tests

self.log.info(f'Configuring workspace permissions for {self.username}')
subprocess.run(['sudo', 'chown', '-R', f'{self.username}:{group_name}', self.user_dir], check=True)
subprocess.run(['sudo', 'chmod', '-R', '750', self.user_dir], check=True)

Check warning on line 114 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L112-L114

Added lines #L112 - L114 were not covered by tests

def _init_logger(self):
"""
Initializes a logger for tracking the operations performed for the user.
"""
logger = logging.getLogger(self.username)
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
return logger

Check warning on line 126 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L120-L126

Added lines #L120 - L126 were not covered by tests


def main():
print("Starting hub_singleuser.py...")

Check warning on line 130 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L130

Added line #L130 was not covered by tests

username = os.getenv('JUPYTERHUB_USER')

Check warning on line 132 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L132

Added line #L132 was not covered by tests

if not username:
raise ValueError('JUPYTERHUB_USER environment variable not set')

Check warning on line 135 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L134-L135

Added lines #L134 - L135 were not covered by tests

single_user_manager = SingleUserEnvManager(username)
single_user_manager.setup_environment()
single_user_manager.start_single_user_server()

Check warning on line 139 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L137-L139

Added lines #L137 - L139 were not covered by tests


if __name__ == "__main__":
main()

Check warning on line 143 in src/jupyterhub_config/hub_singleuser.py

View check run for this annotation

Codecov / codecov/patch

src/jupyterhub_config/hub_singleuser.py#L143

Added line #L143 was not covered by tests
4 changes: 4 additions & 0 deletions test/src/jupyterhub_config/hub_singleuser_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
from jupyterhub_config.hub_singleuser import *

def test_noop():
pass
Loading