diff --git a/python/packages/jumpstarter/jumpstarter/client/lease.py b/python/packages/jumpstarter/jumpstarter/client/lease.py index 1e431a10b..81dba4960 100644 --- a/python/packages/jumpstarter/jumpstarter/client/lease.py +++ b/python/packages/jumpstarter/jumpstarter/client/lease.py @@ -58,6 +58,7 @@ class Lease(ContextManagerMixin, AsyncContextManagerMixin): acquisition_timeout: int = field(default=7200) # Timeout in seconds for lease acquisition, polled in 5s intervals dial_timeout: float = field(default=30.0) # Timeout in seconds for Dial retry loop when exporter not ready exporter_name: str = field(default="remote", init=False) # Populated during acquisition + exporter_labels: dict[str, str] = field(default_factory=dict, init=False) # Populated during acquisition lease_ending_callback: Callable[[Self, timedelta], None] | None = field( default=None, init=False ) # Called when lease is ending @@ -143,6 +144,15 @@ async def request_async(self): return await self._acquire() + async def _fetch_exporter_labels(self): + """Fetch the exporter's labels after lease acquisition.""" + try: + exporter = await self.svc.GetExporter(name=self.exporter_name) + self.exporter_labels = exporter.labels + except Exception as e: + self.exporter_labels = {} + logger.warning("Could not fetch labels for exporter %s: %s", self.exporter_name, e) + def _update_spinner_status(self, spinner, result): """Update spinner with appropriate status message based on lease conditions.""" if condition_true(result.conditions, "Pending"): @@ -183,6 +193,7 @@ async def _acquire(self): logger.debug("Lease %s acquired", self.name) spinner.update_status(f"Lease {self.name} acquired successfully!", force=True) self.exporter_name = result.exporter + await self._fetch_exporter_labels() break # lease unsatisfiable diff --git a/python/packages/jumpstarter/jumpstarter/common/utils.py b/python/packages/jumpstarter/jumpstarter/common/utils.py index 7f63c4eac..55f432ad3 100644 --- a/python/packages/jumpstarter/jumpstarter/common/utils.py +++ b/python/packages/jumpstarter/jumpstarter/common/utils.py @@ -10,14 +10,14 @@ from anyio.from_thread import BlockingPortal, start_blocking_portal from jumpstarter.client import client_from_path -from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JUMPSTARTER_HOST +from jumpstarter.config.env import JMP_DRIVERS_ALLOW, JMP_EXPORTER, JMP_EXPORTER_LABELS, JMP_LEASE, JUMPSTARTER_HOST from jumpstarter.exporter import Session -from jumpstarter.utils.env import env +from jumpstarter.utils.env import ExporterMetadata, env, env_with_metadata if TYPE_CHECKING: from jumpstarter.driver import Driver -__all__ = ["env"] +__all__ = ["ExporterMetadata", "env", "env_with_metadata"] @asynccontextmanager @@ -84,6 +84,19 @@ def _run_process( return process.wait() +def _lease_env_vars(lease) -> dict[str, str]: + """Extract environment variables from a lease object.""" + env_vars: dict[str, str] = {} + env_vars[JMP_EXPORTER] = lease.exporter_name + if lease.name: + env_vars[JMP_LEASE] = lease.name + if lease.exporter_labels: + env_vars[JMP_EXPORTER_LABELS] = ",".join( + f"{k}={v}" for k, v in sorted(lease.exporter_labels.items()) + ) + return env_vars + + def launch_shell( host: str, context: str, @@ -118,6 +131,9 @@ def launch_shell( "_JMP_SUPPRESS_DRIVER_WARNINGS": "1", # Already warned during client initialization } + if lease is not None: + common_env.update(_lease_env_vars(lease)) + if command: return _run_process(list(command), common_env, lease) diff --git a/python/packages/jumpstarter/jumpstarter/common/utils_test.py b/python/packages/jumpstarter/jumpstarter/common/utils_test.py index 86cd52dfd..b555318fd 100644 --- a/python/packages/jumpstarter/jumpstarter/common/utils_test.py +++ b/python/packages/jumpstarter/jumpstarter/common/utils_test.py @@ -1,6 +1,8 @@ import shutil +from types import SimpleNamespace from .utils import launch_shell +from jumpstarter.utils.env import ExporterMetadata def test_launch_shell(tmp_path, monkeypatch): @@ -22,3 +24,52 @@ def test_launch_shell(tmp_path, monkeypatch): use_profiles=False ) assert exit_code == 1 + + +def test_launch_shell_sets_lease_env(tmp_path, monkeypatch): + monkeypatch.setenv("SHELL", shutil.which("env")) + lease = SimpleNamespace( + exporter_name="my-exporter", + name="lease-123", + exporter_labels={"board": "rpi4", "location": "lab-1"}, + lease_ending_callback=None, + ) + exit_code = launch_shell( + host=str(tmp_path / "test.sock"), + context="my-exporter", + allow=["*"], + unsafe=False, + use_profiles=False, + lease=lease, + ) + assert exit_code == 0 + + +def test_exporter_metadata_from_env(monkeypatch): + monkeypatch.setenv("JMP_EXPORTER", "my-board") + monkeypatch.setenv("JMP_LEASE", "lease-abc") + monkeypatch.setenv("JMP_EXPORTER_LABELS", "board=rpi4,location=lab-1,team=qa") + + meta = ExporterMetadata.from_env() + assert meta.name == "my-board" + assert meta.lease == "lease-abc" + assert meta.labels == {"board": "rpi4", "location": "lab-1", "team": "qa"} + + +def test_exporter_metadata_from_env_empty(monkeypatch): + monkeypatch.delenv("JMP_EXPORTER", raising=False) + monkeypatch.delenv("JMP_LEASE", raising=False) + monkeypatch.delenv("JMP_EXPORTER_LABELS", raising=False) + + meta = ExporterMetadata.from_env() + assert meta.name == "" + assert meta.lease is None + assert meta.labels == {} + + +def test_exporter_metadata_from_env_labels_with_equals_in_value(monkeypatch): + monkeypatch.setenv("JMP_EXPORTER", "board") + monkeypatch.setenv("JMP_EXPORTER_LABELS", "key=val=123,other=ok") + + meta = ExporterMetadata.from_env() + assert meta.labels == {"key": "val=123", "other": "ok"} diff --git a/python/packages/jumpstarter/jumpstarter/config/env.py b/python/packages/jumpstarter/jumpstarter/config/env.py index 145966d27..68f4162a3 100644 --- a/python/packages/jumpstarter/jumpstarter/config/env.py +++ b/python/packages/jumpstarter/jumpstarter/config/env.py @@ -8,6 +8,8 @@ JMP_DRIVERS_ALLOW = "JMP_DRIVERS_ALLOW" JUMPSTARTER_HOST = "JUMPSTARTER_HOST" JMP_LEASE = "JMP_LEASE" +JMP_EXPORTER = "JMP_EXPORTER" +JMP_EXPORTER_LABELS = "JMP_EXPORTER_LABELS" JMP_DISABLE_COMPRESSION = "JMP_DISABLE_COMPRESSION" JMP_OIDC_CALLBACK_PORT = "JMP_OIDC_CALLBACK_PORT" diff --git a/python/packages/jumpstarter/jumpstarter/utils/env.py b/python/packages/jumpstarter/jumpstarter/utils/env.py index c6977b7e9..0ae8f38d2 100644 --- a/python/packages/jumpstarter/jumpstarter/utils/env.py +++ b/python/packages/jumpstarter/jumpstarter/utils/env.py @@ -1,12 +1,38 @@ import os from contextlib import ExitStack, asynccontextmanager, contextmanager +from dataclasses import dataclass, field from anyio.from_thread import start_blocking_portal from jumpstarter.client import client_from_path from jumpstarter.common.exceptions import EnvironmentVariableNotSetError from jumpstarter.config.client import ClientConfigV1Alpha1Drivers -from jumpstarter.config.env import JUMPSTARTER_HOST +from jumpstarter.config.env import JMP_EXPORTER, JMP_EXPORTER_LABELS, JMP_LEASE, JUMPSTARTER_HOST + + +@dataclass(frozen=True) +class ExporterMetadata: + """Metadata about the exporter connected to the current shell session.""" + + name: str + labels: dict[str, str] = field(default_factory=dict) + lease: str | None = None + + @classmethod + def from_env(cls) -> "ExporterMetadata": + """Build metadata from JMP_EXPORTER, JMP_EXPORTER_LABELS, and JMP_LEASE env vars.""" + name = os.environ.get(JMP_EXPORTER, "") + lease = os.environ.get(JMP_LEASE) or None + + labels: dict[str, str] = {} + raw_labels = os.environ.get(JMP_EXPORTER_LABELS, "") + if raw_labels: + for pair in raw_labels.split(","): + if "=" in pair: + k, v = pair.split("=", 1) + labels[k] = v + + return cls(name=name, labels=labels, lease=lease) @asynccontextmanager @@ -38,6 +64,19 @@ async def env_async(portal, stack): client.close() +@asynccontextmanager +async def env_with_metadata_async(portal, stack): + """Provide a client and exporter metadata for an existing Jumpstarter shell. + + Async version of env_with_metadata() + + Yields a (client, ExporterMetadata) tuple. The metadata is read from environment + variables set by ``jmp shell``: JMP_EXPORTER, JMP_EXPORTER_LABELS, and JMP_LEASE. + """ + async with env_async(portal, stack) as client: + yield client, ExporterMetadata.from_env() + + @contextmanager def env(): """Provide a client for an existing JUMPSTARTER_HOST environment variable. @@ -49,3 +88,23 @@ def env(): with ExitStack() as stack: with portal.wrap_async_context_manager(env_async(portal, stack)) as client: yield client + + +@contextmanager +def env_with_metadata(): + """Provide a client and exporter metadata for an existing Jumpstarter shell. + + This is useful when you need both the client and information about the connected + exporter (name, labels, lease ID). + + Example:: + + with env_with_metadata() as (client, metadata): + print(metadata.name) + print(metadata.labels) + print(metadata.lease) + """ + with start_blocking_portal() as portal: + with ExitStack() as stack: + with portal.wrap_async_context_manager(env_with_metadata_async(portal, stack)) as result: + yield result