Close HDF5 handles eagerly when reading and writing MPS and MPO

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorMPS"
 uuid = "0d1a4710-d33b-49a5-8f18-73bdf49b47e2"
-version = "0.4.1"
+version = "0.4.2"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>", "Miles Stoudenmire <mstoudenmire@flatironinstitute.org>"]
 
 [workspace]
@@ -38,11 +38,11 @@ Adapt = "4.1"
 ChainRulesCore = "1.10"
 Compat = "4.16"
 HDF5 = "0.14, 0.15, 0.16, 0.17"
-ITensors = "0.8, 0.9"
+ITensors = "0.9.32"
 IsApprox = "2"
 KrylovKit = "0.8.1, 0.9, 0.10"
 LinearAlgebra = "1.10"
-NDTensors = "0.4"
+NDTensors = "0.4.32"
 Observers = "0.2"
 PackageCompiler = "1, 2"
 Printf = "1.10"

ext/ITensorMPSHDF5Ext/ITensorMPSHDF5Ext.jl

@@ -1,4 +1,15 @@
 module ITensorMPSHDF5Ext
+
+# HDF5.jl has no do-block form of `open_group` or `create_group`, so without this
+# the handles opened below stay alive until their finalizers run.
+function closeafter(f, obj)
+    try
+        return f(obj)
+    finally
+        close(obj)
+    end
+end
+
 include("mps.jl")
 include("mpo.jl")
 end

ext/ITensorMPSHDF5Ext/mpo.jl

@@ -1,29 +1,31 @@
-using HDF5: HDF5, attributes, create_group, open_group, read, write
+using HDF5: HDF5, attributes, create_group, open_group, read, read_attribute, write
 using ITensorMPS: MPO
 using ITensors: ITensor
 
 function HDF5.write(parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, M::MPO)
-    g = create_group(parent, name)
-    attributes(g)["type"] = "MPO"
-    attributes(g)["version"] = 1
-    N = length(M)
-    write(g, "rlim", M.rlim)
-    write(g, "llim", M.llim)
-    write(g, "length", N)
-    for n in 1:N
-        write(g, "MPO[$(n)]", M[n])
+    closeafter(create_group(parent, name)) do g
+        attributes(g)["type"] = "MPO"
+        attributes(g)["version"] = 1
+        N = length(M)
+        write(g, "rlim", M.rlim)
+        write(g, "llim", M.llim)
+        write(g, "length", N)
+        for n in 1:N
+            write(g, "MPO[$(n)]", M[n])
+        end
     end
     return
 end
 
 function HDF5.read(parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{MPO})
-    g = open_group(parent, name)
-    if read(attributes(g)["type"]) != "MPO"
-        error("HDF5 group or file does not contain MPO data")
+    return closeafter(open_group(parent, name)) do g
+        if read_attribute(g, "type") != "MPO"
+            error("HDF5 group or file does not contain MPO data")
+        end
+        N = read(g, "length")
+        rlim = read(g, "rlim")
+        llim = read(g, "llim")
+        v = [read(g, "MPO[$(i)]", ITensor) for i in 1:N]
+        return MPO(v, llim, rlim)
     end
-    N = read(g, "length")
-    rlim = read(g, "rlim")
-    llim = read(g, "llim")
-    v = [read(g, "MPO[$(i)]", ITensor) for i in 1:N]
-    return MPO(v, llim, rlim)
 end

ext/ITensorMPSHDF5Ext/mps.jl

@@ -1,29 +1,31 @@
-using HDF5: HDF5, attributes, create_group, open_group, read, write
+using HDF5: HDF5, attributes, create_group, open_group, read, read_attribute, write
 using ITensorMPS: MPS
 using ITensors: ITensor
 
 function HDF5.write(parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, M::MPS)
-    g = create_group(parent, name)
-    attributes(g)["type"] = "MPS"
-    attributes(g)["version"] = 1
-    N = length(M)
-    write(g, "length", N)
-    write(g, "rlim", M.rlim)
-    write(g, "llim", M.llim)
-    for n in 1:N
-        write(g, "MPS[$(n)]", M[n])
+    closeafter(create_group(parent, name)) do g
+        attributes(g)["type"] = "MPS"
+        attributes(g)["version"] = 1
+        N = length(M)
+        write(g, "length", N)
+        write(g, "rlim", M.rlim)
+        write(g, "llim", M.llim)
+        for n in 1:N
+            write(g, "MPS[$(n)]", M[n])
+        end
     end
     return
 end
 
 function HDF5.read(parent::Union{HDF5.File, HDF5.Group}, name::AbstractString, ::Type{MPS})
