Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
…-hacks which were merged in by accident
  • Loading branch information
o-smirnov committed Dec 2, 2024
1 parent 79d91e7 commit c8c47c5
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 34 deletions.
36 changes: 6 additions & 30 deletions stimela/task_stats.py
Original file line number Diff line number Diff line change
Expand Up @@ -237,7 +237,7 @@ def update_stats(now: datetime, sample: TaskStatsDatum):
_taskstats[key][0] = (now - start).total_seconds()


def update_process_status(stimela_process=None, child_processes=None):
def update_process_status():
# current subtask info
ti = _task_stack[-1] if _task_stack else None

Expand All @@ -247,27 +247,11 @@ def update_process_status(stimela_process=None, child_processes=None):

# form up sample datum
s = TaskStatsDatum(num_samples=1)

# If we have the process objects, we don't need to block.
interval = 0 if child_processes else 1

# Grab the stimela process and its children (recursively).
stimela_process = stimela_process or psutil.Process()
stimela_children = child_processes or stimela_process.children(recursive=True)

# Assume that all child processes belong to the same task.
# TODO: Handling of children is rudimentary at present.
# How would this work for scattered/parallel steps?
if stimela_children and ti:
processes = stimela_children
else:
processes = [stimela_process]

# CPU and memory
s.cpu = sum(p.cpu_percent(interval=interval) for p in processes)
system_memory = psutil.virtual_memory().total
s.mem_used = round(sum(p.memory_info().rss for p in processes) / 2**30)
s.mem_total = round(system_memory / 2**30)
s.cpu = psutil.cpu_percent()
mem = psutil.virtual_memory()
s.mem_used = round(mem.total*mem.percent/100 / 2**30)
s.mem_total = round(mem.total / 2**30)
# load
s.load, _, _ = psutil.getloadavg()

Expand Down Expand Up @@ -335,16 +319,8 @@ def update_process_status(stimela_process=None, child_processes=None):
async def run_process_status_update():
if progress_bar:
with contextlib.suppress(asyncio.CancelledError):
stimela_process = psutil.Process()
child_processes = set()
while True:
new_children = {c for c in stimela_process.children(recursive=True) if c not in child_processes}
old_children = {c for c in child_processes if c not in stimela_process.children(recursive=True)}
child_processes = child_processes.union(new_children) - old_children
try:
update_process_status(stimela_process, child_processes)
except psutil.NoSuchProcess:
continue
update_process_status()
await asyncio.sleep(1)

_printed_stats = dict(
Expand Down
8 changes: 4 additions & 4 deletions tests/stimela_tests/test_recipe.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,22 +89,22 @@ def test_test_loop_recipe():
assert retcode == 0

print("===== expecting no errors now =====")
retcode = os.system("stimela exec test_loop_recipe.yml same_as_cubical_image_loop ms=foo")
retcode = os.system("stimela -v -b native exec test_loop_recipe.yml same_as_cubical_image_loop ms=foo")
assert retcode == 0

for name in "abc":
msname = f"test-{name}.ms"
if not os.path.exists(msname):
os.mkdir(msname)
print("===== expecting no errors now =====")
retcode = os.system("stimela exec test_loop_recipe.yml loop_recipe")
retcode = os.system("stimela -v -b native exec test_loop_recipe.yml loop_recipe")
assert retcode == 0

def test_scatter():
print("===== expecting no errors now =====")
retcode = os.system("stimela exec test_scatter.yml basic_loop")
retcode = os.system("stimela -v -b native exec test_scatter.yml basic_loop")
assert retcode == 0

print("===== expecting no errors now =====")
retcode = os.system("stimela exec test_scatter.yml nested_loop")
retcode = os.system("stimela -v -b native exec test_scatter.yml nested_loop")
assert retcode == 0

0 comments on commit c8c47c5

Please sign in to comment.