Repository: GoogleCloudPlatform/PerfKitBenchmarker

Issue:
RobustRemoteCommand fails is run after the VM is restarted
`RobustRemoteCommand` fails if it is run on a VM that has been rebooted, because the `/tmp/pkb` directory has been deleted. `linux_virtual_machine.py` should recreate the tmp directory if necessary. 


Discussion:

Resolving patch:
diff --git a/perfkitbenchmarker/linux_virtual_machine.py b/perfkitbenchmarker/linux_virtual_machine.py
--- a/perfkitbenchmarker/linux_virtual_machine.py
+++ b/perfkitbenchmarker/linux_virtual_machine.py
@@ -108,6 +108,9 @@ def __init__(self):
     self._remote_command_script_upload_lock = threading.Lock()
     self._has_remote_command_script = False
 
+  def _CreateVmTmpDir(self):
+        self.RemoteCommand('mkdir -p %s' % vm_util.VM_TMP_DIR)
+
   def _PushRobustCommandScripts(self):
     """Pushes the scripts required by RobustRemoteCommand to this VM.
 
@@ -208,7 +211,7 @@ def SetupPackageManager(self):
 
   def PrepareVMEnvironment(self):
     self.SetupProxy()
-    self.RemoteCommand('mkdir -p %s' % vm_util.VM_TMP_DIR)
+    self._CreateVmTmpDir()
     if FLAGS.setup_remote_firewall:
       self.SetupRemoteFirewall()
     if self.install_packages:
@@ -409,6 +412,17 @@ def RemoteHostCommand(self, command,
 
     return stdout, stderr
 
+  def _Reboot(self):
+    """OS-specific implementation of reboot command"""
+    self.RemoteCommand('sudo reboot', ignore_failure=True)
+
+  def _AfterReboot(self):
+    """Performs any OS-specific setup on the VM following reboot.
+
+    This will be called after every call to Reboot().
+    """
+    self._CreateVmTmpDir()
+
   def MoveFile(self, target, source_path, remote_path=''):
     self.MoveHostFile(target, source_path, remote_path)
 
@@ -888,7 +902,7 @@ def _CheckDockerExists(self):
 
   def PrepareVMEnvironment(self):
     """Initializes docker before proceeding with preparation."""
-    self.RemoteHostCommand('mkdir -p %s' % vm_util.VM_TMP_DIR)
+    self._CreateVmTmpDir()
     if not self._CheckDockerExists():
       self.Install('docker')
     self.InitDocker()
diff --git a/perfkitbenchmarker/virtual_machine.py b/perfkitbenchmarker/virtual_machine.py
--- a/perfkitbenchmarker/virtual_machine.py
+++ b/perfkitbenchmarker/virtual_machine.py
@@ -395,6 +395,24 @@ def TryRemoteCommand(self, command, **kwargs):
     except:
       raise
 
+  def Reboot(self):
+    """Reboot the VM."""
+    self._Reboot()
+    self.WaitForBootCompletion()
+    self._AfterReboot()
+
+  @abc.abstractmethod
+  def _Reboot(self):
+    """OS-specific implementation of reboot command"""
+    raise NotImplementedError()
+
+  def _AfterReboot(self):
+    """Performs any OS-specific setup on the VM following reboot.
+
+    This will be called after every call to Reboot().
+    """
+    pass
+
   @abc.abstractmethod
   def RemoteCopy(self, file_path, remote_path='', copy_to=True):
     """Copies a file to or from the VM.
diff --git a/perfkitbenchmarker/windows_virtual_machine.py b/perfkitbenchmarker/windows_virtual_machine.py
--- a/perfkitbenchmarker/windows_virtual_machine.py
+++ b/perfkitbenchmarker/windows_virtual_machine.py
@@ -165,6 +165,17 @@ def OnStartup(self):
     self.RemoteCommand('mkdir %s' % self.temp_dir)
     self.DisableGuestFirewall()
 
+  def _Reboot(self):
+    """OS-specific implementation of reboot command"""
+    self.RemoteCommand('shutdown -t 0 -r -f', ignore_failure=True)
+
+  def _AfterReboot(self):
+    """Performs any OS-specific setup on the VM following reboot.
+
+    This will be called after every call to Reboot().
+    """
+    pass
+
   def Install(self, package_name):
     """Installs a PerfKit package on the VM."""
     if not self.install_packages:


Test patch: