def test_resume_workflow(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.resume_workflow_scheduling.return_value = WorkflowInfo(
workflow_name='test_workflow')
client = SchedulerClient("localhost:{}".format(_PORT))
workflow = client.resume_workflow_scheduling(
namespace='namespace', workflow_name='test_workflow')
self.assertTrue('test_workflow', workflow.name)
def test_list_workflow_executions(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.list_workflow_executions.return_value \
= [WorkflowExecutionInfo(workflow_execution_id='id_1', status=Status.INIT),
WorkflowExecutionInfo(workflow_execution_id='id_2', status=Status.INIT)]
client = SchedulerClient("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_submit_workflow(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
instance.submit_workflow.return_value = WorkflowInfo(
workflow_name='test_workflow')
client = SchedulerClient("localhost:{}".format(_PORT))
with self.assertRaises(Exception) as context:
workflow = client.submit_workflow_to_scheduler(
namespace='namespace',
workflow_name='test_workflow',
workflow_json='')
self.assertTrue('workflow json is empty' in str(context.exception))
def test_get_workflow_execution(self):
with mock.patch(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.get_workflow_execution.return_value \
= WorkflowExecutionInfo(workflow_execution_id='id', status=Status.INIT)
client = SchedulerClient("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(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.list_job_executions.return_value \
= [JobExecutionInfo(job_name='job_name_1',
status=Status.RUNNING,
workflow_execution=WorkflowExecutionInfo(workflow_execution_id='id',
status=Status.INIT)),
JobExecutionInfo(job_name='job_name_2',
status=Status.RUNNING,
workflow_execution=WorkflowExecutionInfo(workflow_execution_id='id',
status=Status.INIT))]
client = SchedulerClient("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(SCHEDULER_CLASS) as mockScheduler:
instance = mockScheduler.return_value
self.server.scheduler_service._scheduler = instance
instance.get_job_executions.return_value \
= [JobExecutionInfo(job_name='job_name',
status=Status.RUNNING,
workflow_execution=WorkflowExecutionInfo(workflow_execution_id='id',
status=Status.INIT))]
client = SchedulerClient("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)
MetricClient.__init__(self, server_uri)
SchedulerClient.__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_server_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()