-    g = open_group(parent, name)
-    if read(attributes(g)["type"]) != "MPS"
-        error("HDF5 group or file does not contain MPS data")
+    return closeafter(open_group(parent, name)) do g
+        if read_attribute(g, "type") != "MPS"
+            error("HDF5 group or file does not contain MPS data")
+        end
+        N = read(g, "length")
+        rlim = read(g, "rlim")
+        llim = read(g, "llim")
+        v = [read(g, "MPS[$(i)]", ITensor) for i in 1:N]
+        return MPS(v, llim, rlim)
     end
-    N = read(g, "length")
-    rlim = read(g, "rlim")
-    llim = read(g, "llim")
-    v = [read(g, "MPS[$(i)]", ITensor) for i in 1:N]
-    return MPS(v, llim, rlim)
 end

test/base/test_readwrite.jl

@@ -7,37 +7,54 @@ using Test
 include(joinpath(@__DIR__, "utils", "util.jl"))
 
 @testset "HDF5 Read and Write" begin
-    @testset "MPO/MPS" begin
-        N = 6
-        sites = siteinds("S=1/2", N)
-
-        # MPO
-        mpo = makeRandomMPO(sites)
-
-        h5open("data.h5", "w") do fo
-            return write(fo, "mpo", mpo)
-        end
-
-        h5open("data.h5", "r") do fi
-            rmpo = read(fi, "mpo", MPO)
-            @test prod([norm(rmpo[i] - mpo[i]) / norm(mpo[i]) < 1.0e-10 for i in 1:N])
+    mktempdir() do dir
+        fn = joinpath(dir, "data.h5")
+        @testset "MPO/MPS" begin
+            N = 6
+            sites = siteinds("S=1/2", N)
+
+            # MPO
+            mpo = makeRandomMPO(sites)
+
+            h5open(fn, "w") do fo
+                return write(fo, "mpo", mpo)
+            end
+
+            h5open(fn, "r") do fi
+                rmpo = read(fi, "mpo", MPO)
+                @test prod([norm(rmpo[i] - mpo[i]) / norm(mpo[i]) < 1.0e-10 for i in 1:N])
+            end
+
+            # MPS
+            mps = makeRandomMPS(sites)
+            h5open(fn, "w") do fo
+                return write(fo, "mps", mps)
+            end
+
+            h5open(fn, "r") do fi
+                rmps = read(fi, "mps", MPS)
+                @test prod([norm(rmps[i] - mps[i]) / norm(mps[i]) < 1.0e-10 for i in 1:N])
+            end
         end
 
-        # MPS
-        mps = makeRandomMPS(sites)
-        h5open("data.h5", "w") do fo
-            return write(fo, "mps", mps)
-        end
-
-        h5open("data.h5", "r") do fi
-            rmps = read(fi, "mps", MPS)
-            @test prod([norm(rmps[i] - mps[i]) / norm(mps[i]) < 1.0e-10 for i in 1:N])
+        @testset "No leaked HDF5 handles" begin
+            nopen(file) = HDF5.API.h5f_get_obj_count(file, HDF5.API.H5F_OBJ_ALL)
+            sites = siteinds("S=1/2", 20)
+            mps = makeRandomMPS(sites)
+            mpo = makeRandomMPO(sites)
+            GC.gc()
+            h5open(fn, "w") do fo
+                write(fo, "mps", mps)
+                write(fo, "mpo", mpo)
+                # `h5f_get_obj_count` includes the file handle itself.
+                @test nopen(fo) == 1
+            end
+            h5open(fn, "r") do fi
+                read(fi, "mps", MPS)
+                read(fi, "mpo", MPO)
+                @test nopen(fi) == 1
+            end
         end
     end
-
-    #
-    # Clean up the test hdf5 file
-    #
-    rm("data.h5"; force = true)
 end
 end