mlops-131 user managed notebooks sync and push for all notebooks

src/wanna/cli/plugins/common_options.py

@@ -20,9 +20,8 @@
     PushMode.all,
     "--mode",
     "-m",
-    help="Pipeline push mode, due to CI/CD not "
-    "allowing to push to docker registry from "
-    "GCP Agent, we need to split it. "
+    help="Push mode, this is useful if you want to "
+    "push containers in one step and deploy instances in other."
     "Use all for dev",
 )
 

src/wanna/cli/plugins/managed_notebook_plugin.py

@@ -4,10 +4,19 @@
 import typer
 
 from wanna.cli.plugins.base_plugin import BasePlugin
-from wanna.cli.plugins.common_options import instance_name_option, profile_name_option, wanna_file_option
+from wanna.cli.plugins.common_options import (
+    instance_name_option,
+    profile_name_option,
+    push_mode_option,
+    wanna_file_option,
+)
+from wanna.core.deployment.models import PushMode
+from wanna.core.loggers.wanna_logger import get_logger
 from wanna.core.services.managed_notebook import ManagedNotebookService
 from wanna.core.utils.config_loader import load_config_from_yaml
 
+logger = get_logger(__name__)
+
 
 class ManagedNotebookPlugin(BasePlugin):
     """
@@ -23,6 +32,7 @@ def __init__(self) -> None:
                 self.sync,
                 self.report,
                 self.build,
+                self.push,
             ]
         )
 
@@ -45,6 +55,7 @@ def create(
         file: Path = wanna_file_option,
         profile_name: str = profile_name_option,
         instance_name: str = instance_name_option("managed_notebook", "create"),
+        mode: PushMode = push_mode_option,
     ) -> None:
         """
         Create a Managed Workbench Notebook.
@@ -57,13 +68,14 @@ def create(
         config = load_config_from_yaml(file, gcp_profile_name=profile_name)
         workdir = pathlib.Path(file).parent.resolve()
         nb_service = ManagedNotebookService(config=config, workdir=workdir)
-        nb_service.create(instance_name)
+        nb_service.create(instance_name, push_mode=mode)
 
     @staticmethod
     def sync(
         file: Path = wanna_file_option,
         profile_name: str = profile_name_option,
         force: bool = typer.Option(False, "--force", help="Synchronisation without prompt"),
+        mode: PushMode = push_mode_option,
     ) -> None:
         """
         Synchronize existing Managed Notebooks with wanna.yaml
@@ -76,7 +88,7 @@ def sync(
         config = load_config_from_yaml(file, gcp_profile_name=profile_name)
         workdir = pathlib.Path(file).parent.resolve()
         nb_service = ManagedNotebookService(config=config, workdir=workdir)
-        nb_service.sync(force)
+        nb_service.sync(force=force, push_mode=mode)
 
     @staticmethod
     def report(
@@ -111,3 +123,30 @@ def build(
         workdir = pathlib.Path(file).parent.resolve()
         nb_service = ManagedNotebookService(config=config, workdir=workdir)
         nb_service.build()
+
+    @staticmethod
+    def push(
+        file: Path = wanna_file_option,
+        profile_name: str = profile_name_option,
+        instance_name: str = instance_name_option("notebook", "push"),
+        mode: PushMode = typer.Option(
+            PushMode.containers,
+            "--mode",
+            "-m",
+            help="Managed-Notebook push mode, due to CI/CD not "
+            "allowing to push to docker registry from "
+            "GCP Agent, we need to split it. "
+            "Notebooks currently support only containers, as we do not create manifests as of now.",
+        ),
+    ) -> None:
+        """
+        Push docker containers. This command also builds the images.
+        """
+        if mode != PushMode.containers:
+            logger.user_error("Only containers are supported push mode as of now.")
+            typer.Exit(1)
+
+        config = load_config_from_yaml(file, gcp_profile_name=profile_name)
+        workdir = pathlib.Path(file).parent.resolve()
+        nb_service = ManagedNotebookService(config=config, workdir=workdir)
+        nb_service.push(instance_name=instance_name)

src/wanna/cli/plugins/notebook_plugin.py

@@ -5,10 +5,19 @@
 import typer
 
 from wanna.cli.plugins.base_plugin import BasePlugin
-from wanna.cli.plugins.common_options import instance_name_option, profile_name_option, wanna_file_option
+from wanna.cli.plugins.common_options import (
+    instance_name_option,
+    profile_name_option,
+    push_mode_option,
+    wanna_file_option,
+)
+from wanna.core.deployment.models import PushMode
+from wanna.core.loggers.wanna_logger import get_logger
 from wanna.core.services.notebook import NotebookService
 from wanna.core.utils.config_loader import load_config_from_yaml
 
+logger = get_logger(__name__)
+
 
 class NotebookPlugin(BasePlugin):
     """
@@ -24,6 +33,8 @@ def __init__(self) -> None:
                 self.ssh,
                 self.report,
                 self.build,
+                self.push,
+                self.sync,
             ]
         )
 
@@ -47,6 +58,7 @@ def create(
         profile_name: str = profile_name_option,
         instance_name: str = instance_name_option("notebook", "create"),
         owner: Optional[str] = typer.Option(None, "--owner", "-o", help=""),
+        mode: PushMode = push_mode_option,
     ) -> None:
         """
         Create a User-Managed Workbench Notebook.
@@ -59,7 +71,7 @@ def create(
         config = load_config_from_yaml(file, gcp_profile_name=profile_name)
         workdir = pathlib.Path(file).parent.resolve()
         nb_service = NotebookService(config=config, workdir=workdir, owner=owner)
-        nb_service.create(instance_name)
+        nb_service.create(instance_name, push_mode=mode)
 
     @staticmethod
     def ssh(
@@ -133,3 +145,50 @@ def build(
         workdir = pathlib.Path(file).parent.resolve()
         nb_service = NotebookService(config=config, workdir=workdir)
         nb_service.build()
+
+    @staticmethod
+    def push(
+        file: Path = wanna_file_option,
+        profile_name: str = profile_name_option,
+        instance_name: str = instance_name_option("notebook", "push"),
+        mode: PushMode = typer.Option(
+            PushMode.containers,
+            "--mode",
+            "-m",
+            help="Notebook push mode, due to CI/CD not "
+            "allowing to push to docker registry from "
+            "GCP Agent, we need to split it. "
+            "Notebooks currently support only containers, as we do not create manifests as of now.",
+        ),
+    ) -> None:
+        """
+        Push docker containers. This command also builds the images.
+        """
+        if mode != PushMode.containers:
+            logger.user_error("Only containers are supported push mode as of now.")
+            typer.Exit(1)
+
+        config = load_config_from_yaml(file, gcp_profile_name=profile_name)
+        workdir = pathlib.Path(file).parent.resolve()
+        nb_service = NotebookService(config=config, workdir=workdir)
+        nb_service.push(instance_name=instance_name)
+
+    @staticmethod
+    def sync(
+        file: Path = wanna_file_option,
+        profile_name: str = profile_name_option,
+        force: bool = typer.Option(False, "--force", help="Synchronisation without prompt"),
+        mode: PushMode = push_mode_option,
+    ) -> None:
+        """
+        Synchronize existing User-managed Notebooks with wanna.yaml
+
+        1. Reads current notebooks where label is defined per field wanna_project.name in wanna.yaml
+        2. Does a diff between what is on GCP and what is on yaml
+        3. Create the ones defined in yaml and missing in GCP
+        4. Delete the ones in GCP that are not in wanna.yaml
+        """
+        config = load_config_from_yaml(file, gcp_profile_name=profile_name)
+        workdir = pathlib.Path(file).parent.resolve()
+        nb_service = NotebookService(config=config, workdir=workdir)
+        nb_service.sync(force=force, push_mode=mode)

src/wanna/core/models/base_instance.py

@@ -16,7 +16,7 @@ class BaseInstanceModel(BaseModel, extra=Extra.ignore, validate_assignment=True)
     network: Optional[str]
     bucket: Optional[str]
     tags: Optional[List[str]]
-    metadata: Optional[List[Dict[str, Any]]]
+    metadata: Optional[Dict[str, Any]]
 
     _project_id = validator("project_id", allow_reuse=True)(validators.validate_project_id)
     _zone = validator("zone", allow_reuse=True)(validators.validate_zone)

src/wanna/core/models/notebook.py

@@ -34,6 +34,7 @@ class NotebookModel(BaseInstanceModel):
     network: Optional[str]
     subnet: Optional[str]
     tensorboard_ref: Optional[str]
+    enable_monitoring: bool = True
     no_public_ip: bool = True
     no_proxy_access: bool = False
 
@@ -46,7 +47,7 @@ class ManagedNotebookModel(BaseInstanceModel):
     machine_type: Optional[str] = "n1-standard-4"
     gpu: Optional[GPU]
     data_disk: Optional[Disk]
-    kernels: Optional[List[str]]
+    kernel_docker_image_refs: Optional[List[str]]
     tensorboard_ref: Optional[str]
     network: Optional[str]
     subnet: Optional[str]

src/wanna/core/services/base.py

@@ -1,5 +1,7 @@
 from abc import ABC
-from typing import Generic, List, Optional, TypeVar
+from typing import Generic, List, Optional, Tuple, TypeVar
+
+import typer
 
 from wanna.core.deployment.models import PushMode
 from wanna.core.loggers.wanna_logger import get_logger
@@ -180,3 +182,36 @@ def _get_resource_subnet(self, network: Optional[str], subnet: Optional[str], re
                 return f"projects/{network_project_id}/regions/{region}/subnetworks/{subnet}"
         else:
             return None
+
+    def _return_diff(self) -> Tuple[List[T], List[T]]:
+        """
+        Abstract class.
+        """
+        ...
+
+    def sync(self, force: bool, push_mode: PushMode = PushMode.all) -> None:
+        """
+        1. Reads current instances where label is defined per field wanna_project.name in wanna.yaml
+        2. Does a diff between what is on GCP and what is on yaml
+        3. Delete the ones in GCP that are not in wanna.yaml
+        4. Create the ones defined in yaml and missing in GCP
+        """
+        to_be_deleted, to_be_created = self._return_diff()  # pylint: disable=assignment-from-no-return
+
+        if to_be_deleted:
+            to_be_deleted_str = "\n".join(["- " + item.name for item in to_be_deleted])
+            logger.user_info(f"{self.instance_type.capitalize()}s to be deleted:\n{to_be_deleted_str}")
+            should_delete = True if force else typer.confirm("Are you sure you want to delete them?")
+            if should_delete:
+                for notebook in to_be_deleted:
+                    self._delete_one_instance(notebook)
+
+        if to_be_created:
+            to_be_created_str = "\n".join(["- " + item.name for item in to_be_created])
+            logger.user_info(f"{self.instance_type.capitalize()}s to be created:\n{to_be_created_str}")
+            should_create = True if force else typer.confirm("Are you sure you want to create them?")
+            if should_create:
+                for notebook in to_be_created:
+                    self._create_one_instance(notebook, push_mode=push_mode)
+
+        logger.user_info(f"{self.instance_type.capitalize()}s on GCP are in sync with wanna.yaml")

src/wanna/core/services/docker.py

@@ -11,6 +11,7 @@
 from google.protobuf.duration_pb2 import Duration  # pylint: disable=no-name-in-module
 from python_on_whales import Image, docker
 
+from wanna.core.deployment.models import PushMode
 from wanna.core.loggers.wanna_logger import get_logger
 from wanna.core.models.docker import (
     DockerBuildConfigModel,
@@ -194,14 +195,8 @@ def _get_image(
         build_args = self.build_config.dict() if self.build_config else {}
 
         if isinstance(docker_image_model, NotebookReadyImageModel):
-            image_name = f"{self.wanna_project_name}/{docker_image_model.name}"
             tags = self.construct_image_tag(
-                registry=self.docker_registry,
-                project=self.docker_project_id,
-                repository=self.docker_repository,
-                image_name=image_name,
-                versions=[self.version, "latest"],
-                registry_suffix=self.docker_registry_suffix,
+                image_name=docker_image_model.name,
             )
             template_path = Path("notebook_template.Dockerfile")
             shutil.copy2(
@@ -214,14 +209,8 @@ def _get_image(
                 context_dir, file_path=file_path, tags=tags, docker_image_ref=docker_image_ref, **build_args
             )
         elif isinstance(docker_image_model, LocalBuildImageModel):
-            image_name = f"{self.wanna_project_name}/{docker_image_model.name}"
             tags = self.construct_image_tag(
-                registry=self.docker_registry,
-                project=self.docker_project_id,
-                repository=self.docker_repository,
-                image_name=image_name,
-                versions=[self.version, "latest"],
-                registry_suffix=self.docker_registry_suffix,
+                image_name=docker_image_model.name,
             )
             file_path = self.work_dir / docker_image_model.dockerfile
             context_dir = self.work_dir / docker_image_model.context_dir
@@ -338,29 +327,36 @@ def remove_image(image: Image, force=False, prune=True) -> None:
         """
         docker.image.remove(image, force=force, prune=prune)
 
