star

venv/lib/python3.9/site-packages/mlflow/server/graphql/autogenerated_graphql_schema.py

@@ -0,0 +1,288 @@
+# GENERATED FILE. PLEASE DON'T MODIFY.
+# Run python3 ./dev/proto_to_graphql/code_generator.py to regenerate.
+import graphene
+import mlflow
+from mlflow.server.graphql.graphql_custom_scalars import LongString
+from mlflow.server.graphql.graphql_errors import ApiError
+from mlflow.utils.proto_json_utils import parse_dict
+
+
+class MlflowModelVersionStatus(graphene.Enum):
+    PENDING_REGISTRATION = 1
+    FAILED_REGISTRATION = 2
+    READY = 3
+
+
+class MlflowRunStatus(graphene.Enum):
+    RUNNING = 1
+    SCHEDULED = 2
+    FINISHED = 3
+    FAILED = 4
+    KILLED = 5
+
+
+class MlflowViewType(graphene.Enum):
+    ACTIVE_ONLY = 1
+    DELETED_ONLY = 2
+    ALL = 3
+
+
+class MlflowModelVersionTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowModelVersion(graphene.ObjectType):
+    name = graphene.String()
+    version = graphene.String()
+    creation_timestamp = LongString()
+    last_updated_timestamp = LongString()
+    user_id = graphene.String()
+    current_stage = graphene.String()
+    description = graphene.String()
+    source = graphene.String()
+    run_id = graphene.String()
+    status = graphene.Field(MlflowModelVersionStatus)
+    status_message = graphene.String()
+    tags = graphene.List(graphene.NonNull(MlflowModelVersionTag))
+    run_link = graphene.String()
+    aliases = graphene.List(graphene.String)
+
+
+class MlflowSearchModelVersionsResponse(graphene.ObjectType):
+    model_versions = graphene.List(graphene.NonNull(MlflowModelVersion))
+    next_page_token = graphene.String()
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowDatasetSummary(graphene.ObjectType):
+    experiment_id = graphene.String()
+    name = graphene.String()
+    digest = graphene.String()
+    context = graphene.String()
+
+
+class MlflowSearchDatasetsResponse(graphene.ObjectType):
+    dataset_summaries = graphene.List(graphene.NonNull(MlflowDatasetSummary))
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowMetricWithRunId(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.Float()
+    timestamp = LongString()
+    step = LongString()
+    run_id = graphene.String()
+
+
+class MlflowGetMetricHistoryBulkIntervalResponse(graphene.ObjectType):
+    metrics = graphene.List(graphene.NonNull(MlflowMetricWithRunId))
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowFileInfo(graphene.ObjectType):
+    path = graphene.String()
+    is_dir = graphene.Boolean()
+    file_size = LongString()
+
+
+class MlflowListArtifactsResponse(graphene.ObjectType):
+    root_uri = graphene.String()
+    files = graphene.List(graphene.NonNull(MlflowFileInfo))
+    next_page_token = graphene.String()
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowDataset(graphene.ObjectType):
+    name = graphene.String()
+    digest = graphene.String()
+    source_type = graphene.String()
+    source = graphene.String()
+    schema = graphene.String()
+    profile = graphene.String()
+
+
+class MlflowInputTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowDatasetInput(graphene.ObjectType):
+    tags = graphene.List(graphene.NonNull(MlflowInputTag))
+    dataset = graphene.Field(MlflowDataset)
+
+
+class MlflowRunInputs(graphene.ObjectType):
+    dataset_inputs = graphene.List(graphene.NonNull(MlflowDatasetInput))
+
+
+class MlflowRunTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowParam(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowMetric(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.Float()
+    timestamp = LongString()
+    step = LongString()
+
+
+class MlflowRunData(graphene.ObjectType):
+    metrics = graphene.List(graphene.NonNull(MlflowMetric))
+    params = graphene.List(graphene.NonNull(MlflowParam))
+    tags = graphene.List(graphene.NonNull(MlflowRunTag))
+
+
+class MlflowRunInfo(graphene.ObjectType):
+    run_id = graphene.String()
+    run_uuid = graphene.String()
+    run_name = graphene.String()
+    experiment_id = graphene.String()
+    user_id = graphene.String()
+    status = graphene.Field(MlflowRunStatus)
+    start_time = LongString()
+    end_time = LongString()
+    artifact_uri = graphene.String()
+    lifecycle_stage = graphene.String()
+
+
+class MlflowRun(graphene.ObjectType):
+    info = graphene.Field(MlflowRunInfo)
+    data = graphene.Field(MlflowRunData)
+    inputs = graphene.Field(MlflowRunInputs)
+
+
+class MlflowSearchRunsResponse(graphene.ObjectType):
+    runs = graphene.List(graphene.NonNull('mlflow.server.graphql.graphql_schema_extensions.MlflowRunExtension'))
+    next_page_token = graphene.String()
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowGetRunResponse(graphene.ObjectType):
+    run = graphene.Field('mlflow.server.graphql.graphql_schema_extensions.MlflowRunExtension')
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowExperimentTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowExperiment(graphene.ObjectType):
+    experiment_id = graphene.String()
+    name = graphene.String()
+    artifact_location = graphene.String()
+    lifecycle_stage = graphene.String()
+    last_update_time = LongString()
+    creation_time = LongString()
+    tags = graphene.List(graphene.NonNull(MlflowExperimentTag))
+
+
+class MlflowGetExperimentResponse(graphene.ObjectType):
+    experiment = graphene.Field(MlflowExperiment)
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowSearchModelVersionsInput(graphene.InputObjectType):
+    filter = graphene.String()
+    max_results = LongString()
+    order_by = graphene.List(graphene.String)
+    page_token = graphene.String()
+
+
+class MlflowSearchDatasetsInput(graphene.InputObjectType):
+    experiment_ids = graphene.List(graphene.String)
+
+
+class MlflowGetMetricHistoryBulkIntervalInput(graphene.InputObjectType):
+    run_ids = graphene.List(graphene.String)
+    metric_key = graphene.String()
+    start_step = graphene.Int()
+    end_step = graphene.Int()
+    max_results = graphene.Int()
+
+
+class MlflowListArtifactsInput(graphene.InputObjectType):
+    run_id = graphene.String()
+    run_uuid = graphene.String()
+    path = graphene.String()
+    page_token = graphene.String()
+
+
+class MlflowSearchRunsInput(graphene.InputObjectType):
+    experiment_ids = graphene.List(graphene.String)
+    filter = graphene.String()
+    run_view_type = graphene.Field(MlflowViewType)
+    max_results = graphene.Int()
+    order_by = graphene.List(graphene.String)
+    page_token = graphene.String()
+
+
+class MlflowGetRunInput(graphene.InputObjectType):
+    run_id = graphene.String()
+    run_uuid = graphene.String()
+
+
+class MlflowGetExperimentInput(graphene.InputObjectType):
+    experiment_id = graphene.String()
+
+
+class QueryType(graphene.ObjectType):
+    mlflow_get_experiment = graphene.Field(MlflowGetExperimentResponse, input=MlflowGetExperimentInput())
+    mlflow_get_metric_history_bulk_interval = graphene.Field(MlflowGetMetricHistoryBulkIntervalResponse, input=MlflowGetMetricHistoryBulkIntervalInput())
+    mlflow_get_run = graphene.Field(MlflowGetRunResponse, input=MlflowGetRunInput())
+    mlflow_list_artifacts = graphene.Field(MlflowListArtifactsResponse, input=MlflowListArtifactsInput())
+    mlflow_search_model_versions = graphene.Field(MlflowSearchModelVersionsResponse, input=MlflowSearchModelVersionsInput())
+
+    def resolve_mlflow_get_experiment(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.GetExperiment()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.get_experiment_impl(request_message)
+
+    def resolve_mlflow_get_metric_history_bulk_interval(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.GetMetricHistoryBulkInterval()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.get_metric_history_bulk_interval_impl(request_message)
+
+    def resolve_mlflow_get_run(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.GetRun()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.get_run_impl(request_message)
+
+    def resolve_mlflow_list_artifacts(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.ListArtifacts()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.list_artifacts_impl(request_message)
+
+    def resolve_mlflow_search_model_versions(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.model_registry_pb2.SearchModelVersions()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.search_model_versions_impl(request_message)
+
+
+class MutationType(graphene.ObjectType):
+    mlflow_search_datasets = graphene.Field(MlflowSearchDatasetsResponse, input=MlflowSearchDatasetsInput())
+    mlflow_search_runs = graphene.Field(MlflowSearchRunsResponse, input=MlflowSearchRunsInput())
+
+    def resolve_mlflow_search_datasets(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.SearchDatasets()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.search_datasets_impl(request_message)
+
+    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)

venv/lib/python3.9/site-packages/mlflow/server/graphql/graphql_custom_scalars.py

@@ -0,0 +1,24 @@
+import graphene
+from graphql.language.ast import IntValueNode
+
+
+class LongString(graphene.Scalar):
+    """
+    LongString Scalar type to prevent truncation to max integer in JavaScript.
+    """
+
+    description = "Long converted to string to prevent truncation to max integer in JavaScript"
+
+    @staticmethod
+    def serialize(long):
+        return str(long)
+
+    @staticmethod
+    def parse_literal(node):
+        if isinstance(node, IntValueNode):
+            return int(node.value)
+        return None
+
+    @staticmethod
+    def parse_value(value):
+        return int(value)

venv/lib/python3.9/site-packages/mlflow/server/graphql/graphql_errors.py

@@ -0,0 +1,15 @@
+import graphene
+
+
+class ErrorDetail(graphene.ObjectType):
+    # NOTE: This is not an exhaustive list, might need to add more things in the future if needed.
+    field = graphene.String()
+    message = graphene.String()
+
+
+class ApiError(graphene.ObjectType):
+    code = graphene.String()
+    message = graphene.String()
+    help_url = graphene.String()
+    trace_id = graphene.String()
+    error_details = graphene.List(ErrorDetail)

venv/lib/python3.9/site-packages/mlflow/server/graphql/graphql_schema_extensions.py

@@ -0,0 +1,74 @@
+import graphene
+from graphql import (
+    DirectiveLocation,
+    GraphQLArgument,
+    GraphQLDirective,
+    GraphQLNonNull,
+    GraphQLString,
+)
+
+import mlflow
+from mlflow.server.graphql.autogenerated_graphql_schema import (
+    MlflowExperiment,
+    MlflowModelVersion,
+    MlflowRun,
+    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 Query(QueryType):
+    test = graphene.Field(Test, input_string=graphene.String(), description="Simple echoing field")
+
+    def resolve_test(self, info, input_string):
+        return {"output": input_string}
+
+
+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])