wehub-resource-sync
96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
import math
|
|
|
|
import graphene
|
|
from graphql import (
|
|
DirectiveLocation,
|
|
GraphQLArgument,
|
|
GraphQLDirective,
|
|
GraphQLNonNull,
|
|
GraphQLString,
|
|
)
|
|
|
|
import mlflow
|
|
from mlflow.server.graphql.autogenerated_graphql_schema import (
|
|
MlflowExperiment,
|
|
MlflowMetric,
|
|
MlflowModelVersion,
|
|
MlflowRun,
|
|
MlflowSearchRunsInput,
|
|
MlflowSearchRunsResponse,
|
|
MutationType,
|
|
QueryType,
|
|
)
|
|
from mlflow.utils.proto_json_utils import parse_dict
|
|
|
|
# Component identifier, to keep compatible with Databricks in-house implementations.
|
|
ComponentDirective = GraphQLDirective(
|
|
name="component",
|
|
locations=[
|
|
DirectiveLocation.QUERY,
|
|
DirectiveLocation.MUTATION,
|
|
],
|
|
args={"name": GraphQLArgument(GraphQLNonNull(GraphQLString))},
|
|
)
|
|
|
|
|
|
class Test(graphene.ObjectType):
|
|
output = graphene.String(description="Echoes the input string")
|
|
|
|
|
|
class TestMutation(graphene.ObjectType):
|
|
output = graphene.String(description="Echoes the input string")
|
|
|
|
|
|
class MlflowRunExtension(MlflowRun):
|
|
experiment = graphene.Field(MlflowExperiment)
|
|
model_versions = graphene.List(graphene.NonNull(MlflowModelVersion))
|
|
|
|
def resolve_experiment(self, info):
|
|
experiment_id = self.info.experiment_id
|
|
input_dict = {"experiment_id": experiment_id}
|
|
request_message = mlflow.protos.service_pb2.GetExperiment()
|
|
parse_dict(input_dict, request_message)
|
|
return mlflow.server.handlers.get_experiment_impl(request_message).experiment
|
|
|
|
def resolve_model_versions(self, info):
|
|
run_id = self.info.run_id
|
|
input_dict = {"filter": f"run_id='{run_id}'"}
|
|
request_message = mlflow.protos.model_registry_pb2.SearchModelVersions()
|
|
parse_dict(input_dict, request_message)
|
|
return mlflow.server.handlers.search_model_versions_impl(request_message).model_versions
|
|
|
|
|
|
class MlflowMetricExtension(MlflowMetric):
|
|
value = graphene.Float()
|
|
|
|
# metric values that are NaN will cause an error in graphQL validation as
|
|
# the type is Float. as a workaround, we return None if the value is NaN.
|
|
def resolve_value(self, info):
|
|
return None if math.isnan(self.value) else self.value
|
|
|
|
|
|
class Query(QueryType):
|
|
test = graphene.Field(Test, input_string=graphene.String(), description="Simple echoing field")
|
|
mlflow_search_runs = graphene.Field(MlflowSearchRunsResponse, input=MlflowSearchRunsInput())
|
|
|
|
def resolve_test(self, info, input_string):
|
|
return {"output": input_string}
|
|
|
|
def resolve_mlflow_search_runs(self, info, input):
|
|
input_dict = vars(input)
|
|
request_message = mlflow.protos.service_pb2.SearchRuns()
|
|
parse_dict(input_dict, request_message)
|
|
return mlflow.server.handlers.search_runs_impl(request_message)
|
|
|
|
|
|
class Mutation(MutationType):
|
|
testMutation = graphene.Field(
|
|
TestMutation, input_string=graphene.String(), description="Simple echoing field"
|
|
)
|
|
|
|
def resolve_test_mutation(self, info, input_string):
|
|
return {"output": input_string}
|
|
|
|
|
|
schema = graphene.Schema(query=Query, mutation=Mutation, directives=[ComponentDirective])
|