Coverage for streamwise/job_manager.py: 72%

1""" 

2Submit jobs to services in the cluster. 

3""" 

4 

5import sys 

6import logging 

7 

8from aiohttp import ClientTimeout 

9from asyncio import TimeoutError 

10 

11from http import HTTPStatus 

12 

13from quart import request 

14from quart import jsonify 

15from quart import Response 

16 

17import http_session_manager 

18 

19sys.path.append("..") 

20from quart_utils import QuartReturn 

21 

22 

23# HTTP clients 

24JSON_HEADER = { 

25 "Content-Type": "application/json" 

26} 

27GENERATION_TIMEOUT = ClientTimeout( 

28 total=60.0, 

29 connect=5.0) 

30 

31 

32async def submit_job( 

33 service_name: str, 

34 container_ip: str, 

35 container_port: int 

36) -> QuartReturn: 

37 """API interface to submit a job to the specified service.""" 

38 if not service_name or not container_ip or not container_port: 

39 return jsonify({"error": "Service name, container IP and port are required"}), HTTPStatus.BAD_REQUEST 

40 try: 

41 payload_json = await request.get_json() 

42 if not payload_json: 

43 return jsonify({"error": "No job data provided"}), HTTPStatus.BAD_REQUEST 

44 

45 url = f"{http_session_manager.SERVICE_SCHEME}://{container_ip}:{container_port}/{service_name}" 

46 session = await http_session_manager.get_global_session() 

47 async with session.post(url, json=payload_json, headers=JSON_HEADER, timeout=GENERATION_TIMEOUT) as response: 

48 content_type = response.headers.get("Content-Type", "application/octet-stream") 

49 data = await response.read() 

50 headers = {} 

51 content_disposition = response.headers.get("Content-Disposition") 

52 if content_disposition: 

53 headers["Content-Disposition"] = content_disposition 

54 return Response( 

55 response=data, 

56 status=response.status, 

57 headers=headers, 

58 content_type=content_type) 

59 except TimeoutError: 

60 logging.error(f"Timeout submitting job to {service_name} at {container_ip}:{container_port}.") 

61 return jsonify({"error": "Request timed out"}), HTTPStatus.GATEWAY_TIMEOUT 

62 except Exception as ex: 

63 logging.error(f"Error submitting job to {service_name} at {container_ip}:{container_port}: {ex}") 

64 return jsonify({"error": str(ex)}), HTTPStatus.INTERNAL_SERVER_ERROR