Finishing with logbatch functionality

Author: pebeto

src/api.jl

@@ -9,7 +9,7 @@ MLFlowClient.uri(mlf, "experiments/get", Dict(:experiment_id=>10))
 ```
 """
 uri(mlf::MLFlow, endpoint::String;
-    parameters::Dict{Symbol, <:Any}=Dict{Symbol, IntOrString}()) =
+    parameters::Dict{Symbol, <:Any}=Dict{Symbol, NumberOrString}()) =
     URI("$(mlf.apiroot)/$(mlf.apiversion)/mlflow/$(endpoint)";
         query=parameters)
 

src/services/experiment.jl

@@ -23,7 +23,8 @@ function createexperiment(instance::MLFlow, name::String;
     tags::MLFlowUpsertData{Tag}=Tag[])::String
     try
         result = mlfpost(instance, "experiments/create"; name=name,
-            artifact_location=artifact_location, tags=(tags |> parse))
+            artifact_location=artifact_location,
+            tags=parse(Tag, tags))
         return result["experiment_id"]
     catch e
         if isa(e, HTTP.ExceptionRequest.StatusError) && e.status == 400

src/services/loggers.jl

@@ -28,7 +28,9 @@ logmetric(instance::MLFlow, run::Run, key::String, value::Float64;
     logmetric(instance, run.info.run_id, key, value; timestamp=timestamp, step=step)
 
 """
-    logbatch(instance::MLFlow, run_id::String, metrics::Array{Metric},
+    logbatch(instance::MLFlow, run_id::String; metrics::Array{Metric},
+        params::Array{Param}, tags::Array{Tag})
+    logbatch(instance::MLFlow, run::Run; metrics::Array{Metric},
         params::Array{Param}, tags::Array{Tag})
 
 Log a batch of metrics, params, and tags for a run. In case of error, partial
@@ -37,8 +39,15 @@ data may be written.
 For more information about this function, check [MLFlow official documentation](https://mlflow.org/docs/latest/rest-api.html#log-batch).
 """
 function logbatch(instance::MLFlow, run_id::String;
-    metrics::Array{Metric}=Metric[], params::Array{Param}=Param[],
-        tags::MLFlowUpsertData{Tag}=Tag[])
-    mlfpost(instance, "runs/log-batch"; run_id=run_id, metrics=metrics,
-            params=params, tags=(tags |> parse))
+    metrics::MLFlowUpsertData{Metric}=Metric[],
+    params::MLFlowUpsertData{Param}=Param[], tags::MLFlowUpsertData{Tag}=Tag[])
+    mlfpost(instance, "runs/log-batch"; run_id=run_id,
+        metrics=parse(Metric, metrics), params=parse(Param, params),
+        tags=parse(Tag, tags))
 end
+logbatch(instance::MLFlow, run::Run;
+    metrics::MLFlowUpsertData{Metric}=Metric[],
+    params::MLFlowUpsertData{Param}=Param[],
+    tags::MLFlowUpsertData{Tag}=Tag[]) =
+    logbatch(instance, run.info.run_id; metrics=metrics, params=params,
+        tags=tags)

src/services/run.jl

@@ -23,7 +23,7 @@ function createrun(instance::MLFlow, experiment_id::String;
     tags::MLFlowUpsertData{Tag}=Tag[])
     try
         result = mlfpost(instance, "runs/create"; experiment_id=experiment_id,
-            run_name=run_name, start_time=start_time, tags=(tags |> parse))
+            run_name=run_name, start_time=start_time, tags=parse(Tag, tags))
         return result["run"] |> Run
     catch e
         throw(e)

src/types/run.jl

@@ -1,26 +1,26 @@
 """
-    Metric
+    Metric <: LoggingData
 
 Metric associated with a run, represented as a key-value pair.
 
 # Fields
 - `key::String`: Key identifying this metric.
 - `value::Float64`: Value associated with this metric.
 - `timestamp::Int64`: The timestamp at which this metric was recorded.
-- `step::Int64`: Step at which to log the metric.
+- `step::Union{Int64, Missing}`: Step at which to log the metric.
 """
 struct Metric <: LoggingData
     key::String
     value::Float64
     timestamp::Int64
-    step::Int64
+    step::Union{Int64, Missing}
 end
 Metric(data::Dict{String, Any}) = Metric(data["key"], data["value"],
     data["timestamp"], data["step"])
 Base.show(io::IO, t::Metric) = show(io, ShowCase(t, new_lines=true))
 
 """
-    Param
+    Param <: LoggingData
 
 Param associated with a run.
 

src/types/tag.jl

@@ -1,5 +1,5 @@
 """
-    Tag
+    Tag <: LoggingData
 
 Generic tag type for MLFlow entities.
 

src/utils.jl

@@ -1,50 +1,101 @@
-const IntOrString = Union{Int, String}
-const MLFlowUpsertData{T} = Union{Array{T}, Dict{String, String},
-    Array{Pair{String, String}}, Array{Dict{String, String}}}
+const NumberOrString = Union{Number, String}
+const MLFlowUpsertData{T} = Union{
+    Array{T},
+    Array{<:Dict{String, <:Any}},
+    Dict{String, <:NumberOrString},
+    Array{<:Pair{String, <:NumberOrString}},
+    Array{<:Tuple{String, <:NumberOrString}}
+}
 
 const MLFLOW_ERROR_CODES = (;
     RESOURCE_ALREADY_EXISTS = "RESOURCE_ALREADY_EXISTS",
     RESOURCE_DOES_NOT_EXIST = "RESOURCE_DOES_NOT_EXIST",
 )
 
-function dict_to_array(dict::Dict{String, String})::MLFlowUpsertData
-    tags = Tag[]
+function dict_to_T_array(::Type{T}, dict::Dict{String, <:NumberOrString}) where T<:LoggingData
+    entities = T[]
     for (key, value) in dict
-        push!(tags, Tag(key, value))
+        if T<:Metric
+            push!(entities, Metric(key, Float64(value),
+                round(Int, now() |> datetime2unix), missing))
+        else
+            push!(entities, T(key, value |> string))
+        end
     end
 
-    return tags
+    return entities
 end
 
-function pairsarray_to_array(pair_array::Array{<:Pair})::MLFlowUpsertData
-    entity_array = Tag[]
+function pairarray_to_T_array(::Type{T}, pair_array::Array{<:Pair}) where T<:LoggingData
+    entities = T[]
     for pair in pair_array
-        println(pair)
         key = pair.first |> string
-        value = pair.second |> string
-        push!(entity_array, Tag(key, value))
+        if T<:Metric
+            value = pair.second
+            push!(entities, Metric(key, Float64(value),
+                round(Int, now() |> datetime2unix), missing))
+        else
+            value = pair.second |> string
+            push!(entities, T(key, value))
+        end
     end
 
-    return entity_array
+    return entities
 end
 
-function dictarray_to_array(dict_array::Array{Dict{String, String}})::MLFlowUpsertData
-    tags = Tag[]
+function tuplearray_to_T_array(::Type{T},
+    tuple_array::Array{<:Tuple{String, <:NumberOrString}}) where T<:LoggingData
+    entities = T[]
+    for tuple in tuple_array
+        if length(tuple) != 2
+            error("Tuple must have exactly two elements (format: (key, value))")
+        end
+
+        key = tuple |> first |> string
+        if T<: Metric
+            value = tuple |> last
+            push!(entities, Metric(key, Float64(value),
+                round(Int, now() |> datetime2unix), missing))
+        else
+            value = tuple |> last |> string
+            push!(entities, T(key, value))
+        end
+    end
+
+    return entities
+end
+
+function dictarray_to_T_array(::Type{T},
+    dict_array::Array{<:Dict{String, <:Any}}) where T<:LoggingData
+    entities = T[]
     for dict in dict_array
-        push!(tags, Tag(dict["key"], dict["value"]))
+        key = dict["key"] |> string
+        if T<:Metric
+            value = Float64(dict["value"])
+            if haskey(dict, "timestamp")
+                timestamp = dict["timestamp"]
+            else
+                timestamp = round(Int, now() |> datetime2unix)
+            end
+            push!(entities, Metric(key, value, timestamp, missing))
+        else
+            value = dict["value"] |> string
+            push!(entities, T(key, value))
+        end
     end
 
-    return tags
+    return entities
 end
 
-function parse(entities::MLFlowUpsertData{T}) where T<:LoggingData
-    println(typeof(entities))
-    if entities isa Dict{String, String}
-        return entities |> dict_to_array
-    elseif entities isa Array{Pair{String, String}}
-        return entities |> pairsarray_to_array
-    elseif entities isa Array{Dict{String, String}}
-        return entities |> dictarray_to_array
+function parse(::Type{T}, entities::MLFlowUpsertData{T}) where T<:LoggingData
+    if entities isa Dict{String, <:NumberOrString}
+        return dict_to_T_array(T, entities)
+    elseif entities isa Array{<:Dict{String, <:Any}}
+        return dictarray_to_T_array(T, entities)
+    elseif entities isa Array{<:Pair{String, <:NumberOrString}}
+        return pairarray_to_T_array(T, entities)
+    elseif entities isa Array{<:Tuple{String, <:NumberOrString}}
+        return tuplearray_to_T_array(T, entities)
     end
     return entities
 end

test/runtests.jl

@@ -5,5 +5,5 @@ end
 include("base.jl")
 
 include("services/experiment.jl")
-# include("services/run.jl")
-# include("services/loggers.jl")
+include("services/run.jl")
+include("services/loggers.jl")

test/services/experiment.jl

@@ -14,29 +14,29 @@
         deleteexperiment(mlf, experiment.experiment_id)
     end
 
-    # @testset "with tags as array of tags" begin
-    #     experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
-    #         tags=[Tag("test_key", "test_value")])
-    #     deleteexperiment(mlf, experiment_id)
-    # end
+    @testset "with tags as array of tags" begin
+        experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
+            tags=[Tag("test_key", "test_value")])
+        deleteexperiment(mlf, experiment_id)
+    end
 
     @testset "with tags as array of pairs" begin
         experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
             tags=["test_key" => "test_value"])
         deleteexperiment(mlf, experiment_id)
     end
 
-    # @testset "with tags as array of dicts" begin
-    #     experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
-    #         tags=[Dict("key" => "test_key", "value" => "test_value")])
-    #     deleteexperiment(mlf, experiment_id)
-    # end
+    @testset "with tags as array of dicts" begin
+        experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
+            tags=[Dict("key" => "test_key", "value" => "test_value")])
+        deleteexperiment(mlf, experiment_id)
+    end
 
-    # @testset "with tags as dict" begin
-    #     experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
-    #         tags=Dict("test_key" => "test_value"))
-    #     deleteexperiment(mlf, experiment_id)
-    # end
+    @testset "with tags as dict" begin
+        experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string;
+            tags=Dict("test_key" => "test_value"))
+        deleteexperiment(mlf, experiment_id)
+    end
 end
 
 @testset verbose = true "get experiment" begin