Combination of forward references, and changing objects `__module__` attribute leads to odd bug

[EthanMarx]

🐛 Bug report

I found an interesting edge case when trying to build out a CLI for the ray.tune.Tuner.

Not sure if this should be considered a bug on the jsonargparse side, or an odd (mis?)use of python flexibility by Ray.

The issue stems when trying to build a CLI for the ray.train.RunConfig. It is the combination of the following that leads to problems:

1. The the actual RunConfig class definition lives in ray.air.config
2. ray changes the __module__ attribute of the ray.train.RunConfig import to be ray.train.
3. The actual RunConfig class uses ForwardRefs (e.g. for the SyncConfig)

Due to the above, the following parser isn't able to add any forward referenced arguments in RunConfig

# test.py
from jsonargparse import ArgumentParser
from ray.train import RunConfig

parser = ArgumentParser()
parser.add_class_arguments(RunConfig, "run_config")

if __name__ == "__main__":
     parser.parse_args()

Specifically, the SyncConfig and other attributes that use forward references aren't able to be resolved, and thus aren't added to the parser:

python test.py --help

usage: test.py [-h] [--test CONFIG] [--test.name NAME] [--test.storage_path STORAGE_PATH]
               [--test.storage_filesystem.help CLASS_PATH_OR_NAME] [--test.storage_filesystem STORAGE_FILESYSTEM]
               [--test.failure_config FAILURE_CONFIG] [--test.checkpoint_config CHECKPOINT_CONFIG] [--test.verbose VERBOSE]
               [--test.stop STOP] [--test.log_to_file LOG_TO_FILE] [--test.local_dir LOCAL_DIR] [--print_shtab {bash,zsh,tcsh}]

options:
  -h, --help            Show this help message and exit.
  --print_shtab {bash,zsh,tcsh}
                        Print shtab shell completion script.

Runtime configuration for training and tuning runs:
  --test CONFIG         Path to a configuration file.
  --test.name NAME      Name of the trial or experiment. If not provided, will be deduced from the Trainable. (type:
                        Optional[str], default: null)
  --test.storage_path STORAGE_PATH
                        [Beta] Path where all results and checkpoints are persisted. Can be a local directory or a destination on
                        cloud storage. For multi-node training/tuning runs, this must be set to a shared storage location (e.g.,
                        S3, NFS). This defaults to the local ``~/ray_results`` directory. (type: Optional[str], default: null)
  --test.storage_filesystem.help CLASS_PATH_OR_NAME
                        Show the help for the given subclass of FileSystem and exit.
  --test.storage_filesystem STORAGE_FILESYSTEM
                        [Beta] A custom filesystem to use for storage. If this is provided, `storage_path` should be a path with
                        its prefix stripped (e.g., `s3://bucket/path` -> `bucket/path`). (type: Optional[FileSystem], default:
                        null)
  --test.failure_config FAILURE_CONFIG
                        Failure mode configuration. (type: Optional[FailureConfig], default: null)
  --test.checkpoint_config CHECKPOINT_CONFIG
                        Checkpointing configuration. (type: Optional[CheckpointConfig], default: null)
  --test.verbose VERBOSE
                        0, 1, or 2. Verbosity mode. 0 = silent, 1 = default, 2 = verbose. Defaults to 1. If the
                        ``RAY_AIR_NEW_OUTPUT=1`` environment variable is set, uses the old verbosity settings: 0 = silent, 1 =
                        only status updates, 2 = status and brief results, 3 = status and detailed results. (type: Optional[int],
                        default: null)
  --test.stop STOP      Stop conditions to consider. Refer to ray.tune.stopper.Stopper for more info. Stoppers should be
                        serializable. (type: Union[Mapping, Callable[[str, Mapping], bool], null], default: null)
  --test.log_to_file LOG_TO_FILE
                        [DeveloperAPI] Log stdout and stderr to files in trial directories. If this is `False` (default), no
                        files are written. If `true`, outputs are written to `trialdir/stdout` and `trialdir/stderr`,
                        respectively. If this is a single string, this is interpreted as a file relative to the trialdir, to
                        which both streams are written. If this is a Sequence (e.g. a Tuple), it has to have length 2 and the
                        elements indicate the files to which stdout and stderr are written, respectively. (type: Union[bool, str,
                        Tuple[str, str]], default: False)
  --test.local_dir LOCAL_DIR
                        (type: Optional[str], default: null)

After some debugging, the main issue is that, since the __module__ attribute has been changed by ray, jsonargparse can't import the proper objects here when resolving references.

Environment

• jsonargparse version: 4.33.2
• Python version: 3.10
• How jsonargparse was installed: pip
• OS: liux

[mauvilsa]

Thanks for reporting!

I looked at it and the problem is not because of ray's __module__ attribute modification. Can be reproduced with:

from jsonargparse import ArgumentParser
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional

if TYPE_CHECKING:
    from ray.train import SyncConfig

@dataclass
class RunConfig:
    name: Optional[str] = None
    storage_path: Optional[str] = None
    sync_config: Optional["SyncConfig"] = None

parser = ArgumentParser()
parser.add_class_arguments(RunConfig, "run_config")

if __name__ == "__main__":
    parser.parse_args()

The problem is the forward ref combined with a dataclass. Since SyncConfig is imported inside a TYPE_CHECKING block, this can be fixed. I will look at it later.

[mauvilsa]

But there is a second problem here ray/air/config.py#L652-L654. That is just an invalid type which does not make sense to support. It is just weird that they did that.

[EthanMarx]

Thanks for taking a look at this. Yeah I did notice that invalid type as well. Curious why thats there.

[mauvilsa]

It should be fixed with #611. Also created a pull request in ray ray-project/ray#48439, which I don´t think there would be a reason to not accept.

[mauvilsa]

I forgot to mention that it doesn't work in python 3.8. Since it already reached end-of-life, I did not bother in fixing it.

[EthanMarx]

@mauvilsa While the above example now works on v3.44.0, still having trouble parsing the following:

from ray import tune
import jsonargparse
config = {
    "tuner": {
        "run_config": {
            "name": "boring-demo",
            "storage_path": "s3://my-bucket",
            "sync_config":
            {
                "sync_period": 100,
            },
        },
    },
}

parser = jsonargparse.ArgumentParser()
parser.add_argument("--config", action=jsonargparse.ActionConfigFile)
parser.add_class_arguments(tune.Tuner, "tuner")

cfg = parser.dump(config)
parser.parse_args([f"--config={cfg}"])

error: Parser key "tuner.run_config":
  Does not validate against any of the Union subtypes
  Subtypes: (<class 'ray.train.RunConfig'>, <class 'NoneType'>)
  Errors:
    - Validation failed: Key 'sync_config.sync_period' is not expected
    - Expected a <class 'NoneType'>
  Given value type: <class 'dict'>
  Given value: {'name': 'boring-demo', 'storage_path': 's3://my-bucket', 'sync_config': {'sync_period': 100}}

Note that removing the sync_config from the dictionary will parse.

from ray import tune
import jsonargparse
config = {
    "tuner": {
        "run_config": {
            "name": "boring-demo",
            "storage_path": "s3://my-bucket",
        },
    },
}

parser = jsonargparse.ArgumentParser()
parser.add_argument("--config", action=jsonargparse.ActionConfigFile)
parser.add_class_arguments(tune.Tuner, "tuner")

cfg = parser.dump(config)
parser.parse_args([f"--config={cfg}"])