From fb16cd06e3f00b9c2419cbb2fe5fe64cc8d0de37 Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Mon, 27 Dec 2021 16:34:06 +0900
Subject: [PATCH] machine: update behavior for exp ssh executors

- log startup script finished message
- add setup_script option for runtime env setup
---
 dvc/config.py           | 7 ++++++-
 dvc/config_schema.py    | 1 +
 dvc/machine/__init__.py | 7 +++++++
 3 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/dvc/config.py b/dvc/config.py
index ad8a98b32..caa62e785 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -266,7 +266,12 @@ class Config(dict):
                     "key_path": func,
                 }
             },
-            "machine": {str: {"startup_script": func}},
+            "machine": {
+                str: {
+                    "startup_script": func,
+                    "setup_script": func,
+                }
+            },
         }
         return Schema(dirs_schema, extra=ALLOW_EXTRA)(conf)
 
diff --git a/dvc/config_schema.py b/dvc/config_schema.py
index 098115f24..aa495a890 100644
--- a/dvc/config_schema.py
+++ b/dvc/config_schema.py
@@ -257,6 +257,7 @@ SCHEMA = {
             "instance_gpu": Lower,
             "ssh_private": str,
             "startup_script": str,
+            "setup_script": str,
         },
     },
     # section for experimental features
diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py
index ccf767f46..d93c3e7b9 100644
--- a/dvc/machine/__init__.py
+++ b/dvc/machine/__init__.py
@@ -37,6 +37,8 @@ sudo wget https://dvc.org/deb/dvc.list
 sudo apt-get update
 sudo apt-get install --yes dvc
 popd
+
+sudo echo "OK" > /var/log/dvc-machine-init.log
 """
 
 
@@ -188,6 +190,7 @@ class MachineManager:
         else:
             startup_script = DEFAULT_STARTUP_SCRIPT
         config["startup_script"] = startup_script
+        config.pop("setup_script", None)
         return backend.create(**config)
 
     def destroy(self, name: Optional[str]):
@@ -215,3 +218,7 @@ class MachineManager:
     def get_executor_kwargs(self, name: Optional[str]):
         config, backend = self.get_config_and_backend(name)
         return backend.get_executor_kwargs(**config)
+
+    def get_setup_script(self, name: Optional[str]) -> Optional[str]:
+        config, _backend = self.get_config_and_backend(name)
+        return config.get("setup_script")