1import os
2import tempfile
3from os.path import basename, dirname
4from shutil import make_archive, unpack_archive
5from subprocess import run
6
7from omegaml.backends.basedata import BaseDataBackend
8from omegaml.backends.package.packager import RunnablePackageMixin
9
10
[docs]
11class MLFlowProjectBackend(RunnablePackageMixin, BaseDataBackend):
12 """
13 Backend to support storage of MLFlow projects
14
15 Usage::
16
17 om.scripts.put('mlflow://path/to/MLProject', 'myproject')
18 om.scripts.get('myproject')
19
20 See Also:
21 https://www.mlflow.org/docs/latest/projects.html#project-directories
22 """
23
24 KIND = 'mlflow.project'
25 MLFLOW_PREFIX = 'mlflow://'
26
[docs]
27 @classmethod
28 def supports(self, obj, name, **kwargs):
29 is_mlflow_prefix = isinstance(obj, str) and obj.startswith(self.MLFLOW_PREFIX)
30 is_mlflow_kind = kwargs.get('kind') == self.KIND
31 return is_mlflow_kind or is_mlflow_prefix
32
[docs]
33 def put(self, obj, name, attributes=None, **kwargs):
34 """
35 save a MLFlow project
36
37 :param obj: full path to the MLFlow project directory
38 mlflow://path/to/MLProject
39 :param name: name to store
40 :return: the Metadata object
41 """
42 pkgsrc = obj.split(self.MLFLOW_PREFIX)[-1]
43 if os.path.exists(pkgsrc):
44 distdir = os.path.join(pkgsrc, 'dist')
45 os.makedirs(distdir, exist_ok=True)
46 tarfn = os.path.join(distdir, f'{name}')
47 pkgdist = make_archive(tarfn, 'gztar', root_dir=pkgsrc, base_dir='.')
48 else:
49 raise FileNotFoundError(pkgsrc)
50 filename = self.data_store.object_store_key(name, 'pkg', hashed=True)
51 gridfile = self._store_to_file(self.data_store, pkgdist, filename)
52 return self.data_store._make_metadata(
53 name=name,
54 prefix=self.data_store.prefix,
55 bucket=self.data_store.bucket,
56 kind=self.KIND,
57 attributes=attributes,
58 gridfile=gridfile,
59 ).save()
60
[docs]
61 def get(self, name, localpath=None, **kwargs):
62 """
63 Load MLFlow project from store
64
65 This copies the projects's .tar.gz file from om.scripts to a local temp
66 path and returns the MLFlowProject to it
67
68 :param name: the name of the package
69 :param keep: keep the packages load path in sys.path, defaults to False
70 :param localpath: the local path to store the package
71 :param install: if True call pip install on the retrieved package
72 :param kwargs:
73 :return: the loaded module
74 """
75 pkgname = basename(name)
76 dstdir = localpath or self.data_store.tmppath
77 packagefname = f'{os.path.join(localpath or self.packages_path, pkgname)}.tar.gz'
78 os.makedirs(dirname(packagefname), exist_ok=True)
79 meta = self.data_store.metadata(name)
80 outf = meta.gridfile
81 with open(packagefname, 'wb') as pkgf:
82 pkgf.write(outf.read())
83 unpack_archive(packagefname, dstdir)
84 if localpath:
85 mod = localpath
86 else:
87 mod = MLFlowProject(dstdir)
88 return mod
89
90 @property
91 def packages_path(self):
92 return os.path.join(self.data_store.tmppath, 'packages')
93
94
95class MLFlowProject:
96 """a proxy to the MLFlow project that runs a script
97
98 This provides the mod.run() interface for scripts so that
99 we can use the same semantics for mlflow projects and pypi
100 packages
101 """
102
103 def __init__(self, uri):
104 self.uri = uri
105
106 def run(self, om, pure_python=False, **kwargs):
107 kwargs.setdefault('env-manager', 'local')
108 options = ' '.join(f'--{k.replace("_", "-")} {v}' for k, v in kwargs.items())
109 tmpdir = tempfile.mkdtemp()
110 # fix issue
111 with open(os.path.join(tmpdir, 'pyenv'), 'w') as fout:
112 fout.write('#/bin/bash')
113 fout.write('conda activate $1')
114 cmd = rf'PATH={tmpdir}:$PATH; cd {tmpdir}; chmod +x ./pyenv; mlflow run {options} {self.uri}'
115 print(cmd)
116 output = run(cmd, capture_output=True, shell=True)
117 print(output)
118 if output.stderr:
119 output = {
120 'stdout': output.stdout.decode('utf8'),
121 'stderr': output.stderr.decode('utf8'),
122 }
123 else:
124 output = output.stdout.decode('utf8')
125 return {'output': output}