[core][experimental] Add multi-GPU CI tests for accelerated DAG

.buildkite/core.rayci.yml

@@ -327,3 +327,17 @@ steps:
  - forge
  - raycpubase
  - corebuild
+
+ - label: ":ray: core: multi gpu tests"
+ tags:
+ - accelerated_dag
+ - gpu
+ instance_type: gpu-large
+ commands:
+ # This machine has 4 GPUs, and we need 2 GPUs, so allow 2 tests to run in
+ # parallel.
+ - bazel run //ci/ray_ci:test_in_docker -- //python/ray/tests/... //python/ray/dag/... core
+ --parallelism-per-worker 2 --gpus 2
+ --build-name coregpubuild
+ --only-tags multi_gpu
+ depends_on: coregpubuild

ci/docker/core.build.Dockerfile

@@ -18,11 +18,7 @@ RUN <<EOF
 
 set -euo pipefail
 
-pip install -U --ignore-installed \
- -c python/requirements_compiled.txt \
- -r python/requirements.txt \
- -r python/requirements/test-requirements.txt \
- -r python/requirements/ml/dl-cpu-requirements.txt
+DL=1 ./ci/env/install-dependencies.sh
 
 if [[ "$RAYCI_IS_GPU_BUILD" == "true" ]]; then
  pip install -Ur ./python/requirements/ml/dl-gpu-requirements.txt

ci/docker/core.build.wanda.yaml

@@ -2,11 +2,14 @@ name: "corebuild-py$PYTHON"
 froms: ["cr.ray.io/rayproject/oss-ci-base_build-py$PYTHON"]
 dockerfile: ci/docker/core.build.Dockerfile
 srcs:
+ - ci/env/install-dependencies.sh
  - python/requirements.txt
  - python/requirements_compiled.txt
  - python/requirements/test-requirements.txt
  - python/requirements/ml/dl-cpu-requirements.txt
+ - python/requirements/ml/dl-gpu-requirements.txt
 build_args:
  - DOCKER_IMAGE_BASE_BUILD=cr.ray.io/rayproject/oss-ci-base_build-py$PYTHON
+ - RAYCI_IS_GPU_BUILD
 tags:
  - cr.ray.io/rayproject/corebuild-py$PYTHON

ci/pipeline/determine_tests_to_run.py

@@ -108,6 +108,7 @@ def get_commit_range():
  RAY_CI_WORKFLOW_AFFECTED = 0
  RAY_CI_RELEASE_TESTS_AFFECTED = 0
  RAY_CI_COMPILED_PYTHON_AFFECTED = 0
+ RAY_CI_ACCELERATED_DAG_AFFECTED = 0
 
  if is_pull_request():
  commit_range = get_commit_range()
@@ -258,6 +259,14 @@ def get_commit_range():
  if changed_file.endswith(compiled_extension):
  RAY_CI_COMPILED_PYTHON_AFFECTED = 1
  break
+
+ # Some accelerated DAG tests require GPUs so we only run them
+ # if Ray DAGs or experimental.channels were affected.
+ if changed_file.startswith("python/ray/dag") or changed_file.startswith(
+ "python/ray/experimental/channel"
+ ):
+ RAY_CI_ACCELERATED_DAG_AFFECTED = 1
+
  elif changed_file == ".buildkite/core.rayci.yml":
  RAY_CI_PYTHON_AFFECTED = 1
  RAY_CI_CORE_CPP_AFFECTED = 1
@@ -377,6 +386,7 @@ def get_commit_range():
  RAY_CI_MACOS_WHEELS_AFFECTED = 1
  RAY_CI_DASHBOARD_AFFECTED = 1
  RAY_CI_RELEASE_TESTS_AFFECTED = 1
+ RAY_CI_ACCELERATED_DAG_AFFECTED = 1
  else:
  print(
  "Unhandled source code change: {changed_file}".format(
@@ -458,6 +468,9 @@ def get_commit_range():
  "RAY_CI_COMPILED_PYTHON_AFFECTED={}".format(
  RAY_CI_COMPILED_PYTHON_AFFECTED
  ),
+ "RAY_CI_ACCELERATED_DAG_AFFECTED={}".format(
+ RAY_CI_ACCELERATED_DAG_AFFECTED
+ ),
  ]
  )
 

python/ray/dag/BUILD

@@ -2,7 +2,7 @@ load("//bazel:python.bzl", "doctest")
 load("//bazel:python.bzl", "py_test_module_list")
 
 doctest(
- files = glob(["**/*.py"]),
+ files = glob(["**/*.py"], exclude=["**/experimental/**/*.py"]),
  tags = ["team:core"],
  deps = [":dag_lib"]
 )

python/ray/experimental/channel/torch_tensor_type.py

@@ -31,9 +31,18 @@ def __init__(
  def register_custom_serializer(outer: Any) -> None:
  # Helper method to run on the DAG driver and actors to register custom
  # serializers.
+ import torch
+
  from ray.air._internal import torch_utils
 
- default_device = torch_utils.get_devices()[0]
+ if ray.get_gpu_ids():
+ default_device = torch_utils.get_devices()[0]
+ else:
+ # torch_utils defaults to returning GPU 0 if no
+ # GPU IDs were assigned by Ray. We instead want
+ # the default to be CPU.
+ default_device = torch.device("cpu")
+
  torch_tensor_serializer = _TorchTensorSerializer(default_device)
 
  CUSTOM_SERIALIZERS = (