adding environment setup in init command

MANIFEST.in

@@ -2,7 +2,7 @@ include datmo/core/controller/environment/driver/templates/baseDockerfile
 include datmo/core/controller/environment/driver/templates/python2Dockerfile
 include datmo/core/controller/environment/driver/templates/python3Dockerfile
 include datmo/core/util/validation/schemas.yml
-include datmo/core/controller/environment/driver/docker.json
+include datmo/core/controller/environment/driver/config/docker.json
 include datmo/VERSION
 prune examples
 prune docs

datmo/cli/command/project.py

@@ -3,6 +3,7 @@
 from datmo.cli.command.base import BaseCommand
 from datmo.core.util.misc_functions import mutually_exclusive
 from datmo.core.controller.project import ProjectController
+from datmo.core.controller.environment.environment import EnvironmentController, EnvironmentDoesNotExist
 from datmo.core.controller.task import TaskController
 
 class ProjectCommand(BaseCommand):
@@ -98,7 +99,40 @@ def init(self, name, description):
         for k, v in self.project_controller.model.to_dictionary().items():
             if k != "config":
                 self.cli_helper.echo(str(k) + ": " + str(v))
+        # Ask question if the user would like to setup environment
+        environment_setup = self.cli_helper.prompt_bool(
+            __("prompt", "cli.project.environment.setup"))
+        if environment_setup:
+            # Setting up the environment definition file
+            self.environment_controller = EnvironmentController(home=self.home)
+            available_environments = self.environment_controller.get_supported_environments()
+            available_environment_names, available_environment_description = zip(*available_environments)
+            for idx, environment_name_description in enumerate(available_environments):
+                environment_name = environment_name_description[0]
+                environment_description = environment_name_description[1]
+                self.cli_helper.echo("(%s) %s: %s" % (idx + 1, environment_name, environment_description))
+            input_environment_name = self.cli_helper.prompt(
+                __("prompt", "cli.environment.setup.name"))
+            try:
+                name_environment_index = int(input_environment_name)
+            except ValueError:
+                name_environment_index = 0
+            if 0 < name_environment_index < len(available_environments):
+                input_environment_name = available_environment_names[name_environment_index - 1]
+            else:
+                self.cli_helper.echo(
+                    __("error", "cli.environment.setup.argument", input_environment_name))
 
+            try:
+                options = {"name": input_environment_name}
+                environment_obj = self.environment_controller.setup(
+                    options=options)
+                self.cli_helper.echo(
+                    __("info", "cli.environment.setup.success",
+                       (environment_obj.name, environment_obj.id)))
+            except EnvironmentDoesNotExist:
+                self.cli_helper.echo(
+                    __("error", "cli.environment.setup.argument", input_environment_name))
         return self.project_controller.model
 
     def version(self):

datmo/cli/command/tests/test_environment.py

@@ -49,7 +49,11 @@ def __set_variables(self):
         self.project = ProjectCommand(self.temp_dir, self.cli_helper)
         self.project.parse(
             ["init", "--name", "foobar", "--description", "test model"])
-        self.project.execute()
+        @self.project.cli_helper.input("\n")
+        def dummy(self):
+            return self.project.execute()
+
+        dummy(self)
         self.environment_command = EnvironmentCommand(self.temp_dir,
                                                       self.cli_helper)
 

datmo/cli/command/tests/test_project.py

@@ -48,21 +48,57 @@ def setup_method(self):
     def teardown_method(self):
         pass
 
-    def test_init_create_success(self):
+    def test_init_create_success_no_environment(self):
         test_name = "foobar"
         test_description = "test model"
         self.project_command.parse(
             ["init", "--name", test_name, "--description", test_description])
-        result = self.project_command.execute()
+
+        # Test when environment is not created
+        @self.project_command.cli_helper.input("\n")
+        def dummy_no_environment(self):
+            return self.project_command.execute()
+
+        result = dummy_no_environment(self)
+
+        definition_filepath = os.path.join(self.temp_dir, "datmo_environment", "Dockerfile")
+
+        assert result
+        assert not os.path.isfile(definition_filepath)
+        assert os.path.exists(os.path.join(self.temp_dir, '.datmo'))
+        assert result.name == test_name
+        assert result.description == test_description
+
+    def test_init_create_success_environment(self):
+        test_name = "foobar"
+        test_description = "test model"
+        self.project_command.parse(
+            ["init", "--name", test_name, "--description", test_description])
+
+        # Test when environment is created
+        @self.project_command.cli_helper.input("y\n1\n")
+        def dummy_environment(self):
+            return self.project_command.execute()
+
+        result = dummy_environment(self)
+
+        definition_filepath = os.path.join(self.temp_dir, "datmo_environment", "Dockerfile")
+
+        assert result
+        assert os.path.isfile(definition_filepath)
+        assert "FROM datmo/xgboost:cpu" in open(definition_filepath,
+                                                "r").read()
+
         # test for desired side effects
         assert os.path.exists(os.path.join(self.temp_dir, '.datmo'))
         assert result.name == test_name
         assert result.description == test_description
 
+
     def test_init_create_failure(self):
         self.project_command.parse(["init", "--name", "", "--description", ""])
         # test if prompt opens
-        @self.project_command.cli_helper.input("\n\n")
+        @self.project_command.cli_helper.input("\n\n\n")
         def dummy(self):
             return self.project_command.execute()
 
@@ -74,14 +110,20 @@ def test_init_update_success(self):
         test_description = "test model"
         self.project_command.parse(
             ["init", "--name", test_name, "--description", test_description])
