Public API Reference

Public / stable. Import the symbols on this page from opensportslib.apis or opensportslib.tools. Other package modules are not compatibility-stable unless a developer page labels them an extension point.

Task wrappers

All task wrappers accept a canonical config path or Config, optional local or Hugging Face weights, and optional remote-server connection parameters. train() uses configured or supplied manifests; infer() returns an in-memory prediction payload; evaluate() returns a metrics dictionary; and save_predictions() is the explicit disk-write operation.

Class	Use	Important inputs and side effects
ClassificationModel	Clip, frame-array, and supported tracking classification.	Creates run artifacts under the configured save directory; may load local/Hub weights.
LocalizationModel	Action spotting/localization for configured video, feature, tracking, and HDF5 paths.	Model family and loader backend determine the execution route.
VQAModel	X-VARS/Qwen visual question answering.	Requires the matching VQA dependency profile; direct inference accepts video_path and question.
BaseTaskModel	Shared abstract wrapper contract.	Provides local/remote request lifecycle; subclass methods implement task execution.

from opensportslib.apis import ClassificationModel

model = ClassificationModel(config="/path/to/classification.yaml")
predictions = model.infer(test_set="/path/to/test.json")
model.save_predictions("/path/to/predictions.json", predictions)

Exact wrapper signatures

Bases: ConfigurationMixin, ABC

Thin shared contract for task-level OpenSportsLib wrappers.

Source code in opensportslib/apis/base_task_model.py

