From f7a82f6fc01f4738cf333e0993c326f0805aa9a5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?=
 <peter@pmrowla.com>
Date: Tue, 7 Sep 2021 06:07:22 +0000
Subject: [PATCH] machine: support instance size & startup script config
 (#6537)

* machine: install dvc via startup_script by default

* machine: support configuring instance type and startup script

* accept startup_script as a config relpath

* use pushd/popd in default install script
---
 dvc/config.py           |  1 +
 dvc/config_schema.py    |  2 ++
 dvc/machine/__init__.py | 18 ++++++++++++++++++
 3 files changed, 21 insertions(+)

diff --git a/dvc/config.py b/dvc/config.py
index ac927c5c63..98804bf225 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -263,6 +263,7 @@ def _map_dirs(conf, func):
                     "key_path": func,
                 }
             },
+            "machine": {str: {"startup_script": func}},
         }
         return Schema(dirs_schema, extra=ALLOW_EXTRA)(conf)
 
diff --git a/dvc/config_schema.py b/dvc/config_schema.py
index 7882a63c74..1b55d45dc8 100644
--- a/dvc/config_schema.py
+++ b/dvc/config_schema.py
@@ -229,6 +229,8 @@ class RelPath(str):
     "machine": {
         str: {
             "cloud": All(Lower, Choices("aws", "azure")),
+            "instance_type": Lower,
+            "startup_script": str,
         },
     },
     # section for experimental features
diff --git a/dvc/machine/__init__.py b/dvc/machine/__init__.py
index 8d87b1c393..8197df4ba9 100644
--- a/dvc/machine/__init__.py
+++ b/dvc/machine/__init__.py
@@ -25,6 +25,18 @@
 
 RESERVED_NAMES = {"local", "localhost"}
 
+DEFAULT_STARTUP_SCRIPT = """#!/bin/bash
+sudo apt-get update
+sudo apt-get install --yes python3 python3-pip
+python3 -m pip install --upgrade pip
+
+pushd /etc/apt/sources.list.d
+sudo wget https://dvc.org/deb/dvc.list
+sudo apt-get update
+sudo apt-get install --yes dvc
+popd
+"""
+
 
 def validate_name(name: str):
     from dvc.exceptions import InvalidArgumentError
@@ -168,6 +180,12 @@ def _get_backend(self, cloud: str) -> "BaseMachineBackend":
     def create(self, name: Optional[str]):
         """Create and start the specified machine instance."""
         config, backend = self.get_config_and_backend(name)
+        if "startup_script" in config:
+            with open(config["startup_script"]) as fobj:
+                startup_script = fobj.read()
+        else:
+            startup_script = DEFAULT_STARTUP_SCRIPT
+        config["startup_script"] = startup_script
         return backend.create(**config)
 
     def destroy(self, name: Optional[str]):