Parallelize download of artifacts and improve performance of downloading large models.

[mehtayogita]

What changes are proposed in this pull request?

Parallelize download function for ArtifactRepository.

How is this patch tested?

Unit testing
Manually testing with models having large artifacts. For a model with ~200 artifacts, this change provided 2x improvement. The code version at HEAD took 109 secs to load the model using ModelsArtifactRepository, while the updated code in this PR took 66.56 secs.

Does this PR change the documentation?

• No. You can skip the rest of this section.
• Yes. Make sure the changed pages / sections render correctly by following the steps below.

1. Check the status of the ci/circleci: build_doc check. If it's successful, proceed to the
   next step, otherwise fix it.
2. Click Details on the right to open the job page of CircleCI.
3. Click the Artifacts tab.
4. Click docs/build/html/index.html.
5. Find the changed pages / sections and make sure they render correctly.

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.

Support parallel download of artifacts. Fix issue: #5338

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: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• 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

[mehtayogita]

@ankit-db @dbczumar sending this for early review to confirm if it is ok to add the parallel download functionality directly in the ArtifactRepository class. If this approach looks good, will proceed with further testing and adding more unit tests.

[dbczumar]

@mehtayogita This is awesome! Thanks for bringing download parallelism to all of the artifact repositories across MLflow. LGTM; feel free to merge once you've supplemented unit tests & conducted manual testing and updated the PR description with a manual test plan.

[ankit-db]

Really excited about the contribution! 2 small questions

[ankit-db]

Also, it'd be good to add some unit testing here to verify edge cases - how do we plan on doing that?

[mehtayogita]

Thanks for the quick review! Will confirm this is working as intended with manual testing, add tests and request another round of review. Thanks!

Also, it'd be good to add some unit testing here to verify edge cases - how do we plan on doing that?

Thanks for the quick review! Yes the CL needs to be polished more and have unit tests. The parallelism is currently tested using tests/store/artifact/test_databricks_artifact_repo.py, while it is better to be part of tests/store/artifact/test_artifact_repo.py. The edge cases will be tested in there. I will first confirm this change is working as intended with manual testing and will then update the PR with tests and address TODOs left in the code. Thanks!

[mehtayogita]

Updated the PR to include the unit tests and confirmed performance improvements with manual testing.

[ankit-db]

Looks good! Thanks for the change - great first contribution to MLflow!

[ankit-db]

Btw, could you add a bit more detail about the performance test you did and include a reference to the performance improvement in the release note?

[dbczumar]

LGTM! Thanks @mehtayogita !

[mehtayogita]

Btw, could you add a bit more detail about the performance test you did and include a reference to the performance improvement in the release note?

Added details.