freesolo-flash-dev 1.2.63__tar.gz → 1.2.64__tar.gz

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/CHANGELOG.md

@@ -54,6 +54,10 @@ This file starts at 1.1.35. Earlier releases are not reconstructed here; use
 
 ### Fixed
 
+- Completed verl children now finish each captured W&B run before the W&B service shuts down,
+  avoiding teardown tracebacks, and `flash env setup` now gives every starter config explicit W&B
+  project and folder-derived run metadata.
+
 - An ordered `[gpu] type` list is now priced over every class it declares acceptable, not just the
   first. Allocation cost-ranks the whole set, so quoting the head alone priced a shape the run may
   never be given — an authored `["B200", "H100"]` quoted roughly 3x the H100 the allocator would

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: freesolo-flash-dev
-Version: 1.2.63
+Version: 1.2.64
 Summary: Flash — managed LoRA post-training (SFT/GRPO/OPD) for Freesolo environments, driven by the `flash` CLI
 Project-URL: Homepage, https://github.com/freesolo-co/flash
 Project-URL: Repository, https://github.com/freesolo-co/flash

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/cli/commands/env/setup.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import json
 import sys
 import tomllib
+from collections.abc import Mapping
 from pathlib import Path
 
 from flash._internal.channel import CLI_NAME
@@ -13,6 +14,7 @@ from flash.cli.commands.env.retained import (
     _warn_if_environment_form_disagrees,
     _warn_if_retained_starter_files_describe_another_plane,
 )
+from flash.cli.commands.env.setup_wandb import _wandb_block, _wandb_project
 from flash.cli.scaffold import TRAINING_MD
 from flash.cli.ui import render
 from flash.envs.evaluations import _DEFAULT_EVALUATIONS_PATH
