flyteplugins-trackio 2.6.11__py3-none-any.whl

flyteplugins/trackio/__init__.py

@@ -0,0 +1,52 @@
+"""
+Flyte Trackio plugin.
+
+Provides seamless Trackio experiment tracking for Flyte tasks through the
+`@trackio_init` decorator.
+
+Basic usage:
+
+    from flyteplugins.trackio import (
+        get_trackio_run,
+        trackio_init,
+    )
+
+    @trackio_init(project="my-project")
+    @env.task
+    async def train():
+        run = get_trackio_run()
+        run.log({"loss": 0.123})
+        return run.id
+
+Configuration can also be provided via `trackio_config()`:
+
+    r = flyte.with_runcontext(
+        custom_context=trackio_config(
+            project="my-project",
+            group="baseline",
+            config={"learning_rate": 1e-3},
+        )
+    ).run(train)
+"""
+
+from __future__ import annotations
+
+from ._context import (
+    get_trackio_context,
+    get_trackio_run,
+    trackio_config,
+)
+from ._decorator import (
+    trackio_init,
+)
+from ._link import Trackio
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Trackio",
+    "get_trackio_context",
+    "get_trackio_run",
+    "trackio_config",
+    "trackio_init",
+]

flyteplugins/trackio/_context.py

@@ -0,0 +1,241 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from typing import Any, Optional
+
+import flyte
+
+import trackio
+
+_TRACKIO_RUN_KEY = "_trackio_run"
+_TRACKIO_CONTEXT_PREFIX = "trackio"
+
+
+def _to_dict_helper(obj, prefix: str) -> dict[str, str]:
+    """Serialize a dataclass into Flyte custom_context."""
+
+    result: dict[str, str] = {}
+
+    for key, value in asdict(obj).items():
+        if value is None:
+            continue
+
+        if isinstance(value, (dict, list, bool)):
+            result[f"{prefix}_{key}"] = json.dumps(value)
+        else:
+            result[f"{prefix}_{key}"] = str(value)
+
+    return result
+
+
+def _from_dict_helper(cls, d: dict[str, str], prefix: str):
+    """Deserialize a dataclass from Flyte custom_context."""
+
+    kwargs = {}
+
+    prefix = f"{prefix}_"
+
+    for key, value in d.items():
+        if not key.startswith(prefix):
+            continue
+
+        field = key[len(prefix) :]
+
+        try:
+            kwargs[field] = json.loads(value)
+        except Exception:
+            kwargs[field] = value
+
+    return cls(**kwargs)
+
+
+def _context_manager_enter(obj, prefix: str):
+    ctx = flyte.ctx()
+
+    saved = {}
+
+    if ctx and ctx.custom_context:
+        for key in list(ctx.custom_context.keys()):
+            if key.startswith(f"{prefix}_"):
+                saved[key] = ctx.custom_context[key]
+
+    ctx_mgr = flyte.custom_context(**obj)
+
+    ctx_mgr.__enter__()
+
+    return saved, ctx_mgr
+
+
+def _context_manager_exit(ctx_mgr, saved: dict, prefix: str, *args):
+    if ctx_mgr:
+        ctx_mgr.__exit__(*args)
+
+    ctx = flyte.ctx()
+
+    if ctx and ctx.custom_context:
+        for key in list(ctx.custom_context.keys()):
+            if key.startswith(f"{prefix}_"):
+                del ctx.custom_context[key]
+
+        ctx.custom_context.update(saved)
+
+
+@dataclass
+class _TrackioConfig:
+    """
+    Trackio configuration stored inside the Flyte custom_context.
+    Mirrors the supported subset of `trackio.init()`.
+    """
+
+    project: Optional[str] = None
+
+    name: Optional[str] = None
+
+    group: Optional[str] = None
+
+    space_id: Optional[str] = None
+
+    dataset_id: Optional[str] = None
+
+    bucket_id: Optional[str] = None
+
+    server_url: Optional[str] = None
+
+    config: Optional[dict[str, Any]] = None
+
+    resume: str = "never"
+
+    auto_log_gpu: Optional[bool] = None
+
+    gpu_log_interval: float = 10.0
+
+    auto_log_cpu: Optional[bool] = None
+
+    cpu_log_interval: float = 10.0
+
+    def to_trackio_init(self) -> dict[str, Any]:
+        """Convert to arguments for `trackio.init()`."""
+
+        return {k: v for k, v in asdict(self).items() if v is not None}
+
+    def to_dict(self) -> dict[str, str]:
+        return _to_dict_helper(self, _TRACKIO_CONTEXT_PREFIX)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, str]):
+        return _from_dict_helper(cls, d, _TRACKIO_CONTEXT_PREFIX)
+
+    def __enter__(self):
+        self._saved, self._ctx = _context_manager_enter(
+            self,
+            _TRACKIO_CONTEXT_PREFIX,
+        )
+        return self
+
+    def __exit__(self, *args):
+        _context_manager_exit(
+            self._ctx,
+            self._saved,
+            _TRACKIO_CONTEXT_PREFIX,
+            *args,
+        )
+
+
+def get_trackio_context() -> Optional[_TrackioConfig]:
+    """
+    Return the current Trackio configuration.
+    """
+
+    ctx = flyte.ctx()
+
+    if not ctx or not ctx.custom_context:
+        return None
+
+    prefix = f"{_TRACKIO_CONTEXT_PREFIX}_"
+
+    if not any(k.startswith(prefix) for k in ctx.custom_context):
+        return None
+
+    return _TrackioConfig.from_dict(ctx.custom_context)
+
+
+def trackio_config(
+    *,
+    project: str | None = None,
+    name: str | None = None,
+    group: str | None = None,
+    space_id: str | None = None,
+    dataset_id: str | None = None,
+    bucket_id: str | None = None,
+    server_url: str | None = None,
+    config: dict[str, Any] | None = None,
+    resume: str = "never",
+    auto_log_gpu: bool | None = None,
+    gpu_log_interval: float = 10.0,
+    auto_log_cpu: bool | None = None,
+    cpu_log_interval: float = 10.0,
+) -> _TrackioConfig:
+    """
+    Create Trackio configuration for Flyte.
+    """
+
+    return _TrackioConfig(
+        project=project,
+        name=name,
+        group=group,
+        space_id=space_id,
+        dataset_id=dataset_id,
+        bucket_id=bucket_id,
+        server_url=server_url,
+        config=config,
+        resume=resume,
+        auto_log_gpu=auto_log_gpu,
+        gpu_log_interval=gpu_log_interval,
+        auto_log_cpu=auto_log_cpu,
+        cpu_log_interval=cpu_log_interval,
+    )
+
+
+def set_trackio_run(run) -> None:
+    """Store the active Trackio run in the Flyte context."""
+
+    ctx = flyte.ctx()
+
+    if not ctx:
+        return
+
+    if ctx.data is None:
+        ctx.data = {}
+
+    ctx.data[_TRACKIO_RUN_KEY] = run
+
+
+def get_trackio_run():
+    """
+    Return the active Trackio run.
+
+    If called inside a `@trackio_init` decorated Flyte task, this returns the
+    Trackio run managed by the plugin. Otherwise it falls back to Trackio's
+    globally active run (if one exists).
+
+    Returns:
+        trackio.Run | None: The active Trackio run.
+    """
+    ctx = flyte.ctx()
+
+    if ctx and ctx.data:
+        run = ctx.data.get(_TRACKIO_RUN_KEY)
+        if run is not None:
+            return run
+
+    return getattr(trackio, "run", None)
+
+
+def clear_trackio_run() -> None:
+    """Remove the Trackio run from the Flyte context."""
+
+    ctx = flyte.ctx()
+
+    if ctx and ctx.data:
+        ctx.data.pop(_TRACKIO_RUN_KEY, None)

