Автоматическое сохранение переменных рабочего пространства для `DrWatson.jl`

Пакет Юлии DrWatson.jl направлен на научные вычисления и анализ данных. Он описывает себя как «программное обеспечение для помощников по научным проектам». Недавно я решил попробовать, и мне пока нравится функциональность.

К полезным функциям, на мой взгляд, относятся: safesave метод, который сохраняет словари в JLD файлы, убедившись, что данные не потеряны, а @tag! макрос, который добавляет пары ключ-значение в словарь, которые включают, например, текущий хэш фиксации git, а также предупреждение вместе с «gitpatch», показывающим различия, в случае грязного репозитория. Он также предоставляет инструменты для сбора результатов всех этих файлов вместе в DataFrame объекты.

Я использую сценарии с множеством настраиваемых параметров, которые необходимо запускать на кластерных машинах в течение нескольких дней для вывода результатов. Вскоре после запуска данной симуляции я хочу быть как можно более уверенным в том, что (при условии, что симуляция работает нормально) моя программа не выйдет из строя в конце при попытке сохранить данные, и что легко отслеживать, где находятся файлы результатов пошел, глядя на консольный вывод запуска.

Мой текущий рабочий процесс склоняется к использованию DrWatson просто небрежно сбросить все входные параметры и все результаты моих симуляций в .jld файлы как Dict объекты. Может быть лучше использовать какой-нибудь формат на основе JSON для сохранения всех параметров моделирования. Однако у меня есть много разных симуляций, которые в основном имеют непересекающиеся наборы входных параметров и результатов. В целом, время, потраченное на разработку какой-то системы для сохранения всех метаданных, кажется серьезным отвлечением от моей реальной работы (это то, что она сказала, когда писала сообщение CodeReview об инструментах для упрощения и автоматизации рабочего процесса …).

Проблема

Пока DrWatson кажется полезным, но, приняв его в свою кодовую базу, я закончил тем, что создал безумное количество шаблонного кода. Репозиторий, содержащий сценарии моделирования, не содержит кода моделирования и ссылается на другие репозитории (пакеты Julia). Таким образом, мой проект «моделирования» состоит из множества полностью независимых и непересекающихся скриптов, каждый из которых выглядел следующим образом (обратите внимание на шаблон):

using DrWatson
@quickactivate "PROJECTNAME"

println(repeat("-", 80))
println(repeat("-", 10), " Script $(@__FILE__)")
println(repeat("-", 80))

using SparseArrays
using JLD

using Pkg
@info("Status of custom simulation packages used:")
Pkg.status("<SOMESTUFF>")
# ...

function main()
    p = 32  # blahblahblah
    q = 4  # ...
    # many more parameters...
    # some of these parameters get organized into `struct`'s
    
    simul_params = DrWatson.@strdict p q # a LOT more stuff.
          # Multiple lines of listed parameters, easy to forget something...
    @tag!(simul_params) # add git commit information, etc.
    output_file_name = DrWatson.savename(<important parameters>, ".jld")
    println("Output file complete results: '$output_file_name'")

    # ... simulations can take days to complete
    # some of the simulations actually internally save intermediate results
    # to paths that are part of the parameters. 
    result1 = external_library.some_simulation(parameters)
    result2 = external_library.other_simulation(result1, other_parameters)
    
    simul_params["result1"] = result1
    simul_params["result2"] = result2
    # ...
    # you get the idea

    output_file_path = DrWatson.datadir("sims", "script_name", filename), simul_params)
    DrWatson.safesave(output_file_path, simul_params)
end


main()

Предложенное решение

То, что я в итоге сделал, немного похоже на этот пост с обзором кода (на самом деле я не совсем понимаю, что там делается). Однако, используя DrWatson Я получаю гораздо больше метаданных, и мой подход кажется гораздо более радикальным, поскольку я экономлю все доступные переменные. Возможно, такой подход, как связанный пост, на самом деле объективно лучше, я не уверен, насколько опасны / недостатки экономии абсолютно все вернется, чтобы укусить меня.

