Add some PyTorch CUDA tests.

This also adds a repro case for issue #9827, although it is commented out
for now since it doesn't work yet.

PiperOrigin-RevId: 595238934

Makefile

@@ -293,17 +293,19 @@ cos-gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
 # This is a superset of those needed for smoke tests.
 # It includes non-GPU images that are used as part of GPU tests,
 # e.g. busybox and python.
-gpu-images: gpu-smoke-images load-gpu_ollama load-basic_busybox load-basic_python
+gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-basic_busybox load-basic_python
 .PHONY: gpu-images
 
 gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true --nvproxy-docker=true)
+	@$(call sudo,test/gpu:pytorch_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-all-tests
 
 cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true)
+	@$(call sudo,test/gpu:pytorch_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-all-tests

images/gpu/pytorch/Dockerfile

@@ -0,0 +1,40 @@
+FROM nvidia/cuda:12.2.0-devel-ubuntu20.04
+
+RUN apt-get update && apt-get install --yes \
+      python3 \
+      python3-distutils \
+      python3-pip \
+      clang \
+      wget \
+      vim \
+      git
+
+RUN python3 -m pip install --ignore-installed \
+      "clang~=$(clang --version | grep -oP '10\.[^-]+')" \
+      torch \
+      torchvision \
+      lightning \
+      numpy \
+      memory_profiler
+
+ENV PYTORCH_DATASETS_DIR=/pytorch-data
+ENV TORCH_HOME=/pytorch-home
+COPY download_pytorch_datasets.py /tmp/
+# Some PyTorch examples hardcode the data directory to "data", so
+# make a symlink for that too.
+RUN mkdir "$PYTORCH_DATASETS_DIR" && \
+    python3 /tmp/download_pytorch_datasets.py && \
+    rm /tmp/download_pytorch_datasets.py
+
+RUN PYTORCH_EXAMPLES_COMMIT=30b310a977a82dbfc3d8e4a820f3b14d876d3bd2 && \
+    mkdir /pytorch-examples && \
+    cd /pytorch-examples && \
+    git init && \
+    git remote add origin https://github.com/pytorch/examples && \
+    git fetch --depth 1 origin "$PYTORCH_EXAMPLES_COMMIT" && \
+    git checkout FETCH_HEAD && \
+    sed -ri "s~(datasets.*)\\(['\"](../)?data['\"],~\\1('$PYTORCH_DATASETS_DIR',~g" **/*.py && \
+    sed -ri 's/download=True/download=False/' **/*.py
+
+COPY *.py /
+RUN rm /download_pytorch_datasets.py

images/gpu/pytorch/download_pytorch_datasets.py

@@ -0,0 +1,31 @@
+# Copyright 2023 The gVisor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Download PyTorch datasets used in tests."""
+
+import os
+
+from torchvision import datasets
+from torchvision import models
+
+datasets_dir = os.environ["PYTORCH_DATASETS_DIR"]
+for dataset in (
+    datasets.MNIST,
+    datasets.CIFAR100,
+):
+  dataset(datasets_dir, train=True, download=True)
+  dataset(datasets_dir, train=False, download=True)
+
+# Download resnet50 weights to TORCH_HOME:
+models.resnet50(weights=models.ResNet50_Weights.DEFAULT)

images/gpu/pytorch/is_cuda_available.py

@@ -0,0 +1,25 @@
+# Copyright 2023 The gVisor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Checks if CUDA is recognized as available by PyTorch."""
+
+import sys
+
+import torch
+
+if not torch.cuda.is_available():
+  print("CUDA is not available.", file=sys.stderr)
+  sys.exit(1)
+
+print("CUDA is available.", file=sys.stderr)

images/gpu/pytorch/issue_9827.py

@@ -0,0 +1,94 @@
+# Copyright 2023 The gVisor Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Reproduction case for https://github.com/google/gvisor/issues/9827."""
+
+import os
+import time
+
+import lightning as L
+import psutil
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import models
+from torchvision import transforms
+from torchvision.datasets import CIFAR100
+
+current_process = psutil.Process()
+parent_process = current_process.parent()
+print(f"Processes: {current_process=} {parent_process=}")
+
+
+class NeuralNet(L.LightningModule):
+  """NeuralNet is the neural network used in this test."""
+
+  def __init__(self, nbr_cat):
+    super().__init__()
+
+    module = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+    module.fc = nn.Linear(2048, nbr_cat)
+
+    self.module = module
+
+  def forward(self, x):
+    return self.module(x)
+
+  def training_step(self, batch, batch_idx):
+    x, y = batch
+    y_hat = self(x)
+    loss = F.cross_entropy(y_hat, y)
+    return loss
+
+  def configure_optimizers(self):
+    return torch.optim.Adam(self.parameters(), lr=0.02)
+
+
+def prepare_data():
+  """prepare_data prepares the data to feed to the training pipeline."""
+  pipeline = transforms.Compose([
+      transforms.Resize((224, 224)),
+      transforms.ToTensor(),
+  ])
+
+  train_ds = CIFAR100(os.environ["PYTORCH_DATASETS_DIR"],
+                      train=True,
+                      download=False,
+                      transform=pipeline)
+  train_dataloader = DataLoader(train_ds, batch_size=128, num_workers=4)
+
+  val_ds = CIFAR100(os.environ["PYTORCH_DATASETS_DIR"],
+                    train=False,
+                    download=False,
+                    transform=pipeline)
+  val_dataloader = DataLoader(val_ds, batch_size=128, num_workers=4)
+
+  return train_dataloader, val_dataloader
+
+
+torch.set_float32_matmul_precision("medium")
+train_dl, val_dl = prepare_data()
+model = NeuralNet(100)
+trainer = L.Trainer(max_epochs=1, strategy="ddp_notebook")
+
+start = time.time()
+# TODO(gvisor.dev/issue/9827): Make this not take forever.
+trainer.fit(model, train_dl, val_dl)
+time.sleep(20)
+end = time.time()
+
+training_duration = end - start
+
+print(f"Training duration (seconds): {training_duration}")

test/gpu/BUILD

@@ -17,6 +17,18 @@ go_test(
     deps = ["//pkg/test/dockerutil"],
 )
 
+go_test(
+    name = "pytorch_test",
+    srcs = ["pytorch_test.go"],
+    tags = [
+        "local",
+        "noguitar",
+        "notap",
+    ],
+    visibility = ["//:sandbox"],
+    deps = ["//pkg/test/dockerutil"],
+)
+
 go_test(
     name = "textgen_test",
     srcs = ["textgen_test.go"],

test/gpu/pytorch_test.go

@@ -0,0 +1,61 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package pytorch_test tests basic PyTorch workloads.
+package pytorch_test
+
+import (
+	"context"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+)
+
+// runPytorch runs the given script and command in a PyTorch container.
+func runPytorch(ctx context.Context, t *testing.T, scriptPath string, args ...string) {
+	t.Helper()
+	c := dockerutil.MakeContainer(ctx, t)
+	opts := dockerutil.GPURunOpts()
+	opts.Image = "gpu/pytorch"
+	cmd := append([]string{"python3", scriptPath}, args...)
+	out, err := c.Run(ctx, opts, cmd...)
+	if err != nil {
+		t.Errorf("Failed: %v\nContainer output:\n%s", err, out)
+	} else {
+		t.Logf("Container output:\n%s", out)
+	}
+}
+
+// TestCUDAIsAvailable checks that PyTorch recognizes that CUDA is available.
+func TestCUDAIsAvailable(t *testing.T) {
+	runPytorch(context.Background(), t, "/is_cuda_available.py")
+}
+
+// TestLinearRegressionModel runs a simple linear regression model.
+func TestLinearRegressionModel(t *testing.T) {
+	runPytorch(context.Background(), t, "/pytorch-examples/regression/main.py", "--cuda")
+}
+
+// TestMNIST runs an MNIST model.
+func TestMNIST(t *testing.T) {
+	runPytorch(context.Background(), t, "/pytorch-examples/mnist/main.py", "--epochs=1", "--dry-run")
+}
+
+// TestIssue9827 verifies that issue 9827 is fixed.
+func TestIssue9827(t *testing.T) {
+	// TODO(gvisor.dev/issue/9827): Don't skip this once the
+	// test works and doesn't run forever:
+	t.Skip("TODO(gvisor.dev/issue/9827): Issue 9827 is not yet fixed.")
+	runPytorch(context.Background(), t, "/issue_9827.py")
+}