From 275d66bc3d9d9885942f7b14d7d5ca1944768cee Mon Sep 17 00:00:00 2001 From: JoyboyBrian Date: Sun, 30 Aug 2026 02:49:58 -0700 Subject: [PATCH 1/3] [misc] refactor: simplify internal implementation --- osmosis_ai/cli/output/__init__.py | 29 +-- osmosis_ai/cli/output/display.py | 10 +- osmosis_ai/platform/auth/credentials.py | 16 +- osmosis_ai/platform/auth/flow.py | 7 +- osmosis_ai/platform/cli/dataset.py | 73 ++---- osmosis_ai/rollout/controller/listener.py | 1 - osmosis_ai/rollout/controller/llm_bridge.py | 14 +- osmosis_ai/rollout/trajectory/save.py | 114 ++++---- osmosis_ai/rollout/utils/concurrency.py | 21 +- osmosis_ai/rollout/utils/http.py | 23 +- osmosis_ai/templates/catalog.py | 28 +- osmosis_ai/templates/registry.py | 7 +- tests/unit/auth/test_credentials.py | 2 +- tests/unit/cli/output/test_console_facade.py | 2 + tests/unit/cli/output/test_error.py | 19 +- tests/unit/cli/output/test_renderer_json.py | 36 +-- tests/unit/cli/output/test_serializers.py | 2 +- tests/unit/cli/test_eval_rubric_json.py | 59 ++++- tests/unit/eval/local/test_results.py | 42 +-- tests/unit/eval/local/test_runner_e2e.py | 22 -- tests/unit/eval/local/test_runner_units.py | 15 +- .../unit/platform/cli/test_dataset_upload.py | 45 ++-- tests/unit/rollout/test_utils_concurrency.py | 12 + tests/unit/test_cli.py | 144 +++------- tests/unit/test_rubric_cli_command.py | 245 ++++++------------ 25 files changed, 361 insertions(+), 627 deletions(-) diff --git a/osmosis_ai/cli/output/__init__.py b/osmosis_ai/cli/output/__init__.py index 5c1255a3..b0a24c24 100644 --- a/osmosis_ai/cli/output/__init__.py +++ b/osmosis_ai/cli/output/__init__.py @@ -51,34 +51,19 @@ ) _SERIALIZER_EXPORTS: dict[str, tuple[str, str]] = { - "serialize_benchmark_run": ( - "osmosis_ai.cli.output.serializers", + name: ("osmosis_ai.cli.output.serializers", name) + for name in ( "serialize_benchmark_run", - ), - "serialize_checkpoint": ( - "osmosis_ai.cli.output.serializers", "serialize_checkpoint", - ), - "serialize_dataset": ("osmosis_ai.cli.output.serializers", "serialize_dataset"), - "serialize_dev_rollout_server": ( - "osmosis_ai.cli.output.serializers", + "serialize_dataset", "serialize_dev_rollout_server", - ), - "serialize_environment_secret": ( - "osmosis_ai.cli.output.serializers", "serialize_environment_secret", - ), - "serialize_eval_run": ("osmosis_ai.cli.output.serializers", "serialize_eval_run"), - "serialize_lora_model": ( - "osmosis_ai.cli.output.serializers", + "serialize_eval_run", "serialize_lora_model", - ), - "serialize_model": ("osmosis_ai.cli.output.serializers", "serialize_model"), - "serialize_rollout": ("osmosis_ai.cli.output.serializers", "serialize_rollout"), - "serialize_training_run": ( - "osmosis_ai.cli.output.serializers", + "serialize_model", + "serialize_rollout", "serialize_training_run", - ), + ) } diff --git a/osmosis_ai/cli/output/display.py b/osmosis_ai/cli/output/display.py index 0ffff39f..636dcd0c 100644 --- a/osmosis_ai/cli/output/display.py +++ b/osmosis_ai/cli/output/display.py @@ -20,12 +20,6 @@ def _parse_iso_datetime(value: str | None) -> datetime | None: return parsed -def _localize(dt: datetime, *, tz: tzinfo | None = None) -> datetime: - if tz is not None: - return dt.astimezone(tz) - return dt.astimezone() - - def _twelve_hour_time(dt: datetime, *, with_seconds: bool = False) -> str: """12-hour clock time with AM/PM and no leading-zero hour (e.g. ``6:16 PM``).""" hour = dt.hour % 12 or 12 @@ -40,7 +34,7 @@ def format_local_date( parsed = _parse_iso_datetime(value) if parsed is None: return "" if value is None else str(value)[:10] - local = _localize(parsed, tz=tz) + local = parsed.astimezone(tz) return f"{local.strftime('%Y-%m-%d')} {_twelve_hour_time(local)} {local.strftime('%Z')}" @@ -50,7 +44,7 @@ def format_local_datetime( parsed = _parse_iso_datetime(value) if parsed is None: return "" if value is None else str(value) - local = _localize(parsed, tz=tz) + local = parsed.astimezone(tz) return ( f"{local.strftime('%Y-%m-%d')} " f"{_twelve_hour_time(local, with_seconds=True)} {local.strftime('%Z')}" diff --git a/osmosis_ai/platform/auth/credentials.py b/osmosis_ai/platform/auth/credentials.py index c478ae77..6fbf064f 100644 --- a/osmosis_ai/platform/auth/credentials.py +++ b/osmosis_ai/platform/auth/credentials.py @@ -239,16 +239,6 @@ def _is_default_platform_url(platform_url: str) -> bool: return normalize_platform_url(platform_url) == _default_platform_url() -def _dedupe(values: list[str]) -> list[str]: - result: list[str] = [] - seen: set[str] = set() - for value in values: - if value and value not in seen: - result.append(value) - seen.add(value) - return result - - def _is_platform_registry(data: dict[str, Any]) -> bool: return "platforms" in data @@ -354,7 +344,7 @@ def _keyring_accounts_for_entry( if is_default_platform: accounts.append(KEYRING_ACCOUNT) - return _dedupe(accounts) + return list(dict.fromkeys(account for account in accounts if account)) def _cleanup_platform_keyring_entries( @@ -644,7 +634,7 @@ def save_credentials( del registry["platforms"][old_key] registry["platforms"][platform_url] = data try: - atomic_write_json(CREDENTIALS_FILE, registry, mode=0o600) + atomic_write_json(CREDENTIALS_FILE, registry) except Exception: if keyring_account not in old_keyring_accounts: try: @@ -783,7 +773,7 @@ def delete_credentials( del registry["platforms"][platform_key] if registry["platforms"]: - atomic_write_json(CREDENTIALS_FILE, registry, mode=0o600) + atomic_write_json(CREDENTIALS_FILE, registry) return True try: diff --git a/osmosis_ai/platform/auth/flow.py b/osmosis_ai/platform/auth/flow.py index cbeb591a..128fc62d 100644 --- a/osmosis_ai/platform/auth/flow.py +++ b/osmosis_ai/platform/auth/flow.py @@ -355,11 +355,10 @@ def poll_device_token( except HTTPError as e: if e.code in (426, 429, 500, 502, 503, 504): raise _login_error_from_http(e, "Polling failed") from e - try: - error_data = json.loads(e.read().decode()) - error_code = error_data.get("error", "") - except Exception: + error_data = _read_error_body(e) + if not error_data: raise LoginError(f"Polling failed: HTTP {e.code}") from e + error_code = error_data.get("error", "") if error_code == "authorization_pending": if on_poll: diff --git a/osmosis_ai/platform/cli/dataset.py b/osmosis_ai/platform/cli/dataset.py index ab2b14bb..3a5a24bf 100644 --- a/osmosis_ai/platform/cli/dataset.py +++ b/osmosis_ai/platform/cli/dataset.py @@ -303,66 +303,45 @@ def _perform_upload( ctx, progress_cb = make_progress_bar(file_size) - if is_multipart: - try: - with ctx: + parts = None + try: + with ctx: + if is_multipart: parts = upload_file_multipart( file_path, upload_info, progress_callback=progress_cb ) - except KeyboardInterrupt: - _abort_upload( - client, - dataset.id, - credentials=credentials, - git_identity=git_identity, - ) - raise - except Exception as e: - _abort_upload( - client, - dataset.id, - credentials=credentials, - git_identity=git_identity, - ) - raise CLIError(f"Upload failed: {e}") from e - completed = _complete_with_retry( + else: + upload_file_simple( + file_path, upload_info, progress_callback=progress_cb + ) + except KeyboardInterrupt: + if not is_multipart: + console.print("\nUpload interrupted.") + _abort_upload( client, dataset.id, - parts=parts, - file_extension=ext, credentials=credentials, git_identity=git_identity, ) - else: - try: - with ctx: - upload_file_simple( - file_path, upload_info, progress_callback=progress_cb - ) - except KeyboardInterrupt: - console.print("\nUpload interrupted.") - _abort_upload( - client, - dataset.id, - credentials=credentials, - git_identity=git_identity, - ) - raise - except Exception as e: - _abort_upload( - client, - dataset.id, - credentials=credentials, - git_identity=git_identity, - ) - raise CLIError(f"Upload failed: {e}") from e - completed = _complete_with_retry( + raise + except Exception as e: + _abort_upload( client, dataset.id, - file_extension=ext, credentials=credentials, git_identity=git_identity, ) + raise CLIError(f"Upload failed: {e}") from e + + complete_kwargs = {"parts": parts} if is_multipart else {} + completed = _complete_with_retry( + client, + dataset.id, + **complete_kwargs, + file_extension=ext, + credentials=credentials, + git_identity=git_identity, + ) return completed or dataset diff --git a/osmosis_ai/rollout/controller/listener.py b/osmosis_ai/rollout/controller/listener.py index 7eb1a074..21681692 100644 --- a/osmosis_ai/rollout/controller/listener.py +++ b/osmosis_ai/rollout/controller/listener.py @@ -234,7 +234,6 @@ def __init__( advertised_base_url: str | None = None, bridge_keepalive: bool = False, ) -> None: - _assert_non_empty_auth_token(auth_token) app = create_callback_app(store, auth_token=auth_token) if bridge is not None: if bridge_token is None: diff --git a/osmosis_ai/rollout/controller/llm_bridge.py b/osmosis_ai/rollout/controller/llm_bridge.py index 8f89e7e5..4d4b530b 100644 --- a/osmosis_ai/rollout/controller/llm_bridge.py +++ b/osmosis_ai/rollout/controller/llm_bridge.py @@ -27,9 +27,13 @@ from osmosis_ai._imports import raise_optional_dependency_error from osmosis_ai.rollout.controller.openai_responses import ( + _compact_json, build_responses_kwargs, to_chat_response, ) +from osmosis_ai.rollout.controller.openai_responses import ( + _field as _getattr_or_key, +) try: from fastapi import APIRouter, Depends, HTTPException, Request @@ -97,12 +101,6 @@ def _get_litellm() -> Any: return litellm -def _getattr_or_key(obj: Any, name: str, default: Any = None) -> Any: - if isinstance(obj, dict): - return obj.get(name, default) - return getattr(obj, name, default) - - def _unsupported_openai_sampling_param(exc: Exception, *, litellm: Any) -> str | None: """Return a rejected sampling parameter from a structured OpenAI 400.""" bad_request_error = getattr(litellm, "BadRequestError", None) @@ -260,10 +258,6 @@ def _model_response_to_payload( return payload -def _compact_json(data: Any) -> str: - return json.dumps(data, separators=(",", ":")) - - class LiteLLMBridge: """Convert OpenAI-format chat requests to any litellm provider in-process.""" diff --git a/osmosis_ai/rollout/trajectory/save.py b/osmosis_ai/rollout/trajectory/save.py index 8109a451..adda0396 100644 --- a/osmosis_ai/rollout/trajectory/save.py +++ b/osmosis_ai/rollout/trajectory/save.py @@ -57,81 +57,61 @@ async def save_trajectory( ``diagnostics`` overrides ``result.extra_fields`` for the sidecar. """ try: - await _save( + root = artifact_root or default_artifact_root() + # Written before the sample-None early return so failures leave a record. + payload = diagnostics if diagnostics is not None else result.extra_fields + if payload is not None: + diagnostics_dest = root / rollout_id / "diagnostics.json" + diagnostics_data = json.dumps( + payload, ensure_ascii=False, indent=2, sort_keys=True, default=str + ).encode() + await asyncio.to_thread(_write_document, diagnostics_dest, diagnostics_data) + logger.info( + "Saved rollout diagnostics for %s -> %s", + rollout_id, + diagnostics_dest, + ) + + sample = result.sample + if sample is None: + return + if sample.trajectory_messages is None: + # Explicit opt-out, or an upstream conversion/snapshot failure + # that already warned with a traceback -- not worth a warning here. + logger.info( + "Skipping trajectory for rollout %s: no trajectory messages " + "(persistence disabled or conversion failed upstream)", + rollout_id, + ) + return + + matched_report, unmatched_reports = _resolve_sample_report(report) + if unmatched_reports: + logger.warning( + "Trajectory report for rollout %s has %d entries but the rollout " + "produced one sample; preserving them under " + "extra.osmosis.unmatched_sample_reports", + rollout_id, + len(unmatched_reports), + ) + trajectory = convert_sample_to_trajectory( + sample, rollout_id=rollout_id, - result=result, request_label=request_label, request_metadata=request_metadata, request_extra_fields=request_extra_fields, - report=report, - artifact_root=artifact_root or default_artifact_root(), - diagnostics=diagnostics, + report=matched_report, + default_model_name=report.model_name if report else None, + unmatched_sample_reports=unmatched_reports or None, ) + dest = root / rollout_id / "trajectory.json" + # Keep large token-id/logprob arrays compact inside the pretty document. + data = format_trajectory_json(trajectory.to_json_dict()).encode() + await asyncio.to_thread(_write_document, dest, data) + logger.info("Saved trajectory document for rollout %s -> %s", rollout_id, dest) except Exception: logger.warning( "Failed to save the trajectory for rollout %s (best-effort)", rollout_id, exc_info=True, ) - - -async def _save( - *, - rollout_id: str, - result: ExecutionResult, - request_label: str | None, - request_metadata: dict[str, Any] | None, - request_extra_fields: dict[str, Any] | None, - report: TrajectoryReport | None, - artifact_root: Path, - diagnostics: dict[str, Any] | None = None, -) -> None: - # Written before the sample-None early return so failures leave a record. - payload = diagnostics if diagnostics is not None else result.extra_fields - if payload is not None: - diagnostics_dest = artifact_root / rollout_id / "diagnostics.json" - diagnostics_data = json.dumps( - payload, ensure_ascii=False, indent=2, sort_keys=True, default=str - ).encode() - await asyncio.to_thread(_write_document, diagnostics_dest, diagnostics_data) - logger.info( - "Saved rollout diagnostics for %s -> %s", rollout_id, diagnostics_dest - ) - - sample = result.sample - if sample is None: - return - if sample.trajectory_messages is None: - # Explicit opt-out, or an upstream conversion/snapshot failure - # that already warned with a traceback -- not worth a warning here. - logger.info( - "Skipping trajectory for rollout %s: no trajectory messages " - "(persistence disabled or conversion failed upstream)", - rollout_id, - ) - return - - matched_report, unmatched_reports = _resolve_sample_report(report) - if unmatched_reports: - logger.warning( - "Trajectory report for rollout %s has %d entries but the rollout " - "produced one sample; preserving them under " - "extra.osmosis.unmatched_sample_reports", - rollout_id, - len(unmatched_reports), - ) - trajectory = convert_sample_to_trajectory( - sample, - rollout_id=rollout_id, - request_label=request_label, - request_metadata=request_metadata, - request_extra_fields=request_extra_fields, - report=matched_report, - default_model_name=report.model_name if report else None, - unmatched_sample_reports=unmatched_reports or None, - ) - dest = artifact_root / rollout_id / "trajectory.json" - # Keep large token-id/logprob arrays compact inside the pretty document. - data = format_trajectory_json(trajectory.to_json_dict()).encode() - await asyncio.to_thread(_write_document, dest, data) - logger.info("Saved trajectory document for rollout %s -> %s", rollout_id, dest) diff --git a/osmosis_ai/rollout/utils/concurrency.py b/osmosis_ai/rollout/utils/concurrency.py index bcb38b45..e32ed783 100644 --- a/osmosis_ai/rollout/utils/concurrency.py +++ b/osmosis_ai/rollout/utils/concurrency.py @@ -14,28 +14,21 @@ def __init__(self, *, max_concurrent: int | None) -> None: @asynccontextmanager async def acquire(self) -> AsyncIterator[None]: - if self._semaphore is None: - self.running += 1 + semaphore = self._semaphore + if semaphore is not None: + self.queued += 1 try: - yield + await semaphore.acquire() finally: - self.running -= 1 - return + self.queued -= 1 - self.queued += 1 - try: - await self._semaphore.acquire() - except BaseException: - self.queued -= 1 - raise - - self.queued -= 1 self.running += 1 try: yield finally: self.running -= 1 - self._semaphore.release() + if semaphore is not None: + semaphore.release() def snapshot(self) -> dict[str, int | None]: return { diff --git a/osmosis_ai/rollout/utils/http.py b/osmosis_ai/rollout/utils/http.py index 90b4244a..83357683 100644 --- a/osmosis_ai/rollout/utils/http.py +++ b/osmosis_ai/rollout/utils/http.py @@ -36,8 +36,6 @@ async def post_json_with_retry( raise ValueError("max_attempts must be >= 1") client = get_shared_client(timeout_seconds) - last_exception: Exception | None = None - for attempt in range(1, max_attempts + 1): try: response = await client.post(url, json=payload, headers=headers) @@ -51,10 +49,7 @@ async def post_json_with_retry( response.raise_for_status() return response except (httpx.RequestError, httpx.HTTPStatusError) as exc: - last_exception = exc - is_last_attempt = attempt >= max_attempts - should_retry = _is_retryable_exception(exc) - if is_last_attempt or not should_retry: + if attempt >= max_attempts or not _is_retryable_exception(exc): raise delay = min(base_delay_seconds * (2 ** (attempt - 1)), max_delay_seconds) @@ -70,17 +65,11 @@ async def post_json_with_retry( await asyncio.sleep(delay) # This should be unreachable, but keeps type-checkers satisfied. - raise RuntimeError( - "POST request failed without raising an exception" - ) from last_exception + raise RuntimeError("POST request failed without raising an exception") def _is_retryable_exception(exc: Exception) -> bool: - if isinstance(exc, httpx.RequestError): - return True - - if isinstance(exc, httpx.HTTPStatusError): - status_code = exc.response.status_code - return status_code in {429, 500, 502, 503, 504} - - return False + return isinstance(exc, httpx.RequestError) or ( + isinstance(exc, httpx.HTTPStatusError) + and exc.response.status_code in {429, 500, 502, 503, 504} + ) diff --git a/osmosis_ai/templates/catalog.py b/osmosis_ai/templates/catalog.py index 4192a834..5c4119e0 100644 --- a/osmosis_ai/templates/catalog.py +++ b/osmosis_ai/templates/catalog.py @@ -30,11 +30,7 @@ class ScaffoldEntry: official: bool = False -def _path(value: str) -> Path: - return Path(value) - - -_MULTIPLY_DATA_PATH = _path("data/multiply.jsonl") +_MULTIPLY_DATA_PATH = Path("data/multiply.jsonl") def _recipe(name: str, description: str) -> TemplateRecipe: @@ -42,12 +38,12 @@ def _recipe(name: str, description: str) -> TemplateRecipe: name=name, description=description, files=( - _path(f"rollouts/{name}/**"), - _path(f"configs/eval/{name}.toml"), - _path(f"configs/training/{name}.toml"), + Path(f"rollouts/{name}/**"), + Path(f"configs/eval/{name}.toml"), + Path(f"configs/training/{name}.toml"), _MULTIPLY_DATA_PATH, ), - owned_dirs=(_path(f"rollouts/{name}"),), + owned_dirs=(Path(f"rollouts/{name}"),), next_steps=( f"pip install -e rollouts/{name}", "git push", @@ -66,16 +62,16 @@ def _recipe(name: str, description: str) -> TemplateRecipe: OFFICIAL_AGENT_SCAFFOLD_PATHS: tuple[Path, ...] = ( - _path("AGENTS.md"), - _path("CLAUDE.md"), - _path("configs/AGENTS.md"), + Path("AGENTS.md"), + Path("CLAUDE.md"), + Path("configs/AGENTS.md"), ) REQUIRED_WORKSPACE_DIRS: tuple[Path, ...] = ( - _path("rollouts"), - _path("configs/training"), - _path("configs/eval"), - _path("data"), + Path("rollouts"), + Path("configs/training"), + Path("configs/eval"), + Path("data"), ) diff --git a/osmosis_ai/templates/registry.py b/osmosis_ai/templates/registry.py index 38c91c5f..d6f77d09 100644 --- a/osmosis_ai/templates/registry.py +++ b/osmosis_ai/templates/registry.py @@ -42,11 +42,8 @@ def _expand_catalog_files(root: Path, patterns: tuple[Path, ...]) -> list[Path]: rel_paths: set[Path] = set() for pattern in patterns: pattern_text = pattern.as_posix() - if any(part in {"*", "**"} or "*" in part for part in pattern.parts): - if pattern.parts[-1] == "**": - matches = sorted((root / Path(*pattern.parts[:-1])).rglob("*")) - else: - matches = sorted(root.glob(pattern_text)) + if pattern.parts[-1] == "**": + matches = sorted((root / Path(*pattern.parts[:-1])).rglob("*")) file_matches = [path for path in matches if path.is_file()] if not file_matches: raise CLIError( diff --git a/tests/unit/auth/test_credentials.py b/tests/unit/auth/test_credentials.py index 97fd7fd6..80f22e6c 100644 --- a/tests/unit/auth/test_credentials.py +++ b/tests/unit/auth/test_credentials.py @@ -203,7 +203,7 @@ def fake_set(account: str, token: str) -> bool: events.append(("set", account)) return True - def fake_write(path, data, *, mode): + def fake_write(path, data): events.append( ("metadata", data["platforms"][DEFAULT_PLATFORM]["keyring_account"]) ) diff --git a/tests/unit/cli/output/test_console_facade.py b/tests/unit/cli/output/test_console_facade.py index 27935de9..461bbf88 100644 --- a/tests/unit/cli/output/test_console_facade.py +++ b/tests/unit/cli/output/test_console_facade.py @@ -122,6 +122,8 @@ def test_console_print_warning_is_structured_json_on_stderr() -> None: assert out.getvalue() == "" payload = json.loads(err.getvalue()) + assert list(payload) == ["schema_version", "cli_version", "warning"] + assert list(payload["warning"]) == ["code", "message"] assert payload["schema_version"] == 1 assert "cli_version" in payload assert payload["warning"] == { diff --git a/tests/unit/cli/output/test_error.py b/tests/unit/cli/output/test_error.py index 712b8885..755a7df5 100644 --- a/tests/unit/cli/output/test_error.py +++ b/tests/unit/cli/output/test_error.py @@ -41,13 +41,13 @@ def _capture_envelope(err: CLIError) -> dict[str, Any]: def test_envelope_keys_match_golden() -> None: envelope = _capture_envelope(CLIError("Bad input.", code="VALIDATION")) expected = json.loads((GOLDEN / "error_envelope.json").read_text(encoding="utf-8")) - assert sorted(envelope.keys()) == sorted(expected["keys"]) + assert list(envelope) == expected["keys"] assert envelope["schema_version"] == 1 assert envelope["command"] == "dataset list" assert envelope["cli_version"] assert envelope["error"]["code"] == "VALIDATION" assert envelope["error"]["details"] == {} - assert sorted(envelope["error"].keys()) == sorted(expected["error_keys"]) + assert list(envelope["error"]) == expected["error_keys"] assert "request_id" not in envelope["error"] @@ -113,11 +113,6 @@ def test_billing_required_maps_to_billing_code() -> None: assert cli_err.code == "BILLING_REQUIRED" -def test_generic_403_stays_platform_error() -> None: - cli_err = classify_error(PlatformAPIError("forbidden", status_code=403)) - assert cli_err.code == "PLATFORM_ERROR" - - def test_authentication_expired_error_maps_to_auth_required() -> None: cli_err = classify_error(AuthenticationExpiredError("expired")) assert cli_err.code == "AUTH_REQUIRED" @@ -370,16 +365,6 @@ def test_command_registry_matches_registered_app() -> None: assert command_names == STANDALONE_COMMANDS -def test_command_path_uses_click_context_when_available() -> None: - parent = Context(typer.core.TyperCommand(name="osmosis")) - parent.info_name = "osmosis" - middle = Context(typer.core.TyperCommand(name="dataset"), parent=parent) - middle.info_name = "dataset" - nested = Context(typer.core.TyperCommand(name="list"), parent=middle) - nested.info_name = "list" - assert command_path_for_error(nested) == "dataset list" - - def test_command_path_root_when_argv_empty(monkeypatch) -> None: monkeypatch.setattr("sys.argv", ["osmosis"]) assert command_path_for_error(None) == "" diff --git a/tests/unit/cli/output/test_renderer_json.py b/tests/unit/cli/output/test_renderer_json.py index fb34bcd8..1541a740 100644 --- a/tests/unit/cli/output/test_renderer_json.py +++ b/tests/unit/cli/output/test_renderer_json.py @@ -47,7 +47,7 @@ def test_list_envelope_required_keys() -> None: expected_keys = json.loads( (GOLDEN / "list_envelope.json").read_text(encoding="utf-8") )["keys"] - assert sorted(payload.keys()) == sorted(expected_keys) + assert list(payload) == expected_keys assert payload["schema_version"] == 1 assert payload["items"] == [{"id": "ds_1"}] assert payload["next_offset"] is None @@ -121,18 +121,9 @@ def test_sectioned_list_envelope_required_keys() -> None: golden = json.loads( (GOLDEN / "sectioned_list_envelope.json").read_text(encoding="utf-8") ) - assert sorted(payload.keys()) == sorted(golden["keys"]) + assert list(payload) == golden["keys"] for section_key in ("base_models", "lora_models"): - assert sorted(payload[section_key].keys()) == sorted(golden["section_keys"]) - assert payload["schema_version"] == 1 - assert stderr == "" - - -def test_sectioned_list_envelope_keys_each_section_with_own_pagination() -> None: - payload, stderr = _render_to_json(_sectioned_list_result()) - assert sorted(payload.keys()) == sorted( - ["schema_version", "base_models", "lora_models"] - ) + assert list(payload[section_key]) == golden["section_keys"] assert payload["schema_version"] == 1 assert payload["base_models"] == { "items": [{"id": "m_1"}], @@ -179,7 +170,7 @@ def test_detail_envelope_required_keys() -> None: expected_keys = json.loads( (GOLDEN / "detail_envelope.json").read_text(encoding="utf-8") )["keys"] - assert sorted(payload.keys()) == sorted(expected_keys) + assert list(payload) == expected_keys assert payload["data"] == {"id": "ds_1"} @@ -188,12 +179,20 @@ def test_operation_envelope_required_keys() -> None: operation="deploy", status="success", resource={"id": "dep_1", "checkpoint_name": "run-step-40", "status": "active"}, + message="Deployed.", + next_steps_structured=[{"action": "model_info", "model_name": "run-step-40"}], + extra={"workspace": "ws-a"}, ) payload, _ = _render_to_json(result) expected_keys = json.loads( (GOLDEN / "operation_envelope.json").read_text(encoding="utf-8") )["keys"] - assert set(expected_keys["required"]).issubset(payload.keys()) + expected_order = [ + *expected_keys["required"], + *(key for key in expected_keys["optional"] if key in payload), + "workspace", + ] + assert list(payload) == expected_order assert payload["status"] == "success" assert payload["operation"] == "deploy" @@ -246,15 +245,6 @@ def test_no_ansi_or_rich_box_on_json_stdout() -> None: json.loads(raw) -def test_render_marks_output_emitted() -> None: - result = OperationResult(operation="logout", status="success", message="ok") - out = io.StringIO() - with override_output_context(format=OutputFormat.json) as ctx: - with redirect_stdout(out): - render(result, ctx) - assert ctx.output_emitted is True - - def test_json_envelope_rejects_nonfinite_floats() -> None: result = DetailResult(title="Metrics", data={"value": float("nan")}) with override_output_context(format=OutputFormat.json) as ctx: diff --git a/tests/unit/cli/output/test_serializers.py b/tests/unit/cli/output/test_serializers.py index 2b915508..7af05c0e 100644 --- a/tests/unit/cli/output/test_serializers.py +++ b/tests/unit/cli/output/test_serializers.py @@ -29,7 +29,7 @@ def _assert_keys_match_golden(payload: dict, golden_name: str) -> None: expected = json.loads((GOLDEN_DIR / golden_name).read_text(encoding="utf-8")) - assert sorted(payload.keys()) == sorted(expected["keys"]) + assert list(payload) == expected["keys"] def test_serialize_dataset_keys() -> None: diff --git a/tests/unit/cli/test_eval_rubric_json.py b/tests/unit/cli/test_eval_rubric_json.py index 7506aa42..28c672c7 100644 --- a/tests/unit/cli/test_eval_rubric_json.py +++ b/tests/unit/cli/test_eval_rubric_json.py @@ -52,8 +52,46 @@ def test_eval_rubric_json_returns_operation_result( assert payload["schema_version"] == 1 assert payload["status"] == "success" assert payload["operation"] == "eval.rubric" - assert payload["resource"]["statistics"]["average"] == pytest.approx(0.8) - assert payload["resource"]["records"][0]["scores"] == [0.8] + resource = payload["resource"] + assert list(resource) == [ + "model", + "data_path", + "number", + "statistics", + "record_count", + "error_count", + "records", + ] + assert resource == { + "model": "openai/gpt-5.4", + "data_path": str(data_path), + "number": 1, + "statistics": { + "average": 0.8, + "variance": 0.0, + "stdev": 0.0, + "min": 0.8, + "max": 0.8, + }, + "record_count": 1, + "error_count": 0, + "records": [ + { + "index": 1, + "label": "record[1]", + "scores": [0.8], + "explanations": ["good"], + "errors": [], + "statistics": { + "average": 0.8, + "variance": 0.0, + "stdev": 0.0, + "min": 0.8, + "max": 0.8, + }, + } + ], + } def test_eval_rubric_json_output_file_omits_records( @@ -91,9 +129,21 @@ def test_eval_rubric_json_output_file_omits_records( captured = capsys.readouterr() assert exit_code == 0 payload = json.loads(captured.out) - assert payload["resource"]["output_path"] == str(output_path) - assert "records" not in payload["resource"] + resource = payload["resource"] + assert list(resource) == [ + "model", + "data_path", + "number", + "statistics", + "record_count", + "error_count", + "output_path", + ] + assert resource["output_path"] == str(output_path) + assert "records" not in resource assert output_path.exists() + written = json.loads(output_path.read_text(encoding="utf-8")) + assert written["records"][0]["label"] == "row-1" def test_eval_rubric_json_suppresses_tty_progress( @@ -169,6 +219,7 @@ def test_eval_rubric_plain_allows_tty_progress( captured = capsys.readouterr() assert exit_code == 0 + assert captured.out == "Rubric evaluation completed.\n" assert captured.err != "" diff --git a/tests/unit/eval/local/test_results.py b/tests/unit/eval/local/test_results.py index 6e3448ea..b572f17b 100644 --- a/tests/unit/eval/local/test_results.py +++ b/tests/unit/eval/local/test_results.py @@ -208,23 +208,6 @@ def _row( return payload -def test_pass_rate_excludes_skipped_from_scored() -> None: - summary = aggregate_metrics( - [ - _row(0, 0, "success", 1.0), - _row(1, 0, "success", 0.0), - _row(2, 0, "skipped", None), - ], - pass_threshold=1.0, - ) - assert summary["total_samples"] == 3 - assert summary["skipped"] == 1 - assert summary["completed_samples"] == 2 - assert summary["graded"] == 2 - assert summary["passed"] == 1 - assert summary["pass_rate"] == 0.5 - - def test_passed_uses_a_greater_or_equal_threshold() -> None: summary = aggregate_metrics([_row(0, 0, "success", 0.7)], pass_threshold=0.7) assert summary["passed"] == 1 @@ -545,28 +528,23 @@ def test_a_cancelled_attempt_is_never_projected_as_success(tmp_path: Path) -> No # --------------------------------------------------------------------------- # -def test_reward_less_failures_stay_in_the_pass_rate_denominator() -> None: - # Excluding them would report pass_rate 1.0 for a run where half the rows - # failed -- the metric would hide exactly what the user needs to see. - summary = aggregate_metrics( - [_row(0, 0, "success", 1.0), _row(1, 0, "failed", None)], pass_threshold=1.0 - ) - assert summary["graded"] == 1 - assert summary["completed_samples"] == 2 - assert summary["pass_rate"] == 0.5 - - def test_skipped_rows_are_the_only_thing_excluded_from_scored() -> None: summary = aggregate_metrics( [ _row(0, 0, "success", 1.0), - _row(1, 0, "failed", None), - _row(2, 0, "skipped", None), + _row(1, 0, "success", 0.0), + _row(2, 0, "failed", None), + _row(3, 0, "skipped", None), ], pass_threshold=1.0, ) - assert summary["completed_samples"] == 2 - assert summary["pass_rate"] == 0.5 + assert summary["total_samples"] == 4 + assert summary["skipped"] == 1 + assert summary["failed"] == 1 + assert summary["completed_samples"] == 3 + assert summary["graded"] == 2 + assert summary["passed"] == 1 + assert summary["pass_rate"] == 1 / 3 def test_pass_at_k_counts_a_reward_less_failure_as_a_non_pass() -> None: diff --git a/tests/unit/eval/local/test_runner_e2e.py b/tests/unit/eval/local/test_runner_e2e.py index 15d5af9a..c76934f8 100644 --- a/tests/unit/eval/local/test_runner_e2e.py +++ b/tests/unit/eval/local/test_runner_e2e.py @@ -376,28 +376,6 @@ async def test_a_missing_entrypoint_fails_before_dispatch( # --------------------------------------------------------------------------- # -# Smoke coverage only: these prove the knobs are accepted end to end and every -# work item still lands. The bound itself is unit-tested against -# ``_resolve_concurrency`` in test_runner_units.py. - - -async def test_batch_size_is_accepted_and_the_run_completes( - harness: RunnerHarness, -) -> None: - summary = await harness.runner(spec=harness.spec(batch_size=2)).run() - assert summary.dispatched == 4 - - -async def test_max_in_flight_and_batch_size_together_complete_the_run( - harness: RunnerHarness, -) -> None: - summary = await harness.runner( - spec=harness.spec(batch_size=1), - options=LocalEvalOptions(name="run-1", max_in_flight=4), - ).run() - assert summary.dispatched == 4 - - async def test_confirmation_receives_the_pending_count(harness: RunnerHarness) -> None: hooks = RecordingHooks() await harness.runner(hooks=hooks).run() diff --git a/tests/unit/eval/local/test_runner_units.py b/tests/unit/eval/local/test_runner_units.py index efea2fbc..71b9dedd 100644 --- a/tests/unit/eval/local/test_runner_units.py +++ b/tests/unit/eval/local/test_runner_units.py @@ -37,7 +37,6 @@ compute_source_digest, reserve_free_port, ) -from osmosis_ai.eval.local.state import digest_of from osmosis_ai.rollout.controller import TerminalCallbackResult from osmosis_ai.rollout.types import GraderCompleteRequest, GraderStatus, RolloutSample @@ -167,7 +166,7 @@ def test_semantic_changes_change_the_fingerprint( ) -> None: baseline = _inputs(_spec(), _dataset(), tmp_path) changed = _inputs(_spec(**{field_name: value}), _dataset(), tmp_path) - assert digest_of(changed) != digest_of(baseline) + assert changed != baseline # The refusal message has to be able to name what moved. assert changed_input_keys(baseline, changed) != [] @@ -181,7 +180,7 @@ def test_pass_threshold_does_not_change_the_fingerprint(tmp_path: Path) -> None: def test_dataset_bytes_change_the_fingerprint(tmp_path: Path) -> None: baseline = _inputs(_spec(), _dataset(), tmp_path) changed = _inputs(_spec(), _dataset(sha="c" * 64), tmp_path) - assert digest_of(changed) != digest_of(baseline) + assert changed != baseline def test_source_digest_changes_the_fingerprint(tmp_path: Path) -> None: @@ -194,13 +193,13 @@ def test_source_digest_changes_the_fingerprint(tmp_path: Path) -> None: second = build_run_inputs( _spec(), dataset=_dataset(), selection=selection, rollout_source_digest="2" * 64 ) - assert digest_of(first) != digest_of(second) + assert first != second def test_row_selection_changes_the_fingerprint(tmp_path: Path) -> None: baseline = _inputs(_spec(), _dataset(), tmp_path) narrowed = _inputs(_spec(), _dataset(), tmp_path, row_selector=(0,)) - assert digest_of(narrowed) != digest_of(baseline) + assert narrowed != baseline assert baseline["dataset"]["selected_source_rows"] == "0-1" assert narrowed["dataset"]["selected_source_rows"] == "0" @@ -216,13 +215,13 @@ def test_throughput_and_display_fields_are_excluded( } baseline = _inputs(_spec(), _dataset(), tmp_path) changed = _inputs(_spec(**{field_name: values[field_name]}), _dataset(), tmp_path) - assert digest_of(changed) == digest_of(baseline) + assert changed == baseline def test_secret_names_are_included_and_values_are_never_present(tmp_path: Path) -> None: baseline = _inputs(_spec(), _dataset(), tmp_path) with_secret = _inputs(_spec(secret_names=("OPENAI_API_KEY",)), _dataset(), tmp_path) - assert digest_of(with_secret) != digest_of(baseline) + assert with_secret != baseline assert with_secret["secret_names"] == ["OPENAI_API_KEY"] assert "OPENAI_API_KEY" not in str(with_secret).replace("'OPENAI_API_KEY'", "") @@ -230,7 +229,7 @@ def test_secret_names_are_included_and_values_are_never_present(tmp_path: Path) def test_env_ordering_does_not_change_the_fingerprint(tmp_path: Path) -> None: first = _inputs(_spec(env={"A": "1", "B": "2"}), _dataset(), tmp_path) second = _inputs(_spec(env={"B": "2", "A": "1"}), _dataset(), tmp_path) - assert digest_of(first) == digest_of(second) + assert first == second def test_changed_input_keys_names_the_changed_top_level_keys() -> None: diff --git a/tests/unit/platform/cli/test_dataset_upload.py b/tests/unit/platform/cli/test_dataset_upload.py index 1a84914f..bf991d1c 100644 --- a/tests/unit/platform/cli/test_dataset_upload.py +++ b/tests/unit/platform/cli/test_dataset_upload.py @@ -3,9 +3,11 @@ from __future__ import annotations from contextlib import nullcontext +from io import StringIO import pytest +from osmosis_ai.cli.console import Console from osmosis_ai.cli.output.context import OutputFormat, override_output_context from osmosis_ai.platform.api.models import DatasetFile, UploadInfo @@ -72,7 +74,8 @@ def fake_platform_call(message, call): assert result.id == "dataset-1" assert messages == ["Uploading dataset..."] - def test_simple_upload_flow(self, monkeypatch, tmp_path): + @pytest.mark.parametrize("method", ["simple", "multipart"]) + def test_upload_flow(self, monkeypatch, tmp_path, method): """_perform_upload creates dataset, uploads to S3, and completes.""" import osmosis_ai.platform.api.client as api_client_module import osmosis_ai.platform.api.upload as upload_module @@ -83,7 +86,10 @@ def test_simple_upload_flow(self, monkeypatch, tmp_path): calls: dict[str, bool] = {} fake_credentials = object() - fake_dataset = _make_fake_dataset(file_size=file_size) + fake_dataset = _make_fake_dataset(file_size=file_size, method=method) + expected_parts = ( + [{"PartNumber": 1, "ETag": "etag"}] if method == "multipart" else None + ) class FakeClient: def create_dataset( @@ -107,7 +113,7 @@ def complete_upload( ): calls["complete"] = True assert file_id == "dataset-1" - assert parts is None + assert parts == expected_parts assert git_identity == GIT_IDENTITY assert credentials is fake_credentials return DatasetFile( @@ -123,13 +129,12 @@ def complete_upload( "make_progress_bar", lambda _size: (nullcontext(), lambda _done, _total: None), ) - monkeypatch.setattr( - upload_module, - "upload_file_simple", - lambda _fp, _info, progress_callback=None: calls.update( - {"s3_upload": True} - ), - ) + + def fake_upload(_fp, _info, progress_callback=None): + calls["s3_upload"] = True + return expected_parts + + monkeypatch.setattr(upload_module, f"upload_file_{method}", fake_upload) from osmosis_ai.platform.cli.dataset import _perform_upload @@ -238,21 +243,29 @@ def create_dataset(self, *args, git_identity, **kwargs): credentials=None, ) - def test_abort_on_keyboard_interrupt(self, monkeypatch, tmp_path): + @pytest.mark.parametrize( + ("method", "shows_interrupted_message"), + [("simple", True), ("multipart", False)], + ) + def test_abort_on_keyboard_interrupt( + self, monkeypatch, tmp_path, method, shows_interrupted_message + ): """_perform_upload aborts and re-raises KeyboardInterrupt.""" import osmosis_ai.platform.api.client as api_client_module import osmosis_ai.platform.api.upload as upload_module + import osmosis_ai.platform.cli.dataset as dataset_module file_path = tmp_path / "data.jsonl" file_path.write_text("{}") file_size = file_path.stat().st_size aborted = {} + output = StringIO() class FakeClient: def create_dataset(self, *args, git_identity, **kwargs): assert git_identity == GIT_IDENTITY - return _make_fake_dataset(file_size=file_size) + return _make_fake_dataset(file_size=file_size, method=method) def abort_upload(self, file_id, *, git_identity, credentials=None): assert git_identity == GIT_IDENTITY @@ -262,6 +275,7 @@ def complete_upload(self, *args, **kwargs): pass monkeypatch.setattr(api_client_module, "OsmosisClient", FakeClient) + monkeypatch.setattr(dataset_module, "console", Console(file=output)) monkeypatch.setattr( upload_module, "make_progress_bar", @@ -269,16 +283,14 @@ def complete_upload(self, *args, **kwargs): ) monkeypatch.setattr( upload_module, - "upload_file_simple", + f"upload_file_{method}", lambda _fp, _info, progress_callback=None: (_ for _ in ()).throw( KeyboardInterrupt ), ) - from osmosis_ai.platform.cli.dataset import _perform_upload - with pytest.raises(KeyboardInterrupt): - _perform_upload( + dataset_module._perform_upload( file_path=file_path, ext="jsonl", file_size=file_size, @@ -287,6 +299,7 @@ def complete_upload(self, *args, **kwargs): ) assert aborted.get("called") + assert ("Upload interrupted." in output.getvalue()) is shows_interrupted_message def test_duplicate_name_without_overwrite_raises_guided_conflict( self, monkeypatch, tmp_path diff --git a/tests/unit/rollout/test_utils_concurrency.py b/tests/unit/rollout/test_utils_concurrency.py index 1893f2c4..00ef2edd 100644 --- a/tests/unit/rollout/test_utils_concurrency.py +++ b/tests/unit/rollout/test_utils_concurrency.py @@ -32,6 +32,18 @@ async def task(n: int) -> None: await asyncio.gather(task(1), task(2)) assert sorted(order) == [1, 2] + async def test_cancelled_waiter_restores_queue_count(self): + limiter = ConcurrencyLimiter(max_concurrent=1) + + async with limiter.acquire(): + waiter = asyncio.create_task(limiter.acquire().__aenter__()) + await asyncio.sleep(0) + assert limiter.queued == 1 + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert limiter.queued == 0 + async def test_snapshot_reflects_state(self): limiter = ConcurrencyLimiter(max_concurrent=2) assert limiter.snapshot()["running"] == 0 diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index aeb879b2..ff780c7b 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -1,142 +1,64 @@ import json import sys -from pathlib import Path from unittest.mock import AsyncMock -import pytest - -PROJECT_ROOT = Path(__file__).resolve().parents[1] -if str(PROJECT_ROOT) not in sys.path: - sys.path.insert(0, str(PROJECT_ROOT)) - from osmosis_ai.cli import main as cli +from osmosis_ai.cli.console import Console from osmosis_ai.eval.rubric.types import RubricResult -# ============================================================================= -# eval rubric — happy-path -# ============================================================================= - -def test_eval_rubric_basic(tmp_path, monkeypatch, capsys): - """eval rubric runs successfully with mocked evaluate_rubric.""" +def test_eval_rubric_rich_output_contract(tmp_path, monkeypatch, capsys): data_path = tmp_path / "records.jsonl" - record = { - "messages": [ - {"role": "user", "content": "Help me"}, - {"role": "assistant", "content": "Sure, I can help."}, - ] - } - data_path.write_text(json.dumps(record) + "\n", encoding="utf-8") - - monkeypatch.setenv("OPENAI_API_KEY", "test-key") - - mock_eval = AsyncMock( - return_value=RubricResult(score=0.85, explanation="Good response") + data_path.write_text( + json.dumps({"solution_str": "answer"}) + "\n", encoding="utf-8" ) - monkeypatch.setattr("osmosis_ai.eval.rubric.cli.evaluate_rubric", mock_eval) - - exit_code = cli.main( - [ - "eval", - "rubric", - "-d", - str(data_path), - "--rubric", - "Score quality of the assistant response.", - "--model", - "openai/gpt-5.4", - ] + output_path = tmp_path / "results.json" + monkeypatch.setattr( + "osmosis_ai.eval.rubric.cli.evaluate_rubric", + AsyncMock(return_value=RubricResult(score=0.85, explanation="Good response")), ) - - capsys.readouterr() - assert exit_code == 0 - mock_eval.assert_called_once() - - -def test_eval_rubric_with_output(tmp_path, monkeypatch, capsys): - """eval rubric writes JSON output when --output is specified.""" - data_path = tmp_path / "records.jsonl" - record = { - "messages": [ - {"role": "user", "content": "Help me"}, - {"role": "assistant", "content": "Sure, I can help."}, - ] - } - data_path.write_text(json.dumps(record) + "\n", encoding="utf-8") - - monkeypatch.setenv("OPENAI_API_KEY", "test-key") - - mock_eval = AsyncMock( - return_value=RubricResult(score=0.85, explanation="Good response") + capsys.readouterr() # Ignore optional-dependency import diagnostics. + monkeypatch.setattr( + "osmosis_ai.cli.console.console", + Console(file=sys.stdout, force_terminal=False, width=1000), ) - monkeypatch.setattr("osmosis_ai.eval.rubric.cli.evaluate_rubric", mock_eval) - output_path = tmp_path / "results.json" exit_code = cli.main( [ "eval", "rubric", "-d", str(data_path), - "--rubric", - "Score quality of the assistant response.", + "-r", + "Score quality.", "--model", "openai/gpt-5.4", + "-n", + "1", "-o", str(output_path), ] ) - capsys.readouterr() + captured = capsys.readouterr() assert exit_code == 0 - assert output_path.exists() - - payload = json.loads(output_path.read_text(encoding="utf-8")) - assert "overall_statistics" in payload - assert "records" in payload - assert payload["overall_statistics"]["average"] == pytest.approx(0.85, rel=1e-6) - - -def test_eval_rubric_multiple_runs(tmp_path, monkeypatch, capsys): - """eval rubric correctly handles --number for multiple runs per record.""" - data_path = tmp_path / "records.jsonl" - record = { - "messages": [ - {"role": "user", "content": "Help me"}, - {"role": "assistant", "content": "Sure, I can help."}, - ] - } - data_path.write_text(json.dumps(record) + "\n", encoding="utf-8") - - monkeypatch.setenv("OPENAI_API_KEY", "test-key") - - call_count = 0 - - async def mock_eval_fn(**kwargs): - nonlocal call_count - call_count += 1 - score = 0.4 + 0.1 * call_count - return RubricResult(score=score, explanation=f"run-{call_count}") - - monkeypatch.setattr("osmosis_ai.eval.rubric.cli.evaluate_rubric", mock_eval_fn) - - exit_code = cli.main( - [ - "eval", - "rubric", - "-d", - str(data_path), - "--rubric", - "Score quality.", - "--model", - "openai/gpt-5.4", - "-n", - "3", - ] + assert captured.err == "" + assert captured.out == ( + "Model: openai/gpt-5.4\n" + f"Evaluated 1 record(s) from {data_path}\n" + "Runs per record: 1\n" + "\n" + "[record[1]]\n" + " Run 01: score=0.8500\n" + " explanation: Good response\n" + "\n" + "Overall Statistics:\n" + " average: 0.8500\n" + " stdev: 0.0000\n" + " min/max: 0.8500 / 0.8500\n" + f"Wrote results to {output_path}\n" ) - - assert exit_code == 0 - assert call_count == 3 + assert output_path.is_file() # ============================================================================= diff --git a/tests/unit/test_rubric_cli_command.py b/tests/unit/test_rubric_cli_command.py index 88143213..fb6af376 100644 --- a/tests/unit/test_rubric_cli_command.py +++ b/tests/unit/test_rubric_cli_command.py @@ -10,11 +10,9 @@ import pytest from osmosis_ai.cli.errors import CLIError +from osmosis_ai.cli.output import OutputFormat, override_output_context from osmosis_ai.eval.rubric.cli import RubricCommand -from osmosis_ai.eval.rubric.dataset import ( - RubricRecord, - load_rubric_dataset, -) +from osmosis_ai.eval.rubric.dataset import RubricRecord, load_rubric_dataset from osmosis_ai.eval.rubric.report import ( ConsoleReportRenderer, JsonReportWriter, @@ -81,16 +79,6 @@ def test_messages_format_loads_correctly( assert records[0].metadata is None assert records[0].record_id is None - def test_solution_str_format_auto_converts(self, tmp_path: Path): - data_file = tmp_path / "data.jsonl" - record = {"solution_str": "The answer is 42."} - data_file.write_text(json.dumps(record) + "\n", encoding="utf-8") - - records = load_rubric_dataset(data_file) - - assert len(records) == 1 - assert records[0].solution_str == "The answer is 42." - def test_missing_messages_and_solution_str_raises(self, tmp_path: Path): data_file = tmp_path / "data.jsonl" record = {"some_other_field": "value"} @@ -108,13 +96,6 @@ def test_invalid_json_raises(self, tmp_path: Path): with pytest.raises(CLIError, match="Invalid JSON on line 1"): load_rubric_dataset(data_file) - def test_empty_file_raises(self, tmp_path: Path): - data_file = tmp_path / "data.jsonl" - data_file.write_text("", encoding="utf-8") - - with pytest.raises(CLIError, match="No JSON records found"): - load_rubric_dataset(data_file) - def test_blank_lines_only_raises(self, tmp_path: Path): data_file = tmp_path / "data.jsonl" data_file.write_text("\n\n\n", encoding="utf-8") @@ -163,17 +144,6 @@ def test_id_takes_precedence_over_conversation_id(self, tmp_path: Path): records = load_rubric_dataset(data_file) assert records[0].record_id == "primary-id" - def test_multiple_records(self, tmp_path: Path): - data_file = tmp_path / "data.jsonl" - lines = [ - json.dumps({"messages": [{"role": "assistant", "content": f"Answer {i}"}]}) - for i in range(3) - ] - data_file.write_text("\n".join(lines) + "\n", encoding="utf-8") - - records = load_rubric_dataset(data_file) - assert len(records) == 3 - def test_non_dict_json_raises(self, tmp_path: Path): data_file = tmp_path / "data.jsonl" data_file.write_text("[1, 2, 3]\n", encoding="utf-8") @@ -223,33 +193,13 @@ def test_blank_lines_skipped(self, tmp_path: Path): assert len(records) == 2 -# ============================================================================= -# RubricRecord.label Tests -# ============================================================================= - - -class TestRubricRecordLabel: - """Tests for RubricRecord.label() method.""" - - def test_with_record_id_returns_record_id(self): - record = RubricRecord( - solution_str="test", - ground_truth=None, - original_input=None, - metadata=None, - record_id="my-record-id", - ) - assert record.label(5) == "my-record-id" - - def test_without_record_id_returns_indexed_label(self): - record = RubricRecord( - solution_str="test", - ground_truth=None, - original_input=None, - metadata=None, - record_id=None, - ) - assert record.label(3) == "record[3]" +@pytest.mark.parametrize( + ("record_id", "index", "expected"), + [("my-record-id", 5, "my-record-id"), (None, 3, "record[3]")], +) +def test_rubric_record_label(record_id: str | None, index: int, expected: str) -> None: + record = RubricRecord("test", None, None, None, record_id) + assert record.label(index) == expected # ============================================================================= @@ -287,52 +237,17 @@ def test_multiple_scores(self): # ============================================================================= -# ConsoleReportRenderer Tests +# Console report Tests # ============================================================================= -class TestConsoleReportRenderer: - """Tests for the console report renderer.""" - - def test_renders_basic_report(self, tmp_path: Path): - lines: list[str] = [] - renderer = ConsoleReportRenderer(printer=lines.append) - +class TestConsoleReport: + def test_renders_exact_output(self, tmp_path: Path): + data_path = tmp_path / "data.jsonl" report = RubricReport( model="openai/gpt-5.4", - rubric_text="Score quality", - data_path=tmp_path / "data.jsonl", - number=1, - results=[ - RecordResult( - record_index=1, - label="rec-1", - scores=[0.85], - explanations=["Good"], - errors=[], - statistics=calculate_statistics([0.85]), - ) - ], - overall_statistics=calculate_statistics([0.85]), - ) - renderer.render(report) - - output = "\n".join(lines) - assert "Model: openai/gpt-5.4" in output - assert "Evaluated 1 record(s)" in output - assert "[rec-1]" in output - assert "score=0.8500" in output - assert "explanation: Good" in output - assert "Overall Statistics:" in output - - def test_renders_multi_run_summary(self, tmp_path: Path): - lines: list[str] = [] - renderer = ConsoleReportRenderer(printer=lines.append) - - report = RubricReport( - model="openai/gpt-5.4", - rubric_text="Score quality", - data_path=tmp_path / "data.jsonl", + rubric_text="评分质量", + data_path=data_path, number=2, results=[ RecordResult( @@ -342,57 +257,55 @@ def test_renders_multi_run_summary(self, tmp_path: Path): explanations=["Good", "Better"], errors=[], statistics=calculate_statistics([0.8, 0.9]), - ) - ], - overall_statistics=calculate_statistics([0.8, 0.9]), - ) - renderer.render(report) - - output = "\n".join(lines) - assert "Summary: avg=" in output - - def test_renders_errors(self, tmp_path: Path): - lines: list[str] = [] - renderer = ConsoleReportRenderer(printer=lines.append) - - report = RubricReport( - model="openai/gpt-5.4", - rubric_text="Score quality", - data_path=tmp_path / "data.jsonl", - number=1, - results=[ + ), RecordResult( - record_index=1, - label="rec-1", + record_index=2, + label="rec-2", scores=[], explanations=[], errors=["Something went wrong"], statistics=calculate_statistics([]), - ) + ), ], - overall_statistics=calculate_statistics([]), + overall_statistics=calculate_statistics([0.8, 0.9]), ) - renderer.render(report) - - output = "\n".join(lines) - assert "ERROR: Something went wrong" in output + lines: list[str] = [] + ConsoleReportRenderer(lines.append).render(report) + + assert lines == [ + "Model: openai/gpt-5.4", + f"Evaluated 2 record(s) from {data_path}", + "Runs per record: 2", + "", + "[rec-1]", + " Run 01: score=0.8000", + " explanation: Good", + " Run 02: score=0.9000", + " explanation: Better", + " Summary: avg=0.8500 stdev=0.0500 min=0.8000 max=0.9000", + "", + "[rec-2]", + " ERROR: Something went wrong", + "", + "Overall Statistics:", + " average: 0.8500", + " stdev: 0.0500", + " min/max: 0.8000 / 0.9000", + ] # ============================================================================= -# JsonReportWriter Tests +# JSON report Tests # ============================================================================= -class TestJsonReportWriter: - """Tests for the JSON report writer.""" - +class TestJsonReport: def test_writes_valid_json(self, tmp_path: Path): - writer = JsonReportWriter() output_path = tmp_path / "output" / "result.json" report = RubricReport( model="openai/gpt-5.4", - rubric_text="Score quality", + rubric_text="评分质量", data_path=tmp_path / "data.jsonl", number=1, results=[ @@ -408,19 +321,39 @@ def test_writes_valid_json(self, tmp_path: Path): overall_statistics=calculate_statistics([0.85]), ) - result_path = writer.write(report, output_path) + result_path = JsonReportWriter().write(report, output_path) assert result_path == output_path assert output_path.exists() - data = json.loads(output_path.read_text(encoding="utf-8")) + raw = output_path.read_text(encoding="utf-8") + data = json.loads(raw) + assert raw == json.dumps(data, indent=2, ensure_ascii=False) + assert list(data) == [ + "generated_at", + "model", + "rubric", + "data_path", + "number", + "overall_statistics", + "records", + ] + assert data["generated_at"].endswith("+00:00") assert data["model"] == "openai/gpt-5.4" - assert data["rubric"] == "Score quality" + assert data["rubric"] == "评分质量" + assert data["data_path"] == str(report.data_path) assert data["number"] == 1 + assert data["overall_statistics"] == calculate_statistics([0.85]) assert len(data["records"]) == 1 - assert data["records"][0]["scores"] == [0.85] - assert "generated_at" in data - assert "overall_statistics" in data + assert data["records"][0] == { + "index": 1, + "label": "rec-1", + "scores": [0.85], + "explanations": ["Good"], + "errors": [], + "statistics": calculate_statistics([0.85]), + } + assert data["records"][0] == report.results[0].to_payload() # ============================================================================= @@ -435,13 +368,6 @@ def test_inline_text_returned_stripped(self): result = RubricCommand._resolve_rubric_text(" Score quality ") assert result == "Score quality" - def test_file_reference_reads_content(self, tmp_path: Path): - rubric_file = tmp_path / "rubric.txt" - rubric_file.write_text(" Score factual accuracy. \n", encoding="utf-8") - - result = RubricCommand._resolve_rubric_text(f"@{rubric_file}") - assert result == "Score factual accuracy." - def test_nonexistent_file_raises(self): with pytest.raises(CLIError, match="does not exist"): RubricCommand._resolve_rubric_text("@/nonexistent/rubric.txt") @@ -458,27 +384,10 @@ def test_nonexistent_file_raises(self): class TestRubricCommandRun: """End-to-end tests for RubricCommand.run.""" - def test_run_succeeds(self, tmp_path: Path): - data_file = tmp_path / "data.jsonl" - record = { - "messages": [ - {"role": "user", "content": "What is 2+2?"}, - {"role": "assistant", "content": "4"}, - ] - } - data_file.write_text(json.dumps(record) + "\n", encoding="utf-8") - - mock_result = RubricResult(score=0.9, explanation="Correct") - - with patch( - _EVALUATE_RUBRIC_PATCH, new_callable=AsyncMock, return_value=mock_result - ): - RubricCommand().run( - data=str(data_file), - rubric="Score accuracy", - model="openai/gpt-5.4", - api_key="test-key", - ) + @pytest.fixture(autouse=True) + def _json_output(self): + with override_output_context(format=OutputFormat.json): + yield def test_run_with_output_writes_json(self, tmp_path: Path): data_file = tmp_path / "data.jsonl" @@ -540,7 +449,7 @@ def test_run_with_rubric_file(self, tmp_path: Path): data_file.write_text(json.dumps(record) + "\n", encoding="utf-8") rubric_file = tmp_path / "rubric.txt" - rubric_file.write_text("Score the response quality.", encoding="utf-8") + rubric_file.write_text(" Score the response quality. \n", encoding="utf-8") mock_result = RubricResult(score=0.5, explanation="Average") From 24db5d998dd1e0d8469b3f9669a7458c46d0b0e6 Mon Sep 17 00:00:00 2001 From: JoyboyBrian Date: Sun, 30 Aug 2026 03:05:18 -0700 Subject: [PATCH 2/3] [misc] fix: preserve template wildcard expansion --- osmosis_ai/templates/registry.py | 7 +++++-- tests/unit/templates/test_registry.py | 10 ++++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/osmosis_ai/templates/registry.py b/osmosis_ai/templates/registry.py index d6f77d09..38c91c5f 100644 --- a/osmosis_ai/templates/registry.py +++ b/osmosis_ai/templates/registry.py @@ -42,8 +42,11 @@ def _expand_catalog_files(root: Path, patterns: tuple[Path, ...]) -> list[Path]: rel_paths: set[Path] = set() for pattern in patterns: pattern_text = pattern.as_posix() - if pattern.parts[-1] == "**": - matches = sorted((root / Path(*pattern.parts[:-1])).rglob("*")) + if any(part in {"*", "**"} or "*" in part for part in pattern.parts): + if pattern.parts[-1] == "**": + matches = sorted((root / Path(*pattern.parts[:-1])).rglob("*")) + else: + matches = sorted(root.glob(pattern_text)) file_matches = [path for path in matches if path.is_file()] if not file_matches: raise CLIError( diff --git a/tests/unit/templates/test_registry.py b/tests/unit/templates/test_registry.py index d4164c94..be202cf1 100644 --- a/tests/unit/templates/test_registry.py +++ b/tests/unit/templates/test_registry.py @@ -9,6 +9,7 @@ from osmosis_ai.cli.errors import CLIError from osmosis_ai.templates.registry import ( TemplateNotFoundError, + _expand_catalog_files, iter_template_files, list_templates, ) @@ -89,6 +90,15 @@ def test_iter_template_files_returns_catalog_relative_paths( assert not rel.is_absolute() +def test_expand_catalog_files_supports_non_terminal_wildcards( + workspace_template: Path, +) -> None: + assert _expand_catalog_files(workspace_template, (Path("configs/*/*.toml"),)) == [ + Path("configs/eval/multiply-local-strands.toml"), + Path("configs/training/multiply-local-strands.toml"), + ] + + def test_iter_template_files_missing_catalog_file_uses_user_facing_template_terms( workspace_template: Path, ) -> None: From eda200fbef80b11a2a10c9a27fcf07f8a803bae7 Mon Sep 17 00:00:00 2001 From: JoyboyBrian Date: Sun, 30 Aug 2026 03:18:56 -0700 Subject: [PATCH 3/3] [rollout] fix: preserve trajectory failure logs --- osmosis_ai/rollout/trajectory/save.py | 114 +++++++++++++++----------- 1 file changed, 67 insertions(+), 47 deletions(-) diff --git a/osmosis_ai/rollout/trajectory/save.py b/osmosis_ai/rollout/trajectory/save.py index adda0396..8109a451 100644 --- a/osmosis_ai/rollout/trajectory/save.py +++ b/osmosis_ai/rollout/trajectory/save.py @@ -57,61 +57,81 @@ async def save_trajectory( ``diagnostics`` overrides ``result.extra_fields`` for the sidecar. """ try: - root = artifact_root or default_artifact_root() - # Written before the sample-None early return so failures leave a record. - payload = diagnostics if diagnostics is not None else result.extra_fields - if payload is not None: - diagnostics_dest = root / rollout_id / "diagnostics.json" - diagnostics_data = json.dumps( - payload, ensure_ascii=False, indent=2, sort_keys=True, default=str - ).encode() - await asyncio.to_thread(_write_document, diagnostics_dest, diagnostics_data) - logger.info( - "Saved rollout diagnostics for %s -> %s", - rollout_id, - diagnostics_dest, - ) - - sample = result.sample - if sample is None: - return - if sample.trajectory_messages is None: - # Explicit opt-out, or an upstream conversion/snapshot failure - # that already warned with a traceback -- not worth a warning here. - logger.info( - "Skipping trajectory for rollout %s: no trajectory messages " - "(persistence disabled or conversion failed upstream)", - rollout_id, - ) - return - - matched_report, unmatched_reports = _resolve_sample_report(report) - if unmatched_reports: - logger.warning( - "Trajectory report for rollout %s has %d entries but the rollout " - "produced one sample; preserving them under " - "extra.osmosis.unmatched_sample_reports", - rollout_id, - len(unmatched_reports), - ) - trajectory = convert_sample_to_trajectory( - sample, + await _save( rollout_id=rollout_id, + result=result, request_label=request_label, request_metadata=request_metadata, request_extra_fields=request_extra_fields, - report=matched_report, - default_model_name=report.model_name if report else None, - unmatched_sample_reports=unmatched_reports or None, + report=report, + artifact_root=artifact_root or default_artifact_root(), + diagnostics=diagnostics, ) - dest = root / rollout_id / "trajectory.json" - # Keep large token-id/logprob arrays compact inside the pretty document. - data = format_trajectory_json(trajectory.to_json_dict()).encode() - await asyncio.to_thread(_write_document, dest, data) - logger.info("Saved trajectory document for rollout %s -> %s", rollout_id, dest) except Exception: logger.warning( "Failed to save the trajectory for rollout %s (best-effort)", rollout_id, exc_info=True, ) + + +async def _save( + *, + rollout_id: str, + result: ExecutionResult, + request_label: str | None, + request_metadata: dict[str, Any] | None, + request_extra_fields: dict[str, Any] | None, + report: TrajectoryReport | None, + artifact_root: Path, + diagnostics: dict[str, Any] | None = None, +) -> None: + # Written before the sample-None early return so failures leave a record. + payload = diagnostics if diagnostics is not None else result.extra_fields + if payload is not None: + diagnostics_dest = artifact_root / rollout_id / "diagnostics.json" + diagnostics_data = json.dumps( + payload, ensure_ascii=False, indent=2, sort_keys=True, default=str + ).encode() + await asyncio.to_thread(_write_document, diagnostics_dest, diagnostics_data) + logger.info( + "Saved rollout diagnostics for %s -> %s", rollout_id, diagnostics_dest + ) + + sample = result.sample + if sample is None: + return + if sample.trajectory_messages is None: + # Explicit opt-out, or an upstream conversion/snapshot failure + # that already warned with a traceback -- not worth a warning here. + logger.info( + "Skipping trajectory for rollout %s: no trajectory messages " + "(persistence disabled or conversion failed upstream)", + rollout_id, + ) + return + + matched_report, unmatched_reports = _resolve_sample_report(report) + if unmatched_reports: + logger.warning( + "Trajectory report for rollout %s has %d entries but the rollout " + "produced one sample; preserving them under " + "extra.osmosis.unmatched_sample_reports", + rollout_id, + len(unmatched_reports), + ) + trajectory = convert_sample_to_trajectory( + sample, + rollout_id=rollout_id, + request_label=request_label, + request_metadata=request_metadata, + request_extra_fields=request_extra_fields, + report=matched_report, + default_model_name=report.model_name if report else None, + unmatched_sample_reports=unmatched_reports or None, + ) + dest = artifact_root / rollout_id / "trajectory.json" + # Keep large token-id/logprob arrays compact inside the pretty document. + data = format_trajectory_json(trajectory.to_json_dict()).encode() + await asyncio.to_thread(_write_document, dest, data) + logger.info("Saved trajectory document for rollout %s -> %s", rollout_id, dest)