From ac2d9ae5a3069d04f88df7648dbe2d5bf877d3f4 Mon Sep 17 00:00:00 2001 From: John Date: Thu, 26 Dec 2024 15:04:04 -0700 Subject: [PATCH 1/4] Chore: Sanitize container requests and add stub id to container events --- pkg/repository/events.go | 22 ++++++++++++++++++---- pkg/types/event.go | 6 ++++-- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/pkg/repository/events.go b/pkg/repository/events.go index 42edddce8..f74b46ac9 100644 --- a/pkg/repository/events.go +++ b/pkg/repository/events.go @@ -111,7 +111,8 @@ func (t *TCPEventClientRepo) PushContainerRequestedEvent(request *types.Containe types.EventContainerStatusRequestedSchemaVersion, types.EventContainerStatusRequestedSchema{ ContainerID: request.ContainerId, - Request: *request, + Request: sanitizeContainerRequest(request), + StubID: request.StubId, Status: types.EventContainerLifecycleRequested, }, ) @@ -124,7 +125,8 @@ func (t *TCPEventClientRepo) PushContainerScheduledEvent(containerID string, wor types.EventContainerLifecycleSchema{ ContainerID: containerID, WorkerID: workerID, - Request: *request, + Request: sanitizeContainerRequest(request), + StubID: request.StubId, Status: types.EventContainerLifecycleScheduled, }, ) @@ -137,7 +139,8 @@ func (t *TCPEventClientRepo) PushContainerStartedEvent(containerID string, worke types.EventContainerLifecycleSchema{ ContainerID: containerID, WorkerID: workerID, - Request: *request, + Request: sanitizeContainerRequest(request), + StubID: request.StubId, Status: types.EventContainerLifecycleStarted, }, ) @@ -150,7 +153,8 @@ func (t *TCPEventClientRepo) PushContainerStoppedEvent(containerID string, worke types.EventContainerLifecycleSchema{ ContainerID: containerID, WorkerID: workerID, - Request: *request, + Request: sanitizeContainerRequest(request), + StubID: request.StubId, Status: types.EventContainerLifecycleStopped, }, ) @@ -295,3 +299,13 @@ func (t *TCPEventClientRepo) PushStubStateUnhealthy(workspaceId string, stubId s }, ) } + +func sanitizeContainerRequest(request *types.ContainerRequest) types.ContainerRequest { + requestCopy := *request + requestCopy.Env = nil + requestCopy.EntryPoint = nil + requestCopy.Stub = types.StubWithRelated{} + requestCopy.Mounts = nil + requestCopy.PoolSelector = "" + return requestCopy +} diff --git a/pkg/types/event.go b/pkg/types/event.go index 6aaa841a4..9111b13c2 100644 --- a/pkg/types/event.go +++ b/pkg/types/event.go @@ -57,11 +57,12 @@ var ( // Schema versions should be in ISO 8601 format -var EventContainerLifecycleSchemaVersion = "1.0" +var EventContainerLifecycleSchemaVersion = "1.1" type EventContainerLifecycleSchema struct { ContainerID string `json:"container_id"` WorkerID string `json:"worker_id"` + StubID string `json:"stub_id"` Status string `json:"status"` Request ContainerRequest `json:"request"` } @@ -95,11 +96,12 @@ type EventContainerMetricsData struct { GPUType string `json:"gpu_type"` } -var EventContainerStatusRequestedSchemaVersion = "1.0" +var EventContainerStatusRequestedSchemaVersion = "1.1" type EventContainerStatusRequestedSchema struct { ContainerID string `json:"container_id"` Request ContainerRequest `json:"request"` + StubID string `json:"stub_id"` Status string `json:"status"` } From c80890bfe7328949e683fb7d60c99b00c768f4c0 Mon Sep 17 00:00:00 2001 From: John Date: Thu, 26 Dec 2024 16:18:54 -0700 Subject: [PATCH 2/4] Chore: Need to filter out signing key as well --- pkg/repository/events.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/repository/events.go b/pkg/repository/events.go index f74b46ac9..2d3b41f77 100644 --- a/pkg/repository/events.go +++ b/pkg/repository/events.go @@ -307,5 +307,6 @@ func sanitizeContainerRequest(request *types.ContainerRequest) types.ContainerRe requestCopy.Stub = types.StubWithRelated{} requestCopy.Mounts = nil requestCopy.PoolSelector = "" + requestCopy.Workspace.SigningKey = nil return requestCopy } From b04c02a25a14dc39aa11efa84a8e9c3269fdae1f Mon Sep 17 00:00:00 2001 From: John Date: Thu, 26 Dec 2024 16:24:34 -0700 Subject: [PATCH 3/4] Chore: Prune more data from the SanitizedContainerRequest --- pkg/repository/events.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/repository/events.go b/pkg/repository/events.go index 2d3b41f77..0edfc5fc7 100644 --- a/pkg/repository/events.go +++ b/pkg/repository/events.go @@ -307,6 +307,6 @@ func sanitizeContainerRequest(request *types.ContainerRequest) types.ContainerRe requestCopy.Stub = types.StubWithRelated{} requestCopy.Mounts = nil requestCopy.PoolSelector = "" - requestCopy.Workspace.SigningKey = nil + requestCopy.Workspace = types.Workspace{} return requestCopy } From 59f8cbce855b32766f7647f88466d4215365a93b Mon Sep 17 00:00:00 2001 From: John Date: Thu, 26 Dec 2024 16:24:41 -0700 Subject: [PATCH 4/4] Chore: Prune more data from the SanitizedContainerRequest --- pkg/repository/events.go | 26 +++++++++++++++++--------- pkg/types/event.go | 35 ++++++++++++++++++++++++++--------- 2 files changed, 43 insertions(+), 18 deletions(-) diff --git a/pkg/repository/events.go b/pkg/repository/events.go index 0edfc5fc7..5b4278fa6 100644 --- a/pkg/repository/events.go +++ b/pkg/repository/events.go @@ -300,13 +300,21 @@ func (t *TCPEventClientRepo) PushStubStateUnhealthy(workspaceId string, stubId s ) } -func sanitizeContainerRequest(request *types.ContainerRequest) types.ContainerRequest { - requestCopy := *request - requestCopy.Env = nil - requestCopy.EntryPoint = nil - requestCopy.Stub = types.StubWithRelated{} - requestCopy.Mounts = nil - requestCopy.PoolSelector = "" - requestCopy.Workspace = types.Workspace{} - return requestCopy +func sanitizeContainerRequest(request *types.ContainerRequest) types.SanitizedContainerRequest { + return types.SanitizedContainerRequest{ + ContainerId: request.ContainerId, + Cpu: request.Cpu, + Memory: request.Memory, + Gpu: request.Gpu, + GpuRequest: request.GpuRequest, + GpuCount: request.GpuCount, + ImageId: request.ImageId, + StubId: request.StubId, + WorkspaceId: request.WorkspaceId, + RetryCount: request.RetryCount, + PoolSelector: request.PoolSelector, + Preemptable: request.Preemptable, + CheckpointEnabled: request.CheckpointEnabled, + BuildOptions: request.BuildOptions, + } } diff --git a/pkg/types/event.go b/pkg/types/event.go index 9111b13c2..7df46be86 100644 --- a/pkg/types/event.go +++ b/pkg/types/event.go @@ -55,16 +55,33 @@ var ( EventWorkerLifecycleStopped = "stopped" ) +type SanitizedContainerRequest struct { + ContainerId string `json:"container_id"` + Cpu int64 `json:"cpu"` + Memory int64 `json:"memory"` + Gpu string `json:"gpu"` + GpuRequest []string `json:"gpu_request"` + GpuCount uint32 `json:"gpu_count"` + ImageId string `json:"image_id"` + StubId string `json:"stub_id"` + WorkspaceId string `json:"workspace_id"` + RetryCount int `json:"retry_count"` + PoolSelector string `json:"pool_selector"` + Preemptable bool `json:"preemptable"` + CheckpointEnabled bool `json:"checkpoint_enabled"` + BuildOptions BuildOptions `json:"build_options"` +} + // Schema versions should be in ISO 8601 format var EventContainerLifecycleSchemaVersion = "1.1" type EventContainerLifecycleSchema struct { - ContainerID string `json:"container_id"` - WorkerID string `json:"worker_id"` - StubID string `json:"stub_id"` - Status string `json:"status"` - Request ContainerRequest `json:"request"` + ContainerID string `json:"container_id"` + WorkerID string `json:"worker_id"` + StubID string `json:"stub_id"` + Status string `json:"status"` + Request SanitizedContainerRequest `json:"request"` } var EventContainerMetricsSchemaVersion = "1.0" @@ -99,10 +116,10 @@ type EventContainerMetricsData struct { var EventContainerStatusRequestedSchemaVersion = "1.1" type EventContainerStatusRequestedSchema struct { - ContainerID string `json:"container_id"` - Request ContainerRequest `json:"request"` - StubID string `json:"stub_id"` - Status string `json:"status"` + ContainerID string `json:"container_id"` + Request SanitizedContainerRequest `json:"request"` + StubID string `json:"stub_id"` + Status string `json:"status"` } var EventWorkerLifecycleSchemaVersion = "1.0"