[misc] refactor: simplify internal implementation

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",
-    ),
+    )
 }
 
 

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')}"

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:

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:

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
 

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:

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."""
 

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 {

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}
+    )