From cee40c6f11652e166da3b1efbd60cecbdde30fdc Mon Sep 17 00:00:00 2001 From: Matt Bookman Date: Mon, 27 Mar 2023 20:29:32 +0000 Subject: [PATCH] POC demonstrating how to use guest attributes to improve notebook VM startup output. A good implementation would probably include getting attributes through the GCP API (through a wrapper in the terra-cloud-resource-lib). Would anticipate having command-line flags for people to choose whether they want to fire and forget or wait for completion. --- .../cli/app/utils/LocalProcessLauncher.java | 13 ++ .../command/resource/create/GcpNotebook.java | 115 ++++++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/src/main/java/bio/terra/cli/app/utils/LocalProcessLauncher.java b/src/main/java/bio/terra/cli/app/utils/LocalProcessLauncher.java index 374da6587..b5fece5c4 100644 --- a/src/main/java/bio/terra/cli/app/utils/LocalProcessLauncher.java +++ b/src/main/java/bio/terra/cli/app/utils/LocalProcessLauncher.java @@ -6,6 +6,7 @@ import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; +import java.io.OutputStream; import java.io.PrintStream; import java.nio.charset.StandardCharsets; import java.nio.file.Path; @@ -96,4 +97,16 @@ public int waitForTerminate() { throw new SystemException("Error waiting for child process to terminate", intEx); } } + + public InputStream getErrorStream() { + return process.getErrorStream(); + } + + public InputStream getInputStream() { + return process.getInputStream(); + } + + public OutputStream getOutputStream() { + return process.getOutputStream(); + } } diff --git a/src/main/java/bio/terra/cli/command/resource/create/GcpNotebook.java b/src/main/java/bio/terra/cli/command/resource/create/GcpNotebook.java index 07228feee..a91fa97d9 100644 --- a/src/main/java/bio/terra/cli/command/resource/create/GcpNotebook.java +++ b/src/main/java/bio/terra/cli/command/resource/create/GcpNotebook.java @@ -1,5 +1,6 @@ package bio.terra.cli.command.resource.create; +import bio.terra.cli.app.utils.LocalProcessLauncher; import bio.terra.cli.command.shared.WsmBaseCommand; import bio.terra.cli.command.shared.options.ControlledResourceCreation; import bio.terra.cli.command.shared.options.Format; @@ -11,9 +12,17 @@ import bio.terra.workspace.model.AccessScope; import bio.terra.workspace.model.CloudPlatform; import bio.terra.workspace.model.StewardshipType; +import java.io.BufferedReader; +import java.io.IOException; +import java.io.InputStream; +import java.io.InputStreamReader; +import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.Optional; +import org.json.JSONObject; import picocli.CommandLine; /** This class corresponds to the fourth-level "terra resource create gcp-notebook" command. */ @@ -167,6 +176,112 @@ protected void execute() { bio.terra.cli.businessobject.resource.GcpNotebook createdResource = bio.terra.cli.businessobject.resource.GcpNotebook.createControlled(createParams.build()); formatOption.printReturnValue(new UFGcpNotebook(createdResource), GcpNotebook::printText); + + boolean doWaitForNotebooksReady = true; // this would be a flag + if (doWaitForNotebooksReady) { + waitForNotebooksReady(createdResource.getInstanceId()); + } + } + + private String getNotebookGuestAttribute(String instanceId, String guestAttribute) { + String command = + "gcloud compute instances get-guest-attributes " + + instanceId + + " " + + "--zone us-central1-a " + + "--query-path '" + + guestAttribute + + "' " + + "--format 'json(value)'"; + + List processCommand = new ArrayList<>(); + processCommand.add("bash"); + processCommand.add("-ce"); + processCommand.add(command); + + Map envVars = new HashMap(); + + LocalProcessLauncher localProcessLauncher = new LocalProcessLauncher(); + localProcessLauncher.launchProcess(processCommand, envVars); + + InputStream stdout = localProcessLauncher.getInputStream(); + int exitCode = localProcessLauncher.waitForTerminate(); + if (exitCode != 0) { + // This is actually expected for several iterations immediately + // after provisioning. + ERR.println(String.format("ERROR getting notebook startup status: %d", exitCode)); + return ""; + } + + final BufferedReader reader = + new BufferedReader(new InputStreamReader(localProcessLauncher.getInputStream())); + + String output; + try { + StringBuilder sb = new StringBuilder(); + String line = null; + while ((line = reader.readLine()) != null) { + sb.append(line); + } + reader.close(); + + output = sb.toString().trim(); + } catch (IOException ex) { + // TODO: + return null; + } + + if (output.isEmpty()) { + return ""; + } else { + // Remove the starting and trailing square brackets + // and then parse as JSON. + String jsonstr = output.substring(1, output.length() - 1); + JSONObject obj = new JSONObject(jsonstr); + return obj.getString("value").trim(); + } + } + + private String getNotebookReadyStatus(String instanceId) { + return getNotebookGuestAttribute(instanceId, "startup_script/status"); + } + + private String getNotebookReadyMessage(String instanceId) { + return getNotebookGuestAttribute(instanceId, "startup_script/message"); + } + + private void waitForNotebooksReady(String instanceId) { + String lastStatus = ""; + while (true) { + String currStatus = getNotebookReadyStatus(instanceId); + + if (currStatus.isEmpty()) { + OUT.println("The post-startup script has not started."); + } else if (currStatus.equals(lastStatus)) { + OUT.print("."); + } else { + OUT.println(); + OUT.print(currStatus); + + if (currStatus.equals("COMPLETE")) { + OUT.println(); + break; + } else if (currStatus.equals("ERROR")) { + OUT.println(); + OUT.println(getNotebookReadyMessage(instanceId)); + break; + } + + lastStatus = currStatus; + } + + try { + Thread.sleep(5000); + } catch (InterruptedException ex) { + ERR.println(ex); + break; + } + } } static class VmOrContainerImage {