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

vantage6/cli/sandbox/new.py

@@ -0,0 +1,207 @@
+from pathlib import Path
+
+import click
+from colorama import Fore, Style
+
+from vantage6.common import error
+from vantage6.common.globals import Ports
+
+from vantage6.cli.common.utils import select_context_and_namespace
+from vantage6.cli.context.server import ServerContext
+from vantage6.cli.sandbox.config.core import CoreSandboxConfigManager
+from vantage6.cli.sandbox.start import execute_sandbox_start
+from vantage6.cli.server.common import get_server_context
+from vantage6.cli.utils import prompt_config_name
+
+
+@click.command()
+@click.option(
+    "-n", "--name", default=None, type=str, help="Name for your development setup"
+)
+@click.option(
+    "--num-nodes",
+    type=int,
+    default=3,
+    help="Generate this number of nodes in the development network",
+)
+@click.option(
+    "-p",
+    "--server-port",
+    type=int,
+    default=Ports.DEV_SERVER.value,
+    help=f"Port to run the server on. Default is {Ports.DEV_SERVER}.",
+)
+@click.option(
+    "--ui-port",
+    type=int,
+    default=Ports.DEV_UI.value,
+    help=f"Port to run the UI on. Default is {Ports.DEV_UI}.",
+)
+@click.option(
+    "--algorithm-store-port",
+    type=int,
+    default=Ports.DEV_ALGO_STORE.value,
+    help=f"Port to run the algorithm store on. Default is {Ports.DEV_ALGO_STORE}.",
+)
+@click.option(
+    "--server-image",
+    type=str,
+    default=None,
+    help="Server docker image to use when setting up resources for "
+    "the development server",
+)
+@click.option(
+    "--ui-image",
+    type=str,
+    default=None,
+    help="UI docker image to specify in configuration files. Will be used on startup of"
+    " the network",
+)
+@click.option(
+    "--store-image",
+    type=str,
+    default=None,
+    help="Algorithm store docker image to use when setting up resources for "
+    "the development algorithm store",
+)
+@click.option(
+    "--node-image",
+    type=str,
+    default=None,
+    help="Node docker image to use when setting up resources for the development node",
+)
+@click.option(
+    "--extra-server-config",
+    type=click.Path(exists=True),
+    default=None,
+    help="YAML File with additional server "
+    "configuration. This will be appended to the server "
+    "configuration file",
+)
+@click.option(
+    "--extra-node-config",
+    type=click.Path("rb"),
+    default=None,
+    help="YAML File with additional node configuration. This will be"
+    " appended to each of the node configuration files",
+)
+@click.option(
+    "--extra-store-config",
+    type=click.Path("rb"),
+    default=None,
+    help="YAML File with additional algorithm store configuration. This will be"
+    " appended to the algorithm store configuration file",
+)
+@click.option(
+    "--extra-auth-config",
+    type=click.Path("rb"),
+    default=None,
+    help="YAML File with additional auth configuration. This will be"
+    " appended to the auth configuration file",
+)
+@click.option(
+    "--add-dataset",
+    type=(str, click.Path()),
+    default=None,
+    multiple=True,
+    help="Add a dataset to the nodes. The first argument is the label of the database, "
+    "the second is the path to the dataset file.",
+)
+@click.option("--context", default=None, help="Kubernetes context to use")
+@click.option("--namespace", default=None, help="Kubernetes namespace to use")
+@click.option(
+    "--k8s-node-name", default="docker-desktop", help="Kubernetes node name to use"
+)
+@click.option(
+    "--data-dir",
+    type=click.Path(exists=True),
+    default=None,
+    help="Path to a custom data directory to use. This option is especially useful "
+    "on WSL because of mount issues for default directories",
+)
+@click.option(
+    "--local-chart-dir",
+    type=click.Path(exists=True),
+    default=None,
+    help="Local chart repository to use.",
+)
+@click.pass_context
+def cli_new_sandbox(
+    click_ctx: click.Context,
+    name: str,
+    num_nodes: int,
+    server_port: int,
+    ui_port: int,
+    algorithm_store_port: int,
+    server_image: str | None,
+    ui_image: str | None,
+    store_image: str | None,
+    node_image: str | None,
+    extra_server_config: Path | None,
+    extra_node_config: Path | None,
+    extra_store_config: Path | None,
+    extra_auth_config: Path | None,
+    add_dataset: tuple[str, Path] | None,
+    context: str | None,
+    namespace: str | None,
+    k8s_node_name: str,
+    data_dir: str | None,
+    local_chart_dir: Path | None,
+) -> None:
+    """
+    Create a sandbox environment.
+    """
+
+    # Prompt for the k8s namespace and context
+    context, namespace = select_context_and_namespace(
+        context=context,
+        namespace=namespace,
+    )
+
+    if data_dir is not None:
+        data_dir = Path(data_dir)
+        if not data_dir.exists():
+            error(f"Data directory {data_dir} does not exist!")
+            exit(1)
+
+    server_name = prompt_config_name(name)
+    if ServerContext.config_exists(server_name, False, is_sandbox=True):
+        error(f"Configuration {Fore.RED}{server_name}{Style.RESET_ALL} already exists!")
+        exit(1)
+
+    sb_config_manager = CoreSandboxConfigManager(
+        server_name=server_name,
+        server_port=server_port,
+        ui_port=ui_port,
+        algorithm_store_port=algorithm_store_port,
+        server_image=server_image,
+        ui_image=ui_image,
+        store_image=store_image,
+        extra_server_config=extra_server_config,
+        extra_store_config=extra_store_config,
+        extra_auth_config=extra_auth_config,
+        context=context,
+        namespace=namespace,
+        k8s_node_name=k8s_node_name,
+        custom_data_dir=data_dir,
+    )
+
+    sb_config_manager.generate_server_configs()
+
+    ctx = get_server_context(server_name, False, ServerContext, is_sandbox=True)
+
+    execute_sandbox_start(
+        click_ctx=click_ctx,
+        ctx=ctx,
+        server_name=server_name,
+        context=context,
+        namespace=namespace,
+        num_nodes=num_nodes,
+        initialize=True,
+        node_image=node_image,
+        k8s_node_name=k8s_node_name,
+        extra_node_config=extra_node_config,
+        add_dataset=add_dataset,
+        custom_data_dir=data_dir,
+        local_chart_dir=local_chart_dir,
+    )

