def setup_from_config(config_file=None, fallback=None):
""" setup from a cloud configuration file
If the configuration files does not contain OMEGA_USERID and OMEGA_APIKEY,
will use given fallback.
"""
from omegaml import _base_config
config_file = config_file or _base_config.OMEGA_CONFIG_FILE
AuthenticationEnv.secure()
if isinstance(config_file, str) and os.path.exists(config_file):
with open(config_file, 'r') as fin:
userconfig = yaml.safe_load(fin)
if isinstance(userconfig, dict) and 'OMEGA_USERID' in userconfig:
try:
omega = setup(userid=userconfig['OMEGA_USERID'],
apikey=userconfig['OMEGA_APIKEY'],
qualifier=userconfig.get('OMEGA_QUALIFIER'),
api_url=userconfig.get('OMEGA_RESTAPI_URL'))
except Exception as e:
# TODO make this a SystemError so that OmegaDeferredIstance.setup reverts to proper defaults
raise SystemError(
f'Could not login using config file {config_file}, error={e}'
)
elif fallback:
# if alternative loading was provided, use that
omega = fallback()
else:
raise SystemError(
'No cloud cloud userid/apikey found in config file {}.'.
format(config_file))
return omega
raise SystemError('Config file {} does not exist'.format(config_file))
def setup(userid=None, apikey=None, api_url=None, qualifier=None, bucket=None):
import omegaml as om_mod
api_url = ensure_api_url(api_url, om_mod._base_config)
view = getattr(om_mod._base_config, 'OMEGA_SERVICES_INCLUSTER', False)
auth_env = AuthenticationEnv.secure()
om = auth_env.get_omega_from_apikey(userid=userid,
apikey=apikey,
api_url=api_url,
qualifier=qualifier,
view=view)
# prepare om to be linked to default
om.Omega = OmegaCloud
om.setup = lambda *args, **kwargs: setup(
**{
**dict(userid=userid,
apikey=apikey,
qualifier=qualifier,
bucket=bucket),
**kwargs
})
om._om = om
# ensure link to deferred instance
om_mod.Omega = OmegaCloud
om_mod.link(om)
return om[bucket]
def login(self):
userid = self.args.get('<userid>') or self.args.get('--userid')
apikey = self.args.get('<apikey>') or self.args.get('--apikey')
qualifier = self.args.get('<qualifier>') or 'default'
api_url = self.args.get('--apiurl')
configfile = self.args.get('--config') or 'config.yml'
if not userid:
userid = self.ask('Userid:')
if not apikey:
apikey = self.ask('Apikey:')
auth_env = AuthenticationEnv().secure()
auth_env.save_userconfig_from_apikey(configfile,
userid,
apikey,
qualifier=qualifier,
api_url=api_url)
def _save_userconfig_from_apikey(configfile,
userid,
apikey,
api_url=None,
requested_userid=None,
view=False,
keys=None,
qualifier=None):
from omegaml import settings
defaults = settings()
api_url = ensure_api_url(api_url, defaults)
required_keys = [
'OMEGA_USERID', 'OMEGA_APIKEY', 'OMEGA_RESTAPI_URL', 'OMEGA_QUALIFIER'
]
keys = keys or []
auth_env = AuthenticationEnv.secure()
with open(configfile, 'w') as fconfig:
configs = auth_env.get_userconfig_from_api(
api_url=api_url,
userid=userid,
apikey=apikey,
requested_userid=requested_userid,
qualifier=qualifier,
view=view)
config = configs['objects'][0]['data']
config['OMEGA_RESTAPI_URL'] = api_url
config['OMEGA_QUALIFIER'] = qualifier or 'default'
config['OMEGA_USERID'] = userid
config['OMEGA_APIKEY'] = apikey
config = {
k: v
for k, v in config.items() if k in (required_keys + keys)
}
yaml.safe_dump(config, fconfig, default_flow_style=False)
print("Config is in {configfile}".format(**locals()))
def _get_omega_from_apikey(userid,
apikey,
api_url=None,
requested_userid=None,
qualifier=None,
view=False):
"""
setup an Omega instance from userid and apikey
:param userid: the userid
:param apikey: the apikey
:param api_url: the api URL
:param requested_userid: the userid to request config for. in this case userid
and apikey must for a staff user for the request to succeed
:param qualifier: the database qualifier requested. defaults to 'default'
:returns: OmegaCloud instance configured for the given user
Returns:
OmegaCloud
"""
from omegaml.client.cloud import OmegaCloud
from omegaml import settings, _base_config
defaults = settings(reload=True)
qualifier = qualifier or 'default'
defaults.OMEGA_USERID = userid
defaults.OMEGA_APIKEY = apikey
defaults.OMEGA_QUALIFIER = qualifier
api_url = ensure_api_url(api_url, defaults)
auth_env = AuthenticationEnv.secure()
if api_url.startswith('http') or any('test' in v for v in sys.argv):
configs = auth_env.get_userconfig_from_api(
requested_userid=requested_userid, view=view, defaults=defaults)
configs = configs['objects'][0]['data']
elif api_url == 'local':
configs = {
k: getattr(defaults, k)
for k in dir(defaults) if k.startswith('OMEGA')
}
else:
raise ValueError('invalid api_url {}'.format(api_url))
config = configs.get(qualifier, configs)
# update
_base_config.update_from_dict(config, attrs=defaults)
_base_config.update_from_config(defaults)
_base_config.load_framework_support(defaults)
_base_config.load_user_extensions(defaults)
auth = auth_env.get_runtime_auth(defaults)
om = OmegaCloud(defaults=defaults, auth=auth)
# update config to reflect request
om.defaults.OMEGA_RESTAPI_URL = api_url
om.defaults.OMEGA_USERID = userid
om.defaults.OMEGA_APIKEY = apikey
om.defaults.OMEGA_QUALIFIER = qualifier
return om
def __init__(self, *args, **kwargs):
super(OmegamlTask, self).__init__(*args, **kwargs)
self.auth_env = AuthenticationEnv().active()
class OmegamlTask(EagerSerializationTaskMixin, Task):
abstract = True
def __init__(self, *args, **kwargs):
super(OmegamlTask, self).__init__(*args, **kwargs)
self.auth_env = AuthenticationEnv().active()
@property
def om(self):
# TODO do some more intelligent caching, i.e. by client/auth
kwargs = self.request.kwargs or {}
if not hasattr(self.request, '_om'):
self.request._om = None
if self.request._om is None:
bucket = kwargs.get('__bucket')
self.request._om = self.auth_env.get_omega_for_task(self)[bucket]
return self.request._om
def get_delegate(self, name, kind='models'):
get_delegate_provider = getattr(self.om, kind)
self.enable_delegate_tracking(name, kind, get_delegate_provider)
result = get_delegate_provider.get_backend(name,
data_store=self.om.datasets,
tracking=self.tracking)
return result
def enable_delegate_tracking(self, name, kind, delegate_provider):
exp = self.tracking.experiment
meta = delegate_provider.metadata(name)
exp.log_artifact(meta, 'related')
tracking = meta.attributes.setdefault('tracking', {})
tracking['experiments'] = set(
tracking.get('experiments', []) + [exp._experiment])
meta.save()
return meta
@property
def delegate_args(self):
return self.request.args
@property
def delegate_kwargs(self):
kwargs = self.request.kwargs or {}
return {k: v for k, v in kwargs.items() if not k.startswith('__')}
@property
def system_kwargs(self):
kwargs = self.request.kwargs or {}
return {k: v for k, v in kwargs.items() if k.startswith('__')}
@property
def tracking(self):
if not hasattr(self.request, '_om_tracking'):
self.request._om_tracking = None
if self.request._om_tracking is None:
kwargs = self.request.kwargs or {}
experiment = kwargs.get('__experiment')
if experiment is not None:
# we reuse implied_run=False to use the currently active run,
# i.e. with block will NOT call exp.start()
tracker = self.om.runtime.experiment(experiment,
implied_run=False)
self.request._om_tracking = tracker
else:
self.request._om_tracking = self.om.runtime.experiment(
'.notrack', provider='notrack')
return self.request._om_tracking
@property
def logging(self):
kwargs = self.request.kwargs or {}
logging = kwargs.get('__logging', False)
if isinstance(logging, tuple):
logname, level = logging
if logname is True:
logname = 'root'
elif isinstance(logging, str):
if logging.lower() in ('info', 'warn', 'fatal', 'error', 'debug',
'critical'):
logname, level = 'root', logging.upper()
else:
logname, level = logging, 'INFO'
elif logging is True:
logname, level = 'root', 'INFO'
else:
logname, level = None, 'NOTSET'
return logname, level
def __call__(self, *args, **kwargs):
import logging
@contextmanager
def task_logging():
logname, level = self.logging
if logname:
self.om.logger.setLevel(level)
logger = logging.getLogger(
name=logname if logname != 'root' else None)
logger.setLevel(level)
save_stdout, save_stderr = sys.stdout, sys.stderr
self.app.log.redirect_stdouts_to_logger(logger, loglevel=level)
handler = OmegaLoggingHandler.setup(store=self.om.datasets,
logger=logger,
exit_hook=True,
level=level)
else:
logger = None
try:
yield
finally:
if logger:
handler.flush()
handler.close()
logger.removeHandler(handler)
# reset stdout redirects
sys.stdout, sys.stderr = save_stdout, save_stderr
with task_logging():
with self.tracking as exp:
exp.log_event(f'task_call', self.name, {
'args': args,
'kwargs': kwargs
})
if self.request.id is not None:
# PERFTUNED
# if we have a request, avoid super().__call__()
# calling super() removes the request.id, which means our request._om caching does not work
# Rationale: https://github.com/celery/celery/issues/2633#issuecomment-286694356
result = self.run(*args, **self.delegate_kwargs)
else:
# PERFISSUE new task cannot access the request cache, new omega instance is created
# if we don't have a request.id it means we don't have a proper callstack
# so we need to create one using super().__call__
# this may happen when calling tasks directly, such as in run_omega_callback_script
result = super().__call__(*args, **kwargs)
return result
def reset(self):
# ensure next call will start over and get a new om instance
self.request._om = None
def on_failure(self, exc, task_id, args, kwargs, einfo):
with self.tracking as exp:
exp.log_event(f'task_failure', self.name, {
'exception': repr(exc),
'task_id': task_id,
})
self.reset()
return super().on_failure(exc, task_id, args, kwargs, einfo)
def on_retry(self, exc, task_id, args, kwargs):
with self.tracking as exp:
exp.log_event(f'task_retry', self.name, {
'exception': repr(exc),
'task_id': task_id,
})
self.reset()
return super().on_retry(exc, task_id, args, kwargs)
def on_success(self, retval, task_id, args, kwargs):
with self.tracking as exp:
exp.log_event(f'task_success', self.name, {
'result': sanitized(retval),
'task_id': task_id,
})
self.reset()
return super().on_success(retval, task_id, args, kwargs)
def _restapi_auth(self):
return AuthenticationEnv.active().get_restapi_auth(om=self.om)