Source code for omegaml.mixins.store.promotion
1from functools import partial
2from pathlib import Path
3
4from omegaml.documents import MDREGISTRY
5from omegaml.util import dict_merge
6
7
[docs]
8class PromotionMixin:
9 """Promote objects from one bucket to another
10
11 Promotion Methods:
12 * `getput` - performs target.put(source.get()) and copies metadata
13 attributes by merging target.metadata's .attributes and .kind_meta
14 * `metadata` - creates a new metadata data entry in target, copying
15 metadata attributes and kind_meta. Does not get/put the object itself
16 (i.e. no associated data is promoted).
17 * `data` - like `getput`, but does not merge metadata
18 * `export` - performs .to_archive() and .from_archive(), effectively
19 copying metadata, the associated gridfile (if available) and collection
20 data (if available). This is equivalent to `om runtime export`.
21
22 The default promotion method is getput(), or the object's backend.PROMOTE
23 method, if specified.
24
25 Some object backends provide a default promotion other than getput:
26
27 * sqlalchemy.conx - uses the `metadata` promotion, effectively copying only
28 metadata. Use promote(..., method='metadata,data') to also promote data
29
30 * virtualobj.dill - uses the `export` promotion, effectively copying all
31 metadata and the associated @virtualobj function. If the source object
32 is a versioned model, this copies the current version and metadata. To
33 copy a specific version, use promote('model@version'). To create a
34 new version in the target bucket use promote(..., method='getput').
35 """
36
37 DEFAULT_DROP = {
38 'models/': False # models are versioned, don't drop
39 }
40
41 @classmethod
42 def supports(cls, store, **kwargs):
43 return store.prefix in ('data/', 'jobs/', 'models/', 'scripts/', 'streams/')
44
[docs]
45 def promote(self, name, other, asname=None, drop=None, method='default', get=None, put=None, **kwargs):
46 """Promote object to another store.
47
48 This effectively copies the object. If the objects exists in the
49 target it will be replaced.
50
51 Args:
52 name: The name of the object
53 other: the OmegaStore instance to promote to
54 asname: the name to use in other, defaults to .metadata(name).name
55 drop (bool): if True calls other.drop(force=True) before promoting, defaults to False
56 method (str|list): specify the method or multiple methods in sequence, available methods
57 are 'default', 'getput', 'metadata', 'data'. For 'default', the object backend's
58 .PROMOTE property is used, defaulting to 'getput'
59 get (dict): optional, specifies the store.get(**kwargs)
60 put (dict): optional, specifies the other.put(**kwargs)
61 kwargs: additional kwargs are passed to the initial other.put(), for metadata promotion
62
63 Returns:
64 The Metadata of the new object
65 """
66 from omegaml.store import OmegaStore
67
68 assert isinstance(other, OmegaStore), f"specify promote(..., other, ...) as om.models, not {type(other)}"
69 drop = drop if drop is not None else self.DEFAULT_DROP.get(self.prefix, True)
70 # sanity checks
71 # -- promotion by same name requires a different store
72 meta = self.metadata(name)
73 asname = asname or meta.name
74 if name == asname and self == other:
75 raise ValueError(f'must specify asname= different from {meta.name}')
76 # see if the backend supports explicit promotion
77 backend = self.get_backend(name)
78 if hasattr(backend, 'promote'):
79 return backend.promote(name, other, asname=asname, drop=drop, get=get, put=put, method=method, **kwargs)
80 # run all promotion methods as requested or provided by the backend
81 methods = method.split(',') if isinstance(method, str) else method
82 promotion_methods = self.promotion_methods(
83 name, other, asname=asname, drop=drop, get=get, put=put, backend=backend, **kwargs
84 )
85 [promotion_methods[m]() for m in methods]
86 return other.metadata(asname)
87
88 def promotion_methods(self, name, other, asname=None, drop=None, get=None, put=None, backend=None, **kwargs):
89 # do default promotion, i.e. get()/put()
90 PROMOTION_METHODS = {
91 'getput': partial(self._get_put_promotion, name, other, asname=asname, drop=drop, get=get, put=put),
92 'data': partial(self._data_promotion, name, other, asname=asname, get=get, put=put, **kwargs),
93 'metadata': partial(
94 self._metadata_promotion, name, other, asname=None, drop=False, get=get, put=put, **kwargs
95 ),
96 'export': partial(self._export_promotion, name, other, asname=None, drop=False, get=get, put=put, **kwargs),
97 }
98 default_method = PROMOTION_METHODS[getattr(backend, 'PROMOTE', 'getput')]
99 PROMOTION_METHODS['default'] = lambda *args, **kwargs: default_method(*args, **kwargs)
100 return PROMOTION_METHODS
101
102 def _get_put_promotion(self, name, other, asname=None, drop=None, get=None, put=None, **kwargs):
103 get_kwargs = get or kwargs
104 put_kwargs = put or kwargs
105 meta = self.metadata(name)
106 obj = self.get(name, **get_kwargs)
107 asname = asname or meta.name
108 other.drop(asname, force=True) if drop else None
109 # TODO refactor to promote of python native data backend
110 if meta.kind == MDREGISTRY.PYTHON_DATA:
111 # in case of native python data we get back a list of
112 # all previously inserted objects. do the same in other
113 [other.put(o, asname) for o in obj]
114 other_meta = other.metadata(asname)
115 else:
116 other_meta = other.put(obj, asname, **put_kwargs)
117 # promote metadata, exception versions
118 # TODO: move versions to kind_meta, which is not promoted
119 meta.attributes.pop('versions', None)
120 dict_merge(other_meta.attributes, meta.attributes)
121 other_meta.save()
122 return other_meta
123
124 def _metadata_promotion(self, name, other, asname=None, drop=False, **kwargs):
125 """promote metadata
126
127 This is called by PromotionMixin.promote() to promote metadata
128 to another store, not data.
129
130 Args:
131 name (str): the name of the object
132 other (om.datasets): the other om.datasets
133
134 Returns:
135 Metadata of the promoted dataset as the result of other.put()
136 """
137 meta = self.get.metadata(name)
138 asname = asname or meta.name
139 other.drop(asname, force=True) if drop else None
140 other_meta = other.make_metadata(asname, meta.kind, bucket=other.bucket, prefix=other.prefix, **kwargs)
141 other_meta.kind_meta.update(meta.kind_meta)
142 other_meta.attributes.update(meta.attributes)
143 other_meta.save()
144 return other_meta
145
146 def _data_promotion(self, name, other, asname=None, get=None, put=None, **kwargs):
147 """promote data
148
149 This is called by PromotionMixin.promote() to promote an object's data
150 to another store. If source.get() returns an interator, target.put() will
151 be called for every iteration.
152 """
153 asname = asname or name
154 get_kwargs = get or kwargs
155 put_kwargs = put or kwargs
156 to_copy = self.get(name, **get_kwargs)
157 if hasattr(to_copy, '__iter__'):
158 for chunk in to_copy:
159 other.put(chunk, asname, **put_kwargs)
160 else:
161 other.put(to_copy, asname, **put_kwargs)
162 other_meta = other.metadata(asname)
163 return other_meta.save()
164
165 def _export_promotion(self, name, other, asname=None, get=None, put=None, **kwargs):
166 """promote as export/import
167
168 This is called by PromotionMixin.promote() to promote an object's data
169 to another store in the way that the `om runtime export/import` commands do.
170 This includes exporting all data and metadata associated with an object.
171 """
172 asname = asname or name
173 get_kwargs = get or kwargs
174 put_kwargs = put or kwargs
175 arc = self.to_archive(name, Path(self.tmppath) / asname, **get_kwargs)
176 other_meta = other.from_archive(arc, asname, **put_kwargs)
177 return other_meta