Add return type annotations to key MLflow APIs

[harupy]

Signed-off-by: harupy 17039389+harupy@users.noreply.github.com

What changes are proposed in this pull request?

Put type annotations on key MLflow APIs.

Which API should be annotated?

• Fluent APIs (e.g. mlflow.start_run)
• Client APIs (e.g. mlflow.tracking.MlflowClient.get_run)

How is this patch tested?

Manual testing

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/projects: MLproject format, project running backends
• area/scoring: Local serving, model deployment tools, spark UDFs
• area/server-infra: MLflow server, JavaScript dev server
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, JavaScript, plotting
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[harupy]

A workaround for the skinny client (which doesn't install pandas).

[harupy]

Now mlflow.start_run is type-annotated and looks like this:

def start_run(...) -> ActiveRun:
    ...

Sphinx tries to create a link for ActiveRun using its fully qualified class name (mlflow.tracking.fluent.ActiveRun) as a reference target, but our doc doesn't contain the reference mlflow.tracking.fluent.ActiveRun so it will result in the following warning (which is treated as an error due to the -W option):

/home/circleci/project/mlflow/__init__.py:docstring of mlflow.tracking.fluent.active_run:: WARNING: py:class reference target not found: mlflow.tracking.fluent.ActiveRun

As a workaround, map mlflow.tracking.fluent.ActiveRun to mlflow.ActiveRun.

[WeichenXu123]

Seems missing some similar cases, such as:

mlflow.deployments.BaseDeploymentClient
mflow.models.{Model, ModelSignature, ModelInputExample, infer_signature, FlavorBackend}
mflow.projects.SubmittedRun

mlflow.tracking.MlflowClient
mflow.types.{DataType, ColSpec, Schema, TensorSpec}

[WeichenXu123]

Maybe we can automatically generate the reference_map, by searching all __all__ exports in every sub-module, and for each item in __all__ , if it is not a function but an Class, then we can add it into reference_map, we can generate the original ref by https://github.com/mlflow/mlflow/pull/4242/files#r614140548

[harupy]

One challenge is:

{"mlflow.tracking.fluent.ActiveRun": "< valid_ref >"}

How can we dynamically get < valid_ref > ?

[WeichenXu123]

How can we dynamically get < valid_ref > ?

Using module (include the all list) name + item name?
we can search 2 depth level modules

[harupy]

Offline discussion: Ideally, we should put both argument and return type annotations because both are benefitial for auto-completion in Databricks notebook, but we can address them separately. In this PR, I'm adding return type annotations.

[smurching]

Could we add some context (e.g. basically just your description in https://github.com/mlflow/mlflow/pull/4242/files#r613161818) to why this is necessary in the docstring, for future readers? Namely that Sphinx computes references using fully-qualified classnames, so references in undocumented modules (even if the referenced object is exposed via a different module from the one it's defined in) are considered invalid by Sphinx (you can use the ActiveRun example)

[harupy]

Thanks for the comment! Added some context!

[smurching]

Basically LGTM with minor comments

[harupy]

@smurching Thanks for the review! I'll merge the PR once all the CI checks pass.