mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-08-03 07:21:13 +08:00
Support for async node functions (#8830)
* Support for async execution functions This commit adds support for node execution functions defined as async. When a node's execution function is defined as async, we can continue executing other nodes while it is processing. Standard uses of `await` should "just work", but people will still have to be careful if they spawn actual threads. Because torch doesn't really have async/await versions of functions, this won't particularly help with most locally-executing nodes, but it does work for e.g. web requests to other machines. In addition to the execute function, the `VALIDATE_INPUTS` and `check_lazy_status` functions can also be defined as async, though we'll only resolve one node at a time right now for those. * Add the execution model tests to CI * Add a missing file It looks like this got caught by .gitignore? There's probably a better place to put it, but I'm not sure what that is. * Add the websocket library for automated tests * Add additional tests for async error cases Also fixes one bug that was found when an async function throws an error after being scheduled on a task. * Add a feature flags message to reduce bandwidth We now only send 1 preview message of the latest type the client can support. We'll add a console warning when the client fails to send a feature flags message at some point in the future. * Add async tests to CI * Don't actually add new tests in this PR Will do it in a separate PR * Resolve unit test in GPU-less runner * Just remove the tests that GHA can't handle * Change line endings to UNIX-style * Avoid loading model_management.py so early Because model_management.py has a top-level `logging.info`, we have to be careful not to import that file before we call `setup_logging`. If we do, we end up having the default logging handler registered in addition to our custom one.
This commit is contained in:
77
tests-unit/websocket_feature_flags_test.py
Normal file
77
tests-unit/websocket_feature_flags_test.py
Normal file
@@ -0,0 +1,77 @@
|
||||
"""Simplified tests for WebSocket feature flags functionality."""
|
||||
from comfy_api import feature_flags
|
||||
|
||||
|
||||
class TestWebSocketFeatureFlags:
|
||||
"""Test suite for WebSocket feature flags integration."""
|
||||
|
||||
def test_server_feature_flags_response(self):
|
||||
"""Test server feature flags are properly formatted."""
|
||||
features = feature_flags.get_server_features()
|
||||
|
||||
# Check expected server features
|
||||
assert "supports_preview_metadata" in features
|
||||
assert features["supports_preview_metadata"] is True
|
||||
assert "max_upload_size" in features
|
||||
assert isinstance(features["max_upload_size"], (int, float))
|
||||
|
||||
def test_progress_py_checks_feature_flags(self):
|
||||
"""Test that progress.py checks feature flags before sending metadata."""
|
||||
# This simulates the check in progress.py
|
||||
client_id = "test_client"
|
||||
sockets_metadata = {"test_client": {"feature_flags": {}}}
|
||||
|
||||
# The actual check would be in progress.py
|
||||
supports_metadata = feature_flags.supports_feature(
|
||||
sockets_metadata, client_id, "supports_preview_metadata"
|
||||
)
|
||||
|
||||
assert supports_metadata is False
|
||||
|
||||
def test_multiple_clients_different_features(self):
|
||||
"""Test handling multiple clients with different feature support."""
|
||||
sockets_metadata = {
|
||||
"modern_client": {
|
||||
"feature_flags": {"supports_preview_metadata": True}
|
||||
},
|
||||
"legacy_client": {
|
||||
"feature_flags": {}
|
||||
}
|
||||
}
|
||||
|
||||
# Check modern client
|
||||
assert feature_flags.supports_feature(
|
||||
sockets_metadata, "modern_client", "supports_preview_metadata"
|
||||
) is True
|
||||
|
||||
# Check legacy client
|
||||
assert feature_flags.supports_feature(
|
||||
sockets_metadata, "legacy_client", "supports_preview_metadata"
|
||||
) is False
|
||||
|
||||
def test_feature_negotiation_message_format(self):
|
||||
"""Test the format of feature negotiation messages."""
|
||||
# Client message format
|
||||
client_message = {
|
||||
"type": "feature_flags",
|
||||
"data": {
|
||||
"supports_preview_metadata": True,
|
||||
"api_version": "1.0.0"
|
||||
}
|
||||
}
|
||||
|
||||
# Verify structure
|
||||
assert client_message["type"] == "feature_flags"
|
||||
assert "supports_preview_metadata" in client_message["data"]
|
||||
|
||||
# Server response format (what would be sent)
|
||||
server_features = feature_flags.get_server_features()
|
||||
server_message = {
|
||||
"type": "feature_flags",
|
||||
"data": server_features
|
||||
}
|
||||
|
||||
# Verify structure
|
||||
assert server_message["type"] == "feature_flags"
|
||||
assert "supports_preview_metadata" in server_message["data"]
|
||||
assert server_message["data"]["supports_preview_metadata"] is True
|
Reference in New Issue
Block a user