environment workspace

datmo/cli/command/environment.py

@@ -28,7 +28,7 @@ def setup(self, **kwargs):
         # TODO: remove business logic from here and create common helper
         # environment types
         environment_types = self.environment_controller.get_environment_types()
-        if not environment_type or environment_type not in environment_types:
+        if environment_types and environment_type not in environment_types:
             environment_type = self.cli_helper.prompt_available_options(
                 environment_types, option_type="type")
         # environment frameworks
@@ -37,13 +37,13 @@ def setup(self, **kwargs):
         available_frameworks = [
             item[0] for item in available_framework_details
         ]
-        if not environment_framework or environment_framework not in available_frameworks:
+        if available_frameworks and environment_framework not in available_frameworks:
             environment_framework = self.cli_helper.prompt_available_options(
                 available_framework_details, option_type="framework")
         # environment languages
         available_environment_languages = self.environment_controller.get_supported_languages(
             environment_type, environment_framework)
-        if available_environment_languages and not environment_language or environment_language not in available_environment_languages:
+        if available_environment_languages and environment_language not in available_environment_languages:
             environment_language = self.cli_helper.prompt_available_options(
                 available_environment_languages, option_type="language")
         try:

datmo/cli/command/run.py

@@ -404,7 +404,8 @@ def rerun(self, **kwargs):
             "ports": task_obj.ports,
             "interactive": task_obj.interactive,
             "mem_limit": task_obj.mem_limit,
-            "command_list": command
+            "command_list": command,
+            "workspace": task_obj.workspace
         }
         # Run task and return Task object result
         new_task_obj = self.task_run_helper(task_dict, snapshot_dict,

datmo/cli/command/tests/test_environment.py

@@ -8,6 +8,7 @@
 import os
 import glob
 import time
+import pytest
 import uuid
 import tempfile
 import shutil
@@ -138,6 +139,19 @@ def dummy(self):
         assert "FROM datmo/data-analytics:cpu-py27" in open(
             definition_filepath, "r").read()
 
+        # Test success with correct prompt input using numbers
+        self.environment_command.parse(["environment", "setup"])
+
+        @self.environment_command.cli_helper.input(
+            "cpu\ncaffe2\n")
+        def dummy(self):
+            return self.environment_command.execute()
+        result = dummy(self)
+        assert result
+        assert os.path.isfile(definition_filepath)
+        assert "FROM datmo/caffe2:cpu" in open(
+            definition_filepath, "r").read()
+
         # Test success with correct prompt input using string
         self.environment_command.parse(["environment", "setup"])
 
@@ -166,16 +180,38 @@ def dummy(self):
         # Test with prompt input string incorrect
         self.environment_command.parse(["environment", "setup"])
 
-        @self.environment_command.cli_helper.input("random\n\n\n")
+        # quit at environment type
+        @self.environment_command.cli_helper.input("quit\n")
         def dummy(self):
             return self.environment_command.execute()
 
-        result = dummy(self)
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            dummy(self)
 
-        assert result
-        assert os.path.isfile(definition_filepath)
-        assert "FROM datmo/python-base:cpu-py27" in open(
-            definition_filepath, "r").read()
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
+
+        # quit at environment framework
+        @self.environment_command.cli_helper.input("\nquit\n")
+        def dummy(self):
+            return self.environment_command.execute()
+
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            dummy(self)
+
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
+
+        # quit at environment language
+        @self.environment_command.cli_helper.input("\n\nquit\n")
+        def dummy(self):
+            return self.environment_command.execute()
+
+        with pytest.raises(SystemExit) as pytest_wrapped_e:
+            dummy(self)
+
+        assert pytest_wrapped_e.type == SystemExit
+        assert pytest_wrapped_e.value.code == 0
 
     def test_environment_create(self):
         # 1) Environment definition file in project environment directory (with name / description)

datmo/cli/driver/helper.py

@@ -157,6 +157,10 @@ def get_command_choices(self):
 
     def prompt_available_options(self, available_options, option_type):
         """Prompt user to choose an available environment. Returns the environment name"""
+        # If there are no options
+        if not available_options:
+            return None
+        # If there exists list of options
         if option_type == "framework":
             available_options_info = available_options
             available_options = []
@@ -171,6 +175,9 @@ def prompt_available_options(self, available_options, option_type):
         option_environment_index = None
         valid_option = False
         while input_environment_option is not None and not valid_option:
+            # exit when user wants to quit
+            if input_environment_option in ["q", "quit"]:
+                sys.exit(0)
             try:
                 option_environment_index = int(input_environment_option)
                 if 0 < option_environment_index <= len(available_options):

datmo/core/controller/environment/environment.py

@@ -125,11 +125,17 @@ def setup(self, options, save_hardware_file=True):
                 definition_path=self.file_driver.environment_directory)
         except Exception:
             raise
+        name = options.get('name', None)
+        if name is None:
+            environment_framework = options['environment_framework']
+            environment_type = options['environment_type']
+            environment_language = options['environment_language']
+            if environment_language:
+                name = "%s:%s-%s" % (environment_framework, environment_type, environment_language)
+            else:
+                name = "%s:%s" % (environment_framework, environment_type)
         create_dict = {
-            "name":
-                options['name'] if options.get('name', None) else "%s:%s-%s" %
-                (options['environment_framework'], options['environment_type'],
-                 options['environment_language']),
+            "name": name,
             "description":
                 "supported environment created by datmo"
         }