Buckets:

tracinginsights
/

2026

Files
xet
tracinginsights/2026 / experiments /exp3_workers.py
tracinginsights
2 months ago
download
raw
51.4 kB
"""
Season Session Telemetry Extraction Script (parallel)
=====================================================
Parallel variant of R.py. Drivers are processed with a process pool
(see MAX_WORKERS / TEL_MAX_WORKERS).
On platforms with os.fork (Linux), workers inherit the parent's already-loaded
session via copy-on-write (no per-worker reload). Elsewhere, each worker loads
the session once from cache and reuses it across its assigned drivers.
Laps without usable car/pos telemetry are skipped with a cheap pre-check
instead of paying for a failing FastF1 get_telemetry() merge.
Output directory:
{event_name}/{session_name}/
Per session also writes session_laptimes.json (merged from all drivers'
laptimes.json, alphabetical driver order).
Standalone cache:
cache
"""

import gc
import logging
import multiprocessing as mp
import os
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from typing import Dict, List, Optional, Tuple

import fastf1
import numpy as np
import orjson
import pandas as pd
import psutil
import requests

# ---------------------------------------------------------------------------
# Constants & Configuration
# ---------------------------------------------------------------------------

DEFAULT_YEAR = 2026
CACHE_DIR = "cache"
# With fork inheritance workers share one session; spawn still reloads once
# per worker so keep the default bounded. Override with TEL_MAX_WORKERS.
_MAX_WORKERS_ENV = os.environ.get("TEL_MAX_WORKERS", "").strip()
MAX_WORKERS = (
    max(1, int(_MAX_WORKERS_ENV))
    if _MAX_WORKERS_ENV.isdigit()
    else max(1, min(32, (os.cpu_count() or 2)))
)
# Uncomment one or more events to process.
TARGET_EVENT_NAMES_LIST = [

#    "Australian Grand Prix",
#     "Chinese Grand Prix",
#     "Japanese Grand Prix",

#     "Miami Grand Prix",

    # "Monaco Grand Prix",
    # "Barcelona Grand Prix",

    # "Canadian Grand Prix",
    "Austrian Grand Prix",
    # "British Grand Prix",
    # "Belgian Grand Prix",
    # "Hungarian Grand Prix",

    # "Bahrain Grand Prix",
    # "Saudi Arabian Grand Prix",
    # "Spanish Grand Prix",
    # "Dutch Grand Prix",
    # "Italian Grand Prix",
    # "Azerbaijan Grand Prix",
    # "Singapore Grand Prix",
    # "United States Grand Prix",
    # "Mexico City Grand Prix",
    # "São Paulo Grand Prix",
    # "Las Vegas Grand Prix",
    # "Qatar Grand Prix",
    # "Abu Dhabi Grand Prix",
]
TARGET_EVENT_NAMES = [e.strip() for e in TARGET_EVENT_NAMES_LIST if e.strip()]
if not TARGET_EVENT_NAMES:
    raise ValueError("Set at least one active event in TARGET_EVENT_NAMES_LIST.")
AVAILABLE_SESSIONS = [
    "Practice 1",
    "Practice 2",
    "Practice 3",
    "Qualifying",
    "Sprint Qualifying",
    "Sprint",
    "Race",
]
# Select one or more sessions from AVAILABLE_SESSIONS.
TARGET_SESSIONS = [
    # "Practice 1",
    # "Practice 2",
    # "Practice 3",
    # "Qualifying",
    # "Sprint Qualifying",
    # "Sprint",
    "Race",
]
invalid_target_sessions = sorted(set(TARGET_SESSIONS) - set(AVAILABLE_SESSIONS))
if invalid_target_sessions:
    raise ValueError(
        "Invalid TARGET_SESSIONS value(s): "
        + ", ".join(invalid_target_sessions)
    )
PROTO = "https"
HOST = "api.multiviewer.app"
HEADERS = {"User-Agent": "FastF1/"}
ORJSON_OPTS = orjson.OPT_SERIALIZE_NUMPY | orjson.OPT_NON_STR_KEYS
EPS = np.finfo(float).eps

