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

Code before:
            model_params: The model parameters to use for inference.
            generation_params: The generation parameters to use for inference.
            remote_params: The remote parameters to use for inference.
            project_id_env_key: The environment variable key name for the project ID.
            region_env_key: The environment variable key name for the region.
            project_id: The project ID to use for inference.
            region: The region to use for inference.
        """
        super().__init__(
            model_params=model_params,
            generation_params=generation_params,
            remote_params=remote_params,
        )
        self._project_id_env_key = (
            project_id_env_key or self._DEFAULT_PROJECT_ID_ENV_KEY
        )
        self._region_env_key = region_env_key or self._DEFAULT_REGION_ENV_KEY
        self._project_id = project_id
        self._region = region

    @override
    def _set_required_fields_for_inference(self, remote_params: RemoteParams) -> None:
        """Set required fields for inference."""
        if (
            not remote_params.api_url
            and not self._remote_params.api_url
            and not self.base_url
        ):
            if self._project_id and self._region:
                project_id = self._project_id
                region = self._region
            elif os.getenv(self._project_id_env_key) and os.getenv(
                self._region_env_key
            ):
                project_id = os.getenv(self._project_id_env_key)
                region = os.getenv(self._region_env_key)
            else:
                raise ValueError(
                    "The GCP inference engine requires that either the `api_url` is "
                    "set or that both the `project_id` and `region` are set. "
                    "You can set the `api_url` in `RemoteParams`. "
                    "Alternatively, you can set the `project_id` and `region` when "
                    "constructing the inference engine (GoogleVertexInferenceEngine), "
                    f"or as environment variables: `{self._project_id_env_key}` and "
                    f"`{self._region_env_key}`, respectively."
                )

            remote_params.api_url = self._API_URL_TEMPLATE.format(
                project_id=project_id,
                region=region,
            )

Review comment (suggestion):
We should probably return an error if the user supplies both `project_id_env_key` and `project_id`, or both `region_env_key` and `region`

Code after:
            model_params: The model parameters to use for inference.
            generation_params: The generation parameters to use for inference.
            remote_params: The remote parameters to use for inference.
            project_id_env_key: The environment variable key name for the project ID.
            region_env_key: The environment variable key name for the region.
            project_id: The project ID to use for inference.
            region: The region to use for inference.
        """
        super().__init__(
            model_params=model_params,
            generation_params=generation_params,
            remote_params=remote_params,
        )
        if project_id and project_id_env_key:
            raise ValueError(
                "You cannot set both `project_id` and `project_id_env_key`."
            )
        if region and region_env_key:
            raise ValueError("You cannot set both `region` and `region_env_key`.")

        self._project_id_env_key = (
            project_id_env_key or self._DEFAULT_PROJECT_ID_ENV_KEY
        )
        self._region_env_key = region_env_key or self._DEFAULT_REGION_ENV_KEY
        self._project_id = project_id
        self._region = region

    @override
    def _set_required_fields_for_inference(self, remote_params: RemoteParams) -> None:
        """Set required fields for inference."""
        if (
            not remote_params.api_url
            and not self._remote_params.api_url
            and not self.base_url
        ):
            if self._project_id and self._region:
                project_id = self._project_id
                region = self._region
            elif os.getenv(self._project_id_env_key) and os.getenv(
                self._region_env_key
            ):
                project_id = os.getenv(self._project_id_env_key)
                region = os.getenv(self._region_env_key)
            else:
                raise ValueError(
                    "This inference engine requires that either `api_url` is set in "
                    "`RemoteParams` or that both `project_id` and `region` are set. "
                    "You can set the `project_id` and `region` when "
                    "constructing a GoogleVertexInferenceEngine, "
                    f"or as environment variables: `{self._project_id_env_key}` and "
                    f"`{self._region_env_key}`."