add trace decorators

neuro-inc · Apr 28, 2021 · 0442b7c · 0442b7c
1 parent 84a6ced
commit 0442b7c
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/platform_api/orchestrator/job_policy_enforcer.py b/platform_api/orchestrator/job_policy_enforcer.py
@@ -20,6 +20,7 @@
 
 from neuro_auth_client import AuthClient
 from notifications_client import Client as NotificationsClient, CreditsWillRunOutSoon
+from platform_logging import new_trace, trace
 
 from platform_api.admin_client import AdminClient
 from platform_api.config import JobPolicyEnforcerConfig
@@ -73,6 +74,7 @@ async def _notify_user_if_needed(
         )
         self._sent[notification_key] = credits
 
+    @trace
     async def enforce(self) -> None:
         running_jobs = self._jobs_service.iter_all_jobs(
             job_filter=JobFilter(
@@ -103,6 +105,7 @@ class RuntimeLimitEnforcer(JobPolicyEnforcer):
     def __init__(self, service: JobsService):
         self._service = service
 
+    @trace
     async def enforce(self) -> None:
         active_jobs = await self._service.get_all_jobs(
             job_filter=JobFilter(
@@ -144,6 +147,7 @@ def _groupby(
             res[key(item)].append(item)
         return res
 
+    @trace
     async def enforce(self) -> None:
         jobs = await self._service.get_all_jobs(
             job_filter=JobFilter(
@@ -184,6 +188,7 @@ def __init__(self, service: JobsService, admin_client: AdminClient):
         self._service = service
         self._admin_client = admin_client
 
+    @trace
     async def enforce(self) -> None:
         coros = [
             self._bill_single_wrapper(job)
@@ -258,6 +263,7 @@ async def _run(self) -> None:
                 delay = 0
             await asyncio.sleep(delay)
 
+    @new_trace
     async def _run_once(self) -> None:
         for enforcer in self._enforcers:
             try:

diff --git a/platform_api/orchestrator/jobs_poller.py b/platform_api/orchestrator/jobs_poller.py
@@ -7,6 +7,7 @@
 import aiohttp
 from iso8601 import iso8601
 from multidict import MultiDict
+from platform_logging import new_trace
 from yarl import URL
 
 from ..cluster import SingleClusterUpdater
@@ -280,6 +281,7 @@ async def _run(self) -> None:
             await self._run_once()
             await self._wait()
 
+    @new_trace
     async def _run_once(self) -> None:
         try:
             await self._cluster_updater.do_update()