removing environment artifacts after the run and updating the tests

datmo/cli/command/run.py

@@ -265,6 +265,8 @@ def run(self, **kwargs):
 
         # Create the task object
         task_obj = self.task_controller.create()
+
+        updated_task_obj = task_obj
         try:
             # Pass in the task to run
             updated_task_obj = self.task_controller.run(
@@ -274,9 +276,13 @@ def run(self, **kwargs):
             self.cli_helper.echo("%s" % e)
             self.cli_helper.echo(__("error", "cli.task.run", task_obj.id))
             return False
+        finally:
+            self.cli_helper.echo(
+                __("info", "cli.task.run.stop"))
+            self.task_controller.stop(updated_task_obj.id)
+            self.cli_helper.echo(
+                __("info", "cli.task.run.complete", updated_task_obj.id))
 
-        self.cli_helper.echo(
-            __("info", "cli.task.run.complete", updated_task_obj.id))
         return updated_task_obj
 
     @Helper.notify_no_project_found

datmo/cli/command/tests/test_run.py

@@ -160,7 +160,7 @@ def test_run(self):
         assert "accuracy" in result.logs
         assert result.results
         assert result.results == {"accuracy": "0.45"}
-        assert result.status == "SUCCESS"
+        assert result.status in ['SUCCESS','STOPPED']
 
         # teardown
         self.task_command.parse(["task", "stop", "--all"])
@@ -195,7 +195,7 @@ def test_run_string_command(self):
         assert "accuracy" in result.logs
         assert result.results
         assert result.results == {"accuracy": "0.45"}
-        assert result.status == "SUCCESS"
+        assert result.status in ["SUCCESS", "STOPPED"]
 
         # teardown
         self.task_command.parse(["task", "stop", "--all"])
@@ -242,7 +242,7 @@ def test_run_string_command(self):
     #         assert "accuracy" in result.logs
     #         assert result.results
     #         assert result.results == {"accuracy": "0.45"}
-    #         assert result.status == "SUCCESS"
+    #         assert result.status in ['SUCCESS','STOPPED']
 
     @pytest_docker_environment_failed_instantiation(test_datmo_dir)
     def test_run_notebook(self):
@@ -285,7 +285,7 @@ def test_run_notebook(self):
         assert isinstance(result, CoreTask)
         assert result.logs
         assert "Currently running servers" in result.logs
-        assert result.status == "SUCCESS"
+        assert result.status in ['SUCCESS','STOPPED']
 
         # teardown
         self.task_command.parse(["task", "stop", "--all"])
@@ -323,7 +323,7 @@ def test_run_ls(self):
         self.run_command.parse(["ls"])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
 
         test_session_id = 'test_session_id'
         self.run_command.parse(["ls", "--session-id", test_session_id])
@@ -351,13 +351,13 @@ def test_run_ls(self):
         self.run_command.parse(["ls", "--format", "csv"])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
 
         # Test success format csv, download default
         self.run_command.parse(["ls", "--format", "csv", "--download"])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
         test_wildcard = os.path.join(os.getcwd(), "run_ls_*")
         paths = [n for n in glob.glob(test_wildcard) if os.path.isfile(n)]
         assert paths
@@ -371,7 +371,7 @@ def test_run_ls(self):
         ])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
         assert os.path.isfile(test_path)
         assert open(test_path, "r").read()
         os.remove(test_path)
@@ -380,13 +380,13 @@ def test_run_ls(self):
         self.run_command.parse(["ls"])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
 
         # Test success format table, download default
         self.run_command.parse(["ls", "--download"])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
         test_wildcard = os.path.join(os.getcwd(), "run_ls_*")
         paths = [n for n in glob.glob(test_wildcard) if os.path.isfile(n)]
         assert paths
@@ -399,7 +399,7 @@ def test_run_ls(self):
             ["ls", "--download", "--download-path", test_path])
         run_objs = self.run_command.execute()
         assert run_objs
-        assert run_objs[0].status == 'SUCCESS'
+        assert run_objs[0].status in ['SUCCESS','STOPPED']
         assert os.path.isfile(test_path)
         assert open(test_path, "r").read()
         os.remove(test_path)

datmo/cli/command/tests/test_task.py

@@ -135,7 +135,7 @@ def test_task_run(self):
         assert "accuracy" in result.logs
         assert result.results
         assert result.results == {"accuracy": "0.45"}
-        assert result.status == "SUCCESS"
+        assert result.status in ['SUCCESS','STOPPED']
 
         # teardown
         self.task_command.parse(["task", "stop", "--all"])
@@ -170,7 +170,7 @@ def test_task_run_string_command(self):
         assert "accuracy" in result.logs
         assert result.results
         assert result.results == {"accuracy": "0.45"}
-        assert result.status == "SUCCESS"
+        assert result.status in ['SUCCESS','STOPPED']
 
         # teardown
         self.task_command.parse(["task", "stop", "--all"])
@@ -217,7 +217,7 @@ def test_task_run_string_command(self):
     #         assert "accuracy" in result.logs
     #         assert result.results
     #         assert result.results == {"accuracy": "0.45"}
-    #         assert result.status == "SUCCESS"
+    #         assert result.status in ['SUCCESS','STOPPED']
 
     @pytest_docker_environment_failed_instantiation(test_datmo_dir)
     def test_task_run_notebook(self):
@@ -259,7 +259,7 @@ def test_task_run_notebook(self):
         assert isinstance(result, CoreTask)
         assert result.logs
         assert "Currently running servers" in result.logs
-        assert result.status == "SUCCESS"
+        assert result.status in ['SUCCESS','STOPPED']
 
         # teardown
         self.task_command.parse(["task", "stop", "--all"])

datmo/core/controller/code/driver/git.py

@@ -492,7 +492,7 @@ def check_unstaged_changes(self):
                     __("error", "controller.code.driver.git.status",
                        str(stderr)))
             stdout = stdout.decode().strip()
-            if "working tree clean" not in stdout:
+            if "clean" not in stdout:
                 raise UnstagedChanges()
         except subprocess.CalledProcessError as e:
             raise GitExecutionError(

datmo/core/controller/environment/environment.py

@@ -337,7 +337,7 @@ def delete(self, environment_id):
         return file_collection_deleted and environment_artifacts_removed and \
                delete_success
 
-    def stop(self, run_id=None, match_string=None, all=False):
+    def stop(self, run_id=None, match_string=None, environment_id=None, all=False):
         """Stop the trace of running environment
 
         Parameters
@@ -348,6 +348,8 @@ def stop(self, run_id=None, match_string=None, all=False):
         match_string : str, optional
             stop environment with a string to match the environment name
             (default is None, which means it is not used)
+        environment_id : str
+            environment object id to remove the artifacts
         all : bool, optional
             stop all environments
 
@@ -384,6 +386,9 @@ def stop(self, run_id=None, match_string=None, all=False):
             all_match_string = "datmo-task-" + self.model.id
             stop_success = self.environment_driver.stop_remove_containers_by_term(
                 term=all_match_string, force=True)
+        # Remove artifacts associated with the environment_driver
+        if environment_id:
+            self.environment_driver.remove(environment_id, force=True)
         return stop_success
 
     def exists(self, environment_id=None, environment_unique_hash=None):

datmo/core/controller/environment/tests/test_environment.py

@@ -35,7 +35,7 @@ def to_bytes(val):
 from datmo.core.util.exceptions import (
     EntityNotFound, RequiredArgumentMissing, TooManyArgumentsFound,
     FileAlreadyExistsError, UnstagedChanges, EnvironmentDoesNotExist,
-    ProjectNotInitialized)
+    ProjectNotInitialized, EnvironmentExecutionError)
 from datmo.core.util.misc_functions import pytest_docker_environment_failed_instantiation
 
 # provide mountable tmp directory for docker
@@ -815,6 +815,7 @@ def test_stop_success(self):
         # 1) Test run_id input to stop
         # 2) Test match_string input to stop
         # 3) Test all input to stop
+        # 4) Test if the image was removed by stop
         self.project_controller.init("test5", "test description")
         self.environment_controller = EnvironmentController()
 
@@ -867,19 +868,26 @@ def test_stop_success(self):
 
         _, run_id, _ = \
             self.environment_controller.run(environment_obj.id, run_options, log_filepath)
-        return_code = self.environment_controller.stop(run_id=run_id)
+        return_code = self.environment_controller.stop(run_id=run_id, environment_id=environment_obj.id)
 
         assert return_code
 
         # 2) Test option 2
+
+        # Rebuild environment in the project
+        _ = self.environment_controller.build(environment_obj.id)
+
         _, _, _ = \
             self.environment_controller.run(environment_obj.id, run_options, log_filepath)
         return_code = self.environment_controller.stop(
-            match_string="datmo-task-" + self.environment_controller.model.id)
+            match_string="datmo-task-" + self.environment_controller.model.id, environment_id=environment_obj.id)
 
         assert return_code
 
         # 3) Test option 3
+
+        # Rebuild environment in the project
+        _ = self.environment_controller.build(environment_obj.id)
         _, _, _ = \
             self.environment_controller.run(environment_obj.id, run_options, log_filepath)
         run_options_2 = {
@@ -903,10 +911,19 @@ def test_stop_success(self):
         }
         _, _, _ = \
             self.environment_controller.run(environment_obj.id, run_options_2, log_filepath)
-        return_code = self.environment_controller.stop(all=True)
+        return_code = self.environment_controller.stop(all=True, environment_id=environment_obj.id)
 
         assert return_code
 
+        # 4) Test option 4
+        failed = False
+        try:
+            _, _, _ = \
+                self.environment_controller.run(environment_obj.id, run_options_2, log_filepath)
+        except EnvironmentExecutionError:
+            failed = True
+        assert failed
+
         # teardown
         self.environment_controller.delete(environment_obj.id)
 

datmo/core/controller/task.py

@@ -473,10 +473,16 @@ def stop(self, task_id=None, all=False):
         if task_id:
             _ = self.dal.task.get_by_id(task_id)  # verify if task_id exists
             task_match_string = "datmo-task-" + self.model.id + "-" + task_id
-            return_code = self.environment.stop(match_string=task_match_string)
+            # Get the environment id associated with the task
+            kwargs = {'match_string': task_match_string}
+            task_obj = self.get(task_id)
+            before_snapshot_id = task_obj.before_snapshot_id
+            if before_snapshot_id:
+                before_snapshot_obj = self.snapshot.get(before_snapshot_id)
+                kwargs['environment_id'] = before_snapshot_obj.environment_id
+            return_code = self.environment.stop(**kwargs)
         if all:
             return_code = self.environment.stop(all=True)
-
         # Set stopped task statuses to STOPPED if return success
         if return_code:
             if task_id: