Skip to content

feat(client): unified ArcticRLClient + on-prem Ray transport - #45

Merged
sfc-gh-mwyatt merged 2 commits into
mainfrom
mwyatt/unified-client-1
Jul 28, 2026
Merged

sfc-gh-mwyatt merged 2 commits into
mainfrom
mwyatt/unified-client-1

Conversation

@sfc-gh-mwyatt

Copy link
Copy Markdown
Collaborator

Add the unified RL client frontend: a single op API (client.py) that lowers every op to a canonical Request(op, job_id, body), a Transport ABC with JobHandles/Request (transport.py), a backend config (config.py), and the shared OnPremTransport base with the in-process Ray transport (onprem.py, onprem_ray.py). ray_server.py accepts (job_id, body) on step/save_checkpoint/reset_prefix_cache so the client can forward uniformly.

Includes tests/client/test_client_ops.py covering the op->Request mapping via a FakeTransport (no GPUs/backends needed).

Add the unified RL client frontend: a single op API (client.py) that lowers
every op to a canonical Request(op, job_id, body), a Transport ABC with
JobHandles/Request (transport.py), a backend config (config.py), and the shared
OnPremTransport base with the in-process Ray transport (onprem.py, onprem_ray.py).
ray_server.py accepts (job_id, body) on step/save_checkpoint/reset_prefix_cache
so the client can forward uniformly.

Includes tests/client/test_client_ops.py covering the op->Request mapping via a
FakeTransport (no GPUs/backends needed).

Co-authored-by: Cursor <cursoragent@cursor.com>
@sfc-gh-kganesan

Copy link
Copy Markdown
Collaborator

Reviewed the whole three-PR stack end-to-end; full write-up on #47. Notes specific to this PR below.

Ran the tests on this branch: 13/13 pass.

Design

The three-layer split — client owns op vocab + job resolution, Transport ABC delivers Request(op, job_id, body, binary), OnPremTransport base owns creation/ordering/teardown, concretes implement only _start / _rpc / _destroy / _wait_running — is exactly the seam that lets #46 (HTTP) and #47 (Cortex) drop in without cross-cutting changes. JOB_CREATE_ORDER = ("sampling", "log_prob", "training") with training last as NCCL root is correct and documented. JobHandles.require() guard is a nice defensive fence. Docstrings across transport.py / onprem.py / client.py clearly state design intent — reviewer-friendly.

Plan-level work this PR opens (details on #47 as P3 / P4)

Merging this stack adds a second client with a different calling convention rather than replacing the first. AST scan:

arctic_platform/rl/http_client.py (ArcticRLHTTPClient) — every op is `async def`
arctic_platform/client/client.py  (ArcticRLClient)     — every op is `sync`

Every existing in-repo caller (the integrations/verl/ RemoteBackend, the recipes/rl/skyrl/ recipes) has to be rewritten before arctic_platform/rl/client.py can be deleted. Please add to UNIFICATION_NOTES.md: which caller migrates on what PR, when rl/*_client.py gets removed. Otherwise "one client" becomes "one-plus-legacy" indefinitely.

Related: ops present in arctic_platform/rl/*_client.py but missing from ArcticRLClient: sleep_/wake_{inference,training,log_prob}, empty_training_cache, weight_norm. Not on the immediate critical path but they block deleting the async clients. Enumerate in UNIFICATION_NOTES.md under a "not yet in unified client" section.

Branch hygiene

main has moved past #36 (the verl integration adapter); all three PRs in the stack are behind. This branch also carries stray commit e697039 (recipes/rl/skyrl … (#32)) which isn't in main and inflates the client-only diff (~866 LOC advertised, +4468 LOC actually on the branch). Rebase before merge, or land #32 first.

Not blocking

  • RayTransport._rpc creates a new event loop per op — asyncio.run(method(*args)) at arctic_platform/client/transports/onprem_ray.py:65. Every fwd_bwd / generate / step / log_probs spins up and tears down a fresh loop; on the hot path this is real overhead. Fix: hold one loop on the transport.
  • OnPremTransport.shutdown() iterates JOB_TYPES in insertion order rather than reverse creation order. Usually you want to destroy training before its samplers.
  • No OnPremTransport unit tests — everything is exercised through FakeTransport. A test that pins _init_payload composition, initialize() ordering, and shutdown() traversal would prevent regressions in feat(client): add on-prem HTTP transport #46/feat(client): add Cortex transport + DSSST1 wire codec #47.
  • No op-registry validation — if a client op is renamed and a transport misses the update, AttributeError at runtime instead of import time.
  • ArcticRLClientConfig.checkpoint_path overlaps with save_checkpoint(path=...); docstring should clarify which wins when both are set.

None of these gate merge.

@sfc-gh-truwase

Copy link
Copy Markdown
Collaborator

@sfc-gh-mwyatt, I think the gsm8k recipe of either verl or skyrl would be good gating tests for these PRs.

Comment thread arctic_platform/client/config.py Resolved
@sfc-gh-mwyatt

Copy link
Copy Markdown
Collaborator Author

@sfc-gh-mwyatt, I think the gsm8k recipe of either verl or skyrl would be good gating tests for these PRs.

Let's add that in a fast followup after getting #45, #46, and #47 merged.

@sfc-gh-mwyatt

Copy link
Copy Markdown
Collaborator Author

main has moved past #36 (the verl integration adapter); all three PRs in the stack are behind. This branch also carries stray commit e697039 (recipes/rl/skyrl … (#32)) which isn't in main and inflates the client-only diff (~866 LOC advertised, +4468 LOC actually on the branch). Rebase before merge, or land #32 first.

@sfc-gh-kganesan I think this is your agent doing something wrong because this is not the case. Ignored, but please confirm manually on your side if you think I'm wrong.

@sfc-gh-mwyatt

Copy link
Copy Markdown
Collaborator Author

OnPremTransport.shutdown() iterates JOB_TYPES in insertion order rather than reverse creation order. Usually you want to destroy training before its samplers.

@sfc-gh-kganesan Is there a reason to do reverse order? Has this been a problem with onprem deployments previously?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants