Add sampling support for Tinker SkyRL backend

[tyler-griggs]

Summary

Implements sampling/inference support for the SkyRL Tinker backend so we can start enabling RL training workflows via the Tinker API.

Key Changes

• Implements sample() method in SkyRLTrainBackend using existing InferenceEngineClient
• Adds weight synchronization via save_weights_for_sampler(), which just passes it along to the WorkerDispatch

Limitations

• Prompt logprobs not yet supported (returns None with warning)
• Checkpoint paths in requests ignored (always uses current weights)
• Single model/checkpoint at a time

[tyler-griggs]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds sampling and checkpointing support for the SkyRL Tinker backend, a crucial step for enabling RL training workflows. However, it introduces Path Traversal vulnerabilities in the new save_checkpoint and load_checkpoint methods, as these rely on unsanitized user-supplied identifiers (model_id, checkpoint_id), potentially allowing an attacker to read or write files outside the intended checkpoint directory. Additionally, there is a critical performance issue in the sampling implementation and a couple of medium-severity suggestions for improving configuration and maintainability.

[tyler-griggs]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces significant and well-structured functionality for sampling and inference to the SkyRL backend, which is a crucial step for enabling RL workflows. The code is generally of high quality with good documentation and clear separation of concerns, particularly between in-memory weight synchronization and disk-based checkpointing. My review focuses on a critical compatibility/security issue related to tarfile, improving error handling during sampling, enhancing code readability, and adhering to Python's standard style guidelines.

I am having trouble creating individual review comments. Click here to see my feedback.

skyrl-tx/tx/tinker/backends/skyrl_train.py (462)

The filter argument for tarfile.extractall was introduced in Python 3.12 to mitigate path traversal vulnerabilities. If this project supports Python versions older than 3.12, this call will raise a TypeError. Please ensure the project's minimum Python version is 3.12, or provide a version-conditional fallback for older versions. For example:

import sys

if sys.version_info >= (3, 12):
    tar.extractall(temp_dir, filter="data")
else:
    # This is insecure, but maintains compatibility.
    # Consider adding your own checks to prevent path traversal.
    tar.extractall(temp_dir)

skyrl-tx/tx/tinker/backends/skyrl_train.py (340-342)

Converting exceptions to None here causes the loss of valuable error information. It would be better to pass the sample_outputs list (which may contain Exception objects) directly to _aggregate_sample_results and handle the exceptions there to generate more informative error messages.

For example, in _aggregate_sample_results, you could check for exceptions and format the error message accordingly:

if isinstance(output, Exception):
    has_error = True
    error_msg = f"Sampling failed for sample {i} with error: {output!r}"
    break

skyrl-tx/tx/tinker/backends/skyrl_train.py (370)

This line is a bit complex and can be hard to read. Simplifying it would improve maintainability.

                response_logprobs = (output.get("response_logprobs") or [[]])[0]

skyrl-tx/tx/tinker/engine.py (495)

According to PEP 8, imports should usually be at the top of the file. Please move import asyncio to the top of the file to improve code style and consistency.

[tyler-griggs]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds sampling support to the SkyRL Tinker backend, a crucial feature for enabling reinforcement learning workflows. The changes are well-structured, introducing a new sample method, helpers for creating inference engines, and an async save_weights_for_sampler method for in-memory weight updates. The implementation aligns well with the detailed design in the pull request description. I've identified one critical issue in the sampling logic that would cause a runtime error. Addressing this will ensure the new functionality works as intended.

[gemini-code-assist]

The code attempts to access sampling_params.stop, but this attribute does not exist on the types.SamplingParams model, which will cause a runtime AttributeError. The correct attributes for stop criteria are stop_strings and stop_tokens. This suggestion fixes the handling of stop sequences while preserving the correct logic for top_k and top_p.

                if sampling_params.top_k is not None:
                    params_dict["top_k"] = sampling_params.top_k
                if sampling_params.top_p is not None:
                    params_dict["top_p"] = sampling_params.top_p
                if sampling_params.stop_strings:
                    params_dict["stop"] = sampling_params.stop_strings
                if sampling_params.stop_tokens:
                    params_dict["stop_token_ids"] = sampling_params.stop_tokens

[gemini-code-assist]

Code Review

