Enable autocompletion on Python language servers

[BlueskyFR]

Hi!

It seems that the Determined Python package is not compatible with Python language servers because of the way it exports classes/members in __init__.py files.

Issue in VSCode (PyLance/PyRight language server)

Fix example

Example fix proposed for this file: https://github.com/determined-ai/determined/blob/728078b5d0581c73b66392c9723ac0c800d30b80/harness/determined/keras/__init__.py#L13-L20

from determined.keras import callbacks
from determined.keras._data import (
    _ArrayLikeAdapter,
    _adapt_data_from_data_loader,
    _adapt_data_from_fit_args,
-   ArrayLike,
-   SequenceAdapter,
-   InputData,
+   ArrayLike as ArrayLike,
+   SequenceAdapter as SequenceAdapter,
+   InputData as InputData,
)
from determined.keras._enqueuer import _Enqueuer, _Sampler, _build_enqueuer
from determined.keras._tensorboard_callback import TFKerasTensorBoard
from determined.keras._tf_keras_context import (
-   TFKerasExperimentalContext,
-   TFKerasTrainConfig,
-   TFKerasTrialContext,
+   TFKerasExperimentalContext as TFKerasExperimentalContext,
+   TFKerasTrainConfig as TFKerasTrainConfig,
+   TFKerasTrialContext as TFKerasTrialContext,
)
from determined.keras._tf_keras_multi_gpu import (
    _check_if_aggregation_frequency_will_work,
)
-from determined.keras._tf_keras_trial import TFKerasTrial, TFKerasTrialController
-from determined.keras._load import load_model_from_checkpoint_path
+from determined.keras._tf_keras_trial import (
+   TFKerasTrial as TFKerasTrial,
+   TFKerasTrialController as TFKerasTrialController,
+)
+from determined.keras._load import (
+   load_model_from_checkpoint_path as load_model_from_checkpoint_path,
+)

Result after fix

Please note that this should be applied to the entire repo as it does not only concern this file.

[rb-determined-ai]

Interesting... do you know, does an __all__ = ["TFKerasTrial"] mechanism also meet this need?

[rb-determined-ai]

Also, I am not able to reproduce the error you are seeing either by running pyright directly on a main.py like yours, or by installing the python plugin to vscode and opening such a file.

In vscode I can go-to-def so it seems to understand what determined files to look at, but I don't see any warnings or errors.

Is there a setting or something I'm missing? Can you run pyright directly and get the same error message?

[BlueskyFR]

Hey @rb-determined-ai! I use the Pylance language server version 2023.5.50 through the Python VSCode extension (version v2023.8.0).

• First, yes, using __all__ seems to have the same behavior as adding "X as X"
• Second, for me too go to definition works but the autocompletion is still not working:

I checked Pylance's logs but they did not provide any relevant information

Also I am running the latest pip install determined + Python 3.11 but the behavior is the same using Python 3.10.9