mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-03-25 13:30:45 -07:00
feat(offload): remove auto-deletion and transition to persistent one-worker model
This commit is contained in:
@@ -25,17 +25,23 @@ async function listWorkers() {
|
||||
}
|
||||
|
||||
async function provisionWorker() {
|
||||
const instanceId = Math.floor(Date.now() / 1000);
|
||||
const name = `${INSTANCE_PREFIX}-${instanceId}`;
|
||||
const name = INSTANCE_PREFIX;
|
||||
const zone = 'us-west1-a';
|
||||
|
||||
console.log(`🔍 Checking if worker ${name} already exists...`);
|
||||
const existCheck = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'describe', name,
|
||||
'--project', PROJECT_ID,
|
||||
'--zone', zone
|
||||
], { stdio: 'pipe' });
|
||||
|
||||
if (existCheck.status === 0) {
|
||||
console.log(`✅ Worker ${name} already exists and is ready for use.`);
|
||||
return;
|
||||
}
|
||||
|
||||
console.log(`🚀 Provisioning secure offload worker: ${name}...`);
|
||||
|
||||
// Hardened Metadata: Enable OS Login and 72h Self-Deletion
|
||||
const startupScript = `#!/bin/bash
|
||||
echo "gcloud compute instances delete ${name} --zone ${zone} --project ${PROJECT_ID} --quiet" | at now + 72 hours
|
||||
`;
|
||||
|
||||
const result = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'create', name,
|
||||
'--project', PROJECT_ID,
|
||||
@@ -43,7 +49,7 @@ async function provisionWorker() {
|
||||
'--machine-type', 'n2-standard-8',
|
||||
'--image-family', 'gcli-maintainer-worker',
|
||||
'--image-project', PROJECT_ID,
|
||||
'--metadata', `enable-oslogin=TRUE,startup-script=${startupScript}`,
|
||||
'--metadata', `enable-oslogin=TRUE`,
|
||||
'--labels', `owner=${USER.replace(/[^a-z0-9_-]/g, '_')},type=offload-worker`,
|
||||
'--tags', `gcli-offload-${USER}`,
|
||||
'--scopes', 'https://www.googleapis.com/auth/cloud-platform'
|
||||
|
||||
@@ -36,29 +36,30 @@ export async function runOrchestrator(args: string[], env: NodeJS.ProcessEnv = p
|
||||
}
|
||||
|
||||
const { projectId, zone, terminalType, syncAuth } = config;
|
||||
const userPrefix = `gcli-offload-${env.USER || 'mattkorwel'}`;
|
||||
const targetVM = `gcli-offload-${env.USER || 'mattkorwel'}`;
|
||||
|
||||
console.log(`🔍 Finding active fleet workers for ${userPrefix}...`);
|
||||
console.log(`🔍 Connecting to offload worker: ${targetVM}...`);
|
||||
|
||||
// 2. Discover Worker VM
|
||||
const gcloudList = spawnSync(`gcloud compute instances list --project ${projectId} --filter="name~^${userPrefix} AND status=RUNNING" --format="json"`, { shell: true });
|
||||
// 2. Verify Worker is RUNNING
|
||||
const statusCheck = spawnSync('gcloud', [
|
||||
'compute', 'instances', 'describe', targetVM,
|
||||
'--project', projectId,
|
||||
'--zone', zone,
|
||||
'--format', 'get(status)'
|
||||
], { stdio: 'pipe' });
|
||||
|
||||
let instances = [];
|
||||
try {
|
||||
instances = JSON.parse(gcloudList.stdout.toString());
|
||||
} catch (e) {
|
||||
console.error('❌ Failed to parse gcloud output. Ensure you are logged in.');
|
||||
return 1;
|
||||
const status = statusCheck.stdout.toString().trim();
|
||||
if (status !== 'RUNNING') {
|
||||
if (status === '') {
|
||||
console.log(`⚠️ Worker ${targetVM} does not exist. Please run "npm run offload:fleet provision" first.`);
|
||||
} else {
|
||||
console.log(`⚠️ Worker ${targetVM} is ${status}. Starting it now...`);
|
||||
spawnSync('gcloud', ['compute', 'instances', 'start', targetVM, '--project', projectId, '--zone', zone], { stdio: 'inherit' });
|
||||
}
|
||||
}
|
||||
|
||||
if (instances.length === 0) {
|
||||
console.log('⚠️ No active workers found. Please run "npm run offload:fleet provision" first.');
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Default to the first found worker
|
||||
const targetVM = instances[0].name;
|
||||
const remoteWorkDir = '/home/ubuntu/.offload/workspace';
|
||||
|
||||
const sessionName = `offload-${prNumber}-${action}`;
|
||||
|
||||
// Fetch Metadata (local)
|
||||
|
||||
@@ -50,7 +50,4 @@ npm -v
|
||||
echo "Installing Gemini CLI..."
|
||||
npm install -g @google/gemini-cli@nightly
|
||||
|
||||
# 5. Self-Deletion Cron (Safety)
|
||||
(crontab -u $USER -l 2>/dev/null; echo "0 0 * * * gcloud compute instances delete $(hostname) --zone $(curl -H Metadata-Flavor:Google http://metadata.google.internal/computeMetadata/v1/instance/zone | cut -d/ -f4) --quiet") | crontab -u $USER -
|
||||
|
||||
echo "✅ Provisioning Complete!"
|
||||
|
||||
@@ -31,15 +31,16 @@ describe('Offload Orchestration (GCE)', () => {
|
||||
vi.mocked(fs.mkdirSync).mockReturnValue(undefined as any);
|
||||
vi.mocked(fs.writeFileSync).mockReturnValue(undefined as any);
|
||||
vi.mocked(fs.createWriteStream).mockReturnValue({ pipe: vi.fn() } as any);
|
||||
vi.spyOn(process, 'chdir').mockImplementation(() => {});
|
||||
|
||||
vi.spyOn(process, 'chdir').mockImplementation(() => {});
|
||||
vi.spyOn(process, 'cwd').mockReturnValue('/test-cwd');
|
||||
|
||||
// Default mock for gcloud instance describe
|
||||
vi.mocked(spawnSync).mockImplementation((cmd: any, args: any) => {
|
||||
const callInfo = JSON.stringify({ cmd, args });
|
||||
// 1. Mock GCloud Instance List
|
||||
if (callInfo.includes('gcloud') && callInfo.includes('instances') && callInfo.includes('list')) {
|
||||
return { status: 0, stdout: Buffer.from(JSON.stringify([{ name: 'gcli-offload-test-worker' }])), stderr: Buffer.from('') } as any;
|
||||
if (callInfo.includes('compute') && callInfo.includes('describe')) {
|
||||
return { status: 0, stdout: Buffer.from('RUNNING\n'), stderr: Buffer.from('') } as any;
|
||||
}
|
||||
// 2. Mock GH Metadata Fetching (local or remote)
|
||||
if (callInfo.includes('gh') && callInfo.includes('view')) {
|
||||
return { status: 0, stdout: Buffer.from('test-meta\n'), stderr: Buffer.from('') } as any;
|
||||
}
|
||||
@@ -57,8 +58,8 @@ describe('Offload Orchestration (GCE)', () => {
|
||||
});
|
||||
|
||||
describe('orchestrator.ts', () => {
|
||||
it('should discover active workers and use gcloud compute ssh', async () => {
|
||||
await runOrchestrator(['123'], {});
|
||||
it('should connect to the deterministic worker and use gcloud compute ssh', async () => {
|
||||
await runOrchestrator(['123'], { USER: 'testuser' });
|
||||
|
||||
const spawnCalls = vi.mocked(spawnSync).mock.calls;
|
||||
const sshCall = spawnCalls.find(call =>
|
||||
@@ -66,7 +67,8 @@ describe('Offload Orchestration (GCE)', () => {
|
||||
);
|
||||
|
||||
expect(sshCall).toBeDefined();
|
||||
expect(JSON.stringify(sshCall)).toContain('gcli-offload-test-worker');
|
||||
// Match the new deterministic name: gcli-offload-<USER>
|
||||
expect(JSON.stringify(sshCall)).toContain('gcli-offload-testuser');
|
||||
expect(JSON.stringify(sshCall)).toContain('test-project');
|
||||
});
|
||||
|
||||
|
||||
Reference in New Issue
Block a user