ronantakizawa/codereview-bench huggingface dataset:
Repository: vllm-project/aibrix
File: pkg/plugins/gateway/algorithms/pd_disaggregation.go (Go)

Code before:
*/

package routingalgorithms

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"strconv"
	"time"

	"github.com/vllm-project/aibrix/pkg/cache"
	"github.com/vllm-project/aibrix/pkg/types"
	"github.com/vllm-project/aibrix/pkg/utils"
	"github.com/vllm-project/aibrix/pkg/utils/prefixcacheindexer"
	"github.com/vllm-project/aibrix/pkg/utils/tokenizer"
	v1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

const (
	RouterPD                    types.RoutingAlgorithm = "pd"
	PDRoleIdentifier            string                 = "model.aibrix.ai/role"
	defaultDecodeOnlyTokenLimit int                    = 256
)

var (
	decodeOnlyTokenLimit int = utils.LoadEnvInt("AIBRIX_DECODE_ONLY_TOKEN_LENGTH", defaultDecodeOnlyTokenLimit)
)

func init() {
	Register(RouterPD, NewPDRouter)
}

type pdRouter struct {
	cache              cache.Cache
	tokenizer          tokenizer.Tokenizer
	prefixCacheIndexer *prefixcacheindexer.PrefixHashTable
}

func NewPDRouter() (types.Router, error) {
	var tokenizerObj tokenizer.Tokenizer
	if tokenizerType == "tiktoken" {
		tokenizerObj = tokenizer.NewTiktokenTokenizer()
	} else {
		tokenizerObj = tokenizer.NewCharacterTokenizer()
	}


Review comment (performance):
There's no need to introduce new roles. cache should be aware of stormservice pod labels/envs. 

```
Labels:           role-name=decode
                  role-template-hash=7dcf585456
                  roleset-name=llm-xpyd-roleset-dc8bz
                  storm-service-name=llm-xpyd
                  stormservice.orchestration.aibrix.ai/pod-group-index=3
                  stormservice.orchestration.aibrix.ai/role-replica-index=0
```


Code after:
*/

package routingalgorithms

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"strconv"
	"time"

	"github.com/vllm-project/aibrix/pkg/cache"
	"github.com/vllm-project/aibrix/pkg/types"
	"github.com/vllm-project/aibrix/pkg/utils"
	"github.com/vllm-project/aibrix/pkg/utils/prefixcacheindexer"
	"github.com/vllm-project/aibrix/pkg/utils/tokenizer"
	v1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
)

const (
	RouterPD                    types.RoutingAlgorithm = "pd"
	LLMEngineLabel              string                 = "model.aibrix.ai/engine"
	PDRoleIdentifier            string                 = "role-name"
	RoleReplicaIndex            string                 = "stormservice.orchestration.aibrix.ai/role-replica-index"
	PodGroupIndex               string                 = "stormservice.orchestration.aibrix.ai/pod-group-index"
	defaultDecodeOnlyTokenLimit int                    = 256
)

func init() {
	Register(RouterPD, NewPDRouter)
}

type pdRouter struct {
	cache              cache.Cache
	tokenizer          tokenizer.Tokenizer
	prefixCacheIndexer *prefixcacheindexer.PrefixHashTable
}

func NewPDRouter() (types.Router, error) {
	var tokenizerObj tokenizer.Tokenizer
	if tokenizerType == "tiktoken" {
		tokenizerObj = tokenizer.NewTiktokenTokenizer()
	} else {
		tokenizerObj = tokenizer.NewCharacterTokenizer()
	}

	c, err := cache.Get()
Repository: vllm-project/aibrix
File: pkg/plugins/gateway/algorithms/least_load_test.go (Go)

Code before:
			},
			expectAddr: "2.2.2.2:8000",
			expectErr:  false,
		},
		{
			name: "utilization error",
			pods: []*v1.Pod{
				newPod("pod1", "1.1.1.1", true, map[string]string{"model.aibrix.ai/port": "8000"}),
			},
			utilErr:   cache.ErrorLoadCapacityReached,
			expectErr: true,
		},
		{
			name: "some pods with utilization error",
			pods: []*v1.Pod{
				newPod("pod1", "1.1.1.1", true, map[string]string{"model.aibrix.ai/port": "8000"}),
				newPod("pod2", "2.2.2.2", true, map[string]string{"model.aibrix.ai/port": "8000"}),
			},
			utilizations: map[string]float64{
				"pod2": 0.5,
			},
			expectAddr: "1.1.1.1:8000",
			expectErr:  false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := &mockLoadProvider{
				utilizations: tt.utilizations,
				utilErr:      tt.utilErr,
			}

			router, err := NewLeastLoadRouter(provider)
			assert.NoError(t, err)

			ctx := types.NewRoutingContext(context.Background(), "test", "model1", "", "", "")
			podList := newMockPodList(tt.pods, tt.indexes)

			addr, err := router.Route(ctx, podList)
			if tt.expectErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
				assert.Equal(t, tt.expectAddr, addr) // Use exact match instead of Contains
			}
		})
	}
}

func TestLeastLoadRouter_PullMode(t *testing.T) {

Review comment (suggestion):
The "some pods with utilization error" case seems a bit strange. Why is there no metric for p1 but expectErr is nil and the expected pod is p1? 🤔 

Code after:
			name:      "no pods available",
			pods:      []*v1.Pod{},
			expectErr: true,
		},
		{
			name: "pod without IP",
			pods: []*v1.Pod{
				newPod("pod1", "", true, map[string]string{"model.aibrix.ai/port": "8000"}),
				newPod("pod2", "2.2.2.2", true, map[string]string{"model.aibrix.ai/port": "8000"}),
			},
			utilizations: map[string]float64{
				"pod1": 0.3,
				"pod2": 0.5,
			},
			expectAddr: "2.2.2.2:8000",
			expectErr:  false,
		},
		{
			name: "utilization error",
			pods: []*v1.Pod{
				newPod("pod1", "1.1.1.1", true, map[string]string{"model.aibrix.ai/port": "8000"}),
			},
			globalUtilizationError: cache.ErrorLoadCapacityReached,
			expectErr:              true,
		},
		{
			name: "some pods with utilization error",
			pods: []*v1.Pod{
				newPod("pod1", "1.1.1.1", true, map[string]string{"model.aibrix.ai/port": "8000"}),
				newPod("pod2", "2.2.2.2", true, map[string]string{"model.aibrix.ai/port": "8000"}),
				newPod("pod3", "3.3.3.3", true, map[string]string{"model.aibrix.ai/port": "8000"}),
			},
			utilizations: map[string]float64{
				"pod2": 0.7,
				"pod3": 0.3,
			},
			perPodUtilizationError: map[string]error{
				"pod1": cache.ErrorLoadCapacityReached,
			},
			expectAddr: "3.3.3.3:8000", // Should select pod3 (lowest util among healthy pods)
			expectErr:  false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			provider := &mockLoadProvider{
				utilizations:           tt.utilizations,
				perPodUtilizationError: tt.perPodUtilizationError,
				globalUtilizationError: tt.globalUtilizationError,
			}
Repository: vllm-project/aibrix
File: pkg/controller/modelrouter/modelrouter_controller.go (Go)

Code before:
	})
	if err != nil {
		return err
	}

	_, err = fleetInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    modelRouter.addRouteFromRayClusterFleet,
		DeleteFunc: modelRouter.deleteRouteFromRayClusterFleet,
	})
	if err != nil {
		return err
	}

	// add dynamic informer for all workloads
	for _, gvk := range watchedWorkloads {
		exists, err := utils.GVKCheckExists(mgr.GetConfig(), gvk)
		if err != nil {
			klog.ErrorS(err, "Failed to check CRD", "GVK", gvk)
			continue // continue to skip error in suit_test when crd is not found
		}
		if !exists {
			klog.Infof("%s not found, skipping add informer for it. "+
				"This is optional - install %s CRD if you need ", gvk, gvk)
		} else {
			err = addInformerForGVK(mgr, modelRouter, gvk)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

type ModelRouter struct {
	client.Client
	Scheme        *runtime.Scheme
	RuntimeConfig config.RuntimeConfig
}

func (m *ModelRouter) addRouteFromDeployment(obj interface{}) {
	deployment := obj.(*appsv1.Deployment)
	m.createHTTPRoute(deployment.Namespace, deployment.Labels, deployment.Annotations)
}

func (m *ModelRouter) deleteRouteFromDeployment(obj interface{}) {
	deployment, ok := obj.(*appsv1.Deployment)
	if !ok {
		tombstone, ok := obj.(cache.DeletedFinalStateUnknown)
		if !ok {
			return

Review comment (suggestion):
how about: 
```go
          exists, err := utils.GVKCheckExists(mgr.GetConfig(), gvk)
	if err != nil || !exists {
		klog.InfoS("skip informer (optional CRD)", "GVK", gvk, "exists", exists, "error", err)
		continue
	}
	if err := addInformerForGVK(mgr, modelRouter, gvk); err != nil {
		return err
	}
```

Code after:
		AddFunc:    modelRouter.addRouteFromModelAdapter,
		DeleteFunc: modelRouter.deleteRouteFromModelAdapter,
	})
	if err != nil {
		return err
	}

	_, err = fleetInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    modelRouter.addRouteFromRayClusterFleet,
		DeleteFunc: modelRouter.deleteRouteFromRayClusterFleet,
	})
	if err != nil {
		return err
	}

	// add dynamic informer for all workloads
	for _, gvk := range watchedWorkloads {
		exists, err := utils.GVKCheckExists(mgr.GetConfig(), gvk)
		if err != nil || !exists {
			klog.InfoS("skip informer (optional CRD)", "GVK", gvk, "exists", exists, "error", err)
			continue
		}
		if err := addInformerForGVK(mgr, modelRouter, gvk); err != nil {
			return err
		}
	}

	return nil
}

