def dag_pipe(session: Session, config: dict, config_text: str = None):
assert 'interfaces' in config, 'interfaces missed'
assert 'pipes' in config, 'pipe missed'
info = config['info']
storage = Storage(session)
dag_provider = DagProvider(session)
folder = os.getcwd()
project = ProjectProvider(session).by_name(info['project']).id
dag = dag_provider.add(
Dag(
config=config_text,
project=project,
name=info['name'],
docker_img=info.get('docker_img'),
type=DagType.Pipe.value
)
)
storage.upload(folder, dag)
# Change model dags which have the same name
ModelProvider(session
).change_dag(project=project, name=info['name'], to=dag.id)
def __init__(self,
session: Session,
task: Task,
layout: str,
part: str = 'valid',
name: str = 'img_segment',
max_img_size: Tuple[int, int] = None,
stack_type: str = 'vertical',
main_metric: str = 'dice',
plot_count: int = 0,
colors: List[Tuple] = None):
self.session = session
self.task = task
self.layout = layout
self.part = part
self.name = name or 'img_segment'
self.max_img_size = max_img_size
self.stack_type = stack_type
self.main_metric = main_metric
self.colors = colors
self.plot_count = plot_count
self.dag_provider = DagProvider(session)
self.report_provider = ReportProvider(session)
self.layout_provider = ReportLayoutProvider(session)
self.task_provider = TaskProvider(session)
self.report_img_provider = ReportImgProvider(session)
self.report_task_provider = ReportTasksProvider(session)
self.report_series_provider = ReportSeriesProvider(session)
self.project = self.task_provider.project(task.id).id
self.layout = self.layout_provider.by_name(layout)
self.layout_dict = yaml_load(self.layout.content)
self.create_base()
def __init__(self,
session: Session,
task: Task,
layout: str,
part: str = 'valid',
name: str = 'img_classify',
max_img_size: Tuple[int, int] = None,
main_metric: str = 'accuracy',
plot_count: int = 0):
self.session = session
self.task = task
self.layout = layout
self.part = part
self.name = name or 'img_classify'
self.max_img_size = max_img_size
self.main_metric = main_metric
self.plot_count = plot_count
self.dag_provider = DagProvider(session)
self.report_provider = ReportProvider(session)
self.layout_provider = ReportLayoutProvider(session)
self.task_provider = TaskProvider(session)
self.report_img_provider = ReportImgProvider(session)
self.report_task_provider = ReportTasksProvider(session)
self.report_series_provider = ReportSeriesProvider(session)
self.project = self.task_provider.project(task.id).id
self.layout = self.layout_provider.by_name(layout)
self.layout_dict = yaml_load(self.layout.content)
def create_providers(self):
self.log_info('create_providers')
self.dag_provider = DagProvider(self.session)
self.task_provider = TaskProvider(self.session)
self.file_provider = FileProvider(self.session)
self.dag_storage_provider = DagStorageProvider(self.session)
def dag_stop():
data = request_data()
provider = DagProvider(_write_session)
id = int(data['id'])
dag = provider.by_id(id, joined_load=['tasks'])
for t in dag.tasks:
celery_tasks.stop(logger, _write_session, t, dag)
return {'dag': provider.get({'id': id})['data'][0]}
def dag_model_start(session: Session, data: dict):
provider = ModelProvider(session)
model = provider.by_id(data['model_id'])
dag_provider = DagProvider(session)
dag = dag_provider.by_id(data['dag'], joined_load=[Dag.project_rel])
project = dag.project_rel
src_config = Config.from_yaml(dag.config)
pipe = src_config['pipes'][data['pipe']['name']]
equations = yaml_load(model.equations)
versions = data['pipe']['versions']
if len(versions) > 0:
version = data['pipe']['version']
pipe_equations = yaml_load(version['equations'])
found_version = versions[0]
for v in versions:
if v['name'] == version['name']:
found_version = v
break
found_version['used'] = now()
if len(pipe) == 1:
pipe[list(pipe)[0]].update(pipe_equations)
else:
pipe.update(pipe_equations)
equations[data['pipe']['name']] = versions
model.equations = yaml_dump(equations)
for v in pipe.values():
v['model_id'] = model.id
v['model_name'] = model.name
config = {
'info': {
'name': data['pipe']['name'],
'project': project.name
},
'executors': pipe
}
if model.dag:
old_dag = dag_provider.by_id(model.dag)
if old_dag.name != dag.name:
model.dag = dag.id
else:
model.dag = dag.id
provider.commit()
dag_standard(session=session,
config=config,
debug=False,
upload_files=False,
copy_files_from=data['dag'])
def create_providers(self):
self.provider = TaskProvider(self.session)
self.report_provider = ReportProvider(self.session)
self.report_tasks_provider = ReportTasksProvider(self.session)
self.report_layout_provider = ReportLayoutProvider(self.session)
self.project_provider = ProjectProvider(self.session)
self.storage = Storage(self.session)
self.dag_provider = DagProvider(self.session)
def dag_stop():
data = request_data()
provider = TaskProvider(_write_session)
id = int(data['id'])
tasks = provider.by_dag(id)
supervisor.stop_tasks(tasks)
dag_provider = DagProvider(_write_session)
return {'dag': dag_provider.get({'id': id})['data'][0]}
def describe_dag(dag, axis):
provider = DagProvider()
graph = provider.graph(dag)
status_colors = {
'not_ran': '#808080',
'queued': '#add8e6',
'in_progress': '#bfff00',
'failed': '#e83217',
'stopped': '#cb88ea',
'skipped': '#ffa500',
'success': '#006400'
}
node_color = []
edge_color = []
G = nx.DiGraph()
labels = dict()
for n in graph['nodes']:
G.add_node(n['id'])
labels[n['id']] = n['id']
node_color.append(status_colors[n['status']])
edges = []
for e in graph['edges']:
G.add_edge(e['from'], e['to'])
edges.append((e['from'], e['to']))
edge_color.append(status_colors[e['status']])
pos = nx.spring_layout(G, seed=0)
nx.draw_networkx_nodes(G,
pos,
node_color=node_color,
ax=axis,
node_size=2000)
nx.draw_networkx_labels(G,
pos,
labels,
ax=axis,
with_labels=True,
font_color='orange',
font_weight='bold',
font_size=18)
nx.draw_networkx_edges(G,
pos,
edgelist=edges,
edge_color=edge_color,
arrows=True,
arrowsize=80,
ax=axis)
axis.set_xticks([])
axis.axis('off')
axis.set_title('Graph')
def create_providers(self):
self.log_info('create_providers')
self.provider = TaskProvider(self.session)
self.report_provider = ReportProvider(self.session)
self.report_tasks_provider = ReportTasksProvider(self.session)
self.report_layout_provider = ReportLayoutProvider(self.session)
self.project_provider = ProjectProvider(self.session)
self.storage = Storage(self.session,
logger=self.logger,
component=self.component)
self.dag_provider = DagProvider(self.session)
def __init__(self, session: Session, logger=None,
component: ComponentType = None,
max_file_size: int = 10 ** 5, max_count=10 ** 3):
self.file_provider = FileProvider(session)
self.provider = DagStorageProvider(session)
self.task_provider = TaskProvider(session)
self.library_provider = DagLibraryProvider(session)
self.dag_provider = DagProvider(session)
self.logger = logger
self.component = component
self.max_file_size = max_file_size
self.max_count = max_count
def create_base(self):
self.session.commit()
self.provider = TaskProvider(self.session)
self.computer_provider = ComputerProvider(self.session)
self.docker_provider = DockerProvider(self.session)
self.auxiliary_provider = AuxiliaryProvider(self.session)
self.dag_provider = DagProvider(self.session)
self.queues = [
f'{d.computer}_{d.name}' for d in self.docker_provider.all()
if d.last_activity >= now() - datetime.timedelta(seconds=15)
]
self.auxiliary['queues'] = self.queues
def work(self):
project = ProjectProvider(self.session).by_id(self.project)
self.info(f'Task = {self.train_task} child_task: {self.child_task}')
model = Model(
created=now(),
name=self.name,
project=self.project,
equations='',
fold=self.fold
)
provider = ModelProvider(self.session)
if self.train_task:
task_provider = TaskProvider(self.session)
dag_provider = DagProvider(self.session)
task = task_provider.by_id(self.train_task)
dag = dag_provider.by_id(task.dag)
task_dir = join(TASK_FOLDER, str(self.child_task or task.id))
# get log directory
config = yaml_load(dag.config)
executor_config = config['executors'][task.executor]
catalyst_config_file = executor_config['args']['config']
catalyst_config_file = join(task_dir, catalyst_config_file)
catalyst_config = yaml_load(file=catalyst_config_file)
catalyst_logdir = catalyst_config['args']['logdir']
model.score_local = task.score
src_log = f'{task_dir}/{catalyst_logdir}'
models_dir = join(MODEL_FOLDER, project.name)
os.makedirs(models_dir, exist_ok=True)
model_path_tmp = f'{src_log}/traced.pth'
traced = trace_model_from_checkpoint(src_log, self, file=self.file)
model_path = f'{models_dir}/{model.name}.pth'
model_weight_path = f'{models_dir}/{model.name}_weight.pth'
torch.jit.save(traced, model_path_tmp)
shutil.copy(model_path_tmp, model_path)
file = self.file = 'best_full'
shutil.copy(f'{src_log}/checkpoints/{file}.pth',
model_weight_path)
provider.add(model)
def code_download():
id = int(request.args['id'])
storage = Storage(_read_session)
dag = DagProvider().by_id(id)
folder = os.path.join(TMP_FOLDER, f'{dag.id}({dag.name})')
try:
storage.download_dag(id, folder)
file_name = f'{dag.id}({dag.name}).zip'
dst = os.path.join(TMP_FOLDER, file_name)
zip_folder(folder, dst)
res = send_from_directory(TMP_FOLDER, file_name)
os.remove(dst)
return res
finally:
shutil.rmtree(folder, ignore_errors=True)
def dag_model_start(session: Session, data: dict):
provider = ModelProvider(session)
model = provider.by_id(data['model_id'])
dag = DagProvider(session
).by_id(data['dag'], joined_load=[Dag.project_rel])
project = dag.project_rel
src_config = Config.from_yaml(dag.config)
pipe = src_config['pipes'][data['pipe']]
for k, v in pipe.items():
if v.get('slot') != data['slot']:
continue
params = yaml_load(data['interface_params'])
slot = {
'interface': data['interface'],
'interface_params': params,
'slot': k,
'name': model.name,
'id': data['model_id']
}
v['slot'] = slot
config = {
'info': {
'name': data['pipe'],
'project': project.name
},
'executors': pipe
}
dag_standard(
session=session,
config=config,
debug=False,
upload_files=False,
copy_files_from=data['dag']
)
model.dag = data['dag']
model.interface = data['interface']
model.interface_params = data['interface_params']
model.slot = data['slot']
provider.commit()
def work(self):
task_provider = TaskProvider(self.session)
task = task_provider.by_id(self.train_task)
dag = DagProvider(self.session).by_id(self.dag_pipe,
joined_load=[Dag.project_rel])
task_dir = join(TASK_FOLDER, str(self.child_task or task.id))
src_log = f'{task_dir}/log'
models_dir = join(MODEL_FOLDER, dag.project_rel.name)
os.makedirs(models_dir, exist_ok=True)
self.info(f'Task = {self.task} child_task: {self.child_task}')
model_path_tmp = f'{src_log}/traced.pth'
traced = trace_model_from_checkpoint(src_log, self)
model = Model(dag=self.dag_pipe,
interface=self.interface,
slot=self.slot,
score_local=task.score,
created=now(),
name=self.name,
project=dag.project,
interface_params=yaml_dump(self.interface_params))
provider = ModelProvider(self.session)
provider.add(model, commit=False)
try:
model_path = f'{models_dir}/{model.name}.pth'
model_weight_path = f'{models_dir}/{model.name}_weight.pth'
torch.jit.save(traced, model_path_tmp)
shutil.copy(model_path_tmp, model_path)
shutil.copy(f'{src_log}/checkpoints/best.pth', model_weight_path)
interface_params = yaml_load(model.interface_params)
interface_params['file'] = join('models', model.name + '.pth')
model.interface_params = yaml_dump(interface_params)
provider.update()
except Exception as e:
provider.rollback()
raise e
def dag_tag_add():
data = request_data()
provider = DagProvider(_write_session)
tag = DagTag(dag=data['dag'], tag=data['tag'])
provider.add(tag)
def dag_remove():
id = request_data()['id']
celery_tasks.remove_dag(_write_session, id)
DagProvider(_write_session).remove(id)
def dag_start():
data = request_data()
provider = DagProvider(_write_session)
task_provider = TaskProvider(_write_session)
id = int(data['id'])
dag = provider.by_id(id, joined_load=['tasks'])
can_start_statuses = [
TaskStatus.Failed.value, TaskStatus.Skipped.value,
TaskStatus.Stopped.value
]
tasks = list(dag.tasks)
def find_resume(task):
children = task_provider.children(task.id)
children = sorted(children, key=lambda x: x.id, reverse=True)
if len(children) > 0:
for c in children:
if c.parent != task.id:
continue
info = yaml_load(c.additional_info)
if 'distr_info' not in info:
continue
if info['distr_info']['rank'] == 0:
return {
'master_computer': c.computer_assigned,
'master_task_id': c.id,
'load_last': True
}
raise Exception('Master task not found')
else:
return {
'master_computer': task.computer_assigned,
'master_task_id': task.id,
'load_last': True
}
for t in tasks:
if t.status not in can_start_statuses:
continue
if t.parent:
continue
info = yaml_load(t.additional_info)
info['resume'] = find_resume(t)
t.additional_info = yaml_dump(info)
t.status = TaskStatus.NotRan.value
t.pid = None
t.started = None
t.finished = None
t.computer_assigned = None
t.celery_id = None
t.worker_index = None
t.docker_assigned = None
provider.commit()
def file_before_create(mapper, connection, target):
provider = DagProvider(_session)
dag = provider.by_id(target.dag)
dag.file_size += target.size
provider.commit()
def dags():
data = request_data()
options = PaginatorOptions(**data['paginator'])
provider = DagProvider(_read_session)
res = provider.get(data, options)
return res
def graph():
id = request_data()
res = DagProvider(_read_session).graph(id)
return res
def config():
id = request_data()
res = DagProvider(_read_session).config(id)
return {'data': res}
def remove_all_dags():
data = request_data()
provider = DagProvider(_write_session)
dags = provider.by_project(data['project'])
provider.remove_all([d.id for d in dags])
def dag_tags():
data = request_data()
provider = DagProvider(_write_session)
return provider.tags(data['name'])
def dag_tag_remove():
data = request_data()
provider = DagProvider(_write_session)
provider.remove_tag(dag=data['dag'], tag=data['tag'])
def statuses(folder: str = None):
rows = []
folder_status = 'OK'
folder_comment = ''
folders = [
ROOT_FOLDER, DATA_FOLDER, MODEL_FOLDER, TASK_FOLDER, LOG_FOLDER,
CONFIG_FOLDER, DB_FOLDER, REPORT_FOLDER, TMP_FOLDER
]
for f in folders:
if not exists(f):
folder_status = 'ERROR'
folder_comment = f'folder {f} does not exist'
files = [join(CONFIG_FOLDER, '.env')]
for f in files:
if not exists(f):
folder_status = 'ERROR'
folder_comment = f'file {f} does not exist'
rows.append({
'name': 'Folders',
'status': folder_status,
'comment': folder_comment
})
database_status = 'OK'
database_comment = f'DB_TYPE = {DB_TYPE}'
try:
provider = DagProvider()
provider.count()
except Exception:
database_status = 'ERROR'
database_comment += ' ' + traceback.format_exc()
rows.append({
'name': 'Database',
'status': database_status,
'comment': database_comment
})
redis_status = 'OK'
redis_comment = f''
try:
celery.backend.client.echo(1)
except Exception:
redis_status = 'ERROR'
redis_comment += ' ' + traceback.format_exc()
rows.append({
'name': 'Redis',
'status': redis_status,
'comment': redis_comment
})
if database_status == 'OK':
migrate_status = 'OK'
repository_folder = join(dirname(__file__), 'migration')
repository_version = api.version(repository_folder)
db_version = api.db_version(SA_CONNECTION_STRING, repository_folder)
if db_version != repository_version:
migrate_status = 'ERROR'
migrate_comment = f'Repository version = {repository_version} ' \
f'Db version = {db_version}'
else:
migrate_comment = f'version: {db_version}'
rows.append({
'name': 'Migrate',
'status': migrate_status,
'comment': migrate_comment
})
df = pd.DataFrame(rows)
if folder is not None:
print('Statuses:')
print(df)
df.to_csv(join(folder, 'statuses.csv'), index=False)
return df