vantage6/cli/sandbox/populate/__init__.py

@@ -0,0 +1,173 @@
+"""
+Script to populate the server with basic fixtures.
+"""
+
+import time
+import traceback
+from logging import info
+
+from vantage6.common import error
+
+from vantage6.client import Client
+from vantage6.client.utils import LogLevel
+
+from vantage6.cli.sandbox.populate.helpers.connect_store import connect_store
+from vantage6.cli.sandbox.populate.helpers.delete_fixtures import delete_fixtures
+from vantage6.cli.sandbox.populate.helpers.load_fixtures import create_fixtures
+from vantage6.cli.sandbox.populate.helpers.utils import NodeConfigCreationDetails
+
+
+def populate_server_dev(
+    server_url: str,
+    auth_url: str,
+    number_of_nodes: int,
+    node_config_creation_details: NodeConfigCreationDetails,
+) -> str | None:
+    """
+    Populate the server with basic fixtures.
+
+    Parameters
+    ----------
+    server_url : str
+        The URL of the server to connect to.
+    auth_url : str
+        The URL of the auth service to connect to.
+    number_of_nodes : int
+        The number of nodes to create.
+    node_config_creation_details : NodeConfigCreationDetails | None
+        The details of the node config creation. If None, the node configs will not be
+        created.
+    clear_dev_folders : bool
+        Whether to clear everything inthe dev directory before creating the new
+        resources.
+    create_node_configs: bool
+        Whether to create the node configs. Default is True.
+
+    Returns
+    -------
+    str | None
+        A report of the creation process, or None if an error occurred.
+
+    Raises
+    ------
+    Exception
+        If connection/authentication to the server fails.
+    """
+    client = _initalize_client(server_url, auth_url)
+
+    # Create new resources in the server
+    # Delete existing resources in the server first, before creating new ones.
+    try:
+        report_deletion = delete_fixtures(client)
+        report_creation = create_fixtures(
+            client,
+            number_of_nodes,
+            node_config_creation_details=node_config_creation_details,
+            clear_dev_folders=True,
+        )
+        report_store = connect_store(client)
+
+        return report_deletion + "\n" + report_creation + "\n" + report_store
+
+    except Exception:
+        error("=" * 80)
+        error("Failed to populate server")
+        error(traceback.format_exc())
+        error("=" * 80)
+
+
+def populate_server_sandbox(
+    server_url: str,
+    auth_url: str,
+    number_of_nodes: int,
+) -> dict:
+    """
+    Populate sandbox server with basic resources.
+
+    Parameters
+    ----------
+    server_url : str
+        The URL of the server to connect to.
+    auth_url : str
+        The URL of the auth service to connect to.
+    number_of_nodes : int
+        The number of nodes to create.
+
+    Returns
+    -------
+    dict
+        A dictionary containing the report of the creation process.
+
+    Raises
+    ------
+    Exception
+        If connection/authentication to the server fails.
+    """
+    client = _initalize_client(server_url, auth_url)
+
+    try:
+        delete_fixtures(client)
+        report_creation = create_fixtures(
+            client,
+            number_of_nodes,
+            clear_dev_folders=False,
+            return_as_dict=True,
+        )
+        connect_store(client)
+    except Exception:
+        error("=" * 80)
+        error("Failed to populate server")
+        error(traceback.format_exc())
+        error("=" * 80)
+        exit(1)
+
+    # return the details of the created nodes so that config files can be created
+    return report_creation["nodes"]["created"]
+
+
+def _initalize_client(server_url, auth_url) -> Client:
+    """
+    Initialize an authenticated client to the server.
+
+    The server may not be ready yet, so we retry until it is.
+
+    Parameters
+    ----------
+    server_url : str
+        The URL of the server to connect to.
+    auth_url : str
+        The URL of the auth service to connect to.
+
+    Returns
+    -------
+    Client
+        An authenticated client to the server.
+    """
+    client = Client(
+        server_url=server_url,
+        auth_url=auth_url,
+        log_level=LogLevel.ERROR,
+    )
+
+    info("Waiting for authentication...")
+    max_attempts = 120
+    attempt = 1
+
+    while attempt <= max_attempts:
+        try:
+            print(".", end="", flush=True)
+            client.authenticate()
+            info("Successfully authenticated with server!")
+            break
+        except Exception as e:
+            if attempt == max_attempts:
+                error(
+                    f"Failed to authenticate after {max_attempts} attempts. "
+                    "Server may not be online."
+                )
+                raise e
+
+            time.sleep(5)
+            attempt += 1
+
+    return client

