Merge pull request #115 from datmo/flow-issues

Combined basic flow issues resolution

datmo/cli/command/snapshot.py

@@ -25,7 +25,6 @@ def snapshot(self):
 
     def create(self, **kwargs):
         self.cli_helper.echo(__("info", "cli.snapshot.create"))
-
         task_id = kwargs.get("task_id", None)
         # creating snapshot with task id if it exists
         if task_id is not None:
@@ -45,32 +44,50 @@ def create(self, **kwargs):
             # Create a new core snapshot object
             snapshot_task_obj = self.snapshot_controller.create_from_task(
                 message, task_id, label=label)
+            self.cli_helper.echo(
+                "Created snapshot id: %s" % snapshot_task_obj.id)
             return snapshot_task_obj.id
         else:
             # creating snapshot without task id
             snapshot_dict = {"visible": True}
 
             # Code
-            mutually_exclusive_args = ["code_id", "commit_id"]
-            mutually_exclusive(mutually_exclusive_args, kwargs, snapshot_dict)
+            if kwargs.get("code_id", None) or kwargs.get("commit_id", None):
+                mutually_exclusive_args = ["code_id", "commit_id"]
+                mutually_exclusive(mutually_exclusive_args, kwargs,
+                                   snapshot_dict)
 
             # Environment
-            mutually_exclusive_args = [
-                "environment_id", "environment_definition_filepath"
-            ]
-            mutually_exclusive(mutually_exclusive_args, kwargs, snapshot_dict)
+            if kwargs.get("environment_id", None) or kwargs.get(
+                    "environment_definition_filepath", None):
+                mutually_exclusive_args = [
+                    "environment_id", "environment_definition_filepath"
+                ]
+                mutually_exclusive(mutually_exclusive_args, kwargs,
+                                   snapshot_dict)
 
             # File
-            mutually_exclusive_args = ["file_collection_id", "filepaths"]
-            mutually_exclusive(mutually_exclusive_args, kwargs, snapshot_dict)
+            if kwargs.get("file_collection_id", None) or kwargs.get(
+                    "filepaths", None):
+                mutually_exclusive_args = ["file_collection_id", "filepaths"]
+                mutually_exclusive(mutually_exclusive_args, kwargs,
+                                   snapshot_dict)
 
             # Config
-            mutually_exclusive_args = ["config_filepath", "config_filename"]
-            mutually_exclusive(mutually_exclusive_args, kwargs, snapshot_dict)
+            if kwargs.get("config_filepath", None) or kwargs.get(
+                    "config_filename", None):
+                mutually_exclusive_args = [
+                    "config_filepath", "config_filename"
+                ]
+                mutually_exclusive(mutually_exclusive_args, kwargs,
+                                   snapshot_dict)
 
             # Stats
-            mutually_exclusive_args = ["stats_filepath", "stats_filename"]
-            mutually_exclusive(mutually_exclusive_args, kwargs, snapshot_dict)
+            if kwargs.get("stats_filepath", None) or kwargs.get(
+                    "stats_filename", None):
+                mutually_exclusive_args = ["stats_filepath", "stats_filename"]
+                mutually_exclusive(mutually_exclusive_args, kwargs,
+                                   snapshot_dict)
 
             optional_args = ["session_id", "message", "label"]
 
@@ -79,12 +96,15 @@ def create(self, **kwargs):
                     snapshot_dict[arg] = kwargs[arg]
 
             snapshot_obj = self.snapshot_controller.create(snapshot_dict)
-
+            self.cli_helper.echo(
+                __("info", "cli.snapshot.create.success", snapshot_obj.id))
             return snapshot_obj.id
 
     def delete(self, **kwargs):
         self.cli_helper.echo(__("info", "cli.snapshot.delete"))
         snapshot_id = kwargs.get("id", None)
+        self.cli_helper.echo(
+            __("info", "cli.snapshot.delete.success", snapshot_id))
         return self.snapshot_controller.delete(snapshot_id)
 
     def ls(self, **kwargs):
@@ -134,4 +154,8 @@ def ls(self, **kwargs):
 
     def checkout(self, **kwargs):
         snapshot_id = kwargs.get("id", None)
+        checkout_success = self.snapshot_controller.checkout(snapshot_id)
+        if checkout_success:
+            self.cli_helper.echo(
+                __("info", "cli.snapshot.checkout.success", snapshot_id))
         return self.snapshot_controller.checkout(snapshot_id)

datmo/cli/command/task.py

@@ -3,15 +3,18 @@
 import shlex
 import platform
 import prettytable
+# https://stackoverflow.com/questions/11301138/how-to-check-if-variable-is-string-with-python-2-and-3-compatibility/11301392?utm_medium=organic&utm_source=google_rich_qa&utm_campaign=google_rich_qa
 try:
-    to_unicode = unicode
+    basestring
 except NameError:
-    to_unicode = str
+    basestring = str
 
 from datmo.core.util.i18n import get as __
+from datmo.core.util.misc_functions import mutually_exclusive
 from datmo.cli.command.project import ProjectCommand
 from datmo.core.controller.task import TaskController
 from datmo.core.util.logger import DatmoLogger
+from datmo.core.util.exceptions import RequiredArgumentMissing
 
 
 class TaskCommand(ProjectCommand):
@@ -32,11 +35,10 @@ def run(self, **kwargs):
         if kwargs['environment_definition_filepath']:
             snapshot_dict["environment_definition_filepath"] =\
                 kwargs['environment_definition_filepath']
-
         if not isinstance(kwargs['cmd'], list):
             if platform.system() == "Windows":
                 kwargs['cmd'] = kwargs['cmd']
-            elif isinstance(kwargs['cmd'], to_unicode):
+            elif isinstance(kwargs['cmd'], basestring):
                 kwargs['cmd'] = shlex.split(kwargs['cmd'])
 
         task_dict = {
@@ -56,33 +58,57 @@ def run(self, **kwargs):
             self.logger.error("%s %s" % (e, task_dict))
             self.cli_helper.echo(__("error", "cli.task.run", task_obj.id))
             return False
+        self.cli_helper.echo("Ran task id: %s" % updated_task_obj.id)
         return updated_task_obj
 
     def ls(self, **kwargs):
         session_id = kwargs.get('session_id',
                                 self.task_controller.current_session.id)
         # Get all snapshot meta information
-        header_list = ["id", "command", "results", "created at"]
+        header_list = ["id", "command", "status", "results", "created at"]
         t = prettytable.PrettyTable(header_list)
         task_objs = self.task_controller.list(
             session_id, sort_key='created_at', sort_order='descending')
         for task_obj in task_objs:
             t.add_row([
-                task_obj.id, task_obj.command, task_obj.results,
+                task_obj.id, task_obj.command, task_obj.status,
+                task_obj.results,
                 task_obj.created_at.strftime("%Y-%m-%d %H:%M:%S")
             ])
         self.cli_helper.echo(t)
 
         return True
 
     def stop(self, **kwargs):
-        task_id = kwargs.get('id', None)
-        self.cli_helper.echo(__("info", "cli.task.stop", task_id))
+        input_dict = {}
+        mutually_exclusive(["id", "all"], kwargs, input_dict)
+        if "id" in input_dict:
+            self.cli_helper.echo(__("info", "cli.task.stop", input_dict['id']))
+        elif "all" in input_dict:
+            self.cli_helper.echo(__("info", "cli.task.stop.all"))
+        else:
+            raise RequiredArgumentMissing()
         try:
-            result = self.task_controller.stop(task_id)
-            if not result:
-                self.cli_helper.echo(__("error", "cli.task.stop", task_id))
+            if "id" in input_dict:
+                result = self.task_controller.stop(task_id=input_dict['id'])
+                if not result:
+                    self.cli_helper.echo(
+                        __("error", "cli.task.stop", input_dict['id']))
+                else:
+                    self.cli_helper.echo(
+                        __("info", "cli.task.stop.success", input_dict['id']))
+            if "all" in input_dict:
+                result = self.task_controller.stop(all=input_dict['all'])
+                if not result:
+                    self.cli_helper.echo(__("error", "cli.task.stop.all"))
+                else:
+                    self.cli_helper.echo(
+                        __("info", "cli.task.stop.all.success"))
             return result
         except:
-            self.cli_helper.echo(__("error", "cli.task.stop", task_id))
+            if "id" in input_dict:
+                self.cli_helper.echo(
+                    __("error", "cli.task.stop", input_dict['id']))
+            if "all" in input_dict:
+                self.cli_helper.echo(__("error", "cli.task.stop.all"))
             return False

datmo/cli/command/tests/test_snapshot.py

@@ -12,6 +12,7 @@
 #     # Python 3
 #     import builtins as __builtin__
 import os
+import shutil
 import tempfile
 import platform
 from io import open
@@ -539,11 +540,14 @@ def test_datmo_snapshot_checkout_invalid_arg(self):
 
     def test_datmo_snapshot_checkout(self):
         self.__set_variables()
-
         # Test when optional parameters are not given
         self.snapshot.parse(["snapshot", "create", "-m", "my test snapshot"])
         snapshot_id = self.snapshot.execute()
 
+        # remove datmo_task folder to have no changes before checkout
+        datmo_tasks_dirpath = os.path.join(self.snapshot.home, "datmo_tasks")
+        shutil.rmtree(datmo_tasks_dirpath)
+
         # Test when optional parameters are not given
         self.snapshot.parse(["snapshot", "checkout", "--id", snapshot_id])
 

datmo/cli/command/tests/test_task.py

@@ -28,7 +28,7 @@
 from datmo.cli.command.project import ProjectCommand
 from datmo.cli.command.task import TaskCommand
 from datmo.core.entity.task import Task as CoreTask
-from datmo.core.util.exceptions import ProjectNotInitializedException
+from datmo.core.util.exceptions import ProjectNotInitializedException, MutuallyExclusiveArguments, RequiredArgumentMissing
 
 
 class TestTaskCommand():
@@ -77,17 +77,16 @@ def test_task_run_should_fail1(self):
         result = self.task.execute()
         assert not result
 
-    def test_task_run_should_fail2(self):
+    def test_task_run(self):
+        # TODO: Adding test with `--interactive` argument and terminate inside container
         self.__set_variables()
-        # Test failure case execute
+
+        # Test failure command execute
         test_command = ["yo", "yo"]
         self.task.parse(["task", "run", test_command])
         result = self.task.execute()
-        assert not result
+        assert result
 
-    def test_task_run(self):
-        # TODO: Adding test with `--interactive` argument and terminate inside container
-        self.__set_variables()
         # Test success case
         test_command = ["sh", "-c", "echo accuracy:0.45"]
         test_ports = ["8888:8888", "9999:9999"]
@@ -135,7 +134,6 @@ def test_task_run_string_command(self):
             "task", "run", "--ports", test_ports[0], "--ports", test_ports[1],
             "--environment-def", test_dockerfile, test_command
         ])
