Skip to content

Commit 080a1cb

Browse files
alexgshawgabeorlanski
alexgshaw
and
gabeorlanski
authored
Simplify trial flow (harbor-framework#1672)
* Refactor trial execution by shape * Clean up trial helper typing * Skip Windows container hello world without Docker * Partial refactor. * Improve artifact handler. * Minor multi step fixes. * Make artifact handler paths operation scoped * Fix CI after trial flow cleanup * Keep download dir excludes explicit * Rename download dir exclusions helper * Address artifact exclusion review comments * Avoid duplicate single-step artifact recovery * Avoid double stop after cancellation --------- Co-authored-by: gabeorlanski <[email protected]>
1 parent a083e1b commit 080a1cb

28 files changed

Lines changed: 2538 additions & 1342 deletions

‎pyproject.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,9 @@ build-backend = "uv_build"
4343
[tool.uv.workspace]
4444
members = ["packages/*"]
4545

46+
[tool.uv.sources]
47+
harbor-rewardkit = { workspace = true }
48+
4649
[project.optional-dependencies]
4750
e2b = ["e2b>=2.4.2", "dockerfile-parse>=2.0.1"]
4851
daytona = ["daytona>=0.165.0"]
@@ -64,6 +67,7 @@ tinker = [
6467
dev = [
6568
"harbor[cloud]",
6669
"harbor[tinker]",
70+
"harbor-rewardkit",
6771
"ipykernel>=6.30.1",
6872
"pytest>=8.4.2",
6973
"pytest-asyncio>=1.2.0",

‎src/harbor/agents/base.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,10 +59,6 @@ def _init_model_info(self):
5959
self._parsed_model_name = self.model_name
6060

6161
def to_agent_info(self) -> AgentInfo:
62-
# Record model_info whenever we have a name — provider is optional
63-
# (``-m gpt-5.4`` without a ``<provider>/`` prefix still yields a
64-
# usable model identity). The DB-side ``model.provider`` column has
65-
# a ``DEFAULT 'unknown'`` that takes over for the None case.
6662
return AgentInfo(
6763
name=self.name(),
6864
version=self.version() or "unknown",
@@ -134,3 +130,11 @@ async def run(
134130
environment: The environment in which to complete the task.
135131
context: The context to populate with the results of the agent execution.
136132
"""
133+
134+
def populate_context_post_run(self, context: AgentContext) -> None:
135+
"""Optionally backfill context after ``run()`` completes.
136+
137+
Agents that write logs or trajectories during execution can override this
138+
to parse those outputs after the trial syncs logs back to the host.
139+
"""
140+
pass

‎src/harbor/agents/installed/base.py‎

Lines changed: 1 addition & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,7 @@
33
from abc import ABC, abstractmethod
44
from dataclasses import dataclass
55
from pathlib import Path
6-
from typing import TYPE_CHECKING, Any, ClassVar, Literal
7-
8-
if TYPE_CHECKING:
9-
from harbor.models.agent.context import AgentContext
6+
from typing import Any, ClassVar, Literal
107

118
from harbor.agents.base import BaseAgent
129
from harbor.environments.base import BaseEnvironment
@@ -255,15 +252,6 @@ def _get_env_prefixed(self, prefix: str) -> dict[str, str]:
255252
result[key[len(prefix) :]] = value
256253
return result
257254

258-
@abstractmethod
259-
def populate_context_post_run(self, context: "AgentContext") -> None:
260-
"""Populate the context with the results of the agent execution.
261-
262-
Called by the trial after ``run()`` completes (even on failure).
263-
Typically involves parsing trajectory files and extracting token counts.
264-
"""
265-
pass
266-
267255
def version(self) -> str | None:
268256
return self._version
269257

‎src/harbor/environments/base.py‎

Lines changed: 59 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
import asyncio
2+
import contextlib
23
import logging
34
import shlex
5+
import tarfile
6+
import tempfile
47
import time
58
import warnings
69
from abc import ABC, abstractmethod
7-
from collections.abc import Sequence
8-
from pathlib import Path, PurePath
10+
from collections.abc import Generator, Sequence
11+
from pathlib import Path, PurePath, PurePosixPath
912

1013
from pydantic import BaseModel
1114

@@ -18,6 +21,8 @@
1821
from harbor.utils.scripts import quote_shell_arg
1922

2023
EnvironmentPath = str | PurePath
24+
_TRANSFER_TAR_FILENAME = ".hb-transfer.tar.gz"
25+
_ENV_TRANSFER_TAR_PATH = str(PurePosixPath("/tmp") / _TRANSFER_TAR_FILENAME)
2126

2227

2328
class HealthcheckError(RuntimeError):
@@ -165,6 +170,19 @@ def _resolve_user(self, user: str | int | None) -> str | int | None:
165170
"""
166171
return user if user is not None else self.default_user
167172

173+
@contextlib.contextmanager
174+
def with_default_user(
175+
self,
176+
user: str | int | None,
177+
) -> Generator[None, None, None]:
178+
"""Temporarily set the default user for environment operations."""
179+
previous = self.default_user
180+
self.default_user = user
181+
try:
182+
yield
183+
finally:
184+
self.default_user = previous
185+
168186
def _merge_env(self, env: dict[str, str] | None) -> dict[str, str] | None:
169187
"""Merge persistent env vars with per-exec env vars.
170188
@@ -468,6 +486,45 @@ async def download_dir(self, source_dir: str, target_dir: Path | str):
468486
target_dir: The local path to which to copy the directory.
469487
"""
470488

489+
async def download_dir_with_exclusions(
490+
self,
491+
*,
492+
source_dir: str,
493+
target_dir: Path | str,
494+
exclude: list[str],
495+
) -> None:
496+
"""Download a directory through a temporary tar archive with excludes."""
497+
target = Path(target_dir)
498+
target.mkdir(parents=True, exist_ok=True)
499+
500+
exclude_flags = " ".join(
501+
f"--exclude={shlex.quote(pattern)}" for pattern in exclude
502+
)
503+
env_tar_path = shlex.quote(_ENV_TRANSFER_TAR_PATH)
504+
source_path = shlex.quote(source_dir)
505+
506+
result = await self.exec(
507+
f"tar czf {env_tar_path} {exclude_flags} -C {source_path} .",
508+
timeout_sec=120,
509+
user="root",
510+
)
511+
if result.return_code != 0:
512+
output = result.stderr or result.stdout or "no output"
513+
raise RuntimeError(
514+
"Failed to create transfer archive for "
515+
f"{source_dir!r} with code {result.return_code}: {output}"
516+
)
517+
518+
with tempfile.TemporaryDirectory() as host_tmp_dir:
519+
host_tar_path = Path(host_tmp_dir) / _TRANSFER_TAR_FILENAME
520+
await self.download_file(
521+
source_path=_ENV_TRANSFER_TAR_PATH,
522+
target_path=host_tar_path,
523+
)
524+
525+
with tarfile.open(host_tar_path, "r:gz") as tf:
526+
tf.extractall(path=target, filter="data")
527+
471528
@abstractmethod
472529
async def exec(
473530
self,

‎src/harbor/models/trial/artifact_manifest.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
from typing import Any, Literal
2+
3+
from pydantic import BaseModel, Field
4+
5+
6+
class ArtifactManifestEntry(BaseModel):
7+
source: str
8+
destination: str
9+
type: Literal["file", "directory"]
10+
status: Literal["ok", "failed", "empty"]
11+
12+
13+
class ArtifactManifest(BaseModel):
14+
entries: list[ArtifactManifestEntry] = Field(default_factory=list)
15+
16+
def to_json_data(self) -> list[dict[str, Any]]:
17+
return [entry.model_dump(mode="json") for entry in self.entries]

0 commit comments

Comments
 (0)