diff --git a/sdk/python/kfp/compiler/_container_builder.py b/sdk/python/kfp/compiler/_container_builder.py
index b1233e9945c..924eed1a63a 100644
--- a/sdk/python/kfp/compiler/_container_builder.py
+++ b/sdk/python/kfp/compiler/_container_builder.py
@@ -71,15 +71,10 @@ def _generate_kaniko_spec(self, context, docker_filename, target_image):
 
   def _wrap_dir_in_tarball(self, tarball_path, dir_name):
     """ _wrap_files_in_tarball creates a tarball for all the files in the directory"""
-    old_wd = os.getcwd()
-    os.chdir(dir_name)
     if not tarball_path.endswith('.tar.gz'):
       raise ValueError('the tarball path should end with .tar.gz')
     with tarfile.open(tarball_path, 'w:gz') as tarball:
-      for f in os.listdir(dir_name):
-        tarball.add(f)
-    os.chdir(old_wd)
-
+      tarball.add(dir_name, arcname='')
 
   def build(self, local_dir, docker_filename, target_image, timeout):
     """
diff --git a/sdk/python/tests/compiler/container_builder_test.py b/sdk/python/tests/compiler/container_builder_test.py
index debebcc7738..13ed3231707 100644
--- a/sdk/python/tests/compiler/container_builder_test.py
+++ b/sdk/python/tests/compiler/container_builder_test.py
@@ -42,9 +42,8 @@ def test_wrap_dir_in_tarball(self):
     self.assertTrue(os.path.exists(temp_tarball))
     with tarfile.open(temp_tarball) as temp_tarball_handle:
       temp_files = temp_tarball_handle.getmembers()
-      self.assertTrue(len(temp_files) == 2)
       for temp_file in temp_files:
-        self.assertTrue(temp_file.name in ['test_data_one.tmp', 'test_data_two.tmp'])
+        self.assertTrue(temp_file.name in ['test_data_one.tmp', 'test_data_two.tmp', ''])
 
     # clean up
     os.remove(temp_tarball)