Укажите тип машины для одного компонента конвейера TFX в Vertex AI

Question

Я использую TFX для создания AI Pipeline на Vertex AI. Я следовал этому руководству, чтобы начать работу, а затем адаптировал конвейер для моих собственных данных, которые содержат более 100 миллионов строк данных временных рядов. Несколько моих компонентов были убиты на полпути из-за проблем с памятью, поэтому я хотел бы установить требования к памяти только для этих компонентов. я использую KubeflowV2DagRunner для оркестровки и запуска конвейера в Vertex AI с помощью следующего кода:

      runner = tfx.orchestration.experimental.KubeflowV2DagRunner(
    config=tfx.orchestration.experimental.KubeflowV2DagRunnerConfig(
        default_image = 'gcr.io/watch-hop/hop-tfx-covid:0.6.2'
    ),
    output_filename=PIPELINE_DEFINITION_FILE)

_ = runner.run(
    create_pipeline(
        pipeline_name=PIPELINE_NAME,
        pipeline_root=PIPELINE_ROOT,
        data_path=DATA_ROOT, metadata_path=METADATA_PATH))

На аналогичный вопрос был дан ответ на Stack Overflow, что привело меня к способу к установки требованийпамяти в платформе AI , но эти конфигурации больше не существуют в KubeflowV2DagRunnerConfig, так что я в тупике.

Любая помощь приветствуется.

Answer 1

Оказывается, на данный момент вы не можете этого сделать, но в соответствии с этой проблемой, эта функция появится.

Альтернативное решение - преобразовать конвейер TFX в конвейер Kubeflow. Конвейеры Vertex AI поддерживают kubeflow, и с их помощью вы можете устанавливать ограничения памяти и ЦП на уровне компонентов.

      @component // imported from kfp.dsl
def MyComponent(Input[Dataset] input_data):
  // ...

@pipeline // imported from kfp.dsl
def MyPipeline(...):
  component = MyComponent(...)
  component.set_memory_request('64G')

Answer 2

Альтернативным вариантом этого решения может быть использование средства управления потоком данных, которое позволяет запускать компоненты кластера потока данных через Vertex. Мне все еще нужно найти способ указать типы машин для пользовательских компонентов.

Пример ввода луча:

      BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS = [
--project=  GOOGLE_CLOUD_PROJECT,
--temp_location= GCS_LOCAITON,
--runner=DataflowRunner

]

К настоящему времени вы бы перешли на Vertex AI.