Customizing Endpoint Resolution

The sample APIs provide an option for developers to customize how the streaming endpoints are resolved.

By default the resolve class is Generic, the sample also includes an AWS Targetgroup specific resolver class.

Both can be found within the container at the following path: /usr/local/lib/python3.11/site-packages/nv/svc/streaming/_csp.py

The class and arguments to load at startup is managed by the following settings:

  • backend.csp.cls

  • backend.csp.args

These can be overwritten via the following settings in the helm values file:

  • streaming.serviceConfig.backend_csp_cls

  • streaming.serviceConfig.backend_csp_args

Below is the Generic class defined in _csp.py. It inherits from the _CSP class. The class always get passed a configured Kubernetes client that uses the service account defined in the helm chart of the streaming session manager. An example of it being used can be seen in the _fetch_resources function inside the _CSP class.

class Generic(_CSP):
    """Generic and default CSP manager."""

    def __init__(
        self,
        k8s_client: KubernetesClient,
        enable_wss: bool = False,
        hostname_annotation_key: str = "external-dns.alpha.kubernetes.io/hostname",
        service_annotations_location: str = "streamingKit.service.annotations",
        base_domain: str = ""
    ) -> Any:
        """Initialize."""
        super().__init__(k8s_client)
        self._enable_wss = enable_wss
        self._service_annotations_location = service_annotations_location
        self._hostname_annotation_key = hostname_annotation_key
        self._base_domain = base_domain

    async def on_create(self, profile_data: Dict, settings: Dict) -> Dict:
        """Process CSP customisations on stream creation."""
        if not self._enable_wss:
            return {}

        values = profile_data.get("settings", {}).get("values", {})
        keys = self._service_annotations_location.split(".")
        data = values
        for key in keys:
            data = data.get(key, {})

        prefix = self._generate_random_dns_prefix()
        data[self._hostname_annotation_key] = f"{prefix}.{self._base_domain}"
        settings[self._service_annotations_location] = data

        logging.debug(f"Generated settings {settings}")

        return settings

    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
        """Resolve the endpoint to connect to for a given session."""
        services = await self._fetch_resources("service", selectors={"sessionId": session_id})
        routes, status = await self._extract_routes(services)
        return routes, status
