Skip to content

Commit

Permalink
PLAT-1019: Refresh the project ID in the gretel-agent
Browse files Browse the repository at this point in the history
GitOrigin-RevId: 148e6d9e9952c99fb0744cf7bc7050dde52ae7ba
  • Loading branch information
misberner authored and mckornfield committed Aug 14, 2023
1 parent 0ecf14f commit bacaee2
Showing 1 changed file with 13 additions and 5 deletions.
18 changes: 13 additions & 5 deletions src/gretel_client/agents/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@

import requests

from backports.cached_property import cached_property

import gretel_client.agents.agent_telemetry as telemetry

from gretel_client.agents.drivers.driver import ComputeUnit, Driver
Expand Down Expand Up @@ -79,6 +77,9 @@ class AgentConfig:
_max_runtime_seconds: Optional[int] = None
"""TODO: implement"""

_project_id: Optional[str] = None
"""Cached project ID."""

def __post_init__(self):
if not self._max_runtime_seconds:
self._max_runtime_seconds = self._lookup_max_runtime()
Expand Down Expand Up @@ -133,13 +134,19 @@ def _lookup_max_jobs_active(self) -> int:
def as_dict(self) -> dict:
return asdict(self)

@cached_property
@property
def project_id(self) -> Optional[str]:
if self.project is None:
return None

project = get_project(name=self.project)
return project.project_id
if self._project_id is None:
project = get_project(name=self.project)
self._project_id = project.project_id

return self._project_id

def invalidate_project_id(self) -> None:
self._project_id = None


@dataclass
Expand Down Expand Up @@ -325,6 +332,7 @@ def __next__(self) -> Optional[Job]:
self._logger.warning(
f"There was a problem calling the jobs endpoint {ex}"
)
self._agent_config.invalidate_project_id()
if job:
return job
self._interrupt.wait(wait_secs)
Expand Down

0 comments on commit bacaee2

Please sign in to comment.