From 49e897933f84fbe3ca47f5863bda9d63582dcc0a Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Mar 2026 19:11:34 +0000 Subject: [PATCH] feat(wandb): store Flyte execution ID in wandb run config for reverse lookups Adds flyte_execution_id, flyte_execution_project, and flyte_execution_domain to wandb run config during remote execution. This enables looking up the corresponding Flyte execution from a WandB run, making the WandB<->Flyte mapping a single lookup in either direction. Co-Authored-By: ben@exa.ai --- .../flytekit-wandb/flytekitplugins/wandb/tracking.py | 12 ++++++++++++ plugins/flytekit-wandb/tests/test_wandb_init.py | 12 ++++++++++++ 2 files changed, 24 insertions(+) diff --git a/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py b/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py index 8bf033f2ab..7048f6bb02 100644 --- a/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py +++ b/plugins/flytekit-wandb/flytekitplugins/wandb/tracking.py @@ -96,6 +96,18 @@ def execute(self, *args, **kwargs): run = wandb.init(project=self.project, entity=self.entity, id=wand_id, **self.init_kwargs) + # Store Flyte execution metadata in wandb run config for reverse lookups. + if not is_local_execution: + exec_id = ctx.user_space_params.execution_id + run.config.update( + { + "flyte_execution_id": exec_id.name, + "flyte_execution_project": exec_id.project, + "flyte_execution_domain": exec_id.domain, + }, + allow_val_change=True, + ) + # If FLYTE_EXECUTION_URL is defined, inject it into wandb to link back to the execution. execution_url = os.getenv("FLYTE_EXECUTION_URL") if execution_url is not None: diff --git a/plugins/flytekit-wandb/tests/test_wandb_init.py b/plugins/flytekit-wandb/tests/test_wandb_init.py index 1db9f5c0fd..f6c1984cf3 100644 --- a/plugins/flytekit-wandb/tests/test_wandb_init.py +++ b/plugins/flytekit-wandb/tests/test_wandb_init.py @@ -67,6 +67,8 @@ def test_non_local_execution(wandb_mock, manager_mock, monkeypatch): ctx_mock.user_space_params.secrets.get.return_value = "this_is_the_secret" ctx_mock.user_space_params.execution_id.name = "my_execution_id" + ctx_mock.user_space_params.execution_id.project = "my_project" + ctx_mock.user_space_params.execution_id.domain = "staging" manager_mock.current_context.return_value = ctx_mock execution_url = "http://execution_url.com/afsdfsafafasdfs" @@ -83,6 +85,16 @@ def test_non_local_execution(wandb_mock, manager_mock, monkeypatch): wandb_mock.login.assert_called_with(key="this_is_the_secret", host="https://api.wandb.ai") assert run_mock.notes == f"[Execution URL]({execution_url})" + # Verify Flyte execution metadata is stored in wandb config + run_mock.config.update.assert_called_once_with( + { + "flyte_execution_id": "my_execution_id", + "flyte_execution_project": "my_project", + "flyte_execution_domain": "staging", + }, + allow_val_change=True, + ) + def test_errors(): with pytest.raises(ValueError, match="project must be set"):