Skip to content

Commit 771bf4c

Browse files
mayinghanmayinghan
committed
support only persisting jsonl to local but no upload
1 parent 495ff63 commit 771bf4c

File tree

2 files changed

+132
-124
lines changed

2 files changed

+132
-124
lines changed

‎eval_protocol/pytest/handle_persist_flow.py‎

Lines changed: 128 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,9 @@
1616
def handle_persist_flow(all_results: list[list[EvaluationRow]], test_func_name: str):
1717
try:
1818
# Default is to save and upload experiment JSONL files, unless explicitly disabled
19-
should_save_and_upload = os.getenv("EP_NO_UPLOAD") != "1"
19+
should_save = os.getenv("EP_NO_PERSIST_RESULTS_JSONL") != "1"
2020

21-
if should_save_and_upload:
21+
if should_save:
2222
current_run_rows = [item for sublist in all_results for item in sublist]
2323
if current_run_rows:
2424
experiments: dict[str, list[EvaluationRow]] = defaultdict(list)
@@ -81,129 +81,133 @@ def handle_persist_flow(all_results: list[list[EvaluationRow]], test_func_name:
8181
json.dump(row_data, f, ensure_ascii=False)
8282
f.write("\n")
8383

84-
def get_auth_value(key: str) -> str | None:
85-
"""Get auth value from config file or environment."""
86-
try:
87-
config_path = Path.home() / ".fireworks" / "auth.ini"
88-
if config_path.exists():
89-
config = configparser.ConfigParser() # noqa: F821
90-
config.read(config_path)
91-
for section in ["DEFAULT", "auth"]:
92-
if config.has_section(section) and config.has_option(section, key):
93-
return config.get(section, key)
94-
except Exception:
95-
pass
96-
return os.getenv(key)
97-
98-
fireworks_api_key = get_auth_value("FIREWORKS_API_KEY")
99-
fireworks_account_id = get_auth_value("FIREWORKS_ACCOUNT_ID")
100-
101-
if not fireworks_api_key and not fireworks_account_id:
102-
store_experiment_link(
103-
experiment_id,
104-
"No Fireworks API key AND account ID found",
105-
"failure",
106-
)
107-
continue
108-
elif not fireworks_api_key:
109-
store_experiment_link(
110-
experiment_id,
111-
"No Fireworks API key found",
112-
"failure",
113-
)
114-
continue
115-
elif not fireworks_account_id:
116-
store_experiment_link(
117-
experiment_id,
118-
"No Fireworks account ID found",
119-
"failure",
120-
)
121-
continue
122-
123-
headers = {"Authorization": f"Bearer {fireworks_api_key}", "Content-Type": "application/json"}
124-
125-
# Make dataset first
126-
dataset_url = f"https://api.fireworks.ai/v1/accounts/{fireworks_account_id}/datasets"
127-
128-
dataset_payload = { # pyright: ignore[reportUnknownVariableType]
129-
"dataset": {
130-
"displayName": dataset_name,
131-
"evalProtocol": {},
132-
"format": "FORMAT_UNSPECIFIED",
133-
"exampleCount": f"{len(exp_rows)}",
134-
},
135-
"datasetId": dataset_name,
136-
}
137-
138-
dataset_response = requests.post(dataset_url, json=dataset_payload, headers=headers) # pyright: ignore[reportUnknownArgumentType]
139-
140-
# Skip if dataset creation failed
141-
if dataset_response.status_code not in [200, 201]:
142-
store_experiment_link(
143-
experiment_id,
144-
f"Dataset creation failed: {dataset_response.status_code} {dataset_response.text}",
145-
"failure",
146-
)
147-
continue
148-
149-
dataset_data: dict[str, Any] = dataset_response.json() # pyright: ignore[reportAny, reportExplicitAny]
150-
dataset_id = dataset_data.get("datasetId", dataset_name) # pyright: ignore[reportAny]
151-
152-
# Upload the JSONL file content
153-
upload_url = (
154-
f"https://api.fireworks.ai/v1/accounts/{fireworks_account_id}/datasets/{dataset_id}:upload"
155-
)
156-
upload_headers = {"Authorization": f"Bearer {fireworks_api_key}"}
157-
158-
with open(exp_file, "rb") as f:
159-
files = {"file": f}
160-
upload_response = requests.post(upload_url, files=files, headers=upload_headers)
161-
162-
# Skip if upload failed
163-
if upload_response.status_code not in [200, 201]:
164-
store_experiment_link(
165-
experiment_id,
166-
f"File upload failed: {upload_response.status_code} {upload_response.text}",
167-
"failure",
168-
)
169-
continue
170-
171-
# Create evaluation job (optional - don't skip experiment if this fails)
172-
eval_job_url = f"https://api.fireworks.ai/v1/accounts/{fireworks_account_id}/evaluationJobs"
173-
# Truncate job ID to fit 63 character limit
174-
job_id_base = f"{dataset_name}-job"
175-
if len(job_id_base) > 63:
176-
# Keep the "-job" suffix and truncate the dataset_name part
177-
max_dataset_name_len = 63 - 4 # 4 = len("-job")
178-
truncated_dataset_name = dataset_name[:max_dataset_name_len]
179-
job_id_base = f"{truncated_dataset_name}-job"
180-
181-
eval_job_payload = {
182-
"evaluationJobId": job_id_base,
183-
"evaluationJob": {
184-
"evaluator": f"accounts/{fireworks_account_id}/evaluators/dummy",
185-
"inputDataset": f"accounts/{fireworks_account_id}/datasets/dummy",
186-
"outputDataset": f"accounts/{fireworks_account_id}/datasets/{dataset_id}",
187-
},
188-
}
189-
190-
eval_response = requests.post(eval_job_url, json=eval_job_payload, headers=headers)
191-
192-
if eval_response.status_code in [200, 201]:
193-
eval_job_data = eval_response.json() # pyright: ignore[reportAny]
194-
job_id = eval_job_data.get("evaluationJobId", job_id_base) # pyright: ignore[reportAny]
195-
196-
store_experiment_link(
197-
experiment_id,
198-
f"https://app.fireworks.ai/dashboard/evaluation-jobs/{job_id}",
199-
"success",
200-
)
201-
else:
202-
store_experiment_link(
203-
experiment_id,
204-
f"Job creation failed: {eval_response.status_code} {eval_response.text}",
205-
"failure",
84+
should_upload = os.getenv("EP_NO_UPLOAD") != "1"
85+
86+
if should_upload:
87+
88+
def get_auth_value(key: str) -> str | None:
89+
"""Get auth value from config file or environment."""
90+
try:
91+
config_path = Path.home() / ".fireworks" / "auth.ini"
92+
if config_path.exists():
93+
config = configparser.ConfigParser() # noqa: F821
94+
config.read(config_path)
95+
for section in ["DEFAULT", "auth"]:
96+
if config.has_section(section) and config.has_option(section, key):
97+
return config.get(section, key)
98+
except Exception:
99+
pass
100+
return os.getenv(key)
101+
102+
fireworks_api_key = get_auth_value("FIREWORKS_API_KEY")
103+
fireworks_account_id = get_auth_value("FIREWORKS_ACCOUNT_ID")
104+
105+
if not fireworks_api_key and not fireworks_account_id:
106+
store_experiment_link(
107+
experiment_id,
108+
"No Fireworks API key AND account ID found",
109+
"failure",
110+
)
111+
continue
112+
elif not fireworks_api_key:
113+
store_experiment_link(
114+
experiment_id,
115+
"No Fireworks API key found",
116+
"failure",
117+
)
118+
continue
119+
elif not fireworks_account_id:
120+
store_experiment_link(
121+
experiment_id,
122+
"No Fireworks account ID found",
123+
"failure",
124+
)
125+
continue
126+
127+
headers = {"Authorization": f"Bearer {fireworks_api_key}", "Content-Type": "application/json"}
128+
129+
# Make dataset first
130+
dataset_url = f"https://api.fireworks.ai/v1/accounts/{fireworks_account_id}/datasets"
131+
132+
dataset_payload = { # pyright: ignore[reportUnknownVariableType]
133+
"dataset": {
134+
"displayName": dataset_name,
135+
"evalProtocol": {},
136+
"format": "FORMAT_UNSPECIFIED",
137+
"exampleCount": f"{len(exp_rows)}",
138+
},
139+
"datasetId": dataset_name,
140+
}
141+
142+
dataset_response = requests.post(dataset_url, json=dataset_payload, headers=headers) # pyright: ignore[reportUnknownArgumentType]
143+
144+
# Skip if dataset creation failed
145+
if dataset_response.status_code not in [200, 201]:
146+
store_experiment_link(
147+
experiment_id,
148+
f"Dataset creation failed: {dataset_response.status_code} {dataset_response.text}",
149+
"failure",
150+
)
151+
continue
152+
153+
dataset_data: dict[str, Any] = dataset_response.json() # pyright: ignore[reportAny, reportExplicitAny]
154+
dataset_id = dataset_data.get("datasetId", dataset_name) # pyright: ignore[reportAny]
155+
156+
# Upload the JSONL file content
157+
upload_url = (
158+
f"https://api.fireworks.ai/v1/accounts/{fireworks_account_id}/datasets/{dataset_id}:upload"
206159
)
160+
upload_headers = {"Authorization": f"Bearer {fireworks_api_key}"}
161+
162+
with open(exp_file, "rb") as f:
163+
files = {"file": f}
164+
upload_response = requests.post(upload_url, files=files, headers=upload_headers)
165+
166+
# Skip if upload failed
167+
if upload_response.status_code not in [200, 201]:
168+
store_experiment_link(
169+
experiment_id,
170+
f"File upload failed: {upload_response.status_code} {upload_response.text}",
171+
"failure",
172+
)
173+
continue
174+
175+
# Create evaluation job (optional - don't skip experiment if this fails)
176+
eval_job_url = f"https://api.fireworks.ai/v1/accounts/{fireworks_account_id}/evaluationJobs"
177+
# Truncate job ID to fit 63 character limit
178+
job_id_base = f"{dataset_name}-job"
179+
if len(job_id_base) > 63:
180+
# Keep the "-job" suffix and truncate the dataset_name part
181+
max_dataset_name_len = 63 - 4 # 4 = len("-job")
182+
truncated_dataset_name = dataset_name[:max_dataset_name_len]
183+
job_id_base = f"{truncated_dataset_name}-job"
184+
185+
eval_job_payload = {
186+
"evaluationJobId": job_id_base,
187+
"evaluationJob": {
188+
"evaluator": f"accounts/{fireworks_account_id}/evaluators/dummy",
189+
"inputDataset": f"accounts/{fireworks_account_id}/datasets/dummy",
190+
"outputDataset": f"accounts/{fireworks_account_id}/datasets/{dataset_id}",
191+
},
192+
}
193+
194+
eval_response = requests.post(eval_job_url, json=eval_job_payload, headers=headers)
195+
196+
if eval_response.status_code in [200, 201]:
197+
eval_job_data = eval_response.json() # pyright: ignore[reportAny]
198+
job_id = eval_job_data.get("evaluationJobId", job_id_base) # pyright: ignore[reportAny]
199+
200+
store_experiment_link(
201+
experiment_id,
202+
f"https://app.fireworks.ai/dashboard/evaluation-jobs/{job_id}",
203+
"success",
204+
)
205+
else:
206+
store_experiment_link(
207+
experiment_id,
208+
f"Job creation failed: {eval_response.status_code} {eval_response.text}",
209+
"failure",
210+
)
207211

208212
except Exception as e:
209213
# Do not fail evaluation if experiment JSONL writing fails

‎eval_protocol/pytest/plugin.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,10 @@ def pytest_configure(config) -> None:
258258
if threshold_env is not None:
259259
os.environ["EP_PASSED_THRESHOLD"] = threshold_env
260260

261+
if config.getoption("--ep-no-persist-results-jsonl"):
262+
# flag to turn off persisting results as jsonl files
263+
os.environ["EP_NO_PERSIST_RESULTS_JSONL"] = "1"
264+
261265
if config.getoption("--ep-no-upload"):
262266
os.environ["EP_NO_UPLOAD"] = "1"
263267

0 commit comments

Comments
 (0)