def test_delete_none_workflow(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.delete_workflow.return_value = None client = SchedulingClient("localhost:{}".format(_PORT)) with self.assertRaises(Exception) as context: workflow = client.delete_workflow( namespace='namespace', workflow_name='test_workflow')
def test_get_workflow(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.get_workflow.return_value = WorkflowInfo( workflow_name='test_workflow') client = SchedulingClient("localhost:{}".format(_PORT)) workflow = client.get_workflow(namespace='namespace', workflow_name='test_workflow') self.assertTrue('test_workflow', workflow.name)
def test_list_workflow_executions(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.list_workflow_executions.return_value \ = [WorkflowExecutionInfo(execution_id='id_1', state=State.INIT), WorkflowExecutionInfo(execution_id='id_2', state=State.INIT)] client = SchedulingClient("localhost:{}".format(_PORT)) workflow_execution_list = client.list_workflow_executions( namespace='namespace', workflow_name='test_workflow') self.assertEqual(2, len(workflow_execution_list))
def test_list_workflows(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.list_workflows.return_value = [ WorkflowInfo(workflow_name='test_workflow_1'), WorkflowInfo(workflow_name='test_workflow_2') ] client = SchedulingClient("localhost:{}".format(_PORT)) workflow_list = client.list_workflows(namespace='namespace') self.assertTrue(2, len(workflow_list))
def test_get_workflow_execution(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.get_workflow_execution.return_value \ = WorkflowExecutionInfo(execution_id='id', state=State.INIT) client = SchedulingClient("localhost:{}".format(_PORT)) workflow_execution = client.get_workflow_execution( execution_id='id') self.assertEqual('id', workflow_execution.execution_id) self.assertEqual(StateProto.INIT, workflow_execution.execution_state)
def test_list_jobs(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.list_jobs.return_value \ = [JobInfo(job_name='job_name_1', state=State.RUNNING, workflow_execution=WorkflowExecutionInfo(execution_id='id', state=State.INIT)), JobInfo(job_name='job_name_2', state=State.RUNNING, workflow_execution=WorkflowExecutionInfo(execution_id='id', state=State.INIT))] client = SchedulingClient("localhost:{}".format(_PORT)) job_list = client.list_jobs(execution_id='id') self.assertEqual(2, len(job_list))
def test_get_job(self): with mock.patch( 'ai_flow.test.scheduler.test_scheduling_service.MockScheduler' ) as mockScheduler: instance = mockScheduler.return_value self.server.scheduling_service._scheduler = instance instance.get_job.return_value \ = JobInfo(job_name='job_name', state=State.RUNNING, workflow_execution=WorkflowExecutionInfo(execution_id='id', state=State.INIT)) client = SchedulingClient("localhost:{}".format(_PORT)) job = client.get_job(job_name='job_name', execution_id='id') self.assertEqual('job_name', job.name) self.assertEqual(StateProto.RUNNING, job.job_state) self.assertEqual('id', job.workflow_execution.execution_id) self.assertEqual(StateProto.INIT, job.workflow_execution.execution_state)
def __init__(self, server_uri=_SERVER_URI, notification_service_uri=None, project_config: ProjectConfig = None): MetadataClient.__init__(self, server_uri) ModelCenterClient.__init__(self, server_uri) DeployClient.__init__(self, server_uri) MetricClient.__init__(self, server_uri) SchedulingClient.__init__(self, server_uri) self.enable_ha = False self.list_member_interval_ms = 5000 self.retry_interval_ms = 1000 self.retry_timeout_ms = 10000 if project_config is not None: if server_uri is None: server_uri = project_config.get_master_uri() if notification_service_uri is None: notification_service_uri = project_config.get_notification_service_uri( ) self.enable_ha = project_config.get_enable_ha() self.list_member_interval_ms = project_config.get_list_member_interval_ms( ) self.retry_interval_ms = project_config.get_retry_interval_ms() self.retry_timeout_ms = project_config.get_retry_timeout_ms() if notification_service_uri is None: NotificationClient.__init__( self, server_uri, enable_ha=self.enable_ha, list_member_interval_ms=self.list_member_interval_ms, retry_interval_ms=self.retry_interval_ms, retry_timeout_ms=self.retry_timeout_ms) else: NotificationClient.__init__( self, notification_service_uri, enable_ha=self.enable_ha, list_member_interval_ms=self.list_member_interval_ms, retry_interval_ms=self.retry_interval_ms, retry_timeout_ms=self.retry_timeout_ms) if self.enable_ha: server_uris = server_uri.split(",") self.living_aiflow_members = [] self.current_aiflow_uri = None last_error = None for server_uri in server_uris: channel = grpc.insecure_channel(server_uri) high_availability_stub = HighAvailabilityManagerStub(channel) try: request = ListMembersRequest(timeout_seconds=0) response = high_availability_stub.listMembers(request) if response.return_code == ReturnStatus.CALL_SUCCESS: self.living_aiflow_members = [ proto_to_member(proto).server_uri for proto in response.members ] else: raise Exception(response.return_msg) self.current_aiflow_uri = server_uri self.high_availability_stub = high_availability_stub break except grpc.RpcError as e: last_error = e if self.current_aiflow_uri is None: raise Exception( "No available aiflow server uri!") from last_error self.aiflow_ha_change_lock = threading.Lock() self.aiflow_ha_running = True self._replace_aiflow_stubs(self.current_aiflow_uri) self.list_aiflow_member_thread = threading.Thread( target=self._list_aiflow_members, daemon=True) self.list_aiflow_member_thread.start()