added MLFlowBackend

.github/workflows/ci.yml

@@ -17,12 +17,25 @@ jobs:
           - windows-latest
         arch:
           - x64
+        python-version: ['3.10']
     steps:
       - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
+      - name: Setup python and mlflow server
+        uses: actions/setup-python@v1
+        with:
+          python-version: ${{ matrix.python-version }}
+          architecture: ${{ matrix.arch }}
+      - run: |
+          python -m pip install mlflow
+          python -m pip show mlflow
+          mlflow server --host localhost --port 5000 &
+          sleep 5
       - uses: julia-actions/julia-buildpkg@latest
       - uses: julia-actions/julia-runtest@latest
+        env:
+          MLFLOW_URI: "http://localhost:5000"
 

Project.toml

@@ -13,6 +13,7 @@ Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
 InlineTest = "bd334432-b1e7-49c7-a2dc-dd9149e4ebd6"
+MLFlowLogger = "a17d1b34-d2df-4d9e-9e11-6289e57bd259"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParameterSchedulers = "d7d3b36b-41b8-4d0d-a2bf-768c6151755e"

docs/features.md

@@ -24,7 +24,7 @@ For logging, use the logging callbacks:
 - [`LogHyperParams`](#)
 - [`LogHistograms`](#)
 
-They each can have multiple logging backends, but right now the only one implemented in *FluxTraining.jl* is [`TensorBoardBackend`](#). See also [`LoggerBackend`](#), [`log_to`](#), and [`Loggables.Loggable`](#).
+They each can have multiple logging backends, but right now the only ones implemented in *FluxTraining.jl* are [`TensorBoardBackend`](#) and [`MLFlowBackend`](#). See also [`LoggerBackend`](#), [`log_to`](#), and [`Loggables.Loggable`](#).
 
 There is also an external package [Wandb.jl](https://github.com/avik-pal/Wandb.jl) that implements a logging backend for [Weights&Biases](https://wandb.ai).
 

src/FluxTraining.jl

@@ -47,6 +47,8 @@ include("./callbacks/execution.jl")
 # logging
 include("./callbacks/logging/Loggables.jl")
 include("./callbacks/logging/logger.jl")
+include("./callbacks/logging/combinename.jl")
+include("./callbacks/logging/mlflow.jl")
 include("./callbacks/logging/tensorboard.jl")
 include("./callbacks/logging/checkpointer.jl")
 
@@ -111,6 +113,7 @@ export AbstractCallback,
     LogHyperParams,
     LogVisualization,
     TensorBoardBackend,
+    MLFlowBackend,
     StopOnNaNLoss,
     LearningRate,
     throttle,

src/callbacks/logging/combinename.jl

@@ -0,0 +1,3 @@
+_combinename(name, group::String) = _combinename((group, name))
+_combinename(name, group::Tuple) = _combinename((group..., name))
+_combinename(strings::Tuple) = join(strings, '/')

src/callbacks/logging/mlflow.jl

@@ -0,0 +1,29 @@
+using MLFlowLogger: MLFLogger, log_metric
+
+"""
+    MLFlowBackend(;
+        tracking_uri=nothing, 
+        experiment_name=nothing,
+        run_id=nothing,
+        start_step=0,
+        step_increment=1,
+        min_level=CoreLogging.Info,
+        kwargs...)
+MLFlow backend for logging callbacks. Takes the same arguments
+as [`MLFlowLogger.MLFlowLogger`](https://github.com/rejuvyesh/MLFlowLogger.jl/blob/master/src/MLFlowLogger.jl).
+"""
+struct MLFlowBackend <: LoggerBackend
+    logger::MLFLogger
+
+    function MLFlowBackend(; kwargs...)
+        return new(MLFLogger(; kwargs...))
+    end
+end
+
+Base.show(io::IO, backend::MLFlowBackend) = print(
+    io, "MLFlowBackend(", backend.logger, ")")
+
+function log_to(backend::MLFlowBackend, value::Loggables.Value, name, i; group = ())
+    name = _combinename(name, group)
+    log_metric(backend.logger, name, value.data; step = i)
+end

src/callbacks/logging/tensorboard.jl

@@ -44,9 +44,3 @@ function log_to(backend::TensorBoardBackend, hist::Loggables.Histogram, name, i;
     name = _combinename(name, group)
     log_histogram(backend.logger, name, hist.data, step=i)
 end
-
-# Utilities
-
-_combinename(name, group::String) = _combinename((group, name))
-_combinename(name, group::Tuple) = _combinename((group..., name))
-_combinename(strings::Tuple) = join(strings, '/')

test/callbacks/logging.jl

@@ -1,11 +1,16 @@
 include("../imports.jl")
 
 tbbackend() = TensorBoardBackend(mktempdir())
+mlflowbackend() = MLFlowBackend(tracking_uri=ENV["MLFLOW_URI"])
 
 @testset "`LogMetrics`" begin
     cb = LogMetrics(tbbackend())
     learner = testlearner(Metrics(accuracy), Recorder(), cb)
     @test_nowarn fit!(learner, 1)
+
+    cb = LogMetrics(mlflowbackend())
+    learner = testlearner(Metrics(accuracy), Recorder(), cb)
+    @test_nowarn fit!(learner, 1)
 end