-
         # test for desired side effects
         assert self.task.args.cmd == test_command
         assert self.task.args.ports == test_ports
@@ -236,8 +234,13 @@ def test_task_run_invalid_arg(self):
 
     def test_task_ls(self):
         self.__set_variables()
-        test_session_id = 'test_session_id'
 
+        self.task.parse(["task", "ls"])
+        task_ls_command = self.task.execute()
+
+        assert task_ls_command == True
+
+        test_session_id = 'test_session_id'
         self.task.parse(["task", "ls", "--session-id", test_session_id])
 
         # test for desired side effects
@@ -255,7 +258,9 @@ def test_task_ls_invalid_arg(self):
             exception_thrown = True
         assert exception_thrown
 
-    def test_task_stop(self):
+    def test_task_stop_success(self):
+        # 1) Test stop with task_id
+        # 2) Test stop with all
         self.__set_variables()
 
         test_command = ["sh", "-c", "echo yo"]
@@ -269,6 +274,7 @@ def test_task_stop(self):
 
         test_task_obj = self.task.execute()
 
+        # 1) Test option 1
         self.task.parse(["task", "stop", "--id", test_task_obj.id])
 
         # test for desired side effects
@@ -278,20 +284,49 @@ def test_task_stop(self):
         task_stop_command = self.task.execute()
         assert task_stop_command == True
 
-    def test_task_stop_invalid_task_id(self):
+        # 2) Test option 2
+        self.task.parse(["task", "stop", "--all"])
+
+        # test when all is passed to stop all
+        task_stop_command = self.task.execute()
+        assert task_stop_command == True
+
+    def test_task_stop_failure_required_args(self):
         self.__set_variables()
         # Passing wrong task id
-        self.task.parse(["task", "stop", "--id", "invalid-task-id"])
+        self.task.parse(["task", "stop"])
+        failed = False
+        try:
+            _ = self.task.execute()
+        except RequiredArgumentMissing:
+            failed = True
+        assert failed
 
-        # test when wrong task id is passed to stop it
-        result = self.task.execute()
-        assert not result
+    def test_task_stop_failure_mutually_exclusive_vars(self):
+        self.__set_variables()
+        # Passing wrong task id
+        self.task.parse(["task", "stop", "--id", "invalid-task-id", "--all"])
+        failed = False
+        try:
+            _ = self.task.execute()
+        except MutuallyExclusiveArguments:
+            failed = True
+        assert failed
 
-    def test_task_stop_invalid_arg(self):
+    def test_task_stop_failure_invalid_arg(self):
         self.__set_variables()
         exception_thrown = False
         try:
             self.task.parse(["task", "stop" "--foobar", "foobar"])
         except Exception:
             exception_thrown = True
         assert exception_thrown
+
+    def test_task_stop_invalid_task_id(self):
+        self.__set_variables()
+        # Passing wrong task id
+        self.task.parse(["task", "stop", "--id", "invalid-task-id"])
+
+        # test when wrong task id is passed to stop it
+        result = self.task.execute()
+        assert not result

datmo/cli/parser.py

@@ -215,5 +215,11 @@ def get_datmo_parser():
     task_stop = task_subcommand_parsers.add_parser("stop", help="stop tasks")
     task_stop.add_argument(
         "--id", dest="id", default=None, type=str, help="task id to stop")
+    task_stop.add_argument(
+        "--all",
+        "-a",
+        dest="all",
+        action="store_true",
+        help="stop all datmo tasks")
 
     return parser

datmo/core/controller/code/code.py

@@ -49,7 +49,6 @@ def create(self, commit_id=None):
         create_dict = {
             "model_id": self.model.id,
         }
-
         ## Required args for Code entity
         required_args = ["driver_type", "commit_id"]
         for required_arg in required_args: