[BUG] fixed memory leak in TimeSeriesDataset by using @cached_property and clean-up of index construction

[Vishnu-Rangiah]

Reference Issues/PRs

Fixes #648

What does this implement/fix? Explain your changes.

• Replaced @property and @lru_cache with @cached_property to fix a self-reference leak: previously, the cache kept strong references to every instance, preventing garbage collection and causing memory growth if many instances were created.
• Improved _construct_index() function to return only essential columns in a consistent format.

What should a reviewer concentrate their feedback on?

Did you add any tests for the change?

Any other comments?

Code to check memory leak and fix:

"""
Save as profile_timeseries.py
Run with:
    mprof run python profile_timeseries.py      # interactive plot after the run
or plain:
    python -m memory_profiler profile_timeseries.py
"""

import gc
import os
import time
import tracemalloc

import numpy as np
import pandas as pd
import psutil
from memory_profiler import profile
from pytorch_forecasting import TimeSeriesDataSet

process = psutil.Process(os.getpid())


def rss_mb() -> float:
    """Return resident set size (physical RAM) in MB."""
    try:
        return process.memory_info().rss / 1024**2
    except psutil.NoSuchProcess:
        # process gone: sampler is too late, just return 0
        return 0.0


# ---------- one-time data build ----------
test_data = pd.DataFrame(
    {
        "value": np.random.rand(3_000_000) - 0.5,
        "group": np.repeat(np.arange(3), 1_000_000),
        "time_idx": np.tile(np.arange(1_000_000), 3),
    }
)
print(f"Base DataFrame  RSS: {rss_mb():.1f} MB")

# ---------- enable tracemalloc ----------
tracemalloc.start()


# ---------- main loop ----------
@profile  # <- memory_profiler hook
def build_datasets(n_iter: int = 100):
    for i in range(n_iter):
        t0 = time.perf_counter()
        
        dataset = TimeSeriesDataSet(
            test_data,
            group_ids=["group"],
            target="value",
            time_idx="time_idx",
            min_encoder_length=5,
            max_encoder_length=5,
            min_prediction_length=2,
            max_prediction_length=2,
            time_varying_unknown_reals=["value"],
            predict_mode=False,
        )

        # ----------- point-in-time stats -----------
        current, peak = tracemalloc.get_traced_memory()
        print(
            f"[{i:03}] dataset built in {time.perf_counter()-t0:5.2f}s | "
            f"RSS {rss_mb():8.1f} MB | "
            f"tracemalloc current {current/1e6:8.1f} MB (peak {peak/1e6:8.1f} MB)"
        )

        # ----------- cleanup for the next loop -----------
        # del dataset
        gc.collect()
        tracemalloc.clear_traces()


build_datasets(10)

# ---------- post-mortem: where did the bytes go? ----------
top_stats = tracemalloc.take_snapshot().statistics("lineno")[:15]
print("\nTop 15 allocation sites:")
for stat in top_stats:
    print(stat)

PR checklist

• The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [BUG] - bugfix, [MNT] - CI, test framework, [ENH] - adding or improving code, [DOC] - writing or improving documentation or docstrings.
• Added/modified tests
• Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
  To run hooks independent of commit, execute pre-commit run --all-files

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@cae3174). Learn more about missing BASE report.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1905   +/-   ##
=======================================
  Coverage        ?   86.25%           
=======================================
  Files           ?       96           
  Lines           ?     7769           
  Branches        ?        0           
=======================================
  Hits            ?     6701           
  Misses          ?     1068           
  Partials        ?        0           

Flag	Coverage Δ
cpu	86.25% <100.00%> (?)
pytest	86.25% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fkiraly]

Looks good - the failures are not related to this PR.

[fnhirwa]

Looks good to me 😊