Конвейеры Kubeflow — обзор проблем GitHub

1. Введение

Kubeflow is a Machine Learning toolkit for Kubernetes . The project is dedicated to making deployments of Machine Learning (ML) workflows on Kubernetes simple, portable, and scalable. The goal is to provide a straightforward way to deploy best-of-breed open-source systems for ML to diverse infrastructures.

Рабочий процесс машинного обучения может включать множество взаимозависимых этапов: от подготовки и анализа данных до обучения, оценки, развертывания и многого другого. Сложно организовать и отслеживать эти процессы произвольно — например, в наборе блокнотов или скриптов — и такие вопросы, как аудит и воспроизводимость, становятся все более проблематичными. Kubeflow Pipelines (KFP) помогает решить эти проблемы, предоставляя способ развертывания надежных, воспроизводимых конвейеров машинного обучения, а также мониторинг, аудит, отслеживание версий и воспроизводимость. Cloud AI Pipelines упрощает настройку установки KFP.

Что вы построите

In this codelab, you will build a web app that summarizes GitHub issues using Kubeflow Pipelines to train and serve a model. It is based on an example in the Kubeflow Examples repo . Upon completion, your infrastructure will contain:

• Кластер Google Kubernetes Engine (GKE) с установленными Kubeflow Pipelines (через Cloud AI Pipelines ).
• Конвейер обработки данных для обучения модели Tensor2Tensor на графических процессорах.
• Контейнер для подачи блюд, отображающий прогнозы, полученные с помощью обученной модели.
• Пользовательский интерфейс, который интерпретирует прогнозы для предоставления сводных данных по задачам GitHub.
• Блокнот, создающий конвейер обработки данных с нуля с использованием SDK Kubeflow Pipelines (KFP).

Что вы узнаете

The pipeline you will build trains a Tensor2Tensor model on GitHub issue data, learning to predict issue titles from issue bodies. It then exports the trained model and deploys the exported model using Tensorflow Serving . The final step in the pipeline launches a web app, which interacts with the TF-Serving instance in order to get model predictions.

• Как установить Kubeflow Pipelines в кластере GKE
• Как создавать и запускать рабочие процессы машинного обучения с помощью конвейеров Kubernetes
• Как определять и запускать конвейеры обработки данных из блокнота платформы ИИ

Что вам понадобится

• Базовое понимание Kubernetes будет полезным, но не обязательным.
• Активный проект GCP, для которого у вас есть права владельца.
• (Необязательно) Аккаунт GitHub
• Доступ к оболочке Google Cloud Shell , доступной в консоли Google Cloud Platform (GCP).

2. Настройка

Облачная оболочка

Откройте консоль GCP в браузере и войдите в систему, используя учетные данные своего проекта:

Откройте консоль GCP.

При необходимости нажмите «Выбрать проект», чтобы начать работу с вашим проектом из Codelab.

Затем нажмите значок «Активировать Cloud Shell» в правом верхнем углу консоли, чтобы запустить Cloud Shell .

При запуске Cloud Shell отобразится название проекта, который будет использоваться. Убедитесь, что эта настройка верна.

Чтобы найти идентификатор своего проекта, перейдите на главную панель консоли GCP. Если экран пуст, нажмите «Да» в появившемся окне, чтобы создать панель мониторинга.

Затем в терминале Cloud Shell при необходимости выполните следующие команды, чтобы настроить gcloud для использования правильного проекта:

export PROJECT_ID=<your_project_id>
gcloud config set project ${PROJECT_ID}

Создайте хранилище (сумку для хранения).

Примечание: имена сегментов должны быть уникальными во всей системе GCP, а не только в вашей организации.

Create a Cloud Storage bucket for storing pipeline files. You'll need to use a globally unique ID, so it is convenient to define a bucket name that includes your project ID. Create the bucket using the gsutil mb (make bucket) command:

export PROJECT_ID=<your_project_id>
export BUCKET_NAME=kubeflow-${PROJECT_ID}
gsutil mb gs://${BUCKET_NAME}

В качестве альтернативы, вы можете создать хранилище через консоль GCP .

Необязательно**: Создайте токен GitHub**

This codelab calls the GitHub API to retrieve publicly available data. To prevent rate-limiting, especially at events where a large number of anonymized requests are sent to the GitHub APIs, set up an access token with no permissions. This is simply to authorize you as an individual rather than anonymous user.

