挂有火把的数据库笔记本

Question

我们有一个数据库笔记本的问题。我们的笔记本细胞中似乎有一个挂着，而驱动程序日志显示笔记本细胞已经被执行。有人知道为什么我们的笔记本电脑一直挂着，而且不完整吗？详见下文。

Situation

• 在Databricks笔记本用户界面中，我们正在训练一个带有py手电筒的ML模型。
• 培训使用mlflow注册模型。
• 在单元格的末尾，我们打印一条“完成训练”的声明。
• 我们使用的是带有的单节点集群。
  • Databricks运行时:10.4LTS ML (包括ApacheSpark3.2.1、GPU、Scala2.12)
  • 节点类型: Standard_NC6s_v3

Observations

• 在Databricks笔记本用户界面中，我们看到了单元格运行pytorch训练和显示训练的中间日志。
• 过了一会儿，模型在mlflow中注册了，但是我们在Databricks笔记本用户界面中没有看到这个日志
• 我们还可以在驱动程序日志中看到打印语句“使用培训完成”。我们在Databricks笔记本用户界面中没有看到这个语句

码

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")

包

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)

Answer 1

我解决了这个问题。要解决这个问题，请调整torch.utils.data.DataLoader的参数

1. 禁用pin_memory
2. 将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有关，是由死锁问题造成的。见这里的细节。

Answer 2

https://stackoverflow.com/a/72473053/10555941

在上面的答案中还有更多的补充。当您设置pin_memory True并使num_workers等于节点上的vCpus总数时，它将使用IPC进行线程通信。这些IPC使用共享内存，并窒息VM的共享内存。

这就导致了过程的悬吊。DataLoader num_workers只是为了帮助使用子线程加载数据。尽管如此，它不一定是一个极值，以加快数据加载。如果它小于vCPUS的30%，就足以加载数据。