Spaces:
Running
Running
| # Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. | |
| """Vespa Cloud backend.""" | |
| import logging | |
| import sys | |
| from pathlib import Path | |
| from cryptography import x509 | |
| from cryptography.hazmat.backends import default_backend | |
| from cryptography.hazmat.primitives import serialization | |
| from vespa.application import Vespa | |
| from vespa.deployment import VespaCloud | |
| from application_package import get_application_root | |
| from config import Config | |
| log = logging.getLogger(__name__) | |
| AUTH_JSON_PATH = Path.home() / ".vespa" / "auth.json" | |
| class _ConfiguredVespaCloud(VespaCloud): | |
| """VespaCloud wrapper that honors explicit data-plane cert/key overrides.""" | |
| def __init__(self, cfg: Config, **kwargs) -> None: | |
| self._cfg = cfg | |
| super().__init__(**kwargs) | |
| def _load_certificate_pair(self, generate_cert: bool = True): | |
| if self._cfg.vespa_cloud_cert_path and self._cfg.vespa_cloud_key_path: | |
| cert_path = self._cfg.get_cloud_cert_path() | |
| key_path = self._cfg.get_cloud_key_path() | |
| with open(key_path, "rb") as key_file: | |
| private_key = serialization.load_pem_private_key( | |
| key_file.read(), | |
| password=None, | |
| backend=default_backend(), | |
| ) | |
| with open(cert_path, "rb") as cert_file: | |
| cert = x509.load_pem_x509_certificate( | |
| cert_file.read(), | |
| default_backend(), | |
| ) | |
| self.data_cert_path = str(cert_path) | |
| self.data_key_path = str(key_path) | |
| return private_key, cert | |
| return super()._load_certificate_pair(generate_cert=generate_cert) | |
| def has_control_plane_auth(cfg: Config) -> bool: | |
| """Return True when control-plane auth is available for Vespa Cloud.""" | |
| return bool(cfg.vespa_cloud_api_key) or AUTH_JSON_PATH.exists() | |
| def create_cloud_client(cfg: Config) -> VespaCloud: | |
| """Create a VespaCloud control-plane client from config.""" | |
| if not has_control_plane_auth(cfg) and not sys.stdin.isatty(): | |
| raise RuntimeError( | |
| "Vespa Cloud control-plane auth is required in non-interactive mode. " | |
| "Set VESPA_CLOUD_API_KEY or run `vespa auth login` first." | |
| ) | |
| api_key_location = None | |
| api_key_content = None | |
| if cfg.vespa_cloud_api_key: | |
| if cfg.get_cloud_api_key_path().exists(): | |
| api_key_location = str(cfg.get_cloud_api_key_path()) | |
| else: | |
| api_key_content = cfg.vespa_cloud_api_key | |
| return _ConfiguredVespaCloud( | |
| cfg, | |
| tenant=cfg.vespa_cloud_tenant, | |
| application=cfg.vespa_cloud_application, | |
| instance=cfg.vespa_cloud_instance, | |
| application_root=str(get_application_root()), | |
| key_location=api_key_location, | |
| key_content=api_key_content, | |
| ) | |
| class CloudBackend: | |
| """Prepare a Vespa Cloud dev deployment for use by HNSWTuner.""" | |
| def prepare(self, cfg: Config) -> Vespa: | |
| cloud = create_cloud_client(cfg) | |
| if cfg.bootstrap_mode == "full": | |
| log.info( | |
| "Deploying application package to Vespa Cloud dev (%s) ...", | |
| cfg.get_cloud_application_id(), | |
| ) | |
| return cloud.deploy( | |
| instance=cfg.vespa_cloud_instance, | |
| environment="dev", | |
| max_wait=cfg.vespa_cloud_max_wait, | |
| ) | |
| log.info( | |
| "Connecting to existing Vespa Cloud dev deployment (%s) ...", | |
| cfg.get_cloud_application_id(), | |
| ) | |
| return cloud.get_application( | |
| instance=cfg.vespa_cloud_instance, | |
| environment="dev", | |
| endpoint_type="mtls", | |
| max_wait=cfg.vespa_cloud_connect_wait, | |
| ) | |