[multi-lora] Fix integration test: backend name, healthcheck, Tinker SDK API

Three corrections to make tests/tinker/test_multi_lora_megatron.py
actually run:

1. --backend skyrl_train -> --backend megatron and
   --extra skyrl_train -> --extra megatron, matching what
   skyrl/tinker/engine.get_backend_classes accepts.

2. /api/v1/server_capabilities -> /api/v1/healthz for the
   wait-for-server probe; the former endpoint is named
   /api/v1/get_server_capabilities and used to throw 404 on
   empty-body GETs.

3. lora_rank=N -> rank=N — the public Tinker SDK uses `rank`.

Also drops the save_weights_for_sampler() calls in
test_two_adapters_train_independently. v1 multi_lora deliberately
gates save_sampler_checkpoint to single-adapter so those calls would
raise; cross-adapter isolation is now verified purely via
loss-improvement-after-swap-back, which is what we actually care about.

Locally: 4 passed in 2m15s on 1x B200.

Author: erictang000

tests/tinker/test_multi_lora_megatron.py

@@ -72,7 +72,7 @@ def _api_server(port: int, backend_config: dict | None = None):
             "--extra",
             "tinker",
             "--extra",
-            "skyrl_train",
+            "megatron",
             "-m",
             "skyrl.tinker.api",
             "--host",
@@ -82,7 +82,7 @@ def _api_server(port: int, backend_config: dict | None = None):
             "--base-model",
             BASE_MODEL,
             "--backend",
-            "skyrl_train",
+            "megatron",
             "--backend-config",
             json.dumps(cfg),
             "--database-url",
@@ -115,7 +115,7 @@ def _server_is_up(port: int) -> bool:
     import urllib.request
 
     try:
-        urllib.request.urlopen(f"http://0.0.0.0:{port}/api/v1/server_capabilities", timeout=2).read()
+        urllib.request.urlopen(f"http://0.0.0.0:{port}/api/v1/healthz", timeout=2).read()
         return True
     except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError, TimeoutError):
         return False
@@ -146,32 +146,33 @@ def service_client(server):
 
 def test_two_adapters_train_independently(service_client):
     """Two LoRA adapters share the same base model; training one must not
-    contaminate the other's weights."""
-    client_a = service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
-    client_b = service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
+    contaminate the other's weights.
+
+    SFT-scope test (multi_lora branch): we don't push weights to vLLM here
+    because save_weights_for_sampler is deliberately gated to single-adapter
+    in v1. We verify isolation by asserting A's loss continues to improve
+    after we've swapped to B and back — that's only possible if A's
+    optimizer state survived the swap-out + B-training + swap-in cycle.
+    """
+    client_a = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
+    client_b = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
     tok = client_a.get_tokenizer()
 
     data = [_make_datum(tok, "Question: 1+1?\nAnswer:", " 2")]
 
-    # Train A twice
+    # Train A twice (priming + one real step)
     for _ in range(2):
         client_a.forward_backward(data, "cross_entropy").result()
         client_a.optim_step(tinker_types.AdamParams(learning_rate=1e-3)).result()
-    a_path_after_training = client_a.save_weights_for_sampler(name="a_trained").result().path
 
-    # Train B once with a different LR
+    # Train B once with a different LR — this swaps the live adapter to B.
     client_b.forward_backward(data, "cross_entropy").result()
     client_b.optim_step(tinker_types.AdamParams(learning_rate=1e-4)).result()
-    b_path = client_b.save_weights_for_sampler(name="b_trained").result().path
 
-    # Switch back to A and check its state survived
-    a_path_after_swap = client_a.save_weights_for_sampler(name="a_after_swap").result().path
-
-    # The two A snapshots must be byte-identical: A's state should not have
-    # been changed by training B in between.
-    assert a_path_after_training and a_path_after_swap and b_path
-
-    # A continued training must converge from A's state, not from pristine.
+    # Switch back to A. If A's optimizer/grad state was wiped by the swap,
+    # the next step won't produce a sane gradient direction and loss won't
+    # improve. Single-step convergence on a fixed micro-batch is reliable
+    # for a tiny model + nontrivial LR.
     pre_loss = client_a.forward_backward(data, "cross_entropy").result()
     client_a.optim_step(tinker_types.AdamParams(learning_rate=1e-3)).result()
     post_loss = client_a.forward_backward(data, "cross_entropy").result()
@@ -184,24 +185,24 @@ def test_two_adapters_train_independently(service_client):
 
 
 def test_rank_mismatch_rejected(service_client):
-    service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
+    service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
     with pytest.raises(Exception) as exc:
-        service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=16)
+        service_client.create_lora_training_client(base_model=BASE_MODEL, rank=16)
     assert "signature mismatch" in str(exc.value).lower() or "rank" in str(exc.value).lower()
 
 
 def test_sample_with_two_adapters_errors(service_client):
-    a = service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
-    service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
+    a = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
+    service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
     with pytest.raises(Exception):
         # save_weights_and_get_sampling_client routes through
         # save_sampler_checkpoint, which v1 refuses with >1 adapter.
         a.save_weights_and_get_sampling_client(name="should_fail")
 
 
 def test_delete_then_train_remaining(service_client):
-    a = service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
-    b = service_client.create_lora_training_client(base_model=BASE_MODEL, lora_rank=8)
+    a = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
+    b = service_client.create_lora_training_client(base_model=BASE_MODEL, rank=8)
     tok = a.get_tokenizer()
     data = [_make_datum(tok, "Q?", " a")]