machine: update behavior for exp ssh executors

- log startup script finished message
- add setup_script option for runtime env setup

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)
 

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

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")