В любом случае приведенный ниже код определяет макросы, которые либо возвращаются как Dict (или сохраните прямо в JLD файл) все глобальные и локальные переменные (кроме функций, модулей, переменных со специальными именами и т. д.). Это делается с помощью Base.@local чтобы получить локальные переменные в области действия вызывающего и getfield для получения глобальных переменных вызывающего модуля. Использование getfield вдохновлен источником JLD2.jl.

"""
This file is to be `include`'d into all simulation scripts.
It makes sure that the DrWatson package and its required imports are available.
"""

using DrWatson
@quickactivate "LDPC4QKD"

using SparseArrays
using JLD

using Pkg
@info("Status of RateAdaptiveLDPC used:")
Pkg.status("RateAdaptiveLDPC")



"""
Provides macros for obtaining `Dict`'s of selected local (potentially also global) variables.
"""
module ScopeCrawler

using DrWatson
using SparseArrays
using JLD

export @all_vars_dict
export @safesave_vars

"""
    macro safesave_vars(out_path="", save_globals=true)

Saves the outputs of `@all_vars_dict` to a `JLD` file as a `Dict{String, Any}`.
"""
macro safesave_vars(out_path="", save_globals=true)
    return quote
        output_file_path = $(esc(out_path))

        # auto-generate path if none given.
        if output_file_path == ""
            output_file_path = datadir("autoname", "$(basename(string(__source__.file)))_line$(__source__.line)_$(__module__).jld")
            @warn("@safesave_locals received empty destination path from $(__source__.file). Auto-generated path:
                    '$output_file_path'")
        end

        # make sure the file path has '.jld' extension.
        if length(output_file_path) < 4 || output_file_path[end - 3:end] != ".jld"
            @warn("Appending extension '.jld' to output file path '$output_file_path'.")
            output_file_path *= ".jld"
        end

        results = @all_vars_dict($(esc(save_globals)))
        DrWatson.@tag!(results)

        if isfile(output_file_path)
            @warn("The requested output path $output_file_path already exists. 
                Final path is chosen by `DrWatson.safesave` and will (probably) be
                $(abspath(DrWatson.recursively_clear_path(output_file_path)))")
        else
            @info("Final output file destination: $(abspath(output_file_path))")
        end
        DrWatson.safesave(
            output_file_path, 
            results
        )
    end
end


"""
    macro all_vars_dict(save_globals=true)

Returns all local variables (with reasonable names, no functions and no modules)
in the caller scope and global variables of the module in a `Dict{String, Any}`
(with some additional meta-information). 


Note: this is slow and not optimized. Meant to be called rarely!
TODO: test properly
TODO specify additional information
"""
macro all_vars_dict(save_globals=true)
    # Note: output_file_path should evaluate to a String.
    sourcefile = "unknown"
    sourceline = "unknown"
    try 
        sourcefile = basename(string(__source__.file))
        sourceline = __source__.line
    catch e
        @warn("macro received invalid __source__: $e")
    end

    # could add additional metadata that the macro has access to
    additional_info = Dict{String,Any}()
    try
        additional_info["#source#"] = repr(__source__)
        additional_info["#module#"] = repr(__module__)
        additional_info["#stacktrace#"] = stacktrace()
    catch e
        @warn("Failed to privide additional info in safesave_locals. Reason: $e")
    end

    # helper functions
    symbol_dict_to_strdict(dict::Dict{Symbol,T} where T) = Dict(string(key) => val for (key, val) in dict)

    function combine_conflicting_keys(x, y) # when globals and conflict, resolve
        @warn("Name conflict found between global/local variables while saving: $x vs $y.")
        return Dict("local" => x, "global" => y)
    end

    # Generate expression.
    # Basic idea: 
    # Evaluate `Base.@locals` and function `get_globals_dict` in the scope of the caller.
    return quote
        globals_dict = Dict{String, Any}()
        if $(esc(save_globals))
            globals_dict = $get_globals_dict(@__MODULE__)
        end

        merge($combine_conflicting_keys,
            # `Any` makes sure the right method is for `merge` is called:
            Dict{String, Any}(),
            $symbol_dict_to_strdict(Base.@locals),  
            $additional_info,
            globals_dict
        )
    end
end


"""
    function get_globals_dict(m::Module)

Returns a `Dict{String, Any}` containing global fields of a module that are likely to be user-defined variables.
Ignores modules, functions and non-standard names.

Note: inspired from JLD2.jl, 
see https://github.com/JuliaIO/JLD2.jl/blob/f7535ab19cb65e703ce4f1f2209ccd32e5b9287f/src/loadsave.jl#L78
"""
function get_globals_dict(m::Module)
    results = Dict{String,Any}()

    try
        for vname in names(m; all=true)
            s = string(vname)
            if (!occursin(r"^_+[0-9]*$", s) # skip IJulia history vars
                && !occursin(r"[@#$%^&*()?|\/,.<>]", s)) # skip variables with special characters
                v = getfield(m, vname)
                if !isa(v, Module) && !isa(v, Function)
                    try
                        results[s] = v
                    catch e
                        if isa(e, PointerException)
                            @warn("Saving globals, skipping $vname because it contains a pointer.")
                        else
                            @warn("Saving globals, skipping $vname because $e")
                        end
                    end
                end
            end
        end
    catch e
        @warn("Saving globals failed because: $e")
    end
    return results
end

end # module ScopeCrawler



"""
poor man's unit tests for the above
Note: running this test at the start of each job using the functionality makes sense 
because if this stuff crashes at the end of a simulation, all the work is lost.
"""
module TestingScopeCrawler
using Test
using ..ScopeCrawler

using JLD

tmptmp_global_var_for_test = "tmptmp_global_var_for_test"

@testset "testing @all_vars_dict" begin
    local_var = 3
    othervar = "asdf"

    @testset "no args" begin
        allvarsdict = @all_vars_dict
        @test allvarsdict["local_var"] == 3
        @test allvarsdict["othervar"] == "asdf"
        @test allvarsdict["tmptmp_global_var_for_test"] == "tmptmp_global_var_for_test"
    end
    @testset "arg: true" begin
        allvarsdict = @all_vars_dict true
        @test allvarsdict["local_var"] == 3
        @test allvarsdict["othervar"] == "asdf"
        @test allvarsdict["tmptmp_global_var_for_test"] == "tmptmp_global_var_for_test"
    end
    @testset "arg: false" begin
        allvarsdict = @all_vars_dict false
        @test allvarsdict["local_var"] == 3
        @test allvarsdict["othervar"] == "asdf"
        @test_throws(KeyError, allvarsdict["tmptmp_global_var_for_test"] == "tmptmp_global_var_for_test")
    end
    @testset "arg: variable name" begin
        variable_containing_false = false
        allvarsdict = @all_vars_dict variable_containing_false
        @test allvarsdict["local_var"] == 3
        @test allvarsdict["othervar"] == "asdf"
        @test_throws(KeyError, allvarsdict["tmptmp_global_var_for_test"] == "tmptmp_global_var_for_test")
    end
end


@testset "testing @safesave_locals" begin
    try
        local_var = 3
        othervar = "asdf"

        dest_path = "tmptmp_test_safesave_locals.jld"
        @safesave_vars dest_path
        @test isfile("tmptmp_test_safesave_locals.jld")

        loadagain = load("tmptmp_test_safesave_locals.jld")
        @test loadagain["local_var"] == 3
        @test loadagain["othervar"] == "asdf"
        @test loadagain["tmptmp_global_var_for_test"] == "tmptmp_global_var_for_test"
    finally
        rm("tmptmp_test_safesave_locals.jld")  # throws exception `IOError` if the file does not exist.
    end
end

end  # module TestingScopeCrawler

Как использовать:

С моим стандартным скриптом приведенный выше код упрощается до

include("dr_watson_boilerplate.jl")

function main()
    p = 32  # blahblahblah
    q = 4  # ...
    # many more parameters...
    # some of these parameters get organized into `struct`'s
    
    output_file_name = DrWatson.savename(<important parameters>, ".jld")
    println("Output file complete results: '$output_file_name'")

    # ... simulations can take days to complete
    # some of the simulations actually internally save intermediate results
    # to paths that are part of the parameters. 
    result1 = external_library.some_simulation(parameters)
    result2 = external_library.other_simulation(result1, other_parameters)

    output_file_path = DrWatson.datadir("sims", "script_name", filename), simul_params)
    ScopeCrawler.@safesave_vars output_file_path
end


main()

Что вы думаете?