KeyError in upgrade_config when running predict

[perliedman]

Hi,
I'm writing a plugin to handle a custom model format, and have stumbled on a somewhat unrelated issue that I just can't figure out.

I have added a new config class, Ml2ModelConfig (inherits from SemanticSegmentationModelConfig) which I have used the @register_config decorator for, and I have also added the plugin registration code in my plugin's __init__.py as described in https://docs.rastervision.io/en/stable/framework/architecture.html#example-1-a-simple-pipeline

When running and using my new config, training and validation seems to work ok, but when I get to Running predict command..., it fails while trying to upgrade the config, presumably since my plugin is not in pipeline-config.json's plugin_versions dict:

Log looks like this:

2023-01-03 07:51:35:rastervision.pytorch_learner.learner: INFO - Unzipping model-bundle to /opt/data/tmp/tmpwqq55p73/model-bundle
Traceback (most recent call last):
  File "/opt/conda/lib/python3.9/runpy.py", line 197, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.9/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 251, in <module>
    _main()
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 247, in _main
    main()
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1055, in main
    rv = self.invoke(ctx)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/opt/conda/lib/python3.9/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 175, in run
    _run_pipeline(cfg, runner, tmp_dir, splits, commands,
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 125, in _run_pipeline
    runner.run(
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/runner/inprocess_runner.py", line 25, in run
    _run_command(cfg_json_uri, command, 0, 1)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/cli.py", line 218, in _run_command
    command_fn()
  File "/opt/src/rastervision_core/rastervision/core/rv_pipeline/rv_pipeline.py", line 181, in predict
    self.backend.load_model()
  File "/opt/src/rastervision_pytorch_backend/rastervision/pytorch_backend/pytorch_learner_backend.py", line 123, in load_model
    self.learner = self._build_learner_from_bundle(training=False)
  File "/opt/src/rastervision_pytorch_backend/rastervision/pytorch_backend/pytorch_learner_backend.py", line 131, in _build_learner_from_bundle
    return Learner.from_model_bundle(
  File "/opt/src/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py", line 969, in from_model_bundle
    config_dict = upgrade_config(config_dict)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/config.py", line 247, in upgrade_config
    return _upgrade_config(config_dict, plugin_versions)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/config.py", line 185, in _upgrade_config
    new_x[k] = _upgrade_config(v, plugin_versions)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/config.py", line 185, in _upgrade_config
    new_x[k] = _upgrade_config(v, plugin_versions)
  File "/opt/src/rastervision_pipeline/rastervision/pipeline/config.py", line 191, in _upgrade_config
    old_version = plugin_versions[plugin]
KeyError: 'rastervision.ml2_pipeline'

I'm sure I missed some detail, but even after digging around the code a bit I just can't find the missing piece (or anyone else having the same issue...) Any ideas what I'm missing?

[AdeelH]

Can you try calling registry.set_plugin_version() inside your register_plugin() function and see if that solves the problem?

Example: https://github.com/azavea/raster-vision/blob/0.20/rastervision_pytorch_learner/rastervision/pytorch_learner/__init__.py#L4-L5

In your case, this would probably be:

def register_plugin(registry):
    registry.set_plugin_version('rastervision.ml2_pipeline', 0)

[perliedman]

Thanks, @AdeelH, that did indeed solve the issue!

Should I expect to call this method in my plugins, or does this indicate some other problem in my setup?

[AdeelH]

I believe this always needs to be called in every plugin. The documentation and example plugins are at fault for not mentioning this.