vantage6 5.0.0a36__py3-none-any.whl → 5.0.0a38__py3-none-any.whl

vantage6/cli/sandbox/config/core.py

@@ -0,0 +1,300 @@
+from pathlib import Path
+
+from vantage6.common.globals import HTTP_LOCALHOST, InstanceType, Ports
+
+from vantage6.cli.common.new import new
+from vantage6.cli.sandbox.config.base import BaseSandboxConfigManager
+
+
+class CoreSandboxConfigManager(BaseSandboxConfigManager):
+    """
+    Class to store the sandbox configurations.
+
+    Parameters
+    ----------
+    server_name : str
+        Name of the server.
+    server_port : int
+        Port of the server.
+    ui_port : int
+        Port of the UI.
+    algorithm_store_port : int
+        Port of the algorithm store.
+    server_image : str | None
+        Image of the server.
+    store_image : str | None
+        Image of the algorithm store.
+    ui_image : str | None
+        Image of the UI.
+    extra_server_config : Path | None
+        Path to the extra server configuration file.
+    extra_store_config : Path | None
+        Path to the extra algorithm store configuration file.
+    context : str | None
+        Kubernetes context.
+    namespace : str | None
+        Kubernetes namespace.
+    k8s_node_name : str
+        Kubernetes node name.
+    custom_data_dir : Path | None
+        Path to the custom data directory. Useful on WSL because of mount issues for
+        default directories.
+    """
+
+    def __init__(
+        self,
+        server_name: str,
+        server_port: int,
+        ui_port: int,
+        algorithm_store_port: int,
+        server_image: str | None,
+        store_image: str | None,
+        ui_image: str | None,
+        extra_server_config: Path | None,
+        extra_store_config: Path | None,
+        extra_auth_config: Path | None,
+        context: str,
+        namespace: str,
+        k8s_node_name: str,
+        custom_data_dir: Path | None = None,
+    ) -> None:
+        super().__init__(server_name, custom_data_dir)
+
+        self.server_port = server_port
+        self.ui_port = ui_port
+        self.algorithm_store_port = algorithm_store_port
+        self.server_image = server_image
+        self.store_image = store_image
+        self.ui_image = ui_image
+        self.extra_server_config = extra_server_config
+        self.extra_store_config = extra_store_config
+        self.extra_auth_config = extra_auth_config
+        self.context = context
+        self.namespace = namespace
+
+        self.server_config_file = None
+        self.store_config_file = None
+        self.auth_config_file = None
+        self.k8s_node_name = k8s_node_name
+
+    def generate_server_configs(self) -> None:
+        """Generates the demo network."""
+
+        self._create_auth_config()
+
+        self._create_vserver_config()
+
+        self._create_algo_store_config()
+
+    def __server_config_return_func(self, extra_config: dict, data_dir: Path) -> dict:
+        """
+        Return a dict with server configuration values to be used in creating the
+        config file.
+
+        Parameters
+        ----------
+        extra_config : dict
+            Extra configuration (parsed from YAML) to be added to the server
+            configuration.
+        data_dir : Path
+            Path to the data directory.
+
+        Returns
+        -------
+        dict
+            Dictionary with server configuration values.
+        """
+        store_service = (
+            f"vantage6-{self.server_name}-store-user-algorithm-store-store-service"
+        )
+        store_address = (
+            f"http://{store_service}.{self.namespace}.svc.cluster.local:"
+            f"{Ports.DEV_ALGO_STORE}"
+        )
+        config = {
+            "server": {
+                "baseUrl": f"{HTTP_LOCALHOST}:{self.server_port}",
+                # TODO: v5+ set to latest v5 image
+                # TODO make this configurable
+                "image": (
+                    self.server_image
+                    or "harbor2.vantage6.ai/infrastructure/server:5.0.0a37"
+                ),
+                "algorithm_stores": [
+                    {
+                        "name": "Local store",
+                        "url": store_address,
+                        "api_path": "/store",
+                    }
+                ],
+                "logging": {
+                    "level": "DEBUG",
+                },
+                "jwt": {
+                    "secret": "development-constant-secret!",
+                },
+                "dev": {
+                    "host_uri": (
+                        "host.docker.internal"
+                        if self.k8s_node_name == "docker-desktop"
+                        else "172.17.0.1"
+                    ),
+                    "store_address": store_address,
+                },
+                "keycloakUrl": (
+                    f"http://vantage6-{self.server_name}-auth-user-auth-keycloak."
+                    f"{self.namespace}.svc.cluster.local"
+                ),
+            },
+            "rabbitmq": {},
+            "database": {
+                "volumePath": str(data_dir),
+                "k8sNodeName": self.k8s_node_name,
+            },
+            "ui": {
+                "port": self.ui_port,
+                # TODO: v5+ set to latest v5 image
+                # TODO: make this configurable
+                "image": (
+                    self.ui_image or "harbor2.vantage6.ai/infrastructure/ui:5.0.0a37"
+                ),
+            },
+        }
+
+        # merge the extra config with the server config
+        if extra_config is not None:
+            config.update(extra_config)
+
+        return config
+
+    def _create_vserver_config(self) -> None:
+        """Creates server configuration file (YAML)."""
+
+        data_dir = self._create_and_get_data_dir(instance_type=InstanceType.SERVER)
+
+        extra_config = self._read_extra_config_file(self.extra_server_config)
+        if self.ui_image is not None:
+            ui_config = extra_config.get("ui", {}) if extra_config is not None else {}
+            ui_config["image"] = self.ui_image
+            extra_config["ui"] = ui_config
+
+        # Create the server config file
+        self.server_config_file = new(
+            config_producing_func=self.__server_config_return_func,
+            config_producing_func_args=(extra_config, data_dir),
+            name=self.server_name,
+            system_folders=False,
+            namespace=self.namespace,
+            context=self.context,
+            type_=InstanceType.SERVER,
+            is_sandbox=True,
+        )
+
+    def _create_algo_store_config(self) -> None:
+        """Create algorithm store configuration file (YAML)."""
+
+        extra_config = self._read_extra_config_file(self.extra_store_config)
+
+        data_dir = self._create_and_get_data_dir(InstanceType.ALGORITHM_STORE)
+
+        self.store_config_file = new(
+            config_producing_func=self.__algo_store_config_return_func,
+            config_producing_func_args=(extra_config, data_dir),
+            name=f"{self.server_name}-store",
+            system_folders=False,
+            namespace=self.namespace,
+            context=self.context,
+            type_=InstanceType.ALGORITHM_STORE,
+            is_sandbox=True,
+        )
+
+    def __algo_store_config_return_func(
+        self, extra_config: dict, data_dir: Path
+    ) -> dict:
+        """
+        Return a dict with algorithm store configuration values to be used in creating
+        the config file.
+
+        Returns
+        -------
+        dict
+            Dictionary with algorithm store configuration values.
+        """
+        config = {
+            "store": {
+                "internal": {
+                    "port": self.algorithm_store_port,
+                },
+                "logging": {
+                    "level": "DEBUG",
+                },
+                "vantage6ServerUri": f"{HTTP_LOCALHOST}:{self.server_port}",
+                "image": (
+                    self.store_image
+                    or "harbor2.vantage6.ai/infrastructure/algorithm-store:5.0.0a37"
+                ),
+                "keycloakUrl": (
+                    f"http://vantage6-{self.server_name}-auth-user-auth-keycloak."
+                    f"{self.namespace}.svc.cluster.local"
+                ),
+                "policies": {
+                    "allowLocalhost": True,
+                    "assignReviewOwnAlgorithm": True,
+                },
+                "dev": {
+                    "host_uri": (
+                        "host.docker.internal"
+                        if self.k8s_node_name == "docker-desktop"
+                        else "172.17.0.1"
+                    ),
+                    "disable_review": True,
+                    "review_own_algorithm": True,
+                },
+            },
+            "database": {
+                "volumePath": str(data_dir),
+                "k8sNodeName": self.k8s_node_name,
+            },
+        }
+
+        # merge the extra config with the algorithm store config
+        if extra_config is not None:
+            config.update(extra_config)
+
+        return config
+
+    def _create_auth_config(self) -> None:
+        """Create auth configuration file (YAML)."""
+        self.auth_config_file = new(
+            config_producing_func=self.__auth_config_return_func,
+            config_producing_func_args=(self.extra_auth_config,),
+            name=f"{self.server_name}-auth",
+            system_folders=False,
+            namespace=self.namespace,
+            context=self.context,
+            type_=InstanceType.AUTH,
+            is_sandbox=True,
+        )
+
+    def __auth_config_return_func(self, extra_config: dict) -> dict:
+        """
+        Return a dict with auth configuration values to be used in creating the
+        config file.
+        """
+
+        config = {
+            "keycloak": {
+                "production": False,
+                "no_password_update_required": True,
+                "redirectUris": [
+                    f"{HTTP_LOCALHOST}:7600",
+                    f"{HTTP_LOCALHOST}:7681",
+                ],
+            },
+        }
+
+        # merge the extra config with the auth config
+        if extra_config is not None:
+            config.update(extra_config)
+
+        return config

