Pytorch and Lightning 2.0 fixes

[shrinath-suresh]

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Pytorch 2.0 and Lightning 2.0 has been released recently.

Lightning 2.0 changes are published here

Following are the impacts on mlflow pytorch library

1. The package name has been changed from pytorch_lightning to lightning
2. validation_epoch_end and test_epoch_end has been removed
3. pl.Trainer.add_argparse_args(parent_parser=parser) has been removed. This line of code combines all the pytorch lightning command line arguments with the user's ArgumentParser. Instead we have to use LightningCLI as mentioned in the release link
4. The training arguments have changed. Previously, we use --max_epochs . Now, its --trainer.max_epochs
5. When installing lightning package, pytorch_lightning gets installed as well. Hence, we dont have to change everything in mllfow.pytorch.autolog library. But, we may need to remove older version checks.

How is this patch tested?

• MNIST
  mnist_cpu_conda.txt
  mnist_cpu_virtualenv.txt
  mnist_gpu_conda.txt

• BERT
  bert_cpu_conda .txt
  bert_gpu_conda.txt

• Captum
  captum_cpu_virtual_env.txt

• AX
  ax_hyper_param_tuning_cpu_conda.txt

• Iterative Pruning
  iterative_pruning_cpu_conda.txt

Observation:
In multi gpu tests, the lightning throws _thread.lock object error while pickling the model and the model is not getting logged into mlflow. However, the issue has been fixed in lightning master branch. Attaching screenshot for reference

screenshots.zip

• Existing unit/integration tests
• New unit/integration tests
• Manual tests (describe details, including test results, below)

Does this PR change the documentation?

• No. You can skip the rest of this section.
• Yes. Make sure the changed pages / sections render correctly in the documentation preview.

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/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• 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

[mlflow-automation]

Documentation preview for 3ade269 will be available here when this CircleCI job completes successfully.

More info

• Ignore this comment if this PR does not change the documentation.
• It takes a few minutes for the preview to be available.
• The preview is updated when a new commit is pushed to this PR.
• This comment was created by https://github.com/mlflow/mlflow/actions/runs/4500176285.

[harupy]

@shrinath-suresh thanks for the PR, ping me when it's ready for review!

[shrinath-suresh]

@shrinath-suresh thanks for the PR, ping me when it's ready for review!

sure @harupy . We are working on the GPU tests. Will let you know once the PR is ready for review

[shrinath-suresh]

@harupy The PR is ready for review

[shrinath-suresh]

@harupy Observing the following error message in CI - Message: 'Could not install packages due to an OSError: [Errno 28] No space left on device\n' ... Some of the test cases are failing during the package installation.

And the cross version tests are failing.. When i check the failure logs.. it has failed in import pytorch_lightning as pl . And in the master branch autolog code, the import is in try except block with pass in except - code link

For running autolog test, pytorch_lightning or lightning is mandatory right ?

[harupy]

Message: 'Could not install packages due to an OSError: [Errno 28] No space left on device\n'

Where did you see this error? Can you share the URL?

For running autolog test, pytorch_lightning or lightning is mandatory right ?

No, it's not mandatory. See #7627 for more details.

[shrinath-suresh]

Message: 'Could not install packages due to an OSError: [Errno 28] No space left on device\n'

Where did you see this error? Can you share the URL?

For running autolog test, pytorch_lightning or lightning is mandatory right ?

No, it's not mandatory. See #7627 for more details.

I was getting the Space error in this run - . But in the latest run i am not seeing it.

Thanks for the reference. Cross functional tests (for pytorch flavor) failed due to the import error.. i fixed it supporting all three options (lightning, pytorch lightning <=1.9.4 and pytorch flavor)

    # For lightning>=2.0.0 flavor
    try:
        import lightning as L
    except ImportError:
        # for pytorch-lightning <= 1.9.4 flavor
        try:
            import pytorch_lightning as pl
        except ImportError:
            # For PyTorch flavor
            pass
        else:
            from mlflow.pytorch._lightning_autolog import patched_fit

            safe_patch(FLAVOR_NAME, pl.Trainer, "fit", patched_fit, manage_run=True)
    else:
        from mlflow.pytorch._lightning_autolog import patched_fit

        safe_patch(FLAVOR_NAME, L.Trainer, "fit", patched_fit, manage_run=True)

[harupy]

@shrinath-suresh What happens if we patch both L.Trainer and pl.Trainer?

[shrinath-suresh]

@shrinath-suresh What happens if we patch both L.Trainer and pl.Trainer?

I haven't tested the case where we patch both of it. Can you please explain a bit more on the context.

As of now, the logic is to patch L.Trainer for lightning package and pl.Trainer for the ptl<=1.9.4. If we patch both, ptl<=1.9.4 will not have L.Trainer and fail.

[harupy]

I haven't tested the case where we patch both of it. Can you please explain a bit more on the context.

Sure. If patching both Trainer is fine, we can simplify the code as I suggested.

    try:
        import pytorch_lightning as pl
    except ImportError:
        pass
    else:
        from mlflow.pytorch._lightning_autolog import patched_fit

        safe_patch(FLAVOR_NAME, pl.Trainer, "fit", patched_fit, manage_run=True)

    try:
        import lightning as L
    except ImportError:
        pass
    else:
        from mlflow.pytorch._lightning_autolog import patched_fit

        safe_patch(FLAVOR_NAME, L.Trainer, "fit", patched_fit, manage_run=True)

Just because lightning can be imported doesn't mean the user uses lightning. The user might use pytorch_lightning even when lightning is available.

[harupy]

Thanks for the update, let's see what happens.

[shrinath-suresh]

Thanks for the update, let's see what happens.

Thanks for the suggestion @harupy . I added both and tested it in cpu and gpu. Until pytorch_lightning support is removed from lightning , we are good. At some point, we might need to move into lightning. Which may involve almost the rewrite of entire library. We can use MlflowLogger at that time and change the library. Can you please review and let me know your comments.

[harupy]

LGTM!