1import os
2import traceback
3from uuid import uuid4
4
5from omegaml.util import inprogress
6
7from ._version import version
8from .mixins.store.requests import CombinedStoreRequestCache
9from .store.combined import CombinedOmegaStoreMixin
10from .store.logging import OmegaSimpleLogger
11
12
[docs]
13class Omega(CombinedStoreRequestCache, CombinedOmegaStoreMixin):
14 """
15 Client API to omegaml
16
17 Provides the following APIs:
18
19 * :code:`datasets` - access to datasets stored in the cluster
20 * :code:`models` - access to models stored in the cluster
21 * :code:`runtimes` - access to the cluster compute resources
22 * :code:`jobs` - access to jobs stored and executed in the cluster
23 * :code:`scripts` - access to lambda modules stored and executed in the cluster
24
25 """
26
27 def __init__(self, defaults=None, mongo_url=None, celeryconf=None, bucket=None, **kwargs):
28 """
29 Initialize the client API
30
31 Without arguments create the client API according to the user's
32 configuration in :code:`~/omegaml/config.yml`.
33
34 Arguments override the user's configuration.
35
36 :param defaults: the DefaultsContext
37 :param mongo_url: the fully qualified URI to the mongo database,
38 of format :code:`mongodb://user:password@host:port/database`
39 :param celeryconf: the celery configuration dictionary
40 """
41 from omegaml.util import settings
42
43 # celery and mongo configuration
44 self.defaults = defaults or settings()
45 self.mongo_url = mongo_url or self.defaults.OMEGA_MONGO_URL
46 self.bucket = bucket
47 # setup storage locations
48 self._dbalias = self._make_dbalias()
49 self.models = self._make_store(prefix='models/')
50 self.datasets = self._make_store(prefix='data/')
51 self.jobs = self._make_store(prefix='jobs/')
52 self.scripts = self._make_store(prefix='scripts/')
53 # minibatch integration
54 self.streams = self._make_streams(prefix='streams/')
55 # runtimes environments
56 self.runtime = self._make_runtime(celeryconf)
57 # logger
58 self.logger = OmegaSimpleLogger(store=self.datasets, defaults=self.defaults)
59 # stores
60 self._stores = [self.models, self.datasets, self.scripts, self.jobs, self.streams]
61 # monitoring
62 self._monitor = None # will be created by .status() on first access
63
64 def __repr__(self):
65 return f'Omega(bucket={self.bucket})'
66
67 def _clone(self, **kwargs):
68 return self.__class__(defaults=self.defaults, mongo_url=self.mongo_url, **kwargs)
69
70 def _make_runtime(self, celeryconf):
71 from omegaml.runtimes import OmegaRuntime
72
73 return OmegaRuntime(self, bucket=self.bucket, defaults=self.defaults, celeryconf=celeryconf)
74
75 def _make_store(self, prefix):
76 from omegaml.store import OmegaStore
77
78 return OmegaStore(
79 mongo_url=self.mongo_url, bucket=self.bucket, prefix=prefix, defaults=self.defaults, dbalias=self._dbalias
80 )
81
82 def _make_dbalias(self):
83 return f'omega-{uuid4().hex}'
84
85 def _make_streams(self, prefix):
86 from omegaml.store.streams import StreamsProxy
87
88 return StreamsProxy(mongo_url=self.mongo_url, bucket=self.bucket, prefix=prefix, defaults=self.defaults)
89
90 def _make_monitor(self):
91 import weakref
92
93 from omegaml.client.lunamon import LunaMonitor, OmegaMonitors
94
95 status_logger = self.runtime.experiment('.system')
96 for_keys = lambda event, keys: {k: event[k] for k in keys if k in event}
97 on_status = lambda event: (
98 status_logger.use().log_event(
99 'monitor', event['check'], for_keys(event, ('status', 'message', 'error', 'elapsed'))
100 )
101 if event['check'] == 'health'
102 else None
103 )
104 monitor = LunaMonitor(checks=OmegaMonitors.on(self), on_status=on_status, interval=15)
105 weakref.finalize(self, monitor.stop)
106 return monitor
107
[docs]
108 def status(self, check=None, data=False, by_status=False, wait=False):
109 """get the status of the omegaml cluster
110
111 Args:
112 check (str): the check to run, e.g. 'storage', 'runtime'
113 data (bool): return the monitoring data for the check
114 by_status (bool): return data by status
115 wait (bool): wait for the check to complete
116
117 Returns:
118 dict
119 """
120 if self._monitor is None:
121 # we defer the creation of the monitor to the first access
122 # -- this is to avoid creating a monitor for every instance
123 # -- e.g. in runtime where the instance is created for short-lived tasks,
124 # the monitor would be created and immediately gc'd
125 self._monitor = self._make_monitor()
126 if wait:
127 self._check_connections()
128 return self._monitor.status(check=check, data=data, by_status=by_status)
129
130 def _check_connections(self):
131 return self._monitor.wait_ok()
132
133 def __getitem__(self, bucket):
134 """
135 return Omega instance configured for the given bucket
136
137 Args:
138 bucket (str): the bucket name. If it does not exist
139 it gets created on first storage of an object.
140 If bucket=None returns self.
141
142 Usage:
143 import omegaml as om
144
145 # om is configured on the default bucket
146 # om_mybucket will use the same settings, but configured for mybucket
147 om_mybucket = om['mybucket']
148
149 Returns:
150 Omega instance configured for the given bucket
151 """
152 return self._get_bucket(bucket)
153
154 @property
155 def buckets(self):
156 from itertools import chain
157
158 return list(set(chain(*[getattr(store, 'buckets', []) for store in self._stores])))
159
160 def _get_bucket(self, bucket):
161 # enable patching in testing
162 if bucket is None or bucket == self.bucket or (bucket == 'default' and self.bucket is None):
163 return self
164 if bucket == 'default':
165 # actual bucket selection is a responsibility of the store, thus we pass None
166 # (it is for good reason: OmegaStore should be instantiatable without giving a specific bucket)
167 bucket = None
168 return self._clone(bucket=bucket)
169
170
171class OmegaDeferredInstance:
172 """
173 A deferred instance of Omega() that is only instantiated on access
174
175 This is to ensure that module-level imports don't trigger instantiation
176 of Omega.
177 """
178
179 def __init__(self, base=None, attribute=None):
180 self.omega = 'not initialized -- call .setup() or access an attribute'
181 self.initialized = False
182 self.base = base
183 self.attribute = attribute
184
185 @inprogress(text='loading ...')
186 def setup(self, *args, **kwargs):
187 """loads omegaml
188
189 Loading order
190 - cloud using environment
191 - cloud using config file
192 - local instance
193
194 Returns:
195 omega instance
196 """
197 from omegaml import _base_config
198
199 @inprogress(text='loading base ...')
200 def setup_base():
201 from omegaml import _base_config
202
203 _base_config.load_framework_support()
204 _base_config.load_user_extensions()
205 return Omega(*args, **kwargs)
206
207 @inprogress(text='loading cloud ...')
208 def setup_cloud():
209 from omegaml.client.cloud import setup
210
211 return setup(*args, **kwargs)
212
213 @inprogress(text='loading cloud from env ...')
214 def setup_env():
215 from omegaml.client.cloud import setup
216
217 return setup(
218 userid=os.environ['OMEGA_USERID'],
219 apikey=os.environ['OMEGA_APIKEY'],
220 qualifier=os.environ.get('OMEGA_QUALIFIER'),
221 )
222
223 @inprogress(text='loading cloud from config ...')
224 def setup_cloud_config():
225 from omegaml.client.cloud import setup_from_config
226
227 return setup_from_config(fallback=setup_base)
228
229 omega = None
230 from_args = len(args) > 0 or any(kw in kwargs for kw in ('userid', 'apikey', 'api_url', 'qualifier'))
231 from_env = {'OMEGA_USERID', 'OMEGA_APIKEY'} < set(os.environ)
232 from_config = _base_config.OMEGA_CONFIG_FILE and os.path.exists(_base_config.OMEGA_CONFIG_FILE)
233 loaders = (
234 (from_args, setup_cloud),
235 (from_env, setup_env),
236 (from_config, setup_cloud_config),
237 (True, setup_base),
238 )
239 must_load = (from_env, setup_env), (from_config, setup_cloud_config)
240 errors = []
241 for condition, loader in loaders:
242 if not condition:
243 continue
244 try:
245 omega = loader()
246 except Exception as e:
247 errors.append((loader.__name__, e, traceback.format_exc()))
248 if any(condition and loader is expected for condition, expected in must_load):
249 raise
250 else:
251 break
252 else:
253 assert omega is not None, f"failed to load omega due to {errors}"
254
255 if not self.initialized:
256 self.initialized = True
257 self.omega = omega
258 return omega
259
260 def __getattr__(self, name):
261 if self.base:
262 base = getattr(self.base, self.attribute)
263 return getattr(base, name)
264 if not self.initialized:
265 self.setup()
266 return getattr(self.omega, name)
267
268 def __getitem__(self, bucket):
269 if self.base:
270 base = getattr(self.base, self.attribute)
271 return base[bucket]
272 if not self.initialized:
273 self.setup()
274 return self.omega[bucket]
275
276 def __repr__(self):
277 if self.base:
278 return repr(getattr(self.base, self.attribute))
279 self.setup()
280 return repr(self.omega)
281
282 def __call__(self, *args, **kwargs):
283 if self.base:
284 base = getattr(self.base, self.attribute)
285 return base(*args, **kwargs)
286 if not self.initialized:
287 self.setup()
288 return self(*args, **kwargs)
289
290 @property
291 def instance(self):
292 return self.base
293
294
295def setup(*args, **kwargs):
296 """
297 configure and return the omega client instance
298 """
299 return _om.setup(*args, **kwargs)
300
301
302# default instance
303# -- these are deferred instanced that is the actual Omega instance
304# is only created on actual attribute access
305_om = OmegaDeferredInstance()
306version = version # ensure we keep imports