From 14c70c0e68eed4231102dfdd2066128e8c07edfd Mon Sep 17 00:00:00 2001 From: Alexey Gorbatovski Date: Fri, 4 Jul 2025 17:13:34 +0300 Subject: [PATCH 1/2] Include run name in wandb initialization in BaseEnv --- atroposlib/envs/base.py | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/atroposlib/envs/base.py b/atroposlib/envs/base.py index 774e4130..45d88332 100644 --- a/atroposlib/envs/base.py +++ b/atroposlib/envs/base.py @@ -394,21 +394,28 @@ class BaseEnv(ABC): # Setup wandb getting the group and project via the server while self.wandb_project is None: async with aiohttp.ClientSession() as session: - async with session.get( - f"{self.config.rollout_server_url}/wandb_info" - ) as resp: + async with session.get(f"{self.config.rollout_server_url}/wandb_info") as resp: data = await parse_http_response(resp, logger) self.wandb_group = data["group"] self.wandb_project = data["project"] + if self.wandb_project is None: await asyncio.sleep(1) - else: - wandb.init( - project=self.wandb_project, - group=self.wandb_group, - config=self.config.model_dump(), - ) - break + continue + + wandb_run_name = None + if self.config.wandb_name: + random_id = "".join(random.choices(string.ascii_lowercase, k=6)) + current_date = datetime.now().strftime("%Y-%m-%d") + wandb_run_name = f"{self.config.wandb_name}-{current_date}-{random_id}" + + wandb.init( + name=wandb_run_name, + project=self.wandb_project, + group=self.wandb_group, + config=self.config.model_dump(), + ) + break @retry( stop=stop_after_attempt(3), From ee5257522a5292fc2985c62d6de794755d783fa6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 4 Jul 2025 14:34:33 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- atroposlib/envs/base.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/atroposlib/envs/base.py b/atroposlib/envs/base.py index 45d88332..28c043a8 100644 --- a/atroposlib/envs/base.py +++ b/atroposlib/envs/base.py @@ -394,7 +394,9 @@ class BaseEnv(ABC): # Setup wandb getting the group and project via the server while self.wandb_project is None: async with aiohttp.ClientSession() as session: - async with session.get(f"{self.config.rollout_server_url}/wandb_info") as resp: + async with session.get( + f"{self.config.rollout_server_url}/wandb_info" + ) as resp: data = await parse_http_response(resp, logger) self.wandb_group = data["group"] self.wandb_project = data["project"] @@ -407,7 +409,9 @@ class BaseEnv(ABC): if self.config.wandb_name: random_id = "".join(random.choices(string.ascii_lowercase, k=6)) current_date = datetime.now().strftime("%Y-%m-%d") - wandb_run_name = f"{self.config.wandb_name}-{current_date}-{random_id}" + wandb_run_name = ( + f"{self.config.wandb_name}-{current_date}-{random_id}" + ) wandb.init( name=wandb_run_name,