Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
Commit ·
829ae5c
1
Parent(s): 00102c4
fix: use user OAuth token for job creation (has job.write on org)
Browse files- agent/tools/jobs_tool.py +12 -29
agent/tools/jobs_tool.py
CHANGED
|
@@ -293,11 +293,8 @@ class HfJobsTool:
|
|
| 293 |
log_callback: Optional[Callable[[str], Awaitable[None]]] = None,
|
| 294 |
session: Any = None,
|
| 295 |
tool_call_id: Optional[str] = None,
|
| 296 |
-
user_token: Optional[str] = None,
|
| 297 |
):
|
| 298 |
self.hf_token = hf_token
|
| 299 |
-
# user_token is injected into job secrets; hf_token is for API calls (job creation)
|
| 300 |
-
self.user_token = user_token or hf_token
|
| 301 |
self.api = HfApi(token=hf_token)
|
| 302 |
self.namespace = namespace
|
| 303 |
self.log_callback = log_callback
|
|
@@ -522,7 +519,7 @@ class HfJobsTool:
|
|
| 522 |
image=image,
|
| 523 |
command=command,
|
| 524 |
env=_add_default_env(args.get("env")),
|
| 525 |
-
secrets=_add_environment_variables(args.get("secrets"), self.
|
| 526 |
flavor=args.get("hardware_flavor", "cpu-basic"),
|
| 527 |
timeout=args.get("timeout", "30m"),
|
| 528 |
namespace=self.namespace,
|
|
@@ -754,7 +751,7 @@ To verify, call this tool with `{{"operation": "inspect", "job_id": "{job_id}"}}
|
|
| 754 |
command=command,
|
| 755 |
schedule=schedule,
|
| 756 |
env=_add_default_env(args.get("env")),
|
| 757 |
-
secrets=_add_environment_variables(args.get("secrets"), self.
|
| 758 |
flavor=args.get("hardware_flavor", "cpu-basic"),
|
| 759 |
timeout=args.get("timeout", "30m"),
|
| 760 |
namespace=self.namespace,
|
|
@@ -1038,34 +1035,20 @@ async def hf_jobs_handler(
|
|
| 1038 |
# If script is a sandbox file path, read it from the sandbox
|
| 1039 |
script = arguments.get("script", "")
|
| 1040 |
sandbox = getattr(session, "sandbox", None) if session else None
|
| 1041 |
-
|
| 1042 |
-
|
| 1043 |
-
|
| 1044 |
-
|
| 1045 |
-
|
| 1046 |
-
|
| 1047 |
-
script
|
| 1048 |
-
or script.startswith("./")
|
| 1049 |
-
or script.startswith("../")
|
| 1050 |
-
)
|
| 1051 |
-
)
|
| 1052 |
-
if is_path:
|
| 1053 |
-
import shlex
|
| 1054 |
-
|
| 1055 |
-
result = await asyncio.to_thread(sandbox.bash, f"cat {shlex.quote(script)}")
|
| 1056 |
-
if not result.success:
|
| 1057 |
-
return f"Failed to read {script} from sandbox: {result.error}", False
|
| 1058 |
-
arguments = {**arguments, "script": result.output}
|
| 1059 |
|
| 1060 |
-
|
| 1061 |
-
|
| 1062 |
-
admin_token = os.environ.get("HF_ADMIN_TOKEN") or user_token
|
| 1063 |
-
namespace = os.environ.get("HF_NAMESPACE") or (HfApi(token=admin_token).whoami().get("name") if admin_token else None)
|
| 1064 |
|
| 1065 |
tool = HfJobsTool(
|
| 1066 |
namespace=namespace,
|
| 1067 |
-
hf_token=
|
| 1068 |
-
user_token=user_token,
|
| 1069 |
log_callback=log_callback if session else None,
|
| 1070 |
session=session,
|
| 1071 |
tool_call_id=tool_call_id,
|
|
|
|
| 293 |
log_callback: Optional[Callable[[str], Awaitable[None]]] = None,
|
| 294 |
session: Any = None,
|
| 295 |
tool_call_id: Optional[str] = None,
|
|
|
|
| 296 |
):
|
| 297 |
self.hf_token = hf_token
|
|
|
|
|
|
|
| 298 |
self.api = HfApi(token=hf_token)
|
| 299 |
self.namespace = namespace
|
| 300 |
self.log_callback = log_callback
|
|
|
|
| 519 |
image=image,
|
| 520 |
command=command,
|
| 521 |
env=_add_default_env(args.get("env")),
|
| 522 |
+
secrets=_add_environment_variables(args.get("secrets"), self.hf_token),
|
| 523 |
flavor=args.get("hardware_flavor", "cpu-basic"),
|
| 524 |
timeout=args.get("timeout", "30m"),
|
| 525 |
namespace=self.namespace,
|
|
|
|
| 751 |
command=command,
|
| 752 |
schedule=schedule,
|
| 753 |
env=_add_default_env(args.get("env")),
|
| 754 |
+
secrets=_add_environment_variables(args.get("secrets"), self.hf_token),
|
| 755 |
flavor=args.get("hardware_flavor", "cpu-basic"),
|
| 756 |
timeout=args.get("timeout", "30m"),
|
| 757 |
namespace=self.namespace,
|
|
|
|
| 1035 |
# If script is a sandbox file path, read it from the sandbox
|
| 1036 |
script = arguments.get("script", "")
|
| 1037 |
sandbox = getattr(session, "sandbox", None) if session else None
|
| 1038 |
+
if sandbox and script:
|
| 1039 |
+
from agent.tools.sandbox_tool import resolve_sandbox_script
|
| 1040 |
+
content, error = await resolve_sandbox_script(sandbox, script)
|
| 1041 |
+
if error:
|
| 1042 |
+
return error, False
|
| 1043 |
+
if content:
|
| 1044 |
+
arguments = {**arguments, "script": content}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1045 |
|
| 1046 |
+
hf_token = session.hf_token if session else None
|
| 1047 |
+
namespace = os.environ.get("HF_NAMESPACE") or (HfApi(token=hf_token).whoami().get("name") if hf_token else None)
|
|
|
|
|
|
|
| 1048 |
|
| 1049 |
tool = HfJobsTool(
|
| 1050 |
namespace=namespace,
|
| 1051 |
+
hf_token=hf_token,
|
|
|
|
| 1052 |
log_callback=log_callback if session else None,
|
| 1053 |
session=session,
|
| 1054 |
tool_call_id=tool_call_id,
|