vantage6/cli/sandbox/populate/helpers/connect_store.py

@@ -0,0 +1,203 @@
+"""
+Development script to connect the server to the local store.
+"""
+
+import time
+from http import HTTPStatus
+
+import requests
+
+from vantage6.common import error, info
+from vantage6.common.enum import AlgorithmStepType
+
+from vantage6.client import Client
+
+
+def _wait_for_store_to_be_online(
+    local_store_url: str, local_store_api_path: str
+) -> None:
+    """
+    Wait for the store to be online.
+
+    Parameters
+    ---------
+    client: Client
+        The client to use to connect to the server.
+    local_store_url: str
+        The URL of the local store.
+    local_store_api_path: str
+        The API path of the local store.
+    """
+    info(
+        f"Waiting for store to be online at {local_store_url}{local_store_api_path}..."
+    )
+    max_retries = 100
+    wait_time = 3
+    ready = False
+    for _ in range(max_retries):
+        try:
+            result = requests.get(f"{local_store_url}{local_store_api_path}/version")
+            if result.status_code == HTTPStatus.OK:
+                ready = True
+                break
+        except Exception:
+            info(f"Store not ready yet, waiting {wait_time} seconds...")
+            time.sleep(wait_time)
+
+    if not ready:
+        error("Store did not become ready in time. Exiting...")
+        exit(1)
+    else:
+        info("Store is online!")
+
+
+def connect_store(client: Client) -> str:
+    """
+    Connect the server to the local store.
+
+    Parameters
+    ---------
+    client: Client
+        The client to use to connect to the server.
+    dev_dir: Path
+        The directory to use to store the development data.
+    """
+
+    existing_stores = client.store.list().get("data", [])
+    summary = "=== Store Connection Summary ===\n"
+
+    # URL should be retrieved from the store, see issue:
+    # https://github.com/vantage6/vantage6/issues/1824
+    local_store_url = "http://localhost:7602"
+    local_store_api_path = "/store"
+    client.store.store_id = 1
+
+    _wait_for_store_to_be_online(local_store_url, local_store_api_path)
+
+    # note that the store is already coupled to the server in the sandbox/devspace
+    # config. To find the store, either check that it is a localhost URL or that it
+    # contains "svc.cluster.local" (which is for local k8s services)
+    try:
+        store = next(s for s in existing_stores if s["url"] == local_store_url)
+        client.store.set(store["id"])
+    except StopIteration:
+        error(
+            "Local algorithm store not found. Please register its resources manually."
+        )
+        return
+
+    # register also the other users in the local store
+    users_in_store = client.store.user.list()["data"]
+    all_users = client.user.list()["data"]
+    for user in all_users:
+        if user["keycloak_id"] not in [u["keycloak_id"] for u in users_in_store]:
+            summary += f"Registering user {user['username']} in local store\n"
+            client.store.user.register(username=user["username"], roles=[1])
+
+    # Remove existing algorithm
+    # This is broken, see issue: https://github.com/vantage6/vantage6/issues/1824
+    algorithms = client.algorithm.list(name="session basic example")["data"]
+    if len(algorithms) > 0:
+        summary += f"Removing existing algorithm {algorithms[0]['name']}\n"
+        client.algorithm.delete(id_=algorithms[0]["id"])
+
+    # Download the algorithm store from the github repo. This ensures that the data is
+    # always up to date.
+    try:
+        url = (
+            "https://raw.githubusercontent.com/vantage6/v6-session-basics/"
+            "main/algorithm_store.json"
+        )
+        response = requests.get(url, timeout=10)
+        response.raise_for_status()
+        algorithm_json = response.json()
+        function_metadata = algorithm_json["functions"]
+    except requests.exceptions.RequestException as e:
+        # Fallback to local file if download fails
+        print(f"Warning: Could not download algorithm store from GitHub: {e}")
+        print("Not putting the algorithm in the store.")
+        function_metadata = []
+
+    summary += "Creating Session Basics algorithm\n"
+    client.algorithm.create(
+        name="Session Basics",
+        description="A set of basic algorithms for a session management",
+        image="harbor2.vantage6.ai/algorithms/session-basics:latest",
+        vantage6_version="5.0.0",
+        code_url="https://github.com/vantage6-ai/v6-session-basics",
+        partitioning="horizontal",
+        functions=function_metadata or [],
+    )
+
+    # TODO: v5+ get this json data by downloading it from the github repo - that ensures
+    # that the data is always up to date.
+    summary += "Creating Network Diagnostics algorithm\n"
+    client.algorithm.create(
+        name="Network Diagnostics",
+        description=(
+            "Functions to diagnose network policies, that is, to what extent the "
+            "jobs running on a V6 Node could have access to the internal k8s network "
+            "or to the outside world."
+        ),
+        image="ghcr.io/hcadavid/v6-sessions-k8s-diagnostics:latest",
+        vantage6_version="5.0.0",
+        code_url="https://github.com/hcadavid/v6-sessions-k8s-diagnostics",
+        partitioning="horizontal",
+        functions=[
+            {
+                "name": "network_status",
+                "display_name": "Get the network diagnostics performed on a given node",
+                "standalone": True,
+                "description": "Get the network diagnostics performed on a given node",
+                "ui_visualizations": [],
+                "step_type": AlgorithmStepType.FEDERATED_COMPUTE.value,
+                "arguments": [
+                    {
+                        "has_default_value": True,
+                        "name": "sleep_time",
+                        "display_name": "Delay ",
+                        "description": (
+                            "Delay in seconds before finishing the the diagnostics "
+                            "task (useful to keep the POD running and access it to "
+                            "perform further tests) "
+                        ),
+                        "type": "integer",
+                        "default_value": "1",
+                        "is_frontend_only": False,
+                    }
+                ],
+                "databases": [],
+            },
+            {
+                "name": "central_network_diagnostics",
+                "display_name": (
+                    "Get the network diagnostics performed on all the nodes"
+                ),
+                "standalone": True,
+                "description": (
+                    "Get the network diagnostics performed on all the nodes within a "
+                    "collaboration"
+                ),
+                "ui_visualizations": [],
+                "step_type": AlgorithmStepType.CENTRAL_COMPUTE.value,
+                "arguments": [
+                    {
+                        "has_default_value": True,
+                        "name": "sleep_time",
+                        "display_name": "Delay ",
+                        "description": (
+                            "Delay in seconds before finishing the the diagnostics "
+                            "task (useful to keep the PODs running on all the nodes "
+                            "to access them and perform further inspections) "
+                        ),
+                        "type": "integer",
+                        "default_value": "1",
+                        "is_frontend_only": False,
+                    }
+                ],
+                "databases": [],
+            },
+        ],
+    )
+
+    return summary

