Skip to content

Commit 1930778

Browse files
deep1401deep1401
committed
fix job not found errors when running with local providers
1 parent 07bcedd commit 1930778

3 files changed

Lines changed: 82 additions & 61 deletions

File tree

‎api/transformerlab/compute_providers/local.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,10 +150,18 @@ def launch_cluster(self, cluster_name: str, config: ClusterConfig) -> Dict[str,
150150

151151
self._ensure_venv_and_sync(venv_path)
152152

153+
# Use a per-job workspace directory as HOME for local runs so tools that
154+
# rely on ~ and $HOME resolve inside the job workspace instead of the
155+
# user's real home directory. This makes it easier to clone and run
156+
# code in an isolated workspace for each job.
157+
workspace_home = job_dir / "workspace"
158+
workspace_home.mkdir(parents=True, exist_ok=True)
159+
153160
venv_bin = venv_path / "bin"
154161
env = os.environ.copy()
155162
env.update(config.env_vars or {})
156163
env["PATH"] = f"{venv_bin}{os.pathsep}{env.get('PATH', '')}"
164+
env["HOME"] = str(workspace_home)
157165

158166
print(f"DEBUG: LocalProvider.launch_cluster: cluster_name={cluster_name}")
159167
print(f"DEBUG: LocalProvider.launch_cluster: job_dir={job_dir}")

‎api/transformerlab/routers/compute_provider.py‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1506,16 +1506,13 @@ async def launch_template_on_provider(
15061506
# env_vars["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key
15071507

15081508
# For local provider, set TFL_WORKSPACE_DIR so the lab SDK in the subprocess can find
1509-
# the job directory (workspace/jobs/<job_id>). Without this, Job.get(id) looks under
1510-
# ~/.transformerlab/workspace and fails with "Directory for Job with id '...' not found".
1509+
# the job directory (workspace/jobs/<job_id>). The organization context for the API
1510+
# request is already set by authentication middleware, so we can rely on
1511+
# get_workspace_dir() without mutating the global org context here.
15111512
if provider.type == ProviderType.LOCAL.value and team_id:
1512-
set_organization_id(team_id)
1513-
try:
1514-
workspace_dir = await get_workspace_dir()
1515-
if workspace_dir and not storage.is_remote_path(workspace_dir):
1516-
env_vars["TFL_WORKSPACE_DIR"] = workspace_dir
1517-
finally:
1518-
set_organization_id(None)
1513+
workspace_dir = await get_workspace_dir()
1514+
if workspace_dir and not storage.is_remote_path(workspace_dir):
1515+
env_vars["TFL_WORKSPACE_DIR"] = workspace_dir
15191516

15201517
# Resolve command (and optional setup override) for interactive sessions from gallery
15211518
base_command = request.command
@@ -1666,6 +1663,7 @@ async def launch_template_on_provider(
16661663
job_id=str(job_id),
16671664
experiment_id=request.experiment_id,
16681665
provider_id=provider.id,
1666+
team_id=team_id,
16691667
cluster_name=formatted_cluster_name,
16701668
cluster_config=cluster_config,
16711669
quota_hold_id=str(quota_hold.id) if quota_hold else None,

‎api/transformerlab/services/local_provider_queue.py‎

Lines changed: 67 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
from transformerlab.services import job_service, quota_service
88
from transformerlab.services.provider_service import get_provider_by_id, get_provider_instance
99
from transformerlab.db.session import async_session
10+
from transformerlab.shared.request_context import set_current_org_id
11+
from lab import dirs as lab_dirs
1012

1113

1214
class LocalLaunchWorkItem(BaseModel):
@@ -15,6 +17,7 @@ class LocalLaunchWorkItem(BaseModel):
1517
job_id: str
1618
experiment_id: str
1719
provider_id: str
20+
team_id: str
1821
cluster_name: str
1922
cluster_config: ClusterConfig
2023
quota_hold_id: Optional[str] = None
@@ -30,6 +33,7 @@ async def enqueue_local_launch(
3033
job_id: str,
3134
experiment_id: str,
3235
provider_id: str,
36+
team_id: str,
3337
cluster_name: str,
3438
cluster_config: ClusterConfig,
3539
quota_hold_id: Optional[str],
@@ -42,6 +46,7 @@ async def enqueue_local_launch(
4246
job_id=str(job_id),
4347
experiment_id=str(experiment_id),
4448
provider_id=str(provider_id),
49+
team_id=str(team_id),
4550
cluster_name=cluster_name,
4651
cluster_config=cluster_config,
4752
quota_hold_id=quota_hold_id,
@@ -62,71 +67,81 @@ async def _local_launch_worker() -> None:
6267
try:
6368
# Use a dedicated DB session inside the worker
6469
async with async_session() as session:
65-
provider = await get_provider_by_id(session, item.provider_id)
66-
if not provider:
67-
await job_service.job_update_status(
68-
item.job_id,
69-
"FAILED",
70-
experiment_id=item.experiment_id,
71-
error_msg="Provider not found for local launch",
72-
session=session,
73-
)
74-
if item.quota_hold_id:
75-
await quota_service.release_quota_hold(session, hold_id=item.quota_hold_id)
76-
await session.commit()
77-
continue
78-
79-
provider_instance = await get_provider_instance(provider)
80-
81-
# Transition from WAITING -> initial_status (LAUNCHING / INTERACTIVE)
82-
await job_service.job_update_status(
83-
item.job_id,
84-
item.initial_status,
85-
experiment_id=item.experiment_id,
86-
session=session,
87-
)
88-
await session.commit()
89-
90-
loop = asyncio.get_running_loop()
70+
# Ensure lab SDK and API request context are scoped to the correct organization
71+
set_current_org_id(item.team_id)
72+
lab_dirs.set_organization_id(item.team_id)
9173
try:
92-
# Ensure only one local launch runs at a time
93-
async with _worker_lock:
94-
launch_result = await loop.run_in_executor(
95-
None,
96-
lambda: provider_instance.launch_cluster(item.cluster_name, item.cluster_config),
74+
provider = await get_provider_by_id(session, item.provider_id)
75+
if not provider:
76+
await job_service.job_update_status(
77+
item.job_id,
78+
"FAILED",
79+
experiment_id=item.experiment_id,
80+
error_msg="Provider not found for local launch",
81+
session=session,
9782
)
98-
except Exception as exc: # noqa: BLE001
99-
# Release quota hold and mark job failed
100-
if item.quota_hold_id:
101-
await quota_service.release_quota_hold(session, hold_id=item.quota_hold_id)
102-
await session.commit()
83+
if item.quota_hold_id:
84+
await quota_service.release_quota_hold(session, hold_id=item.quota_hold_id)
85+
await session.commit()
86+
continue
87+
88+
provider_instance = await get_provider_instance(provider)
10389

90+
# Transition from WAITING -> initial_status (LAUNCHING / INTERACTIVE)
10491
await job_service.job_update_status(
10592
item.job_id,
106-
"FAILED",
93+
item.initial_status,
10794
experiment_id=item.experiment_id,
108-
error_msg=str(exc),
10995
session=session,
11096
)
11197
await session.commit()
112-
continue
11398

114-
# On success, we keep the job in LAUNCHING/INTERACTIVE; status checks will
115-
# complete it when the local process exits.
116-
if isinstance(launch_result, dict):
117-
await job_service.job_update_job_data_insert_key_value(
118-
item.job_id,
119-
"provider_launch_result",
120-
launch_result,
121-
item.experiment_id,
122-
)
123-
request_id = launch_result.get("request_id")
124-
if request_id:
99+
loop = asyncio.get_running_loop()
100+
try:
101+
# Ensure only one local launch runs at a time
102+
async with _worker_lock:
103+
launch_result = await loop.run_in_executor(
104+
None,
105+
lambda: provider_instance.launch_cluster(item.cluster_name, item.cluster_config),
106+
)
107+
except Exception as exc: # noqa: BLE001
108+
# Release quota hold and mark job failed
109+
if item.quota_hold_id:
110+
await quota_service.release_quota_hold(session, hold_id=item.quota_hold_id)
111+
await session.commit()
112+
113+
await job_service.job_update_status(
114+
item.job_id,
115+
"FAILED",
116+
experiment_id=item.experiment_id,
117+
error_msg=str(exc),
118+
session=session,
119+
)
120+
await session.commit()
121+
continue
122+
123+
print(f"DEBUG: LocalProviderQueue._local_launch_worker: launch_result={launch_result}")
124+
125+
# On success, we keep the job in LAUNCHING/INTERACTIVE; status checks will
126+
# complete it when the local process exits.
127+
if isinstance(launch_result, dict):
125128
await job_service.job_update_job_data_insert_key_value(
126129
item.job_id,
127-
"orchestrator_request_id",
128-
request_id,
130+
"provider_launch_result",
131+
launch_result,
129132
item.experiment_id,
130133
)
134+
request_id = launch_result.get("request_id")
135+
if request_id:
136+
await job_service.job_update_job_data_insert_key_value(
137+
item.job_id,
138+
"orchestrator_request_id",
139+
request_id,
140+
item.experiment_id,
141+
)
142+
finally:
143+
# Clear org context after each item
144+
set_current_org_id(None)
145+
lab_dirs.set_organization_id(None)
131146
finally:
132147
_local_launch_queue.task_done()

0 commit comments

Comments
 (0)