sashatrask
1002 lines
44 KiB
Python
1002 lines
44 KiB
Python
import hashlib
|
|
import os
|
|
import re
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import datetime, timedelta
|
|
from types import MappingProxyType, SimpleNamespace
|
|
|
|
from console_runner import (
|
|
_GitResolutionFailure,
|
|
git_resolution_failure_result,
|
|
prepare_scan_options,
|
|
reserve_v2_admission_with_recovery,
|
|
resolve_and_bind_git_claim,
|
|
validate_v2_capacity_model,
|
|
)
|
|
from process_identity import current_process_identity
|
|
from result_bundle import BundleReservation, ResultBundleReader, bundle_ready_path
|
|
from lifecycle_authority import DISCOVERY_PRODUCER_SOURCES
|
|
from scan_execution import (
|
|
PROTOCOL_VERSION,
|
|
QueueDispositionPolicy,
|
|
ScanCompatibility,
|
|
ScanExecutionError,
|
|
WorkerBuildCompatibility,
|
|
canonical_json_sha256,
|
|
normalize_docker_direct_execution_snapshot,
|
|
normalize_docker_direct_execution_target,
|
|
normalize_exact_git_execution_snapshot,
|
|
normalize_huggingface_space_execution_snapshot,
|
|
normalize_huggingface_space_execution_target,
|
|
normalize_remote_scan_policy,
|
|
remote_execution_identity,
|
|
remote_execution_snapshot_sha256,
|
|
stage_scan_result_in_scope,
|
|
validate_scan_kwargs,
|
|
validate_worker_build_compatibility,
|
|
)
|
|
from scanner_db import ScannerDB, ScanEventConflictError
|
|
from scanner import scan_config
|
|
from runtime_security import sha256_file
|
|
from worker_contracts import (
|
|
DIAGNOSTIC_PROJECTION_VERSION,
|
|
ordered_diagnostic_uid_set_sha256,
|
|
)
|
|
from worker_package import (
|
|
KNOWN_WORKER_PACKAGE_CAPABILITIES,
|
|
PACKAGE_DETECTOR_POLICY,
|
|
load_worker_package_manifest,
|
|
normalize_worker_package_manifest,
|
|
worker_package_build_compatibility,
|
|
)
|
|
|
|
|
|
DEFAULT_ASSIGNMENT_TTL_SECONDS = 24 * 60 * 60
|
|
DEFAULT_RESULT_UPLOAD_BODY_TIMEOUT_SECONDS = 30 * 60
|
|
PROTOCOL1_NEW_CLAIM_SOURCES = frozenset(('github', 'gitlab'))
|
|
SUPPORTED_REMOTE_GIT_SOURCES = PROTOCOL1_NEW_CLAIM_SOURCES
|
|
_EXACT_GIT_ASSIGNMENT_FLOW = object()
|
|
_DIRECT_ASSIGNMENT_FLOW = object()
|
|
NO_WORK_REASONS = frozenset((
|
|
'empty_queue', 'assignment_cap', 'dispatch_paused', 'capacity',
|
|
'compatibility',
|
|
))
|
|
_ADMISSION_NO_WORK_REASONS = {
|
|
'no_claimable_target': 'empty_queue',
|
|
'remote_user_quota_closed': 'assignment_cap',
|
|
'remote_global_quota_closed': 'assignment_cap',
|
|
'dispatch_gate_closed': 'dispatch_paused',
|
|
'pipeline_capacity_closed': 'capacity',
|
|
'quarantine_capacity_conflict': 'capacity',
|
|
'quarantine_admission_closed': 'capacity',
|
|
'ingester_not_ready': 'capacity',
|
|
}
|
|
_CANONICAL_CAPABILITY_NAME = re.compile(r'[a-z][a-z0-9_]{0,63}\Z')
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssignmentCapability:
|
|
source: str
|
|
platform: str
|
|
planning_kind: str
|
|
|
|
def as_dict(self):
|
|
return {
|
|
'source': self.source,
|
|
'platform': self.platform,
|
|
'planning_kind': self.planning_kind,
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssignmentSourceAdapter:
|
|
queue_source: str
|
|
worker_platform: str
|
|
planning_kind: str
|
|
package_capability: AssignmentCapability
|
|
snapshot_validator: object
|
|
assignment_flow: object = None
|
|
|
|
def validate_source_args(self, args):
|
|
if str(getattr(args, 'platform', '') or '').strip().lower() != self.worker_platform:
|
|
raise ValueError('remote assignment source platform does not match its adapter')
|
|
if self.assignment_flow is None:
|
|
raise ValueError('remote assignment source is not available for new claims')
|
|
if self.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW and (
|
|
self.planning_kind != 'exact_git_v1'
|
|
or not bool(getattr(args, 'exact_git_planning_enabled', False))
|
|
):
|
|
raise ValueError('remote Git sources require enabled exact Git planning')
|
|
if self.assignment_flow is _DIRECT_ASSIGNMENT_FLOW:
|
|
if self.planning_kind not in {
|
|
'docker_direct_v1', 'huggingface_space_v1',
|
|
} or any(str(getattr(args, name, '') or '') for name in (
|
|
'token', 'docker_username', 'docker_token', 'auth_name',
|
|
)):
|
|
raise ValueError('remote direct sources must be credential-free')
|
|
elif self.assignment_flow is not _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
raise ValueError('remote assignment source flow is invalid')
|
|
|
|
def validate_snapshot(self, value):
|
|
normalized = self.snapshot_validator(value)
|
|
if (
|
|
normalized['execution']['source'] != self.worker_platform
|
|
or normalized['planning']['kind'] != self.planning_kind
|
|
or normalized['credential_ref']['source'] != self.queue_source
|
|
):
|
|
raise ScanExecutionError('remote execution snapshot does not match its source adapter')
|
|
return normalized
|
|
|
|
|
|
def _exact_git_snapshot(value):
|
|
return normalize_exact_git_execution_snapshot(value)
|
|
|
|
|
|
def _adapter(source, platform, planning_kind, *, assignment_flow=None, validator=None):
|
|
capability = AssignmentCapability(source, platform, planning_kind)
|
|
return AssignmentSourceAdapter(
|
|
source, platform, planning_kind, capability,
|
|
validator, assignment_flow,
|
|
)
|
|
|
|
|
|
LEGACY_GITHUB_ASSIGNMENT_ADAPTER = _adapter(
|
|
'github', 'github', 'exact_git_v1',
|
|
assignment_flow=_EXACT_GIT_ASSIGNMENT_FLOW,
|
|
validator=_exact_git_snapshot,
|
|
)
|
|
CORE_ASSIGNMENT_SOURCE_ADAPTERS = MappingProxyType({
|
|
'gitlab': _adapter(
|
|
'gitlab', 'gitlab', 'exact_git_v1',
|
|
assignment_flow=_EXACT_GIT_ASSIGNMENT_FLOW,
|
|
validator=_exact_git_snapshot,
|
|
),
|
|
'dockerhub': _adapter(
|
|
'dockerhub', 'docker', 'docker_direct_v1',
|
|
assignment_flow=_DIRECT_ASSIGNMENT_FLOW,
|
|
validator=normalize_docker_direct_execution_snapshot,
|
|
),
|
|
'huggingface': _adapter(
|
|
'huggingface', 'huggingface', 'huggingface_space_v1',
|
|
assignment_flow=_DIRECT_ASSIGNMENT_FLOW,
|
|
validator=normalize_huggingface_space_execution_snapshot,
|
|
),
|
|
})
|
|
ASSIGNMENT_SOURCE_ADAPTERS = MappingProxyType({
|
|
'github': LEGACY_GITHUB_ASSIGNMENT_ADAPTER,
|
|
**CORE_ASSIGNMENT_SOURCE_ADAPTERS,
|
|
})
|
|
PROTOCOL2_NEW_CLAIM_SOURCES = frozenset(CORE_ASSIGNMENT_SOURCE_ADAPTERS)
|
|
|
|
|
|
def _validate_assignment_source_adapters():
|
|
if tuple(CORE_ASSIGNMENT_SOURCE_ADAPTERS) != tuple(DISCOVERY_PRODUCER_SOURCES):
|
|
raise RuntimeError('core assignment adapters do not match discovery producers')
|
|
pairs = set()
|
|
for key, adapter in ASSIGNMENT_SOURCE_ADAPTERS.items():
|
|
capability = adapter.package_capability
|
|
identities = (
|
|
key, adapter.queue_source, adapter.worker_platform,
|
|
adapter.planning_kind,
|
|
)
|
|
if any(_CANONICAL_CAPABILITY_NAME.fullmatch(value) is None for value in identities):
|
|
raise RuntimeError('remote assignment adapter identity is invalid')
|
|
if key != adapter.queue_source or capability.as_dict() != {
|
|
'source': adapter.queue_source,
|
|
'platform': adapter.worker_platform,
|
|
'planning_kind': adapter.planning_kind,
|
|
}:
|
|
raise RuntimeError('remote assignment adapter capability is inconsistent')
|
|
pair = (adapter.queue_source, adapter.worker_platform)
|
|
if pair in pairs:
|
|
raise RuntimeError('remote assignment adapter source/platform is duplicated')
|
|
pairs.add(pair)
|
|
if {
|
|
source for source, adapter in ASSIGNMENT_SOURCE_ADAPTERS.items()
|
|
if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW
|
|
} != set(PROTOCOL1_NEW_CLAIM_SOURCES):
|
|
raise RuntimeError('protocol-1 assignment source registry is inconsistent')
|
|
if {
|
|
source for source, adapter in ASSIGNMENT_SOURCE_ADAPTERS.items()
|
|
if adapter.assignment_flow is _DIRECT_ASSIGNMENT_FLOW
|
|
} != {'dockerhub', 'huggingface'}:
|
|
raise RuntimeError('direct assignment source registry is inconsistent')
|
|
if {
|
|
(
|
|
adapter.package_capability.source,
|
|
adapter.package_capability.platform,
|
|
adapter.package_capability.planning_kind,
|
|
)
|
|
for adapter in ASSIGNMENT_SOURCE_ADAPTERS.values()
|
|
} != set(KNOWN_WORKER_PACKAGE_CAPABILITIES):
|
|
raise RuntimeError('worker package capabilities do not match source adapters')
|
|
|
|
|
|
_validate_assignment_source_adapters()
|
|
|
|
|
|
def assignment_source_adapter(source):
|
|
source = str(source or '').strip().lower()
|
|
adapter = ASSIGNMENT_SOURCE_ADAPTERS.get(source)
|
|
if adapter is None:
|
|
raise ValueError('unsupported remote assignment source')
|
|
return adapter
|
|
|
|
|
|
def _stable_id(device_id, request_id, label):
|
|
payload = f'truf-worker-v1\0{int(device_id)}\0{request_id}\0{label}'.encode('ascii')
|
|
return hashlib.sha256(payload).hexdigest()[:32]
|
|
|
|
|
|
def _hash_file(path):
|
|
digest = hashlib.sha256()
|
|
with open(path, 'rb', buffering=0) as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b''):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
class RemoteAssignmentBuilder:
|
|
"""Central remote admission builder routed through closed source adapters."""
|
|
|
|
def __init__(
|
|
self, db_url, bundle_root, source_args, compatibility_profiles,
|
|
supervisor_instance_id, *, assignment_ttl_seconds=DEFAULT_ASSIGNMENT_TTL_SECONDS,
|
|
assignment_ttl_seconds_by_source=None,
|
|
result_upload_body_timeout_seconds=DEFAULT_RESULT_UPLOAD_BODY_TIMEOUT_SECONDS,
|
|
credential_refs=None,
|
|
db_factory=ScannerDB, admission=reserve_v2_admission_with_recovery,
|
|
planner=resolve_and_bind_git_claim,
|
|
):
|
|
self.db_url = str(db_url or '')
|
|
self.bundle_root = os.path.abspath(bundle_root)
|
|
self.supervisor_instance_id = str(supervisor_instance_id or '')
|
|
if not self.db_url or not self.supervisor_instance_id:
|
|
raise ValueError('remote assignment builder requires database and supervisor identities')
|
|
if (
|
|
type(assignment_ttl_seconds) is not int
|
|
or not 60 <= assignment_ttl_seconds <= 7 * 24 * 60 * 60
|
|
):
|
|
raise ValueError('remote assignment lifetime must be between one minute and seven days')
|
|
self.assignment_ttl_seconds = assignment_ttl_seconds
|
|
raw_ttl_overrides = assignment_ttl_seconds_by_source
|
|
if raw_ttl_overrides is None:
|
|
raw_ttl_overrides = {}
|
|
if (
|
|
type(raw_ttl_overrides) is not dict
|
|
or set(raw_ttl_overrides) - set(CORE_ASSIGNMENT_SOURCE_ADAPTERS)
|
|
):
|
|
raise ValueError('remote assignment lifetime overrides contain unsupported sources')
|
|
self.assignment_ttl_seconds_by_source = {}
|
|
for source, value in raw_ttl_overrides.items():
|
|
if type(value) is not int or not 60 <= value <= 7 * 24 * 60 * 60:
|
|
raise ValueError('remote assignment lifetime override is outside its bounds')
|
|
self.assignment_ttl_seconds_by_source[source] = value
|
|
if (
|
|
type(result_upload_body_timeout_seconds) is not int
|
|
or not 30 <= result_upload_body_timeout_seconds <= 24 * 60 * 60
|
|
):
|
|
raise ValueError('remote result upload body timeout is outside its bounds')
|
|
self.result_upload_body_timeout_seconds = result_upload_body_timeout_seconds
|
|
self.db_factory = db_factory
|
|
self.admission = admission
|
|
self.planner = planner
|
|
|
|
self.source_args = {}
|
|
self.source_adapters = {}
|
|
for source, args in dict(source_args or {}).items():
|
|
source = str(source or '').strip().lower()
|
|
adapter = assignment_source_adapter(source)
|
|
adapter.validate_source_args(args)
|
|
self.source_args[source] = args
|
|
self.source_adapters[source] = adapter
|
|
if not self.source_args:
|
|
raise ValueError('remote assignment builder has no supported sources')
|
|
raw_credential_refs = dict(credential_refs or {})
|
|
if set(raw_credential_refs) - set(self.source_args):
|
|
raise ValueError('remote credential references contain unsupported sources')
|
|
self.credential_refs = {}
|
|
for source in self.source_args:
|
|
reference = str(raw_credential_refs.get(source) or '')
|
|
if len(reference) > 128 or '\x00' in reference:
|
|
raise ValueError('remote credential reference is invalid')
|
|
if (
|
|
self.source_adapters[source].assignment_flow
|
|
is _DIRECT_ASSIGNMENT_FLOW and reference
|
|
):
|
|
raise ValueError('remote direct source credential reference must be empty')
|
|
self.credential_refs[source] = reference
|
|
|
|
self.compatibility_profiles = {}
|
|
for profile_name, value in dict(compatibility_profiles or {}).items():
|
|
profile_name = str(profile_name or '').strip()
|
|
profile = dict(value or {})
|
|
if set(profile) - {'package_manifest', 'sources'}:
|
|
raise ValueError('remote compatibility profile shape is invalid')
|
|
package_value = profile.get('package_manifest')
|
|
package = (
|
|
load_worker_package_manifest(package_value)
|
|
if isinstance(package_value, (str, os.PathLike))
|
|
else normalize_worker_package_manifest(package_value)
|
|
)
|
|
required_build = WorkerBuildCompatibility.from_mapping(
|
|
worker_package_build_compatibility(package),
|
|
)
|
|
package_capabilities = frozenset(
|
|
(
|
|
item['source'], item['platform'], item['planning_kind'],
|
|
)
|
|
for item in package['capabilities']
|
|
)
|
|
package_sources = frozenset(
|
|
source for source, _platform, _planning in package_capabilities
|
|
)
|
|
allowed_sources = frozenset(
|
|
str(item or '').strip().lower()
|
|
for item in (profile.get('sources') or package_sources)
|
|
)
|
|
if (
|
|
not profile_name or not allowed_sources
|
|
or not allowed_sources <= set(self.source_args)
|
|
or any(
|
|
(
|
|
self.source_adapters[source].package_capability.source,
|
|
self.source_adapters[source].package_capability.platform,
|
|
self.source_adapters[source].package_capability.planning_kind,
|
|
) not in package_capabilities
|
|
for source in allowed_sources
|
|
)
|
|
):
|
|
raise ValueError('remote compatibility profile has invalid sources')
|
|
key = tuple(required_build.as_dict().values())
|
|
if key in self.compatibility_profiles:
|
|
raise ValueError('remote compatibility profile identity is duplicated')
|
|
self.compatibility_profiles[key] = {
|
|
'profile_name': profile_name,
|
|
'build': required_build,
|
|
'package_manifest': package,
|
|
'sources': allowed_sources,
|
|
'capabilities': package_capabilities,
|
|
}
|
|
if not self.compatibility_profiles:
|
|
raise ValueError('remote assignment builder has no compatibility profiles')
|
|
|
|
def compatibility_snapshot(self):
|
|
profiles = []
|
|
for profile in sorted(
|
|
self.compatibility_profiles.values(),
|
|
key=lambda item: item['profile_name'],
|
|
):
|
|
profiles.append({
|
|
'profile_name': profile['profile_name'],
|
|
**profile['build'].as_dict(),
|
|
'sources': sorted(profile['sources']),
|
|
'capabilities': [
|
|
{
|
|
'source': source,
|
|
'platform': platform,
|
|
'planning_kind': planning_kind,
|
|
}
|
|
for source, platform, planning_kind in sorted(
|
|
profile['capabilities'],
|
|
)
|
|
],
|
|
})
|
|
required_capabilities = [
|
|
self.source_adapters[source].package_capability.as_dict()
|
|
for source in sorted(self.source_adapters)
|
|
]
|
|
return {
|
|
'profiles': profiles,
|
|
'required_capabilities': required_capabilities,
|
|
}
|
|
|
|
@staticmethod
|
|
def _queue_policy(args):
|
|
return QueueDispositionPolicy(
|
|
target_retry_max_attempts=int(getattr(args, 'target_retry_max_attempts', 3) or 3),
|
|
target_retry_base_delay_sec=int(getattr(args, 'target_retry_base_delay_sec', 3600) or 3600),
|
|
target_retry_max_delay_sec=int(getattr(args, 'target_retry_max_delay_sec', 86400) or 86400),
|
|
target_timeout_retry_delay_sec=int(getattr(args, 'target_timeout_retry_delay_sec', 21600) or 21600),
|
|
)
|
|
|
|
@staticmethod
|
|
def _capacity(args):
|
|
return {
|
|
'bundle_items': int(getattr(args, 'result_bundle_max_items', 10000)),
|
|
'bundle_bytes': int(getattr(args, 'result_bundle_max_total_bytes', 3 << 30)),
|
|
'projection_items': int(getattr(args, 'projection_backlog_max_items', 10000)),
|
|
'projection_bytes': int(getattr(args, 'projection_backlog_max_bytes', 2 << 30)),
|
|
'projection_headroom_bytes': int(getattr(
|
|
args, 'projection_backlog_headroom_bytes', 0,
|
|
)),
|
|
'keycheck_items': int(getattr(args, 'keycheck_queue_max_items', 100000)),
|
|
'keycheck_bytes': int(getattr(args, 'keycheck_queue_max_bytes', 512 << 20)),
|
|
'quarantine_items': int(getattr(args, 'pipeline_quarantine_max_items', 10000)),
|
|
'quarantine_bytes': int(getattr(args, 'pipeline_quarantine_max_bytes', 1 << 30)),
|
|
}
|
|
|
|
@staticmethod
|
|
def _scan_policy(args):
|
|
def setting(name, default):
|
|
return getattr(args, name, getattr(scan_config, name, default))
|
|
|
|
return normalize_remote_scan_policy({
|
|
'drop_detectors': setting('drop_detectors', ()),
|
|
'strict_git_provider_token_filter': bool(setting(
|
|
'strict_git_provider_token_filter', True,
|
|
)),
|
|
'trufflehog_stdout_max_mb': int(setting('trufflehog_stdout_max_mb', 32)),
|
|
'trufflehog_stderr_max_mb': int(setting('trufflehog_stderr_max_mb', 8)),
|
|
'result_bundle_max_event_bytes': int(setting(
|
|
'result_bundle_max_event_bytes', 64 << 20,
|
|
)),
|
|
'trufflehog_max_findings_per_target': int(setting(
|
|
'trufflehog_max_findings_per_target', 20000,
|
|
)),
|
|
'trufflehog_job_memory_limit_bytes': int(setting(
|
|
'trufflehog_job_memory_limit_bytes', 0,
|
|
)),
|
|
'trufflehog_windows_job_cpu_weight': int(setting(
|
|
'trufflehog_windows_job_cpu_weight', 0,
|
|
)),
|
|
'trufflehog_windows_memory_priority': int(setting(
|
|
'trufflehog_windows_memory_priority', 0,
|
|
)),
|
|
'trufflehog_diagnostic_max_lines': int(setting(
|
|
'trufflehog_diagnostic_max_lines', 2000,
|
|
)),
|
|
'trufflehog_diagnostic_max_line_chars': int(setting(
|
|
'trufflehog_diagnostic_max_line_chars', 8192,
|
|
)),
|
|
'trufflehog_diagnostic_max_line_bytes': int(setting(
|
|
'trufflehog_diagnostic_max_line_bytes', 8192,
|
|
)),
|
|
'trufflehog_diagnostic_max_errors': int(setting(
|
|
'trufflehog_diagnostic_max_errors', 200,
|
|
)),
|
|
'trufflehog_diagnostic_max_warnings': int(setting(
|
|
'trufflehog_diagnostic_max_warnings', 200,
|
|
)),
|
|
'trufflehog_diagnostic_max_unclassified': int(setting(
|
|
'trufflehog_diagnostic_max_unclassified', 20,
|
|
)),
|
|
})
|
|
|
|
def _db(self, application_name):
|
|
db = self.db_factory(db_url=self.db_url, initialize=False)
|
|
if not db.enabled:
|
|
db.close()
|
|
raise RuntimeError('remote assignment PostgreSQL connection is unavailable')
|
|
db.set_application_name(application_name)
|
|
return db
|
|
|
|
def _accept_staged(self, identity, claim, staged):
|
|
ready_path = bundle_ready_path(self.bundle_root, claim['bundle_id'])
|
|
relative_path = os.path.relpath(ready_path, self.bundle_root).replace(os.sep, '/')
|
|
staged_relative_path = str(getattr(staged, 'relative_path', '') or '').replace(
|
|
'\\', '/'
|
|
)
|
|
if staged_relative_path and staged_relative_path != relative_path:
|
|
raise RuntimeError('published remote bundle path conflicts with its reservation')
|
|
digest = _hash_file(ready_path)
|
|
reader = ResultBundleReader(
|
|
ready_path, max_event_bytes=int(claim['declared_bundle_bytes']),
|
|
)
|
|
metadata = reader.validate()
|
|
effective_diagnostics = reader.effective_diagnostics()
|
|
values = metadata.as_dict()
|
|
values['effective_diagnostic_count'] = len(effective_diagnostics)
|
|
values['effective_diagnostic_projection_version'] = (
|
|
DIAGNOSTIC_PROJECTION_VERSION
|
|
)
|
|
values['effective_diagnostic_uids_sha256'] = (
|
|
ordered_diagnostic_uid_set_sha256(effective_diagnostics)
|
|
)
|
|
values['relative_path'] = relative_path
|
|
db = self._db('truf-worker-planning-result')
|
|
try:
|
|
return db.mark_result_bundle_ready(
|
|
int(claim['reservation_id']), values,
|
|
remote_acceptance={
|
|
'device_id': int(identity['device_id']),
|
|
'payload_sha256': digest,
|
|
'token_sha256': str(identity['token_sha256']),
|
|
},
|
|
bundle_capacity_bytes=self._capacity(
|
|
self.source_args[claim['source']]
|
|
)['bundle_bytes'],
|
|
)
|
|
finally:
|
|
db.close()
|
|
|
|
def _recover_published(self, identity, claim):
|
|
path = bundle_ready_path(self.bundle_root, claim['bundle_id'])
|
|
if not os.path.lexists(path):
|
|
return False
|
|
metadata = ResultBundleReader(
|
|
path, max_event_bytes=int(claim['declared_bundle_bytes']),
|
|
).validate()
|
|
if metadata.header != BundleReservation.from_mapping(claim).header():
|
|
raise RuntimeError('published remote bundle conflicts with its reservation')
|
|
self._accept_staged(identity, claim, metadata)
|
|
return True
|
|
|
|
def _bound_plan(self, identity, claim):
|
|
db = self._db('truf-worker-plan-reconcile')
|
|
try:
|
|
return db.remote_bound_git_scan_plan(
|
|
int(claim['reservation_id']), int(identity['device_id']),
|
|
str(claim['claim_lease_token']), str(identity['token_sha256']),
|
|
)
|
|
finally:
|
|
db.close()
|
|
|
|
def _reconcile_request(self, identity, request_id):
|
|
db = self._db('truf-worker-request-reconcile')
|
|
try:
|
|
return db.reconcile_remote_assignment_request(
|
|
request_id, int(identity['device_id']), str(identity['token_sha256']),
|
|
)
|
|
finally:
|
|
db.close()
|
|
|
|
def _snapshot(self, adapter, args, compatibility, execution):
|
|
planning = {'kind': adapter.planning_kind}
|
|
if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
planning.update({
|
|
'git_baseline_depth': max(
|
|
1, int(getattr(args, 'git_baseline_depth', 100) or 100),
|
|
),
|
|
'git_ref_resolution_attempts': max(
|
|
1, int(getattr(args, 'git_ref_resolution_attempts', 2) or 2),
|
|
),
|
|
'git_ref_resolution_timeout_sec': max(
|
|
0.1, float(getattr(args, 'git_ref_resolution_timeout_sec', 10) or 10),
|
|
),
|
|
'git_ref_resolution_max_bytes': max(
|
|
1024, int(getattr(args, 'git_ref_resolution_max_bytes', 1 << 20) or (1 << 20)),
|
|
),
|
|
})
|
|
return adapter.validate_snapshot({
|
|
'schema': 1,
|
|
'compatibility': compatibility.as_dict(),
|
|
'execution': execution,
|
|
'planning': planning,
|
|
'credential_ref': {
|
|
'source': adapter.queue_source,
|
|
'auth_entry': self.credential_refs[adapter.queue_source],
|
|
},
|
|
})
|
|
|
|
def _rehydrate(self, client, claim, snapshot, snapshot_sha256=None):
|
|
adapter = assignment_source_adapter(claim.get('source'))
|
|
if (
|
|
self.source_adapters.get(adapter.queue_source) is not adapter
|
|
or str(claim.get('platform') or '') != adapter.worker_platform
|
|
or adapter.assignment_flow not in {
|
|
_EXACT_GIT_ASSIGNMENT_FLOW, _DIRECT_ASSIGNMENT_FLOW,
|
|
}
|
|
):
|
|
raise RuntimeError('remote assignment source adapter is unavailable')
|
|
snapshot = adapter.validate_snapshot(snapshot)
|
|
if (
|
|
snapshot_sha256 is not None
|
|
and str(snapshot_sha256) != remote_execution_snapshot_sha256(snapshot)
|
|
):
|
|
raise ScanEventConflictError('remote execution snapshot changed after admission')
|
|
required = ScanCompatibility.from_mapping(snapshot['compatibility'])
|
|
required_build = WorkerBuildCompatibility.from_mapping({
|
|
'protocol_version': required.protocol_version,
|
|
'bundle_format_version': required.bundle_format_version,
|
|
'platform_tag': required.platform_tag,
|
|
'code_manifest_sha256': required.code_manifest_sha256,
|
|
'detector_policy_sha256': required.detector_policy_sha256,
|
|
})
|
|
if required.protocol_version not in {1, PROTOCOL_VERSION}:
|
|
raise ScanEventConflictError(
|
|
'remote execution snapshot protocol is unsupported'
|
|
)
|
|
validate_worker_build_compatibility(
|
|
required_build, client,
|
|
expected_protocol_version=required.protocol_version,
|
|
)
|
|
source = adapter.queue_source
|
|
args = self.source_args.get(source)
|
|
if args is None or snapshot['credential_ref'] != {
|
|
'source': source, 'auth_entry': self.credential_refs.get(source),
|
|
}:
|
|
raise RuntimeError('remote assignment credential reference is unavailable')
|
|
values = vars(args).copy()
|
|
values.update(snapshot['planning'])
|
|
values.pop('kind', None)
|
|
planning_args = SimpleNamespace(**values)
|
|
event_scan_options = dict(snapshot['execution']['scan_kwargs'])
|
|
scan_kwargs = dict(event_scan_options)
|
|
if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
scan_kwargs['token'] = str(getattr(args, 'token', '') or '')
|
|
return (
|
|
adapter, planning_args, snapshot, scan_kwargs, event_scan_options,
|
|
QueueDispositionPolicy(**snapshot['execution']['queue_policy']),
|
|
dict(snapshot['execution']['limits']),
|
|
dict(snapshot['execution']['scan_policy']), required,
|
|
)
|
|
|
|
def _complete_claim(
|
|
self, identity, client, claim, snapshot, *, snapshot_sha256=None,
|
|
recovered_execution_plan=None,
|
|
):
|
|
(
|
|
adapter, args, snapshot, scan_kwargs, event_scan_options, queue_policy,
|
|
limits, scan_policy, required,
|
|
) = self._rehydrate(client, claim, snapshot, snapshot_sha256)
|
|
device_id = int(identity['device_id'])
|
|
status_db = self._db('truf-worker-claim-reconcile')
|
|
try:
|
|
status = status_db.remote_assignment_status(
|
|
int(claim['reservation_id']), device_id,
|
|
str(identity['token_sha256']),
|
|
)
|
|
finally:
|
|
status_db.close()
|
|
if status and status.get('receipt_id'):
|
|
return None
|
|
if self._recover_published(identity, claim):
|
|
return None
|
|
|
|
if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
execution_target = str(claim.get('target') or '')
|
|
elif adapter.planning_kind == 'docker_direct_v1':
|
|
execution_target = normalize_docker_direct_execution_target(
|
|
claim.get('target'),
|
|
)['image']
|
|
elif adapter.planning_kind == 'huggingface_space_v1':
|
|
execution_target = normalize_huggingface_space_execution_target(
|
|
claim.get('target'),
|
|
)
|
|
else:
|
|
raise RuntimeError('remote assignment execution plan is unavailable')
|
|
|
|
bound_plan = None
|
|
if recovered_execution_plan is not None:
|
|
if (
|
|
not isinstance(recovered_execution_plan, dict)
|
|
or set(recovered_execution_plan) != {
|
|
'kind', 'execution_target', 'bound_plan',
|
|
}
|
|
or recovered_execution_plan['kind'] != adapter.planning_kind
|
|
or str(recovered_execution_plan['execution_target']) != execution_target
|
|
or (
|
|
adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW
|
|
and not isinstance(recovered_execution_plan['bound_plan'], dict)
|
|
)
|
|
or (
|
|
adapter.assignment_flow is _DIRECT_ASSIGNMENT_FLOW
|
|
and recovered_execution_plan['bound_plan'] is not None
|
|
)
|
|
):
|
|
raise ScanEventConflictError(
|
|
'recovered execution plan conflicts with its assignment'
|
|
)
|
|
bound_plan = recovered_execution_plan['bound_plan']
|
|
if (
|
|
adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW
|
|
and bound_plan is None
|
|
):
|
|
bound_plan = self._bound_plan(identity, claim)
|
|
if (
|
|
adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW
|
|
and bound_plan is None
|
|
):
|
|
try:
|
|
bound_plan = self.planner(
|
|
args, self.db_url, adapter.queue_source, claim, scan_kwargs,
|
|
remote_credential={
|
|
'device_id': device_id,
|
|
'token_sha256': str(identity['token_sha256']),
|
|
},
|
|
)
|
|
except ScanEventConflictError:
|
|
bound_plan = self._bound_plan(identity, claim)
|
|
if bound_plan is None:
|
|
raise
|
|
except _GitResolutionFailure as failure:
|
|
bound_plan = self._bound_plan(identity, claim)
|
|
if bound_plan is None:
|
|
result = git_resolution_failure_result(args, claim, scan_kwargs, failure)
|
|
staged = stage_scan_result_in_scope(
|
|
result, claim, self.bundle_root, event_scan_options, queue_policy,
|
|
attempts=int(claim.get('attempts') or 0),
|
|
candidate_max_items=limits['candidate_max_items'],
|
|
candidate_max_bytes=limits['candidate_max_bytes'],
|
|
require_s_drive=False,
|
|
)
|
|
self._accept_staged(identity, claim, staged)
|
|
return None
|
|
|
|
assigned_scan_kwargs = dict(scan_kwargs)
|
|
if adapter.assignment_flow is _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
assigned_scan_kwargs['git_plan'] = bound_plan
|
|
execution_plan = {
|
|
'kind': adapter.planning_kind,
|
|
'execution_target': execution_target,
|
|
'bound_plan': bound_plan,
|
|
}
|
|
issued_at = str(claim.get('remote_issued_at') or '')
|
|
deadline_at = str(claim.get('remote_expires_at') or '')
|
|
try:
|
|
issued = datetime.fromisoformat(issued_at)
|
|
deadline = datetime.fromisoformat(deadline_at)
|
|
except ValueError as exc:
|
|
raise RuntimeError('remote assignment reservation deadline is invalid') from exc
|
|
if (
|
|
issued.tzinfo is None
|
|
or deadline.tzinfo is None
|
|
or issued.utcoffset() != timedelta(0)
|
|
or deadline.utcoffset() != timedelta(0)
|
|
or issued.isoformat(timespec='seconds') != issued_at
|
|
or deadline.isoformat(timespec='seconds') != deadline_at
|
|
):
|
|
raise RuntimeError('remote assignment reservation deadline is invalid')
|
|
assignment_ttl = (deadline - issued).total_seconds()
|
|
if (
|
|
not assignment_ttl.is_integer()
|
|
or not 60 <= assignment_ttl <= 7 * 24 * 60 * 60
|
|
):
|
|
raise RuntimeError('remote assignment reservation lifetime is invalid')
|
|
target_scan_timeout = scan_kwargs.get('timeout_sec')
|
|
if (
|
|
isinstance(target_scan_timeout, bool)
|
|
or not isinstance(target_scan_timeout, (int, float))
|
|
or not float(target_scan_timeout).is_integer()
|
|
or target_scan_timeout <= 0
|
|
):
|
|
raise RuntimeError('remote assignment target scan timeout is invalid')
|
|
persisted_upload_timeout = claim.get(
|
|
'remote_result_upload_body_timeout_seconds'
|
|
)
|
|
if persisted_upload_timeout is None:
|
|
persisted_upload_timeout = self.result_upload_body_timeout_seconds
|
|
if (
|
|
isinstance(persisted_upload_timeout, bool)
|
|
or not isinstance(persisted_upload_timeout, int)
|
|
or not 30 <= persisted_upload_timeout <= 24 * 60 * 60
|
|
):
|
|
raise RuntimeError('remote assignment result upload timeout is invalid')
|
|
reservation = dict(claim)
|
|
reservation.pop('remote_result_upload_body_timeout_seconds', None)
|
|
return {
|
|
'reservation': reservation,
|
|
'deadlines': {
|
|
'target_scan_timeout_seconds': int(target_scan_timeout),
|
|
'result_upload_body_timeout_seconds': persisted_upload_timeout,
|
|
'assignment_ttl_seconds': int(assignment_ttl),
|
|
'assignment_issued_at': issued_at,
|
|
'assignment_deadline_at': deadline_at,
|
|
},
|
|
'compatibility': required.as_dict(),
|
|
'scan_kwargs': assigned_scan_kwargs,
|
|
'event_scan_options': event_scan_options,
|
|
'queue_policy': asdict(queue_policy),
|
|
'limits': limits,
|
|
'scan_policy': scan_policy,
|
|
'execution_snapshot': snapshot,
|
|
'execution_snapshot_sha256': remote_execution_snapshot_sha256(snapshot),
|
|
'execution_plan': execution_plan,
|
|
}
|
|
|
|
def __call__(self, identity, request_id, client_compatibility):
|
|
client = WorkerBuildCompatibility.from_mapping(client_compatibility)
|
|
reconciled = self._reconcile_request(identity, request_id)
|
|
if reconciled is not None:
|
|
if reconciled.get('state') != 'committed':
|
|
return None
|
|
if reconciled.get('receipt'):
|
|
return {'resolution': dict(reconciled['receipt'])}
|
|
execution_plan = reconciled.get('execution_plan')
|
|
if execution_plan is None:
|
|
if assignment_source_adapter(
|
|
reconciled['claim'].get('source')
|
|
).assignment_flow is not _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
raise ScanEventConflictError(
|
|
'direct assignment execution plan is unavailable'
|
|
)
|
|
execution_plan = {
|
|
'kind': 'exact_git_v1',
|
|
'execution_target': reconciled['claim']['target'],
|
|
'bound_plan': reconciled.get('git_plan'),
|
|
}
|
|
return self._complete_claim(
|
|
identity, client, reconciled['claim'],
|
|
reconciled['execution_snapshot'],
|
|
snapshot_sha256=reconciled['execution_snapshot_sha256'],
|
|
recovered_execution_plan=execution_plan,
|
|
)
|
|
profile = self.compatibility_profiles.get(tuple(client.as_dict().values()))
|
|
if profile is None:
|
|
return {'no_assignment': {'reason': 'compatibility'}}
|
|
required_build = validate_worker_build_compatibility(profile['build'], client)
|
|
|
|
candidates = []
|
|
for source in sorted(profile['sources']):
|
|
adapter = self.source_adapters[source]
|
|
capability = adapter.package_capability
|
|
if (
|
|
capability.source, capability.platform, capability.planning_kind,
|
|
) not in profile['capabilities']:
|
|
raise RuntimeError(
|
|
'remote package capability changed after profile validation'
|
|
)
|
|
values = vars(self.source_args[source]).copy()
|
|
central_policy = str(values.get('trufflehog_config') or '')
|
|
if (
|
|
not central_policy or not os.path.isfile(central_policy)
|
|
or sha256_file(central_policy)
|
|
!= required_build.detector_policy_sha256
|
|
):
|
|
raise RuntimeError('remote package detector policy does not match central authority')
|
|
values['trufflehog_config'] = PACKAGE_DETECTOR_POLICY
|
|
source_args = SimpleNamespace(**values)
|
|
source_scan_kwargs = validate_scan_kwargs(
|
|
adapter.worker_platform,
|
|
prepare_scan_options(source_args, 1, quiet=True),
|
|
)
|
|
source_event_options = {
|
|
key: value for key, value in source_scan_kwargs.items() if key != 'token'
|
|
}
|
|
source_queue_policy = self._queue_policy(source_args)
|
|
source_limits = {
|
|
'candidate_max_items': int(getattr(source_args, 'keycheck_candidates_per_event', 2000)),
|
|
'candidate_max_bytes': int(getattr(source_args, 'keycheck_candidate_bytes_per_event', 2 << 20)),
|
|
}
|
|
source_scan_policy = self._scan_policy(source_args)
|
|
effective_identity, execution = remote_execution_identity(
|
|
adapter.worker_platform, source_scan_kwargs, source_event_options,
|
|
source_queue_policy, source_limits, source_scan_policy,
|
|
)
|
|
required = ScanCompatibility(
|
|
protocol_version=required_build.protocol_version,
|
|
bundle_format_version=required_build.bundle_format_version,
|
|
platform_tag=required_build.platform_tag,
|
|
code_manifest_sha256=required_build.code_manifest_sha256,
|
|
effective_config_sha256=effective_identity,
|
|
detector_policy_sha256=required_build.detector_policy_sha256,
|
|
)
|
|
snapshot = self._snapshot(adapter, source_args, required, execution)
|
|
candidates.append((
|
|
adapter, source_args, source_scan_kwargs, source_event_options,
|
|
source_queue_policy, source_limits, source_scan_policy, required,
|
|
snapshot,
|
|
))
|
|
if not candidates:
|
|
raise RuntimeError('remote compatibility profile has no central source configuration')
|
|
device_id = int(identity['device_id'])
|
|
producer = current_process_identity().as_dict()
|
|
start = int(request_id[:16], 16) % len(candidates)
|
|
ordered_candidates = candidates[start:] + candidates[:start]
|
|
multiple_sources = len(ordered_candidates) > 1
|
|
no_work_reasons = []
|
|
for candidate in ordered_candidates:
|
|
(
|
|
adapter, args, scan_kwargs, event_scan_options, queue_policy, limits,
|
|
scan_policy, required, snapshot,
|
|
) = candidate
|
|
admission_token = (
|
|
_stable_id(
|
|
device_id, request_id,
|
|
f'admission:{adapter.queue_source}',
|
|
)
|
|
if multiple_sources else request_id
|
|
)
|
|
if multiple_sources:
|
|
recovered = self._reconcile_request(identity, admission_token)
|
|
if recovered is not None:
|
|
if recovered.get('state') != 'committed':
|
|
continue
|
|
if recovered.get('receipt'):
|
|
return {'resolution': dict(recovered['receipt'])}
|
|
execution_plan = recovered.get('execution_plan')
|
|
if execution_plan is None:
|
|
if assignment_source_adapter(
|
|
recovered['claim'].get('source')
|
|
).assignment_flow is not _EXACT_GIT_ASSIGNMENT_FLOW:
|
|
raise ScanEventConflictError(
|
|
'direct assignment execution plan is unavailable'
|
|
)
|
|
execution_plan = {
|
|
'kind': 'exact_git_v1',
|
|
'execution_target': recovered['claim']['target'],
|
|
'bound_plan': recovered.get('git_plan'),
|
|
}
|
|
return self._complete_claim(
|
|
identity, client, recovered['claim'],
|
|
recovered['execution_snapshot'],
|
|
snapshot_sha256=recovered[
|
|
'execution_snapshot_sha256'
|
|
],
|
|
recovered_execution_plan=execution_plan,
|
|
)
|
|
max_event_bytes = int(getattr(
|
|
args, 'result_bundle_max_event_bytes', 64 << 20,
|
|
))
|
|
remote_reserve_bytes = int(getattr(
|
|
args, 'remote_assignment_reserve_bytes', 2 << 20,
|
|
))
|
|
remote_max_active = int(getattr(
|
|
args, 'remote_assignment_max_active', 50,
|
|
))
|
|
validate_v2_capacity_model(
|
|
int(getattr(args, 'max_active_scans', 1) or 1),
|
|
max_event_bytes,
|
|
int(getattr(args, 'projection_backlog_max_bytes', 2 << 30)),
|
|
int(getattr(
|
|
args, 'projection_backlog_headroom_bytes',
|
|
max_event_bytes * 2,
|
|
)),
|
|
)
|
|
outcome = self.admission(
|
|
self.db_url, adapter.queue_source, adapter.worker_platform,
|
|
producer, self.supervisor_instance_id,
|
|
max_event_bytes, remote_reserve_bytes,
|
|
limits['candidate_max_items'], limits['candidate_max_bytes'],
|
|
lease_seconds=self.assignment_ttl_seconds_by_source.get(
|
|
adapter.queue_source, self.assignment_ttl_seconds,
|
|
),
|
|
max_attempts=int(getattr(
|
|
args, 'target_retry_max_attempts', 3,
|
|
) or 3),
|
|
capacity_limits=self._capacity(args), run_id=None, cycle_id=None,
|
|
reservation_token=admission_token,
|
|
bundle_id=_stable_id(device_id, admission_token, 'bundle'),
|
|
scan_event_id=_stable_id(device_id, admission_token, 'event'),
|
|
resolution_attempts=int(getattr(
|
|
args, 'admission_resolution_attempts', 8,
|
|
) or 8),
|
|
resolution_seconds=float(getattr(
|
|
args, 'admission_resolution_seconds', 30,
|
|
) or 30),
|
|
retry_delay=float(getattr(
|
|
args, 'admission_resolution_retry_delay_sec', 0.2,
|
|
) or 0.2),
|
|
claim_order=str(getattr(
|
|
args, 'target_claim_order', 'oldest',
|
|
) or 'oldest'),
|
|
final_cutover=True,
|
|
reserved_bundle_bytes=remote_reserve_bytes,
|
|
remote_max_active=remote_max_active,
|
|
db_factory=self.db_factory,
|
|
remote_assignment={
|
|
'user_id': int(identity['user_id']),
|
|
'device_id': device_id,
|
|
'effective_config_sha256': required.effective_config_sha256,
|
|
'client_compat_sha256': canonical_json_sha256(
|
|
client.as_dict()
|
|
),
|
|
'token_sha256': str(identity['token_sha256']),
|
|
'result_upload_body_timeout_seconds': (
|
|
self.result_upload_body_timeout_seconds
|
|
),
|
|
'execution_snapshot': snapshot,
|
|
},
|
|
)
|
|
if outcome.claim is not None:
|
|
return self._complete_claim(
|
|
identity, client, outcome.claim, snapshot,
|
|
)
|
|
reason = _ADMISSION_NO_WORK_REASONS.get(str(
|
|
getattr(outcome, 'reason', None) or ''
|
|
))
|
|
if reason:
|
|
no_work_reasons.append(reason)
|
|
if no_work_reasons:
|
|
priority = (
|
|
'dispatch_paused', 'assignment_cap', 'capacity', 'empty_queue',
|
|
)
|
|
return {'no_assignment': {'reason': next(
|
|
reason for reason in priority if reason in no_work_reasons
|
|
)}}
|
|
return None
|
|
|
|
|
|
RemoteGitAssignmentBuilder = RemoteAssignmentBuilder
|