prime train logs: expose -c trainer / -c inference / -c env-server

Backend's /api/v1/rft/runs/{run_id}/logs now accepts component +
env_name params (dedicated full-FT). Surface them through the CLI:

  prime train logs <run_id> -c trainer
  prime train logs <run_id> -c inference
  prime train logs <run_id> -c env-server --env <name>

Legacy --env <name>/<idx> still routes through the env-server-logs
endpoint (shared-RFT pods, cluster_id-backed lookup). Dedicated
env-server (slug, no slash) goes through the unified /logs route.

Per-rank --pod-index intentionally not exposed yet: the chart's
torchrun --local-ranks-filter=0 already collapses in-pod rank fan-out
to rank 0 stdout, and Loki's pod-label indexing in this tenant
doesn't actually filter the prime-job-* streams — per-pod inspection
on multi-node runs is kubectl + the PVC log files for now.

Author: JannikSt

packages/prime/src/prime_cli/api/rl.py

@@ -384,12 +384,31 @@ def get_run(self, run_id: str) -> RLRun:
                 raise APIError(f"Failed to get RL run: {e.response.text}")
             raise APIError(f"Failed to get RL run: {str(e)}")
 
-    def get_logs(self, run_id: str, tail_lines: int = 1000) -> str:
-        """Get orchestrator logs for an RL run."""
+    def get_logs(
+        self,
+        run_id: str,
+        tail_lines: int = 1000,
+        component: Optional[str] = None,
+        pod_index: int = 0,
+        env_name: Optional[str] = None,
+    ) -> str:
+        """Get logs for one component of an RFT run.
+
+        Defaults to the orchestrator pod. Dedicated full-FT runs additionally
+        expose `trainer`, `inference`, and `env-server` components.
+        `pod_index` narrows to a specific replica for multi-node
+        trainer/inference; `env_name` picks among per-env env-server
+        StatefulSets when `component='env-server'`.
+        """
+        params: Dict[str, Any] = {"tail_lines": tail_lines}
+        if component:
+            params["component"] = component
+        if pod_index:
+            params["pod_index"] = pod_index
+        if env_name:
+            params["env_name"] = env_name
         try:
-            response = self.client.get(
-                f"/rft/runs/{run_id}/logs", params={"tail_lines": tail_lines}
-            )
+            response = self.client.get(f"/rft/runs/{run_id}/logs", params=params)
             return response.get("logs", "")
         except Exception as e:
             if hasattr(e, "response") and hasattr(e.response, "text"):

packages/prime/src/prime_cli/commands/rl.py

@@ -1856,8 +1856,9 @@ def get_logs(
         "--component",
         "-c",
         help=(
-            "Pod to read logs from: 'orchestrator' (default) or 'env-server'. "
-            "Inferred from --env when omitted."
+            "Pod to read logs from: 'orchestrator' (default), 'trainer', "
+            "'inference', or 'env-server'. trainer/inference apply only "
+            "to dedicated full-FT runs. Inferred from --env when omitted."
         ),
     ),
     env: Optional[str] = typer.Option(
@@ -1875,30 +1876,38 @@ def get_logs(
 ) -> None:
     """Get logs for a run.
 
-    Defaults to the orchestrator pod. Pass ``--env <name>`` to read an
-    env-server pod instead — useful when an env-server is crash-looping
-    (e.g. ``ModuleNotFoundError``) and the orchestrator has stalled at
-    "Starting orchestrator step 0".
+    Defaults to the orchestrator pod. Use ``--component`` to pick one of
+    ``trainer`` / ``inference`` / ``env-server`` (dedicated full-FT only).
+    Pass ``--env <name>`` to read an env-server pod by name (shorthand for
+    ``--component=env-server``).
 
     List available pods first with ``prime train components <run_id>``.
 
+    Per-rank narrowing on multi-replica trainer/inference is not yet
+    surfaced here — `--local-ranks-filter=0` in the chart's torchrun
+    invocation already dedupes the in-pod rank fan-out, and per-pod
+    inspection on multi-node runs requires kubectl + the PVC log files.
+
     Examples:
 
         prime train logs <run_id>
         prime train logs <run_id> -f
+        prime train logs <run_id> -c trainer
+        prime train logs <run_id> -c inference
         prime train logs <run_id> --env reverse-text
         prime train logs <run_id> --env reverse-text/1 -f
     """
+    valid_components = ("orchestrator", "trainer", "inference", "env-server")
     if component is None:
         component = "env-server" if env is not None else "orchestrator"
-    elif component not in ("orchestrator", "env-server"):
+    elif component not in valid_components:
         raise typer.BadParameter(
-            f"Invalid component '{component}'. Use 'orchestrator' or 'env-server'.",
+            f"Invalid component '{component}'. Use one of: {', '.join(valid_components)}.",
             param_hint="--component",
         )
-    if component == "orchestrator" and env is not None:
+    if env is not None and component != "env-server":
         raise typer.BadParameter(
-            "--env applies only to env-server logs. Drop --component=orchestrator or drop --env.",
+            f"--env applies only to env-server logs. Drop --component={component} or drop --env.",
             param_hint="--env",
         )
     if component == "env-server" and env is None:
@@ -1912,25 +1921,46 @@ def get_logs(
         api_client = APIClient()
         rl_client = RLClient(api_client)
 
-        if component == "orchestrator":
+        if component == "env-server" and env is not None and "/" in env:
+            # Legacy shared-RFT env-server (`name/index` qualifier) — go
+            # through the dedicated env-server endpoint which uses the
+            # cluster_id-backed pod lookup path. Dedicated full-FT
+            # env-servers use the unified /logs route with
+            # component=env-server + env_name (StatefulSets always run
+            # one pod per env, so no index disambiguation needed).
+            env_name_q, env_index_q = _parse_env_qualifier(env)
+
+            def fetch(t: int) -> str:
+                return rl_client.get_env_server_logs(
+                    run_id,
+                    env_name=env_name_q,
+                    env_index=env_index_q,
+                    tail_lines=t,
+                )
+
+            label = f"env-server {env}"
+        elif component == "orchestrator":
 
             def fetch(t: int) -> str:
                 return rl_client.get_logs(run_id, tail_lines=t)
 
             label = "orchestrator"
         else:
-            assert env is not None  # narrowed by validation above
-            env_name, env_index = _parse_env_qualifier(env)
+            # trainer / inference / dedicated env-server — unified /logs
+            # route. env (no slash) names the dedicated env-server's
+            # StatefulSet.
+            fetch_component = component
+            fetch_env = env if component == "env-server" else None
 
             def fetch(t: int) -> str:
-                return rl_client.get_env_server_logs(
+                return rl_client.get_logs(
                     run_id,
-                    env_name=env_name,
-                    env_index=env_index,
                     tail_lines=t,
+                    component=fetch_component,
+                    env_name=fetch_env,
                 )
 
-            label = f"env-server {env}"
+            label = f"env-server {env}" if component == "env-server" else component
 
         _stream_logs(
             fetch_fn=fetch,