fix: preserve Docker runner exit status

Signed-off-by: Zakhar Dvurechensky <72825626+Zakharden@users.noreply.github.com>

Author: Zakharden

sdk/python/kfp/local/docker_task_handler.py

@@ -140,21 +140,16 @@ def run_docker_container(client: 'docker.DockerClient', image: str,
         stdout=True,
         stderr=True,
         volumes=volumes,
-        auto_remove=True,
+        auto_remove=False,
         **container_run_args)
     try:
         for line in container.logs(stream=True):
             # the inner logs should already have trailing \n
             # we do not need to add another
             print(line.decode(), end='')
-    except docker.errors.NotFound:
-        # Container was auto-removed before we could stream logs
-        # This can happen if the container exits very quickly
-        pass
-    try:
         return container.wait()['StatusCode']
-    except docker.errors.NotFound:
-        # Container was auto-removed after logs completed
-        # This means the container exited successfully (logs streaming completed)
-        # We assume success since we couldn't get the actual status code
-        return 0
+    finally:
+        try:
+            container.remove(force=True)
+        except docker.errors.NotFound:
+            pass

sdk/python/kfp/local/docker_task_handler_test.py

@@ -66,7 +66,7 @@ def test_no_volumes(self):
             stdout=True,
             stderr=True,
             volumes={},
-            auto_remove=True,
+            auto_remove=False,
         )
 
     def test_cwd_volume(self):
@@ -90,8 +90,25 @@ def test_cwd_volume(self):
                 'bind': '/localdir',
                 'mode': 'ro'
             }},
-            auto_remove=True,
+            auto_remove=False,
         )
+        self.mocked_docker_client.containers.run.return_value.remove.assert_called_once_with(
+            force=True)
+
+    def test_returns_container_exit_status_before_removing_container(self):
+        mock_container = self.mocked_docker_client.containers.run.return_value
+        mock_container.wait.return_value = {'StatusCode': 1}
+
+        return_code = docker_task_handler.run_docker_container(
+            docker.from_env(),
+            image='alpine',
+            command=['false'],
+            volumes={},
+        )
+
+        self.assertEqual(return_code, 1)
+        mock_container.wait.assert_called_once()
+        mock_container.remove.assert_called_once_with(force=True)
 
 
 class TestDockerTaskHandler(DockerMockTestCase):
@@ -239,7 +256,7 @@ def test_run(self):
                         'mode': 'rw'
                     }
                 },
-                auto_remove=True,
+                auto_remove=False,
             )
 
     def test_run_passes_env_vars(self):