Skip to content

nos.common.spec

nos.common.spec.FunctionSignature

Function signature that fully describes the remote-model to be executed including inputs, outputs, func_or_cls to be executed, initialization args/kwargs.

Source code in nos/common/spec.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class FunctionSignature(BaseModel):
    """Function signature that fully describes the remote-model to be executed
    including `inputs`, `outputs`, `func_or_cls` to be executed,
    initialization `args`/`kwargs`."""

    func_or_cls: Callable
    """Class instance."""
    method: str
    """Class method name. (e.g. forward, __call__ etc)"""

    init_args: Tuple[Any, ...] = Field(default_factory=tuple)
    """Arguments to initialize the model instance."""
    init_kwargs: Dict[str, Any] = Field(default_factory=dict)
    """Keyword arguments to initialize the model instance."""

    parameters: Dict[str, Any] = Field(init=False, default_factory=dict)
    """Input function signature (as returned by inspect.signature)."""
    return_annotation: Any = Field(init=False, default=None)
    """Output / return function signature (as returned by inspect.signature)."""

    input_annotations: Dict[str, Any] = Field(default_factory=dict)
    """Mapping of input keyword arguments to dtypes."""
    output_annotations: Union[Any, Dict[str, Any], None] = Field(default=None)
    """Mapping of output names to dtypes."""

    def __init__(
        self,
        func_or_cls: Callable,
        method: str,
        init_args: Tuple[Any, ...] = (),
        init_kwargs: Dict[str, Any] = {},  # noqa: B006
        parameters: Dict[str, Any] = {},  # noqa: B006
        return_annotation: Any = None,
        input_annotations: Dict[str, Any] = {},  # noqa: B006
        output_annotations: Union[Any, Dict[str, Any], None] = None,
    ):
        super().__init__(
            func_or_cls=func_or_cls,
            method=method,
            init_args=init_args,
            init_kwargs=init_kwargs,
            parameters=parameters,
            return_annotation=return_annotation,
            input_annotations=input_annotations,
            output_annotations=output_annotations,
        )
        if not callable(self.func_or_cls):
            raise ValueError(f"Invalid function/class provided, func_or_cls={self.func_or_cls}.")

        if not self.method or not hasattr(self.func_or_cls, self.method):
            raise ValueError(f"Invalid method name provided, method={self.method}.")

        # Get the function signature
        sig: Dict[str, inspect.Parameter] = inspect.signature(getattr(self.func_or_cls, self.method))

        # Get the input/output annotations
        self.parameters = sig.parameters
        self.return_annotation = sig.return_annotation
        logger.debug(f"Function signature [method={self.method}, sig={sig}].")

    @staticmethod
    def validate(inputs: Dict[str, Any], sig: Dict[str, Any]) -> Dict[str, Any]:
        """Validate the input dict against the defined signature (input or output)."""
        # TOFIX (spillai): This needs to be able to validate using args/kwargs instead
        if not set(inputs.keys()).issubset(set(sig.keys())):  # noqa: W503
            raise InputValidationException(
                f"Invalid inputs, provided={set(inputs.keys())}, expected={set(sig.keys())}."
            )
        # TODO (spillai): Validate input types and shapes.
        return inputs

    @field_validator("init_args", mode="before")
    @classmethod
    def _validate_init_args(cls, init_args: Union[Tuple[Any, ...], Any]) -> Tuple[Any, ...]:
        """Validate the initialization arguments."""
        # TODO (spillai): Check the function signature of the func_or_cls class and validate
        # the init_args against the signature.
        return init_args

    @field_validator("init_kwargs", mode="before")
    @classmethod
    def _validate_init_kwargs(cls, init_kwargs: Dict[str, Any]) -> Dict[str, Any]:
        """Validate the initialization keyword arguments."""
        # TODO (spillai): Check the function signature of the func_or_cls class and validate
        # the init_kwargs against the signature.
        return init_kwargs

    def _encode_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Encode inputs based on defined signature."""
        inputs = FunctionSignature.validate(inputs, self.parameters)
        return {k: dumps(v) for k, v in inputs.items()}

    def _decode_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Decode inputs based on defined signature."""
        inputs = FunctionSignature.validate(inputs, self.parameters)
        return {k: loads(v) for k, v in inputs.items()}

func_or_cls instance-attribute 

func_or_cls: Callable

Class instance.

method instance-attribute 

method: str

Class method name. (e.g. forward, call etc)

init_args class-attribute instance-attribute 

init_args: Tuple[Any, ...] = Field(default_factory=tuple)

Arguments to initialize the model instance.

init_kwargs class-attribute instance-attribute 

init_kwargs: Dict[str, Any] = Field(default_factory=dict)

Keyword arguments to initialize the model instance.

input_annotations class-attribute instance-attribute 

input_annotations: Dict[str, Any] = Field(default_factory=dict)

Mapping of input keyword arguments to dtypes.

output_annotations class-attribute instance-attribute 

output_annotations: Union[Any, Dict[str, Any], None] = Field(default=None)

Mapping of output names to dtypes.

parameters class-attribute instance-attribute 

parameters: Dict[str, Any] = parameters

Input function signature (as returned by inspect.signature).

return_annotation class-attribute instance-attribute 

return_annotation: Any = return_annotation

Output / return function signature (as returned by inspect.signature).

validate staticmethod 

validate(inputs: Dict[str, Any], sig: Dict[str, Any]) -> Dict[str, Any]

Validate the input dict against the defined signature (input or output).

Source code in nos/common/spec.py 
87
88
89
90
91
92
93
94
95
96
@staticmethod
def validate(inputs: Dict[str, Any], sig: Dict[str, Any]) -> Dict[str, Any]:
    """Validate the input dict against the defined signature (input or output)."""
    # TOFIX (spillai): This needs to be able to validate using args/kwargs instead
    if not set(inputs.keys()).issubset(set(sig.keys())):  # noqa: W503
        raise InputValidationException(
            f"Invalid inputs, provided={set(inputs.keys())}, expected={set(sig.keys())}."
        )
    # TODO (spillai): Validate input types and shapes.
    return inputs

nos.common.spec.ModelSpec

Model specification for the registry.

ModelSpec captures all the relevant information for the instantiation, runtime and execution of a model.

Source code in nos/common/spec.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
class ModelSpec(BaseModel):
    """Model specification for the registry.

    ModelSpec captures all the relevant information for
    the instantiation, runtime and execution of a model.
    """

    id: str
    """Model identifier."""
    signature: Dict[str, FunctionSignature] = Field(default_factory=dict)
    """Model function signatures to export (method -> FunctionSignature)."""
    runtime_env: Union[RuntimeEnv, None] = Field(default=None)
    """Runtime environment with custom packages."""

    def __init__(
        self, id: str, signature: Dict[str, FunctionSignature] = {}, runtime_env: RuntimeEnv = None  # noqa: B006
    ):
        super().__init__(id=id, signature=signature, runtime_env=runtime_env)

    def __repr__(self):
        return f"""ModelSpec(id={self.id}, methods=({', '.join(list(self.signature.keys()))}), tasks=({', '.join([str(self.task(m)) for m in self.signature])}))"""

    @field_validator("id", mode="before")
    def _validate_id(cls, id: str) -> str:
        """Validate the model identifier."""
        regex = re.compile(r"^[a-zA-Z0-9\/._-]+$")  # allow alphanumerics, `/`, `.`, `_`, and `-`
        if not regex.match(id):
            raise ValueError(
                f"Invalid model id, id={id} can only contain alphanumerics characters, `/`, `.`, `_`, and `-`"
            )
        return id

    @field_validator("signature", mode="before")
    def _validate_signature(
        cls, signature: Union[FunctionSignature, Dict[str, FunctionSignature]], **kwargs: Dict[str, Any]
    ) -> Dict[str, FunctionSignature]:
        """Validate the model signature / signatures.

        Checks that the model class `cls` has the function name attribute
        as defined in the signature `function_name`.

        Args:
            signature (Union[FunctionSignature, Dict[str, FunctionSignature]]): Model signature.
            **kwargs: Keyword arguments.
        Returns:
            Dict[str, FunctionSignature]: Model signature.
        """
        if isinstance(signature, (list, tuple)):
            raise TypeError(f"Invalid signature provided, signature={signature}.")
        if isinstance(signature, FunctionSignature):
            signature = {signature.method: signature}
        for method, sig in signature.items():
            if method != sig.method:
                raise ValueError(f"Invalid method name provided, method={method}, sig.method={sig.method}.")
            if sig and sig.func_or_cls:
                model_cls = sig.func_or_cls
                if sig.method and not hasattr(model_cls, sig.method):
                    raise ValueError(f"Model class {model_cls} does not have function {sig.method}.")
        return signature

    @property
    def name(self) -> str:
        """Return the model name (for backwards compatibility)."""
        return self.id

    def task(self, method: str = None) -> TaskType:
        """Return the task type for a given method (or defaults to default method)."""
        if method is None:
            method = self.default_method
        try:
            md = self.metadata(method)
            return md.task
        except Exception:
            logger.debug(f"Model metadata not found, id={self.id}.")
            return None

    def metadata(self, method: str = None) -> ModelSpecMetadata:
        """Return the model spec metadata for a given method (or defaults to default method)."""
        if method is None:
            method = self.default_method
        catalog = ModelSpecMetadataCatalog.get()
        try:
            metadata: ModelSpecMetadata = catalog._metadata_catalog[f"{self.id}/{method}"]
        except KeyError:
            logger.debug(f"Model metadata not found in catalog, id={self.id}, method={method}.")
            return ModelSpecMetadata(id=self.id, method=method)
        return metadata

    @cached_property
    def default_method(self) -> str:
        """Return the default method name."""
        assert len(self.signature) > 0, f"No default signature found, signature={self.signature}."
        return list(self.signature.keys())[0]

    @cached_property
    def default_signature(self) -> FunctionSignature:
        """Return the default function signature.

        Returns:
            FunctionSignature: Default function signature.
        """
        # Note (spillai): For now, we assume that the first
        # signature is the default signature. In the `.from_cls()`
        # method, we add the __call__ method as the first method
        # for this exact reason.
        return self.signature[self.default_method]

    def set_default_method(self, method: str) -> None:
        """Set the default method name."""
        if method not in self.signature:
            raise ValueError(f"Invalid method name provided, method={method}.")

        # Update the default method in the signature
        signature = {}
        signature[method] = self.signature.pop(method)
        signature.update(self.signature)
        self.signature = signature
        # Clear the cached properties to force re-computation
        self.__dict__.pop("default_method", None)
        self.__dict__.pop("default_signature", None)

    def __call__(self, *init_args, **init_kwargs) -> Any:
        """Create a model instance.

        This method allows us to create a model instance directly
        from the model spec. Let's consider the example below

            ```
            class CustomModel:
                ...

            CustomModel = ModelSpec.from_cls(CustomModel)
            model = CustomModel(*init_args, **init_kwargs)
            ```

        Args:
            *init_args: Positional arguments.
            **init_kwargs: Keyword arguments.
        Returns:
            Any: Model instance.
        """
        sig: FunctionSignature = self.default_signature
        return sig.func_or_cls(*init_args, **init_kwargs)

    @classmethod
    def from_yaml(cls, filename: str) -> "ModelSpec":
        raise NotImplementedError()

    @classmethod
    def from_cls(
        cls,
        func_or_cls: Callable,
        init_args: Tuple[Any, ...] = (),
        init_kwargs: Dict[str, Any] = {},  # noqa: B006
        method: str = "__call__",
        runtime_env: RuntimeEnv = None,
        model_id: str = None,
        **kwargs: Any,
    ) -> "ModelSpec":
        """Wrap custom models/classes into a nos-compatible model spec.

        Args:
            func_or_cls (Callable): Model function or class. For now, only classes are supported.
            init_args (Tuple[Any, ...]): Initialization arguments.
            init_kwargs (Dict[str, Any]): Initialization keyword arguments.
            method (str): Method name to be executed.
            runtime_env (RuntimeEnv): Runtime environment specification.
            model_id (str): Optional model identifier.
            **kwargs: Additional keyword arguments.
                These include `init_args` and `init_kwargs` to initialize the model instance.

        Returns:
            ModelSpec: The resulting model specification that fully describes the model execution.
        """
        # Check if the cls is not a function
        if not callable(func_or_cls) or not inspect.isclass(func_or_cls):
            raise ValueError(f"Invalid class `{func_or_cls}` provided, needs to be a class object.")

        # Check if the cls has the method_name
        if not hasattr(func_or_cls, method):
            raise ValueError(f"Invalid method name `{method}` provided.")

        # TODO (spillai): Provide additional RayRuntimeEnvConfig as `config`
        # config = dict(setup_timeout_seconds=10 * 60, eager_install=True)
        if runtime_env:
            logger.debug(f"Using custom runtime_env [env={runtime_env}]")

        # Inspect all the public methods of the class
        # and expose them as model methods
        ignore_methods = ["__init__", method]
        all_methods = [m for m, _ in inspect.getmembers(func_or_cls, predicate=inspect.isfunction)]
        methods = [m for m in all_methods if m not in ignore_methods]

        # Note (spillai): See .default_signature property for why we add
        # the __call__ method as the first method.
        if method in all_methods:
            methods.insert(0, method)  # first method is the default method
        logger.debug(f"Registering methods [methods={methods}].")

        # Add function signature for each method
        model_id: str = func_or_cls.__name__ if model_id is None else model_id
        signature: Dict[str, FunctionSignature] = {}
        metadata: Dict[str, ModelSpecMetadata] = {}
        for method in methods:
            # Add the function signature
            sig = FunctionSignature(
                func_or_cls=func_or_cls,
                method=method,
                init_args=init_args,
                init_kwargs=init_kwargs,
            )
            signature[method] = sig
            metadata[method] = ModelSpecMetadata(model_id, method, task=None)
            logger.debug(f"Added function signature [method={method}, signature={sig}].")

        # Build the model spec from the function signature
        spec = cls(
            model_id,
            signature=signature,
            runtime_env=runtime_env,
        )
        return spec

    def _to_proto(self) -> nos_service_pb2.GenericResponse:
        """Convert the model spec to proto."""
        spec: ModelSpec = loads(dumps(self, protocol=-1))
        # Note (spillai): We only serialize the input/output
        # signatures and method of the spec. Notably, the
        # `func_or_cls` attribute is not serialized to avoid
        # the dependency on torch and other server-side dependencies.
        for method in spec.signature:
            spec.signature[method].func_or_cls = None
            spec.signature[method].init_args = ()
            spec.signature[method].init_kwargs = {}
        return nos_service_pb2.GenericResponse(
            response_bytes=dumps(spec),
        )

    @staticmethod
    def _from_proto(minfo: nos_service_pb2.GenericResponse) -> "ModelSpec":
        """Convert the generic response back to the spec."""
        return loads(minfo.response_bytes)

id instance-attribute 

id: str

Model identifier.

signature class-attribute instance-attribute 

signature: Dict[str, FunctionSignature] = Field(default_factory=dict)

Model function signatures to export (method -> FunctionSignature).

runtime_env class-attribute instance-attribute 

runtime_env: Union[RuntimeEnv, None] = Field(default=None)

Runtime environment with custom packages.

name property 

name: str

Return the model name (for backwards compatibility).

default_method cached property 

default_method: str

Return the default method name.

default_signature cached property 

default_signature: FunctionSignature

Return the default function signature.

Returns:

task 

task(method: str = None) -> TaskType

Return the task type for a given method (or defaults to default method).

Source code in nos/common/spec.py 
454
455
456
457
458
459
460
461
462
463
def task(self, method: str = None) -> TaskType:
    """Return the task type for a given method (or defaults to default method)."""
    if method is None:
        method = self.default_method
    try:
        md = self.metadata(method)
        return md.task
    except Exception:
        logger.debug(f"Model metadata not found, id={self.id}.")
        return None

metadata 

metadata(method: str = None) -> ModelSpecMetadata

Return the model spec metadata for a given method (or defaults to default method).

Source code in nos/common/spec.py 
465
466
467
468
469
470
471
472
473
474
475
def metadata(self, method: str = None) -> ModelSpecMetadata:
    """Return the model spec metadata for a given method (or defaults to default method)."""
    if method is None:
        method = self.default_method
    catalog = ModelSpecMetadataCatalog.get()
    try:
        metadata: ModelSpecMetadata = catalog._metadata_catalog[f"{self.id}/{method}"]
    except KeyError:
        logger.debug(f"Model metadata not found in catalog, id={self.id}, method={method}.")
        return ModelSpecMetadata(id=self.id, method=method)
    return metadata

set_default_method 

set_default_method(method: str) -> None

Set the default method name.

Source code in nos/common/spec.py 
496
497
498
499
500
501
502
503
504
505
506
507
508
def set_default_method(self, method: str) -> None:
    """Set the default method name."""
    if method not in self.signature:
        raise ValueError(f"Invalid method name provided, method={method}.")

    # Update the default method in the signature
    signature = {}
    signature[method] = self.signature.pop(method)
    signature.update(self.signature)
    self.signature = signature
    # Clear the cached properties to force re-computation
    self.__dict__.pop("default_method", None)
    self.__dict__.pop("default_signature", None)

__call__ 

__call__(*init_args, **init_kwargs) -> Any

Create a model instance.

This method allows us to create a model instance directly from the model spec. Let's consider the example below

```
class CustomModel:
    ...

CustomModel = ModelSpec.from_cls(CustomModel)
model = CustomModel(*init_args, **init_kwargs)
```

Parameters:

  • *init_args –

    Positional arguments.

  • **init_kwargs –

    Keyword arguments.

Returns: Any: Model instance.

Source code in nos/common/spec.py 
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
def __call__(self, *init_args, **init_kwargs) -> Any:
    """Create a model instance.

    This method allows us to create a model instance directly
    from the model spec. Let's consider the example below

        ```
        class CustomModel:
            ...

        CustomModel = ModelSpec.from_cls(CustomModel)
        model = CustomModel(*init_args, **init_kwargs)
        ```

    Args:
        *init_args: Positional arguments.
        **init_kwargs: Keyword arguments.
    Returns:
        Any: Model instance.
    """
    sig: FunctionSignature = self.default_signature
    return sig.func_or_cls(*init_args, **init_kwargs)

from_cls classmethod 

from_cls(func_or_cls: Callable, init_args: Tuple[Any, ...] = (), init_kwargs: Dict[str, Any] = {}, method: str = '__call__', runtime_env: RuntimeEnv = None, model_id: str = None, **kwargs: Any) -> ModelSpec

Wrap custom models/classes into a nos-compatible model spec.

Parameters:

  • func_or_cls (Callable) –

    Model function or class. For now, only classes are supported.

  • init_args (Tuple[Any, ...], default: () ) –

    Initialization arguments.

  • init_kwargs (Dict[str, Any], default: {} ) –

    Initialization keyword arguments.

  • method (str, default: '__call__' ) –

    Method name to be executed.

  • runtime_env (RuntimeEnv, default: None ) –

    Runtime environment specification.

  • model_id (str, default: None ) –

    Optional model identifier.

  • **kwargs (Any, default: {} ) –

    Additional keyword arguments. These include init_args and init_kwargs to initialize the model instance.

Returns:

  • ModelSpec ( ModelSpec ) –

    The resulting model specification that fully describes the model execution.

Source code in nos/common/spec.py 
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
@classmethod
def from_cls(
    cls,
    func_or_cls: Callable,
    init_args: Tuple[Any, ...] = (),
    init_kwargs: Dict[str, Any] = {},  # noqa: B006
    method: str = "__call__",
    runtime_env: RuntimeEnv = None,
    model_id: str = None,
    **kwargs: Any,
) -> "ModelSpec":
    """Wrap custom models/classes into a nos-compatible model spec.

    Args:
        func_or_cls (Callable): Model function or class. For now, only classes are supported.
        init_args (Tuple[Any, ...]): Initialization arguments.
        init_kwargs (Dict[str, Any]): Initialization keyword arguments.
        method (str): Method name to be executed.
        runtime_env (RuntimeEnv): Runtime environment specification.
        model_id (str): Optional model identifier.
        **kwargs: Additional keyword arguments.
            These include `init_args` and `init_kwargs` to initialize the model instance.

    Returns:
        ModelSpec: The resulting model specification that fully describes the model execution.
    """
    # Check if the cls is not a function
    if not callable(func_or_cls) or not inspect.isclass(func_or_cls):
        raise ValueError(f"Invalid class `{func_or_cls}` provided, needs to be a class object.")

    # Check if the cls has the method_name
    if not hasattr(func_or_cls, method):
        raise ValueError(f"Invalid method name `{method}` provided.")

    # TODO (spillai): Provide additional RayRuntimeEnvConfig as `config`
    # config = dict(setup_timeout_seconds=10 * 60, eager_install=True)
    if runtime_env:
        logger.debug(f"Using custom runtime_env [env={runtime_env}]")

    # Inspect all the public methods of the class
    # and expose them as model methods
    ignore_methods = ["__init__", method]
    all_methods = [m for m, _ in inspect.getmembers(func_or_cls, predicate=inspect.isfunction)]
    methods = [m for m in all_methods if m not in ignore_methods]

    # Note (spillai): See .default_signature property for why we add
    # the __call__ method as the first method.
    if method in all_methods:
        methods.insert(0, method)  # first method is the default method
    logger.debug(f"Registering methods [methods={methods}].")

    # Add function signature for each method
    model_id: str = func_or_cls.__name__ if model_id is None else model_id
    signature: Dict[str, FunctionSignature] = {}
    metadata: Dict[str, ModelSpecMetadata] = {}
    for method in methods:
        # Add the function signature
        sig = FunctionSignature(
            func_or_cls=func_or_cls,
            method=method,
            init_args=init_args,
            init_kwargs=init_kwargs,
        )
        signature[method] = sig
        metadata[method] = ModelSpecMetadata(model_id, method, task=None)
        logger.debug(f"Added function signature [method={method}, signature={sig}].")

    # Build the model spec from the function signature
    spec = cls(
        model_id,
        signature=signature,
        runtime_env=runtime_env,
    )
    return spec