Implement idle-killer enhancements and swarm management improvements
- Marked critical bugs as resolved in the review documentation, including changes to the `arm_idle_killer` function to raise errors on credential creation failures and ensure proper file permissions for JSON credentials. - Introduced a new `_try_arm_idle_killer` function in `provision.py` to manage idle-killer state more effectively, ensuring it arms correctly during provisioning. - Updated the `swarm_busy` function in `remote/idle_killer.py` to allow idle state after a specified duration of Swarm unavailability, preventing unnecessary billing. - Enhanced performance tuning logic in `tune_swarm_perf.py` to ensure proper handling of pip installation success before applying extra arguments. - Added tests to validate the new idle-killer behavior and swarm management logic, ensuring robustness in handling idle states and error conditions.
This commit is contained in:
@@ -106,17 +106,18 @@ def patch_backends_extra_args(extra: str) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def pip_install_sage(pip: Path) -> None:
|
||||
def pip_install_sage(pip: Path) -> bool:
|
||||
print(f"pip install triton sageattention via {pip}")
|
||||
env = dict(os.environ)
|
||||
env["PIP_DISABLE_PIP_VERSION_CHECK"] = "1"
|
||||
# Best-effort: do not fail whole tune if wheels missing for this torch.
|
||||
cmd = [str(pip), "install", "-U", "triton", "sageattention"]
|
||||
try:
|
||||
subprocess.check_call(cmd, env=env)
|
||||
print("triton + sageattention installed")
|
||||
return True
|
||||
except subprocess.CalledProcessError as exc:
|
||||
print(f"WARN: pip install failed ({exc}) — ExtraArgs may no-op until fixed")
|
||||
print(f"WARN: pip install failed ({exc}) — ExtraArgs NOT applied")
|
||||
return False
|
||||
|
||||
|
||||
def main() -> int:
|
||||
@@ -134,24 +135,25 @@ def main() -> int:
|
||||
return 0
|
||||
|
||||
print(f"perf tune: {plan['name']} tier={plan['tier']} sage={plan['use_sage']}")
|
||||
pip_ok = False
|
||||
pip_ok = not plan["use_sage"]
|
||||
restarted_needed = False
|
||||
if plan["use_sage"]:
|
||||
pip = find_pip()
|
||||
if pip:
|
||||
pip_install_sage(pip)
|
||||
pip_ok = True
|
||||
pip_ok = pip_install_sage(pip)
|
||||
else:
|
||||
print("Comfy venv pip not found yet — will retry next up")
|
||||
if patch_backends_extra_args(plan["extra_args"]):
|
||||
pip_ok = False
|
||||
# Only patch ExtraArgs when wheels installed — otherwise Comfy may break.
|
||||
if pip_ok and patch_backends_extra_args(plan["extra_args"]):
|
||||
restarted_needed = True
|
||||
|
||||
marker = {
|
||||
"uuid": plan["uuid"],
|
||||
"name": plan["name"],
|
||||
"tier": plan["tier"],
|
||||
"extra_args": plan["extra_args"],
|
||||
"pip_ok": pip_ok or not plan["use_sage"],
|
||||
"extra_args": plan["extra_args"] if pip_ok else "",
|
||||
"pip_ok": pip_ok,
|
||||
"restart_needed": restarted_needed,
|
||||
}
|
||||
MARKER.write_text(json.dumps(marker, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
Reference in New Issue
Block a user