Merge pull request #14104 from vespa-engine/tgm/deploy-return-vespa-instance

return Vespa connection when deploying Vespa app.

2 files changed, 37 insertions, 3 deletions

diff --git a/python/vespa/vespa/application.py b/python/vespa/vespa/application.py
index dfb8f181cd4..3ab31c4dd8c 100644
--- a/python/vespa/vespa/application.py
+++ b/python/vespa/vespa/application.py
@@ -10,12 +10,18 @@ from vespa.evaluation import EvalMetric
 
 
 class Vespa(object):
-    def __init__(self, url: str, port: Optional[int] = None) -> None:
+    def __init__(
+        self,
+        url: str,
+        port: Optional[int] = None,
+        deployment_message: Optional[List[str]] = None,
+    ) -> None:
         """
         Establish a connection with a Vespa application.
 
         :param url: URL
         :param port: Port
+        :param deployment_message: Message returned by Vespa engine after deployment.
 
             >>> Vespa(url = "https://cord19.vespa.ai")
             >>> Vespa(url = "http://localhost", port = 8080)
@@ -23,6 +29,7 @@ class Vespa(object):
         """
         self.url = url
         self.port = port
+        self.deployment_message = deployment_message
 
         if port is None:
             self.end_point = self.url
@@ -30,6 +37,12 @@ class Vespa(object):
             self.end_point = str(url).rstrip("/") + ":" + str(port)
         self.search_end_point = self.end_point + "/search/"
 
+    def __repr__(self):
+        if self.port:
+            return "Vespa({}, {})".format(self.url, self.port)
+        else:
+            return "Vespa({})".format(self.url)
+
     def query(
         self,
         body: Optional[Dict] = None,
diff --git a/python/vespa/vespa/package.py b/python/vespa/vespa/package.py
index 7b6b5fae757..4b5d1e701d5 100644
--- a/python/vespa/vespa/package.py
+++ b/python/vespa/vespa/package.py
@@ -1,4 +1,5 @@
 import os
+import re
 from time import sleep
 from typing import List, Mapping, Optional
 from pathlib import Path
@@ -7,6 +8,7 @@ from jinja2 import Environment, PackageLoader, select_autoescape
 import docker
 
 from vespa.json_serialization import ToJson, FromJson
+from vespa.application import Vespa
 
 
 class Field(ToJson, FromJson["Field"]):
@@ -390,6 +392,7 @@ class VespaDocker(object):
         """
         self.application_package = application_package
         self.container = None
+        self.local_port = 8080
 
     def run_vespa_engine_container(self, disk_folder: str, container_memory: str):
         """
@@ -412,7 +415,7 @@ class VespaDocker(object):
                     hostname=self.application_package.name,
                     privileged=True,
                     volumes={disk_folder: {"bind": "/app", "mode": "rw"}},
-                    ports={8080: 8080, 19112: 19112},
+                    ports={self.local_port: self.local_port, 19112: 19112},
                 )
 
     def check_configuration_server(self) -> bool:
@@ -432,6 +435,14 @@ class VespaDocker(object):
         )
 
     def deploy(self, disk_folder: str, container_memory: str = "4G"):
+        """
+        Deploy the application into a Vespa container.
+
+        :param disk_folder: Disk folder to save the required Vespa config files.
+        :param container_memory: Docker container memory available to the application.
+
+        :return: a Vespa connection instance.
+        """
 
         self.application_package.create_application_package_files(dir_path=disk_folder)
 
@@ -446,4 +457,14 @@ class VespaDocker(object):
         deployment = self.container.exec_run(
             "bash -c '/opt/vespa/bin/vespa-deploy prepare /app/application && /opt/vespa/bin/vespa-deploy activate'"
         )
-        return deployment.output.decode("utf-8").split("\n")
+
+        deployment_message = deployment.output.decode("utf-8").split("\n")
+
+        if not any(re.match("Generation: [0-9]+", line) for line in deployment_message):
+            raise RuntimeError(deployment_message)
+
+        return Vespa(
+            url="http://localhost",
+            port=self.local_port,
+            deployment_message=deployment_message,
+        )