1. Перейдите по ссылке https://github.com/settings/tokens и сгенерируйте новый токен без ограничений по области действия (scopes).
2. Сохраните его в надёжном месте. Если вы его потеряете, вам придётся удалить и создать новый.

Если вы пропустите этот шаг, лабораторная работа все равно будет работать – просто ваши возможности по генерации входных данных для тестирования модели будут несколько ограничены.

Дополнительно: Закрепить полезные панели мониторинга

В консоли GCP закрепите панели мониторинга Kubernetes Engine и Storage для более удобного доступа.

Создайте установку AI Platform Pipelines (Hosted Kubeflow Pipelines).

Follow the instructions in the 'Before you begin' and 'Set up your instance' sections here to set up a GKE instance with KFP installed. Be sure to check the Allow access to the following Cloud APIs box as indicated in the documentation. (If you don't, the example pipeline won't run successfully). Leave the installation namespace as default .

Вам потребуется выбрать зону, поддерживающую Nvidia k80. В качестве зоны по умолчанию можно использовать us-central1-a или us-central1-c .

Запомните имя кластера GKE и зону, указанные для вашей установки на панели мониторинга AI Pipelines после завершения установки, и для удобства установите переменные среды в соответствии с этими значениями.

export ZONE=<your zone>
export CLUSTER_NAME=<your cluster name>

Настройте kubectl для использования учетных данных вашего нового кластера GKE.

После создания кластера GKE настройте kubectl для использования учетных данных нового кластера, выполнив следующую команду в Cloud Shell:

