Skip to content

py-gen-ml

Generate typed ML configuration tooling from a protobuf definition.

This command is the core of the py-gen-ml toolbox. You pass it one or more .proto files that you authored. Under the hood it runs protoc with the protoc-gen-py-ml plugin. That is ordinary, deterministic schema-driven code generation. It then writes Pydantic models and JSON Schemas for your project.

By default, generated Python lands in src/pgml_out. If your proto is called example.proto, you get:

  • src/pgml_out/example_base.py. Base model that follows the protobuf definition.
  • src/pgml_out/example_sweep.py. Sweep model for hyperparameter search spaces.
  • src/pgml_out/example_patch.py. Patch model for overlaying a base config.

If you've set the CLI option on a message called Foo, it will also generate:

  • src/pgml_out/example_cli_args.py. CLI argument models.
  • src/pgml_out/foo_entrypoint.py. Entrypoint that combines base config, sweep, and CLI arguments.

It also writes JSON schemas under configs/:

  • configs/base/schemas/<message_name>.json
  • configs/patch/schemas/<message_name>.json
  • configs/sweep/schemas/<message_name>.json

🏗️ Typical project layout

<project_root>/
  src/
    your_package/
      config.proto          # you write this
      train.py              # you write this
    pgml_out/               # generated by py-gen-ml
      config_base.py
      config_patch.py
      config_sweep.py
      ...
  configs/
    base/
      schemas/              # generated JSON schemas
      default.yaml          # you write YAML configs
    patch/
      schemas/
    sweep/
      schemas/

💻 Usage

$ py-gen-ml [OPTIONS] PROTO_FILE...

Arguments:

  • PROTO_FILE...: Path to the protobuf file(s). [required]

Options:

  • --proto-root TEXT: Path to the root of the protobuf files which will be passed to the protoc command as an include path. If not specified, the script will try to infer it from the proto_file arguments by adding the parent of the proto_file arguments.
  • --code-dir TEXT: Path to the generated code directory. [default: src/pgml_out]
  • --source-root TEXT: Path to the root of the source code. [default: src]
  • --configs-dir TEXT: Path to the base directory for configs. [default: configs]
  • --generators TEXT: Comma-separated subset of generators to run, e.g. --generators=base,patch. Defaults to every registered generator with enabled_by_default=True (built-in plus any third-party generators discovered via the py_gen_ml.generators entry-point group). Built-ins: base, patch, sweep, cli_args. Opt-in: lancedb (see LanceDB schemas), bentoml (see BentoML services), litserve (see LitServe services), pydantic_ai (see PydanticAI synthesis), argilla (see Argilla datasets), mlflow (see MLflow tracking), wandb (see Weights & Biases tracking).
  • --install-completion: Install completion for the current shell.
  • --show-completion: Show completion for the current shell, to copy it or customize the installation.
  • --help: Show help and exit.

🔌 Extending generators

Built-in emitters are registered through the py_gen_ml.generators entry-point group. Third-party packages can ship additional GeneratorSpecs under the same group. --generators then lets you select which ones to run.