diff --git a/Makefile b/Makefile index e1599e2a2..01765f730 100644 --- a/Makefile +++ b/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 diff --git a/images/gpu/pytorch/Dockerfile b/images/gpu/pytorch/Dockerfile new file mode 100644 index 000000000..ea29b768e --- /dev/null +++ b/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 diff --git a/images/gpu/pytorch/download_pytorch_datasets.py b/images/gpu/pytorch/download_pytorch_datasets.py new file mode 100644 index 000000000..28a88054c --- /dev/null +++ b/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) diff --git a/images/gpu/pytorch/is_cuda_available.py b/images/gpu/pytorch/is_cuda_available.py new file mode 100644 index 000000000..b5060c33d --- /dev/null +++ b/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) diff --git a/images/gpu/pytorch/issue_9827.py b/images/gpu/pytorch/issue_9827.py new file mode 100644 index 000000000..f0384926c --- /dev/null +++ b/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}") diff --git a/test/gpu/BUILD b/test/gpu/BUILD index 50e658787..9caaa2845 100644 --- a/test/gpu/BUILD +++ b/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"], diff --git a/test/gpu/pytorch_test.go b/test/gpu/pytorch_test.go new file mode 100644 index 000000000..04f617754 --- /dev/null +++ b/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") +}