type ModelRouter struct {
	client.Client
	Scheme        *runtime.Scheme
	RuntimeConfig config.RuntimeConfig
}

func (m *ModelRouter) addRouteFromDeployment(obj interface{}) {
	deployment := obj.(*appsv1.Deployment)
	m.createHTTPRoute(deployment.Namespace, deployment.Labels, deployment.Annotations)
}

func (m *ModelRouter) deleteRouteFromDeployment(obj interface{}) {
	deployment, ok := obj.(*appsv1.Deployment)
	if !ok {
		tombstone, ok := obj.(cache.DeletedFinalStateUnknown)
		if !ok {
			return
		}
		deployment, ok = tombstone.Obj.(*appsv1.Deployment)
		if !ok {
			return
Repository: vllm-project/aibrix
File: pkg/controller/modelrouter/modelrouter_controller.go (Go)

Code before:
			continue
		}
		httpRoute.Spec.Rules[0].Matches = append(httpRoute.Spec.Rules[0].Matches,
			gatewayv1.HTTPRouteMatch{
				Path: &gatewayv1.HTTPPathMatch{
					Type:  ptr.To(gatewayv1.PathMatchPathPrefix),
					Value: ptr.To(path),
				},
				Headers: []gatewayv1.HTTPHeaderMatch{
					modelHeaderMatch,
				},
			})
		klog.InfoS("Added custom model router path", "path", path)
	}
}

func addInformerForGVK(mgr manager.Manager, modelRouter *ModelRouter, gvk schema.GroupVersionKind) error {
	// create dynamic Informer
	uObj := &unstructured.Unstructured{}
	uObj.SetGroupVersionKind(gvk)

	uInformer, err := mgr.GetCache().GetInformer(context.TODO(), uObj)
	if err != nil {
		klog.ErrorS(err, "Failed to get informer for LWS")
		return err
	}

	// add Event Handler
	_, err = uInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    modelRouter.addRouteFromUnstructuredObj,
		DeleteFunc: modelRouter.deleteRouteFromUnstructuredObj,
	})
	if err != nil {
		return err
	}
	klog.Infof("Added model router informer for %s", gvk)
	return nil
}


Review comment (bug):
Should be generic since the function handles any GVK:
`klog.ErrorS(err, "Failed to get informer", "GVK", gvk)`

Code after:
		path = strings.ReplaceAll(path, " ", "")
		if _, exists := pathSet[path]; path == "" || exists {
			continue
		}
		httpRoute.Spec.Rules[0].Matches = append(httpRoute.Spec.Rules[0].Matches,
			gatewayv1.HTTPRouteMatch{
				Path: &gatewayv1.HTTPPathMatch{
					Type:  ptr.To(gatewayv1.PathMatchPathPrefix),
					Value: ptr.To(path),
				},
				Headers: []gatewayv1.HTTPHeaderMatch{
					modelHeaderMatch,
				},
			})
		klog.InfoS("Added custom model router path", "path", path)
	}
}

func addInformerForGVK(mgr manager.Manager, modelRouter *ModelRouter, gvk schema.GroupVersionKind) error {
	// create dynamic Informer
	uObj := &unstructured.Unstructured{}
	uObj.SetGroupVersionKind(gvk)

	uInformer, err := mgr.GetCache().GetInformer(context.TODO(), uObj)
	if err != nil {
		klog.ErrorS(err, "Failed to get informer for LWS")
		return err
	}

	// add Event Handler
	_, err = uInformer.AddEventHandler(cache.ResourceEventHandlerFuncs{
		AddFunc:    modelRouter.addRouteFromUnstructuredObj,
		DeleteFunc: modelRouter.deleteRouteFromUnstructuredObj,
	})
	if err != nil {
		return err
	}
	klog.Infof("Added model router informer for %s", gvk)
	return nil
}
Repository: vllm-project/aibrix
File: pkg/utils/util.go (Go)

Code before:
	if err != nil {
		klog.Warningf("Invalid port value in label %s=%s: %v",
			constants.ModelLabelPort, portStr, err)
		return 0, fmt.Errorf("invalid port format: %w", err)
	}

	if port <= 0 || port > 65535 {
		klog.Warningf("Port %d from label %s is out of valid range",
			port, constants.ModelLabelPort)
		return 0, fmt.Errorf("port %d out of range [1-65535]", port)
	}

	return port, nil
}

// getDataParallelSize extracts data-parallel size from pod environment variables
func getDataParallelSize(pod *v1.Pod) int {
	for _, container := range pod.Spec.Containers {
		for _, env := range container.Env {
			if env.Name == "data-parallel-size" && env.Value != "" {
				size, err := strconv.Atoi(env.Value)
				if err != nil {
					klog.Warningf("Invalid data-parallel-size value '%s': %v",
						env.Value, err)
					return 0
				}

				if size < 0 {
					klog.Warningf("Invalid data-parallel-size %d (must be >= 0)", size)
					return 0
				}

				return size
			}
		}
	}

	return 0
}

// GVKCheckExists check if gvk is support
func GVKCheckExists(cfg *rest.Config, gvk schema.GroupVersionKind) (bool, error) {
	dc, err := discovery.NewDiscoveryClientForConfig(cfg)
	if err != nil {
		return false, err
	}

	// get all resource from GroupVersion
	apiResources, err := dc.ServerResourcesForGroupVersion(gvk.GroupVersion().String())
	if err != nil {
		if discovery.IsGroupDiscoveryFailedError(err) {

Review comment (suggestion):
this is injected and control by users? If so, I suggest we add some comments here, to enable this feature, we require user to pass it in the pod template env etc. Later ,we can move it to the doc to guide user following this setting

Code after:
		klog.Warningf("Invalid port value in label %s=%s: %v",
			constants.ModelLabelPort, portStr, err)
		return 0, fmt.Errorf("invalid port format: %w", err)
	}

	if port <= 0 || port > 65535 {
		klog.Warningf("Port %d from label %s is out of valid range",
			port, constants.ModelLabelPort)
		return 0, fmt.Errorf("port %d out of range [1-65535]", port)
	}

	return port, nil
}

// getDataParallelSize retrieves the data parallelism degree from the pod's environment variables.
// This feature requires the user to explicitly pass the "data-parallel-size"
// variable in the pod template's container env section.
// Example:
//
//		  env:
//		  - name: "data-parallel-size"
//	     value: "2"
func getDataParallelSize(pod *v1.Pod) int {
	for _, container := range pod.Spec.Containers {
		for _, env := range container.Env {
			if env.Name == "data-parallel-size" && env.Value != "" {
				size, err := strconv.Atoi(env.Value)
				if err != nil {
					klog.Warningf("Invalid data-parallel-size value '%s': %v",
						env.Value, err)
					return 0
				}

				if size < 0 {
					klog.Warningf("Invalid data-parallel-size %d (must be >= 0)", size)
					return 0
				}

				return size
			}
		}
	}

	return 0
}

// GVKCheckExists check if gvk is support
func GVKCheckExists(cfg *rest.Config, gvk schema.GroupVersionKind) (bool, error) {
	dc, err := discovery.NewDiscoveryClientForConfig(cfg)
	if err != nil {
		return false, err
Repository: vllm-project/vllm
File: tests/models/decoder_only/language/test_mistral.py (Python)

Code before:
    ) as mistral_format_model:
        mistral_format_outputs = mistral_format_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

    check_logprobs_close(
        outputs_0_lst=hf_format_outputs,
        outputs_1_lst=mistral_format_outputs,
        name_0="hf",
        name_1="mistral",
    )


@pytest.mark.parametrize("dtype", ["bfloat16"])
@pytest.mark.parametrize("model", MODELS[1:])  # v1 can't do func calling
def test_mistral_function_calling(
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    with vllm_runner(model,
                     dtype=dtype,
                     tokenizer_mode="mistral",
                     config_format="mistral",
                     load_format="mistral") as vllm_model:
        outputs = vllm_model.model.chat(MSGS,
                                        tools=TOOLS,
                                        sampling_params=SAMPLING_PARAMS)

        assert outputs[0].outputs[0].text.strip() == EXPECTED_FUNC_CALL


Review comment (suggestion):
> @DarkLight1337 seems like mistral-nemo is too large for the decoder-only tests:

We can set `max_model_len` here to reduce the memory usage.

Code after:
            load_format="mistral",
            config_format="mistral",
    ) as mistral_format_model:
        mistral_format_outputs = mistral_format_model.generate_greedy_logprobs(
            example_prompts, max_tokens, num_logprobs)

    check_logprobs_close(
        outputs_0_lst=hf_format_outputs,
        outputs_1_lst=mistral_format_outputs,
        name_0="hf",
        name_1="mistral",
    )


@pytest.mark.parametrize("dtype", ["bfloat16"])
@pytest.mark.parametrize("model", MODELS[1:])  # v1 can't do func calling
def test_mistral_function_calling(
    vllm_runner,
    model: str,
    dtype: str,
) -> None:
    with vllm_runner(model,
                     dtype=dtype,
                     tokenizer_mode="mistral",
                     config_format="mistral",
                     load_format="mistral") as vllm_model:
        outputs = vllm_model.model.chat(MSGS,
                                        tools=TOOLS,
                                        sampling_params=SAMPLING_PARAMS)

        assert outputs[0].outputs[0].text.strip() == EXPECTED_FUNC_CALL
Repository: vllm-project/vllm
File: tools/pre_commit/check_torch_cuda.py (Python)

Code before:
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys

import regex as re

# --------------------------------------------------------------------------- #
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
# --------------------------------------------------------------------------- #
_TORCH_CUDA_RE = re.compile(r"\btorch\.cuda\.empty_cache\b")


ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}


def scan_file(path: str) -> int:
    with open(path, encoding="utf-8") as f:
        for i, line in enumerate(f, 1):
            if _TORCH_CUDA_RE.search(line):
                print(
                    f"{path}:{i}: "
                    "\033[91merror:\033[0m "  # red color
                    "Found torch.cuda API call"
                )
                return 1
    return 0


def main():
    returncode = 0
    for filename in sys.argv[1:]:
        if any(filename.startswith(prefix) for prefix in ALLOWED_FILES):
            continue
        returncode |= scan_file(filename)
    return returncode


if __name__ == "__main__":
    sys.exit(main())


Review comment (suggestion):
```suggestion
    with open(path, encoding="utf-8") as f:
        content = f.read()
    for match in re.finditer(_TORCH_CUDA_RE, content, re.MULTILINE):
        # Calculate line number from match position
        line_num = content[: match.start() + 1].count("\n") + 1
        print(
            f"{path}:{line_num}: "
            "\033[91merror:\033[0m "  # red color
            "Found torch.cuda API call"
        )
        return 1
```

Code after:
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys

import regex as re

# --------------------------------------------------------------------------- #
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
# --------------------------------------------------------------------------- #
_TORCH_CUDA_PATTERNS = [
    r"\btorch\.cuda\.empty_cache\b",
]

ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}