# Pre-allocated smoothing kernels
_KERNEL_3 = np.ones(3, dtype=np.float64) / 3.0
_KERNEL_9 = np.ones(9, dtype=np.float64) / 9.0

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("session_extraction.log"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger("session_extractor")
logging.getLogger("fastf1").setLevel(logging.WARNING)
logging.getLogger("fastf1").propagate = False

_MISSING_TEXT_VALUES = frozenset({
    "",
    "null",
    "nan",
    "nat",
    "none",
    "inf",
    "-inf",
    "infinity",
    "-infinity",
})
_MISSING_TEXT_LIST = list(_MISSING_TEXT_VALUES)


# ---------------------------------------------------------------------------
# Helper Functions (Copied from main_optimized.py for standalone execution)
# ---------------------------------------------------------------------------
def _write_json(path: str, obj, normalize_missing: bool = False) -> None:
    if normalize_missing:
        obj = _normalize_missing_for_json(obj)
    with open(path, "wb") as f:
        f.write(orjson.dumps(obj, option=ORJSON_OPTS))


def _td_col_to_seconds(series: pd.Series) -> list:
    if series.empty:
        return []
    seconds = series.dt.total_seconds().to_numpy()
    mask = series.isna().to_numpy()
    out = np.round(seconds, 3).astype(object)
    out[mask] = "None"
    return out.tolist()


def _col_to_list_str_or_none(col) -> list:
    if isinstance(col, np.ndarray):
        vals = col
    else:
        if col.empty:
            return []
        vals = col.to_numpy()
    if len(vals) == 0:
        return []
    mask = pd.isna(vals)
    valid = ~mask
    out = np.empty(vals.shape, dtype=object)
    out[mask] = "None"
    valid_vals = vals[valid]
    s_vals = np.array([str(v).strip().lower() for v in valid_vals])
    missing_mask = np.isin(s_vals, _MISSING_TEXT_LIST)
    str_vals = np.array([str(v) for v in valid_vals])
    out[valid] = np.where(missing_mask, "None", str_vals)
    return out.tolist()


def _col_to_list_int_or_none(series: pd.Series) -> list:
    if series.empty:
        return []
    vals = series.to_numpy()
    mask = pd.isna(vals)
    out = np.empty(vals.shape, dtype=object)
    out[mask] = "None"
    out[~mask] = vals[~mask].astype(int)
    return out.tolist()


def _col_to_list_bool_or_none(series: pd.Series) -> list:
    if series.empty:
        return []
    vals = series.to_numpy()
    mask = pd.isna(vals)
    out = np.empty(vals.shape, dtype=object)
    out[mask] = "None"
    out[~mask] = vals[~mask].astype(bool)
    return out.tolist()


def _series_to_json_list(series: pd.Series) -> list:
    if series.empty:
        return []

    if pd.api.types.is_timedelta64_dtype(series.dtype):
        return _td_col_to_seconds(series)

    vals = series.to_numpy()
    if pd.api.types.is_float_dtype(series.dtype):
        vals_f = vals.astype(np.float64, copy=False)
        mask = ~np.isfinite(vals_f)
    else:
        mask = pd.isna(vals)
    out = np.empty(vals.shape, dtype=object)
    out[mask] = "None"

    valid = ~mask
    if not valid.any():
        return out.tolist()

    if pd.api.types.is_bool_dtype(series.dtype):
        out[valid] = vals[valid].astype(bool)
    elif pd.api.types.is_integer_dtype(series.dtype):
        out[valid] = vals[valid].astype(int)
    elif pd.api.types.is_float_dtype(series.dtype):
        out[valid] = vals[valid].astype(float)
    else:
        valid_vals = vals[valid]
        s_vals = np.array([str(v).strip().lower() for v in valid_vals])
        missing_mask = np.isin(s_vals, _MISSING_TEXT_LIST)
        str_vals = np.array([str(v) for v in valid_vals])
        out[valid] = np.where(missing_mask, "None", str_vals)

    return out.tolist()


def _scalar_to_json_primitive_or_none(value):
    if isinstance(value, (float, np.floating)):
        return "None" if not np.isfinite(value) else float(value)
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.bool_):
        return bool(value)
    if isinstance(value, str):
        return "None" if value.strip().lower() in _MISSING_TEXT_VALUES else value
    if pd.isna(value):
        return "None"
    return value


