Fix bug with passing capture_* args to neptune callback

[AleksanderWWW]

What does this PR do?

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[AleksanderWWW]

What does this PR do

This PR aims to fix a bug that appears when using NeptuneCallback with run=None and at least one of the capture_* params set.

Cause of the problem

This is due to the fact, that in one of the methods those params have hardcoded values, but the kwargs passed to the constructor are also passed to that method. Therefore a TypeError like this occurs during evaluation:
TypeError: neptune.metadata_containers.run.Run() got multiple values for keyword argument 'capture_stdout'

Solution

Before invoking the problematic method we want to make sure that the capture_* params are not present in the callback's state by temporarily removing them. Once the method returns, the original kwargs state is restored.

[amyeroberts]

Thanks for opening this PR @AleksanderWWW! Let us know when you want a review 🤗

[AleksanderWWW]

Thank you @amyeroberts !

Let me wait for my colleagues at Neptune to share their thoughts on my proposal 😃

[Raalsky]

I would recommend to override the params as a part of def _initialize_run:

    def _initialize_run(self, **additional_neptune_kwargs):
        ...

        self._stop_run_if_exists()

        try:
            run_params = additional_neptune_kwargs.copy()
            run_params.update(self._init_run_kwargs)
            self._run = init_run(**run_params)
            self._run_id = self._run["sys/id"].fetch()
        except (NeptuneMissingProjectNameException, NeptuneMissingApiTokenException) as e:
            raise NeptuneMissingConfiguration() from e

I'm not sure if I haven't done any typo etc. but I don't see a need to "preserve" these params.

[Raalsky]

Actually by use case mentioned here: neptune-ai/neptune-client#1663 let's override with the user-provided parameters.

[AleksanderWWW]

But setting those capture_* to False is only required in one invocation of _initialize_run. I'm not sure I understand

[Raalsky]

No need for this, right?

[AleksanderWWW]

Hey @amyeroberts
It seems that the integration_utils.py module has become very long and rather hard to navigate.
It's difficult to find what you're looking for and making changes in more than one place in one PR is a nightmare.
Would it be possible to split it into multiple modules separating different integrations?

[amyeroberts]

Hi @AleksanderWWW, what kind of structure would you propose for the files?

[AleksanderWWW]

@amyeroberts
I thought of something similar to what pytorch_lightning has: https://github.com/Lightning-AI/pytorch-lightning/tree/master/src/lightning/pytorch/loggers

The root could be the already existing integrations package, but each provider (e.g. Neptune, CometML, etc.) has its own module.
Does it make sense?

---

And btw. the PR is ready for review from the code owners' side 🚀

[amyeroberts]

@AleksanderWWW OK, sounds like a good idea. We'd have to make all objects importable from the integrations module for backwards compatibility. If you'd like to open a PR I'd be very happy to review :)

[amyeroberts]

Thanks for working on a fix!

Just a question about how this works

[amyeroberts]

This looks to me like it won't work. Because we're updating run_params with self._init_run_kwargs we're guaranteeing that all the keys in self._init_run_kwargs are in run_params i.e. we're passing the same argument twice to init_run, which I'd expect to trigger an error.

[AleksanderWWW]

Yes, you are absolutely correct!. In fact, it should be just **run_params in the init_run call. My bad! Fixing rn :) Thank you for noticing

[AleksanderWWW]

If you'd like to open a PR I'd be very happy to review :)

Will do :)

[AleksanderWWW]

The failing CI test ModelUtilsTest.test_use_safetensors doesn't seem to have anything to do with my changes. Is it possible to maybe rerun it?

[amyeroberts]

Thanks for fixing!

[AleksanderWWW]

@amyeroberts will you merge it once the checks pass? 🚀

[amyeroberts]

@AleksanderWWW A fix has been merged into main which should resolve this. Could you try rebasing to include this and trigger an new CI run?

[AleksanderWWW]

@amyeroberts Done 😃