def scan_file(path: str) -> int:
    with open(path, encoding="utf-8") as f:
        content = f.read()
    for pattern in _TORCH_CUDA_PATTERNS:
        for match in re.finditer(pattern, content, re.MULTILINE):
            # Calculate line number from match position
            line_num = content[: match.start() + 1].count("\n") + 1
            print(
                f"{path}:{line_num}: "
                "\033[91merror:\033[0m "  # red color
                "Found torch.cuda API call"
            )
            return 1
    return 0


def main():
    returncode = 0
    for filename in sys.argv[1:]:
        if any(filename.startswith(prefix) for prefix in ALLOWED_FILES):
            continue
        returncode |= scan_file(filename)
    return returncode


if __name__ == "__main__":
    sys.exit(main())
Repository: vllm-project/vllm
File: tools/pre_commit/check_torch_cuda.py (Python)

Code before:
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys

import regex as re

# --------------------------------------------------------------------------- #
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
# --------------------------------------------------------------------------- #
_TORCH_CUDA_RE = re.compile(r"\btorch\.cuda\.empty_cache\b")


ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}


def scan_file(path: str) -> int:
    with open(path, encoding="utf-8") as f:
        for i, line in enumerate(f, 1):
            if _TORCH_CUDA_RE.search(line):
                print(
                    f"{path}:{i}: "
                    "\033[91merror:\033[0m "  # red color
                    "Found torch.cuda API call"
                )
                return 1
    return 0


def main():
    returncode = 0
    for filename in sys.argv[1:]:
        if any(filename.startswith(prefix) for prefix in ALLOWED_FILES):
            continue
        returncode |= scan_file(filename)
    return returncode

Review comment (performance):
```suggestion
_TORCH_CUDA_RE = r"\btorch\.cuda\.empty_cache\b"
```

Code after:
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys

import regex as re

# --------------------------------------------------------------------------- #
# Regex: match `torch.cuda.xxx` but allow `torch.accelerator.xxx`
# --------------------------------------------------------------------------- #
_TORCH_CUDA_PATTERNS = [
    r"\btorch\.cuda\.empty_cache\b",
]

ALLOWED_FILES = {"vllm/platforms/", "vllm/device_allocator/"}


def scan_file(path: str) -> int:
    with open(path, encoding="utf-8") as f:
        content = f.read()
    for pattern in _TORCH_CUDA_PATTERNS:
        for match in re.finditer(pattern, content, re.MULTILINE):
            # Calculate line number from match position
            line_num = content[: match.start() + 1].count("\n") + 1
            print(
                f"{path}:{line_num}: "
                "\033[91merror:\033[0m "  # red color
                "Found torch.cuda API call"
            )
            return 1
    return 0


def main():
    returncode = 0
    for filename in sys.argv[1:]: