Skip to main content

Using GPUs with IDEs

Conveyor IDEs run on nodes that do not have GPUs. If you want to run jobs that require GPUs, you should use the Conveyor Python SDK to create a job that runs on a GPU node. The Python SDK allows you to specify the instance type and the number of GPUs you want to use. The steps needed to run a job that uses a GPU from IDEs are as follows:

  • Create a Python/Spark job that uses GPUs for processing. In this howto guide we will use the mnist example from PyTorch.
  • Prepare a Dockerfile that installs the necessary dependencies and packages the code for your job. If you need help with writing a Dockerfile with GPU support, take a look at working-with-gpus.
  • Make sure you have the Conveyor Python SDK installed in your IDE, see Python SDK for more information.
  • Trigger the Python/Spark job using the Conveyor Python SDK using a new Jupyter notebook file (ipynb extension). You can use the ContainerTaskRunner or SparkTaskRunner to run your job on a GPU instance. When using the ContainerTaskRunner, the code looks as follows:
from conveyor.project import ProjectBuilder

# Build the project and push the Docker image to the Conveyor registry.
build_id = ProjectBuilder(project_path="<relative-path-project-root>").build()

from conveyor.runner import ContainerTaskRunner
from conveyor.types import InstanceType

# Trigger the job using the ContainerTaskRunner.
with ContainerTaskRunner(
    task_name="<some task name>",
    project_name="<project-name>",
    environment_name=<some-env-name>,
    build_id=build_id, # refers to the build_id of the project we just built
    instance_type=InstanceType.g4dn_xlarge, # specify an instance type with GPU support
    args=["-m", "samplepytorch.mnist"],
    iam_identity=f"<some-iam-identity>"
).run() as runner:
    print(f"The run has {'failed' if runner.has_failed() else 'succeeded'}")

Note: since GPU images are often large, pushing the container image, as well as starting your job quickly takes a couple of minutes.