From 06a0d5d3bb03e619231229890d0ffe2a133cc62c Mon Sep 17 00:00:00 2001 From: c-bata Date: Wed, 24 Jan 2024 10:06:40 +0900 Subject: [PATCH] Remove deprecation warnings --- optuna_dashboard/_serializer.py | 2 +- optuna_dashboard/artifact/_backend.py | 16 +++++++++++----- python_tests/artifact/test_backend.py | 7 +++++-- .../artifact/test_optuna_compatibility.py | 5 ++++- 4 files changed, 21 insertions(+), 9 deletions(-) diff --git a/optuna_dashboard/_serializer.py b/optuna_dashboard/_serializer.py index 7030abec..b280d411 100644 --- a/optuna_dashboard/_serializer.py +++ b/optuna_dashboard/_serializer.py @@ -246,7 +246,7 @@ def serialize_frozen_trial( ], "user_attrs": serialize_attrs(trial.user_attrs), "note": note.get_note_from_system_attrs(study_system_attrs, trial._trial_id), - "artifacts": list_trial_artifacts(study_system_attrs, trial), + "artifacts": list_trial_artifacts(study_system_attrs, trial_system_attrs, trial), "constraints": trial_system_attrs.get(CONSTRAINTS_KEY, []), } diff --git a/optuna_dashboard/artifact/_backend.py b/optuna_dashboard/artifact/_backend.py index a76568e2..941805de 100644 --- a/optuna_dashboard/artifact/_backend.py +++ b/optuna_dashboard/artifact/_backend.py @@ -139,10 +139,12 @@ def register_artifact_route( storage.set_trial_system_attr(trial_id, attr_key, json.dumps(artifact)) response.status = 201 - trial = storage.get_trial(trial_id) # Fetch trial.system_attrs again. + study_system_attrs = storage.get_study_system_attrs(study_id) + trial_system_attrs = storage.get_trial_system_attrs(trial_id) + artifacts = list_trial_artifacts(study_system_attrs, trial_system_attrs, trial) return { "artifact_id": artifact_id, - "artifacts": list_trial_artifacts(storage.get_study_system_attrs(study_id), trial), + "artifacts": artifacts, } @app.post("/api/artifacts/") @@ -322,7 +324,11 @@ def delete_all_artifacts(backend: ArtifactStore, storage: BaseStorage, study_id: study_system_attrs = storage.get_study_system_attrs(study_id) artifact_metas.extend(list_study_artifacts(study_system_attrs)) for trial in storage.get_all_trials(study_id): - trial_artifacts = list_trial_artifacts(study_system_attrs, trial) + trial_system_attrs = getattr(trial, "_system_attrs") + if trial_system_attrs is None: + # This is unreachable line until Optuna v5.0.0 release. + trial_system_attrs = storage.get_trial_system_attrs(trial._trial_id) + trial_artifacts = list_trial_artifacts(study_system_attrs, trial_system_attrs, trial) artifact_metas.extend(trial_artifacts) for meta in artifact_metas: @@ -339,7 +345,7 @@ def list_study_artifacts(study_system_attrs: dict[str, Any]) -> list[ArtifactMet def list_trial_artifacts( - study_system_attrs: dict[str, Any], trial: FrozenTrial + study_system_attrs: dict[str, Any], trial_system_attrs: dict[str, Any], trial: FrozenTrial ) -> list[ArtifactMeta]: # Collect ArtifactMeta from study_system_attrs due to backward compatibility. dashboard_artifact_metas = [ @@ -353,7 +359,7 @@ def list_trial_artifacts( # See https://github.com/optuna/optuna/blob/f827582a8/optuna/artifacts/_upload.py#L16 optuna_artifact_metas = [ json.loads(value) - for key, value in trial.system_attrs.items() + for key, value in trial_system_attrs.items() if key.startswith(ARTIFACTS_ATTR_PREFIX) ] artifact_metas = dashboard_artifact_metas + optuna_artifact_metas diff --git a/python_tests/artifact/test_backend.py b/python_tests/artifact/test_backend.py index 8325343b..90004f11 100644 --- a/python_tests/artifact/test_backend.py +++ b/python_tests/artifact/test_backend.py @@ -84,9 +84,12 @@ def test_delete_all_artifacts(init_storage_with_artifact_meta: MagicMock) -> Non def test_list_trial_artifacts(init_storage_with_artifact_meta: MagicMock) -> None: storage = init_storage_with_artifact_meta - trial = MagicMock(_trial_id=0, system_attrs=storage.get_trial_system_attrs(0)) + trial_system_attrs = storage.get_trial_system_attrs(0) + trial = MagicMock(_trial_id=0, system_attrs={}) - actual = _backend.list_trial_artifacts(storage.get_study_system_attrs(0), trial) + actual = _backend.list_trial_artifacts( + storage.get_study_system_attrs(0), trial_system_attrs, trial + ) assert actual == [ {"artifact_id": "id0", "filename": "foo.txt"}, {"artifact_id": "id1", "filename": "bar.txt"}, diff --git a/python_tests/artifact/test_optuna_compatibility.py b/python_tests/artifact/test_optuna_compatibility.py index 82d36450..1ade64ab 100644 --- a/python_tests/artifact/test_optuna_compatibility.py +++ b/python_tests/artifact/test_optuna_compatibility.py @@ -36,8 +36,11 @@ def test_list_optuna_trial_artifacts() -> None: study.tell(trial, 0.0) study_system_attrs = storage.get_study_system_attrs(study._study_id) + trial_system_attrs = storage.get_trial_system_attrs(trial._trial_id) frozen_trial = storage.get_trial(trial._trial_id) - artifact_meta_list = list_trial_artifacts(study_system_attrs, frozen_trial) + artifact_meta_list = list_trial_artifacts( + study_system_attrs, trial_system_attrs, frozen_trial + ) assert len(artifact_meta_list) == 1 artifact_id = artifact_meta_list[0]["artifact_id"]