121class Generic(_CSP):
122    """Generic and default CSP manager."""
123
124    def __init__(
125        self,
126        k8s_client: KubernetesClient,
127        enable_wss: bool = False,
128        hostname_annotation_key: str = "external-dns.alpha.kubernetes.io/hostname",
129        service_annotations_location: str = "streamingKit.service.annotations",
130        base_domain: str = ""
131    ) -> Any:
132        """Initialize."""
133        super().__init__(k8s_client)
134        self._enable_wss = enable_wss
135        self._service_annotations_location = service_annotations_location
136        self._hostname_annotation_key = hostname_annotation_key
137        self._base_domain = base_domain
138
139    async def on_create(self, profile_data: Dict, settings: Dict) -> Dict:
140        """Process CSP customisations on stream creation."""
141        if not self._enable_wss:
142            return {}
143
144        values = profile_data.get("settings", {}).get("values", {})
145        keys = self._service_annotations_location.split(".")
146        data = values
147        for key in keys:
148            data = data.get(key, {})
149
150        prefix = self._generate_random_dns_prefix()
151        data[self._hostname_annotation_key] = f"{prefix}.{self._base_domain}"
152        settings[self._service_annotations_location] = data
153
154        logging.debug(f"Generated settings {settings}")
155
156        return settings
157
158    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
159        """Resolve the endpoint to connect to for a given session."""
160        services = await self._fetch_resources("service", selectors={"sessionId": session_id})
161        routes, status = await self._extract_routes(services)
162        return routes, status
163
164    def _extract_ports(self, service_spec: Dict) -> Tuple[Dict, bool]:
165        """Extract the port information from the service specification.
166
167        Args:
168            service_spec (Dict): The service specification dictionary.
169
170        Returns:
171            Dict: A dictionary containing the route information.
172        """
173        ports = service_spec.spec.ports
174
175        routes = []
176        for port in ports:
177            routes.append(
178                {
179                    "source_port": port.port,
180                    "description": port.name,
181                    "protocol": port.protocol,
182                    "destination_port": port.node_port
183                }
184            )
185
186        status = True if routes else False
187
188        logging.debug(f"Extracted port mappings {routes}, readiness status {status}")
189        return {"routes": routes}, status
190
191    async def _extract_routes(self, services):
192        routes = {}
193        statuses = []
194
195        for service in services:
196            ports, port_ready = self._extract_ports(service)
197
198            entries = []
199            lb_ready = False
200            if self._enable_wss:
201                hostname, lb_ready = await self._extract_hostname(service)
202                entries.append(hostname)
203            else:
204                ips, lb_ready = await self._extract_lb_ips(service)
205                entries.extend(ips)
206
207            statuses.append(all([port_ready, lb_ready]))
208            if not lb_ready:
209                continue
210
211            for entry in entries:
212                routes[entry] = ports
213
214        status = all(statuses) if statuses else False
215        logging.debug(f"Extracted routes {routes}, readiness status {status}")
216        return routes, status
217
218    async def _extract_lb_ips(self, service) -> Tuple[List, bool]:
219        ips = []
220        hostname = None
221
222        ingress = service.status.load_balancer.ingress
223        if not ingress:
224            return ips, False
225
226        for entry in ingress:
227            if entry.ip:
228                ips.append(entry.ip)
229            hostname = entry.hostname
230
231        if not ips and hostname:
232            logging.debug("No IPs were found attached to the service, trying to resolve hostname")
233            ips = await self._resolve_hostname(hostname)
234
235        status = True if ips else False
236
237        logging.debug(f"Extracted IPs {ips}, readiness status {status}")
238        return ips, status
239
240    async def _extract_hostname(self, service) -> Tuple[str, bool]:
241        annotations = service.metadata.annotations
242        hostname = None
243
244        try:
245            hostname = annotations[self._hostname_annotation_key]
246        except KeyError:
247            logging.error(f"Hostname field `{self._hostname_annotation_key}` not found in annotations")
248            raise HostnameNotFoundError("Unable to find hostname")
249
250        ready = False
251        try:
252            ips = await self._resolve_hostname(hostname)
253            if ips:
254                ready = True
255        except aiodns.error.DNSError as exc:
256            logging.warning(f"Unable to resolve {hostname}: {exc}")
257
258        return hostname, ready
class _CSP(object):
    """Custom CSP behaviour base class."""

    def __init__(self, k8s_client: KubernetesClient) -> None:
        """Initialize."""
        self._k8s_client = k8s_client
        self._default_namespace = "omni-streaming"

    async def on_create(self, profile_data: dict, settings: dict) -> Dict:
        """Process CSP customisations on stream creation."""
        return {}

    async def on_delete(self, data: dict) -> None:
        """Proccess CSP customisations on stream termination."""
        pass

    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
        """Resolve the endpoint to connect to for a given session."""
        return {}, False

    async def _fetch_resources(self, resource_type, selectors=None, args: Dict = None):
        api_class = self._get_k8s_api_class(resource_type)

        args = args or {}
        async with self._k8s_client as api_client:
            api_instance = api_class(api_client.api_http)
            func_name = args.pop('func_name', f"list_namespaced_{resource_type}")
            func = getattr(api_instance, func_name)

            if selectors:
                args['label_selector'] = ",".join(f"{k}={v}" for k, v in selectors.items())

            res = await func(
                namespace=self._default_namespace,
                **args
            )
            return res['items'] if isinstance(res, dict) else res.items
 40class _CSP(object):
 41    """Custom CSP behaviour base class."""
 42
 43    def __init__(self, k8s_client: KubernetesClient) -> None:
 44        """Initialize."""
 45        self._k8s_client = k8s_client
 46        self._default_namespace = "omni-streaming"
 47
 48    def _generate_random_dns_prefix(self, length=6):
 49        letters = string.ascii_lowercase
 50        return ''.join(random.choice(letters) for _ in range(length))
 51
 52    async def on_create(self, profile_data: dict, settings: dict) -> Dict:
 53        """Process CSP customisations on stream creation."""
 54        return {}
 55
 56    async def on_delete(self, data: dict) -> None:
 57        """Proccess CSP customisations on stream termination."""
 58        pass
 59
 60    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
 61        """Resolve the endpoint to connect to for a given session."""
 62        return {}, False
 63
 64    async def _resolve_hostname(self, hostname):
 65        ips = []
 66
 67        resolver = aiodns.DNSResolver()
 68        try:
 69            ipv4_records = await resolver.query(hostname, 'A')
 70        except aiodns.error.DNSError as exc:
 71            logging.warning(f"Failed to resolve DNS for {hostname}. The domain might not have propagated yet if this is a new stream. {exc}")
 72            return []
 73
 74        ips.extend([record.host for record in ipv4_records])
 75        return ips
 76
 77    async def _fetch_resources(self, resource_type, selectors=None, args: Dict = None):
 78        api_class = self._get_k8s_api_class(resource_type)
 79
 80        args = args or {}
 81        async with self._k8s_client as api_client:
 82            api_instance = api_class(api_client.api_http)
 83            func_name = args.pop('func_name', f"list_namespaced_{resource_type}")
 84            func = getattr(api_instance, func_name)
 85
 86            if selectors:
 87                args['label_selector'] = ",".join(f"{k}={v}" for k, v in selectors.items())
 88
 89            res = await func(
 90                namespace=self._default_namespace,
 91                **args
 92            )
 93            return res['items'] if isinstance(res, dict) else res.items
 94
 95    def _get_k8s_api_class(self, resource_type):
 96        resource_to_api_class = {
 97            'pod': CoreV1Api,
 98            'service': CoreV1Api,
 99            'deployment': AppsV1Api,
100            'targetgroupbinding': CustomObjectsApi,
101
102        }
103        return resource_to_api_class.get(resource_type, CoreV1Api)
104
105    async def _tcp_port_ready(self, host: str, port: int) -> bool:
106        try:
107            logging.debug(f"Testing stream readiness on {host}:{port}")
108            reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=2)
109            writer.close()
110            await writer.wait_closed()
111            logging.debug(f"Stream connection on {host}:{port} successful")
112            return True
113        except asyncio.TimeoutError:
114            logging.info(f"Failed to connect to {host}:{port}. Connection not ready.")
115        except Exception as exc:
116            logging.error(f"Failed to connect to {host}:{port}: {exc}")
117
118        return False
class AWS(_CSP):
    """AWS customisations."""

    def __init__(
        self,
        k8s_client: KubernetesClient,
        nlb_mgmt_svc: str = "",
        enable_wss: bool = False
    ) -> Any:
        """Initialize AWS class."""
        super().__init__(k8s_client=k8s_client)
        self._nlb_mgmt_svc_url = nlb_mgmt_svc

        self._port_locations = {
            "media": "streamingKit.service.mediaPort",
            "signaling": "streamingKit.service.signalingPort"
        }

        self._targetgroup_arn_locations = {
            "media": "streamingKit.aws.targetgroups.media",
            "signaling": "streamingKit.aws.targetgroups.signaling",
        }

        self._listeners_arn_locations = {
            "media": "streamingKit.aws.listeners.media",
            "signaling": "streamingKit.aws.listeners.signaling",
        }

        self._nlb_location = "streamingKit.aws.nlb"
        self._alias_location = "streamingKit.aws.alias"
        self._enable_wss = enable_wss

    async def on_create(self, profile_data: dict, settings: dict) -> Dict:
        """Process AWS customisations on stream creation."""
        values = profile_data["settings"]["values"]

        ports = {}
        for port_name, location in self._port_locations.items():
            ports[port_name] = self._lookup_nested_dict(values, location)

        allocations = []
        default_protocol = "TLS" if self._enable_wss else "TCP"

        for name in ports.keys():
            allocations.append(
                {
                    "name": name,
                    "protocol": "UDP" if name == "media" else default_protocol
                }
            )

        url = f"{self._nlb_mgmt_svc_url}/allocation"
        async with aiohttp.ClientSession() as session:
            async with session.post(url, json={"allocations": allocations}) as resp:
                if resp.status not in [200]:
                    detail = await resp.text()
                    error_msg = f"Failed request to {url}: {resp.status}, {detail}"
                    logging.error(error_msg)
                    raise _APIError(status_code=resp.status, details=error_msg)

                arns = await resp.json()

        settings[self._nlb_location] = arns["loadbalancer"]["dnsName"]
        settings[self._alias_location] = arns["loadbalancer"].get("alias", "")

        for key, location in self._targetgroup_arn_locations.items():
            listener_arn = arns['allocations'][key]['listenerArn']
            listener_port = arns['allocations'][key]['listenerPort']
            listener_protocol = arns['allocations'][key]['listenerProtocol']
            settings[location] = arns["allocations"][key]["targetGroupArn"]
            settings[self._listeners_arn_locations[key]] = f"{listener_arn}@{listener_port}@{listener_protocol}@{key}"

        logging.debug(f"Generated settings {settings}")
        return settings

    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
        """Resolve the endpoint to connect to for a given session."""

        args = {
            'func_name': 'list_namespaced_custom_object',
            'group': 'elbv2.k8s.aws',
            'version': 'v1beta1',
            'plural': 'targetgroupbindings',
        }

        tgbs = await self._fetch_resources(
            "targetgroupbinding",
            selectors={"sessionId": session_id},
            args=args
        )

        routes, status = await self._extract_routes(tgbs)

        return routes, status
261class AWS(_CSP):
262    """AWS customisations."""
263
264    def __init__(
265        self,
266        k8s_client: KubernetesClient,
267        nlb_mgmt_svc: str = "",
268        enable_wss: bool = False
269    ) -> Any:
270        """Initialize AWS class."""
271        super().__init__(k8s_client=k8s_client)
272        self._nlb_mgmt_svc_url = nlb_mgmt_svc
273
274        self._port_locations = {
275            "media": "streamingKit.service.mediaPort",
276            "signaling": "streamingKit.service.signalingPort"
277        }
278
279        self._targetgroup_arn_locations = {
280            "media": "streamingKit.aws.targetgroups.media",
281            "signaling": "streamingKit.aws.targetgroups.signaling",
282        }
283
284        self._listeners_arn_locations = {
285            "media": "streamingKit.aws.listeners.media",
286            "signaling": "streamingKit.aws.listeners.signaling",
287        }
288
289        self._nlb_location = "streamingKit.aws.nlb"
290        self._alias_location = "streamingKit.aws.alias"
291        self._enable_wss = enable_wss
292
293    def _lookup_nested_dict(self, nested_dict, key_string):
294        keys = key_string.split('.')
295        value = nested_dict
296        for key in keys:
297            value = value[key]
298        return value
299
300    async def on_create(self, profile_data: dict, settings: dict) -> Dict:
301        """Process AWS customisations on stream creation."""
302        values = profile_data["settings"]["values"]
303
304        ports = {}
305        for port_name, location in self._port_locations.items():
306            ports[port_name] = self._lookup_nested_dict(values, location)
307
308        allocations = []
309        default_protocol = "TLS" if self._enable_wss else "TCP"
310
311        for name in ports.keys():
312            allocations.append(
313                {
314                    "name": name,
315                    "protocol": "UDP" if name == "media" else default_protocol
316                }
317            )
318
319        url = f"{self._nlb_mgmt_svc_url}/allocation"
320        async with aiohttp.ClientSession() as session:
321            async with session.post(url, json={"allocations": allocations}) as resp:
322                if resp.status not in [200]:
323                    detail = await resp.text()
324                    error_msg = f"Failed request to {url}: {resp.status}, {detail}"
325                    logging.error(error_msg)
326                    raise _APIError(status_code=resp.status, details=error_msg)
327
328                arns = await resp.json()
329
330        settings[self._nlb_location] = arns["loadbalancer"]["dnsName"]
331        settings[self._alias_location] = arns["loadbalancer"].get("alias", "")
332
333        for key, location in self._targetgroup_arn_locations.items():
334            listener_arn = arns['allocations'][key]['listenerArn']
335            listener_port = arns['allocations'][key]['listenerPort']
336            listener_protocol = arns['allocations'][key]['listenerProtocol']
337            settings[location] = arns["allocations"][key]["targetGroupArn"]
338            settings[self._listeners_arn_locations[key]] = f"{listener_arn}@{listener_port}@{listener_protocol}@{key}"
339
340        logging.debug(f"Generated settings {settings}")
341        return settings
342
343    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
344        """Resolve the endpoint to connect to for a given session."""
345
346        args = {
347            'func_name': 'list_namespaced_custom_object',
348            'group': 'elbv2.k8s.aws',
349            'version': 'v1beta1',
350            'plural': 'targetgroupbindings',
351        }
352
353        tgbs = await self._fetch_resources(
354            "targetgroupbinding",
355            selectors={"sessionId": session_id},
356            args=args
357        )
358
359        routes, status = await self._extract_routes(tgbs)
360
361        return routes, status
362
363    async def _extract_routes(self, tgbs: List) -> Tuple[Dict, bool] :
364        routes = []
365        statuses = []
366
367        hostnames = []
368
369        for tgb in tgbs:
370            annotations = tgb["metadata"]["annotations"]
371            listener = annotations.get('nvidia.com/omniverse.listener', '')
372            nlb_hostname = annotations.get('nvidia.com/omniverse.nlb', '')
373            alias = annotations.get('nvidia.com/omniverse.alias', '')
374
375            hostname = alias if self._enable_wss else nlb_hostname
376
377            if not listener or not hostname:
378                statuses.append(False)
379                logging.error(f"Invalid targetgroupbinding was found. No listener annotation was found {tgb}")
380                continue
381
382            listener_arn, port, protocol, name = listener.split("@")
383            hostnames.append(hostname)
384            routes.append({
385                "source_port": int(port),
386                "description": name,
387                "protocol": protocol,
388                "destination_port": -1
389            })
390            status = await self._tcp_port_ready(hostname, port) if protocol.lower() == "tcp" else True
391            statuses.append(status)
392
393        status = all(statuses) if statuses else False
394
395        assembled_routes = {}
396
397        ips = []
398        if not self._enable_wss:
399            for hostname in hostnames:
400                ips.extend(await self._resolve_hostname(hostname))
401            hostnames = ips
402
403        for hostname in hostnames:
404            assembled_routes[hostname] = {"routes": routes}
405
406        logging.debug(f"Extracted routes {assembled_routes}, readiness status {status}")
407        return assembled_routes, status

