Spaces:
Runtime error
Runtime error
File size: 2,654 Bytes
f8f02c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import io
import time as _time
from typing import Any, Dict
import pandas as pd
from .client import api_client
from .parser import parse_model_response
from .schema import TransformedOutput
from .utils import wait_for_completion
def execute_model_step(
current_df: pd.DataFrame,
model_name: str,
compliance_type: str = "firco",
user_id: str = "system",
version: str = "latest",
number_of_reasonings: int = 1
) -> TransformedOutput:
"""
Executes a model integration step for a workflow.
Blocks minimally for 60 seconds (or more if model takes longer)
before retrieving results and transforming them.
"""
# 1. Prepare CSV buffer from dataframe
csv_buffer = io.BytesIO()
current_df.to_csv(csv_buffer, index=False)
csv_bytes = csv_buffer.getvalue()
# Use a unique filename per prediction to avoid S3 key collisions on the remote server
unique_filename = f"workflow_input_{int(_time.time() * 1000)}.csv"
# 2. Call prediction initialization
prediction_id = api_client.run_prediction(
csv_bytes=csv_bytes,
filename=unique_filename,
model_name=model_name,
compliance_type=compliance_type,
source_type="file",
version=version,
number_of_reasonings=number_of_reasonings
)
if not prediction_id:
raise ValueError("Failed to get a valid prediction_id from API.")
# 3 & 4 & 5. Wait, poll, and validate status
def check_status() -> Dict[str, Any]:
return api_client.get_prediction_run(prediction_id)
def is_completed(response: Dict[str, Any]) -> bool:
# The actual workflow status is typically nested under "data"
data_block = response.get("data", response)
status = str(data_block.get("status", "")).lower()
if status in ("failed", "error"):
raise RuntimeError(f"Prediction {prediction_id} failed remotely: {data_block.get('error')}")
return status == "completed"
try:
completed_response = wait_for_completion(
check_status_callable=check_status,
is_completed_callable=is_completed,
poll_interval=5,
max_wait_seconds=600
)
except TimeoutError as e:
raise TimeoutError(f"Prediction {prediction_id} timed out: {e}")
# 6. Parse and transform the final output
transformed_output = parse_model_response(current_df.copy(), completed_response)
# Tag it implicitly with the prediction ID for workflow traceability
transformed_output.prediction_id = prediction_id
return transformed_output
|