def _normalize_missing_for_json(value):
    if isinstance(value, dict):
        return {k: _normalize_missing_for_json(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_normalize_missing_for_json(v) for v in value]
    if isinstance(value, tuple):
        return [_normalize_missing_for_json(v) for v in value]
    if isinstance(value, np.ndarray):
        return [_normalize_missing_for_json(v) for v in value.tolist()]
    return _scalar_to_json_primitive_or_none(value)


def _dataframe_to_column_lists(df: pd.DataFrame) -> Dict[str, list]:
    if df is None or df.empty:
        return {}
    return {col: _series_to_json_list(df[col]) for col in df.columns}


_LAP_WEATHER_COL_MAP = (
    ("wT", "Time"),
    ("wAT", "AirTemp"),
    ("wH", "Humidity"),
    ("wP", "Pressure"),
    ("wR", "Rainfall"),
    ("wTT", "TrackTemp"),
    ("wWD", "WindDirection"),
    ("wWS", "WindSpeed"),
)
LAP_WEATHER_KEYS = tuple(k for k, _ in _LAP_WEATHER_COL_MAP)

# Fixed column order for per-driver / session laptimes (matches laps_data).
_LAPTIMES_KEYS = (
    "time", "lap", "compound", "stint",
    "s1", "s2", "s3", "life", "pos", "status", "pb",
    "sesT", "drv", "dNum", "pout", "pin",
    "s1T", "s2T", "s3T", "vi1", "vi2",
    "vfl", "vst", "fresh", "team", "lST",
    "lSD", "del", "delR", "ff1G", "iacc",
    *LAP_WEATHER_KEYS,
)

_RCM_COL_MAP = (
    ("time", "Time"),
    ("cat", "Category"),
    ("msg", "Message"),
    ("status", "Status"),
    ("flag", "Flag"),
    ("scope", "Scope"),
    ("sector", "Sector"),
    ("dNum", "RacingNumber"),
    ("lap", "Lap"),
)


def _session_weather_to_column_lists(weather_df: pd.DataFrame) -> Dict[str, list]:
    if weather_df is None or weather_df.empty:
        return {}

    out: Dict[str, list] = {}
    for short_key, weather_col in _LAP_WEATHER_COL_MAP:
        if weather_col in weather_df.columns:
            out[short_key] = _series_to_json_list(weather_df[weather_col])
    return out


def _session_rcm_to_column_lists(rcm_df: pd.DataFrame) -> Dict[str, list]:
    if rcm_df is None or rcm_df.empty:
        return {}

    out: Dict[str, list] = {}
    for short_key, rcm_col in _RCM_COL_MAP:
        if rcm_col in rcm_df.columns:
            out[short_key] = _series_to_json_list(rcm_df[rcm_col])
    return out


def _lap_weather_to_column_lists(laps: pd.DataFrame, weather_df: pd.DataFrame = None) -> Dict[str, list]:
    n_laps = len(laps)
    if n_laps == 0:
        return {k: [] for k in LAP_WEATHER_KEYS}

    none_row = ["None"] * n_laps
    out = {k: none_row.copy() for k in LAP_WEATHER_KEYS}

    if weather_df is None:
        if not hasattr(laps, "get_weather_data"):
            return out
        try:
            weather_df = laps.get_weather_data()
        except Exception:
            return out

    if weather_df is None:
        return out

    for short_key, weather_col in _LAP_WEATHER_COL_MAP:
        if weather_col not in weather_df.columns:
            continue

        values = _series_to_json_list(weather_df[weather_col])
        if len(values) < n_laps:
            values.extend(["None"] * (n_laps - len(values)))
        elif len(values) > n_laps:
            values = values[:n_laps]
        out[short_key] = values

    return out


def _array_to_list_float_or_none(arr: np.ndarray) -> list:
    if arr.size == 0:
        return []
    valid = np.isfinite(arr)
    if valid.all():
        return arr.tolist()
    out = np.empty(arr.shape, dtype=object)
    out[~valid] = "None"
    out[valid] = arr[valid]
    return out.tolist()


def _array_to_list_int_or_none(arr: np.ndarray) -> list:
    if arr.size == 0:
        return []
    mask = ~np.isfinite(arr)
    if not mask.any():
        return arr.astype(int).tolist()
    out = np.empty(arr.shape, dtype=object)
    out[mask] = "None"
    out[~mask] = arr[~mask].astype(int)
    return out.tolist()


def _smooth_outliers(arr: np.ndarray, threshold: float, use_abs: bool) -> None:
    if use_abs:
        mask = np.abs(arr) > threshold
    else:
        mask = arr > threshold
    if mask.any():
        indices = np.where(mask)[0]
        indices = indices[(indices >= 1) & (indices < len(arr) - 1)]
        if len(indices) > 0:
            arr[indices] = arr[indices - 1]


def _compute_accelerations(
    speed: np.ndarray,
    time_arr: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    z: np.ndarray,
    dist: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    # Convert speed km/h -> m/s as float64
    vx = speed * (1.0 / 3.6)
    if vx.dtype != np.float64:
        vx = vx.astype(np.float64)
    time_f = (time_arr / np.timedelta64(1, "s")).astype(np.float64)

    # Ensure float64 only when needed
    x_f = x if x.dtype == np.float64 else x.astype(np.float64)
    y_f = y if y.dtype == np.float64 else y.astype(np.float64)
    z_f = z if z.dtype == np.float64 else z.astype(np.float64)
    dist_f = dist if dist.dtype == np.float64 else dist.astype(np.float64)

    # --- X acceleration ---
    dtime = np.gradient(time_f)
    ax = np.gradient(vx) / dtime
    _smooth_outliers(ax, 25.0, use_abs=False)
    ax = np.convolve(ax, _KERNEL_3, mode="same")

    # --- Shared gradient for Y and Z ---
    dx = np.gradient(x_f)
    ds = np.gradient(dist_f)

    # --- Y acceleration ---
    dy = np.gradient(y_f)
    theta = np.arctan2(dy, dx + EPS)
    theta[0] = theta[1]
    dtheta = np.gradient(np.unwrap(theta))
    _smooth_outliers(dtheta, 0.5, use_abs=True)
    C = dtheta / (ds + 0.0001)
    ay = np.square(vx) * C
    ay[np.abs(ay) > 150] = 0
    ay = np.convolve(ay, _KERNEL_9, mode="same")

    # --- Z acceleration ---
    dz = np.gradient(z_f)
    z_theta = np.arctan2(dz, dx + EPS)
    z_theta[0] = z_theta[1]
    z_dtheta = np.gradient(np.unwrap(z_theta))
    _smooth_outliers(z_dtheta, 0.5, use_abs=True)
    z_C = z_dtheta / (ds + 0.0001)
    az = np.square(vx) * z_C
    az[np.abs(az) > 150] = 0
    az = np.convolve(az, _KERNEL_9, mode="same")

    return ax, ay, az, time_f


def _process_telemetry_to_dict(telemetry: pd.DataFrame, data_key: str) -> dict:
    time_arr = telemetry["Time"].to_numpy()
    speed = telemetry["Speed"].to_numpy()
    x = telemetry["X"].to_numpy()
    y = telemetry["Y"].to_numpy()
    z = telemetry["Z"].to_numpy()
    dist = telemetry["Distance"].to_numpy()

    ax, ay, az, time_s = _compute_accelerations(speed, time_arr, x, y, z, dist)

    drs_raw = telemetry["DRS"].to_numpy()
    drs = np.isin(drs_raw, [10, 12, 14]).astype(np.int8)
    brake = telemetry["Brake"].to_numpy().astype(bool).astype(np.int8)
    driver_ahead = (
        telemetry["DriverAhead"]
        if "DriverAhead" in telemetry.columns
        else np.full(len(telemetry), np.nan)
    )
    distance_to_driver_ahead = (
        telemetry["DistanceToDriverAhead"].to_numpy()
        if "DistanceToDriverAhead" in telemetry.columns
        else np.full(len(telemetry), np.nan, dtype=np.float64)
    )

    return {
        "tel": {
            "time": _array_to_list_float_or_none(time_s),
            "rpm": _array_to_list_float_or_none(telemetry["RPM"].to_numpy()),
            "speed": _array_to_list_float_or_none(speed),
            "gear": _array_to_list_int_or_none(telemetry["nGear"].to_numpy()),
            "throttle": _array_to_list_float_or_none(telemetry["Throttle"].to_numpy()),
            "brake": _array_to_list_int_or_none(brake),
            "drs": _array_to_list_int_or_none(drs),
            "distance": _array_to_list_float_or_none(dist),
            "rel_distance": _array_to_list_float_or_none(
                telemetry["RelativeDistance"].to_numpy()
                if "RelativeDistance" in telemetry.columns
                else np.full(len(telemetry), np.nan, dtype=np.float64)
            ),
            "DriverAhead": _col_to_list_str_or_none(driver_ahead),
            "DistanceToDriverAhead": _array_to_list_float_or_none(
                distance_to_driver_ahead
            ),
            "acc_x": _array_to_list_float_or_none(ax),
            "acc_y": _array_to_list_float_or_none(ay),
            "acc_z": _array_to_list_float_or_none(az),
            "x": _array_to_list_float_or_none(x),
            "y": _array_to_list_float_or_none(y),
            "z": _array_to_list_float_or_none(z),
            "dataKey": data_key,
        }
    }


def _lap_telemetry_or_none(selected) -> Optional[pd.DataFrame]:
    """Build lap telemetry, or None when car/pos data cannot be merged.
    Mirrors FastF1 ``get_telemetry()`` but bails out before the expensive
    ``add_driver_ahead`` / ``merge_channels`` path when pos or car slices are
    empty (the common Monaco KeyError: Date case).
    """
    try:
        pos_data = selected.get_pos_data(pad=1, pad_side="both")
        if pos_data is None or pos_data.empty or "Date" not in pos_data.columns:
            return None

        car_data = selected.get_car_data(pad=1, pad_side="both")
        if (
            car_data is None
            or car_data.empty
            or "Date" not in car_data.columns
            or len(car_data) < 3
        ):
            return None

        # Same pipeline as fastf1.core.Laps.get_telemetry / Lap.get_telemetry.
        drv_ahead = (
            car_data.iloc[1:-1]
            .add_driver_ahead()
            .loc[
                :,
                (
                    "DriverAhead",
                    "DistanceToDriverAhead",
                    "Date",
                    "Time",
                    "SessionTime",
                ),
            ]
        )
        car_data = car_data.add_distance().add_relative_distance()
        car_data = car_data.merge_channels(drv_ahead, frequency=None)
        merged = pos_data.merge_channels(car_data, frequency=None)
        return merged.slice_by_lap(selected, interpolate_edges=True)
    except Exception:
        return None


def _lap_telemetry_from_merged(merged, selected):
    """Slice pre-merged telemetry for one lap (fast path, no re-merge)."""
    try:
        if merged is None or selected.empty:
            return None
        return merged.slice_by_lap(selected, interpolate_edges=True)
    except Exception:
        return None


def _driver_merged_telemetry(driver_laps):
    """Build the interpolated merged car+pos telemetry once for a driver.
    Mirrors FastF1 ``Laps.get_telemetry`` merge but returns the merged object
    once so that per-lap slicing avoids repeating the expensive merge.
    Returns None when pos/car data cannot be merged.
    """
    try:
        pos_data = driver_laps.get_pos_data(pad=1, pad_side="both")
        if pos_data is None or pos_data.empty or "Date" not in pos_data.columns:
            return None
        car_data = driver_laps.get_car_data(pad=1, pad_side="both")
        if (
            car_data is None
            or car_data.empty
            or "Date" not in car_data.columns
            or len(car_data) < 3
        ):
            return None
        drv_ahead = (
            car_data.iloc[1:-1]
            .add_driver_ahead()
            .loc[:, ("DriverAhead", "DistanceToDriverAhead", "Date", "Time", "SessionTime")]
        )
        car_data = car_data.add_distance().add_relative_distance()
        car_data = car_data.merge_channels(drv_ahead, frequency=None)
        return pos_data.merge_channels(car_data, frequency=None)
    except Exception:
        return None


def check_memory_usage(threshold_percent=80, session_cache=None, circuit_cache=None):
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    memory_percent = process.memory_percent()

    logger.info(
        f"Current memory usage: {memory_percent:.2f}% "
        f"({memory_info.rss / 1024 / 1024:.2f} MB)"
    )

    if memory_percent > threshold_percent:
        logger.warning(
            f"Memory usage exceeds {threshold_percent}% threshold, clearing caches"
        )
        if session_cache is not None:
            session_cache.clear()
        if circuit_cache is not None:
            circuit_cache.clear()
        gc.collect()

        new_pct = psutil.Process(os.getpid()).memory_percent()
        logger.info(f"New memory usage after clearing caches: {new_pct:.2f}%")
        return True

    return False



# ---------------------------------------------------------------------------
# Process-pool workers (one session load per worker, then many drivers)
# ---------------------------------------------------------------------------
_WORKER: Dict[str, object] = {}


def _prefer_fork_workers() -> bool:
    """True when workers can inherit the parent's loaded session via fork."""
    return "fork" in mp.get_all_start_methods()


def _bind_worker_session(
    extractor: "SeasonSessionExtractor",
    f1session: fastf1.core.Session,
    session_weather_df,
    event_name: str,
    session_name: str,
) -> None:
    """Publish session state for fork-inherited workers (and tests)."""
    _WORKER.clear()
    _WORKER.update(
        {
            "extractor": extractor,
            "f1session": f1session,
            "session_weather_df": session_weather_df,
            "event_name": event_name,
            "session_name": session_name,
        }
    )


def _init_driver_worker(
    year: int,
    cache_dir: str,
    event_name: str,
    session_name: str,
) -> None:
    """Spawn-only: load session once per worker; reused across driver tasks."""
    os.makedirs(cache_dir, exist_ok=True)
    fastf1.Cache.enable_cache(cache_dir)
    logging.getLogger("fastf1").setLevel(logging.WARNING)

    extractor = SeasonSessionExtractor(year=year)
    f1session = extractor.get_session(event_name, session_name, load_telemetry=True)

    session_weather_df = None
    if hasattr(f1session.laps, "get_weather_data"):
        try:
            session_weather_df = f1session.laps.get_weather_data()
        except Exception:
            session_weather_df = None

    _bind_worker_session(
        extractor, f1session, session_weather_df, event_name, session_name
    )


def _process_driver_in_worker(
    driver: str, base_dir: str
) -> Tuple[str, Optional[Dict[str, list]]]:
    extractor: SeasonSessionExtractor = _WORKER["extractor"]  # type: ignore[assignment]
    laptimes = extractor.process_driver(
        _WORKER["event_name"],  # type: ignore[arg-type]
        _WORKER["session_name"],  # type: ignore[arg-type]
        driver,
        base_dir,
        _WORKER["f1session"],  # type: ignore[arg-type]
        _WORKER["session_weather_df"],  # type: ignore[arg-type]
    )
    return driver, laptimes


# ---------------------------------------------------------------------------
# Season Session Extractor
# ---------------------------------------------------------------------------
class SeasonSessionExtractor:
    """Extract telemetry from non-testing season sessions."""

    def __init__(self, year: int = DEFAULT_YEAR):
        self.year = year
        self._session_cache: Dict[str, fastf1.core.Session] = {}
        self._circuit_cache: Dict[str, dict] = {}

    @staticmethod
    def _fix_missing_retired_on_track_laps(f1session: fastf1.core.Session) -> int:
        """Add partial last laps for cars that stopped on track.
        FastF1's ``_fix_missing_laps_retired_on_track`` skips generation when
        *any* prior session-status is ``Finished``. That is wrong for
        multi-segment sessions (Qualifying / Sprint Qualifying): Q1/Q2 already
        emit ``Finished`` before a Q3 crash, so the crash lap is dropped
        (e.g. VER lap 11, 2026 Austrian GP Qualifying).
        Mirror FastF1's generation logic, but treat the driver as finished only
        when the *latest* status at/before the last lap end is ``Finished``.
        """
        if not hasattr(f1session, "_laps") or f1session._laps is None:
            return 0
        if not hasattr(f1session, "session_status") or f1session.session_status is None:
            return 0

        laps = f1session._laps
        if laps.empty or "DriverNumber" not in laps.columns:
            return 0

        session_status = f1session.session_status
        car_data = getattr(f1session, "_car_data", None) or {}
        total_laps = getattr(f1session, "_total_laps", None)
        added = 0
        new_rows: List[pd.DataFrame] = []

        for drv in laps["DriverNumber"].unique():
            drv_laps = laps[laps["DriverNumber"] == drv]
            if drv_laps.empty:
                continue

            if len(drv_laps) == 1 and bool(drv_laps["FastF1Generated"].iloc[0]):
                continue

            ref_time = drv_laps["Time"].iloc[-1]
            if pd.isna(ref_time):
                ref_time = drv_laps["LapStartTime"].iloc[-1]
            if pd.isna(ref_time):
                continue

            next_statuses = session_status[session_status["Time"] > ref_time]
            prev_statuses = session_status[session_status["Time"] <= ref_time]

            if total_laps is not None and drv_laps["LapNumber"].max() == total_laps:
                continue

            # FastF1 bug: used .any() on Finished, which fires after Q1/Q2.
            if (
                not prev_statuses.empty
                and prev_statuses["Status"].iloc[-1] == "Finished"
            ):
                continue

            if next_statuses.empty or not (next_statuses["Status"] == "Finished").any():
                continue

            if not pd.isna(drv_laps["PitInTime"].iloc[-1]):
                continue

            if total_laps is not None and drv_laps.shape[0] >= total_laps:
                continue

            if (
                len(drv_laps) >= 2
                and not pd.isna(drv_laps["PitInTime"].iloc[-2])
                and pd.isna(drv_laps["PitOutTime"].iloc[-1])
            ):
                continue

            next_status = next_statuses.iloc[0]
            if next_status["Status"] == "Aborted":
                assumed_end_time = next_status["Time"]
            else:
                assumed_end_time = None
                if drv in car_data:
                    try:
                        assumed_end_time = car_data[drv].loc[
                            (car_data[drv]["SessionTime"] > ref_time)
                            & (car_data[drv]["Speed"] == 0.0)
                        ].iloc[0]["SessionTime"]
                    except (IndexError, KeyError):
                        assumed_end_time = None
                if assumed_end_time is None:
                    assumed_end_time = ref_time + pd.Timedelta(150, "sec")

            new_last = pd.DataFrame(
                {
                    "LapStartTime": [drv_laps["Time"].iloc[-1]],
                    "Time": [assumed_end_time],
                    "Driver": [drv_laps["Driver"].iloc[-1]],
                    "DriverNumber": [drv_laps["DriverNumber"].iloc[-1]],
                    "Team": [drv_laps["Team"].iloc[-1]],
                    "LapNumber": [drv_laps["LapNumber"].iloc[-1] + 1],
                    "Stint": [drv_laps["Stint"].iloc[-1]],
                    "Compound": [drv_laps["Compound"].iloc[-1]],
                    "TyreLife": [drv_laps["TyreLife"].iloc[-1] + 1],
                    "FreshTyre": [drv_laps["FreshTyre"].iloc[-1]],
                    "Position": [np.nan],
                    "FastF1Generated": [True],
                    "IsAccurate": [False],
                }
            )
            if hasattr(f1session, "_add_track_status_to_laps"):
                f1session._add_track_status_to_laps(new_last)
            new_rows.append(new_last)
            added += 1
            logger.info(
                "Added retired-on-track lap %s for driver %s "
                "(FastF1 skipped it due to multi-segment Finished status)",
                int(new_last["LapNumber"].iloc[0]),
                drv_laps["Driver"].iloc[-1],
            )

        if not new_rows:
            return 0

        f1session._laps = (
            pd.concat([f1session._laps, *new_rows], ignore_index=True)
            .__finalize__(f1session._laps)
            .sort_values(by=["DriverNumber", "LapNumber"])
            .reset_index(drop=True)
        )
        return added

    def get_session(
        self, event_name: str, session_name: str, load_telemetry: bool = True
    ) -> fastf1.core.Session:
        cache_key = f"{self.year}-{event_name}-{session_name}"
        cached = self._session_cache.get(cache_key)

        if cached is not None:
            if load_telemetry and not getattr(cached, "_telemetry_loaded", False):
                cached.load(telemetry=True, weather=True, messages=True)
                cached._telemetry_loaded = True
                self._session_cache[cache_key] = cached
            # Idempotent: safe on cache hits; needed after late telemetry load.
            self._fix_missing_retired_on_track_laps(cached)
            return cached

        f1session = fastf1.get_session(self.year, event_name, session_name)
        f1session.load(telemetry=load_telemetry, weather=True, messages=True)
        f1session._telemetry_loaded = load_telemetry
        self._fix_missing_retired_on_track_laps(f1session)
        self._session_cache[cache_key] = f1session
        return f1session

    def session_drivers(
        self, event_name: str, session_name: str, f1session: fastf1.core.Session = None
    ) -> Dict[str, List[Dict[str, str]]]:
        try:
            if f1session is None:
                f1session = self.get_session(event_name, session_name)
            laps = f1session.laps
            driver_cols = ["Driver", "Team"]
            has_driver_number = "DriverNumber" in laps.columns
            if has_driver_number:
                driver_cols.append("DriverNumber")
            driver_team = laps.drop_duplicates(subset="Driver")[driver_cols]

            results = f1session.results
            result_by_abbr = {}
            result_by_number = {}
            if results is not None and not results.empty:
                for row in results.itertuples():
                    abbr = getattr(row, "Abbreviation", None)
                    if pd.notna(abbr):
                        result_by_abbr[str(abbr)] = row

                    driver_number = getattr(row, "DriverNumber", None)
                    if pd.notna(driver_number):
                        result_by_number[str(driver_number)] = row

            drivers = [
                self._build_driver_info(
                    row=row,
                    has_driver_number=has_driver_number,
                    result_by_abbr=result_by_abbr,
                    result_by_number=result_by_number,
                )
                for row in driver_team.itertuples(index=False)
            ]
            return {"drivers": drivers}
        except Exception as e:
            logger.error(
                f"Error getting drivers for {event_name} {session_name}: {e}"
            )
            return {"drivers": []}

    def _build_driver_info(
        self,
        row,
        has_driver_number: bool,
        result_by_abbr: Dict[str, object],
        result_by_number: Dict[str, object],
    ) -> Dict[str, str]:
        driver = _scalar_to_json_primitive_or_none(row.Driver)
        team = _scalar_to_json_primitive_or_none(row.Team)
        driver_number = _scalar_to_json_primitive_or_none(
            row.DriverNumber if has_driver_number else "None"
        )

        result_row = result_by_abbr.get(str(driver))
        if result_row is None and driver_number != "None":
            result_row = result_by_number.get(str(driver_number))

        first_name = "None"
        last_name = "None"
        team_color = "None"
        headshot_url = "None"

        if result_row is not None:
            first_name = _scalar_to_json_primitive_or_none(
                getattr(result_row, "FirstName", "None")
            )
            last_name = _scalar_to_json_primitive_or_none(
                getattr(result_row, "LastName", "None")
            )
            team_color = _scalar_to_json_primitive_or_none(
                getattr(result_row, "TeamColor", "None")
            )
            headshot_url = _scalar_to_json_primitive_or_none(
                getattr(
                    result_row,
                    "HeadshotUrl",
                    getattr(result_row, "HeadShotUrl", "None"),
                )
            )

        return {
            "driver": driver,
            "team": team,
            "dn": driver_number,
            "fn": first_name,
            "ln": last_name,
            "tc": team_color,
            "url": headshot_url,
        }

    def laps_data(
        self,
        driver: str,
        f1session: fastf1.core.Session,
        driver_laps: pd.DataFrame = None,
        session_weather_df: pd.DataFrame = None,
    ) -> Dict[str, list]:
        try:
            if driver_laps is None:
                driver_laps = f1session.laps.pick_drivers(driver)

            lap_weather = _lap_weather_to_column_lists(driver_laps, session_weather_df)

            return {
                "time": _td_col_to_seconds(driver_laps["LapTime"]),
                "lap": _col_to_list_int_or_none(driver_laps["LapNumber"]),
                "compound": _col_to_list_str_or_none(driver_laps["Compound"]),
                "stint": _col_to_list_int_or_none(driver_laps["Stint"]),
                "s1": _td_col_to_seconds(driver_laps["Sector1Time"]),
                "s2": _td_col_to_seconds(driver_laps["Sector2Time"]),
                "s3": _td_col_to_seconds(driver_laps["Sector3Time"]),
                "life": _col_to_list_int_or_none(driver_laps["TyreLife"]),
                "pos": _col_to_list_int_or_none(driver_laps["Position"]),
                "status": _col_to_list_str_or_none(driver_laps["TrackStatus"]),
                "pb": _col_to_list_bool_or_none(driver_laps["IsPersonalBest"]),
                "sesT": _td_col_to_seconds(driver_laps["Time"]),
                "drv": _col_to_list_str_or_none(driver_laps["Driver"]),
                "dNum": _col_to_list_str_or_none(driver_laps["DriverNumber"]),
                "pout": _td_col_to_seconds(driver_laps["PitOutTime"]),
                "pin": _td_col_to_seconds(driver_laps["PitInTime"]),
                "s1T": _td_col_to_seconds(driver_laps["Sector1SessionTime"]),
                "s2T": _td_col_to_seconds(driver_laps["Sector2SessionTime"]),
                "s3T": _td_col_to_seconds(driver_laps["Sector3SessionTime"]),
                "vi1": _array_to_list_float_or_none(driver_laps["SpeedI1"].to_numpy()),
                "vi2": _array_to_list_float_or_none(driver_laps["SpeedI2"].to_numpy()),
                "vfl": _array_to_list_float_or_none(driver_laps["SpeedFL"].to_numpy()),
                "vst": _array_to_list_float_or_none(driver_laps["SpeedST"].to_numpy()),
                "fresh": _col_to_list_bool_or_none(driver_laps["FreshTyre"]),
                "team": _col_to_list_str_or_none(driver_laps["Team"]),
                "lST": _td_col_to_seconds(driver_laps["LapStartTime"]),
                "lSD": _col_to_list_str_or_none(driver_laps["LapStartDate"]),
                "del": _col_to_list_bool_or_none(driver_laps["Deleted"]),
                "delR": _col_to_list_str_or_none(driver_laps["DeletedReason"]),
                "ff1G": _col_to_list_bool_or_none(driver_laps["FastF1Generated"]),
                "iacc": _col_to_list_bool_or_none(driver_laps["IsAccurate"]),
                **lap_weather,
            }
        except Exception as e:
            logger.error(f"Error getting lap data for {driver}: {e}")
            return {k: [] for k in _LAPTIMES_KEYS}

    def get_circuit_info(
        self, event_name: str, session_name: str
    ) -> Optional[Dict]:
        cache_key = f"{self.year}-{event_name}-{session_name}"
        if cache_key in self._circuit_cache:
            return self._circuit_cache[cache_key]

        try:
            f1session = self.get_session(event_name, session_name)
            circuit_key = f1session.session_info["Meeting"]["Circuit"]["Key"]

            try:
                circuit_info = f1session.get_circuit_info()
                corners = circuit_info.corners
                result = {
                    "CornerNumber": _series_to_json_list(corners["Number"]),
                    "X": _series_to_json_list(corners["X"]),
                    "Y": _series_to_json_list(corners["Y"]),
                    "Angle": _series_to_json_list(corners["Angle"]),
                    "Distance": _series_to_json_list(corners["Distance"]),
                    "Rotation": _scalar_to_json_primitive_or_none(circuit_info.rotation),
                }
                self._circuit_cache[cache_key] = result
                return result
            except (AttributeError, KeyError):
                circuit_df, rotation = self._get_circuit_info_from_api(circuit_key)
                if circuit_df is not None:
                    result = {
                        "CornerNumber": _series_to_json_list(circuit_df["Number"]),
                        "X": _series_to_json_list(circuit_df["X"]),
                        "Y": _series_to_json_list(circuit_df["Y"]),
                        "Angle": _series_to_json_list(circuit_df["Angle"]),
                        "Distance": _series_to_json_list(circuit_df["Distance"] / 10),
                        "Rotation": _scalar_to_json_primitive_or_none(rotation),
                    }
                    self._circuit_cache[cache_key] = result
                    return result

            logger.warning(
                f"Could not get corner data for {event_name} {session_name}"
            )
            return None
        except Exception as e:
            logger.error(
                f"Error getting circuit info for {event_name} {session_name}: {e}"
            )
            return None

    def _get_circuit_info_from_api(
        self, circuit_key: int
    ) -> Tuple[Optional[pd.DataFrame], float]:
        url = f"{PROTO}://{HOST}/api/v1/circuits/{circuit_key}/{self.year}"
        try:
            response = requests.get(url, headers=HEADERS)
            if response.status_code != 200:
                logger.debug(f"[{response.status_code}] {response.content.decode()}")
                return None, 0.0

            data = response.json()
            rotation = float(data.get("rotation", 0.0))

            rows = [
                (
                    float(e.get("trackPosition", {}).get("x", 0.0)),
                    float(e.get("trackPosition", {}).get("y", 0.0)),
                    int(e.get("number", 0)),
                    str(e.get("letter", "")),
                    float(e.get("angle", 0.0)),
                    float(e.get("length", 0.0)),
                )
                for e in data["corners"]
            ]

            return (
                pd.DataFrame(
                    rows, columns=["X", "Y", "Number", "Letter", "Angle", "Distance"]
                ),
                rotation,
            )
        except Exception as e:
            logger.error(f"Error fetching circuit data from API: {e}")
            return None, 0.0

    def _process_single_lap(
        self,
        driver: str,
        lap_number: int,
        driver_dir: str,
        driver_laps: pd.DataFrame,
        event_name: str,
        session_name: str,
        merged=None,
    ) -> bool:
        file_path = f"{driver_dir}/{lap_number}_tel.json"
        try:
            selected = driver_laps[driver_laps.LapNumber == lap_number]
            if selected.empty:
                logger.warning(
                    f"No data for {driver} lap {lap_number} in {event_name} {session_name}"
                )
                return False

            telemetry = (
                _lap_telemetry_from_merged(merged, selected)
                if merged is not None
                else _lap_telemetry_or_none(selected)
            )
            if telemetry is None or telemetry.empty:
                return False

            data_key = f"{self.year}-{event_name}-{session_name}-{driver}-{lap_number}"
            tel_data = _process_telemetry_to_dict(telemetry, data_key)
            _write_json(file_path, tel_data)
            return True
        except Exception as e:
            logger.error(f"Error processing lap {lap_number} for {driver}: {e}")
            return False

    @staticmethod
    def _merge_driver_laptimes(
        merged: Optional[Dict[str, list]],
        laptimes: Dict[str, list],
        offset: int,
        total_laps: int,
    ) -> Tuple[Dict[str, list], int]:
        """Copy one driver's column arrays into a preallocated session merge."""
        n = len(laptimes["lap"])
        if merged is None:
            merged = {key: [None] * total_laps for key in _LAPTIMES_KEYS}

        end = offset + n
        for key in _LAPTIMES_KEYS:
            merged[key][offset:end] = laptimes[key]
        return merged, end

    def process_driver(
        self,
        event_name: str,
        session_name: str,
        driver: str,
        base_dir: str,
        f1session: fastf1.core.Session = None,
        session_weather_df: pd.DataFrame = None,
    ) -> Optional[Dict[str, list]]:
        driver_dir = f"{base_dir}/{driver}"
        os.makedirs(driver_dir, exist_ok=True)

        try:
            if f1session is None:
                f1session = self.get_session(
                    event_name, session_name, load_telemetry=True
                )

            driver_laps = f1session.laps.pick_drivers(driver)
            driver_laps = driver_laps.assign(
                LapNumber=driver_laps["LapNumber"].astype(int)
            )

            laptimes = self.laps_data(driver, f1session, driver_laps, session_weather_df)
            _write_json(f"{driver_dir}/laptimes.json", laptimes)

            lap_numbers = driver_laps["LapNumber"].tolist()

            existing = (
                set(os.listdir(driver_dir))
                if os.path.isdir(driver_dir)
                else set()
            )

            merged = _driver_merged_telemetry(driver_laps)
            for lap_number in lap_numbers:
                fname = f"{lap_number}_tel.json"
                if fname in existing:
                    continue
                self._process_single_lap(
                    driver, lap_number, driver_dir, driver_laps,
                    event_name, session_name, merged
                )

            return laptimes

        except Exception as e:
            logger.error(f"Error processing driver {driver}: {e}")
            return None

    def process_event_session(self, event_name: str, session_name: str) -> None:
        label = f"{event_name} - {session_name}"
        logger.info(f"Processing {label}")

        _root = os.environ.get("EXP_ROOT", "").rstrip("/")
        base_dir = (
            f"{_root}/{event_name}/{session_name}"
            if _root
            else f"{event_name}/{session_name}"
        )
        os.makedirs(_root, exist_ok=True) if _root else None
        os.makedirs(base_dir, exist_ok=True)

        try:
            # Warm FastF1 disk cache once in the parent so workers only read.
            f1session = self.get_session(event_name, session_name, load_telemetry=True)

            weather_data = _session_weather_to_column_lists(f1session.weather_data)
            _write_json(f"{base_dir}/weather.json", weather_data)

            session_control_messages = _session_rcm_to_column_lists(
                f1session.race_control_messages
            )
            _write_json(f"{base_dir}/rcm.json", session_control_messages)

            drivers_info = self.session_drivers(event_name, session_name, f1session)
            _write_json(f"{base_dir}/drivers.json", drivers_info)

            corner_info = self.get_circuit_info(event_name, session_name)
            if corner_info:
                _write_json(f"{base_dir}/corners.json", corner_info)

            # Alphabetical order matches session_laptimes concatenation order.
            drivers = sorted(d["driver"] for d in drivers_info.get("drivers", []))

            if not drivers:
                logger.warning(f"No drivers found for {label}")
                return

            # Pre-count laps so session_laptimes can be allocated once.
            driver_set = set(drivers)
            lap_drivers = f1session.laps["Driver"]
            total_laps = int(lap_drivers.isin(driver_set).sum())

            session_weather_df = None
            if hasattr(f1session.laps, "get_weather_data"):
                try:
                    session_weather_df = f1session.laps.get_weather_data()
                except Exception:
                    session_weather_df = None

            workers = max(1, min(MAX_WORKERS, len(drivers)))
            use_fork = _prefer_fork_workers()

            if use_fork:
                # Workers inherit this session via COW — no per-worker reload.
                _bind_worker_session(
                    self,
                    f1session,
                    session_weather_df,
                    event_name,
                    session_name,
                )
                mp_ctx = mp.get_context("fork")
                pool_kwargs = {
                    "max_workers": workers,
                    "mp_context": mp_ctx,
                }
                logger.info(
                    f"{label}: processing {len(drivers)} drivers with "
                    f"{workers} fork workers (shared session)"
                )
            else:
                # Spawn cannot inherit memory; free parent copy then reload
                # once per worker in the initializer.
                cache_key = f"{self.year}-{event_name}-{session_name}"
                self._session_cache.pop(cache_key, None)
                del f1session
                gc.collect()
                pool_kwargs = {
                    "max_workers": workers,
                    "mp_context": mp.get_context("spawn"),
                    "initializer": _init_driver_worker,
                    "initargs": (self.year, CACHE_DIR, event_name, session_name),
                }
                logger.info(
                    f"{label}: processing {len(drivers)} drivers with "
                    f"{workers} spawn workers"
                )

            results: Dict[str, Optional[Dict[str, list]]] = {}
            with ProcessPoolExecutor(**pool_kwargs) as executor:
                futures = {
                    executor.submit(_process_driver_in_worker, driver, base_dir): driver
                    for driver in drivers
                }
                done = 0
                for fut in as_completed(futures):
                    driver = futures[fut]
                    done += 1
                    try:
                        drv, laptimes = fut.result()
                        results[drv] = laptimes
                        logger.info(
                            f"Finished driver {drv} ({done}/{len(drivers)})"
                        )
                    except Exception as e:
                        results[driver] = None
                        logger.error(
                            f"Worker failed for driver {driver} "
                            f"({done}/{len(drivers)}): {e}"
                        )

            session_laptimes: Optional[Dict[str, list]] = None
            offset = 0
            merged_drivers = 0

            # Merge in alphabetical driver order (same as sequential path).
            for driver in drivers:
                laptimes = results.get(driver)
                if laptimes is None:
                    continue
                n = len(laptimes["lap"])
                if n == 0:
                    continue
                session_laptimes, offset = self._merge_driver_laptimes(
                    session_laptimes, laptimes, offset, total_laps
                )
                merged_drivers += 1

            if session_laptimes is None or offset == 0:
                logger.warning(f"No driver laptimes to merge for {label}")
            else:
                # Trim if any drivers were skipped after the pre-count.
                if offset != total_laps:
                    for key in _LAPTIMES_KEYS:
                        del session_laptimes[key][offset:]
                out_path = f"{base_dir}/session_laptimes.json"
                _write_json(out_path, session_laptimes)
                logger.info(
                    f"{label}: wrote {out_path} "
                    f"({merged_drivers} drivers, {offset} total laps)"
                )
        except Exception as e:
            logger.error(f"Error processing {label}: {e}")


    def process_all(self) -> None:
        logger.info(f"Starting season session extraction for {self.year}")
        start_time = time.time()

        if not TARGET_EVENT_NAMES:
            logger.warning("No TARGET_EVENT_NAMES configured — nothing to extract.")
            return

        sessions = [s for s in TARGET_SESSIONS if isinstance(s, str) and s.strip()]
        if not sessions:
            logger.warning("No TARGET_SESSIONS configured — nothing to extract.")
            return

        for event_name in TARGET_EVENT_NAMES:
            logger.info(f"Processing {event_name} ({', '.join(sessions)})")
            for session_name in sessions:
                try:
                    self.process_event_session(event_name, session_name)
                except Exception as e:
                    logger.error(f"Failed {event_name} {session_name}: {e}")
                check_memory_usage(
                    session_cache=self._session_cache,
                    circuit_cache=self._circuit_cache,
                )

        elapsed = time.time() - start_time
        logger.info(f"Season session extraction completed in {elapsed:.2f} seconds")


# ======================================================================
# Data Availability
# ======================================================================




def is_session_data_available(
    year: int,
    events: Optional[List[str]] = None,
    sessions: Optional[List[str]] = None,
) -> bool:
    """Check if data is available for the first specified event/session pair."""
    try:
        if events is None:
            events = list(TARGET_EVENT_NAMES)
        if sessions is None:
            sessions = list(TARGET_SESSIONS)

        if not events or not sessions:
            logger.warning("No events or sessions specified to check")
            return False

        event = events[0]
        session = sessions[0]

        logger.info(f"Checking data availability for {year} {event} {session}...")

        f1session = fastf1.get_session(year, event, session)
        f1session.load(telemetry=False, weather=False, messages=False)

        if f1session.laps.empty:
            logger.info(f"No lap data available yet for {year} {event} {session}")
            return False

        if "Driver" not in f1session.laps.columns:
            logger.info(f"No driver data available yet for {year} {event} {session}")
            return False

        if len(f1session.laps["Driver"].dropna().unique()) == 0:
            logger.info(f"No driver data available yet for {year} {event} {session}")
            return False

        logger.info(f"Data is available for {year} {event} {session}")
        return True

    except Exception as e:
        logger.info(f"Data not yet available: {str(e)}")
        return False


def main():
    os.makedirs("cache", exist_ok=True)
    fastf1.Cache.enable_cache("cache")
    extractor = SeasonSessionExtractor(year=DEFAULT_YEAR)
    extractor.process_all()


if __name__ == "__main__":
    main()
Xet Storage Details

Size:: 51.4 kB
Xet hash:: 75794c42b3b0cc6e74a718af8bca463ffab1f35c2960343877e64ec68a9d687f
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.