Ignore non-event files in deserialized folder (#89)

* Ignore non-event files in deserialized folder

* add test for #88

* bump version

* fix header

Author: PhilipVinc

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorBoardLogger"
 uuid = "899adc3e-224a-11e9-021f-63837185c80f"
 authors = ["Filippo Vicentini <[email protected]>"]
-version = "0.1.13"
+version = "0.1.14"
 
 [deps]
 CRC32c = "8bf52ea8-c179-5cab-976a-9e18b702a9bc"

src/Deserialization/deserialization.jl

@@ -1,5 +1,27 @@
 export summary_iterator
 
+"""
+    is_valid_event(f::IOStream) => Bool
+
+Returns true if the stream points to a valid TensorBoard event, false overwise.
+This is accomplished by checeking the crc checksum on the header (first 8
+bytes) of the event.
+"""
+function is_valid_event(f::IOStream)
+    eof(f) && return false
+
+    header = read(f, 8)
+    length(header) != 8 && return false
+
+    crc_header = read(f, 4)
+    length(crc_header) != 4 && return false
+
+    # check
+    crc_header_ck = reinterpret(UInt8, UInt32[masked_crc32c(header)])
+    return crc_header == crc_header_ck
+end
+
+
 """
     read_event(f::IOStream) => Event
 
@@ -13,7 +35,9 @@ function read_event(f::IOStream)
 
     # check
     crc_header_ck = reinterpret(UInt8, UInt32[masked_crc32c(header)])
-    @assert crc_header == crc_header_ck
+    if crc_header != crc_header_ck
+        error("Invalid event checksum for stream", f)
+    end
 
     # read data
     data_len = first(reinterpret(Int64, header))
@@ -42,12 +66,27 @@ struct TBEventFileCollectionIterator
 
     purge::Bool
 end
-TBEventFileCollectionIterator(path; purge=true) =
-    TBEventFileCollectionIterator(path, sort(readdir(path)), purge)
 
 TBEventFileCollectionIterator(logger::TBLogger; purge=true) =
     TBEventFileCollectionIterator(logdir(logger), purge=true)
 
+function TBEventFileCollectionIterator(path; purge=true)
+    fnames = sort(readdir(path))
+    good_fnames = typeof(fnames)()
+
+    # Only consider files whose first event file would be valid.
+    # So if there are other files in this folder, we ignore them.
+    for fname in fnames
+        open(joinpath(path,fname), "r") do f
+            is_valid_event(f) && push!(good_fnames, fname)
+        end
+    end
+
+    @debug "Valid TensorBoard event files" good_fnames
+
+    TBEventFileCollectionIterator(path, good_fnames, purge)
+end
+
 function Base.iterate(it::TBEventFileCollectionIterator, state=1)
     state > length(it.files) && return nothing
     fstream = open(joinpath(it.dir, it.files[state]))

test/deserialization.jl

@@ -17,6 +17,12 @@ ENV["GKSwstype"] = "100"
         end
     end
 
+    # Issue #88: ignore non -tensorboard files
+    dir = TensorBoardLogger.logdir(logger)
+    f = open(joinpath(dir, "mockdata"), "w")
+    write(f, "baddta")
+    close(f)
+    
     tgs = TensorBoardLogger.tags(logger)
     @test "test/val" ∈ tgs
     @test "test/b" ∈ tgs
@@ -29,4 +35,5 @@ ENV["GKSwstype"] = "100"
     @test all(hist["test/val"].values .== (1-0.5*im).* is)
     @test all(hist["test/b"].values == is .* 2)
     @test all([v == mri for v=hist["test2/mri"].values])
+
 end