greatsky-internal-metaflow

CLI for authenticating with the GreatSky Metaflow platform.

Install

# pip (use a venv)
pip install greatsky-internal-metaflow

# or with uv
uv pip install greatsky-internal-metaflow

Quick Start

# Authenticate (opens GitHub in your browser)
gsm login

# Verify everything is working
gsm validate

# Run Metaflow flows — no further config needed
python my_flow.py run

Commands

Authentication

Service Keys

Admin (org members only)

Hugging Face Integration

The package includes first-class Hugging Face Hub support. When running on the platform, HF_TOKEN and HF_ORG are injected automatically -- no manual setup. All dataset/model names without a / are auto-prefixed with the org (Great-Sky/).

Datasets

from greatsky_internal_metaflow.hf import load_dataset, save_dataset, list_datasets

# Load a dataset with version tracking (org prefix auto-resolved)
ds, version = load_dataset("my-dataset", split="train")
self.dataset_version = version  # store as Metaflow artifact for lineage

# Save a dataset (accepts Dataset, DataFrame, dict, or directory path)
version = save_dataset(my_dataframe, "my-processed-dataset")
version = save_dataset({"text": texts, "label": labels}, "my-dataset")
version = save_dataset(hf_dataset, "my-dataset", commit_message="v2 cleaned")

# List all org datasets
for ds_info in list_datasets():
    print(ds_info["id"], ds_info["downloads"])

The DatasetVersion object captures the exact HF commit SHA, split, timestamp, and Metaflow run context. Query it later via the Metaflow Client API:

from metaflow import Flow
run = Flow("HFTrainingFlow").latest_run
print(run.data.dataset_version)
# DatasetVersion(Great-Sky/imdb@a3b5c8d2, split=train, run=HFTrainingFlow/42)

Models

from greatsky_internal_metaflow.hf import push_model_to_hub, pull_from_hub, sync_s3_to_hf

# Pull a model (org prefix auto-resolved)
local_path = pull_from_hub("my-model")

# Push a trained model to Great-Sky/my-model-finetuned
url = push_model_to_hub("/tmp/output", "my-model-finetuned")

# S3 <-> HF bridge
sync_s3_to_hf("s3://bucket/models/my-model", "my-model")

With metaflow-checkpoint decorators

Install with the [metaflow] extra to get @huggingface_hub, @model, and @checkpoint decorators:

pip install "greatsky-internal-metaflow[metaflow]"
# or: uv pip install "greatsky-internal-metaflow[metaflow]"

from metaflow import FlowSpec, step
from metaflow_extensions.obcheckpoint.plugins.hf_hub_decorator import huggingface_hub

class MyFlow(FlowSpec):
    @huggingface_hub(models=["distilbert-base-uncased"])
    @step
    def train(self):
        ...

Locally

If running outside the platform, set HF_TOKEN and HF_ORG as environment variables, or authenticate with huggingface-cli login.

How It Works

1. gsm login starts a GitHub Device Flow — you get a code to enter at github.com/login/device
2. Once authorized, the CLI exchanges the GitHub token with the GreatSky auth API for a platform API key
3. The API key and Metaflow config are written to a .greatsky_gsm/ directory (project, venv, or global scope)
4. The Metaflow extension reads that config automatically — no wrappers or env vars needed

API keys are valid for 14 days and are automatically revoked if you leave the GitHub org.

Config Resolution

The CLI resolves config from three locations, in priority order:

1. Project — .greatsky_gsm/ in the nearest directory containing pyproject.toml or setup.py
2. Virtualenv — .greatsky_gsm/ inside $VIRTUAL_ENV
3. Global — ~/.greatsky_gsm/

Use --scope with gsm login to control where credentials are stored.

Development

git clone https://github.com/greatsky-ai/greatsky-internal-metaflow.git
cd greatsky-internal-metaflow
uv venv && uv pip install -e ".[dev]"
uv run pytest
uv run ruff check src/ tests/

Versioning

Versions are derived automatically from git tags via hatch-vcs. To release:

git tag v0.2.0
git push origin v0.2.0

The publish workflow runs lint, tests, and smoke tests before pushing to PyPI.