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()
