From be2f99103df6893b6fdd0df0fd2969469422fd7f Mon Sep 17 00:00:00 2001
From: Antoni Ivanov <aivanov@vmware.com>
Date: Fri, 29 Sep 2023 11:54:25 +0300
Subject: [PATCH] vdk-jupyter: print summary output to temp dir (#2715)

Use the newly provided JOB_RUN_SUMMARY_FILE_PATH and JobSummaryParser to
get information about job summary which can be shown in Notebook UI.

This passes the path explicitly to vdk run as environment variable
making sure not hidden depencies would break things in the future.
---
 .../vdk_jupyterlab_extension/vdk_ui.py        | 70 ++++++++++---------
 1 file changed, 36 insertions(+), 34 deletions(-)

diff --git a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/vdk_ui.py b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/vdk_ui.py
index 8c412fc76d..28f48797f0 100644
--- a/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/vdk_ui.py
+++ b/projects/vdk-plugins/vdk-jupyter/vdk-jupyterlab-extension/vdk_jupyterlab_extension/vdk_ui.py
@@ -9,6 +9,7 @@
 import tempfile
 from pathlib import Path
 
+from vdk.internal.builtin_plugins.run.summary_output import JobSummaryParser
 from vdk.internal.control.command_groups.job.create import JobCreate
 from vdk.internal.control.command_groups.job.delete import JobDelete
 from vdk.internal.control.command_groups.job.deploy_cli_impl import JobDeploy
@@ -61,42 +62,43 @@ def run_job(path, arguments=None):
         if len(script_files) == 0:
             return {"message": f"No steps were found in {path}!"}
         with open("vdk_run.log", "w+") as log_file:
-            path = shlex.quote(str(path))
-            cmd: list[str] = ["vdk", "run", f"{path}"]
-            if arguments:
-                arguments = shlex.quote(arguments)
-                cmd.append("--arguments")
-                cmd.append(f"{arguments}")
+            cmd = VdkUI.__prepare_vdk_run_command(arguments, path)
+            with tempfile.NamedTemporaryFile(
+                prefix="jupyter-run-summary-", mode="w+"
+            ) as summary_output_file:
+                environ_copy = os.environ.copy()
+                environ_copy["JOB_RUN_SUMMARY_FILE_PATH"] = summary_output_file.name
+                # We add the PYTHONUNBUFFERED env variable to ensure print statements inside user jobs are correctly
+                # interspersed in between the rest of the job logs instead of being placed at the end
+                environ_copy["PYTHONUNBUFFERED"] = "x"
 
-            # We add the PYTHONUNBUFFERED env variable to ensure print statements inside user jobs are correctly
-            # interspersed in between the rest of the job logs instead of being placed at the end
-            env = os.environ.copy()
-            env["PYTHONUNBUFFERED"] = "x"
+                try:
+                    process = subprocess.Popen(
+                        cmd,
+                        stdout=log_file,
+                        stderr=log_file,
+                        env=environ_copy,
+                        shell=False,
+                    )
+                    process.wait()
+                    if process.returncode != 0:
+                        content = pathlib.Path(summary_output_file.name).read_text()
+                        job_summary = JobSummaryParser.from_json(content)
+                        return {"message": job_summary.details}
+                    else:
+                        return {"message": process.returncode}
+                except Exception as e:
+                    return {"message": f"Unhandled exception: {e}"}
 
-            process = subprocess.Popen(
-                cmd,
-                stdout=log_file,
-                stderr=log_file,
-                env=env,
-                shell=False,
-            )
-            process.wait()
-            if process.returncode != 0:
-                job_name = (
-                    os.path.basename(path[:-1])
-                    if path.endswith("/")
-                    else os.path.basename(path)
-                )
-                error_file = os.path.join(
-                    os.path.dirname(path), f".{job_name}_error.json"
-                )
-                if os.path.exists(error_file):
-                    with open(error_file) as file:
-                        # the json is generated in vdk-notebook plugin
-                        # you can see /vdk-notebook/src/vdk/notebook-plugin.py
-                        error = json.load(file)
-                        return {"message": error["details"]}
-            return {"message": process.returncode}
+    @staticmethod
+    def __prepare_vdk_run_command(arguments, path):
+        path = shlex.quote(str(path))
+        cmd: list[str] = ["vdk", "run", f"{path}"]
+        if arguments:
+            arguments = shlex.quote(arguments)
+            cmd.append("--arguments")
+            cmd.append(f"{arguments}")
+        return cmd
 
     @staticmethod
     def delete_job(name: str, team: str) -> str: