optim: Bunch of micro-optimizations

automl · Dec 10, 2024 · c6a5003 · c6a5003
1 parent a51a286
commit c6a5003
Show file tree

Hide file tree

Showing 16 changed files with 336 additions and 494 deletions.
diff --git a/neps/env.py b/neps/env.py
@@ -40,7 +40,6 @@ def yaml_or_json(e: str) -> Literal["yaml", "json"]:
     parse=str,
     default="lockf",
 )
-
 MAX_RETRIES_GET_NEXT_TRIAL = get_env(
     "NEPS_MAX_RETRIES_GET_NEXT_TRIAL",
     parse=int,
@@ -96,7 +95,6 @@ def yaml_or_json(e: str) -> Literal["yaml", "json"]:
     parse=lambda e: None if is_nullable(e) else float(e),
     default=120,
 )
-
 GLOBAL_ERR_FILELOCK_POLL = get_env(
     "NEPS_GLOBAL_ERR_FILELOCK_POLL",
     parse=float,

diff --git a/neps/runtime.py b/neps/runtime.py
@@ -37,8 +37,10 @@
 from neps.state._eval import evaluate_trial
 from neps.state.neps_state import NePSState
 from neps.state.optimizer import BudgetInfo, OptimizationState, OptimizerInfo
+from neps.state.seed_snapshot import SeedSnapshot
 from neps.state.settings import DefaultReportValues, OnErrorPossibilities, WorkerSettings
 from neps.state.trial import Trial
+from neps.utils.common import gc_disabled
 
 if TYPE_CHECKING:
     from neps.optimizers.base_optimizer import BaseOptimizer
@@ -358,7 +360,9 @@ def _get_next_trial(self) -> Trial | Literal["break"]:
             # NOTE: It's important to release the trial lock before sampling
             # as otherwise, any other service, such as reporting the result
             # of a trial. Hence we do not lock these together with the above.
-            with self.state._trial_lock.lock(worker_id=self.worker_id):
+            # OPTIM: We try to prevent garbage collection from happening in here to
+            # minimize time spent holding on to the lock.
+            with self.state._trial_lock.lock(worker_id=self.worker_id), gc_disabled():
                 # Give the file-system some time to sync if we encountered out-of-order
                 # issues with this worker.
                 if self._GRACE > 0:
@@ -375,7 +379,7 @@ def _get_next_trial(self) -> Trial | Literal["break"]:
                 pending_trials = [
                     trial
                     for trial in trials.values()
-                    if trial.state == Trial.State.PENDING
+                    if trial.metadata.state == Trial.State.PENDING
                 ]
 
                 if len(pending_trials) > 0:
@@ -387,10 +391,7 @@ def _get_next_trial(self) -> Trial | Literal["break"]:
                         time_started=time.time(),
                         worker_id=self.worker_id,
                     )
-                    self.state._trials.update_trial(
-                        earliest_pending,
-                        hints=["metadata", "state"],
-                    )
+                    self.state._trials.update_trial(earliest_pending, hints="metadata")
                     logger.info(
                         "Worker '%s' picked up pending trial: %s.",
                         self.worker_id,
@@ -404,7 +405,7 @@ def _get_next_trial(self) -> Trial | Literal["break"]:
                 trials=trials,
             )
 
-            with self.state._trial_lock.lock(worker_id=self.worker_id):
+            with self.state._trial_lock.lock(worker_id=self.worker_id), gc_disabled():
                 try:
                     sampled_trial.set_evaluating(
                         time_started=time.time(),
@@ -546,7 +547,7 @@ def run(self) -> None:  # noqa: C901, PLR0915
                 "Worker '%s' evaluated trial: %s as %s.",
                 self.worker_id,
                 evaluated_trial.id,
-                evaluated_trial.state,
+                evaluated_trial.metadata.state,
             )
 
             if report.cost is not None:
@@ -610,6 +611,7 @@ def _launch_runtime(  # noqa: PLR0913
                 load_only=False,
                 optimizer_info=OptimizerInfo(optimizer_info),
                 optimizer_state=OptimizationState(
+                    seed_snapshot=SeedSnapshot.new_capture(),
                     budget=(
                         BudgetInfo(
                             max_cost_budget=max_cost_total,
@@ -618,7 +620,7 @@ def _launch_runtime(  # noqa: PLR0913
                             used_evaluations=0,
                         )
                     ),
-                    shared_state={},  # TODO: Unused for the time being...
+                    shared_state=None,  # TODO: Unused for the time being...
                 ),
             )
             break