diff --git a/dataikuapi/dss/savedmodel.py b/dataikuapi/dss/savedmodel.py index ce4f65d3..b2808f98 100644 --- a/dataikuapi/dss/savedmodel.py +++ b/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,7 +121,7 @@ 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. @@ -127,7 +129,7 @@ def import_mlflow_version_from_path(self, version_id, path, code_env_name = "INH :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) diff --git a/dataikuapi/utils.py b/dataikuapi/utils.py index e4512acc..6503fdb8 100644 --- a/dataikuapi/utils.py +++ b/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