hnsw-tuner / backend_cloud.py
raduvespa's picture
Initial HF Space
88f21af
# 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,
)