Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
…8684/files#r1854943374, use `ResourceNotFoundError` if study is not found.

PiperOrigin-RevId: 699334215
  • Loading branch information
xingyousong authored and copybara-github committed Nov 23, 2024
1 parent 2e7dce0 commit 2f43cb7
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 26 deletions.
4 changes: 2 additions & 2 deletions vizier/_src/pyglove/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -287,7 +287,7 @@ def _setup_study(self) -> bool:
)

return is_chief
except KeyError:
except client_abc.ResourceNotFoundError:
# Study does not exist.
if self._run_mode == TunerMode.SECONDARY:
pg.logging.info(
Expand Down Expand Up @@ -481,7 +481,7 @@ def _wait_for_study(
while True:
try:
return self._tuner.load_study(owner, name)
except KeyError:
except client_abc.ResourceNotFoundError:
logging.info(
'Study %s (owner=%s) does not exist. Retrying after 10 seconds.',
name,
Expand Down
39 changes: 15 additions & 24 deletions vizier/_src/service/clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@

from typing import Any, Callable, Iterable, Iterator, List, Mapping, Optional, Type
import attr

from vizier._src.service import constants
from vizier._src.service import resources
from vizier._src.service import vizier_client
Expand Down Expand Up @@ -138,12 +137,10 @@ def _trial_client(self, trial: vz.Trial) -> Trial:
def suggest(
self, *, count: Optional[int] = None, client_id: str = 'default_client_id'
) -> List[Trial]:
return [
self._trial_client(t)
for t in self._client.get_suggestions(
count, client_id_override=client_id
)
]
suggestions = self._client.get_suggestions(
count, client_id_override=client_id
)
return [self._trial_client(t) for t in suggestions]

def delete(self) -> None:
self._client.delete_study()
Expand Down Expand Up @@ -181,7 +178,7 @@ def get_trial(self, trial_id: int) -> Trial:
return self._trial_client(trial)
except KeyError as err:
raise ResourceNotFoundError(
f'Study f{self.resource_name} does not have Trial {trial_id}.'
f'Study {self.resource_name} does not have Trial {trial_id}.'
) from err

def optimal_trials(self, count: Optional[int] = None) -> TrialIterable:
Expand All @@ -207,14 +204,11 @@ def materialize_state(self) -> vz.StudyState:

@classmethod
def from_resource_name(cls: Type['Study'], name: str) -> 'Study':
client = vizier_client.VizierClient(
name,
constants.UNUSED_CLIENT_ID,
)
client = vizier_client.VizierClient(name, constants.UNUSED_CLIENT_ID)
try:
_ = client.get_study_config() # Make sure study exists.
except Exception as err:
raise KeyError(f'Study {name} does not exist.') from err
raise ResourceNotFoundError(f'Study {name} does not exist.') from err
return Study(client)

@classmethod
Expand All @@ -233,10 +227,8 @@ def from_owner_and_id(
Raises:
ResourceNotFoundError.
"""
study_resource_name = resources.StudyResource(
owner_id=owner, study_id=study_id
).name
return cls.from_resource_name(study_resource_name)
resource = resources.StudyResource(owner, study_id)
return cls.from_resource_name(resource.name)

@classmethod
def from_study_config(
Expand All @@ -257,11 +249,10 @@ def from_study_config(
"""
# Use a dummy client_id, since for multi-worker workflows, each worker is
# expected to provide a client_id in the suggest() call.
return Study(
vizier_client.create_or_load_study(
owner_id=owner,
client_id=constants.UNUSED_CLIENT_ID,
study_id=study_id,
study_config=config,
)
client = vizier_client.create_or_load_study(
owner_id=owner,
client_id=constants.UNUSED_CLIENT_ID,
study_id=study_id,
study_config=config,
)
return Study(client)
2 changes: 2 additions & 0 deletions vizier/service/clients/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,9 @@
from __future__ import annotations

"""OSS Client API."""

from vizier._src.service.clients import environment_variables
from vizier._src.service.clients import ResourceNotFoundError
from vizier._src.service.clients import Study
from vizier._src.service.clients import Trial
from vizier._src.service.clients import TrialIterable

0 comments on commit 2f43cb7

Please sign in to comment.