Skip to content

Commit

Permalink
Add error handlings to the workflow
Browse files Browse the repository at this point in the history
  • Loading branch information
@GoldenAnpu
GoldenAnpu committed Jul 12, 2024
1 parent 53686ba commit 267391f
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 54 deletions.
17 changes: 10 additions & 7 deletions supervisely/serve/src/workflow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,13 +34,16 @@ def check_instance_ver_compatibility(self):

@check_compatibility
def add_input(self, checkpoint_url: str):
meta = {"customNodeSettings": {"title": "<h4>Serve Custom Model</h4>"}}
sly.logger.debug(f"Workflow Input: Checkpoint URL - {checkpoint_url}")
if checkpoint_url and self.api.file.exists(sly.env.team_id(), checkpoint_url):
self.api.app.workflow.add_input_file(checkpoint_url, model_weight=True, meta=meta)
else:
sly.logger.debug(f"Checkpoint {checkpoint_url} not found in Team Files. Cannot set workflow input")

try:
meta = {"customNodeSettings": {"title": "<h4>Serve Custom Model</h4>"}}
sly.logger.debug(f"Workflow Input: Checkpoint URL - {checkpoint_url}")
if checkpoint_url and self.api.file.exists(sly.env.team_id(), checkpoint_url):
self.api.app.workflow.add_input_file(checkpoint_url, model_weight=True, meta=meta)
else:
sly.logger.debug(f"Checkpoint {checkpoint_url} not found in Team Files. Cannot set workflow input")
except Exception as e:
sly.logger.error(f"Failed to add input to the workflow: {e}")

@check_compatibility
def add_output(self):
raise NotImplementedError("add_output is not implemented in this workflow")
104 changes: 57 additions & 47 deletions supervisely/train/src/workflow.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,55 +37,65 @@ def check_instance_ver_compatibility(self):

@check_compatibility
def add_input(self, project_info: sly.ProjectInfo, state: dict):
project_version_id = self.api.project.version.create(
project_info, "Train YOLO v5", f"This backup was created automatically by Supervisely before the Train YOLO task with ID: {self.api.task_id}"
)
if project_version_id is None:
project_version_id = project_info.version.get("id", None) if project_info.version else None
self.api.app.workflow.add_input_project(project_info.id, version_id=project_version_id)
file_info = False
if state["weightsInitialization"] is not None and state["weightsInitialization"] == "custom":
file_info = self.api.file.get_info_by_path(sly.env.team_id(), state["_weightsPath"])
self.api.app.workflow.add_input_file(file_info, model_weight=True)
sly.logger.debug(f"Workflow Input: Project ID - {project_info.id}, Project Version ID - {project_version_id}, Input File - {True if file_info else False}")
try:
project_version_id = self.api.project.version.create(
project_info, "Train YOLO v5", f"This backup was created automatically by Supervisely before the Train YOLO task with ID: {self.api.task_id}"
)
except Exception as e:
sly.logger.error(f"Failed to create a project version: {e}")
project_version_id = None

try:
if project_version_id is None:
project_version_id = project_info.version.get("id", None) if project_info.version else None
self.api.app.workflow.add_input_project(project_info.id, version_id=project_version_id)
file_info = False
if state["weightsInitialization"] is not None and state["weightsInitialization"] == "custom":
file_info = self.api.file.get_info_by_path(sly.env.team_id(), state["_weightsPath"])
self.api.app.workflow.add_input_file(file_info, model_weight=True)
sly.logger.debug(f"Workflow Input: Project ID - {project_info.id}, Project Version ID - {project_version_id}, Input File - {True if file_info else False}")
except Exception as e:
sly.logger.error(f"Failed to add input to the workflow: {e}")

@check_compatibility
def add_output(self, state: dict, team_files_dir: str):
weights_dir_in_team_files = os.path.join(team_files_dir, "weights")
files_info = self.api.file.list(sly.env.team_id(), weights_dir_in_team_files, return_type="fileinfo")
best_filename_info = None
for file_info in files_info:
if "best" in file_info.name:
best_filename_info = file_info
break
if best_filename_info:
module_id = self.api.task.get_info_by_id(self.api.task_id).get("meta", {}).get("app", {}).get("id")
if state["weightsInitialization"] is not None and state["weightsInitialization"] == "custom":
model_name = "Custom Model"
else:
model_name = "YOLOv5"

meta = {
"customNodeSettings": {
"title": f"<h4>Train {model_name}</h4>",
"mainLink": {
"url": f"/apps/{module_id}/sessions/{self.api.task_id}" if module_id else f"apps/sessions/{self.api.task_id}",
"title": "Show Results"
}
},
"customRelationSettings": {
"icon": {
"icon": "zmdi-folder",
"color": "#FFA500",
"backgroundColor": "#FFE8BE"
try:
weights_dir_in_team_files = os.path.join(team_files_dir, "weights")
files_info = self.api.file.list(sly.env.team_id(), weights_dir_in_team_files, return_type="fileinfo")
best_filename_info = None
for file_info in files_info:
if "best" in file_info.name:
best_filename_info = file_info
break
if best_filename_info:
module_id = self.api.task.get_info_by_id(self.api.task_id).get("meta", {}).get("app", {}).get("id")
if state["weightsInitialization"] is not None and state["weightsInitialization"] == "custom":
model_name = "Custom Model"
else:
model_name = "YOLOv5"

meta = {
"customNodeSettings": {
"title": f"<h4>Train {model_name}</h4>",
"mainLink": {
"url": f"/apps/{module_id}/sessions/{self.api.task_id}" if module_id else f"apps/sessions/{self.api.task_id}",
"title": "Show Results"
}
},
"title": "<h4>Checkpoints</h4>",
"mainLink": {"url": f"/files/{best_filename_info.id}/true", "title": "Open Folder"}
"customRelationSettings": {
"icon": {
"icon": "zmdi-folder",
"color": "#FFA500",
"backgroundColor": "#FFE8BE"
},
"title": "<h4>Checkpoints</h4>",
"mainLink": {"url": f"/files/{best_filename_info.id}/true", "title": "Open Folder"}
}
}
}
sly.logger.debug(f"Workflow Output: Team Files dir - {team_files_dir}, Best filename - {best_filename_info.name}")
sly.logger.debug(f"Workflow Output: meta \n {meta}")
self.api.app.workflow.add_output_file(best_filename_info, model_weight=True, meta=meta)
else:
sly.logger.debug(f"File with the best weighs not found in Team Files. Cannot set workflow output.")

sly.logger.debug(f"Workflow Output: Team Files dir - {team_files_dir}, Best filename - {best_filename_info.name}")
sly.logger.debug(f"Workflow Output: meta \n {meta}")
self.api.app.workflow.add_output_file(best_filename_info, model_weight=True, meta=meta)
else:
sly.logger.debug(f"File with the best weighs not found in Team Files. Cannot set workflow output.")
except Exception as e:
sly.logger.error(f"Failed to add output to the workflow: {e}")

0 comments on commit 267391f

Please sign in to comment.