@@ -358,22 +360,14 @@ _REASONING_OPTIONS = [
 _SKIP_TRACES = ""  # sentinel option value: scaffold the starter rows instead of importing
 
 
-def _require_setup_project(args) -> str:
+def _require_setup_project(args, *, api_url: str | None, api_key: str) -> Mapping[str, object]:
     """Resolve and validate the one explicit project used by the whole scaffold."""
-    from flash.client import ClientError, list_projects, resolve_project_id
-    from flash.client.config import load_credentials
+    from flash.client import ClientError, list_projects, resolve_project
     from flash.client.http import has_freesolo_backend
 
-    api_url, api_key = load_credentials()
-    if not api_key:
-        raise ClientError(
-            f"not logged in. Run `{CLI_NAME} login` before `{CLI_NAME} env setup` "
-            "so the project can be validated"
-        )
-
     supplied = str(getattr(args, "project", "") or "").strip()
     if supplied:
-        return resolve_project_id(supplied, api_key, api_url)
+        return resolve_project(supplied, api_key, api_url)
     if not _setup_interactive(args):
         raise ClientError(
             "--project PROJECT_UUID is required in noninteractive mode, with redirected stdin, or with --yes"
@@ -395,7 +389,7 @@ def _require_setup_project(args) -> str:
             f"`{CLI_NAME} projects create NAME`"
         )
     selected = render.select_required("Choose the Freesolo project for this environment", options)
-    return resolve_project_id(selected, api_key, api_url)
+    return resolve_project(selected, api_key, api_url)
 
 
 # Guidance the starter .py docstrings carry, per plane kind. The templates hold placeholders
@@ -450,7 +444,7 @@ def _render_starter(template: str, project_id: str, *, can_publish: bool) -> str
     return template.replace("PROJECT_UUID", project_id)
 
 
-def _plane_can_publish_environments() -> bool:
+def _plane_can_publish_environments(api_url: str | None) -> bool:
     """Whether `flash env push` can actually publish for the logged-in plane.
 
     Publishing uploads to the plane, which commits into the managed environment hub
@@ -462,10 +456,8 @@ def _plane_can_publish_environments() -> bool:
     operator-run Freesolo-compatible backend has a project directory (which is what that helper
     answers) but still cannot publish to the hub.
     """
-    from flash.client.config import load_credentials
     from flash.serve.urls import is_freesolo_hosted_url
 
-    api_url, _ = load_credentials()
     # Unset means the built-in default, which is the managed plane.
     return api_url is None or is_freesolo_hosted_url(api_url)
 
@@ -553,7 +545,7 @@ def _existing_reasoning(configs: tuple[Path, ...]) -> bool | None:
     return next(iter(found.values()), None)
 
 
-def _traces_dataset(args, project_id: str) -> str | None:
+def _traces_dataset(args, project_id: str, api_key: str) -> str | None:
     """Optionally export traces from the already-selected project."""
     if not _setup_interactive(args):
         return None
@@ -568,7 +560,7 @@ def _traces_dataset(args, project_id: str) -> str | None:
     if use_traces == _SKIP_TRACES:
         return None
     try:
-        exported = traces.fetch_records(project_id)
+        exported = traces.fetch_records(project_id, api_key)
     except Exception as exc:  # trace import is optional after project validation succeeds
         _warn(f"could not export traces from {project_id} ({exc}); using the starter dataset")
         return None
@@ -692,6 +684,7 @@ def _write_rl_config(
     rl: Path,
     project_line: str,
     thinking_line: str,
+    wandb_block: str,
     env_comment: str,
     max_examples_line: str,
     rl_reasoning_train: str,
@@ -704,6 +697,7 @@ def _write_rl_config(
             'algorithm = "grpo"\n'
             f"{thinking_line}"
             "\n"
+            f"{wandb_block}"
             f"{env_comment}"
             "[train]\n"
             "epochs = 1\n"
@@ -722,6 +716,7 @@ def _write_sft_config(
     sft: Path,
     project_line: str,
     thinking_line: str,
+    wandb_block: str,
     env_comment: str,
     max_examples_line: str,
     sft_reasoning_note: str,
@@ -734,6 +729,7 @@ def _write_sft_config(
             'algorithm = "sft"\n'
             f"{thinking_line}"
             "\n"
+            f"{wandb_block}"
             f"{env_comment}"
             "[train]\n"
             "epochs = 1\n"
@@ -766,6 +762,7 @@ def _write_opd_config(
     project_line: str,
     project_id: str,
     thinking_line: str,
+    wandb_block: str,
     max_examples_line: str,
     *,
     can_publish: bool = True,
@@ -788,6 +785,7 @@ def _write_opd_config(
             'algorithm = "opd"   # on-policy distillation from a managed parasail teacher (default glm 5.2)\n'
             f"{thinking_line}"
             "\n"
+            f"{wandb_block}"
             + _environment_comment(
                 project_id,
                 can_publish=can_publish,
@@ -816,15 +814,34 @@ def _write_training_guide(training: Path, project_id: str) -> None:
         )
 
 
+def _resolve_setup_identity(args) -> tuple[str, str, str, str, str]:
+    """Resolve the credential, project, and folder identity used by every setup phase."""
+    from flash.client import ClientError
+    from flash.client.config import load_credentials
+
+    api_url, api_key = load_credentials()
+    if not api_key:
+        raise ClientError(
+            f"not logged in. Run `{CLI_NAME} login` before `{CLI_NAME} env setup` "
+            "so the project can be validated"
+        )
+    project = _require_setup_project(args, api_url=api_url, api_key=api_key)
+    project_id = str(project["id"])
+    folder_name = Path.cwd().name
+    if not folder_name.strip():
+        raise ClientError("the current environment folder name must be nonblank")
+    return api_url, api_key, project_id, _wandb_project(project), folder_name
+
+
 def cmd_env_setup(args) -> int:
-    project_id = _require_setup_project(args)
+    api_url, api_key, project_id, project_name, folder_name = _resolve_setup_identity(args)
     _validate_existing_config_projects(project_id)
     starter_env = Path("environment.py")
     starter_evaluations = Path(_DEFAULT_EVALUATIONS_PATH)
     dataset = Path("dataset/train.jsonl")
     # trace import is optional and can only read the selected project. an existing dataset is never
     # overwritten, so importing into it would be a silently discarded download.
-    traces_jsonl = None if dataset.exists() else _traces_dataset(args, project_id)
+    traces_jsonl = None if dataset.exists() else _traces_dataset(args, project_id, api_key)
     multi_turn, starter_env_exists = _resolve_turn_mode(args, starter_env, dataset)
 
     rl = Path("configs/rl.toml")
@@ -837,7 +854,7 @@ def cmd_env_setup(args) -> int:
     reasoning_configs = (rl, opd, sft)
     reasoning = _resolve_reasoning_mode(args, reasoning_configs)
 
-    can_publish = _plane_can_publish_environments()
+    can_publish = _plane_can_publish_environments(api_url)
     _warn_if_environment_form_disagrees(reasoning_configs, can_publish=can_publish, warn=_warn)
     env_py = _render_starter(
         _STARTER_ENV_MULTITURN_PY if multi_turn else _STARTER_ENV_PY,
@@ -908,6 +925,7 @@ def cmd_env_setup(args) -> int:
         rl,
         project_line,
         thinking_line,
+        _wandb_block(project_name, f"{folder_name}-grpo"),
         env_comment,
         max_examples_line,
         rl_reasoning_train,
@@ -916,6 +934,7 @@ def cmd_env_setup(args) -> int:
         sft,
         project_line,
         thinking_line,
+        _wandb_block(project_name, f"{folder_name}-sft"),
         env_comment,
         max_examples_line,
         sft_reasoning_note,
@@ -927,6 +946,7 @@ def cmd_env_setup(args) -> int:
         project_line,
         project_id,
         thinking_line,
+        _wandb_block(project_name, f"{folder_name}-opd"),
         max_examples_line,
         can_publish=can_publish,
     )

freesolo_flash_dev-1.2.64/flash/cli/commands/env/setup_wandb.py

@@ -0,0 +1,30 @@
+"""W&B metadata helpers for environment scaffolds."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+
+_WANDB_PROJECT_MAX_LENGTH = 128
+_WANDB_PROJECT_FORBIDDEN = frozenset("/\\#?%:")
+
+
+def _wandb_project(project: Mapping[str, object]) -> str:
+    """Adapt a resolved canonical project to a W&B-safe project name."""
+    project_id = str(project["id"])
+    name = project.get("name")
+    if not isinstance(name, str) or not name.strip():
+        return project_id
+    if len(name) > _WANDB_PROJECT_MAX_LENGTH or any(
+        character in _WANDB_PROJECT_FORBIDDEN for character in name
+    ):
+        return project_id
+    return name
+
+
+def _wandb_block(project_name: str, run_name: str) -> str:
+    return (
+        "[wandb]\n"
+        f"project = {json.dumps(project_name, ensure_ascii=False)}\n"
+        f"run_name = {json.dumps(run_name, ensure_ascii=False)}\n\n"
+    )

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/cli/commands/traces.py

@@ -77,8 +77,8 @@ def fetch_records(
 ) -> dict:
     """A project's traces in the requested shape, converted server-side.
 
-    Omitting `api_key` reads the stored credential, which is right for `env_setup`'s call: it holds
-    no snapshot this key has to agree with.
+    Omitting `api_key` reads the stored credential. Callers that already resolved configuration
+    should pass the key from that snapshot so a concurrent login cannot change the request identity.
     """
     key = _require_api_key(api_key or load_credentials()[1])
     return export_trace_records(project_id, key, export_format=export_format)

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/cli/scaffold/TRAINING.md

@@ -272,8 +272,10 @@ lora_rank = 32              # lora_alpha defaults to 2 x lora_rank; set it to ov
 # All knobs live under [train]. Do not add [sft], [grpo], or [opd] tables.
 
 [wandb]
-# project  = "my-project"   # the table allows exactly `project` and `run_name`;
-# run_name = "sft-run-1"    # `name` is rejected
+# new scaffolds set `project` from the selected Freesolo project and `run_name` from the folder
+# plus algorithm. edit either value as needed; only `project` and `run_name` are allowed.
+# project  = "my-project"
+# run_name = "sft-run-1"
 ```
 
 **Knobs are scoped by algorithm.** `[train]` is one flat table shared by all three algorithms,

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/client/__init__.py

@@ -1,5 +1,7 @@
 """HTTP client for the managed Flash control plane (used by the CLI)."""
 
+from collections.abc import Mapping
+
 from flash._internal.channel import CLI_NAME
 from flash.client.config import load_credentials, save_credentials
 from flash.client.http import (
@@ -20,8 +22,10 @@ from flash.client.http import (
 from flash.core.spec import require_project_id
 
 
-def resolve_project_id(project_id: str, api_key: str, api_url: str | None = None) -> str:
-    """Resolve one project id, turning "not yours / not there" into the actionable CLI refusal.
+def resolve_project(
+    project_id: str, api_key: str, api_url: str | None = None
+) -> Mapping[str, object]:
+    """Resolve one project, turning "not yours / not there" into the actionable CLI refusal.
 
     403 and 404 are the same answer to a CLI user -- this project is not one you can record against
     -- and neither is distinguishable from a typo, so they collapse into one message naming the way
@@ -48,9 +52,9 @@ def resolve_project_id(project_id: str, api_key: str, api_url: str | None = None
     at call time, which is the binding the CLI tests patch.
     """
     if api_url is not None and not has_freesolo_backend(api_url):
-        return require_project_id(project_id)
+        return {"id": require_project_id(project_id)}
     try:
-        return str(get_project(project_id, api_key)["id"])
+        return get_project(project_id, api_key)
     except ApiError as exc:
         if exc.status not in {403, 404}:
             raise
@@ -60,6 +64,11 @@ def resolve_project_id(project_id: str, api_key: str, api_url: str | None = None
         ) from exc
 
 
+def resolve_project_id(project_id: str, api_key: str, api_url: str | None = None) -> str:
+    """Resolve and return one canonical project id."""
+    return str(resolve_project(project_id, api_key, api_url)["id"])
+
+
 __all__ = [
     "ApiClient",
     "ApiError",
@@ -72,6 +81,7 @@ __all__ = [
     "list_projects",
     "list_trace_projects",
     "load_credentials",
+    "resolve_project",
     "resolve_project_id",
     "save_credentials",
     "upload_eval_run",

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/engine/worker/opd_train_runner.py

@@ -124,6 +124,31 @@ def _validate_teacher_transport() -> tuple[str, str]:
     return capability, control_panel_url
 
 
+def _prepare_prompt_messages(
+    example: dict,
+    messages: list[dict],
+    *,
+    multi_turn: bool,
+    package_root: str | None,
+) -> tuple[list[dict], tuple[str, ...]]:
+    from flash.content.multimodal import normalize_prompt_images, record_has_images
+
+    if record_has_images(example, messages):
+        normalized = normalize_prompt_images(example, messages, package_root)
+        if multi_turn:
+            _opd_train.validate_transcript_messages(
+                normalized.messages,
+                source="environment initial prompt",
+                allow_content_blocks=True,
+            )
+        return normalized.messages, tuple(normalized.descriptors)
+    if multi_turn:
+        messages = _opd_train.validate_transcript_messages(
+            messages, source="environment initial prompt"
+        )
+    return messages, ()
+
+
 def _prepare_prompts(
     request: _OpdRequest,
     prompt_rows: list[tuple[Any, Any]],
@@ -131,11 +156,7 @@ def _prepare_prompts(
     capability: str,
     control_panel_url: str,
 ) -> _PromptState:
-    from flash.content.multimodal import (
-        image_teacher_prompt_messages,
-        normalize_prompt_images,
-        record_has_images,
-    )
+    from flash.content.multimodal import image_teacher_prompt_messages
     from flash.engine.worker.teacher.client import TeacherClient
 
     teacher = TeacherClient(capability, control_panel_url, request.knobs.teacher_model)
@@ -187,15 +208,14 @@ def _prepare_prompts(
     with _opd_train.liveness_heartbeat("opd_image_prep", progress=lambda: prepped[0]):
         for example, messages in prompt_rows:
             prepped[0] += 1
-            if request.multi_turn:
-                messages = _opd_train.validate_transcript_messages(
-                    messages, source="environment initial prompt"
-                )
-            if record_has_images(example, messages):
+            student_messages, image_descriptors = _prepare_prompt_messages(
+                example,
+                messages,
+                multi_turn=request.multi_turn,
+                package_root=package_root,
+            )
+            if image_descriptors:
                 assert processor is not None
-                normalized = normalize_prompt_images(example, messages, package_root)
-                student_messages = normalized.messages
-                image_descriptors = tuple(normalized.descriptors)
                 teacher_messages = image_teacher_prompt_messages(
                     student_messages, len(image_descriptors)
                 )
@@ -207,9 +227,7 @@ def _prepare_prompts(
                     enable_thinking=bool(_opd_train._w.THINKING),
                 )
             else:
-                student_messages = messages
-                teacher_messages = messages
-                image_descriptors = ()
+                teacher_messages = student_messages
                 if processor is not None:
                     # mixed job: the verl child tokenizes EVERY row through the multimodal dataset
                     # path (the processor), so text-only rows must freeze via the same path or the
@@ -224,7 +242,7 @@ def _prepare_prompts(
                 else:
                     prompt_ids = _opd_train._normalize_prompt_ids(
                         tokenizer.apply_chat_template(
-                            messages,
+                            student_messages,
                             tokenize=True,
                             add_generation_prompt=True,
                             enable_thinking=_opd_train._w.THINKING,

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/engine/worker/train/core/child/glue.py

@@ -81,6 +81,73 @@ def content_block_text(content, *, source: str, position: int) -> str:
     return "".join(parts)
 
 
+def validate_structured_messages(messages: list[dict], *, source: str) -> list[dict]:
+    """validate and canonicalize role/content messages without dropping media placement."""
+    if not isinstance(messages, list):
+        raise ValueError(f"{source} messages must be a list")
+    normalized: list[dict[str, Any]] = []
+    for position, message in enumerate(messages):
+        if not isinstance(message, dict):
+            raise ValueError(f"{source} message {position} must be an object")
+        extras = sorted(
+            key
+            for key, value in message.items()
+            if key not in _ALLOWED_MESSAGE_KEYS and value is not None
+        )
+        if extras:
+            raise ValueError(
+                f"{source} message {position} carries unsupported transcript metadata {extras}; "
+                "tool names, call ids, tool calls, and other message fields cannot be represented "
+                "in a role/content multi-turn transcript"
+            )
+        role = message.get("role")
+        content = message.get("content")
+        if not isinstance(role, str) or not role.strip():
+            raise ValueError(f"{source} message {position} has an invalid role")
+        if isinstance(content, str):
+            normalized.append({"role": role, "content": content})
+            continue
+        if not isinstance(content, list):
+            raise ValueError(
+                f"{source} message {position} content must be text or content blocks for multi-turn"
+            )
+        blocks: list[dict[str, Any]] = []
+        for block_index, block in enumerate(content):
+            if not isinstance(block, dict):
+                raise ValueError(
+                    f"{source} message {position} has a content block that is not an object"
+                )
+            block_type = block.get("type")
+            allowed_keys = (
+                {"type", block_type} if block_type in _MEDIA_BLOCK_TYPES else {"type", "text"}
+            )
+            extras = sorted(
+                key for key, value in block.items() if key not in allowed_keys and value is not None
+            )
+            if extras:
+                raise ValueError(
+                    f"{source} message {position} content block {block_index} carries unsupported "
+                    f"fields {extras}"
+                )
+            if block_type in _MEDIA_BLOCK_TYPES:
+                blocks.append({"type": block_type})
+            elif block_type == "text" and isinstance(block.get("text"), str):
+                text = block["text"]
+                if not text:
+                    continue
+                if blocks and blocks[-1].get("type") == "text":
+                    blocks[-1]["text"] += text
+                else:
+                    blocks.append({"type": "text", "text": text})
+            else:
+                raise ValueError(
+                    f"{source} message {position} has an unsupported content block at index "
+                    f"{block_index}"
+                )
+        normalized.append({"role": role, "content": blocks})
+    return normalized
+
+
 def validate_transcript_messages(
     messages: list[dict], *, source: str, allow_content_blocks: bool = False
 ) -> list[dict]:
@@ -314,3 +381,94 @@ async def run_executor_call(loop, callback):
         with contextlib.suppress(Exception):
             await asyncio.shield(task)
         raise
+
+
+def multi_modal_image_count(multi_modal_data) -> int:
+    """how many images a verl multimodal payload carries.
+
+    single-turn and multi-turn both report this count to the bridge, which compares it against the
+    frozen parent prompt. two counters would let the two paths disagree about the same payload, so
+    this is the one definition: verl carries "image" (singular) for a single-image row and "images"
+    for a list, and a bare non-sequence value under either key is one image.
+    """
+    if not multi_modal_data:
+        return 0
+    if not isinstance(multi_modal_data, dict):
+        raise TypeError("verl multimodal data must be a mapping")
+    images = multi_modal_data.get("images")
+    if images is None:
+        images = multi_modal_data.get("image")
+    if images is None:
+        return 0
+    if not isinstance(images, (list, tuple)):
+        return 1
+    try:
+        return len(images)
+    except TypeError as error:
+        raise TypeError("verl multimodal images must be a sized collection") from error
+
+
+class EpisodePrompt:
+    """the tokenized initial prompt plus the decoded media every generate call must carry.
+
+    an image-bearing prompt tokenizes to placeholder tokens that carry no pixels. both
+    apply_chat_template and every generate call need the decoded media alongside the ids, or the
+    engine sees placeholders it cannot expand and the rollout either dies on a feature/placeholder
+    mismatch or conditions on nothing. text-only prompts yield {}.
+
+    the media is frozen at episode start and resent unchanged on every turn. later environment
+    replies never add to it: a mid-episode image would change the placeholder/pixel correspondence
+    the prompt ids were built against, which is why both loops reject one upstream.
+    """
+
+    __slots__ = (
+        "audios",
+        "images",
+        "mm_processor_kwargs",
+        "multi_modal_data",
+        "prompt_ids",
+        "structured_messages",
+        "videos",
+    )
+
+    def __init__(self, multi_modal_data, mm_processor_kwargs, prompt_ids, structured_messages):
+        self.multi_modal_data = multi_modal_data
+        self.mm_processor_kwargs = mm_processor_kwargs
+        self.prompt_ids = prompt_ids
+        self.structured_messages = structured_messages
+        self.images = multi_modal_data.get("images")
+        self.videos = multi_modal_data.get("videos")
+        self.audios = multi_modal_data.get("audios")
+
+    def image_count(self) -> int:
+        """how many images the engine actually decoded for this prompt.
+
+        the bridge authenticates this against the frozen parent prompt, so it counts what the child
+        will condition on rather than what the transcript appears to mention. it must agree with the
+        count the single-turn path reports, hence the shared counter rather than a second one.
+        """
+        return multi_modal_image_count(self.multi_modal_data)
+
+
+async def prepare_episode_prompt(loop_self, raw_prompt) -> EpisodePrompt:
+    """extract media and ids before building the canonical structured prompt view.
+
+    order matters: an image prompt arrives as content blocks, so the media and block-rendered ids
+    must be captured before structured validation. the canonical view preserves media placement for
+    bridge authentication and rejects unsupported blocks and message metadata before rollout.
+    """
+    messages = [dict(message) for message in raw_prompt]
+    multi_modal_data = await loop_self.process_multi_modal_info(messages)
+    images = multi_modal_data.get("images")
+    videos = multi_modal_data.get("videos")
+    audios = multi_modal_data.get("audios")
+    mm_processor_kwargs = loop_self._get_mm_processor_kwargs(audios)
+    prompt_ids = await loop_self.apply_chat_template(
+        messages,
+        images=images,
+        videos=videos,
+        audios=audios,
+        mm_processor_kwargs=mm_processor_kwargs,
+    )
+    structured = validate_structured_messages(messages, source="initial prompt")
+    return EpisodePrompt(multi_modal_data, mm_processor_kwargs, prompt_ids, structured)

{freesolo_flash_dev-1.2.63 → freesolo_flash_dev-1.2.64}/flash/engine/worker/train/core/child/runtime.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import atexit
+import contextlib
 import json
 import os
 import sys
@@ -200,9 +202,22 @@ def install_wandb_link_reporting() -> None:
         if getattr(wandb.init, "_flash_wandb_reporting", False):
             return
         original_init = wandb.init
+        finished_run_ids: set[int] = set()
+
+        def finish_run(run) -> None:
+            run_id = id(run)
+            if run_id in finished_run_ids:
+                return
+            try:
+                run.finish(exit_code=0)
+            except Exception:
+                return
+            finished_run_ids.add(run_id)
 
         def init_reporting(*args, **kwargs):
             run = original_init(*args, **kwargs)
+            with contextlib.suppress(Exception):
+                atexit.register(finish_run, run)
             try:
                 url = getattr(run, "url", None)
                 run_id = getattr(run, "id", None)