fix: use trace count as index (#909)

This commit is contained in:
Tim
2025-05-29 16:52:37 +08:00
committed by GitHub
parent 63ebdb5ada
commit 1091d5ea29
@@ -70,7 +70,7 @@ class LimitTimeCKPSelector(CheckpointSelector):
return (-1,) # Continue with latest trial for new sub-trace
# Calculate elapsed time for current sub-trace
elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count]
elapsed_time = current_time - self.sub_trace_start_times[trace.sub_trace_count - 1]
if elapsed_time < self.time_limit_pre_trace:
# Continue with current sub-trace
@@ -81,7 +81,7 @@ class LimitTimeCKPSelector(CheckpointSelector):
return (-1,)
else:
# Check if we've reached the maximum number of traces
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
logger.info(
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
)
@@ -89,11 +89,11 @@ class LimitTimeCKPSelector(CheckpointSelector):
return (-1,)
# Time limit exceeded, start a new sub-trace
self.sub_trace_start_times[trace.sub_trace_count + 1] = current_time
self.sub_trace_start_times[trace.sub_trace_count] = current_time
logger.info(
f"Elapsed time {elapsed_time} exceeds time limit {self.time_limit_pre_trace}, jump to a new sub-trace"
)
logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
logger.info(f"current sub-trace count: {trace.sub_trace_count}")
return tuple() # Empty tuple signals starting a new sub-trace
@@ -129,7 +129,7 @@ class SOTAJumpCKPSelector(CheckpointSelector):
sota_count += 1
if sota_count < self.SOTA_COUNT_THRESHOLD:
# Check if we've reached the maximum number of traces
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
logger.info(
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
)
@@ -139,7 +139,7 @@ class SOTAJumpCKPSelector(CheckpointSelector):
logger.info(
f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump to a new sub-trace"
)
logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
logger.info(f"current sub-trace count: {trace.sub_trace_count}")
return ()
else:
logger.info(
@@ -189,7 +189,7 @@ class BackJumpCKPSelector(CheckpointSelector):
if sota_count < self.SOTA_COUNT_THRESHOLD:
# Check if we've reached the maximum number of traces before creating a new one
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
logger.info(
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
)
@@ -216,7 +216,7 @@ class BackJumpCKPSelector(CheckpointSelector):
return (last_second_sota_idx,)
else:
# Check max trace limit again before creating a new trace
if trace.sub_trace_count + 1 >= self.MAX_TRACE_NUM:
if trace.sub_trace_count >= self.MAX_TRACE_NUM:
logger.info(
f"Reached maximum trace count ({self.MAX_TRACE_NUM}), continuing with the current sub-trace"
)
@@ -226,7 +226,7 @@ class BackJumpCKPSelector(CheckpointSelector):
logger.info(
f"SOTA count {sota_count} is below threshold {self.SOTA_COUNT_THRESHOLD}, jump a new sub-trace"
)
logger.info(f"current sub-trace count: {trace.sub_trace_count + 1}")
logger.info(f"current sub-trace count: {trace.sub_trace_count}")
return () # reboot a new sub-trace
else: