Source code for pycif.plugins.obsoperators.standard.transforms.utils.clean_memory
import code
import resource
import copy
import sys
from logging import debug
import tracemalloc
import pandas as pd
import numpy as np
from .fetch_inoutputs import fetch_inputs_outputs
from ......utils.datastores.compress import compress_datastore
[docs]
def clean_memory(self, transform_pipe, mapper, period_order,
ddi, transform, direction, mode, only_init):
# if self.monitor_memory:
# current, peak = tracemalloc.get_traced_memory()
# debug(f"Memory usage before cleaning is {current / 1024 ** 2}MB; "
# f"Peak was {peak / 1024 ** 2}MB")
transform_id = (ddi, transform, direction)
# Determine whether the transform outputs are needed
# by future transforms
transf_inout = "inputs" if direction == "adjoint" else "outputs"
neighbout_inout = "outputs" if direction == "adjoint" else "inputs"
pipe_links = self.pipe_links_adj if direction == "adjoint" else self.pipe_links_fwd
transform_ds = transform_pipe.datastore[transform][ddi][transf_inout]
used_outputs = {trid: {d: False for d in transform_ds[trid]}
for trid in transform_ds}
for trf_id in pipe_links[transform_id]:
dd_tmp, trf_tmp, direction_tmp = trf_id
# Check whether successor uses inputs/outputs
for trid in mapper[trf_tmp]["subsimus"][dd_tmp][neighbout_inout]:
if trid not in used_outputs:
continue
for d in mapper[trf_tmp]["subsimus"][dd_tmp][neighbout_inout][trid]:
if d in used_outputs[trid]:
used_outputs[trid][d] = True
# Now loop over the transform datastores and delete those not useful anymore
for trid in transform_pipe.datastore[transform][ddi][transf_inout]:
# Skip if data is further needed after initialization
if mapper[transform][transf_inout][trid].get("keep_data_after_init", False) and only_init:
continue
dates_orig = copy.deepcopy(
list(
transform_pipe.datastore[
transform][ddi][transf_inout][trid].keys()
)
)
for d in dates_orig:
if not used_outputs[trid].get(d, True):
clear_ds = transform_pipe.datastore[
transform][ddi][transf_inout][trid][d]
for tr in clear_ds:
clear_ds[tr] = None
clear_ds[tr] = {}
# Clean memory from the inputs in fwd and outputs in adj
# as they are not used anymore
# Exception with only_init as the input structure might be needed
transf_cleanin = "outputs" if direction == "adjoint" else "inputs"
transf_cleanout = "inputs" if direction == "adjoint" else "outputs"
for trid in transform_pipe.datastore[transform][ddi][transf_cleanin]:
needed_data = any([
mapper[transform][transf_cleanout][trid_in].get(
f"need_{transf_cleanin}", False)
for trid_in in mapper[transform][f'{transf_cleanin}2{transf_cleanout}'][trid]
if trid_in in mapper[transform][transf_cleanout]
])
if needed_data and only_init:
continue
dates_orig = copy.deepcopy(
list(
transform_pipe.datastore[
transform][ddi][transf_cleanin][trid].keys()
)
)
for d in dates_orig:
transform_pipe.datastore[
transform][ddi][transf_cleanin][trid][d] = None
transform_pipe.datastore[
transform][ddi][transf_cleanin][trid][d] = {}
# Now clean memory for past transforms
transform_ind = period_order.index(transform_id)
for trf_id in period_order[:transform_ind]:
if trf_id not in pipe_links:
continue
if transform_id not in pipe_links[trf_id]:
continue
dd_tmp, trf_tmp, direction_tmp = trf_id
# Check that other successors of the neighbours are still used
# Can flush inputs otherwise
if np.all(np.array(
[period_order.index(tr) for tr in pipe_links[trf_id]
if tr in period_order]) <= transform_ind):
neighbour_inout = \
"inputs" if direction_tmp == "adjoint" else "outputs"
ds_in = transform_pipe.datastore[trf_tmp][dd_tmp][neighbour_inout]
for trid in ds_in:
# Skip if data is further needed after initialization
if mapper[trf_tmp][neighbour_inout][trid].get("keep_data_after_init", False) and only_init:
continue
for d in ds_in[trid]:
for tr in ds_in[trid][d]:
ds_in[trid][d][tr] = None
ds_in[trid][d][tr] = {}
# # Force compression of datastores for remaining data
# for inout in ["inputs", "outputs"]:
# for trid_in in transform_pipe.datastore[transform][ddi][inout]:
# inoutputs = transform_pipe.datastore[transform][ddi][inout][trid_in]
# for d in inoutputs:
# if not isinstance(inoutputs[d], pd.DataFrame):
# continue
# inoutputs[d] = compress_datastore(inoutputs[d])
# if self.monitor_memory:
# current, peak = tracemalloc.get_traced_memory()
# debug(f"Memory usage after cleaning is {current / 1024 ** 2}MB; "
# f"Peak was {peak / 1024 ** 2}MB")
# for tr in transform_pipe.datastore:
# print(tr)
# print(transform_pipe.datastore[tr])
# print()
# print(__file__)
# import code
# code.interact(local=dict(locals(), **globals()))
return transform_pipe
# transform_direction = "adjoint" if mode == "adj" else "forward"
# transform_id = (ddi, transform, transform_direction)
# transform_id = (ddi, transform, direction)
# # Determine whether the transform outputs are needed
# # by future transforms
# transf_inout = "inputs" if mode == "adj" else "outputs"
# neighbout_inout = "outputs" if mode == "adj" else "inputs"
# pipe_links = self.pipe_links_adj if mode == "adj" else self.pipe_links_fwd
# transform_ds = transform_pipe.datastore[transform][ddi][transf_inout]
# used_outputs = {trid: {d: False for d in transform_ds[trid]}
# for trid in transform_ds}
# print(__file__)
# import code
# code.interact(local=dict(locals(), **globals()))
# for trf_id in pipe_links[transform_id]:
# dd_tmp, trf_tmp, direction_tmp = trf_id
# # Check whether successor uses inputs/outputs
# for trid in mapper[trf_tmp]["subsimus"][dd_tmp][neighbout_inout]:
# if trid not in used_outputs:
# continue
# for d in mapper[trf_tmp]["subsimus"][dd_tmp][neighbout_inout][trid]:
# if d in used_outputs[trid]:
# used_outputs[trid][d] = True
# # Now loop over the transform datastores and delete those not useful anymore
# for trid in transform_pipe.datastore[transform][ddi][transf_inout]:
# dates_orig = copy.deepcopy(
# list(
# transform_pipe.datastore[
# transform][ddi][transf_inout][trid].keys()
# )
# )
# for d in dates_orig:
# if not used_outputs[trid].get(d, True):
# clear_ds = transform_pipe.datastore[
# transform][ddi][transf_inout][trid][d]
# for tr in clear_ds:
# clear_ds[tr] = None
# clear_ds[tr] = {}
# # Clean memory from the inputs in fwd and outputs in adj
# # as they are not used anymore
# transf_cleanin = "outputs" if mode == "adj" else "inputs"
# for trid in transform_pipe.datastore[transform][ddi][transf_cleanin]:
# dates_orig = copy.deepcopy(
# list(
# transform_pipe.datastore[
# transform][ddi][transf_cleanin][trid].keys()
# )
# )
# for d in dates_orig:
# transform_pipe.datastore[
# transform][ddi][transf_cleanin][trid][d] = None
# transform_pipe.datastore[
# transform][ddi][transf_cleanin][trid][d] = {}
# # Now clean memory for past transforms
# transform_ind = period_order.index(transform_id)
# for trf_id in period_order[:transform_ind]:
# if trf_id not in pipe_links:
# continue
# if transform_id not in pipe_links[trf_id]:
# continue
# dd_tmp, trf_tmp, direction_tmp = trf_id
# # Check that other successors of the neighbours are still used
# # can flush inputs otherwise
# if np.all(np.array(
# [period_order.index(tr) for tr in pipe_links[trf_id]
# if tr in period_order]) <= transform_ind):
# neighbour_inout = \
# "inputs" if direction_tmp == "adjoint" else "outputs"
# ds_in = transform_pipe.datastore[trf_tmp][dd_tmp][neighbour_inout]
# for trid in ds_in:
# for d in ds_in[trid]:
# for tr in ds_in[trid][d]:
# ds_in[trid][d][tr] = None
# ds_in[trid][d][tr] = {}
# return transform_pipe
