mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 04:12:36 +00:00
fix(cli): review pass on the docsgpt command
- The image pins DOCSGPT_HOME=/app: it ships no checkout, so the data home no longer depends on the working directory. - api, worker, beat and migrate print the data home and env file they resolved, so an API and a worker started from different directories show it. - The worker passes -Q only when asked; a bare worker consumes every configured queue, which honours EMBEDDINGS_QUEUE and DOCUMENT_PARSE_QUEUE. - The worker runs through celery.start and returns its exit code; click usage errors print usage and exit 2 instead of a traceback. - Windows: solo pool and no embedded scheduler (celery rejects -B there), with a pointer to the new `docsgpt beat` command, which runs the scheduler on its own. - prefetch_models and verify_offline parse their arguments, so --help is help rather than a model name. - A DOCSGPT_ENV_FILE that is not a file raises instead of booting with defaults.
This commit is contained in:
1 parent
f6f0d49452
commit
3036ece3ff
8 files changed
+134
-23
No files matched your search
@@ -80,7 +80,9 @@ docsgpt worker # in a second terminal: the Celery worker, with the scheduler
|
||||
|
||||
The API applies pending migrations when it starts (`AUTO_MIGRATE`), so `docsgpt migrate` is the explicit step for deployments that want the schema in place before the first request or that run the API with a restricted database role.
|
||||
|
||||
The worker is not optional: query embedding runs on it, so search fails without one. `docsgpt worker --help` lists the queue, concurrency and pool options; `--no-beat` starts a worker without the scheduler when another worker already runs it.
|
||||
Both commands print the data home they resolved on start-up. Run them from the same directory, or set `DOCSGPT_HOME` for both, so the worker finds the files the API stores and the API finds the indexes the worker builds.
|
||||
|
||||
The worker is not optional: query embedding runs on it, so search fails without one. `docsgpt worker --help` lists the queue, concurrency and pool options; `--no-beat` starts a worker without the scheduler when another worker already runs it. On Windows the scheduler cannot be embedded, so run `docsgpt beat` in a third terminal.
|
||||
|
||||
Other commands:
|
||||
|
||||
|
||||
+4
-1
@@ -150,7 +150,10 @@ ENV FLASK_APP=app.py
|
||||
# Thread caps. onnxruntime (FastEmbed) ignores OMP_NUM_THREADS and sizes its
|
||||
# pool to the host's core count, which a CPU-limited container still reports;
|
||||
# EMBEDDINGS_THREADS pins it the way OMP_NUM_THREADS pinned torch before.
|
||||
ENV MALLOC_ARENA_MAX=2 \
|
||||
# The image ships no checkout, so pin the data home (inputs, indexes, .env)
|
||||
# instead of depending on the working directory.
|
||||
ENV DOCSGPT_HOME=/app \
|
||||
MALLOC_ARENA_MAX=2 \
|
||||
OMP_NUM_THREADS=4 \
|
||||
MKL_NUM_THREADS=4 \
|
||||
OPENBLAS_NUM_THREADS=4 \
|
||||
|
||||
+45
-10
@@ -16,11 +16,18 @@ from docsgpt.version import __version__
|
||||
|
||||
DEFAULT_HOST = "127.0.0.1"
|
||||
DEFAULT_PORT = 7091
|
||||
DEFAULT_QUEUES = "docsgpt,parsing,embeddings"
|
||||
|
||||
|
||||
def _announce_home() -> None:
|
||||
"""Say where runtime data and the env file come from; the API and the worker must agree."""
|
||||
from docsgpt.core.paths import env_file, home_dir
|
||||
|
||||
print(f"docsgpt: data home {home_dir()} (env file {env_file()})", file=sys.stderr)
|
||||
|
||||
|
||||
def _api(args: argparse.Namespace) -> int:
|
||||
"""Serve the ASGI app: gunicorn with the bounded-drain uvicorn worker, or uvicorn when reloading."""
|
||||
_announce_home()
|
||||
if args.reload or sys.platform == "win32":
|
||||
import uvicorn
|
||||
|
||||
@@ -49,20 +56,43 @@ def _api(args: argparse.Namespace) -> int:
|
||||
return 0
|
||||
|
||||
|
||||
def _worker(args: argparse.Namespace) -> int:
|
||||
"""Run the Celery worker (with the beat scheduler unless ``--no-beat``)."""
|
||||
def _celery(argv: list[str]) -> int:
|
||||
"""Run a celery subcommand on the app and return its exit code (usage errors print usage)."""
|
||||
import click
|
||||
|
||||
from docsgpt.app import celery
|
||||
|
||||
argv = ["worker", "-l", args.loglevel, "-Q", args.queues]
|
||||
try:
|
||||
code = celery.start(argv)
|
||||
except click.ClickException as exc:
|
||||
exc.show()
|
||||
return exc.exit_code
|
||||
return int(code or 0)
|
||||
|
||||
|
||||
def _worker(args: argparse.Namespace) -> int:
|
||||
"""Run the Celery worker, with the beat scheduler embedded unless ``--no-beat`` (or on Windows)."""
|
||||
_announce_home()
|
||||
windows = sys.platform == "win32"
|
||||
argv = ["worker", "-l", args.loglevel]
|
||||
if args.queues:
|
||||
argv += ["-Q", args.queues]
|
||||
if args.concurrency:
|
||||
argv += ["--concurrency", str(args.concurrency)]
|
||||
if args.beat:
|
||||
if args.beat and windows:
|
||||
print("docsgpt: the embedded scheduler is not available on Windows; run `docsgpt beat` separately.", file=sys.stderr)
|
||||
elif args.beat:
|
||||
argv.append("-B")
|
||||
pool = args.pool or ("solo" if sys.platform == "darwin" else None)
|
||||
pool = args.pool or ("solo" if sys.platform in ("darwin", "win32") else None)
|
||||
if pool:
|
||||
argv += ["--pool", pool]
|
||||
celery.worker_main(argv)
|
||||
return 0
|
||||
return _celery(argv)
|
||||
|
||||
|
||||
def _beat(args: argparse.Namespace) -> int:
|
||||
"""Run the beat scheduler on its own (Windows, or a worker started with ``--no-beat``)."""
|
||||
_announce_home()
|
||||
return _celery(["beat", "-l", args.loglevel])
|
||||
|
||||
|
||||
def _migrate(args: argparse.Namespace) -> int:
|
||||
@@ -70,6 +100,7 @@ def _migrate(args: argparse.Namespace) -> int:
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.storage.db.bootstrap import ensure_database_ready
|
||||
|
||||
_announce_home()
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
|
||||
if not settings.POSTGRES_URI:
|
||||
print("POSTGRES_URI is not set; nothing to migrate.", file=sys.stderr)
|
||||
@@ -112,13 +143,17 @@ def build_parser() -> argparse.ArgumentParser:
|
||||
api.set_defaults(func=_api)
|
||||
|
||||
worker = commands.add_parser("worker", help="run the Celery worker (and the scheduler)")
|
||||
worker.add_argument("-Q", "--queues", default=DEFAULT_QUEUES, help=f"queues to consume (default: {DEFAULT_QUEUES})")
|
||||
worker.add_argument("-Q", "--queues", help="queues to consume (default: every configured queue)")
|
||||
worker.add_argument("--concurrency", type=int, help="worker processes (default: one per CPU)")
|
||||
worker.add_argument("--pool", help="celery pool (default: prefork; solo on macOS)")
|
||||
worker.add_argument("--pool", help="celery pool (default: prefork; solo on macOS and Windows)")
|
||||
worker.add_argument("-l", "--loglevel", default="INFO")
|
||||
worker.add_argument("--no-beat", dest="beat", action="store_false", help="do not embed the beat scheduler")
|
||||
worker.set_defaults(func=_worker)
|
||||
|
||||
beat = commands.add_parser("beat", help="run the beat scheduler on its own")
|
||||
beat.add_argument("-l", "--loglevel", default="INFO")
|
||||
beat.set_defaults(func=_beat)
|
||||
|
||||
migrate = commands.add_parser("migrate", help="create the database if needed and run the migrations")
|
||||
migrate.add_argument("--no-create", dest="create_db", action="store_false", help="fail instead of creating a missing database")
|
||||
migrate.set_defaults(func=_migrate)
|
||||
|
||||
@@ -43,5 +43,8 @@ def env_file() -> Path:
|
||||
"""The ``.env`` file settings load: ``DOCSGPT_ENV_FILE``, else ``<home>/.env``."""
|
||||
configured = os.environ.get(ENV_FILE_ENV)
|
||||
if configured:
|
||||
return Path(configured).expanduser()
|
||||
path = Path(configured).expanduser()
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"{ENV_FILE_ENV} is set to {path}, which is not a file")
|
||||
return path
|
||||
return home_dir() / ".env"
|
||||
@@ -97,11 +97,21 @@ def prefetch(names: Sequence[str], cache_dir: Optional[str] = None) -> List[str]
|
||||
return fetched
|
||||
|
||||
|
||||
def _parse(argv: Optional[Sequence[str]], prog: str, description: str) -> list[str]:
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(prog=prog, description=description)
|
||||
parser.add_argument(
|
||||
"models", nargs="*", help=f"embedding model names or aliases (default: {', '.join(DEFAULT_MODELS)})"
|
||||
)
|
||||
return parser.parse_args(argv).models or list(DEFAULT_MODELS)
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
||||
import os
|
||||
|
||||
names = list(argv) if argv else list(DEFAULT_MODELS)
|
||||
names = _parse(argv, "prefetch-models", "Download the embedding models and the tiktoken encodings into the local caches.")
|
||||
fetched = prefetch(names, os.environ.get("EMBEDDINGS_CACHE_DIR"))
|
||||
logger.info("Cached %d model(s): %s", len(fetched), ", ".join(fetched))
|
||||
encodings = prefetch_tiktoken()
|
||||
|
||||
@@ -138,8 +138,16 @@ def _docling_check() -> object:
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
import argparse
|
||||
|
||||
logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(message)s")
|
||||
models: List[str] = list(argv) if argv else list(DEFAULT_MODELS)
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="verify-offline", description="Check that the install serves its defaults without network access."
|
||||
)
|
||||
parser.add_argument(
|
||||
"models", nargs="*", help=f"embedding model names or aliases to check (default: {', '.join(DEFAULT_MODELS)})"
|
||||
)
|
||||
models: List[str] = parser.parse_args(argv).models or list(DEFAULT_MODELS)
|
||||
if _network_reachable():
|
||||
print("note network is reachable; run with --network none to prove the offline path")
|
||||
else:
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from docsgpt.core import paths
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
@@ -32,8 +34,16 @@ class TestEnvFile:
|
||||
assert paths.env_file() == tmp_path.resolve() / ".env"
|
||||
|
||||
def test_the_env_var_wins(self, monkeypatch, tmp_path):
|
||||
monkeypatch.setenv(paths.ENV_FILE_ENV, str(tmp_path / "custom.env"))
|
||||
assert paths.env_file() == tmp_path / "custom.env"
|
||||
custom = tmp_path / "custom.env"
|
||||
custom.write_text("LLM_NAME=x\n")
|
||||
monkeypatch.setenv(paths.ENV_FILE_ENV, str(custom))
|
||||
assert paths.env_file() == custom
|
||||
|
||||
def test_a_missing_configured_file_is_an_error(self, monkeypatch, tmp_path):
|
||||
"""Booting with every default because of a typo in the path must not be silent."""
|
||||
monkeypatch.setenv(paths.ENV_FILE_ENV, str(tmp_path / "nope.env"))
|
||||
with pytest.raises(FileNotFoundError, match="DOCSGPT_ENV_FILE"):
|
||||
paths.env_file()
|
||||
|
||||
|
||||
class TestSettingsFollowTheHome:
|
||||
|
||||
+46
-6
@@ -6,6 +6,7 @@ import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import click
|
||||
import pytest
|
||||
|
||||
from docsgpt import cli
|
||||
@@ -34,7 +35,7 @@ class TestTopLevel:
|
||||
|
||||
|
||||
class TestApi:
|
||||
def test_gunicorn_argv(self, monkeypatch):
|
||||
def test_gunicorn_argv(self, monkeypatch, capsys):
|
||||
run = MagicMock()
|
||||
monkeypatch.setitem(sys.modules, "gunicorn.app.wsgiapp", types.SimpleNamespace(run=run))
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
@@ -47,6 +48,7 @@ class TestApi:
|
||||
assert argv[argv.index("-w") + 1] == "2"
|
||||
assert argv[argv.index("-k") + 1] == "docsgpt.gunicorn_worker.BoundedDrainUvicornWorker"
|
||||
assert argv[argv.index("--config") + 1] == "python:docsgpt.gunicorn_conf"
|
||||
assert "data home" in capsys.readouterr().err
|
||||
|
||||
def test_reload_uses_uvicorn(self, monkeypatch):
|
||||
uvicorn = types.SimpleNamespace(run=MagicMock())
|
||||
@@ -57,31 +59,62 @@ class TestApi:
|
||||
|
||||
class TestWorker:
|
||||
@staticmethod
|
||||
def _celery(monkeypatch):
|
||||
def _celery(monkeypatch, start=None):
|
||||
celery = MagicMock()
|
||||
celery.start = start or MagicMock(return_value=0)
|
||||
monkeypatch.setitem(sys.modules, "docsgpt.app", types.SimpleNamespace(celery=celery))
|
||||
return celery
|
||||
|
||||
def test_defaults(self, monkeypatch):
|
||||
def test_defaults_consume_every_configured_queue(self, monkeypatch, capsys):
|
||||
celery = self._celery(monkeypatch)
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
assert cli.main(["worker"]) == 0
|
||||
argv = celery.worker_main.call_args.args[0]
|
||||
argv = celery.start.call_args.args[0]
|
||||
assert argv[:3] == ["worker", "-l", "INFO"]
|
||||
assert argv[argv.index("-Q") + 1] == cli.DEFAULT_QUEUES
|
||||
assert "-Q" not in argv, "a bare worker honours EMBEDDINGS_QUEUE and DOCUMENT_PARSE_QUEUE"
|
||||
assert "-B" in argv
|
||||
assert "--pool" not in argv
|
||||
assert "data home" in capsys.readouterr().err
|
||||
|
||||
def test_options_and_the_macos_solo_pool(self, monkeypatch):
|
||||
celery = self._celery(monkeypatch)
|
||||
monkeypatch.setattr(sys, "platform", "darwin")
|
||||
assert cli.main(["worker", "--no-beat", "-Q", "embeddings", "--concurrency", "2"]) == 0
|
||||
argv = celery.worker_main.call_args.args[0]
|
||||
argv = celery.start.call_args.args[0]
|
||||
assert argv[argv.index("--pool") + 1] == "solo"
|
||||
assert "-B" not in argv
|
||||
assert argv[argv.index("-Q") + 1] == "embeddings"
|
||||
assert argv[argv.index("--concurrency") + 1] == "2"
|
||||
|
||||
def test_windows_runs_solo_without_the_embedded_scheduler(self, monkeypatch, capsys):
|
||||
celery = self._celery(monkeypatch)
|
||||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
assert cli.main(["worker"]) == 0
|
||||
argv = celery.start.call_args.args[0]
|
||||
assert "-B" not in argv
|
||||
assert argv[argv.index("--pool") + 1] == "solo"
|
||||
assert "docsgpt beat" in capsys.readouterr().err
|
||||
|
||||
def test_the_worker_exit_code_is_returned(self, monkeypatch):
|
||||
self._celery(monkeypatch, start=MagicMock(return_value=1))
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
assert cli.main(["worker"]) == 1
|
||||
|
||||
def test_a_usage_error_prints_usage_instead_of_a_traceback(self, monkeypatch, capsys):
|
||||
self._celery(monkeypatch, start=MagicMock(side_effect=click.UsageError("No such option: --bogus")))
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
assert cli.main(["worker"]) == 2
|
||||
assert "No such option" in capsys.readouterr().err
|
||||
|
||||
|
||||
class TestBeat:
|
||||
def test_runs_the_scheduler_alone(self, monkeypatch):
|
||||
celery = MagicMock()
|
||||
celery.start = MagicMock(return_value=0)
|
||||
monkeypatch.setitem(sys.modules, "docsgpt.app", types.SimpleNamespace(celery=celery))
|
||||
assert cli.main(["beat", "-l", "DEBUG"]) == 0
|
||||
assert celery.start.call_args.args[0] == ["beat", "-l", "DEBUG"]
|
||||
|
||||
|
||||
class TestMigrate:
|
||||
def test_runs_the_bootstrap(self, monkeypatch):
|
||||
@@ -109,3 +142,10 @@ class TestScripts:
|
||||
def test_the_script_exit_code_is_returned(self, monkeypatch):
|
||||
monkeypatch.setattr("docsgpt.scripts.verify_offline.main", MagicMock(return_value=3))
|
||||
assert cli.main(["verify-offline"]) == 3
|
||||
|
||||
@pytest.mark.parametrize("script", ["prefetch-models", "verify-offline"])
|
||||
def test_help_is_help_not_a_model_name(self, script, capsys):
|
||||
with pytest.raises(SystemExit) as exc:
|
||||
cli.main([script, "--help"])
|
||||
assert exc.value.code == 0
|
||||
assert "models" in capsys.readouterr().out
|
||||
Reference in new issue
Block a user