Replace shutil.make_archive by custom zip function

dataikuapi/dss/savedmodel.py

@@ -5,6 +5,8 @@
 from .ml import DSSTrainedClusteringModelDetails
 from .ml import DSSTrainedPredictionModelDetails
 
+from ..utils import make_zipfile
+
 try:
     basestring
 except NameError:
@@ -119,15 +121,15 @@ def get_origin_ml_task(self):
         if fmi is not None:
             return DSSMLTask.from_full_model_id(self.client, fmi, project_key=self.project_key)
 
-    def import_mlflow_version_from_path(self, version_id, path, code_env_name = "INHERIT"):
+    def import_mlflow_version_from_path(self, version_id, path, code_env_name="INHERIT"):
         """
         Create a new version for this saved model from a path containing a MLFlow model.
 
         Requires the saved model to have been created using :meth:`dataikuapi.dss.project.DSSProject.create_mlflow_pyfunc_model`.
 
         :param str version_id: Identifier of the version to create
         :param str path: An absolute path on the local filesystem. Must be a folder, and must contain a MLFlow model
-        :param str code_env_name: Name of the code env to use for this model version. The code env must contain at least 
+        :param str code_env_name: Name of the code env to use for this model version. The code env must contain at least
                                   mlflow and the package(s) corresponding to the used MLFlow-compatible frameworks.
                                   If value is "INHERIT", the default active code env of the project will be used
         :return a :class:MLFlowVersionHandler in order to interact with the new MLFlow model version
@@ -136,13 +138,14 @@ def import_mlflow_version_from_path(self, version_id, path, code_env_name = "INH
         # TODO: cleanup the archive
         import shutil
         import os
+
         archive_temp_dir = tempfile.mkdtemp()
         try:
-            archive_filename = shutil.make_archive(os.path.join(archive_temp_dir, "tmpmodel"), "zip", path) #[, root_dir[, base_dir[, verbose[, dry_run[, owner[, group[, logger]]]]]]])
+            archive_filename = make_zipfile(os.path.join(archive_temp_dir, "tmpmodel.zip"), path)
 
             with open(archive_filename, "rb") as fp:
                 self.client._perform_empty("POST", "/projects/%s/savedmodels/%s/versions/%s?codeEnvName=%s" % (self.project_key, self.sm_id, version_id, code_env_name),
-                    files={"file":(archive_filename, fp)})
+                                           files={"file": (archive_filename, fp)})
             return self.get_mlflow_version_handler(version_id)
         finally:
             shutil.rmtree(archive_temp_dir)

dataikuapi/utils.py

@@ -1,7 +1,8 @@
 import csv, sys
 from dateutil import parser as date_iso_parser
 from contextlib import closing
-
+import os
+import zipfile
 import itertools
 
 if sys.version_info > (3,0):
@@ -101,3 +102,18 @@ def __init__(self, val):
 
     def __call__(self):
         return self.val
+
+
+def make_zipfile(output_filename, source_dir):
+    """Replace shutil.make_archive which adds undesired folders to the archive
+    in python 2.7 in some environments.
+    """
+    relroot = os.path.abspath(os.path.join(source_dir))
+    with zipfile.ZipFile(output_filename, "w", zipfile.ZIP_DEFLATED) as zipfp:
+        for root, dirs, files in os.walk(source_dir):
+            for file in files:
+                filename = os.path.join(root, file)
+                if os.path.isfile(filename):
+                    arcname = os.path.join(os.path.relpath(root, relroot), file)
+                    zipfp.write(filename, arcname)
+    return output_filename