Repository: wandb/wandb

Issue:
[CLI]: W&B service deadlock on KeyboardInterrupt with PyTorch-Lightning and Hydra
### Describe the bug

<!--- Description of the issue below  -->
I have a Pytorch-Lightning-Project where I use [hydra](https://hydra.cc) for configuration. A simplified code example follows below:

<!--- A minimal code snippet between the quotes below  -->
```python
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
import wandb
import hydra

@hydra.main(config_name="config", config_path=".")
def train(cfg : DictConfig) -> None:
    logger = WandbLogger(project='project')

    model = Model(cfg)
    data = DataModule(cfg)

    trainer = pl.Trainer(
        **cfg.trainer,
        logger=logger
    )

    try:
        trainer.fit(model, datamodule=data)
        trainer.test(datamodule=data)
        exit_code = 0
    except Exception as e:
        exit_code = -1
        raise e
    finally:
        wandb.finish(exit_code)

if __name__ == "__main__":
    train()
```
But since Pytorch-Lightning version 1.6.0 my code hangs, if I interrupt the training with a CTRL+C. I used the [hanging-threads](https://pypi.org/project/hanging-threads/2.0.3/) library to find out the hanging thread and the traceback reads:

<!--- A full traceback of the exception in the quotes below -->
```shell
Epoch 0:   12%
Aborted!
/home/user/.local/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py:724: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...
  rank_zero_warn("Detected KeyboardInterrupt, attempting graceful shutdown...")
[...]
wandb: Waiting for W&B process to finish... (failed -1). Press Control-C to abort syncing.
[...]
---------Thread 140402038120640 "MainThread" hangs ---------
        File "/home/user/.local/lib/python3.7/site-packages/wandb/sdk/wandb_manager.py", line 126, in <lambda>
                self._atexit_lambda = lambda: self._atexit_teardown()
        File "/home/user/.local/lib/python3.7/site-packages/wandb/sdk/wandb_manager.py", line 134, in _atexit_teardown
                self._teardown(exit_code)
        File "/home/user/.local/lib/python3.7/site-packages/wandb/sdk/wandb_manager.py", line 141, in _teardown
                result = self._service.join()
        File "/home/user/.local/lib/python3.7/site-packages/wandb/sdk/service/service.py", line 119, in join
                ret = self._internal_proc.wait()
        File "/opt/conda/lib/python3.7/subprocess.py", line 1019, in wait
                return self._wait(timeout=timeout)
        File "/opt/conda/lib/python3.7/subprocess.py", line 1653, in _wait
                (pid, sts) = self._try_wait(0)
        File "/opt/conda/lib/python3.7/subprocess.py", line 1611, in _try_wait
                (pid, sts) = os.waitpid(self.pid, wait_flags)
```
Then I modified the wandb-service code and added a timeout in L121 in [wandb/sdk/service/service.py](https://github.com/wandb/client/blob/master/wandb/sdk/service/service.py#L121) and this solved the issue by killing the service-process. Since the service-process was added to the WandbLogger from PyTorch-Lightning in version 1.6.0 (see [here](https://github.com/wandb/client/blob/master/docs/dev/wandb-service-user.md)) I suspect it to be the problem.

### Additional Files

_No response_

### Environment

WandB version: 0.12.11 (also tested 0.12.14)

OS: Ubuntu 16.04.

Python version: 3.7

Versions of relevant libraries:
PyTorch-Lightning 1.6.0 (also tested 1.6.1)


### Additional Context

_No response_


Discussion:
Hi @kielnino,

Thank you for reporting this and I am sorry you are facing this issue. Could you share a bit more information regarding this issue? Specifically, could you share the `debug.log` and `debug-internal.log` files associated with the run where you see this deadlock? They should be present in a folder called `wandb` in the directory where you ran these experiments.

I'll try to reproduce this issue on my end so that we can get this resolved ASAP.

Thanks,
Ramit
Hi @ramit-wandb,

thank you for the quick response. Here is a complete running code-example and the corresponding log-files:

```shell
pip install hydra-core==1.1.0
pip install pytorch-lightning==1.6.1
pip install wandb==0.12.14
```

First create a file `config.yaml` with the following content:
```yaml
trainer:
  max_epochs: 50
  log_every_n_steps: 1
```
and then run this python-file from the same folder:
```python
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
import wandb
import hydra
import torch
from torch.nn import functional as F
from omegaconf import DictConfig
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader


class MNISTModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


@hydra.main(config_name="config", config_path=".")
def train(cfg : DictConfig) -> None:
    logger = WandbLogger(project='project')

    mnist_model = MNISTModel()
    # Init DataLoader from MNIST Dataset
    train_ds = MNIST('/tmp', train=True, download=True, transform=transforms.ToTensor())
    train_loader = DataLoader(train_ds, batch_size=32)

    trainer = pl.Trainer(
        **cfg.trainer,
        logger=logger
    )

    try:
        trainer.fit(mnist_model, train_loader)
        exit_code = 0
    except Exception as e:
        exit_code = -1
        raise e
    finally:
        wandb.finish(exit_code)

if __name__ == "__main__":
    train()
```
When the training starts, just abort it with CTRL + C and the process will hang...

Here are the debug files:

[debug.log](https://github.com/wandb/client/files/8528393/debug.log)
[debug-internal.log](https://github.com/wandb/client/files/8528394/debug-internal.log)
WandB Internal User commented: 
kielnino commented: 
Hi @ramit-wandb,

thank you for the quick response. Here is a complete running code-example and the corresponding log-files:

```shell
pip install hydra-core==1.1.0
pip install pytorch-lightning==1.6.1
pip install wandb==0.12.14
```

First create a file `config.yaml` with the following content:
```yaml
trainer:
  max_epochs: 50
  log_every_n_steps: 1
```
and then run this python-file from the same folder:
```python
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
import wandb
import hydra
import torch
from torch.nn import functional as F
from omegaconf import DictConfig
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader


class MNISTModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


@hydra.main(config_name="config", config_path=".")
def train(cfg : DictConfig) -> None:
    logger = WandbLogger(project='project')

    mnist_model = MNISTModel()
    # Init DataLoader from MNIST Dataset
    train_ds = MNIST('/tmp', train=True, download=True, transform=transforms.ToTensor())
    train_loader = DataLoader(train_ds, batch_size=32)

    trainer = pl.Trainer(
        **cfg.trainer,
        logger=logger
    )

    try:
        trainer.fit(mnist_model, train_loader)
        exit_code = 0
    except Exception as e:
        exit_code = -1
        raise e
    finally:
        wandb.finish(exit_code)

if __name__ == "__main__":
    train()
```

Here are the debug files:

[debug.log](https://github.com/wandb/client/files/8528393/debug.log)
[debug-internal.log](https://github.com/wandb/client/files/8528394/debug-internal.log)

Hi @kielnino, Thanks a lot for reporting this issue and for the detailed reproduction steps, it was a pleasure to read your issue. We were able to reproduce it and working on a solution. Hopefully this will be included in the upcoming release. 
Hi @kielnino,

This should now be resolved with our latest release of `wandb 0.12.16`. Please let me know if you have any questions with this.

Thanks,
Ramit
Hi,

The issue still exist with the latest versions. I'm using `wandb=0.13.1` and `pytorch-lightning=1.7.1`, and KeyboardInterrupt will deadlock Pytorch Lightning when a WandbLogger is used, but works fine when there is no WandbLogger.
@TheMrZZ Thanks for reporting the issue. mind sharing more details about your code?
(since `service` was enabled in `pytorch-lightning=1.6` where we verified the issue was resolved)

In the meantime we will try to reproduce it in the latest version of `pytorch-lightning`
Thanks. The minimum reproduction example is [available here](https://gist.github.com/TheMrZZ/b1989ea7bf2867e0b2a6bb44e7632ffa).

Removing `logger=...` will allow the program to stop successfully on KeyboardInterrupt.

## Output with logger
(removed clutter from the logs)
```
  | Name   | Type   | Params
----------------------------------
0 | linear | Linear | 2
----------------------------------
2         Trainable params
0         Non-trainable params
2         Total params
0.000     Total estimated model params size (MB)
Epoch 94: 100%|████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 333.36it/s, loss=0, v_num=fkyi]

C:\Users\redacted: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...
  rank_zero_warn("Detected KeyboardInterrupt, attempting graceful shutdown...")

Aborted!
⬜
```
And then, the program deadlocks. I never get back the control of the terminal.
![image](https://user-images.githubusercontent.com/30944236/184986034-c77798e7-8239-4481-8fa8-469a988cd45d.png)

@TheMrZZ Thanks a lot for the repro. I see that you are using windows, I think this is the reason the fix didn't work for you. We are working on a new fix and will update this branch when it is available. Apologies for the inconvenience.

Resolving patch:
diff --git a/wandb/docker/__init__.py b/wandb/docker/__init__.py
--- a/wandb/docker/__init__.py
+++ b/wandb/docker/__init__.py
@@ -141,7 +141,7 @@ def auth_token(registry: str, repo: str) -> Dict[str, str]:
     auth_info = auth_config.resolve_authconfig(registry)
     if auth_info:
         normalized = {k.lower(): v for k, v in auth_info.items()}
-        normalized_auth_info: Optional[Tuple[str, str]] = (  # type: ignore
+        normalized_auth_info = (
             normalized.get("username"),
             normalized.get("password"),
         )
@@ -166,7 +166,7 @@ def auth_token(registry: str, repo: str) -> Dict[str, str]:
             + "?service={}&scope=repository:{}:pull".format(
                 info["bearer"]["service"], repo
             ),
-            auth=normalized_auth_info,
+            auth=normalized_auth_info,  # type: ignore
             timeout=3,
         )
         res.raise_for_status()
diff --git a/wandb/sdk/service/service.py b/wandb/sdk/service/service.py
--- a/wandb/sdk/service/service.py
+++ b/wandb/sdk/service/service.py
@@ -4,6 +4,7 @@
 """
 
 import os
+import platform
 import subprocess
 import sys
 import tempfile
@@ -68,7 +69,12 @@ def _launch_server(self) -> None:
         # - https://github.com/wandb/wandb/blob/archive/old-cli/wandb/__init__.py
         # - https://stackoverflow.com/questions/1196074/how-to-start-a-background-process-in-python
 
-        kwargs: Dict[str, Any] = dict(close_fds=True, start_new_session=True)
+        kwargs: Dict[str, Any] = dict(close_fds=True)
+        # flags to handle keyboard interrupt signal that is causing a hang
+        if platform.system() == "Windows":
+            kwargs.update(creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)  # type: ignore [attr-defined]
+        else:
+            kwargs.update(start_new_session=True)
 
         pid = os.getpid()
 


Test patch: