Method __call__

sgs/models/query_local.py:373–531

(
        self,
        model_config: ModelConfig,
        buffer_size: int = 128,
        gpu_id: Optional[int] = None,
    )

        sys.exit(1)

    def __call__(
        self,
        model_config: ModelConfig,
        buffer_size: int = 128,
        gpu_id: Optional[int] = None,
    ) -> None:
        self.exit_reason = "Unkown"

        if gpu_id is not None:
            # Set CUDA_VISIBLE_DEVICES to only show the specific GPU for this subprocess
            os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
            # Set the GPU for this process
            torch.cuda.set_device(
                0
            )  # Use GPU 0 since CUDA_VISIBLE_DEVICES makes it the only visible GPU

        try:
            self.monitor = ResourceMonitor(
                interval_sec=10,
                track_children=False,
                track_cgroup=False,
                profile_gpu=False,
                timings_only=not self.do_in_depth_monitoring,
            )
            self.monitor.start()

            vllm_model = LLM(
                model=model_config.model_name,
                dtype=model_config.dtype,
                max_model_len=model_config.max_tokens,
            )

            sampling_params = SamplingParams(
                temperature=model_config.temperature,
                max_tokens=model_config.max_tokens,
                top_p=model_config.top_p,
                top_k=model_config.top_k,
                # skip_special_tokens=False,
                # include_stop_str_in_output=True,
                logprobs=10,
            )

            # We keep going while there are still things in the buffer or we are still processing
            status_check_interval = 2
            last_status_check = time.time()
            report_util_interval = 20
            last_report_util = time.time()

            print("=" * 100)
            print("BEGINNING GENERATION LOOP WITH THE SERVER")
            print("=" * 100)

            # We keep going while there are still things in the buffer or we are still processing
            while True:
                # First we do some reporting to the server
                time.sleep(random.uniform(0.5, 1.5))
                now = time.time()
                if now - last_status_check > status_check_interval: