diff options
Diffstat (limited to 'tools/guest_agent')
| -rw-r--r-- | tools/guest_agent/guest_test_agent.py | 332 |
1 files changed, 332 insertions, 0 deletions
diff --git a/tools/guest_agent/guest_test_agent.py b/tools/guest_agent/guest_test_agent.py new file mode 100644 index 0000000..15a578d --- /dev/null +++ b/tools/guest_agent/guest_test_agent.py @@ -0,0 +1,332 @@ +#!/usr/bin/env python3 +""" +guest_test_agent.py - Runs inside guest VM, executes tests, reports results +Communication via vsock, serial, or HTTP +""" + +import sys +import os +import json +import time +import subprocess +import threading +import socket +import argparse +from datetime import datetime +from pathlib import Path + +class GuestTestAgent: + def __init__(self, listen_port=9999, report_url=None, vsock_cid=3, vsock_port=9999): + self.listen_port = listen_port + self.report_url = report_url + self.vsock_cid = vsock_cid + self.vsock_port = vsock_port + self.running = False + self.results = [] + + def start_server(self): + """Start command listener (TCP or vsock)""" + self.running = True + + # Try vsock first (preferred for VM-host communication) + try: + import socket + if hasattr(socket, 'AF_VSOCK'): + self.server = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM) + self.server.bind((self.vsock_cid, self.vsock_port)) + self.server.listen(5) + print(f"[{datetime.now()}] Guest agent listening on vsock CID:{self.vsock_cid} PORT:{self.vsock_port}") + else: + raise ImportError("vsock not available") + except Exception: + # Fallback to TCP + self.server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.server.bind(('0.0.0.0', self.listen_port)) + self.server.listen(5) + print(f"[{datetime.now()}] Guest agent listening on TCP port {self.listen_port}") + + while self.running: + try: + conn, addr = self.server.accept() + threading.Thread(target=self._handle_connection, args=(conn,), daemon=True).start() + except Exception as e: + if self.running: + print(f"[{datetime.now()}] Server error: {e}") + + def _handle_connection(self, conn): + """Handle incoming test commands""" + try: + data = conn.recv(4096).decode('utf-8') + if not data: + return + + cmd = json.loads(data) + action = cmd.get('action') + test_id = cmd.get('test_id', 'unknown') + + print(f"[{datetime.now()}] Received command: {action} (test_id: {test_id})") + + if action == 'ping': + result = {'status': 'ok', 'agent': 'guest_test_agent', 'version': '1.0'} + elif action == 'run_test': + result = self._run_test(cmd.get('test'), cmd.get('args', {})) + elif action == 'run_suite': + result = self._run_suite(cmd.get('suite')) + elif action == 'get_results': + result = {'results': self.results} + elif action == 'shutdown': + self.running = False + result = {'status': 'shutting_down'} + else: + result = {'error': f'Unknown action: {action}'} + + result['test_id'] = test_id + conn.send(json.dumps(result).encode('utf-8')) + + except Exception as e: + conn.send(json.dumps({'error': str(e)}).encode('utf-8')) + finally: + conn.close() + + def _run_test(self, test_name, args): + """Run a single test""" + test_method = getattr(self, f'test_{test_name}', None) + if not test_method: + return {'status': 'error', 'message': f'Unknown test: {test_name}'} + + start = time.time() + try: + result = test_method(**args) + result['duration'] = time.time() - start + result['timestamp'] = datetime.now().isoformat() + self.results.append(result) + return result + except Exception as e: + return {'status': 'error', 'message': str(e), 'duration': time.time() - start} + + def _run_suite(self, suite_name): + """Run a test suite""" + suite_method = getattr(self, f'suite_{suite_name}', None) + if not suite_method: + return {'status': 'error', 'message': f'Unknown suite: {suite_name}'} + + return suite_method() + + # ============================================================ + # INDIVIDUAL TESTS + # ============================================================ + + def test_nvidia_smi(self): + """Run nvidia-smi and verify GPU detection""" + try: + out = subprocess.check_output(['nvidia-smi'], text=True, timeout=30) + return { + 'test': 'nvidia_smi', + 'status': 'pass' if 'NVIDIA' in out else 'fail', + 'output': out[:2000] + } + except subprocess.CalledProcessError as e: + return {'test': 'nvidia_smi', 'status': 'fail', 'error': str(e), 'output': e.output} + except Exception as e: + return {'test': 'nvidia_smi', 'status': 'error', 'message': str(e)} + + def test_cuda_vectoradd(self): + """Compile and run CUDA vectorAdd sample""" + try: + # Find CUDA samples + cuda_samples = Path('/usr/local/cuda/samples') + if not cuda_samples.exists(): + cuda_samples = Path('/opt/cuda/samples') + + vecadd_dir = cuda_samples / '1_Utilities' / 'vectorAdd' + if not vecadd_dir.exists(): + return {'test': 'cuda_vectoradd', 'status': 'skip', 'message': 'vectorAdd sample not found'} + + # Build + subprocess.check_call(['make'], cwd=vecadd_dir, timeout=60) + + # Run + out = subprocess.check_output(['./vectorAdd'], cwd=vecadd_dir, text=True, timeout=30) + + return { + 'test': 'cuda_vectoradd', + 'status': 'pass' if 'Test passed' in out else 'fail', + 'output': out[:2000] + } + except subprocess.CalledProcessError as e: + return {'test': 'cuda_vectoradd', 'status': 'fail', 'error': str(e), 'output': e.output} + except Exception as e: + return {'test': 'cuda_vectoradd', 'status': 'error', 'message': str(e)} + + def test_cuda_devquery(self): + """Run CUDA deviceQuery""" + try: + cuda_samples = Path('/usr/local/cuda/samples') + if not cuda_samples.exists(): + cuda_samples = Path('/opt/cuda/samples') + + devq_dir = cuda_samples / '1_Utilities' / 'deviceQuery' + if not devq_dir.exists(): + return {'test': 'cuda_devquery', 'status': 'skip', 'message': 'deviceQuery sample not found'} + + subprocess.check_call(['make'], cwd=devq_dir, timeout=60) + out = subprocess.check_output(['./deviceQuery'], cwd=devq_dir, text=True, timeout=30) + + return { + 'test': 'cuda_devquery', + 'status': 'pass' if 'Result = PASS' in out else 'fail', + 'output': out[:3000] + } + except Exception as e: + return {'test': 'cuda_devquery', 'status': 'error', 'message': str(e)} + + def test_vulkan_info(self): + """Run vulkaninfo""" + try: + out = subprocess.check_output(['vulkaninfo'], text=True, timeout=30) + return { + 'test': 'vulkan_info', + 'status': 'pass' if 'GPU id' in out else 'fail', + 'output': out[:3000] + } + except Exception as e: + return {'test': 'vulkan_info', 'status': 'error', 'message': str(e)} + + def test_vkcube(self): + """Run vkcube (headless)""" + try: + # vkcube may need display; try with xvfb or skip + out = subprocess.check_output(['vkcube', '--duration', '2'], text=True, timeout=30) + return {'test': 'vkcube', 'status': 'pass', 'output': out[:2000]} + except FileNotFoundError: + return {'test': 'vkcube', 'status': 'skip', 'message': 'vkcube not installed'} + except Exception as e: + return {'test': 'vkcube', 'status': 'error', 'message': str(e)} + + def test_glxgears(self): + """Run glxgears (brief)""" + try: + out = subprocess.check_output(['timeout', '5', 'glxgears'], text=True, timeout=15) + return {'test': 'glxgears', 'status': 'pass', 'output': out[:1000]} + except Exception as e: + return {'test': 'glxgears', 'status': 'error', 'message': str(e)} + + def test_ffmpeg_nvenc(self): + """Test FFmpeg NVENC""" + try: + out = subprocess.check_output(['ffmpeg', '-encoders'], text=True, timeout=10) + has_nvenc = 'h264_nvenc' in out or 'hevc_nvenc' in out + return { + 'test': 'ffmpeg_nvenc', + 'status': 'pass' if has_nvenc else 'fail', + 'output': out[:2000] + } + except Exception as e: + return {'test': 'ffmpeg_nvenc', 'status': 'error', 'message': str(e)} + + def test_pytorch(self): + """Test PyTorch CUDA""" + try: + code = """ +import torch +print(f'CUDA available: {torch.cuda.is_available()}') +if torch.cuda.is_available(): + print(f'Device count: {torch.cuda.device_count()}') + print(f'Device name: {torch.cuda.get_device_name(0)}') + x = torch.randn(100, 100).cuda() + y = torch.randn(100, 100).cuda() + z = torch.mm(x, y) + print(f'Matmul result shape: {z.shape}') + print('PyTorch CUDA test: PASS') +else: + print('PyTorch CUDA test: FAIL - CUDA not available') +""" + out = subprocess.check_output([sys.executable, '-c', code], text=True, timeout=60) + return { + 'test': 'pytorch', + 'status': 'pass' if 'PASS' in out else 'fail', + 'output': out + } + except Exception as e: + return {'test': 'pytorch', 'status': 'error', 'message': str(e)} + + # ============================================================ + # TEST SUITES + # ============================================================ + + def suite_smoke(self): + """Quick smoke test suite""" + tests = ['nvidia_smi', 'cuda_devquery', 'cuda_vectoradd'] + return self._run_test_list(tests) + + def suite_compute(self): + """Compute-focused test suite""" + tests = ['nvidia_smi', 'cuda_devquery', 'cuda_vectoradd', 'pytorch'] + return self._run_test_list(tests) + + def suite_graphics(self): + """Graphics-focused test suite""" + tests = ['nvidia_smi', 'vulkan_info', 'vkcube', 'glxgears'] + return self._run_test_list(tests) + + def suite_video(self): + """Video encode/decode test suite""" + tests = ['nvidia_smi', 'ffmpeg_nvenc'] + return self._run_test_list(tests) + + def suite_full(self): + """Full test suite""" + tests = ['nvidia_smi', 'cuda_devquery', 'cuda_vectoradd', 'vulkan_info', 'vkcube', 'glxgears', 'ffmpeg_nvenc', 'pytorch'] + return self._run_test_list(tests) + + def _run_test_list(self, tests): + """Run a list of tests""" + results = [] + for test in tests: + print(f"[{datetime.now()}] Running suite test: {test}") + result = self._run_test(test, {}) + results.append(result) + print(f" -> {result.get('status', 'unknown')}") + + passed = sum(1 for r in results if r.get('status') == 'pass') + failed = sum(1 for r in results if r.get('status') in ('fail', 'error')) + skipped = sum(1 for r in results if r.get('status') == 'skip') + + return { + 'suite': 'custom', + 'total': len(tests), + 'passed': passed, + 'failed': failed, + 'skipped': skipped, + 'results': results + } + +def main(): + parser = argparse.ArgumentParser(description="Guest Test Agent") + parser.add_argument('--port', type=int, default=9999, help='TCP listen port') + parser.add_argument('--vsock-cid', type=int, default=3, help='VSOCK CID') + parser.add_argument('--vsock-port', type=int, default=9999, help='VSOCK port') + parser.add_argument('--run', help='Run single test and exit (for standalone mode)') + parser.add_argument('--suite', help='Run test suite and exit') + args = parser.parse_args() + + agent = GuestTestAgent(listen_port=args.port, vsock_cid=args.vsock_cid, vsock_port=args.vsock_port) + + if args.run: + # Standalone mode + result = agent._run_test(args.run, {}) + print(json.dumps(result, indent=2)) + elif args.suite: + result = getattr(agent, f'suite_{args.suite}')() + print(json.dumps(result, indent=2)) + else: + # Server mode + try: + agent.start_server() + except KeyboardInterrupt: + print("\nShutting down...") + agent.running = False + +if __name__ == "__main__": + main()
\ No newline at end of file |