Fix gateway/webhook/run bugs, align types with REST API, increase test coverage

Author: pebeto

src/api.jl

@@ -204,26 +204,4 @@ function mlfdelete_v3(mlf, endpoint; kwargs...)
         @error error_message
         throw(ErrorException(error_message))
     end
-end
-
-"""
-    mlfpatch_v3(mlf, endpoint; kwargs...)
-
-Performs a HTTP PATCH to the specified endpoint using API version 3.0. kwargs are converted to JSON and become
-the PATCH body.
-"""
-function mlfpatch_v3(mlf, endpoint; kwargs...)
-    apiuri = uri_v3(mlf, endpoint;)
-    apiheaders = headers(mlf, Dict("Content-Type" => "application/json"))
-    body = JSON.json(kwargs)
-
-    try
-        response = HTTP.patch(apiuri, apiheaders, body)
-        return response.body |> String |> JSON.parse
-    catch e
-        error_response = e.response.body |> String |> JSON.parse
-        error_message = "$(error_response["error_code"]) -  $(error_response["message"])"
-        @error error_message
-        throw(ErrorException(error_message))
-    end
-end
+end

src/services/experiment.jl

@@ -21,9 +21,21 @@ The ID of the newly created [`Experiment`](@ref).
 function createexperiment(instance::MLFlow, name::String;
     artifact_location::Union{String,Missing}=missing,
     tags::MLFlowUpsertData{Tag}=Tag[])::String
-    result = mlfpost(instance, "experiments/create"; name=name,
-        artifact_location=artifact_location, tags=parse(Tag, tags))
-    return result["experiment_id"]
+    try
+        result = mlfpost(instance, "experiments/create"; name=name,
+            artifact_location=artifact_location, tags=parse(Tag, tags))
+        return result["experiment_id"]
+    catch e
+        # MLflow with basic-auth may return RESOURCE_ALREADY_EXISTS due to a
+        # permission auto-creation race condition, even though the experiment was
+        # successfully created. Recover by fetching the experiment by name.
+        if e isa ErrorException && occursin("Experiment permission", e.msg) &&
+                occursin("already exists", e.msg)
+            experiment = getexperimentbyname(instance, name)
+            return experiment.experiment_id
+        end
+        rethrow()
+    end
 end
 
 """
@@ -163,12 +175,19 @@ updateexperiment(instance::MLFlow, experiment::Experiment, new_name::String)::Bo
 - The next page token if there are more results.
 """
 function searchexperiments(instance::MLFlow; max_results::Int64=20000,
-    page_token::String="", filter::String="", order_by::Array{String}=String[],
+    page_token::Union{String,Missing}=missing, filter::Union{String,Missing}=missing,
+    order_by::Array{String}=String[],
     view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY)::Tuple{Array{Experiment},Union{String,Nothing}}
-    parameters = (; max_results, page_token, filter, :view_type => view_type |> Integer)
+    parameters = Dict{Symbol,Any}(:max_results => max_results, :view_type => view_type |> Integer)
 
-    if order_by |> !isempty
-        parameters = (; order_by, parameters...)
+    if !ismissing(page_token) && !isempty(page_token)
+        parameters[:page_token] = page_token
+    end
+    if !ismissing(filter) && !isempty(filter)
+        parameters[:filter] = filter
+    end
+    if !isempty(order_by)
+        parameters[:order_by] = order_by
     end
 
     result = mlfget(instance, "experiments/search"; parameters...)

src/services/gateway.jl

