From 9dacaeece93e59a23a0479790fe528a302df02c1 Mon Sep 17 00:00:00 2001 From: c-bata Date: Fri, 8 Sep 2023 18:28:50 +0900 Subject: [PATCH] Revert changes in upload_artifact api --- optuna_dashboard/artifact/_backend.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/optuna_dashboard/artifact/_backend.py b/optuna_dashboard/artifact/_backend.py index 63fec1c5..10a5b1ee 100644 --- a/optuna_dashboard/artifact/_backend.py +++ b/optuna_dashboard/artifact/_backend.py @@ -166,7 +166,7 @@ def register_artifact_route( def upload_artifact( backend: ArtifactBackend, - study_or_trial: optuna.Trial | optuna.Study, + trial: optuna.Trial, file_path: str, *, mimetype: Optional[str] = None, @@ -202,6 +202,8 @@ def upload_artifact( ) filename = os.path.basename(file_path) + storage = trial.storage + trial_id = trial._trial_id artifact_id = str(uuid.uuid4()) guess_mimetype, guess_encoding = mimetypes.guess_type(filename) artifact: ArtifactMeta = { @@ -211,15 +213,7 @@ def upload_artifact( "filename": filename, } attr_key = ARTIFACTS_ATTR_PREFIX + artifact_id - - if isinstance(study_or_trial, optuna.Study): - storage = study_or_trial._storage - study_id = study_or_trial._study_id - storage.set_study_system_attr(study_id, attr_key, json.dumps(artifact)) - else: - storage = study_or_trial.storage - trial_id = study_or_trial._trial_id - storage.set_trial_system_attr(trial_id, attr_key, json.dumps(artifact)) + storage.set_trial_system_attr(trial_id, attr_key, json.dumps(artifact)) with open(file_path, "rb") as f: backend.write(artifact_id, f)