wherobots-python-sdk 0.1.0__py3-none-any.whl

wherobots/client.py

@@ -0,0 +1,640 @@
+"""Wherobots Jobs API Client.
+
+``WherobotsJob`` is the primary user-facing class.  It orchestrates
+job submission, monitoring, log streaming, and cancellation by
+delegating HTTP work to the ``RunsAPI`` layer.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Iterator
+from typing import Any
+
+from wherobots.api.runs import RunsAPI
+from wherobots.config import WherobotsConfig
+from wherobots.enums import DependencyType, JobStatus, Region, Runtime, is_terminal_status
+from wherobots.exceptions import (
+    WherobotsAPIError,
+    WherobotsJobError,
+    WherobotsTimeoutError,
+    WherobotsValidationError,
+)
+from wherobots.models import (
+    CreateRunPayload,
+    LogsResponse,
+    RunEnvironment,
+    RunJarPayload,
+    RunListPage,
+    RunMetricsResponse,
+    RunPythonPayload,
+    RunView,
+)
+from wherobots.utils.logger import get_logger
+from wherobots.utils.validation import validate_name
+
+logger = get_logger("client")
+
+
+def _status_str(status: JobStatus | str | None) -> str:
+    """Human-readable status string for logging, handling the
+    forward-compat case where *status* is a raw string the enum
+    doesn't know about."""
+    if status is None:
+        return "UNKNOWN"
+    if isinstance(status, JobStatus):
+        return status.value
+    return status
+
+
+class WherobotsJob:
+    """Wherobots Job Run Manager.
+
+    Manages the lifecycle of Wherobots job runs including submission,
+    monitoring, log streaming, and cancellation.
+    """
+
+    def __init__(
+        self,
+        script: str,
+        name: str,
+        runtime: str | Runtime = "tiny",
+        region: str | Region | None = None,
+        api_key: str | None = None,
+        version: str | None = None,
+        timeout_seconds: int = 3600,
+        args: list[str] | None = None,
+        spark_configs: dict[str, str] | None = None,
+        dependencies: list[dict[str, Any]] | None = None,
+        spark_driver_disk_gb: int | None = None,
+        spark_executor_disk_gb: int | None = None,
+        s3_bucket: str | None = None,
+        s3_prefix: str | None = None,
+        jar_main_class: str | None = None,
+        auto_upload: bool = True,
+        base_url: str | None = None,
+        request_timeout_seconds: int | None = None,
+        config: WherobotsConfig | None = None,
+    ):
+        """Initialize a Wherobots Job.
+
+        Args:
+            script: Path to Python script (.py) or JAR file (.jar).
+                   Can be local path or S3 URI (s3://bucket/key)
+            name: Job name (8-255 chars, alphanumeric, _, -, .)
+            runtime: Compute runtime size (default: "tiny")
+            region: AWS region (default: from config or env)
+            api_key: Wherobots API key (or set WHEROBOTS_API_KEY env var)
+            version: Wherobots version ("latest" or "preview")
+            timeout_seconds: Job timeout in seconds (default: 3600)
+            args: Optional list of arguments to pass to script
+            spark_configs: Optional Spark configuration dict
+            dependencies: Optional list of PyPI or file dependencies
+            spark_driver_disk_gb: Optional driver disk size
+            spark_executor_disk_gb: Optional executor disk size
+            s3_bucket: Deprecated. Ignored; presigned uploads are used
+                instead. Will be removed in a future release.
+            s3_prefix: Deprecated. Ignored; presigned uploads are used
+                instead. Will be removed in a future release.
+            jar_main_class: Main class for JAR files (required for JARs)
+            auto_upload: Automatically upload local scripts via presigned
+                URL (default: True)
+            base_url: Override the API base URL
+            request_timeout_seconds: HTTP request timeout in seconds
+            config: Optional WherobotsConfig to override defaults
+
+        Raises:
+            WherobotsValidationError: If script is None/empty, disk sizes
+                are negative, or jar_main_class is missing for JARs.
+        """
+        if not script:
+            raise WherobotsValidationError("script must not be None or empty")
+
+        if spark_driver_disk_gb is not None and spark_driver_disk_gb < 0:
+            raise WherobotsValidationError(
+                f"spark_driver_disk_gb must be non-negative, got {spark_driver_disk_gb}"
+            )
+        if spark_executor_disk_gb is not None and spark_executor_disk_gb < 0:
+            raise WherobotsValidationError(
+                f"spark_executor_disk_gb must be non-negative, got {spark_executor_disk_gb}"
+            )
+
+        self.script = script
+        self.name = validate_name(name)
+        self.runtime = runtime.value if isinstance(runtime, Runtime) else runtime
+        region_value = region.value if isinstance(region, Region) else region
+
+        # Deprecation warnings for s3_bucket / s3_prefix are emitted by
+        # ``WherobotsConfig.from_env`` so they fire exactly once per
+        # construction regardless of whether the user went through
+        # ``WherobotsJob`` or ``WherobotsConfig`` directly.
+        self._config = config or WherobotsConfig.from_env(
+            api_key=api_key,
+            region=region_value,
+            s3_bucket=s3_bucket,
+            s3_prefix=s3_prefix,
+            base_url=base_url,
+            version=version,
+            request_timeout_seconds=request_timeout_seconds,
+        )
+
+        self.region = region_value or self._config.region or "aws-us-west-2"
+        self.version = version or self._config.version
+        self.timeout_seconds = timeout_seconds
+        self.args = args or []
+        self.spark_configs = spark_configs or {}
+        self.dependencies = dependencies or []
+        self.spark_driver_disk_gb = spark_driver_disk_gb
+        self.spark_executor_disk_gb = spark_executor_disk_gb
+        self.s3_bucket = s3_bucket or self._config.s3_bucket
+        self.s3_prefix = s3_prefix or self._config.s3_prefix
+        self.jar_main_class = jar_main_class
+        self.auto_upload = auto_upload
+
+        self.run_id: str | None = None
+        # Status may be a string when the server returns a value the
+        # SDK's JobStatus enum doesn't recognize yet (forward-compat).
+        self.status: JobStatus | str | None = None
+        self._last_log_cursor: int | str = 0
+        self._script_uri: str | None = None
+
+        self.is_jar = script.lower().endswith(".jar")
+        if self.is_jar and not jar_main_class:
+            raise WherobotsValidationError("jar_main_class is required for JAR files")
+
+        # Build the API layer
+        self._api = RunsAPI.from_config(self._config)
+
+    # ------------------------------------------------------------------ #
+    # Upload helpers
+    # ------------------------------------------------------------------ #
+
+    @staticmethod
+    def _is_s3_uri(path: str) -> bool:
+        return path.startswith("s3://")
+
+    def _prepare_script_uri(self) -> str:
+        """Resolve the script to an S3 URI, uploading if necessary.
+
+        Decision tree:
+
+        1. Already resolved → return cached URI.
+        2. Script is already an ``s3://`` URI → use directly.
+        3. Local file + ``auto_upload=True`` → upload via presigned URL
+           (``FilesAPI``).  Only an API key is needed.
+        4. ``auto_upload=False`` → error.
+        """
+        if self._script_uri:
+            return self._script_uri
+
+        if self._is_s3_uri(self.script):
+            self._script_uri = self.script
+        elif self.auto_upload:
+            self._script_uri = self._upload_via_presigned(self.script)
+        else:
+            raise WherobotsValidationError(
+                f"Script path '{self.script}' is not an S3 URI. "
+                "Set auto_upload=True or provide an S3 URI."
+            )
+
+        assert self._script_uri is not None  # all branches above set or raise
+        return self._script_uri
+
+    def _upload_via_presigned(self, local_path: str) -> str:
+        """Upload a local file via presigned URL (FilesAPI).
+
+        Args:
+            local_path: Path to the local file to upload.
+
+        Returns:
+            The ``s3://`` URI of the uploaded file.
+
+        Raises:
+            WherobotsAPIError: If the presigned upload fails.
+        """
+        from wherobots.api.files import FilesAPI
+
+        files_api = FilesAPI.from_config(self._config)
+        try:
+            s3_uri = files_api.upload_file(local_path)
+            logger.info("Presigned upload succeeded: %s", s3_uri)
+            return s3_uri
+        finally:
+            files_api.close()
+
+    # ------------------------------------------------------------------ #
+    # Payload construction
+    # ------------------------------------------------------------------ #
+
+    def _build_payload(self) -> CreateRunPayload:
+        script_uri = self._prepare_script_uri()
+
+        run_python: RunPythonPayload | None = None
+        run_jar: RunJarPayload | None = None
+
+        if self.is_jar:
+            run_jar = RunJarPayload(
+                uri=script_uri,
+                main_class=self.jar_main_class or "",
+                args=self.args,
+            )
+        else:
+            run_python = RunPythonPayload(uri=script_uri, args=self.args)
+
+        environment: RunEnvironment | None = None
+        if (
+            self.spark_configs
+            or self.dependencies
+            or self.spark_driver_disk_gb is not None
+            or self.spark_executor_disk_gb is not None
+        ):
+            environment = RunEnvironment(
+                spark_configs=self.spark_configs or None,
+                dependencies=self.dependencies or None,
+                spark_driver_disk_gb=self.spark_driver_disk_gb,
+                spark_executor_disk_gb=self.spark_executor_disk_gb,
+            )
+
+        return CreateRunPayload(
+            runtime=self.runtime,
+            name=self.name,
+            version=self.version or "latest",
+            timeout_seconds=self.timeout_seconds,
+            run_python=run_python,
+            run_jar=run_jar,
+            environment=environment,
+        )
+
+    # ------------------------------------------------------------------ #
+    # Lifecycle
+    # ------------------------------------------------------------------ #
+
+    def close(self) -> None:
+        """Close the underlying API session.
+
+        Should be called when the job manager is no longer needed,
+        or use as a context manager instead.
+        """
+        self._api.close()
+
+    def __enter__(self) -> WherobotsJob:
+        return self
+
+    def __exit__(self, *exc: Any) -> None:
+        self.close()
+
+    # ------------------------------------------------------------------ #
+    # Public API
+    # ------------------------------------------------------------------ #
+
+    def submit(self) -> str:
+        """Submit the job run.
+
+        Returns:
+            Run ID
+        """
+        if self.run_id:
+            logger.warning("Job already submitted with run_id: %s", self.run_id)
+            return self.run_id
+
+        payload = self._build_payload()
+        logger.info("Submitting job '%s' to %s", self.name, self.region)
+
+        run_view = self._api.create(payload, region=self.region)
+        self.run_id = run_view.id
+        self.status = run_view.status
+
+        logger.info("Job submitted successfully. Run ID: %s", self.run_id)
+        if self.status:
+            logger.info("Status: %s", _status_str(self.status))
+
+        return self.run_id
+
+    def get_status(self) -> RunView:
+        """Get current job status and details.
+
+        Returns:
+            RunView with full job details
+        """
+        if not self.run_id:
+            raise WherobotsJobError("Job not submitted. Call submit() first.")
+
+        run_view = self._api.get(self.run_id)
+        self.status = run_view.status
+        return run_view
+
+    def get_logs(
+        self,
+        cursor: int | str = 0,
+        size: int = 100,
+    ) -> LogsResponse:
+        """Get job logs.
+
+        Args:
+            cursor: Position to start reading logs. Accepts integer byte
+                offsets (legacy) or string cursors returned in
+                ``LogsResponse.next_page``. Defaults to ``0``.
+            size: Maximum number of log entries to retrieve (default: 100)
+
+        Returns:
+            LogsResponse with items and pagination info
+        """
+        if not self.run_id:
+            raise WherobotsJobError("Job not submitted. Call submit() first.")
+
+        return self._api.get_logs(self.run_id, cursor=cursor, size=size)
+
+    def get_metrics(self) -> RunMetricsResponse:
+        """Get job metrics for the run.
+
+        Returns:
+            RunMetricsResponse with series and instant metrics
+        """
+        if not self.run_id:
+            raise WherobotsJobError("Job not submitted. Call submit() first.")
+
+        return self._api.get_metrics(self.run_id)
+
+    def iter_logs(
+        self,
+        cursor: int | str = 0,
+        size: int = 100,
+        max_pages: int = 10_000,
+    ) -> Iterator[dict[str, Any]]:
+        """Iterate over log entries until pagination ends.
+
+        Args:
+            cursor: Initial pagination cursor.
+            size: Page size passed to :meth:`get_logs`.
+            max_pages: Hard ceiling on pages fetched, to guard against a
+                buggy server that always returns a distinct
+                ``next_page`` cursor. Defaults to ``10_000``.
+
+        Raises:
+            WherobotsJobError: If the page ceiling is hit.
+        """
+        next_cursor: int | str = cursor
+        for _ in range(max_pages):
+            logs = self.get_logs(cursor=next_cursor, size=size)
+            for item in logs.items:
+                yield item.to_dict()
+            if logs.next_page is None or logs.next_page == next_cursor:
+                return
+            next_cursor = logs.next_page
+        raise WherobotsJobError(
+            f"iter_logs exceeded max_pages={max_pages} without reaching end of log stream"
+        )
+
+    def poll_for_logs(
+        self,
+        follow: bool = True,
+        interval: float = 2.0,
+        log_handler: Callable[[dict[str, Any]], None] | None = None,
+        max_errors: int = 10,
+        _deadline: float | None = None,
+    ) -> None:
+        """Poll and print job logs.
+
+        Args:
+            follow: If True, continue polling until job completes.
+            interval: Polling interval in seconds (default: 2.0).
+            log_handler: Optional handler for each log item dict.
+            max_errors: Max consecutive transient errors before giving up
+                (default: 10). Non-transient errors (4xx except 429)
+                are raised immediately.
+            _deadline: Internal. Monotonic clock deadline for timeout
+                enforcement.  Callers should use
+                ``wait_for_completion(max_wait_seconds=...)`` instead.
+
+        Raises:
+            WherobotsJobError: If the job has not been submitted.
+            WherobotsTimeoutError: If *_deadline* is set and exceeded.
+            WherobotsAPIError: On non-transient HTTP errors.
+        """
+        if not self.run_id:
+            raise WherobotsJobError("Job not submitted. Call submit() first.")
+
+        def default_handler(item: dict[str, Any]) -> None:
+            print(item.get("raw", ""), flush=True)
+
+        handler = log_handler or default_handler
+        consecutive_errors = 0
+
+        while True:
+            # Enforce timeout when called from wait_for_completion
+            if _deadline is not None and time.monotonic() >= _deadline:
+                raise WherobotsTimeoutError(
+                    f"Timed out waiting for job {self.run_id} "
+                    f"(deadline exceeded during log streaming)"
+                )
+
+            try:
+                logs = self.get_logs(cursor=self._last_log_cursor)
+
+                for item in logs.items:
+                    handler(item.to_dict())
+
+                if logs.next_page is not None:
+                    self._last_log_cursor = logs.next_page
+
+                consecutive_errors = 0  # Reset on success
+
+                if not follow:
+                    break
+
+                run_view = self.get_status()
+                if is_terminal_status(run_view.status):
+                    # Drain remaining logs
+                    final_logs = self.get_logs(cursor=self._last_log_cursor, size=1000)
+                    for item in final_logs.items:
+                        handler(item.to_dict())
+                    break
+
+                time.sleep(interval)
+
+            except KeyboardInterrupt:
+                logger.info("Log streaming interrupted by user")
+                break
+            except WherobotsAPIError as exc:
+                # Non-transient HTTP errors should not be retried
+                if exc.status_code and 400 <= exc.status_code < 500 and exc.status_code != 429:
+                    raise
+                consecutive_errors += 1
+                logger.error("Error polling logs (%d/%d): %s", consecutive_errors, max_errors, exc)
+                if consecutive_errors >= max_errors or not follow:
+                    raise
+                time.sleep(interval)
+            except Exception as exc:
+                consecutive_errors += 1
+                logger.error("Error polling logs (%d/%d): %s", consecutive_errors, max_errors, exc)
+                if consecutive_errors >= max_errors or not follow:
+                    raise
+                time.sleep(interval)
+
+    def cancel(self) -> bool:
+        """Cancel the job run.
+
+        Returns:
+            True if cancellation request was successful
+        """
+        if not self.run_id:
+            raise WherobotsJobError("Job not submitted. Call submit() first.")
+
+        self._api.cancel(self.run_id)
+
+        logger.info("Cancellation requested for run %s", self.run_id)
+        logger.info("It may take a minute for status to change to CANCELLED")
+
+        return True
+
+    def wait_for_completion(
+        self,
+        poll_interval: float = 5.0,
+        stream_logs: bool = True,
+        log_interval: float = 2.0,
+        max_wait_seconds: float | None = None,
+    ) -> JobStatus | str:
+        """Wait for job to complete, optionally streaming logs.
+
+        Args:
+            poll_interval: Status check interval in seconds (default: 5.0)
+            stream_logs: Stream logs while waiting (default: True)
+            log_interval: Log polling interval in seconds (default: 2.0)
+            max_wait_seconds: Maximum time to wait before raising
+                ``WherobotsTimeoutError``.  ``None`` means wait
+                indefinitely (default: None).
+
+        Returns:
+            Final job status
+
+        Raises:
+            WherobotsTimeoutError: If max_wait_seconds is exceeded.
+        """
+        if not self.run_id:
+            raise WherobotsJobError("Job not submitted. Call submit() first.")
+
+        logger.info("Waiting for job %s to complete...", self.run_id)
+        start = time.monotonic()
+
+        if stream_logs:
+            deadline = start + max_wait_seconds if max_wait_seconds is not None else None
+            self.poll_for_logs(follow=True, interval=log_interval, _deadline=deadline)
+        else:
+            while True:
+                if max_wait_seconds is not None:
+                    elapsed = time.monotonic() - start
+                    if elapsed >= max_wait_seconds:
+                        raise WherobotsTimeoutError(
+                            f"Timed out waiting for job {self.run_id} "
+                            f"after {elapsed:.0f}s (limit: {max_wait_seconds}s)"
+                        )
+                run_view = self.get_status()
+                if is_terminal_status(run_view.status):
+                    break
+                if run_view.status:
+                    logger.info("Status: %s", _status_str(run_view.status))
+                time.sleep(poll_interval)
+
+        final_view = self.get_status()
+        if not final_view.status:
+            raise WherobotsJobError(f"Job {self.run_id} returned an unrecognized status")
+
+        self.status = final_view.status
+        logger.info("Job %s finished with status: %s", self.run_id, _status_str(self.status))
+
+        if self.status == JobStatus.FAILED:
+            logger.error("Job failed. Check logs for details.")
+        elif self.status == JobStatus.CANCELLED:
+            logger.warning("Job was cancelled.")
+
+        return self.status
+
+    # ------------------------------------------------------------------ #
+    # Class-level helpers (no instance required)
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    def list_runs(
+        cls,
+        api_key: str | None = None,
+        region: str | None = None,
+        name_pattern: str | None = None,
+        created_after: str | None = None,
+        status: list[str | JobStatus] | None = None,
+        cursor: str | None = None,
+        size: int = 50,
+        base_url: str | None = None,
+        request_timeout_seconds: int | None = None,
+        config: WherobotsConfig | None = None,
+    ) -> RunListPage:
+        """List job runs in the organization.
+
+        Args:
+            api_key: Wherobots API key (or use WHEROBOTS_API_KEY env var)
+            region: Filter by region
+            name_pattern: Filter by name pattern (supports * wildcard)
+            created_after: Filter runs created after this ISO timestamp
+            status: Filter by status (list of statuses)
+            cursor: Pagination cursor
+            size: Number of results per page (default: 50, max: 100)
+            base_url: Override API base URL
+            request_timeout_seconds: HTTP request timeout in seconds
+            config: Optional WherobotsConfig to override defaults
+
+        Returns:
+            RunListPage with items, total, and pagination cursor
+        """
+        cfg = config or WherobotsConfig.from_env(
+            api_key=api_key,
+            region=region,
+            base_url=base_url,
+            request_timeout_seconds=request_timeout_seconds,
+        )
+
+        with RunsAPI.from_config(cfg) as api:
+            return api.list(
+                region=region or cfg.region,
+                name_pattern=name_pattern,
+                created_after=created_after,
+                status=status,
+                cursor=cursor,
+                size=size,
+            )
+
+    @classmethod
+    def add_pypi_dependency(cls, library_name: str, library_version: str) -> dict[str, str]:
+        """Create a PyPI dependency object.
+
+        Args:
+            library_name: PyPI package name
+            library_version: Package version
+
+        Returns:
+            Dependency dictionary
+        """
+        return {
+            "sourceType": DependencyType.PYPI.value,
+            "libraryName": library_name,
+            "libraryVersion": library_version,
+        }
+
+    @classmethod
+    def add_file_dependency(cls, file_path: str) -> dict[str, str]:
+        """Create a file dependency object.
+
+        Args:
+            file_path: Path to dependency file (.jar, .whl, .zip, .json)
+
+        Returns:
+            Dependency dictionary
+        """
+        valid_extensions = [".jar", ".whl", ".zip", ".json"]
+        if not any(file_path.endswith(ext) for ext in valid_extensions):
+            raise WherobotsValidationError(
+                f"File must have one of these extensions: {valid_extensions}"
+            )
+
+        return {
+            "sourceType": DependencyType.FILE.value,
+            "filePath": file_path,
+        }