feat(launcher): add vLLM container support for data synthesis

- core.py: clear Docker ENTRYPOINT so containers with vllm serve as
  entrypoint (e.g. vllm/vllm-openai:qwen3_5-cu130) run correctly
- vllm/query.sh: use python3, install datasets+openai before running
  query.py
- query.py: add --max-tokens arg to cap response length; reinitialize
  OpenAI client after fork to avoid connection-pool corruption in
  datasets.map(num_proc>1)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: chenhany <chenhany@nvidia.com>

Author: ChenhanYu

tools/launcher/common/query.py

@@ -44,16 +44,30 @@ def _strip_thinking(content: str) -> str:
 class LLM:
     def __init__(self, args):
         self.args = args
+        self._pid = os.getpid()
         self.client = OpenAI(base_url=args.base_url)
         self.generate(messages=[{"role": "user", "content": "Hello! /no_think"}], verbose=True)
 
+    def _ensure_client(self):
+        """Reinitialize the HTTP client if we've been forked into a new process.
+
+        datasets.map(num_proc>1) forks worker processes that inherit the parent's
+        connection pool.  Reusing inherited sockets across processes causes
+        "Invalid HTTP request" errors.  Creating a fresh client per-process avoids this.
+        """
+        if os.getpid() != self._pid:
+            self._pid = os.getpid()
+            self.client = OpenAI(base_url=self.args.base_url)
+
     def generate(self, messages, verbose=False, **chat_template_kwargs):
         global early_termination
+        self._ensure_client()
         try:
             completion = self.client.chat.completions.create(
                 model=self.args.model,
                 messages=messages,
                 temperature=self.args.temperature,
+                max_tokens=self.args.max_tokens,
             )
             new_message = completion.choices[0].message.content
             if verbose:
@@ -88,6 +102,7 @@ def generate(self, messages, verbose=False, **chat_template_kwargs):
 )
 parser.add_argument("--num-proc", type=int, default=32, help="number of processes (concurrency).")
 parser.add_argument("--temperature", type=float, default=0.0, help="temperature.")
+parser.add_argument("--max-tokens", type=int, default=None, help="maximum tokens to generate per response.")
 args = parser.parse_args()
 
 llm = LLM(args)
@@ -162,7 +177,11 @@ def synthesize(data):
     return {"conversations": current_messages}
 
 
-dataset = load_dataset(args.data, split=args.data_split)
+# Support both HF Hub repo IDs and local file paths (.jsonl, .json, .parquet, etc.)
+if os.path.isfile(args.data):
+    dataset = load_dataset("json", data_files=args.data, split=args.data_split)
+else:
+    dataset = load_dataset(args.data, split=args.data_split)
 
 if args.num_shards * 100 > len(dataset):
     args.num_shards = min(16, len(dataset) // 100)

tools/launcher/common/vllm/query.sh

@@ -118,7 +118,8 @@ while true; do
     sleep 10
 done
 
-cmd="python common/query.py http://localhost:8000/v1 ${MODEL} ${QUERY_ARGS[*]}"
+pip3 install -q datasets openai 2>/dev/null || true
+cmd="python3 common/query.py http://localhost:8000/v1 ${MODEL} ${QUERY_ARGS[*]}"
 echo "Running command: $cmd"
 eval $cmd
 echo "Main process exit"

tools/launcher/core.py

@@ -322,7 +322,7 @@ def build_docker_executor(
         ipc_mode="host",
         container_image=slurm_config.container,
         volumes=container_mounts,
-        additional_kwargs={"user": f"{os.getuid()}:{os.getgid()}"},
+        additional_kwargs={"user": f"{os.getuid()}:{os.getgid()}", "entrypoint": ""},
         packager=packager,
     )
     return executor