github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/ext/etl/emd.go (about)

     1  // Package etl provides utilities to initialize and use transformation pods.
     2  /*
     3   * Copyright (c) 2021-2022, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package etl
     6  
     7  import (
     8  	"encoding/json"
     9  	"fmt"
    10  
    11  	"github.com/NVIDIA/aistore/cmn"
    12  	"github.com/NVIDIA/aistore/cmn/cos"
    13  	"github.com/NVIDIA/aistore/cmn/debug"
    14  	"github.com/NVIDIA/aistore/cmn/jsp"
    15  	jsoniter "github.com/json-iterator/go"
    16  )
    17  
    18  type (
    19  	ETLs map[string]InitMsg
    20  
    21  	// ETL metadata
    22  	MD struct {
    23  		Version int64
    24  		ETLs    ETLs
    25  		Ext     any
    26  	}
    27  
    28  	jsonETL struct {
    29  		Type string              `json:"type,string"`
    30  		Msg  jsoniter.RawMessage `json:"msg"`
    31  	}
    32  	jsonMD struct {
    33  		Version int64              `json:"version"`
    34  		ETLs    map[string]jsonETL `json:"etls"`
    35  		Ext     any                `json:"ext,omitempty"` // within meta-version extensions
    36  	}
    37  )
    38  
    39  var etlMDJspOpts = jsp.CCSign(cmn.MetaverEtlMD)
    40  
    41  // interface guard
    42  var (
    43  	_ json.Marshaler   = (*MD)(nil)
    44  	_ json.Unmarshaler = (*MD)(nil)
    45  
    46  	_ jsp.Opts = (*MD)(nil)
    47  )
    48  
    49  ////////
    50  // MD //
    51  ////////
    52  
    53  func (e *MD) Init(l int)         { e.ETLs = make(ETLs, l) }
    54  func (e *MD) Add(msg InitMsg)    { e.ETLs[msg.Name()] = msg }
    55  func (*MD) JspOpts() jsp.Options { return etlMDJspOpts }
    56  
    57  func (e *MD) Get(id string) (msg InitMsg, present bool) {
    58  	if e == nil {
    59  		return
    60  	}
    61  	msg, present = e.ETLs[id]
    62  	return
    63  }
    64  
    65  func (e *MD) Del(id string) (deleted bool) {
    66  	if _, present := e.ETLs[id]; !present {
    67  		return
    68  	}
    69  	delete(e.ETLs, id)
    70  	return true
    71  }
    72  
    73  func (e *MD) String() string {
    74  	if e == nil {
    75  		return "EtlMD <nil>"
    76  	}
    77  	return fmt.Sprintf("EtlMD v%d(%d)", e.Version, len(e.ETLs))
    78  }
    79  
    80  func (e *MD) MarshalJSON() ([]byte, error) {
    81  	jsonMD := jsonMD{
    82  		Version: e.Version,
    83  		ETLs:    make(map[string]jsonETL, len(e.ETLs)),
    84  		Ext:     e.Ext,
    85  	}
    86  	for k, v := range e.ETLs {
    87  		jsonMD.ETLs[k] = jsonETL{v.MsgType(), cos.MustMarshal(v)}
    88  	}
    89  	return jsoniter.Marshal(jsonMD)
    90  }
    91  
    92  func (e *MD) UnmarshalJSON(data []byte) (err error) {
    93  	jsonMD := &jsonMD{}
    94  	if err = jsoniter.Unmarshal(data, jsonMD); err != nil {
    95  		return
    96  	}
    97  	e.Version, e.Ext = jsonMD.Version, jsonMD.Ext
    98  	e.ETLs = make(ETLs, len(jsonMD.ETLs))
    99  	for k, v := range jsonMD.ETLs {
   100  		switch v.Type {
   101  		case Code:
   102  			e.ETLs[k] = &InitCodeMsg{}
   103  		case Spec:
   104  			e.ETLs[k] = &InitSpecMsg{}
   105  		default:
   106  			err = fmt.Errorf("invalid InitMsg type %q", v.Type)
   107  			debug.AssertNoErr(err)
   108  			return
   109  		}
   110  		if err = jsoniter.Unmarshal(v.Msg, e.ETLs[k]); err != nil {
   111  			break
   112  		}
   113  	}
   114  	return
   115  }