From 18e60c7b28d20a977f79311b3dea712935c9015c Mon Sep 17 00:00:00 2001
From: Jose Esparza <28990958+pebeto@users.noreply.github.com>
Date: Mon, 17 Jul 2023 02:11:51 -0500
Subject: [PATCH] uploading tests and including mlflow on CI

---
 .github/workflows/ci.yml   |  5 ++++
 Project.toml               |  2 +-
 ext/LoggersExt/mlflow.jl   |  5 ++--
 src/MLJBase.jl             |  9 +++---
 test/extensions/loggers.jl | 60 ++++++++++++++++++++++++++++++++++++++
 test/runtests.jl           |  4 +++
 6 files changed, 78 insertions(+), 7 deletions(-)
 create mode 100644 test/extensions/loggers.jl

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ac1e885e..568d8e48 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -15,6 +15,11 @@ env:
   TEST_MLJBASE: "true"
 jobs:
   test:
+    services:
+      mlflow:
+        image: adacotechjp/mlflow:2.3.1
+        ports:
+          - 5000:5000
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 60
diff --git a/Project.toml b/Project.toml
index a5829d3c..792a7ca5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -68,4 +68,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 
 [targets]
-test = ["DataFrames", "DecisionTree", "Distances", "Logging", "MultivariateStats", "NearestNeighbors", "StableRNGs", "Test", "TypedTables"]
+test = ["DataFrames", "DecisionTree", "Distances", "Logging", "MultivariateStats", "NearestNeighbors", "StableRNGs", "Test", "TypedTables", "MLFlowClient"]
diff --git a/ext/LoggersExt/mlflow.jl b/ext/LoggersExt/mlflow.jl
index 69e1c158..a793c982 100644
--- a/ext/LoggersExt/mlflow.jl
+++ b/ext/LoggersExt/mlflow.jl
@@ -18,8 +18,8 @@ function _logmodelparams(client::MLFlow, run::MLFlowRun, model::Model)
     end
 end
 
-function _logmachinemeasures(client::MLFlow, run::MLFlowRun, measures::Vector{Measure},
-                             measurements::Vector{Float64})
+function _logmachinemeasures(client::MLFlow, run::MLFlowRun, measures::Vector{T},
+                             measurements::Vector{Float64}) where T<:Measure
     measure_names = measures .|> info .|> x -> x.name
     for (name, value) in zip(measure_names, measurements)
         logmetric(client, run, name, value)
@@ -60,4 +60,5 @@ function save(logger::MLFlowLogger, mach::Machine)
     _logmodelparams(logger.client, run, mach.model)
     fname = "$(model_name).jls"
     logartifact(logger.client, run, fname, io)
+    updaterun(logger.client, run, "FINISHED")
 end
diff --git a/src/MLJBase.jl b/src/MLJBase.jl
index d0fbf78e..ae895a19 100644
--- a/src/MLJBase.jl
+++ b/src/MLJBase.jl
@@ -294,11 +294,15 @@ export coerce, coerce!, autotype, schema, info
 export UnivariateFiniteArray, UnivariateFiniteVector
 
 # -----------------------------------------------------------------------
-# abstract model types defined in MLJModelInterface.jl and extended here:
+# re-export from MLJModelInterface.jl
+
+#abstract model types defined in MLJModelInterface.jl and extended here:
 for T in EXTENDED_ABSTRACT_MODEL_TYPES
     @eval(export $T)
 end
 
+export params
+
 # -------------------------------------------------------------------
 # exports from this module, MLJBase
 
@@ -308,9 +312,6 @@ export default_resource
 # one_dimensional_ranges.jl:
 export ParamRange, NumericRange, NominalRange, iterator, scale
 
-# parameter_inspection.jl:
-export params # note this is *not* an extension of StatsBase.params
-
 # data.jl:
 export partition, unpack, complement, restrict, corestrict
 
diff --git a/test/extensions/loggers.jl b/test/extensions/loggers.jl
new file mode 100644
index 00000000..64d2830d
--- /dev/null
+++ b/test/extensions/loggers.jl
@@ -0,0 +1,60 @@
+module TestLoggers
+
+using Test
+using MLJBase
+using ..Models
+
+@testset "mlflow logger" begin
+    artifact_directory = "mlj-test"
+    experiment_name = "mlflow logger tests"
+
+    @testset "outside extension tests" begin
+        @test_throws ErrorException mlflow_logger()
+
+        using MLFlowClient
+        logger = mlflow_logger(; experiment_name=experiment_name, artifact_location=artifact_directory)
+
+        @test logger.client isa MLFlow
+        @test logger.experiment_name == experiment_name
+        @test logger.artifact_location == artifact_directory
+    end # @testset "outside extension tests"
+
+    @testset "extension tests" begin
+        X = (x=rand(4),)
+        y = ["Chenta", "Missy", "Gala", "Wendy"] |> categorical
+
+        mach = machine(ConstantClassifier(), X, y)
+        fit!(mach, verbosity=0)
+
+        logger = mlflow_logger(; experiment_name=experiment_name, artifact_location=artifact_directory)
+
+        @testset "save" begin
+            run = MLJBase.save(logger, mach)
+            experiment = getexperiment(logger.client, run.info.experiment_id)
+            @test run isa MLFlowRun
+            @test experiment isa MLFlowExperiment
+
+            deleterun(logger.client, run)
+            deleteexperiment(logger.client, experiment)
+        end # @testset "save"
+
+        @testset "evaluate!" begin
+            evaluate!(mach, resampling=Holdout(), logger=logger)
+
+            experiments = searchexperiments(logger.client)
+            experiments_ids = experiments .|> (e -> e.experiment_id)
+            runs = searchruns(logger.client, experiments_ids)
+
+            # it's 2 because of the default experiment
+            @test length(experiments_ids) == 2
+            @test length(runs) == 1
+
+            deleterun(logger.client, runs[1])
+            deleteexperiment(logger.client, experiments[2])
+        end # @testset "evaluate!"
+    end # @testset "extension tests"
+end # @testset "mlflow logger"
+
+end # module
+
+true
diff --git a/test/runtests.jl b/test/runtests.jl
index 8b07929e..a55c5d19 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -80,3 +80,7 @@ end
     @test include("hyperparam/one_dimensional_ranges.jl")
     @test include("hyperparam/one_dimensional_range_methods.jl")
 end
+
+@conditional_testset "extensions" begin
+    @test include("extensions/loggers.jl")
+end