Adding loginputs function

Author: pebeto

src/MLFlowClient.jl

@@ -71,6 +71,7 @@ export
 include("services/loggers.jl")
 export
     logmetric,
-    logbatch
+    logbatch,
+    loginputs
 
 end

src/services/loggers.jl

@@ -28,26 +28,41 @@ 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},
-        params::Array{Param}, tags::Array{Tag})
+    logbatch(instance::MLFlow, run_id::String;
+        metrics::MLFlowUpsertData{Metric}, params::MLFlowUpsertData{Param},
+        tags::MLFlowUpsertData{Tag})
     logbatch(instance::MLFlow, run::Run; metrics::Array{Metric},
-        params::Array{Param}, tags::Array{Tag})
+        params::MLFlowUpsertData{Param}, tags::MLFlowUpsertData{Tag})
 
 Log a batch of metrics, params, and tags for a run. In case of error, partial
 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;
+logbatch(instance::MLFlow, run_id::String;
     metrics::MLFlowUpsertData{Metric}=Metric[],
-    params::MLFlowUpsertData{Param}=Param[], tags::MLFlowUpsertData{Tag}=Tag[])
+    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)
+
+"""
+    loginputs(instance::MLFlow, run_id::String; datasets::Array{DatasetInput})
+    loginputs(instance::MLFlow, run::Run; datasets::Array{DatasetInput})
+
+# Arguments
+- `instance`: [`MLFlow`](@ref) configuration.
+- `run_id`: ID of the run to log under This field is required.
+- `datasets`: Dataset inputs.
+"""
+loginputs(instance::MLFlow, run_id::String, datasets::Array{DatasetInput}) =
+    mlfpost(instance, "runs/log-inputs"; run_id=run_id, datasets=datasets)
+loginputs(instance::MLFlow, run::Run, datasets::Array{DatasetInput}) =
+    loginputs(instance, run.info.run_id, datasets)

src/types/run.jl

@@ -7,13 +7,13 @@ Metric associated with a run, represented as a key-value pair.
 - `key::String`: Key identifying this metric.
 - `value::Float64`: Value associated with this metric.
 - `timestamp::Int64`: The timestamp at which this metric was recorded.
