github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/internal/runner/cuda_test.go (about)

     1  // +build !NO_CUDA
     2  
     3  // Copyright 2018-2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     4  
     5  package runner
     6  
     7  import (
     8  	"bufio"
     9  	"fmt"
    10  	"os"
    11  	"strconv"
    12  	"strings"
    13  	"testing"
    14  
    15  	"github.com/go-stack/stack"
    16  	"github.com/jjeffery/kv" // MIT License
    17  
    18  	nvml "github.com/karlmutch/go-nvml"
    19  
    20  	"github.com/leaf-ai/studio-go-runner/pkg/studio"
    21  )
    22  
    23  var (
    24  	errFormatIssue = kv.NewError("unexpected format, lines should be in the format x=y")
    25  )
    26  
    27  func init() {
    28  	CudaInTest = true
    29  }
    30  
    31  // This file contains an integration test implementation that submits a studio runner
    32  // task across an SQS queue and then validates is has completed successfully by
    33  // the go runner this test is running within
    34  
    35  func readIni(fn string) (items map[string]string, err kv.Error) {
    36  
    37  	items = map[string]string{}
    38  
    39  	fh, errGo := os.Open(fn)
    40  	if errGo != nil {
    41  		return items, kv.Wrap(errGo).With("filename", fn).With("stack", stack.Trace().TrimRuntime())
    42  	}
    43  	defer fh.Close()
    44  
    45  	scanner := bufio.NewScanner(fh)
    46  	for scanner.Scan() {
    47  		aLine := scanner.Text()
    48  		kv := strings.SplitN(aLine, "=", 2)
    49  		if len(kv) != 2 {
    50  			return items, errFormatIssue.With("line", aLine).With("filename", fn).With("stack", stack.Trace().TrimRuntime())
    51  		}
    52  		items[kv[0]] = kv[1]
    53  	}
    54  	if errGo := scanner.Err(); errGo != nil {
    55  		return items, kv.Wrap(errGo).With("filename", fn).With("stack", stack.Trace().TrimRuntime())
    56  	}
    57  
    58  	return items, nil
    59  }
    60  
    61  // TestCUDAActive checks that at least one GPU is available before any other GPU tests are used
    62  //
    63  func TestCUDAActive(t *testing.T) {
    64  	logger := studio.NewLogger("cuda_active_test")
    65  	defer logger.Warn("completed")
    66  
    67  	if !*UseGPU {
    68  		logger.Warn("TestCUDAActive not run")
    69  		t.Skip("no GPUs present for testing")
    70  	}
    71  
    72  	devs, errGo := nvml.GetAllGPUs()
    73  	if errGo != nil {
    74  		t.Fatal(errGo)
    75  	}
    76  	if len(devs) < 1 {
    77  		t.Fatal("no CUDA capable devices found during the CUDA testing")
    78  	}
    79  
    80  	annotations, err := readIni(k8sAnnotations)
    81  	if err != nil {
    82  		logger.Warn("test appears to be running without k8s specifications")
    83  
    84  		if *useK8s {
    85  			t.Fatal("Kubernetes cluster present for testing, however the downward API files are missing " + err.Error())
    86  		}
    87  		return
    88  	}
    89  
    90  	if gpus, isPresent := annotations["gpus"]; isPresent {
    91  		gpus = strings.Trim(gpus, "\"'")
    92  		expected, errGo := strconv.Atoi(gpus)
    93  		if errGo != nil {
    94  			t.Fatal(errGo.Error())
    95  
    96  		}
    97  		if len(devs) != expected {
    98  			t.Fatal(fmt.Sprintln("expected ", expected, " gpus got ", len(devs)))
    99  		}
   100  	}
   101  }