vantage6 4.10.2__py3-none-any.whl → 4.11.0__py3-none-any.whl

vantage6/cli/globals.py

@@ -42,6 +42,18 @@ DIAGNOSTICS_IMAGE = "harbor2.vantage6.ai/algorithms/diagnostic"
 # Address of community algorithm store
 COMMUNITY_STORE = "https://store.cotopaxi.vantage6.ai/api"
 
+DEFAULT_PROMETHEUS_IMAGE = "prom/prometheus"
+PROMETHEUS_CONFIG = "prometheus.yml"
+PROMETHEUS_DIR = "prometheus"
+
+
+# datasets included in the nodes of the dev network
+class DefaultDatasets(str, Enum):
+    """Enum containing default datasets"""
+
+    OLYMPIC_ATHLETES = "olympic_athletes_2016.csv"
+    KAPLAN_MEIER_TEST = "km_dataset.csv"
+
 
 class ServerType(str, Enum):
     """Enum containing server types"""

vantage6/cli/prometheus/monitoring_manager.py

@@ -0,0 +1,146 @@
+import yaml
+import docker
+from pathlib import Path
+
+from vantage6.common import info, error
+from vantage6.common.docker.network_manager import NetworkManager
+from vantage6.common.globals import DEFAULT_PROMETHEUS_EXPORTER_PORT
+from vantage6.cli.context.server import ServerContext
+from vantage6.cli.globals import (
+    DEFAULT_PROMETHEUS_IMAGE,
+    PROMETHEUS_CONFIG,
+)
+
+
+class PrometheusServer:
+    """
+    Manages the Prometheus Docker container
+    """
+
+    def __init__(
+        self, ctx: ServerContext, network_mgr: NetworkManager, image: str = None
+    ):
+        """
+        Initialize the PrometheusServer instance.
+
+        Parameters
+        ----------
+        ctx : ServerContext
+            The server context containing configuration and paths.
+        network_mgr : NetworkManager
+            The network manager responsible for managing Docker networks.
+        image : str, optional
+            The Docker image to use for the Prometheus container. If not provided,
+            the default Prometheus image will be used.
+        """
+        self.ctx = ctx
+        self.network_mgr = network_mgr
+        self.docker = docker.from_env()
+        self.image = image if image else DEFAULT_PROMETHEUS_IMAGE
+        self.config_file = Path(__file__).parent / PROMETHEUS_CONFIG
+        self.data_dir = self.ctx.prometheus_dir
+
+    def start(self):
+        """
+        Start a Docker container running Prometheus
+        """
+        self._prepare_config()
+
+        volumes = {
+            str(self.config_file): {
+                "bind": "/etc/prometheus/prometheus.yml",
+                "mode": "ro",
+            },
+            str(self.data_dir): {"bind": "/prometheus", "mode": "rw"},
+        }
+        ports = {"9090/tcp": 9090}
+
+        if self._is_container_running():
+            info("Prometheus is already running!")
+            return
+
+        self.docker.containers.run(
+            name=self.ctx.prometheus_container_name,
+            image=self.image,
+            volumes=volumes,
+            ports=ports,
+            detach=True,
+            restart_policy={"Name": "unless-stopped"},
+            network=self.network_mgr.network_name,
+        )
+        info("Prometheus container started successfully!")
+
+    def _prepare_config(self):
+        """
+        Prepare the Prometheus configuration and data directories
+        """
+        if not self.config_file.exists():
+            error(f"Prometheus configuration file {self.config_file} not found!")
+            raise FileNotFoundError(f"{self.config_file} not found!")
+
+        if not self.data_dir.exists():
+            self.data_dir.mkdir(parents=True, exist_ok=True)
+
+        self._update_prometheus_config()
+
+    def _update_prometheus_config(self):
+        """
+        Update the Prometheus configuration file with the server address.
+        """
+
+        try:
+            prometheus_exporter_port = self.ctx.config.get("prometheus", {}).get(
+                "exporter_port", DEFAULT_PROMETHEUS_EXPORTER_PORT
+            )
+            server_hostname = self.ctx.prometheus_container_name
+            server_address = f"{server_hostname}:{prometheus_exporter_port}"
+
+            info(
+                f"Using Docker container hostname '{server_hostname}' for Prometheus "
+                "target. Ensure Prometheus is in the same Docker network to resolve "
+                "this address."
+            )
+
+            with open(self.config_file, "r") as f:
+                config = yaml.safe_load(f)
+
+            job_name = "vantage6_server_metrics"
+            job_exists = any(
+                job.get("job_name") == job_name
+                for job in config.get("scrape_configs", [])
+            )
+
+            if not job_exists:
+                new_job = {
+                    "job_name": job_name,
+                    "static_configs": [{"targets": [server_address]}],
+                }
+                config.setdefault("scrape_configs", []).append(new_job)
+            else:
+                for job in config["scrape_configs"]:
+                    if job.get("job_name") == job_name:
+                        job["static_configs"] = [{"targets": [server_address]}]
+
+            with open(self.config_file, "w") as f:
+                yaml.dump(config, f)
+
+            info(f"Prometheus configuration updated with target: {server_address}")
+
+        except Exception as e:
+            error(f"Failed to update Prometheus configuration: {e}")
+            raise
+
+    def _is_container_running(self) -> bool:
+        """
+        Check if a Prometheus container is already running.
+
+        Returns
+        -------
+        bool
+            True if the Prometheus container is running, False otherwise.
+        """
+        try:
+            container = self.docker.containers.get(self.ctx.prometheus_container_name)
+            return container.status == "running"
+        except docker.errors.NotFound:
+            return False

vantage6/cli/prometheus/prometheus.yml

@@ -0,0 +1,5 @@
+global:
+  # How frequently to scrape targets by default.
+  scrape_interval: 1m
+
+scrape_configs: []

vantage6/cli/server/start.py

@@ -11,11 +11,12 @@ from vantage6.common.globals import (
     InstanceType,
 )
 
-from vantage6.common.globals import Ports
+from vantage6.common.globals import Ports, DEFAULT_PROMETHEUS_EXPORTER_PORT
 from vantage6.cli.context.server import ServerContext
 from vantage6.cli.rabbitmq.queue_manager import RabbitMQManager
 from vantage6.cli.server.common import stop_ui
 from vantage6.cli.common.decorator import click_insert_context
+from vantage6.cli.prometheus.monitoring_manager import PrometheusServer
 from vantage6.cli.common.start import (
     attach_logs,
     check_for_start,
@@ -49,6 +50,14 @@ from vantage6.cli.common.start import (
     "container - use in development only",
 )
 @click.option("--rabbitmq-image", default=None, help="RabbitMQ docker image to use")
+@click.option(
+    "--with-prometheus",
+    "start_prometheus",
+    flag_value=True,
+    default=False,
+    help="Start Prometheus monitoring as a local container",
+)
+@click.option("--prometheus-image", default=None, help="Prometheus docker image to use")
 @click.option(
     "--keep/--auto-remove",
     default=False,
@@ -75,6 +84,8 @@ def cli_server_start(
     ui_port: int,
     start_rabbitmq: bool,
     rabbitmq_image: str,
+    start_prometheus: bool,
+    prometheus_image: str,
     keep: bool,
     mount_src: str,
     attach: bool,
@@ -135,6 +146,24 @@ def cli_server_start(
             "cannot be scaled horizontally!"
         )
 
+    if (
+        start_prometheus
+        or ctx.config.get("prometheus")
+        and ctx.config["prometheus"].get("start_with_server", False)
+    ):
+        info("Starting Prometheus container")
+        _start_prometheus(ctx, prometheus_image, server_network_mgr)
+    elif ctx.config.get("prometheus"):
+        info(
+            "Prometheus is provided in the config file as external service. "
+            "Assuming this service is up and running."
+        )
+    else:
+        warning(
+            "Monitoring is not set up! This means that the vantage6 server "
+            "cannot be monitored with Prometheus!"
+        )
+
     # start the UI if requested
     if start_ui or ctx.config.get("ui") and ctx.config["ui"].get("enabled"):
         _start_ui(docker_client, ctx, ui_port)
@@ -152,6 +181,9 @@ def cli_server_start(
 
     info("Run Docker container")
     port_ = str(port or ctx.config["port"] or Ports.DEV_SERVER.value)
+    prometheus_exporter_port = ctx.config.get("prometheus", {}).get(
+        "exporter_port", DEFAULT_PROMETHEUS_EXPORTER_PORT
+    )
     container = docker_client.containers.run(
         image,
         command=cmd,
@@ -162,7 +194,10 @@ def cli_server_start(
             "name": ctx.config_file_name,
         },
         environment=environment_vars,
-        ports={f"{internal_port}/tcp": (ip, port_)},
+        ports={
+            f"{internal_port}/tcp": (ip, port_),  # API port
+            f"{prometheus_exporter_port}/tcp": prometheus_exporter_port,
+        },
         name=ctx.docker_container_name,
         auto_remove=not keep,
         tty=True,
@@ -202,6 +237,27 @@ def _start_rabbitmq(
     rabbit_mgr.start()
 
 
+def _start_prometheus(
+    ctx: ServerContext, prometheus_image: str, network_mgr: NetworkManager
+) -> None:
+    """
+    Start the Prometheus container if it is not already running.
+
+    Parameters
+    ----------
+    ctx : ServerContext
+        Server context object
+    prometheus_image : str
+        Prometheus image to use
+    network_mgr : NetworkManager
+        Network manager object
+    """
+    prometheus_server = PrometheusServer(
+        ctx=ctx, network_mgr=network_mgr, image=prometheus_image
+    )
+    prometheus_server.start()
+
+
 def _start_ui(client: DockerClient, ctx: ServerContext, ui_port: int) -> None:
     """
     Start the UI container.

vantage6/cli/server/stop.py

@@ -117,3 +117,10 @@ def _stop_server_containers(
                 f"Stopped the {Fore.GREEN}{rabbit_container_name}"
                 f"{Style.RESET_ALL} container."
             )
+
+    if ctx.config.get("prometheus", {}).get("start_with_server"):
+        remove_container_if_exists(client, name=ctx.prometheus_container_name)
+        info(
+            f"Stopped the {Fore.GREEN}{ctx.prometheus_container_name}"
+            f"{Style.RESET_ALL} container."
+        )

vantage6/cli/template/node_config.j2

@@ -1,11 +1,13 @@
 api_key: {{ api_key }}
 api_path: /api
 databases:
-  {% for label, path in databases.items() %}
+  {% for db in databases %}
+  {% for label, path in db.items() %}
   - label: {{ label }}
     uri: {{ path }}
     type: csv
   {% endfor %}
+  {% endfor %}
 encryption:
   enabled: false
   private_key: null
@@ -31,4 +33,4 @@ logging:
 port: {{ port }}
 server_url: {{ server_url }}
 task_dir: {{ task_dir}}
-{{ user_provided_config }}
+{{- user_provided_config -}}

vantage6/cli/test/algo_test_scripts/algo_test_arguments.py

@@ -0,0 +1,33 @@
+average = {
+    "collaboration": 1,
+    "organizations": [1],
+    "name": "test_average_task",
+    "image": "harbor2.vantage6.ai/demo/average",
+    "description": "",
+    "input_": {
+        "method": "central_average",
+        "args": [],
+        "kwargs": {"column_name": "Age"},
+    },
+    "databases": [{"label": "olympic_athletes"}],
+}
+
+kaplan_meier = {
+    "collaboration": 1,
+    "organizations": [1],
+    "name": "test_average_task",
+    "image": "harbor2.vantage6.ai/algorithms/kaplan-meier",
+    "description": "",
+    "input_": {
+        "method": "kaplan_meier_central",
+        "args": [],
+        "kwargs": {
+            "time_column_name": "days",
+            "censor_column_name": "censor",
+            "organizations_to_include": [1, 2, 3],
+        },
+    },
+    "databases": [{"label": "kaplan_meier_test"}],
+}
+
+args = {"average": average, "kaplan_meier": kaplan_meier}

vantage6/cli/test/algo_test_scripts/algo_test_script.py

@@ -0,0 +1,88 @@
+import logging
+import algo_test_arguments as arguments
+import json
+import sys
+
+import vantage6.common.task_status as task_status
+
+from vantage6.client import Client
+from vantage6.common import error
+from vantage6.common.globals import Ports
+
+
+def create_and_run_task(client: Client, task_args: dict, algo_name: str = "algorithm"):
+    """
+    Create and run a task using the provided client and task arguments.
+
+    Parameters
+    ----------
+        client: Client
+            The client instance to use for creating and running the task.
+        task_args: dict
+            The arguments to pass to the task creation method.
+        algo_name: str, optional
+            The name of the algorithm for logging purposes. Default is "algorithm".
+
+    Raises
+    ------
+        AssertionError: If the task fails.
+    """
+    task = client.task.create(**task_args)
+    task_id = task["id"]
+    client.wait_for_results(task_id)
+
+    try:
+        # check if the task has failed
+        assert not task_status.has_task_failed(client.task.get(task_id)["status"])
+
+        logging.info(f"Task for {algo_name} completed successfully.")
+
+    except AssertionError:
+        error(
+            f"Task for {algo_name} failed. Check the log file of the task "
+            f"{task_id} for more information."
+        )
+        exit(1)
+
+
+def run_test(custom_args: dict | None = None):
+    """
+    Run a test by creating and running tasks using the provided arguments.
+
+    Parameters
+    ----------
+    custom_args: dict, optional
+        The arguments to pass to the task creation method. If not provided,
+        the arguments from the `arguments` module will be used.
+    """
+    # Create a client and authenticate
+    client = Client("http://localhost", Ports.DEV_SERVER.value, "/api")
+    try:
+        client.authenticate("dev_admin", "password")
+    except ConnectionError:
+        error(
+            "Could not connect to the server. Please check if a dev network is running."
+        )
+        exit(1)
+
+    # if custom arguments are provided, use them for running the task
+    if custom_args:
+        create_and_run_task(client, custom_args)
+
+    else:
+        # Run the task for each algorithm in the arguments file
+        for algo, task_args in arguments.args.items():
+            logging.info("Running task for %s", algo)
+            logging.info("Task arguments: %s", task_args)
+            create_and_run_task(client, task_args, algo)
+
+
+if __name__ == "__main__":
+    # check if arguments are provided
+    if len(sys.argv) > 1:
+        input_string = sys.argv[1].replace("'", '"')
+        json_input = json.loads(input_string)
+    else:
+        json_input = None
+
+    run_test(json_input)

vantage6/cli/test/client_script.py

@@ -0,0 +1,150 @@
+import click
+import json
+import subprocess
+import sys
+
+from pathlib import Path
+
+from rich.console import Console
+
+from vantage6.cli.dev.create import create_demo_network
+from vantage6.cli.dev.remove import remove_demo_network
+from vantage6.cli.dev.start import start_demo_network
+from vantage6.cli.dev.stop import stop_demo_network
+from vantage6.cli.utils import prompt_config_name
+from vantage6.common.globals import Ports
+
+TEST_FILE_PATH = Path(__file__).parent / "algo_test_scripts" / "algo_test_script.py"
+
+
+@click.command()
+@click.option(
+    "--script",
+    type=click.Path(),
+    default=TEST_FILE_PATH,
+    help="Path of the script to test the algorithm. If a script is not provided, the default script is used.",
+)
+@click.option(
+    "--task-arguments",
+    type=str,
+    default=None,
+    help="Arguments to be provided to Task.create function. If --script is provided, this should not be set.",
+)
+@click.option(
+    "--create-dev-network",
+    is_flag=True,
+    help="Create a new dev network to run the test",
+)
+@click.option(
+    "--start-dev-network",
+    is_flag=True,
+    help="Start a dev network to run the test",
+)
+@click.option(
+    "-n", "--name", default=None, type=str, help="Name for your development setup"
+)
+@click.option(
+    "--server-url",
+    type=str,
+    default="http://host.docker.internal",
+    help="Server URL to point to. If you are using Docker Desktop, "
+    "the default http://host.docker.internal should not be changed.",
+)
+@click.option(
+    "-i", "--image", type=str, default=None, help="Server Docker image to use"
+)
+@click.option(
+    "--keep",
+    type=bool,
+    default=False,
+    help="Keep the dev network after finishing the test",
+)
+@click.option(
+    "--add-dataset",
+    type=(str, click.Path()),
+    default=[],
+    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.pass_context
+def cli_test_client_script(
+    click_ctx: click.Context,
+    script: Path | None,
+    task_arguments: str | None,
+    name: str,
+    server_url: str,
+    create_dev_network: bool,
+    start_dev_network: bool,
+    image: str,
+    keep: bool,
+    add_dataset: list[tuple[str, Path]] = (),
+) -> int:
+    """
+    Run a script for testing an algorithm on a dev network.
+    The path to the script must be provided as an argument.
+    """
+    if not (script or task_arguments):
+        raise click.UsageError("--script or --task-arguments must be set.")
+    elif script != TEST_FILE_PATH and task_arguments:
+        raise click.UsageError("--script and --task-arguments cannot be set together.")
+
+    # Check if the task_arguments is a valid JSON string
+    if task_arguments:
+        try:
+            json.loads(task_arguments.replace("'", '"'))
+        except json.JSONDecodeError:
+            raise click.UsageError("task-arguments must be a valid JSON string.")
+
+    name = prompt_config_name(name)
+
+    # create the network
+    if create_dev_network:
+        click_ctx.invoke(
+            create_demo_network,
+            name=name,
+            num_nodes=3,
+            server_url=server_url,
+            server_port=Ports.DEV_SERVER.value,
+            image=image,
+            extra_server_config=None,
+            extra_node_config=None,
+            add_dataset=add_dataset,
+        )
+
+    # start the server and nodes
+    if create_dev_network or start_dev_network:
+        click_ctx.invoke(
+            start_demo_network,
+            name=name,
+            server_image=image,
+            node_image=image,
+        )
+
+    # run the test script and get the result
+    if not task_arguments:
+        subprocess_args = ["python", script]
+    else:
+        subprocess_args = ["python", script, task_arguments]
+
+    result = subprocess.run(subprocess_args, stdout=sys.stdout, stderr=sys.stderr)
+
+    # check the exit code. If the test passed, it should be 0
+    if result.returncode == 0:
+        msg = ":heavy_check_mark: [green]Test passed[/green]"
+    else:
+        msg = ":x: [red]Test failed[/red]"
+
+    console = Console()
+    console.print(msg)
+
+    # clean up the test resources. Keep the network if --keep is set, or if the network
+    # was created for this test. If the network was started for this test, stop it but
+    # do not remove it.
+    if not keep:
+        if create_dev_network or start_dev_network:
+            click_ctx.invoke(stop_demo_network, name=name)
+        if create_dev_network:
+            click_ctx.invoke(remove_demo_network, name=name)
+
+    return result.returncode

{vantage6-4.10.2.dist-info → vantage6-4.11.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vantage6
-Version: 4.10.2
+Version: 4.11.0
 Summary: vantage6 command line interface
 Home-page: https://github.com/vantage6/vantage6
 Requires-Python: >=3.10
@@ -16,8 +16,8 @@ Requires-Dist: questionary==1.10.0
 Requires-Dist: rich==13.5.2
 Requires-Dist: schema==0.7.5
 Requires-Dist: SQLAlchemy==1.4.46
-Requires-Dist: vantage6-common==4.10.2
-Requires-Dist: vantage6-client==4.10.2
+Requires-Dist: vantage6-common==4.11.0
+Requires-Dist: vantage6-client==4.11.0
 Provides-Extra: dev
 Requires-Dist: coverage==6.4.4; extra == "dev"
 Requires-Dist: black; extra == "dev"