This pull request successfully adds sampling support for the SkyRL Tinker backend, enabling RL training workflows. However, two significant security issues were identified: a flawed concurrency model in the sample method that leads to deadlocks in asynchronous environments, and a lack of resource constraints when processing large batches of samples, which can be exploited for Denial of Service attacks. The implementation of the sample() method using InferenceEngineClient is otherwise well-done, including robust handling of asynchronous operations and batch processing, and save_weights_for_sampler is a key feature. Minor suggestions for code style and readability have also been provided.

[gemini-code-assist]

The sample method contains logic that will cause a permanent deadlock if called from a thread where an asyncio event loop is already running. The implementation uses asyncio.get_running_loop() and asyncio.run_coroutine_threadsafe(sample_all(), loop).result(), which blocks the event loop thread, leading to a permanent hang of the process. To remediate this, avoid blocking the event loop thread; consider making the sample method asynchronous or offloading work to a separate thread pool. Additionally, for better code organization and PEP 8 adherence, the concurrent.futures import should be moved to the top of the file.

[gemini-code-assist]

The sample method is vulnerable to resource exhaustion leading to a Denial of Service. It iterates through all prompts in prepared_batch.all_prompts (the size of which is controlled by the user-supplied num_samples parameter) and creates a task for each one, then executes all tasks in parallel using asyncio.gather. There is no upper bound on num_samples nor any chunking of the tasks. An attacker can provide an extremely large value for num_samples, causing the application to consume excessive memory and network resources, potentially leading to an Out-Of-Memory (OOM) crash or system instability.

Implement a maximum limit for num_samples and process the sampling tasks in smaller, fixed-size batches (e.g., using a semaphore or by chunking the input list).

[gemini-code-assist]

This line is a bit dense and could be simplified for better readability. Using the or operator can make handling the default case more concise and clear.

Suggested change

	response_logprobs = output.get("response_logprobs", [[]])[0] if output.get("response_logprobs") else []
	response_logprobs = (output.get("response_logprobs") or [[]])[0]

[gemini-code-assist]

To follow standard Python style guidelines (PEP 8), please move this import to the top of the file. This improves readability and makes dependencies clear.

[pcmoritz]

Should we remove this before merging? It would probably be better to expose the .yaml so it can be set in there, right? You could add an argument to SkyRLTrainBackendConfig which specifies the path to the .yaml and then load it from here -- and then it can be used like `--backend-config '{"config_path": "/path/to/yaml"}' or similar.

[tyler-griggs]

Yeah I'll remove! I'm hoping to rely on Sumanth's changes to use pythonic configs very soon..

[pcmoritz]

I realize trying to avoid importing skyrl_train will get uglier and uglier over time. We need to figure out a better way, we can make this a function

SkyRL/skyrl-tx/tx/tinker/engine.py

Line 133 in afa60b7

BACKENDS = {

that returns the available backends, and it will only import skyrl_train if the dependencies are installed, so everything is in one centralized place. E.g.

def get_backends() -> dict[str, tuple[...]]:
    backends = {}
    try:
        from tx.tinker.backends.skyrl_train import SkyRLTrainBackend, SkyRLTrainBackendConfig
        backends["skyrl_train"] = (SkyRLTrainBackend, SkyRLTrainBackendConfig)
    except ImportError:
        # print warning

    try:
        from tx.tinker.backends.jax import JaxBackend, JaxBackendConfig
        backends["jax"] = (JaxBackend, JaxBackendConfig)
    except:
        # print warning

    return backends

Maybe this can be made even nicer...

[pcmoritz]

Could this code be structured such that cfg.generator is mostly passed through? That would make it easier going forward e.g. when new fields are added. These fields are currently implementation details that this file should better not be concerned with (the more minimal we can keep this file, the better).

[pcmoritz]

It is worth checking if these are actually needed (e.g. https://github.com/vllm-project/vllm/blob/8869cd8ec1b2586df96e6556f2d2e3b3760bfb93/vllm/sampling_params.py#L159). In external_inference.py we are just passing it through and it has been working well afaik.

[pcmoritz]

Do these need to be handled differently? E.g. stop vs stop_token_ids https://github.com/vllm-project/vllm/blob/8869cd8ec1b2586df96e6556f2d2e3b3760bfb93/vllm/sampling_params.py#L171