-    @staticmethod
     def construct_image_tag(
-        registry: str,
-        project: str,
-        repository: str,
+        self,
         image_name: str,
-        registry_suffix: str,
-        versions: List[str] = ["latest"],
     ):
         """
         Construct full image tag.
         Args:
-            registry:
-            project:
-            repository:
             image_name:
-            versions:
 
         Returns:
             List of full image tag
         """
-
-        return [f"{registry}/{registry_suffix}{project}/{repository}/{image_name}:{version}" for version in versions]
+        versions = [self.version, "latest"]
+        return [
+            f"{self.docker_registry}/{self.docker_registry_suffix}{self.docker_project_id}/"
+            f"{self.docker_repository}/{self.wanna_project_name}/{image_name}:{version}"
+            for version in versions
+        ]
+
+    def build_container_and_get_image_url(self, docker_image_ref: str, push_mode: PushMode = PushMode.all) -> str:
+        if push_mode == PushMode.quick:
+            docker_image_model = self.find_image_model_by_name(docker_image_ref)
+            tags = self.construct_image_tag(image_name=docker_image_model.name)
+            image_url = tags[0]
+        else:
+            image_tag = self.get_image(docker_image_ref=docker_image_ref)
+            if len(image_tag) > 1 and image_tag[1]:
+                self.push_image(image_tag[1])
+            image_url = image_tag[2]
+        return image_url
 
     @staticmethod
     def _jinja_render_dockerfile(