Databricks 笔记本挂着 pytorch

我们遇到 Databricks 笔记本问题。我们的一个笔记本单元似乎挂起，而驱动程序日志确实显示该笔记本单元已被执行。有谁知道为什么我们的笔记本单元一直挂起并且无法完成？请参阅下面的详细信息。

情况

• 我们正在训练 ML 模型pytorch在 Databricks 笔记本 UI 中
• 训练使用mlflow注册模型
• 在单元格的末尾，我们打印一条语句“完成训练”
• We are using a single node cluster with
  • Databricks 运行时：10.4 LTS ML（包括 Apache Spark 3.2.1、GPU、Scala 2.12）
  • 节点类型：Standard_NC6s_v3

观察结果

• 在 Databricks 笔记本 UI 中，我们看到单元正在运行pytorch训练并显示训练的中间日志
• 一段时间后，模型被注册在mlflow但我们在 Databricks 笔记本 UI 中看不到此日志
• 我们还可以在驱动程序日志中看到打印语句“完成训练”。我们在 Databricks 笔记本 UI 中没有看到此语句

Code

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks.early_stopping import EarlyStopping

trainer = Trainer(gpus=-1, num_sanity_val_steps=0, logger = logger, callbacks=[EarlyStopping(monitor="test_loss", patience = 2, mode = "min", verbose=True)])

with mlflow.start_run(experiment_id = experiment_id) as run:
  trainer.fit(model, train_loader, val_loader)
  mlflow.log_param("param1", param1)    
  mlflow.log_param("param2", param2)    
  mlflow.pytorch.log_model(model._image_model, artifact_path="model", registered_model_name="image_model")
  mlflow.pytorch.log_state_dict(model._image_model.state_dict(), "model")
  
print("Done with training")

Packages

mlflow-skinny==1.25.1
torch==1.10.2+cu111 
torchvision==0.11.3+cu111

我尝试过但不起作用的解决方案

• 尝试添加缓存删除，但没有成功

# Cleaning up to avoid any open processes...  
del trainer
torch.cuda.empty_cache()
# force garbage collection
gc.collect()

• 尝试强制退出笔记本，但还是不行

parameters = json.dumps({"Status": "SUCCESS", "Message": "DONE"})
dbutils.notebook.exit(parameters)

我想通了这个问题。为了解决这个问题，调整参数torch.utils.data.DataLoader

1. Disable pin_memory
2. Set num_workers占总 vCPU 的 30%（例如，Standard_NC6s_v3 为 1 或 2）

例如：

train_loader = DataLoader(
    train_dataset,
    batch_size=32,
    num_workers=1,
    pin_memory=False,
    shuffle=True,
)

这个问题似乎与 PyTorch 有关，是由于死锁问题造成的。请参阅详细信息在这里 https://pytorch.org/docs/stable/notes/multiprocessing.html#avoiding-and-fighting-deadlocks.