Skip to content

Commit

Permalink
avoid unnecessary heap init
Browse files Browse the repository at this point in the history
  • Loading branch information
ankona committed Jul 24, 2024
1 parent 6f9ad9c commit d1a26c0
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions smartsim/_core/launcher/dragon/pqueue.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,9 +81,13 @@ def _initialize_reference_counters(self, nodes: t.List[Node]) -> None:
self._ref_map[node.hostname] = tracking_info

if node.num_gpus:
heapq.heappush(self._gpu_refs, tracking_info)
self._gpu_refs.append(tracking_info)
else:
heapq.heappush(self._cpu_refs, tracking_info)
self._cpu_refs.append(tracking_info)

# no need to reorder anything during initialization
# heapq.heapify(self._gpu_refs)
# heapq.heapify(self._cpu_refs)

def increment(self, host: str) -> None:
"""Directly increment the reference count of a given node and ensure the
Expand Down Expand Up @@ -191,7 +195,7 @@ def _get_next_available_node(
if ref_count == 0:
# increment the ref count before putting back onto heap
tracking_info[0] = ref_count + 1

if tracking_info is not None:
heapq.heappush(heap, tracking_info)

Expand Down

0 comments on commit d1a26c0

Please sign in to comment.