1import datetime
2import re
3from io import BytesIO, StringIO
4from pathlib import Path
5from uuid import uuid4
6
7import gridfs
8import yaml
9from croniter import croniter
10from jupyter_client import AsyncKernelManager
11from nbformat import NotebookNode
12from nbformat import read as nbread
13from nbformat import v4 as nbv4
14from nbformat import write as nbwrite
15
16from omegaml.backends.basedata import BaseDataBackend
17from omegaml.documents import MDREGISTRY
18from omegaml.notebook.jobschedule import JobSchedule
19
20
21class NotebookBackend(BaseDataBackend):
22 """
23 Omega Jobs API
24 """
25
26 # TODO this class should be a proper backend class with a mixin for ipynb
27
28 KIND = MDREGISTRY.OMEGAML_JOBS
29
30 @classmethod
31 def supports(cls, obj, name, data_store=None, **kwargs):
32 return data_store.prefix == 'jobs/' and (isinstance(obj, (str, NotebookNode)) or name.endswith('.ipynb'))
33
34 @property
35 def store(self):
36 return self.data_store
37
38 @property
39 def _db(self):
40 return self.store.mongodb
41
42 @property
43 def _fs(self):
44 return self.store.fs
45
46 def put(self, obj, name, attributes=None, **kwargs):
47 """store a notebook job
48
49 Args:
50 obj (str|NotebookNode|filelike): the notebook object or the notebook's code as a string
51 name (str): the name of the job
52 attributes (dict): the attributes to store with the job
53
54 Returns:
55 Metadata
56 """
57 if not name.endswith('.ipynb'):
58 name += '.ipynb'
59 if isinstance(obj, (str, Path)) and Path(obj).is_file():
60 with Path(obj).open() as fin:
61 value = fin.read()
62 elif isinstance(obj, str):
63 return self.store.create(obj, name)
64 elif hasattr(obj, 'read') and hasattr(obj, 'seek'):
65 obj.seek(0)
66 value = obj.read()
67 elif isinstance(obj, NotebookNode):
68 # nbwrite expects string, fs.put expects bytes
69 sbuf = StringIO()
70 nbwrite(obj, sbuf, version=4)
71 sbuf.seek(0)
72 value = sbuf.getvalue()
73 else:
74 raise ValueError('cannote store {type(obj}}, require obj:str|NotebookNode|filelike')
75 bbuf = BytesIO()
76 bbuf.write(value.encode('utf8'))
77 bbuf.seek(0)
78 # see if we have a file already, if so replace the gridfile
79 meta = self.store.metadata(name)
80 if not meta:
81 filename = uuid4().hex
82 fileid = self._store_to_file(self.store, bbuf, filename)
83 meta = self.store._make_metadata(
84 name=name,
85 prefix=self.store.prefix,
86 bucket=self.store.bucket,
87 kind=self.KIND,
88 attributes=attributes,
89 gridfile=fileid,
90 )
91 meta = meta.save()
92 else:
93 filename = uuid4().hex
94 meta.gridfile = self._store_to_file(self.store, bbuf, filename)
95 meta = meta.save()
96 # set config
97 nb_config = self.store.get_notebook_config(name)
98 meta_config = meta.attributes.get(
99 'config',
100 {},
101 )
102 if nb_config:
103 meta_config.update(dict(**nb_config))
104 meta.attributes['config'] = meta_config
105 meta = meta.save()
106 if not name.startswith('results') and ('run-at' in meta_config):
107 meta = self.store.schedule(name)
108 return meta
109
110 def get(self, name, version=-1, force_python=False, lazy=False, **kwargs):
111 """
112 Retrieve a notebook and return a NotebookNode
113 """
114 if not name.endswith('.ipynb'):
115 name += '.ipynb'
116 meta = self.store.metadata(name)
117 if meta:
118 try:
119 outf = meta.gridfile
120 except gridfs.errors.NoFile as e:
121 raise e
122 # nbwrite wants a string, outf is bytes
123 sbuf = StringIO()
124 data = outf.read()
125 if data is None:
126 msg = 'Expected content in {name}, got None'.format(**locals())
127 raise ValueError(msg)
128 sbuf.write(data.decode('utf8'))
129 sbuf.seek(0)
130 nb = nbread(sbuf, as_version=4)
131 return nb
132 else:
133 raise gridfs.errors.NoFile(f">{name}< does not exist in jobs bucket '{self.store.bucket}'")
134
135
136class NotebookMixin:
137 """om.jobs storage methods
138
139 Mixin to OmegaStore to provide om.jobs-specific methods
140
141 Usage:
142 # create a new notebook from code
143 code = \"""
144 print('hello')
145 \"""
146 om.jobs.create(code, 'name')
147
148 # schedule a notebook
149 om.jobs.schedule('name', run_at='<cron spec>')
150
151 # run a notebook immediately
152 om.jobs.run('name')
153 """
154
155 _nb_config_magic = 'omega-ml', 'schedule', 'run-at', 'cron'
156 _dir_placeholder = '_placeholder.ipynb'
157
158 def _init_mixin(self, *args, **kwargs):
159 self._include_dir_placeholder = True
160 # convenience so you can do om.jobs.schedule(..., run_at=om.jobs.Schedule(....))
161 self.Schedule = JobSchedule
162
163 @classmethod
164 def supports(cls, store, *args, **kwargs):
165 return store.prefix == 'jobs/'
166
167 def __repr__(self):
168 return f'OmegaJobs(store={super().__repr__()},)'
169
170 def collection(self, name):
171 if not name.endswith('.ipynb'):
172 name += '.ipynb'
173 return super().collection(name)
174
175 def metadata(self, name, **kwargs):
176 """retrieve metadata of a notebook
177
178 Args:
179 name (str): the name of the notebook
180
181 Returns:
182 Metadata
183 """
184 meta = super().metadata(name, **kwargs)
185 if meta is None and not name.endswith('.ipynb'):
186 name += '.ipynb'
187 meta = super().metadata(name)
188 return meta
189
190 def exists(self, name):
191 """check if the notebook exists
192
193 Args:
194 name (str): the name of the notebook
195
196 Returns:
197 Metadata
198 """
199 return len(super().list(name)) + len(super().list(name + '.ipynb')) > 0
200
201 def create(self, code, name, **kwargs):
202 """
203 create a notebook from code
204
205 :param code: the code as a string
206 :param name: the name of the job to create
207 :return: the metadata object created
208 """
209 cells = []
210 cells.append(nbv4.new_code_cell(source=code))
211 notebook = nbv4.new_notebook(cells=cells)
212 # put the notebook
213 meta = self.put(notebook, name, **kwargs)
214 return meta
215
216 def get_fs(self, collection=None):
217 # legacy support
218 return self._fs
219
220 def get_collection(self, collection):
221 """
222 returns the collection object
223 """
224 # FIXME this should use store.collection
225 return getattr(self.mongodb, collection)
226
227 def list(self, pattern=None, regexp=None, raw=False, **kwargs):
228 """
229 list all jobs matching filter.
230 filter is a regex on the name of the ipynb entry.
231 The default is all, i.e. `.*`
232 """
233 job_list = super().list(pattern=pattern, regexp=regexp, raw=raw, **kwargs)
234 name = lambda v: v.name if raw else v
235 if not self._include_dir_placeholder:
236 job_list = [v for v in job_list if not name(v).endswith(self._dir_placeholder)]
237 return job_list
238
239 def get_notebook_config(self, nb_filename):
240 """
241 returns the omegaml script config on the notebook's first cell
242
243 The config cell is in any of the following formats:
244
245 *Option 1*::
246
247 # omega-ml:
248 # run-at: "<cron schedule>"
249
250 *Option 2*::
251
252 # omega-ml:
253 # schedule: "weekday(s), hour[, month]"
254
255 *Option 3*::
256
257 # omega-ml:
258 # cron: "<cron schedule>"
259
260 You may optionally specify only ``run-at``, ``schedule`` or ``cron``,
261 i.e. without the ``omega-ml`` header::
262
263 # run-at: "<cron schedule>"
264 # cron: "<cron schedule>"
265 # schedule: "weekday(s), hour[, month]"
266
267 Args:
268 nb_filename (str): the name of the notebook
269
270 Returns:
271 config(dict) => ``{ 'run-at': <specifier> }``
272
273 Raises:
274 ValueError, if there is no config cell or the config cell is invalid
275
276 See Also:
277 * JobSchedule
278 """
279 notebook = self.get(nb_filename)
280 config_cell = None
281 config_magic = [f'# {kw}' for kw in self._nb_config_magic]
282 for cell in notebook.get('cells'):
283 if any(cell.source.startswith(kw) for kw in config_magic):
284 config_cell = cell
285 if not config_cell:
286 return {}
287 yaml_conf = '\n'.join([re.sub('#', '', x, 1) for x in str(config_cell.source).splitlines()])
288 try:
289 yaml_conf = yaml.safe_load(yaml_conf)
290 config = yaml_conf.get(self._nb_config_magic[0], yaml_conf)
291 except Exception:
292 raise ValueError('Notebook configuration cannot be parsed')
293 # translate config to canonical form
294 # TODO refactor to seperate method / mapped translation functions
295 if 'schedule' in config:
296 config['run-at'] = JobSchedule(text=config.get('schedule', '')).cron
297 if 'cron' in config:
298 config['run-at'] = JobSchedule.from_cron(config.get('cron')).cron
299 if 'run-at' in config:
300 config['run-at'] = JobSchedule.from_cron(config.get('run-at')).cron
301 return config
302
303 def run(self, name, event=None, timeout=None):
304 """
305 Run a job immediately
306
307 The job is run and the results are stored in om.jobs('results/name <timestamp>') and
308 the result's Metadata is returned.
309
310 ``Metadata.attributes`` of the original job as given by name is updated:
311
312 * ``attributes['job_runs']`` (list) - list of status of each run. Status is
313 a dict as below
314 * ``attributes['job_results']`` (list) - list of results job names in same
315 index-order as job_runs
316 * ``attributes['trigger']`` (list) - list of triggers
317
318 The status of each job run is a dict with keys:
319
320 * ``status`` (str): the status of the job run, OK or ERROR
321 * ``ts`` (datetime): time of execution
322 * ``message`` (str): error mesasge in case of ERROR, else blank
323 * ``results`` (str): name of results in case of OK, else blank
324
325 Usage::
326
327 # directly (sync)
328 meta = om.jobs.run('mynb')
329
330 # via runtime (async)
331 job = om.runtime.job('mynb')
332 result = job.run()
333
334 Args:
335 name (str): the name of the jobfile
336 event (str): an event name
337 timeout (int): timeout in seconds, None means no timeout
338
339 Returns:
340 Metadata of the results entry
341
342 See Also:
343 * OmegaJobs.run_notebook
344 """
345 return self.run_notebook(name, event=event, timeout=timeout)
346
347 def run_notebook(self, name, event=None, timeout=None):
348 """run a given notebook immediately.
349
350 Args:
351 name (str): the name of the jobfile
352 event (str): an event name
353 timeout (int): timeout in seconds
354
355 Returns:
356 Metadata of results
357
358 Notes:
359 The notebook is run using nbconvert.
360
361 The job's metadata is updated to reflect all job runs and their results
362 as follows::
363
364 {
365 'job_runs': [{
366 'status': '<execution status>',
367 'ts': <datetime>,
368 'message': '<blank | exception>,
369 'results': 'results/<name>_<datetime>.ipynb',
370 } ...],
371 'job_results': [
372 'results/<name>_<datetime>.ipynb',
373 ]
374 }
375
376 `job_runs` is a list representing result status of each time the job was
377 executed by the runtime, either by `om.runtime.job('<name>').run()`, or
378 as a scheduled execution. In case of an exception, the message is truncated
379 to contain an excerpt of str(exception).
380
381 `job_results` is a list of members in `om.jobs` representing the result
382 each job execution. Each result is equivalent of an interactive execution
383 in jupyter.
384
385 See Also:
386 * nbconvert https://nbconvert.readthedocs.io/en/latest/execute_api.html
387
388 .. versionchanged:: 0.18.0
389 `job_runs[].results` lists the name of the respective `job_results` item
390 regardless of status. In prior versions `job_runs[].results` was null for
391 failed jobs.
392
393 .. versionchanged:: 0.18.0
394 `job_runs[].message` is truncated to contain the first and last 80
395 characters of an exception. The actual exception is preserved in the
396 job_results notebook, stored as a separate object.
397 """
398 from nbconvert.preprocessors import ClearOutputPreprocessor
399 from nbconvert.preprocessors.execute import ExecutePreprocessor
400
401 notebook = self.get(name)
402 meta_job = self.metadata(name)
403 assert meta_job is not None, f"Cannot run non-existent jobs {name}"
404 ts = datetime.datetime.now()
405 # execute kwargs
406 # -- see ExecuteProcessor class
407 # -- see https://nbconvert.readthedocs.io/en/latest/execute_api.html
408 ep_kwargs = {
409 # avoid timeouts to stop kernel
410 'timeout': timeout,
411 # avoid kernel at exit functions
412 # -- this stops ipykernel AttributeError 'send_multipart'
413 'shutdown_kernel': 'immediate',
414 # set kernel name, blank is default
415 # -- e.g. python3, ir
416 # -- see https://stackoverflow.com/a/47053020/890242
417 'kernel_name': '',
418 # always use the async kernel manager
419 # -- see https://github.com/jupyter/nbconvert/issues/1964
420 'kernel_manager_class': AsyncKernelManager,
421 }
422 # overrides from metadata
423 ep_kwargs.update(
424 meta_job.kind_meta.get(
425 'ep_kwargs',
426 {},
427 )
428 )
429 try:
430 resources = {
431 'metadata': {'path': self.defaults.OMEGA_TMP},
432 }
433 if not meta_job.kind_meta.get('keep_output', False):
434 # https://nbconvert.readthedocs.io/en/latest/api/preprocessors.html
435 cp = ClearOutputPreprocessor()
436 cp.preprocess(notebook, resources)
437 ep = ExecutePreprocessor(**ep_kwargs)
438 ep.preprocess(notebook, resources)
439 except (BaseException, Exception) as e:
440 status = 'ERROR'
441 message = str(e)
442 else:
443 status = 'OK'
444 message = ''
445 finally:
446 # limit size of metadata.attributes #496
447 message = f'{message[0:80]}...{message[-80:]}'
448 del ep
449 # record results
450 meta_results = self.put(notebook, 'results/{name}_{ts}'.format(**locals()))
451 meta_results.attributes['source_job'] = name
452 meta_results.save()
453 job_results = meta_job.attributes.get('job_results', [])
454 job_results.append(meta_results.name)
455 meta_job.attributes['job_results'] = job_results
456 # record final job status
457 job_runs = meta_job.attributes.get('job_runs', [])
458 runstate = {
459 'status': status,
460 'ts': ts,
461 'message': message,
462 'results': meta_results.name,
463 }
464 job_runs.append(runstate)
465 meta_job.attributes['job_runs'] = job_runs
466 # set event run state if event was specified
467 if event:
468 attrs = meta_job.attributes
469 triggers = attrs['triggers'] = attrs.get('triggers', [])
470 scheduled = (trigger for trigger in triggers if trigger['event-kind'] == 'scheduled')
471 for trigger in scheduled:
472 if event == trigger['event']:
473 trigger['status'] = status
474 trigger['ts'] = ts
475 meta_job.save()
476 return meta_results
477
478 def schedule(self, nb_file, run_at=None, last_run=None):
479 """
480 Schedule a processing of a notebook as per the interval
481 specified on the job script
482
483 Notes:
484 This updates the notebook's Metadata entry by adding the
485 next scheduled run in ``attributes['triggers']```
486
487 Args:
488 nb_file (str): the name of the notebook
489 run_at (str|dict|JobSchedule): the schedule specified in a format
490 suitable for JobSchedule. If not specified, this value is
491 extracted from the first cell of the notebook
492 last_run (datetime|None): the last time this job was run, use this to
493 reschedule the job for the next run. Defaults to the last
494 timestamp listed in ``attributes['job_runs']``, or datetime.utcnow()
495 if no previous run exists.
496
497 See Also:
498 * croniter.get_next()
499 * JobSchedule
500 * OmegaJobs.get_notebook_config
501 * cron expression - https://en.wikipedia.org/wiki/Cron#CRON_expression
502 """
503 meta = self.metadata(nb_file)
504 attrs = meta.attributes
505 # get/set run-at spec
506 config = attrs.get(
507 'config',
508 {},
509 )
510 # see what we have as a schedule
511 # -- a dictionary of JobSchedule
512 if isinstance(run_at, dict):
513 run_at = self.Schedule(**run_at).cron
514 # -- a cron string
515 elif isinstance(run_at, str):
516 run_at = self.Schedule(run_at).cron
517 # -- a JobSchedule
518 elif isinstance(run_at, self.Schedule):
519 run_at = run_at.cron
520 # -- nothing, may we have it on the job's config already
521 if not run_at:
522 interval = config.get('run-at')
523 else:
524 # store the new schedule
525 config['run-at'] = run_at
526 interval = run_at
527 if not interval:
528 # if we don't have a run-spec, return without scheduling
529 raise ValueError('no run-at specification found, cannot schedule')
530 # get last time the job was run
531 if last_run is None:
532 job_runs = attrs.get('job_runs')
533 if job_runs:
534 last_run = job_runs[-1]['ts']
535 else:
536 last_run = datetime.datetime.utcnow()
537 # calculate next run time
538 iter_next = croniter(interval, last_run)
539 run_at = iter_next.get_next(datetime.datetime)
540 # store next scheduled run
541 triggers = attrs['triggers'] = attrs.get('triggers', [])
542 # set up a schedule event
543 scheduled_run = {
544 'event-kind': 'scheduled',
545 'event': run_at.isoformat(),
546 'run-at': run_at,
547 'status': 'PENDING',
548 }
549 # remove all pending triggers, add new triggers
550 past_triggers = [cur for cur in triggers if cur.get('status') != 'PENDING']
551 triggers.clear()
552 triggers.extend(past_triggers)
553 triggers.append(scheduled_run)
554 attrs['config'] = config
555 return meta.save()
556
557 def get_schedule(self, name, only_pending=False):
558 """
559 return the cron schedule and corresponding triggers
560
561 Args:
562 name (str): the name of the job
563
564 Returns:
565 tuple of (run_at, triggers)
566
567 run_at (str): the cron spec, None if not scheduled
568 triggers (list): the list of triggers
569 """
570 meta = self.metadata(name)
571 attrs = meta.attributes
572 config = attrs.get('config')
573 triggers = attrs.get('triggers', [])
574 if only_pending:
575 triggers = [trigger for trigger in triggers if trigger['status'] == 'PENDING']
576 if config and 'run-at' in config:
577 run_at = config.get('run-at')
578 else:
579 run_at = None
580 return run_at, triggers
581
582 def drop_schedule(self, name):
583 """
584 Drop an existing schedule, if any
585
586 This will drop any existing schedule and any pending triggers of
587 event-kind 'scheduled'.
588
589 Args:
590 name (str): the name of the job
591
592 Returns:
593 Metadata
594 """
595 meta = self.metadata(name)
596 attrs = meta.attributes
597 config = attrs.get('config')
598 triggers = attrs.get('triggers')
599 if 'run-at' in config:
600 del config['run-at']
601 for trigger in triggers:
602 if trigger['event-kind'] == 'scheduled' and trigger['status'] == 'PENDING':
603 trigger['status'] = 'CANCELLED'
604 return meta.save()
605
606 def export(self, name, localpath='memory', format='html'):
607 """
608 Export a job or result file to HTML
609
610 The job is exported in the given format.
611
612 :param name: the name of the job, as in jobs.get
613 :param localpath: the path of the local file to write. If you
614 specify an empty path or 'memory' a tuple of (body, resource)
615 is returned instead
616 :param format: the output format. currently only :code:`'html'` is supported
617 :return: the (data, resources) tuple as returned by nbconvert. For
618 format html data is the HTML's body, for PDF it is the pdf file contents
619 """
620 from nbconvert import SlidesExporter
621 from nbconvert.exporters.html import HTMLExporter
622 from nbconvert.exporters.pdf import PDFExporter
623 from traitlets.config import Config
624
625 # https://nbconvert.readthedocs.io/en/latest/nbconvert_library.html
626 # (exporter class, filemode, config-values
627 EXPORTERS = {
628 'html': (
629 HTMLExporter,
630 '',
631 {},
632 ),
633 'htmlbody': (
634 HTMLExporter,
635 '',
636 {},
637 ),
638 'pdf': (
639 PDFExporter,
640 'b',
641 {},
642 ),
643 'slides': (
644 SlidesExporter,
645 '',
646 {'RevealHelpPreprocessor.url_prefix': 'https://cdnjs.cloudflare.com/ajax/libs/reveal.js/3.6.0/'},
647 ),
648 }
649 # get exporter according to format
650 if format not in EXPORTERS:
651 raise ValueError(f'format {format} is invalid. Choose one of {EXPORTERS.keys()}')
652 exporter_cls, fmode, configkw = EXPORTERS[format]
653 # prepare config
654 # http://nbconvert.readthedocs.io/en/latest/nbconvert_library.html#Using-different-preprocessors
655 c = Config()
656 for k, v in configkw.items():
657 context, key = k.split('.')
658 setattr(c[context], key, v)
659 # get configured exporter
660 exporter = exporter_cls(config=c)
661 # get notebook, convert and store in file if requested
662 notebook = self.get(name)
663 (data, resources) = exporter.from_notebook_node(notebook)
664 if localpath and localpath != 'memory':
665 with open(localpath, 'w' + fmode) as fout:
666 fout.write(data)
667 return data, resources
668
669 def help(self, name_or_obj=None, kind=None, raw=False):
670 """get help for a notebook
671
672 Args:
673 name_or_obj (str|obj): the name or actual object to get help for
674 kind (str): optional, if specified forces retrieval of backend for the given kind
675 raw (bool): optional, if True forces help to be the backend type of the object.
676 If False returns the attributes[docs] on the object's metadata, if available.
677 Defaults to False
678
679 Returns:
680 * help(obj) if python is in interactive mode
681 * text(str) if python is in not interactive mode
682 """
683 return super().help(name_or_obj, kind=kind, raw=raw)
684
685 def to_archive(self, name, path, **kwargs):
686 if not name.endswith('.ipynb'):
687 name += '.ipynb'
688 return super().to_archive(name, path, **kwargs)
689
690 def from_archive(self, path, name, **kwargs):
691 if not name.endswith('.ipynb'):
692 name += '.ipynb'
693 return super().from_archive(path, name, **kwargs)
694
695 def promote(self, name, other, **kwargs):
696 nb = self.get(name)
697 return other.put(nb, name)
698
699 def summary(self, *args, **kwargs):
700 return super().summary(*args, **kwargs)
701
702
703# sphinx docs compatibility
[docs]
704class OmegaJobs(NotebookMixin, NotebookBackend):
705 pass