Repository: oumi-ai/oumi
File: src/oumi/inference/remote_inference_engine.py (Python)

Code before:
            generation_params = self._generation_params
            output_path = None
        else:
            remote_params = inference_config.remote_params or self._remote_params
            generation_params = inference_config.generation or self._generation_params
            output_path = inference_config.output_path

        self._set_required_fields_for_inference(remote_params)
        if not remote_params.api_url:
            raise ValueError("API URL is required for remote inference.")
        async with semaphore:
            api_input = self._convert_conversation_to_api_input(
                conversation, generation_params
            )
            headers = self._get_request_headers(remote_params)
            retries = 0
            failure_reason = None
            # Retry the request if it fails.
            for _ in range(remote_params.max_retries + 1):
                async with session.post(
                    remote_params.api_url,
                    json=api_input,
                    headers=headers,
                    timeout=remote_params.connection_timeout,
                ) as response:
                    print("response", response)
                    response_json = await response.json()
                    if response.status == 200:
                        result = self._convert_api_output_to_conversation(
                            response_json, conversation
                        )
                        if output_path:
                            # Write what we have so far to our scratch directory.
                            self._save_conversation(
                                result,
                                self._get_scratch_filepath(output_path),
                            )
                        await asyncio.sleep(remote_params.politeness_policy)
                        return result
                    else:
                        if isinstance(response_json, list):
                            # If the response is a list, it is likely an error message.
                            response_json = response_json[0]
                        failure_reason = (
                            response_json.get("error").get("message")
                            if response_json and response_json.get("error")
                            else None
                        )
                        retries += 1
                        await asyncio.sleep(remote_params.politeness_policy)
            raise RuntimeError(

Review comment (bug):
In general: try to use the logger (from oumi.utils.logging import logger), if you want to print something meaningful. That said, specifically for this case, maybe printing every output is not something useful for the user. It is useful for debugging purposes (not for the final PR)

Code after:
            generation_params = self._generation_params
            output_path = None
        else:
            remote_params = inference_config.remote_params or self._remote_params
            generation_params = inference_config.generation or self._generation_params
            output_path = inference_config.output_path

        self._set_required_fields_for_inference(remote_params)
        if not remote_params.api_url:
            raise ValueError("API URL is required for remote inference.")
        async with semaphore:
            api_input = self._convert_conversation_to_api_input(
                conversation, generation_params
            )
            headers = self._get_request_headers(remote_params)
            retries = 0
            failure_reason = None
            # Retry the request if it fails.
            for _ in range(remote_params.max_retries + 1):
                async with session.post(
                    remote_params.api_url,
                    json=api_input,
                    headers=headers,
                    timeout=remote_params.connection_timeout,
                ) as response:
                    response_json = await response.json()
                    if response.status == 200:
                        result = self._convert_api_output_to_conversation(
                            response_json, conversation
                        )
                        if output_path:
                            # Write what we have so far to our scratch directory.
                            self._save_conversation(
                                result,
                                self._get_scratch_filepath(output_path),
                            )
                        await asyncio.sleep(remote_params.politeness_policy)
                        return result
                    else:
                        if isinstance(response_json, list):
                            # If the response is a list, it is likely an error message.
                            response_json = response_json[0]
                        failure_reason = (
                            response_json.get("error").get("message")
                            if response_json and response_json.get("error")
                            else None
                        )
                        retries += 1
                        await asyncio.sleep(remote_params.politeness_policy)
            raise RuntimeError(
                f"Failed to query API after {remote_params.max_retries} retries. "