vantage6/cli/sandbox/populate/helpers/delete_fixtures.py

@@ -0,0 +1,67 @@
+"""
+Development script to delete all entities from the server
+"""
+
+from vantage6.client import Client
+
+
+def delete_fixtures(client: Client) -> str:
+    """
+    Delete all entities from the server.
+
+    Arguments
+    ---------
+    client: Client
+        The client to use to delete the entities.
+    """
+
+    # Track deletion counts
+    deletion_counts = {}
+
+    # Delete sessions & dataframes
+    sessions = client.session.list(per_page=999)["data"]
+    for session in sessions:
+        client.session.delete(session["id"], delete_dependents=True)
+    deletion_counts["sessions"] = len(sessions)
+
+    # Delete tasks and nodes
+    for client_subclass_name in ("task", "node"):
+        client_subclass = getattr(client, client_subclass_name)
+        entities = client_subclass.list(per_page=999)["data"]
+        for entity in entities:
+            client_subclass.delete(entity["id"])
+        deletion_counts[f"{client_subclass_name}s"] = len(entities)
+
+    # Delete collaborations
+    collabs = client.collaboration.list(per_page=999, scope="global")["data"]
+    for collab in collabs:
+        client.collaboration.delete(collab["id"])
+    deletion_counts["collaborations"] = len(collabs)
+
+    # Delete users (excluding admin)
+    users = client.user.list(per_page=999)["data"]
+    deleted_users = 0
+    for user in users:
+        if user["username"] == "admin":
+            continue
+        client.user.delete(user["id"])
+        deleted_users += 1
+    deletion_counts["users"] = deleted_users
+
+    # Delete organizations (excluding the default organization with id 1)
+    orgs = client.organization.list(per_page=999)["data"]
+    deleted_orgs = 0
+    for org in orgs:
+        if org["id"] == 1:
+            continue
+        client.organization.delete(org["id"])
+        deleted_orgs += 1
+    deletion_counts["organizations"] = deleted_orgs
+
+    # Create summary string
+    summary = "=== Deletion Summary ===\n"
+    for entity_type, count in deletion_counts.items():
+        summary += f"Deleted {count} {entity_type}\n"
+    summary += "======================="
+    print(summary)
+    return summary