fix: correctly handle method parameter counting in function_extra_arg

[benieric]

Issue #, if available:

• Custom Inference Code - model_fn() takes more positional argument #126
• https://discuss.huggingface.co/t/typeerror-model-fn-takes-1-positional-argument-but-2-were-given/76391/5

Description of changes:

• This change should address error like model_fn() takes x positional argument but y were given
• This error occurs under a race condition where validate_and_initialize_user_module() is called by 2 workers and extra arg is calculated incorrectly

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[arjkesh]

Can you add a test to check for regression against the race condition you described? Besides this, LGTM

[davidthomas426]

Can you add a test to check for regression against the race condition you described? Besides this, LGTM

+1.

[benieric]

Thanks for taking a look, I'll work on getting a test for this in

[davidthomas426]

Can you clarify why there are two threads calling initialize twice in a single python process?

[benieric]

Yup, it is in the handle() method -

sagemaker-huggingface-inference-toolkit/src/sagemaker_huggingface_inference_toolkit/handler_service.py

Line 234 in 92b57dd

def handle(self, data, context):

specifically this block:

 if not self.initialized:
                if self.attempted_init:
                    logger.warn(
                        "Model is not initialized, will try to load model again.\n"
                        "Please consider increase wait time for model loading.\n"
                    )
                self.initialize(context)

The test just assumes the fail condition already occurred

[benieric]

From some personal testing I was able to see that that model gets loaded and later on fails when attempting to load again:

2025-03-25T18:41:48,682 [INFO ] W-9000-model-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - Model model loaded io_fd=7a3c6bfffe7cf36a-0000007c-00000000-9116165d8c9c5504-f2c269a2

...

2025-03-25T18:42:34,468 [INFO ] W-model-1-stdout com.amazonaws.ml.mms.wlm.WorkerLifeCycle - mms.service.PredictionException: model_fn() takes 1 positional argument but 2 were given : 400

[benieric]

After testing with fix, did not see such error

[davidthomas426]

LGTM.

Though I'm still slightly confused about how a race condition will happen in this code. But I see that if the initialization function runs twice, it causes this problem, and I can see how this fixes that.