Drive memory utilization cleanup

Nate-Haines
New Contributor III

Issue Summary:

When running multiple jobs on the same compute cluster, over time, I see an increase in memory utilization that is seemingly never fully released, even when jobs finish. This eventually leads to some jobs stalling out as memory hits the upper limit, as well as cluster crashes. I have replicated the issue by running a single job in a notebook environment (connected to the cluster), detaching the notebook, then running again, and I indeed see the consistent uptick in memory with each subsequent run. From this test it seems that the cluster is never fully releasing memory after a job (or notebook) is finished/detached. 


Solutions Attempted:

  1. Manually deleting objects produced throughout the job:

    • del object

    Outcome:
    Objects removed, but memory usage in the cluster did not decrease significantly.

  2. Garbage Collection:

    • Triggered Python’s garbage collection using gc.collect() after deleting objects.

    Outcome:
    No noticeable reduction in memory usage.

  3. Notebook Detach and Reattach:

    • Replicated the job in a notebook, detaching the notebook once job is done.

    Outcome:
    Produces a small decrease in memory util, but not a return to what util looked like before attaching + running the notebook.

  4. Restarting the Cluster:

    • This works as intended, but is not viable given the cluster will be running concurrent jobs in practice.

Additional Context:

  • My workflow is specifically using Ibis with the Polars backend in the Databricks job. It involves creating many memtables dynamically for temporary use, which I assume should just all be cleaned up once the job is finished. The Pyspark backend does not work well for my use case so is not a solution here.
  • The Databricks cluster is configured with Shared access mode with a Shared Compute policy. The driver node is the node getting the high memory util that is not released. The driver has 128 GB RAM, 16 cores (r5d.4xlarge)
  • Restarting the cluster between jobs is not an option due to concurrent runs.

Request:

I am looking for:

  1. A reliable way to return memory utilization back to normal once a job is finished, whether configured as a notebook or Python script.
  2. General guidance on managing memory in shared compute Databricks clusters. 
  3. Any insights into how memory management works with Databricks clusters, and if particular configurations would be better suited for my use case.