@@ -191,24 +191,31 @@ end
 
 """
     updategatewaymodeldefinition(instance::MLFlow, model_definition_id::String;
-        model_name::Union{String,Missing}=missing, updated_by::Union{String,Missing}=missing)
+        name::Union{String,Missing}=missing, secret_id::Union{String,Missing}=missing,
+        model_name::Union{String,Missing}=missing, provider::Union{String,Missing}=missing,
+        updated_by::Union{String,Missing}=missing)
 
 Update a model definition.
 
 # Arguments
 - `instance`: [`MLFlow`](@ref) configuration.
 - `model_definition_id`: ID of the model definition to update.
+- `name`: Optional new name.
+- `secret_id`: Optional new secret ID.
 - `model_name`: Optional new model name.
+- `provider`: Optional new provider.
 - `updated_by`: Username of the updater.
 
 # Returns
 An instance of type [`GatewayModelDefinition`](@ref).
 """
 function updategatewaymodeldefinition(instance::MLFlow, model_definition_id::String;
-    model_name::Union{String,Missing}=missing, updated_by::Union{String,Missing}=missing)::GatewayModelDefinition
+    name::Union{String,Missing}=missing, secret_id::Union{String,Missing}=missing,
+    model_name::Union{String,Missing}=missing, provider::Union{String,Missing}=missing,
+    updated_by::Union{String,Missing}=missing)::GatewayModelDefinition
     result = mlfpost_v3(instance, "gateway/model-definitions/update";
-        model_definition_id=model_definition_id, model_name=model_name,
-        updated_by=updated_by)
+        model_definition_id=model_definition_id, name=name, secret_id=secret_id,
+        model_name=model_name, provider=provider, updated_by=updated_by)
     return result["model_definition"] |> GatewayModelDefinition
 end
 
@@ -289,26 +296,47 @@ end
 
 """
     updategatewayendpoint(instance::MLFlow, endpoint_id::String;
-        name::Union{String,Missing}=missing, config=missing,
+        name::Union{String,Missing}=missing,
+        model_configs::Union{Array,Missing}=missing,
+        routing_strategy::Union{String,Missing}=missing,
+        fallback_config::Union{Dict,Missing}=missing,
+        experiment_id::Union{String,Missing}=missing,
+        usage_tracking::Union{Bool,Missing}=missing,
         updated_by::Union{String,Missing}=missing)
 
 Update a gateway endpoint.
 
 # Arguments
 - `instance`: [`MLFlow`](@ref) configuration.
 - `endpoint_id`: ID of the endpoint to update.
-- `name`: Optional new name.
-- `config`: Optional new configuration.
+- `name`: Optional new name for the endpoint.
+- `model_configs`: Optional new list of model configurations (replaces all existing).
+- `routing_strategy`: Optional new routing strategy.
+- `fallback_config`: Optional fallback configuration.
+- `experiment_id`: Optional experiment ID for tracing.
+- `usage_tracking`: Whether to enable usage tracking.
 - `updated_by`: Username of the updater.
 
 # Returns
 An instance of type [`GatewayEndpoint`](@ref).
 """
 function updategatewayendpoint(instance::MLFlow, endpoint_id::String;
-    name::Union{String,Missing}=missing, config=missing,
+    name::Union{String,Missing}=missing,
+    model_configs::Union{Array,Missing}=missing,
+    routing_strategy::Union{String,Missing}=missing,
+    fallback_config::Union{Dict,Missing}=missing,
+    experiment_id::Union{String,Missing}=missing,
+    usage_tracking::Union{Bool,Missing}=missing,
     updated_by::Union{String,Missing}=missing)::GatewayEndpoint
-    result = mlfpost_v3(instance, "gateway/endpoints/update";
-        endpoint_id=endpoint_id, name=name, config=config, updated_by=updated_by)
+    params = Dict{Symbol,Any}(:endpoint_id => endpoint_id)
+    !ismissing(name) && (params[:name] = name)
+    !ismissing(model_configs) && (params[:model_configs] = model_configs)
+    !ismissing(routing_strategy) && (params[:routing_strategy] = routing_strategy)
+    !ismissing(fallback_config) && (params[:fallback_config] = fallback_config)
+    !ismissing(experiment_id) && (params[:experiment_id] = experiment_id)
+    !ismissing(usage_tracking) && (params[:usage_tracking] = usage_tracking)
+    !ismissing(updated_by) && (params[:updated_by] = updated_by)
+    result = mlfpost_v3(instance, "gateway/endpoints/update"; params...)
     return result["endpoint"] |> GatewayEndpoint
 end
 
@@ -367,14 +395,14 @@ Attach a model to a gateway endpoint.
 - `created_by`: Username of the creator.
 
 # Returns