_csp.py

Creating a custom CSP resolver

Because the class and arguments to load can be done dynamically it is possible for a developer to provide their own class and arguments as long as the python file containing the class is within the container and within a python package (ie: an __init__.py file is present).

There are three functions that allow custom behavior:

  • on_create

    • Allows customizing the chart values that are passed to the rmcp service before the session gets instantiated.

    • Takes as arguments the chart values that come from resolving the profile, this gives access to the resolved values and settings, the overrides specified within a version.

    • It returns an updated settings dictionary that are resolved to update the chart values.

  • on_delete

    • Allows additional actions to be taken when a session is terminated

    • Takes as an argument a dictionary which currently only contains the session_id

    • Returns nothing.

  • resolve_endpoints

    • Determines the routes passed back to the web client, and if the session is ready to connect to

    • Takes as an argument the session_id.

    • Returns a tuple of a dictionary with routes and a boolean indicating if the session is ready.

      • If the boolean is false, a 202 will be returned to the client.

Example custom endpoint resolver

Here is an example of creating a custom endpoint resolver and deploying it to an Omniverse Application Streaming API instance.

Example python class that extracts a streaming session’s POD IP and ports
 1from typing import Tuple, Dict, List
 2
 3from nv.svc.streaming._csp import Generic
 4
 5class NodePortResolver(Generic):
 6
 7    async def resolve_endpoints(self, session_id: str) -> Tuple[Dict, bool]:
 8        """Resolve the endpoint to connect to for a given session."""
 9        services = await self._fetch_resources("service", selectors={"sessionId": session_id})
