Add ZoneProfilers support (#40)

* Add ZoneProfilers support

* Apply suggestions from code review

Co-authored-by: Copilot <[email protected]>

* fix copilot typofix

* Improve zone text

* Add zone for io

---------

Co-authored-by: Copilot <[email protected]>

Author: nhz2

.github/workflows/CI.yml

@@ -21,7 +21,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.11'
+          - '1.12'
           - 'pre'
           - 'nightly'
         os:

Project.toml

@@ -1,7 +1,10 @@
 name = "MEDYANSimRunner"
 uuid = "b58a3b99-22e3-44d1-b5ea-258f082a6fe8"
+version = "0.6.2-dev"
 authors = ["nhz2 <[email protected]>"]
-version = "0.6.1"
+
+[workspace]
+projects = ["test"]
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
@@ -16,6 +19,7 @@ OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
 SmallZarrGroups = "d423b6e5-1c84-4ae2-8d2d-b903aee15ac7"
+ZoneProfilers = "402d8ed9-92af-47cb-aae4-ae215e44b74b"
 
 [compat]
 ArgCheck = "2"
@@ -30,7 +34,5 @@ OrderedCollections = "1"
 Random = "1"
 SHA = "0.7, 1"
 SmallZarrGroups = "0.8.3"
-julia = "1.11"
-
-[workspace]
-projects = ["test"]
+ZoneProfilers = "0.1"
+julia = "1.12"

README.md

@@ -51,6 +51,40 @@ end
 
 To run the simulation if `main.jl` is called as a julia script.
 
+## Performance Profiling with ZoneProfilers
+
+MEDYANSimRunner integrates with [ZoneProfilers.jl](https://github.com/nhz2/ZoneProfilers.jl) to provide detailed performance instrumentation of your simulations. This allows you to visualize where time is spent in your `setup`, `loop`, `save`, `load`, and `done` functions, as well as in file I/O operations.
+
+### Using the Profiler
+
+To enable profiling, pass a profiler instance to the `run` function:
+
+```julia
+if abspath(PROGRAM_FILE) == @__FILE__
+    # Open and connect to the tracy GUI
+    using ZoneProfilerTracy: TracyProfiler
+    import TracyProfiler_jll
+    profiler = TracyProfiler(TracyProfiler_jll)
+    MEDYANSimRunner.run(ARGS; jobs, setup, loop, load, save, done, profiler)
+end
+```
+
+### Adding Custom Profiling
+
+You can add additional profiling zones within your user functions by using the `profiler` keyword argument:
+
+```julia
+function loop(step::Int, state; output, profiler=NullProfiler())
+    state = @zone profiler compute_physics(state)
+    @zone profiler collect_data!(output, state)
+    return state
+end
+```
+
+For production runs without profiling overhead, simply omit the `profiler` parameter (defaults to `NullProfiler()` which has zero runtime cost).
+
+See the [ZoneProfilers.jl documentation](https://github.com/nhz2/ZoneProfilers.jl) for more advanced profiling features.
+
 #### Standard input parameters.
  - `step::Int`: starts out at 0 after setup and is incremented right before every call to `loop`.
 

src/MEDYANSimRunner.jl

@@ -1,5 +1,19 @@
 module MEDYANSimRunner
 
+using ZoneProfilers: NullProfiler, @zone, frame_mark!, zone_text!, zone_color!, zone_active, message!, app_info!
+using ArgCheck: ArgCheck, @argcheck
+using Logging: Logging, current_logger, with_logger
+using SmallZarrGroups: SmallZarrGroups, ZGroup, attrs
+import InteractiveUtils
+import LoggingExtras
+import JSON3
+import Dates
+using SHA: sha256
+import FileWatching
+import OrderedCollections
+using Random: Random, RandomDevice
+import DeepDiffs
+
 include("constants.jl")
 include("rng-load-save.jl")
 include("file-saving.jl")

src/file-saving.jl

@@ -4,14 +4,6 @@ Helper functions to try to reliably save trajectories on distributed file system
 None of these functions try to be thread safe, or prevent file system attacks.
 =#
 
-
-using SHA: sha256
-using Random: RandomDevice
-import Dates
-import LoggingExtras
-using Logging
-using ArgCheck
-
 const DATE_FORMAT = Dates.dateformat"yyyy-mm-ddTHH:MM:SS"
 
 function timestamp_logger(file, level) 
@@ -44,6 +36,7 @@ end
 
 
 """
+Return the sha256 hash of the data.
 Ensure a file with the contents `data` exists at the path `joinpath(dir_name, file_name)`
 
 This function can fail or be interrupted, in that case, an existing file 
@@ -58,63 +51,52 @@ If interrupted, there may be a temporary file left behind in `dir_name`.
 function write_traj_file(
         dir_name::String,
         file_name::String,
-        data::AbstractVector{UInt8},
-    )::Nothing
-    new_hash = sha256(data)
-    file_path = joinpath(dir_name, file_name)
-    if isfile(file_path)
-        if filesize(file_path) == length(data)
-            existing_hash = open(sha256, file_path)
-            if new_hash == existing_hash
-                # file exists and is correct, return
-                return
+        data::AbstractVector{UInt8};
+        profiler= NullProfiler(),
+    )::String
+    @zone profiler name="write_traj_file" begin
+        new_hash = sha256(data)
+        file_path = joinpath(dir_name, file_name)
+        zone_active(profiler) && zone_text!(profiler, repr(file_path))
+        zone_active(profiler) && zone_text!(profiler, "$(round(length(data)*1E-3; sigdigits=4)) kB")
+        if isfile(file_path)
+            if filesize(file_path) == length(data)
+                existing_hash = open(sha256, file_path)
+                if new_hash == existing_hash
+                    zone_text!(profiler, "file exists and is correct")
+                    return bytes2hex(new_hash)
+                end
             end
+            zone_text!(profiler, "file exists but does not match")
         end
-    end
-    # safely create the new file.
-    mktemp(dir_name) do temp_path, temp_out
-        nb = write(temp_out, data)
-        if nb != length(data)
-            error("short write of $(repr(file_name)) data")
-        end
-        close(temp_out)
-        # TODO Drop support for 1.11
-        @static if VERSION > v"1.12.0-rc"
+        # safely create the new file.
+        mktemp(dir_name) do temp_path, temp_out
+            nb = write(temp_out, data)
+            if nb != length(data)
+                error("short write of $(repr(file_name)) data")
+            end
+            close(temp_out)
             try
                 Base.rename(temp_path, file_path)
             catch err
+                zone_text!(profiler, "rename error")
+                zone_color!(profiler, :red)
                 err isa Base.IOError || rethrow()
                 # on error, check if file was made by another process, and is still valid.
                 if isfile(file_path)
                     if filesize(file_path) == length(data)
                         existing_hash = open(sha256, file_path)
                         if new_hash == existing_hash
-                            # file exists and is correct, return
-                            return
-                        end
-                    end
-                end
-                # otherwise error
-                error("$(repr(file_path)) is corrupted")
-            end
-        else
-            err = ccall(:jl_fs_rename, Int32, (Cstring, Cstring), temp_path, file_path)
-            # on error, check if file was made by another process, and is still valid.
-            if err < 0
-                if isfile(file_path)
-                    if filesize(file_path) == length(data)
-                        existing_hash = open(sha256, file_path)
-                        if new_hash == existing_hash
-                            # file exists and is correct, return
-                            return
+                            zone_text!(profiler, "file exists and is correct")
+                            return bytes2hex(new_hash)
                         end
                     end
                 end
                 # otherwise error
                 error("$(repr(file_path)) is corrupted")
             end
+            return bytes2hex(new_hash)
         end
-        nothing
     end
 end
 

src/outputdiff.jl

@@ -1,9 +1,5 @@
 # functions to show the difference between two job outputs
 
-import DeepDiffs
-using SmallZarrGroups
-import JSON3
-
 
 """
     print_json_diff(io::IO, json1::AbstractString, json2::AbstractString)