Skip to content

Commit

Permalink
feat(core): shell completion for sessions (#3450)
Browse files Browse the repository at this point in the history
  • Loading branch information
mohammad-alisafaee authored May 24, 2023
1 parent 41927a1 commit 9fa63dd
Show file tree
Hide file tree
Showing 25 changed files with 311 additions and 196 deletions.
6 changes: 2 additions & 4 deletions conftest.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,5 @@
# -*- coding: utf-8 -*-
#
# Copyright 2017-2023 Swiss Data Science Center (SDSC)
# A partnership between École Polytechnique Fédérale de Lausanne (EPFL) and
# Copyright Swiss Data Science Center (SDSC). A partnership between
# École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down
3 changes: 1 addition & 2 deletions design/003-interactive-session/003-interactive-session.md
Original file line number Diff line number Diff line change
Expand Up @@ -160,10 +160,9 @@ class ISessionProvider:
"""
pass

def session_list(self, project_name: str, config: Optional[Dict[str, Any]]) -> List[Session]:
def session_list(self, project_name: str) -> List[Session]:
"""Lists all the sessions currently running by the given session provider.
:param project_name: Renku project name.
:param config: Path to the session provider specific configuration YAML.
:returns: a list of sessions.
"""
pass
Expand Down
1 change: 1 addition & 0 deletions docs/how-to-guides/shell-integration.rst
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ To activate tab completion for your supported shell run the following command af
$ eval "$(_RENKU_COMPLETE=zsh_source renku)"
You can put the same command in your shell's startup script to enable completion by default.
After this not only sub-commands of ``renku`` will be auto-completed using tab, but for example
in case of ``renku workflow execute`` the available ``Plans`` are going to be listed.

Expand Down
28 changes: 22 additions & 6 deletions renku/command/session.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
#
# Copyright 2018-2023- Swiss Data Science Center (SDSC)
# A partnership between École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
# Copyright Swiss Data Science Center (SDSC). A partnership between
# École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
Expand All @@ -16,9 +15,26 @@
# limitations under the License.
"""Renku session commands."""


from renku.command.command_builder.command import Command
from renku.core.session.session import session_list, session_open, session_start, session_stop, ssh_setup
from renku.core.session.session import (
search_session_providers,
search_sessions,
session_list,
session_open,
session_start,
session_stop,
ssh_setup,
)


def search_sessions_command():
"""Get all the session names that match a pattern."""
return Command().command(search_sessions).require_migration().with_database(write=False)


def search_session_providers_command():
"""Get all the session provider names that match a pattern."""
return Command().command(search_session_providers).require_migration().with_database(write=False)


def session_list_command():
Expand Down
4 changes: 3 additions & 1 deletion renku/core/plugin/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,4 +39,6 @@ def get_supported_session_providers() -> List[ISessionProvider]:
from renku.core.plugin.pluginmanager import get_plugin_manager

pm = get_plugin_manager()
return pm.hook.session_provider()
providers = pm.hook.session_provider()

return sorted(providers, key=lambda p: p.priority)
7 changes: 3 additions & 4 deletions renku/core/session/__init__.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
#
# Copyright 2018-2023 - Swiss Data Science Center (SDSC)
# A partnership between École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
# Copyright Swiss Data Science Center (SDSC). A partnership between
# École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
Expand Down
44 changes: 27 additions & 17 deletions renku/core/session/docker.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
#
# Copyright 2018-2023 - Swiss Data Science Center (SDSC)
# A partnership between École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
# Copyright Swiss Data Science Center (SDSC). A partnership between
# École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -33,7 +32,7 @@
from renku.core.plugin import hookimpl
from renku.core.util import communication
from renku.domain_model.project_context import project_context
from renku.domain_model.session import ISessionProvider, Session
from renku.domain_model.session import ISessionProvider, Session, SessionStopStatus

if TYPE_CHECKING:
from renku.core.dataset.providers.models import ProviderParameter
Expand All @@ -43,7 +42,7 @@ class DockerSessionProvider(ISessionProvider):
"""A docker based interactive session provider."""

JUPYTER_PORT = 8888
# NOTE: Give the docker provider a higher priority so that it's checked first
# NOTE: Give the docker provider the highest priority so that it's checked first
priority: ProviderPriority = ProviderPriority.HIGHEST

def __init__(self):
Expand All @@ -54,7 +53,7 @@ def docker_client(self) -> docker.client.DockerClient:
Note:
This is not a @property, even though it should be, because ``pluggy``
will call it in that case in unrelated parts of the code that will
will call it in that case in unrelated parts of the code.
Raises:
errors.DockerError: Exception when docker is not available.
Returns:
Expand Down Expand Up @@ -133,7 +132,7 @@ def get_open_parameters(self) -> List["ProviderParameter"]:
"""Returns parameters that can be set for session open."""
return []

def session_list(self, project_name: str, config: Optional[Dict[str, Any]]) -> List[Session]:
def session_list(self, project_name: str) -> List[Session]:
"""Lists all the sessions currently running by the given session provider.
Returns:
Expand Down Expand Up @@ -297,29 +296,36 @@ def session_start_helper(consider_disk_request: bool):
else:
return result, ""

def session_stop(self, project_name: str, session_name: Optional[str], stop_all: bool) -> bool:
def session_stop(self, project_name: str, session_name: Optional[str], stop_all: bool) -> SessionStopStatus:
"""Stops all or a given interactive session."""
try:
docker_containers = (
self._get_docker_containers(project_name)
if stop_all
else self.docker_client().containers.list(filters={"id": session_name})
if session_name
else self.docker_client().containers.list()
)

if len(docker_containers) == 0:
return False
n_docker_containers = len(docker_containers)

if n_docker_containers == 0:
return SessionStopStatus.FAILED if session_name else SessionStopStatus.NO_ACTIVE_SESSION
elif not session_name and len(docker_containers) > 1:
return SessionStopStatus.NAME_NEEDED

[c.stop() for c in docker_containers]
return True
except docker.errors.APIError as error:
raise errors.DockerError(error.msg)
else:
return SessionStopStatus.SUCCESSFUL

def session_open(self, project_name: str, session_name: str, **kwargs) -> bool:
def session_open(self, project_name: str, session_name: Optional[str], **kwargs) -> bool:
"""Open a given interactive session.
Args:
project_name(str): Renku project name.
session_name(str): The unique id of the interactive session.
session_name(Optional[str]): The unique id of the interactive session.
"""
url = self.session_url(session_name)

Expand All @@ -329,10 +335,14 @@ def session_open(self, project_name: str, session_name: str, **kwargs) -> bool:
webbrowser.open(url)
return True

def session_url(self, session_name: str) -> Optional[str]:
def session_url(self, session_name: Optional[str]) -> Optional[str]:
"""Get the URL of the interactive session."""
for c in self.docker_client().containers.list():
if c.short_id == session_name and f"{DockerSessionProvider.JUPYTER_PORT}/tcp" in c.ports:
sessions = self.docker_client().containers.list()

for c in sessions:
if (
c.short_id == session_name or (not session_name and len(sessions) == 1)
) and f"{DockerSessionProvider.JUPYTER_PORT}/tcp" in c.ports:
host = c.ports[f"{DockerSessionProvider.JUPYTER_PORT}/tcp"][0]
return f'http://{host["HostIp"]}:{host["HostPort"]}/?token={c.labels["jupyter_token"]}'
return None
Expand Down
62 changes: 43 additions & 19 deletions renku/core/session/renkulab.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
#
# Copyright 2018-2023 - Swiss Data Science Center (SDSC)
# A partnership between École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
# Copyright Swiss Data Science Center (SDSC). A partnership between
# École Polytechnique Fédérale de Lausanne (EPFL) and
# Eidgenössische Technische Hochschule Zürich (ETHZ).
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
Expand All @@ -26,6 +25,7 @@

from renku.core import errors
from renku.core.config import get_value
from renku.core.constant import ProviderPriority
from renku.core.login import read_renku_token
from renku.core.plugin import hookimpl
from renku.core.session.utils import get_renku_project_name, get_renku_url
Expand All @@ -34,7 +34,7 @@
from renku.core.util.jwt import is_token_expired
from renku.core.util.ssh import SystemSSHConfig
from renku.domain_model.project_context import project_context
from renku.domain_model.session import ISessionProvider, Session
from renku.domain_model.session import ISessionProvider, Session, SessionStopStatus

if TYPE_CHECKING:
from renku.core.dataset.providers.models import ProviderParameter
Expand All @@ -44,6 +44,8 @@ class RenkulabSessionProvider(ISessionProvider):
"""A session provider that uses the notebook service API to launch sessions."""

DEFAULT_TIMEOUT_SECONDS = 300
# NOTE: Give the renkulab provider the lowest priority so that it's checked last
priority: ProviderPriority = ProviderPriority.LOWEST

def __init__(self):
self.__renku_url: Optional[str] = None
Expand Down Expand Up @@ -187,7 +189,7 @@ def _cleanup_ssh_connection_configs(
gotten from the server.
"""
if not running_sessions:
running_sessions = self.session_list("", None, ssh_garbage_collection=False)
running_sessions = self.session_list(project_name="", ssh_garbage_collection=False)

system_config = SystemSSHConfig()

Expand All @@ -199,7 +201,8 @@ def _cleanup_ssh_connection_configs(
if path not in session_config_paths:
path.unlink()

def _remote_head_hexsha(self):
@staticmethod
def _remote_head_hexsha():
remote = get_remote(repository=project_context.repository)

if remote is None:
Expand All @@ -221,7 +224,8 @@ def _send_renku_request(self, req_type: str, *args, **kwargs):
)
return res

def _project_name_from_full_project_name(self, project_name: str) -> str:
@staticmethod
def _project_name_from_full_project_name(project_name: str) -> str:
"""Get just project name of project name if in owner/name form."""
if "/" not in project_name:
return project_name
Expand Down Expand Up @@ -282,9 +286,7 @@ def get_open_parameters(self) -> List["ProviderParameter"]:
ProviderParameter("ssh", help="Open a remote terminal through SSH.", is_flag=True),
]

def session_list(
self, project_name: str, config: Optional[Dict[str, Any]], ssh_garbage_collection: bool = True
) -> List[Session]:
def session_list(self, project_name: str, ssh_garbage_collection: bool = True) -> List[Session]:
"""Lists all the sessions currently running by the given session provider.
Returns:
Expand Down Expand Up @@ -398,45 +400,67 @@ def session_start(
)
raise errors.RenkulabSessionError("Cannot start session via the notebook service because " + res.text)

def session_stop(self, project_name: str, session_name: Optional[str], stop_all: bool) -> bool:
def session_stop(self, project_name: str, session_name: Optional[str], stop_all: bool) -> SessionStopStatus:
"""Stops all sessions (for the given project) or a specific interactive session."""
responses = []
sessions = self.session_list(project_name=project_name)
n_sessions = len(sessions)

if n_sessions == 0:
return SessionStopStatus.NO_ACTIVE_SESSION

if stop_all:
sessions = self.session_list(project_name=project_name, config=None)
for session in sessions:
responses.append(
self._send_renku_request(
"delete", f"{self._notebooks_url()}/servers/{session.id}", headers=self._auth_header()
)
)
self._wait_for_session_status(session.id, "stopping")
else:
elif session_name:
responses.append(
self._send_renku_request(
"delete", f"{self._notebooks_url()}/servers/{session_name}", headers=self._auth_header()
)
)
self._wait_for_session_status(session_name, "stopping")
elif n_sessions == 1:
responses.append(
self._send_renku_request(
"delete", f"{self._notebooks_url()}/servers/{sessions[0].id}", headers=self._auth_header()
)
)
self._wait_for_session_status(sessions[0].id, "stopping")
else:
return SessionStopStatus.NAME_NEEDED

self._cleanup_ssh_connection_configs(project_name)

return all([response.status_code == 204 for response in responses]) if responses else False
n_successfully_stopped = len([r for r in responses if r.status_code == 204])

def session_open(self, project_name: str, session_name: str, ssh: bool = False, **kwargs) -> bool:
return SessionStopStatus.SUCCESSFUL if n_successfully_stopped == n_sessions else SessionStopStatus.FAILED

def session_open(self, project_name: str, session_name: Optional[str], ssh: bool = False, **kwargs) -> bool:
"""Open a given interactive session.
Args:
project_name(str): Renku project name.
session_name(str): The unique id of the interactive session.
session_name(Optional[str]): The unique id of the interactive session.
ssh(bool): Whether to open an SSH connection or a normal browser interface.
"""
sessions = self.session_list("", None)
sessions = self.session_list(project_name="")
system_config = SystemSSHConfig()
name = self._project_name_from_full_project_name(project_name)
ssh_prefix = f"{system_config.renku_host}-{name}-"

if not session_name:
if len(sessions) == 1:
session_name = sessions[0].id
else:
return False

if session_name.startswith(ssh_prefix):
# NOTE: use passed in ssh connection name instead of session id by accident
# NOTE: User passed in ssh connection name instead of session id by accident
session_name = session_name.replace(ssh_prefix, "", 1)

if not any(s.id == session_name for s in sessions):
Expand Down
Loading

0 comments on commit 9fa63dd

Please sign in to comment.