class BaseTaskModel(ConfigurationMixin, ABC):
    """Thin shared contract for task-level OpenSportsLib wrappers."""

    def __init__(
        self,
        config=None,
        weights=None,
        remote: str | None = None,
        remote_timeout: float = 30.0,
        remote_poll_interval: float = 1.0,
        remote_result_timeout: float = 7200.0,
        remote_model_id: str | None = None,
    ):
        self._configure_logging()
        self.remote = remote.rstrip("/") if remote else None
        self.remote_timeout = float(remote_timeout)
        self.remote_poll_interval = float(remote_poll_interval)
        self.remote_result_timeout = float(remote_result_timeout)
        self.remote_model_id = remote_model_id
        self.remote_session_id: str | None = None
        if self.remote_timeout <= 0 or self.remote_poll_interval <= 0 or self.remote_result_timeout <= 0:
            raise ValueError("Remote timeout values must be positive.")

        self._config_editor = copy.deepcopy(config) if isinstance(config, Config) else None
        if self._config_editor is not None:
            weights = weights if weights is not None else self._config_editor.weights
            config = self._config_editor.source
            if self.is_remote:
                self._config_editor.remote_overrides()

        remote_only = config is None and self._config_editor is None and self.is_remote and weights is None
        if remote_only:
            # Remote workers own model configuration and weights. Keep a
            # minimal namespace so shared decorators and remote-only methods
            # can operate without downloading or constructing a local model.
            self.config_path = None
            self.config = dict_to_namespace({})
        elif config is None and self._config_editor is None:
            from huggingface_hub.utils import HFValidationError, validate_repo_id

            if (not isinstance(weights, str) or os.path.exists(expand(weights))
                    or weights.endswith((".pt", ".pth", ".tar"))):
                raise ValueError("config path is required unless weights is a Hugging Face model ID")
            try:
                validate_repo_id(weights)
            except HFValidationError as exc:
                raise ValueError("config path is required unless weights is a Hugging Face model ID") from exc
            try:
                config = resolve_config_path(weights)
            except Exception as exc:
                raise ValueError(
                    f"Could not load OpenSportsLib config.yaml from {weights!r}; "
                    "provide config explicitly or publish a compatible config.yaml."
                ) from exc

        if self._config_editor is not None:
            self.config_path = self._config_editor.source
            self.config = dict_to_namespace(self._config_editor.get_config())
        elif self._config_editor is None and not remote_only:
            self.config_path = resolve_config_path(config)
            self._config_editor = Config.from_file(self.config_path)
            self.config = dict_to_namespace(self._config_editor.get_config())
        self.last_loaded_weights = None
        self.best_checkpoint = None

        if weights is not None and not self.is_remote:
            self.config = fetch_and_merge_config_from_HF(self.config, weights, merge_policy="compatibility")
            self.config = resolve_inference_class_metadata(self.config)
            self.last_loaded_weights = weights
            self.best_checkpoint = weights
        elif weights is not None:
            # The remote worker owns model loading; retain this only as caller metadata.
            self.last_loaded_weights = weights
            self.best_checkpoint = weights

        self.config = self._effective_config(self.config)
        self.train_flag = False  # Flag to indicate whether we're in training mode (affects checkpoint loading behavior)

        data_cfg = getattr(self.config, "DATA", None)
        if data_cfg is not None and hasattr(data_cfg, "data_dir"):
            data_cfg.data_dir = expand(data_cfg.data_dir)
            logging.info(f"Data directory: {data_cfg.data_dir}")

        self.run_id = os.environ.get("RUN_ID") or str(uuid.uuid4())[:8]
        os.environ["RUN_ID"] = self.run_id

        system_cfg = getattr(self.config, "SYSTEM", None)
        if system_cfg is not None:
            system_paths = getattr(system_cfg, "paths", None)
            base_save_dir = expand(
                getattr(system_paths, "save_dir", None)
                or getattr(system_cfg, "save_dir", None)
                or "./checkpoints"
            )
            model_name = get_component_name_by_kind(self.config, "encoder") or "model"
            run_save_dir = os.path.join(base_save_dir, model_name, self.run_id)
            self.save_dir = run_save_dir
            if system_paths is not None:
                system_paths.save_dir = run_save_dir
                if hasattr(system_paths, "work_dir"):
                    system_paths.work_dir = run_save_dir
            else:
                system_cfg.save_dir = run_save_dir
                if hasattr(system_cfg, "work_dir"):
                    system_cfg.work_dir = run_save_dir
            os.makedirs(run_save_dir, exist_ok=True)
        else:
            self.save_dir = expand("./checkpoints")
            os.makedirs(self.save_dir, exist_ok=True)

        logging.info(f"Save directory: {self.save_dir}")

        self.model = None
        self.processor = None
        self.trainer = None
        self.last_remote_failures: list[dict[str, Any]] = []

        if weights is not None and not self.is_remote:
            self.load_weights(weights=weights)

    @property
    def is_remote(self) -> bool:
        """Whether inference requests are sent to an OpenSportsLib server."""

        return self.remote is not None

    def submit_inference(
        self,
        *,
        task_type: str,
        test_set: str,
        model_id: str | None = None,
        task_options: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Upload a complete JSON manifest and its referenced local media as one job."""

        if not self.remote:
            raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
        manifest_path = Path(expand(test_set)).resolve()
        if not manifest_path.is_file():
            raise FileNotFoundError(f"Test manifest not found: {manifest_path}")
        if manifest_path.suffix.lower() != ".json":
            raise ValueError("Remote test-set uploads currently require a JSON manifest.")

        with tempfile.TemporaryDirectory(prefix="opensportslib-remote-") as temp_dir:
            staged_manifest, archive_path = self._stage_manifest_upload(manifest_path, Path(temp_dir))
            logging.info(
                "Submitting remote test-set inference | task=%s manifest=%s manifest_bytes=%d archive=%s archive_bytes=%d",
                task_type,
                manifest_path,
                staged_manifest.stat().st_size,
                archive_path,
                archive_path.stat().st_size,
            )
            response = self._post_multipart(
                "/predict",
                fields={
                    "task_type": task_type,
                    "model_id": model_id or self.remote_model_id or "",
                    "task_options": json.dumps(task_options or {}),
                },
                files={
                    "test_set_file": staged_manifest,
                    "media_archive": archive_path,
                },
            )
            return response

    def submit_per_sample_inference(
        self,
        *,
        task_type: str,
        test_set: str,
        model_id: str | None = None,
        task_options: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Submit one asynchronous remote job for every sample in an OSL manifest."""

        if not self.remote:
            raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
        manifest_path, payload = self._load_remote_manifest(test_set)
        template = copy.deepcopy(payload)
        samples = template.pop("data")
        jobs: list[dict[str, Any]] = []
        submission_errors: list[dict[str, Any]] = []

        for sample_index, sample in enumerate(samples):
            sample_id = str(sample.get("id") or sample_index) if isinstance(sample, dict) else str(sample_index)
            record: dict[str, Any] = {"sample_index": sample_index, "sample_id": sample_id}
            try:
                with tempfile.TemporaryDirectory(prefix="opensportslib-remote-sample-") as temp_dir:
                    sample_payload = copy.deepcopy(template)
                    sample_payload["data"] = [copy.deepcopy(sample)]
                    staged_manifest, archive_path = self._stage_manifest_payload(
                        sample_payload,
                        manifest_path.parent,
                        Path(temp_dir),
                    )
                    response = self._post_multipart(
                        "/predict",
                        fields={
                            "task_type": task_type,
                            "model_id": model_id or self.remote_model_id or "",
                            "task_options": json.dumps(task_options or {}),
                        },
                        files={"test_set_file": staged_manifest, "media_archive": archive_path},
                    )
                record.update({"job_id": response.get("job_id"), "status": response.get("status", "queued")})
                logging.info("Submitted remote sample job | sample_id=%s job_id=%s", sample_id, record["job_id"])
            except Exception as exc:
                record.update({"status": "submission_failed", "error": str(exc)})
                submission_errors.append(dict(record))
                logging.error("Remote sample submission failed | sample_id=%s error=%s", sample_id, exc)
            jobs.append(record)

        return {
            "task_type": task_type,
            "source_manifest": str(manifest_path),
            "remote_mode": "per_sample",
            "manifest_template": template,
            "jobs": jobs,
            "submission_errors": submission_errors,
        }

    def submit_video_inference(
        self,
        *,
        task_type: str,
        video_path: str,
        question: str | None = None,
        session_id: str | None = None,
        model_id: str | None = None,
        task_options: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Submit one video for direct task inference."""

        if not self.remote:
            raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
        source = Path(expand(video_path)).resolve()
        if not source.is_file():
            raise FileNotFoundError(f"Video file not found: {source}")
        logging.info(
            "Submitting remote video inference | task=%s video=%s video_bytes=%d",
            task_type,
            source,
            source.stat().st_size,
        )
        fields = {
            "task_type": task_type,
            "model_id": model_id or self.remote_model_id or "",
            "task_options": json.dumps(task_options or {}),
        }
        if question is not None:
            fields["question"] = question
        if session_id is not None:
            fields["session_id"] = session_id
        response = self._post_multipart("/predict", fields=fields, files={"upload_file": source})
        self._remember_remote_session(response)
        return response

    def submit_session_inference(
        self,
        *,
        task_type: str,
        session_id: str,
        model_id: str | None = None,
        task_options: dict[str, Any] | None = None,
        question: str | None = None,
    ) -> dict[str, Any]:
        if not self.remote:
            raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
        fields = {
            "task_type": task_type,
            "model_id": model_id or self.remote_model_id or "",
            "session_id": session_id,
            "task_options": json.dumps(task_options or {}),
        }
        if question is not None:
            fields["question"] = question
        response = self._post_multipart("/predict", fields=fields, files={})
        self._remember_remote_session(response)
        return response

    def _remember_remote_session(self, response: dict[str, Any]) -> None:
        session_id = response.get("session_id")
        if session_id:
            self.remote_session_id = str(session_id)

    @property
    def last_remote_session_id(self) -> str | None:
        return self.remote_session_id

    def clear_remote_session(self) -> None:
        self.remote_session_id = None

    def get_remote_job(self, job_id: str) -> dict[str, Any]:
        return self._request_json(f"/jobs/{job_id}")

    def get_remote_result(self, job_id: str) -> dict[str, Any]:
        return self._request_json(f"/jobs/{job_id}/result")

    def wait_for_remote_result(self, job_id: str, timeout: float | None = None) -> dict[str, Any]:
        """Poll a submitted remote job until it succeeds, fails, or times out."""

        deadline = time.monotonic() + (self.remote_result_timeout if timeout is None else float(timeout))
        while True:
            status = self.get_remote_job(job_id)
            state = str(status.get("status", "")).lower()
            logging.info("Remote job status | job_id=%s status=%s", job_id, state or "unknown")
            if state == "succeeded":
                return self.get_remote_result(job_id)
            if state == "failed":
                raise RuntimeError(f"Remote inference job `{job_id}` failed: {status.get('error', 'unknown error')}")
            if time.monotonic() >= deadline:
                raise TimeoutError(f"Timed out waiting for remote inference job `{job_id}`.")
            time.sleep(self.remote_poll_interval)

    def wait_for_remote_batch(self, batch: dict[str, Any], timeout: float | None = None) -> dict[str, Any]:
        """Collect a per-sample batch without discarding successful sample predictions."""

        if batch.get("remote_mode") != "per_sample":
            raise ValueError("`wait_for_remote_batch()` requires a `remote_mode=per_sample` submission payload.")
        deadline = time.monotonic() + (self.remote_result_timeout if timeout is None else float(timeout))
        pending = [record for record in batch.get("jobs", []) if record.get("job_id")]
        completed: dict[str, dict[str, Any]] = {}
        failures = list(batch.get("submission_errors") or [])

        while pending:
            next_pending = []
            for record in pending:
                job_id = str(record["job_id"])
                try:
                    status = self.get_remote_job(job_id)
                    state = str(status.get("status", "")).lower()
                    logging.info("Remote batch job status | sample_id=%s job_id=%s status=%s", record.get("sample_id"), job_id, state)
                    if state == "succeeded":
                        completed[job_id] = self.get_remote_result(job_id)
                    elif state == "failed":
                        failures.append({**record, "error": status.get("error", "Remote job failed.")})
                    else:
                        next_pending.append(record)
                except Exception as exc:
                    failures.append({**record, "error": str(exc)})
            pending = next_pending
            if pending:
                if time.monotonic() >= deadline:
                    failures.extend({**record, "error": "Timed out waiting for remote job."} for record in pending)
                    break
                time.sleep(self.remote_poll_interval)

        predictions = copy.deepcopy(batch.get("manifest_template") or {})
        predictions["data"] = []
        for record in batch.get("jobs", []):
            result = completed.get(str(record.get("job_id")))
            if result is None:
                continue
            result_payload = (result.get("result") or {}).get("predictions") or {}
            predictions["data"].extend(result_payload.get("data") or [])
        return {
            "task_type": batch.get("task_type"),
            "source_manifest": batch.get("source_manifest"),
            "predictions": predictions,
            "failures": failures,
            "jobs": batch.get("jobs") or [],
        }

    def _stage_manifest_upload(self, manifest_path: Path, temp_dir: Path) -> tuple[Path, Path]:
        manifest_path, payload = self._load_remote_manifest(manifest_path)
        return self._stage_manifest_payload(payload, manifest_path.parent, temp_dir)

    def _load_remote_manifest(self, test_set: str | Path) -> tuple[Path, dict[str, Any]]:
        manifest_path = Path(expand(str(test_set))).resolve()
        if not manifest_path.is_file():
            raise FileNotFoundError(f"Test manifest not found: {manifest_path}")
        if manifest_path.suffix.lower() != ".json":
            raise ValueError("Remote test-set uploads currently require a JSON manifest.")
        payload = json.loads(manifest_path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict) or not isinstance(payload.get("data"), list):
            raise ValueError("Test manifest must be an OSL JSON object with a `data` list.")
        return manifest_path, payload

    def _stage_manifest_payload(
        self,
        payload: dict[str, Any],
        source_root: Path,
        temp_dir: Path,
    ) -> tuple[Path, Path]:
        temp_dir.mkdir(parents=True, exist_ok=True)
        payload = copy.deepcopy(payload)

        files: dict[Path, str] = {}
        for value_ref in _manifest_media_references(payload):
            raw_path = str(value_ref[0])
            source = Path(raw_path)
            if not source.is_absolute():
                source = source_root / source
            source = source.resolve()
            if not source.is_file():
                raise FileNotFoundError(f"Manifest media file not found: {source}")
            archive_name = files.setdefault(source, f"media/{len(files):05d}_{source.name}")
            value_ref[1](archive_name)

        staged_manifest = temp_dir / "test_set.json"
        staged_manifest.write_text(json.dumps(payload), encoding="utf-8")
        archive_path = temp_dir / "media.zip"
        with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
            for source, archive_name in files.items():
                archive.write(source, archive_name)
        logging.info(
            "Prepared remote test-set upload | manifest=%s media_files=%d archive_bytes=%d",
            staged_manifest,
            len(files),
            archive_path.stat().st_size,
        )
        return staged_manifest, archive_path

    def _request_json(self, endpoint: str) -> dict[str, Any]:
        if not self.remote:
            raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
        request = urlrequest.Request(f"{self.remote}{endpoint}", method="GET")
        return self._open_request(request)

    def _post_multipart(self, endpoint: str, *, fields: dict[str, str], files: dict[str, Path]) -> dict[str, Any]:
        if endpoint == "/predict":
            fields = self._config_request_fields(fields)
        boundary = f"----OpenSportsLib{uuid.uuid4().hex}"
        chunks: list[bytes] = []
        for name, value in fields.items():
            if value == "":
                continue
            chunks.extend((
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="{name}"\r\n\r\n'.encode(),
                value.encode(),
                b"\r\n",
            ))
        for name, path in files.items():
            chunks.extend((
                f"--{boundary}\r\n".encode(),
                f'Content-Disposition: form-data; name="{name}"; filename="{path.name}"\r\n'.encode(),
                b"Content-Type: application/octet-stream\r\n\r\n",
                path.read_bytes(),
                b"\r\n",
            ))
        chunks.append(f"--{boundary}--\r\n".encode())
        request = urlrequest.Request(
            f"{self.remote}{endpoint}",
            data=b"".join(chunks),
            method="POST",
            headers={
                "Content-Type": f"multipart/form-data; boundary={boundary}",
                "Content-Length": str(sum(len(chunk) for chunk in chunks)),
            },
        )
        return self._open_request(request)

    def _open_request(self, request: urlrequest.Request) -> dict[str, Any]:
        started_at = time.monotonic()
        content_length = request.headers.get("Content-length", "0")
        logging.info(
            "Remote server request | method=%s url=%s content_bytes=%s",
            request.get_method(),
            request.full_url,
            content_length,
        )
        try:
            with urlrequest.urlopen(request, timeout=self.remote_timeout) as response:
                body = response.read()
                logging.info(
                    "Remote server response | method=%s url=%s status=%s response_bytes=%d elapsed_s=%.2f",
                    request.get_method(),
                    request.full_url,
                    response.status,
                    len(body),
                    time.monotonic() - started_at,
                )
                return json.loads(body.decode("utf-8"))
        except urlerror.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            logging.error(
                "Remote server error | method=%s url=%s status=%s elapsed_s=%.2f detail=%s",
                request.get_method(),
                request.full_url,
                exc.code,
                time.monotonic() - started_at,
                detail,
            )
            raise RuntimeError(f"Remote server returned HTTP {exc.code}: {detail}") from exc
        except urlerror.URLError as exc:
            logging.error(
                "Remote server connection error | method=%s url=%s elapsed_s=%.2f reason=%s",
                request.get_method(),
                request.full_url,
                time.monotonic() - started_at,
                exc.reason,
            )
            raise ConnectionError(f"Could not reach remote server `{self.remote}`: {exc.reason}") from exc

    @staticmethod
    def _configure_logging() -> None:
        root_logger = logging.getLogger()
        if not root_logger.handlers:
            logging.basicConfig(
                level=logging.INFO,
                format="%(asctime)s | %(levelname)s | %(message)s",
            )
        elif root_logger.level > logging.INFO:
            root_logger.setLevel(logging.INFO)

    @abstractmethod
    def load_weights(
        self,
        weights: str | None = None,
        **kwargs,
    ) -> None:
        raise NotImplementedError

    @abstractmethod
    def train(
        self,
        train_set: str | None = None,
        valid_set: str | None = None,
        weights: str | None = None,
        use_wandb: bool = True,
        **kwargs,
    ) -> str | None:
        raise NotImplementedError

    @abstractmethod
    def infer(
        self,
        test_set: str | None = None,
        weights: str | None = None,
        use_wandb: bool = True,
        **kwargs,
    ) -> dict:
        raise NotImplementedError

    @abstractmethod
    def evaluate(
        self,
        test_set: str | None = None,
        weights: str | None = None,
        predictions: str | dict[str, Any] | None = None,
        use_wandb: bool = True,
        **kwargs,
    ) -> dict | str | None:
        raise NotImplementedError

    def save_predictions(
        self,
        output_path: str,
        predictions: dict,
    ) -> str:
        """Persist in-memory prediction JSON payload to a target file path."""

        dst = expand(output_path)
        os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)

        if not isinstance(predictions, dict):
            raise TypeError(
                f"Unsupported predictions type: {type(predictions).__name__}. "
                "Expected dict."
            )

        with open(dst, "w", encoding="utf-8") as f:
            json.dump(predictions, f)
        return dst

is_remote property

Whether inference requests are sent to an OpenSportsLib server.

submit_inference(*, task_type, test_set, model_id=None, task_options=None)

Upload a complete JSON manifest and its referenced local media as one job.

Source code in opensportslib/apis/base_task_model.py

def submit_inference(
    self,
    *,
    task_type: str,
    test_set: str,
    model_id: str | None = None,
    task_options: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Upload a complete JSON manifest and its referenced local media as one job."""

    if not self.remote:
        raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
    manifest_path = Path(expand(test_set)).resolve()
    if not manifest_path.is_file():
        raise FileNotFoundError(f"Test manifest not found: {manifest_path}")
    if manifest_path.suffix.lower() != ".json":
        raise ValueError("Remote test-set uploads currently require a JSON manifest.")

    with tempfile.TemporaryDirectory(prefix="opensportslib-remote-") as temp_dir:
        staged_manifest, archive_path = self._stage_manifest_upload(manifest_path, Path(temp_dir))
        logging.info(
            "Submitting remote test-set inference | task=%s manifest=%s manifest_bytes=%d archive=%s archive_bytes=%d",
            task_type,
            manifest_path,
            staged_manifest.stat().st_size,
            archive_path,
            archive_path.stat().st_size,
        )
        response = self._post_multipart(
            "/predict",
            fields={
                "task_type": task_type,
                "model_id": model_id or self.remote_model_id or "",
                "task_options": json.dumps(task_options or {}),
            },
            files={
                "test_set_file": staged_manifest,
                "media_archive": archive_path,
            },
        )
        return response

submit_per_sample_inference(*, task_type, test_set, model_id=None, task_options=None)

Submit one asynchronous remote job for every sample in an OSL manifest.

Source code in opensportslib/apis/base_task_model.py

def submit_per_sample_inference(
    self,
    *,
    task_type: str,
    test_set: str,
    model_id: str | None = None,
    task_options: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Submit one asynchronous remote job for every sample in an OSL manifest."""

    if not self.remote:
        raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
    manifest_path, payload = self._load_remote_manifest(test_set)
    template = copy.deepcopy(payload)
    samples = template.pop("data")
    jobs: list[dict[str, Any]] = []
    submission_errors: list[dict[str, Any]] = []

    for sample_index, sample in enumerate(samples):
        sample_id = str(sample.get("id") or sample_index) if isinstance(sample, dict) else str(sample_index)
        record: dict[str, Any] = {"sample_index": sample_index, "sample_id": sample_id}
        try:
            with tempfile.TemporaryDirectory(prefix="opensportslib-remote-sample-") as temp_dir:
                sample_payload = copy.deepcopy(template)
                sample_payload["data"] = [copy.deepcopy(sample)]
                staged_manifest, archive_path = self._stage_manifest_payload(
                    sample_payload,
                    manifest_path.parent,
                    Path(temp_dir),
                )
                response = self._post_multipart(
                    "/predict",
                    fields={
                        "task_type": task_type,
                        "model_id": model_id or self.remote_model_id or "",
                        "task_options": json.dumps(task_options or {}),
                    },
                    files={"test_set_file": staged_manifest, "media_archive": archive_path},
                )
            record.update({"job_id": response.get("job_id"), "status": response.get("status", "queued")})
            logging.info("Submitted remote sample job | sample_id=%s job_id=%s", sample_id, record["job_id"])
        except Exception as exc:
            record.update({"status": "submission_failed", "error": str(exc)})
            submission_errors.append(dict(record))
            logging.error("Remote sample submission failed | sample_id=%s error=%s", sample_id, exc)
        jobs.append(record)

    return {
        "task_type": task_type,
        "source_manifest": str(manifest_path),
        "remote_mode": "per_sample",
        "manifest_template": template,
        "jobs": jobs,
        "submission_errors": submission_errors,
    }

submit_video_inference(*, task_type, video_path, question=None, session_id=None, model_id=None, task_options=None)

Submit one video for direct task inference.

Source code in opensportslib/apis/base_task_model.py

def submit_video_inference(
    self,
    *,
    task_type: str,
    video_path: str,
    question: str | None = None,
    session_id: str | None = None,
    model_id: str | None = None,
    task_options: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Submit one video for direct task inference."""

    if not self.remote:
        raise RuntimeError("Remote inference is not configured. Pass `remote=` to the model constructor.")
    source = Path(expand(video_path)).resolve()
    if not source.is_file():
        raise FileNotFoundError(f"Video file not found: {source}")
    logging.info(
        "Submitting remote video inference | task=%s video=%s video_bytes=%d",
        task_type,
        source,
        source.stat().st_size,
    )
    fields = {
        "task_type": task_type,
        "model_id": model_id or self.remote_model_id or "",
        "task_options": json.dumps(task_options or {}),
    }
    if question is not None:
        fields["question"] = question
    if session_id is not None:
        fields["session_id"] = session_id
    response = self._post_multipart("/predict", fields=fields, files={"upload_file": source})
    self._remember_remote_session(response)
    return response

get_remote_job(job_id)

Source code in opensportslib/apis/base_task_model.py

def get_remote_job(self, job_id: str) -> dict[str, Any]:
    return self._request_json(f"/jobs/{job_id}")

get_remote_result(job_id)

Source code in opensportslib/apis/base_task_model.py

def get_remote_result(self, job_id: str) -> dict[str, Any]:
    return self._request_json(f"/jobs/{job_id}/result")

wait_for_remote_result(job_id, timeout=None)

Poll a submitted remote job until it succeeds, fails, or times out.

Source code in opensportslib/apis/base_task_model.py

def wait_for_remote_result(self, job_id: str, timeout: float | None = None) -> dict[str, Any]:
    """Poll a submitted remote job until it succeeds, fails, or times out."""

    deadline = time.monotonic() + (self.remote_result_timeout if timeout is None else float(timeout))
    while True:
        status = self.get_remote_job(job_id)
        state = str(status.get("status", "")).lower()
        logging.info("Remote job status | job_id=%s status=%s", job_id, state or "unknown")
        if state == "succeeded":
            return self.get_remote_result(job_id)
        if state == "failed":
            raise RuntimeError(f"Remote inference job `{job_id}` failed: {status.get('error', 'unknown error')}")
        if time.monotonic() >= deadline:
            raise TimeoutError(f"Timed out waiting for remote inference job `{job_id}`.")
        time.sleep(self.remote_poll_interval)

wait_for_remote_batch(batch, timeout=None)

Collect a per-sample batch without discarding successful sample predictions.

Source code in opensportslib/apis/base_task_model.py

def wait_for_remote_batch(self, batch: dict[str, Any], timeout: float | None = None) -> dict[str, Any]:
    """Collect a per-sample batch without discarding successful sample predictions."""

    if batch.get("remote_mode") != "per_sample":
        raise ValueError("`wait_for_remote_batch()` requires a `remote_mode=per_sample` submission payload.")
    deadline = time.monotonic() + (self.remote_result_timeout if timeout is None else float(timeout))
    pending = [record for record in batch.get("jobs", []) if record.get("job_id")]
    completed: dict[str, dict[str, Any]] = {}
    failures = list(batch.get("submission_errors") or [])

    while pending:
        next_pending = []
        for record in pending:
            job_id = str(record["job_id"])
            try:
                status = self.get_remote_job(job_id)
                state = str(status.get("status", "")).lower()
                logging.info("Remote batch job status | sample_id=%s job_id=%s status=%s", record.get("sample_id"), job_id, state)
                if state == "succeeded":
                    completed[job_id] = self.get_remote_result(job_id)
                elif state == "failed":
                    failures.append({**record, "error": status.get("error", "Remote job failed.")})
                else:
                    next_pending.append(record)
            except Exception as exc:
                failures.append({**record, "error": str(exc)})
        pending = next_pending
        if pending:
            if time.monotonic() >= deadline:
                failures.extend({**record, "error": "Timed out waiting for remote job."} for record in pending)
                break
            time.sleep(self.remote_poll_interval)

    predictions = copy.deepcopy(batch.get("manifest_template") or {})
    predictions["data"] = []
    for record in batch.get("jobs", []):
        result = completed.get(str(record.get("job_id")))
        if result is None:
            continue
        result_payload = (result.get("result") or {}).get("predictions") or {}
        predictions["data"].extend(result_payload.get("data") or [])
    return {
        "task_type": batch.get("task_type"),
        "source_manifest": batch.get("source_manifest"),
        "predictions": predictions,
        "failures": failures,
        "jobs": batch.get("jobs") or [],
    }

clear_remote_session()

Source code in opensportslib/apis/base_task_model.py

def clear_remote_session(self) -> None:
    self.remote_session_id = None

Bases: BaseTaskModel

Top-level task wrapper for classification.

Source code in opensportslib/apis/classification.py

class ClassificationModel(BaseTaskModel):
    """Top-level task wrapper for classification."""

    def _resolve_split_path(self, split: str, override: str | None = None) -> str:
        if not self.is_remote:
            from opensportslib.datasets.hf_tracking import (
                hf_tracking_source,
                prepare_hf_tracking_split,
            )

            if hf_tracking_source(self.config):
                if override is not None:
                    raise ValueError(
                        "hf_webdataset uses its configured Hub splits; "
                        "train_set, valid_set, and test_set overrides are unsupported."
                    )
                return str(prepare_hf_tracking_split(self.config, split).annotations_path)
        if override is not None:
            return expand(override)

        path = get_split_annotation_path(self.config, split)
        if path:
            return expand(path)

        raise ValueError(
            f"Could not resolve path for split '{split}'. "
            f"Expected DATA.common.splits.{split}.annotation_path."
        )

    # -----------------------------------------------------------------
    # internal DDP worker
    # -----------------------------------------------------------------
    @staticmethod
    def _worker_ddp(
        rank,
        world_size,
        mode,
        config_path,
        config,
        return_queue=None,
        train_set=None,
        valid_set=None,
        test_set=None,
        weights=None,
        use_wandb=False,
    ):
        """Execute one training/inference job on a single process."""
        import torch
        import wandb
        from opensportslib.core.trainer.classification_trainer import Trainer_Classification
        from opensportslib.core.utils.ddp import ddp_cleanup, ddp_setup
        from opensportslib.core.utils.wandb import init_wandb
        from opensportslib.core.utils.seed import set_reproducibility
        from opensportslib.datasets.builder import build_dataset
        from opensportslib.models.builder import build_model

        logging.basicConfig(
            level=logging.INFO,
            format=f"[RANK {rank}] %(asctime)s | %(levelname)s | %(message)s",
            force=True,
        )
        if rank != 0:
            logging.getLogger().setLevel(logging.ERROR)

        if rank == 0:
            init_wandb(
                config_path,
                config,
                run_id=os.environ["RUN_ID"],
                use_wandb=use_wandb,
            )

        if get_system_use_seed(config):
            set_reproducibility(get_system_seed(config))

        is_ddp = world_size > 1
        if is_ddp:
            torch.cuda.set_device(rank)
            ddp_setup(rank, world_size)
            device = torch.device(f"cuda:{rank}")
        else:
            from opensportslib.core.utils.config import select_device

            device = select_device(config.SYSTEM)

        trainer = Trainer_Classification(config)
        trainer.device = device

        if weights:
            model, processor, _, _ = trainer.load(weights)
        else:
            model, processor = build_model(config, device)

        trainer.model = model

        modality = get_data_modality(config)
        use_tracking_collate = _is_tracking_graph_modality(modality)
        logging.info(
            "Worker setup | mode=%s | modality=%s | tracking_collate=%s",
            mode,
            modality,
            use_tracking_collate,
        )

        try:
            if mode == "train":
                train_data = build_dataset(config, train_set, processor, split="train")
                valid_data = build_dataset(config, valid_set, processor, split="valid")
                best_ckpt = trainer.train(
                    model,
                    train_data,
                    valid_data,
                    rank=rank,
                    world_size=world_size,
                )
                if rank == 0 and return_queue is not None:
                    best_ckpt = best_ckpt or getattr(trainer.trainer, "best_checkpoint_path", None)
                    return_queue.put(best_ckpt)

            elif mode == "infer":
                test_data = build_dataset(config, test_set, processor, split="test")
                _ = trainer.infer(
                    test_data,
                    rank=rank,
                    world_size=world_size,
                )
                if rank == 0 and return_queue is not None:
                    if world_size > 1:
                        return_queue.put(getattr(trainer, "predictions_path", None))
                    else:
                        return_queue.put(getattr(trainer, "predictions_payload", None))
        finally:
            if rank == 0 and use_wandb and getattr(wandb, "run", None) is not None:
                try:
                    wandb.finish(quiet=True)
                except Exception:
                    pass
            if is_ddp:
                ddp_cleanup()

    def load_weights(
        self,
        weights: str | None = None,
        **kwargs,
    ) -> None:
        from opensportslib.core.trainer.classification_trainer import Trainer_Classification

        del kwargs
        if weights is None:
            raise ValueError("`weights` must be provided to load_weights().")

        self.trainer = Trainer_Classification(self.config)
        loaded = self.trainer.load(weights)
        self.model = loaded[0]

        if get_component_provider_by_kind(self.config, "encoder") == "huggingface":
            self.processor = loaded[1]

        self.last_loaded_weights = weights
        self.best_checkpoint = weights

    # -----------------------------------------------------------------
    # public training interface
    # -----------------------------------------------------------------

    @config_operation
    def train(
        self,
        train_set=None,
        valid_set=None,
        test_set=None,
        weights=None,
        use_ddp=False,
        use_wandb=True,
        **kwargs,
    ):
        """Run full training and return best checkpoint path."""
        import torch
        import torch.multiprocessing as mp
        from opensportslib.core.utils.config import resolve_config_omega

        del test_set  # retained for API compatibility

        train_set = self._resolve_split_path("train", train_set)
        valid_set = self._resolve_split_path("valid", valid_set)

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        logging.info("Configuration:")
        logging.info(self.config)

        del kwargs

        effective_weights = weights if weights is not None else self.last_loaded_weights

        world_size = torch.cuda.device_count() or get_system_gpu_count(self.config)
        use_ddp = use_ddp and world_size > 1

        ctx = mp.get_context("spawn")
        queue = ctx.SimpleQueue()

        if use_ddp:
            logging.info(f"Launching DDP on {world_size} GPUs")
            mp.spawn(
                ClassificationModel._worker_ddp,
                args=(
                    world_size,
                    "train",
                    self.config_path,
                    self.config,
                    queue,
                    train_set,
                    valid_set,
                    None,
                    effective_weights,
                    use_wandb,
                ),
                nprocs=world_size,
            )
        else:
            logging.info("Single GPU training")
            ClassificationModel._worker_ddp(
                rank=0,
                world_size=1,
                mode="train",
                config_path=self.config_path,
                config=self.config,
                return_queue=queue,
                train_set=train_set,
                valid_set=valid_set,
                weights=effective_weights,
                use_wandb=use_wandb,
            )

        self.best_checkpoint = queue.get()
        self.last_loaded_weights = self.best_checkpoint
        return self.best_checkpoint

    @config_operation
    def infer(
        self,
        test_set=None,
        weights=None,
        use_ddp=False,
        use_wandb=True,
        video_path: str | None = None,
        session_id: str | None = None,
        **kwargs,
    ):
        """Run model inference and return predictions in OSL JSON format."""
        if test_set is not None and video_path is not None:
            raise ValueError("Provide either `test_set` or `video_path`, not both.")
        remote_mode_provided = "remote_mode" in kwargs
        remote_mode = kwargs.pop("remote_mode", "full_test_set")
        if self.is_remote:
            remote_model_id = kwargs.pop("remote_model_id", None)
            remote_task_options = kwargs.pop("remote_task_options", None)
            if kwargs:
                raise TypeError(f"Unsupported remote inference options: {', '.join(kwargs)}")
            active_session_id = (
                session_id
                if session_id is not None
                else (self.remote_session_id if video_path is None else None)
            )
            if video_path is not None:
                if remote_mode != "full_test_set":
                    raise ValueError("`remote_mode=per_sample` requires `test_set`, not direct video input.")
                job = self.submit_video_inference(
                    task_type="classification",
                    video_path=video_path,
                    session_id=active_session_id,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                self.last_remote_failures = []
                return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
            if test_set is None and active_session_id is not None:
                job = self.submit_session_inference(
                    task_type="classification",
                    session_id=active_session_id,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                self.last_remote_failures = []
                return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
            test_set = self._resolve_split_path("test", test_set)
            if remote_mode == "per_sample":
                batch = self.submit_per_sample_inference(
                    task_type="classification",
                    test_set=test_set,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                collected = self.wait_for_remote_batch(batch)
                self.last_remote_failures = collected["failures"]
                if self.last_remote_failures:
                    logging.warning("Remote per-sample inference completed with %d failures.", len(self.last_remote_failures))
                return collected["predictions"]
            if remote_mode != "full_test_set":
                raise ValueError("`remote_mode` must be `full_test_set` or `per_sample`.")
            job = self.submit_inference(
                task_type="classification",
                test_set=test_set,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if remote_mode_provided:
            raise ValueError("`remote_mode` is available only when `remote` is configured.")
        del kwargs

        if video_path is not None:
            from opensportslib.core.utils.direct_video import direct_video_manifest
            from opensportslib.core.utils.config import resolve_config_omega

            manifest_config = self._effective_config(resolve_config_omega(self.config, weights=weights))
            manifest_config = resolve_inference_class_metadata(manifest_config)
            with direct_video_manifest(manifest_config, video_path, "classification") as manifest:
                return self.infer(
                    test_set=manifest,
                    weights=weights,
                    use_ddp=use_ddp,
                    use_wandb=use_wandb,
                )

        import torch
        import torch.multiprocessing as mp
        from opensportslib.core.utils.config import resolve_config_omega

        test_set = self._resolve_split_path("test", test_set)

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        self.config = resolve_inference_class_metadata(self.config)
        logging.info("Configuration:")
        logging.info(self.config)

        effective_weights = weights if weights is not None else self.last_loaded_weights

        world_size = torch.cuda.device_count()
        use_ddp = use_ddp and world_size > 1

        ctx = mp.get_context("spawn")
        queue = ctx.Queue()

        if use_ddp:
            mp.spawn(
                ClassificationModel._worker_ddp,
                args=(
                    world_size,
                    "infer",
                    self.config_path,
                    self.config,
                    queue,
                    None,
                    None,
                    test_set,
                    effective_weights,
                    use_wandb,
                ),
                nprocs=world_size,
            )
        else:
            ClassificationModel._worker_ddp(
                rank=0,
                world_size=1,
                mode="infer",
                config_path=self.config_path,
                config=self.config,
                return_queue=queue,
                test_set=test_set,
                weights=effective_weights,
                use_wandb=use_wandb,
            )

        predictions = queue.get()
        if use_ddp and isinstance(predictions, str):
            with open(predictions, encoding="utf-8") as f:
                predictions = json.load(f)
        return predictions

    @config_operation
    def evaluate(
        self,
        test_set=None,
        weights=None,
        predictions=None,
        use_ddp=False,
        use_wandb=True,
        **kwargs,
    ):
        """Run inference on test set and return evaluation metrics."""
        del kwargs

        from opensportslib.datasets.builder import build_dataset
        from opensportslib.core.trainer.classification_trainer import Trainer_Classification
        from opensportslib.core.utils.config import resolve_config_omega

        test_set = self._resolve_split_path("test", test_set)

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        self.config = resolve_inference_class_metadata(self.config)
        logging.info("Configuration:")
        logging.info(self.config)
        if predictions is None:
            predictions = self.infer(
                test_set=test_set,
                weights=weights,
                use_ddp=use_ddp,
                use_wandb=use_wandb,
            )

        self.trainer = self.trainer or Trainer_Classification(self.config)
        test_data = build_dataset(self.config, test_set, None, split="test")
        metrics = self.trainer.evaluate(
            pred_path=predictions,
            gt_path=test_set,
            class_names=test_data.label_map,
            exclude_labels=test_data.exclude_labels,
        )

        logging.info(f"TEST METRICS : {metrics}")
        return metrics

load_weights(weights=None, **kwargs)

Source code in opensportslib/apis/classification.py

def load_weights(
    self,
    weights: str | None = None,
    **kwargs,
) -> None:
    from opensportslib.core.trainer.classification_trainer import Trainer_Classification

    del kwargs
    if weights is None:
        raise ValueError("`weights` must be provided to load_weights().")

    self.trainer = Trainer_Classification(self.config)
    loaded = self.trainer.load(weights)
    self.model = loaded[0]

    if get_component_provider_by_kind(self.config, "encoder") == "huggingface":
        self.processor = loaded[1]

    self.last_loaded_weights = weights
    self.best_checkpoint = weights

train(train_set=None, valid_set=None, test_set=None, weights=None, use_ddp=False, use_wandb=True, **kwargs)

Run full training and return best checkpoint path.

Source code in opensportslib/apis/classification.py

@config_operation
def train(
    self,
    train_set=None,
    valid_set=None,
    test_set=None,
    weights=None,
    use_ddp=False,
    use_wandb=True,
    **kwargs,
):
    """Run full training and return best checkpoint path."""
    import torch
    import torch.multiprocessing as mp
    from opensportslib.core.utils.config import resolve_config_omega

    del test_set  # retained for API compatibility

    train_set = self._resolve_split_path("train", train_set)
    valid_set = self._resolve_split_path("valid", valid_set)

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    logging.info("Configuration:")
    logging.info(self.config)

    del kwargs

    effective_weights = weights if weights is not None else self.last_loaded_weights

    world_size = torch.cuda.device_count() or get_system_gpu_count(self.config)
    use_ddp = use_ddp and world_size > 1

    ctx = mp.get_context("spawn")
    queue = ctx.SimpleQueue()

    if use_ddp:
        logging.info(f"Launching DDP on {world_size} GPUs")
        mp.spawn(
            ClassificationModel._worker_ddp,
            args=(
                world_size,
                "train",
                self.config_path,
                self.config,
                queue,
                train_set,
                valid_set,
                None,
                effective_weights,
                use_wandb,
            ),
            nprocs=world_size,
        )
    else:
        logging.info("Single GPU training")
        ClassificationModel._worker_ddp(
            rank=0,
            world_size=1,
            mode="train",
            config_path=self.config_path,
            config=self.config,
            return_queue=queue,
            train_set=train_set,
            valid_set=valid_set,
            weights=effective_weights,
            use_wandb=use_wandb,
        )

    self.best_checkpoint = queue.get()
    self.last_loaded_weights = self.best_checkpoint
    return self.best_checkpoint

infer(test_set=None, weights=None, use_ddp=False, use_wandb=True, video_path=None, session_id=None, **kwargs)

Run model inference and return predictions in OSL JSON format.

Source code in opensportslib/apis/classification.py

@config_operation
def infer(
    self,
    test_set=None,
    weights=None,
    use_ddp=False,
    use_wandb=True,
    video_path: str | None = None,
    session_id: str | None = None,
    **kwargs,
):
    """Run model inference and return predictions in OSL JSON format."""
    if test_set is not None and video_path is not None:
        raise ValueError("Provide either `test_set` or `video_path`, not both.")
    remote_mode_provided = "remote_mode" in kwargs
    remote_mode = kwargs.pop("remote_mode", "full_test_set")
    if self.is_remote:
        remote_model_id = kwargs.pop("remote_model_id", None)
        remote_task_options = kwargs.pop("remote_task_options", None)
        if kwargs:
            raise TypeError(f"Unsupported remote inference options: {', '.join(kwargs)}")
        active_session_id = (
            session_id
            if session_id is not None
            else (self.remote_session_id if video_path is None else None)
        )
        if video_path is not None:
            if remote_mode != "full_test_set":
                raise ValueError("`remote_mode=per_sample` requires `test_set`, not direct video input.")
            job = self.submit_video_inference(
                task_type="classification",
                video_path=video_path,
                session_id=active_session_id,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if test_set is None and active_session_id is not None:
            job = self.submit_session_inference(
                task_type="classification",
                session_id=active_session_id,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        test_set = self._resolve_split_path("test", test_set)
        if remote_mode == "per_sample":
            batch = self.submit_per_sample_inference(
                task_type="classification",
                test_set=test_set,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            collected = self.wait_for_remote_batch(batch)
            self.last_remote_failures = collected["failures"]
            if self.last_remote_failures:
                logging.warning("Remote per-sample inference completed with %d failures.", len(self.last_remote_failures))
            return collected["predictions"]
        if remote_mode != "full_test_set":
            raise ValueError("`remote_mode` must be `full_test_set` or `per_sample`.")
        job = self.submit_inference(
            task_type="classification",
            test_set=test_set,
            model_id=remote_model_id,
            task_options=remote_task_options,
        )
        self.last_remote_failures = []
        return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
    if remote_mode_provided:
        raise ValueError("`remote_mode` is available only when `remote` is configured.")
    del kwargs

    if video_path is not None:
        from opensportslib.core.utils.direct_video import direct_video_manifest
        from opensportslib.core.utils.config import resolve_config_omega

        manifest_config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        manifest_config = resolve_inference_class_metadata(manifest_config)
        with direct_video_manifest(manifest_config, video_path, "classification") as manifest:
            return self.infer(
                test_set=manifest,
                weights=weights,
                use_ddp=use_ddp,
                use_wandb=use_wandb,
            )

    import torch
    import torch.multiprocessing as mp
    from opensportslib.core.utils.config import resolve_config_omega

    test_set = self._resolve_split_path("test", test_set)

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    self.config = resolve_inference_class_metadata(self.config)
    logging.info("Configuration:")
    logging.info(self.config)

    effective_weights = weights if weights is not None else self.last_loaded_weights

    world_size = torch.cuda.device_count()
    use_ddp = use_ddp and world_size > 1

    ctx = mp.get_context("spawn")
    queue = ctx.Queue()

    if use_ddp:
        mp.spawn(
            ClassificationModel._worker_ddp,
            args=(
                world_size,
                "infer",
                self.config_path,
                self.config,
                queue,
                None,
                None,
                test_set,
                effective_weights,
                use_wandb,
            ),
            nprocs=world_size,
        )
    else:
        ClassificationModel._worker_ddp(
            rank=0,
            world_size=1,
            mode="infer",
            config_path=self.config_path,
            config=self.config,
            return_queue=queue,
            test_set=test_set,
            weights=effective_weights,
            use_wandb=use_wandb,
        )

    predictions = queue.get()
    if use_ddp and isinstance(predictions, str):
        with open(predictions, encoding="utf-8") as f:
            predictions = json.load(f)
    return predictions

evaluate(test_set=None, weights=None, predictions=None, use_ddp=False, use_wandb=True, **kwargs)

Run inference on test set and return evaluation metrics.

Source code in opensportslib/apis/classification.py

@config_operation
def evaluate(
    self,
    test_set=None,
    weights=None,
    predictions=None,
    use_ddp=False,
    use_wandb=True,
    **kwargs,
):
    """Run inference on test set and return evaluation metrics."""
    del kwargs

    from opensportslib.datasets.builder import build_dataset
    from opensportslib.core.trainer.classification_trainer import Trainer_Classification
    from opensportslib.core.utils.config import resolve_config_omega

    test_set = self._resolve_split_path("test", test_set)

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    self.config = resolve_inference_class_metadata(self.config)
    logging.info("Configuration:")
    logging.info(self.config)
    if predictions is None:
        predictions = self.infer(
            test_set=test_set,
            weights=weights,
            use_ddp=use_ddp,
            use_wandb=use_wandb,
        )

    self.trainer = self.trainer or Trainer_Classification(self.config)
    test_data = build_dataset(self.config, test_set, None, split="test")
    metrics = self.trainer.evaluate(
        pred_path=predictions,
        gt_path=test_set,
        class_names=test_data.label_map,
        exclude_labels=test_data.exclude_labels,
    )

    logging.info(f"TEST METRICS : {metrics}")
    return metrics

Bases: BaseTaskModel

Top-level task wrapper for localization / spotting.

Source code in opensportslib/apis/localization.py

class LocalizationModel(BaseTaskModel):
    """Top-level task wrapper for localization / spotting."""

    _HF_BACKEND_SPLIT_TYPES = {
        "dali": {
            "VideoGameWithOpencv": "VideoGameWithDali",
            "VideoGameWithOpencvVideo": "VideoGameWithDaliVideo",
        },
        "opencv": {
            "VideoGameWithDali": "VideoGameWithOpencv",
            "VideoGameWithDaliVideo": "VideoGameWithOpencvVideo",
        },
    }

    # def __init__(self, config=None, weights=None):
    #     super().__init__(config=config, weights=None)
    #     if weights is not None:
    #         self.last_loaded_weights = weights
    #         self.best_checkpoint = weights

    #     self.train_flag = False  # Flag to indicate whether we're in training mode (affects checkpoint loading behavior)

    def _resolve_split_path(self, split: str, override: str | None = None) -> str:
        if override is not None:
            return expand(override)

        from opensportslib.datasets.hf_json import hf_json_source, prepare_hf_json_split

        if hf_json_source(self.config):
            prepared = prepare_hf_json_split(self.config, split)
            set_split_source_path(self.config, split, str(prepared.source_path))
            if split == "valid":
                set_split_source_path(self.config, "valid_data_frames", str(prepared.source_path))
            return str(prepared.annotation_path)

        path = get_split_annotation_path(self.config, split)
        if path:
            return expand(path)

        raise ValueError(
            f"Could not resolve path for split '{split}'. "
            f"Expected DATA.common.splits.{split}.annotation_path."
        )

    def _set_split_path(self, split: str, value: str) -> str:
        resolved = expand(value)
        set_split_annotation_path(self.config, split, resolved)
        return resolved

    def _gate_multi_gpu_by_device(self, device) -> None:
        """Disable TRAIN.execution.multi_gpu when effective device is CPU."""
        execution = getattr(getattr(self.config, "TRAIN", None), "execution", None)
        if execution is None:
            return

        multi_gpu = bool(getattr(execution, "multi_gpu", False))
        if device.type == "cpu" and multi_gpu:
            execution.multi_gpu = False
            logging.warning(
                "Detected SYSTEM.device=%s; forcing TRAIN.execution.multi_gpu=false for localization runtime.",
                device,
            )

    @staticmethod
    def _device_type(device) -> str:
        device_type = getattr(device, "type", device)
        return str(device_type).split(":", 1)[0].lower()

    @staticmethod
    def _is_hf_repo_weights(weights: str | None) -> bool:
        if not weights:
            return False
        from opensportslib.core.utils.config import is_local_path

        return not is_local_path(weights)

    @staticmethod
    def _iter_split_items(splits):
        if splits is None:
            return []
        if isinstance(splits, dict):
            return list(splits.items())
        return list(vars(splits).items())

    @staticmethod
    def _ensure_namespace(parent, attr: str):
        current = getattr(parent, attr, None)
        if current is None:
            current = SimpleNamespace()
            setattr(parent, attr, current)
        return current

    def _normalize_opencv_dataloader(self, split_name: str, split_cfg) -> None:
        dataloader = getattr(split_cfg, "dataloader", None)
        if dataloader is None:
            dataloader = SimpleNamespace()
            setattr(split_cfg, "dataloader", dataloader)

        if getattr(dataloader, "batch_size", None) is None:
            dataloader.batch_size = 1
        if getattr(dataloader, "shuffle", None) is None:
            dataloader.shuffle = split_name == "train"
        if getattr(dataloader, "num_workers", None) is None:
            dataloader.num_workers = 0
        if getattr(dataloader, "pin_memory", None) is None:
            dataloader.pin_memory = False

    def _adapt_hf_backend_for_device(self, weights: str | None) -> None:
        if not self._is_hf_repo_weights(weights):
            return

        from opensportslib.core.utils.config import select_device

        device = select_device(self.config.SYSTEM)
        backend = "dali" if self._device_type(device) == "cuda" and _dali_available() else "opencv"
        current_backend = get_loader_backend(self.config)
        set_loader_backend(self.config, backend)

        common = getattr(getattr(self.config, "DATA", None), "common", None)
        if common is None:
            return
        splits = self._ensure_namespace(common, "splits")
        remap = self._HF_BACKEND_SPLIT_TYPES[backend]

        for split_name, split_cfg in self._iter_split_items(splits):
            split_type = getattr(split_cfg, "type", None)
            if split_type in remap:
                setattr(split_cfg, "type", remap[split_type])
            if backend == "opencv":
                self._normalize_opencv_dataloader(split_name, split_cfg)

        if current_backend != backend:
            logging.info(
                "HF localization runtime backend override: %s -> %s for weights=%s",
                current_backend,
                backend,
                weights,
            )

    def _configure_test_time_adaptation(self) -> None:
        """Start a fresh adaptation session for one localization inference call."""
        model_cfg = getattr(self.config, "MODEL", None)
        policies = getattr(model_cfg, "policies", None)
        if isinstance(policies, dict):
            adaptation_cfg = policies.get("test_time_adaptation")
        else:
            adaptation_cfg = getattr(policies, "test_time_adaptation", None)

        if adaptation_cfg is None:
            if hasattr(self.model, "configure_test_time_adaptation"):
                self.model.configure_test_time_adaptation(None)
            return

        enabled = (
            adaptation_cfg.get("enabled", False)
            if isinstance(adaptation_cfg, dict)
            else getattr(adaptation_cfg, "enabled", False)
        )
        if enabled and str(get_model_family(self.config)).strip().lower() != "e2e":
            raise ValueError("SpoTTA is currently integrated only for the E2ESpot family.")
        test_split_type = getattr(get_split_cfg(self.config, "test"), "type", None)
        if enabled and test_split_type != "VideoGameWithOpencvVideo":
            raise ValueError(
                "The reproducible SpoTTA E2ESpot recipe requires a "
                "VideoGameWithOpencvVideo test split."
            )
        if enabled:
            set_loader_backend(self.config, "opencv")
            import random

            import numpy as np
            import torch

            seed = get_system_seed(self.config)
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            if torch.cuda.is_available():
                torch.cuda.manual_seed(seed)
                torch.cuda.manual_seed_all(seed)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
        if not hasattr(self.model, "configure_test_time_adaptation"):
            if enabled:
                raise ValueError("This localization model does not support test-time adaptation.")
            return
        self.model.configure_test_time_adaptation(adaptation_cfg)

    def load_weights(
        self,
        weights: str | None = None,
        **kwargs,
    ) -> None:
        from opensportslib.models.builder import build_model
        from opensportslib.core.utils.config import is_local_path, select_device
        from opensportslib.core.utils.checkpoint import (
            load_checkpoint,
            localization_remap,
        )
        from opensportslib.core.optimizer.builder import build_optimizer
        from opensportslib.core.scheduler.builder import build_scheduler
        default_args = kwargs.get("default_args", None)
        del kwargs
        if weights is None:
            raise ValueError("`weights` must be provided to load_weights().")

        self._adapt_hf_backend_for_device(weights)
        device = select_device(self.config.SYSTEM)
        self._gate_multi_gpu_by_device(device)
        if self.model is None:
            self.model = build_model(self.config, device=device)

        if hasattr(self.model, "configure_test_time_adaptation"):
            self.model.configure_test_time_adaptation(None)

        inner_model = getattr(self.model, "_model", None)
        if inner_model is None:
            inner_model = getattr(self.model, "model", self.model)

        if is_local_path(weights):
            set_system_path(
                self.config,
                "work_dir",
                os.path.dirname(os.path.abspath(weights)),
            )

        if default_args is not None:
            logging.info("Building optimizer + scaler for checkpoint restore...")
            optimizer, scaler = build_optimizer(
                inner_model.parameters(),  # or _get_params() if required
                self.config.TRAIN.optimizer
            )

            logging.info("Building scheduler for checkpoint restore...")
            scheduler = build_scheduler(
                optimizer,
                self.config.TRAIN.scheduler,
                default_args
            )
        else:
            optimizer = scheduler = scaler = None

        inner_model, optimizer, scheduler, scaler, epoch, checkpoint = load_checkpoint(
            model=inner_model,
            path=weights,
            optimizer=optimizer,
            scheduler=scheduler,
            scaler=scaler,
            device=device,
            key_remap_fn=localization_remap,
        )

        if hasattr(self.model, "_model"):
            self.model._model = inner_model
        elif hasattr(self.model, "model"):
            self.model.model = inner_model
        else:
            self.model = inner_model

        self.last_loaded_weights = weights
        self.best_checkpoint = weights

        best_epoch = checkpoint.get("best_epoch", 0)

        best_criterion_valid = checkpoint.get(
            "best_criterion_valid",
            0 if get_train_execution(self.config).get("criterion_valid") == "map" else float("inf")
        )
        self._resume_state = {
            "optimizer": optimizer,
            "scheduler": scheduler,
            "scaler": scaler,
            "epoch": epoch if epoch is not None else 0,
            "best_epoch": best_epoch,
            "best_criterion_valid": best_criterion_valid,
        }

    @config_operation
    def train(
        self,
        train_set=None,
        valid_set=None,
        weights=None,
        use_wandb=True,
        **kwargs,
    ):
        from opensportslib.datasets.builder import build_dataset
        from opensportslib.models.builder import build_model
        from opensportslib.core.trainer.localization_trainer import build_trainer
        from opensportslib.core.utils.default_args import (
            get_default_args_train,
            get_default_args_trainer,
        )
        from opensportslib.core.utils.config import resolve_config_omega, select_device
        from opensportslib.core.utils.load_annotations import check_config
        from opensportslib.core.utils.wandb import init_wandb
        import random
        import numpy as np
        import torch

        del kwargs

        if str(get_model_family(self.config)).lower() == "rulebased":
            raise NotImplementedError(
                "RuleBased localization models are inference-only; call infer() instead of train()."
            )

        train_set = self._resolve_split_path("train", train_set)
        valid_set = self._resolve_split_path("valid", valid_set)
        self._set_split_path("train", train_set)
        self._set_split_path("valid", valid_set)
        # E2E validation mAP uses the `valid_data_frames` split; keep it in sync
        # with explicit valid annotation overrides.
        self._set_split_path("valid_data_frames", valid_set)

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        self.config = resolve_inference_class_metadata(self.config)
        effective_weights = weights if weights is not None else self.last_loaded_weights
        self._adapt_hf_backend_for_device(effective_weights)
        check_config(self.config, split="train")
        init_wandb(
            self.config_path,
            self.config,
            run_id=os.environ["RUN_ID"],
            use_wandb=use_wandb,
        )

        logging.info("Configuration:")
        logging.info(self.config)

        def set_seed(seed):
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
            torch.use_deterministic_algorithms(True, warn_only=True)

        set_seed(get_system_seed(self.config))

        start = time.time()

        data_obj_train = build_dataset(self.config, split="train")
        dataset_train = data_obj_train.building_dataset(
            cfg=data_obj_train.cfg,
            gpu=get_system_gpu_count(self.config),
            default_args=data_obj_train.default_args,
        )
        train_loader = data_obj_train.building_dataloader(
            dataset_train,
            cfg=data_obj_train.cfg.dataloader,
            gpu=get_system_gpu_count(self.config),
            dali=(get_loader_backend(self.config) == "dali"),
        )

        data_obj_valid = build_dataset(self.config, split="valid")
        dataset_valid = data_obj_valid.building_dataset(
            cfg=data_obj_valid.cfg,
            gpu=get_system_gpu_count(self.config),
            default_args=data_obj_valid.default_args,
        )
        valid_loader = data_obj_valid.building_dataloader(
            dataset_valid,
            cfg=data_obj_valid.cfg.dataloader,
            gpu=get_system_gpu_count(self.config),
            dali=(get_loader_backend(self.config) == "dali"),
        )

        default_args = get_default_args_trainer(self.config, len(train_loader))

        self.train_flag = True  # Set flag to indicate training mode for checkpoint loading
        if effective_weights is not None:
            if self.model is None or self.last_loaded_weights != effective_weights:
                self.load_weights(weights=effective_weights, default_args=default_args)
        elif self.model is None:
            device = select_device(self.config.SYSTEM)
            self._gate_multi_gpu_by_device(device)
            self.model = build_model(self.config, device=device)

        self.trainer = build_trainer(
            cfg=self.config,
            model=self.model,
            default_args=default_args,
            resume_from=self._resume_state if hasattr(self, "_resume_state") else None,
        )

        logging.info("Start training")

        self.trainer.train(
            **get_default_args_train(
                self.model,
                train_loader,
                valid_loader,
                get_data_classes(self.config),
                get_train_trainer_type(self.config),
            )
        )

        self.best_checkpoint = self.trainer.best_checkpoint_path
        self.last_loaded_weights = self.best_checkpoint

        logging.info(f"Total Execution Time is {time.time()-start} seconds")
        return self.best_checkpoint

    @config_operation
    def infer(
        self,
        test_set=None,
        weights=None,
        use_wandb=True,
        video_path: str | None = None,
        session_id: str | None = None,
        **kwargs,
    ):
        """Run model inference and return predictions in OSL JSON format."""
        if test_set is not None and video_path is not None:
            raise ValueError("Provide either `test_set` or `video_path`, not both.")
        remote_mode_provided = "remote_mode" in kwargs
        remote_mode = kwargs.pop("remote_mode", "full_test_set")
        if self.is_remote:
            remote_model_id = kwargs.pop("remote_model_id", None)
            remote_task_options = kwargs.pop("remote_task_options", None)
            if kwargs:
                raise TypeError(f"Unsupported remote inference options: {', '.join(kwargs)}")
            active_session_id = (
                session_id
                if session_id is not None
                else (self.remote_session_id if video_path is None else None)
            )
            if video_path is not None:
                if remote_mode != "full_test_set":
                    raise ValueError("`remote_mode=per_sample` requires `test_set`, not direct video input.")
                job = self.submit_video_inference(
                    task_type="localization",
                    video_path=video_path,
                    session_id=active_session_id,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                self.last_remote_failures = []
                return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
            if test_set is None and active_session_id is not None:
                job = self.submit_session_inference(
                    task_type="localization",
                    session_id=active_session_id,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                self.last_remote_failures = []
                return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
            test_set = self._resolve_split_path("test", test_set)
            if remote_mode == "per_sample":
                batch = self.submit_per_sample_inference(
                    task_type="localization",
                    test_set=test_set,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                collected = self.wait_for_remote_batch(batch)
                self.last_remote_failures = collected["failures"]
                if self.last_remote_failures:
                    logging.warning("Remote per-sample inference completed with %d failures.", len(self.last_remote_failures))
                return collected["predictions"]
            if remote_mode != "full_test_set":
                raise ValueError("`remote_mode` must be `full_test_set` or `per_sample`.")
            job = self.submit_inference(
                task_type="localization",
                test_set=test_set,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if remote_mode_provided:
            raise ValueError("`remote_mode` is available only when `remote` is configured.")
        if video_path is not None:
            from opensportslib.core.utils.direct_video import direct_video_manifest
            from opensportslib.core.utils.config import resolve_config_omega

            manifest_config = self._effective_config(resolve_config_omega(self.config, weights=weights))
            manifest_config = resolve_inference_class_metadata(manifest_config)
            with direct_video_manifest(manifest_config, video_path, "localization") as manifest:
                return self.infer(
                    test_set=manifest,
                    weights=weights,
                    use_wandb=use_wandb,
                )
        from opensportslib.datasets.builder import build_dataset
        from opensportslib.models.builder import build_model
        from opensportslib.core.trainer.localization_trainer import build_inferer
        from opensportslib.core.utils.config import resolve_config_omega, select_device
        from opensportslib.core.utils.load_annotations import (
            check_config,
            whether_infer_split,
        )
        from opensportslib.core.utils.wandb import init_wandb

        del kwargs

        test_set = self._resolve_split_path("test", test_set)
        self._set_split_path("test", test_set)

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        self.config = resolve_inference_class_metadata(self.config)
        effective_weights = weights if weights is not None else self.last_loaded_weights
        self._adapt_hf_backend_for_device(effective_weights)
        check_config(self.config, split="test")
        self.config = resolve_inference_class_metadata(self.config)
        self.config.infer_split = whether_infer_split(get_split_cfg(self.config, "test"))

        init_wandb(
            self.config_path,
            self.config,
            run_id=os.environ["RUN_ID"],
            use_wandb=use_wandb,
        )

        logging.info("Configuration:")
        logging.info(self.config)

        start = time.time()

        if effective_weights is not None:
            if self.model is None or self.last_loaded_weights != effective_weights:
                self.load_weights(weights=effective_weights)
        elif self.model is None:
            device = select_device(self.config.SYSTEM)
            self._gate_multi_gpu_by_device(device)
            self.model = build_model(self.config, device=device)

        self._configure_test_time_adaptation()

        data_obj_test = build_dataset(self.config, split="test")
        dataset_test = data_obj_test.building_dataset(
            cfg=data_obj_test.cfg,
            gpu=get_system_gpu_count(self.config),
            default_args=data_obj_test.default_args,
        )
        test_loader = data_obj_test.building_dataloader(
            dataset_test,
            cfg=data_obj_test.cfg.dataloader,
            gpu=get_system_gpu_count(self.config),
            dali=(get_loader_backend(self.config) == "dali"),
        )

        inferer = build_inferer(cfg=self.config, model=self.model)
        predictions = inferer.infer(
            cfg=self.config,
            data=dataset_test,
            dataloader=test_loader,
        )

        adaptation_stats = getattr(self.model, "test_time_adaptation_stats", None)
        if adaptation_stats is not None:
            logging.info("SpoTTA inference stats: %s", adaptation_stats)

        logging.info(f"Total Execution Time is {time.time()-start} seconds")
        return predictions

    @config_operation
    def evaluate(
        self,
        test_set=None,
        weights=None,
        predictions=None,
        use_wandb=True,
        **kwargs,
    ):
        from opensportslib.core.trainer.localization_trainer import build_evaluator
        from opensportslib.core.utils.config import resolve_config_omega
        from opensportslib.core.utils.load_annotations import (
            check_config,
            has_localization_events,
            whether_infer_split,
        )
        from opensportslib.core.utils.wandb import init_wandb

        del kwargs

        test_set = self._resolve_split_path("test", test_set)
        self._set_split_path("test", test_set)
        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        self.config = resolve_inference_class_metadata(self.config)
        effective_weights = weights if weights is not None else self.last_loaded_weights
        self._adapt_hf_backend_for_device(effective_weights)
        check_config(self.config, split="test")
        self.config = resolve_inference_class_metadata(self.config)
        self.config.infer_split = whether_infer_split(get_split_cfg(self.config, "test"))

        init_wandb(
            self.config_path,
            self.config,
            run_id=os.environ["RUN_ID"],
            use_wandb=use_wandb,
        )

        if predictions is None:
            predictions = self.infer(
                test_set=test_set,
                weights=weights,
                use_wandb=use_wandb,
            )

        metrics = None

        test_path = get_split_annotation_path(self.config, "test")
        if has_localization_events(test_path):
            logging.info("Ground truth labels detected -> running evaluation")
            evaluator = build_evaluator(cfg=self.config)
            eval_input = (
                getattr(get_split_cfg(self.config, "test"), "results")
                if isinstance(predictions, dict)
                else predictions
            )
            metrics = evaluator.evaluate(
                cfg_testset=get_split_cfg(self.config, "test"),
                json_gz_file=eval_input,
            )
        else:
            logging.info("No labels found in annotation file -> skipping evaluation")

        return metrics

load_weights(weights=None, **kwargs)

Source code in opensportslib/apis/localization.py

def load_weights(
    self,
    weights: str | None = None,
    **kwargs,
) -> None:
    from opensportslib.models.builder import build_model
    from opensportslib.core.utils.config import is_local_path, select_device
    from opensportslib.core.utils.checkpoint import (
        load_checkpoint,
        localization_remap,
    )
    from opensportslib.core.optimizer.builder import build_optimizer
    from opensportslib.core.scheduler.builder import build_scheduler
    default_args = kwargs.get("default_args", None)
    del kwargs
    if weights is None:
        raise ValueError("`weights` must be provided to load_weights().")

    self._adapt_hf_backend_for_device(weights)
    device = select_device(self.config.SYSTEM)
    self._gate_multi_gpu_by_device(device)
    if self.model is None:
        self.model = build_model(self.config, device=device)

    if hasattr(self.model, "configure_test_time_adaptation"):
        self.model.configure_test_time_adaptation(None)

    inner_model = getattr(self.model, "_model", None)
    if inner_model is None:
        inner_model = getattr(self.model, "model", self.model)

    if is_local_path(weights):
        set_system_path(
            self.config,
            "work_dir",
            os.path.dirname(os.path.abspath(weights)),
        )

    if default_args is not None:
        logging.info("Building optimizer + scaler for checkpoint restore...")
        optimizer, scaler = build_optimizer(
            inner_model.parameters(),  # or _get_params() if required
            self.config.TRAIN.optimizer
        )

        logging.info("Building scheduler for checkpoint restore...")
        scheduler = build_scheduler(
            optimizer,
            self.config.TRAIN.scheduler,
            default_args
        )
    else:
        optimizer = scheduler = scaler = None

    inner_model, optimizer, scheduler, scaler, epoch, checkpoint = load_checkpoint(
        model=inner_model,
        path=weights,
        optimizer=optimizer,
        scheduler=scheduler,
        scaler=scaler,
        device=device,
        key_remap_fn=localization_remap,
    )

    if hasattr(self.model, "_model"):
        self.model._model = inner_model
    elif hasattr(self.model, "model"):
        self.model.model = inner_model
    else:
        self.model = inner_model

    self.last_loaded_weights = weights
    self.best_checkpoint = weights

    best_epoch = checkpoint.get("best_epoch", 0)

    best_criterion_valid = checkpoint.get(
        "best_criterion_valid",
        0 if get_train_execution(self.config).get("criterion_valid") == "map" else float("inf")
    )
    self._resume_state = {
        "optimizer": optimizer,
        "scheduler": scheduler,
        "scaler": scaler,
        "epoch": epoch if epoch is not None else 0,
        "best_epoch": best_epoch,
        "best_criterion_valid": best_criterion_valid,
    }

train(train_set=None, valid_set=None, weights=None, use_wandb=True, **kwargs)

Source code in opensportslib/apis/localization.py

@config_operation
def train(
    self,
    train_set=None,
    valid_set=None,
    weights=None,
    use_wandb=True,
    **kwargs,
):
    from opensportslib.datasets.builder import build_dataset
    from opensportslib.models.builder import build_model
    from opensportslib.core.trainer.localization_trainer import build_trainer
    from opensportslib.core.utils.default_args import (
        get_default_args_train,
        get_default_args_trainer,
    )
    from opensportslib.core.utils.config import resolve_config_omega, select_device
    from opensportslib.core.utils.load_annotations import check_config
    from opensportslib.core.utils.wandb import init_wandb
    import random
    import numpy as np
    import torch

    del kwargs

    if str(get_model_family(self.config)).lower() == "rulebased":
        raise NotImplementedError(
            "RuleBased localization models are inference-only; call infer() instead of train()."
        )

    train_set = self._resolve_split_path("train", train_set)
    valid_set = self._resolve_split_path("valid", valid_set)
    self._set_split_path("train", train_set)
    self._set_split_path("valid", valid_set)
    # E2E validation mAP uses the `valid_data_frames` split; keep it in sync
    # with explicit valid annotation overrides.
    self._set_split_path("valid_data_frames", valid_set)

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    self.config = resolve_inference_class_metadata(self.config)
    effective_weights = weights if weights is not None else self.last_loaded_weights
    self._adapt_hf_backend_for_device(effective_weights)
    check_config(self.config, split="train")
    init_wandb(
        self.config_path,
        self.config,
        run_id=os.environ["RUN_ID"],
        use_wandb=use_wandb,
    )

    logging.info("Configuration:")
    logging.info(self.config)

    def set_seed(seed):
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        torch.use_deterministic_algorithms(True, warn_only=True)

    set_seed(get_system_seed(self.config))

    start = time.time()

    data_obj_train = build_dataset(self.config, split="train")
    dataset_train = data_obj_train.building_dataset(
        cfg=data_obj_train.cfg,
        gpu=get_system_gpu_count(self.config),
        default_args=data_obj_train.default_args,
    )
    train_loader = data_obj_train.building_dataloader(
        dataset_train,
        cfg=data_obj_train.cfg.dataloader,
        gpu=get_system_gpu_count(self.config),
        dali=(get_loader_backend(self.config) == "dali"),
    )

    data_obj_valid = build_dataset(self.config, split="valid")
    dataset_valid = data_obj_valid.building_dataset(
        cfg=data_obj_valid.cfg,
        gpu=get_system_gpu_count(self.config),
        default_args=data_obj_valid.default_args,
    )
    valid_loader = data_obj_valid.building_dataloader(
        dataset_valid,
        cfg=data_obj_valid.cfg.dataloader,
        gpu=get_system_gpu_count(self.config),
        dali=(get_loader_backend(self.config) == "dali"),
    )

    default_args = get_default_args_trainer(self.config, len(train_loader))

    self.train_flag = True  # Set flag to indicate training mode for checkpoint loading
    if effective_weights is not None:
        if self.model is None or self.last_loaded_weights != effective_weights:
            self.load_weights(weights=effective_weights, default_args=default_args)
    elif self.model is None:
        device = select_device(self.config.SYSTEM)
        self._gate_multi_gpu_by_device(device)
        self.model = build_model(self.config, device=device)

    self.trainer = build_trainer(
        cfg=self.config,
        model=self.model,
        default_args=default_args,
        resume_from=self._resume_state if hasattr(self, "_resume_state") else None,
    )

    logging.info("Start training")

    self.trainer.train(
        **get_default_args_train(
            self.model,
            train_loader,
            valid_loader,
            get_data_classes(self.config),
            get_train_trainer_type(self.config),
        )
    )

    self.best_checkpoint = self.trainer.best_checkpoint_path
    self.last_loaded_weights = self.best_checkpoint

    logging.info(f"Total Execution Time is {time.time()-start} seconds")
    return self.best_checkpoint

infer(test_set=None, weights=None, use_wandb=True, video_path=None, session_id=None, **kwargs)

Run model inference and return predictions in OSL JSON format.

Source code in opensportslib/apis/localization.py

@config_operation
def infer(
    self,
    test_set=None,
    weights=None,
    use_wandb=True,
    video_path: str | None = None,
    session_id: str | None = None,
    **kwargs,
):
    """Run model inference and return predictions in OSL JSON format."""
    if test_set is not None and video_path is not None:
        raise ValueError("Provide either `test_set` or `video_path`, not both.")
    remote_mode_provided = "remote_mode" in kwargs
    remote_mode = kwargs.pop("remote_mode", "full_test_set")
    if self.is_remote:
        remote_model_id = kwargs.pop("remote_model_id", None)
        remote_task_options = kwargs.pop("remote_task_options", None)
        if kwargs:
            raise TypeError(f"Unsupported remote inference options: {', '.join(kwargs)}")
        active_session_id = (
            session_id
            if session_id is not None
            else (self.remote_session_id if video_path is None else None)
        )
        if video_path is not None:
            if remote_mode != "full_test_set":
                raise ValueError("`remote_mode=per_sample` requires `test_set`, not direct video input.")
            job = self.submit_video_inference(
                task_type="localization",
                video_path=video_path,
                session_id=active_session_id,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if test_set is None and active_session_id is not None:
            job = self.submit_session_inference(
                task_type="localization",
                session_id=active_session_id,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        test_set = self._resolve_split_path("test", test_set)
        if remote_mode == "per_sample":
            batch = self.submit_per_sample_inference(
                task_type="localization",
                test_set=test_set,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            collected = self.wait_for_remote_batch(batch)
            self.last_remote_failures = collected["failures"]
            if self.last_remote_failures:
                logging.warning("Remote per-sample inference completed with %d failures.", len(self.last_remote_failures))
            return collected["predictions"]
        if remote_mode != "full_test_set":
            raise ValueError("`remote_mode` must be `full_test_set` or `per_sample`.")
        job = self.submit_inference(
            task_type="localization",
            test_set=test_set,
            model_id=remote_model_id,
            task_options=remote_task_options,
        )
        self.last_remote_failures = []
        return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
    if remote_mode_provided:
        raise ValueError("`remote_mode` is available only when `remote` is configured.")
    if video_path is not None:
        from opensportslib.core.utils.direct_video import direct_video_manifest
        from opensportslib.core.utils.config import resolve_config_omega

        manifest_config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        manifest_config = resolve_inference_class_metadata(manifest_config)
        with direct_video_manifest(manifest_config, video_path, "localization") as manifest:
            return self.infer(
                test_set=manifest,
                weights=weights,
                use_wandb=use_wandb,
            )
    from opensportslib.datasets.builder import build_dataset
    from opensportslib.models.builder import build_model
    from opensportslib.core.trainer.localization_trainer import build_inferer
    from opensportslib.core.utils.config import resolve_config_omega, select_device
    from opensportslib.core.utils.load_annotations import (
        check_config,
        whether_infer_split,
    )
    from opensportslib.core.utils.wandb import init_wandb

    del kwargs

    test_set = self._resolve_split_path("test", test_set)
    self._set_split_path("test", test_set)

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    self.config = resolve_inference_class_metadata(self.config)
    effective_weights = weights if weights is not None else self.last_loaded_weights
    self._adapt_hf_backend_for_device(effective_weights)
    check_config(self.config, split="test")
    self.config = resolve_inference_class_metadata(self.config)
    self.config.infer_split = whether_infer_split(get_split_cfg(self.config, "test"))

    init_wandb(
        self.config_path,
        self.config,
        run_id=os.environ["RUN_ID"],
        use_wandb=use_wandb,
    )

    logging.info("Configuration:")
    logging.info(self.config)

    start = time.time()

    if effective_weights is not None:
        if self.model is None or self.last_loaded_weights != effective_weights:
            self.load_weights(weights=effective_weights)
    elif self.model is None:
        device = select_device(self.config.SYSTEM)
        self._gate_multi_gpu_by_device(device)
        self.model = build_model(self.config, device=device)

    self._configure_test_time_adaptation()

    data_obj_test = build_dataset(self.config, split="test")
    dataset_test = data_obj_test.building_dataset(
        cfg=data_obj_test.cfg,
        gpu=get_system_gpu_count(self.config),
        default_args=data_obj_test.default_args,
    )
    test_loader = data_obj_test.building_dataloader(
        dataset_test,
        cfg=data_obj_test.cfg.dataloader,
        gpu=get_system_gpu_count(self.config),
        dali=(get_loader_backend(self.config) == "dali"),
    )

    inferer = build_inferer(cfg=self.config, model=self.model)
    predictions = inferer.infer(
        cfg=self.config,
        data=dataset_test,
        dataloader=test_loader,
    )

    adaptation_stats = getattr(self.model, "test_time_adaptation_stats", None)
    if adaptation_stats is not None:
        logging.info("SpoTTA inference stats: %s", adaptation_stats)

    logging.info(f"Total Execution Time is {time.time()-start} seconds")
    return predictions

evaluate(test_set=None, weights=None, predictions=None, use_wandb=True, **kwargs)

Source code in opensportslib/apis/localization.py

@config_operation
def evaluate(
    self,
    test_set=None,
    weights=None,
    predictions=None,
    use_wandb=True,
    **kwargs,
):
    from opensportslib.core.trainer.localization_trainer import build_evaluator
    from opensportslib.core.utils.config import resolve_config_omega
    from opensportslib.core.utils.load_annotations import (
        check_config,
        has_localization_events,
        whether_infer_split,
    )
    from opensportslib.core.utils.wandb import init_wandb

    del kwargs

    test_set = self._resolve_split_path("test", test_set)
    self._set_split_path("test", test_set)
    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    self.config = resolve_inference_class_metadata(self.config)
    effective_weights = weights if weights is not None else self.last_loaded_weights
    self._adapt_hf_backend_for_device(effective_weights)
    check_config(self.config, split="test")
    self.config = resolve_inference_class_metadata(self.config)
    self.config.infer_split = whether_infer_split(get_split_cfg(self.config, "test"))

    init_wandb(
        self.config_path,
        self.config,
        run_id=os.environ["RUN_ID"],
        use_wandb=use_wandb,
    )

    if predictions is None:
        predictions = self.infer(
            test_set=test_set,
            weights=weights,
            use_wandb=use_wandb,
        )

    metrics = None

    test_path = get_split_annotation_path(self.config, "test")
    if has_localization_events(test_path):
        logging.info("Ground truth labels detected -> running evaluation")
        evaluator = build_evaluator(cfg=self.config)
        eval_input = (
            getattr(get_split_cfg(self.config, "test"), "results")
            if isinstance(predictions, dict)
            else predictions
        )
        metrics = evaluator.evaluate(
            cfg_testset=get_split_cfg(self.config, "test"),
            json_gz_file=eval_input,
        )
    else:
        logging.info("No labels found in annotation file -> skipping evaluation")

    return metrics

Bases: BaseTaskModel

Top-level task wrapper for VQA.

Source code in opensportslib/apis/vqa.py

class VQAModel(BaseTaskModel):
    """Top-level task wrapper for VQA."""

    def _init_wandb(self, use_wandb: bool) -> None:
        from opensportslib.core.utils.wandb import init_wandb

        init_wandb(
            self.config_path,
            self.config,
            run_id=os.environ["RUN_ID"],
            use_wandb=use_wandb,
        )

    @staticmethod
    def _worker_ddp(
        rank,
        world_size,
        config_path,
        config,
        return_queue=None,
        train_set=None,
        valid_set=None,
        resume_from_checkpoint=None,
        use_wandb=False,
    ):
        import torch
        from opensportslib.core.trainer.vqa_trainer import Trainer_VQA
        from opensportslib.core.utils.ddp import ddp_cleanup, ddp_setup
        from opensportslib.core.utils.wandb import init_wandb
        from opensportslib.datasets.builder import build_dataset

        logging.basicConfig(
            level=logging.INFO,
            format=f"[RANK {rank}] %(asctime)s | %(levelname)s | %(message)s",
            force=True,
        )
        if rank != 0:
            logging.getLogger().setLevel(logging.ERROR)
        # Keep INFO on all ranks while debugging multi-GPU startup/hangs.
        # This makes DDP failures visible instead of appearing as a silent stall.
        #logging.getLogger().setLevel(logging.INFO)

        is_ddp = world_size > 1
        if is_ddp:
            os.environ["RANK"] = str(rank)
            os.environ["WORLD_SIZE"] = str(world_size)
            os.environ["LOCAL_RANK"] = str(rank)
            if "TORCH_DISTRIBUTED_DEBUG" not in os.environ:
                os.environ["TORCH_DISTRIBUTED_DEBUG"] = "INFO"
            torch.cuda.set_device(rank)
            ddp_setup(rank, world_size)
            logging.info(
                "Initialized VQA DDP worker | rank=%s | world_size=%s | torch_distributed_debug=%s",
                rank,
                world_size,
                os.environ.get("TORCH_DISTRIBUTED_DEBUG"),
            )

        try:
            if rank == 0:
                init_wandb(
                    config_path,
                    config,
                    run_id=os.environ["RUN_ID"],
                    use_wandb=use_wandb,
                )
            train_data = build_dataset(config, train_set, None, split="train")
            valid_data = build_dataset(config, valid_set, None, split="valid")
            trainer = Trainer_VQA(config)
            train_kwargs = {"rank": rank, "world_size": world_size, "use_wandb": use_wandb}
            if resume_from_checkpoint is not None:
                train_kwargs["resume_from_checkpoint"] = resume_from_checkpoint
            ckpt = trainer.train(None, train_data, valid_data, **train_kwargs)
            if rank == 0 and return_queue is not None:
                return_queue.put(ckpt)
        finally:
            if is_ddp:
                ddp_cleanup()

    def _resolve_split_path(self, split: str, override: str | None = None) -> str:
        if override is not None:
            return expand(override)
        path = get_split_annotation_path(self.config, split)
        if not path:
            raise ValueError(
                f"Could not resolve path for split '{split}'. "
                f"Expected DATA.common.splits.{split}.annotation_path."
            )
        return expand(path)

    def load_weights(self, weights: str | None = None, **kwargs) -> None:
        del kwargs
        if weights is None:
            raise ValueError("`weights` must be provided to load_weights().")
        from opensportslib.core.trainer.vqa_trainer import Trainer_VQA

        _set_model_checkpoint_path(self.config, weights)
        self.trainer = Trainer_VQA(self.config)
        self.trainer.load(weights)
        self.last_loaded_weights = weights
        self.best_checkpoint = weights

    @config_operation
    def train(
        self,
        train_set: str | None = None,
        valid_set: str | None = None,
        weights: str | None = None,
        resume_from_checkpoint: str | None = None,
        use_ddp: bool | None = None,
        use_wandb: bool = True,
        **kwargs,
    ) -> str | None:
        del kwargs

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        execution = get_train_execution(self.config)
        backend = str(execution.get("training_backend", "placeholder")).lower()
        vqa_backend = get_vqa_backend(self.config)
        if vqa_backend == "qwen_xvars_infer" and backend != "qwen_xvars_lora":
            raise ValueError(
                "The 'qwen_xvars_infer' backend requires TRAIN.execution.training_backend='qwen_xvars_lora' for train()."
            )
        if vqa_backend == "qwen_vl_native_infer" and backend != "qwen_vl_native_lora":
            raise ValueError(
                "The 'qwen_vl_native_infer' backend requires TRAIN.execution.training_backend='qwen_vl_native_lora' for train()."
            )
        import torch
        import torch.multiprocessing as mp

        train_set = self._resolve_split_path("train", train_set)
        valid_set = self._resolve_split_path("valid", valid_set)
        if backend in {"xvars_videochatgpt_lora", "qwen_xvars_lora", "qwen_vl_native_lora"}:
            requested_gpus = get_system_gpu_count(self.config)
            available_gpus = int(torch.cuda.device_count() or 0)
            if requested_gpus > 0:
                world_size = min(available_gpus, int(requested_gpus)) if available_gpus > 0 else int(requested_gpus)
            else:
                world_size = available_gpus
            effective_world_size = max(int(world_size), 1)
            if use_ddp is None:
                launch_ddp = effective_world_size > 1
            else:
                launch_ddp = bool(use_ddp) and effective_world_size > 1
            if requested_gpus > 0 and available_gpus > 0 and requested_gpus != available_gpus:
                logging.info(
                    "VQA train GPU selection | requested=%s | available=%s | launching=%s",
                    requested_gpus,
                    available_gpus,
                    effective_world_size if launch_ddp else 1,
                )
            logging.info(
                "VQA train launch | mode=%s | requested_gpus=%s | available_gpus=%s | world_size=%s",
                "ddp" if launch_ddp else "single",
                requested_gpus,
                available_gpus,
                effective_world_size if launch_ddp else 1,
            )

            ctx = mp.get_context("spawn")
            queue = ctx.SimpleQueue()
            if launch_ddp:
                mp.spawn(
                    VQAModel._worker_ddp,
                    args=(
                        effective_world_size,
                        self.config_path,
                        self.config,
                        queue,
                        train_set,
                        valid_set,
                        resume_from_checkpoint,
                        use_wandb,
                    ),
                    nprocs=effective_world_size,
                )
            else:
                VQAModel._worker_ddp(
                    rank=0,
                    world_size=1,
                    config_path=self.config_path,
                    config=self.config,
                    return_queue=queue,
                    train_set=train_set,
                    valid_set=valid_set,
                    resume_from_checkpoint=resume_from_checkpoint,
                    use_wandb=use_wandb,
                )
            ckpt = queue.get()
            self.best_checkpoint = ckpt
            self.last_loaded_weights = ckpt
            return ckpt

        raise ValueError(
            f"Unsupported VQA training backend '{backend}'. "
            "Only 'xvars_videochatgpt_lora' and 'qwen_xvars_lora' are supported."
        )

    @config_operation
    def infer(
        self,
        test_set: str | None = None,
        weights: str | None = None,
        use_wandb: bool = True,
        video_path: str | None = None,
        question: str | None = None,
        session_id: str | None = None,
        **kwargs,
    ) -> dict:
        remote_mode_provided = "remote_mode" in kwargs
        remote_mode = kwargs.pop("remote_mode", "full_test_set")
        if self.is_remote:
            remote_model_id = kwargs.pop("remote_model_id", None)
            remote_task_options = kwargs.pop("remote_task_options", None)
            if kwargs:
                raise TypeError(f"Unsupported remote inference options: {', '.join(kwargs)}")
            if test_set is not None:
                if video_path is not None or question is not None or session_id is not None:
                    raise ValueError("Remote VQA manifest inference accepts only `test_set`.")
                if remote_mode == "per_sample":
                    batch = self.submit_per_sample_inference(
                        task_type="vqa",
                        test_set=self._resolve_split_path("test", test_set),
                        model_id=remote_model_id,
                        task_options=remote_task_options,
                    )
                    collected = self.wait_for_remote_batch(batch)
                    self.last_remote_failures = collected["failures"]
                    if self.last_remote_failures:
                        logging.warning("Remote per-sample inference completed with %d failures.", len(self.last_remote_failures))
                    return collected["predictions"]
                if remote_mode != "full_test_set":
                    raise ValueError("`remote_mode` must be `full_test_set` or `per_sample`.")
                job = self.submit_inference(
                    task_type="vqa",
                    test_set=self._resolve_split_path("test", test_set),
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                self.last_remote_failures = []
                return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
            if remote_mode != "full_test_set":
                raise ValueError("`remote_mode=per_sample` requires `test_set`, not direct VQA input.")
            # A stored session is for question-only follow-ups. A new video
            # must start a fresh server session; explicitly supplied session
            # IDs remain caller-controlled and are forwarded as requested.
            active_session_id = (
                session_id
                if session_id is not None
                else (self.remote_session_id if video_path is None else None)
            )
            if active_session_id is not None and video_path is None and question:
                job = self.submit_session_inference(
                    task_type="vqa",
                    session_id=active_session_id,
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                    question=str(question),
                )
                return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
            if not video_path or not str(question or "").strip():
                raise ValueError("Remote direct VQA inference requires `video_path` and a non-empty `question`.")
            job = self.submit_video_inference(
                task_type="vqa",
                video_path=video_path,
                question=str(question),
                session_id=active_session_id,
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if remote_mode_provided:
            raise ValueError("`remote_mode` is available only when `remote` is configured.")
        del kwargs
        from opensportslib.core.trainer.vqa_trainer import Trainer_VQA
        from opensportslib.datasets.builder import build_dataset
        from opensportslib.models.builder import build_model
        from opensportslib.core.utils.config import select_device

        direct_requested = video_path is not None or question is not None
        if direct_requested and test_set is not None:
            raise ValueError("Provide either `test_set` or `video_path`/`question`, not both.")
        if direct_requested and (not video_path or not str(question or "").strip()):
            raise ValueError("Direct VQA inference requires both `video_path` and a non-empty `question`.")

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        backend = get_vqa_backend(self.config)
        effective_weights = weights if weights is not None else self.last_loaded_weights
        _set_model_checkpoint_path(self.config, effective_weights)
        self.trainer = Trainer_VQA(self.config)
        if effective_weights is not None:
            # Validate OpenSportsLib adapter metadata before allocating the base model.
            self.trainer.load(effective_weights)
        resolved_video_path = None
        if direct_requested:
            resolved_video_path = expand(str(video_path))
            if not os.path.isfile(resolved_video_path):
                raise FileNotFoundError(f"Video file not found: {resolved_video_path}")
        device = select_device(self.config.SYSTEM)
        model, _ = build_model(self.config, device)
        if direct_requested:
            test_data = [
                {
                    "id": os.path.splitext(os.path.basename(resolved_video_path))[0],
                    "question": str(question).strip(),
                    "references": [],
                    "video_path": resolved_video_path,
                    "frame_paths": [],
                    "video_frames": [],
                    "video_spatio_temporal_features": None,
                    "prior_prediction_text": "",
                    "labels": {},
                    "metadata": {},
                    "_xvars_demo_parity_direct_infer": True,
                }
            ]
        else:
            test_set = self._resolve_split_path("test", test_set)
            test_data = build_dataset(self.config, test_set, None, split="test")
        self._init_wandb(use_wandb=use_wandb)
        return self.trainer.infer(model, test_data, use_wandb=use_wandb)

    @config_operation
    def evaluate(
        self,
        test_set: str | None = None,
        weights: str | None = None,
        predictions: str | dict | None = None,
        use_wandb: bool = True,
        **kwargs,
    ) -> dict | str | None:
        del kwargs
        from opensportslib.core.trainer.vqa_trainer import Trainer_VQA
        from opensportslib.datasets.builder import build_dataset

        self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
        test_set = self._resolve_split_path("test", test_set)
        test_data = build_dataset(self.config, test_set, None, split="test")
        self._init_wandb(use_wandb=use_wandb)
        if predictions is None:
            predictions = self.infer(test_set=test_set, weights=weights, use_wandb=use_wandb)
        elif isinstance(predictions, str):
            import json

            with open(expand(predictions), encoding="utf-8") as f:
                predictions = json.load(f)
        self.trainer = self.trainer or Trainer_VQA(self.config)
        return self.trainer.evaluate(predictions, test_data, use_wandb=use_wandb)

    def save_predictions(
        self,
        output_path: str,
        predictions: dict,
        output_format: str = "osl",
    ) -> str:
        """Persist VQA predictions, optionally as X-VARS-compatible rows."""

        if str(output_format).lower() != "xvars":
            return super().save_predictions(output_path, predictions)

        import json

        payload = self._to_xvars_prediction_rows(predictions)
        dst = expand(output_path)
        os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)
        with open(dst, "w", encoding="utf-8") as f:
            json.dump(payload, f)
        return dst

    @staticmethod
    def _to_xvars_prediction_rows(predictions: dict[str, Any]) -> list[dict[str, Any]]:
        rows = []
        for item in predictions.get("data", []) if isinstance(predictions, dict) else []:
            video_path = str(item.get("video_path") or "")
            video_name = os.path.splitext(os.path.basename(video_path))[0] if video_path else str(item.get("id"))
            rows.append(
                {
                    "id": item.get("id"),
                    "video_name": video_name,
                    "Q": item.get("question"),
                    "pred": item.get("answer_text"),
                }
            )
        return rows

load_weights(weights=None, **kwargs)

Source code in opensportslib/apis/vqa.py

def load_weights(self, weights: str | None = None, **kwargs) -> None:
    del kwargs
    if weights is None:
        raise ValueError("`weights` must be provided to load_weights().")
    from opensportslib.core.trainer.vqa_trainer import Trainer_VQA

    _set_model_checkpoint_path(self.config, weights)
    self.trainer = Trainer_VQA(self.config)
    self.trainer.load(weights)
    self.last_loaded_weights = weights
    self.best_checkpoint = weights

train(train_set=None, valid_set=None, weights=None, resume_from_checkpoint=None, use_ddp=None, use_wandb=True, **kwargs)

Source code in opensportslib/apis/vqa.py

@config_operation
def train(
    self,
    train_set: str | None = None,
    valid_set: str | None = None,
    weights: str | None = None,
    resume_from_checkpoint: str | None = None,
    use_ddp: bool | None = None,
    use_wandb: bool = True,
    **kwargs,
) -> str | None:
    del kwargs

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    execution = get_train_execution(self.config)
    backend = str(execution.get("training_backend", "placeholder")).lower()
    vqa_backend = get_vqa_backend(self.config)
    if vqa_backend == "qwen_xvars_infer" and backend != "qwen_xvars_lora":
        raise ValueError(
            "The 'qwen_xvars_infer' backend requires TRAIN.execution.training_backend='qwen_xvars_lora' for train()."
        )
    if vqa_backend == "qwen_vl_native_infer" and backend != "qwen_vl_native_lora":
        raise ValueError(
            "The 'qwen_vl_native_infer' backend requires TRAIN.execution.training_backend='qwen_vl_native_lora' for train()."
        )
    import torch
    import torch.multiprocessing as mp

    train_set = self._resolve_split_path("train", train_set)
    valid_set = self._resolve_split_path("valid", valid_set)
    if backend in {"xvars_videochatgpt_lora", "qwen_xvars_lora", "qwen_vl_native_lora"}:
        requested_gpus = get_system_gpu_count(self.config)
        available_gpus = int(torch.cuda.device_count() or 0)
        if requested_gpus > 0:
            world_size = min(available_gpus, int(requested_gpus)) if available_gpus > 0 else int(requested_gpus)
        else:
            world_size = available_gpus
        effective_world_size = max(int(world_size), 1)
        if use_ddp is None:
            launch_ddp = effective_world_size > 1
        else:
            launch_ddp = bool(use_ddp) and effective_world_size > 1
        if requested_gpus > 0 and available_gpus > 0 and requested_gpus != available_gpus:
            logging.info(
                "VQA train GPU selection | requested=%s | available=%s | launching=%s",
                requested_gpus,
                available_gpus,
                effective_world_size if launch_ddp else 1,
            )
        logging.info(
            "VQA train launch | mode=%s | requested_gpus=%s | available_gpus=%s | world_size=%s",
            "ddp" if launch_ddp else "single",
            requested_gpus,
            available_gpus,
            effective_world_size if launch_ddp else 1,
        )

        ctx = mp.get_context("spawn")
        queue = ctx.SimpleQueue()
        if launch_ddp:
            mp.spawn(
                VQAModel._worker_ddp,
                args=(
                    effective_world_size,
                    self.config_path,
                    self.config,
                    queue,
                    train_set,
                    valid_set,
                    resume_from_checkpoint,
                    use_wandb,
                ),
                nprocs=effective_world_size,
            )
        else:
            VQAModel._worker_ddp(
                rank=0,
                world_size=1,
                config_path=self.config_path,
                config=self.config,
                return_queue=queue,
                train_set=train_set,
                valid_set=valid_set,
                resume_from_checkpoint=resume_from_checkpoint,
                use_wandb=use_wandb,
            )
        ckpt = queue.get()
        self.best_checkpoint = ckpt
        self.last_loaded_weights = ckpt
        return ckpt

    raise ValueError(
        f"Unsupported VQA training backend '{backend}'. "
        "Only 'xvars_videochatgpt_lora' and 'qwen_xvars_lora' are supported."
    )

infer(test_set=None, weights=None, use_wandb=True, video_path=None, question=None, session_id=None, **kwargs)

Source code in opensportslib/apis/vqa.py

@config_operation
def infer(
    self,
    test_set: str | None = None,
    weights: str | None = None,
    use_wandb: bool = True,
    video_path: str | None = None,
    question: str | None = None,
    session_id: str | None = None,
    **kwargs,
) -> dict:
    remote_mode_provided = "remote_mode" in kwargs
    remote_mode = kwargs.pop("remote_mode", "full_test_set")
    if self.is_remote:
        remote_model_id = kwargs.pop("remote_model_id", None)
        remote_task_options = kwargs.pop("remote_task_options", None)
        if kwargs:
            raise TypeError(f"Unsupported remote inference options: {', '.join(kwargs)}")
        if test_set is not None:
            if video_path is not None or question is not None or session_id is not None:
                raise ValueError("Remote VQA manifest inference accepts only `test_set`.")
            if remote_mode == "per_sample":
                batch = self.submit_per_sample_inference(
                    task_type="vqa",
                    test_set=self._resolve_split_path("test", test_set),
                    model_id=remote_model_id,
                    task_options=remote_task_options,
                )
                collected = self.wait_for_remote_batch(batch)
                self.last_remote_failures = collected["failures"]
                if self.last_remote_failures:
                    logging.warning("Remote per-sample inference completed with %d failures.", len(self.last_remote_failures))
                return collected["predictions"]
            if remote_mode != "full_test_set":
                raise ValueError("`remote_mode` must be `full_test_set` or `per_sample`.")
            job = self.submit_inference(
                task_type="vqa",
                test_set=self._resolve_split_path("test", test_set),
                model_id=remote_model_id,
                task_options=remote_task_options,
            )
            self.last_remote_failures = []
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if remote_mode != "full_test_set":
            raise ValueError("`remote_mode=per_sample` requires `test_set`, not direct VQA input.")
        # A stored session is for question-only follow-ups. A new video
        # must start a fresh server session; explicitly supplied session
        # IDs remain caller-controlled and are forwarded as requested.
        active_session_id = (
            session_id
            if session_id is not None
            else (self.remote_session_id if video_path is None else None)
        )
        if active_session_id is not None and video_path is None and question:
            job = self.submit_session_inference(
                task_type="vqa",
                session_id=active_session_id,
                model_id=remote_model_id,
                task_options=remote_task_options,
                question=str(question),
            )
            return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
        if not video_path or not str(question or "").strip():
            raise ValueError("Remote direct VQA inference requires `video_path` and a non-empty `question`.")
        job = self.submit_video_inference(
            task_type="vqa",
            video_path=video_path,
            question=str(question),
            session_id=active_session_id,
            model_id=remote_model_id,
            task_options=remote_task_options,
        )
        return self.wait_for_remote_result(job["job_id"])["result"]["predictions"]
    if remote_mode_provided:
        raise ValueError("`remote_mode` is available only when `remote` is configured.")
    del kwargs
    from opensportslib.core.trainer.vqa_trainer import Trainer_VQA
    from opensportslib.datasets.builder import build_dataset
    from opensportslib.models.builder import build_model
    from opensportslib.core.utils.config import select_device

    direct_requested = video_path is not None or question is not None
    if direct_requested and test_set is not None:
        raise ValueError("Provide either `test_set` or `video_path`/`question`, not both.")
    if direct_requested and (not video_path or not str(question or "").strip()):
        raise ValueError("Direct VQA inference requires both `video_path` and a non-empty `question`.")

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    backend = get_vqa_backend(self.config)
    effective_weights = weights if weights is not None else self.last_loaded_weights
    _set_model_checkpoint_path(self.config, effective_weights)
    self.trainer = Trainer_VQA(self.config)
    if effective_weights is not None:
        # Validate OpenSportsLib adapter metadata before allocating the base model.
        self.trainer.load(effective_weights)
    resolved_video_path = None
    if direct_requested:
        resolved_video_path = expand(str(video_path))
        if not os.path.isfile(resolved_video_path):
            raise FileNotFoundError(f"Video file not found: {resolved_video_path}")
    device = select_device(self.config.SYSTEM)
    model, _ = build_model(self.config, device)
    if direct_requested:
        test_data = [
            {
                "id": os.path.splitext(os.path.basename(resolved_video_path))[0],
                "question": str(question).strip(),
                "references": [],
                "video_path": resolved_video_path,
                "frame_paths": [],
                "video_frames": [],
                "video_spatio_temporal_features": None,
                "prior_prediction_text": "",
                "labels": {},
                "metadata": {},
                "_xvars_demo_parity_direct_infer": True,
            }
        ]
    else:
        test_set = self._resolve_split_path("test", test_set)
        test_data = build_dataset(self.config, test_set, None, split="test")
    self._init_wandb(use_wandb=use_wandb)
    return self.trainer.infer(model, test_data, use_wandb=use_wandb)

evaluate(test_set=None, weights=None, predictions=None, use_wandb=True, **kwargs)

Source code in opensportslib/apis/vqa.py

@config_operation
def evaluate(
    self,
    test_set: str | None = None,
    weights: str | None = None,
    predictions: str | dict | None = None,
    use_wandb: bool = True,
    **kwargs,
) -> dict | str | None:
    del kwargs
    from opensportslib.core.trainer.vqa_trainer import Trainer_VQA
    from opensportslib.datasets.builder import build_dataset

    self.config = self._effective_config(resolve_config_omega(self.config, weights=weights))
    test_set = self._resolve_split_path("test", test_set)
    test_data = build_dataset(self.config, test_set, None, split="test")
    self._init_wandb(use_wandb=use_wandb)
    if predictions is None:
        predictions = self.infer(test_set=test_set, weights=weights, use_wandb=use_wandb)
    elif isinstance(predictions, str):
        import json

        with open(expand(predictions), encoding="utf-8") as f:
            predictions = json.load(f)
    self.trainer = self.trainer or Trainer_VQA(self.config)
    return self.trainer.evaluate(predictions, test_data, use_wandb=use_wandb)

save_predictions(output_path, predictions, output_format='osl')

Persist VQA predictions, optionally as X-VARS-compatible rows.

Source code in opensportslib/apis/vqa.py

def save_predictions(
    self,
    output_path: str,
    predictions: dict,
    output_format: str = "osl",
) -> str:
    """Persist VQA predictions, optionally as X-VARS-compatible rows."""

    if str(output_format).lower() != "xvars":
        return super().save_predictions(output_path, predictions)

    import json

    payload = self._to_xvars_prediction_rows(predictions)
    dst = expand(output_path)
    os.makedirs(os.path.dirname(dst) or ".", exist_ok=True)
    with open(dst, "w", encoding="utf-8") as f:
        json.dump(payload, f)
    return dst

Configuration API

Config loads, composes, migrates, validates, and edits a configuration before model allocation. Config.from_pretrained() downloads config.yaml from a Hub model repository, so it needs network access and any required Hugging Face authentication. Use options() before update() to discover supported, validated settings.

Editable OSL config. Loading this object never allocates model weights.

Source code in opensportslib/core/config/editable.py

class Config:
    """Editable OSL config. Loading this object never allocates model weights."""

    def __init__(self, document, *, source=None, weights=None, provenance=None):
        self._document = deepcopy(migrate_config(document))
        self.source = str(source) if source is not None else None
        self.weights = weights
        self._provenance = dict(provenance or {})
        self._updates: dict[str, Any] = {}
        self._friendly: dict[str, dict[str, Any]] = {}
        self._has_dotted_overrides = False
        self._validate(self._document)

    @classmethod
    def from_file(cls, path: str | Path):
        from omegaconf import OmegaConf
        from .schemas.schema_canonical import is_canonical_schema

        path = Path(path).expanduser().resolve()
        layers = _compose_yaml_layers(path) or [path]
        document = OmegaConf.create({})
        provenance = {}
        for layer in layers:
            raw = OmegaConf.load(layer)
            for key, _ in _leaves(OmegaConf.to_container(raw, resolve=False)):
                provenance[key] = str(layer)
            document = OmegaConf.merge(document, raw)
        raw = OmegaConf.to_container(document, resolve=False)
        if not is_canonical_schema(raw):
            raw = migrate_config(_resolve(raw))
            provenance = {key: str(path) for key, _ in _leaves(raw)}
        return cls(raw, source=path, provenance=provenance)

    @classmethod
    def from_pretrained(cls, model_id: str):
        from huggingface_hub import hf_hub_download
        from huggingface_hub.utils import validate_repo_id

        validate_repo_id(model_id)
        try:
            path = hf_hub_download(repo_id=model_id, filename="config.yaml")
        except Exception as exc:
            raise ValueError(
                f"Cannot load OSL config.yaml from {model_id!r}: {exc}. "
                "Check the model ID, network access, authentication for private models, "
                "or provide a config file."
            ) from exc
        config = cls.from_file(path)
        config.weights = model_id
        config._provenance = {key: model_id for key in config._provenance}
        return config

    def get_config(self) -> dict[str, Any]:
        """Return a detached, resolved canonical document."""
        return deepcopy(_resolve(self._document))

    @staticmethod
    def _validate(document):
        resolved = _resolve(document)
        validate_config(resolved)
        for split, settings in _get(resolved, "DATA.common.splits", {}).items():
            overlap = settings.get("overlap_len")
            clip = _get(resolved, "DATA.inputs.video.sampling.clip_len")
            if overlap is not None and clip is not None and not 0 <= overlap < clip:
                raise ValueError(f"DATA.common.splits.{split}.overlap_len must be smaller than clip_len")
        native = _get(resolved, "TRAIN.execution.native_vl", {})
        if "min_pixels" in native and "max_pixels" in native and native["min_pixels"] > native["max_pixels"]:
            raise ValueError("TRAIN.execution.native_vl.min_pixels must not exceed max_pixels")
        return resolved

    def _locked_paths(self, doc):
        locked = set()
        for key, comp in _get(doc, "MODEL.components", {}).items():
            name = _get(comp, "source.name", "") or ""
            if name.startswith("h5_header_"):
                from .rule_variants import HEADER_RULE_VARIANTS, SKELETON_RULE_VARIANTS

                variant = {**HEADER_RULE_VARIANTS, **SKELETON_RULE_VARIANTS}.get(name, {})
                locked.update(f"MODEL.components.{key}.params.{param}" for param in variant)
        return locked

    def _registry(self):
        doc = self.get_config()
        result = {}
        missing = object()

        def add(name, paths, description, minimum=None, maximum=None, choices=None, init=False):
            if isinstance(paths, str):
                paths = (paths,)
            paths = tuple(p for p in paths if _get(doc, p, missing) is not missing)
            if paths:
                result[name] = Option(paths, description, minimum, maximum, choices, init)

        splits = _get(doc, "DATA.common.splits", {})
        for split in ("train", "valid", "test"):
            add(f"data.{split}_set", f"DATA.common.splits.{split}.annotation_path", f"{split.title()} annotation manifest")
        add("data.data_root", "DATA.common.data_root", "Root directory for data and referenced split paths")
        add("runtime.device", "SYSTEM.device", "Execution device", choices=("auto", "cpu", "cuda"), init=True)
        add("runtime.seed", "SYSTEM.reproducibility.seed", "Random seed (enables reproducibility)", 0, 2**32 - 1)
        add("runtime.output_dir", "SYSTEM.paths.save_dir", "Root directory for new run artifacts", init=True)
        task = str(doc.get("TASK", "")).lower()
        trainable = _get(doc, "TRAIN.execution.enabled", True) and _get(doc, "MODEL.metadata.family") != "RuleBased"
        backend = _get(doc, "MODEL.metadata.backend", _get(doc, "TRAIN.execution.backend", "baseline"))
        if trainable:
            for name, path, minimum in (("epochs", "epochs", 1), ("learning_rate", "optimizer.lr", 0), ("weight_decay", "optimizer.weight_decay", 0)):
                add(f"training.{name}", f"TRAIN.{path}", name.replace("_", " ").capitalize(), minimum)
            for field in ("batch_size", "num_workers"):
                add(f"training.{field}", f"DATA.common.splits.train.dataloader.{field}", f"Training {field.replace('_', ' ')}", 1 if field == "batch_size" else 0)
        # VQA engines generate one sample at a time without a DataLoader.
        if task != "vqa" and _get(doc, "MODEL.metadata.family") != "RuleBased":
            for field in ("batch_size", "num_workers"):
                add(f"inference.{field}", tuple(f"DATA.common.splits.{s}.dataloader.{field}" for s in splits if s in {"valid", "test", "valid_data_frames"}), f"Validation and test {field.replace('_', ' ')}", 1 if field == "batch_size" else 0)
        if task == "vqa" and backend in {"qwen_xvars_infer", "qwen_vl_native_infer", "xvars_videochatgpt"}:
            for name in ("max_new_tokens", "temperature", "repetition_penalty", "no_repeat_ngram_size"):
                if backend == "qwen_vl_native_infer" and name in {"repetition_penalty", "no_repeat_ngram_size"}:
                    continue
                add(f"inference.{name}", f"TRAIN.execution.generation.{name}", f"Generation {name.replace('_', ' ')}", 1 if name in {"max_new_tokens", "repetition_penalty"} else 0)
            result["inference.do_sample"] = Option(
                ("TRAIN.execution.generation.do_sample",), "Enable sampling during generation"
            )
            result["inference.top_p"] = Option(
                ("TRAIN.execution.generation.top_p",), "Nucleus sampling probability", 0, 1
            )
        sections = {"sampling": "DATA.inputs.video.sampling", "augmentation": "DATA.inputs.video.augmentations"}
        if trainable:
            sections.update({"scheduler": "TRAIN.scheduler", "training_sampling": "TRAIN.sampling"})
        if task == "vqa":
            sections["prompt"] = "TRAIN.execution.prompt"
            if trainable:
                sections.update({"lora": "TRAIN.execution.lora", "sft": "TRAIN.execution.sft"})
            sections["native_vl"] = "TRAIN.execution.native_vl"
        if _get(doc, "MODEL.policies.test_time_adaptation.enabled", False):
            sections["adaptation"] = "MODEL.policies.test_time_adaptation"
        for key, comp in _get(doc, "MODEL.components", {}).items():
            if comp.get("kind") in {"algorithm", "postprocessor"}:
                sections[key] = f"MODEL.components.{key}.params"
        locked = self._locked_paths(doc)
        for group, root in sections.items():
            for suffix, value in _leaves(_get(doc, root, {})):
                if not suffix or suffix.split(".")[-1] in {"type", "name", "label", "head_name", "trainable_parameters"}:
                    continue
                path = f"{root}.{suffix}"
                if path in locked or isinstance(value, dict):
                    continue
                add(f"{group}.{suffix}", path, f"{group.replace('_', ' ').title()}: {suffix.replace('_', ' ')}", init=group not in {"scheduler", "training_sampling", "sft", "prompt"})
        return result

    def options(self) -> dict[str, dict[str, Any]]:
        doc = self.get_config()
        options = {}
        for name, spec in self._registry().items():
            defaults = {"inference.do_sample": False, "inference.top_p": 1.0}
            values = {path: deepcopy(_get(doc, path, defaults.get(name))) for path in spec.paths}
            current = list(values.values())
            options[name] = {
                "value": current[0] if all(v == current[0] for v in current) else values,
                "type": type(current[0]).__name__, "description": spec.description,
                "minimum": spec.minimum, "maximum": spec.maximum,
                "choices": list(spec.choices) if spec.choices else None,
                "paths": list(spec.paths),
                "sources": {p: self._provenance.get(p, self.source or "runtime") for p in spec.paths},
                "requires_initialization": spec.requires_initialization,
            }
        return options

    @staticmethod
    def _check_value(path, value, old, spec=None):
        if isinstance(old, bool):
            valid = isinstance(value, bool)
        elif isinstance(old, int):
            valid = isinstance(value, int) and not isinstance(value, bool)
        elif isinstance(old, float):
            valid = isinstance(value, (int, float)) and not isinstance(value, bool)
        elif old is None:
            valid = value is None or isinstance(value, (str, int, float, bool, list, dict))
        else:
            valid = isinstance(value, type(old))
        if not valid:
            raise ValueError(f"{path}: expected {type(old).__name__}, got {type(value).__name__}")
        if isinstance(value, (float, int)) and not isinstance(value, bool) and not math.isfinite(value):
            raise ValueError(f"{path}: value must be finite")
        if spec is not None:
            if spec.choices and value not in spec.choices:
                raise ValueError(f"{path}: choose one of {spec.choices}")
            if (spec.minimum is not None and value < spec.minimum) or (spec.maximum is not None and value > spec.maximum):
                raise ValueError(f"{path}: value outside [{spec.minimum}, {spec.maximum}]")

    def update(self, *, data: DataOptions | None = None, runtime: RuntimeOptions | None = None,
               training: TrainingOptions | None = None, inference: InferenceOptions | None = None,
               overrides: dict[str, Any] | None = None, **groups):
        requested = {**groups, **{k: v for k, v in {"data": data, "runtime": runtime, "training": training, "inference": inference}.items() if v is not None}}
        registry = self._registry()
        doc = self.get_config()
        changes = {}
        friendly = deepcopy(self._friendly)
        for group, settings in requested.items():
            if not isinstance(settings, dict):
                raise ValueError(f"{group}: expected a dictionary of options")
            if not any(name.startswith(group + ".") for name in registry):
                raise ValueError(f"Unsupported option group {group!r}; inspect config.options()")
            for name, value in settings.items():
                key = f"{group}.{name}"
                if key not in registry:
                    raise ValueError(f"Unsupported option {key!r}; inspect config.options()")
                spec = registry[key]
                for path in spec.paths:
                    old = _get(doc, path, {"inference.do_sample": False, "inference.top_p": 1.0}.get(key))
                    self._check_value(key, value, old, spec)
                    if path in changes:
                        raise ValueError(f"Conflicting options for {path}")
                    changes[path] = value
                friendly.setdefault(group, {})[name] = deepcopy(value)
        if overrides is not None and not isinstance(overrides, dict):
            raise ValueError("overrides: expected a dictionary of canonical dotted paths")
        missing = object()
        for path, value in (overrides or {}).items():
            if not isinstance(path, str) or _get(doc, path, missing) is missing:
                raise ValueError(f"Unknown config path {path!r}; inspect config.get_config()")
            if any(_overlap(path, other) for other in changes):
                raise ValueError(f"Conflicting overrides for {path}")
            self._check_value(path, value, _get(doc, path))
            changes[path] = value
        candidate = deepcopy(self._document)
        for path, value in changes.items():
            _set(candidate, path, value)
        for locked in self._locked_paths(_resolve(candidate)):
            if any(_overlap(path, locked) for path in changes):
                raise ValueError(f"{locked} is controlled by the selected rule variant")
        resolved = _resolve(candidate)
        generation = _get(resolved, "TRAIN.execution.generation", {})
        if generation.get("do_sample") and float(generation.get("temperature", 0)) <= 0:
            raise ValueError("inference.temperature must be positive when inference.do_sample is true")
        for spec in registry.values():
            for path in spec.paths:
                defaults = {
                    "TRAIN.execution.generation.do_sample": False,
                    "TRAIN.execution.generation.top_p": 1.0,
                }
                old = _get(doc, path, defaults.get(path, _get(resolved, path)))
                value = _get(resolved, path, defaults.get(path))
                self._check_value(path, value, old, spec)
        self._synchronize(candidate, changes)
        self._validate(candidate)
        updates = deepcopy(self._updates)
        for path, value in changes.items():
            for old in list(updates):
                if _overlap(old, path):
                    del updates[old]
            updates[path] = deepcopy(value)
        self._document = candidate
        self._updates = updates
        self._friendly = friendly
        self._has_dotted_overrides |= bool(overrides)
        self._provenance.update({path: "user" for path in changes})
        return self

    @staticmethod
    def _synchronize(doc, changes):
        """Keep aliases consumed by existing backends in agreement."""
        task = str(doc.get("TASK", "")).lower()
        for path, value in list(changes.items()):
            if path == "DATA.common.data_root":
                for split, split_cfg in _get(doc, "DATA.common.splits", {}).items():
                    source_path = split_cfg.get("source_path")
                    if source_path is not None and "${" not in str(source_path):
                        _set(doc, f"DATA.common.splits.{split}.source_path", value)
            if path == "SYSTEM.reproducibility.seed":
                _set(doc, "SYSTEM.reproducibility.use_seed", True)
            if path == "SYSTEM.device":
                _set(doc, "MODEL.runtime.device", value)
                if value != "auto":
                    _set(doc, "TRAIN.execution.hf.prefer_cuda", value == "cuda")
            if path.endswith(".dataloader.num_workers") and value == 0:
                root = path.rsplit(".", 1)[0]
                _set(doc, root + ".persistent_workers", False)
                _set(doc, root + ".prefetch_factor", None)
            if task == "vqa" and path.endswith(".dataloader.batch_size"):
                split = path.split(".")[3]
                field = "per_device_train_batch_size" if split == "train" else "per_device_eval_batch_size"
                _set(doc, "TRAIN.execution.sft." + field, value)
            if task == "vqa" and path == "DATA.common.splits.train.dataloader.num_workers":
                _set(doc, "TRAIN.execution.sft.dataloader_num_workers", value)

    def apply_to(self, config):
        """Reapply explicit updates after a runtime/checkpoint config merge."""
        doc = deepcopy(namespace_to_plain_dict(config))
        source_values = self.get_config()
        for path, expression in _leaves(self._document):
            if isinstance(expression, str) and "${" in expression:
                refs = re.findall(r"\$\{([^}]+)\}", expression)
                if _get(doc, path) == _get(source_values, path) or any(any(_overlap(ref, changed) for changed in self._updates) for ref in refs):
                    _set(doc, path, expression)
        for path, value in self._updates.items():
            _set(doc, path, value)
        self._synchronize(doc, self._updates)
        return _normalize_cpu_loader_backend(self._validate(doc))

    def remote_overrides(self):
        if self._has_dotted_overrides:
            raise ValueError("Remote execution does not accept dotted config overrides")
        forbidden = set(self._friendly) - {"data", "inference"}
        if forbidden or "num_workers" in self._friendly.get("inference", {}):
            raise ValueError("Remote tuning accepts advertised inference options only; runtime and workers are server-controlled")
        return deepcopy(self._friendly.get("inference", {}))

from_file(path) classmethod

Source code in opensportslib/core/config/editable.py

@classmethod
def from_file(cls, path: str | Path):
    from omegaconf import OmegaConf
    from .schemas.schema_canonical import is_canonical_schema

    path = Path(path).expanduser().resolve()
    layers = _compose_yaml_layers(path) or [path]
    document = OmegaConf.create({})
    provenance = {}
    for layer in layers:
        raw = OmegaConf.load(layer)
        for key, _ in _leaves(OmegaConf.to_container(raw, resolve=False)):
            provenance[key] = str(layer)
        document = OmegaConf.merge(document, raw)
    raw = OmegaConf.to_container(document, resolve=False)
    if not is_canonical_schema(raw):
        raw = migrate_config(_resolve(raw))
        provenance = {key: str(path) for key, _ in _leaves(raw)}
    return cls(raw, source=path, provenance=provenance)

from_pretrained(model_id) classmethod

Source code in opensportslib/core/config/editable.py

@classmethod
def from_pretrained(cls, model_id: str):
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import validate_repo_id

    validate_repo_id(model_id)
    try:
        path = hf_hub_download(repo_id=model_id, filename="config.yaml")
    except Exception as exc:
        raise ValueError(
            f"Cannot load OSL config.yaml from {model_id!r}: {exc}. "
            "Check the model ID, network access, authentication for private models, "
            "or provide a config file."
        ) from exc
    config = cls.from_file(path)
    config.weights = model_id
    config._provenance = {key: model_id for key in config._provenance}
    return config

get_config()

Return a detached, resolved canonical document.

Source code in opensportslib/core/config/editable.py

def get_config(self) -> dict[str, Any]:
    """Return a detached, resolved canonical document."""
    return deepcopy(_resolve(self._document))

options()

Source code in opensportslib/core/config/editable.py

def options(self) -> dict[str, dict[str, Any]]:
    doc = self.get_config()
    options = {}
    for name, spec in self._registry().items():
        defaults = {"inference.do_sample": False, "inference.top_p": 1.0}
        values = {path: deepcopy(_get(doc, path, defaults.get(name))) for path in spec.paths}
        current = list(values.values())
        options[name] = {
            "value": current[0] if all(v == current[0] for v in current) else values,
            "type": type(current[0]).__name__, "description": spec.description,
            "minimum": spec.minimum, "maximum": spec.maximum,
            "choices": list(spec.choices) if spec.choices else None,
            "paths": list(spec.paths),
            "sources": {p: self._provenance.get(p, self.source or "runtime") for p in spec.paths},
            "requires_initialization": spec.requires_initialization,
        }
    return options

update(*, data=None, runtime=None, training=None, inference=None, overrides=None, **groups)

Source code in opensportslib/core/config/editable.py

def update(self, *, data: DataOptions | None = None, runtime: RuntimeOptions | None = None,
           training: TrainingOptions | None = None, inference: InferenceOptions | None = None,
           overrides: dict[str, Any] | None = None, **groups):
    requested = {**groups, **{k: v for k, v in {"data": data, "runtime": runtime, "training": training, "inference": inference}.items() if v is not None}}
    registry = self._registry()
    doc = self.get_config()
    changes = {}
    friendly = deepcopy(self._friendly)
    for group, settings in requested.items():
        if not isinstance(settings, dict):
            raise ValueError(f"{group}: expected a dictionary of options")
        if not any(name.startswith(group + ".") for name in registry):
            raise ValueError(f"Unsupported option group {group!r}; inspect config.options()")
        for name, value in settings.items():
            key = f"{group}.{name}"
            if key not in registry:
                raise ValueError(f"Unsupported option {key!r}; inspect config.options()")
            spec = registry[key]
            for path in spec.paths:
                old = _get(doc, path, {"inference.do_sample": False, "inference.top_p": 1.0}.get(key))
                self._check_value(key, value, old, spec)
                if path in changes:
                    raise ValueError(f"Conflicting options for {path}")
                changes[path] = value
            friendly.setdefault(group, {})[name] = deepcopy(value)
    if overrides is not None and not isinstance(overrides, dict):
        raise ValueError("overrides: expected a dictionary of canonical dotted paths")
    missing = object()
    for path, value in (overrides or {}).items():
        if not isinstance(path, str) or _get(doc, path, missing) is missing:
            raise ValueError(f"Unknown config path {path!r}; inspect config.get_config()")
        if any(_overlap(path, other) for other in changes):
            raise ValueError(f"Conflicting overrides for {path}")
        self._check_value(path, value, _get(doc, path))
        changes[path] = value
    candidate = deepcopy(self._document)
    for path, value in changes.items():
        _set(candidate, path, value)
    for locked in self._locked_paths(_resolve(candidate)):
        if any(_overlap(path, locked) for path in changes):
            raise ValueError(f"{locked} is controlled by the selected rule variant")
    resolved = _resolve(candidate)
    generation = _get(resolved, "TRAIN.execution.generation", {})
    if generation.get("do_sample") and float(generation.get("temperature", 0)) <= 0:
        raise ValueError("inference.temperature must be positive when inference.do_sample is true")
    for spec in registry.values():
        for path in spec.paths:
            defaults = {
                "TRAIN.execution.generation.do_sample": False,
                "TRAIN.execution.generation.top_p": 1.0,
            }
            old = _get(doc, path, defaults.get(path, _get(resolved, path)))
            value = _get(resolved, path, defaults.get(path))
            self._check_value(path, value, old, spec)
    self._synchronize(candidate, changes)
    self._validate(candidate)
    updates = deepcopy(self._updates)
    for path, value in changes.items():
        for old in list(updates):
            if _overlap(old, path):
                del updates[old]
        updates[path] = deepcopy(value)
    self._document = candidate
    self._updates = updates
    self._friendly = friendly
    self._has_dotted_overrides |= bool(overrides)
    self._provenance.update({path: "user" for path in changes})
    return self

apply_to(config)

Reapply explicit updates after a runtime/checkpoint config merge.

Source code in opensportslib/core/config/editable.py

def apply_to(self, config):
    """Reapply explicit updates after a runtime/checkpoint config merge."""
    doc = deepcopy(namespace_to_plain_dict(config))
    source_values = self.get_config()
    for path, expression in _leaves(self._document):
        if isinstance(expression, str) and "${" in expression:
            refs = re.findall(r"\$\{([^}]+)\}", expression)
            if _get(doc, path) == _get(source_values, path) or any(any(_overlap(ref, changed) for changed in self._updates) for ref in refs):
                _set(doc, path, expression)
    for path, value in self._updates.items():
        _set(doc, path, value)
    self._synchronize(doc, self._updates)
    return _normalize_cpu_loader_backend(self._validate(doc))

remote_overrides()

Source code in opensportslib/core/config/editable.py

def remote_overrides(self):
    if self._has_dotted_overrides:
        raise ValueError("Remote execution does not accept dotted config overrides")
    forbidden = set(self._friendly) - {"data", "inference"}
    if forbidden or "num_workers" in self._friendly.get("inference", {}):
        raise ValueError("Remote tuning accepts advertised inference options only; runtime and workers are server-controlled")
    return deepcopy(self._friendly.get("inference", {}))

Remote model registry

RemoteModelRegistry is an administrative client for the separately deployed server. register_model, set_default, unregister_model, and reconcile_runtime mutate server state; use reconcile_runtime(dry_run=True) to inspect safely. Local registration requires a server API key; Hub operations may require an HF token.

Bases: RuntimeError

Source code in opensportslib/remote_registry.py

class RemoteRegistryError(RuntimeError):
    def __init__(self, status_code: int, detail: Any):
        self.status_code = status_code
        self.detail = detail
        super().__init__(f"Remote model registry returned HTTP {status_code}: {detail}")

Administrative client for an OpenSportsLib server model registry.

Source code in opensportslib/remote_registry.py

class RemoteModelRegistry:
    """Administrative client for an OpenSportsLib server model registry."""

    def __init__(
        self,
        remote: str,
        api_key: str | None = None,
        timeout: float = 30.0,
        *,
        admin_token: str | None = None,
    ):
        if not remote:
            raise ValueError("remote is required.")
        if admin_token is not None:
            warnings.warn("admin_token is deprecated; use api_key.", DeprecationWarning, stacklevel=2)
            if api_key is not None and api_key != admin_token:
                raise ValueError("api_key and admin_token must match when both are provided.")
            api_key = admin_token
        self.remote = remote.rstrip("/")
        self.api_key = api_key
        self.timeout = float(timeout)
        if self.timeout <= 0:
            raise ValueError("timeout must be positive.")

    def register_model(
        self,
        *,
        task_type: str,
        huggingface_model_id: str | None = None,
        weights_path: str | Path | None = None,
        config_path: str | Path | None = None,
        model_id: str | None = None,
        hf_token: str | None = None,
    ) -> dict[str, Any]:
        """Register exactly one Hugging Face or server-local model source."""
        if (huggingface_model_id is None) == (weights_path is None):
            raise ValueError("Provide exactly one of huggingface_model_id or weights_path.")
        if huggingface_model_id is not None:
            if config_path is not None:
                raise ValueError("config_path applies only to local models.")
            if model_id is not None and model_id != huggingface_model_id:
                raise ValueError("A Hugging Face model uses its repository ID as model_id.")
            source = {"type": "huggingface", "model_id": huggingface_model_id}
            if hf_token is not None:
                source["hf_token"] = hf_token
            payload = {
                "task_type": task_type,
                "source": source,
            }
        else:
            payload = {
                "task_type": task_type,
                "model_id": model_id,
                "source": {
                    "type": "local",
                    "weights_path": str(weights_path),
                    "config_path": str(config_path) if config_path is not None else None,
                },
            }
        return self._request("POST", "/models", payload, authenticated=weights_path is not None)

    def list_models(self) -> list[dict[str, Any]]:
        return self._request("GET", "/models")["models"]

    def get_model(self, model_id: str) -> dict[str, Any]:
        query = parse.urlencode({"model_id": model_id})
        return self._request("GET", f"/models/status?{query}")

    def get_operation(self, operation_id: str) -> dict[str, Any]:
        return self._request("GET", f"/model-operations/{parse.quote(operation_id, safe='')}")

    def wait_for_operation(
        self,
        operation_id: str,
        *,
        timeout: float = 7200.0,
        poll_interval: float = 1.0,
    ) -> dict[str, Any]:
        deadline = time.monotonic() + timeout
        while True:
            operation = self.get_operation(operation_id)
            if operation["status"] == "succeeded":
                return operation
            if operation["status"] == "failed":
                raise RuntimeError(operation.get("error") or "Model operation failed.")
            if time.monotonic() >= deadline:
                raise TimeoutError(f"Timed out waiting for model operation `{operation_id}`.")
            time.sleep(poll_interval)

    def set_default(self, task_type: str, model_id: str) -> dict[str, Any]:
        task = parse.quote(task_type, safe="")
        return self._request("PUT", f"/models/defaults/{task}", {"model_id": model_id}, authenticated=True)

    def unregister_model(self, model_id: str, hf_token: str | None = None) -> dict[str, Any]:
        encoded = parse.quote(model_id, safe="")
        return self._request("DELETE", f"/models/{encoded}", authenticated=True, hf_token=hf_token)

    def reconcile_runtime(
        self,
        *,
        dry_run: bool = True,
        include_active: bool = False,
    ) -> dict[str, Any]:
        """Inspect or recover stale server jobs; dry-run is the safe default."""
        return self._request(
            "POST",
            "/admin/runtime/reconcile",
            {"dry_run": bool(dry_run), "include_active": bool(include_active)},
            authenticated=True,
        )

    def _request(
        self,
        method: str,
        endpoint: str,
        payload: dict[str, Any] | None = None,
        authenticated: bool = False,
        hf_token: str | None = None,
    ) -> dict[str, Any]:
        body = json.dumps(payload).encode() if payload is not None else None
        headers = {"Content-Type": "application/json"}
        if authenticated and self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        if hf_token:
            headers["X-HF-Token"] = hf_token
        outgoing = request.Request(self.remote + endpoint, data=body, method=method, headers=headers)
        try:
            with request.urlopen(outgoing, timeout=self.timeout) as response:
                return json.loads(response.read().decode())
        except error.HTTPError as exc:
            detail = exc.read().decode(errors="replace")
            try:
                parsed = json.loads(detail)
                detail = parsed.get("detail", parsed)
            except json.JSONDecodeError:
                pass
            raise RemoteRegistryError(exc.code, detail) from exc
        except error.URLError as exc:
            raise ConnectionError(f"Could not reach remote server `{self.remote}`: {exc.reason}") from exc

register_model(*, task_type, huggingface_model_id=None, weights_path=None, config_path=None, model_id=None, hf_token=None)

Register exactly one Hugging Face or server-local model source.

Source code in opensportslib/remote_registry.py

def register_model(
    self,
    *,
    task_type: str,
    huggingface_model_id: str | None = None,
    weights_path: str | Path | None = None,
    config_path: str | Path | None = None,
    model_id: str | None = None,
    hf_token: str | None = None,
) -> dict[str, Any]:
    """Register exactly one Hugging Face or server-local model source."""
    if (huggingface_model_id is None) == (weights_path is None):
        raise ValueError("Provide exactly one of huggingface_model_id or weights_path.")
    if huggingface_model_id is not None:
        if config_path is not None:
            raise ValueError("config_path applies only to local models.")
        if model_id is not None and model_id != huggingface_model_id:
            raise ValueError("A Hugging Face model uses its repository ID as model_id.")
        source = {"type": "huggingface", "model_id": huggingface_model_id}
        if hf_token is not None:
            source["hf_token"] = hf_token
        payload = {
            "task_type": task_type,
            "source": source,
        }
    else:
        payload = {
            "task_type": task_type,
            "model_id": model_id,
            "source": {
                "type": "local",
                "weights_path": str(weights_path),
                "config_path": str(config_path) if config_path is not None else None,
            },
        }
    return self._request("POST", "/models", payload, authenticated=weights_path is not None)

list_models()

Source code in opensportslib/remote_registry.py

def list_models(self) -> list[dict[str, Any]]:
    return self._request("GET", "/models")["models"]

get_model(model_id)

Source code in opensportslib/remote_registry.py

def get_model(self, model_id: str) -> dict[str, Any]:
    query = parse.urlencode({"model_id": model_id})
    return self._request("GET", f"/models/status?{query}")

get_operation(operation_id)

Source code in opensportslib/remote_registry.py

def get_operation(self, operation_id: str) -> dict[str, Any]:
    return self._request("GET", f"/model-operations/{parse.quote(operation_id, safe='')}")

wait_for_operation(operation_id, *, timeout=7200.0, poll_interval=1.0)

Source code in opensportslib/remote_registry.py

def wait_for_operation(
    self,
    operation_id: str,
    *,
    timeout: float = 7200.0,
    poll_interval: float = 1.0,
) -> dict[str, Any]:
    deadline = time.monotonic() + timeout
    while True:
        operation = self.get_operation(operation_id)
        if operation["status"] == "succeeded":
            return operation
        if operation["status"] == "failed":
            raise RuntimeError(operation.get("error") or "Model operation failed.")
        if time.monotonic() >= deadline:
            raise TimeoutError(f"Timed out waiting for model operation `{operation_id}`.")
        time.sleep(poll_interval)

set_default(task_type, model_id)

Source code in opensportslib/remote_registry.py

def set_default(self, task_type: str, model_id: str) -> dict[str, Any]:
    task = parse.quote(task_type, safe="")
    return self._request("PUT", f"/models/defaults/{task}", {"model_id": model_id}, authenticated=True)

unregister_model(model_id, hf_token=None)

Source code in opensportslib/remote_registry.py

def unregister_model(self, model_id: str, hf_token: str | None = None) -> dict[str, Any]:
    encoded = parse.quote(model_id, safe="")
    return self._request("DELETE", f"/models/{encoded}", authenticated=True, hf_token=hf_token)

reconcile_runtime(*, dry_run=True, include_active=False)

Inspect or recover stale server jobs; dry-run is the safe default.

Source code in opensportslib/remote_registry.py

def reconcile_runtime(
    self,
    *,
    dry_run: bool = True,
    include_active: bool = False,
) -> dict[str, Any]:
    """Inspect or recover stale server jobs; dry-run is the safe default."""
    return self._request(
        "POST",
        "/admin/runtime/reconcile",
        {"dry_run": bool(dry_run), "include_active": bool(include_active)},
        authenticated=True,
    )

Conversion and Hugging Face tools

These exports are stable, but many perform disk and network work. Conversion functions create Parquet/WebDataset or JSON output. Hugging Face download and upload functions access remote repositories and can write local media or remote repository content. Inspect their signatures before use and use the dataset transfer guide for complete flows.

Convert an OpenSportsLib-style JSON annotation file into

1) a Parquet table with flattened metadata 2) WebDataset TAR shards containing the referenced input files + sample metadata

Each sample typically contains

• id
• inputs: [{"type": "...", "path": "...", ...}, ...]
• optional task-specific fields: events / captions / dense_captions / labels / metadata

Public entry point

convert_json_to_parquet(...)

parse_shard_size(value)

Parse byte counts such as 500MB, 1GB, or 1024MiB.

Source code in opensportslib/tools/osl_json_to_parquet.py

def parse_shard_size(value: int | str) -> int:
    """Parse byte counts such as 500MB, 1GB, or 1024MiB."""
    if isinstance(value, int):
        size = value
    else:
        text = str(value or "").strip()
        if not text:
            raise ValueError("shard_size is required.")
        match = re.fullmatch(r"(\d+)(?:\s*([a-zA-Z]+))?", text)
        if not match:
            raise ValueError(f"Invalid shard size: {value!r}")
        amount = int(match.group(1))
        unit = (match.group(2) or "B").lower()
        multipliers = {
            "b": 1,
            "byte": 1,
            "bytes": 1,
            "kb": 1_000,
            "mb": 1_000_000,
            "gb": 1_000_000_000,
            "tb": 1_000_000_000_000,
            "kib": 1024,
            "mib": 1024**2,
            "gib": 1024**3,
            "tib": 1024**4,
        }
        if unit not in multipliers:
            raise ValueError(f"Invalid shard size unit: {match.group(2)!r}")
        size = amount * multipliers[unit]
    if size < 1:
        raise ValueError("shard_size must be >= 1.")
    return size

convert_json_to_parquet(json_path, media_root, output_dir, *, shard_mode='size', shard_size=DEFAULT_SHARD_SIZE, samples_per_shard=100, compression='zstd', shard_prefix='shard', missing_policy='raise', keep_relative_paths_in_parquet=True, overwrite=False)

Convert an OSL-style JSON file to

• metadata.parquet (flattened, queryable per-sample table)
• shards/*.tar (WebDataset TAR shards with input files + sidecar JSON)
• shard_manifest.parquet

Parameters

json_path: Path to the JSON annotation file. media_root: Root directory where files referenced in inputs[].path live. output_dir: Destination directory. shard_mode: "size" groups samples by target TAR shard size. "samples" groups samples by samples_per_shard. shard_size: Target on-disk TAR shard size for shard_mode="size". Accepts bytes or strings such as "500MB", "1GB", or "1024MiB". samples_per_shard: Number of samples grouped in one TAR shard when shard_mode="samples". compression: Parquet compression codec ("zstd", "snappy", "gzip", or None). shard_prefix: Prefix for TAR shard file names. missing_policy: "raise" — abort if a referenced input file is missing. "skip" — keep sample in Parquet but omit missing file from shard. keep_relative_paths_in_parquet: If True (default), Parquet stores original relative paths from JSON. If False, Parquet stores resolved absolute paths. overwrite: Whether to remove and recreate output_dir if it already exists.

Returns

dict Summary with counts of samples, shards, and input files processed.

Source code in opensportslib/tools/osl_json_to_parquet.py

def convert_json_to_parquet(
    json_path: str | Path,
    media_root: str | Path,
    output_dir: str | Path,
    *,
    shard_mode: Literal["size", "samples"] = "size",
    shard_size: int | str = DEFAULT_SHARD_SIZE,
    samples_per_shard: int = 100,
    compression: Optional[str] = "zstd",
    shard_prefix: str = "shard",
    missing_policy: str = "raise",
    keep_relative_paths_in_parquet: bool = True,
    overwrite: bool = False,
) -> Dict[str, Any]:
    """
    Convert an OSL-style JSON file to:
      - metadata.parquet  (flattened, queryable per-sample table)
      - shards/*.tar      (WebDataset TAR shards with input files + sidecar JSON)
      - shard_manifest.parquet

    Parameters
    ----------
    json_path:
        Path to the JSON annotation file.
    media_root:
        Root directory where files referenced in ``inputs[].path`` live.
    output_dir:
        Destination directory.
    shard_mode:
        ``"size"`` groups samples by target TAR shard size.
        ``"samples"`` groups samples by ``samples_per_shard``.
    shard_size:
        Target on-disk TAR shard size for ``shard_mode="size"``. Accepts bytes
        or strings such as ``"500MB"``, ``"1GB"``, or ``"1024MiB"``.
    samples_per_shard:
        Number of samples grouped in one TAR shard when ``shard_mode="samples"``.
    compression:
        Parquet compression codec (``"zstd"``, ``"snappy"``, ``"gzip"``, or ``None``).
    shard_prefix:
        Prefix for TAR shard file names.
    missing_policy:
        ``"raise"`` — abort if a referenced input file is missing.
        ``"skip"``  — keep sample in Parquet but omit missing file from shard.
    keep_relative_paths_in_parquet:
        If ``True`` (default), Parquet stores original relative paths from JSON.
        If ``False``, Parquet stores resolved absolute paths.
    overwrite:
        Whether to remove and recreate ``output_dir`` if it already exists.

    Returns
    -------
    dict
        Summary with counts of samples, shards, and input files processed.
    """
    json_path = Path(json_path)
    media_root = Path(media_root)
    output_dir = Path(output_dir)
    shards_dir = output_dir / "shards"
    parquet_path = output_dir / "metadata.parquet"

    if output_dir.exists():
        if not overwrite:
            raise FileExistsError(f"{output_dir} already exists. Use overwrite=True.")
        shutil.rmtree(output_dir)

    shards_dir.mkdir(parents=True, exist_ok=True)

    doc = _load_json(json_path)
    samples = doc.get("data", [])
    if not isinstance(samples, list):
        raise ValueError("JSON format error: top-level 'data' must be a list.")
    cleaned_shard_mode = str(shard_mode or "size").strip().lower()
    if cleaned_shard_mode not in {"size", "samples"}:
        raise ValueError("shard_mode must be either 'size' or 'samples'.")
    cleaned_shard_size = parse_shard_size(shard_size)
    cleaned_samples_per_shard = int(samples_per_shard or 100)
    if cleaned_samples_per_shard < 1:
        raise ValueError("samples_per_shard must be >= 1.")

    parquet_rows: List[Dict[str, Any]] = []
    shard_manifest: List[Dict[str, Any]] = []
    total_input_files_added = 0
    total_missing_input_files = 0
    header = {k: v for k, v in doc.items() if k != "data"}

    if cleaned_shard_mode == "size":
        shard_ranges = _build_shard_ranges_by_size(
            samples,
            media_root,
            shard_size=cleaned_shard_size,
            missing_policy=missing_policy,
        )
    else:
        shard_ranges = _build_shard_ranges_by_samples(samples, cleaned_samples_per_shard)
    num_shards = len(shard_ranges)

    for shard_idx, (start, end) in enumerate(shard_ranges):
        shard_samples = samples[start:end]

        shard_name = f"{shard_prefix}-{shard_idx:06d}.tar"
        shard_path = shards_dir / shard_name

        with tarfile.open(shard_path, mode="w") as tar:
            for local_idx, sample in enumerate(shard_samples):
                global_idx = start + local_idx
                sample_id = sample.get("id", f"sample_{global_idx:06d}")
                key = f"{global_idx:09d}"

                row = _flatten_sample_for_parquet(sample, global_idx, shard_pattern=f"{shard_prefix}-%06d.tar")
                row["shard_name"] = shard_name
                row["sample_index"] = global_idx
                row["header"] = json_dumps_compact(header)
                parquet_rows.append(row)

                _add_bytes_to_tar(tar, f"{key}.json", _build_sidecar_metadata(sample))

                for input_idx, input_item in enumerate(_extract_inputs_with_path(sample)):
                    rel_path = str(input_item["path"])
                    resolved = _resolve_media_path(media_root, rel_path, missing_policy=missing_policy)
                    input_type = str(input_item.get("type", "")).strip()

                    if resolved is None:
                        total_missing_input_files += 1
                        shard_manifest.append({
                            "sample_id": sample_id,
                            "shard_name": shard_name,
                            "input_index": input_idx,
                            "file_role": "primary",
                            "input_type": input_type,
                            "relative_path": rel_path,
                            "resolved_path": None,
                            "status": "missing",
                        })
                    else:
                        ext = resolved.suffix.lstrip(".").lower() or "bin"
                        arcname = f"{key}.{input_idx}.{ext}"
                        _add_file_to_tar(tar, resolved, arcname)
                        shard_manifest.append({
                            "sample_id": sample_id,
                            "shard_name": shard_name,
                            "input_index": input_idx,
                            "file_role": "primary",
                            "input_type": input_type,
                            "relative_path": rel_path,
                            "resolved_path": str(resolved if not keep_relative_paths_in_parquet else rel_path),
                            "status": "ok",
                            "wds_member": arcname,
                        })
                        total_input_files_added += 1

                    # player_joints_h5 / player_centroids_h5 inputs may carry a
                    # sidecar ball_path pointing at a separate ball-tracking h5
                    # file; bundle it into the same shard, tagged with a distinct
                    # arcname (".ball.") so it doesn't collide with the primary
                    # input's tar member.
                    ball_rel_path = str(input_item.get("ball_path") or "").strip()
                    if not ball_rel_path:
                        continue
                    ball_resolved = _resolve_media_path(media_root, ball_rel_path, missing_policy=missing_policy)
                    if ball_resolved is None:
                        total_missing_input_files += 1
                        shard_manifest.append({
                            "sample_id": sample_id,
                            "shard_name": shard_name,
                            "input_index": input_idx,
                            "file_role": "ball",
                            "input_type": input_type,
                            "relative_path": ball_rel_path,
                            "resolved_path": None,
                            "status": "missing",
                        })
                        continue

                    ball_ext = ball_resolved.suffix.lstrip(".").lower() or "bin"
                    ball_arcname = f"{key}.{input_idx}.ball.{ball_ext}"
                    _add_file_to_tar(tar, ball_resolved, ball_arcname)
                    shard_manifest.append({
                        "sample_id": sample_id,
                        "shard_name": shard_name,
                        "input_index": input_idx,
                        "file_role": "ball",
                        "input_type": input_type,
                        "relative_path": ball_rel_path,
                        "resolved_path": str(ball_resolved if not keep_relative_paths_in_parquet else ball_rel_path),
                        "status": "ok",
                        "wds_member": ball_arcname,
                    })
                    total_input_files_added += 1

    df = pd.DataFrame(parquet_rows)

    if not keep_relative_paths_in_parquet:
        manifest_df = pd.DataFrame(shard_manifest)
        if not manifest_df.empty:
            ok_manifest = manifest_df[manifest_df["status"] == "ok"].copy()
            if not ok_manifest.empty:
                ok_manifest["input_index"] = ok_manifest["input_index"].astype(int)
                primary_manifest = ok_manifest[ok_manifest["file_role"] == "primary"]
                ball_manifest = ok_manifest[ok_manifest["file_role"] == "ball"]

                by_sample_input_paths: Dict[str, Dict[int, str]] = {}
                for sample_id, sample_manifest in primary_manifest.groupby("sample_id", sort=False):
                    sample_manifest = sample_manifest.sort_values("input_index")
                    by_sample_input_paths[str(sample_id)] = {
                        int(rec["input_index"]): str(rec["resolved_path"])
                        for _, rec in sample_manifest.iterrows()
                    }

                by_sample_ball_paths: Dict[str, Dict[int, str]] = {}
                for sample_id, sample_manifest in ball_manifest.groupby("sample_id", sort=False):
                    sample_manifest = sample_manifest.sort_values("input_index")
                    by_sample_ball_paths[str(sample_id)] = {
                        int(rec["input_index"]): str(rec["resolved_path"])
                        for _, rec in sample_manifest.iterrows()
                    }

                def _resolved_payload_for_row(row: pd.Series) -> str:
                    payload = maybe_json_loads(row.get("sample_payload"), {})
                    if not isinstance(payload, dict):
                        return json_dumps_compact(payload)

                    payload_copy = dict(payload)
                    inputs_value = payload_copy.get("inputs", [])
                    if not isinstance(inputs_value, list):
                        return json_dumps_compact(payload_copy)

                    resolved_by_index = by_sample_input_paths.get(str(row.get("sample_id")), {})
                    resolved_ball_by_index = by_sample_ball_paths.get(str(row.get("sample_id")), {})
                    with_path_idx = 0
                    rewritten: List[Any] = []
                    for inp in inputs_value:
                        if not isinstance(inp, dict):
                            rewritten.append(inp)
                            continue
                        inp_copy = dict(inp)
                        if inp_copy.get("path"):
                            resolved = resolved_by_index.get(with_path_idx)
                            if resolved is not None:
                                inp_copy["path"] = resolved
                            if inp_copy.get("ball_path"):
                                resolved_ball = resolved_ball_by_index.get(with_path_idx)
                                if resolved_ball is not None:
                                    inp_copy["ball_path"] = resolved_ball
                            with_path_idx += 1
                        rewritten.append(inp_copy)
                    payload_copy["inputs"] = rewritten
                    return json_dumps_compact(payload_copy)

                df["sample_payload"] = df.apply(_resolved_payload_for_row, axis=1)

    df.to_parquet(parquet_path, index=False, compression=compression)

    manifest_path = output_dir / "shard_manifest.parquet"
    pd.DataFrame(shard_manifest).to_parquet(manifest_path, index=False, compression=compression)

    return {
        "json_path": str(json_path),
        "media_root": str(media_root),
        "output_dir": str(output_dir),
        "parquet_path": str(parquet_path),
        "manifest_path": str(manifest_path),
        "shards_dir": str(shards_dir),
        "num_samples": len(samples),
        "num_shards": num_shards,
        "shard_mode": cleaned_shard_mode,
        "shard_size": cleaned_shard_size,
        "samples_per_shard": cleaned_samples_per_shard,
        "input_files_added": total_input_files_added,
        "missing_input_files": total_missing_input_files,
    }

Convert a dataset stored as

• metadata.parquet
• WebDataset TAR shards (created by osl_json_to_parquet.py)

back into an OpenSportsLib-style JSON file.

Notes

• Full reconstruction prefers the per-sample sidecar JSON stored inside each TAR shard.
• Current metadata.parquet files also store a full sample_payload copy, so annotation-only reconstruction does not need to fetch any shard.
• By default, reconstructed inputs[].path values remain the original relative paths. Pass extract_media=True to also extract the input files from the shards.

Public entry point

convert_parquet_to_json(...)

convert_parquet_metadata_to_json(metadata_path, output_json_path, *, json_indent=2)

Reconstruct OSL JSON using only a WebDataset metadata.parquet file.

Metadata-only reconstruction requires the current schema's full-fidelity sample_payload column. No TAR shards are opened or required.

Source code in opensportslib/tools/parquet_to_osl_json.py

def convert_parquet_metadata_to_json(
    metadata_path: str | Path,
    output_json_path: str | Path,
    *,
    json_indent: int = 2,
) -> Dict[str, Any]:
    """Reconstruct OSL JSON using only a WebDataset ``metadata.parquet`` file.

    Metadata-only reconstruction requires the current schema's full-fidelity
    ``sample_payload`` column.  No TAR shards are opened or required.
    """
    metadata_path = Path(metadata_path)
    output_json_path = Path(output_json_path)
    if not metadata_path.is_file():
        raise FileNotFoundError(f"Missing metadata file: {metadata_path}")

    df = (
        pd.read_parquet(metadata_path)
        .sort_values("sample_index")
        .reset_index(drop=True)
    )
    if len(df) == 0:
        raise ValueError("metadata.parquet is empty.")

    required_cols = {"sample_id", "sample_index", "header", "sample_payload"}
    missing = sorted(required_cols.difference(df.columns))
    if missing:
        raise ValueError(
            "metadata.parquet cannot be converted without shards; missing required "
            f"columns: {', '.join(missing)}"
        )

    top_doc = maybe_json_loads(df.iloc[0].get("header"), None)
    if not isinstance(top_doc, dict):
        raise ValueError(
            "metadata.parquet has invalid 'header' payload; expected JSON object."
        )
    top_doc = dict(top_doc)
    top_doc.pop("data", None)
    top_doc["data"] = [
        _reconstruct_sample_from_row(row, sidecar=None)
        for _, row in df.iterrows()
    ]

    output_json_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_json_path, "w", encoding="utf-8") as handle:
        json.dump(top_doc, handle, ensure_ascii=False, indent=json_indent)
        handle.write("\n")

    return {
        "metadata_path": str(metadata_path),
        "output_json_path": str(output_json_path),
        "num_samples": len(top_doc["data"]),
        "extract_media": False,
        "extracted_input_files": 0,
        "extracted_media_files": 0,
    }

convert_parquet_to_json(dataset_dir, output_json_path, *, extract_media=False, output_media_root=None, overwrite_media=False, json_indent=2)

Convert a Parquet + WebDataset directory back to an OSL-style JSON file.

The input directory must contain metadata.parquet and a shards/ sub-directory produced by :func:opensportslib.tools.convert_json_to_parquet.

Parameters

dataset_dir: Directory produced by the forward converter. output_json_path: Destination JSON file path. extract_media: If True, extract media files from TAR shards into output_media_root. inputs[].path values are preserved as-is (original relative paths). output_media_root: Root directory for extracted media. Required when extract_media=True. overwrite_media: Whether to overwrite already-extracted media files. json_indent: Indentation level for the output JSON.

Returns

dict Summary with sample count and extraction statistics.

Source code in opensportslib/tools/parquet_to_osl_json.py

def convert_parquet_to_json(
    dataset_dir: str | Path,
    output_json_path: str | Path,
    *,
    extract_media: bool = False,
    output_media_root: Optional[str | Path] = None,
    overwrite_media: bool = False,
    json_indent: int = 2,
) -> Dict[str, Any]:
    """
    Convert a Parquet + WebDataset directory back to an OSL-style JSON file.

    The input directory must contain ``metadata.parquet`` and a ``shards/``
    sub-directory produced by :func:`opensportslib.tools.convert_json_to_parquet`.

    Parameters
    ----------
    dataset_dir:
        Directory produced by the forward converter.
    output_json_path:
        Destination JSON file path.
    extract_media:
        If ``True``, extract media files from TAR shards into ``output_media_root``.
        ``inputs[].path`` values are preserved as-is (original relative paths).
    output_media_root:
        Root directory for extracted media. Required when ``extract_media=True``.
    overwrite_media:
        Whether to overwrite already-extracted media files.
    json_indent:
        Indentation level for the output JSON.

    Returns
    -------
    dict
        Summary with sample count and extraction statistics.
    """
    dataset_dir = Path(dataset_dir)
    output_json_path = Path(output_json_path)
    metadata_path = dataset_dir / "metadata.parquet"
    shards_dir = dataset_dir / "shards"

    if not metadata_path.exists():
        raise FileNotFoundError(f"Missing metadata file: {metadata_path}")
    if not shards_dir.exists():
        raise FileNotFoundError(f"Missing shards directory: {shards_dir}")
    if extract_media and output_media_root is None:
        raise ValueError("output_media_root must be provided when extract_media=True")

    output_media_root_path = Path(output_media_root) if output_media_root is not None else None
    if output_media_root_path is not None:
        output_media_root_path.mkdir(parents=True, exist_ok=True)

    df = pd.read_parquet(metadata_path).sort_values("sample_index").reset_index(drop=True)
    if len(df) == 0:
        raise ValueError("metadata.parquet is empty.")

    required_cols = {"sample_id", "sample_index", "shard_name", "header"}
    missing = sorted(required_cols.difference(df.columns))
    if missing:
        if "header" in missing:
            raise ValueError(
                "metadata.parquet is missing required 'header' column. "
                "Legacy schemas are no longer supported."
            )
        raise ValueError(f"metadata.parquet is missing required columns: {', '.join(missing)}")

    top_doc = maybe_json_loads(df.iloc[0].get("header"), None)
    if not isinstance(top_doc, dict):
        raise ValueError("metadata.parquet has invalid 'header' payload; expected JSON object.")
    top_doc = dict(top_doc)
    top_doc.pop("data", None)

    data: List[Dict[str, Any]] = []
    extracted_media_count = 0

    for _, row in df.iterrows():
        sample_index = int(row["sample_index"])
        shard_name = str(row["shard_name"]) if row.get("shard_name") is not None else ""
        if not shard_name:
            raise ValueError(f"Invalid shard_name for sample_index={sample_index}.")

        tar_path = shards_dir / shard_name
        if not tar_path.exists():
            raise FileNotFoundError(f"Missing shard file: {tar_path}")

        sidecar = _read_sidecar_json_from_tar(tar_path, sample_index)
        sample = _reconstruct_sample_from_row(row, sidecar=sidecar)

        if extract_media:
            inputs = sample.get("inputs", []) if isinstance(sample, dict) else []
            original_input_paths: List[str] = []
            original_ball_paths: List[Optional[str]] = []
            for inp in inputs:
                if not isinstance(inp, dict) or not inp.get("path"):
                    continue
                original_input_paths.append(str(inp["path"]))
                ball_path = inp.get("ball_path")
                original_ball_paths.append(str(ball_path) if ball_path else None)
            extracted_media_count += _extract_sample_media_from_tar(
                tar_path=tar_path,
                sample_index=sample_index,
                output_media_root=output_media_root_path,
                original_paths=original_input_paths,
                original_ball_paths=original_ball_paths,
                overwrite=overwrite_media,
            )

        data.append(sample)

    top_doc["data"] = data

    output_json_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_json_path, "w", encoding="utf-8") as f:
        json.dump(top_doc, f, ensure_ascii=False, indent=json_indent)

    return {
        "dataset_dir": str(dataset_dir),
        "output_json_path": str(output_json_path),
        "num_samples": len(data),
        "extract_media": extract_media,
        "output_media_root": str(output_media_root_path) if output_media_root_path else None,
        "extracted_input_files": extracted_media_count,
        "extracted_media_files": extracted_media_count,
    }

HfTransferCancelled

Bases: RuntimeError

Source code in opensportslib/tools/hf_transfer.py

class HfTransferCancelled(RuntimeError):
    pass

MissingDatasetInputsError

Bases: FileNotFoundError

Raised before upload when dataset references are absent locally.

Source code in opensportslib/tools/hf_transfer.py

class MissingDatasetInputsError(FileNotFoundError):
    """Raised before upload when dataset references are absent locally."""

    def __init__(self, missing_inputs: list[dict[str, str]]):
        self.missing_inputs = list(missing_inputs)
        preview = ", ".join(
            item["path"] for item in self.missing_inputs[:5]
        )
        if len(self.missing_inputs) > 5:
            preview += f", and {len(self.missing_inputs) - 5} more"
        super().__init__(
            f"Dataset upload blocked: {len(self.missing_inputs)} referenced input "
            f"file(s) are missing locally: {preview}"
        )

download_dataset_split_from_hf(repo_id, revision, split, output_dir, *, download_format='parquet', annotations_only=False, dry_run=False, token=None, progress_cb=None, byte_progress_cb=None, file_plan_cb=None, file_completed_cb=None, json_ready_cb=None, is_cancelled=None)

Source code in opensportslib/tools/hf_transfer.py

def download_dataset_split_from_hf(
    repo_id: str,
    revision: str,
    split: str,
    output_dir: str,
    *,
    download_format: str = "parquet",
    annotations_only: bool = False,
    dry_run: bool = False,
    token: str | None = None,
    progress_cb: ProgressCallback | None = None,
    byte_progress_cb: ByteProgressCallback | None = None,
    file_plan_cb: FilePlanCallback | None = None,
    file_completed_cb: FileCompletedCallback | None = None,
    json_ready_cb: JsonReadyCallback | None = None,
    is_cancelled: CancelCheck | None = None,
) -> dict[str, Any]:
    cleaned_repo_id = str(repo_id or "").strip()
    cleaned_revision = str(revision or "").strip() or "main"
    cleaned_split = _clean_hf_split(split)
    cleaned_format = str(download_format or "parquet").strip().lower()
    if cleaned_format not in {"json", "parquet"}:
        raise ValueError("download_format must be 'json' or 'parquet'.")
    if not cleaned_repo_id:
        raise ValueError("repo_id is required.")
    if annotations_only and dry_run:
        raise ValueError("dry_run cannot be combined with annotations_only.")
    split_output_dir = _build_split_output_dir(output_dir, cleaned_revision, cleaned_split)
    resolved_commit = _resolve_hf_commit(cleaned_repo_id, cleaned_revision, token)

    if cleaned_format == "parquet":
        if dry_run:
            raise ValueError("dry_run is only supported for JSON split downloads.")
        return _download_parquet_split_and_convert(
            cleaned_repo_id,
            cleaned_revision,
            cleaned_split,
            split_output_dir,
            commit=resolved_commit,
            annotations_only=annotations_only,
            token=token,
            progress_cb=progress_cb,
            byte_progress_cb=byte_progress_cb,
            file_plan_cb=file_plan_cb,
            file_completed_cb=file_completed_cb,
            json_ready_cb=json_ready_cb,
            is_cancelled=is_cancelled,
        )

    return _download_json_path_from_hf(
        cleaned_repo_id,
        cleaned_revision,
        cleaned_split,
        split_output_dir,
        commit=resolved_commit,
        annotations_only=annotations_only,
        dry_run=dry_run,
        token=token,
        progress_cb=progress_cb,
        byte_progress_cb=byte_progress_cb,
        file_plan_cb=file_plan_cb,
        file_completed_cb=file_completed_cb,
        json_ready_cb=json_ready_cb,
        is_cancelled=is_cancelled,
    )

download_dataset_sample_inputs_from_hf(dataset_json_path, sample_id, *, input_path=None, overwrite=False, token=None, progress_cb=None, byte_progress_cb=None, is_cancelled=None)

Download one sample or input from the immutable HF source of a local JSON.

Source code in opensportslib/tools/hf_transfer.py

def download_dataset_sample_inputs_from_hf(
    dataset_json_path: str,
    sample_id: str,
    *,
    input_path: str | None = None,
    overwrite: bool = False,
    token: str | None = None,
    progress_cb: ProgressCallback | None = None,
    byte_progress_cb: ByteProgressCallback | None = None,
    is_cancelled: CancelCheck | None = None,
) -> dict[str, Any]:
    """Download one sample or input from the immutable HF source of a local JSON."""
    cleaned_json_path = Path(str(dataset_json_path or "").strip()).resolve()
    if not cleaned_json_path.is_file():
        raise ValueError(f"Dataset JSON does not exist: {cleaned_json_path}")
    with open(cleaned_json_path, "r", encoding="utf-8") as handle:
        dataset_json = json.load(handle)
    if not isinstance(dataset_json, dict):
        raise ValueError("Invalid dataset JSON: expected root object.")

    source = read_hf_source_metadata_from_dataset(dataset_json)
    missing_source = [key for key in ("repo_id", "split", "format", "commit") if not source.get(key)]
    if missing_source:
        raise ValueError(
            "Dataset lacks selective-download provenance (missing: "
            f"{', '.join(missing_source)}). Re-download its annotations with a newer OpenSportsLib checkout."
        )
    if source["format"] not in {"json", "parquet"}:
        raise ValueError(f"Unsupported Hugging Face dataset format: {source['format']}")

    requested_paths = _sample_requested_asset_paths(dataset_json, sample_id, input_path)
    common = {
        "repo_id": source["repo_id"],
        "split": source["split"],
        "commit": source["commit"],
        "sample_id": str(sample_id),
        "input_path": input_path,
        "overwrite": bool(overwrite),
        "token": token,
        "progress_cb": progress_cb,
        "byte_progress_cb": byte_progress_cb,
        "is_cancelled": is_cancelled,
    }
    _ensure_not_cancelled(is_cancelled)
    if source["format"] == "json":
        return _download_json_sample_inputs(cleaned_json_path.parent, requested_paths, **common)
    return _download_parquet_sample_inputs(cleaned_json_path.parent, requested_paths, **common)

download_dataset_missing_inputs_from_hf(dataset_json_path, *, token=None, progress_cb=None, byte_progress_cb=None, is_cancelled=None)

Hydrate all locally missing inputs from the JSON's pinned HF source.

Source code in opensportslib/tools/hf_transfer.py

def download_dataset_missing_inputs_from_hf(
    dataset_json_path: str,
    *,
    token: str | None = None,
    progress_cb: ProgressCallback | None = None,
    byte_progress_cb: ByteProgressCallback | None = None,
    is_cancelled: CancelCheck | None = None,
) -> dict[str, Any]:
    """Hydrate all locally missing inputs from the JSON's pinned HF source."""
    initial_missing = find_missing_dataset_inputs(dataset_json_path)
    sample_ids = list(
        dict.fromkeys(item["sample_id"] for item in initial_missing if item["sample_id"])
    )
    if initial_missing and len(sample_ids) != len(
        {item["sample_id"] for item in initial_missing}
    ):
        raise ValueError("Every sample with a missing input must have an id.")

    results = []
    total = len(sample_ids)
    for index, sample_id in enumerate(sample_ids, start=1):
        _ensure_not_cancelled(is_cancelled)

        def _scoped_progress(
            message: str,
            _index: int = index,
            _sample_id: str = sample_id,
        ) -> None:
            _emit_progress(
                progress_cb,
                f"[{_index}/{total}] {_sample_id}: {message}",
            )

        results.append(
            download_dataset_sample_inputs_from_hf(
                dataset_json_path,
                sample_id,
                overwrite=False,
                token=token,
                progress_cb=_scoped_progress,
                byte_progress_cb=byte_progress_cb,
                is_cancelled=is_cancelled,
            )
        )

    remaining_missing = find_missing_dataset_inputs(dataset_json_path)
    return {
        "operation": "missing_assets",
        "dataset_json_path": str(Path(dataset_json_path).resolve()),
        "initial_missing": initial_missing,
        "initial_missing_count": len(initial_missing),
        "remaining_missing": remaining_missing,
        "remaining_missing_count": len(remaining_missing),
        "sample_results": results,
        "requested_downloaded_count": sum(
            int(result.get("requested_downloaded_count") or 0) for result in results
        ),
        "opportunistic_downloaded_count": sum(
            int(result.get("opportunistic_downloaded_count") or 0)
            for result in results
        ),
        "failed_count": sum(int(result.get("failed_count") or 0) for result in results),
    }

find_missing_dataset_inputs(dataset_json_path)

Return referenced primary and companion inputs that are absent locally.

Source code in opensportslib/tools/hf_transfer.py

def find_missing_dataset_inputs(dataset_json_path: str) -> list[dict[str, str]]:
    """Return referenced primary and companion inputs that are absent locally."""
    cleaned_json_path = Path(str(dataset_json_path or "").strip()).resolve()
    if not cleaned_json_path.is_file():
        raise ValueError(f"Dataset JSON does not exist: {cleaned_json_path}")
    with open(cleaned_json_path, "r", encoding="utf-8") as handle:
        payload = json.load(handle)
    data_items = payload.get("data", []) if isinstance(payload, dict) else []
    if not isinstance(data_items, list):
        raise ValueError("Invalid dataset JSON: expected top-level 'data' list.")

    missing: list[dict[str, str]] = []
    seen: set[tuple[str, str, str]] = set()
    for sample in data_items:
        if not isinstance(sample, dict):
            continue
        sample_id = str(sample.get("id") or "").strip()
        for input_item in sample.get("inputs", []) or []:
            if not isinstance(input_item, dict):
                continue
            primary_path = str(input_item.get("path") or "").strip()
            for role, raw_path in (
                ("primary", primary_path),
                ("ball", str(input_item.get("ball_path") or "").strip()),
            ):
                if not raw_path:
                    continue
                local_path = (
                    Path(raw_path)
                    if os.path.isabs(raw_path)
                    else cleaned_json_path.parent / raw_path
                ).resolve(strict=False)
                if local_path.is_file():
                    continue
                key = (sample_id, role, raw_path.replace("\\", "/"))
                if key in seen:
                    continue
                seen.add(key)
                missing.append(
                    {
                        "sample_id": sample_id,
                        "input_path": primary_path,
                        "path": raw_path,
                        "role": role,
                        "local_path": str(local_path),
                    }
                )
    return missing

upload_dataset_inputs_from_json_to_hf(repo_id, json_path, *, revision='main', split=None, commit_message=None, token=None, progress_cb=None, is_cancelled=None)

Source code in opensportslib/tools/hf_transfer.py

def upload_dataset_inputs_from_json_to_hf(
    repo_id: str,
    json_path: str,
    *,
    revision: str | None = "main",
    split: str | None = None,
    commit_message: str | None = None,
    token: str | None = None,
    progress_cb: ProgressCallback | None = None,
    is_cancelled: CancelCheck | None = None,
) -> dict[str, Any]:
    HfApi, _, _ = _import_hf_hub()
    CommitOperationAdd = _import_hf_commit_operation_add()

    cleaned_repo_id = str(repo_id or "").strip()
    cleaned_json_path = os.path.abspath(str(json_path or "").strip())
    if not cleaned_repo_id:
        raise ValueError("repo_id is required.")
    if not cleaned_json_path:
        raise ValueError("json_path is required.")
    if not os.path.isfile(cleaned_json_path):
        raise ValueError(f"JSON file does not exist: {cleaned_json_path}")
    cleaned_revision = str(revision or "").strip() or "main"

    effective_commit_message = (commit_message or "").strip() or "Upload dataset inputs from JSON"
    input_upload_entries, missing_input_entries = _collect_local_input_upload_entries_from_json(
        cleaned_json_path
    )
    missing_input_paths = list(
        dict.fromkeys(entry["path_in_repo"] for entry in missing_input_entries)
    )
    if missing_input_paths:
        _emit_progress(
            progress_cb,
            (
                f"Skipping {len(missing_input_paths)} referenced input files that "
                "are not available locally."
            ),
        )
    unique_input_entries: list[dict[str, str]] = []
    input_entry_by_repo_path: dict[str, dict[str, str]] = {}
    duplicate_input_refs = 0
    for entry in input_upload_entries:
        path_in_repo = entry["path_in_repo"]
        existing = input_entry_by_repo_path.get(path_in_repo)
        if existing is None:
            input_entry_by_repo_path[path_in_repo] = entry
            unique_input_entries.append(entry)
            continue

        duplicate_input_refs += 1
        if os.path.abspath(existing["local_path"]) != os.path.abspath(entry["local_path"]):
            raise ValueError(
                "Conflicting local files mapped to the same repo path "
                f"'{path_in_repo}': '{existing['local_path']}' vs '{entry['local_path']}'."
            )

    if duplicate_input_refs:
        _emit_progress(
            progress_cb,
            (
                f"Deduplicated {duplicate_input_refs} repeated input references "
                f"into {len(unique_input_entries)} unique repo paths."
            ),
        )

    cleaned_split = _clean_hf_split(split) if split else ""
    json_path_in_repo = f"{cleaned_split}.json" if cleaned_split else _normalize_repo_path(os.path.basename(cleaned_json_path))
    if not json_path_in_repo:
        json_path_in_repo = "dataset.json"

    existing_repo_paths = {entry["path_in_repo"] for entry in unique_input_entries}
    json_already_listed_in_inputs = json_path_in_repo in existing_repo_paths
    upload_entries = [
        {
            "local_path": cleaned_json_path,
            "path_in_repo": json_path_in_repo,
        }
    ]
    # Always place dataset JSON first in the commit operations list.
    upload_entries.extend(
        entry for entry in unique_input_entries if entry["path_in_repo"] != json_path_in_repo
    )

    _ensure_not_cancelled(is_cancelled)
    api = HfApi(token=token or None)
    _emit_progress(
        progress_cb,
        (
            f"Preparing batched upload of {len(upload_entries)} files to {cleaned_repo_id}@{cleaned_revision} "
            f"(dataset JSON + {len(unique_input_entries)} unique inputs) from {cleaned_json_path}"
        ),
    )

    operations = []
    for idx, entry in enumerate(upload_entries, start=1):
        _ensure_not_cancelled(is_cancelled)
        _emit_progress(progress_cb, f"[{idx}/{len(upload_entries)}] Queueing {entry['path_in_repo']}")
        operations.append(
            CommitOperationAdd(
                path_in_repo=entry["path_in_repo"],
                path_or_fileobj=entry["local_path"],
            )
        )

    _ensure_not_cancelled(is_cancelled)
    _emit_progress(progress_cb, f"Submitting one Hugging Face commit with {len(operations)} files...")
    commit_info = api.create_commit(
        repo_id=cleaned_repo_id,
        repo_type="dataset",
        revision=cleaned_revision,
        operations=operations,
        commit_message=effective_commit_message,
    )
    commit_ref = (
        str(getattr(commit_info, "oid", "") or "").strip()
        or str(getattr(commit_info, "commit_id", "") or "").strip()
        or str(getattr(commit_info, "commit_url", "") or "").strip()
        or str(commit_info)
    )

    _emit_progress(progress_cb, f"Upload completed in one commit. Uploaded {len(upload_entries)} files.")
    return {
        "repo_id": cleaned_repo_id,
        "repo_type": "dataset",
        "upload_kind": "json",
        "json_path": cleaned_json_path,
        "revision": cleaned_revision,
        "split": cleaned_split,
        "json_path_in_repo": json_path_in_repo,
        "input_file_count": len(input_upload_entries),
        "unique_input_file_count": len(unique_input_entries),
        "skipped_missing_input_count": len(missing_input_paths),
        "skipped_missing_input_paths": missing_input_paths,
        "uploaded_file_count": len(upload_entries),
        "uploaded_json_separately": not json_already_listed_in_inputs,
        "commit_message": effective_commit_message,
        "commit_ref": commit_ref,
    }

upload_dataset_as_parquet_to_hf(repo_id, json_path, *, revision='main', split=None, commit_message=None, shard_mode='size', shard_size=DEFAULT_SHARD_SIZE, samples_per_shard=100, token=None, progress_cb=None, is_cancelled=None)

Convert an OSL JSON dataset to Parquet + WebDataset format and upload the result to a HuggingFace dataset repository under a folder named after the JSON file stem.

For example, annotations_test.json is converted and uploaded to the annotations_test/ folder on the repository.

A temporary directory is used for the conversion output and removed when done.

Source code in opensportslib/tools/hf_transfer.py

def upload_dataset_as_parquet_to_hf(
    repo_id: str,
    json_path: str,
    *,
    revision: str | None = "main",
    split: str | None = None,
    commit_message: str | None = None,
    shard_mode: str = "size",
    shard_size: int | str = DEFAULT_SHARD_SIZE,
    samples_per_shard: int = 100,
    token: str | None = None,
    progress_cb: ProgressCallback | None = None,
    is_cancelled: CancelCheck | None = None,
) -> dict[str, Any]:
    """
    Convert an OSL JSON dataset to Parquet + WebDataset format and upload the result
    to a HuggingFace dataset repository under a folder named after the JSON file stem.

    For example, ``annotations_test.json`` is converted and uploaded to the
    ``annotations_test/`` folder on the repository.

    A temporary directory is used for the conversion output and removed when done.
    """
    cleaned_repo_id = str(repo_id or "").strip()
    cleaned_json_path = os.path.abspath(str(json_path or "").strip())
    if not cleaned_repo_id:
        raise ValueError("repo_id is required.")
    if not os.path.isfile(cleaned_json_path):
        raise ValueError(f"JSON file does not exist: {cleaned_json_path}")
    missing_inputs = find_missing_dataset_inputs(cleaned_json_path)
    if missing_inputs:
        raise MissingDatasetInputsError(missing_inputs)

    HfApi, _, _ = _import_hf_hub()
    CommitOperationAdd = _import_hf_commit_operation_add()
    CommitOperationDelete = _import_hf_commit_operation_delete()

    cleaned_revision = str(revision or "").strip() or "main"
    effective_commit_message = (commit_message or "").strip() or "Upload dataset as Parquet + WebDataset"
    cleaned_shard_mode = str(shard_mode or "size").strip().lower()
    if cleaned_shard_mode not in {"size", "samples"}:
        raise ValueError("shard_mode must be either 'size' or 'samples'.")
    cleaned_shard_size = parse_shard_size(shard_size)
    cleaned_samples_per_shard = int(samples_per_shard or 100)
    if cleaned_samples_per_shard < 1:
        raise ValueError("samples_per_shard must be >= 1.")
    cleaned_split = _clean_hf_split(split) if split else ""
    folder_name = cleaned_split or Path(cleaned_json_path).stem
    media_root = Path(cleaned_json_path).parent

    _ensure_not_cancelled(is_cancelled)
    shard_desc = (
        f"shard_size={cleaned_shard_size}"
        if cleaned_shard_mode == "size"
        else f"samples_per_shard={cleaned_samples_per_shard}"
    )
    _emit_progress(
        progress_cb,
        f"Converting {cleaned_json_path} to Parquet + WebDataset ({shard_desc})...",
    )

    conversion_result: dict[str, Any] = {}
    total = 0
    commit_ref = ""
    tmp_dir = tempfile.mkdtemp(prefix="hf_parquet_ul_")
    try:
        parquet_output = Path(tmp_dir) / folder_name
        conversion_result = convert_json_to_parquet(
            json_path=cleaned_json_path,
            media_root=media_root,
            output_dir=parquet_output,
            shard_mode=cleaned_shard_mode,
            shard_size=cleaned_shard_size,
            samples_per_shard=cleaned_samples_per_shard,
            missing_policy="raise",
            overwrite=True,
        )

        _ensure_not_cancelled(is_cancelled)

        # Collect all files to upload, preserving sub-paths under folder_name/
        upload_entries: list[dict[str, str]] = []
        for local_file in sorted(parquet_output.rglob("*")):
            if not local_file.is_file():
                continue
            rel = local_file.relative_to(tmp_dir).as_posix()
            upload_entries.append({"local_path": str(local_file), "path_in_repo": rel})

        api = HfApi(token=token or None)
        total = len(upload_entries)
        _emit_progress(
            progress_cb,
            f"Preparing batched parquet upload of {total} files to {cleaned_repo_id}@{cleaned_revision} under '{folder_name}/'..."
        )

        generated_repo_paths = {entry["path_in_repo"] for entry in upload_entries}
        existing_repo_paths = set(
            api.list_repo_files(
                cleaned_repo_id,
                revision=cleaned_revision,
                repo_type="dataset",
            )
        )
        target_prefix = f"{folder_name}/"
        stale_repo_paths = sorted(
            path
            for path in existing_repo_paths
            if path.startswith(target_prefix)
            and path not in generated_repo_paths
            and (
                path in {
                    f"{folder_name}/metadata.parquet",
                    f"{folder_name}/shard_manifest.parquet",
                }
                or path.startswith(f"{folder_name}/shards/")
            )
        )

        operations = [
            CommitOperationDelete(path_in_repo=path) for path in stale_repo_paths
        ]
        for idx, entry in enumerate(upload_entries, start=1):
            _ensure_not_cancelled(is_cancelled)
            _emit_progress(progress_cb, f"[{idx}/{total}] Queueing {entry['path_in_repo']}")
            operations.append(
                CommitOperationAdd(
                    path_in_repo=entry["path_in_repo"],
                    path_or_fileobj=entry["local_path"],
                )
            )

        _ensure_not_cancelled(is_cancelled)
        _emit_progress(
            progress_cb,
            "Submitting one Hugging Face commit with "
            f"{total} generated file(s) and {len(stale_repo_paths)} obsolete "
            "shard deletion(s)...",
        )
        commit_info = api.create_commit(
            repo_id=cleaned_repo_id,
            repo_type="dataset",
            revision=cleaned_revision,
            operations=operations,
            commit_message=effective_commit_message,
        )
        commit_ref = (
            str(getattr(commit_info, "oid", "") or "").strip()
            or str(getattr(commit_info, "commit_id", "") or "").strip()
            or str(getattr(commit_info, "commit_url", "") or "").strip()
            or str(commit_info)
        )
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)

    _emit_progress(progress_cb, f"Parquet upload completed. Uploaded {total} files.")
    input_file_count = int(conversion_result.get("input_files_added") or 0)
    return {
        "repo_id": cleaned_repo_id,
        "revision": cleaned_revision,
        "split": cleaned_split,
        "upload_kind": "parquet",
        "json_path": cleaned_json_path,
        "folder_name": folder_name,
        "shard_mode": cleaned_shard_mode,
        "shard_size": cleaned_shard_size,
        "samples_per_shard": cleaned_samples_per_shard,
        "num_shards": int(conversion_result.get("num_shards") or 0),
        "num_samples": int(conversion_result.get("num_samples") or 0),
        "input_file_count": input_file_count,
        "uploaded_file_count": total,
        "deleted_file_count": len(stale_repo_paths),
        "commit_message": effective_commit_message,
        "commit_ref": commit_ref,
    }

create_dataset_repo_on_hf(repo_id, *, token=None, private=False, progress_cb=None)

Source code in opensportslib/tools/hf_transfer.py

def create_dataset_repo_on_hf(
    repo_id: str,
    *,
    token: str | None = None,
    private: bool = False,
    progress_cb: ProgressCallback | None = None,
) -> dict[str, Any]:
    cleaned_repo_id = str(repo_id or "").strip()
    if not cleaned_repo_id:
        raise ValueError("repo_id is required.")

    HfApi, _, _ = _import_hf_hub()
    api = HfApi(token=token or None)

    _emit_progress(progress_cb, f"Creating Hugging Face dataset repository: {cleaned_repo_id}")
    repo_url = api.create_repo(
        repo_id=cleaned_repo_id,
        repo_type="dataset",
        private=bool(private),
        exist_ok=True,
    )
    _emit_progress(progress_cb, f"Repository is ready: {cleaned_repo_id}")

    return {
        "repo_id": cleaned_repo_id,
        "repo_type": "dataset",
        "repo_url": str(repo_url),
    }

dataset_repo_exists_on_hf(repo_id, *, token=None)

Source code in opensportslib/tools/hf_transfer.py

def dataset_repo_exists_on_hf(
    repo_id: str,
    *,
    token: str | None = None,
) -> bool:
    cleaned_repo_id = str(repo_id or "").strip()
    if not cleaned_repo_id:
        raise ValueError("repo_id is required.")

    HfApi, _, _ = _import_hf_hub()
    api = HfApi(token=token or None)
    try:
        api.repo_info(repo_id=cleaned_repo_id, repo_type="dataset")
        return True
    except Exception as exc:
        if is_hf_repo_not_found_error(str(exc)):
            return False
        raise

create_dataset_branch_on_hf(repo_id, branch, *, source_revision='main', token=None, progress_cb=None)

Source code in opensportslib/tools/hf_transfer.py

def create_dataset_branch_on_hf(
    repo_id: str,
    branch: str,
    *,
    source_revision: str = "main",
    token: str | None = None,
    progress_cb: ProgressCallback | None = None,
) -> dict[str, Any]:
    cleaned_repo_id = str(repo_id or "").strip()
    cleaned_branch = str(branch or "").strip()
    cleaned_source_revision = str(source_revision or "").strip() or "main"
    if not cleaned_repo_id:
        raise ValueError("repo_id is required.")
    if not cleaned_branch:
        raise ValueError("branch is required.")

    HfApi, _, _ = _import_hf_hub()
    api = HfApi(token=token or None)

    # Resolve the actual commit to branch from.
    # Always use the oldest (initial) commit so the new branch is never empty —
    # branching from HEAD would create an empty ref on a repo with no files yet.
    _emit_progress(
        progress_cb,
        f"Resolving initial commit for {cleaned_repo_id} to use as branch base...",
    )
    commits = api.list_repo_commits(cleaned_repo_id, repo_type="dataset")
    initial_commit_id = commits[-1].commit_id if commits else cleaned_source_revision

    _emit_progress(
        progress_cb,
        (
            f"Creating Hugging Face dataset branch '{cleaned_branch}' "
            f"from initial commit {initial_commit_id!r} in {cleaned_repo_id}"
        ),
    )
    api.create_branch(
        repo_id=cleaned_repo_id,
        repo_type="dataset",
        branch=cleaned_branch,
        revision=initial_commit_id,
        exist_ok=True,
    )
    _emit_progress(progress_cb, f"Branch is ready: {cleaned_repo_id}@{cleaned_branch}")

    return {
        "repo_id": cleaned_repo_id,
        "repo_type": "dataset",
        "branch": cleaned_branch,
        "source_revision": initial_commit_id,
    }

See extension APIs for builders and runtime modules, and internal catalog for concrete implementation classes.