-        result_1 = self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        result_1 = dummy(self)
         updated_name = "foobar2"
         updated_description = "test model 2"
         self.project_command.parse([
             "init", "--name", updated_name, "--description",
             updated_description
         ])
-        result_2 = self.project_command.execute()
+
+        result_2 = dummy(self)
         # test for desired side effects
         assert os.path.exists(os.path.join(self.temp_dir, '.datmo'))
         assert result_2.id == result_1.id
@@ -93,12 +135,17 @@ def test_init_update_success_only_name(self):
         test_description = "test model"
         self.project_command.parse(
             ["init", "--name", test_name, "--description", test_description])
-        result_1 = self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        result_1 = dummy(self)
         updated_name = "foobar2"
         self.project_command.parse(
             ["init", "--name", updated_name, "--description", ""])
 
-        @self.project_command.cli_helper.input("\n")
+        @self.project_command.cli_helper.input("\n\n")
         def dummy(self):
             return self.project_command.execute()
 
@@ -114,12 +161,17 @@ def test_init_update_success_only_description(self):
         test_description = "test model"
         self.project_command.parse(
             ["init", "--name", test_name, "--description", test_description])
-        result_1 = self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        result_1 = dummy(self)
         updated_description = "test model 2"
         self.project_command.parse(
             ["init", "--name", "", "--description", updated_description])
 
-        @self.project_command.cli_helper.input("\n")
+        @self.project_command.cli_helper.input("\n\n")
         def dummy(self):
             return self.project_command.execute()
 
@@ -157,7 +209,12 @@ def test_status(self):
         test_description = "test model"
         self.project_command.parse(
             ["init", "--name", test_name, "--description", test_description])
-        _ = self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        _ = dummy(self)
 
         self.project_command.parse(["status"])
         result = self.project_command.execute()
@@ -178,11 +235,16 @@ def test_cleanup(self):
         test_description = "test model"
         self.project_command.parse(
             ["init", "--name", test_name, "--description", test_description])
-        _ = self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        _ = dummy(self)
 
         self.project_command.parse(["cleanup"])
 
-        @self.project_command.cli_helper.input("y\n")
+        @self.project_command.cli_helper.input("y\n\n")
         def dummy(self):
             return self.project_command.execute()
 
@@ -203,7 +265,12 @@ def test_cleanup_invalid_arg(self):
     def test_notebook(self):
         self.project_command.parse(
             ["init", "--name", "foobar", "--description", "test model"])
-        self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            self.project_command.execute()
+
+        dummy(self)
 
         # Create environment_driver definition
         env_def_path = os.path.join(self.temp_dir, "Dockerfile")

datmo/cli/command/tests/test_session.py

@@ -36,7 +36,12 @@ def __set_variables(self):
         self.project_command = ProjectCommand(self.temp_dir, self.cli_helper)
         self.project_command.parse(
             ["init", "--name", "foobar", "--description", "test model"])
-        self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        dummy(self)
         self.session_command = SessionCommand(self.temp_dir, self.cli_helper)
 
     def test_session_project_not_init(self):

datmo/cli/command/tests/test_snapshot.py

@@ -59,7 +59,12 @@ def __set_variables(self):
         self.project_command = ProjectCommand(self.temp_dir, self.cli_helper)
         self.project_command.parse(
             ["init", "--name", "foobar", "--description", "test model"])
-        self.project_command.execute()
+
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        dummy(self)
         self.snapshot_command = SnapshotCommand(self.temp_dir, self.cli_helper)
 
         # Create environment_driver definition

datmo/cli/command/tests/test_task.py

@@ -60,7 +60,11 @@ def __set_variables(self):
         self.project_command = ProjectCommand(self.temp_dir, self.cli_helper)
         self.project_command.parse(
             ["init", "--name", "foobar", "--description", "test model"])
-        self.project_command.execute()
+        @self.project_command.cli_helper.input("\n")
+        def dummy(self):
+            return self.project_command.execute()
+
+        dummy(self)
 
         self.task_command = TaskCommand(self.temp_dir, self.cli_helper)
 

datmo/core/controller/environment/driver/dockerenv.py

@@ -140,7 +140,8 @@ def setup(self, options, definition_path):
         available_environment_names, available_environment_description = zip(*available_environments)
         # Validate that the name exists
         if not name or name not in available_environment_names:
-            raise EnvironmentDoesNotExist()
+            raise EnvironmentDoesNotExist(__("error", "controller.environment.driver.docker.setup.dne", name))
+
         # Validate the given definition path exists
         if not os.path.isdir(definition_path):
             raise PathDoesNotExist()

datmo/core/util/lang/en.py

@@ -108,7 +108,8 @@
         "cli.project.notebook":
             "Error while running the notebook with task: %s",
         "cli.environment.setup.argument":
-            "This name or index does not match any supported environments: %s",
+            "This name or index does not match any supported environments: %s, "
+            "Update Dockerfile with image name in environment folder",
         "cli.task.run":
             "Error while running the task: %s",
         "cli.task.run.already_running":
@@ -215,6 +216,8 @@
             "Error in getting tags: %s",
         "controller.environment.driver.docker.create.dne":
             "path does not exist: %s",
+        "controller.environment.driver.docker.setup.dne":
+            "input environment does not exist: %s, Update Dockerfile with image name in environment folder",
         "controller.environment.driver.docker.create.exists":
             "output path already exists: %s",
         "controller.environment.driver.docker.build_image":
@@ -350,6 +353,8 @@
             "Is it okay?",
         "cli.project.cleanup.confirm":
             "Are you sure you want to delete all datmo project information? [yN]",
+        "cli.project.environment.setup":
+            "Would you like to setup the environment? [yN]"
     },
     "argparser": {
         "cli.datmo.usage":