-An instance of type [`GatewayEndpoint`](@ref).
+An instance of type [`GatewayEndpointModelMapping`](@ref).
 """
 function attachmodeltogatewayendpoint(instance::MLFlow, endpoint_id::String,
-    model_config::Dict{String,Any}; created_by::Union{String,Missing}=missing)::GatewayEndpoint
+    model_config::Dict{String,Any}; created_by::Union{String,Missing}=missing)::GatewayEndpointModelMapping
     params = Dict{Symbol,Any}(:endpoint_id => endpoint_id, :model_config => model_config)
     !ismissing(created_by) && (params[:created_by] = created_by)
     result = mlfpost_v3(instance, "gateway/endpoints/models/attach"; params...)
-    return result["endpoint"] |> GatewayEndpoint
+    return result["mapping"] |> GatewayEndpointModelMapping
 end
 
 """
@@ -393,7 +421,7 @@ Detach a model from a gateway endpoint.
 """
 function detachmodelfromgatewayendpoint(instance::MLFlow, endpoint_id::String,
     model_definition_id::String)::Bool
-    mlfdelete_v3(instance, "gateway/endpoints/models/detach";
+    mlfpost_v3(instance, "gateway/endpoints/models/detach";
         endpoint_id=endpoint_id, model_definition_id=model_definition_id)
     return true
 end

src/services/model_version.jl

@@ -112,13 +112,19 @@ end
 - Vector of [`ModelVersion`](@ref) that were found in the [`MLFlow`](@ref) instance.
 - The next page token if there are more results.
 """
-function searchmodelversions(instance::MLFlow; filter::String="",
+function searchmodelversions(instance::MLFlow; filter::Union{String,Missing}=missing,
     max_results::Int64=200000, order_by::Array{String}=String[],
-    page_token::String="")::Tuple{Array{ModelVersion},Union{String,Nothing}}
-    parameters = (; max_results, page_token, filter)
+    page_token::Union{String,Missing}=missing)::Tuple{Array{ModelVersion},Union{String,Nothing}}
+    parameters = Dict{Symbol,Any}(:max_results => max_results)
 
-    if order_by |> !isempty
-        parameters = (; order_by, parameters...)
+    if !ismissing(filter) && !isempty(filter)
+        parameters[:filter] = filter
+    end
+    if !ismissing(page_token) && !isempty(page_token)
+        parameters[:page_token] = page_token
+    end
+    if !isempty(order_by)
+        parameters[:order_by] = order_by
     end
 
     result = mlfget(instance, "model-versions/search"; parameters...)

src/services/registered_model.jl

@@ -108,13 +108,19 @@ end
 - Vector of [`RegisteredModel`](@ref) that were found in the [`MLFlow`](@ref) instance.
 - The next page token if there are more results.
 """
-function searchregisteredmodels(instance::MLFlow; filter::String="",
+function searchregisteredmodels(instance::MLFlow; filter::Union{String,Missing}=missing,
     max_results::Int64=100, order_by::Array{String}=String[],
-    page_token::String="")::Tuple{Array{RegisteredModel},Union{String,Nothing}}
-    parameters = (; max_results, page_token, filter)
+    page_token::Union{String,Missing}=missing)::Tuple{Array{RegisteredModel},Union{String,Nothing}}
+    parameters = Dict{Symbol,Any}(:max_results => max_results)
 
-    if order_by |> !isempty
-        parameters = (; order_by, parameters...)
+    if !ismissing(filter) && !isempty(filter)
+        parameters[:filter] = filter
+    end
+    if !ismissing(page_token) && !isempty(page_token)
+        parameters[:page_token] = page_token
+    end
+    if !isempty(order_by)
+        parameters[:order_by] = order_by
     end
 
     result = mlfget(instance, "registered-models/search"; parameters...)

src/services/run.jl

@@ -170,14 +170,24 @@ Search for runs that satisfy expressions. Search expressions can use [`Metric`](
 - The next page token if there are more results.
 """
 function searchruns(instance::MLFlow; experiment_ids::Array{String}=String[],
-    filter::String="", run_view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY,
+    filter::Union{String,Missing}=missing,
+    run_view_type::ViewType.ViewTypeEnum=ViewType.ACTIVE_ONLY,
     max_results::Int=1000, order_by::Array{String}=String[],
-    page_token::String="")::Tuple{Array{Run},Union{String,Nothing}}
-    parameters = (; experiment_ids, filter, :run_view_type => run_view_type |> Integer,
-        max_results, page_token)
-
-    if order_by |> !isempty
-        parameters = (; order_by, parameters...)
+    page_token::Union{String,Missing}=missing)::Tuple{Array{Run},Union{String,Nothing}}
+    parameters = Dict{Symbol,Any}(
+        :experiment_ids => experiment_ids,
+        :run_view_type => run_view_type |> Integer,
+        :max_results => max_results
+    )
+
+    if !ismissing(filter) && !isempty(filter)
+        parameters[:filter] = filter
+    end
+    if !ismissing(page_token) && !isempty(page_token)
+        parameters[:page_token] = page_token
+    end
+    if !isempty(order_by)
+        parameters[:order_by] = order_by
     end
 
     result = mlfpost(instance, "runs/search"; parameters...)
