fix: Handle KeyError in W&B migration script and add validation for missing Neptune workspace

[SiddhantSadangi]

Description

Include a summary of the changes and the related issue.

Related to: <ClickUp/JIRA task name>

Any expected test failures?

---

Add a [X] to relevant checklist items

❔ This change

• adds a new feature
• fixes breaking code
• is cosmetic (refactoring/reformatting)

---

✔️ Pre-merge checklist

• Refactored code (sourcery)
• Tested code locally
• Precommit installed and run before pushing changes
• Added code to GitHub tests (notebooks, scripts)
• Updated GitHub README
• Updated the projects overview page on Notion

---

🧪 Test Configuration

• OS:
• Python version:
• Neptune version:
• Affected libraries with version:

Summary by Sourcery

Validate Neptune workspace input and extend exception handling in the W&B-to-Neptune migration script to catch KeyError and avoid crashes.

Bug Fixes:

• Add validation to prevent empty Neptune workspace name
• Handle KeyError when copying metrics from W&B to Neptune

[sourcery-ai]

Reviewer's guide (collapsed on small PRs)

Reviewer's Guide

Enhance robustness of the W&B to Neptune migration script by validating the Neptune workspace input and extending error handling in metric copying to include KeyError.

Sequence diagram for Neptune workspace name validation in migration script

sequenceDiagram
    actor User
    participant Script
    User->>Script: Start migration script
    Script->>User: Prompt for Neptune workspace name
    User->>Script: Enter workspace name
    Script->>Script: Validate workspace name
    alt Workspace name is empty
        Script->>User: Print error message
        Script->>Script: Exit with error
    else Workspace name is valid
        Script->>Script: Continue execution
    end

Loading

Class diagram for updated error handling in copy_metrics function

classDiagram
    class copy_metrics {
        +copy_metrics(neptune_run, wandb_run)
        - Handles IndexError, TypeError, KeyError when copying metrics
    }

Loading

File-Level Changes

Change	Details	Files
Validate Neptune workspace input to prevent empty values	Trim user input for workspace name Print an error message and exit with status code 1 on empty input	utils/migration_tools/from_wandb/wandb_to_neptune.py
Extend error handling in copy_metrics to catch KeyError	Add KeyError to the exception tuple alongside IndexError and TypeError	utils/migration_tools/from_wandb/wandb_to_neptune.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey there - I've reviewed your changes and they look great!

Prompt for AI Agents

Please address the comments from this code review:

## Individual Comments

### Comment 1
<location> `utils/migration_tools/from_wandb/wandb_to_neptune.py:236` </location>
<code_context>
                 ):  # Not uploading W&B table-file to Neptune
                     continue
-            except (IndexError, TypeError):
+            except (IndexError, TypeError, KeyError):
                 if epoch:
                     neptune_run[key].append(value, step=epoch)
</code_context>

<issue_to_address>
**issue (bug_risk):** Adding KeyError to the exception handling may mask unexpected issues.

If missing keys are expected, handle them explicitly or log them for debugging.
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}