vantage6/cli/sandbox/config/node.py

@@ -0,0 +1,311 @@
+from dataclasses import dataclass
+from importlib import resources as impresources
+from pathlib import Path
+
+import pandas as pd
+from colorama import Fore, Style
+
+from vantage6.common import error, info
+from vantage6.common.globals import InstanceType
+
+import vantage6.cli.sandbox.data as node_datafiles_dir
+from vantage6.cli.common.new import new
+from vantage6.cli.context.node import NodeContext
+from vantage6.cli.globals import (
+    DefaultDatasets,
+)
+from vantage6.cli.sandbox.config.base import BaseSandboxConfigManager
+from vantage6.cli.sandbox.populate.helpers.utils import replace_wsl_path
+
+
+@dataclass
+class NodeDataset:
+    label: str
+    path: Path
+
+
+class NodeSandboxConfigManager(BaseSandboxConfigManager):
+    """
+    Class to store the node sandbox configurations.
+
+    Parameters
+    ----------
+    server_name : str
+        Name of the server.
+    api_keys : list[str]
+        List of API keys.
+    node_names : list[str]
+        List of node names.
+    server_port : int
+        Port of the server.
+    node_image : str | None
+        Image of the node.
+    extra_node_config : Path | None
+        Path to the extra node configuration file.
+    extra_dataset : NodeDataset | None
+        List of tuples with the label and path to the dataset file.
+    context : str
+        Kubernetes context.
+    namespace : str
+        Kubernetes namespace.
+    k8s_node_name : str
+        Kubernetes node name.
+    custom_data_dir : Path | None
+        Path to the custom data directory. Useful on WSL because of mount issues for
+        default directories.
+    """
+
+    def __init__(
+        self,
+        server_name: str,
+        api_keys: list[str],
+        node_names: list[str],
+        server_port: int,
+        node_image: str | None,
+        extra_node_config: Path | None,
+        extra_dataset: NodeDataset | None,
+        context: str,
+        namespace: str,
+        k8s_node_name: str,
+        custom_data_dir: Path | None,
+    ) -> None:
+        super().__init__(server_name, custom_data_dir)
+        self.api_keys = api_keys
+        self.node_names = node_names
+        self.num_nodes = len(api_keys)
+        self.server_port = server_port
+        self.node_image = node_image
+        self.extra_node_config = extra_node_config
+        if extra_dataset:
+            self.node_datasets = [extra_dataset]
+        else:
+            self.node_datasets = []
+        self.context = context
+        self.namespace = namespace
+        self.k8s_node_name = k8s_node_name
+
+        self.node_configs = []
+        self.node_config_files = []
+        self.node_config_names = []
+        self.extra_config = None
+
+    def generate_node_configs(self) -> None:
+        """
+        Generates ``num_nodes`` node configuration files.
+        """
+        node_data_files = []
+        self.extra_config = self._read_extra_config_file(self.extra_node_config)
+
+        data_directory = impresources.files(node_datafiles_dir)
+
+        # Add default datasets to the list of dataset provided
+        for default_dataset in DefaultDatasets:
+            # note that the label is the name of the dataset with the underscores
+            # replace with hyphens, to make it valid for k8s
+            self.node_datasets.append(
+                NodeDataset(
+                    label=default_dataset.name.lower().replace("_", "-"),
+                    path=data_directory / default_dataset.value,
+                )
+            )
+
+        # Check for duplicate dataset labels
+        seen_labels = set()
+        duplicates = [
+            label
+            for label in [dataset.label for dataset in self.node_datasets]
+            if (label in seen_labels or seen_labels.add(label))
+        ]
+
+        if len(duplicates) > 0:
+            error(
+                f"Duplicate dataset labels found: {duplicates}. "
+                f"Please make sure all dataset labels are unique."
+            )
+            exit(1)
+
+        # create the data files for the nodes and get the path and label for each
+        # dataset
+        for dataset in self.node_datasets:
+            node_data_files.append(self._create_node_data_files(dataset))
+
+        for idx, api_key in enumerate(self.api_keys):
+            config = {
+                "org_id": idx + 1,
+                "api_key": api_key,
+                "node_name": self.node_names[idx],
+            }
+            config_file = self._create_node_config_file(
+                config,
+                [files[idx] for files in node_data_files],
+            )
+            self.node_configs.append(config)
+            self.node_config_files.append(config_file)
+        info(
+            f"Created {Fore.GREEN}{len(self.node_config_files)}{Style.RESET_ALL} node "
+            f"configuration(s), attaching them to {Fore.GREEN}{self.server_name}"
+            f"{Style.RESET_ALL}."
+        )
+
+    def _create_node_data_files(
+        self, node_dataset: NodeDataset
+    ) -> list[tuple[str, Path]]:
+        """
+        Create data files for nodes.
+
+        Parameters
+        ----------
+        node_dataset : NodeDataset
+            Tuple with the label and path to the dataset file.
+
+        Returns
+        -------
+        list[tuple[str, Path]]
+            List of tuples with the label and path to the dataset file.
+        """
+        info(
+            f"Creating data files using dataset '{node_dataset.label}' for "
+            f"{self.num_nodes} nodes"
+        )
+        data_files = []
+        full_df = pd.read_csv(node_dataset.path)
+        length_df = len(full_df)
+        for i in range(self.num_nodes):
+            node_name = f"{self.server_name}_node_{i + 1}"
+            path_to_dev_dir = self._create_and_get_data_dir(
+                InstanceType.NODE, is_data_folder=False
+            )
+
+            # Split the data over the nodes
+            start = i * length_df // self.num_nodes
+            end = (i + 1) * length_df // self.num_nodes
+            data = full_df[start:end]
+            data_file = (
+                replace_wsl_path(path_to_dev_dir, to_mnt_wsl=True)
+                / f"df_{node_dataset.label}_{node_name}.csv"
+            )
+
+            # write data to file
+            data.to_csv(data_file, index=False)
+            data_files.append(
+                (node_dataset.label, replace_wsl_path(data_file, to_mnt_wsl=False))
+            )
+        return data_files
+
+    def _create_node_config_file(
+        self, config: dict, datasets: list[tuple[str, Path]]
+    ) -> Path:
+        """
+        Create a node configuration file (YAML).
+
+        Parameters
+        ----------
+        config : dict
+            Configuration dictionary.
+        datasets : list[tuple[str, Path]]
+            List of tuples with the label and path to the dataset file.
+
+        Returns
+        -------
+        Path
+            Path to the node configuration file.
+        """
+        node_name = config["node_name"]
+        config_name = f"{self.server_name}-{node_name}"
+
+        path_to_data_dir = self._create_and_get_data_dir(
+            InstanceType.NODE, is_data_folder=True, node_name=node_name
+        )
+
+        # delete old node config if it exists
+        NodeContext.remove_config_file_if_exists(
+            InstanceType.NODE, config_name, False, is_sandbox=True
+        )
+
+        # create new node config
+        node_config = new(
+            config_producing_func=self.__node_config_return_func,
+            config_producing_func_args=(
+                config,
+                self.extra_config,
+                datasets,
+                path_to_data_dir,
+            ),
+            name=config_name,
+            system_folders=False,
+            namespace=self.namespace,
+            context=self.context,
+            type_=InstanceType.NODE,
+            is_sandbox=True,
+        )
+
+        self.node_config_names.append(config_name)
+
+        return node_config
+
+    def __node_config_return_func(
+        self,
+        node_specific_config: dict,
+        extra_config: dict,
+        datasets: list[tuple[str, Path]],
+        path_to_data_dir: Path,
+    ) -> dict:
+        """
+        Return a dict with node configuration values to be used in creating the
+        config file.
+        """
+        config = {
+            "node": {
+                "proxyPort": 7676 + int(node_specific_config["org_id"]),
+                "apiKey": node_specific_config["api_key"],
+                "name": node_specific_config["node_name"],
+                "image": (
+                    self.node_image
+                    # TODO v5+ update
+                    or "harbor2.vantage6.ai/infrastructure/node:5.0.0a37"
+                ),
+                "logging": {
+                    "level": "DEBUG",
+                    "file": f"{node_specific_config['node_name']}.log",
+                },
+                "keycloakUrl": (
+                    f"http://vantage6-{self.server_name}-auth-user-auth-keycloak."
+                    f"{self.namespace}.svc.cluster.local"
+                ),
+                "persistence": {
+                    "tasks": {
+                        "hostPath": str(path_to_data_dir),
+                        "size": "1Gi",
+                    },
+                    "database": {
+                        "size": "1Gi",
+                    },
+                },
+                "k8sNodeName": self.k8s_node_name,
+                "databases": {
+                    "fileBased": [
+                        {
+                            "name": dataset[0],
+                            "uri": dataset[1],
+                            "type": "csv",
+                            "volumePath": Path(dataset[1]).parent,
+                            "originalName": dataset[0],
+                        }
+                        for dataset in [datasets[0]]
+                    ]
+                },
+                "server": {
+                    "url": (
+                        f"http://vantage6-{self.server_name}-user-server-"
+                        f"vantage6-server-service.{self.namespace}.svc.cluster.local"
+                    ),
+                    "port": self.server_port,
+                },
+            },
+        }
+
+        # merge the extra config with the node config
+        if extra_config is not None:
+            config.update(extra_config)
+
+        return config