Skip to content

Threading support for legacy loading of checkpoints #12814

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Apr 28, 2022
Merged

Threading support for legacy loading of checkpoints #12814

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
203ce91
Add lock to pl_legacy_patch context manager; avoiding race condition …
krshrimali Apr 18, 2022
99acdf2
Add tests, make lock global after discussion with Adrian
krshrimali Apr 28, 2022
a8fc820
Imports
krshrimali Apr 28, 2022
ee21793
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Apr 28, 2022
a67dc9b
make lock protected
krshrimali Apr 28, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions pytorch_lightning/utilities/migration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,14 @@
from __future__ import annotations

import sys
import threading
from types import ModuleType, TracebackType

import pytorch_lightning.utilities.argparse

# Create a global lock to ensure no race condition with deleting sys modules
_lock = threading.Lock()


class pl_legacy_patch:
"""Registers legacy artifacts (classes, methods, etc.) that were removed but still need to be included for
Expand All @@ -35,6 +39,7 @@ class pl_legacy_patch:
"""

def __enter__(self) -> None:
_lock.acquire()
# `pl.utilities.argparse_utils` was renamed to `pl.utilities.argparse`
legacy_argparse_module = ModuleType("pytorch_lightning.utilities.argparse_utils")
sys.modules["pytorch_lightning.utilities.argparse_utils"] = legacy_argparse_module
Expand All @@ -49,3 +54,4 @@ def __exit__(
if hasattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default"):
delattr(pytorch_lightning.utilities.argparse, "_gpus_arg_default")
del sys.modules["pytorch_lightning.utilities.argparse_utils"]
_lock.release()
23 changes: 23 additions & 0 deletions tests/checkpointing/test_legacy_checkpoints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
import glob
import os
import sys
import threading
from unittest.mock import patch

import pytest
Expand Down Expand Up @@ -60,6 +61,28 @@ def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningMo
trainer.should_stop = True


@pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
def test_legacy_ckpt_threading(tmpdir, pl_version: str):
def load_model():
import torch

from pytorch_lightning.utilities.migration import pl_legacy_patch

with pl_legacy_patch():
_ = torch.load(PATH_LEGACY)

PATH_LEGACY = os.path.join(LEGACY_CHECKPOINTS_PATH, pl_version)
with patch("sys.path", [PATH_LEGACY] + sys.path):
t1 = threading.Thread(target=load_model)
t2 = threading.Thread(target=load_model)

t1.start()
t2.start()

t1.join()
t2.join()


@pytest.mark.parametrize("pl_version", LEGACY_BACK_COMPATIBLE_PL_VERSIONS)
def test_resume_legacy_checkpoints(tmpdir, pl_version: str):
PATH_LEGACY = os.path.join(LEGACY_CHECKPOINTS_PATH, pl_version)
Expand Down