33import resource
44import sys
55from enum import Enum
6- from typing import TYPE_CHECKING , List , Optional , Sequence , Union , cast
6+ from pathlib import Path
7+ from typing import TYPE_CHECKING , List , Literal , Optional , Sequence , Union , cast
78
8- from seml .database import get_collection
9+ from pymongo .collection import Collection
10+
11+ from seml .database import States , get_collection
12+ from seml .document import ExperimentDoc
913from seml .experiment .observers import create_mongodb_observer
1014from seml .settings import SETTINGS
1115from seml .utils .multi_process import is_main_process
2428 recursive_fill_in ,
2529 undogmatize ,
2630)
31+ from sacred .utils import SacredInterrupt
2732
2833
2934class LoggerOptions (Enum ):
@@ -32,7 +37,13 @@ class LoggerOptions(Enum):
3237 RICH = 'rich'
3338
3439
40+ class RescheduleInterrupt (SacredInterrupt ):
41+ STATUS = States .PENDING [0 ]
42+
43+
3544class Experiment (ExperimentBase ):
45+ _reschedule_watch_path : Path | None | Literal ['INVALID' ]
46+
3647 def __init__ (
3748 self ,
3849 name : Optional [str ] = None ,
@@ -65,6 +76,8 @@ def __init__(
6576 if collect_stats :
6677 self .post_run_hook (lambda _run : _collect_exp_stats (_run ))
6778
79+ self ._reschedule_watch_path = None
80+
6881 def run (
6982 self ,
7083 command_name : Optional [str ] = None ,
@@ -87,6 +100,138 @@ def run(
87100 options = options ,
88101 )
89102
103+ def reschedule_hook (self , f ):
104+ """Decorator to register the reschedule hook.
105+
106+ In case of a required rescheduling, the decorated function will be called.
107+ It should return a config dictionary that will be used to update the current
108+ configuration before submitting the rescheduled job.
109+
110+ Parameters
111+ ----------
112+ f: Callable[[..., dict]]
113+ User-defined function to be called when rescheduling is triggered.
114+
115+ Returns
116+ -------
117+ Callable
118+ Wrapped function that checks for rescheduling and calls the user-defined function.
119+
120+
121+ Example
122+ -------
123+
124+ ```python
125+ from seml import Experiment
126+ ex = Experiment()
127+
128+ @ex.reschedule_hook
129+ def reschedule(step: int):
130+ print(f'Reschedule triggered at step {step}.')
131+ return {'checkpoint': step}
132+
133+ @ex.automain
134+ def run(n_steps: int, checkpoint: int | None = None):
135+ if checkpoint is not None:
136+ print(f'Resuming from checkpoint: {checkpoint}')
137+ for step in range(checkpoint or 0, n_steps):
138+ reschedule(step)
139+ print(f'Processing step {step + 1}/{n_steps}')
140+ print('Experiment completed successfully.')
141+ ```
142+ """
143+
144+ def _reschedule_hook (* args , ** kwargs ):
145+ # Check whether we must reschedule
146+ if self ._reschedule_watch_path is None :
147+ self ._reschedule_watch_path = (
148+ self ._get_reschedule_watch_path () or 'INVALID'
149+ )
150+ if self ._reschedule_watch_path == 'INVALID' or not is_main_process ():
151+ return
152+
153+ _must_reschedule = Path (self ._reschedule_watch_path ).exists ()
154+ if not _must_reschedule :
155+ return
156+
157+ # If yes, call the user-defined function
158+ logging .info ('Caught reschedule signal, calling reschedule_hook.' )
159+ new_config = f (* args , ** kwargs )
160+ assert isinstance (new_config , dict ), (
161+ 'Reschedule hook must return a configuration dictionary.'
162+ f' Got type: + { type (new_config )}
163+ )
164+ logging .info ('Reschedule hook returned new configuration: ' f'{ new_config } )
165+
166+ # Merge config update with experiment doc
167+ exp = self ._get_exp_document_from_db ()
168+ assert exp is not None , (
169+ 'Could not retrieve experiment document from database.'
170+ ' Rescheduling not possible.'
171+ )
172+ collection = self ._get_db_collection ()
173+ assert collection is not None
174+ # TODO: TIGRU: Update config of experiment document in the database
175+
176+ raise RescheduleInterrupt
177+
178+ return _reschedule_hook
179+
180+ def _get_db_collection (self ) -> Optional [Collection [ExperimentDoc ]]:
181+ assert self .current_run is not None
182+ db_collection = self .current_run .config .get ('db_collection' )
183+ if db_collection is None :
184+ logging .warning (
185+ 'Reschedule hook called outside of a SEML-managed experiment.'
186+ ' Rescheduling hook is ignored.'
187+ )
188+ return None
189+ db_collection = get_collection (db_collection )
190+ return db_collection
191+
192+ def _get_exp_id (self ) -> Optional [int ]:
193+ assert self .current_run is not None
194+ exp_id = self .current_run .config .get ('overwrite' )
195+ if exp_id is None :
196+ logging .warning (
197+ 'Reschedule hook called without a SEML experiment ID.'
198+ ' Is this execution unobserved? Rescheduling hook is ignored.'
199+ )
200+ return exp_id
201+
202+ def _get_exp_document_from_db (self ) -> Optional [ExperimentDoc ]:
203+ db_collection = self ._get_db_collection ()
204+ exp_id = self ._get_exp_id ()
205+ if db_collection is None or exp_id is None :
206+ return None
207+ exp = db_collection .find_one (exp_id )
208+ if exp is None :
209+ logging .warning (
210+ f'No experiment with ID { exp_id }
211+ ' Cannot retrieve experiment document.'
212+ )
213+ return exp
214+
215+ def _get_reschedule_watch_path (self ) -> Optional [Path ]:
216+ exp = self ._get_exp_document_from_db ()
217+ if exp is None :
218+ return
219+
220+ if exp .get ('execution' , {}).get ('cluster' , None ) == 'local' :
221+ logging .info ('Experiment is executed locally. Reschedule hook is ignored.' )
222+ return
223+
224+ reschedule_path = exp .get ('execution' , {}).get ('reschedule_file' , None )
225+ assert (
226+ reschedule_path is not None
227+ ), 'No SLURM reschedule file recorded for this experiment in the database.'
228+ reschedule_path = Path (reschedule_path )
229+ logging .info (
230+ 'Found signal path for reschedule hook.'
231+ f' Watching for file: { reschedule_path .as_posix ()}
232+ )
233+ return reschedule_path
234+
90235
91236class MongoDbObserverConfig :
92237 def __init__ (self , experiment : Experiment ):
0 commit comments