Dask：任务停滞？

Question

我有一个 dask 批量推理工作负载，它对 50 万行特征应用机器学习预测。我在 AWS 上的 EMR 实例上使用 dask-yarn。似乎有一些掉队的人阻碍了达斯克，我不知道该如何对付他们。对于分配任务，它一直卡在99531/100000，我一直在等待。当我使用 Dask 时，有时它会“卡在”更大的数据集上。

代码示例：

def model_predict(row, model_config, bert_model, bert_tokenizer, xgb_model):

    input_doc = row.to_dict()
    try:
        return predict_one(model_config, bert_model, bert_tokenizer, xgb_model, input_doc)
    except:
        return {'error_code': 500}


all_features_df = dd.read_csv(file)
all_features_df = all_features_df.repartition(npartitions=100000).persist()
all_features_df['prediction_response'] = all_features_df.apply(
        model_predict,
        model_config=base_config,
        bert_model=bert,
        bert_tokenizer=tokenizer,
        xgb_model=xgb,
        axis=1, 
        meta=('prediction_response', 'object'))
all_features_df = all_features_df.persist()
all_features_df.repartition(partition_size="100MB").to_parquet(f's3://my_data/predictions/batch={batch}/',
                                                       engine='pyarrow',
                                                       schema='infer')