def setUp(self):
super().setUp()
workflow_0 = Workflow(id=0, name='test-workflow-0', project_id=0)
workflow_1 = Workflow(id=1, name='test-workflow-1', project_id=0)
db.session.add_all([workflow_0, workflow_1])
config = workflow_definition_pb2.JobDefinition(
name='test-job').SerializeToString()
job_0 = Job(id=0,
name='raw_data_0',
job_type=JobType.RAW_DATA,
state=JobState.STARTED,
workflow_id=0,
project_id=0,
config=config)
job_1 = Job(id=1,
name='raw_data_1',
job_type=JobType.RAW_DATA,
state=JobState.COMPLETED,
workflow_id=0,
project_id=0,
config=config)
job_2 = Job(id=2,
name='data_join_0',
job_type=JobType.DATA_JOIN,
state=JobState.WAITING,
workflow_id=0,
project_id=0,
config=config)
job_3 = Job(id=3,
name='data_join_1',
job_type=JobType.DATA_JOIN,
state=JobState.COMPLETED,
workflow_id=1,
project_id=0,
config=config)
job_4 = Job(id=4,
name='train_job_0',
job_type=JobType.NN_MODEL_TRANINING,
state=JobState.WAITING,
workflow_id=1,
project_id=0,
config=config)
db.session.add_all([job_0, job_1, job_2, job_3, job_4])
job_dep_0 = JobDependency(src_job_id=job_0.id,
dst_job_id=job_2.id,
dep_index=0)
job_dep_1 = JobDependency(src_job_id=job_1.id,
dst_job_id=job_2.id,
dep_index=1)
job_dep_2 = JobDependency(src_job_id=job_3.id,
dst_job_id=job_4.id,
dep_index=0)
db.session.add_all([job_dep_0, job_dep_1, job_dep_2])
db.session.commit()
def setUp(self, mock_get_checkpoint_path):
super().setUp()
self.model_manager = ModelManager()
self.train_job = Job(name='train-job',
job_type=JobType.TREE_MODEL_TRAINING)
self.eval_job = Job(name='eval-job',
job_type=JobType.TREE_MODEL_EVALUATION)
mock_get_checkpoint_path.return_value = 'output'
self.model_manager.create(parent_id=None, job=self.train_job)
model = Model.query.filter_by(job_name=self.train_job.name).one()
self.model_manager.create(parent_id=model.id, job=self.eval_job)
def _setup_jobs(self):
if self.forked_from is not None:
trunk = Workflow.query.get(self.forked_from)
assert trunk is not None, \
'Source workflow %d not found' % self.forked_from
trunk_job_defs = trunk.get_config().job_definitions
trunk_name2index = {
job.name: i
for i, job in enumerate(trunk_job_defs)
}
job_defs = self.get_config().job_definitions
flags = self.get_create_job_flags()
assert len(job_defs) == len(flags), \
'Number of job defs does not match number of create_job_flags ' \
'%d vs %d'%(len(job_defs), len(flags))
jobs = []
for i, (job_def, flag) in enumerate(zip(job_defs, flags)):
if flag == common_pb2.CreateJobFlag.REUSE:
assert job_def.name in trunk_name2index, \
f'Job {job_def.name} not found in base workflow'
j = trunk.get_job_ids()[trunk_name2index[job_def.name]]
job = Job.query.get(j)
assert job is not None, \
'Job %d not found' % j
# TODO: check forked jobs does not depend on non-forked jobs
else:
job = Job(
name=f'{self.uuid}-{job_def.name}',
job_type=JobType(job_def.job_type),
config=job_def.SerializeToString(),
workflow_id=self.id,
project_id=self.project_id,
state=JobState.NEW,
is_disabled=(flag == common_pb2.CreateJobFlag.DISABLED))
db.session.add(job)
jobs.append(job)
db.session.flush()
name2index = {job.name: i for i, job in enumerate(job_defs)}
for i, (job, flag) in enumerate(zip(jobs, flags)):
if flag == common_pb2.CreateJobFlag.REUSE:
continue
for j, dep_def in enumerate(job.get_config().dependencies):
dep = JobDependency(
src_job_id=jobs[name2index[dep_def.source]].id,
dst_job_id=job.id,
dep_index=j)
db.session.add(dep)
self.set_job_ids([job.id for job in jobs])
if Features.FEATURE_MODEL_WORKFLOW_HOOK:
for job in jobs:
ModelService(db.session).workflow_hook(job)
def _setup_jobs(self):
if self.forked_from is not None:
trunk = Workflow.query.get(self.forked_from)
assert trunk is not None, \
'Source workflow %d not found'%self.forked_from
trunk_job_defs = trunk.get_config().job_definitions
trunk_name2index = {
job.name: i
for i, job in enumerate(trunk_job_defs)
}
else:
assert not self.get_reuse_job_names()
job_defs = self.get_config().job_definitions
jobs = []
reuse_jobs = set(self.get_reuse_job_names())
for i, job_def in enumerate(job_defs):
if job_def.name in reuse_jobs:
assert job_def.name in trunk_name2index, \
"Job %s not found in base workflow"%job_def.name
j = trunk.get_job_ids()[trunk_name2index[job_def.name]]
job = Job.query.get(j)
assert job is not None, \
'Job %d not found'%j
# TODO: check forked jobs does not depend on non-forked jobs
else:
job = Job(name=f'{self.name}-{job_def.name}',
job_type=JobType(job_def.type),
config=job_def.SerializeToString(),
workflow_id=self.id,
project_id=self.project_id,
state=JobState.STOPPED)
job.set_yaml_template(job_def.yaml_template)
db.session.add(job)
jobs.append(job)
db.session.commit()
name2index = {job.name: i for i, job in enumerate(job_defs)}
for i, job in enumerate(jobs):
if job.name in reuse_jobs:
continue
for j, dep_def in enumerate(job.get_config().dependencies):
dep = JobDependency(
src_job_id=jobs[name2index[dep_def.source]].id,
dst_job_id=job.id,
dep_index=j)
db.session.add(dep)
self.set_job_ids([job.id for job in jobs])
db.session.commit()
def test_get_waiting_jobs(self):
db.session.add(
Job(name='testtes',
state=JobState.STOPPED,
job_type=JobType.DATA_JOIN,
workflow_id=1,
project_id=1))
db.session.add(
Job(name='testtest',
state=JobState.WAITING,
job_type=JobType.DATA_JOIN,
workflow_id=1,
project_id=1))
db.session.commit()
self.assertEqual(_get_waiting_jobs(), [2])
def test_patch_create_job_flags(self):
wd = WorkflowDefinition()
jd = wd.job_definitions.add()
workflow = Workflow(
name='test-workflow',
project_id=123,
config=wd.SerializeToString(),
forkable=False,
state=WorkflowState.READY,
)
db.session.add(workflow)
db.session.flush()
job = Job(name='test_job',
job_type=JobType(1),
config=jd.SerializeToString(),
workflow_id=workflow.id,
project_id=123,
state=JobState.STOPPED,
is_disabled=False)
db.session.add(job)
db.session.flush()
workflow.job_ids = str(job.id)
db.session.commit()
response = self.patch_helper(f'/api/v2/workflows/{workflow.id}',
data={'create_job_flags': [3]})
self.assertEqual(response.status_code, HTTPStatus.OK)
patched_job = Job.query.get(job.id)
self.assertEqual(patched_job.is_disabled, True)
response = self.patch_helper(f'/api/v2/workflows/{workflow.id}',
data={'create_job_flags': [1]})
self.assertEqual(response.status_code, HTTPStatus.OK)
patched_job = Job.query.get(job.id)
self.assertEqual(patched_job.is_disabled, False)
def setUpClass(cls):
grpc_spec = GrpcSpec(
authority=cls._TEST_AUTHORITY,
extra_headers={cls._X_HOST_HEADER_KEY: cls._TEST_X_HOST})
participant = Participant(name=cls._TEST_RECEIVER_NAME,
domain_name='fl-test.com',
grpc_spec=grpc_spec)
project_config = Project(name=cls._TEST_PROJECT_NAME,
token='test-auth-token',
participants=[participant],
variables=[{
'name': 'EGRESS_URL',
'value': cls._TEST_URL
}])
job = Job(name='test-job')
cls._participant = participant
cls._project_config = project_config
cls._project = ProjectModel(name=cls._TEST_PROJECT_NAME)
cls._project.set_config(project_config)
cls._job = job
# Inserts the project entity
cls._DB.create_all()
cls._DB.session.add(cls._project)
cls._DB.session.commit()
def follower_test_peer_metrics(self):
self.setup_project('follower',
JobMetricsBuilderTest.Config.GRPC_LISTEN_PORT)
workflow = Workflow(name='test-workflow',
project_id=1,
metric_is_public=True)
workflow.set_job_ids([1])
db.session.add(workflow)
job = Job(name='automl-2782410011',
job_type=JobType.NN_MODEL_TRANINING,
workflow_id=1,
project_id=1,
config=workflow_definition_pb2.JobDefinition(
name='test-job').SerializeToString())
db.session.add(job)
db.session.commit()
while True:
time.sleep(1)
def add_fake_workflow(session):
wd = WorkflowDefinition()
jd = wd.job_definitions.add()
workflow = Workflow(
name='test-workflow',
project_id=123,
config=wd.SerializeToString(),
forkable=False,
state=WorkflowState.READY,
)
session.add(workflow)
session.flush()
job = Job(name='test_job',
job_type=JobType(1),
config=jd.SerializeToString(),
workflow_id=workflow.id,
project_id=123,
state=JobState.STOPPED,
is_disabled=False)
session.add(job)
session.flush()
workflow.job_ids = str(job.id)
session.commit()
return workflow, job
def test_generate_self_dict(self):
config = {
'variables': [
{
'name': 'namespace',
'value': 'leader'
},
{
'name': 'basic_envs',
'value': '{}'
},
{
'name': 'storage_root_dir',
'value': '/'
},
{
'name': 'EGRESS_URL',
'value': '127.0.0.1:1991'
}
]
}
job = Job(name='aa', project_id=1, workflow_id=1, state=JobState.NEW)
job.set_config(ParseDict(config, JobDefinition()))
self.assertEqual(generate_self_dict(job),
{'id': None, 'name': 'aa',
'job_type': None, 'state': 'NEW', 'config':
{'expert_mode': False,
'variables': [
{
'name': 'namespace',
'value': 'leader',
'access_mode': 'UNSPECIFIED',
'widget_schema': '',
'value_type': 'STRING'},
{
'name': 'basic_envs',
'value': '{}',
'access_mode': 'UNSPECIFIED',
'widget_schema': '',
'value_type': 'STRING'},
{
'name': 'storage_root_dir',
'value': '/',
'access_mode': 'UNSPECIFIED',
'widget_schema': '',
'value_type': 'STRING'},
{
'name': 'EGRESS_URL',
'value': '127.0.0.1:1991',
'access_mode': 'UNSPECIFIED',
'widget_schema': '',
'value_type': 'STRING'}],
'name': '',
'job_type': 'UNSPECIFIED',
'is_federated': False,
'dependencies': [],
'yaml_template': ''},
'is_disabled': None, 'workflow_id': 1, 'project_id': 1, 'flapp_snapshot': None,
'pods_snapshot': None, 'error_message': None, 'created_at': None, 'updated_at': None,
'deleted_at': None, 'pods': [], 'complete_at': None,
'variables': {'namespace': 'leader', 'basic_envs': '{}', 'storage_root_dir': '/',
'EGRESS_URL': '127.0.0.1:1991'}}
)
def test_nn_metrics(self):
job = Job(name='automl-2782410011',
job_type=JobType.NN_MODEL_TRANINING)
print(JobMetricsBuilder(job).plot_metrics())
def test_data_join_metrics(self):
job = Job(name='multi-indices-test27', job_type=JobType.DATA_JOIN)
import json
print(json.dumps(JobMetricsBuilder(job).plot_metrics()))