10        pods = await self._fetch_resources("pod", selectors={"sessionId": session_id})
11        routes, status = await self._extract_routes(services, pods)
12        return routes, status
13
14    def _extract_ports(self, service_spec: Dict) -> Tuple[Dict, bool]:
15        """Extract the port information from the service specification.
16
17        Args:
18            service_spec (Dict): The service specification dictionary.
19
20        Returns:
21            Dict: A dictionary containing the route information.
22        """
23        ports = service_spec.spec.ports
24
25        routes = []
26        for port in ports:
27            routes.append(
28                {
29                    "source_port": port.node_port,
30                    "description": port.name,
31                    "protocol": port.protocol,
32                    "destination_port": 0 # this field is no longer used.
33                }
34            )
35
36        status = True if routes else False
37
38        return {"routes": routes}, status
39
40    async def _extract_routes(self, services, pods):
41        routes = {}
42        statuses = []
43
44        for service in services:
45            ports, port_ready = self._extract_ports(service)
46            ips, ip_ready = await self._extract_pod_node_ips(pods)
47            statuses.append(port_ready and ip_ready)
48
49            if not ip_ready:
50                continue
51
52            for ip in ips:
53                routes[ip] = ports
54
55        status = all(statuses) if statuses else False
56        return routes, status
57
58    async def _extract_pod_node_ips(self, pods) -> Tuple[List[str], bool]:
59        """Extract the host IPs of the nodes hosting the pods of the stream."""
60        node_ips = set()
61        for pod in pods.items:  # .items to iterate through pod list
62            if pod.status.host_ip:
63                node_ips.add(pod.status.host_ip)
64
65        return list(node_ips), bool(node_ips)
  1. node_port.py

  2. Create a dockerfile that copies node_port.py into the streaming session manager container:

    FROM nvcr.io/nvidia/omniverse/kit-appstreaming-manager:1.4.0
    
    # Create the extensions directory
    RUN mkdir -p /resolvers
    
    # Copy the NodePortResolver.py into /extensions
    COPY node_port.py /resolvers/
    
    # Create an empty __init__.py to make it a package
    RUN touch /resolvers/__init__.py
    
  3. Build and push the container, replacing the registry, name and version:

    docker build -t {my-registry}/{kit-appstreaming-manager-node-port}:{1.0.0}
    docker push {my-registry}/{kit-appstreaming-manager-node-port}:{1.0.0}
    
  4. Create or update the values file to use the new container and instruct it to use the new resolver class (this is a partial values file with just the required fields changed):

    streaming:
       image:
          # -- Image repository.
          repository: "my-registry/kit-appstreaming-manager-node-port"
          # -- Image pull policy.
          pullPolicy: Always
          # -- Image tag.
          tag: 1.0.0
    
       serviceConfig:
          # -- CSP customisation manager
          backend_csp_cls: "resolvers.node_port:NodePortResolver"
    
          # -- CSP Customation manager arguments
          backend_csp_args: {}
    
  5. Apply the helm chart.

At this point, the NodePortResolver will be loaded and called to resolve the routes based on NodePort and HostIP.