flyteplugins/trackio/_decorator.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import functools
+import logging
+from contextlib import contextmanager
+from inspect import iscoroutinefunction
+from typing import Any, Callable, Optional, TypeVar, cast
+
+import flyte
+from flyte._task import AsyncFunctionTaskTemplate
+
+import trackio
+
+from ._context import (
+    _TRACKIO_RUN_KEY,
+    clear_trackio_run,
+    get_trackio_context,
+    set_trackio_run,
+)
+from ._link import Trackio
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def _build_init_kwargs(
+    decorator_kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    """
+    Build arguments for `trackio.init()`.
+
+    Values from the current Trackio configuration are used as defaults and are
+    overridden by non-`None` decorator arguments.
+
+    Args:
+        decorator_kwargs:
+            Keyword arguments supplied to `@trackio_init`.
+
+    Returns:
+        A dictionary of resolved keyword arguments to pass to
+        `trackio.init()`.
+    """
+    ctx = get_trackio_context()
+
+    init_kwargs = ctx.to_trackio_init() if ctx else {}
+
+    init_kwargs.update({k: v for k, v in decorator_kwargs.items() if v is not None})
+
+    return init_kwargs
+
+
+@contextmanager
+def _trackio_run(**decorator_kwargs):
+    """
+    Context manager responsible for Trackio run lifecycle.
+
+    If a parent Trackio run already exists, it is reused.
+    Otherwise a new Trackio run is created for the duration
+    of the task.
+    """
+
+    flyte_ctx = flyte.ctx()
+
+    if not flyte_ctx:
+        run = trackio.init(**decorator_kwargs)
+
+        try:
+            yield run
+        finally:
+            run.finish()
+
+        return
+
+    if flyte_ctx.data is None:
+        flyte_ctx.data = {}
+
+    saved_run = flyte_ctx.data.get(_TRACKIO_RUN_KEY)
+
+    if saved_run is not None:
+        yield saved_run
+        return
+
+    init_kwargs = _build_init_kwargs(decorator_kwargs)
+
+    run = trackio.init(**init_kwargs)
+
+    set_trackio_run(run)
+
+    try:
+        yield run
+    finally:
+        try:
+            run.finish()
+        except Exception:
+            logger.exception("Failed to finish Trackio run.")
+        finally:
+            clear_trackio_run()
+
+
+def trackio_init(
+    _func: Optional[F] = None,
+    **decorator_kwargs: Any,
+) -> F:
+    """
+    Initialize a Trackio run around a Flyte task.
+
+    Usage
+    -----
+
+    @trackio_init
+    @env.task
+    async def train():
+        ...
+
+    @trackio_init(
+        project="vision",
+        space_id="user/demo",
+    )
+    @env.task
+    async def train():
+        ...
+    """
+
+    def decorator(task: F) -> F:
+
+        if isinstance(task, AsyncFunctionTaskTemplate):
+            existing_links = getattr(task, "links", ())
+
+            task = task.override(
+                links=(
+                    *existing_links,
+                    Trackio(
+                        project=decorator_kwargs.get("project"),
+                        server_url=decorator_kwargs.get("server_url"),
+                        space_id=decorator_kwargs.get("space_id"),
+                    ),
+                )
+            )
+
+            original_execute = task.execute
+
+            async def wrapped_execute(*args, **kwargs):
+
+                with _trackio_run(**decorator_kwargs):
+                    return await original_execute(*args, **kwargs)
+
+            task.execute = wrapped_execute
+
+            return cast(F, task)
+
+        if iscoroutinefunction(task):
+
+            @functools.wraps(task)
+            async def async_wrapper(*args, **kwargs):
+
+                with _trackio_run(**decorator_kwargs):
+                    return await task(*args, **kwargs)
+
+            return cast(F, async_wrapper)
+
+        @functools.wraps(task)
+        def sync_wrapper(*args, **kwargs):
+
+            with _trackio_run(**decorator_kwargs):
+                return task(*args, **kwargs)
+
+        return cast(F, sync_wrapper)
+
+    if _func is None:
+        return decorator
+
+    return decorator(_func)

flyteplugins/trackio/_link.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from flyte import Link
+
+from ._context import get_trackio_context
+
+
+@dataclass
+class Trackio(Link):
+    """
+    Generates a Trackio dashboard link for Flyte.
+
+    The link resolution order is:
+
+        1. Explicit server_url (self or context)
+        2. Hugging Face Space (space_id)
+        3. Hugging Face Trackio documentation
+
+    Args:
+        host:
+            Base Hugging Face host.
+
+        project:
+            Trackio project name.
+
+        server_url:
+            Base URL of a self-hosted Trackio instance.
+
+        space_id:
+            Hugging Face Space hosting the Trackio dashboard.
+
+        name:
+            Display name in the Flyte UI.
+    """
+
+    host: str = "https://huggingface.co"
+
+    project: Optional[str] = None
+
+    server_url: Optional[str] = None
+
+    space_id: Optional[str] = None
+
+    name: str = "Trackio"
+
+    def get_link(
+        self,
+        run_name: str,
+        project: str,
+        domain: str,
+        context: Dict[str, str],
+        parent_action_name: str,
+        action_name: str,
+        pod_name: str,
+        **kwargs,
+    ) -> str:
+        """
+        Resolve the Trackio dashboard URL.
+        """
+
+        cfg = get_trackio_context()
+
+        project_name = self.project
+        server_url = self.server_url
+        space_id = self.space_id
+
+        if cfg is not None:
+            project_name = project_name or cfg.project
+            server_url = server_url or cfg.server_url
+            space_id = space_id or cfg.space_id
+
+        if server_url:
+            server_url = server_url.rstrip("/")
+
+            if project_name:
+                return f"{server_url}/projects/{project_name}"
+
+            return server_url
+
+        if space_id:
+            return f"{self.host}/spaces/{space_id}"
+
+        return f"{self.host}/docs/trackio"

flyteplugins_trackio-2.6.11.dist-info/METADATA

@@ -0,0 +1,124 @@
+Metadata-Version: 2.4
+Name: flyteplugins-trackio
+Version: 2.6.11
+Summary: Trackio plugin for Flyte
+Author-email: Parag Ekbote <23150020.dypsst@dpu.edu.in>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: trackio
+Requires-Dist: flyte
+
+# Flyte Trackio Plugin
+
+Native Flyte support for Trackio experiment tracking.
+
+This plugin makes it easy to manage Trackio runs inside Flyte tasks, log metrics from within task code or callbacks, and expose a Trackio dashboard link in the Flyte UI.
+
+## Deploying a Trackio Server
+
+This plugin works with Trackio's local-first design. By default, Trackio stores experiment data locally and can optionally send metrics to a self-hosted Trackio server or a Hugging Face Space by configuring the appropriate `server_url` or `space_id`.
+
+For deployment instructions, please refer to the official Trackio documentation:
+
+- Self-host a Trackio server: https://huggingface.co/docs/trackio/self_hosted_server
+- Deploy and embed Trackio dashboards: https://huggingface.co/docs/trackio/deploy_embed
+
+## What it provides
+
+- `@trackio_init` decorator to initialize and manage a Trackio run for a Flyte task.
+- `trackio_config(...)` helper to store Trackio initialization settings in Flyte `custom_context`.
+- `get_trackio_run()` to access the active Trackio run from task code or callback code.
+- Optional Flyte task link support via the `Trackio` link class.
+
+## Installation
+
+```bash
+pip install flyteplugins-trackio
+```
+
+## Quick start
+
+```python
+import flyte
+from flyteplugins.trackio import (
+    get_trackio_run,
+    trackio_config,
+    trackio_init,
+)
+
+env = flyte.TaskEnvironment(name="trackio")
+
+@trackio_init
+@env.task
+def train() -> dict[str, float]:
+    run = get_trackio_run()
+    run.log({"accuracy": 0.92})
+    return {"accuracy": 0.92}
+
+cfg = trackio_config(
+    project="my-project",
+    space_id="my-org/my-trackio-space",
+    bucket_id="my-bucket",
+    auto_log_cpu=True,
+)
+
+flyte.with_runcontext(custom_context=cfg.to_dict()).run(train)
+```
+
+## Key APIs
+
+### `trackio_init`
+
+Decorates a Flyte task or plain Python function to create a Trackio run for the decorated execution.
+
+- Works with Flyte task functions.
+- Works with plain synchronous or asynchronous functions outside of Flyte.
+- Reuses an existing Trackio run if one is already present in the Flyte context.
+
+### `trackio_config(...)`
+
+Returns a configuration object with Trackio initialization options such as:
+
+- `project`
+- `name`
+- `group`
+- `space_id`
+- `dataset_id`
+- `bucket_id`
+- `server_url`
+- `config`
+- `resume`
+- `auto_log_gpu`
+- `gpu_log_interval`
+- `auto_log_cpu`
+- `cpu_log_interval`
+
+Use `cfg.to_dict()` to store these settings in Flyte `custom_context`, then run the task with `flyte.with_runcontext(...)`.
+
+### `get_trackio_run()`
+
+Returns the active Trackio run from the current Flyte context, or falls back to Trackio's global active run when running outside of Flyte.
+
+### `Trackio`
+
+A Flyte `Link` implementation that resolves a Trackio dashboard URL based on:
+
+- explicit `server_url`
+- `space_id` for Hugging Face Space deployments
+- plugin `project`/context values
+
+This link can be attached automatically by `trackio_init` for Flyte tasks.
+
+## Usage notes
+
+- The plugin manages the Trackio run lifecycle automatically: `trackio_init` creates and finishes the run around the decorated execution.
+- Custom context values from `trackio_config` are merged with decorator-level overrides.
+- Use `get_trackio_run()` in callbacks or training loops to log metrics incrementally.
+
+## Examples
+
+See [distilbert_text_classification.py](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/trackio/examples) and [vit_image_classification.py](https://github.com/flyteorg/flyte-sdk/tree/main/plugins/trackio/examples) for full usage patterns with Hugging Face training and callback-based metric logging. Before running the examples, please install the required dependencies:
+
+```bash
+pip install datasets transformers accelerate evaluate psutil
+```

flyteplugins_trackio-2.6.11.dist-info/RECORD

@@ -0,0 +1,8 @@
+flyteplugins/trackio/__init__.py,sha256=H9YHuhiBHQedE7q_7AUATlMKlEGwtnzNdMR5lNYg45s,1017
+flyteplugins/trackio/_context.py,sha256=ANJdp_lSVwcXVxTFKrajs2byl084A95RC-QMsyktc2g,5547
+flyteplugins/trackio/_decorator.py,sha256=EYELNYEMIjPi7QvFw5krVF1ikvxLGfNZMYUPTUjAVM0,3938
+flyteplugins/trackio/_link.py,sha256=K-Z22Riy0a0e2x66B3J_j2sueeiPHY9yxE6cUyLY9QM,1906
+flyteplugins_trackio-2.6.11.dist-info/METADATA,sha256=xYl5KpG7rA5b_IijgOiLZv9sV7OTsxiXKDTVIxVR-E0,3995
+flyteplugins_trackio-2.6.11.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flyteplugins_trackio-2.6.11.dist-info/top_level.txt,sha256=cgd779rPu9EsvdtuYgUxNHHgElaQvPn74KhB5XSeMBE,13
+flyteplugins_trackio-2.6.11.dist-info/RECORD,,

flyteplugins_trackio-2.6.11.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flyteplugins_trackio-2.6.11.dist-info/top_level.txt

@@ -0,0 +1 @@
+flyteplugins