Fix: long import times

[oKatanaaa]

Problem Description

This PR addresses two related issues with Unsloth's module importing mechanism:

1. Slow/hanging imports when called from entry points: When importing Unsloth via Python entry points (defined in pyproject.toml), import times can exceed 60 seconds or even hang indefinitely (Slow Import Times (60+ seconds) When Using Unsloth via Entry Points unsloth#1859) due to an improper retry mechanism. This occurs because the module import system fails differently depending on how Python is invoked.

2. Inconsistent cache locations: Previously, the UNSLOTH_COMPILE_LOCATION was set as a relative path ("unsloth_compiled_cache"), which would create different cache directories depending on the working directory where a script was launched.

Root Cause Analysis

The core issue was in the create_new_function method's module loading mechanism:

1. It tried to import modules using importlib.import_module(UNSLOTH_COMPILE_LOCATION + "." + name), but this fails when the directory is not in sys.path (common when running through entry points).

2. The fallback mechanism used a while True loop that could potentially run indefinitely, with each iteration adding another sleep delay (Slow Import Times (60+ seconds) When Using Unsloth via Entry Points unsloth#1859 (comment)).

Solution

1. Standardized cache location using appdirs.user_data_dir() for consistency across sessions.

2. Fixed module importing by temporarily adding the cache directory to sys.path.

3. Eliminated the infinite loop in the fallback mechanism with proper exception handling.

4. Added clearer error messages to simplify troubleshooting.

Testing

Tested the changes with both direct module invocation (python -m) and through entry points, confirming:

• Consistent import times (5-7 seconds) regardless of invocation method
• Consistent cache location across multiple runs and directories

[danielhanchen]

Oh wait I just did some changes which adds a /tmp folder and some other changes - is it possible to directly add it to sys.path instead of using a new dependency?

[oKatanaaa]

Oh wait I just did some changes which adds a /tmp folder and some other changes - is it possible to directly add it to sys.path instead of using a new dependency?

Good point! We could use tempfile.gettempdir (a standard python lib) instead of appdirs.user_data_dir. That would keep all of the benefits of the fix while keeping it cross-platform.

I've implemented the change, take a look.

Edit: tempfile.gettempdir returns the /tmp folder (or its Windows analog), so it's aligned with your idea.

[oKatanaaa]

Looking at the current implementation:

UNSLOTH_COMPILE_LOCATION = "unsloth_compiled_cache"
UNSLOTH_CREATED_FUNCTIONS = []

UNSLOTH_COMPILE_LOCATION_USE_TEMP = False

# Try creating a directory for cache, or else use a temporary folder
try:
    os.makedirs(UNSLOTH_COMPILE_LOCATION, exist_ok = True)
    if not os.path.exists(UNSLOTH_COMPILE_LOCATION): raise
except:
    from tempfile import TemporaryDirectory
    UNSLOTH_COMPILE_LOCATION_USE_TEMP = True
    UNSLOTH_COMPILE_LOCATION = TemporaryDirectory(ignore_cleanup_errors = True).name
    print(f"Unsloth: We can't create folders, so as a hack, we used a temporary directory = {UNSLOTH_COMPILE_LOCATION}")
pass

The problem would still persist, since the cache folder will be created in the directory where python code is launched (unless current directory editing rights are to restrictive to create the cache dir).

Using the temp dir approach should be the go to at all times.

[danielhanchen]

As an update - I revamped everything! I incorporated some of your ideas - hope this makes it better - essentially it tries to create a folder, and if i fails, it utilizes your idea of a tempfile directory