feat(cli): add --force-build flag to docker session start (#3369)

Author: Panaetius

renku/core/session/docker.py

@@ -125,6 +125,7 @@ def get_start_parameters(self) -> List["ProviderParameter"]:
 
         return [
             ProviderParameter("port", help="Local port to use (random if not specified).", type=int),
+            ProviderParameter("force-build", help="Always build image and don't check if it exists.", is_flag=True),
         ]
 
     def get_open_parameters(self) -> List["ProviderParameter"]:
@@ -322,3 +323,7 @@ def session_url(self, session_name: str) -> Optional[str]:
                 host = c.ports[f"{DockerSessionProvider.JUPYTER_PORT}/tcp"][0]
                 return f'http://{host["HostIp"]}:{host["HostPort"]}/?token={c.labels["jupyter_token"]}'
         return None
+
+    def force_build_image(self, force_build: bool = False, **kwargs) -> bool:
+        """Whether we should force build the image directly or check for an existing image first."""
+        return force_build

renku/core/session/renkulab.py

@@ -45,8 +45,9 @@ class RenkulabSessionProvider(ISessionProvider):
     DEFAULT_TIMEOUT_SECONDS = 300
 
     def __init__(self):
-        self.__renku_url = None
-        self.__notebooks_url = None
+        self.__renku_url: Optional[str] = None
+        self.__notebooks_url: Optional[str] = None
+        self._force_build: bool = False
 
     def _renku_url(self) -> str:
         """Get the URL of the renku instance."""
@@ -164,13 +165,15 @@ def pre_start_checks(self, ssh: bool = False, **kwargs):
                     "Your system is not set up for SSH connections to Renkulab. Would you like to set it up?"
                 ):
                     ssh_setup()
+                    self._force_build = True
                 else:
                     raise errors.RenkulabSessionError(
                         "Can't run ssh session without setting up Renku SSH support. Run without '--ssh' or "
                         "run 'renku session ssh-setup'."
                     )
 
-            system_config.setup_session_keys()
+            if system_config.setup_session_keys():
+                self._force_build = True
 
     def _cleanup_ssh_connection_configs(
         self, project_name: str, running_sessions: Optional[List[Session]] = None
@@ -475,3 +478,7 @@ def session_url(self, session_name: str) -> str:
             session_name,
         ]
         return urllib.parse.urljoin(self._renku_url(), "/".join(session_url_parts))
+
+    def force_build_image(self, **kwargs) -> bool:
+        """Whether we should force build the image directly or check for an existing image first."""
+        return self._force_build

renku/core/session/session.py

@@ -130,17 +130,19 @@ def session_start(
         if repo_host:
             image_name = f"{repo_host}/{image_name}"
 
-        if not provider_api.find_image(image_name, config):
-            communication.confirm(
-                f"The container image '{image_name}' does not exist. Would you like to build it using {provider}?",
-                abort=True,
-            )
-            with communication.busy(msg=f"Building image {image_name}"):
-                provider_api.build_image(project_context.dockerfile_path.parent, image_name, config)
-            communication.echo(f"Image {image_name} built successfully.")
-    else:
-        if not provider_api.find_image(image_name, config):
-            raise errors.ParameterError(f"Cannot find the provided container image '{image_name}'!")
+    force_build_image = provider_api.force_build_image(**kwargs)
+
+    if not force_build_image and not provider_api.find_image(image_name, config):
+        communication.confirm(
+            f"The container image '{image_name}' does not exist. Would you like to build it using {provider}?",
+            abort=True,
+        )
+        force_build_image = True
+
+    if force_build_image:
+        with communication.busy(msg=f"Building image {image_name}"):
+            provider_api.build_image(project_context.dockerfile_path.parent, image_name, config)
+        communication.echo(f"Image {image_name} built successfully.")
 
     # set resource settings
     cpu_limit = cpu_request or get_value("interactive", "cpu_request")

renku/core/util/ssh.py

@@ -133,7 +133,7 @@ def session_config_path(self, project_name: str, session_name: str) -> Path:
         """
         return self.renku_ssh_root / f"00-{project_name}-{session_name}.conf"
 
-    def setup_session_keys(self):
+    def setup_session_keys(self) -> bool:
         """Add a users key to a project."""
         project_context.ssh_authorized_keys_path.parent.mkdir(parents=True, exist_ok=True)
         project_context.ssh_authorized_keys_path.touch(mode=0o600, exist_ok=True)
@@ -142,7 +142,7 @@ def setup_session_keys(self):
             raise errors.SSHNotSetupError()
 
         if self.public_key_string in project_context.ssh_authorized_keys_path.read_text():
-            return
+            return False
 
         communication.info("Adding SSH public key to project.")
         with project_context.ssh_authorized_keys_path.open("at") as f:
@@ -153,6 +153,7 @@ def setup_session_keys(self):
         communication.info(
             "Added public key. Changes need to be pushed and remote image built for changes to take effect."
         )
+        return True
 
     def setup_session_config(self, project_name: str, session_name: str) -> str:
         """Setup local SSH config for connecting to a session.

renku/domain_model/session.py

@@ -182,3 +182,7 @@ def pre_start_checks(self, **kwargs):
         make sure that the session launches successfully. By default this method does not do any checks.
         """
         return None
+
+    def force_build_image(self, **kwargs) -> bool:
+        """Whether we should force build the image directly or check for an existing image first."""
+        return False

renku/ui/cli/session.py

@@ -54,8 +54,10 @@
 
     $ renku session start -p docker
 
-The command first looks for a local image to use. If a local image isn't found, it searches the remote Renku deployment
-(if any) and pulls the image if it exists. Finally, it prompts the user to build the image locally if no image is found.
+The command first looks for a local image to use. If a local image isn't found, it
+searches the remote Renku deployment (if any) and pulls the image if it exists.
+Finally, it prompts the user to build the image locally if no image is found. You
+can force the image to always be built by using the ``--force-build`` flag.
 
 Renkulab provider
 ~~~~~~~~~~~~~~~~~

tests/core/plugins/test_session.py

@@ -68,6 +68,10 @@ def fake_pre_start_checks(self, **kwargs):
     pass
 
 
+def fake_force_build_image(self, **kwargs):
+    return kwargs.get("force_build", False)
+
+
 @pytest.mark.parametrize(
     "provider_name,session_provider,provider_patches",
     [
@@ -80,7 +84,8 @@ def fake_pre_start_checks(self, **kwargs):
     [
         ({}, "0xdeadbeef"),
         ({"image_name": "fixed_image"}, "0xdeadbeef"),
-        ({"image_name": "missing_image"}, ParameterError),
+        ({"image_name": "missing_image"}, click.Abort),
+        (({"image_name": "missing_image", "force_build": True}, "0xdeadbeef")),
     ],
 )
 def test_session_start(
@@ -101,6 +106,7 @@ def test_session_start(
         find_image=fake_find_image,
         build_image=fake_build_image,
         pre_start_checks=fake_pre_start_checks,
+        force_build_image=fake_force_build_image,
         **provider_patches,
     ):
         provider_implementation = next(

tests/service/views/test_dataset_views.py

@@ -1343,7 +1343,13 @@ def test_edit_datasets_view_unset_values(svc_client_with_repo):
 
     assert_rpc_response(response)
     assert {"warnings", "edited", "remote_branch"} == set(response.json["result"])
-    assert {"keywords": [], "custom_metadata": None, "images": []} == response.json["result"]["edited"]
+    assert {
+        "keywords": [],
+        "custom_metadata": None,
+        "images": [],
+    } == response.json[
+        "result"
+    ]["edited"]
 
     params_list = {
         "project_id": project_id,

tests/service/views/test_project_views.py

@@ -143,7 +143,12 @@ def test_edit_project_view_unset(svc_client_with_repo):
 
     assert_rpc_response(response)
     assert {"warning", "edited", "remote_branch"} == set(response.json["result"])
-    assert {"keywords": None, "custom_metadata": None} == response.json["result"]["edited"]
+    assert {
+        "keywords": None,
+        "custom_metadata": None,
+    } == response.json[
+        "result"
+    ]["edited"]
 
 
 @pytest.mark.service