@@ -208,13 +218,14 @@ Update [`Run`](@ref) metadata.
 """
 function updaterun(instance::MLFlow, run_id::String;
     status::Union{RunStatus.RunStatusEnum,Missing}=missing,
-    end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing})::RunInfo
-    result = mlfpost(instance, "runs/update"; run_id=run_id, status=(status |> Integer),
+    end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing}=missing)::RunInfo
+    result = mlfpost(instance, "runs/update"; run_id=run_id,
+        status=ismissing(status) ? missing : (status |> Integer),
         end_time=end_time, run_name=run_name)
     return result["run_info"] |> RunInfo
 end
 updaterun(instance::MLFlow, run::Run;
     status::Union{RunStatus.RunStatusEnum,Missing}=missing,
-    end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing})::RunInfo =
+    end_time::Union{Int64,Missing}=missing, run_name::Union{String,Missing}=missing)::RunInfo =
     updaterun(instance, run.info.run_id; status=status, end_time=end_time,
         run_name=run_name)

src/services/webhook.jl

@@ -60,8 +60,12 @@ List all webhooks.
 - The next page token if there are more results.
 """
 function listwebhooks(instance::MLFlow; max_results::Int64=100,
-    page_token::String="")::Tuple{Array{Webhook},Union{String,Nothing}}
-    result = mlfget(instance, "webhooks"; max_results=max_results, page_token=page_token)
+    page_token::Union{String,Missing}=missing)::Tuple{Array{Webhook},Union{String,Nothing}}
+    parameters = Dict{Symbol,Any}(:max_results => max_results)
+    if !ismissing(page_token) && !isempty(page_token)
+        parameters[:page_token] = page_token
+    end
+    result = mlfget(instance, "webhooks"; parameters...)
     webhooks = get(result, "webhooks", []) |> (x -> [Webhook(y) for y in x])
     next_page_token = get(result, "next_page_token", nothing)
     return webhooks, next_page_token
@@ -174,5 +178,5 @@ function testwebhook(instance::MLFlow, webhook_id::String;
     end
 
     result = mlfpost(instance, "webhooks/$(webhook_id)/test"; kwargs...)
-    return result["test_result"] |> WebhookTestResult
+    return result["result"] |> WebhookTestResult
 end

src/types/artifact.jl

@@ -11,6 +11,8 @@ struct FileInfo
     is_dir::Bool
     file_size::Int64
 end
+FileInfo(data::AbstractDict) = FileInfo(
+    data["path"], get(data, "is_dir", false), get(data, "file_size", 0))
 Base.show(io::IO, t::FileInfo) = show(io, ShowCase(t, new_lines=true))
 
 """
@@ -28,6 +30,9 @@ struct MultipartUploadCredential
     upload_url::String
     headers::Dict{String,String}
 end
+MultipartUploadCredential(data::AbstractDict) = MultipartUploadCredential(
+    get(data, "part_number", 0), get(data, "url", ""),
+    Dict{String,String}(k => v for (k, v) in get(data, "headers", Dict())))
 
 """
     MultipartUploadPart