有没有办法在不轮询 REST API 的情况下通知 Google AI Platform 训练作业的状态变化？

Posted 2023-03-25

【中文标题】有没有办法在不轮询 REST API 的情况下通知 Google AI Platform 训练作业的状态变化？

【英文标题】：Is there a way to be notified of status changes in Google AI Platform training jobs without polling the REST API?

【发布时间】：2020-01-24 08:39:15

【问题描述】：

现在，我通过轮询 job REST API 来监控我在 Google AI Platform（以前称为 ml 引擎）上提交的作业。我不喜欢这个解决方案有几个原因：

如果状态更改之间的间隔小于监控轮询率，则通常会延迟或完全错过对状态更改的感知 大量不必要的网络流量 大量不必要的函数调用

我希望在我的培训工作完成后立即收到通知。如果有某种方法可以分配钩子或回调以在作业状态更改时运行。

我还考虑直接在 AI Platform 上运行的训练任务 python 包中添加对云函数的调用。但是，我认为这些函数调用不会发生在训练作业意外关闭的情况下，例如当作业被 GCP 取消或强制结束时。

有没有更好的方法来解决这个问题？

【参考方案1】：

您可以使用 Stackdriver 接收器来读取日志并将其发送到 Pub/Sub。通过 Pub/Sub，您可以连接到许多其他提供商：

1。设置 Pub/Sub 接收器

在开始之前，请确保您拥有 access to the logs 和 publish rights to the topic you desire。按照说明设置 Stackdriver -> Pub/Sub 接收器。您将希望使用此查询将事件限制为仅培训作业：

resource.type = "ml_job"
resource.labels.task_name = "service"

请注意Stackdriver can further limit down the query。例如，您可以通过添加像 resource.labels.job_id = "..." 这样的条件或使用像 jsonPayload.message : "..." 这样的过滤器来限制特定事件。

2。回复 Pub/Sub 消息

为了知道发生了什么变化，Pub/Sub 消息的接收者可以从 ml.googleapis.com API 查询作业状态或阅读消息的文本

从 ml.googleapis.com 读取状态

收到消息后，调用https://ml.googleapis.com/v1/<project_id>/jobs/<job_id>获取Job信息，将URL中的[project_id]和[job_id]替换为resource.label.project_id和resource.label.job_id的值分别发布/订阅消息。

返回的 Job 对象包含一个字段 state，它自然地表明了作业的状态。

从消息文本中读取状态

Pub/Sub 消息将包含一个字符串，说明作业发生了什么。您可能希望工作结束时的行为。在 jsonPayload.message 中查找这些字符串：

“作业成功完成。” “作业已取消。” “作业失败。”

【讨论】：

感谢详细的解决方案。阅读日志并发布到 pub/sub 肯定是解决此问题的一种方法。我将把它理解为 ml-engine 没有可以由状态更改触发的内置钩子。

@EdmundRosewright 你有没有想出任何其他解决方案，比如让作业本身在完成时发布一条 pubsub 消息？我正在尝试做完全相同的事情，但还没有开始工作。

到目前为止，我认为@htappen 的解决方案是我们现在能做的最好的。它对我来说效果很好（尽管实施起来很痛苦）。我以前试过你的建议。不幸的是，实际作业中的任何代码都不能保证在作业取消或突然崩溃的情况下运行。最重要的是，它只允许我监控我的应用程序的内部状态，而不是实际的 AI Platform job state

【参考方案2】：

正如@htappen 所说，我实现了一个 Terraform 模块。如果对你有帮助，我很高兴。但我真正希望 Google 更新 AI Platform 以提供相同的功能。

https://github.com/sfujiwara/terraform-google-ai-platform-notification

【参考方案3】：

我认为您可以在培训作业代码的末尾以编程方式发布 PubSub 消息。像这样的：

from google.cloud import pubsub_v1

# publish job complete message
client = pubsub_v1.PublisherClient()
topic = client.topic_path(args.gcp_project_id, 'topic-name')
data = 
    'ACTION': 'JOB_COMPLETE',
    'SAVED_MODEL_DIR': args.job_dir

data_bytes = json.dumps(data).encode('utf-8')
client.publish(topic, data_bytes)

然后您可以设置一个云功能以由相同的 pubsub 主题触发。

【讨论】：

我考虑过这样做，但不幸的是，它只有在训练代码能够一直执行到最后的情况下才有效。训练工作突然结束的方法有很多。例如，如果您取消作业、崩溃，或者作业被 gcp 中断/抢占。在这些情况下，pubsub 调用可能永远无法执行。此外，这只会让我监控应用的内部状态，而不是实际的 AI Platform 作业状态。

@EdmundRosewright。这就说得通了。我想您可以将您的培训代码包装在 try/except 中，但这并不能提供工作状态的细粒度通知。出于这个原因，我希望 AI Platform Training 能够自动将作业状态事件发布到预定义的 PubSub 主题。

是的，我同意如果他们自动将作业状态事件发布到 PubSub 主题，那将是一个巨大的改进。并将整个训练代码包装在 try/except 中有助于减少由于程序内部错误而导致的崩溃。但它对外部崩溃没有帮助，例如内存不足崩溃或作业取消。我还注意到，在我们的 AI Platform Job 开始和我们的实际训练代码开始运行之间以及我们的代码完成执行和 Job 实际标记为完成之间可能会有很长的延迟。

【参考方案4】：

您可以通过将 LamdbaCallback 添加到 fit() 调用来解决自定义 TF 训练作业中缺少来自服务的回调的问题。在 on_epoch 方法中，您可以在工作完成时向自己发送有关工作进度和 on_train_end 的通知。

https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/LambdaCallback

【讨论】：

不幸的是，由于大卫布朗解决方案的评论线程中讨论的原因，此解决方案不起作用。