gcloud container clusters get-credentials ${CLUSTER_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

Alternatively, click on the name of the cluster in the AI Pipelines dashboard to visit its GKE page, then click "Connect" at the top of the page. From the popup, paste the command into your Cloud Shell.

Это настраивает контекст kubectl , позволяя вам взаимодействовать с кластером. Для проверки конфигурации выполните следующую команду:

kubectl get nodes -o wide

Вы должны увидеть узлы со статусом " Ready ", а также другую информацию о возрасте узла, версии, внешнем IP-адресе, образе ОС, версии ядра и среде выполнения контейнера.

Настройте кластер для установки драйвера Nvidia на пулы узлов с поддержкой графических процессоров.

Далее мы применим к кластеру DaemonSet , который установит драйвер Nvidia на все узлы кластера с поддержкой графических процессоров:

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nvidia-driver-installer/cos/daemonset-preloaded.yaml

Затем выполните следующую команду, которая предоставит компонентам KFP разрешение на создание новых ресурсов Kubernetes:

kubectl create clusterrolebinding sa-admin --clusterrole=cluster-admin --serviceaccount=kubeflow:pipeline-runner

Создайте пул узлов GPU.

Затем мы настроим пул узлов GPU размером 1:

gcloud container node-pools create gpu-pool \
    --cluster=${CLUSTER_NAME} \
    --zone ${ZONE} \
    --num-nodes=1 \
    --machine-type n1-highmem-8 \
    --scopes cloud-platform --verbosity error \
    --accelerator=type=nvidia-tesla-k80,count=1

Примечание: Это может занять около 5 минут. Если при создании пула узлов вы видите ошибку квоты на использование графического процессора, вы можете запросить дополнительную квоту .

3. Запустите конвейер с панели управления «Конвейеры».

Откройте панель управления «Конвейеры».

In the Cloud Console, visit the Pipelines panel if you're not already there. Then click on " OPEN PIPELINES DASHBOARD " for your installation, and click on Pipelines in the left menu bar. If you get a load error, refresh the tab. You should see a new page like this:

Описание конвейера

Конвейер выполнения, который вы будете запускать, состоит из нескольких этапов (подробности см. в Приложении к данному практическому заданию):

1. Существующая контрольная точка модели копируется в ваш бакет.
2. Модель Tensor2Tensor обучается с использованием предварительно обработанных данных.

• Обучение начинается с существующей контрольной точки модели, скопированной на первом шаге, а затем продолжается еще несколько сотен шагов. (Полное обучение модели в рамках практического занятия заняло бы слишком много времени).
• После завершения обучения на этапе конвейера модель экспортируется в формате, подходящем для развертывания с помощью TensorFlow Serving .

3. Экземпляр TensorFlow-Serving развертывается с использованием этой модели.
4. Для взаимодействия с предоставленной моделью и получения прогнозов запускается веб-приложение.

Загрузите и скомпилируйте конвейер обработки данных.

В этом разделе мы рассмотрим, как скомпилировать определение конвейера. Первым делом нам нужно установить KFP SDK. Выполните следующие действия в Cloud Shell:

pip3 install -U kfp

Чтобы загрузить файл определения конвейера, выполните следующую команду в Cloud Shell:

curl -O https://raw.githubusercontent.com/amygdala/kubeflow-examples/ghsumm/github_issue_summarization/pipelines/example_pipelines/gh_summ_hosted_kfp.py

Затем скомпилируйте файл определения конвейера, запустив его следующим образом:

python3 gh_summ_hosted_kfp.py

В результате вы увидите файл gh_summ_hosted_kfp.py.tar.gz .

Примечание: Если вы получили ошибку, убедитесь, что у вас установлен Pipelines SDK и используется Python 3.

Загрузите скомпилированный конвейер.

In the Kubeflow Pipelines web UI, click on Upload pipeline , and select Import by URL . Copy, then paste in the following URL, which points to the same pipeline that you just compiled. (It's a few extra steps to upload a file from Cloud Shell, so we're taking a shortcut).

https://storage.googleapis.com/aju-dev-demos-codelabs/KF/compiled_pipelines/gh_summ_hosted_kfp.py.tar.gz

Присвойте конвейеру имя (например gh_summ ).

Запустите конвейер

Click on the uploaded pipeline in the list —this lets you view the pipeline's static graph— then click on Create experiment to create a new Experiment using the pipeline. An Experiment is a way to group together semantically related runs.

Присвойте эксперименту имя (например, такое же, как у конвейера, gh_summ ), затем нажмите «Далее» , чтобы создать его.

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

Для заполнения параметров вам может потребоваться выполнить следующие команды в Cloud Shell.

gcloud config get-value project
echo "gs://${BUCKET_NAME}/codelab"

Имя для запуска будет заполнено автоматически, но при желании вы можете задать другое имя.

Затем заполните три поля параметров:

• project
• (необязательный) github-token
• working-dir

For the working-dir, enter some path under the GCS bucket you created. Include the ' gs:// ' prefix. For the github-token field, enter either the token that you optionally generated earlier, or leave the placeholder string as is if you did not generate a token.

After filling in the fields, click Start , then click on the listed run to view its details. While a given pipeline step is running, you can click on it to get more information about it, including viewing its pod logs. (You can also view the logs for a pipeline step via the link to its Cloud Logging (Stackdriver) logs, even if the cluster node has been torn down).

Примечание: Завершение работы конвейера займет приблизительно 15 минут.

Просмотрите определение конвейера

Во время работы конвейера вы можете внимательнее изучить его структуру и выполняемые функции. Более подробная информация приведена в разделе «Приложение» к практическому руководству.

Просмотрите информацию об обучении модели в TensorBoard.

После завершения этапа обучения выберите вкладку «Визуализация» и нажмите синюю кнопку «Запустить TensorBoard» , а затем, когда он будет готов, нажмите «Открыть Tensorboard» .

Примечание: Если через минуту или около того вкладка TensorBoard не загружается, попробуйте перезагрузить страницу.

Изучите панель мониторинга артефактов и выполнений.

Kubeflow Pipelines automatically logs metadata about the pipeline steps as a pipeline executes. Both Artifact and Execution information is recorded. Click these entries in the left nav bar of the dashboard to explore further.

Для артефактов вы можете просмотреть как обзорную панель, так и панель «Исследователь родословной».

Запустите веб-приложение, созданное конвейером обработки данных, и сделайте несколько прогнозов.

На последнем этапе конвейера развертывается веб-приложение, которое предоставляет пользовательский интерфейс для запроса к обученной модели — предоставляемой через TF Serving — для выполнения прогнозов.

After the pipeline completes, connect to the web app by port-forwarding to its service (we're port-forwarding because, for this codelab, the webapp service is not set up to have an external endpoint).

Чтобы узнать имя сервиса, выполните следующую команду в Cloud Shell:

kubectl get services

Найдите в списке сервис с именем, похожим на: ghsumm-*-webappsvc .

Затем в Cloud Shell настройте переадресацию портов на этот сервис следующим образом, изменив следующую команду, чтобы использовать имя вашего webappsvc :

kubectl port-forward svc/ghsumm-xxxxx-webappsvc 8080:80

После запуска переадресации портов щелкните значок «предварительный просмотр» над панелью Cloud Shell, а в раскрывающемся списке выберите «Предварительный просмотр на порту 8080».

В новой вкладке должна открыться страница, похожая на эту:

Нажмите кнопку «Заполнить случайный выпуск» , чтобы получить блок текста. Нажмите кнопку «Сгенерировать заголовок» , чтобы вызвать обученную модель и отобразить прогноз.

Примечание: отображение сводки может занять несколько секунд, особенно при первом запросе — в рамках этого семинара мы не используем графические процессоры для экземпляра TensorFlow Serving.

If your pipeline parameters included a valid GitHub token, you can alternately try entering a GitHub URL in the second field, then clicking "Generate Title". If you did not set up a valid GitHub token, use only the "Populate Random Issue" field.

4. Запустите конвейер обработки данных из блокнота платформы ИИ.

You can also interactively define and run Kubeflow Pipelines from a Jupyter notebook using the KFP SDK. AI Platform Notebooks , which we'll use for this codelab, makes this very straightforward.

Создайте экземпляр блокнота.

Мы создадим экземпляр блокнота из Cloud Shell, используя его API. (В качестве альтернативы вы можете создать блокнот через Cloud Console . Дополнительную информацию см. в документации).

В оболочке Cloud Shell установите следующие переменные среды:

export INSTANCE_NAME="kfp-ghsumm"
export VM_IMAGE_PROJECT="deeplearning-platform-release"
export VM_IMAGE_FAMILY="tf2-2-3-cpu"
export MACHINE_TYPE="n1-standard-4"
export LOCATION="us-central1-c"

Затем в оболочке Cloud Shell выполните команду для создания экземпляра блокнота:

gcloud beta notebooks instances create $INSTANCE_NAME \
  --vm-image-project=$VM_IMAGE_PROJECT \
  --vm-image-family=$VM_IMAGE_FAMILY \
  --machine-type=$MACHINE_TYPE --location=$LOCATION

При первом запуске этой команды вас могут попросить включить API notebooks для вашего проекта. В этом случае ответьте « y ».

Через несколько минут ваш сервер Notebook будет запущен и начнет работу. Вы сможете увидеть список своих экземпляров Notebook в консоли Cloud .

Загрузите блокнот Codelab.

После создания экземпляра блокнота нажмите эту ссылку , чтобы загрузить блокнот Jupyter из учебного модуля. Выберите нужный экземпляр блокнота. Блокнот откроется автоматически.

Запустите блокнот

Follow the instructions in the notebook for the remainder of the lab. Note that in the "Setup" part of the notebook, you will need to fill in your own values before running the rest of the notebook.

(Если вы используете собственный проект, не забудьте вернуться и выполнить раздел «Уборка» этой лабораторной работы).

5. Уборка

Если вы используете временную учетную запись Codelab, вам это делать не нужно, но если вы используете собственный проект, возможно, вам стоит удалить установку Pipelines и Notebook.

Вывести из строя кластер Pipelines GKE

Вы можете удалить кластер Pipelines из Cloud Console . (Если вы хотите повторно использовать кластер GKE, у вас есть возможность удалить только установку Pipelines).

Удалите экземпляр AI Notebook.

Если вы выполняли часть практического задания, посвященную "блокноту", вы можете удалить или остановить экземпляр блокнота из облачной консоли .

(Необязательно) Удалите токен GitHub

Перейдите по ссылке https://github.com/settings/tokens и удалите сгенерированный токен.

6. Приложения

Взгляд на код

Определение конвейера

Конвейер обработки данных, используемый в этом практическом занятии, описан здесь .

Давайте рассмотрим, как это определяется, а также как определяются его компоненты (шаги). Мы затронем некоторые основные моменты, но для получения более подробной информации обратитесь к документации .

Note: While not illustrated in this example, you may also be interested to explore KFP lightweight python components . Lightweight python components do not require you to build a new container image for every code change, and are helpful for fast iteration and prototyping.

Kubeflow Pipeline steps are container-based. When you're building a pipeline, you can use pre-built components , with already-built container images, or build your own components. For this codelab, we've built our own.

Four of the pipeline steps are defined from reusable components , accessed via their component definition files . In this first code snippet, we're accessing these component definition files via their URL, and using these definitions to create 'ops' that we'll use to create a pipeline step.

import kfp.dsl as dsl
import kfp.gcp as gcp
import kfp.components as comp

...

copydata_op = comp.load_component_from_url(
  'https://raw.githubusercontent.com/kubeflow/examples/master/github_issue_summarization/pipelines/components/t2t/datacopy_component.yaml'
  )

train_op = comp.load_component_from_url(
  'https://raw.githubusercontent.com/kubeflow/examples/master/github_issue_summarization/pipelines/components/t2t/train_component.yaml'
  )

Ниже приведено одно из определений компонента для операции обучения в формате YAML. Вы можете видеть, что определены его входные и выходные данные, образ контейнера и аргументы точки входа контейнера.

name: Train T2T model
description: |
  A Kubeflow Pipeline component to train a Tensor2Tensor
  model
metadata:
  labels:
    add-pod-env: 'true'
inputs:
  - name: train_steps
    description: '...'
    type: Integer
    default: 2019300
  - name: data_dir
    description: '...'
    type: GCSPath
  - name: model_dir
    description: '...'
    type: GCSPath
  - name: action
    description: '...'
    type: String
  - name: deploy_webapp
    description: '...'
    type: String
outputs:
  - name: launch_server
    description: '...'
    type: String
  - name: train_output_path
    description: '...'
    type: GCSPath
  - name: MLPipeline UI metadata
    type: UI metadata
implementation:
  container:
    image: gcr.io/google-samples/ml-pipeline-t2ttrain:v3ap
    args: [
      --data-dir, {inputValue: data_dir},
      --action, {inputValue: action},
      --model-dir, {inputValue: model_dir},
      --train-steps, {inputValue: train_steps},
      --deploy-webapp, {inputValue: deploy_webapp},
      --train-output-path, {outputPath: train_output_path}
    ]
    env:
      KFP_POD_NAME: "{{pod.name}}"
    fileOutputs:
      launch_server: /tmp/output
      MLPipeline UI metadata: /mlpipeline-ui-metadata.json

Также можно определить шаг конвейера с помощью конструктора dsl.ContainerOp , как мы увидим ниже.

Below is the bulk of the pipeline definition. We're defining the pipeline inputs (and their default values). Then we define the pipeline steps. For most we're using the 'ops' defined above, but we're also defining a 'serve' step inline via ContainerOp , specifying the container image and entrypoint arguments directly.

Как видите, этапы train , log_model и serve используют выходные данные предыдущих этапов в качестве входных. Подробнее о том, как это задается, можно прочитать здесь .

@dsl.pipeline(
 name='Github issue summarization',
 description='Demonstrate Tensor2Tensor-based training and TF-Serving'
)
def gh_summ(  #pylint: disable=unused-argument
 train_steps: 'Integer' = 2019300,
 project: str = 'YOUR_PROJECT_HERE',
 github_token: str = 'YOUR_GITHUB_TOKEN_HERE',
 working_dir: 'GCSPath' = 'gs://YOUR_GCS_DIR_HERE',
 checkpoint_dir: 'GCSPath' = 'gs://aju-dev-demos-codelabs/kubecon/model_output_tbase.bak2019000/',
 deploy_webapp: str = 'true',
 data_dir: 'GCSPath' = 'gs://aju-dev-demos-codelabs/kubecon/t2t_data_gh_all/'
 ):


 copydata = copydata_op(
   data_dir=data_dir,
   checkpoint_dir=checkpoint_dir,
   model_dir='%s/%s/model_output' % (working_dir, dsl.RUN_ID_PLACEHOLDER),
   action=COPY_ACTION,
   )


 train = train_op(
   data_dir=data_dir,
   model_dir=copydata.outputs['copy_output_path'],
   action=TRAIN_ACTION, train_steps=train_steps,
   deploy_webapp=deploy_webapp
   )

 serve = dsl.ContainerOp(
     name='serve',
     image='gcr.io/google-samples/ml-pipeline-kubeflow-tfserve:v6',
     arguments=["--model_name", 'ghsumm-%s' % (dsl.RUN_ID_PLACEHOLDER,),
         "--model_path", train.outputs['train_output_path']
         ]
     )

 train.set_gpu_limit(1)

Обратите внимание, что для выполнения этапа «обучение» требуется узел в кластере, имеющий как минимум 1 доступный графический процессор.

  train.set_gpu_limit(1)

The final step in the pipeline— also defined inline— is conditional. It will run after the ' serve ' step is finished, only if the training step launch_server output is the string 'true'. It launches the 'prediction web app', that we used to request issue summaries from the trained T2T model.

 with dsl.Condition(train.outputs['launch_server'] == 'true'):
   webapp = dsl.ContainerOp(
       name='webapp',
       image='gcr.io/google-samples/ml-pipeline-webapp-launcher:v1',
       arguments=["--model_name", 'ghsumm-%s' % (dsl.RUN_ID_PLACEHOLDER,),
           "--github_token", github_token]

       )
   webapp.after(serve)

Определения образов контейнеров компонентов

The Kubeflow Pipeline documentation describes some best practices for building your own components. As part of this process, you will need to define and build a container image. You can see the component steps for this codelab's pipeline here . The Dockerfile definitions are in the containers subdirectories, eg here .

Используйте виртуальные машины с возможностью вытеснения с помощью графических процессоров для обучения.

Preemptible VMs are Compute Engine VM instances that last a maximum of 24 hours and provide no availability guarantees. The pricing of preemptible VMs is lower than that of standard Compute Engine VMs.

With Google Kubernetes Engine (GKE) , it is easy to set up a cluster or node pool that uses preemptible VMs . You can set up such a node pool with GPUs attached to the preemptible instances . These work the same as regular GPU-enabled nodes, but the GPUs persist only for the life of the instance.

You can set up a preemptible, GPU-enabled node pool for your cluster by running a command similar to the following, editing the following command with your cluster name and zone, and adjusting the accelerator type and count according to your requirements. You can optionally define the node pool to autoscale based on current workloads.

gcloud container node-pools create preemptible-gpu-pool \
    --cluster=<your-cluster-name> \
    --zone <your-cluster-zone> \
    --enable-autoscaling --max-nodes=4 --min-nodes=0 \
    --machine-type n1-highmem-8 \
    --preemptible \
    --node-taints=preemptible=true:NoSchedule \
    --scopes cloud-platform --verbosity error \
    --accelerator=type=nvidia-tesla-k80,count=4

Вы также можете настроить пул узлов через консоль Cloud Console .

Определение конвейера Kubeflow, использующего вытесняемые узлы GKE.

Если вы используете Kubeflow в GKE, теперь легко определять и запускать конвейеры Kubeflow, в которых один или несколько шагов (компонентов) конвейера выполняются на вытесняемых узлах , что снижает стоимость выполнения задания. Для корректного использования вытесняемых виртуальных машин шаги, которые вы определяете как вытесняемые, должны быть либо идемпотентными (то есть, если вы запустите шаг несколько раз, он даст тот же результат), либо должны иметь функцию контрольных точек, позволяющую шагу продолжить выполнение с того места, где он был прерван.

При определении конвейера Kubeflow можно указать, что данный шаг должен выполняться на узле, который можно вытеснить, изменив операцию следующим образом:

your_pipelines_op.apply(gcp.use_preemptible_nodepool())

Подробности см. в документации .

Вероятно, вам также потребуется повторить этот шаг несколько раз, если узел будет вытеснен. Это можно сделать следующим образом — здесь мы указываем 5 повторных попыток.

your_pipelines_op.set_gpu_limit(1).apply(gcp.use_preemptible_nodepool()).set_retry(5)

Попробуйте отредактировать конвейер Kubeflow, который мы использовали в этом практическом занятии, чтобы запустить этап обучения на виртуальной машине с возможностью прерывания .

Измените следующую строку в спецификации конвейера, чтобы дополнительно использовать пул узлов с возможностью прерывания (убедитесь, что вы создали его, как указано выше) и повторить попытку 5 раз:

  train.set_gpu_limit(1)

Затем перекомпилируйте конвейер, загрузите новую версию (присвойте ей новое имя) и запустите новую версию конвейера.