-- `step::Union{Int64, Missing}`: Step at which to log the metric.
+- `step::Union{Int64, Nothing}`: Step at which to log the metric.
 """
 struct Metric <: LoggingData
     key::String
     value::Float64
     timestamp::Int64
-    step::Union{Int64, Missing}
+    step::Union{Int64, Nothing}
 end
 Metric(data::Dict{String, Any}) = Metric(data["key"], data["value"],
     data["timestamp"], data["step"])

src/utils.jl

@@ -17,7 +17,7 @@ function dict_to_T_array(::Type{T}, dict::Dict{String, <:NumberOrString}) where
     for (key, value) in dict
         if T<:Metric
             push!(entities, Metric(key, Float64(value),
-                round(Int, now() |> datetime2unix), missing))
+                round(Int, now() |> datetime2unix), nothing))
         else
             push!(entities, T(key, value |> string))
         end
@@ -33,7 +33,7 @@ function pairarray_to_T_array(::Type{T}, pair_array::Array{<:Pair}) where T<:Log
         if T<:Metric
             value = pair.second
             push!(entities, Metric(key, Float64(value),
-                round(Int, now() |> datetime2unix), missing))
+                round(Int, now() |> datetime2unix), nothing))
         else
             value = pair.second |> string
             push!(entities, T(key, value))
@@ -55,7 +55,7 @@ function tuplearray_to_T_array(::Type{T},
         if T<: Metric
             value = tuple |> last
             push!(entities, Metric(key, Float64(value),
-                round(Int, now() |> datetime2unix), missing))
+                round(Int, now() |> datetime2unix), nothing))
         else
             value = tuple |> last |> string
             push!(entities, T(key, value))
@@ -77,7 +77,7 @@ function dictarray_to_T_array(::Type{T},
             else
                 timestamp = round(Int, now() |> datetime2unix)
             end
-            push!(entities, Metric(key, value, timestamp, missing))
+            push!(entities, Metric(key, value, timestamp, nothing))
         else
             value = dict["value"] |> string
             push!(entities, T(key, value))

test/services/loggers.jl

@@ -36,7 +36,7 @@ end
 
     @testset "with run id as string" begin
         run = createrun(mlf, experiment_id)
-        logbatch(mlf, run.info.run_id, metrics=[("gala", 0.1)])
+        logbatch(mlf, run.info.run_id; metrics=[("gala", 0.1)])
 
         run = refresh(mlf, run)
         last_metric = run.data.metrics |> last
@@ -49,7 +49,7 @@ end
 
     @testset "with run" begin
         run = createrun(mlf, experiment_id)
-        logbatch(mlf, run, metrics=[("missy", 0.9)])
+        logbatch(mlf, run; metrics=[("missy", 0.9)])
 
         run = refresh(mlf, run)
         last_metric = run.data.metrics |> last
@@ -62,7 +62,7 @@ end
 
     @testset "with metrics, params and tags as dict" begin
         run = createrun(mlf, experiment_id)
-        logbatch(mlf, run, metrics=Dict("ana" => 0.5),
+        logbatch(mlf, run; metrics=Dict("ana" => 0.5),
             params=Dict("test_param" => "0.9"),
             tags=Dict("test_tag" => "gala"))
 
@@ -88,7 +88,7 @@ end
 
     @testset "with metrics, params and tags as pair array" begin
         run = createrun(mlf, experiment_id)
-        logbatch(mlf, run, metrics=["ana" => 0.5],
+        logbatch(mlf, run; metrics=["ana" => 0.5],
             params=["test_param" => "0.9"], tags=["test_tag" => "gala"])
 
         run = refresh(mlf, run)
@@ -113,7 +113,7 @@ end
 
     @testset "with metrics, params and tags as tuple array" begin
         run = createrun(mlf, experiment_id)
-        logbatch(mlf, run, metrics=[("ana", 0.5)],
+        logbatch(mlf, run; metrics=[("ana", 0.5)],
             params=[("test_param", "0.9")], tags=[("test_tag", "gala")])
 
         run = refresh(mlf, run)
@@ -138,7 +138,7 @@ end
 
     @testset "with metrics, params and tags as dict array" begin
         run = createrun(mlf, experiment_id)
-        logbatch(mlf, run,
+        logbatch(mlf, run;
             metrics=[Dict("key" => "ana", "value" => 0.5, "timestamp" => 123)],
             params=[Dict("key" => "test_param", "value" => "0.9")],
             tags=[Dict("key" => "test_tag", "value" => "gala")])
@@ -166,3 +166,38 @@ end
 
     deleteexperiment(mlf, experiment_id)
 end
+
+@testset verbose = true "log inputs" begin
+    @ensuremlf
+
+    experiment_id = createexperiment(mlf, UUIDs.uuid4() |> string)
+
+    @testset "with run id as string" begin
+        run = createrun(mlf, experiment_id)
+        inputs = [DatasetInput([Tag("tag_key", "tag_value")],
+            Dataset("dataset_name", "dataset_digest", "dataset_source_type",
+                "dataset_source", nothing, nothing))]
+        loginputs(mlf, run.info.run_id, inputs)
+
+        run = refresh(mlf, run)
+
+        @test run.inputs.dataset_inputs |> length == 1
+
+        dataset_input = run.inputs.dataset_inputs |> first
+        dataset_input_tag = dataset_input.tags |> first
+
+        @test dataset_input_tag isa Tag
+        @test dataset_input_tag.key == "tag_key"
+        @test dataset_input_tag.value == "tag_value"
+
+        @test dataset_input.dataset isa Dataset
+        @test dataset_input.dataset.name == "dataset_name"
+        @test dataset_input.dataset.digest == "dataset_digest"
+        @test dataset_input.dataset.source_type == "dataset_source_type"
+        @test dataset_input.dataset.source == "dataset_source"
+        @test dataset_input.dataset.schema |> isnothing
+        @test dataset_input.dataset.profile |> isnothing
+    end
+
+    deleteexperiment(mlf, experiment_id)
+end