11import getpass
2-
3- import dill
42import os
5- import pandas as pd
6- import pkg_resources
73import platform
84from base64 import b64encode , b64decode
95from datetime import datetime
106from itertools import product
117from uuid import uuid4
128
9+ import dill
10+ import pandas as pd
11+ import pkg_resources
12+ import pymongo
13+
1314from omegaml .backends .basemodel import BaseModelBackend
1415from omegaml .documents import Metadata
15- from omegaml .util import _raise , settings
16+ from omegaml .util import _raise , settings , ensure_index
1617
1718
1819class ExperimentBackend (BaseModelBackend ):
@@ -183,12 +184,24 @@ def active_run(self):
183184 def status (self , run = None ):
184185 return 'STOPPED'
185186
186- def start (self ):
187+ def start (self , run = None ):
187188 raise NotImplementedError
188189
189190 def stop (self ):
190191 raise NotImplementedError
191192
193+ def start_runtime (self ):
194+ # hook to signal the runtime is starting a task inside a worker
195+ # this is unlike the .start() method which is called to start a run
196+ # which can happen in the client or in the runtime
197+ pass
198+
199+ def stop_runtime (self ):
200+ # hook to signal the runtime has completed a task inside a worker
201+ # this is unlike the .stop() method which is called to stop a run
202+ # which can happen in the client or in the runtime
203+ self .flush ()
204+
192205 def log_event (self , event , key , value , step = None , ** extra ):
193206 raise NotImplementedError
194207
@@ -211,6 +224,9 @@ def tensorflow_callback(self):
211224 def data (self , experiment = None , run = None , event = None , step = None , key = None , raw = False ):
212225 raise NotImplementedError
213226
227+ def flush (self ):
228+ pass
229+
214230 @property
215231 def _data_name (self ):
216232 return f'.experiments/{ self ._experiment }
@@ -219,7 +235,7 @@ def _data_name(self):
219235class NoTrackTracker (TrackingProvider ):
220236 """ A default tracker that does not record anything """
221237
222- def start (self ):
238+ def start (self , run = None ):
223239 pass
224240
225241 def stop (self ):
@@ -261,6 +277,12 @@ class OmegaSimpleTracker(TrackingProvider):
261277 _ensure_active = lambda self , r : r if r is not None else _raise (
262278 ValueError ('no active run, call .start() or .use() ' ))
263279
280+ def __init__ (self , * args , ** kwargs ):
281+ super ().__init__ (* args , ** kwargs )
282+ self .log_buffer = []
283+ self .max_buffer = 10
284+ self ._initialize_dataset ()
285+
264286 def active_run (self , run = None ):
265287 """ set the lastest run as the active run
266288
@@ -293,8 +315,9 @@ def use(self, run=None):
293315
294316 @property
295317 def _latest_run (self ):
296- data = self .data (event = 'start' , raw = True )
297- run = data [- 1 ]['run' ] if data is not None and len (data ) > 0 else None
318+ cursor = self .data (event = 'start' , lazy = True )
319+ data = list (cursor .sort ('data.run' , - 1 ).limit (1 )) if cursor else None
320+ run = data [- 1 ].get ('data' , {}).get ('run' ) if data is not None and len (data ) > 0 else None
298321 return run
299322
300323 def status (self , run = None ):
@@ -307,7 +330,7 @@ def status(self, run=None):
307330 status in 'STARTED', 'STOPPED'
308331 """
309332 self ._run = run or self ._run or self ._latest_run
310- data = self .data (event = ( 'start' , 'stop' ) , run = self ._run , raw = True )
333+ data = self .data (event = [ 'start' , 'stop' ] , run = self ._run , raw = True )
311334 no_runs = data is None or len (data ) == 0
312335 has_stop = sum (1 for row in (data or []) if row .get ('event' ) == 'stop' )
313336 return 'PENDING' if no_runs else 'STOPPED' if has_stop else 'STARTED'
@@ -320,7 +343,7 @@ def start(self, run=None):
320343 self ._run = run or (self ._latest_run or 0 ) + 1
321344 self ._startdt = datetime .utcnow ()
322345 data = self ._common_log_data ('start' , key = None , value = None , step = None , dt = self ._startdt )
323- self ._write_log (data )
346+ self ._write_log (data , immediate = True )
324347 return self ._run
325348
326349 def stop (self ):
@@ -331,6 +354,14 @@ def stop(self):
331354 self ._stopdt = datetime .utcnow ()
332355 data = self ._common_log_data ('stop' , key = None , value = None , step = None , dt = self ._stopdt )
333356 self ._write_log (data )
357+ self .flush ()
358+
359+ def flush (self ):
360+ # passing list of list forces insert_many
361+ if self .log_buffer :
362+ self ._store .put (self .log_buffer , self ._data_name ,
363+ noversion = True , as_many = True )
364+ self .log_buffer .clear ()
334365
335366 def _common_log_data (self , event , key , value , step = None , dt = None , ** extra ):
336367 if isinstance (value , dict ):
@@ -358,8 +389,10 @@ def _common_log_data(self, event, key, value, step=None, dt=None, **extra):
358389 data .update (self ._extra_log ) if self ._extra_log else None
359390 return data
360391
361- def _write_log (self , data ):
362- self ._store .put (data , self ._data_name , noversion = True )
392+ def _write_log (self , data , immediate = False ):
393+ self .log_buffer .append (data )
394+ if immediate or len (self .log_buffer ) > self .max_buffer :
395+ self .flush ()
363396
364397 def log_artifact (self , obj , name , step = None , ** extra ):
365398 """ log any object to the current run
@@ -396,7 +429,7 @@ def log_artifact(self, obj, name, step=None, **extra):
396429 meta = self ._model_store .put (obj , f'.experiments/.artefacts/{ objname } )
397430 format = 'model'
398431 rawdata = meta .name
399- elif self ._store .get_backend_by_obj (obj ) is not None :
432+ elif self ._store .get_backend_byobj (obj ) is not None :
400433 objname = uuid4 ().hex
401434 meta = self ._store .put (obj , f'.experiments/.artefacts/{ objname } )
402435 format = 'dataset'
@@ -485,16 +518,17 @@ def log_extra(self, remove=False, **kwargs):
485518 consume (deletions , maxlen = 0 )
486519
487520 def data (self , experiment = None , run = None , event = None , step = None , key = None , raw = False ,
488- ** extra ):
521+ lazy = False , ** extra ):
489522 """ build a dataframe of all stored data
490523
491524 Args:
492- experiment (str): the name of the experiment, defaults to its current value
525+ experiment (str|list ): the name of the experiment, defaults to its current value
493526 run (int|list): the run(s) to get data back, defaults to current run, use 'all' for all
494527 event (str|list): the event(s) to include
495528 step (int|list): the step(s) to include
496529 key (str|list): the key(s) to include
497530 raw (bool): if True returns the raw data instead of a DataFrame
531+ lazy (bool): if True returns the Cursor instead of data, ignores raw
498532
499533 Returns:
500534 * data (DataFrame) if raw == False
@@ -519,12 +553,19 @@ def data(self, experiment=None, run=None, event=None, step=None, key=None, raw=F
519553 for k , v in extra .items ():
520554 if valid (k ):
521555 filter [f'data.{ k } ] = op (v )
522- data = self ._store .get (self ._data_name , filter = filter )
523- if data is not None and not raw :
556+ data = self ._store .get (self ._data_name , filter = filter , lazy = lazy )
557+ if data is not None and not raw and not lazy :
524558 data = pd .DataFrame .from_records (data )
525559 data .sort_values ('dt' , inplace = True ) if 'dt' in data .columns else None
526560 return data
527561
562+ def _initialize_dataset (self , force = False ):
563+ # create indexes when the dataset is first created
564+ if not force and self ._store .exists (self ._data_name ):
565+ return
566+ coll = self ._store .collection (self ._data_name )
567+ ensure_index (coll , {'data.run' : pymongo .ASCENDING , 'data.event' : pymongo .ASCENDING })
568+
528569 def restore_artifact (self , key = None , experiment = None , run = None , step = None , value = None ):
529570 """ restore a logged artificat
530571
@@ -596,7 +637,7 @@ class OmegaProfilingTracker(OmegaSimpleTracker):
596637 def __init__ (self , * args , ** kwargs ):
597638 super ().__init__ (* args , ** kwargs )
598639 self .profile_logs = []
599- self .max_buffer = 6
640+ self .max_buffer = 10
600641
601642 def log_profile (self , data ):
602643 """ the callback for BackgroundProfiler """
@@ -605,27 +646,32 @@ def log_profile(self, data):
605646 self .flush ()
606647
607648 def flush (self ):
608- for step , data in enumerate (self .profile_logs ):
609- # record the actual time instead of logging time (avoid buffering delays)
610- dt = data .get ('profile_dt' )
611- for k , v in data .items ():
612- self .log_event ('profile' , k , v , step = step , dt = dt )
613- self .profile_logs = []
614-
615- def start (self ):
649+ def log_items ():
650+ for step , data in enumerate (self .profile_logs ):
651+ # record the actual time instead of logging time (avoid buffering delays)
652+ dt = data .get ('profile_dt' )
653+ for k , v in data .items ():
654+ item = self ._common_log_data ('profile' , k , v , step = step , dt = dt )
655+ yield item
656+ if self .profile_logs :
657+ self ._store .put ([item for item in log_items ()], self ._data_name ,
658+ index = ['event' ], as_many = True , noversion = True )
659+ self .profile_logs = []
660+
661+ def start_runtime (self ):
616662 self .profiler = BackgroundProfiler (callback = self .log_profile )
617663 self .profiler .start ()
618- super ().start ()
664+ super ().start_runtime ()
619665
620- def stop (self ):
666+ def stop_runtime (self ):
621667 self .profiler .stop ()
622668 self .flush ()
623- super ().stop ()
669+ super ().stop_runtime ()
624670
625671
626672try :
627673 from tensorflow import keras
628- except :
674+ except Exception :
629675 pass
630676else :
631677 class TensorflowCallback (keras .callbacks .Callback ):
0 commit comments