github.com/mweagle/Sparta@v1.15.0/aws/step/sagemaker.go (about)

     1  package step
     2  
     3  import (
     4  	"math/rand"
     5  
     6  	gocf "github.com/mweagle/go-cloudformation"
     7  )
     8  
     9  // SageMakerTag represents a tag for a SageMaker task
    10  type SageMakerTag struct {
    11  	Key   string          `json:",omitempty"`
    12  	Value gocf.Stringable `json:",omitempty"`
    13  }
    14  
    15  // SageMakerTrainingJobParameters are the parameters for a SageMaker Training
    16  // Job as defined by https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTrainingJob.html
    17  type SageMakerTrainingJobParameters struct {
    18  	AlgorithmSpecification map[string]interface{} `json:",omitempty"`
    19  	HyperParameters        map[string]interface{} `json:",omitempty"`
    20  	InputDataConfig        []interface{}          `json:",omitempty"`
    21  	OutputDataConfig       map[string]interface{} `json:",omitempty"`
    22  	ResourceConfig         map[string]interface{} `json:",omitempty"`
    23  	RoleArn                gocf.Stringable        `json:",omitempty"`
    24  	StoppingCondition      interface{}            `json:",omitempty"`
    25  	Tags                   []SageMakerTag         `json:",omitempty"`
    26  	TrainingJobName        string                 `json:",omitempty"`
    27  	VpcConfig              map[string]interface{} `json:",omitempty"`
    28  }
    29  
    30  ////////////////////////////////////////////////////////////////////////////////
    31  /*
    32    _____         _      _
    33   |_   _| _ __ _(_)_ _ (_)_ _  __ _
    34     | || '_/ _` | | ' \| | ' \/ _` |
    35     |_||_| \__,_|_|_||_|_|_||_\__, |
    36                               |___/
    37  */
    38  ////////////////////////////////////////////////////////////////////////////////
    39  
    40  // SageMakerTrainingJob represents bindings for
    41  // https://docs.aws.amazon.com/step-functions/latest/dg/connectors-sqs.html
    42  type SageMakerTrainingJob struct {
    43  	BaseTask
    44  	parameters SageMakerTrainingJobParameters
    45  }
    46  
    47  // MarshalJSON for custom marshalling, since this will be stringified and we need it
    48  // to turn into a stringified
    49  // Ref: https://docs.aws.amazon.com/step-functions/latest/dg/connectors-sqs.html
    50  func (smtj *SageMakerTrainingJob) MarshalJSON() ([]byte, error) {
    51  	return smtj.BaseTask.marshalMergedParams("arn:aws:states:::sagemaker:createTrainingJob.sync",
    52  		&smtj.parameters)
    53  }
    54  
    55  // NewSageMakerTrainingJob returns an initialized SQSTaskState
    56  func NewSageMakerTrainingJob(stateName string,
    57  	parameters SageMakerTrainingJobParameters) *SageMakerTrainingJob {
    58  	smtj := &SageMakerTrainingJob{
    59  		BaseTask: BaseTask{
    60  			baseInnerState: baseInnerState{
    61  				name: stateName,
    62  				id:   rand.Int63(),
    63  			},
    64  		},
    65  		parameters: parameters,
    66  	}
    67  	return smtj
    68  }
    69  
    70  ////////////////////////////////////////////////////////////////////////////////
    71  /*
    72  _____                  __
    73  |_   _| _ __ _ _ _  ___/ _|___ _ _ _ __
    74  	| || '_/ _` | ' \(_-<  _/ _ \ '_| '  \
    75  	|_||_| \__,_|_||_/__/_| \___/_| |_|_|_|
    76  */
    77  ////////////////////////////////////////////////////////////////////////////////
    78  
    79  // SageMakerTransformJobParameters are the parameters for a SageMaker Training
    80  // Job as defined by https://docs.aws.amazon.com/sagemaker/latest/dg/API_CreateTransformJob.html
    81  type SageMakerTransformJobParameters struct {
    82  	BatchStrategy           string                 `json:",omitempty"`
    83  	Environment             map[string]string      `json:",omitempty"`
    84  	MaxConcurrentTransforms int                    `json:",omitempty"`
    85  	MaxPayloadInMB          int                    `json:",omitempty"`
    86  	ModelName               string                 `json:",omitempty"`
    87  	Tags                    []SageMakerTag         `json:",omitempty"`
    88  	TransformInput          map[string]interface{} `json:",omitempty"`
    89  	TransformJobName        string                 `json:",omitempty"`
    90  	TransformOutput         map[string]interface{} `json:",omitempty"`
    91  	TransformResources      map[string]interface{} `json:",omitempty"`
    92  }
    93  
    94  // SageMakerTransformJob represents bindings for
    95  // https://docs.aws.amazon.com/step-functions/latest/dg/connectors-sqs.html
    96  type SageMakerTransformJob struct {
    97  	BaseTask
    98  	parameters SageMakerTransformJobParameters
    99  }
   100  
   101  // MarshalJSON for custom marshalling, since this will be stringified and we need it
   102  // to turn into a stringified
   103  // Ref: https://docs.aws.amazon.com/step-functions/latest/dg/connectors-sqs.html
   104  func (smtj *SageMakerTransformJob) MarshalJSON() ([]byte, error) {
   105  	return smtj.BaseTask.marshalMergedParams("arn:aws:states:::sagemaker:createTransformJob.sync",
   106  		&smtj.parameters)
   107  }
   108  
   109  // NewSageMakerTransformJob returns an initialized SQSTaskState
   110  func NewSageMakerTransformJob(stateName string,
   111  	parameters SageMakerTransformJobParameters) *SageMakerTransformJob {
   112  	sns := &SageMakerTransformJob{
   113  		BaseTask: BaseTask{
   114  			baseInnerState: baseInnerState{
   115  				name: stateName,
   116  				id:   rand.Int63(),
   117  			},
   118  		},
   119  		parameters: parameters,
   120  	}
   121  	return sns
   122  }